scripts/run_all.py

"""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()