File size: 4,493 Bytes
6438a35 | 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 84 85 86 87 | #!/usr/bin/env python3
"""Evaluate QEV2 multitask records by language, origin, and question type."""
from __future__ import annotations
import argparse, json, math, time
from collections import Counter, defaultdict
from pathlib import Path
import torch
from qev.model import QevModel
from qev.train import load_records, normalize
from qev.records import load_jsonl
def load_records_with_meta(path):
"""Normalize records while retaining source metadata for language splits."""
out = []
for raw in load_jsonl(path):
row = normalize(raw)
meta = dict(raw.get("_meta", {})) if isinstance(raw, dict) else {}
# Some corpora place language at the top level instead of _meta.
if "language" not in meta and isinstance(raw, dict) and raw.get("language"):
meta["language"] = raw["language"]
row["_meta"] = meta
out.append(row)
return out
def add(cm, y, p):
cm["n"] += 1
cm["correct"] += int(y == p)
if y and p: cm["tp"] += 1
elif y and not p: cm["fn"] += 1
elif not y and p: cm["fp"] += 1
else: cm["tn"] += 1
def metric(cm):
n = cm["n"]
acc = cm["correct"] / n if n else 0.0
prec = cm["tp"] / (cm["tp"] + cm["fp"]) if cm["tp"] + cm["fp"] else 0.0
rec = cm["tp"] / (cm["tp"] + cm["fn"]) if cm["tp"] + cm["fn"] else 0.0
f1 = 2 * prec * rec / (prec + rec) if prec + rec else 0.0
f2 = 5 * prec * rec / (4 * prec + rec) if 4 * prec + rec else 0.0
return {"n": n, "accuracy": acc, "precision": prec, "recall": rec, "f1": f1, "f2": f2,
"tp": cm["tp"], "fp": cm["fp"], "fn": cm["fn"], "tn": cm["tn"]}
def evaluate(args):
dtype = torch.bfloat16 if args.dtype == "bf16" and args.device.startswith("cuda") else torch.float32
model = QevModel.load_checkpoint(args.checkpoint, args.device, dtype, trainable=False)
model.eval()
records = load_records_with_meta(args.data)
if args.max_records: records = records[:args.max_records]
overall = defaultdict(lambda: Counter()); by_lang = defaultdict(lambda: defaultdict(Counter)); by_origin = defaultdict(lambda: defaultdict(Counter))
type_counts = Counter(); mismatches = Counter(); started = time.time()
with torch.inference_mode():
for i, row in enumerate(records, 1):
logits_many = model.logits_many(row, row["questions"], max_state=args.max_state, max_branch=args.max_branch)
preds = {}
for q, logits in zip(row["questions"], logits_many):
pred = int(torch.argmax(logits).item())
y = int(q["label"])
p = bool(pred)
truth = bool(y)
qid = q["id"]
add(overall[qid], truth, p); add(by_lang[row.get("_meta", {}).get("language")][qid], truth, p); add(by_origin[row.get("_meta", {}).get("origin")][qid], truth, p)
if qid != "is_dlp": type_counts[qid] += 1
preds[qid] = p
if "is_dlp" in preds and any(k != "is_dlp" and v for k, v in preds.items()) != preds["is_dlp"]:
lang = row.get("_meta", {}).get("language") or "unknown"; mismatches[lang] += 1
if i % args.progress == 0: print(json.dumps({"processed": i, "seconds": round(time.time()-started,1)}), flush=True)
report = {"checkpoint": args.checkpoint, "data": args.data, "records": len(records), "seconds": round(time.time()-started, 1),
"overall": {k: metric(v) for k,v in overall.items()},
"by_language": {lang: {k: metric(v) for k,v in qs.items()} for lang, qs in by_lang.items()},
"by_origin": {origin: {k: metric(v) for k,v in qs.items()} for origin, qs in by_origin.items()},
"logical_mismatch": dict(mismatches), "type_question_counts": dict(type_counts)}
Path(args.out).write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
print(json.dumps(report, ensure_ascii=False, indent=2))
if __name__ == "__main__":
p = argparse.ArgumentParser()
p.add_argument("--checkpoint", required=True); p.add_argument("--data", required=True); p.add_argument("--out", required=True)
p.add_argument("--device", default="cuda"); p.add_argument("--dtype", choices=["fp32", "bf16"], default="bf16")
p.add_argument("--max_records", type=int); p.add_argument("--max_state", type=int, default=12288); p.add_argument("--max_branch", type=int, default=16384); p.add_argument("--progress", type=int, default=250)
evaluate(p.parse_args())
|