diff --git a/README.md b/README.md index 6dc35c6..5a6f260 100644 --- a/README.md +++ b/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 diff --git a/examples/retrieval.py b/examples/retrieval.py index 0573469..b23b58d 100644 --- a/examples/retrieval.py +++ b/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}") diff --git a/pyproject.toml b/pyproject.toml index a1915a5..3d75f26 100644 --- a/pyproject.toml +++ b/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 diff --git a/retri_eval/evaluation/pyserini_retriever.py b/retri_eval/evaluation/pyserini_retriever.py new file mode 100644 index 0000000..ddb80d0 --- /dev/null +++ b/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")