File size: 2,173 Bytes
22b7d63 | 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 | """
lexical.py — BM25+ retrieval over the keyword-dense `lexical_text` field.
Uses rank_bm25's BM25Plus (the '+' lower-bounds term-frequency saturation via a
delta term, which helps long documents from being unfairly penalized). The
per-document frequency dicts that BM25Plus already keeps are reused to explain
*which* query terms actually hit a given movie — surfaced in the UI as
`match_terms`.
"""
from __future__ import annotations
import numpy as np
from rank_bm25 import BM25Plus
from app.engine.preprocess import preprocessor
class LexicalIndex:
def __init__(self, tokenized_corpus: list[list[str]], k1: float, b: float, delta: float):
self.k1, self.b, self.delta = k1, b, delta
self.bm25 = BM25Plus(tokenized_corpus, k1=k1, b=b, delta=delta)
@classmethod
def build(cls, texts: list[str], k1: float, b: float, delta: float) -> "LexicalIndex":
tokenized = [preprocessor.tokenize(t) for t in texts]
return cls(tokenized, k1, b, delta)
def scores(self, query_tokens: list[str]) -> np.ndarray:
if not query_tokens:
return np.zeros(self.bm25.corpus_size, dtype=np.float32)
return np.asarray(self.bm25.get_scores(query_tokens), dtype=np.float32)
def search(self, query_tokens: list[str], top_n: int) -> list[tuple[int, float]]:
scores = self.scores(query_tokens)
if scores.size == 0 or not query_tokens:
return []
top_n = min(top_n, scores.size)
# argpartition for speed, then sort just the top slice
idx = np.argpartition(-scores, top_n - 1)[:top_n]
idx = idx[np.argsort(-scores[idx])]
return [(int(i), float(scores[i])) for i in idx if scores[i] > 0]
def match_terms(self, doc_idx: int, query_tokens: list[str], limit: int = 6) -> list[str]:
"""Query terms (lemmas) that actually occur in the document."""
doc_freq = self.bm25.doc_freqs[doc_idx]
seen, out = set(), []
for t in query_tokens:
if t in doc_freq and t not in seen:
seen.add(t)
out.append(t)
if len(out) >= limit:
break
return out
|