src/evaluate_retrieval.py

#!/usr/bin/env python3
"""
Retrieval evaluation: the chunking comparison table.

THIS IS THE DELIVERABLE FOR REQUIREMENT 2 ("chunking must be vast").
Nine strategies with measured numbers beats any amount of prose about chunking.

WHAT IT MEASURES
----------------
Relevance is judged at PASSAGE level (`is_selected`), but retrieval happens at
CHUNK level. A retrieved chunk counts as relevant if it overlaps any positive
passage -- which is why every Chunk carries `block_ids`. That mapping is what
makes chunking strategies comparable on a corpus whose labels are per passage.

  nDCG@5  Hit@5  MRR@5  P@1        standard ranking quality
  zero_hit_rate                    fraction of queries with NO positive in top-k
                                   -- the quality-side analogue of P100, and the
                                   metric that exposes catastrophic failures that
                                   a good mean hides
  per_query_std                    robustness; arXiv:2603.06976 found fine-grained
                                   chunkers have high variance and more zero-hits
  n_chunks, index_mb, embed_s      the efficiency side of the Pareto frontier

GPU REQUIRED (embeddings). Run on jupyter-pod, not the head node.

  python src/evaluate_retrieval.py --max-queries 1000 --langs hi,ta,bn
  python src/evaluate_retrieval.py --max-queries 1000            # all 14
"""
from __future__ import annotations

import argparse
import json
import math
import sys
import time
from collections import defaultdict
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from src.chunkers.base import Document, chunk_document, word_budget  # noqa: E402
from src.chunkers.strategies import build_portfolio  # noqa: E402
from src.golden_set import canonical_id  # noqa: E402
from src.schema_utils import LANG_NAMES, default_root, iter_passages, norm_lang, load_report  # noqa: E402

K = 5


# ------------------------------------------------------------------ metrics

def dcg(rels: list[int]) -> float:
    return sum(r / math.log2(i + 2) for i, r in enumerate(rels))


def score_ranking(retrieved_ids: list[set[str]], positives: set[str], k: int = K) -> dict:
    """retrieved_ids[i] = canonical_ids covered by the chunk at rank i."""
    rels = [1 if (ids & positives) else 0 for ids in retrieved_ids[:k]]
    ideal = sorted(rels, reverse=True)
    hit = max(rels) if rels else 0
    rr = next((1.0 / (i + 1) for i, r in enumerate(rels) if r), 0.0)
    idcg = dcg(ideal)
    return {
        "ndcg": dcg(rels) / idcg if idcg > 0 else 0.0,
        "hit": float(hit),
        "mrr": rr,
        "p1": float(rels[0]) if rels else 0.0,
        "zero_hit": 1.0 - float(hit),
    }


def score_at_budget(ranked_ids: list[set[str]], ranked_words: list[int],
                    positives: set[str], budget_words: int) -> dict:
    """Take chunks in rank order until `budget_words` of context is filled.

    WHY THIS EXISTS
    ---------------
    Fixed top-k is not a fair comparison across chunking strategies. A strategy
    that emits 3.8 chunks per document has all of them retrieved by k=5, so its
    Hit@5 is free; one emitting 9.4 chunks faces a real selection. Measured on
    this corpus, corr(chunks_per_doc, Hit@5) = -0.885 -- the ranking was largely
    reproducing chunk count, not chunk quality.

    A reader has a CONTEXT budget, not a chunk budget. Equalising words delivered
    is the comparison that matches how the system is actually used.
    """
    used, taken = 0, []
    for ids, w in zip(ranked_ids, ranked_words):
        if used and used + w > budget_words:
            break
        taken.append(ids)
        used += w
    rels = [1 if (ids & positives) else 0 for ids in taken]
    hit = max(rels) if rels else 0
    rr = next((1.0 / (i + 1) for i, r in enumerate(rels) if r), 0.0)
    idcg = dcg(sorted(rels, reverse=True))
    return {
        "b_ndcg": dcg(rels) / idcg if idcg > 0 else 0.0,
        "b_hit": float(hit),
        "b_mrr": rr,
        "b_zero_hit": 1.0 - float(hit),
        "b_chunks_used": float(len(taken)),
        "b_words_used": float(used),
    }


