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