|
|
| import math
|
| import random
|
| from nltk.corpus import wordnet
|
| from scrapy.crawler import CrawlerRunner
|
| from scrapy.utils.log import configure_logging
|
| from scrapy.utils.project import get_project_settings
|
| from twisted.internet import reactor, defer
|
| from scrapy import signals
|
| import logging
|
| from my_search_engine.my_search_engine.spiders.search_spider import SearchSpider
|
| from sentence_transformers import SentenceTransformer, util
|
| from ranking import train_ranking_model
|
| import time
|
|
|
| logger = logging.getLogger(__name__)
|
|
|
| class MCTSNode:
|
| def __init__(self, state, parent=None, action=None):
|
| self.state = state
|
| self.parent = parent
|
| self.action = action
|
| self.children = []
|
| self.visits = 0
|
| self.value = 0
|
| self.ucb_score = float('inf')
|
|
|
| def is_leaf(self):
|
| return len(self.children) == 0
|
|
|
| def add_child(self, child_state, action=None):
|
| child_node = MCTSNode(child_state, parent=self, action=action)
|
| self.children.append(child_node)
|
| return child_node
|
|
|
| def update(self, reward):
|
| self.visits += 1
|
| self.value += reward
|
| if self.parent:
|
| self.ucb_score = self.calculate_ucb()
|
|
|
| def calculate_ucb(self, exploration_weight=1.41):
|
| if self.visits == 0 or not self.parent:
|
| return float('inf')
|
| exploitation = self.value / self.visits
|
| exploration = exploration_weight * math.sqrt(math.log(self.parent.visits) / self.visits)
|
| return exploitation + exploration
|
|
|
| class MCTS:
|
| def __init__(self, initial_state, num_simulations=20, exploration_weight=1.41):
|
| self.root = MCTSNode(initial_state)
|
| self.num_simulations = num_simulations
|
| self.exploration_weight = exploration_weight
|
| self.query_model = SentenceTransformer('all-MiniLM-L6-v2')
|
| self.results = []
|
| self.crawler_runner = CrawlerRunner(get_project_settings())
|
| self.initial_state = initial_state
|
| self.num_iterations = 5
|
|
|
| def select(self, node):
|
| while not node.is_leaf():
|
| if not node.children:
|
| return node
|
| node = max(node.children, key=lambda c: c.calculate_ucb(self.exploration_weight))
|
| return node
|
|
|
| def expand(self, node):
|
| if node.visits == 0:
|
| return node
|
| possible_refinements = self.get_possible_refinements(node.state)
|
| for refinement in possible_refinements:
|
| node.add_child(refinement)
|
| return random.choice(node.children) if node.children else node
|
|
|
| def calculate_combined_reward(self, ranking_score, state):
|
| state_length_reward = len(state) / 100
|
| if state:
|
| query_complexity = len(set(state.split())) / len(state.split())
|
| else:
|
| query_complexity = 0
|
| semantic_similarity = self.calculate_semantic_similarity(state, self.root.state)
|
|
|
| combined_reward = (
|
| 0.5 * ranking_score +
|
| 0.2 * state_length_reward +
|
| 0.2 * query_complexity +
|
| 0.1 * semantic_similarity
|
| )
|
| return combined_reward
|
|
|
| def calculate_semantic_similarity(self, query1, query2):
|
| embedding1 = self.query_model.encode(query1)
|
| embedding2 = self.query_model.encode(query2)
|
| return util.pytorch_cos_sim(embedding1, embedding2).item()
|
|
|
| def backpropagate(self, node, reward):
|
| while node is not None:
|
| node.update(reward)
|
| node = node.parent
|
|
|
| def best_action(self):
|
| if not self.root.children:
|
| return self.root
|
|
|
| def score(node):
|
| if node.visits == 0:
|
| return float('-inf')
|
| return node.value / node.visits
|
|
|
| return max(self.root.children, key=score)
|
|
|
| def refine_query(self, query):
|
| words = query.split()
|
| refined_query = []
|
|
|
| for word in words:
|
| if word.lower() not in {"how", "to", "get", "an", "the", "and", "or", "of", "build"}:
|
| synonyms = wordnet.synsets(word)
|
| if synonyms:
|
| synonym_words = [lemma.name() for lemma in synonyms[0].lemmas()
|
| if len(lemma.name().split()) == 1 and word != lemma.name()]
|
| if synonym_words:
|
| refined_query.append(random.choice(synonym_words))
|
| else:
|
| refined_query.append(word)
|
| else:
|
| refined_query.append(word)
|
| else:
|
| refined_query.append(word)
|
|
|
| possible_intent_keywords = ['guide', 'tutorial', 'LLM', 'language model', 'NLP', 'GPT']
|
| refined_query.append(random.choice(possible_intent_keywords))
|
|
|
| return ' '.join(refined_query)
|
|
|
| def get_related_queries(self, query):
|
| query_embedding = self.query_model.encode(query)
|
| refined_query_variations = [query]
|
| words_to_avoid = {'how', 'to', 'get'}
|
| words = query.split()
|
|
|
| for word in words:
|
| if word.lower() not in words_to_avoid:
|
| synonyms = wordnet.synsets(word)
|
| if synonyms:
|
| synonym_words = [lemma.name() for lemma in synonyms[0].lemmas() if lemma.name() != word]
|
| if synonym_words:
|
| refined_query = query.replace(word, random.choice(synonym_words))
|
| refined_query_variations.append(refined_query)
|
|
|
| refined_query_variations = list(set(refined_query_variations))
|
| refined_query_embeddings = [self.query_model.encode(variation) for variation in refined_query_variations]
|
| similarity_scores = util.pytorch_cos_sim(query_embedding, refined_query_embeddings).tolist()[0]
|
|
|
| similarity_threshold = 0.8
|
| filtered_queries = [variation for idx, variation in enumerate(refined_query_variations)
|
| if similarity_scores[idx] > similarity_threshold]
|
|
|
| return filtered_queries[:2] if filtered_queries else [query]
|
|
|
| def get_possible_refinements(self, query):
|
| refined_queries = self.get_related_queries(query)
|
| return refined_queries + [self.refine_query(query)]
|
|
|
| @defer.inlineCallbacks
|
| def web_search(self, query, search_sites=None):
|
| if not query.strip():
|
| logger.error("Cannot perform web search with an empty query.")
|
| defer.returnValue([])
|
|
|
| logger.info(f"Starting web search for query: {query}")
|
| configure_logging(install_root_handler=False)
|
| logging.basicConfig(level=logging.INFO)
|
|
|
| results = []
|
|
|
| def crawler_results(item, response, spider):
|
| logger.info(f"Received result: {item['title']}")
|
| results.append(item)
|
|
|
| try:
|
| crawler = self.crawler_runner.create_crawler(SearchSpider)
|
| crawler.signals.connect(crawler_results, signal=signals.item_scraped)
|
|
|
|
|
| yield self.crawler_runner.crawl(crawler, query=query, search_sites=search_sites)
|
| except Exception as e:
|
| logger.error(f"Error during web search: {str(e)}")
|
| defer.returnValue([])
|
|
|
| logger.info(f"Web search completed. Found {len(results)} results.")
|
| defer.returnValue(results)
|
|
|
| @defer.inlineCallbacks
|
| def run(self):
|
| logger.info(f"Starting MCTS run with {self.num_iterations} iterations")
|
| for i in range(self.num_iterations):
|
| logger.debug(f"Iteration {i+1}/{self.num_iterations}")
|
| leaf = self.select(self.root)
|
| child = self.expand(leaf)
|
| reward = yield self.simulate(child)
|
| self.backpropagate(child, reward)
|
|
|
| best_child = self.best_action()
|
| logger.info(f"MCTS run completed. Best action: {best_child.state}")
|
| defer.returnValue(best_child.state if best_child != self.root else self.root.state)
|
|
|
| @defer.inlineCallbacks
|
| def simulate(self, node):
|
| query_results = yield self.web_search(node.state)
|
| ranked_results = train_ranking_model(node.state, query_results)
|
|
|
| if ranked_results:
|
| top_score = ranked_results[0]['predicted_score']
|
| else:
|
| top_score = 0
|
|
|
| reward = self.calculate_combined_reward(top_score, node.state)
|
| defer.returnValue(reward)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
