| """Run THX-01 on the exact nibzard decision-model-benchmark suites (hash-verified), same state & criteria as Jev.""" |
| import json, sys, time, collections |
| import numpy as np |
| sys.path.insert(0, ".") |
| import thx01 |
| S = "data/nibzard_suites/" |
| INS = ("Choose the single best option for the decision state. You must commit to a single option even if the state does not " |
| "determine the answer; in that case choose the least unreasonable option and report low confidence. Answer confidence " |
| "separately. The confidence is your probability that your chosen option is correct.") |
|
|
|
|
| def ece(conf, cor, bins=10): |
| conf, cor = np.asarray(conf, float), np.asarray(cor, float) |
| idx = np.digitize(conf, np.linspace(0, 1, bins + 1)[1:-1], right=True) |
| return float(sum((idx == b).sum() / len(conf) * abs(cor[idx == b].mean() - conf[idx == b].mean()) for b in range(bins) if (idx == b).any())) |
|
|
|
|
| MP = sys.argv[1] if len(sys.argv) > 1 else "." |
| ag = thx01.load(MP, device="cuda") |
| out = {} |
| for suite in ["s1_intent77", "s2_spam", "s3_cardinality", "s4_order", "s5_confidence"]: |
| items = [json.loads(l) for l in open(S + suite + ".jsonl")] |
| qs = [{"decision": {"type": "choice", "instructions": INS, "criteria": {o: o for o in it["options"]}}} for it in items] |
| t0 = time.perf_counter() |
| res = ag.predict_batch([it["state"].strip() for it in items], qs, batch_size=32) |
| wall = time.perf_counter() - t0 |
| ch, cf = [], [] |
| for it, r in zip(items, res): |
| a = r["answers"]["decision"]; ch.append(it["options"].index(a["choice"])); cf.append(float(a["probabilities"][a["choice"]])) |
| gold = [int(it["gold_index"]) for it in items] |
| g = [i for i, x in enumerate(gold) if x >= 0] |
| cor = [ch[i] == gold[i] for i in g] |
| m = {"n": len(items), "accuracy": float(np.mean(cor)), "ece": ece([cf[i] for i in g], cor), |
| "brier": float(np.mean([(cf[i] - float(ch[i] == gold[i])) ** 2 for i in g])), "ms_per_item": wall / len(items) * 1000} |
| if suite == "s4_order": |
| by = collections.defaultdict(list) |
| for it, c in zip(items, ch): |
| by[it["base_item_id"]].append(it["options"][c]) |
| m["flip_rate"] = sum(len(set(v)) > 1 for v in by.values()) / len(by) |
| if suite == "s5_confidence": |
| ng = [cf[i] for i, x in enumerate(gold) if x < 0] |
| m["admits_ignorance"] = float(np.mean([c <= 0.5 for c in ng])); m["mean_conf_no_good"] = float(np.mean(ng)) |
| if suite == "s3_cardinality": |
| byn = collections.defaultdict(list) |
| for it, c, gg in zip(items, ch, gold): |
| byn[json.loads(it["meta"].replace("'", '"'))["N"] if isinstance(it["meta"], str) else it["meta"]["N"]].append(c == gg) |
| m["by_N"] = {int(k): float(np.mean(v)) for k, v in sorted(byn.items())} |
| out[suite] = m |
| print(suite, {k: (round(v, 4) if isinstance(v, float) else v) for k, v in m.items()}, flush=True) |
| json.dump(out, open(sys.argv[2] if len(sys.argv) > 2 else "./results/jev_suites_thx01.json", "w"), indent=1) |
|
|
