# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import argparse import os from nemo.collections.nlp.data.machine_translation.preproc_mt_data import MTDataPreproc if __name__ == '__main__': parser = argparse.ArgumentParser(description='NMT dataset pre-processing') parser.add_argument('--shared_tokenizer', action="store_true", help='Whether to share encoder/decoder tokenizers') parser.add_argument('--clean', action="store_true", help='Whether to clean dataset based on length diff') parser.add_argument('--tar_file_prefix', type=str, default='parallel', help='Prefix for tar files') parser.add_argument('--src_fname', type=str, required=True, help='Path to the source file') parser.add_argument('--tgt_fname', type=str, required=True, help='Path to the target file') parser.add_argument('--out_dir', type=str, required=True, help='Path to store dataloader and tokenizer models') parser.add_argument('--encoder_model_name', type=str, default=None, help='For use with pretrained encoders') parser.add_argument( '--decoder_model_name', type=str, default=None, help='For use with pretrained decoders (not yet supported)' ) parser.add_argument( '--encoder_tokenizer_model', type=str, default='None', help='Path to pre-trained encoder tokenizer model' ) parser.add_argument( '--encoder_tokenizer_name', type=str, default='yttm', help='Encoder BPE Tokenizer Name, Options: [yttm, sentencepiece]', ) parser.add_argument('--encoder_tokenizer_vocab_size', type=int, default=32000, help='Encoder Vocab size after BPE') parser.add_argument( '--encoder_tokenizer_coverage', type=float, default=0.999, help='Encoder Character coverage for BPE' ) parser.add_argument('--encoder_tokenizer_bpe_dropout', type=float, default=0.1, help='Encoder BPE dropout prob') parser.add_argument( '--encoder_tokenizer_r2l', action="store_true", help='Whether to return encoded sequence from right to left' ) parser.add_argument( '--encoder_tokenizer_legacy', action="store_true", help='Whether to use legacy tokenizer implementation of sentencepiece', ) parser.add_argument( '--decoder_tokenizer_model', type=str, default='None', help='Path to pre-trained decoder tokenizer model' ) parser.add_argument( '--decoder_tokenizer_name', type=str, default='yttm', help='Encoder BPE Tokenizer Name, Options: [yttm, sentencepiece]', ) parser.add_argument('--decoder_tokenizer_vocab_size', type=int, default=32000, help='Encoder Vocab size after BPE') parser.add_argument( '--decoder_tokenizer_coverage', type=float, default=0.999, help='Encoder Character coverage for BPE' ) parser.add_argument('--decoder_tokenizer_bpe_dropout', type=float, default=0.1, help='Encoder BPE dropout prob') parser.add_argument( '--decoder_tokenizer_r2l', action="store_true", help='Whether to return encoded sequence from right to left' ) parser.add_argument( '--decoder_tokenizer_legacy', action="store_true", help='Whether to use legacy tokenizer implementation of sentencepiece', ) parser.add_argument('--max_seq_length', type=int, default=512, help='Max Sequence Length') parser.add_argument('--min_seq_length', type=int, default=1, help='Min Sequence Length') parser.add_argument('--tokens_in_batch', type=int, default=16000, help='# Tokens per batch per GPU') parser.add_argument('--coverage', type=float, default=0.999, help='BPE character coverage [0-1]') parser.add_argument( '--lines_per_dataset_fragment', type=int, default=1000000, help='Number of lines to consider for bucketing and padding', ) parser.add_argument( '--num_batches_per_tarfile', type=int, default=1000, help='Number of batches (pickle files) within each tarfile', ) parser.add_argument( '--n_preproc_jobs', type=int, default=-2, help='Number of processes to use for creating the tarred dataset.', ) parser.add_argument( '--byte_fallback', action="store_true", help='Whether to use byte fallback with sentencepiece for BPE tokenization.', ) parser.add_argument( '--split_digits', action="store_true", help='Whether to split digits while tokenizing with sentencepiece.' ) parser.add_argument( '--no_split_by_whitespace', action="store_true", help='If True, this will not respect whitepsaces while learning BPE merges.', ) args = parser.parse_args() if not os.path.exists(args.out_dir): os.mkdir(args.out_dir) if ( args.encoder_tokenizer_model != 'None' and args.decoder_tokenizer_model == 'None' or args.decoder_tokenizer_model != 'None' and args.encoder_tokenizer_model == 'None' ): if args.shared_tokenizer: raise ValueError( ''' If using a pre-trained shared tokenizer, both encoder and decoder tokenizers must be the same ''' ) else: raise ValueError('Both encoder and decoder pre-trained tokenizer models must be specified') if args.encoder_tokenizer_model == 'None' and args.decoder_tokenizer_model == 'None': encoder_tokenizer_model, decoder_tokenizer_model = MTDataPreproc.train_tokenizers( out_dir=args.out_dir, src_fname=args.src_fname, tgt_fname=args.tgt_fname, shared_tokenizer=args.shared_tokenizer, encoder_tokenizer_name=args.encoder_tokenizer_name, encoder_tokenizer_vocab_size=args.encoder_tokenizer_vocab_size, encoder_tokenizer_coverage=args.encoder_tokenizer_coverage, decoder_tokenizer_name=args.decoder_tokenizer_name, decoder_tokenizer_vocab_size=args.decoder_tokenizer_vocab_size, decoder_tokenizer_coverage=args.decoder_tokenizer_coverage, global_rank=0, byte_fallback=args.byte_fallback, split_digits=args.split_digits, split_by_whitespace=not args.no_split_by_whitespace, ) else: encoder_tokenizer_model, decoder_tokenizer_model = args.encoder_tokenizer_model, args.decoder_tokenizer_model encoder_tokenizer, decoder_tokenizer = MTDataPreproc.get_enc_dec_tokenizers( encoder_tokenizer_name=args.encoder_tokenizer_name, encoder_tokenizer_model=encoder_tokenizer_model, encoder_bpe_dropout=args.encoder_tokenizer_bpe_dropout, encoder_r2l=args.encoder_tokenizer_r2l, decoder_tokenizer_name=args.decoder_tokenizer_name, decoder_tokenizer_model=decoder_tokenizer_model, decoder_bpe_dropout=args.decoder_tokenizer_bpe_dropout, decoder_r2l=args.decoder_tokenizer_r2l, encoder_tokenizer_legacy=args.encoder_tokenizer_legacy, decoder_tokenizer_legacy=args.decoder_tokenizer_legacy, ) _, _ = MTDataPreproc.preprocess_parallel_dataset( clean=args.clean, src_fname=args.src_fname, tgt_fname=args.tgt_fname, out_dir=args.out_dir, encoder_tokenizer_name=args.encoder_tokenizer_name, encoder_model_name=args.encoder_model_name, encoder_tokenizer_model=encoder_tokenizer_model, encoder_bpe_dropout=args.encoder_tokenizer_bpe_dropout, encoder_tokenizer_r2l=args.encoder_tokenizer_r2l, decoder_tokenizer_name=args.decoder_tokenizer_name, decoder_model_name=args.decoder_model_name, decoder_tokenizer_model=decoder_tokenizer_model, decoder_tokenizer_r2l=args.decoder_tokenizer_r2l, decoder_bpe_dropout=args.decoder_tokenizer_bpe_dropout, max_seq_length=args.max_seq_length, min_seq_length=args.min_seq_length, tokens_in_batch=args.tokens_in_batch, lines_per_dataset_fragment=args.lines_per_dataset_fragment, num_batches_per_tarfile=args.num_batches_per_tarfile, tar_file_prefix=args.tar_file_prefix, global_rank=0, world_size=1, n_jobs=args.n_preproc_jobs, encoder_tokenizer_legacy=args.encoder_tokenizer_legacy, decoder_tokenizer_legacy=args.decoder_tokenizer_legacy, )