Update simple_search_engine/search_engine.py

simple_search_engine/search_engine.py

@@ -4,16 +4,28 @@ from typing import List, Dict
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import linear_kernel
 from rank_bm25 import BM25Okapi
+import re
+import numpy as np
+
+CORPUS_DIR = "../corpus"
+
+def _tokenize(text: str):
+    return re.findall(r"\b\w+\b", text.lower())
 
-CORPUS_DIR = "../corpus"  # Now points to the unified directory
 
 class SimpleSearchEngine:
     def __init__(self, corpus_dir: str = CORPUS_DIR):
         self.corpus_dir = corpus_dir
         self.documents: List[Dict] = []
+
+        # TF-IDF
         self.vectorizer: TfidfVectorizer | None = None
         self.doc_tfidf = None
 
+        # BM25
+        self.bm25: BM25Okapi | None = None
+        self._bm25_tokens = []
+
     def _load_documents(self):
         docs = []
         doc_id = 0
@@ -32,7 +44,6 @@ class SimpleSearchEngine:
                     with open(path, "r", encoding="utf-8") as f:
                         data = json.load(f)
 
-                    # Extract TF-IDF specific text
                     text = data.get("tf_idf_text", "").strip()
                     url = data.get("url", "")
                     title = data.get("title", os.path.splitext(fname)[0])
@@ -45,13 +56,12 @@ class SimpleSearchEngine:
                         "title": title,
                         "url": url,
                         "path": path,
-                        "text": text,  # Used for indexing
+                        "text": text,
                     })
                     doc_id += 1
 
                 except Exception as e:
                     print(f"Error reading {path}: {e}")
-                    continue
 
         self.documents = docs
         print(f"Loaded {len(self.documents)} documents from {self.corpus_dir}")
@@ -64,63 +74,57 @@ class SimpleSearchEngine:
             print("No documents found to index.")
             return
 
+        # TF-IDF
         self.vectorizer = TfidfVectorizer(analyzer="word", token_pattern=r"\b\w+\b")
         self.doc_tfidf = self.vectorizer.fit_transform(texts)
         print("TF-IDF index built.")
 
+        # BM25
+        self._bm25_tokens = [_tokenize(t) for t in texts]
+        self.bm25 = BM25Okapi(self._bm25_tokens)
+        print("BM25 index built.")
+
+    def search_bm25(self, query: str, top_k: int = 5):
+        if self.bm25 is None:
+            raise RuntimeError("BM25 index not built. Call build_index() first.")
+
+        q_tokens = _tokenize(query)
+        scores = np.array(self.bm25.get_scores(q_tokens), dtype=float)
+
+        ranked_idx = scores.argsort()[::-1][:top_k]
+
+        results = []
+        for idx in ranked_idx:
+            doc = self.documents[int(idx)]
+            snippet = doc["text"][:200] + ("..." if len(doc["text"]) > 200 else "")
+            results.append({
+                "score": float(scores[int(idx)]),
+                "title": doc["title"],
+                "url": doc["url"],
+                "path": doc["path"],
+                "snippet": snippet,
+            })
+
+        return results
+
     def search(self, query: str, top_k: int = 5):
         if self.vectorizer is None or self.doc_tfidf is None:
             raise RuntimeError("Index not built. Call build_index() first.")
-    
-        # 1) Retrieve a larger candidate set with TF-IDF
-        initial_k = max(top_k * 5, 20)
-    
+
         query_vec = self.vectorizer.transform([query])
         similarities = linear_kernel(query_vec, self.doc_tfidf).flatten()
-        top_idx = similarities.argsort()[::-1][:initial_k]
-    
-        candidates = []
+        top_idx = similarities.argsort()[::-1][:top_k]
+
+        results = []
         for idx in top_idx:
             doc = self.documents[idx]
-            candidates.append({
-                "score": float(similarities[idx]),  # TF-IDF score
-                "title": doc["title"],
-                "url": doc["url"],
-                "path": doc["path"],
-                "text": doc["text"],  # full text for reranking
-            })
-            
-        # 2) Return final top_k results
-        results = []
-        for doc in candidates[:top_k]:
             snippet = doc["text"][:200] + ("..." if len(doc["text"]) > 200 else "")
             results.append({
-                "score": doc["score"],
+                "score": float(similarities[idx]),
                 "title": doc["title"],
                 "url": doc["url"],
                 "path": doc["path"],
                 "snippet": snippet,
             })
-        
-        return results
 
-if __name__ == "__main__":
-    engine = SimpleSearchEngine(corpus_dir=CORPUS_DIR)
-    engine.build_index()
-
-    while True:
-        query = input("\nEnter your query (or 'quit'): ").strip()
-        if query.lower() in {"quit", "exit"}:
-            break
-
-        try:
-            hits = engine.search(query, top_k=5)
-            print(f"\nTop results for: {query!r}")
-            for i, h in enumerate(hits, start=1):
-                print(f"\n[{i}] {h['title']} (score={h['score']:.4f})")
-                print(f"    Source: {h['url']}")
-                print(f"    Snippet: {h['snippet']}")
-        except RuntimeError as e:
-            print(e)
-
-# end of code
+        return results