#!/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())