Introduce code from Rust monorepo. This is a mess

tts/NeMo/scripts/asr_language_modeling/neural_rescorer/create_tarred_transformer_lm_dataset.py

@@ -0,0 +1,318 @@
+# 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.
+
+"""
+Execution :
+
+python create_tarred_tokenized_text_lm_dataset.py \
+    --text_path=<comma seperated text filepaths> \
+    --data_root=<path to output directory> \
+    --tokenizer_name="bert-base-cased" \
+    --tokenizer_vocab_file=<path to vocab file for tokenizer> \
+    --num_shards=64 \
+    --chunk_size=8192 \
+    --chunk_write_buffer=512 \
+    --lower_case \
+    --log
+"""
+
+
+import argparse
+import glob
+import json
+import logging
+import os
+import tarfile
+
+import joblib
+import numpy as np
+from tqdm import tqdm
+
+from nemo.collections.nlp.modules.common.tokenizer_utils import get_tokenizer
+
+parser = argparse.ArgumentParser(description='Tarred Tokenized dataset for text language modelling')
+
+# Data path arguments
+parser.add_argument('--text_path', required=True, default=None, type=str, help='Text paths, seperated by commas')
+parser.add_argument('--data_root', required=True, default=None, type=str, help='Output directory')
+
+# General arguments
+parser.add_argument(
+    '--chunk_write_buffer',
+    default=128,
+    type=int,
+    help='Number of chunks of `chunk_size` to buffer for parallel tokenization and serial write to disk',
+)
+parser.add_argument('--lower_case', action='store_true', help='Whether to lower case the corpus')
+parser.add_argument('--log', action='store_true', help='Whether to print logs to terminal')
+
+# Tokenizer arguments
+parser.add_argument('--tokenizer_name', required=False, default=None, type=str, help='Tokenizer name for resolution')
+parser.add_argument(
+    '--tokenizer_model', required=False, default=None, type=str, help='Path to tokenizer model for sentencepiece'
+)
+parser.add_argument('--tokenizer_vocab_file', required=False, type=str, default=None, help='Path to a vocab file')
+parser.add_argument(
+    '--tokenizer_special_tokens', default=None, type=str, nargs='+', help='List of special tokens for the tokenizer'
+)
+
+# Tarred dataset arguments
+parser.add_argument('--num_shards', default=1, type=int, help='Number of shards for the tarfile')
+parser.add_argument('--chunk_size', default=8192, type=int, help='Number of rows of data concatenated into a vector')
+
+parser.set_defaults(log=False, lower_case=False)
+args = parser.parse_args()
+
+
+def __build_dataset_from_text(texts: str, lower_case: bool, chunk_size: int):
+    if ',' in texts:
+        texts = texts.split(',')
+    else:
+        texts = [texts]
+
+    num_lines = 0
+    text_dataset = []
+
+    for text in texts:
+        with open(text, 'r') as in_reader:
+            reader = tqdm(iter(lambda: in_reader.readline(), ''), desc="Read 0 lines", unit=' lines')
+
+            for i, line in enumerate(reader):
+                # Clean text line
+                line = line.replace("\n", "").strip()
+
+                if lower_case:
+                    line = line.lower()
+
+                if line:
+                    text_dataset.append(line)
+
+                    num_lines += 1
+                    if num_lines % 100000 == 0:
+                        reader.set_description(f"Read {num_lines} lines")
+
+                    if num_lines % chunk_size == 0:
+                        yield text_dataset, num_lines
+
+                        # Empty cache
+                        text_dataset = []
+
+            logging.info(f"Finished extracting manifest : {text}")
+
+        logging.info("Finished extracting all manifests ! Number of sentences : {}".format(num_lines))
+
+    if len(text_dataset) != 0:
+        yield text_dataset, num_lines
+
+
+def __tokenize_str(texts, tokenizer):
+    tokenized_text = []
+    for text in texts:
+        tok_text = tokenizer.text_to_ids(text)
+        tokenized_text.extend(tok_text)
+    return tokenized_text
+
+
+def __tokenize_text(
+    text_paths, tokenizer, tokenized_cachedir, lower_case: bool = False, chunk_size=8192, write_buffer: int = -1
+):
+    if write_buffer < 1:
+        write_buffer = max(os.cpu_count() - write_buffer, 1)
+
+    logging.info(f"Using write chunk buffer of size {write_buffer}")
+
+    if not os.path.exists(tokenized_cachedir):
+        os.makedirs(tokenized_cachedir)
+
+    # global parameters
+    global_chunk_idx = 0
+    chunk_paths = []
+    chunk_lens = []
+
+    # buffer parameters
+    data_cache = []
+    chunk_idx = 0
+
+    text_generator = iter(__build_dataset_from_text(text_paths, lower_case=lower_case, chunk_size=chunk_size))
+    global_num_lines = 0
+    last_batch = False
+
+    with joblib.Parallel(n_jobs=-2, verbose=10) as parallel:
+
+        while True:
+            try:
+                data, num_lines = next(text_generator)
+                data_cache.append(data)
+
+                global_num_lines = num_lines
+            except StopIteration:
+                last_batch = True
+
+            # Update counters
+            chunk_idx += 1
+
+            if (chunk_idx == write_buffer) or last_batch:
+                # write the chunks into disk after parallel tokenization
+                tokenized_data_list = parallel(
+                    joblib.delayed(__tokenize_str)(chunk, tokenizer) for chunk in data_cache
+                )
+
+                # Sequential write cache
+                for chunk in tokenized_data_list:
+                    fp = os.path.join(tokenized_cachedir, f"chunk_{global_chunk_idx}.npy")
+                    chunk = np.asarray(chunk, dtype=np.int64)
+                    np.save(fp, chunk, allow_pickle=False)
+
+                    chunk_paths.append(fp)
+                    chunk_lens.append(len(chunk))
+
+                    global_chunk_idx += 1
+
+                logging.info(f"Wrote a total of {global_chunk_idx} chunks to file...")
+
+                # reset buffers
+                data_cache.clear()
+                del data_cache
+
+                data_cache = []
+                chunk_idx = 0
+
+                if last_batch:
+                    logging.info("Finished tokenizing last chunk")
+                    break
+
+    logging.info(
+        f"Chunking {global_num_lines} rows into {global_num_lines // chunk_size} tasks (each chunk contains {chunk_size} elements)"
+    )
+    return chunk_paths, chunk_lens
+
+
+def __create_chunk(data_root, chunk_path, shard_id, compute_metrics=False):
+    """Creates a tarball containing the tokenized text chunks.
+       """
+    tar = tarfile.open(os.path.join(data_root, f'text_{shard_id}.tar'), mode='a')
+
+    # We squash the filename since we do not preserve directory structure of tokenized text in the tarball.
+    base, ext = os.path.splitext(chunk_path)
+    base = base.replace(os.pathsep, '_')
+    # Need the following replacement as long as WebDataset splits on first period
+    base = base.replace('.', '_')
+    squashed_filename = f'{base}{ext}'
+    tar.add(chunk_path, arcname=squashed_filename)
+
+    tar.close()
+
+    if compute_metrics:
+        data = np.load(chunk_path, allow_pickle=False)
+        chunk_len = len(data)
+        return (chunk_len,)
+    else:
+        return None
+
+
+def __write_tarred_tokenized_text_dataset(data_root, num_shards, chunk_paths, chunk_lens):
+    num_chunks = len(chunk_paths)
+
+    if chunk_lens is not None:
+        num_text = sum(chunk_lens)
+        shard_counts = {chunk_id: chunk_len for chunk_id, chunk_len in enumerate(chunk_lens)}
+        compute_metrics = False
+
+    else:
+        num_text = 0
+        shard_counts = {}
+        compute_metrics = True
+
+    for chunk_id, chunk_path in enumerate(tqdm(chunk_paths, desc='Writing chunk ', unit=' chunks')):
+        shard_id = chunk_id % num_shards
+        metrics = __create_chunk(data_root, chunk_path, shard_id, compute_metrics=compute_metrics)
+
+        if metrics is not None:
+            num_text += metrics[0]
+            shard_counts[chunk_id] = metrics[0]
+
+    # write metadata
+    metadata_path = os.path.join(data_root, 'metadata.json')
+    with open(metadata_path, 'w') as f:
+        metadata = {'num_chunks': num_chunks, 'num_text': num_text, 'shard_count': shard_counts}
+        json.dump(metadata, f, indent=4)
+        logging.info("Metadata writen..")
+
+
+def main():
+    text_path = args.text_path
+    data_root = args.data_root
+
+    if args.log:
+        logging.basicConfig(level=logging.INFO)
+
+    tokenized_cachedir = os.path.join(data_root, '_tokenized_dataset_cachedir')
+
+    if os.path.exists(tokenized_cachedir):
+        logging.warning(
+            f'Tokenized cache directory {tokenized_cachedir} already potentially contains files.'
+            f'In such a case, please be aware that the tarfiles will be **appended** instead of overridden!'
+        )
+
+    if not os.path.exists(data_root):
+        os.makedirs(data_root)
+
+    chunk_paths = None
+    chunk_lens = None
+
+    if os.path.exists(tokenized_cachedir):
+        paths = glob.glob(os.path.join(tokenized_cachedir, "*.npy"))
+        if len(paths) > 0:
+            logging.info("Cached tokenized numpy files found, skipping re-tokenization of dataset")
+
+            chunk_paths = paths
+            chunk_lens = None
+
+    if chunk_paths is None:
+        if args.tokenizer_name is None:
+            raise ValueError("`tokenizer_name` name is required when tokenizing the dataset for the first time.")
+
+        if args.tokenizer_vocab_file is None:
+            raise ValueError("`tokenizer_vocab_file` is required when constructing the tokenized dataset")
+
+        tokenizer = get_tokenizer(
+            tokenizer_name=args.tokenizer_name,
+            tokenizer_model=args.tokenizer_model,
+            vocab_file=args.tokenizer_vocab_file,
+            special_tokens=args.tokenizer_special_tokens,
+        )
+
+        logging.info("Built tokenizer")
+
+        # tokenize text data into sub-words
+        chunk_paths, chunk_lens = __tokenize_text(
+            text_paths=text_path,
+            tokenizer=tokenizer,
+            tokenized_cachedir=tokenized_cachedir,
+            lower_case=args.lower_case,
+            chunk_size=args.chunk_size,
+            write_buffer=args.chunk_write_buffer,
+        )
+        logging.info(f"Tokenized dataset into sub-words and serialized cache at {tokenized_cachedir}")
+
+    # Write tarred dataset
+    __write_tarred_tokenized_text_dataset(
+        data_root, num_shards=args.num_shards, chunk_paths=chunk_paths, chunk_lens=chunk_lens
+    )
+
+    logging.info('Done preparing tokenized dataset!')
+
+
+if __name__ == "__main__":
+    main()

