File size: 5,296 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 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 | """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()
|