process.py

import json
from sklearn.model_selection import train_test_split

# 假设你的数据在 output.json 文件中
with open('output.json', 'r') as file:
    data = json.load(file)

# 确保数据样本是一个列表
if not isinstance(data, list):
    raise ValueError("数据格式错误，output.json 应该包含一个列表！")

# 使用 80% 训练集，10% 验证集，10% 测试集
train_data, temp_data = train_test_split(data, test_size=0.2, random_state=42)  # 80% 训练集
val_data, test_data = train_test_split(temp_data, test_size=0.5, random_state=42)  # 10% 验证集，10% 测试集

# 输出数据集大小
print(f"总数据样本: {len(data)}")
print(f"训练集大小: {len(train_data)}")
print(f"验证集大小: {len(val_data)}")
print(f"测试集大小: {len(test_data)}")

# 保存到文件
with open('train_data.json', 'w') as file:
    json.dump(train_data, file, indent=4)

with open('val_data.json', 'w') as file:
    json.dump(val_data, file, indent=4)

with open('test_data.json', 'w') as file:
    json.dump(test_data, file, indent=4)