File size: 11,121 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 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 | #!/usr/bin/env python3
"""通用 Telegram Desktop HTML 群聊导出解析器（区别于单发送者的频道导出）。
群导出与频道导出的关键差异：
- 多发送者：同发送者连续消息会折叠成 class 含 joined 的条目，没有 from_name，
需要沿用上一个显示的发送者，否则发送者级信息全部错位；
- 回复结构 reply_to；
- 标签不能走"频道先验=全广告"，改用本地模型打分 + 置信度分层，
高置信两端的可直接当伪标签，中间带标 needs_review 等人工裁决。
"""
from __future__ import annotations
import argparse
import collections
import datetime as dt
import hashlib
import json
import os
import re
import sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from build_channel_dataset import DomBuilder, parse_date # noqa: E402
PHOTO_RE = re.compile(r"^photos/photo_(\d+)@", re.IGNORECASE)
def parse_file(path: str, last_sender: dict) -> list[dict]:
with open(path, encoding="utf-8", errors="replace") as fh:
builder = DomBuilder()
builder.feed(fh.read())
out = []
for node in builder.root.iter():
if node.tag != "div" or "message" not in node.classes:
continue
if "service" in node.classes:
out.append({"kind": "service", "text": node.text().strip(), "service": True})
continue
raw_id = node.attrs.get("id", "")
mid = int(raw_id.replace("message", "")) if raw_id.startswith("message") else None
joined = "joined" in node.classes
date_node = next((n for n in node.iter() if "date" in n.classes and n.attrs.get("title")), None)
from_node = next((n for n in node.iter() if "from_name" in n.classes), None)
sender = from_node.text().strip() if from_node else (last_sender.get("name") if joined else None)
if from_node:
last_sender["name"] = sender
fwd = next((n for n in node.iter() if "forwarded" in n.classes), None)
reply = next((n for n in node.iter() if "reply_to" in n.classes), None)
text_node = next((n for n in node.iter() if "text" in n.classes and n.tag == "div"), None)
photos, media_kind = [], "none"
for n in node.iter():
if n.tag == "a" and "photo_wrap" in n.classes and n.attrs.get("href"):
photos.append(n.attrs["href"])
for mk in ("media_photo", "media_video", "media_file", "media_music", "media_sticker"):
if mk in n.classes and media_kind == "none":
media_kind = mk.replace("media_", "")
out.append({
"kind": "message", "message_id": mid, "source_file": os.path.basename(path),
"date": parse_date(date_node.attrs["title"]) if date_node else None,
"from_name": sender, "joined": joined,
"forwarded": fwd.text().strip()[:120] if fwd else None,
"is_reply": reply is not None,
"text": text_node.text().strip() if text_node else "",
"photos": photos, "media_kind": "photo" if photos else media_kind,
})
return out
def collect_photos(base: str) -> dict[str, dict]:
out: dict[str, dict] = {}
pdir = os.path.join(base, "photos")
if not os.path.isdir(pdir):
return out
for name in sorted(os.listdir(pdir)):
m = re.match(r"^photo_(\d+)@", name)
if not m:
continue
rec = out.setdefault(m.group(1), {"full": [], "thumbs": []})
(rec["thumbs"] if "_thumb" in name else rec["full"]).append(name)
return out
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--src", default="ChatExport_2026-10-27")
ap.add_argument("--out", default="data/processed/chatexport_2026-10-27")
ap.add_argument("--model", default="outputs/ad-clf/best")
ap.add_argument("--batch", type=int, default=128)
ap.add_argument("--max-len", type=int, default=512)
ap.add_argument("--low", type=float, default=0.3, help="低于此判 normal 的高置信伪标签")
ap.add_argument("--high", type=float, default=0.7, help="高于此判 ad 的高置信伪标签")
args = ap.parse_args()
src = os.path.abspath(args.src)
htmls = sorted((p for p in os.listdir(src) if re.match(r"^messages\d*\.html$", p)),
key=lambda p: int(re.search(r"(\d+)", p).group(1)) if re.search(r"\d+", p) else 0)
# 按导出编号顺序解析，joined 的发送者归属才正确
msgs, service = [], 0
last_sender = {"name": None}
for name in htmls:
parsed = parse_file(os.path.join(src, name), last_sender)
service += sum(1 for m in parsed if m.get("service"))
msgs.extend(m for m in parsed if not m.get("service"))
print(f"[parse] {len(htmls)} 个 HTML：消息 {len(msgs)}，service {service}")
# 去重 + 日期修正（同 id 的编辑版本取日期更晚的）
by_id: dict[int, dict] = {}
for m in msgs:
if m["message_id"] is None:
continue
prev = by_id.get(m["message_id"])
if prev is None or (m["date"] or 0) >= (prev["date"] or 0):
by_id[m["message_id"]] = m
print(f"[dedup] 保留 {len(by_id)}（重复 {len(msgs) - len(by_id)}）")
photos = collect_photos(src)
# ---------- 打分 ----------
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
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"])
items = sorted((m for m in by_id.values() if m["text"]), key=lambda m: len(m["text"]))
print(f"[score] 有文本 {len(items)} 条，设备 {device}")
probs: dict[int, float] = {}
t0 = dt.datetime.now()
with torch.inference_mode():
for i in range(0, len(items), args.batch):
chunk = items[i:i + args.batch]
enc = tok([m["text"] for m in chunk], truncation=True, max_length=args.max_len,
padding=True, return_tensors="pt").to(device)
p = torch.softmax(model(**enc).logits.float(), dim=-1)[:, ad_id].tolist()
for m, v in zip(chunk, p):
probs[m["message_id"]] = v
if (i // args.batch) % 20 == 0:
print(f" {i + len(chunk)}/{len(items)} {(dt.datetime.now() - t0).total_seconds():.0f}s")
# ---------- 落盘 ----------
os.makedirs(args.out, exist_ok=True)
rows, img_rows = [], []
seen_hash: set[str] = set()
stats = collections.Counter()
for mid, m in sorted(by_id.items()):
p_ad = probs.get(mid)
if p_ad is None:
label, source, conf = None, "unlabeled", None
elif p_ad >= args.high:
label, source, conf = "ad", "laya_pseudo", round(p_ad, 4)
elif p_ad <= args.low:
label, source, conf = "normal", "laya_pseudo", round(1 - p_ad, 4)
else:
label, source, conf = None, "needs_review", round(p_ad, 4)
imgs = []
for href in m["photos"]:
mm = PHOTO_RE.match(href)
rec = photos.get(mm.group(1)) if mm else None
if not rec or not rec["full"]:
stats["photo_missing"] += 1
continue
fname, path = rec["full"][0], os.path.join(src, "photos", rec["full"][0])
try:
from PIL import Image
with Image.open(path) as im:
im.load()
w, hgt, fmt = im.size[0], im.size[1], im.format or ""
except Exception:
stats["photo_corrupt"] += 1
continue
if w < 64 or hgt < 64:
stats["photo_too_small"] += 1
continue
with open(path, "rb") as fh:
digest = hashlib.sha256(fh.read()).hexdigest()
if digest in seen_hash:
stats["photo_dup"] += 1
continue
seen_hash.add(digest)
imgs.append({"path": os.path.relpath(path, os.path.dirname(src)).replace("\\", "/"),
"sha256": digest, "width": w, "height": hgt, "format": fmt,
"bytes": os.path.getsize(path)})
stats["photo_kept"] += 1
rows.append({
"chat_id": "TGCN机场联盟", "message_id": mid, "date": m["date"] or 0,
"iso_time": dt.datetime.fromtimestamp(m["date"]).isoformat(timespec="seconds") if m["date"] else None,
"from_name": m["from_name"], "forwarded": m["forwarded"], "is_reply": m["is_reply"],
"text": m["text"], "text_len": len(m["text"]), "media_kind": m["media_kind"],
"images": [i["path"] for i in imgs],
"p_ad": round(p_ad, 6) if p_ad is not None else None,
"label": label, "label_source": source, "confidence": conf,
"needs_review": source == "needs_review",
})
for i in imgs:
img_rows.append({**i, "message_id": mid, "date": m["date"] or 0,
"p_ad": round(p_ad, 6) if p_ad is not None else None,
"label": label, "label_source": source, "text": m["text"][:300]})
for name, data in (("messages.jsonl", rows), ("images.jsonl", img_rows)):
with open(os.path.join(args.out, name), "w", encoding="utf-8") as fh:
for r in data:
fh.write(json.dumps(r, ensure_ascii=False) + "\n")
lab = [r for r in rows if r["label"]]
report = {
"src": src, "html_files": len(htmls), "service": service,
"messages": len(rows), "with_text": sum(1 for r in rows if r["text"]),
"media_kind": dict(collections.Counter(r["media_kind"] for r in rows)),
"images": len(img_rows), "image_stats": dict(stats),
"label_dist": dict(collections.Counter(r["label_source"] for r in rows)),
"label_value": dict(collections.Counter(r["label"] for r in lab)),
"scored_by_threshold": {
"ad>=0.7": sum(1 for r in rows if r["p_ad"] is not None and r["p_ad"] >= args.high),
"normal<=0.3": sum(1 for r in rows if r["p_ad"] is not None and r["p_ad"] <= args.low),
"review_0.3-0.7": sum(1 for r in rows if r["p_ad"] is not None and args.low < r["p_ad"] < args.high),
},
"date_range": [dt.datetime.fromtimestamp(min(r["date"] for r in rows if r["date"])).isoformat(timespec="seconds"),
dt.datetime.fromtimestamp(max(r["date"] for r in rows if r["date"])).isoformat(timespec="seconds")],
"senders": len({r["from_name"] for r in rows if r["from_name"]}),
"out": os.path.abspath(args.out),
}
with open(os.path.join(args.out, "report.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))
if __name__ == "__main__":
sys.exit(main())
|