From 88177bd7a680490ac85d244c5fff90f21a3be27c Mon Sep 17 00:00:00 2001 From: Leoy Date: Sat, 6 Jun 2026 06:15:44 +0800 Subject: [PATCH] feat(relevance): weight BM25 score_batch by corpus IDF (#646) BM25Scorer.score_batch() ranks a real corpus of documents but weighted every matched term with a constant idf=log(2.0), so a ubiquitous noise word counted the same as a discriminative UUID. The _compute_idf() helper needed to do this properly already existed (and was unit-tested) but was never wired into scoring. - Implement _compute_idf() with the standard floored BM25 IDF its docstring documents: log((N - n + 0.5) / (n + 0.5) + 1). - Thread an optional per-term idf_map through _bm25_score(); single-document score() keeps the neutral log(2.0) weight (no corpus to estimate from). - score_batch() now computes document frequency across the batch and builds the IDF map, so rare/discriminative terms outrank corpus-wide terms in the ranking that CompressionStore.search() and HybridScorer consume. Adds tests covering the IDF formula and the batch ranking behaviour. --- headroom/relevance/bm25.py | 47 ++++++++++++++++++++++++++++++-------- tests/test_relevance.py | 41 +++++++++++++++++++++++++++++++++ 2 files changed, 79 insertions(+), 9 deletions(-) diff --git a/headroom/relevance/bm25.py b/headroom/relevance/bm25.py index 97c96362c..2e83204f8 100644 --- a/headroom/relevance/bm25.py +++ b/headroom/relevance/bm25.py @@ -103,22 +103,26 @@ class BM25Scorer(RelevanceScorer): Uses the standard BM25 IDF formula: IDF = log((N - n + 0.5) / (n + 0.5) + 1) - Where N = total docs, n = docs containing term. + Where N = total docs (``doc_count``) and n = docs containing the + term (``doc_freq``). The ``+ 1`` keeps the result non-negative even + when a term appears in more than half the corpus, which is the + floored variant of BM25 used by Lucene/Elasticsearch. - For single-document scoring, we use a simplified version. + A term that occurs in few documents (low ``doc_freq``) is more + discriminative and earns a higher IDF; a term that occurs in nearly + every document earns an IDF approaching ``log(1) = 0``. """ - if doc_freq == 0: + if doc_freq <= 0: return 0.0 - # Simplified IDF for single-document case - # Term present = higher IDF, term absent = 0 - return math.log(2.0) # Constant since we have single document + return math.log((doc_count - doc_freq + 0.5) / (doc_freq + 0.5) + 1.0) def _bm25_score( self, doc_tokens: list[str], query_tokens: list[str], avg_doc_len: float | None = None, + idf_map: dict[str, float] | None = None, ) -> tuple[float, list[str]]: """Compute BM25 score between document and query. @@ -126,6 +130,13 @@ class BM25Scorer(RelevanceScorer): doc_tokens: Tokenized document. query_tokens: Tokenized query. avg_doc_len: Average document length (optional). + idf_map: Pre-computed corpus IDF per term. When supplied (batch + scoring, where a real corpus exists) each term is weighted by + its inverse document frequency, so a discriminative term such + as a UUID outranks a term that is common across the corpus. + When ``None`` (single-document scoring, where there is no + corpus to estimate IDF from) every term falls back to the + neutral ``log(2.0)`` weight, preserving the original behaviour. Returns: Tuple of (score, matched_terms). @@ -149,8 +160,9 @@ class BM25Scorer(RelevanceScorer): f = doc_freq[term] matched_terms.append(term) - # BM25 term score - idf = math.log(2.0) # Simplified for single doc + # BM25 term score. Use the corpus IDF when available; otherwise + # fall back to the neutral single-document weight. + idf = idf_map.get(term, math.log(2.0)) if idf_map is not None else math.log(2.0) numerator = f * (self.k1 + 1) denominator = f + self.k1 * (1 - self.b + self.b * doc_len / avgdl) @@ -223,9 +235,26 @@ class BM25Scorer(RelevanceScorer): all_tokens = [self._tokenize(item) for item in items] avg_len = sum(len(t) for t in all_tokens) / max(len(items), 1) + # Compute corpus IDF per query term. Unlike single-item scoring, + # a batch is a real corpus, so document frequency is meaningful: + # terms that appear in many items are down-weighted while rare, + # discriminative terms (IDs, UUIDs) are boosted. This is what makes + # the ranking BM25 rather than plain term-frequency weighting. + n_docs = len(all_tokens) + doc_freq_across: Counter[str] = Counter() + for tokens in all_tokens: + doc_freq_across.update(set(tokens)) + idf_map = { + term: self._compute_idf(term, n_docs, doc_freq_across[term]) + for term in set(context_tokens) + if term in doc_freq_across + } + results = [] for item_tokens in all_tokens: - raw_score, matched = self._bm25_score(item_tokens, context_tokens, avg_doc_len=avg_len) + raw_score, matched = self._bm25_score( + item_tokens, context_tokens, avg_doc_len=avg_len, idf_map=idf_map + ) # Normalize if self.normalize_score: diff --git a/tests/test_relevance.py b/tests/test_relevance.py index 6da26e1a0..f5941c42f 100644 --- a/tests/test_relevance.py +++ b/tests/test_relevance.py @@ -107,6 +107,47 @@ class TestBM25Scorer: """BM25Scorer is always available.""" assert BM25Scorer.is_available() + def test_compute_idf_follows_standard_formula(self): + """IDF rewards rare terms and decays toward zero for common terms.""" + scorer = BM25Scorer() + + # Absent term contributes nothing. + assert scorer._compute_idf("x", doc_count=10, doc_freq=0) == 0.0 + + # A term in 1/10 docs is more discriminative than one in 9/10 docs. + rare = scorer._compute_idf("x", doc_count=10, doc_freq=1) + common = scorer._compute_idf("x", doc_count=10, doc_freq=9) + assert rare > common > 0 + + def test_batch_idf_downweights_common_terms(self): + """A discriminative term outranks one shared across the whole corpus. + + ``shared`` appears in every item, so its corpus IDF approaches zero, + while ``zeta`` appears in a single item and stays discriminative. An + item matched only on the rare term must therefore outrank an item + matched only on the ubiquitous term. + """ + scorer = BM25Scorer() + items = [ + "shared zeta", # matches both query terms, one of them rare + "shared alpha", # matches only the ubiquitous term + "shared beta", + "shared gamma", + ] + context = "shared zeta" + + scores = scorer.score_batch(items, context) + assert scores[0].score > scores[1].score + assert scores[0].score > scores[2].score + + def test_batch_idf_does_not_change_matched_terms(self): + """Corpus IDF affects ranking only, not which terms are reported.""" + scorer = BM25Scorer() + items = ["alpha", "alpha beta"] + scores = scorer.score_batch(items, "alpha beta") + assert scores[0].matched_terms == ["alpha"] + assert sorted(scores[1].matched_terms) == ["alpha", "beta"] + class TestEmbeddingScorer: """Tests for embedding-based semantic scorer."""