| |
| """把 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() |
|
|
