Feature/scoring

src/backend/search_engine/models/index.py

@@ -6,6 +6,7 @@ class SearchResult(BaseModel):
     url: HttpUrl
     title: str
     snippet: str
+    score: float
 
 
 class SearchResults(BaseModel):

src/backend/search_engine/query/query_engine.py

@@ -1,4 +1,5 @@
 import time
+import heapq
 from backend.search_engine.index.index_loader import get_index
 from backend.search_engine.models.index import SearchResult, SearchResults
 from cpp_utils import (  # type: ignore [import-untyped]
@@ -19,6 +20,7 @@
 
 from backend.search_engine.spell_correction.spell_corrector import get_spell_corrector
 from backend.search_engine.spell_correction.spell_correction import repl
+from backend.search_engine.scoring.bm25 import bm25_score_docs, BM25Config
 
 logger = get_logger(__name__)
 
@@ -175,10 +177,48 @@ def search_results(self, limit: int = 10) -> SearchResults:
             f"Found {len(result.postings)} results in {time.perf_counter() - start:.6f} seconds"
         )
 
-        top_n_results = result  # TODO will be done by BM25 ranking later
+        # candidates: bool/phrase search returns doc_ids in result.postings
+        candidate_doc_ids = list(result.postings)
+        metadata = self.inverted_index.metadata
+
+        # score which query terms
+        # - bool queries: qt.unique_terms
+        # - AND-auto-query w/o operators: normalized_tokens
+        query_terms: list[str] = (
+            list(qt.unique_terms)
+            if getattr(qt, "unique_terms", None)
+            else list(dict.fromkeys(normalized_tokens))
+        )
+
+        t_score = time.perf_counter()
+
+        scores = bm25_score_docs(
+            query_terms=query_terms,
+            postings_by_term=self.inverted_index.index,  # term -> PostingList
+            candidate_doc_ids=candidate_doc_ids,
+            num_docs=metadata.num_docs,
+            avgdl=metadata.avg_doc_length,
+            get_doc_length=metadata.get_doc_length,
+            cfg=BM25Config(k1=1.2, b=0.75, idf_threshold=0.0, clamp_negative_idf=True),
+        )
+
+        logger.debug(f"BM25 scoring time: {time.perf_counter() - t_score:.6f}s")
+
+        t_sort = time.perf_counter()
+
+        # sort doc_ids acc to score
+        ranked_top = heapq.nlargest(
+            limit,
+            ((doc_id, scores.get(doc_id, 0.0)) for doc_id in candidate_doc_ids),
+            key=lambda x: x[1],
+        )
+
+        logger.debug(f"Ranking sort time: {time.perf_counter() - t_sort:.6f}s")
+
+        t_top = time.perf_counter()
 
         search_results = []
-        for doc_id in top_n_results.postings[:limit]:
+        for doc_id, score in ranked_top:
             doc_data = self.inverted_index.doc_store.get(doc_id)
             if doc_data is None:
                 continue
@@ -195,12 +235,17 @@ def search_results(self, limit: int = 10) -> SearchResults:
                     url=url,  # type: ignore[arg-type]
                     title=title,
                     snippet=snippet,
+                    score=score,
                 )
                 search_results.append(search_result)
             except Exception as e:
                 logger.error(f"Error creating SearchResult for doc_id {doc_id}: {e}")
                 continue
 
