File size: 5,598 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 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | #!/usr/bin/env python3
"""用本地微调的 Laya 模型给 archive 存档补标，并统计与生产判定的一致率。
用途有两个：
1. 质量校验：在已有生产判定（systemone/llm/manual）的消息上跑一遍，
算一致率 —— 这衡量的是"自训模型 vs 线上系统"的口径差。
2. 补标：未判定的消息用模型概率补出来，单独落文件并标记 label_source，
不与生产判定混在一起使用。
产物：
messages_scored.jsonl 每条有文本的消息 + p_ad
agreement.json 按 decider / 群 / 内容类型拆分的一致率
"""
from __future__ import annotations
import argparse
import collections
import json
import os
import sys
import time
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
BATCH = 128
MAX_LEN = 512
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--inp", default="data/processed/archive/messages.jsonl")
ap.add_argument("--model", default="outputs/ad-clf/best")
ap.add_argument("--out", default="data/processed/archive")
ap.add_argument("--max-len", type=int, default=MAX_LEN)
ap.add_argument("--batch", type=int, default=BATCH)
args = ap.parse_args()
device = "cuda" if torch.cuda.is_available() else "cpu"
tok = AutoTokenizer.from_pretrained(args.model)
model = AutoModelForSequenceClassification.from_pretrained(args.model).eval().to(device)
if device == "cuda":
model = model.to(torch.bfloat16)
ad_id = int(model.config.label2id["ad"])
t0 = time.time()
total = scored = 0
rows_meta: list[dict] = []
probs: list[float] = []
# 按长度排序后分批，减少 padding 浪费；批次内顺序在写回前还原
items = []
with open(args.inp, encoding="utf-8") as fh:
for line in fh:
r = json.loads(line)
total += 1
if not r.get("text"):
continue
items.append(r)
items.sort(key=lambda r: len(r["text"]))
print(f"[score] 待打分 {len(items)} 条（总 {total}）")
with torch.inference_mode():
for i in range(0, len(items), args.batch):
chunk = items[i:i + args.batch]
enc = tok([r["text"] for r in chunk], truncation=True, max_length=args.max_len,
padding=True, return_tensors="pt").to(device)
logits = model(**enc).logits.float()
p = torch.softmax(logits, dim=-1)[:, ad_id].tolist()
probs.extend(p)
rows_meta.extend(chunk)
scored += len(chunk)
if scored % 20000 < args.batch:
print(f" {scored}/{len(items)} {time.time() - t0:.0f}s")
# 用 (chat_id, message_id) 作稳定键：重新读文件会生成新对象，不能靠 id()
scored_map: dict[tuple[str, int], float] = {
(r["chat_id"], r["message_id"]): probs[k] for k, r in enumerate(rows_meta)
}
meta_by_key = {(r["chat_id"], r["message_id"]): r for r in rows_meta}
def agreement(rows: list[dict]) -> dict:
ok = collections.Counter()
n = collections.Counter()
for r in rows:
if r["label"] not in ("ad", "clean"):
continue
p_ad = scored_map.get((r["chat_id"], r["message_id"]))
if p_ad is None:
continue
pred = "ad" if p_ad >= 0.5 else "clean"
key_all = "all"
n[key_all] += 1
ok[key_all] += pred == r["label"]
for key in (f"decider={r.get('decider')}", f"chat={r.get('chat_id')}",
f"type={r.get('content_type')}"):
n[key] += 1
ok[key] += pred == r["label"]
return {k: {"n": n[k], "agree": round(ok[k] / n[k], 4)} for k in n if n[k] >= 10}
out_dir = args.out
out_path = os.path.join(out_dir, "messages_scored.jsonl")
pseudo = []
with open(out_path, "w", encoding="utf-8") as fh:
for line in open(args.inp, encoding="utf-8"):
r = json.loads(line)
p_ad = scored_map.get((r["chat_id"], r["message_id"]))
if p_ad is not None:
r["p_ad"] = round(p_ad, 6)
fh.write(json.dumps(r, ensure_ascii=False) + "\n")
if r.get("text") and r["label"] not in ("ad", "clean"):
r2 = dict(r)
r2.update({"label": "ad" if p_ad >= 0.5 else "clean",
"label_source": "laya_pseudo", "p_ad": round(p_ad, 6),
"label_tier": 9, "confidence": round(max(p_ad, 1 - p_ad), 4),
"decider": "laya_pseudo", "ad_kind": None})
pseudo.append(r2)
pseudo_path = os.path.join(out_dir, "messages_pseudo.jsonl")
with open(pseudo_path, "w", encoding="utf-8") as fh:
for r in pseudo:
fh.write(json.dumps(r, ensure_ascii=False) + "\n")
agree = agreement(list(meta_by_key.values()))
report = {
"model": args.model,
"device": device,
"scored": scored,
"pseudo_labeled": len(pseudo),
"pseudo_label_dist": dict(collections.Counter(r["label"] for r in pseudo)),
"agreement": agree,
"elapsed_s": round(time.time() - t0, 1),
}
with open(os.path.join(out_dir, "agreement.json"), "w", encoding="utf-8") as fh:
json.dump(report, fh, ensure_ascii=False, indent=2)
print(json.dumps(report, ensure_ascii=False, indent=2))
print(f"[out] {out_path}\n[out] {pseudo_path}")
if __name__ == "__main__":
sys.exit(main())
|