| """One retriever, three modes β the configurations compared in SQ1. |
| |
| from retrieval import Retriever |
| r = Retriever(model="bge", mode="hybrid") # mode: dense | bm25 | hybrid ; model: any key in embed_models.py |
| hits = r.search("αα»α±α¬ααΊαΈααα¬ αα
α¬", k=5) # -> list of (score, meta) |
| |
| dense cosine similarity between the query vector and every passage vector (the embedding model's view of meaning) |
| bm25 keyword overlap on Burmese syllables / syllable bigrams (exact words, names, numbers) |
| hybrid both lists merged by reciprocal-rank fusion: score = 1/(60 + dense rank) + w Β· 1/(60 + keyword rank) |
| (Cormack et al., 2009; w = bm25_weight, default 0.5). |
| Defaults (model="sealion", mode="hybrid") are the SQ1 winner of 2 Nov 2026: SEA-LION-E5 + hybrid, R@5 0.84, MRR 0.74 on 57 questions. |
| The keyword side is used only for queries that contain Burmese words. |
| |
| per_speech=N caps how many passages from one speech may appear in the top-k (the pilot's "one long speech |
| crowded out the others"). The dense index and the BM25 index hold the same passages in the same order, so |
| positions are shared. |
| """ |
| import numpy as np |
| from search import Index, load_meta |
| from bm25_index import BM25 |
|
|
| RRF_K = 60 |
| POOL = 100 |
|
|
| class Retriever: |
| def __init__(self, model="sealion", mode="hybrid", device=None, meta=None, dense=None, bm25=None, bm25_weight=0.5): |
| self.mode = mode; self.model_key = model; self.bm25_weight = bm25_weight |
| self.meta = meta if meta is not None else load_meta() |
| self.dense = dense if dense is not None else (Index(model=model, device=device, meta=self.meta) if mode != "bm25" else None) |
| self.bm25 = bm25 if bm25 is not None else (BM25() if mode != "dense" else None) |
| if self.dense is not None and len(self.dense) != len(self.meta): raise RuntimeError("dense index / metadata size mismatch") |
|
|
| @property |
| def name(self): return f"{self.model_key}-{self.mode}" if self.mode != "bm25" else "bm25" |
|
|
| def ranking(self, query, pool=POOL): |
| """Fused (or single) ranking: list of (score, passage position), best first.""" |
| if self.mode == "dense": |
| s = self.dense.scores(query); top = np.argpartition(-s, pool)[:pool] |
| return sorted(((float(s[j]), int(j)) for j in top), reverse=True) |
| if self.mode == "bm25": |
| return self.bm25.search(query, pool) |
| d = self.dense.scores(query) |
| |
| |
| use_bm25 = any("\u1000" <= ch <= "\u109f" for ch in query) |
| fused = {} |
| top = np.argpartition(-d, pool)[:pool] |
| for rank, j in enumerate(sorted(top, key=lambda j: -d[j]), 1): |
| fused[int(j)] = fused.get(int(j), 0.0) + 1.0 / (RRF_K + rank) |
| if use_bm25: |
| b = self.bm25.scores(query); top = np.argpartition(-b, pool)[:pool] |
| for rank, j in enumerate(sorted(top, key=lambda j: -b[j]), 1): |
| if b[j] <= 0: break |
| fused[int(j)] = fused.get(int(j), 0.0) + self.bm25_weight / (RRF_K + rank) |
| return sorted(((v, j) for j, v in fused.items()), reverse=True) |
|
|
| def ranking_multi(self, queries, pool=POOL): |
| """Fuse the rankings of several phrasings of one question (e.g. the English original and its Burmese |
| translation) by reciprocal rank fusion, so the keyword index can contribute for a translated query.""" |
| if len(queries) == 1: return self.ranking(queries[0], pool) |
| fused = {} |
| for q in queries: |
| for rank, (_, j) in enumerate(self.ranking(q, pool), 1): |
| fused[j] = fused.get(j, 0.0) + 1.0 / (RRF_K + rank) |
| return sorted(((v, j) for j, v in fused.items()), reverse=True) |
|
|
| def search(self, query, k=5, per_speech=None, filter_fn=None, extra_queries=()): |
| out, seen = [], {} |
| queries = [query] + [q for q in extra_queries if q and q.strip() and q.strip() != query.strip()] |
| for score, j in self.ranking_multi(queries): |
| m = self.meta[j] |
| if filter_fn and not filter_fn(m): continue |
| if per_speech: |
| c = seen.get(m["speech_id"], 0) |
| if c >= per_speech: continue |
| seen[m["speech_id"]] = c + 1 |
| out.append((score, m)) |
| if len(out) >= k: break |
| return out |
|
|
| def all_configs(models=("bge",), modes=("dense", "hybrid"), device=None): |
| """Build every (model Γ mode) retriever once, sharing metadata and the BM25 index.""" |
| meta = load_meta(); bm = BM25() if any(m != "dense" for m in modes) else None |
| out = {} |
| for mk in models: |
| idx = Index(model=mk, device=device, meta=meta) if any(m != "bm25" for m in modes) else None |
| for mode in modes: |
| r = Retriever(model=mk, mode=mode, meta=meta, dense=idx, bm25=bm) |
| out[r.name] = r |
| return out |
|
|
