From 936cb8d98fbaefd1d8586bf57e222e510bb3eee8 Mon Sep 17 00:00:00 2001 From: Viet Date: Mon, 12 Jan 2026 13:55:59 +0100 Subject: [PATCH 1/2] base bm25 + idf thresholding implementation --- src/backend/search_engine/models/index.py | 1 + .../search_engine/query/query_engine.py | 34 ++++- src/backend/search_engine/scoring/__init__.py | 0 src/backend/search_engine/scoring/bm25.py | 116 ++++++++++++++++++ src/backend/search_engine/scripts/query.py | 2 +- 5 files changed, 150 insertions(+), 3 deletions(-) create mode 100644 src/backend/search_engine/scoring/__init__.py create mode 100644 src/backend/search_engine/scoring/bm25.py 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..e8d76f5 100755 --- a/src/backend/search_engine/query/query_engine.py +++ b/src/backend/search_engine/query/query_engine.py @@ -19,6 +19,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 +176,38 @@ 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 = result.postings + + # 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)) + ) + + metadata = self.inverted_index.metadata + 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), + ) + + # sort doc_ids acc to score + ranked = sorted( + ((doc_id, scores.get(doc_id, 0.0)) for doc_id in candidate_doc_ids), + key=lambda x: x[1], + reverse=True, + ) search_results = [] - for doc_id in top_n_results.postings[:limit]: + for doc_id, score in ranked[:limit]: doc_data = self.inverted_index.doc_store.get(doc_id) if doc_data is None: continue @@ -195,6 +224,7 @@ 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: 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..f6ee6c3 --- /dev/null +++ b/src/backend/search_engine/scoring/bm25.py @@ -0,0 +1,116 @@ +# 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's influence is + b: float = 0.75 # level of document normalization + # idf thresholding -> ignore terms with idf < idf_threshold + idf_threshold: float = 0.0 + # clamp negative idf to 0 + clamp_negative_idf: bool = True + + +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 + # df can be 0, if term unknown + 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]: + """ + scores only the provided candidate_doc_ids (typically: boolean result set) + """ + # precompute term idf + apply thresholding + term_idf: dict[str, float] = {} + 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) + if idf >= cfg.idf_threshold: + term_idf[t] = idf + + # fallback: if thresholding removes everything, take all terms w/o threshold + if not 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) + term_idf[t] = bm25_idf( + num_docs, int(df), clamp_negative=cfg.clamp_negative_idf + ) + + scores: dict[int, float] = {int(d): 0.0 for d in candidate_doc_ids} + + for doc_id in list(scores.keys()): + dl = int(get_doc_length(doc_id)) + s = 0.0 + for t, idf in term_idf.items(): + pl = postings_by_term.get(t) + if pl is None: + continue + tf = int(pl.term_frequencies.get(doc_id, 0)) + s += bm25_term_contribution( + tf=tf, + doc_len=dl, + avgdl=avgdl, + idf=idf, + k1=cfg.k1, + b=cfg.b, + ) + scores[doc_id] = s + + 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") From bd72c6ee9ab871cde31abb82fbdfa173f3dd0dd3 Mon Sep 17 00:00:00 2001 From: Viet Date: Tue, 24 Feb 2026 13:26:35 +0100 Subject: [PATCH 2/2] refactor + real idf thresholding --- .../search_engine/query/query_engine.py | 25 +++++-- src/backend/search_engine/scoring/bm25.py | 75 ++++++++++--------- 2 files changed, 61 insertions(+), 39 deletions(-) diff --git a/src/backend/search_engine/query/query_engine.py b/src/backend/search_engine/query/query_engine.py index e8d76f5..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] @@ -177,7 +178,8 @@ def search_results(self, limit: int = 10) -> SearchResults: ) # candidates: bool/phrase search returns doc_ids in result.postings - candidate_doc_ids = result.postings + candidate_doc_ids = list(result.postings) + metadata = self.inverted_index.metadata # score which query terms # - bool queries: qt.unique_terms @@ -188,7 +190,8 @@ def search_results(self, limit: int = 10) -> SearchResults: else list(dict.fromkeys(normalized_tokens)) ) - metadata = self.inverted_index.metadata + t_score = time.perf_counter() + scores = bm25_score_docs( query_terms=query_terms, postings_by_term=self.inverted_index.index, # term -> PostingList @@ -199,15 +202,23 @@ def search_results(self, limit: int = 10) -> SearchResults: 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 = sorted( + 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], - reverse=True, ) + logger.debug(f"Ranking sort time: {time.perf_counter() - t_sort:.6f}s") + + t_top = time.perf_counter() + search_results = [] - for doc_id, score in ranked[:limit]: + for doc_id, score in ranked_top: doc_data = self.inverted_index.doc_store.get(doc_id) if doc_data is None: continue @@ -231,6 +242,10 @@ def search_results(self, limit: int = 10) -> SearchResults: 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/bm25.py b/src/backend/search_engine/scoring/bm25.py index f6ee6c3..398b761 100644 --- a/src/backend/search_engine/scoring/bm25.py +++ b/src/backend/search_engine/scoring/bm25.py @@ -11,13 +11,16 @@ @dataclass(frozen=True) class BM25Config: - k1: float = 1.2 # how strong tf's influence is + k1: float = 1.2 # how strong tf influence is b: float = 0.75 # level of document normalization - # idf thresholding -> ignore terms with idf < idf_threshold + + # 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: """ @@ -25,7 +28,6 @@ def bm25_idf(num_docs: int, df: int, *, clamp_negative: bool = True) -> float: """ if num_docs <= 0: return 0.0 - # df can be 0, if term unknown if df <= 0: return 0.0 @@ -63,11 +65,12 @@ def bm25_score_docs( get_doc_length: Callable[[int], int], cfg: BM25Config = BM25Config(), ) -> dict[int, float]: - """ - scores only the provided candidate_doc_ids (typically: boolean result set) - """ - # precompute term idf + apply thresholding - term_idf: dict[str, 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: @@ -77,40 +80,44 @@ def bm25_score_docs( df = len(pl.postings) idf = bm25_idf(num_docs, int(df), clamp_negative=cfg.clamp_negative_idf) - if idf >= cfg.idf_threshold: - term_idf[t] = idf - - # fallback: if thresholding removes everything, take all terms w/o threshold - if not 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) - term_idf[t] = 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) - scores: dict[int, float] = {int(d): 0.0 for d in candidate_doc_ids} + # thresholding + term_idf: dict[str, float] = { + t: idf for (t, idf) in term_idf_all if idf >= cfg.idf_threshold + } - for doc_id in list(scores.keys()): - dl = int(get_doc_length(doc_id)) - s = 0.0 - for t, idf in term_idf.items(): - pl = postings_by_term.get(t) - if pl is None: + # 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 - tf = int(pl.term_frequencies.get(doc_id, 0)) - s += bm25_term_contribution( - tf=tf, + + 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, ) - scores[doc_id] = s return scores