"""micro-F1 of an HF checkpoint on any jsonl set, greedy, subnet scorer.
usage: evalset.py <model_dir> <jsonl> [N] [B]"""
import json, os, sys, time
from pathlib import Path
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
H = Path(os.path.expanduser("~/mt92")); sys.path.insert(0, str(H/"repo"))
from microtensor.scoring.extraction import parse_entities, micro_f1
MODEL = str(H/sys.argv[1]); SET = H/sys.argv[2]
N = int(sys.argv[3]) if len(sys.argv) > 3 else 10**9
B = int(sys.argv[4]) if len(sys.argv) > 4 else 64
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]
# sort by length for efficient batching
order = sorted(range(len(rows)), key=lambda i: len(rows[i]["prompt"]))
preds = [None]*len(rows); golds = [None]*len(rows); ntok = []
t0 = time.time()
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=256, 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()))
        preds[i] = parse_entities(gen); golds[i] = parse_entities(rows[i]["completion"]) or set()
mal = sum(1 for p in preds if p is None)
tp = sum(len((p or set()) & g) for p, g in zip(preds, golds))
fp = sum(len((p or set()) - g) for p, g in zip(preds, golds))
fn = sum(len(g - (p or set())) for p, g in zip(preds, golds))
t = sorted(ntok)
print(f"=== {sys.argv[1]} on {SET.name} ({len(rows)} docs, {time.time()-t0:.0f}s) ===")
print(f"  MICRO-F1 {micro_f1(preds, golds):.4f} | P {tp/(tp+fp) if tp+fp else 0:.4f} R {tp/(tp+fn) if tp+fn else 0:.4f} | tp/fp/fn {tp}/{fp}/{fn} | malformed {mal} | out tok med {t[len(t)//2]} max {t[-1]}")