File size: 2,607 Bytes
71dabc7 | 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 | #!/usr/bin/env python3
"""在 jsonl({text,label}) 上评估已训练的广告分类模型。
输出 accuracy / per-class P-R-F1 / 混淆矩阵，并把逐条概率落盘便于挑坏例。
用法:
python scripts/eval_ad.py # 默认评估 val 集
python scripts/eval_ad.py --data data/raw/my_test.jsonl
"""
import argparse
import json
import numpy as np
import torch
from sklearn.metrics import classification_report, confusion_matrix
from transformers import AutoModelForSequenceClassification, AutoTokenizer
def to_id(lab, label2id):
"""标签兼容：字符串名 / 旧数字 id 均可。"""
if isinstance(lab, int):
return lab
if lab in label2id:
return label2id[lab]
return int(lab)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", default="outputs/ad-clf/best")
ap.add_argument("--data", default="data/processed/val.jsonl")
ap.add_argument("--batch_size", type=int, default=64)
ap.add_argument("--max_len", type=int, default=512)
args = ap.parse_args()
tok = AutoTokenizer.from_pretrained(args.model)
model = AutoModelForSequenceClassification.from_pretrained(
args.model, attn_implementation="sdpa", dtype=torch.bfloat16).eval().cuda()
rows = [json.loads(l) for l in open(args.data, encoding="utf-8") if l.strip()]
y = np.array([to_id(r["label"], model.config.label2id) for r in rows])
preds, probs_all = [], []
with torch.inference_mode():
for i in range(0, len(rows), args.batch_size):
chunk = rows[i:i + args.batch_size]
enc = tok([r["text"] for r in chunk], truncation=True,
max_length=args.max_len, padding=True,
return_tensors="pt").to("cuda")
probs = torch.softmax(model(**enc).logits.float(), dim=-1).cpu().numpy()
preds.extend(probs.argmax(1).tolist())
probs_all.extend(probs.tolist())
names = [model.config.id2label[i] for i in range(model.config.num_labels)]
print(classification_report(y, preds, target_names=names, digits=4))
print(f"混淆矩阵（行=真实，列=预测，顺序 = {names}）:")
print(confusion_matrix(y, preds))
dump = args.model.rstrip("/\\") + "_val_probs.jsonl"
with open(dump, "w", encoding="utf-8") as f:
for r, p in zip(rows, probs_all):
f.write(json.dumps({"text": r["text"], "label": r["label"], "probs": p},
ensure_ascii=False) + "\n")
print("逐条概率已存:", dump)
if __name__ == "__main__":
main()
|