"""exact-label accuracy of an HF checkpoint on a classify jsonl (greedy). usage: evalcl.py <model> <jsonl> [N]"""
import json, os, re, sys, time
from pathlib import Path
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
H = Path(os.path.expanduser("~/mt92"))
MODEL = str(H/sys.argv[1]); SET = H/sys.argv[2]; N = int(sys.argv[3]) if len(sys.argv) > 3 else 10**9
tok = AutoTokenizer.from_pretrained(MODEL); tok.padding_side = "left"
if tok.pad_token_id is None: tok.pad_token_id = 151643
model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).cuda().eval()
rows = [json.loads(l) for l in SET.read_text(encoding="utf-8").splitlines() if l.strip()][:N]
order = sorted(range(len(rows)), key=lambda i: len(rows[i]["prompt"]))
norm = lambda s: re.sub(r"\s+", " ", str(s).strip().lower())
correct = 0; ntok = []; t0 = time.time(); B = 128
for s in range(0, len(order), B):
    idx = order[s:s+B]; chunk = [rows[i] for i in idx]
    enc = tok([r["prompt"] for r in chunk], return_tensors="pt", padding=True, add_special_tokens=False).to("cuda")
    with torch.no_grad():
        o = model.generate(**enc, max_new_tokens=16, do_sample=False, pad_token_id=151643, eos_token_id=151643)
    L = enc["input_ids"].shape[1]
    for j, i in enumerate(idx):
        gen = tok.decode(o[j][L:], skip_special_tokens=True); ntok.append(int((o[j][L:] != 151643).sum()))
        correct += norm(gen) == norm(rows[i]["completion"])
print(f"=== {sys.argv[1]} on {SET.name} ({len(rows)} docs, {time.time()-t0:.0f}s) ===")
print(f"  ACCURACY {correct/len(rows):.4f} | out tok med {sorted(ntok)[len(ntok)//2]} max {max(ntok)}")