executor1389/modern-search-engine

1import os2import json3from whoosh.index import create_in, open_dir4from whoosh.fields import Schema, TEXT, ID5from sklearn.feature_extraction.text import TfidfVectorizer6from sklearn.metrics.pairwise import cosine_similarity7import numpy as np8import pickle9 10class Indexer:11    def __init__(self, data_dir="data", index_dir="index"):12        self.data_dir = data_dir13        self.index_dir = index_dir14        self.whoosh_dir = os.path.join(index_dir, "whoosh")15        self.tfidf_path = os.path.join(index_dir, "tfidf_model.pkl")16        self.vectors_path = os.path.join(index_dir, "vectors.npy")17        self.meta_path = os.path.join(index_dir, "metadata.pkl")18        19        # Whoosh Schema20        self.schema = Schema(21            url=ID(stored=True, unique=True),22            title=TEXT(stored=True),23            content=TEXT(stored=True)24        )25 26        if not os.path.exists(self.whoosh_dir):27            os.makedirs(self.whoosh_dir)28 29    def build_inverted_index(self):30        print("Building inverted index...")31        ix = create_in(self.whoosh_dir, self.schema)32        writer = ix.writer()33        34        files = [f for f in os.listdir(self.data_dir) if f.endswith('.json')]35        for filename in files:36            with open(os.path.join(self.data_dir, filename), 'r', encoding='utf-8') as f:37                data = json.load(f)38                writer.add_document(39                    url=data['url'],40                    title=data['title'],41                    content=data['content']42                )43        writer.commit()44        print(f"Inverted index built with {len(files)} documents.")45 46    def build_vector_index(self):47        print("Building vector index (using TF-IDF)...")48        files = [f for f in os.listdir(self.data_dir) if f.endswith('.json')]49        documents = []50        urls = []51        titles = []52        53        for filename in files:54            with open(os.path.join(self.data_dir, filename), 'r', encoding='utf-8') as f:55                data = json.load(f)56                documents.append(data['content'])57                urls.append(data['url'])58                titles.append(data['title'])59        60        if not documents:61            print("No documents found to index.")62            return63 64        # Use TfidfVectorizer as a lightweight alternative to neural embeddings65        vectorizer = TfidfVectorizer(stop_words='english', max_features=5000)66        tfidf_matrix = vectorizer.fit_transform(documents)67        68        # Save vectorizer and matrix69        with open(self.tfidf_path, 'wb') as f:70            pickle.dump(vectorizer, f)71        72        np.save(self.vectors_path, tfidf_matrix.toarray())73        74        # Save metadata75        metadata = {76            "urls": urls,77            "titles": titles78        }79        with open(self.meta_path, 'wb') as f:80            pickle.dump(metadata, f)81            82        print(f"Vector index built with {len(documents)} vectors.")83 84    def run_all(self):85        self.build_inverted_index()86        self.build_vector_index()87 88if __name__ == "__main__":89    indexer = Indexer()90    indexer.run_all()91