| |
| """通用 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 |
|
|
| 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) |
| |
| 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}") |
|
|
| |
| 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()) |
|
|
