#!/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())