RASMUS/Finnish-ASR-Canary-v2

1# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15# USAGE: python process_asr_text_tokenizer.py --manifest=<path to train manifest files, seperated by commas> \16#         --data_root="<output directory>" \17#         --vocab_size=<number of tokens in vocabulary> \18#         --tokenizer=<"spe" or "wpe"> \19#         --log20# where <manifest> can be: train_clean_100, train_clean_360, train_other_50021# You can also put more than one data_set comma-separated:22# --manifest="train_clean_100,train_clean_360,train_other_500"23# or24#       python process_asr_text_tokenizer.py --data_file=<path to train text file> \25#         --data_root="<output directory>" \26#         --vocab_size=<number of tokens in vocabulary> \27#         --tokenizer=<"bpe" or "wpe"> \28#         --log29# where <manifest> can be: train_clean_100, train_clean_360, train_other_50030# You can also put more than one data_set comma-separated:31# --manifest="train_clean_100,train_clean_360,train_other_500"32#33# Args:34#   --manifest or --data_file: If your text data lies inside of an ASR manifest file,35#       then use the --manifest path. If instead the text data is inside a file with separate lines36#       corresponding to different text lines, then use --data_file.37#       In either case, you can add commas to concatenate different manifests or different data files.38#39#   --data_root: The output directory (whose subdirectories will be created if not present) where40#       the tokenizers will be placed.41#42#   --vocab_size: The size of the tokenizer vocabulary. Larger vocabularies can accommodate almost entire,43#       words but the decoder size of any model will grow proportionally.44#45#   --tokenizer: Can be either spe or wpe . spe refers to the Google sentencepiece library tokenizer.46#       wpe refers to the HuggingFace BERT Word Piece tokenizer.47#48#   --no_lower_case: When this flag is passed, it will force the tokenizer to create seperate tokens for49#       upper and lower case characters. By default, the script will turn all the text to lower case50#       before tokenization (and if upper case characters are passed during training/inference, the51#       tokenizer will emit a token equivalent to Out-Of-Vocabulary). Used primarily for the52#       English language.53#54#    --spe_type: The sentencepiece library has a few implementations of the tokenization technique, and55#       spe_type refers to these implementations. Currently supported types are unigram, bpe, char, word.56#       Defaults to bpe.57#58#   --spe_character_coverage: The sentencepiece library considers how much of the original vocabulary it59#       should cover in its "base set" of tokens (akin to the lower and upper case characters of the60#       English language). For almost all languages with small base token sets (<1000 tokens), this61#       should be kept at its default of 1.0. For languages with larger vocabularies (say Japanese,62#       Mandarin, Korean etc), the suggested value is 0.9995.63#64#   --spe_user_defined_symbols: The sentencepiece library allows you to add your own tokens to the base set.65#      This flag allows you to pass a space separated list of tokens that you want to add to the base set.66#      These tokens remain in the decoded text and are encoded automatically when present in the input text.67#68#   --spe_control_symbols: The sentencepiece library allows you to add your own tokens to the base set.69#      This flag allows you to pass a space separated list of tokens that you want to add to the base set.70#      These tokens get removed at decode time and are not encoded from the text - can only be added to the71#      input programatically.72#73#   --spe_byte_fallback: If <unk>, fallback to a byte sequence of the characters.74#75#   --spe_split_digits: If true, digits are split into individual tokens.76#77#   --spe_sample_size: If the dataset is too large, consider using a sampled dataset indicated by a78#       positive integer. By default, any negative value (default = -1) will use the entire dataset.79#80#   --spe_train_extremely_large_corpus: When training a sentencepiece tokenizer on very large amounts of text,81#       sometimes the tokenizer will run out of memory or wont be able to process so much data on RAM.82#       At some point you might receive the following error - "Input corpus too large, try with83#       train_extremely_large_corpus=true". If your machine has large amounts of RAM, it might still be possible84#       to build the tokenizer using the above flag. Will silently fail if it runs out of RAM.85#86#   --spe_max_sentencepiece_length: Limits the maximum length that any any SentencePiece subword can be.87#       Using this will change the subword tokens generated.88#89#   --spe_pad: Adds <pad> as special token.90#91#   --spe_bos: Adds <s> as Begining-of-Sentence special token.92#93#   --spe_eos: Adds </s> as End-of-Sentence special token.94#95#   --log: Whether the script should display log messages96 97 98import argparse99import json100import logging101import os102from typing import List, Optional103 104import tokenizers105 106from nemo.collections.common.tokenizers.sentencepiece_tokenizer import create_spt_model107from nemo.utils.data_utils import DataStoreObject108 109parser = argparse.ArgumentParser(description='Create tokenizer')110group = parser.add_mutually_exclusive_group(required=True)111group.add_argument("--manifest", default=None, type=str, help='Comma separated list of manifest files')112group.add_argument("--data_file", default=None, help='data file from which to create tokenizer model')113parser.add_argument("--data_root", required=True, default=None, type=str, help='Output directory')114parser.add_argument("--vocab_size", default=1024, type=int, help='Vocabulary size')115parser.add_argument("--tokenizer", default="wpe", choices=["spe", "wpe"], help='Type of tokenization to perform')116parser.add_argument(117    "--spe_type",118    default="bpe",119    choices=['bpe', 'unigram', 'char', 'word'],120    help='Type of the SentencePiece model. Can be `bpe`, `unigram`, `char` or `word`.'121    'Used only if --tokenizer == `spe`',122)123parser.add_argument(124    '--spe_character_coverage',125    type=float,126    default=1.0,127    help="Character coverage percentage for SentencePiece tokenization. For languages "128    "with large vocabulary, should be close to 0.9995, otherwise kept as 1.0",129)130parser.add_argument('--spe_bos', action='store_true', help='Add <s> token to SentencePiece Tokenizer.')131parser.add_argument('--spe_eos', action='store_true', help='Add </s> token to SentencePiece Tokenizer.')132parser.add_argument('--spe_pad', action='store_true', help='Add <pad> token to SentencePiece Tokenizer.')133parser.add_argument(134    '--spe_user_defined_symbols', default=None, type=str, nargs='+', help='User defined symbols for SentencePiece'135)136parser.add_argument(137    '--spe_control_symbols', default=None, type=str, nargs='+', help='Control symbols for SentencePiece'138)139parser.add_argument('--spe_split_digits', action='store_true', help='Split digits into separate tokens.')140parser.add_argument(141    '--spe_remove_extra_whitespaces',142    action='store_true',143    help='Remove leading, trailing, and duplicate internal whitespace.',144)145 146parser.add_argument(147    '--spe_sample_size',148    type=int,149    default=-1,150    help="Samples the dataset by `sample_size` if positive integer, otherwise uses whole dataset",151)152parser.add_argument('--spe_train_extremely_large_corpus', action='store_true', help='')153parser.add_argument(154    '--spe_max_sentencepiece_length',155    type=int,156    default=-1,157    help='Limit the maximum number of tokens in each SentencePiece subword. '158    'Must be a positive integer > 0. By default places no limit on subword length.',159)160parser.add_argument(161    '--spe_no_split_by_unicode_script',162    dest='spe_split_by_unicode_script',163    action='store_false',164    help="Don't use Unicode script to split sentence pieces.",165)166parser.add_argument(167    '--spe_byte_fallback',168    dest='spe_byte_fallback',169    action='store_true',170    help="If <unk>, fallback to a byte sequence of the characters.",171)172parser.add_argument('--no_lower_case', dest='lower_case', action='store_false')173parser.add_argument("--log", action='store_true')174parser.set_defaults(log=False, lower_case=True, spe_train_extremely_large_corpus=False)175args = parser.parse_args()176 177 178def __build_document_from_manifests(179    data_root: str,180    manifests: str,181):182    if ',' in manifests:183        manifests = manifests.split(',')184    else:185        manifests = [manifests]186 187    document_dir = os.path.join(data_root, 'text_corpus')188    if not os.path.exists(document_dir):189        os.makedirs(document_dir)190 191    document_path = os.path.join(document_dir, 'document.txt')192 193    if os.path.exists(document_path):194        logging.info('Corpus already exists at path : %s', document_path)195        return document_path196 197    num_lines = 0198    with open(document_path, 'w') as out_writer:199        for manifest in manifests:200            with open(DataStoreObject(manifest).get(), 'r') as in_reader:201                for line in in_reader:202                    item = json.loads(line)203                    text = item['text']204 205                    out_writer.write(text + '\n')206                    out_writer.flush()207 208                    num_lines += 1209 210            logging.info(f"Finished extracting manifest : {manifest}")211 212        logging.info("Finished extracting all manifests ! Number of sentences : {}".format(num_lines))213    return document_path214 215 216def __process_data(217    text_path: str,218    dst_folder: str,219    vocab_size: int,220    tokenizer_type: str,221    spe_type: str,222    spe_character_coverage: float,223    spe_train_extremely_large_corpus: bool,224    spe_sample_size: int,225    spe_max_sentencepiece_length: int,226    spe_split_by_unicode_script: bool,227    spe_bos: bool,228    spe_eos: bool,229    spe_pad: bool,230    spe_control_symbols: Optional[List[str]],231    spe_user_defined_symbols: Optional[List[str]],232    spe_byte_fallback: bool,233    spe_split_digits: bool,234    spe_remove_extra_whitespaces: bool,235    lower_case: bool,236):237    """238    Converts flac to wav and build manifests's json239    Args:240        text_path: source with text lines241        dst_folder: where wav files will be stored242        vocab_size: vocabular size used in encoding the text243        tokenizer_type: type of tokenization to perform - wpe or spe244        spe_type: type of tokenization model used for spe.245        spe_character_coverage: float value between 0 and 1 (as a percentage). For languages with a vast charset,246            can be < 1.0, but for all other languages, it should be set as 1.0247        spe_sample_size: int, default of -1. If positive integer is used, samples the dataset248            by given sample size.249        spe_train_extremely_large_corpus: bool. If dataset is too large, and user has sufficient RAM,250            this flag can be set to try to trained the tokenizer. Will silently fail if it runs out of RAM.251        spe_max_sentencepiece_length: Limits the maximum length of the SentencePiece subword that can be constructed.252            By default, no limit is placed.253        spe_bos: Bool flag, whether to add <s> to SentencePiece tokenizer vocabulary.254        spe_eos: Bool flag, whether to add </s> to SentencePiece tokenizer vocabulary.255        spe_pad: Bool flag, whether to add <pad> to SentencePiece tokenizer vocabulary.256        spe_control_symbols: control symbols to add to tokenizer, as defined by sentencepiece.257            These tokens get removed at decode time and are not encoded from the text - can only be added to the input programatically.258        spe_user_defined_symbols: user symbols to add to tokenizer, as defined by sentencepiece.259            These tokens remain in the decoded text and are encoded automatically when present in the input text.260        spe_byte_fallback: If <unk>, fallback to a byte sequence of the character.261        spe_split_digits: If true, digits are split into individual tokens.262        spe_remove_extra_whitespaces: If true, removes leading, trailing, and duplicate internal whitespace.263        lower_case: whether to tokenize with lower case character set only (for english)264 265    Returns:266    """267    if tokenizer_type == 'spe':268 269        # Prepare directory of tokenizer270        if spe_max_sentencepiece_length > 0:271            tokenizer_dir = os.path.join(dst_folder, 'tokenizer_{}_{}_v{}_max_{}').format(272                tokenizer_type, spe_type, vocab_size, spe_max_sentencepiece_length273            )274        else:275            tokenizer_dir = os.path.join(dst_folder, 'tokenizer_{}_{}_v{}').format(276                tokenizer_type, spe_type, vocab_size277            )278 279        if spe_pad:280            tokenizer_dir = f'{tokenizer_dir}_pad'281        if spe_bos:282            tokenizer_dir = f'{tokenizer_dir}_bos'283        if spe_eos:284            tokenizer_dir = f'{tokenizer_dir}_eos'285 286        if not os.path.exists(tokenizer_dir):287            os.makedirs(tokenizer_dir)288 289        if os.path.exists(os.path.join(tokenizer_dir, 'tokenizer.model')):290            logging.warning("Model file already exists, overriding old model file !")291            os.remove(os.path.join(tokenizer_dir, 'tokenizer.model'))292 293        # Build tokenizer294        tokenizer_path, vocab_path = create_spt_model(295            data_file=text_path,296            vocab_size=vocab_size,297            sample_size=spe_sample_size,298            do_lower_case=lower_case,299            output_dir=tokenizer_dir,300            tokenizer_type=spe_type,301            character_coverage=spe_character_coverage,302            train_extremely_large_corpus=spe_train_extremely_large_corpus,303            max_sentencepiece_length=spe_max_sentencepiece_length,304            split_by_unicode_script=spe_split_by_unicode_script,305            bos=spe_bos,306            eos=spe_eos,307            pad=spe_pad,308            control_symbols=spe_control_symbols,309            user_defined_symbols=spe_user_defined_symbols,310            byte_fallback=spe_byte_fallback,311            split_digits=spe_split_digits,312            remove_extra_whitespaces=spe_remove_extra_whitespaces,313        )314 315    else:316        tokenizer_dir = os.path.join(dst_folder, 'tokenizer_{}_v{}').format(tokenizer_type, vocab_size)317 318        if not os.path.exists(tokenizer_dir):319            os.makedirs(tokenizer_dir)320 321        tokenizer = tokenizers.BertWordPieceTokenizer(lowercase=lower_case)322 323        tokenizer.train(text_path, vocab_size=vocab_size)324        tokenizer.save_model(tokenizer_dir)325 326    return tokenizer_dir327 328 329def main():330    data_root = args.data_root331    manifests = args.manifest332    data_file = args.data_file333    vocab_size = args.vocab_size334    tokenizer = args.tokenizer335    spe_type = args.spe_type336    spe_character_coverage = args.spe_character_coverage337    spe_sample_size = args.spe_sample_size338    spe_train_extremely_large_corpus = args.spe_train_extremely_large_corpus339    spe_max_sentencepiece_length = args.spe_max_sentencepiece_length340    spe_split_by_unicode_script = args.spe_split_by_unicode_script341    spe_bos, spe_eos, spe_pad = args.spe_bos, args.spe_eos, args.spe_pad342    spe_control_symbols = args.spe_control_symbols343    spe_user_defined_symbols = args.spe_user_defined_symbols344    spe_byte_fallback = args.spe_byte_fallback345    spe_split_digits = args.spe_split_digits346    spe_remove_extra_whitespaces = args.spe_remove_extra_whitespaces347    lower_case = args.lower_case348 349    if not os.path.exists(data_root):350        os.makedirs(data_root)351 352    if args.log:353        logging.basicConfig(level=logging.INFO)354 355    if manifests:356        text_corpus_path = __build_document_from_manifests(data_root, manifests)357    else:358        text_corpus_path = data_file359    tokenizer_path = __process_data(360        text_corpus_path,361        data_root,362        vocab_size,363        tokenizer,364        spe_type,365        lower_case=lower_case,366        spe_character_coverage=spe_character_coverage,367        spe_sample_size=spe_sample_size,368        spe_train_extremely_large_corpus=spe_train_extremely_large_corpus,369        spe_max_sentencepiece_length=spe_max_sentencepiece_length,370        spe_split_by_unicode_script=spe_split_by_unicode_script,371        spe_bos=spe_bos,372        spe_eos=spe_eos,373        spe_pad=spe_pad,374        spe_control_symbols=spe_control_symbols,375        spe_user_defined_symbols=spe_user_defined_symbols,376        spe_byte_fallback=spe_byte_fallback,377        spe_split_digits=spe_split_digits,378        spe_remove_extra_whitespaces=spe_remove_extra_whitespaces,379    )380 381    print("Serialized tokenizer at location :", tokenizer_path)382    logging.info('Done!')383 384 385if __name__ == "__main__":386    main()387