File size: 4,083 Bytes
dd18e84 f8d2531 dd18e84 f8d2531 dd18e84 f8d2531 dd18e84 f8d2531 dd18e84 f8d2531 dd18e84 f8d2531 dd18e84 f8d2531 dd18e84 f8d2531 dd18e84 f8d2531 dd18e84 | 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 | #!/usr/bin/env python3
"""四头模型推理：一次前向输出 is_ad / kind / scope / severity（0–3 期望分数）。"""
import argparse
import json
from pathlib import Path
import torch
from transformers import AutoTokenizer
from dual_head_model import load_dual_head_any
KINDS = ["crypto", "gambling", "porn", "porn_bait", "promo", "scam", "spam_flood"]
DEMO = [
"加微信 abc123 免费领取涨粉秘籍，名额有限！",
"今晚合唱团排练改到七点半，别走错了",
"这款面霜绝了，用了一个月肉眼可见变好，链接放评论区",
]
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", default="outputs/ad-4head-teacher-heads/best")
ap.add_argument("--text", default=None)
ap.add_argument("--file", default=None, help="jsonl({text}) 或 txt，每行一条")
ap.add_argument("--threshold", type=float, default=0.5)
ap.add_argument("--kind_threshold", type=float, default=0.0,
help=">0 时 kind 置信度低于该值输出 ad_uncertain")
ap.add_argument("--batch_size", type=int, default=64)
ap.add_argument("--max_len", type=int, default=384)
args = ap.parse_args()
if args.text:
texts = [args.text]
elif args.file:
if args.file.endswith(".jsonl"):
texts = [json.loads(l)["text"] for l in open(args.file, encoding="utf-8") if l.strip()]
else:
texts = [l.rstrip("\n") for l in open(args.file, encoding="utf-8") if l.strip()]
else:
texts = DEMO
cfg_path = Path(args.model) / "dual_config.json"
dc = json.loads(cfg_path.read_text(encoding="utf-8")) if cfg_path.is_file() else {}
kind_labels = dc.get("kind_labels") or KINDS
scope_labels = dc.get("scope_labels") or []
severity_labels = dc.get("severity_labels") or []
device = "cuda" if torch.cuda.is_available() else "cpu"
model = load_dual_head_any(args.model, device=device, dtype=torch.bfloat16)
tok = AutoTokenizer.from_pretrained(args.model)
with torch.inference_mode():
for i in range(0, len(texts), args.batch_size):
chunk = texts[i:i + args.batch_size]
enc = tok(chunk, truncation=True, max_length=args.max_len,
padding=True, return_tensors="pt").to(device)
with torch.autocast(device_type=device, dtype=torch.bfloat16, enabled=device == "cuda"):
blog, klog, slog, vlog, _ = model.forward_all(enc["input_ids"], enc["attention_mask"])
p_ad = torch.softmax(blog.float(), -1)[:, 0].cpu()
pk = torch.softmax(klog.float(), -1).cpu()
ps = torch.softmax(slog.float(), -1).cpu() if slog is not None else None
pv = torch.softmax(vlog.float(), -1).cpu() if vlog is not None else None
for j, (t, pa) in enumerate(zip(chunk, p_ad)):
is_ad = bool(pa >= args.threshold)
k = int(pk[j].argmax())
kind = kind_labels[k] if is_ad else "normal"
if is_ad and args.kind_threshold > 0 and float(pk[j][k]) < args.kind_threshold:
kind = "ad_uncertain"
out = {"is_ad": is_ad, "p_ad": round(float(pa), 6),
"kind": kind, "kind_conf": round(float(pk[j][k]), 6),
"kind_probs": {n: round(float(pk[j][x]), 6) for x, n in enumerate(kind_labels)}}
if is_ad and ps is not None:
s = int(ps[j].argmax())
out["scope"] = scope_labels[s] if scope_labels else str(s)
out["scope_probs"] = {n: round(float(ps[j][x]), 6) for x, n in enumerate(scope_labels)}
if is_ad and pv is not None:
probs = [float(x) for x in pv[j]]
out["severity_score"] = round(sum(x * p for x, p in enumerate(probs)), 4)
out["severity_probs"] = {n: round(probs[x], 6) for x, n in enumerate(severity_labels)}
print(json.dumps({"text": t, **out}, ensure_ascii=False))
if __name__ == "__main__":
main()
|