1"""Phase 5 — private hybrid index. Units = structural annotations, summaries, headings and short excerpts (≤ 400 chars, rights.yaml).2Backend: BM25 (rank_bm25) + local sentence-transformers embeddings stored as .npy; metadata in parquet. No full texts are indexed."""3from __future__ import annotations4import json, re5from pathlib import Path6import numpy as np, pandas as pd7from . import config, segmentation8from .provenance import stamp, now_iso910INDEX_DIR = config.ROOT / "data" / "derived" / "index"11EXCERPT = 4001213def units_for_doc(doc_id: str, ann: dict | None) -> list[dict]:14 d = segmentation.load_parsed(doc_id); m = d["metadata"]; base = {"doc_id": doc_id, "sha256": d["source_sha256"], "year": m["year"], "title": m["title"], "article_type": m["article_type"], "doi": m["doi"]}15 a = (ann or {}).get("annotation") or {}16 out = []17 if a:18 out.append({**base, "unit": "document_summary", "section": None, "page": 1, "role": "document", "genre": a["genre"], "subfield": a["subfield"], "method": a["primary_method"],19 "text": f"{m['title']}. Genre: {a['genre']}; subfield: {a['subfield']}; method: {a['primary_method']}. {a['structural_summary']} Keywords: {', '.join(a['writing_profile_keywords'])}", "moves": None, "confidence": a["confidence"]})20 out.append({**base, "unit": "identification", "section": None, "page": 1, "role": "empirical_design", "genre": a["genre"], "subfield": a["subfield"], "method": a["primary_method"],21 "text": f"Estimand: {a['identification']['estimand']}. Strategy: {a['identification']['strategy']}. Assumptions: {'; '.join(a['identification']['key_assumptions'])}. Threats addressed: {'; '.join(a['identification']['threats_addressed'])}. Causal language strength {a['identification']['causal_language_strength']}/5.", "moves": None, "confidence": a["confidence"]})22 for ce in a["claim_evidence"]:23 out.append({**base, "unit": "claim_evidence", "section": "introduction", "page": 1, "role": "introduction", "genre": a["genre"], "subfield": a["subfield"], "method": a["primary_method"],24 "text": f"Claim: {ce['claim']} Evidence type: {ce['evidence_type']} ({ce['evidence_ref']}).", "moves": None, "confidence": a["confidence"]})25 moves_by_para = {}26 for mv in a["introduction"]["moves"]: moves_by_para.setdefault(mv["paragraph_index"], []).append(mv["move"])27 intro = next((s for s in d["sections"] if s.get("is_intro")), None)28 if intro:29 for i, p in enumerate(segmentation.paragraphs(intro)[:12]):30 if i in moves_by_para:31 out.append({**base, "unit": "intro_move", "section": "introduction", "page": p["page"], "role": "introduction", "genre": a["genre"], "subfield": a["subfield"], "method": a["primary_method"],32 "text": f"[{', '.join(moves_by_para[i])}] " + segmentation.clip(p["text"], EXCERPT), "moves": "|".join(moves_by_para[i]), "confidence": a["confidence"]})33 roles = {x["idx"]: x["role"] for x in a["section_roles"]}34 else:35 roles = {}36 for i, s in enumerate(d["sections"]):37 ps = segmentation.paragraphs(s)38 if not ps: continue39 role = roles.get(i) or segmentation.section_role(s["heading"], s.get("is_appendix", False), s.get("is_intro", False))40 out.append({**base, "unit": "section_opening", "section": s["heading"], "page": s["page"], "role": role, "genre": a.get("genre"), "subfield": a.get("subfield"), "method": a.get("primary_method"),41 "text": f"Section '{s['heading']}' ({role}, {s['n_words']} words; subsections: {', '.join(x['heading'] for x in s['subsections'][:6]) or 'none'}). Opening: " + segmentation.clip(ps[0]["text"], EXCERPT), "moves": None, "confidence": d["confidence"]["overall"]})42 for t in d["tables"]:43 if t.get("notes_excerpt"):44 out.append({**base, "unit": "table_notes", "section": t["label"], "page": t["page"], "role": "results", "genre": a.get("genre"), "subfield": a.get("subfield"), "method": a.get("primary_method"),45 "text": f"{t['label']}: {t['title'] or ''}. Notes: " + segmentation.clip(t["notes_excerpt"], EXCERPT), "moves": None, "confidence": 0.9})46 for f in d["figures"][:8]:47 out.append({**base, "unit": "figure_caption", "section": f["label"], "page": f["page"], "role": "results", "genre": a.get("genre"), "subfield": a.get("subfield"), "method": a.get("primary_method"),48 "text": segmentation.clip(f["caption"], EXCERPT), "moves": None, "confidence": 0.9})49 return out5051def build_index(model_name: str | None = None) -> dict:52 from sentence_transformers import SentenceTransformer53 INDEX_DIR.mkdir(parents=True, exist_ok=True)54 man = pd.read_parquet(config.manifests_dir() / "corpus_manifest.parquet"); ids = man[man.decision == "include"].doc_id.tolist()55 adir = config.ROOT / "data" / "annotations" / "documents"; units = []56 for i in ids:57 ap = adir / f"{i}.json"; ann = json.loads(ap.read_text()) if ap.exists() else None58 if ann and not ann.get("schema_valid"): ann = None59 units += units_for_doc(i, ann)60 df = pd.DataFrame(units); df["unit_id"] = range(len(df))61 sp = json.loads((config.ROOT / "data" / "splits" / "heldout_v1.json").read_text()); df["split"] = np.where(df.doc_id.isin(sp["heldout"]), "heldout", "train")62 model_name = model_name or config.load("models")["embedding_model"]63 model = SentenceTransformer(model_name, device="cpu")64 emb = model.encode(df["text"].tolist(), batch_size=64, normalize_embeddings=True, show_progress_bar=False).astype(np.float32)65 np.save(INDEX_DIR / "embeddings.npy", emb); df.to_parquet(INDEX_DIR / "units.parquet", index=False)66 meta = {"built": stamp(), "n_units": int(len(df)), "n_docs": int(df.doc_id.nunique()), "n_docs_with_annotation": int(df[df.unit == "document_summary"].doc_id.nunique()), "embedding_model": model_name, "dim": int(emb.shape[1]), "units_by_type": {k: int(v) for k, v in df.unit.value_counts().items()}, "excerpt_cap_chars": EXCERPT}67 (INDEX_DIR / "index_meta.json").write_text(json.dumps(meta, indent=1)); return meta68