| """Reproduce AUBIN Mini's numbers on Kev's public locked suites (jaredpalmer/kev-suites, Apache-2.0). |
| |
| pip install ./code && python code/eval_kev.py --model emrevrg/AUBIN-Mini --variant phone --out kev_phone.json |
| python code/eval_kev.py --shard 0/4 ... # split the work across processes / machines, then --merge a.json b.json ... |
| |
| Same items, same metrics as the AUBIN/Kev protocol: accuracy, Brier, NLL, 15-bin ECE (overall and per source). |
| decision-v7 test = sources Kev trained on; transfer-v4 test = 764 records from sources Kev never trained on. |
| """ |
| import argparse |
| import json |
| import math |
| import time |
|
|
| import numpy as np |
|
|
| from aubin_mini import AubinMini, options_of |
|
|
| SUITES = {"kev_test": "v7/decision-v7/test.jsonl", "kev_transfer_test": "v4/transfer-v4/test.jsonl", |
| "kev_dev": "v7/decision-v7/development.jsonl", "kev_transfer_dev": "v4/transfer-v4/development.jsonl"} |
|
|
|
|
| def load(name): |
| from huggingface_hub import hf_hub_download |
| p = hf_hub_download("jaredpalmer/kev-suites", SUITES[name], repo_type="dataset") |
| items = [] |
| for i, line in enumerate(open(p, encoding="utf-8")): |
| r = json.loads(line) |
| for qid, q in r["questions"].items(): |
| keys, texts = options_of(q) |
| lab = q["label"] |
| lab = str(lab).lower() if isinstance(lab, bool) else str(lab) |
| skeys = [str(k).lower() if isinstance(k, bool) else str(k) for k in keys] |
| if lab not in skeys: |
| continue |
| items.append({"src": q.get("src") or r.get("_meta", {}).get("source") or name, "state": r["state"], |
| "q": q.get("instructions") or qid, "texts": texts, "y": skeys.index(lab)}) |
| return items |
|
|
|
|
| def metrics(rows): |
| by, allr = {}, [] |
| for src, y, lp in rows: |
| p = np.exp(np.asarray(lp)); oh = np.zeros(len(p)); oh[y] = 1 |
| r = (int(p.argmax()) == y, float(((p - oh) ** 2).sum()), float(p.max()), float(-math.log(max(p[y], 1e-12)))) |
| allr.append(r); by.setdefault(src, []).append(r) |
|
|
| def m(v): |
| ece = 0.0 |
| for j in range(15): |
| bb = [x for x in v if j / 15 < x[2] <= (j + 1) / 15] |
| if bb: |
| ece += len(bb) / len(v) * abs(np.mean([x[2] for x in bb]) - np.mean([x[0] for x in bb])) |
| return {"n": len(v), "accuracy": round(float(np.mean([x[0] for x in v])), 4), "brier": round(float(np.mean([x[1] for x in v])), 4), |
| "nll": round(float(np.mean([x[3] for x in v])), 4), "ece15": round(float(ece), 4)} |
| return {"all": m(allr), "by_source": {k: m(v) for k, v in sorted(by.items())}} |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--model", default="emrevrg/AUBIN-Mini"); ap.add_argument("--variant", default="phone") |
| ap.add_argument("--suites", default="kev_test,kev_transfer_test"); ap.add_argument("--threads", type=int, default=None) |
| ap.add_argument("--shard", default="0/1"); ap.add_argument("--out", default="aubin_mini_kev.json") |
| ap.add_argument("--merge", nargs="*", default=None, help="merge shard outputs instead of running") |
| a = ap.parse_args() |
| if a.merge: |
| rows = {} |
| for f in a.merge: |
| for s, v in json.load(open(f))["rows"].items(): |
| rows.setdefault(s, []).extend(v) |
| R = {s: metrics(v) for s, v in rows.items()} |
| json.dump(R, open(a.out, "w"), indent=1); print(json.dumps({s: R[s]["all"] for s in R})) |
| return |
| k, n = map(int, a.shard.split("/")) |
| m = AubinMini.from_pretrained(a.model, a.variant, threads=a.threads) |
| out, ms = {"variant": a.variant, "rows": {}}, [] |
| for s in a.suites.split(","): |
| its = load(s)[k::n] |
| rows = [] |
| for it in its: |
| t = time.perf_counter() |
| lp = m.log_probs(it["q"], it["texts"], it["state"]) |
| ms.append((time.perf_counter() - t) * 1000) |
| rows.append((it["src"], it["y"], [round(float(x), 6) for x in lp])) |
| out["rows"][s] = rows |
| print(s, json.dumps(metrics(rows)["all"]), flush=True) |
| out["ms_per_question"] = {"median": round(float(np.median(ms)), 1), "p95": round(float(np.percentile(ms, 95)), 1)} |
| json.dump(out, open(a.out, "w")) |
| print("ms/question", out["ms_per_question"], flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
