File size: 2,133 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 | #!/usr/bin/env python3
"""广告识别预测演示：单条文本 / 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()
|