| """Evaluate a topic & format classifier (48 sigmoid outputs topic_* / format_*) on webtax/evals: en_llama405b (WebOrganizer Llama-3.1-405B |
| labels, English) and fw2_judged (65 languages, Gemma-4-31B judge). python webtax/eval_webtax.py --models A --out R.json""" |
| import argparse, json, os, sys |
| import numpy as np, pandas as pd, torch |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer |
| HERE = os.path.dirname(os.path.abspath(__file__)); sys.path.insert(0, HERE) |
| from metrics import acc_f1 |
| ap = argparse.ArgumentParser(); ap.add_argument("--models", required=True); ap.add_argument("--out", required=True); args = ap.parse_args() |
| EV = {n: pd.read_parquet(os.path.join(HERE, "evals", f"{n}.parquet")) for n in ("en_llama405b", "fw2_judged")} |
| res = json.load(open(args.out)) if os.path.exists(args.out) else {} |
| for spec in args.models.split(","): |
| tok = AutoTokenizer.from_pretrained(spec); m = AutoModelForSequenceClassification.from_pretrained(spec, torch_dtype=torch.bfloat16).cuda().eval() |
| labels = [m.config.id2label[i] for i in range(m.config.num_labels)] |
| ti = [i for i, l in enumerate(labels) if l.startswith("topic_")]; fi = [i for i, l in enumerate(labels) if l.startswith("format_")] |
| r = {} |
| for name, df in EV.items(): |
| P = np.zeros((len(df), len(labels))); order = np.argsort(df.text.str.len().values) |
| with torch.no_grad(): |
| for s in range(0, len(df), 32): |
| b = order[s:s + 32]; e = tok(df.text.iloc[b].tolist(), truncation=True, max_length=512, padding=True, return_tensors="pt") |
| P[b] = m(input_ids=e["input_ids"].cuda(), attention_mask=e["attention_mask"].cuda()).logits.float().cpu().numpy() |
| pt = np.array([labels[ti[j]][6:] for j in P[:, ti].argmax(1)]); pf = np.array([labels[fi[j]][7:] for j in P[:, fi].argmax(1)]) |
| r[name] = dict(topic=acc_f1(pt, df.ref_topic.values, df.lang.values), format=acc_f1(pf, df.ref_format.values, df.lang.values)) |
| print(spec, name, {k: {kk: round(vv, 4) for kk, vv in v.items()} for k, v in r[name].items()}, flush=True) |
| res[spec] = r; json.dump(res, open(args.out, "w"), indent=1) |
|
|
