| """Accuracy / NLL / Brier / ECE of a checkpoint on a JSONL with gold labels (and agreement with teacher).
|
| If the data has `box` questions with a gold box, also reports a "box" block: mean IoU, IoU>=0.5 / >=0.75 rates,
|
| mean confidence and ECE(15) of the confidence vs the event IoU>=0.5 (plus the max-prob fallback's ECE)."""
|
| import os; os.environ.setdefault("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "1")
|
| import argparse, json, math, sys
|
| from pathlib import Path
|
| import torch
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
| from glint.model import load_checkpoint, default_device
|
| from glint.data import load_jsonl, to_record, normalize_teacher, normalize_box_gold, derive_presence_gold
|
| from glint.joint_schema_model import collate_records
|
| from glint.schema import encode_glint_record
|
| from glint.box import box_iou, box_metrics, decode_box, max_prob_confidence
|
|
|
| def ece(confs, hits, bins=15):
|
| tot, n = 0.0, len(confs)
|
| for b in range(bins):
|
| lo, hi = b / bins, (b + 1) / bins
|
| idx = [i for i, c in enumerate(confs) if lo < c <= hi or (b == 0 and c == 0)]
|
| if idx:
|
| tot += len(idx) / n * abs(sum(hits[i] for i in idx) / len(idx) - sum(confs[i] for i in idx) / len(idx))
|
| return tot
|
|
|
| @torch.inference_mode()
|
| def main():
|
| ap = argparse.ArgumentParser(); ap.add_argument("--ckpt", required=True); ap.add_argument("--data", required=True)
|
| args = ap.parse_args(); dev = default_device()
|
| model, proc, _ = load_checkpoint(args.ckpt, dev)
|
| confs, hits, nll, brier, agree, nt = [], [], 0.0, 0.0, 0, 0
|
| ious, bconfs, bmaxp, any_box = [], [], [], False
|
| for ex in load_jsonl(args.data):
|
| enc = encode_glint_record(proc.tokenizer, to_record(ex), processor=proc)
|
| all_logits, boxes = model.decide(collate_records([enc], proc.tokenizer.pad_token_id, torch.device(dev)))
|
| logits = all_logits[0]
|
| teacher = normalize_teacher(ex)
|
| if boxes[0] is not None:
|
| any_box = True; gold_boxes = normalize_box_gold(ex); pred = boxes[0]
|
| edges = decode_box(pred.logits).cpu(); conf = pred.conf_logit.float().sigmoid().cpu()
|
| maxp = max_prob_confidence(pred.logits).cpu()
|
| for i, qid in enumerate(pred.question_ids):
|
| g = gold_boxes.get(qid)
|
| if g is not None:
|
| ious.append(float(box_iou(edges[i], torch.tensor(g)))); bconfs.append(float(conf[i])); bmaxp.append(float(maxp[i]))
|
| ex = {**ex, "gold": derive_presence_gold(ex)}
|
| for q, z in zip(enc.questions, logits):
|
| p = z.float().softmax(-1).cpu(); pred = int(p.argmax())
|
| g = (ex.get("gold") or {}).get(q.question_id)
|
| if g is not None:
|
| y = q.option_ids.index(str(g)); confs.append(float(p[pred])); hits.append(int(pred == y))
|
| nll -= math.log(max(float(p[y]), 1e-9)); brier += float(((p - torch.nn.functional.one_hot(torch.tensor(y), len(p))) ** 2).sum())
|
| if q.question_id in teacher:
|
| t = teacher[q.question_id]; nt += 1; agree += int(q.option_ids[pred] == max(t, key=t.get))
|
| n = max(len(hits), 1)
|
| print(json.dumps({"n_questions": len(hits), "accuracy": sum(hits) / n, "nll": nll / n, "brier": brier / n,
|
| "ece_15bin": ece(confs, hits), "teacher_agreement": agree / max(nt, 1),
|
| **({"box": box_metrics(ious, bconfs, bmaxp)} if any_box else {})}, indent=1))
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
