File size: 4,494 Bytes
e2492ec
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
#!/usr/bin/env python3
"""构建广告细类(7 类)数据集,复用 v1.1 的 train/val split,避免泄漏。

类别:crypto / porn / porn_bait / gambling / scam / promo / spam_flood
映射:contact → porn(样本均为色情招揽/迷药);spam / impersonate / unspecified 丢弃(样本极少且语义杂)
测试集:
  gold_val     = processed_v2/val ∩ gold.kind(广告行)
  archive_test = processed_v2/val ∩ archive.strong(ad 行)
训练池:
  processed_v2/train ∩ (gold.kind 广告行 ∪ archive.strong ad 行),85/15 分层切分 train/val
"""
import argparse
import collections
import json
from pathlib import Path

import numpy as np
from sklearn.model_selection import train_test_split

KINDS = ["crypto", "gambling", "porn", "porn_bait", "promo", "scam", "spam_flood"]
MAP = {"contact": "porn"}
DROP = {"spam", "impersonate", "unspecified", "normal", "none", None, ""}


def load(path):
    return [json.loads(l) for l in open(path, encoding="utf-8") if l.strip()]


def norm(t):
    return t.strip()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--gold", default="data/raw/dataset/converted/gold.kind.jsonl")
    ap.add_argument("--archive", default="data/raw/archive/converted/archive.strong.binary.jsonl")
    ap.add_argument("--split_train", default="data/processed_v2/train.jsonl")
    ap.add_argument("--split_val", default="data/processed_v2/val.jsonl")
    ap.add_argument("--out", default="data/processed_kind")
    ap.add_argument("--val_ratio", type=float, default=0.15)
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

    def clean(lab):
        if lab in MAP:
            lab = MAP[lab]
        return lab if lab in KINDS else None

    gold_kind = {}
    for r in load(args.gold):
        lab = clean(r["label"])
        if lab:
            gold_kind.setdefault(norm(r["text"]), lab)
    arch_kind = {}
    for r in load(args.archive):
        if r.get("label") != "ad":
            continue
        lab = clean(r.get("ad_kind"))
        if lab:
            arch_kind.setdefault(norm(r["text"]), lab)

    tr_texts = {norm(r["text"]) for r in load(args.split_train)}
    va_texts = {norm(r["text"]) for r in load(args.split_val)}

    pool = ([{"text": t, "label": lab} for t, lab in gold_kind.items() if t in tr_texts] +
            [{"text": t, "label": lab} for t, lab in arch_kind.items()
             if t in tr_texts and t not in gold_kind])
    gold_val = [{"text": t, "label": lab} for t, lab in gold_kind.items() if t in va_texts]
    archive_test = [{"text": t, "label": lab} for t, lab in arch_kind.items()
                    if t in va_texts and t not in gold_kind]

    assert not ({norm(r["text"]) for r in pool} & {norm(r["text"]) for r in gold_val})
    assert not ({norm(r["text"]) for r in pool} & {norm(r["text"]) for r in archive_test})

    labels = np.array([r["label"] for r in pool])
    idx = np.arange(len(pool))
    tr_idx, va_idx = train_test_split(idx, test_size=args.val_ratio,
                                      stratify=labels, random_state=args.seed)
    train = [pool[i] for i in tr_idx]
    val = [pool[i] for i in va_idx]

    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    for name, rows in [("train", train), ("val", val),
                       ("gold_val", gold_val), ("archive_test", archive_test)]:
        with open(out / f"{name}.jsonl", "w", encoding="utf-8") as f:
            for r in rows:
                f.write(json.dumps(r, ensure_ascii=False) + "\n")
        print(f"{name:<14} n={len(rows):5d} {dict(collections.Counter(r['label'] for r in rows))}")
    label_map = {k: i for i, k in enumerate(sorted(KINDS))}
    (out / "label_map.json").write_text(
        json.dumps(label_map, ensure_ascii=False, indent=2), encoding="utf-8")
    (out / "report.json").write_text(json.dumps({
        "kinds": KINDS, "label_map": label_map,
        "train": len(train), "val": len(val),
        "gold_val": len(gold_val), "archive_test": len(archive_test),
        "train_dist": dict(collections.Counter(r["label"] for r in train)),
        "val_dist": dict(collections.Counter(r["label"] for r in val)),
        "gold_val_dist": dict(collections.Counter(r["label"] for r in gold_val)),
        "archive_test_dist": dict(collections.Counter(r["label"] for r in archive_test)),
    }, ensure_ascii=False, indent=2), encoding="utf-8")
    print("输出:", out.resolve())


if __name__ == "__main__":
    main()