| """BM25 lexical index rebuilt deterministically from the manifest. |
| |
| The corpus statistics cover every READY chunk in the single shared corpus; |
| scope restrictions filter *candidates* before ranking, so an empty scope |
| yields nothing and a document filter never leaks other documents. This is the |
| supported isolation model (one corpus, optional per-request document scope); |
| it is not multi-tenant isolation. |
| |
| Scoring is BM25 (k1=1.5, b=0.75) with the Lucene IDF variant |
| ``log(1 + (N - n + 0.5) / (n + 0.5))``, which stays positive for any corpus |
| size. (The ``rank_bm25`` package's Okapi variant yields non-positive scores |
| when a term appears in every document, which makes one-document corpora |
| unsearchable.) Scores are comparable only within one query. |
| |
| A short English stopword list is removed from documents and queries alike. |
| Without it, on a small corpus, function words such as "what" or "does" carry |
| enough IDF that a question containing no content term still produces lexical |
| hits, and rank fusion then promotes that noise above genuine dense matches. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
| import re |
| import threading |
| from collections import Counter |
| from dataclasses import dataclass |
| from typing import Dict, Iterable, List, Sequence, Tuple |
|
|
| from .types import QueryScope |
|
|
| _TOKEN_RE = re.compile(r"\w+", re.UNICODE) |
|
|
| STOPWORDS = frozenset( |
| """ |
| a an and are as at be been by can could did do does for from had has have he her |
| his how i if in into is it its may might of on or our shall she should that the |
| their them then there these they this those to us was we were what when where |
| which who whom why will with would you your |
| """.split() |
| ) |
|
|
|
|
| def tokenize(text: str) -> List[str]: |
| """Lowercased word tokens with stopwords removed (documents and queries alike).""" |
| return [t for t in _TOKEN_RE.findall(text.lower()) if t not in STOPWORDS] |
|
|
|
|
| @dataclass(frozen=True) |
| class LexicalEntry: |
| chunk_id: str |
| doc_id: str |
| text: str |
|
|
|
|
| @dataclass(frozen=True) |
| class LexicalHit: |
| chunk_id: str |
| doc_id: str |
| score: float |
| rank: int |
|
|
|
|
| class BM25: |
| def __init__(self, documents: Sequence[Sequence[str]], k1: float = 1.5, b: float = 0.75): |
| self.k1 = k1 |
| self.b = b |
| self.n_docs = len(documents) |
| self.doc_freqs: List[Counter] = [Counter(tokens) for tokens in documents] |
| self.doc_lengths = [len(tokens) for tokens in documents] |
| self.avg_length = (sum(self.doc_lengths) / self.n_docs) if self.n_docs else 0.0 |
| df: Counter = Counter() |
| for freqs in self.doc_freqs: |
| df.update(freqs.keys()) |
| self.idf: Dict[str, float] = { |
| term: math.log(1.0 + (self.n_docs - n + 0.5) / (n + 0.5)) for term, n in df.items() |
| } |
|
|
| def scores(self, query_tokens: Sequence[str]) -> List[float]: |
| out = [0.0] * self.n_docs |
| terms = [t for t in dict.fromkeys(query_tokens) if t in self.idf] |
| if not terms: |
| return out |
| for index, freqs in enumerate(self.doc_freqs): |
| length_norm = ( |
| 1.0 |
| - self.b |
| + self.b * (self.doc_lengths[index] / self.avg_length if self.avg_length else 0.0) |
| ) |
| score = 0.0 |
| for term in terms: |
| tf = freqs.get(term, 0) |
| if tf: |
| score += self.idf[term] * (tf * (self.k1 + 1.0)) / (tf + self.k1 * length_norm) |
| out[index] = score |
| return out |
|
|
|
|
| class LexicalIndex: |
| """In-memory BM25 index. ``rebuild`` is O(corpus); fine for a local demo.""" |
|
|
| def __init__(self) -> None: |
| self._lock = threading.RLock() |
| self._entries: List[LexicalEntry] = [] |
| self._bm25: BM25 | None = None |
| self.generation: int = -1 |
|
|
| @property |
| def size(self) -> int: |
| return len(self._entries) |
|
|
| def rebuild(self, entries: Iterable[LexicalEntry], generation: int) -> None: |
| entries = list(entries) |
| bm25 = BM25([tokenize(e.text) for e in entries]) if entries else None |
| with self._lock: |
| self._entries = entries |
| self._bm25 = bm25 |
| self.generation = generation |
|
|
| def search( |
| self, |
| query: str, |
| n_results: int, |
| scope: QueryScope = QueryScope.whole_corpus(), |
| ) -> List[LexicalHit]: |
| if n_results <= 0 or scope.is_empty: |
| return [] |
| tokens = tokenize(query) |
| if not tokens: |
| return [] |
| with self._lock: |
| entries = self._entries |
| bm25 = self._bm25 |
| if bm25 is None: |
| return [] |
| scores = bm25.scores(tokens) |
| candidates: List[Tuple[float, str, str]] = [] |
| for entry, score in zip(entries, scores): |
| if score <= 0.0 or not scope.allows(entry.doc_id): |
| continue |
| candidates.append((float(score), entry.chunk_id, entry.doc_id)) |
| candidates.sort(key=lambda c: (-c[0], c[1])) |
| return [ |
| LexicalHit(chunk_id=cid, doc_id=doc, score=score, rank=i + 1) |
| for i, (score, cid, doc) in enumerate(candidates[:n_results]) |
| ] |
|
|
| def contains_doc(self, doc_id: str) -> bool: |
| with self._lock: |
| return any(e.doc_id == doc_id for e in self._entries) |
|
|
