| |
| """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 {} |
| |
| 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()) |
|
|
