File size: 2,448 Bytes
e6f9ec5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | import json
def read_jsonl(file_path):
"""读取JSONL文件，返回数据列表"""
data_list = []
with open(file_path, 'r', encoding='utf-8') as file:
lines = file.readlines()
for line in lines:
temp_json = json.loads(line)
data_list.append(temp_json)
return data_list
def write_jsonl(data_list, output_path):
"""将数据列表写入JSONL文件"""
with open(output_path, 'w', encoding='utf-8') as file:
for item in data_list:
file.write(json.dumps(item, ensure_ascii=False))
file.write('\n')
print(f"文件已生成至 {output_path}，共 {len(data_list)} 条数据。")
def main():
# 输入文件路径
input_dir = "./data"
base_filename = "gsm8k_math500_test_addbox_"
suffix_range = range(8) # 0到7
# 输出文件路径
output_dir = "./data"
output_base = "math500_"
# 存储所有math500数据
math500_data = []
# 遍历所有文件
for suffix in suffix_range:
input_file = f"{input_dir}/{base_filename}{suffix}.json"
print(f"正在读取: {input_file}")
try:
data_list = read_jsonl(input_file)
# 提取data_source为"math500"的数据
for item in data_list:
if item.get('data_source') == 'math500':
math500_data.append(item)
print(f" - 从文件中找到 {len([item for item in data_list if item.get('data_source') == 'math500'])} 条math500数据")
except FileNotFoundError:
print(f" - 警告: 文件 {input_file} 不存在，跳过")
except Exception as e:
print(f" - 错误: 读取文件 {input_file} 时出错: {e}")
# 将数据切片为5个文件
if math500_data:
num_files = 5
data_per_file = len(math500_data) // num_files
print(f"\n共提取 {len(math500_data)} 条数据，将切片为 {num_files} 个文件，每个文件约 {data_per_file} 条数据\n")
for i in range(num_files):
start_idx = i * data_per_file
end_idx = start_idx + data_per_file if i < num_files - 1 else len(math500_data)
chunk = math500_data[start_idx:end_idx]
output_path = f"{output_dir}/{output_base}{i}.json"
write_jsonl(chunk, output_path)
else:
print("警告: 没有找到任何math500数据")
if __name__ == "__main__":
main()
|