+        logger.debug(
+            f"Build top-{limit} results time: {time.perf_counter() - t_top:.6f}s"
+        )
+
         end = time.perf_counter()
         logger.debug(
             f"Returned {len(search_results)} results. "

src/backend/search_engine/scoring/bm25.py

@@ -0,0 +1,123 @@
+# backend/search_engine/scoring/bm25.py
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Iterable, Mapping, Sequence
+from collections.abc import Callable
+
+from cpp_utils import PostingList  # type: ignore [import-untyped]
+
+
+@dataclass(frozen=True)
+class BM25Config:
+    k1: float = 1.2  # how strong tf influence is
+    b: float = 0.75  # level of document normalization
+
+    # ignore terms with idf < idf_threshold
+    idf_threshold: float = 0.0
+    # clamp negative idf to 0
+    clamp_negative_idf: bool = True
+
+    min_terms_after_threshold: int = 1
+
+
+def bm25_idf(num_docs: int, df: int, *, clamp_negative: bool = True) -> float:
+    """
+    idf: ln((N - df + 0.5) / (df + 0.5))
+    """
+    if num_docs <= 0:
+        return 0.0
+    if df <= 0:
+        return 0.0
+
+    val = math.log((num_docs - df + 0.5) / (df + 0.5))
+    if clamp_negative and val < 0.0:
+        return 0.0
+    return val
+
+
+def bm25_term_contribution(
+    tf: int,
+    doc_len: int,
+    avgdl: float,
+    *,
+    idf: float,
+    k1: float,
+    b: float,
+) -> float:
+    if tf <= 0:
+        return 0.0
+    if avgdl <= 0:
+        avgdl = 1.0
+
+    denom = tf + k1 * (1.0 - b + b * (doc_len / avgdl))
+    return idf * (tf * (k1 + 1.0) / denom)
+
+
+def bm25_score_docs(
+    query_terms: Sequence[str],
+    *,
+    postings_by_term: Mapping[str, PostingList],
+    candidate_doc_ids: Iterable[int],
+    num_docs: int,
+    avgdl: float,
+    get_doc_length: Callable[[int], int],
+    cfg: BM25Config = BM25Config(),
+) -> dict[int, float]:
+    # materialize candidates once
+    cand_list = list(candidate_doc_ids)
+    cand_set = set(cand_list)
+
+    # compute idf for al query terms first
+    term_idf_all: list[tuple[str, float]] = []  # (term, idf)
+    for t in query_terms:
+        pl = postings_by_term.get(t)
+        if pl is None:
+            continue
+        df = getattr(pl, "doc_frequency", None)
+        if df is None:
+            df = len(pl.postings)
+
+        idf = bm25_idf(num_docs, int(df), clamp_negative=cfg.clamp_negative_idf)
+        term_idf_all.append((t, idf))
+
+    # order by decreasing idf
+    term_idf_all.sort(key=lambda x: x[1], reverse=True)
+
+    # thresholding
+    term_idf: dict[str, float] = {
+        t: idf for (t, idf) in term_idf_all if idf >= cfg.idf_threshold
+    }
+
+    # ensure keep at least n best terms
+    min_keep = max(1, int(cfg.min_terms_after_threshold))
+    if len(term_idf) < min_keep:
+        term_idf = {t: idf for (t, idf) in term_idf_all[:min_keep]}
+
+    # scoring only docs that actually appear in term postings (tf>0)
+    scores: dict[int, float] = {int(d): 0.0 for d in cand_list}
+
+    for t, idf in term_idf.items():
+        pl = postings_by_term.get(t)
+        if pl is None:
+            continue
+
+        tf_map = dict(pl.term_frequencies)
+
+        for doc_id, tf in tf_map.items():
+            doc_id = int(doc_id)
+            if doc_id not in cand_set:
+                continue
+
+            dl = int(get_doc_length(doc_id))
+            scores[doc_id] += bm25_term_contribution(
+                tf=int(tf),
+                doc_len=dl,
+                avgdl=avgdl,
+                idf=idf,
+                k1=cfg.k1,
+                b=cfg.b,
+            )
+
+    return scores

src/backend/search_engine/scripts/query.py

@@ -21,7 +21,7 @@ def main():
     end = time.time()
 
     for r in results.search_results:
-        print(f"[{r.document_id}] {r.title} — {r.url}")
+        print(f"[{r.document_id}] {r.title} — {r.url} — {r.score}")
 
     print(f"Total time: {end - start}s")