Spaces:
Runtime error
Runtime error
| """In-memory BM25 keyword index (the 'keyword' half of hybrid search).""" | |
| import heapq | |
| import math | |
| import re | |
| import threading | |
| from collections import Counter | |
| STOP = set( | |
| "a an and are as at be but by for from has have he her his i in is it its of on or she that the their " | |
| "them they this to was we were what when where which who why will with you your do does did how can " | |
| "kya hai hain ka ki ke ko se me mai par pe aur tha thi the".split() | |
| ) | |
| def tokenize(text: str) -> list[str]: | |
| return [t for t in re.findall(r"\w+", text.lower()) if len(t) > 1 and t not in STOP] | |
| class BM25Index: | |
| k1 = 1.5 | |
| b = 0.75 | |
| def __init__(self): | |
| self._lock = threading.RLock() | |
| self._dirty = True | |
| self._clear() | |
| def _clear(self): | |
| self.chunks: list[dict] = [] | |
| self.lens: list[int] = [] | |
| self.postings: dict[str, list[tuple[int, int]]] = {} | |
| self.avgdl = 1.0 | |
| def mark_dirty(self) -> None: | |
| self._dirty = True | |
| def build(self, chunks: list[dict]) -> None: | |
| postings: dict[str, list[tuple[int, int]]] = {} | |
| lens = [] | |
| for i, c in enumerate(chunks): | |
| toks = tokenize(c["text"]) | |
| lens.append(len(toks)) | |
| for term, tf in Counter(toks).items(): | |
| postings.setdefault(term, []).append((i, tf)) | |
| with self._lock: | |
| self.chunks, self.lens, self.postings = chunks, lens, postings | |
| self.avgdl = (sum(lens) / len(lens)) if lens else 1.0 | |
| def ensure_fresh(self) -> None: | |
| if not self._dirty: | |
| return | |
| with self._lock: | |
| if not self._dirty: | |
| return | |
| self._dirty = False # cleared first: a concurrent mark_dirty() during the rebuild is not lost | |
| from . import store | |
| self.build(store.scroll_all()) | |
| def search(self, query: str, doc_ids=None, pages=None, limit: int = 20) -> list[dict]: | |
| self.ensure_fresh() | |
| doc_set = set(doc_ids) if doc_ids else None | |
| page_set = set(pages) if pages else None | |
| with self._lock: | |
| n = len(self.chunks) | |
| if n == 0: | |
| return [] | |
| scores: dict[int, float] = {} | |
| for term in set(tokenize(query)): | |
| plist = self.postings.get(term) | |
| if not plist: | |
| continue | |
| df = len(plist) | |
| idf = math.log(1 + (n - df + 0.5) / (df + 0.5)) | |
| for i, tf in plist: | |
| c = self.chunks[i] | |
| if doc_set and c["doc_id"] not in doc_set: | |
| continue | |
| if page_set and c["page"] not in page_set: | |
| continue | |
| denom = tf + self.k1 * (1 - self.b + self.b * self.lens[i] / self.avgdl) | |
| scores[i] = scores.get(i, 0.0) + idf * tf * (self.k1 + 1) / denom | |
| top = heapq.nlargest(limit, scores.items(), key=lambda kv: kv[1]) | |
| return [{**self.chunks[i], "score": s} for i, s in top] | |
| bm25 = BM25Index() | |