Update simple_search_engine/search_engine.py

simple_search_engine/search_engine.py

@@ -9,6 +9,7 @@ import numpy as np
 
 CORPUS_DIR = "../corpus"
 
+
 def _tokenize(text: str):
     return re.findall(r"\b\w+\b", text.lower())
 
@@ -19,11 +20,11 @@ class SimpleSearchEngine:
         self.documents: List[Dict] = []
 
         # TF-IDF
-        self.vectorizer: TfidfVectorizer | None = None
+        self.vectorizer = None
         self.doc_tfidf = None
 
         # BM25
-        self.bm25: BM25Okapi | None = None
+        self.bm25 = None
         self._bm25_tokens = []
 
     def _load_documents(self):
@@ -45,16 +46,13 @@ class SimpleSearchEngine:
                         data = json.load(f)
 
                     text = data.get("tf_idf_text", "").strip()
-                    url = data.get("url", "")
-                    title = data.get("title", os.path.splitext(fname)[0])
-
                     if not text:
                         continue
 
                     docs.append({
                         "id": doc_id,
-                        "title": title,
-                        "url": url,
+                        "title": data.get("title", os.path.splitext(fname)[0]),
+                        "url": data.get("url", ""),
                         "path": path,
                         "text": text,
                     })
@@ -75,7 +73,7 @@ class SimpleSearchEngine:
             return
 
         # TF-IDF
-        self.vectorizer = TfidfVectorizer(analyzer="word", token_pattern=r"\b\w+\b")
+        self.vectorizer = TfidfVectorizer(token_pattern=r"\b\w+\b")
         self.doc_tfidf = self.vectorizer.fit_transform(texts)
         print("TF-IDF index built.")
 
@@ -91,51 +89,32 @@ class SimpleSearchEngine:
             return np.zeros_like(scores)
         return (scores - min_s) / (max_s - min_s)
 
-    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)
-        print(f"[DEBUG][BM25] query={query}")
-        print("[DEBUG][BM25] top5_scores=", scores[scores.argsort()[::-1][:5]])
+    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.")
 
+        query_vec = self.vectorizer.transform([query])
+        scores = linear_kernel(query_vec, self.doc_tfidf).flatten()
         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,
-            })
+            doc = self.documents[idx]
+            results.append(doc)
 
         return results
 
-    # TF-IDF search
-    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.")
+    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.")
 
-        query_vec = self.vectorizer.transform([query])
-        similarities = linear_kernel(query_vec, self.doc_tfidf).flatten()
-        top_idx = similarities.argsort()[::-1][:top_k]
+        scores = np.array(self.bm25.get_scores(_tokenize(query)))
+        ranked_idx = scores.argsort()[::-1][:top_k]
 
         results = []
-        for idx in top_idx:
+        for idx in ranked_idx:
             doc = self.documents[idx]
-            snippet = doc["text"][:200] + ("..." if len(doc["text"]) > 200 else "")
-            results.append({
-                "score": float(similarities[idx]),
-                "title": doc["title"],
-                "url": doc["url"],
-                "path": doc["path"],
-                "snippet": snippet,
-            })
+            results.append(doc)
 
         return results
 
@@ -148,8 +127,10 @@ class SimpleSearchEngine:
         tfidf_scores = linear_kernel(query_vec, self.doc_tfidf).flatten()
 
         # BM25 scores
-        q_tokens = _tokenize(query)
-        bm25_scores = np.array(self.bm25.get_scores(q_tokens), dtype=float)
+        bm25_scores = np.array(
+            self.bm25.get_scores(_tokenize(query)),
+            dtype=float
+        )
 
         # Normalize
         tfidf_norm = self._minmax_normalize(tfidf_scores)
@@ -157,24 +138,11 @@ class SimpleSearchEngine:
 
         # Combine
         hybrid_scores = alpha * tfidf_norm + (1 - alpha) * bm25_norm
-        print(f"[DEBUG][HYBRID] query={query} alpha={alpha}")
-        print(
-            "[DEBUG][HYBRID] top5_hybrid=",
-            hybrid_scores[hybrid_scores.argsort()[::-1][:5]]
-        )
-
         ranked_idx = hybrid_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(hybrid_scores[int(idx)]),
-                "title": doc["title"],
-                "url": doc["url"],
-                "path": doc["path"],
-                "snippet": snippet,
-            })
-
-        return results
+            doc = self.documents[idx]
+            results.append(doc)
+
+        return results