File size: 3,600 Bytes
e379257 | 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 | """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") # ROCm: faster SDPA (see RESULTS.md); set to 0 to disable
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()
|