(input_folder,output_folder)
| 206 | |
| 207 | |
| 208 | def gen_code_github(input_folder,output_folder): ##预处理github 代码数据集 |
| 209 | os.makedirs(output_folder, exist_ok=True) |
| 210 | eos_token = "[EOS]" |
| 211 | for filename in tqdm(os.listdir(input_folder)): |
| 212 | if filename.endswith(".parquet"): |
| 213 | doc_ids=[] |
| 214 | origin_file = os.path.join(input_folder,filename) |
| 215 | output_file = os.path.join(output_folder,filename) |
| 216 | print('当前正在处理文件：{}'.format(origin_file)) |
| 217 | table = pq.read_table(origin_file) |
| 218 | tmp_df = table.to_pandas() |
| 219 | code_text = tmp_df['code'].values |
| 220 | for code in tqdm(code_text): |
| 221 | text = code+eos_token |
| 222 | doc_ids.append(text) |
| 223 | if doc_ids: |
| 224 | chunk_data = split_txt_cropus_to_chunk_data(doc_ids) |
| 225 | |
| 226 | tb = pa.Table.from_arrays([pa.array(chunk_data)], names=["text"]) |
| 227 | del chunk_data |
| 228 | pq.write_table( |
| 229 | table=tb, |
| 230 | where=output_file, |
| 231 | row_group_size=20000, |
| 232 | # data_page_size=50000, |
| 233 | ) |
| 234 | print(f"处理原文件{origin_file} 保存至{output_file}") |
| 235 | |
| 236 | |
| 237 | if __name__ =='__main__': |
nothing calls this directly
no test coverage detected