| """End-to-end pipeline: data -> runs -> per-query metrics -> golden-set simulation. | |
| Latency is measured separately (scripts/latency.py) because it needs a fixed thread count. | |
| Usage: uv run python scripts/run_all.py [--datasets nfcorpus] [--R 200] [--out results] | |
| """ | |
| import argparse | |
| import hashlib | |
| import json | |
| import platform | |
| import sys | |
| import time | |
| from importlib import metadata | |
| from pathlib import Path | |
| import faiss | |
| import numpy as np | |
| import pandas as pd | |
| import yaml | |
| from gatestats import data, metrics, retrieve, simulate | |
| ROOT = Path(__file__).resolve().parents[1] | |
| CACHE = ROOT / "cache" | |
| def sha256(path): | |
| return hashlib.sha256(Path(path).read_bytes()).hexdigest() | |
| def cached(path, fn): | |
| path = Path(path) | |
| if path.exists(): | |
| return np.load(path) | |
| arr = fn() | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| np.save(path, arr) | |
| return arr | |
| def chunked(stem, texts, fn, size=4096): | |
| """Apply fn to texts in chunks saved as stem.partNNNN.npy, so an interrupted run resumes where it stopped.""" | |
| parts = [] | |
| for i, start in enumerate(range(0, len(texts), size)): | |
| part = stem.with_name(f"{stem.name}.part{i:04d}.npy") | |
| parts.append(cached(part, lambda: fn(texts[start:start + size]))) | |
| out = np.concatenate(parts) | |
| for p in stem.parent.glob(f"{stem.name}.part*.npy"): | |
| p.unlink() | |
| return out | |
| def environment(cfg_path): | |
| import torch | |
| try: | |
| import psutil | |
| ram = round(psutil.virtual_memory().total / 2**30, 1) | |
| except ImportError: | |
| ram = None | |
| libs = ["torch", "sentence-transformers", "transformers", "faiss-cpu", "bm25s", "datasets", "numpy", "scipy"] | |
| return { | |
| "started": time.strftime("%Y-%m-%dT%H:%M:%S%z"), | |
| "python": sys.version, | |
| "platform": platform.platform(), | |
| "processor": platform.processor(), | |
| "cpu_count": __import__("os").cpu_count(), | |
| "ram_gb": ram, | |
| "torch_threads": torch.get_num_threads(), | |
| "faiss_threads": faiss.omp_get_max_threads(), | |
| "libs": {name: metadata.version(name) for name in libs}, | |
| "prereg_sha256": sha256(ROOT / "prereg.md"), | |
| "config_sha256": sha256(cfg_path), | |
| } | |
| def build_runs(ds, cfg, ds_cache, keep_hnsw_s0): | |
| k = cfg["top_k"] | |
| runs = {"bm25": cached(ds_cache / "run_bm25.npy", lambda: retrieve.bm25_run( | |
| retrieve.bm25_index(ds.doc_texts), ds.query_texts, k))} | |
| emb = {} | |
| for key, mcfg in cfg["models"].items(): | |
| model = None | |
| def encode(texts, prefix): | |
| nonlocal model | |
| model = model or retrieve.load_model(mcfg, cfg["max_seq_length"]) | |
| return retrieve.embed(model, texts, prefix, cfg["batch_size"]) | |
| docs = cached(ds_cache / f"emb_{key}_docs.npy", lambda: chunked(ds_cache / f"emb_{key}_docs", ds.doc_texts, | |
| lambda t: encode(t, mcfg["doc_prefix"]))) | |
| queries = cached(ds_cache / f"emb_{key}_queries.npy", lambda: encode(ds.query_texts, mcfg["query_prefix"])) | |
| emb[key] = (docs, queries) | |
| runs[key] = cached(ds_cache / f"run_{key}.npy", lambda: retrieve.exact_run(docs, queries, k)) | |
| h = cfg["hnsw"] | |
| docs, queries = emb[h["model"]] | |
| for seed in h["seeds"]: | |
| paths = {ef: ds_cache / f"run_hnsw_ef{ef}_s{seed}.npy" for ef in h["ef_search"]} | |
| if not all(p.exists() for p in paths.values()): | |
| index, perm = retrieve.hnsw_index(docs, h["M"], h["ef_construction"], seed) | |
| for ef, p in paths.items(): | |
| np.save(p, retrieve.hnsw_run(index, perm, queries, ef, h["k"])) | |
| if seed == h["seeds"][0] and keep_hnsw_s0: | |
| faiss.write_index(index, str(ds_cache / "hnsw_s0.faiss")) | |
| np.save(ds_cache / "hnsw_s0_perm.npy", perm) | |
| for ef, p in paths.items(): | |
| runs[f"hnsw_ef{ef}_s{seed}"] = np.load(p) | |
| rng = np.random.default_rng([cfg["seed"], 11]) | |
| for rate in cfg["deletion_rates"]: | |
| deleted = rng.choice(len(ds.doc_ids), int(rate * len(ds.doc_ids)), replace=False) | |
| runs[f"del{round(rate * 100):02d}"] = retrieve.drop_docs(runs["bge"], deleted, 10) | |
| return runs | |
| def main(): | |
| p = argparse.ArgumentParser() | |
| p.add_argument("--config", default=str(ROOT / "configs" / "default.yaml")) | |
| p.add_argument("--datasets", nargs="*") | |
| p.add_argument("--R", type=int) | |
| p.add_argument("--out", default=str(ROOT / "results")) | |
| args = p.parse_args() | |
| cfg = yaml.safe_load(Path(args.config).read_text()) | |
| out = Path(args.out) | |
| out.mkdir(parents=True, exist_ok=True) | |
| env = environment(args.config) | |
| env["timings_s"], env["datasets"] = {}, {} | |
| per_query, rows_t, rows_p, rows_c = [], [], [], [] | |
| names = args.datasets or cfg["datasets"] | |
| for name in names: | |
| ds_i = cfg["datasets"].index(name) | |
| t0 = time.time() | |
| ds = data.load(name) | |
| runs = build_runs(ds, cfg, CACHE / name, keep_hnsw_s0=name == cfg["latency"]["dataset"]) | |
| systems = {s: metrics.per_query(run, ds.qrels) for s, run in runs.items()} | |
| systems.update(simulate.controlled(systems["bge"], cfg["controlled_deltas"], | |
| np.random.default_rng([cfg["seed"], ds_i, 7]))) | |
| env["timings_s"][f"{name}_runs"] = round(time.time() - t0, 1) | |
| env["datasets"][name] = {"queries": len(ds.query_ids), "docs": len(ds.doc_ids)} | |
| print(f"[{name}] N={len(ds.query_ids)} docs={len(ds.doc_ids)} systems={len(systems)} " | |
| f"bge recall@5={systems['bge']['recall@5'].mean():.4f}", flush=True) | |
| for s, vals in systems.items(): | |
| per_query.append(pd.DataFrame({"dataset": name, "system": s, "qid": ds.query_ids, **vals})) | |
| t0 = time.time() | |
| t, pr, c = simulate.simulate(name, ds_i, systems, cfg, R=args.R) | |
| rows_t += t | |
| rows_p += pr | |
| rows_c += c | |
| env["timings_s"][f"{name}_simulation"] = round(time.time() - t0, 1) | |
| print(f"[{name}] simulation {env['timings_s'][f'{name}_simulation']}s", flush=True) | |
| pd.concat(per_query).to_parquet(out / "per_query.parquet", index=False) | |
| pd.DataFrame(rows_t).to_parquet(out / "threshold.parquet", index=False) | |
| pd.DataFrame(rows_p).to_parquet(out / "pairs.parquet", index=False) | |
| pd.DataFrame(rows_c).to_parquet(out / "contrasts.parquet", index=False) | |
| env["R"] = args.R or cfg["simulation"]["R"] | |
| (out / "env.json").write_text(json.dumps(env, indent=2)) | |
| if __name__ == "__main__": | |
| main() | |