Files
jspace/scripts/gate_threshold_curve.py

75 lines
2.6 KiB
Python

"""Item 20 phase 3 (offline, no GPU): threshold sweep + oracle bound.
Inputs: gate_probs.json (per-item halt probabilities, phase 1) and
eval_code_merge_lut.json (per-item outcomes at k=0/1/2/4, phase 2).
For each threshold theta: k*(item) = first point where cumulative halt
mass >= theta (pre-loop consult included, remainder to kmax), outcome
looked up per item; reports overall/easy/hard accuracy and E[k].
Also: the oracle gate (best k per item) — the gating ceiling.
"""
import json
import os
from pathlib import Path
OUT = Path(os.environ.get("LOOP_OUT",
Path(__file__).resolve().parent.parent / "results-loop"))
probs = {r["task_id"]: r for r in json.load(open(OUT / "gate_probs.json"))}
lut = json.load(open(OUT / "eval_code_merge_lut.json"))
# lut["ks"][k]["per_item"] -> task_id, ok
ok = {}
for k, v in lut["ks"].items():
for r in v["per_item"]:
ok.setdefault(r["task_id"], {})[int(k)] = r["ok"]
KS = sorted(int(k) for k in lut["ks"]) # e.g. [0, 1, 2, 4]
KMAX = max(KS)
def kstar(p, theta):
"""p = [p0, p1, ..., pkmax]; deploy semantics of halted_k_per_item."""
cum, keep = 0.0, 1.0
for i, pi in enumerate(p):
cum += keep * pi
keep *= (1 - pi)
if cum >= theta:
return 0 if i == 0 else min(i, KMAX)
return KMAX
def nearest_k(k):
"""Map k* to the nearest evaluated k (lut has only KS)."""
return min(KS, key=lambda x: (abs(x - k), x))
def score(assign):
n, hit = {}, {}
ek = 0.0
for tid, k in assign.items():
lbl = probs[tid]["label"]
n[lbl] = n.get(lbl, 0) + 1
hit[lbl] = hit.get(lbl, 0) + int(ok[tid][nearest_k(k)])
ek += k
tot_n = sum(n.values())
tot = sum(hit.values()) / tot_n
return tot, {l: hit[l] / n[l] for l in n}, ek / tot_n
print(f"{'theta':>6} {'overall':>8} {'easy':>6} {'hard':>6} {'E[k]':>5}")
best = []
for theta in (0.3, 0.5, 0.7, 0.8, 0.9, 0.95, 0.98, 0.99):
assign = {tid: kstar(r["p"], theta) for tid, r in probs.items()}
tot, by, ek = score(assign)
print(f"{theta:>6} {tot:>8.3f} {by.get('easy',0):>6.3f} "
f"{by.get('hard',0):>6.3f} {ek:>5.2f}")
best.append({"theta": theta, "overall": tot, "by_label": by, "ek": ek})
oracle = {tid: max(KS, key=lambda k: (ok[tid][k], -k)) for tid in probs}
tot, by, ek = score(oracle)
print(f"{'oracle':>6} {tot:>8.3f} {by.get('easy',0):>6.3f} "
f"{by.get('hard',0):>6.3f} {ek:>5.2f}")
json.dump({"curve": best, "oracle": {"overall": tot, "by_label": by,
"ek": ek}},
open(OUT / "gate_threshold_curve.json", "w"), indent=1)
print("wrote", OUT / "gate_threshold_curve.json")