Views
No views yet
math-classifier to retrieve math-related content from fineweb-edu, dclm, ... to upsample math-related content1import json
2import os
3import time
4from concurrent.futures import ProcessPoolExecutor, wait, ALL_COMPLETED
5from time import sleep
6
7import fasttext
8import numpy as np
9import pandas as pd
10import pyarrow.parquet as pq
11from tqdm import tqdm
12
13
14def print_error(value):
15 print("error: ", value)
16
17
18def data_process(index, file, saved_dir):
19
20 try:
21 model_path = "math_score.bin"
22 model = fasttext.load_model(model_path)
23
24 # saved_dir: fineweb-edu/data/CC...-math/
25 filename = file.split('/')[-1].replace('.parquet', '.jsonl')
26 path90 = os.path.join(saved_dir, "09_10", filename)
27 if os.path.exists(path90):
28 print("exist", path90, flush=True)
29 return
30
31 sleep(index * 3)
32 os.makedirs(saved_dir, exist_ok=True)
33
34 label_list = []
35 s67_list = []
36 s78_list = []
37 s89_list = []
38 s90_list = []
39
40 st = time.time()
41 print("reading parquet", file, flush=True)
42 df = pd.read_parquet(file)
43 ed = time.time()
44 print("read parquet time: ", ed - st, flush=True)
45 for _, row_orginal in tqdm(
46 df.iterrows(),
47 total=len(df),
48 position=index,
49 desc=filename,
50 ):
51 row = row_orginal.to_dict()
52 text = row['text'].replace('\n', ' ')
53
54 pred = model.predict(text)
55 label, score = pred[0][0], pred[1][0]
56 label_list.append(pred)
57 if label == '__label__positive':
58 if 0.6 <= score < 0.7:
59 s67_list.append(row)
60 if 0.7 <= score < 0.8:
61 s78_list.append(row)
62 elif 0.8 <= score < 0.9:
63 s89_list.append(row)
64 elif 0.9 <= score <= 1.0:
65 s90_list.append(row)
66 else:
67 continue
68 except Exception as e:
69 print_error(e)
70 return None
71
72 os.makedirs(os.path.join(saved_dir, "labeled"), exist_ok=True)
73
74 print("writing to file", flush=True)
75
76 with open(
77 os.path.join(saved_dir, "labeled",
78 filename.replace('.jsonl', '.txt')), 'w') as f:
79 f.write("\n".join(str(pred) for pred in label_list))
80
81 for dir_name in [ "07_08", "08_09", "09_10"]:
82 os.makedirs(os.path.join(saved_dir, dir_name), exist_ok=True)
83
84 with open(os.path.join(saved_dir, "06_07", filename), 'w') as f:
85 f.write("\n".join(json.dumps(line_now) for line_now in s67_list))
86
87 with open(os.path.join(saved_dir, "07_08", filename), 'w') as f:
88 f.write("\n".join(json.dumps(line_now) for line_now in s78_list))
89
90 with open(os.path.join(saved_dir, "08_09", filename), 'w') as f:
91 f.write("\n".join(json.dumps(line_now) for line_now in s89_list))
92
93 with open(os.path.join(saved_dir, "09_10", filename), 'w') as f:
94 f.write("\n".join(json.dumps(line_now) for line_now in s90_list))
95
96 return None
97
98
99if __name__ == '__main__':
100
101 num_process = 5
102 start_time = time.time()
103 file_paths = []
104 base = "fineweb-edu/data/"
105
106 coun=0
107 for file_name in [
108 'CC-MAIN-2017-04','CC-MAIN-2017-09','CC-MAIN-2017-13',
109 'CC-MAIN-2017-17','CC-MAIN-2017-22','CC-MAIN-2017-26',
110 'CC-MAIN-2017-30','CC-MAIN-2017-34','CC-MAIN-2017-39',
111 'CC-MAIN-2017-43','CC-MAIN-2017-47','CC-MAIN-2017-51',
112
113 "CC-MAIN-2018-05","CC-MAIN-2018-09","CC-MAIN-2018-13",
114 "CC-MAIN-2018-17","CC-MAIN-2018-22","CC-MAIN-2018-26",
115 "CC-MAIN-2018-30","CC-MAIN-2018-34","CC-MAIN-2018-39",
116 "CC-MAIN-2018-43","CC-MAIN-2018-47","CC-MAIN-2018-51",
117
118 "CC-MAIN-2019-04","CC-MAIN-2019-09","CC-MAIN-2019-13",
119 "CC-MAIN-2019-18","CC-MAIN-2019-22","CC-MAIN-2019-26",
120 "CC-MAIN-2019-30","CC-MAIN-2019-35","CC-MAIN-2019-39",
121 "CC-MAIN-2019-43","CC-MAIN-2019-47","CC-MAIN-2019-51",
122
123 ]:
124
125 print("Walking:", file_name)
126 original_file_path = base + file_name
127 math_dir = original_file_path + "-math"
128 print(math_dir)
129
130 for root, dirs, files in os.walk(original_file_path):
131 for file in files:
132 if file.endswith(".parquet"): # 只处理Parquet文件
133 file_path = os.path.abspath(os.path.join(root, file))
134 coun+=1
135 saved_dir = math_dir + "/" + file_path.split("/")[-1][:-8]
136 print(saved_dir)
137 file_paths.append((file_path, saved_dir))
138 print(coun)
139 print(len(lines))
140
141 print("total file paths", len(file_paths))
142 num_process = min(num_process, len(file_paths))
143 print("num_process", num_process)
144
145 futures = []
146 with ProcessPoolExecutor(num_process) as executor:
147 for index, (file_path, saved_dir) in enumerate(file_paths):
148 futures.append(
149 executor.submit(data_process, index % num_process, file_path,
150 saved_dir))
151 done, not_done = wait(futures, return_when=ALL_COMPLETED)
152
153 end_time = time.time()
154
155 # 计算并打印所用时间
156 elapsed_time = end_time - start_time
157 print(f"Time taken: {elapsed_time} seconds")
158 print("=" * 100)
159@article{hu2024yulan,
title={YuLan-Mini: An Open Data-efficient Language Model},
author={Hu, Yiwen and Song, Huatong and Deng, Jia and Wang, Jiapeng and Chen, Jie and Zhou, Kun and Zhu, Yutao and Jiang, Jinhao and Dong, Zican and Zhao, Wayne Xin and others},
journal={arXiv preprint arXiv:2412.17743},
year={2024}
}