| import os |
| import os.path as osp |
| import time |
| import itertools |
| import shutil |
| import glob |
| import argparse |
|
|
| import tqdm |
| import numpy as np |
| import threading |
|
|
| def save_lines(lines, filename): |
| os.makedirs(osp.dirname(filename), exist_ok=True) |
| with open(filename, 'w') as f: |
| f.writelines(lines) |
| del lines |
|
|
| def get_part_jsonls(filepath, total_line_number, parts=512): |
| dirname, filename, ext = osp.dirname(filepath), osp.splitext(osp.basename(filepath))[0], osp.splitext(osp.basename(filepath))[1] |
| if parts == 1: |
| return False, {1: filepath} |
| save_dir = osp.join(dirname, f'{parts:04d}_parts') |
| chunk_id2save_files = {} |
| missing = False |
| chunk_size = int(total_line_number/parts) |
| for chunk_id in range(1, parts+1): |
| if chunk_id == parts: |
| num_of_lines = total_line_number - chunk_size * (parts-1) |
| else: |
| num_of_lines = chunk_size |
| chunk_id2save_files[chunk_id] = osp.join(save_dir, f'{filename}_{chunk_id:04d}_{parts:04d}_{num_of_lines:09d}{ext}') |
| if not osp.exists(chunk_id2save_files[chunk_id]): |
| missing = True |
| return missing, chunk_id2save_files |
|
|
| def split_large_txt_files(filepath, chunk_id2save_files): |
| thread_list = [] |
| chunk_id = 1 |
| with open(filepath, 'r') as f: |
| chunk = [] |
| pbar = tqdm.tqdm(total=len(chunk_id2save_files)) |
| for line in f: |
| chunk.append(line) |
| cur_chunk_size = int(osp.splitext(osp.basename(chunk_id2save_files[chunk_id]))[0].split('_')[-1]) |
| if len(chunk) >= cur_chunk_size: |
| pbar.update(1) |
| thread_list.append(threading.Thread(target=save_lines, args=(chunk, chunk_id2save_files[chunk_id]))) |
| thread_list[-1].start() |
| chunk = [] |
| chunk_id += 1 |
| if len(chunk): |
| import ipdb; ipdb.set_trace() |
| assert not len(chunk) |
| for thread in thread_list: |
| thread.join() |
|
|
| if __name__ == '__main__': |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--jsonl_folder', type=str, default='') |
| parser.add_argument('--parts', type=int, default=600) |
| args = parser.parse_args() |
| for jsonl_filepath in sorted(glob.glob(osp.join(args.jsonl_folder, '*.jsonl'))): |
| print(jsonl_filepath) |
| t1 = time.time() |
| line_num = int(jsonl_filepath.split('_')[-1].split('.')[0]) |
| missing, chunk_id2save_files = get_part_jsonls(jsonl_filepath, line_num, parts=args.parts) |
| split_large_txt_files(jsonl_filepath, chunk_id2save_files) |
| t2 = time.time() |
| print(f'split takes {t2-t1}s') |
|
|
