File size: 2,613 Bytes
f8d2531 | 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 | #!/usr/bin/env python3
"""把 teacher 清洗后的干净数据转成四头训练格式，并分层切分。
输入: data/cleaned/archive_latest_full_clean.jsonl（build_cleaned_dataset.py 输出）
输出: data/processed_teacher/{train,dev,val,holdout827}.jsonl
字段: text / label_bin / label_kind / label_scope / label_sev
- normal: kind=-100, scope=1(message), sev=0
- ad: kind 0..6（none -> -100）; scope account=0/message=1; sev bucket 0..3
"""
import argparse
import json
from collections import Counter
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"]
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--input", default="data/cleaned/archive_latest_full_clean.jsonl")
ap.add_argument("--out", default="data/processed_teacher")
ap.add_argument("--seed", type=int, default=42)
args = ap.parse_args()
rows = []
for line in open(args.input, encoding="utf-8"):
r = json.loads(line)
is_ad = r.get("label") == "ad"
kind = r.get("ad_kind")
rows.append({
"text": r.get("text"),
"label_bin": 0 if is_ad else 1,
"label_kind": KINDS.index(kind) if (is_ad and kind in KINDS) else -100,
"label_scope": (0 if (is_ad and r.get("scope") == "account") else 1),
"label_sev": (int(r.get("severity_bucket", 0)) if is_ad else 0),
})
labels = np.array([r["label_bin"] for r in rows])
idx = np.arange(len(rows))
tr, rest = train_test_split(idx, test_size=0.10, stratify=labels, random_state=args.seed)
dev, rest2 = train_test_split(rest, test_size=0.50, stratify=labels[rest], random_state=args.seed)
val, hold = train_test_split(rest2, test_size=0.50, stratify=labels[rest2], random_state=args.seed)
out = Path(args.out)
out.mkdir(parents=True, exist_ok=True)
for name, ids in [("train", tr), ("dev", dev), ("val", val), ("holdout827", hold)]:
data = [rows[i] for i in ids]
with open(out / f"{name}.jsonl", "w", encoding="utf-8") as f:
for x in data:
f.write(json.dumps(x, ensure_ascii=False) + "\n")
print(name, len(data),
"bin", dict(Counter(x["label_bin"] for x in data)),
"kind", dict(Counter(x["label_kind"] for x in data)),
"scope", dict(Counter(x["label_scope"] for x in data)),
"sev", dict(Counter(x["label_sev"] for x in data)))
print("out", out.resolve())
if __name__ == "__main__":
main()
|