File size: 4,301 Bytes
995bd7d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | """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()
|