File size: 3,470 Bytes
68a82f9 | 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 | """Evaluate a trained run: top-1 vs jev_topic, macro-F1, per-item dump, confidence buckets.
Mirrors the eval/ schema promised in jev-topic-head (per_item jsonl + headline numbers).
No rules_table.json exists in this example, so show/hide is skipped (it's a fixed
lookup table downstream of the topic, not part of the head).
"""
import argparse, json, os
import numpy as np
import torch
import torch.nn.functional as F
import yaml
from sklearn.metrics import f1_score
from transformers import AutoModel, AutoTokenizer
from peft import PeftModel
from train_head import PrototypeHead, encode, read_jsonl
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--config", default="config.yaml")
ap.add_argument("--run", required=True, help="e.g. runs/seed0")
args = ap.parse_args()
cfg = yaml.safe_load(open(args.config))
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if cfg["train"]["dtype"] == "bf16" else torch.float32
hs = torch.load(os.path.join(args.run, "head", "head.pt"), map_location="cpu")
topics = hs["topics"]
topic2id = {t: i for i, t in enumerate(topics)}
tok = AutoTokenizer.from_pretrained(cfg["base_model"])
tok.padding_side = "right"
base = AutoModel.from_pretrained(hs.get("base_model", cfg["base_model"]),
torch_dtype=dtype).to(device)
base.config.use_cache = False
model = PeftModel.from_pretrained(base, os.path.join(args.run, "lora"))
head = PrototypeHead(hs["protos"], float(hs["scale"]))
head.load_state_dict({"protos": hs["protos"], "scale": hs["scale"]})
head = head.to(device).eval()
rows = read_jsonl(os.path.join(cfg["data"]["dir"], cfg["data"]["test"]))
texts = [r["text"] for r in rows]
logits, BATCH = [], 32
with torch.no_grad():
for i in range(0, len(texts), BATCH):
v = encode(model, tok, texts[i:i + BATCH], cfg["max_length"], device)
logits.append(head.logits(v).float().cpu())
logits = torch.cat(logits)
probs = F.softmax(logits, dim=-1)
conf, pred = probs.max(-1)
pred = pred.numpy()
gold = np.array([topic2id[r["jev_topic"]] for r in rows])
per_item = []
for r, p, c in zip(rows, pred, conf.numpy()):
per_item.append({"id": r["id"], "model_topic": topics[p],
"jev_topic": r["jev_topic"], "confidence": round(float(c), 4),
"agree": bool(p == topic2id[r["jev_topic"]])})
os.makedirs(os.path.join(args.run, "eval"), exist_ok=True)
with open(os.path.join(args.run, "eval", "per_item_test.jsonl"), "w") as f:
for r in per_item:
f.write(json.dumps(r) + "\n")
top1 = float((pred == gold).mean())
mf1 = float(f1_score(gold, pred, average="macro", zero_division=0,
labels=list(range(len(topics)))))
conf_all = conf.numpy()
summary = {"n": len(rows), "top1_vs_jev": round(top1, 4),
"macro_f1_vs_jev": round(mf1, 4),
"conf_lt_05": int((conf_all < 0.5).sum()),
"conf_gte_09": int((conf_all >= 0.9).sum()),
"errors_lt_05": int(((pred != gold) & (conf_all < 0.5)).sum()),
"errors_gte_09": int(((pred != gold) & (conf_all >= 0.9)).sum())}
json.dump(summary, open(os.path.join(args.run, "eval", "summary.json"), "w"), indent=2)
print(json.dumps(summary, indent=2))
if __name__ == "__main__":
main()
|