diff --git a/src/backend/search_engine/models/index.py b/src/backend/search_engine/models/index.py index 3728ead..6ee1836 100644 --- a/src/backend/search_engine/models/index.py +++ b/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): diff --git a/src/backend/search_engine/query/query_engine.py b/src/backend/search_engine/query/query_engine.py index 10e7212..ac9aef6 100755 --- a/src/backend/search_engine/query/query_engine.py +++ b/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. " diff --git a/src/backend/search_engine/scoring/__init__.py b/src/backend/search_engine/scoring/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/backend/search_engine/scoring/bm25.py b/src/backend/search_engine/scoring/bm25.py new file mode 100644 index 0000000..398b761 --- /dev/null +++ b/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 diff --git a/src/backend/search_engine/scripts/query.py b/src/backend/search_engine/scripts/query.py index f2905b5..7555eca 100755 --- a/src/backend/search_engine/scripts/query.py +++ b/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")