code/lora/02_split_cleaned_medical_dataset.py

"""将清洗后的完整数据去重，并按 8:2 切分训练集和验证集。"""

import json
import random
from pathlib import Path


ROOT = Path(__file__).resolve().parent
SOURCE_FILE = ROOT / "medical_o1_sft_mix_Chinese_cleaned.json"
OUTPUT_DIR = ROOT / "medical_o1_sft_mix_Chinese_cleaned_dedup_split_80_20"
RANDOM_SEED = 42


def question_key(record):
    """统一问题首尾及连续空白，用于问题级去重。"""
    return " ".join(record["Question"].split())


def save_json(path, records):
    path.write_text(
        json.dumps(records, ensure_ascii=False, indent=4) + "\n",
        encoding="utf-8",
    )


def main():
    records = json.loads(SOURCE_FILE.read_text(encoding="utf-8"))
    required_fields = {"Question", "Complex_CoT", "Response"}
    if not isinstance(records, list) or any(
        not isinstance(record, dict) or not required_fields <= record.keys()
        for record in records
    ):
        raise ValueError(f"每条记录必须包含字段：{required_fields}")

    # 同一个规范化问题只保留第一次出现的完整记录。
    unique = {}
    for record in records:
        unique.setdefault(question_key(record), record)
    unique_records = list(unique.values())

    random.Random(RANDOM_SEED).shuffle(unique_records)
    train_size = len(unique_records) * 8 // 10
    train = unique_records[:train_size]
    validation = unique_records[train_size:]

    train_keys = {question_key(record) for record in train}
    validation_keys = {question_key(record) for record in validation}
    if len(train_keys) != len(train) or len(validation_keys) != len(validation):
        raise RuntimeError("训练集或验证集内部仍然存在重复问题。")
    if train_keys & validation_keys:
        raise RuntimeError("训练集与验证集之间仍然存在重复问题。")

    if OUTPUT_DIR.exists() and any(OUTPUT_DIR.iterdir()):
        raise FileExistsError(f"为避免覆盖已有结果，请先确认目录：{OUTPUT_DIR}")
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    save_json(OUTPUT_DIR / "train.json", train)
    save_json(OUTPUT_DIR / "validation.json", validation)

    print(f"清洗后原始记录：{len(records)}")
    print(f"问题去重后记录：{len(unique_records)}")
    print(f"删除重复记录：{len(records) - len(unique_records)}")
    print(f"训练集：{len(train)}")
    print(f"验证集：{len(validation)}")
    print(f"训练/验证问题交集：{len(train_keys & validation_keys)}")
    print(f"随机种子：{RANDOM_SEED}")
    print(f"保存目录：{OUTPUT_DIR}")


if __name__ == "__main__":
    main()