Add optional Pyserini BM25 retriever backend

README.md

@@ -19,6 +19,12 @@ Our goal is to be complimentary to your existing pipelines!
 ```bash
 pip install retri-evals
 ```
+
+To enable the optional Pyserini BM25 baseline retriever:
+
+```bash
+pip install "retri-evals[pyserini]"
+```
 ### Define your data type
 We use Pydantic to make sure that the index receives the expected data.
 
@@ -163,6 +169,7 @@ results:
 retri-eval is still in active development. We're planning to add the following functionality:
 
 - [ ] Support reranking models
+- [x] Add lexical BM25 baseline via Pyserini (optional extra)
 - [ ] Add support for hybrid retrieval baselines
 - [ ] Support for automatic dataset generation
 - [ ] Support parallel execution

examples/retrieval.py

@@ -4,6 +4,7 @@
 
 from retri_eval.evaluation.mteb_tasks import CQADupstackEnglishRetrieval, Touche2020
 from retri_eval.evaluation.retriever import DenseRetriever
+from retri_eval.evaluation.pyserini_retriever import PyseriniBM25Retriever
 from retri_eval.indexes.qdrant_index import QdrantIndex, QdrantDocument
 from retri_eval.processing.basic_query_processor import QueryProcessor
 from retri_eval.processing.beir_title_processor import BeirTitleProcessor
@@ -33,11 +34,15 @@ def encode_queries(self, batch, **kwargs):
     doc_processor = BeirTitleProcessor(model, name=model_name)
     query_processor = QueryProcessor(model, name=model_name)
 
-    retriever = DenseRetriever(
-        index=index,
-        query_processor=query_processor,
-        doc_processor=doc_processor,
-    )
+    # Choose a retriever backend:
+    # - DenseRetriever: bring your own vector index (e.g., Qdrant)
+    # - PyseriniBM25Retriever: lexical BM25 baseline (indexes the BEIR corpus on-the-fly)
+    retriever = PyseriniBM25Retriever()  # "pyserini-bm25"
+    # retriever = DenseRetriever(
+    #     index=index,
+    #     query_processor=query_processor,
+    #     doc_processor=doc_processor,
+    # )
 
     id = f"{doc_processor.id}-{query_processor.id}"
     print(f"evaluation id: {id}")

pyproject.toml

@@ -42,6 +42,11 @@ transformers = "^4.36.2"
 sentence-transformers = "^2.2.2"
 usearch = "^2.8.15"
 openai = "<=0.28.1"
+pyserini = {version = ">=0.24.0", optional = true}
+pytrec-eval = {version = ">=0.5", optional = true}
+
+[tool.poetry.extras]
+pyserini = ["pyserini", "pytrec-eval"]
 
 [tool.poetry.group.examples]
 optional = true

retri_eval/evaluation/pyserini_retriever.py

@@ -0,0 +1,102 @@
+"""Pyserini-backed lexical retriever for retri-evals.
+
+This integrates as a BEIR/MTEB BaseSearch model so it can be evaluated using the
+existing IndexedTask/EvaluateRetrieval path.
+
+Design goals:
+- Minimal surface area: keep it optional via extras.
+- Index once per evaluation run.
+- Return BEIR-style results dict: {qid: {doc_id: score}}.
+
+Notes:
+- Pyserini BM25 scores are not directly comparable to dense similarity scores;
+  BEIR metrics only require ranking.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, List, Tuple
+
+from beir.retrieval.search import BaseSearch
+
+
+@dataclass
+class PyseriniRetrieverMetrics:
+    latencies: List[float]
+
+
+class PyseriniBM25Retriever(BaseSearch):
+    def __init__(
+        self,
+        k1: float = 0.9,
+        b: float = 0.4,
+        threads: int = 1,
+    ):
+        super().__init__()
+        self.k1 = k1
+        self.b = b
+        self.threads = threads
+
+    def search(
+        self,
+        corpus: Dict[str, Dict[str, str]],
+        queries: Dict[str, str],
+        top_k: int,
+        score_function: str = "bm25",
+        **kwargs,
+    ) -> Tuple[Dict[str, Dict[str, float]], PyseriniRetrieverMetrics]:
+        # Lazy imports so this module is optional.
+        import json
+        import os
+        import tempfile
+        import time
+
+        from pyserini.index.lucene import LuceneIndexer
+        from pyserini.search.lucene import LuceneSearcher
+
+        if score_function not in {"bm25", "pyserini"}:
+            raise ValueError(
+                f"PyseriniBM25Retriever only supports score_function='bm25'/'pyserini' (got {score_function})"
+            )
+
+        # Build a temporary jsonl collection and Lucene index.
+        # Pyserini expects each doc to have at least: {id, contents}.
+        with tempfile.TemporaryDirectory(prefix="retri-evals-pyserini-") as tmp:
+            collection_path = os.path.join(tmp, "collection.jsonl")
+            index_dir = os.path.join(tmp, "index")
+
+            with open(collection_path, "w", encoding="utf-8") as f:
+                for doc_id, doc in corpus.items():
+                    # BEIR corpora often have title/text.
+                    title = (doc.get("title") or "").strip()
+                    text = (doc.get("text") or "").strip()
+                    contents = (title + "\n\n" + text).strip() if title else text
+                    f.write(json.dumps({"id": doc_id, "contents": contents}, ensure_ascii=False) + "\n")
+
+            LuceneIndexer(index_dir=index_dir).index(collection_path)
+
+            searcher = LuceneSearcher(index_dir)
+            searcher.set_bm25(k1=self.k1, b=self.b)
+
+            results: Dict[str, Dict[str, float]] = {qid: {} for qid in queries.keys()}
+            latencies: List[float] = []
+
+            for qid, query in queries.items():
+                start = time.perf_counter()
+                hits = searcher.search(query, k=top_k, threads=self.threads)
+                latencies.append(time.perf_counter() - start)
+
+                for hit in hits:
+                    # hit.docid is the string id we provided.
+                    results[qid][hit.docid] = float(hit.score)
+
+            return results, PyseriniRetrieverMetrics(latencies=latencies)
+
+    # BEIR BaseSearch interface includes encode_* for dense retrievers.
+    # For BM25 these are unused.
+    def encode_queries(self, queries: List[str], batch_size: int, **kwargs):
+        raise NotImplementedError("PyseriniBM25Retriever is lexical; no query encoding")
+
+    def encode_corpus(self, corpus: List[Dict[str, str]], batch_size: int, **kwargs):
+        raise NotImplementedError("PyseriniBM25Retriever is lexical; no corpus encoding")