62 lines
1.7 KiB
Python
62 lines
1.7 KiB
Python
import json
|
||
import random
|
||
|
||
def split_data(input_file, train_file, test_file, split_ratio=0.7, seed=42):
|
||
# 读取原始的jsonl文件
|
||
# with open(input_file, 'r', encoding='utf-8') as f:
|
||
# data = [json.loads(line.strip()) for line in f]
|
||
|
||
with open(input_file, 'rt', encoding='utf-8') as file:
|
||
data = json.load(file)
|
||
|
||
# 设置随机种子，确保每次划分结果一致
|
||
random.seed(seed)
|
||
# 打乱数据
|
||
random.shuffle(data)
|
||
|
||
# 计算训练集和测试集的划分点
|
||
split_point = int(len(data) * split_ratio)
|
||
|
||
# 划分数据
|
||
train_data = data[:split_point]
|
||
test_data = data[split_point:]
|
||
|
||
# # 保存训练集
|
||
# with open(train_file, 'w', encoding='utf-8') as f_train:
|
||
# for d in train_data:
|
||
# json.dump(d, f_train, ensure_ascii=False)
|
||
# f_train.write('\n')
|
||
|
||
# # 保存测试集
|
||
# with open(test_file, 'w', encoding='utf-8') as f_test:
|
||
# for d in test_data:
|
||
# json.dump(d, f_test, ensure_ascii=False)
|
||
# f_test.write('\n')
|
||
|
||
|
||
with open(train_file, 'w', encoding='utf-8') as file:
|
||
json.dump(train_data, file, ensure_ascii=False, indent=4)
|
||
|
||
with open(test_file, 'w', encoding='utf-8') as file:
|
||
json.dump(test_data, file, ensure_ascii=False, indent=4)
|
||
|
||
print(len(data),len(train_data),len(test_data),)
|
||
|
||
|
||
def main():
|
||
|
||
# 输入文件路径和输出文件路径
|
||
input_jsonl_path = 'data/ruozhiba_format_emo.jsonl'
|
||
train_jsonl_path = 'data/train_emo.jsonl'
|
||
test_jsonl_path = 'data/test_emo.jsonl'
|
||
|
||
# 划分数据集，按照7:3的比例
|
||
split_data(input_jsonl_path, train_jsonl_path, test_jsonl_path, split_ratio=0.8)
|
||
|
||
|
||
if __name__== "__main__" :
|
||
main()
|
||
|
||
|
||
|