Spaces:
Running on Zero
Running on Zero
| """Tune every method's knobs on dev splits only (scifact.train, nfcorpus.dev, fiqa.dev; one global | |
| setting per method), then evaluate on the five BEIR test splits. | |
| Usage: python run_eval.py [--quick] | |
| """ | |
| import itertools | |
| import json | |
| import os | |
| import sys | |
| import time | |
| import numpy as np | |
| sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")) | |
| from common import OUT, evaluate, load_qrels, order_to_run # noqa: E402 | |
| from infomax_rerank import by_score, fused, infomax, mmr, zfuse # noqa: E402 | |
| DEV = ["scifact.train", "nfcorpus.dev", "fiqa.dev"] | |
| TEST = ["scifact.test", "nfcorpus.test", "fiqa.test", "scidocs.test", "trec-covid.test"] | |
| def load(tag): | |
| name, split = tag.split(".") | |
| z = np.load(os.path.join(OUT, f"{tag}.scores.npz")) | |
| qs = [] | |
| for q in z["qids"]: | |
| qs.append({k: z[f"{q}|{k}"] for k in ("ids", "bm25", "ce", "emb", "qemb")} | {"qid": str(q)}) | |
| return qs, load_qrels(name, split) | |
| def score(Q, qrels, orders): | |
| run = {x["qid"]: order_to_run([x["ids"][j] for j in o]) for x, o in zip(Q, orders)} | |
| m, per = evaluate(run, qrels) | |
| return m, {q: v["ndcg_cut_10"] for q, v in per.items()} | |
| def method_orders(Q, method, p): | |
| if method == "BM25": | |
| return [by_score(x["bm25"]) for x in Q] | |
| if method == "Dense": | |
| return [by_score(x["emb"] @ x["qemb"]) for x in Q] | |
| if method == "CE": | |
| return [by_score(x["ce"]) for x in Q] | |
| if method == "MMR": | |
| return [mmr(x["ce"], x["emb"], **p) for x in Q] | |
| if method == "Fusion": | |
| return [zfuse(x["ce"], x["emb"], x["qemb"], **p) for x in Q] | |
| if method == "InfoMax": | |
| return [infomax(x["ce"], x["emb"], x["qemb"], **p) for x in Q] | |
| if method == "InfoMax+F": # same selector, relevance/anchor = the tuned CE+dense fusion score | |
| p = dict(p) | |
| a = p.pop("alpha") | |
| return [infomax(fused(x["ce"], x["emb"], x["qemb"], a), x["emb"], x["qemb"], **p) for x in Q] | |
| raise ValueError(method) | |
| GRIDS = { | |
| "MMR": [{"lam": l} for l in (0.5, 0.6, 0.7, 0.8, 0.9, 0.95)], | |
| "Fusion": [{"alpha": a} for a in (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9)], | |
| "InfoMax": [dict(rank=r, gamma=g, rho=rh, sigma2=s, beta=b, delta=d) | |
| for r, g, rh, s, b, d in itertools.product( | |
| (1, 4, 16), (1.0, 2.0, 4.0), (0.3, 1.0), (0.03, 0.3), (0.0, 0.3, 3.0), | |
| (0.0, 0.5, 2.0, np.inf))], | |
| "InfoMax+F": [dict(rank=r, gamma=g, rho=1.0, sigma2=s, beta=b, delta=d) | |
| for r, g, s, b, d in itertools.product( | |
| (1, 4, 16), (1.0, 2.0, 4.0, 8.0), (0.03, 0.3), (0.0, 0.3, 3.0), (0.0, 0.5, 2.0, np.inf))], | |
| } | |
| def boot_ci(a, b, n=10000, seed=0): | |
| d = np.asarray(a) - np.asarray(b) | |
| rng = np.random.default_rng(seed) | |
| m = d[rng.integers(0, len(d), (n, len(d)))].mean(1) | |
| return float(d.mean()), float(np.percentile(m, 2.5)), float(np.percentile(m, 97.5)) | |
| def main(): | |
| quick = "--quick" in sys.argv | |
| best = {} | |
| bp = os.path.join(OUT, "best_params.json") | |
| if "--test" in sys.argv: | |
| best = {k: {kk: (np.inf if vv is None else vv) for kk, vv in v.items()} | |
| for k, v in json.load(open(bp)).items()} | |
| grids = {} | |
| else: | |
| dev = {t: load(t) for t in DEV} | |
| grids = GRIDS | |
| for meth, grid in grids.items(): | |
| if meth == "InfoMax+F": | |
| grid = [p | {"alpha": best["Fusion"]["alpha"]} for p in grid] | |
| if quick: | |
| grid = grid[:4] | |
| t0 = time.time() | |
| res = [] | |
| for p in grid: | |
| v = np.mean([score(Q, qr, method_orders(Q, meth, p))[0]["nDCG@10"] for Q, qr in dev.values()]) | |
| res.append((v, p)) | |
| res.sort(key=lambda t: -t[0]) | |
| best[meth] = res[0][1] | |
| ce_dev = np.mean([score(Q, qr, method_orders(Q, "CE", {}))[0]["nDCG@10"] for Q, qr in dev.values()]) | |
| print(f"[tune] {meth}: best dev nDCG@10 {res[0][0]:.4f} (CE {ce_dev:.4f}) with {res[0][1]} " | |
| f"[{len(grid)} settings, {time.time()-t0:.0f}s]", flush=True) | |
| if meth.startswith("InfoMax"): | |
| for v, p in res[:5]: | |
| print(f" {v:.4f} {p}") | |
| if "--test" not in sys.argv: | |
| json.dump({k: {kk: (None if vv == np.inf else vv) for kk, vv in v.items()} for k, v in best.items()}, | |
| open(bp, "w"), indent=1) | |
| return | |
| methods = ["BM25", "Dense", "CE", "MMR", "Fusion", "InfoMax", "InfoMax+F"] | |
| table, perq = {}, {} | |
| for tag in TEST: | |
| Q, qr = load(tag) | |
| for meth in methods: | |
| o = method_orders(Q, meth, best.get(meth, {})) | |
| m, pq = score(Q, qr, o) | |
| table[(tag, meth)] = m | |
| perq[(tag, meth)] = pq | |
| print("\nnDCG@10 on test (BM25 top-50 pools; knobs tuned on dev only)") | |
| print(f"{'dataset':16s}" + "".join(f"{m:>9s}" for m in methods)) | |
| for tag in TEST: | |
| print(f"{tag:16s}" + "".join(f"{table[(tag, m)]['nDCG@10']:9.4f}" for m in methods)) | |
| print("\nInfoMax+F minus X, nDCG@10, paired bootstrap 95% CI") | |
| rows = {} | |
| for tag in TEST: | |
| qids = sorted(perq[(tag, "CE")]) | |
| for other in ("CE", "MMR", "Fusion"): | |
| a = [perq[(tag, "InfoMax+F")][q] for q in qids] | |
| b = [perq[(tag, other)][q] for q in qids] | |
| rows[(tag, other)] = boot_ci(a, b) | |
| print(f"{tag:16s}" + " ".join(f"vs {o}: {rows[(tag, o)][0]:+.4f} [{rows[(tag, o)][1]:+.4f},{rows[(tag, o)][2]:+.4f}]" | |
| for o in ("CE", "MMR", "Fusion"))) | |
| json.dump({"best": {k: {kk: (None if vv == np.inf else vv) for kk, vv in v.items()} for k, v in best.items()}, | |
| "test": {f"{t}|{m}": v for (t, m), v in table.items()}, | |
| "ci_infomax_minus": {f"{t}|{o}": v for (t, o), v in rows.items()}}, | |
| open(os.path.join(OUT, "results.json"), "w"), indent=1) | |
| if __name__ == "__main__": | |
| main() | |