Reranker adjustments to account for title spam and in text spam of entities (ie Luffy)

simple_search_engine/search_engine.py

@@ -8,89 +8,37 @@ import re
 CORPUS_DIR = "../corpus"  # Now points to the unified directory
 
 class HeuristicReranker:
-    """
-    A simple, classic, pre-neural reranker for search results.
-
-    This reranker is intentionally:
-    - Non-AI (no embeddings, no transformers)
-    - Deterministic (same input → same output)
-    - Explainable (every score adjustment has a clear reason)
-
-    It is designed to run AFTER an initial TF-IDF retrieval step.
-
-    Conceptually:
-        TF-IDF answers: "Is this document related to the query?"
-        This reranker answers: "Is this document actually a good match?"
-    """
-
     def __init__(
         self,
         coverage_weight: float = 1.0,
-        # title_weight: float = 0.8,
         length_penalty_weight: float = 0.2,
         ideal_length: int = 1500,
     ):
-        """
-        Initialize the reranker and its tuning parameters.
-
-        Each weight controls how much influence a particular heuristic has
-        on the final ranking.
-
-        - coverage_weight:
-            How strongly we reward documents that contain *all* query terms.
-            This fixes a common TF-IDF failure mode where one repeated word
-            dominates the score.
-
-        Not USED
-        - title_weight:
-            How strongly we reward query matches in the document title.
-            Titles are usually concise and high-signal.
-
-        - length_penalty_weight:
-            How strongly we penalize very long documents.
-            Long documents often mention many topics but answer none clearly.
-
-        - ideal_length:
-            The document length (in characters) that we consider "reasonable".
-            Documents longer than this are softly penalized.
-        """
         self.coverage_weight = coverage_weight
-        # self.title_weight = title_weight
         self.length_penalty_weight = length_penalty_weight
         self.ideal_length = ideal_length
 
     def _tokenize(self, text: str) -> list[str]:
-        """
-        Convert text into a list of lowercase word tokens.
-
-        This is a deliberately simple tokenizer:
-        - lowercase
-        - split on word boundaries
-        - no stemming or lemmatization
-
-        Why so simple?
-        Because this reranker is meant to be:
-        - predictable
-        - fast
-        - easy to reason about
-        """
         return re.findall(r"\b\w+\b", text.lower())
 
-    def _coverage_score(self, query_tokens, doc_tokens) -> float:
-        """
-        Compute how many query terms appear in the document.
+    def _is_entity_query(self, query_tokens: list[str], query: str) -> bool:
+        if len(query_tokens) > 2:
+            return False
+
+        lowered = query.lower()
 
-        Example:
-            Query tokens: ["python", "error", "handling"]
-            Document tokens contain: ["python", "handling"]
+        forbidden_terms = {
+            "vs", "episode", "season", "arc",
+            "fight", "battle", "when", "what",
+            "how", "why", "where"
+        }
 
-            Coverage score = 2 / 3 ≈ 0.67
+        if any(term in lowered for term in forbidden_terms):
+            return False
 
-        Why this matters:
-        TF-IDF can over-score documents that repeat a single query term.
-        Coverage ensures documents that mention *all* query terms
-        are ranked higher.
-        """
+        return True
+
+    def _coverage_score(self, query_tokens, doc_tokens) -> float:
         if not query_tokens:
             return 0.0
 
@@ -98,99 +46,87 @@ class HeuristicReranker:
         matched = sum(1 for t in query_tokens if t in doc_token_set)
         return matched / len(query_tokens)
 
-    #def _title_boost(self, query_tokens, title: str) -> float:
-        #"""
-        #Compute a boost based on query terms appearing in the title.
+    def _length_penalty(self, doc_length: int) -> float:
+        if doc_length <= self.ideal_length:
+            return 0.0
 
-        #Titles are usually:
-        #- concise
-        #- intentional
-        #- descriptive of the main topic
+        return (doc_length - self.ideal_length) / self.ideal_length
 
-        #Therefore, a document whose title matches the query
-        #is often more relevant than one where the match is buried
-        #deep in the body text.
-        #"""
-        #if not title or not query_tokens:
-        #    return 0.0
+    def _episodic_penalty(self, title: str, url: str) -> float:
+        text = f"{title} {url}".lower()
 
-        #title_tokens = set(self._tokenize(title))
-        #matched = sum(1 for t in query_tokens if t in title_tokens)
-        #return matched / len(query_tokens)
+        patterns = [
+            "season",
+            "episode",
+            "transliteration",
+            "list_of",
+        ]
 
-    def _length_penalty(self, doc_length: int) -> float:
-        """
-        Compute a penalty for very long documents.
+        penalty = 0.0
+        for p in patterns:
+            if p in text:
+                penalty += 1.0
 
-        Why penalize length?
-        - Long documents tend to match many queries accidentally
-        - They often contain background, boilerplate, or tangential info
+        digit_count = sum(c.isdigit() for c in text)
+        penalty += digit_count * 0.1
 
-        This penalty is:
-        - soft (never removes a document)
-        - linear (easy to reason about)
-        """
-        if doc_length <= self.ideal_length:
+        return penalty
+
+    def _entity_title_match(self, query_tokens: list[str], title: str) -> float:
+        if not title:
             return 0.0
 
-        return (doc_length - self.ideal_length) / self.ideal_length
+        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]:
-        """
-        Rerank a list of candidate documents.
-
-        Inputs:
-        - query:
-            The original user query string.
-        - candidates:
-            A list of documents already retrieved by TF-IDF.
-            Each candidate is expected to contain:
-                - "score" (TF-IDF score)
-                - "text"  (document body)
-                - "title" (optional, but recommended)
-
-        Output:
-        - The same documents, reordered by a refined relevance score.
-
-        Important:
-        This method does NOT retrieve new documents.
-        It only reorders existing candidates.
-        """
         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)
 
-            # Start with the original TF-IDF score.
-            # This represents basic lexical relevance.
-            score = c["score"]
-
-            # Add a boost if most or all query terms appear in the document.
-            score += self.coverage_weight * self._coverage_score(
-                query_tokens, doc_tokens
-            )
-
-            # Add a boost if query terms appear in the title.
-            #score += self.title_weight * self._title_boost(
-            #    query_tokens, c.get("title", "")
-            #)
+            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)
+                )
 
-            # Subtract a penalty if the document is very long.
-            score -= self.length_penalty_weight * self._length_penalty(
-                len(doc_text)
-            )
-
-            # Store the adjusted score alongside the original document data.
             reranked.append({**c, "score": score})
 
-        # Sort documents by final score, highest first.
         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