From 908c0bca2cf50e128cf8db1cf4424c7fd52d9c87 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stardust=C2=B7=E5=87=8F?= Date: Tue, 5 Sep 2023 20:17:48 +0800 Subject: [PATCH] Update bert_gen.py --- bert_gen.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/bert_gen.py b/bert_gen.py index 9df3658..42dd5ad 100644 --- a/bert_gen.py +++ b/bert_gen.py @@ -44,6 +44,10 @@ def process_line(line): if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument('-c', '--config', type=str, default="configs/config.json") + parser.add_argument( + '--num_processes', type=int, default=1, help='You are advised to set the number of processes to the same as the number of CPU cores' + ) + args = parser.parse_args() config_path = args.config hps = utils.get_hparams_from_file(config_path) lines = [] @@ -53,9 +57,6 @@ if __name__ == "__main__": with open(hps.data.validation_files, encoding="utf-8") as f: lines.extend(f.readlines()) - parser.add_argument( - '--num_processes', type=int, default=1, help='You are advised to set the number of processes to the same as the number of CPU cores' - ) num_processes = args.num_processes with Pool(processes=num_processes) as pool: