Update bert_gen.py

This commit is contained in:
Stardust·减
2023-09-05 20:17:48 +08:00
committed by GitHub
parent a716d8ef82
commit 908c0bca2c

View File

@@ -44,6 +44,10 @@ def process_line(line):
if __name__ == "__main__": if __name__ == "__main__":
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('-c', '--config', type=str, default="configs/config.json") 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 config_path = args.config
hps = utils.get_hparams_from_file(config_path) hps = utils.get_hparams_from_file(config_path)
lines = [] lines = []
@@ -53,9 +57,6 @@ if __name__ == "__main__":
with open(hps.data.validation_files, encoding="utf-8") as f: with open(hps.data.validation_files, encoding="utf-8") as f:
lines.extend(f.readlines()) 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 num_processes = args.num_processes
with Pool(processes=num_processes) as pool: with Pool(processes=num_processes) as pool: