| |
| """广告识别预测演示：单条文本 / jsonl / 纯文本文件，输出各类别概率。""" |
| import argparse |
| import json |
|
|
| import torch |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer |
|
|
| DEMO_TEXTS = [ |
| "加微信 xxFREE 免费领取涨粉秘籍，名额有限！", |
| "今晚合唱团排练改到七点半，别走错了", |
| "这款面霜绝了，用了一个月肉眼可见变好，链接放评论区", |
| "明天降温10度，记得给孩子加衣服", |
| ] |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--model", default="outputs/ad-clf/best") |
| ap.add_argument("--text", default=None, help="单条文本") |
| ap.add_argument("--file", default=None, help="jsonl({text}) 或 txt，每行一条") |
| ap.add_argument("--batch_size", type=int, default=32) |
| args = ap.parse_args() |
|
|
| tok = AutoTokenizer.from_pretrained(args.model) |
| model = AutoModelForSequenceClassification.from_pretrained( |
| args.model, attn_implementation="sdpa", dtype=torch.bfloat16).eval().cuda() |
| names = [model.config.id2label[i] for i in range(model.config.num_labels)] |
|
|
| 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_TEXTS |
|
|
| 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=512, padding=True, |
| return_tensors="pt").to("cuda") |
| probs = torch.softmax(model(**enc).logits.float(), dim=-1).cpu() |
| for t, p in zip(chunk, probs): |
| top = int(p.argmax()) |
| dist = " ".join(f"{n}={p[j]:.3f}" for j, n in enumerate(names)) |
| print(f"[{names[top]:>7}] {t[:60]} ({dist})") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
