From dfdb07f261405d1c9248088e831487e9ece04a30 Mon Sep 17 00:00:00 2001 From: George Cooper Date: Mon, 2 Nov 2020 17:37:02 -0500 Subject: [PATCH 1/2] Save vocab and merge files in the models dir --- megatron/arguments.py | 6 ++++-- megatron/tokenizer/tokenizer.py | 30 ++++++++++++++++++++++++------ 2 files changed, 28 insertions(+), 8 deletions(-) diff --git a/megatron/arguments.py b/megatron/arguments.py index c4555af1104..b350994ca3d 100644 --- a/megatron/arguments.py +++ b/megatron/arguments.py @@ -405,9 +405,11 @@ def _add_data_args(parser): '`90,5,5` will use 90% of data for training, 5% for ' 'validation and 5% for test.') group.add_argument('--vocab-file', type=str, default=None, - help='Path to the vocab file.') + help='Path to the vocab file. Will be copied to the model directory. If omitted and a vocab ' + 'file already exists in the model directory, that file will be used.') group.add_argument('--merge-file', type=str, default=None, - help='Path to the BPE merge file.') + help='Path to the BPE merge file. Will be copied to the model directory. If omitted and a merge ' + 'file already exists in the model directory, that file will be used.') group.add_argument('--seq-length', type=int, default=None, help="Maximum sequence length to process.") group.add_argument('--mask-prob', type=float, default=0.15, diff --git a/megatron/tokenizer/tokenizer.py b/megatron/tokenizer/tokenizer.py index 9c4f9d003aa..8200e17ee5b 100644 --- a/megatron/tokenizer/tokenizer.py +++ b/megatron/tokenizer/tokenizer.py @@ -14,7 +14,8 @@ # limitations under the License. """Megatron tokenizers.""" - +import os +import shutil from abc import ABC from abc import abstractmethod @@ -22,6 +23,22 @@ from .gpt2_tokenization import GPT2Tokenizer +def copy_file_to_model_dir(args, file_from_args, file_name_in_model_dir): + """Copy a file to the model directory and return the path to the file that should be used.""" + if args.save and file_from_args and args.rank == 0: + file_in_model_dir = os.path.join(args.save, file_name_in_model_dir) + print(f"copying vocab file from {file_from_args} to {file_in_model_dir}") + os.makedirs(args.save, exist_ok=True) + try: + shutil.copyfile(file_from_args, file_in_model_dir) + except shutil.SameFileError: + pass + + file_to_use = file_from_args if file_from_args else os.path.join(args.load, file_name_in_model_dir) + assert os.path.exists(file_to_use) + return file_to_use + + def build_tokenizer(args): """Initialize tokenizer.""" if args.rank == 0: @@ -29,16 +46,17 @@ def build_tokenizer(args): flush=True) # Select and instantiate the tokenizer. - assert args.vocab_file is not None + vocab_file = copy_file_to_model_dir(args, args.vocab_file, "vocab.json") + if args.tokenizer_type == 'BertWordPieceLowerCase': - tokenizer = _BertWordPieceTokenizer(vocab_file=args.vocab_file, + tokenizer = _BertWordPieceTokenizer(vocab_file=vocab_file, lower_case=True) elif args.tokenizer_type == 'BertWordPieceCase': - tokenizer = _BertWordPieceTokenizer(vocab_file=args.vocab_file, + tokenizer = _BertWordPieceTokenizer(vocab_file=vocab_file, lower_case=False) elif args.tokenizer_type == 'GPT2BPETokenizer': - assert args.merge_file is not None - tokenizer = _GPT2BPETokenizer(args.vocab_file, args.merge_file) + merge_file = copy_file_to_model_dir(args, args.merge_file, "merges.txt") + tokenizer = _GPT2BPETokenizer(vocab_file, merge_file) else: raise NotImplementedError('{} tokenizer is not ' 'implemented.'.format(args.tokenizer_type)) From 68edab7b03d71b7aa8eaa07655ba2fa2ddae6670 Mon Sep 17 00:00:00 2001 From: George Cooper Date: Mon, 16 Nov 2020 12:49:39 -0500 Subject: [PATCH 2/2] When attempting to save vocab/merge files to the model directory, don't fail if args.save is not set --- megatron/tokenizer/tokenizer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/tokenizer/tokenizer.py b/megatron/tokenizer/tokenizer.py index 8200e17ee5b..1580a00afa7 100644 --- a/megatron/tokenizer/tokenizer.py +++ b/megatron/tokenizer/tokenizer.py @@ -25,7 +25,7 @@ def copy_file_to_model_dir(args, file_from_args, file_name_in_model_dir): """Copy a file to the model directory and return the path to the file that should be used.""" - if args.save and file_from_args and args.rank == 0: + if hasattr(args, "save") and args.save and file_from_args and args.rank == 0: file_in_model_dir = os.path.join(args.save, file_name_in_model_dir) print(f"copying vocab file from {file_from_args} to {file_in_model_dir}") os.makedirs(args.save, exist_ok=True)