| |
| """四头模型推理：一次前向输出 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() |
|
|
