| """Run a student checkpoint over a JSONL and dump per-record predictions (no metrics here; see student_report.py).
|
|
|
| Each output line: id, per-question probs + logits (choice/noul/score), per-box-question decoded box
|
| [top, bottom, left, right] + IoU-confidence + max-prob confidence, plus gold, normalized teacher, scenario and box_pairs
|
| copied from the input so the report is standalone. --sample N evaluates a seeded random subset (e.g. train-gap check).
|
| --resume appends and skips ids already in the output.
|
|
|
| usage: python scripts/eval_student.py --ckpt checkpoints/x --data d.jsonl --out runs/preds.jsonl [--sample 300 --seed 0] [--resume]
|
| """
|
| import os; os.environ.setdefault("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "1")
|
| import argparse, json, random, sys, time
|
| 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
|
| from glint.joint_schema_model import collate_records
|
| from glint.schema import encode_glint_record
|
| from glint.box import box_iou, decode_box, max_prob_confidence
|
|
|
|
|
| @torch.inference_mode()
|
| def main():
|
| ap = argparse.ArgumentParser()
|
| ap.add_argument("--ckpt", required=True); ap.add_argument("--data", required=True); ap.add_argument("--out", required=True)
|
| ap.add_argument("--sample", type=int, default=None); ap.add_argument("--seed", type=int, default=0)
|
| ap.add_argument("--resume", action="store_true"); ap.add_argument("--stats", default=None)
|
| args = ap.parse_args(); dev = default_device()
|
| rows = load_jsonl(args.data)
|
| if args.sample is not None and args.sample < len(rows):
|
| keep = set(random.Random(args.seed).sample(range(len(rows)), args.sample))
|
| rows = [r for i, r in enumerate(rows) if i in keep]
|
| done = set()
|
| if args.resume and Path(args.out).exists():
|
| for line in Path(args.out).read_text(encoding="utf-8").splitlines():
|
| try:
|
| done.add(json.loads(line)["id"])
|
| except (ValueError, KeyError):
|
| pass
|
| todo = [r for r in rows if r["id"] not in done]
|
| print(f"{len(rows)} records, {len(done)} already done, {len(todo)} to go", flush=True)
|
| t0 = time.time(); model, proc, meta = load_checkpoint(args.ckpt, dev); load_s = time.time() - t0
|
| print(f"loaded {args.ckpt} in {load_s:.1f}s", flush=True)
|
| Path(args.out).parent.mkdir(parents=True, exist_ok=True)
|
| lat = []
|
| if dev == "cuda":
|
| torch.cuda.reset_peak_memory_stats()
|
| with open(args.out, "a" if args.resume else "w", encoding="utf-8", newline="\n") as f:
|
| for i, ex in enumerate(todo):
|
| ts = time.time()
|
| 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)))
|
| out = {"id": ex["id"], "probs": {}, "logits": {}, "boxes": {}}
|
| for q, z in zip(enc.questions, all_logits[0]):
|
| z = z.float().cpu()
|
| out["logits"][q.question_id] = dict(zip(q.option_ids, [round(v, 5) for v in z.tolist()]))
|
| out["probs"][q.question_id] = dict(zip(q.option_ids, [round(v, 6) for v in z.softmax(-1).tolist()]))
|
| if boxes[0] is not None:
|
| b = boxes[0]; edges = decode_box(b.logits).cpu(); conf = b.conf_logit.float().sigmoid().cpu()
|
| maxp = max_prob_confidence(b.logits).cpu(); gb = normalize_box_gold(ex)
|
| for j, qid in enumerate(b.question_ids):
|
| g = gb.get(qid)
|
| out["boxes"][qid] = {"box": [round(v, 5) for v in edges[j].tolist()], "conf": round(float(conf[j]), 5),
|
| "maxp": round(float(maxp[j]), 5),
|
| "iou": None if g is None else round(float(box_iou(edges[j], torch.tensor(g))), 5)}
|
| if dev == "cuda":
|
| torch.cuda.synchronize()
|
| dt = time.time() - ts; lat.append(dt)
|
| m = ex.get("meta") or {}
|
| out.update({"sec": round(dt, 3), "gold": ex.get("gold") or {}, "teacher": normalize_teacher(ex),
|
| "meta": {"scenario": m.get("scenario"), "box_pairs": m.get("box_pairs"), "lang": m.get("lang"),
|
| "theme": m.get("theme"), "style_look": m.get("style_look")}})
|
| f.write(json.dumps(out, ensure_ascii=False) + "\n"); f.flush()
|
| if i % 25 == 0:
|
| print(json.dumps({"i": i, "id": ex["id"], "sec": round(dt, 3), "elapsed_min": round((time.time() - t0) / 60, 1)}), flush=True)
|
| srt = sorted(lat[1:] or lat)
|
| stats = {"ckpt": args.ckpt, "data": args.data, "n": len(lat), "load_seconds": load_s,
|
| "median_sec": srt[len(srt) // 2] if srt else None, "total_seconds": time.time() - t0,
|
| "peak_mem_gb": torch.cuda.max_memory_allocated() / 2**30 if dev == "cuda" else None}
|
| print(json.dumps(stats), flush=True)
|
| if args.stats:
|
| Path(args.stats).write_text(json.dumps(stats, indent=1), encoding="utf-8")
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
