Download src/gatestats/simulate.py from bunyamineymen/rag-gate-stats-code: direct link, hf CLI and curl.
- Browser
- Download file 5.2 kB
-
https://huggingface.co/bunyamineymen/rag-gate-stats-code/resolve/main/src/gatestats/simulate.py
- Command line
-
hf download hf://bunyamineymen/rag-gate-stats-code/src/gatestats/simulate.py
-
curl -L -o simulate.py https://huggingface.co/bunyamineymen/rag-gate-stats-code/resolve/main/src/gatestats/simulate.py
5.2 kB
| """Golden-set subsampling simulation: gate decisions on R golden sets vs the population truth.""" | |
| import numpy as np | |
| from . import gates, stats | |
| EXACT = ["bm25", "bge", "minilm", "e5"] | |
| def scenarios(cfg): | |
| """(family, A, B) system pairs; A is the baseline, B the candidate.""" | |
| seeds = cfg["hnsw"]["seeds"] | |
| pairs = [("noop", f"hnsw_ef{ef}_s{seeds[0]}", f"hnsw_ef{ef}_s{s}") for ef in (16, 128) for s in seeds[1:]] | |
| pairs += [("natural", "bge", "minilm"), ("natural", "bge", "e5"), ("natural", "bge", "bm25"), | |
| ("natural", "hnsw_ef128_s0", "hnsw_ef16_s0")] | |
| pairs += [("natural", "bge", f"del{round(r * 100):02d}") for r in cfg["deletion_rates"]] | |
| pairs += [("controlled", "bge", f"ctl{round(d * 100):02d}") for d in cfg["controlled_deltas"]] | |
| return pairs | |
| def controlled(base, deltas, rng): | |
| """Zero all metrics of randomly ordered queries until population recall@5 drops by each delta.""" | |
| N = len(base["recall@5"]) | |
| order = rng.permutation(N) | |
| order = order[base["recall@5"][order] > 0] | |
| cum = np.cumsum(base["recall@5"][order]) / N | |
| out = {} | |
| for delta in deltas: | |
| chosen = order[: np.searchsorted(cum, delta) + 1] | |
| out[f"ctl{round(delta * 100):02d}"] = {m: np.where(np.isin(np.arange(N), chosen), 0.0, v) for m, v in base.items()} | |
| return out | |
| def _contrast(rows, key, x): | |
| mean, lo, hi = stats.mean_ci(x) | |
| rows.append({**key, "mean": mean, "lo": lo, "hi": hi}) | |
| def simulate(name, ds_i, systems, cfg, R=None): | |
| sim, seed = cfg["simulation"], cfg["seed"] | |
| R, B, alpha = R or sim["R"], sim["B"], sim["alpha"] | |
| N = len(systems["bge"]["recall@5"]) | |
| split = np.random.default_rng([seed, ds_i, 99]).permutation(N) | |
| truths = {"population": (np.arange(N), np.arange(N)), "split_half": (split[: N // 2], split[N // 2:])} | |
| rows_t, rows_p, rows_c = [], [], [] | |
| for ti, (truth, (pool, tidx)) in enumerate(truths.items()): | |
| for n in [n for n in sim["n"] if n <= N // 2]: | |
| rng = np.random.default_rng([seed, ds_i, ti, n]) | |
| S = rng.permuted(np.tile(pool, (R, 1)), axis=1)[:, :n] | |
| Wa = stats.bootstrap_weights(n, B, rng).astype(np.float32) | |
| Wb = stats.bootstrap_weights(n, B, rng).astype(np.float32) | |
| flips = rng.choice(np.array([-1.0, 1.0], np.float32), size=(B, n)) | |
| for metric in sim["metrics"]: | |
| for sysname in EXACT: | |
| v = systems[sysname][metric] | |
| true, est = v[tidx].mean(), v[S].mean(1) | |
| for d in sim["d"]: | |
| rows_t.append(dict(dataset=name, truth=truth, system=sysname, metric=metric, n=n, d=d, | |
| true_value=true, truth_pass=d >= 0, | |
| k_pass=int(gates.g1_threshold(v[S], true - d).sum()), R=R)) | |
| family = {} | |
| for fam, A, Bn in scenarios(cfg): | |
| a, b = systems[A][metric], systems[Bn][metric] | |
| XA, XB = a[S].astype(np.float32), b[S].astype(np.float32) | |
| delta, sd_d = (b - a)[tidx].mean(), (b - a)[pool].std(ddof=1) | |
| lb_p = stats.lower_bound((XB - XA) @ Wa.T, alpha) | |
| lb_u = stats.lower_bound(XB @ Wb.T - XA @ Wa.T, alpha) | |
| key = dict(dataset=name, truth=truth, family=fam, A=A, B=Bn, metric=metric, n=n, R=R) | |
| base = {**key, "delta_true": delta, "sd_d": sd_d} | |
| g1 = {h: gates.g1_threshold(XB, a[tidx].mean() - h) for h in sim["headroom"]} | |
| for h, x in g1.items(): | |
| rows_p.append({**base, "gate": "G1", "param": h, "k_pass": int(x.sum())}) | |
| for m in sim["margins"]: | |
| dec = {"G2": gates.g2_naive_delta(XA, XB, m), "G3": lb_p >= -m, "G3u": lb_u >= -m, | |
| "G3signflip": gates.g3_sign_flip(XA, XB, m, flips, alpha)} | |
| for g, x in dec.items(): | |
| rows_p.append({**base, "gate": g, "param": m, "k_pass": int(x.sum()), | |
| "analytic_pass": stats.noninferiority_pass_prob(delta, m, sd_d, n, alpha) | |
| if g == "G3" else np.nan}) | |
| h2 = g1[max(sim["headroom"])].astype(float) - dec["G3"] | |
| h3 = (~dec["G3u"]).astype(float) - (~dec["G3"]) | |
| _contrast(rows_c, {**key, "param": m, "contrast": "G1pass-G3pass"}, h2) | |
| _contrast(rows_c, {**key, "param": m, "contrast": "G3ufail-G3fail"}, h3) | |
| if fam == "noop": | |
| family[m] = family.get(m, 0) + h3 | |
| n_noop = sum(f == "noop" for f, _, _ in scenarios(cfg)) | |
| for m, total in family.items(): | |
| _contrast(rows_c, dict(dataset=name, truth=truth, family="noop", A="*", B="*", metric=metric, | |
| n=n, R=R, param=m, contrast="G3ufail-G3fail"), total / n_noop) | |
| return rows_t, rows_p, rows_c | |