tts/NeMo/scripts/asr_language_modeling/neural_rescorer/eval_neural_rescorer.py

@@ -0,0 +1,288 @@
+# Copyright (c) 2021, 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.
+
+"""
+This script would evaluate a neural language model (Transformer) trained with
+`examples/nlp/language_modeling/transformer_lm.py' as a rescorer for ASR systems.
+Given a trained TransformerLMModel `.nemo` file, this script can be used to re-score the beams obtained from a beam
+search decoder of an ASR model.
+
+USAGE:
+1. Obtain `.tsv` file with beams and their corresponding scores. Scores can be from a regular beam search decoder or
+   in fusion with an N-gram LM scores. For a given beam size `beam_size` and a number of examples
+   for evaluation `num_eval_examples`, it should contain (`beam_size` x `num_eval_examples`) lines of
+   form `beam_candidate_text \t score`. This file can be generated by `scripts/asr_language_modeling/ngram_lm/eval_beamsearch_ngram.py`.
+
+2. Rescore the candidates:
+    python eval_neural_rescorer.py
+        --lm_model=[path to .nemo file of the LM]
+        --beams_file=[path to beams .tsv file]
+        --beam_size=[size of the beams]
+        --eval_manifest=[path to eval manifest .json file]
+        --batch_size=[batch size used for inference on the LM model]
+        --alpha=[the value for the parameter rescorer_alpha]
+        --beta=[the value for the parameter rescorer_beta]
+
+You may find more info on how to use this script at:
+https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/asr_language_modeling.html
+
+"""
+
+import contextlib
+import json
+from argparse import ArgumentParser
+
+import editdistance
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import torch
+import tqdm
+
+from nemo.collections.nlp.models.language_modeling import TransformerLMModel
+from nemo.utils import logging
+
+
+class BeamScoresDataset(torch.utils.data.Dataset):
+    """
+    Dataset to read the score file containing the beams and their score
+
+    Args:
+        data_path: path to the beams file
+        tokenizer: tokenizer of the LM model
+        manifest_path: manifest `.json` file which contains the ground truths transcripts
+        beam_size: the number of beams per sample
+        max_seq_length: the maximum length of sequences
+    """
+
+    def __init__(self, data_path, tokenizer, manifest_path, beam_size=128, max_seq_length=256):
+        self.data = pd.read_csv(data_path, delimiter="\t", header=None)
+        self.tokenizer = tokenizer
+        self.ground_truths = []
+        with open(manifest_path, 'r') as f_orig:
+            for line in f_orig:
+                item = json.loads(line)
+                self.ground_truths.append(item['text'])
+        self.beam_size = beam_size
+        self.max_seq_length = max_seq_length
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        text = str(self.data[0][idx])
+        tokens = [self.tokenizer.bos_id] + self.tokenizer.text_to_ids(text) + [self.tokenizer.eos_id]
+        input_ids = [self.tokenizer.pad_id] * self.max_seq_length
+        input_ids[: len(tokens)] = tokens
+        input_ids = np.array(input_ids)
+        input_mask = (input_ids != self.tokenizer.pad_id).astype(np.float32)
+        acoustic_score = self.data[1][idx]
+        dist = editdistance.eval(text.split(), self.ground_truths[idx // self.beam_size].split())
+        ref_len = len(self.ground_truths[idx // self.beam_size].split())
+        len_in_chars = len(str(self.data[0][idx]))
+        return input_ids, input_mask, acoustic_score, dist, ref_len, len_in_chars, idx
+
+
+def linear_search_wer(
+    dists, scores1, scores2, total_len, coef_range=[0, 10], coef_steps=10000, param_name='parameter'
+):
+    """
+    performs linear search to find the best coefficient when two set of scores are getting linearly fused.
+
+    Args:
+        dists: Tesnor of the distances between the ground truth and the candidates with shape of [number of samples, beam size]
+        scores1: Tensor of the first set of scores with shape of [number of samples, beam size]
+        scores2: Tensor of the second set of scores with shape of [number of samples, beam size]
+        total_len: The total length of all samples
+        coef_range: the search range for the coefficient
+        coef_steps: the number of steps that the search range would get divided into
+        param_name: the name of the parameter to be used in the figure
+
+    Output:
+        (best coefficient found, best WER achived)
+    """
+    scale = scores1.mean().abs().item() / scores2.mean().abs().item()
+    left = coef_range[0] * scale
+    right = coef_range[1] * scale
+    coefs = np.linspace(left, right, coef_steps)
+
+    best_wer = 10000
+    best_coef = left
+    wers = []
+    for coef in coefs:
+        scores = scores1 + coef * scores2
+        wer = compute_wer(dists, scores, total_len)
+        wers.append(wer)
+        if wer < best_wer:
+            best_wer = wer
+            best_coef = coef
+
+    plt.plot(coefs, wers)
+    plt.title(f'WER% after rescoring with different values of {param_name}')
+    plt.ylabel('WER%')
+    plt.xlabel(param_name)
+    plt.show()
+    return best_coef, best_wer
+
+
+def compute_wer(dists, scores, total_len):
+    """
+    Sorts the candidates based on the scores and calculates the WER with the new top candidates.
+
+    Args:
+        dists: Tensor of the distances between the ground truth and the candidates with shape of [number of samples, beam size]
+        scores: Tensor of the scores for candidates with shape of [number of samples, beam size]
+        total_len: The total length of all samples
+
+    Output:
+        WER with the new scores
+    """
+    indices = scores.max(dim=1, keepdim=True)[1]
+    wer = dists.gather(dim=1, index=indices).sum() / total_len
+    wer = wer.item()
+    return wer
+
+
+def main():
+    parser = ArgumentParser()
+    parser.add_argument("--lm_model_file", type=str, required=True, help="path to LM model .nemo file")
+    parser.add_argument("--beams_file", type=str, required=True, help="path to beams .tsv file")
+    parser.add_argument(
+        "--eval_manifest", type=str, required=True, help="path to the evaluation `.json` manifest file"
+    )
+    parser.add_argument("--beam_size", type=int, required=True, help="number of beams per candidate")
+    parser.add_argument("--batch_size", type=int, default=256, help="inference batch size")
+    parser.add_argument("--alpha", type=float, default=None, help="parameter alpha of the fusion")
+    parser.add_argument("--beta", type=float, default=None, help="parameter beta of the fusion")
+    parser.add_argument(
+        "--scores_output_file", default=None, type=str, help="The optional path to store the rescored beams"
+    )
+    parser.add_argument(
+        "--device", default="cuda", type=str, help="The device to load the model onto to calculate the scores"
+    )
+    parser.add_argument(
+        "--use_amp", action="store_true", help="Whether to use AMP if available to calculate the scores"
+    )
+    args = parser.parse_args()
+
+    device = args.device
+    if device.startswith("cuda") and not torch.cuda.is_available():
+        logging.info(f"cuda is not available! switched to cpu.")
+        device = "cpu"
+
+    if args.lm_model_file.endswith(".nemo"):
+        logging.info("Attempting to initialize from .nemo file")
+        model = TransformerLMModel.restore_from(
+            restore_path=args.lm_model_file, map_location=torch.device(device)
+        ).eval()
+    else:
+        raise NotImplementedError(f"Only supports .nemo files, but got: {args.model}")
+
+    max_seq_length = model.encoder._embedding.position_embedding.pos_enc.shape[0]
+    dataset = BeamScoresDataset(args.beams_file, model.tokenizer, args.eval_manifest, args.beam_size, max_seq_length)
+    data_loader = torch.utils.data.DataLoader(dataset=dataset, batch_size=args.batch_size)
+
+    if args.use_amp:
+        if torch.cuda.is_available() and hasattr(torch.cuda, 'amp') and hasattr(torch.cuda.amp, 'autocast'):
+            logging.info("AMP is enabled!\n")
+            autocast = torch.cuda.amp.autocast
+    else:
+
+        @contextlib.contextmanager
+        def autocast():
+            yield
+
+    logging.info(f"Rescoring with beam_size: {args.beam_size}")
+    logging.info("Calculating the scores...")
+    with autocast():
+        with torch.no_grad():
+            am_scores, lm_scores, dists, ref_lens, lens_in_chars = [], [], [], [], []
+            for batch in tqdm.tqdm(data_loader):
+                input_ids, input_mask, acoustic_score, dist, ref_len, len_in_chars, idx = batch
+
+                max_len_in_batch = input_mask.sum(dim=0).argmin().item()
+                input_ids, input_mask = input_ids[:, :max_len_in_batch], input_mask[:, :max_len_in_batch]
+                if torch.cuda.is_available():
+                    input_ids, input_mask = input_ids.to(device), input_mask.to(device)
+                    dist, acoustic_score, len_in_chars = (
+                        dist.to(device),
+                        acoustic_score.to(device),
+                        len_in_chars.to(device),
+                    )
+
+                log_probs = model.forward(input_ids[:, :-1], input_mask[:, :-1])
+                target_log_probs = log_probs.gather(2, input_ids[:, 1:].unsqueeze(2)).squeeze(2)
+                neural_lm_score = torch.sum(target_log_probs * input_mask[:, 1:], dim=-1)
+
+                am_scores.append(acoustic_score)
+                lm_scores.append(neural_lm_score)
+                dists.append(dist)
+                ref_lens.append(ref_len)
+                lens_in_chars.append(len_in_chars)
+
+    am_scores = torch.cat(am_scores).view(-1, args.beam_size)
+    lm_scores = torch.cat(lm_scores).view(-1, args.beam_size)
+    dists = torch.cat(dists).view(-1, args.beam_size)
+    ref_lens = torch.cat(ref_lens).view(-1, args.beam_size)
+    lens_in_chars = torch.cat(lens_in_chars).view(-1, args.beam_size).to(am_scores.dtype)
+
+    total_len = ref_lens[:, 0].sum()
+    model_wer = dists[:, 0].sum() / total_len
+    ideal_wer = dists.min(dim=1)[0].sum() / total_len
+
+    if args.alpha is None:
+        logging.info("Linear search for alpha...")
+        coef1, _ = linear_search_wer(
+            dists=dists, scores1=am_scores, scores2=lm_scores, total_len=total_len, param_name='alpha'
+        )
+        coef1 = np.round(coef1, 3)
+        logging.info(f"alpha={coef1} achieved the best WER.")
+        logging.info(f"------------------------------------------------")
+    else:
+        coef1 = args.alpha
+
+    scores = am_scores + coef1 * lm_scores
+
+    if args.beta is None:
+        logging.info("Linear search for beta...")
+        coef2, _ = linear_search_wer(
+            dists=dists, scores1=scores, scores2=lens_in_chars, total_len=total_len, param_name='beta'
+        )
+        coef2 = np.round(coef2, 3)
+        logging.info(f"beta={coef2} achieved the best WER.")
+        logging.info(f"------------------------------------------------")
+    else:
+        coef2 = args.beta
+
+    new_scores = am_scores + coef1 * lm_scores + coef2 * lens_in_chars
+    rescored_wer = compute_wer(dists, new_scores, total_len)
+
+    logging.info(f"Input beams WER: {np.round(model_wer.item() * 100, 2)}%")
+    logging.info(f"------------------------------------------------")
+    logging.info(f"  +LM rescoring WER: {np.round(rescored_wer * 100, 2)}%")
+    logging.info(f"  with alpha={coef1}, beta={coef2}")
+    logging.info(f"------------------------------------------------")
+    logging.info(f"Best possible WER: {np.round(ideal_wer.item() * 100, 2)}%")
+    logging.info(f"------------------------------------------------")
+
+    new_scores_flatten = new_scores.flatten()
+    if args.scores_output_file is not None:
+        logging.info(f'Saving the candidates with their new scores at `{args.scores_output_file}`...')
+        with open(args.scores_output_file, "w") as fout:
+            for sample_id in range(len(dataset)):
+                fout.write(f"{dataset.data[0][sample_id]}\t{new_scores_flatten[sample_id]}\n")
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/scripts/asr_language_modeling/ngram_lm/eval_beamsearch_ngram.py

@@ -0,0 +1,345 @@
+# 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.
+#
+
+# This script would evaluate an N-gram language model trained with KenLM library (https://github.com/kpu/kenlm) in
+# fusion with beam search decoders on top of a trained ASR model. NeMo's beam search decoders are capable of using the
+# KenLM's N-gram models to find the best candidates. This script supports both character level and BPE level
+# encodings and models which is detected automatically from the type of the model.
+# You may train the LM model with 'scripts/ngram_lm/train_kenlm.py'.
+#
+# USAGE: python eval_beamsearch_ngram.py --nemo_model_file <path to the .nemo file of the model> \
+#                                         --input_manifest <path to the evaluation JSON manifest file \
+#                                         --kenlm_model_file <path to the binary KenLM model> \
+#                                         --beam_width <list of the beam widths> \
+#                                         --beam_alpha <list of the beam alphas> \
+#                                         --beam_width <list of the beam betas> \
+#                                         --preds_output_folder <optional folder to store the predictions> \
+#                                         --decoding_mode beam_search_ngram
+#                                         ...
+#
+# You may find more info on how to use this script at:
+# https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/asr_language_modeling.html
+
+
+# Please check train_kenlm.py to find out why we need TOKEN_OFFSET for BPE-based models
+TOKEN_OFFSET = 100
+
+import argparse
+import contextlib
+import json
+import os
+import pickle
+
+import editdistance
+import kenlm_utils
+import numpy as np
+import torch
+from sklearn.model_selection import ParameterGrid
+from tqdm.auto import tqdm
+
+import nemo
+import nemo.collections.asr as nemo_asr
+from nemo.utils import logging
+
+
+def beam_search_eval(
+    all_probs,
+    target_transcripts,
+    vocab,
+    ids_to_text_func=None,
+    preds_output_file=None,
+    lm_path=None,
+    beam_alpha=1.0,
+    beam_beta=0.0,
+    beam_width=128,
+    beam_batch_size=128,
+    progress_bar=True,
+):
+    # creating the beam search decoder
+    beam_search_lm = nemo_asr.modules.BeamSearchDecoderWithLM(
+        vocab=vocab,
+        beam_width=beam_width,
+        alpha=beam_alpha,
+        beta=beam_beta,
+        lm_path=lm_path,
+        num_cpus=max(os.cpu_count(), 1),
+        input_tensor=False,
+    )
+
+    wer_dist_first = cer_dist_first = 0
+    wer_dist_best = cer_dist_best = 0
+    words_count = 0
+    chars_count = 0
+    sample_idx = 0
+    if preds_output_file:
+        out_file = open(preds_output_file, 'w')
+
+    if progress_bar:
+        it = tqdm(
+            range(int(np.ceil(len(all_probs) / beam_batch_size))),
+            desc=f"Beam search decoding with width={beam_width}, alpha={beam_alpha}, beta={beam_beta}",
+            ncols=120,
+        )
+    else:
+        it = range(int(np.ceil(len(all_probs) / beam_batch_size)))
+    for batch_idx in it:
+        # disabling type checking
+        with nemo.core.typecheck.disable_checks():
+            probs_batch = all_probs[batch_idx * beam_batch_size : (batch_idx + 1) * beam_batch_size]
+            beams_batch = beam_search_lm.forward(log_probs=probs_batch, log_probs_length=None,)
+
+        for beams_idx, beams in enumerate(beams_batch):
+            target = target_transcripts[sample_idx + beams_idx]
+            target_split_w = target.split()
+            target_split_c = list(target)
+            words_count += len(target_split_w)
+            chars_count += len(target_split_c)
+            wer_dist_min = cer_dist_min = 10000
+            for candidate_idx, candidate in enumerate(beams):
+                if ids_to_text_func is not None:
+                    # For BPE encodings, need to shift by TOKEN_OFFSET to retrieve the original sub-word ids
+                    pred_text = ids_to_text_func([ord(c) - TOKEN_OFFSET for c in candidate[1]])
+                else:
+                    pred_text = candidate[1]
+                pred_split_w = pred_text.split()
+                wer_dist = editdistance.eval(target_split_w, pred_split_w)
+                pred_split_c = list(pred_text)
+                cer_dist = editdistance.eval(target_split_c, pred_split_c)
+
+                wer_dist_min = min(wer_dist_min, wer_dist)
+                cer_dist_min = min(cer_dist_min, cer_dist)
+
+                if candidate_idx == 0:
+                    # first candidate
+                    wer_dist_first += wer_dist
+                    cer_dist_first += cer_dist
+
+                score = candidate[0]
+                if preds_output_file:
+                    out_file.write('{}\t{}\n'.format(pred_text, score))
+            wer_dist_best += wer_dist_min
+            cer_dist_best += cer_dist_min
+        sample_idx += len(probs_batch)
+
+    if preds_output_file:
+        out_file.close()
+        logging.info(f"Stored the predictions of beam search decoding at '{preds_output_file}'.")
+
+    if lm_path:
+        logging.info(
+            'WER/CER with beam search decoding and N-gram model = {:.2%}/{:.2%}'.format(
+                wer_dist_first / words_count, cer_dist_first / chars_count
+            )
+        )
+    else:
+        logging.info(
+            'WER/CER with beam search decoding = {:.2%}/{:.2%}'.format(
+                wer_dist_first / words_count, cer_dist_first / chars_count
+            )
+        )
+    logging.info(
+        'Best WER/CER in candidates = {:.2%}/{:.2%}'.format(wer_dist_best / words_count, cer_dist_best / chars_count)
+    )
+    logging.info(f"=================================================================================")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description='Evaluate an ASR model with beam search decoding and n-gram KenLM language model.'
+    )
+    parser.add_argument(
+        "--nemo_model_file", required=True, type=str, help="The path of the '.nemo' file of the ASR model"
+    )
+    parser.add_argument(
+        "--kenlm_model_file", required=False, default=None, type=str, help="The path of the KenLM binary model file"
+    )
+    parser.add_argument("--input_manifest", required=True, type=str, help="The manifest file of the evaluation set")
+    parser.add_argument(
+        "--preds_output_folder", default=None, type=str, help="The optional folder where the predictions are stored"
+    )
+    parser.add_argument(
+        "--probs_cache_file", default=None, type=str, help="The cache file for storing the outputs of the model"
+    )
+    parser.add_argument(
+        "--acoustic_batch_size", default=16, type=int, help="The batch size to calculate log probabilities"
+    )
+    parser.add_argument(
+        "--device", default="cuda", type=str, help="The device to load the model onto to calculate log probabilities"
+    )
+    parser.add_argument(
+        "--use_amp", action="store_true", help="Whether to use AMP if available to calculate log probabilities"
+    )
+    parser.add_argument(
+        "--decoding_mode",
+        choices=["greedy", "beamsearch", "beamsearch_ngram"],
+        default="beamsearch_ngram",
+        type=str,
+        help="The decoding scheme to be used for evaluation.",
+    )
+    parser.add_argument(
+        "--beam_width",
+        required=True,
+        type=int,
+        nargs="+",
+        help="The width or list of the widths for the beam search decoding",
+    )
+    parser.add_argument(
+        "--beam_alpha",
+        required=True,
+        type=float,
+        nargs="+",
+        help="The alpha parameter or list of the alphas for the beam search decoding",
+    )
+    parser.add_argument(
+        "--beam_beta",
+        required=True,
+        type=float,
+        nargs="+",
+        help="The beta parameter or list of the betas for the beam search decoding",
+    )
+    parser.add_argument(
+        "--beam_batch_size", default=128, type=int, help="The batch size to be used for beam search decoding"
+    )
+    args = parser.parse_args()
+
+    asr_model = nemo_asr.models.EncDecCTCModelBPE.restore_from(
+        args.nemo_model_file, map_location=torch.device(args.device)
+    )
+
+    target_transcripts = []
+    with open(args.input_manifest, 'r') as manifest_file:
+        audio_file_paths = []
+        for line in tqdm(manifest_file, desc=f"Reading Manifest {args.input_manifest} ...", ncols=120):
+            data = json.loads(line)
+            target_transcripts.append(data['text'])
+            audio_file_paths.append(data['audio_filepath'])
+
+    if args.probs_cache_file and os.path.exists(args.probs_cache_file):
+        logging.info(f"Found a pickle file of probabilities at '{args.probs_cache_file}'.")
+        logging.info(f"Loading the cached pickle file of probabilities from '{args.probs_cache_file}' ...")
+        with open(args.probs_cache_file, 'rb') as probs_file:
+            all_probs = pickle.load(probs_file)
+
+        if len(all_probs) != len(audio_file_paths):
+            raise ValueError(
+                f"The number of samples in the probabilities file '{args.probs_cache_file}' does not "
+                f"match the manifest file. You may need to delete the probabilities cached file."
+            )
+    else:
+        if args.use_amp:
+            if torch.cuda.is_available() and hasattr(torch.cuda, 'amp') and hasattr(torch.cuda.amp, 'autocast'):
+                logging.info("AMP is enabled!\n")
+                autocast = torch.cuda.amp.autocast
+        else:
+
+            @contextlib.contextmanager
+            def autocast():
+                yield
+
+        with autocast():
+            with torch.no_grad():
+                all_logits = asr_model.transcribe(audio_file_paths, batch_size=args.acoustic_batch_size, logprobs=True)
+        all_probs = [kenlm_utils.softmax(logits) for logits in all_logits]
+        if args.probs_cache_file:
+            logging.info(f"Writing pickle files of probabilities at '{args.probs_cache_file}'...")
+            with open(args.probs_cache_file, 'wb') as f_dump:
+                pickle.dump(all_probs, f_dump)
+
+    wer_dist_greedy = 0
+    cer_dist_greedy = 0
+    words_count = 0
+    chars_count = 0
+    for batch_idx, probs in enumerate(all_probs):
+        preds = np.argmax(probs, axis=1)
+        preds_tensor = torch.tensor(preds, device='cpu').unsqueeze(0)
+        pred_text = asr_model._wer.ctc_decoder_predictions_tensor(preds_tensor)[0]
+
+        pred_split_w = pred_text.split()
+        target_split_w = target_transcripts[batch_idx].split()
+        pred_split_c = list(pred_text)
+        target_split_c = list(target_transcripts[batch_idx])
+
+        wer_dist = editdistance.eval(target_split_w, pred_split_w)
+        cer_dist = editdistance.eval(target_split_c, pred_split_c)
+
+        wer_dist_greedy += wer_dist
+        cer_dist_greedy += cer_dist
+        words_count += len(target_split_w)
+        chars_count += len(target_split_c)
+
+    logging.info('Greedy WER/CER = {:.2%}/{:.2%}'.format(wer_dist_greedy / words_count, cer_dist_greedy / chars_count))
+
+    encoding_level = kenlm_utils.SUPPORTED_MODELS.get(type(asr_model).__name__, None)
+    if not encoding_level:
+        logging.warning(
+            f"Model type '{type(asr_model).__name__}' may not be supported. Would try to train a char-level LM."
+        )
+        encoding_level = 'char'
+
+    vocab = asr_model.decoder.vocabulary
+    ids_to_text_func = None
+    if encoding_level == "subword":
+        vocab = [chr(idx + TOKEN_OFFSET) for idx in range(len(vocab))]
+        ids_to_text_func = asr_model.tokenizer.ids_to_text
+    # delete the model to free the memory
+    del asr_model
+
+    if args.decoding_mode == "beamsearch_ngram":
+        if not os.path.exists(args.kenlm_model_file):
+            raise FileNotFoundError(f"Could not find the KenLM model file '{args.kenlm_model_file}'.")
+        lm_path = args.kenlm_model_file
+    else:
+        lm_path = None
+
+    # 'greedy' decoding_mode would skip the beam search decoding
+    if args.decoding_mode in ["beamsearch_ngram", "beamsearch"]:
+
+        params = {'beam_width': args.beam_width, 'beam_alpha': args.beam_alpha, 'beam_beta': args.beam_beta}
+        hp_grid = ParameterGrid(params)
+        hp_grid = list(hp_grid)
+
+        logging.info(f"==============================Starting the beam search decoding===============================")
+        logging.info(f"Grid search size: {len(hp_grid)}")
+        logging.info(f"It may take some time...")
+        logging.info(f"==============================================================================================")
+
+        if args.preds_output_folder and not os.path.exists(args.preds_output_folder):
+            os.mkdir(args.preds_output_folder)
+        for hp in hp_grid:
+            if args.preds_output_folder:
+                preds_output_file = os.path.join(
+                    args.preds_output_folder,
+                    f"preds_out_width{hp['beam_width']}_alpha{hp['beam_alpha']}_beta{hp['beam_beta']}.tsv",
+                )
+            else:
+                preds_output_file = None
+
+            beam_search_eval(
+                all_probs=all_probs,
+                target_transcripts=target_transcripts,
+                vocab=vocab,
+                ids_to_text_func=ids_to_text_func,
+                preds_output_file=preds_output_file,
+                lm_path=lm_path,
+                beam_width=hp["beam_width"],
+                beam_alpha=hp["beam_alpha"],
+                beam_beta=hp["beam_beta"],
+                beam_batch_size=args.beam_batch_size,
+                progress_bar=True,
+            )
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/scripts/asr_language_modeling/ngram_lm/install_beamsearch_decoders.sh

@@ -0,0 +1,6 @@
+git clone https://github.com/NVIDIA/OpenSeq2Seq -b ctc-decoders
+mv OpenSeq2Seq/decoders .
+rm -rf OpenSeq2Seq
+cd decoders
+./setup.sh
+cd ..

tts/NeMo/scripts/asr_language_modeling/ngram_lm/kenlm_utils.py

@@ -0,0 +1,112 @@
+# 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.
+#
+
+"""
+Utility methods to be used for training N-gram LM with KenLM in train_kenlm.py
+"""
+
+import json
+import os
+
+import numpy as np
+from joblib import Parallel, delayed
+from tqdm.auto import tqdm
+
+# List of the supported models to be used with N-gram LM and beam search decoding
+SUPPORTED_MODELS = {
+    'EncDecCTCModelBPE': 'subword',
+    'EncDecCTCModel': 'char',
+}
+
+
+def softmax(x):
+    e = np.exp(x - np.max(x))
+    return e / e.sum(axis=-1).reshape([x.shape[0], 1])
+
+
+def read_train_file(path, lowercase: bool = False):
+    lines_read = 0
+    text_dataset = []
+
+    with open(path, 'r') as f:
+        reader = tqdm(iter(lambda: f.readline(), ''), desc="Read 0 lines", unit=' lines')
+        for i, line in enumerate(reader):
+            if path.endswith('.json'):
+                line = json.loads(line)['text']
+
+            line = line.replace("\n", "").strip()
+            if lowercase:
+                line = line.lower()
+
+            if line:
+                text_dataset.append(line)
+
+                lines_read += 1
+                if lines_read % 100000 == 0:
+                    reader.set_description(f"Read {lines_read} lines")
+
+    return text_dataset
+
+
+def tokenize_str(texts, tokenizer, offset):
+    tokenized_text = []
+    for text in texts:
+        tok_text = tokenizer.text_to_ids(text)
+        tok_text = [chr(token + offset) for token in tok_text]
+        tokenized_text.append(tok_text)
+    return tokenized_text
+
+
+def tokenize_text(data, tokenizer, path, chunk_size=8192, buffer_size=32, token_offset=100):
+    dataset_len = len(data)
+    print(
+        f"Chunking {dataset_len} rows into {dataset_len / float(chunk_size):0.4f} tasks (each chunk contains {chunk_size} elements)"
+    )
+
+    current_step = 0
+    if os.path.exists(path):
+        print(f"Deleting previous file : {path}")
+        os.remove(path)
+
+    with Parallel(n_jobs=-2, verbose=10) as parallel:
+        while True:
+            start = current_step * chunk_size
+            end = min((current_step + buffer_size) * chunk_size, dataset_len)
+
+            tokenized_data = parallel(
+                delayed(tokenize_str)(data[start : start + chunk_size], tokenizer, token_offset)
+                for start in range(start, end, chunk_size)
+            )
+
+            # Write dataset
+            write_dataset(tokenized_data, path)
+            current_step += len(tokenized_data)
+            print(f"Finished writing {len(tokenized_data)} chunks to {path}. Current chunk index = {current_step}")
+            del tokenized_data
+            if end >= dataset_len:
+                break
+
+
+def write_dataset(chunks, path):
+    basedir = os.path.dirname(path)
+
+    if not os.path.exists(basedir):
+        os.makedirs(basedir, exist_ok=True)
+
+    with open(path, 'a+', encoding='utf-8') as f:
+        for chunk_idx in tqdm(range(len(chunks)), desc='Chunk ', total=len(chunks), unit=' chunks'):
+            for text in chunks[chunk_idx]:
+                line = ' '.join(text)
+                f.write(f"{line}\n")

tts/NeMo/scripts/asr_language_modeling/ngram_lm/train_kenlm.py

@@ -0,0 +1,154 @@
+# 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.
+#
+
+
+# This script would train an N-gram language model with KenLM library (https://github.com/kpu/kenlm) which can be used
+# with the beam search decoders on top of the ASR models. This script supports both character level and BPE level
+# encodings and models which is detected automatically from the type of the model.
+# After the N-gram model is trained, and stored in the binary format, you may use
+# 'scripts/ngram_lm/eval_beamsearch_ngram.py' to evaluate it on an ASR model.
+#
+# You need to install the KenLM library and also the beam search decoders to use this feature. Please refer
+# to 'scripts/ngram_lm/install_beamsearch_decoders.sh' on how to install them.
+#
+# USAGE: python train_kenlm.py --nemo_model_file <path to the .nemo file of the model> \
+#                              --train_file <path to the training text or JSON manifest file \
+#                              --kenlm_bin_path <path to the bin folder of KenLM library> \
+#                              --kenlm_model_file <path to store the binary KenLM model> \
+#                              --ngram_length <order of N-gram model>
+#
+# After training is done, the binary LM model is stored at the path specified by '--kenlm_model_file'.
+# You may find more info on how to use this script at:
+# https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/asr_language_modeling.html
+
+import argparse
+import logging
+import os
+import subprocess
+import sys
+
+import kenlm_utils
+import torch
+
+import nemo.collections.asr as nemo_asr
+from nemo.utils import logging
+
+"""
+NeMo's beam search decoders only support char-level encodings. In order to make it work with BPE-level encodings, we
+use a trick to encode the sub-word tokens of the training data as unicode characters and train a char-level KenLM. 
+TOKEN_OFFSET is the offset in the unicode table to be used to encode the BPE sub-words. This encoding scheme reduces 
+the required memory significantly, and the LM and its binary blob format require less storage space.
+"""
+TOKEN_OFFSET = 100
+
+CHUNK_SIZE = 8192
+CHUNK_BUFFER_SIZE = 512
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description='Train an N-gram language model with KenLM to be used with beam search decoder of ASR models.'
+    )
+    parser.add_argument(
+        "--train_file",
+        required=True,
+        type=str,
+        help="Path to the training file, it can be a text file or JSON manifest",
+    )
+    parser.add_argument(
+        "--nemo_model_file", required=True, type=str, help="The path of the '.nemo' file of the ASR model"
+    )
+    parser.add_argument(
+        "--kenlm_model_file", required=True, type=str, help="The path to store the KenLM binary model file"
+    )
+    parser.add_argument("--ngram_length", required=True, type=int, help="The order of N-gram LM")
+    parser.add_argument("--kenlm_bin_path", required=True, type=str, help="The path to the bin folder of KenLM")
+    parser.add_argument(
+        "--do_lowercase", action='store_true', help="Whether to apply lower case conversion on the training text"
+    )
+    args = parser.parse_args()
+
+    """ TOKENIZER SETUP """
+    logging.info(f"Loading nemo model '{args.nemo_model_file}' ...")
+    model = nemo_asr.models.ASRModel.restore_from(args.nemo_model_file, map_location=torch.device('cpu'))
+
+    encoding_level = kenlm_utils.SUPPORTED_MODELS.get(type(model).__name__, None)
+    if not encoding_level:
+        logging.warning(
+            f"Model type '{type(model).__name__}' may not be supported. Would try to train a char-level LM."
+        )
+        encoding_level = 'char'
+
+    """ DATASET SETUP """
+    logging.info(f"Encoding the train file '{args.train_file}' ...")
+    dataset = kenlm_utils.read_train_file(args.train_file, lowercase=args.do_lowercase)
+    encoded_train_file = f"{args.kenlm_model_file}.tmp.txt"
+    if encoding_level == "subword":
+        kenlm_utils.tokenize_text(
+            dataset,
+            model.tokenizer,
+            path=encoded_train_file,
+            chunk_size=CHUNK_SIZE,
+            buffer_size=CHUNK_BUFFER_SIZE,
+            token_offset=TOKEN_OFFSET,
+        )
+        # --discount_fallback is needed for training KenLM for BPE-based models
+        discount_arg = "--discount_fallback"
+    else:
+        with open(encoded_train_file, 'w', encoding='utf-8') as f:
+            for line in dataset:
+                f.write(f"{line}\n")
+
+        discount_arg = ""
+
+    del model
+
+    arpa_file = f"{args.kenlm_model_file}.tmp.arpa"
+    """ LMPLZ ARGUMENT SETUP """
+    kenlm_args = [
+        os.path.join(args.kenlm_bin_path, 'lmplz'),
+        "-o",
+        f"{args.ngram_length}",
+        "--text",
+        encoded_train_file,
+        "--arpa",
+        arpa_file,
+        discount_arg,
+    ]
+
+    ret = subprocess.run(kenlm_args, capture_output=False, text=True, stdout=sys.stdout, stderr=sys.stderr)
+    if ret.returncode != 0:
+        raise RuntimeError("Training KenLM was not successful!")
+    """ BINARY BUILD """
+    logging.info(f"Running binary_build command \n\n{' '.join(kenlm_args)}\n\n")
+    kenlm_args = [
+        os.path.join(args.kenlm_bin_path, "build_binary"),
+        "trie",
+        arpa_file,
+        args.kenlm_model_file,
+    ]
+    ret = subprocess.run(kenlm_args, capture_output=False, text=True, stdout=sys.stdout, stderr=sys.stderr)
+
+    if ret.returncode != 0:
+        raise RuntimeError("Training KenLM was not successful!")
+
+    os.remove(encoded_train_file)
+    logging.info(f"Deleted the temporary encoded training file '{encoded_train_file}'.")
+    os.remove(arpa_file)
+    logging.info(f"Deleted the arpa file '{arpa_file}'.")
+
+
+if __name__ == '__main__':
+    main()