Spaces:
Running
Running
File size: 3,747 Bytes
ed65693 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | """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]] = []
@staticmethod
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)
|