executor1389/modern-search-engine

1import os2import pickle3import numpy as np4from whoosh.index import open_dir5from whoosh.qparser import QueryParser6from sklearn.metrics.pairwise import cosine_similarity7import nltk8from nltk.stem import PorterStemmer9from nltk.tokenize import word_tokenize10 11# Ensure nltk resources are available12try:13    nltk.data.find('tokenizers/punkt')14    nltk.data.find('tokenizers/punkt_tab')15except LookupError:16    nltk.download('punkt')17    nltk.download('punkt_tab')18 19class QueryProcessor:20    def __init__(self, index_dir="index"):21        self.index_dir = index_dir22        self.whoosh_dir = os.path.join(index_dir, "whoosh")23        self.tfidf_path = os.path.join(index_dir, "tfidf_model.pkl")24        self.vectors_path = os.path.join(index_dir, "vectors.npy")25        self.meta_path = os.path.join(index_dir, "metadata.pkl")26        27        self.stemmer = PorterStemmer()28        29        # Load indexes30        self.ix = open_dir(self.whoosh_dir)31        32        with open(self.tfidf_path, 'rb') as f:33            self.vectorizer = pickle.load(f)34            35        self.tfidf_matrix = np.load(self.vectors_path)36        37        with open(self.meta_path, 'rb') as f:38            self.metadata = pickle.load(f)39 40    def process_query(self, query):41        tokens = word_tokenize(query.lower())42        stemmed = [self.stemmer.stem(t) for t in tokens]43        return " ".join(stemmed)44 45    def search_bm25(self, query, limit=10):46        results = []47        with self.ix.searcher() as searcher:48            parser = QueryParser("content", self.ix.schema)49            parsed_query = parser.parse(query)50            whoosh_results = searcher.search(parsed_query, limit=limit)51            for r in whoosh_results:52                results.append({53                    "url": r['url'],54                    "title": r['title'],55                    "score": r.score56                })57        return results58 59    def search_vector(self, query, limit=10):60        # Transform query to vector61        query_vec = self.vectorizer.transform([query]).toarray()62        63        # Compute cosine similarity64        similarities = cosine_similarity(query_vec, self.tfidf_matrix).flatten()65        66        # Get top-k indices67        top_indices = np.argsort(similarities)[::-1][:limit]68        69        results = []70        for idx in top_indices:71            if similarities[idx] > 0:72                results.append({73                    "url": self.metadata['urls'][idx],74                    "title": self.metadata['titles'][idx],75                    "score": float(similarities[idx])76                })77        return results78 79    def reciprocal_rank_fusion(self, bm25_results, vector_results, k=60):80        scores = {}81        82        def update_scores(results):83            for rank, res in enumerate(results):84                url = res['url']85                if url not in scores:86                    scores[url] = {"score": 0.0, "title": res['title']}87                scores[url]["score"] += 1.0 / (k + rank + 1)88        89        update_scores(bm25_results)90        update_scores(vector_results)91        92        # Sort by fused score93        fused = sorted(94            [{"url": url, "title": data["title"], "score": data["score"]} for url, data in scores.items()],95            key=lambda x: x["score"],96            reverse=True97        )98        return fused99 100    def hybrid_search(self, query, limit=10):101        processed_query = self.process_query(query)102        print(f"Searching for: {processed_query}")103        104        bm25_res = self.search_bm25(query, limit=limit*2)105        vec_res = self.search_vector(query, limit=limit*2)106        107        fused_res = self.reciprocal_rank_fusion(bm25_res, vec_res)108        return fused_res[:limit]109 110if __name__ == "__main__":111    qp = QueryProcessor()112    test_query = "Python programming tutorials"113    print("\nHybrid Search Results:")114    for res in qp.hybrid_search(test_query):115        print(f"- {res['title']} ({res['url']}) [Score: {res['score']:.4f}]")116