File size: 3,058 Bytes
7aab86a | 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 | 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
# Full source pages preserve nearby conditions that a small chunk can omit.
return [self.pages[key] for key in selected]
|