| |
|
| | 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)
|
| |
|
| |
|
| |
|
| |
|
| |
|
| |
|
| |
|
