| 46 | |
| 47 | |
| 48 | def main(): |
| 49 | parser = argparse.ArgumentParser( |
| 50 | formatter_class=argparse.RawTextHelpFormatter) |
| 51 | parser.add_argument('--input-dir', |
| 52 | '-input_dir', |
| 53 | '-i', |
| 54 | help='folder name of checkpoint files', |
| 55 | required=True) |
| 56 | |
| 57 | parser.add_argument('--output-dir', |
| 58 | '-output_dir', |
| 59 | '-o', |
| 60 | help='folder name of checkpoint files', |
| 61 | required=True) |
| 62 | |
| 63 | parser.add_argument('--num-processes', |
| 64 | '-num_processes', |
| 65 | '-p', |
| 66 | type=int, |
| 67 | default=None, |
| 68 | help='Number of processes') |
| 69 | |
| 70 | args = parser.parse_args() |
| 71 | po = multiprocessing.Pool(args.num_processes) |
| 72 | |
| 73 | if not os.path.exists(args.output_dir): |
| 74 | os.makedirs(args.output_dir) |
| 75 | |
| 76 | for input_file in tqdm(glob(args.input_dir + '/*.json')): |
| 77 | fn = input_file.split('/')[-1] |
| 78 | output_file = os.path.join(args.output_dir, fn) |
| 79 | po.apply_async(func=run_preprocess, args=(input_file, output_file)) |
| 80 | po.close() |
| 81 | po.join() |
| 82 | print('done') |
| 83 | |
| 84 | |
| 85 | if __name__ == '__main__': |