File size: 9,150 Bytes
dfe1826 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 | """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/<key>/ (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+<key>/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_<split>)")
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()
|