| """Evaluate a checkpoint directory with laya's own PyTorch forward (the code path consumers load), |
| and fit calibration temperatures. |
| |
| python -m xdecision.evaluate --ckpt DIR [--calibrate] [--suites typed_decisions,emotion,...] |
| |
| --calibrate fits one temperature per (question type, option-count bucket) on data/built/calib.pt |
| by held-out NLL and writes them into DIR/rl_agent_config.json before evaluating. |
| """ |
| import argparse |
| import json |
| import math |
| import os |
| import time |
| from collections import defaultdict |
|
|
| import numpy as np |
| import torch |
|
|
| from laya.common import QTYPE_NAMES, clamp_temperature, ece_score, temp_bucket |
|
|
| from .model_io import load_config, load_model, load_tokenizer |
| from .train import collate |
| from .devices import pick_device |
|
|
| SUITES = ["typed_decisions", "emotion", "sst5", "prompt_injection", "xcopa", "xwinograd", "belebele", "xstory", |
| "wd_heldout", "wiki_heldout", "hans_formats", "unseen_workflows", "market_heldout", "uncertainty", "fitness_unseen", |
| "unseen_workflows_extra"] |
|
|
|
|
| @torch.no_grad() |
| def logits_for(model, items, pad_id, device, max_tokens=16384): |
| order = sorted(range(len(items)), key=lambda i: len(items[i]["ids"])) |
| out = [None] * len(items) |
| act_out = [None] * len(items) |
| i = 0 |
| while i < len(order): |
| L = len(items[order[i]]["ids"]) |
| bs = max(1, min(128, max_tokens // max(64, ((L + 63) // 64) * 64))) |
| idx = order[i:i + bs] |
| |
| while len(idx) > 1 and len(items[idx[-1]]["ids"]) * len(idx) > max_tokens * 1.5: |
| idx = idx[:-1] |
| b = collate([items[j] for j in idx], pad_id) |
| b = {k: v.to(device) for k, v in b.items()} |
| lg, act = model(b["input_ids"], b["attention_mask"], b["marker_pos"], b["marker_mask"], b["qtype"]) |
| lg = lg.float().cpu().numpy() |
| act = torch.softmax(act.float(), -1)[:, 0].cpu().numpy() |
| for r, j in enumerate(idx): |
| out[j] = lg[r, : len(items[j]["markers"])] |
| act_out[j] = float(act[r]) |
| i += len(idx) |
| return out, act_out |
|
|
|
|
| def temps_from_cfg(cfg): |
| t = [clamp_temperature(x) for x in cfg.get("temperature", [1.0, 1.0, 1.0])] |
| tb = {k: clamp_temperature(v) for k, v in (cfg.get("temperature_by_options") or {}).items()} |
| return t, tb |
|
|
|
|
| def probs(z, qtype, t, tb): |
| T = tb.get(temp_bucket(qtype, len(z)), t[qtype]) |
| z = np.asarray(z, dtype=np.float64) / T |
| p = np.exp(z - z.max()) |
| return p / p.sum() |
|
|
|
|
| def fit_temperature(pairs, lo=0.5, hi=5.0): |
| """1-D NLL minimisation over log T (golden section; the NLL is unimodal in log T).""" |
| def nll(logT): |
| T = math.exp(logT) |
| s = 0.0 |
| for z, tgt in pairs: |
| z = np.asarray(z, dtype=np.float64) / T |
| lse = z.max() + math.log(np.exp(z - z.max()).sum()) |
| s -= float(np.dot(tgt, z - lse)) |
| return s / len(pairs) |
| a, b = math.log(lo), math.log(hi) |
| g = (math.sqrt(5) - 1) / 2 |
| c, d = b - g * (b - a), a + g * (b - a) |
| fc, fd = nll(c), nll(d) |
| for _ in range(40): |
| if fc < fd: |
| b, d, fd = d, c, fc |
| c = b - g * (b - a) |
| fc = nll(c) |
| else: |
| a, c, fc = c, d, fd |
| d = a + g * (b - a) |
| fd = nll(d) |
| return math.exp((a + b) / 2) |
|
|
|
|
| def calibrate(model, tok, device, calib_items, min_n=40): |
| lg, _ = logits_for(model, calib_items, tok.pad_token_id, device) |
| by_bucket, by_type = defaultdict(list), defaultdict(list) |
| for z, it in zip(lg, calib_items): |
| pair = (z, np.asarray(it["target"], dtype=np.float64)) |
| by_bucket[temp_bucket(it["qtype"], len(z))].append(pair) |
| by_type[it["qtype"]].append(pair) |
| t = [round(fit_temperature(by_type[q]), 4) if len(by_type[q]) >= min_n else 1.0 for q in range(3)] |
| tb = {k: round(fit_temperature(v), 4) for k, v in sorted(by_bucket.items()) if len(v) >= min_n} |
| return t, tb, {k: len(v) for k, v in by_bucket.items()} |
|
|
|
|
| def auroc(scores, labels): |
| s, y = np.asarray(scores), np.asarray(labels).astype(bool) |
| if y.all() or (~y).all(): |
| return float("nan") |
| ranks = s.argsort().argsort() + 1 |
| return float((ranks[y].sum() - y.sum() * (y.sum() + 1) / 2) / (y.sum() * (~y).sum())) |
|
|
|
|
| def score_items(items, lg, act, t, tb): |
| rows = [] |
| for it, z, a in zip(items, lg, act): |
| q = it["qtype"] |
| p = probs(z, q, t, tb) |
| tgt = np.asarray(it["target"]) |
| gold = int(it.get("label", int(tgt.argmax()))) |
| pred = int(p.argmax()) |
| row = {"src": it["src"], "lang": it.get("lang", "en"), "t": QTYPE_NAMES[q], "correct": pred == gold, |
| "conf": float(p.max()), "soft": float((p * tgt).sum()), "act": a, "pred": pred, "gold": gold, |
| "gold_prob": float(p[gold]), |
| "brier": float(np.sum((p - np.eye(len(p), dtype=np.float64)[gold]) ** 2)), |
| "determinate": bool(it.get("determinate", bool(tgt.max() >= 0.99) or "label" in it))} |
| if q == 1: |
| lv = np.arange(len(p)) |
| row["mae"] = abs(float((p * lv).sum()) - gold) |
| row["first"] = pred == 0 |
| if "workflow" in it: |
| row["workflow"] = it["workflow"] |
| rows.append(row) |
| return rows |
|
|
|
|
| def summarize(rows): |
| def agg(rs): |
| det = [r for r in rs if r["determinate"]] |
| if not det: |
| conf = np.asarray([r["conf"] for r in rs], dtype=np.float64) |
| act = np.asarray([r["act"] for r in rs], dtype=np.float64) |
| return {"n": len(rs), "mean_conf": round(float(conf.mean()), 4), |
| "p95_conf": round(float(np.percentile(conf, 95)), 4), |
| "high_conf_rate_075": round(float(np.mean(conf >= 0.75)), 4), |
| "mean_act": round(float(act.mean()), 4), |
| "p95_act": round(float(np.percentile(act, 95)), 4)} |
| conf = np.asarray([r["conf"] for r in det], dtype=np.float64) |
| correct = np.asarray([r["correct"] for r in det], dtype=bool) |
| accepted = conf >= 0.8 |
| order = np.argsort(-conf) |
| cumulative_risk = np.cumsum(~correct[order]) / np.arange(1, len(det) + 1) |
| d = {"n": len(det), "acc": round(float(np.mean([r["correct"] for r in det])), 4), |
| "soft_acc": round(float(np.mean([r["soft"] for r in det])), 4), |
| "ece": round(ece_score(np.array([r["conf"] for r in det]), np.array([r["correct"] for r in det], float)), 4), |
| "act_auroc": round(auroc([r["act"] for r in det], [r["correct"] for r in det]), 4), |
| "nll": round(float(np.mean([-math.log(max(r["gold_prob"], 1e-12)) for r in det])), 4), |
| "brier": round(float(np.mean([r["brier"] for r in det])), 4), |
| "wrong_high_conf_rate_080": round(float(np.mean((~correct) & accepted)), 4), |
| "coverage_at_080": round(float(accepted.mean()), 4), |
| "accuracy_at_080": round(float(correct[accepted].mean()), 4) if accepted.any() else None, |
| "aurc": round(float(cumulative_risk.mean()), 4)} |
| sc = [r for r in det if "mae" in r] |
| if sc: |
| d["score_mae"] = round(float(np.mean([r["mae"] for r in sc])), 4) |
| d["first_level_rate"] = round(float(np.mean([r["first"] for r in sc])), 4) |
| d["first_level_gold_rate"] = round(float(np.mean([r["gold"] == 0 for r in sc])), 4) |
| return d |
| out = {"all": agg(rows)} |
| for key in ("t", "lang", "workflow", "src"): |
| groups = defaultdict(list) |
| for r in rows: |
| if key in r: |
| groups[r[key]].append(r) |
| if len(groups) > 1: |
| out["by_" + key] = {k: agg(v) for k, v in sorted(groups.items())} |
| und = [r for r in rows if not r["determinate"]] |
| if und: |
| conf = np.asarray([r["conf"] for r in und], dtype=np.float64) |
| act = np.asarray([r["act"] for r in und], dtype=np.float64) |
| out["undetermined_mean_conf"] = round(float(conf.mean()), 4) |
| out["undetermined_p95_conf"] = round(float(np.percentile(conf, 95)), 4) |
| out["undetermined_high_conf_rate_075"] = round(float(np.mean(conf >= 0.75)), 4) |
| out["undetermined_mean_act"] = round(float(act.mean()), 4) |
| out["undetermined_p95_act"] = round(float(np.percentile(act, 95)), 4) |
| return out |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--ckpt", required=True) |
| ap.add_argument("--data", default="data/built") |
| ap.add_argument("--suites", default=",".join(SUITES + ["val"])) |
| ap.add_argument("--calibrate", action="store_true") |
| ap.add_argument("--calib-file", default="calib.pt") |
| ap.add_argument("--device", default="auto", help="auto, cuda[:index], mps or cpu; evaluation remains FP32") |
| ap.add_argument("--out", default=None) |
| ap.add_argument("--max-per-suite", type=int, default=0) |
| ap.add_argument("--threads", type=int, default=0) |
| args = ap.parse_args() |
| device = pick_device(args.device) |
| print(f"device={device} precision=fp32", flush=True) |
| if args.threads: |
| torch.set_num_threads(args.threads) |
| tok = load_tokenizer(args.ckpt) |
| model, cfg = load_model(args.ckpt, device) |
| model.eval() |
| t0 = time.time() |
| shipped = temps_from_cfg(cfg) |
| if args.calibrate: |
| calib = torch.load(os.path.join(args.data, args.calib_file), weights_only=False) |
| t, tb, counts = calibrate(model, tok, device, calib) |
| cfg["temperature"] = t |
| cfg["temperature_by_options"] = tb |
| cfg["calibration"] = {"items": len(calib), "per_bucket": counts, "method": "held-out NLL, T in [0.5, 5]"} |
| json.dump(cfg, open(os.path.join(args.ckpt, "rl_agent_config.json"), "w"), indent=2, ensure_ascii=False) |
| print(f"calibrated in {time.time()-t0:.0f}s: temperature={t} by_options={tb}", flush=True) |
| t, tb = temps_from_cfg(cfg) |
| report = {"ckpt": os.path.abspath(args.ckpt), "device": str(device), "precision": "fp32", |
| "temperature": t, "temperature_by_options": tb, "suites": {}, |
| "suites_as_shipped": {}} |
| for name in args.suites.split(","): |
| path = os.path.join(args.data, f"{name}.pt" if name == "val" else f"eval_{name}.pt") |
| if not os.path.exists(path): |
| continue |
| items = torch.load(path, weights_only=False) |
| if args.max_per_suite: |
| items = items[: args.max_per_suite] |
| lg, act = logits_for(model, items, tok.pad_token_id, device) |
| report["suites"][name] = summarize(score_items(items, lg, act, t, tb)) |
| if args.calibrate: |
| report["suites_as_shipped"][name] = summarize(score_items(items, lg, act, *shipped)) |
| a = report["suites"][name]["all"] |
| print(f"[{time.time()-t0:5.0f}s] {name:17s} n={a.get('n')} acc={a.get('acc')} ece={a.get('ece')} " |
| f"act_auroc={a.get('act_auroc')}", flush=True) |
| out = args.out or os.path.join(args.ckpt, "eval_report.json") |
| json.dump(report, open(out, "w"), indent=1, ensure_ascii=False) |
| print("wrote", out) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
