laya-ad / code /scripts /build_archive_dataset.py
x6nux's picture
laya-ad v1.1:六精度权重 + 模型卡 + 项目代码快照
71dabc7 verified
Raw History Blame Contribute Delete
22.2 kB
#!/usr/bin/env python3
"""从 archive.tar.gz 解包出的群聊存档构建清洗后的带标签数据集。
输入目录结构(bot 存档导出):
<archive>/<chat_id>/messages-YYYY-MM-DD.jsonl 消息
<archive>/<chat_id>/judgments-YYYY-MM-DD.jsonl 判定(含 verdict/ad_kind)
<archive>/<chat_id>/images/YYYY-MM-DD/<msgid>_{tg|hist_*}.{jpg,webp}
三份产物:
messages.jsonl 全量消息 + 裁决后的标签(ad/clean)
image_clf.jsonl 通过校验的图片 + 图文对齐后的标签
desc_clf.jsonl 图片/贴纸的 VLM 描述文本 + 标签(喂文本模型用)
清洗口径见 report.md。原始存档只读,产物全部写到 --out。
"""
from __future__ import annotations
import argparse
import collections
import datetime as dt
import hashlib
import json
import os
import re
import sys
import unicodedata
from typing import Any, Iterable
# 判定仲裁优先级:数字越小越可信
TIER_NAMES = {
"manual": 0, "manual-ban": 0,
"systemone+llm": 1,
"systemone": 2,
"llm": 3, "llm+llm": 3,
"phash": 5, "hash": 5, "hash+llm": 5,
}
DROP_VERDICTS = {"skipped", "error"}
# 反检测用的不可见字符:零宽空格/连接符、词连接符、BOM、软连字符
INVISIBLE = dict.fromkeys(map(ord, "\u200b\u200c\u200d\u2060\u2061\u2062\u2063\u2064\ufeff\u00ad"), None)
MARK_IMAGE = "[图片]"
MARK_STICKER = "[贴纸]"
MARK_QUOTE = "[引用]"
# 群机器人指令/签到类短文本:样本量大但不含判定信息
COMMAND_WORDS = {"签到", "专享", "订阅", "积分", "积分排名", "积分兑换", "1", "?", ".", "。", "+1"}
# 联系方式/落地页特征:短文本里带这些字样的 ad 才算内容级广告
CONTACT_RE = re.compile(
r"(https?://|t\.me/|www\.|\+\d{5,}|\d{3,}|@[A-Za-z0-9_]{4,}|微信|威信|vx|VX|QQ|qq|telegram|电报|飞机|私聊|加我)",
re.IGNORECASE,
)
def label_trust(label: str | None, decider: str | None, text: str) -> str:
"""标注可信度分层。
存档里的判定不全是内容级:manual 与"极短文本被判 ad"多数是发送者级处置
(用户因历史行为被封,其所有消息都被打成 ad),拿来训练会教模型把"嗯""哈喽"学成广告。
"""
if label not in ("ad", "clean"):
return "none"
if decider == "manual":
return "suspect_sender_level"
if label == "ad" and len(text) < 6 and not CONTACT_RE.search(text):
return "suspect_sender_level"
if decider and (decider.startswith("rule") or decider in ("phash", "hash", "hash+llm")):
return "medium_heuristic"
return "high"
IMG_RE = re.compile(r"^(\d+)_(tg|hist|\d+)[^/]*\.(jpg|jpeg|png|webp)$", re.IGNORECASE)
def iter_jsonl(pattern_dir: str, prefix: str) -> Iterable[dict]:
for root, _dirs, files in os.walk(pattern_dir):
for name in files:
if not (name.startswith(prefix) and name.endswith(".jsonl")):
continue
path = os.path.join(root, name)
with open(path, encoding="utf-8") as fh:
for lineno, line in enumerate(fh, 1):
line = line.strip()
if not line:
continue
try:
yield json.loads(line)
except json.JSONDecodeError:
print(f"[warn] 坏行 {path}:{lineno}", file=sys.stderr)
def tier_of(decider: str | None) -> int:
if decider is None:
return 9
if decider in TIER_NAMES:
return TIER_NAMES[decider]
if decider.startswith("rule"):
return 4
return 6
def load_messages(base: str) -> dict[tuple[str, int], dict]:
"""按 (chat_id, message_id) 去重,保留 edit_date 最新的一条。"""
msgs: dict[tuple[str, int], dict] = {}
total = dup = 0
for d in iter_jsonl(base, "messages-"):
total += 1
key = (str(d.get("chat_id")), int(d.get("message_id", -1)))
prev = msgs.get(key)
if prev is None:
msgs[key] = d
continue
dup += 1
if int(d.get("edit_date") or 0) >= int(prev.get("edit_date") or 0):
msgs[key] = d
print(f"[load] messages 原始 {total} 条,去重后 {len(msgs)}(重复 {dup})")
return msgs
def load_judgments(base: str) -> dict[tuple[str, int], list[dict]]:
groups: dict[tuple[str, int], list[dict]] = collections.defaultdict(list)
total = 0
for d in iter_jsonl(base, "judgments-"):
total += 1
groups[(str(d.get("chat_id")), int(d.get("message_id", -1)))].append(d)
print(f"[load] judgments {total} 条,覆盖 {len(groups)} 个消息")
return dict(groups)
def resolve_label(js: list[dict] | None) -> dict[str, Any]:
"""多判定仲裁:先按 decider 可信度分层,层内按置信度、再按时间取最新。
同层出现 ad/clean 对立即记为冲突(保留高置信度那条,但标记出来供人工复核)。
"""
empty = {"label": None, "ad_kind": None, "confidence": None, "decider": None,
"label_tier": None, "label_conflict": False, "label_votes": 0}
if not js:
return empty
valid = [j for j in js if j.get("verdict") not in DROP_VERDICTS and j.get("verdict")]
if not valid:
return empty
best_tier = min(tier_of(j.get("decider")) for j in valid)
top = [j for j in valid if tier_of(j.get("decider")) == best_tier]
votes = {j["verdict"] for j in top}
top.sort(key=lambda j: (float(j.get("confidence") or 0), int(j.get("created_at") or 0)), reverse=True)
win = top[0]
return {
"label": win["verdict"],
"ad_kind": win.get("ad_kind"),
"confidence": win.get("confidence"),
"decider": win.get("decider"),
"label_tier": best_tier,
"label_conflict": len(votes) > 1,
"label_votes": len(valid),
}
def clean_text(raw: str) -> dict[str, Any]:
"""剥离不可见字符、拆出引用段、判定内容类型。
标记([图片]等)必须在 NFKC 归一化前识别,否则全角括号会被改写成半角。
"""
raw = raw or ""
had_invisible = any(ch in raw for ch in "\u200b\u200c\u200d\u2060\u2061\u2062\u2063\u2064\ufeff\u00ad")
content_type = "text"
if raw.startswith(MARK_IMAGE):
content_type = "image_desc"
elif raw.startswith(MARK_STICKER):
content_type = "sticker_desc"
body, quote = raw, ""
if MARK_QUOTE in raw:
head, _, tail = raw.partition(MARK_QUOTE)
body, quote = (head if head.strip() else raw), tail.strip()
def norm(s: str) -> str:
s = unicodedata.normalize("NFKC", s)
s = s.translate(INVISIBLE)
return re.sub(r"[ \t\u3000]+", " ", re.sub(r"\n{3,}", "\n\n", s)).strip()
return {
"text": norm(body),
"quote": norm(quote),
"content_type": content_type,
"had_invisible": had_invisible,
"had_quote": bool(quote),
}
def dhash(path: str, size: int = 8) -> int | None:
"""8x8 差值哈希,用于跨消息的近似重复图去重。"""
from PIL import Image
try:
with Image.open(path) as im:
im = im.convert("L").resize((size + 1, size), Image.Resampling.LANCZOS)
px = list(im.getdata())
except Exception:
return None
bits = 0
for row in range(size):
for col in range(size):
left = px[row * (size + 1) + col]
right = px[row * (size + 1) + col + 1]
bits = (bits << 1) | (1 if left > right else 0)
return bits
def hamming(a: int, b: int) -> int:
return (a ^ b).bit_count()
def collect_images(base: str) -> tuple[dict[tuple[str, int], list[dict]], dict[str, int]]:
"""枚举图片文件,校验可解码,按消息归组。"""
from PIL import Image
stats = collections.Counter()
per_msg: dict[tuple[str, int], list[dict]] = collections.defaultdict(list)
for root, _dirs, files in os.walk(base):
if os.path.basename(os.path.dirname(root)) != "images":
continue
chat_id = os.path.basename(os.path.dirname(os.path.dirname(root)))
for name in files:
stats["files"] += 1
m = IMG_RE.match(name)
if not m:
stats["unparsable_name"] += 1
continue
path = os.path.join(root, name)
partial = ".partial." in name.lower()
if partial:
stats["partial"] += 1
try:
with Image.open(path) as im:
im.load()
w, h = im.size
fmt = im.format
except Exception:
stats["corrupt"] += 1
continue
if w < 64 or h < 64:
stats["too_small"] += 1
continue
with open(path, "rb") as fh:
digest = hashlib.sha256(fh.read()).hexdigest()
stats["valid"] += 1
per_msg[(chat_id, int(m.group(1)))].append({
"path": path, "sha256": digest, "w": w, "h": h, "format": fmt,
"partial": partial, "variant": m.group(2), "bytes": os.path.getsize(path),
})
print(f"[img] {dict(stats)}")
return per_msg, dict(stats)
def pick_one(cands: list[dict]) -> dict:
"""同一消息多份文件时择优:非 partial > 面积大 > 体积大。"""
return sorted(cands, key=lambda c: (c["partial"], -c["w"] * c["h"], -c["bytes"]))[0]
def build_splits(records: list[dict], val_ratio: float = 0.1, group_keys=None) -> None:
"""时间切分:按 date 排序,尾部 val_ratio 进 val。
同组记录强制同侧(同一文本 / 同一相册),避免近似样本跨集泄漏。
三份数据集各自按自身时间分布切分,否则图片这类集中在新日期的子集会全掉进 val。
"""
dates = sorted(r["date"] for r in records)
if not dates:
return
cutoff = dates[int(len(dates) * (1 - val_ratio))]
for r in records:
r["split"] = "val" if r["date"] >= cutoff else "train"
if group_keys is None:
return
anchor: dict[str, str] = {}
for r in sorted(records, key=lambda x: x["date"]):
keys = group_keys(r)
first = next((anchor[k] for k in keys if k in anchor), None)
if first is None:
for k in keys:
anchor[k] = r["split"]
else:
r["split"] = first
for k in keys:
anchor[k] = first
def msg_group_keys(r: dict) -> list[str]:
keys = [f"t:{r['text_hash']}"] if r.get("text") else []
if r.get("media_group_id") is not None:
keys.append(f"g:{r['chat_id']}:{r['media_group_id']}")
return keys
def desc_group_keys(r: dict) -> list[str]:
return [f"t:{hashlib.sha1(r['text'].encode()).hexdigest()[:16]}"] if r.get("text") else []
def write_jsonl(path: str, rows: Iterable[dict]) -> int:
os.makedirs(os.path.dirname(path), exist_ok=True)
n = 0
with open(path, "w", encoding="utf-8") as fh:
for r in rows:
fh.write(json.dumps(r, ensure_ascii=False) + "\n")
n += 1
return n
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--archive", default="data/raw/archive/archive")
ap.add_argument("--out", default="data/processed/archive")
ap.add_argument("--val-ratio", type=float, default=0.1)
ap.add_argument("--phash-radius", type=int, default=4, help="跨消息近似重复图的汉明距离阈值")
args = ap.parse_args()
msgs = load_messages(args.archive)
judgments = load_judgments(args.archive)
# ---------- 消息:清洗 + 打标 ----------
rows: list[dict] = []
text_seen: dict[str, tuple[str, int]] = {}
conflicts: list[dict] = []
for key, m in msgs.items():
chat_id, mid = key
meta = clean_text(m.get("text") or "")
lab = resolve_label(judgments.get(key))
text_hash = hashlib.sha1(meta["text"].encode("utf-8")).hexdigest()[:16]
# 首次出现决定该文本的 split 锚点;重复文本标记出来便于降采样
is_dup = text_hash in text_seen and bool(meta["text"])
if meta["text"] and not is_dup:
text_seen[text_hash] = ("train", 0)
if lab["label_conflict"]:
conflicts.append({"chat_id": chat_id, "message_id": mid, "text": meta["text"][:120]})
rows.append({
"chat_id": chat_id,
"message_id": mid,
"date": int(m.get("date") or 0),
"iso_time": dt.datetime.fromtimestamp(int(m.get("date") or 0)).isoformat(timespec="seconds"),
"user_id": m.get("user_id"),
"kind": m.get("kind"),
"media_group_id": m.get("media_group_id"),
"punished": bool(m.get("punished")),
"text_raw": m.get("text") or "",
"text": meta["text"],
"quote": meta["quote"],
"content_type": meta["content_type"],
"had_invisible": meta["had_invisible"],
"text_len": len(meta["text"]),
"low_info": len(meta["text"]) < 4 or meta["text"] in COMMAND_WORDS,
"text_hash": text_hash,
"dup_text": is_dup,
"label_trust": label_trust(lab["label"], lab["decider"], meta["text"]),
**lab,
})
labeled = [r for r in rows if r["label"] in ("ad", "clean")]
build_splits(labeled, args.val_ratio, msg_group_keys)
# ---------- 图片:校验 + 去重 + 对齐标签 ----------
per_msg, img_stats = collect_images(args.archive)
by_key = {(r["chat_id"], r["message_id"]): r for r in rows}
image_rows: list[dict] = []
seen_hash: dict[str, dict] = {}
seen_phash: list[tuple[int, dict]] = []
img_dedup = collections.Counter()
for key, cands in sorted(per_msg.items(), key=lambda kv: (kv[0][0], kv[0][1])):
rec = by_key.get(key)
if rec is None:
img_dedup["no_message"] += 1
continue
if rec["label"] not in ("ad", "clean"):
img_dedup["no_label"] += 1
continue
best = pick_one(cands)
if best["sha256"] in seen_hash:
img_dedup["exact_dup"] += 1
continue
ph = dhash(best["path"])
if ph is not None and any(hamming(ph, p) <= args.phash_radius for p, _ in seen_phash):
img_dedup["near_dup"] += 1
continue
out = {
"path": os.path.relpath(best["path"], args.archive).replace("\\", "/"),
"sha256": best["sha256"],
"dhash": f"{ph:016x}" if ph is not None else None,
"width": best["w"], "height": best["h"], "bytes": best["bytes"],
"chat_id": key[0], "message_id": key[1],
"date": rec["date"], "split": rec["split"],
"label": rec["label"], "ad_kind": rec["ad_kind"],
"confidence": rec["confidence"], "decider": rec["decider"],
"content_type": rec["content_type"], "text": rec["text"],
"variants": len(cands),
}
image_rows.append(out)
seen_hash[best["sha256"]] = out
if ph is not None:
seen_phash.append((ph, out))
# 图片按自身时间分布切分:它们集中在近期,跟着消息切分会让 val 反超 train
build_splits(image_rows, args.val_ratio)
# ---------- 描述文本(图片/贴纸的 VLM 描述):喂文本模型的桥接训练集 ----------
img_keys = {(r["chat_id"], r["message_id"]) for r in image_rows}
desc_rows = []
for r in rows:
if r["content_type"] not in ("image_desc", "sticker_desc"):
continue
if r["label"] not in ("ad", "clean"):
continue
desc_rows.append({
"chat_id": r["chat_id"], "message_id": r["message_id"], "date": r["date"],
"split": r["split"], "text": r["text"], "quote": r["quote"],
"content_type": r["content_type"], "label": r["label"],
"ad_kind": r["ad_kind"], "confidence": r["confidence"],
"decider": r["decider"], "has_image": (r["chat_id"], r["message_id"]) in img_keys,
})
build_splits(desc_rows, args.val_ratio, desc_group_keys)
# ---------- 落盘 + 报告 ----------
out = args.out
n_msg = write_jsonl(os.path.join(out, "messages.jsonl"), rows)
n_img = write_jsonl(os.path.join(out, "image_clf.jsonl"), image_rows)
n_desc = write_jsonl(os.path.join(out, "desc_clf.jsonl"), desc_rows)
n_lab = write_jsonl(os.path.join(out, "messages_labeled.jsonl"),
(r for r in rows if r["label"] in ("ad", "clean")))
# 去重视图:同一文本只保留首条("签到"这类群指令单条就能刷到 7420 次)
seen_h: set[str] = set()
dedup_rows = []
for r in labeled:
if r["text"] and r["text_hash"] in seen_h:
continue
seen_h.add(r["text_hash"])
dedup_rows.append(r)
n_dedup = write_jsonl(os.path.join(out, "messages_labeled_dedup.jsonl"), dedup_rows)
# 可信子集:剔除发送者级处置造成的伪标签,供训练直接用
trusted = [r for r in rows
if r["label"] in ("ad", "clean") and r["label_trust"] in ("high", "medium_heuristic")]
n_trusted = write_jsonl(os.path.join(out, "messages_trusted.jsonl"), trusted)
report = {
"generated_at": dt.datetime.now().isoformat(timespec="seconds"),
"archive": os.path.abspath(args.archive),
"messages": {
"total": n_msg,
"labeled": n_lab,
"label_dist": dict(collections.Counter(r["label"] for r in rows)),
"ad_kind_dist": dict(collections.Counter(r["ad_kind"] for r in rows if r["label"] == "ad").most_common()),
"content_type_dist": dict(collections.Counter(r["content_type"] for r in rows)),
"kind_dist": dict(collections.Counter(r["kind"] for r in rows)),
"labeled_content_type_dist": dict(collections.Counter(r["content_type"] for r in labeled)),
"low_info": sum(1 for r in labeled if r["low_info"]),
"label_trust": dict(collections.Counter(r["label_trust"] for r in labeled)),
"trusted": {
"rows": n_trusted,
"label_dist": dict(collections.Counter(r["label"] for r in trusted)),
"split": dict(collections.Counter(r["split"] for r in trusted)),
},
"labeled_dedup": {
"rows": n_dedup,
"label_dist": dict(collections.Counter(r["label"] for r in dedup_rows)),
"split": dict(collections.Counter(r["split"] for r in dedup_rows)),
},
"conflicts": len(conflicts),
"had_invisible": sum(1 for r in rows if r["had_invisible"]),
"dup_text": sum(1 for r in rows if r["dup_text"]),
"split": dict(collections.Counter(r["split"] for r in labeled)),
"date_range": [dt.datetime.fromtimestamp(labeled[0]["date"]).isoformat(timespec="seconds"),
dt.datetime.fromtimestamp(labeled[-1]["date"]).isoformat(timespec="seconds")] if labeled else None,
},
"images": {
"files_scanned": img_stats.get("files", 0),
"invalid_or_dropped": {k: img_stats.get(k, 0) for k in ("corrupt", "partial", "too_small", "unparsable_name")},
"dropped_in_join": dict(img_dedup),
"kept": n_img,
"label_dist": dict(collections.Counter(r["label"] for r in image_rows)),
"content_type_dist": dict(collections.Counter(r["content_type"] for r in image_rows)),
"ad_kind_dist": dict(collections.Counter(r["ad_kind"] for r in image_rows if r["label"] == "ad").most_common()),
"split": dict(collections.Counter(r["split"] for r in image_rows)),
},
"desc": {
"kept": n_desc,
"label_dist": dict(collections.Counter(r["label"] for r in desc_rows)),
"with_image_file": sum(1 for r in desc_rows if r["has_image"]),
"split": dict(collections.Counter(r["split"] for r in desc_rows)),
},
"conflict_samples": conflicts[:20],
}
with open(os.path.join(out, "report.json"), "w", encoding="utf-8") as fh:
json.dump(report, fh, ensure_ascii=False, indent=2)
m, im, de = report["messages"], report["images"], report["desc"]
md = [
"# archive 存档清洗报告", "",
f"- 生成时间:{report['generated_at']}", f"- 源目录:`{report['archive']}`", "",
"## 消息", "",
f"- 总消息 {m['total']},其中带标签 {m['labeled']}",
f"- 标签分布:`{m['label_dist']}`",
f"- 已标内容类型:`{m['labeled_content_type_dist']}`",
f"- 判定冲突(同层 ad/clean 对立):{m['conflicts']}",
f"- 含不可见字符(反检测):{m['had_invisible']}",
f"- 完全重复文本:{m['dup_text']};低信息短文本:{m['low_info']}",
f"- 切分:`{m['split']}`",
f"- 标注可信度:`{m['label_trust']}`",
f"- 可信训练集 `messages_trusted.jsonl`:{m['trusted']['rows']} 条,"
f"标签 `{m['trusted']['label_dist']}`,切分 `{m['trusted']['split']}`",
f"- 去重视图 `messages_labeled_dedup.jsonl`:{m['labeled_dedup']['rows']} 条,"
f"标签 `{m['labeled_dedup']['label_dist']}`", "",
"## 图片", "",
f"- 扫描 {im['files_scanned']} 个文件,丢弃 `{im['invalid_or_dropped']}`,join 丢弃 `{im['dropped_in_join']}`",
f"- 保留 {im['kept']},标签分布 `{im['label_dist']}`,内容类型 `{im['content_type_dist']}`",
f"- 切分:`{im['split']}`", "",
"## 图片描述文本", "",
f"- 保留 {de['kept']},其中 {de['with_image_file']} 条有对应图片文件",
f"- 标签分布:`{de['label_dist']}`",
]
with open(os.path.join(out, "report.md"), "w", encoding="utf-8") as fh:
fh.write("\n".join(md) + "\n")
print("\n".join(md))
print(f"\n[out] {os.path.abspath(out)}")
if __name__ == "__main__":
main()