1"""Hybrid retrieval over the private unit index: BM25 + cosine (local embeddings), metadata filters, provenance on every hit."""2from __future__ import annotations3import json, re4from functools import lru_cache5import numpy as np, pandas as pd6from rank_bm25 import BM25Okapi7from .embeddings import INDEX_DIR8from . import config910def _tok(s: str) -> list[str]: return re.findall(r"[a-z0-9]+", s.lower())1112@lru_cache(maxsize=1)13def _load():14 df = pd.read_parquet(INDEX_DIR / "units.parquet"); emb = np.load(INDEX_DIR / "embeddings.npy")15 bm = BM25Okapi([_tok(t) for t in df["text"]])16 from sentence_transformers import SentenceTransformer17 model = SentenceTransformer(json.loads((INDEX_DIR / "index_meta.json").read_text())["embedding_model"], device="cpu")18 return df, emb, bm, model1920def search(query: str, k: int = 8, section: str | None = None, year: int | None = None, unit: str | None = "annotation", genre: str | None = None, subfield: str | None = None,21 split: str | None = "train", alpha: float = 0.5, mode: str = "hybrid") -> list[dict]:22 """unit: 'annotation' = structural units only (summary, identification, intro_move, claim_evidence); 'raw' = section openings/table notes/captions23 (conventional RAG over short excerpts); None = everything. mode: hybrid | bm25 | dense."""24 df, emb, bm, model = _load()25 mask = np.ones(len(df), bool)26 if split: mask &= (df["split"] == split).values27 if year: mask &= (df["year"] == year).values28 if genre: mask &= (df["genre"] == genre).values29 if subfield: mask &= (df["subfield"] == subfield).values30 if section: mask &= (df["role"] == section).values | df["section"].fillna("").str.contains(section, case=False).values31 if unit == "annotation": mask &= df["unit"].isin(["document_summary", "identification", "intro_move", "claim_evidence"]).values32 elif unit == "raw": mask &= df["unit"].isin(["section_opening", "table_notes", "figure_caption"]).values33 elif unit: mask &= (df["unit"] == unit).values34 idx = np.where(mask)[0]35 if len(idx) == 0: return []36 b = np.array(bm.get_scores(_tok(query)))[idx]; b = (b - b.min()) / (b.max() - b.min() + 1e-9)37 qv = model.encode([query], normalize_embeddings=True)[0]; c = emb[idx] @ qv; c = (c - c.min()) / (c.max() - c.min() + 1e-9)38 score = b if mode == "bm25" else c if mode == "dense" else alpha * b + (1 - alpha) * c39 order = idx[np.argsort(-score)][:k * 3]40 out, seen = [], {}41 for j in order:42 r = df.iloc[j]43 if seen.get(r.doc_id, 0) >= 2: continue # at most 2 units per document44 seen[r.doc_id] = seen.get(r.doc_id, 0) + 145 out.append({"score": round(float(score[np.where(idx == j)[0][0]]), 4), "unit": r.unit, "text": r.text, "provenance": {"doc_id": r.doc_id, "doi": r.doi, "title": r.title, "year": int(r.year) if pd.notna(r.year) else None, "page": int(r.page) if pd.notna(r.page) else None, "section": r.section, "sha256": r.sha256[:16], "role": r.role, "genre": r.genre, "subfield": r.subfield, "confidence": float(r.confidence) if pd.notna(r.confidence) else None, "split": r.split}})46 if len(out) >= k: break47 return out48