File size: 6,882 Bytes
0f5ecaf | 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 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 | """
bm25_backends.py (Improved Version)
Just what it sounds like
"""
from typing import List, Dict, Tuple
import math
import time
import numpy as np
from collections import defaultdict, Counter
# Check for bm25s availability
try:
import bm25s as _bm25s
_BM25S_AVAILABLE = True
_BM25S_ERR = ""
except Exception as _e:
_bm25s = None
_BM25S_AVAILABLE = False
_BM25S_ERR = str(_e)
class AbstractBM25Backend:
"""Abstract base class for a BM25 implementation."""
def __init__(self, tokenizer):
self.tokenizer = tokenizer
self.doc_ids: List[str] = []
def build(self, ids: List[str], texts: List[str]):
raise NotImplementedError
def search(self, query: str, topk: int = 300) -> List[str]:
raise NotImplementedError
@property
def name(self) -> str:
return self.__class__.__name__
class BM25SBackend(AbstractBM25Backend):
"""
High-performance and reliable wrapper for the 'bm25s' library.
- Uses the fast `retrieve` method for efficient top-k search.
- Uses stable sorting (`lexsort`) for deterministic tie-breaking.
- Allows configurable k1 and b parameters.
"""
def __init__(self, tokenizer, k1: float = 1.3, b: float = 0.7):
if not _BM25S_AVAILABLE:
raise ImportError(f"bm25s library not available: {_BM25S_ERR}")
super().__init__(tokenizer)
self.k1 = k1
self.b = b
self._bm25 = None
@property
def name(self) -> str:
return f"BM25SBackend(k1={self.k1}, b={self.b})"
def build(self, ids: List[str], texts: List[str]):
from bm25s import BM25
self.doc_ids = list(ids)
t0 = time.time()
tokenized_corpus = [self.tokenizer(t) for t in texts]
self._bm25 = BM25(k1=self.k1, b=self.b)
self._bm25.index(tokenized_corpus)
print(f"[{self.name}] Indexed {len(self.doc_ids):,} documents in {time.time() - t0:.2f}s")
def search(self, query: str, topk: int = 300) -> List[str]:
tokenized_query = self.tokenizer(query)
if not tokenized_query or self._bm25 is None:
return []
k = min(topk, len(self.doc_ids))
if k == 0:
return []
# bm25s API compatibility: newer accepts positional list + k; older may need positional only
try:
doc_indices, scores = self._bm25.retrieve([tokenized_query], k=k)
except TypeError:
try:
doc_indices, scores = self._bm25.retrieve([tokenized_query], k)
except TypeError:
# very old API uses 'topk' name
doc_indices, scores = self._bm25.retrieve([tokenized_query], topk=k)
doc_indices, scores = doc_indices[0], scores[0]
mask = np.isfinite(scores) & (scores > 0)
doc_indices = doc_indices[mask]
scores = scores[mask]
if len(doc_indices) == 0:
return []
order = np.lexsort((doc_indices, -scores)) # stable: by -score then doc idx
final_indices = doc_indices[order]
return [self.doc_ids[int(i)] for i in final_indices]
# The pure-Python fallback remains the same, as it was already reliable.
class DeterministicBM25Backend(AbstractBM25Backend):
"""Pure-Python deterministic BM25. Slower but a good reference."""
def __init__(self, tokenizer, k1: float = 1.3, b: float = 0.7):
super().__init__(tokenizer)
self.k1 = k1
self.b = b
self.N = 0
self.avgdl = 0.0
self.doc_lens = None
self.vocab = {}
self.postings = {}
self.idf = None
@property
def name(self) -> str:
return f"DeterministicBM25Backend(k1={self.k1}, b={self.b})"
def build(self, ids: List[str], texts: List[str]):
self.doc_ids=list(ids)
self.N=len(ids)
lens=np.zeros(self.N,dtype=np.int32)
tmp=defaultdict(list)
t0=time.time()
for i, text in enumerate(texts):
terms=self.tokenizer(text); lens[i]=len(terms)
if not terms: continue
ctr=Counter(terms)
for t,tf in ctr.items():
tid=self.vocab.setdefault(t, len(self.vocab))
tmp[tid].append((i, tf))
self.doc_lens=lens
self.avgdl=float(np.maximum(1,lens).mean())
V=len(self.vocab)
self.idf=np.zeros(V,dtype=np.float32)
self.postings={}
for tid, pairs in tmp.items():
docs=np.array([d for d,_ in pairs],dtype=np.int32)
tfs =np.array([tf for _,tf in pairs],dtype=np.float32)
df=float(len(docs))
idf=math.log((self.N-df+0.5)/(df+0.5)+1.0)
self.idf[tid]=idf
self.postings[tid]=(docs,tfs)
print(f"[{self.name}] Indexed {self.N:,} documents in {time.time() - t0:.2f}s")
def search(self, query: str, topk: int = 300) -> List[str]:
terms=self.tokenizer(query)
if not terms: return []
seen: Dict[int,float] = {}
for t in terms:
tid=self.vocab.get(t)
if tid is None: continue
idf=self.idf[tid]
docs,tfs=self.postings[tid]
denom=tfs + self.k1*(1-self.b + self.b*(self.doc_lens[docs]/self.avgdl))
contrib = idf * (tfs*(self.k1+1)) / denom
for d, c in zip(docs, contrib):
seen[d]=seen.get(d,0.0)+float(c)
if not seen: return []
idx=np.fromiter(seen.keys(),dtype=np.int32)
scs=np.fromiter(seen.values(),dtype=np.float32)
k=min(topk,len(scs))
# stable top-k: argsort with secondary key by doc index
order = np.lexsort((idx, -scs)) # sort by -score, then doc idx
order = order[:k]
idx = idx[order]
return [self.doc_ids[i] for i in idx]
def get_bm25_backend(use_bm25s: bool, tokenizer, k1=1.3, b=0.7, logger=print) -> AbstractBM25Backend:
"""
Factory function to get the best available BM25 backend.
Prefers the fast and reliable BM25SBackend, with a pure-Python fallback.
"""
if use_bm25s:
if _BM25S_AVAILABLE:
try:
be = BM25SBackend(tokenizer, k1=k1, b=b)
if logger: logger(f"[BM25] Using high-performance BM25S backend.")
return be
except Exception as e:
if logger: logger(f"[BM25] BM25S failed to initialize ({e}); falling back to DeterministicBM25.")
else:
if logger: logger(f"[BM25] bm25s library not installed; falling back to DeterministicBM25.")
if logger: logger(f"[BM25] Using pure-Python DeterministicBM25 backend.")
return DeterministicBM25Backend(tokenizer, k1=k1, b=b) |