def aggregate(per_query: list[dict]) -> dict:
    if not per_query:
        return {}
    out = {}
    keys = [k for k in ("ndcg", "hit", "mrr", "p1", "zero_hit",
                        "b_ndcg", "b_hit", "b_mrr", "b_zero_hit",
                        "b_chunks_used", "b_words_used") if k in per_query[0]]
    for key in keys:
        vals = [q[key] for q in per_query]
        mean = sum(vals) / len(vals)
        var = sum((v - mean) ** 2 for v in vals) / max(1, len(vals) - 1)
        out[key] = round(mean, 4)
        if key == "ndcg":
            out["ndcg_std"] = round(math.sqrt(var), 4)
    out["n_queries"] = len(per_query)
    return out


# ------------------------------------------------------------------ data

def load_eval_data(root: Path, lang: str, max_queries: int):
    """Return (queries, docs) for one language, from the validation split."""
    import polars as pl

    rep = load_report(root)
    fmap, pmap = rep["field_mapping"], rep["passage_mapping"]
    pcol = fmap["passages"]
    qid_c, q_c, lang_c = fmap["query_id"], fmap["query"], fmap.get("lang")
    t_key, en_key, sel_key = pmap["text"], pmap.get("text_en"), pmap.get("is_selected")

    target = None
    for fp in rep["files"]:
        if "val" not in Path(fp).name:
            continue
        try:
            probe = pl.read_parquet(fp, columns=[lang_c], n_rows=1)
        except Exception:
            continue
        if norm_lang(probe[lang_c][0]) == lang:
            target = fp
            break
    if target is None:
        return [], []

    df = pl.read_parquet(target, columns=[qid_c, q_c, pcol], n_rows=max_queries * 3)
    queries, docs = [], []
    for qid, q, plist in zip(df[qid_c].to_list(), df[q_c].to_list(), df[pcol].to_list()):
        if len(queries) >= max_queries:
            break
        texts, ids, pos = [], [], set()
        for idx, text, _en, sel, _u in iter_passages(plist, t_key, en_key, sel_key, None):
            if not isinstance(text, str) or not text.strip():
                continue
            cid = canonical_id(qid, idx)
            texts.append(text)
            ids.append(cid)
            if sel == 1:
                pos.add(cid)
        if not texts or not pos or not isinstance(q, str) or not q.strip():
            continue          # unanswerable queries are excluded from RANKING metrics
        queries.append({"query_id": str(qid), "text": q, "positives": pos})
        docs.append(Document.from_blocks(f"q{qid}", lang, texts, ids))
    return queries, docs


# ------------------------------------------------------------------ embedding

class Embedder:
    def __init__(self, model_path: str, device: str, batch: int, max_len: int):
        import torch
        from transformers import AutoModel, AutoTokenizer
        self.torch = torch
        self.tok = AutoTokenizer.from_pretrained(model_path)
        self.model = AutoModel.from_pretrained(
            model_path, torch_dtype=torch.float16 if "cuda" in device else torch.float32
        ).to(device).eval()
        self.device, self.batch, self.max_len = device, batch, max_len

    @property
    def dim(self) -> int:
        return int(self.model.config.hidden_size)

    def encode(self, texts: list[str]):
        torch = self.torch
        outs = []
        with torch.inference_mode():
            for i in range(0, len(texts), self.batch):
                enc = self.tok(texts[i:i + self.batch], padding=True, truncation=True,
                               max_length=self.max_len, return_tensors="pt").to(self.device)
                h = self.model(**enc).last_hidden_state
                # CLS pooling — bge-m3's dense head
                v = h[:, 0]
                outs.append(torch.nn.functional.normalize(v, dim=-1).to(torch.float16))
        return torch.cat(outs) if outs else torch.zeros((0, self.dim), device=self.device)


# ------------------------------------------------------------------ eval

