Rollback to remove reranker

simple_search_engine/search_engine.py

@@ -3,130 +3,9 @@ import json
 from typing import List, Dict
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import linear_kernel
-import re
 
 CORPUS_DIR = "../corpus"  # Now points to the unified directory
 
-class HeuristicReranker:
-    def __init__(
-        self,
-        coverage_weight: float = 1.0,
-        length_penalty_weight: float = 0.2,
-        ideal_length: int = 1500,
-    ):
-        self.coverage_weight = coverage_weight
-        self.length_penalty_weight = length_penalty_weight
-        self.ideal_length = ideal_length
-
-    def _tokenize(self, text: str) -> list[str]:
-        return re.findall(r"\b\w+\b", text.lower())
-
-    def _is_entity_query(self, query_tokens: list[str], query: str) -> bool:
-        if len(query_tokens) > 2:
-            return False
-
-        lowered = query.lower()
-
-        forbidden_terms = {
-            "vs", "episode", "season", "arc",
-            "fight", "battle", "when", "what",
-            "how", "why", "where"
-        }
-
-        if any(term in lowered for term in forbidden_terms):
-            return False
-
-        return True
-
-    def _coverage_score(self, query_tokens, doc_tokens) -> float:
-        if not query_tokens:
-            return 0.0
-
-        doc_token_set = set(doc_tokens)
-        matched = sum(1 for t in query_tokens if t in doc_token_set)
-        return matched / len(query_tokens)
-
-    def _length_penalty(self, doc_length: int) -> float:
-        if doc_length <= self.ideal_length:
-            return 0.0
-
-        return (doc_length - self.ideal_length) / self.ideal_length
-
-    def _episodic_penalty(self, title: str, url: str) -> float:
-        text = f"{title} {url}".lower()
-
-        patterns = [
-            "season",
-            "episode",
-            "transliteration",
-            "list_of",
-        ]
-
-        penalty = 0.0
-        for p in patterns:
-            if p in text:
-                penalty += 1.0
-
-        digit_count = sum(c.isdigit() for c in text)
-        penalty += digit_count * 0.1
-
-        return penalty
-
-    def _entity_title_match(self, query_tokens: list[str], title: str) -> float:
-        if not title:
-            return 0.0
-
-        title_tokens = self._tokenize(title)
-
-        if title_tokens == query_tokens:
-            return 3.0
-
-        if all(t in title_tokens for t in query_tokens):
-            return 1.5
-
-        return 0.0
-
-    def _mention_spam_penalty(self, query_tokens: list[str], doc_tokens: list[str]) -> float:
-        if len(query_tokens) != 1:
-            return 0.0
-
-        term = query_tokens[0]
-        freq = doc_tokens.count(term)
-
-        if freq <= 10:
-            return 0.0
-
-        return (freq - 10) * 0.05
-
-    def rerank(self, query: str, candidates: list[dict]) -> list[dict]:
-        query_tokens = self._tokenize(query)
-        is_entity = self._is_entity_query(query_tokens, query)
-
-        reranked = []
-        for c in candidates:
-            doc_text = c["text"]
-            doc_tokens = self._tokenize(doc_text)
-
-            if is_entity:
-                score = 0.2 * c["score"]
-                score += self._entity_title_match(query_tokens, c.get("title", ""))
-                score -= self._episodic_penalty(c.get("title", ""), c.get("url", ""))
-                score -= self._mention_spam_penalty(query_tokens, doc_tokens)
-                score -= 0.1 * self._length_penalty(len(doc_text))
-            else:
-                score = c["score"]
-                score += self.coverage_weight * self._coverage_score(
-                    query_tokens, doc_tokens
-                )
-                score -= self.length_penalty_weight * self._length_penalty(
-                    len(doc_text)
-                )
-
-            reranked.append({**c, "score": score})
-
-        reranked.sort(key=lambda x: x["score"], reverse=True)
-        return reranked
-
 class SimpleSearchEngine:
     def __init__(self, corpus_dir: str = CORPUS_DIR):
         self.corpus_dir = corpus_dir
@@ -209,14 +88,10 @@ class SimpleSearchEngine:
                 "path": doc["path"],
                 "text": doc["text"],  # full text for reranking
             })
-    
-        # 2) Rerank using classic heuristics
-        reranker = HeuristicReranker()
-        reranked = reranker.rerank(query, candidates)
-    
-        # 3) Return final top_k results (same output format as before)
+            
+        # 2) Return final top_k results
         results = []
-        for doc in reranked[:top_k]:
+        for doc in candidates[:top_k]:
             snippet = doc["text"][:200] + ("..." if len(doc["text"]) > 200 else "")
             results.append({
                 "score": doc["score"],
@@ -225,7 +100,7 @@ class SimpleSearchEngine:
                 "path": doc["path"],
                 "snippet": snippet,
             })
-    
+        
         return results
 
 if __name__ == "__main__":