model: Add Bayesian BM25 (bb25) baseline model (#4084)

jaepil · web-flow · commit d32b358ed84d · 2026-02-13T11:03:54.000+01:00
* model: Add Bayesian BM25 (bb25) baseline model

Add baseline/bb25, a Bayesian BM25 model that transforms BM25 scores
into calibrated probabilities via sigmoid likelihood and optional
composite prior. Uses bm25s as the BM25 backend.

With default prior_weight=0.0, rankings are identical to bm25s while
outputting calibrated probabilities in [0, 1] for hybrid search fusion.

* test: Compare bb25 against bm25s instead of hardcoded values

* docs: Clarify alpha scaling comment with paper context

* style: Apply ruff format to bb25.py

* refactor: Rename encode to _encode to avoid EncoderProtocol confusion

* refactor: Rename encode to _encode in bm25.py for consistency

* test: Merge bb25 test into parametrized bm25 test
diff --git a/mteb/models/model_implementations/bb25.py b/mteb/models/model_implementations/bb25.py
@@ -0,0 +1,311 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from mteb._create_dataloaders import _create_text_queries_dataloader
+from mteb._requires_package import requires_package
+from mteb.models.model_meta import ModelMeta
+
+if TYPE_CHECKING:
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models.models_protocols import SearchProtocol
+    from mteb.types import (
+        CorpusDatasetType,
+        EncodeKwargs,
+        QueryDatasetType,
+        RetrievalOutputType,
+        TopRankedDocumentsType,
+    )
+
+logger = logging.getLogger(__name__)
+
+
+def _composite_prior(
+    query_tfs: np.ndarray,
+    doc_lengths: np.ndarray,
+    avg_dl: float,
+) -> np.ndarray:
+    """Composite prior from Bayesian BM25 paper Section 4.2.
+
+    Combines term frequency prior (Def 4.2.1) and field norm prior
+    (Def 4.2.2) to produce document-level prior probabilities.
+
+    Args:
+        query_tfs: Total query term frequency per candidate document.
+        doc_lengths: Token count per candidate document.
+        avg_dl: Average document length across the corpus.
+
+    Returns:
+        Prior probabilities clipped to [0.1, 0.9].
+    """
+    p_tf = 0.2 + 0.7 * np.minimum(1.0, query_tfs / 10.0)
+
+    if avg_dl > 0:
+        norm_ratio = doc_lengths / avg_dl
+    else:
+        norm_ratio = np.ones_like(doc_lengths, dtype=np.float64)
+    p_norm = 0.3 + 0.6 * (1.0 - np.minimum(1.0, np.abs(norm_ratio - 1.0) * 0.5))
+
+    prior = 0.7 * p_tf + 0.3 * p_norm
+    return np.clip(prior, 0.1, 0.9)
+
+
+def bb25_loader(model_name, **kwargs) -> SearchProtocol:
+    requires_package(bb25_loader, "bm25s", model_name, "pip install mteb[bm25s]")
+    import bm25s
+    import Stemmer
+
+    class BB25Search:
+        """Bayesian BM25 search using bm25s as the BM25 backend.
+
+        Bayesian BM25 transforms traditional BM25 scores into calibrated
+        probability estimates in [0, 1] through Bayesian inference with a
+        sigmoid likelihood model and composite prior design.
+
+        With the default prior_weight=0.0, the prior is a flat 0.5 for
+        all documents. Since the sigmoid likelihood is strictly monotonic,
+        this preserves BM25 rankings exactly while outputting calibrated
+        probabilities suitable for hybrid search score fusion.
+
+        Setting prior_weight > 0 enables the Composite Prior (Section 4.2
+        of the paper), which re-adjusts rankings based on document-level
+        evidence from query term frequency and document length signals.
+
+        Architecture:
+            1. bm25s handles fast BM25 indexing and top-k retrieval.
+            2. Retrieved candidates are re-scored with Bayesian posterior:
+               - Dynamic beta = median(BM25 scores) per query
+               - Sigmoid likelihood: sigma(alpha * (score - beta))
+               - prior = 0.5 + prior_weight * (composite_prior - 0.5)
+               - Posterior via Bayes' rule
+        """
+
+        retriever: bm25s.BM25
+        corpus_idx_to_id: dict[int, str]
+
+        def __init__(
+            self,
+            previous_results: str | None = None,
+            stopwords: str = "en",
+            stemmer_language: str | None = "english",
+            k1: float = 1.5,
+            b: float = 0.75,
+            alpha: float = 1.0,
+            prior_weight: float = 0.0,
+            **kwargs,
+        ):
+            self.k1 = k1
+            self.b = b
+            self.alpha = alpha
+            self.prior_weight = prior_weight
+            self.stopwords = stopwords
+            self.stemmer = (
+                Stemmer.Stemmer(stemmer_language) if stemmer_language else None
+            )
+
+        def _encode(self, texts: list[str]):
+            """Tokenize texts using bm25s. Not to be confused with EncoderProtocol.encode()."""
+            return bm25s.tokenize(texts, stopwords=self.stopwords, stemmer=self.stemmer)
+
+        def index(
+            self,
+            corpus: CorpusDatasetType,
+            *,
+            task_metadata: TaskMetadata,
+            hf_split: str,
+            hf_subset: str,
+            encode_kwargs: EncodeKwargs,
+            num_proc: int | None = None,
+        ) -> None:
+            logger.info("Encoding Corpus...")
+            corpus_texts = [
+                "\n".join([doc.get("title", ""), doc["text"]]) for doc in corpus
+            ]
+            encoded_corpus = self._encode(corpus_texts)
+
+            logger.info(
+                f"Indexing Corpus... {len(encoded_corpus.ids):,} documents, "
+                f"{len(encoded_corpus.vocab):,} vocab"
+            )
+
+            self.retriever = bm25s.BM25(k1=self.k1, b=self.b)
+            self.retriever.index(encoded_corpus)
+            self.corpus_idx_to_id = {i: row["id"] for i, row in enumerate(corpus)}
+
+            if self.prior_weight > 0:
+                # Per-document token IDs for composite prior computation.
+                # Stored as compact numpy int32 arrays to minimize memory.
+                self.corpus_token_ids = [
+                    np.array(doc_ids, dtype=np.int32) for doc_ids in encoded_corpus.ids
+                ]
+                self.corpus_vocab = dict(encoded_corpus.vocab)
+                self.doc_lengths = np.array(
+                    [len(ids) for ids in encoded_corpus.ids], dtype=np.float64
+                )
+                self.avg_dl = (
+                    float(self.doc_lengths.mean()) if len(self.doc_lengths) > 0 else 0.0
+                )
+
+            logger.info(f"Indexed {len(self.corpus_idx_to_id):,} documents")
+
+        def search(
+            self,
+            queries: QueryDatasetType,
+            *,
+            task_metadata: TaskMetadata,
+            hf_split: str,
+            hf_subset: str,
+            top_k: int,
+            encode_kwargs: EncodeKwargs,
+            top_ranked: TopRankedDocumentsType | None = None,
+            num_proc: int | None = None,
+        ) -> RetrievalOutputType:
+            logger.info("Encoding Queries...")
+            query_ids = list(queries["id"])
+            results: RetrievalOutputType = {qid: {} for qid in query_ids}
+            queries_loader = _create_text_queries_dataloader(queries)
+            queries_texts = [text for batch in queries_loader for text in batch["text"]]
+
+            query_tokenized = self._encode(queries_texts)
+
+            logger.info(f"Retrieving Results... {len(queries):,} queries")
+
+            queries_results, queries_scores = self.retriever.retrieve(
+                query_tokenized,
+                k=min(top_k, len(self.corpus_idx_to_id)),
+            )
+
+            use_prior = self.prior_weight > 0
+            if use_prior:
+                query_id_to_str = {v: k for k, v in query_tokenized.vocab.items()}
+
+            for qi, qid in enumerate(query_ids):
+                doc_indices = queries_results[qi]
+                bm25_scores = queries_scores[qi].astype(np.float64)
+
+                query_documents = (
+                    top_ranked[qid] if top_ranked and qid in top_ranked else None
+                )
+
+                doc_id_to_score: dict[str, float] = {}
+
+                # Separate positive-score candidates for Bayesian re-scoring
+                positive_mask = bm25_scores > 0
+                positive_indices = np.where(positive_mask)[0]
+
+                if len(positive_indices) > 0:
+                    cand_doc_indices = doc_indices[positive_indices]
+                    cand_bm25_scores = bm25_scores[positive_indices]
+
+                    # Dynamic beta = median of BM25 scores for this query
+                    beta = float(np.median(cand_bm25_scores))
+
+                    # Dynamic alpha scaling for query-level score distribution
+                    # invariance. The paper defines alpha as sigmoid steepness;
+                    # dividing by std(scores) keeps the effective steepness
+                    # consistent across queries whose BM25 ranges vary widely,
+                    # preventing sigmoid saturation on high-scoring queries.
+                    score_std = float(np.std(cand_bm25_scores))
+                    alpha_eff = (
+                        self.alpha / score_std if score_std > 1e-10 else self.alpha
+                    )
+
+                    # Sigmoid likelihood (monotonic -- preserves BM25 ranking)
+                    x = np.clip(alpha_eff * (cand_bm25_scores - beta), -500, 500)
+                    likelihood = 1.0 / (1.0 + np.exp(-x))
+
+                    if use_prior:
+                        # Map query token IDs to corpus vocab IDs
+                        query_term_strs = [
+                            query_id_to_str[tid] for tid in query_tokenized.ids[qi]
+                        ]
+                        query_corpus_ids = np.array(
+                            [
+                                self.corpus_vocab[t]
+                                for t in query_term_strs
+                                if t in self.corpus_vocab
+                            ],
+                            dtype=np.int32,
+                        )
+
+                        n_cand = len(cand_doc_indices)
+                        query_tfs = np.zeros(n_cand, dtype=np.float64)
+                        cand_doc_lengths = np.zeros(n_cand, dtype=np.float64)
+
+                        has_query_terms = len(query_corpus_ids) > 0
+                        for ci in range(n_cand):
+                            doc_idx = cand_doc_indices[ci]
+                            cand_doc_lengths[ci] = self.doc_lengths[doc_idx]
+                            if has_query_terms:
+                                doc_ids = self.corpus_token_ids[doc_idx]
+                                query_tfs[ci] = np.isin(doc_ids, query_corpus_ids).sum()
+
+                        composite = _composite_prior(
+                            query_tfs, cand_doc_lengths, self.avg_dl
+                        )
+                        # Interpolate between flat prior (0.5) and composite
+                        prior = 0.5 + self.prior_weight * (composite - 0.5)
+                        posterior = (likelihood * prior) / (
+                            likelihood * prior + (1.0 - likelihood) * (1.0 - prior)
+                        )
+                    else:
+                        # prior=0.5 => posterior = likelihood
+                        posterior = likelihood
+
+                    for ci in range(len(cand_doc_indices)):
+                        doc_id = self.corpus_idx_to_id[cand_doc_indices[ci]]
+                        if query_documents is None or doc_id in query_documents:
+                            doc_id_to_score[doc_id] = float(posterior[ci])
+
+                # Include zero-score documents with score 0.0
+                for vi in np.where(~positive_mask)[0]:
+                    doc_id = self.corpus_idx_to_id[doc_indices[vi]]
+                    if query_documents is None or doc_id in query_documents:
+                        doc_id_to_score[doc_id] = 0.0
+
+                results[qid] = doc_id_to_score
+
+            return results
+
+    return BB25Search(**kwargs)
+
+
+bb25_model = ModelMeta(
+    loader=bb25_loader,
+    name="baseline/bb25",
+    model_type=["sparse"],
+    languages=None,
+    open_weights=True,
+    revision="0_1_1",
+    release_date="2026-02-06",
+    n_parameters=None,
+    n_embedding_parameters=None,
+    memory_usage_mb=None,
+    embed_dim=None,
+    license=None,
+    max_tokens=None,
+    reference="https://github.com/instructkr/bb25",
+    similarity_fn_name=None,
+    framework=[],
+    use_instructions=False,
+    public_training_code="https://github.com/instructkr/bb25",
+    public_training_data=None,
+    training_datasets=None,
+    citation="""@software{jeong2026bayesianbm25,
+  title={Bayesian BM25: A Probabilistic Framework for Hybrid Text and Vector Search},
+  author={Jeong, Jaepil},
+  year={2026},
+  doi={10.5281/zenodo.18414941},
+  url={https://doi.org/10.5281/zenodo.18414941},
+}
+@software{jeong2026neural,
+  title={From Bayesian Inference to Neural Computation: The Analytical Emergence of Neural Network Structure from Probabilistic Relevance Estimation},
+  author={Jeong, Jaepil},
+  year={2026},
+  doi={10.5281/zenodo.18512411},
+  url={https://doi.org/10.5281/zenodo.18512411},
+}""",
+)
diff --git a/mteb/models/model_implementations/bm25.py b/mteb/models/model_implementations/bm25.py
@@ -60,7 +60,7 @@ def index(
             corpus_texts = [
                 "\n".join([doc.get("title", ""), doc["text"]]) for doc in corpus
             ]  # concatenate all document values (title, text, ...)
-            encoded_corpus = self.encode(corpus_texts)
+            encoded_corpus = self._encode(corpus_texts)
 
             logger.info(
                 f"Indexing Corpus... {len(encoded_corpus.ids):,} documents, {len(encoded_corpus.vocab):,} vocab"
@@ -89,7 +89,7 @@ def search(
             queries_loader = _create_text_queries_dataloader(queries)
             queries_texts = [text for batch in queries_loader for text in batch["text"]]
 
-            query_token_strs = self.encode(queries_texts)
+            query_token_strs = self._encode(queries_texts)
 
             logger.info(f"Retrieving Results... {len(queries):,} queries")
 
@@ -120,8 +120,8 @@ def search(
 
             return results
 
-        def encode(self, texts: list[str]):
-            """Encode input text as term vectors"""
+        def _encode(self, texts: list[str]):
+            """Tokenize texts using bm25s. Not to be confused with EncoderProtocol.encode()."""
             return bm25s.tokenize(texts, stopwords=self.stopwords, stemmer=self.stemmer)
 
     return BM25Search(**kwargs)
diff --git a/tests/test_models/test_bm25.py b/tests/test_models/test_bm25.py
@@ -5,12 +5,13 @@
 import mteb
 
 
-def test_bm25s_e2e():
+@pytest.mark.parametrize("model_name", ["bm25s", "baseline/bb25"])
+def test_bm25_e2e(model_name):
     # fails for dataset smaller than 1000
     pytest.importorskip("bm25s", reason="bm25s not installed")
     pytest.importorskip("Stemmer", reason="PyStemmer not installed")
 
-    model = mteb.get_model("bm25s")
+    model = mteb.get_model(model_name)
     task = mteb.get_task("NFCorpus", eval_splits=["test"])
 
     results = mteb.evaluate(model, task, cache=None)
diff --git a/tests/test_models/test_model_meta.py b/tests/test_models/test_model_meta.py
@@ -30,6 +30,7 @@
     "TomoroAI/tomoro-colqwen3-embed-8b",
     "VPLabs/SearchMap_Preview",
     "ai-sage/Giga-Embeddings-instruct",
+    "baseline/bb25",
     "baseline/bm25s",
     "codefuse-ai/C2LLM-0.5B",
     "codefuse-ai/C2LLM-7B",