def evaluate_lang(emb: Embedder, lang: str, queries, docs, strategies, fertility: float,
                  pool: str = "corpus", budget_words: int = 400):
    torch = emb.torch
    lo, hi = word_budget(200, fertility)
    results = []

    t0 = time.perf_counter()
    qv = emb.encode([q["text"] for q in queries])
    q_secs = time.perf_counter() - t0

    for ch in strategies:
        t0 = time.perf_counter()
        chunks, owner = [], []          # owner[i] = index of the query this chunk belongs to
        for qi, doc in enumerate(docs):
            for c in chunk_document(doc, ch, lo, hi, post=True):
                chunks.append(c)
                owner.append(qi)
        chunk_s = time.perf_counter() - t0

        if not chunks:
            continue
        t0 = time.perf_counter()
        cv = emb.encode([c.text for c in chunks])
        embed_s = time.perf_counter() - t0

        # POOL CHOICE MATTERS ENORMOUSLY.
        #   query  : rank only the chunks from this query's own ~10 passages.
        #            Easy (Hit@5 ~ 0.99) and confounded -- a strategy emitting
        #            fewer than k chunks gets every one of them retrieved for free.
        #   corpus : rank against EVERY chunk in the language. Realistic, and the
        #            pool no longer depends on how many chunks a strategy makes.
        by_query = defaultdict(list)
        for i, qi in enumerate(owner):
            by_query[qi].append(i)

        n_words = [c.n_words for c in chunks]
        depth = max(K, 40)          # deep enough to fill the word budget
        per_query = []
        t0 = time.perf_counter()
        with torch.inference_mode():
            if pool == "corpus":
                for qi, q in enumerate(queries):
                    sims = (cv @ qv[qi]).float()
                    top = torch.topk(sims, min(depth, len(chunks))).indices.tolist()
                    ranked = [set(chunks[t].block_ids) for t in top]
                    words = [n_words[t] for t in top]
                    m = score_ranking(ranked, q["positives"])
                    m.update(score_at_budget(ranked, words, q["positives"], budget_words))
                    per_query.append(m)
            else:
                for qi, q in enumerate(queries):
                    idxs = by_query.get(qi)
                    if not idxs:
                        continue
                    sims = (cv[idxs] @ qv[qi]).float()
                    top = torch.topk(sims, min(depth, len(idxs))).indices.tolist()
                    ranked = [set(chunks[idxs[t]].block_ids) for t in top]
                    words = [n_words[idxs[t]] for t in top]
                    m = score_ranking(ranked, q["positives"])
                    m.update(score_at_budget(ranked, words, q["positives"], budget_words))
                    per_query.append(m)
        search_s = time.perf_counter() - t0

        agg = aggregate(per_query)
        agg.update({
            "strategy": ch.name, "family": ch.family, "lang": lang,
            "pool": pool, "budget_words": budget_words,
            "n_chunks": len(chunks),
            "mean_chunk_words": round(sum(c.n_words for c in chunks) / len(chunks), 1),
            "chunks_per_doc": round(len(chunks) / max(1, len(docs)), 2),
            "block_integrity": round(
                1 - sum(c.split_blocks for c in chunks) / max(1, sum(len(c.block_ids) for c in chunks)), 4),
            "index_mb": round(len(chunks) * emb.dim * 2 / 2**20, 1),
            "chunk_s": round(chunk_s, 2), "embed_s": round(embed_s, 2),
            "search_ms_per_query": round(1000 * search_s / max(1, len(per_query)), 3),
        })
        results.append(agg)
        print(f"    {ch.name:6s} nDCG@5 {agg['ndcg']:.4f}  P@1 {agg['p1']:.3f}  "
              f"zero {agg['zero_hit']:.3f}  |  @{budget_words}w: nDCG {agg['b_ndcg']:.4f} "
              f"hit {agg['b_hit']:.3f} ({agg['b_chunks_used']:.1f} chunks)  "
              f"|  {agg['chunks_per_doc'] if 'chunks_per_doc' in agg else len(chunks)/max(1,len(docs)):.1f}/doc  {embed_s:.1f}s")
    return results, q_secs


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", type=Path, default=None)
    ap.add_argument("--model", default=None, help="path or hub id; default bge-m3 from the local cache")
    ap.add_argument("--langs", default=None, help="comma-separated ISO-2; default all")
    ap.add_argument("--max-queries", type=int, default=1000)
    ap.add_argument("--batch", type=int, default=64)
    ap.add_argument("--max-len", type=int, default=192, help="passages are ~55 words; 192 is ample and 2.6x faster than 512")
    ap.add_argument("--device", default=None)
    ap.add_argument("--pool", choices=["corpus", "query"], default="corpus",
                    help="corpus = rank against all chunks in the language (realistic); "
                         "query = only this query's own passages (easy, and confounded "
                         "by chunks-per-doc vs k)")
    ap.add_argument("--budget-words", type=int, default=400,
                    help="context budget for the fair, chunk-count-independent metric")
    ap.add_argument("--allow-cpu", action="store_true",
                    help="run without a GPU anyway (very slow; for debugging only)")
    args = ap.parse_args()

    root = args.root.expanduser().resolve() if args.root else default_root()
    print(f"==> data root: {root}")

    # Fail fast and helpfully: this step needs a GPU and transformers, both of
    # which live on jupyter-pod. The head node has neither.
    import socket
    missing = []
    for mod in ("torch", "transformers"):
        try:
            __import__(mod)
        except ImportError:
            missing.append(mod)
    if missing:
        raise SystemExit(
            f"\nMissing: {', '.join(missing)}  (host: {socket.gethostname()})\n\n"
            "This step needs a GPU and the NVIDIA container's python packages.\n"
            "kls-headnode has neither; jupyter-pod has both.\n\n"
            "  -> Open 03_evaluate.ipynb on jupyter-pod and run it there.\n"
            "     Both hosts share the same NFS, so nothing needs copying.\n\n"
            "  (If you must run here, create an isolated venv:\n"
            f"     python3 -m venv --system-site-packages {root.parent}/.venv\n"
            f"     source {root.parent}/.venv/bin/activate && pip install {' '.join(missing)}\n"
            "   but it will still be CPU-only and very slow.)\n"
        )

    import torch
    device = pick_device(args.device)
    if device == "cpu":
        print(f"  !! no GPU visible on {socket.gethostname()} — embedding will be ~50x slower.")
        print("     Run 03_evaluate.ipynb on jupyter-pod instead.")
        if not args.allow_cpu:
            raise SystemExit("     (pass --allow-cpu to override)")
    else:
        p = torch.cuda.get_device_properties(0)
        print(f"==> device: {device}  ({p.name}, {p.total_memory/2**30:.0f} GiB, {p.multi_processor_count} SMs)")

    model = args.model
    if model is None:
        hits = list((root / "hf_cache" / "hub").glob("models--BAAI--bge-m3/snapshots/*"))
        model = str(hits[0]) if hits else "BAAI/bge-m3"
    print(f"==> model: {model}")

    fert_path = root / "results" / "fertility.json"
    fert = {}
    if fert_path.exists():
        fert = {k: v["fertility_vs_english"]
                for k, v in json.loads(fert_path.read_text())["per_language"].items()}
        print(f"==> fertility loaded for {len(fert)} languages (per-language word budgets)")
    else:
        print("==> fertility.json not found — using fertility=1.0 for every language")

    langs = args.langs.split(",") if args.langs else list(LANG_NAMES)
    langs = [l for l in langs if l != "en"]

    emb = Embedder(model, device, args.batch, args.max_len)
    strategies = build_portfolio()
    print(f"==> {len(strategies)} strategies: {[s.name for s in strategies]}\n")

    all_rows = []
    for lang in langs:
        queries, docs = load_eval_data(root, lang, args.max_queries)
        if not queries:
            print(f"  {lang}: no data, skipping")
            continue
        print(f"  {LANG_NAMES.get(lang, lang)} ({lang}): {len(queries):,} answerable queries, "
              f"{sum(len(d.blocks) for d in docs):,} passages, fertility {fert.get(lang, 1.0):.2f}")
        rows, _ = evaluate_lang(emb, lang, queries, docs, strategies, fert.get(lang, 1.0),
                                pool=args.pool, budget_words=args.budget_words)
        all_rows.extend(rows)
        print()

    if not all_rows:
        raise SystemExit("no results — check --langs and that the validation split is present")

    # ---- the table ----
    by_strat = defaultdict(list)
    for r in all_rows:
        by_strat[r["strategy"]].append(r)

    print("=" * 96)
    print("CHUNKING COMPARISON  (mean across languages)")
    print("=" * 96)
    hdr = (f"{'strategy':8s}{'family':12s}{'nDCG@5':>8}{'Hit@5':>8}{'MRR@5':>8}{'P@1':>7}"
           f"{'zero':>7}{'std':>7}{'chunks':>9}{'BI':>7}{'MB':>7}")
    print(hdr); print("-" * 96)
    table = []
    for name, rows in sorted(by_strat.items(),
                             key=lambda kv: -sum(r["ndcg"] for r in kv[1]) / len(kv[1])):
        m = lambda k: sum(r[k] for r in rows) / len(rows)  # noqa: E731
        table.append({"strategy": name, "family": rows[0]["family"],
                      **{k: round(m(k), 4) for k in
                         ("ndcg", "hit", "mrr", "p1", "zero_hit", "ndcg_std",
                          "block_integrity", "chunks_per_doc",
                          "b_ndcg", "b_hit", "b_mrr", "b_zero_hit",
                          "b_chunks_used", "b_words_used")},
                      "n_chunks": int(m("n_chunks")), "index_mb": round(m("index_mb"), 1),
                      "n_langs": len(rows)})
        print(f"{name:8s}{rows[0]['family']:12s}{m('ndcg'):>8.4f}{m('hit'):>8.3f}{m('mrr'):>8.4f}"
              f"{m('p1'):>7.3f}{m('zero_hit'):>7.3f}{m('ndcg_std'):>7.3f}"
              f"{int(m('n_chunks')):>9,}{m('block_integrity'):>7.3f}{m('index_mb'):>7.1f}")

    best, worst = table[0], table[-1]
    print("-" * 96)
    if worst["p1"] > 0:
        print(f"  best/worst P@1 ratio: {best['p1']/worst['p1']:.2f}x "
              f"({best['strategy']} {best['p1']:.3f} vs {worst['strategy']} {worst['p1']:.3f})")
    print(f"  nDCG spread across {len(table)} strategies: "
          f"{best['ndcg']-worst['ndcg']:.4f} "
          f"({100*(best['ndcg']-worst['ndcg'])/max(1e-9,best['ndcg']):.1f}% of the best)")
    print("  ^ a SMALL spread is itself the finding — it is what the ANOVA will quantify.")

    # --- confound diagnostic: is the ranking just reproducing chunk count? ---
    def _corr(x, y):
        mx = sum(x) / len(x); my = sum(y) / len(y)
        num = sum((a - mx) * (b - my) for a, b in zip(x, y))
        den = (sum((a - mx) ** 2 for a in x) * sum((b - my) ** 2 for b in y)) ** 0.5
        return num / den if den else 0.0

    cpd = [r["chunks_per_doc"] for r in table]
    print(f"\n  corr(chunks_per_doc, nDCG@{K})    = {_corr(cpd, [r['ndcg'] for r in table]):+.3f}")
    print(f"  corr(chunks_per_doc, nDCG@{args.budget_words}w) = "
          f"{_corr(cpd, [r['b_ndcg'] for r in table]):+.3f}   <- should be much weaker")
    free = [r["strategy"] for r in table if r["chunks_per_doc"] <= K]
    if free and args.pool == "query":
        print(f"  !! {free} emit <= k={K} chunks/doc, so top-k retrieves ALL of them.")
        print(f"     Their Hit@{K} is free. Trust the @{args.budget_words}w columns instead.")

    out = root / "results" / "retrieval_eval.json"
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps({
        "config": {"max_queries": args.max_queries, "k": K, "model": model,
                   "max_len": args.max_len, "langs": langs},
        "summary": table, "per_language": all_rows,
    }, indent=2))
    print(f"\n==> wrote {out}")
    print("    per_language rows feed the ANOVA (strategy x language x ...)")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())