| import re |
| import unicodedata |
| import numpy as np |
| from rank_bm25 import BM25Okapi |
| from .schema import Evidence |
|
|
|
|
| def normalized(text): |
| return re.sub(r"\s+", " ", unicodedata.normalize("NFC", text)).strip() |
|
|
|
|
| def terms(text): |
| text = "".join(str(unicodedata.decimal(c)) if c.isdecimal() else c for c in text) |
| return re.findall(r"\w+(?:[./-]\w+)*", text.lower()) or ["empty"] |
|
|
|
|
| def evidence_for(doc): |
| return [Evidence(id=f"{doc.id}:p{p.number}", document=doc.name, page=p.number, |
| text=p.english, original=p.original) for p in doc.pages |
| if p.english.strip()] |
|
|
|
|
| class HybridIndex: |
| """MiniLM cosine similarity + BM25, joined with reciprocal rank fusion.""" |
| def __init__(self, documents, embedder): |
| self.embedder = embedder |
| self.pages = {e.id: e for d in documents if d.ready for e in evidence_for(d)} |
| self.rows = [] |
| for page in self.pages.values(): |
| offsets = embedder.tokenizer(page.text, add_special_tokens=False, |
| return_offsets_mapping=True, truncation=False, verbose=False)["offset_mapping"] |
| for start in range(0, len(offsets), 150): |
| stop = min(start + 180, len(offsets)) |
| text = page.text[offsets[start][0]:offsets[stop - 1][1]] |
| self.rows.append((page.id, text)) |
| if stop == len(offsets): |
| break |
| texts = [row[1] for row in self.rows] |
| self.vectors = embedder.encode(texts, normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=False) |
| self.bm25 = BM25Okapi([terms(text) for text in texts]) |
|
|
| def search(self, question, document_ids, k=4): |
| if not document_ids: |
| return [] |
| positions = [i for i, row in enumerate(self.rows) if row[0].split(":")[0] in document_ids] |
| if not positions: |
| return [] |
| vector = self.embedder.encode([question], normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=False)[0] |
| semantic_scores = self.vectors @ vector |
| keyword_scores = self.bm25.get_scores(terms(question)) |
| query_terms = set(terms(question)) |
| semantic = sorted(positions, key=lambda i: float(semantic_scores[i]), reverse=True)[:20] |
| keyword = sorted(positions, key=lambda i: (float(keyword_scores[i]), |
| len(query_terms.intersection(terms(self.rows[i][1])))), reverse=True)[:20] |
| scores = {} |
| for ranking in [semantic, keyword]: |
| for rank, index in enumerate(ranking, 1): |
| scores[index] = scores.get(index, 0) + 1 / (60 + rank) |
| selected = [] |
| for index in sorted(scores, key=scores.get, reverse=True): |
| page_id = self.rows[index][0] |
| if page_id not in selected: |
| selected.append(page_id) |
| if len(selected) == k: |
| break |
| |
| return [self.pages[key] for key in selected] |
|
|
