Spaces:
Running
Running
| """BM25 sparse retrieval using rank_bm25. | |
| Provides lexical keyword-based retrieval to complement dense FAISS search. | |
| Supports full-corpus scoring for building corpus-level CDFs. | |
| """ | |
| import pickle | |
| import re | |
| from pathlib import Path | |
| from typing import Any | |
| import numpy as np | |
| from rank_bm25 import BM25Okapi | |
| class BM25Index: | |
| """BM25Okapi-based sparse retrieval index with disk persistence.""" | |
| def __init__(self) -> None: | |
| self.bm25: BM25Okapi | None = None | |
| self.tokenized_corpus: list[list[str]] = [] | |
| self.doc_metadata: list[dict[str, Any]] = [] | |
| def tokenize(text: str) -> list[str]: | |
| """Lowercase, strip punctuation, split on whitespace.""" | |
| text = text.lower() | |
| text = re.sub(r"[^\w\s]", " ", text) | |
| return [tok for tok in text.split() if len(tok) > 1] | |
| def add_documents(self, documents: list[dict[str, Any]]) -> None: | |
| """Add documents to the BM25 index. | |
| Args: | |
| documents: List of dicts with keys 'text', 'source', 'chunk'. | |
| """ | |
| for doc in documents: | |
| tokens = self.tokenize(doc["text"]) | |
| self.tokenized_corpus.append(tokens) | |
| self.doc_metadata.append( | |
| { | |
| "source": doc["source"], | |
| "chunk": doc["chunk"], | |
| "text": doc["text"], | |
| } | |
| ) | |
| self.bm25 = BM25Okapi(self.tokenized_corpus) | |
| def search(self, query: str, top_k: int = 20) -> list[dict[str, Any]]: | |
| """Retrieve top-k documents by BM25 score. | |
| Returns list of dicts with 'source', 'chunk', 'text', 'score'. | |
| """ | |
| if self.bm25 is None: | |
| return [] | |
| tokens = self.tokenize(query) | |
| scores = self.bm25.get_scores(tokens) | |
| top_indices = np.argsort(scores)[::-1][:top_k] | |
| results = [] | |
| for idx in top_indices: | |
| if scores[idx] > 0: | |
| results.append( | |
| { | |
| "source": self.doc_metadata[idx]["source"], | |
| "chunk": self.doc_metadata[idx]["chunk"], | |
| "text": self.doc_metadata[idx]["text"], | |
| "score": float(scores[idx]), | |
| } | |
| ) | |
| return results | |
| def score_all(self, query: str) -> np.ndarray: | |
| """Score ALL documents for a query. Used for corpus-level CDF construction.""" | |
| if self.bm25 is None: | |
| return np.array([]) | |
| tokens = self.tokenize(query) | |
| return self.bm25.get_scores(tokens) | |
| def corpus_mean_score(self, query: str) -> float: | |
| """Mean BM25 score across the entire corpus for a query. | |
| Approximation of s(q, C) from Shtok et al. (2012) for NQC computation. | |
| """ | |
| scores = self.score_all(query) | |
| return float(np.mean(scores)) if len(scores) > 0 else 0.0 | |
| def save(self, path: Path) -> None: | |
| """Persist tokenized corpus and metadata to disk.""" | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| with open(path, "wb") as f: | |
| pickle.dump( | |
| { | |
| "tokenized_corpus": self.tokenized_corpus, | |
| "doc_metadata": self.doc_metadata, | |
| }, | |
| f, | |
| ) | |
| def load(self, path: Path) -> None: | |
| """Load tokenized corpus and rebuild BM25 index.""" | |
| if not path.exists(): | |
| return | |
| with open(path, "rb") as f: | |
| data = pickle.load(f) # noqa: S301 | |
| self.tokenized_corpus = data["tokenized_corpus"] | |
| self.doc_metadata = data["doc_metadata"] | |
| if self.tokenized_corpus: | |
| self.bm25 = BM25Okapi(self.tokenized_corpus) | |