| """Dense retrieval over a shard index (BGE-M3 by default; any model in embed_models.py). |
| |
| from search import Index |
| idx = Index() # BGE-M3, data/index |
| idx = Index(model="sealion") # SEA-LION-E5, data/index_sealion (metadata shared from data/index) |
| hits = idx.search("αα»α±α¬ααΊαΈααα¬ αα
α¬", k=5) # -> list of (score, meta) |
| |
| Shards: emb_NNNNN.npy (float16, L2-normalised) per model directory; meta_NNNNN.jsonl lives once, in data/index, |
| because every model embeds the same passages in the same order (make_embed_input.py fixed the order). |
| """ |
| import glob, json, os |
| import numpy as np |
| from common import DATA, INDEX_DIR, load_embedder, header_of, citation_of |
| from embed_models import spec |
|
|
| def load_meta(meta_dir=INDEX_DIR): |
| metas = [] |
| for mp in sorted(glob.glob(os.path.join(meta_dir, "meta_*.jsonl"))): |
| metas.extend(json.loads(l) for l in open(mp, encoding="utf-8")) |
| return metas |
|
|
| class Index: |
| def __init__(self, model="bge", index_dir=None, meta_dir=INDEX_DIR, device=None, embedder=None, meta=None): |
| self.key = model |
| self.spec = spec(model) |
| self.index_dir = index_dir or os.path.join(DATA, self.spec["dir"]) |
| paths = sorted(glob.glob(os.path.join(self.index_dir, "emb_*.npy"))) |
| if not paths: raise RuntimeError(f"no shards found in {self.index_dir} β build the '{model}' index first") |
| self.n_shards = len(paths) |
| self.emb = np.concatenate([np.load(p) for p in paths]).astype(np.float32) |
| self.meta = meta if meta is not None else load_meta(meta_dir)[: len(self.emb)] |
| if len(self.meta) != len(self.emb): |
| raise RuntimeError(f"{len(self.emb):,} embeddings but {len(self.meta):,} metadata rows β index incomplete?") |
| self.model = embedder or load_embedder(device, model=model)[0] |
|
|
| def __len__(self): return len(self.meta) |
|
|
| def encode_query(self, q): |
| q = self.spec["query"] + q |
| return self.model.encode([q], normalize_embeddings=True, convert_to_numpy=True)[0].astype(np.float32) |
|
|
| def scores(self, query): |
| """Cosine score of every passage for the query (used by hybrid fusion).""" |
| return self.emb @ self.encode_query(query) |
|
|
| def search(self, query, k=5, filter_fn=None, per_speech=None): |
| """Top-k passages. per_speech=N caps how many passages of the same speech may appear (pilot: top-k |
| was filled by one long speech).""" |
| sc = self.scores(query) |
| order = np.argsort(-sc) |
| out, seen = [], {} |
| for j in order: |
| 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((float(sc[j]), m)) |
| if len(out) >= k: break |
| return out |
|
|
| def format_hit(rank, score, m, n_chars=220): |
| return (f"{rank}. [{score:.3f}] {citation_of(m)}\n" |
| f" {m['text'][:n_chars].replace(chr(10), ' ')}{'β¦' if len(m['text']) > n_chars else ''}") |
|
|
