chaitiao_and_ASR/4.结构化json-Ai标注/train and test.py

import json
import random
from collections import defaultdict
import pandas as pd


def split_dataset_by_source_id(input_file, test_size=150, random_seed=42):
    """
    根据source_id随机划分数据集为训练集和测试集

    Args:
        input_file: 输入的JSON文件路径
        test_size: 测试集中source_id的数量
        random_seed: 随机种子，确保结果可重现

    Returns:
        train_data, test_data: 训练集和测试集数据
    """

    # 设置随机种子
    random.seed(random_seed)

    print(f"正在读取文件: {input_file}")

    try:
        # 读取JSON文件
        with open(input_file, 'r', encoding='utf-8') as f:
            all_data = json.load(f)

        print(f"✓ 成功读取文件，总记录数: {len(all_data)}")

        # 按source_id分组
        source_id_groups = defaultdict(list)
        for item in all_data:
            source_id_groups[item['source_id']].append(item)

        # 获取所有unique的source_id
        all_source_ids = list(source_id_groups.keys())
        total_source_ids = len(all_source_ids)

        print(f"✓ 发现 {total_source_ids} 个不同的source_id")

        # 检查测试集大小是否合理
        if test_size >= total_source_ids:
            print(f"✗ 错误：测试集大小 ({test_size}) 大于等于总source_id数量 ({total_source_ids})")
            print(f"建议将测试集大小设置为小于 {total_source_ids}")
            return None, None

        # 随机选择测试集的source_id
        test_source_ids = random.sample(all_source_ids, test_size)
        train_source_ids = [sid for sid in all_source_ids if sid not in test_source_ids]

        print(f"✓ 随机选择了 {len(test_source_ids)} 个source_id作为测试集")
        print(f"✓ 剩余 {len(train_source_ids)} 个source_id作为训练集")

        # 构建训练集和测试集
        train_data = []
        test_data = []

        for source_id in train_source_ids:
            train_data.extend(source_id_groups[source_id])

        for source_id in test_source_ids:
            test_data.extend(source_id_groups[source_id])

        print(f"\n=== 数据集划分结果 ===")
        print(f"训练集:")
        print(f"  - Source ID数量: {len(train_source_ids)}")
        print(f"  - 记录数量: {len(train_data)}")

        print(f"测试集:")
        print(f"  - Source ID数量: {len(test_source_ids)}")
        print(f"  - 记录数量: {len(test_data)}")

        # 统计标签分布
        def get_label_distribution(data, dataset_name):
            label_counts = defaultdict(int)
            for item in data:
                label_counts[item['label']] += 1

            print(f"\n{dataset_name}标签分布:")
            for label, count in sorted(label_counts.items()):
                percentage = (count / len(data) * 150) if len(data) > 0 else 0
                print(f"  标签 {label}: {count} 条 ({percentage:.2f}%)")

            return label_counts

        train_labels = get_label_distribution(train_data, "训练集")
        test_labels = get_label_distribution(test_data, "测试集")

        # 显示选中的source_id
        print(f"\n=== 测试集Source ID列表 ===")
        print(f"测试集source_id: {sorted(test_source_ids)}")

        print(f"\n=== 训练集Source ID列表 ===")
        print(f"训练集source_id: {sorted(train_source_ids)}")

        return train_data, test_data, train_source_ids, test_source_ids

    except FileNotFoundError:
        print(f"✗ 错误：找不到文件 {input_file}")
        return None, None, None, None
    except json.JSONDecodeError as e:
        print(f"✗ 错误：JSON文件格式错误 - {str(e)}")
        return None, None, None, None
    except Exception as e:
        print(f"✗ 错误：处理文件时出现异常 - {str(e)}")
        return None, None, None, None


def save_dataset(data, filename, description):
    """
    保存数据集到JSON文件

    Args:
        data: 要保存的数据
        filename: 输出文件名
        description: 数据集描述
    """
    try:
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        print(f"✓ {description}已保存到: {filename}")
        return True
    except Exception as e:
        print(f"✗ 保存{description}时出错: {str(e)}")
        return False


def create_summary_report(train_data, test_data, train_source_ids, test_source_ids):
    """
    创建数据集划分的详细报告
    """
    summary = {
        "split_info": {
            "total_source_ids": len(train_source_ids) + len(test_source_ids),
            "train_source_ids": len(train_source_ids),
            "test_source_ids": len(test_source_ids),
            "total_records": len(train_data) + len(test_data),
            "train_records": len(train_data),
            "test_records": len(test_data)
        },
        "train_source_id_list": sorted(train_source_ids),
        "test_source_id_list": sorted(test_source_ids),
        "label_distribution": {
            "train": {},
            "test": {}
        }
    }

    # 计算标签分布
    for dataset_name, data in [("train", train_data), ("test", test_data)]:
        label_counts = defaultdict(int)
        for item in data:
            label_counts[item['label']] += 1

        summary["label_distribution"][dataset_name] = dict(label_counts)

    # 保存报告
    with open('dataset_split_summary.json', 'w', encoding='utf-8') as f:
        json.dump(summary, f, ensure_ascii=False, indent=2)

    print(f"✓ 数据集划分报告已保存到: dataset_split_summary.json")


# 主程序执行
if __name__ == "__main__":
    print("=" * 60)
    print("数据集划分程序")
    print("=" * 60)

    # 输入文件名
    input_file = "segmentation_results_from_7_retried.json"

    # 执行数据集划分
    train_data, test_data, train_source_ids, test_source_ids = split_dataset_by_source_id(
        input_file=input_file,
        test_size=150,
        random_seed=42
    )

    if train_data is not None and test_data is not None:
        print(f"\n{'=' * 60}")
        print("开始保存数据集文件")
        print(f"{'=' * 60}")

        # 保存训练集
        train_success = save_dataset(train_data, "train_dataset.json", "训练集")

        # 保存测试集
        test_success = save_dataset(test_data, "test_dataset.json", "测试集")

        if train_success and test_success:
            # 创建详细报告
            create_summary_report(train_data, test_data, train_source_ids, test_source_ids)

            print(f"\n{'=' * 60}")
            print("数据集划分完成！")
            print(f"{'=' * 60}")
            print("生成的文件:")
            print("1. train_dataset.json - 训练集数据")
            print("2. test_dataset.json - 测试集数据")
            print("3. dataset_split_summary.json - 划分报告")

            # 验证数据完整性
            print(f"\n=== 数据完整性验证 ===")
            original_count = len(train_data) + len(test_data)
            print(f"原始数据总数: {original_count}")
            print(f"训练集 + 测试集: {len(train_data)} + {len(test_data)} = {len(train_data) + len(test_data)}")

            if len(train_data) + len(test_data) == original_count:
                print("✓ 数据完整性验证通过")
            else:
                print("✗ 数据完整性验证失败")

        else:
            print("✗ 保存文件时出现错误")
    else:
        print("✗ 数据集划分失败")