"""Evaluate retrieval on the reference set Q (Phase 1.4, Eq. 7): Recall@k and MRR. Methods: BM25 (lexical baseline), every dense model with an index under data/index/eval// (with and without its query instruction), and RRF hybrids BM25 + dense. Scores are computed against both gold definitions: (c) `gold`, the main one, and (b) `gold_b`, as a sensitivity check. The evaluation index is built from corpus_eval.jsonl (corpus minus the queries in Q, so a query cannot retrieve itself); the system's final index uses the full corpus.jsonl. For the selected model (models.yaml → embeddings.selected) the script also scores `rrf(bm25+/instr)+generic`: the Retriever's `generic_query` method, i.e. RRF over BM25 and dense rankings of the query and of `generalize_query(query)`. The generalised queries are encoded on CPU (fp32); the original ones reuse the vectors embedded with the corpus. Only the dev split is used for model and k selection. `--split test` is for the single final report and requires `--final`. Usage: python scripts/eval_retrieval.py # dev, all available indexes python scripts/eval_retrieval.py --models bge-base e5-base """ from __future__ import annotations import argparse import json import sys from datetime import datetime, timezone from pathlib import Path ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT / "src")) from formal_mas.config import load_models # noqa: E402 from formal_mas.retrieval.bm25 import BM25 # noqa: E402 from formal_mas.retrieval.corpus import load_jsonl # noqa: E402 from formal_mas.retrieval.evaluate import evaluate, rrf # noqa: E402 from formal_mas.retrieval.retriever import generalize_query # noqa: E402 KS = (4, 8, 16, 32) GOLDS = {"c": "gold", "b": "gold_b"} def query_vectors(folder: Path, key: str, q: list[dict], instr: bool): """Query vectors embedded with the corpus (same run and precision) if present, else encode on CPU.""" import numpy as np label = "instr" if instr else "no_instr" if (folder / f"queries_{label}.npy").exists(): pos = {n: i for i, n in enumerate(json.loads((folder / "queries.json").read_text(encoding="utf-8")))} return np.load(folder / f"queries_{label}.npy")[[pos[r["name"]] for r in q]].astype(np.float32) from formal_mas.retrieval.encoders import Encoder, load_specs spec = load_specs()[key] return Encoder(spec).encode([spec.query_text(r["query"], instr) for r in q]) def dense_rankings(key: str, index_root: Path, q: list[dict], depth: int) -> dict[str, list[list[str]]]: import faiss from formal_mas.retrieval.encoders import load_specs folder = index_root / key index = faiss.read_index(str(folder / "index.faiss")) names = json.loads((folder / "names.json").read_text(encoding="utf-8")) settings = {"no_instr": False} if load_specs()[key].query_instruction: settings["instr"] = True out = {} for label, instr in settings.items(): _, ids = index.search(query_vectors(folder, key, q, instr), depth) out[f"{key}/{label}"] = [[names[i] for i in row if i >= 0] for row in ids] return out def generic_dense_rankings(key: str, index_root: Path, generic: list[str | None], depth: int) -> list[list[str]]: """Dense rankings of the generalised queries (instruction on); [] where there is none.""" import faiss from formal_mas.retrieval.encoders import Encoder, load_specs folder = index_root / key index = faiss.read_index(str(folder / "index.faiss")) names = json.loads((folder / "names.json").read_text(encoding="utf-8")) todo = [i for i, g in enumerate(generic) if g] out: list[list[str]] = [[] for _ in generic] if todo: spec = load_specs()[key] _, ids = index.search(Encoder(spec).encode([spec.query_text(generic[i], True) for i in todo]), depth) for i, row in zip(todo, ids): out[i] = [names[j] for j in row if j >= 0] return out def to_markdown(results: dict, gold: str) -> str: cols = [f"R@{k}" for k in KS] + ["MRR"] lines = [f"| Method ({gold}) | " + " | ".join(cols) + " |", "|---|" + "---|" * len(cols)] for method, res in results.items(): lines.append(f"| {method} | " + " | ".join(f"{res[gold][c]:.3f}" for c in cols) + " |") return "\n".join(lines) def to_latex(results: dict, gold: str, split: str, n: int) -> str: cols = [f"R@{k}" for k in KS] + ["MRR"] best = {c: max(r[gold][c] for r in results.values()) for c in cols} rows = [] for method, res in results.items(): cells = [(f"\\textbf{{{res[gold][c]:.3f}}}" if res[gold][c] == best[c] else f"{res[gold][c]:.3f}") for c in cols] rows.append(f"{method.replace('_', chr(92) + '_')} & " + " & ".join(cells) + " \\\\") header = " & ".join(["Method"] + [c.replace("@", "@") for c in cols]) return "\n".join([ "\\begin{table}[t]", "\\centering", f"\\caption{{Retrieval on $Q_{{\\text{{{split}}}}}$ ($n={n}$), gold ({gold}).}}", "\\begin{tabular}{l" + "c" * len(cols) + "}", "\\toprule", header + " \\\\", "\\midrule", *rows, "\\bottomrule", "\\end{tabular}", "\\end{table}", ]) def main(argv: list[str] | None = None) -> None: parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument("--split", choices=["dev", "test"], default="dev") parser.add_argument("--final", action="store_true", help="required to evaluate on the test split") parser.add_argument("--models", nargs="*", help="dense model keys (default: every built eval index)") parser.add_argument("--corpus", type=Path, default=ROOT / "data" / "index" / "corpus_eval.jsonl") parser.add_argument("--depth", type=int, default=100, help="list length for MRR and RRF") parser.add_argument("--out", type=Path, default=ROOT / "data" / "retrieval_eval") parser.add_argument("--stem", help="output file stem (default: results_)") parser.add_argument("--note", help="free-text note stored in the JSON and Markdown outputs") args = parser.parse_args(argv) if args.split == "test" and not args.final: parser.error("Q-test is reserved for the single final report: pass --final") q = load_jsonl(args.out / f"q_{args.split}.jsonl") queries = [r["query"] for r in q] index_root = args.corpus.parent / "eval" keys = args.models if args.models is not None else sorted( p.name for p in index_root.glob("*") if (p / "index.faiss").exists() and "_limit" not in p.name) corpus = load_jsonl(args.corpus) bm25 = BM25([r["text"] for r in corpus]) names = [r["name"] for r in corpus] del corpus rankings = {"bm25": [[names[i] for i, _ in bm25.top(t, args.depth)] for t in queries]} generic = [generalize_query(t) for t in queries] bm25_generic = [[names[i] for i, _ in bm25.top(g, args.depth)] if g else [] for g in generic] del bm25 selected = load_models()["embeddings"]["selected"] for key in keys: dense = dense_rankings(key, index_root, q, args.depth) rankings.update(dense) for label, ranked in dense.items(): rankings[f"rrf(bm25+{label})"] = [rrf([b, d], args.depth) for b, d in zip(rankings["bm25"], ranked)] if key == selected and f"{key}/instr" in dense: dense_generic = generic_dense_rankings(key, index_root, generic, args.depth) rankings[f"rrf(bm25+{key}/instr)+generic"] = [ rrf([lst for lst in four if lst], args.depth) for four in zip(rankings["bm25"], dense[f"{key}/instr"], bm25_generic, dense_generic)] results = {m: {g: evaluate(r, [x[field] for x in q], KS) for g, field in GOLDS.items()} for m, r in rankings.items()} args.out.mkdir(parents=True, exist_ok=True) stem = args.out / (args.stem or f"results_{args.split}") payload = { "created_utc": datetime.now(timezone.utc).isoformat(timespec="seconds"), "split": args.split, "n": len(q), "ks": list(KS), "depth": args.depth, "corpus": args.corpus.name, "n_generalized": sum(g is not None for g in generic), **({"note": args.note} if args.note else {}), "results": results, } stem.with_suffix(".json").write_text(json.dumps(payload, indent=2, ensure_ascii=False), encoding="utf-8") md = "\n\n".join(to_markdown(results, g) for g in GOLDS) if args.note: md = f"> {args.note}\n\n{md}" stem.with_suffix(".md").write_text(md + "\n", encoding="utf-8") stem.with_suffix(".tex").write_text("\n\n".join(to_latex(results, g, args.split, len(q)) for g in GOLDS) + "\n", encoding="utf-8") # Per-query rankings, for error analysis and the lean_leansearch comparison. with stem.with_name(stem.name + "_rankings.jsonl").open("w", encoding="utf-8") as f: for i, item in enumerate(q): f.write(json.dumps({"name": item["name"], "top": {m: r[i][:32] for m, r in rankings.items()}}, ensure_ascii=False) + "\n") sys.stdout.write(md + "\n") if __name__ == "__main__": main()