| """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] |
| |
| 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]}") |
|
|
