chore: import upstream snapshot with attribution

scripts/asr_language_modeling/neural_rescorer/create_tarred_transformer_lm_dataset.py

@@ -0,0 +1,317 @@
+# 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.common.tokenizers.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', encoding='utf-8') 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', encoding='utf-8')
+
+    # 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()

scripts/asr_language_modeling/neural_rescorer/eval_neural_rescorer.py

@@ -0,0 +1,326 @@
+# 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 inspect
+import json
+from abc import ABC
+from argparse import ArgumentParser
+
+import numpy as np
+import pandas as pd
+import torch
+import tqdm
+from kaldialign import edit_distance
+from transformers import AutoModelForCausalLM
+
+try:
+    from nemo.collections.nlp.models.language_modeling import TransformerLMModel
+except (ImportError, ModuleNotFoundError):
+    TransformerLMModel = ABC
+
+from nemo.collections.common.tokenizers.tokenizer_utils import get_tokenizer
+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', encoding='utf-8') 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
+
+        if self.tokenizer.pad_id is not None:
+            self.pad_id = self.tokenizer.pad_id
+        elif self.tokenizer.eos_id is not None:
+            self.pad_id = self.tokenizer.eos_id
+        else:
+            logging.warning(f"Using 0 as pad_id as the tokenizer has no pad_id or eos_id.")
+            self.pad_id = 0
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        text = str(self.data[0][idx])
+        tokens = self.tokenizer.text_to_ids(text)
+        if self.tokenizer.bos_id is not None:
+            tokens = [self.tokenizer.bos_id] + tokens
+        if self.tokenizer.eos_id is not None:
+            tokens = tokens + [self.tokenizer.eos_id]
+        input_ids = [self.pad_id] * self.max_seq_length
+        input_ids[: len(tokens)] = tokens
+        input_ids = np.array(input_ids)
+        input_mask = np.zeros(self.max_seq_length)
+        input_mask[: len(tokens)] = 1
+        acoustic_score = self.data[1][idx]
+        dist = edit_distance(self.ground_truths[idx // self.beam_size].split(), text.split())['total']
+        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 achieved)
+    """
+    import matplotlib.pyplot as plt
+
+    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 or the name of a HuggingFace pretrained models like 'transfo-xl-wt103' or 'gpt2'",
+    )
+    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("--max_seq_length", default=512, help="Maximum sequence length (in tokens) for the input")
+    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"):
+        nemo_model = True
+        logging.info("Attempting to initialize from .nemo file...")
+        model = TransformerLMModel.restore_from(
+            restore_path=args.lm_model_file, map_location=torch.device(device)
+        ).eval()
+        model_tokenizer = model.tokenizer
+    else:
+        nemo_model = False
+        logging.info("Attempting to initialize from a pretrained model from HuggingFace...")
+        model = (
+            AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path=args.lm_model_file, is_decoder=True)
+            .to(device)
+            .eval()
+        )
+        model_tokenizer = get_tokenizer(tokenizer_name=args.lm_model_file)
+
+    max_seq_length = args.max_seq_length
+    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 "attention_mask" in inspect.getfullargspec(model.forward).args:
+        support_att_mask = True
+    else:
+        support_att_mask = False
+    logging.info(f"Rescoring with beam_size: {args.beam_size}")
+    logging.info("Calculating the scores...")
+    with torch.amp.autocast(model.device.type, enabled=args.use_amp):
+        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),
+                    )
+                # some models like Transformer-XL don't need attention_mask as input
+                if support_att_mask:
+                    log_probs = model(input_ids=input_ids, attention_mask=input_mask)
+                else:
+                    log_probs = model(input_ids=input_ids)
+
+                if not nemo_model:
+                    log_probs = torch.nn.functional.log_softmax(log_probs.logits, dim=-1)
+
+                target_log_probs = log_probs[:, :-1].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"Oracle 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", encoding='utf-8') 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()

scripts/asr_language_modeling/ngram_lm/create_lexicon_from_arpa.py

@@ -0,0 +1,79 @@
+# Copyright (c) 2023, 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.
+#
+# Use this file to create a lexicon file for Flashlight decoding from an existing KenLM arpa file
+# A lexicon file is required for Flashlight decoding in most cases, as it acts as a map from the words
+# in you arpa file to the representation used by your ASR AM.
+# For more details, see: https://github.com/flashlight/flashlight/tree/main/flashlight/app/asr#data-preparation
+#
+# Usage: python create_lexicon_from_arpa.py --arpa /path/to/english.arpa --model /path/to/model.nemo --lower
+#
+#
+
+
+import argparse
+import os
+import re
+
+from nemo.utils import logging
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Utility script for generating lexicon file from a KenLM arpa file")
+    parser.add_argument("--arpa", required=True, help="path to your arpa file")
+    parser.add_argument("--dst", help="directory to store generated lexicon", default=None)
+    parser.add_argument("--lower", action='store_true', help="Whether to lowercase the arpa vocab")
+    parser.add_argument("--model", default=None, help="path to Nemo model for its tokeniser")
+
+    args = parser.parse_args()
+
+    if not os.path.exists(args.arpa):
+        logging.critical(f"ARPA file [ {args.arpa} ] not detected on disk, aborting!")
+        exit(255)
+
+    if args.dst is not None:
+        save_path = args.dst
+    else:
+        save_path = os.path.dirname(args.arpa)
+    os.makedirs(save_path, exist_ok=True)
+
+    tokenizer = None
+    if args.model is not None:
+        from nemo.collections.asr.models import ASRModel
+
+        model = ASRModel.restore_from(restore_path=args.model, map_location='cpu')
+        if hasattr(model, 'tokenizer'):
+            tokenizer = model.tokenizer
+        else:
+            logging.warning('Supplied Nemo model does not contain a tokenizer')
+
+    lex_file = os.path.join(save_path, os.path.splitext(os.path.basename(args.arpa))[0] + '.lexicon')
+
+    logging.info(f"Writing Lexicon file to: {lex_file}...")
+    with open(lex_file, "w", encoding='utf_8', newline='\n') as f:
+        with open(args.arpa, "r", encoding='utf_8') as arpa:
+            for line in arpa:
+                # verify if the line corresponds to unigram
+                if not re.match(r"[-]*[0-9\.]+\t\S+\t*[-]*[0-9\.]*$", line):
+                    continue
+                word = line.split("\t")[1]
+                word = word.strip().lower() if args.lower else word.strip()
+                if word == "<UNK>" or word == "<unk>" or word == "<s>" or word == "</s>":
+                    continue
+
+                if tokenizer is None:
+                    f.write("{w}\t{s}\n".format(w=word, s=" ".join(word)))
+                else:
+                    w_ids = tokenizer.text_to_ids(word)
+                    if tokenizer.unk_id not in w_ids:
+                        f.write("{w}\t{s}\n".format(w=word, s=" ".join(tokenizer.text_to_tokens(word))))

scripts/asr_language_modeling/ngram_lm/eval_beamsearch_ngram_ctc.py

@@ -0,0 +1,406 @@
+# 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 with CTC decoder. To evaluate a model with
+# Transducer (RNN-T) decoder use another script 'scripts/asr_language_modeling/ngram_lm/eval_beamsearch_ngram_transducer.py'.
+# 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/asr_language_modeling/ngram_lm/train_kenlm.py'.
+
+# Config Help
+
+To discover all arguments of the script, please run :
+python eval_beamsearch_ngram_ctc.py --help
+python eval_beamsearch_ngram_ctc.py --cfg job
+
+# USAGE
+
+python eval_beamsearch_ngram_ctc.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, separated with commas>] \
+           beam_alpha=[<list of the beam alphas, separated with commas>] \
+           beam_beta=[<list of the beam betas, separated with commas>] \
+           preds_output_folder=<optional folder to store the predictions> \
+           probs_cache_file=null \
+           decoding_mode=beamsearch_ngram
+           ...
+
+
+# Grid Search for Hyper parameters
+
+For grid search, you can provide a list of arguments as follows -
+
+           beam_width=[4,8,16,....] \
+           beam_alpha=[-2.0,-1.0,...,1.0,2.0] \
+           beam_beta=[-1.0,-0.5,0.0,...,1.0] \
+
+# 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
+import os
+from dataclasses import dataclass, field, is_dataclass
+from pathlib import Path
+from typing import List, Optional
+
+import msgpack
+import numpy as np
+import torch
+from kaldialign import edit_distance
+from omegaconf import MISSING, OmegaConf
+from sklearn.model_selection import ParameterGrid
+from tqdm.auto import tqdm
+
+import nemo.collections.asr as nemo_asr
+from nemo.collections.asr.models import EncDecHybridRNNTCTCModel
+from nemo.collections.asr.parts.submodules import ctc_beam_decoding
+from nemo.collections.asr.parts.utils.transcribe_utils import PunctuationCapitalization, TextProcessingConfig
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+
+# fmt: off
+
+
+@dataclass
+class EvalBeamSearchNGramConfig:
+    """
+    Evaluate an ASR model with beam search decoding and n-gram KenLM language model.
+    """
+    # # The path of the '.nemo' file of the ASR model or the name of a pretrained model (ngc / huggingface)
+    nemo_model_file: str = MISSING
+
+    # File paths
+    input_manifest: str = MISSING  # The manifest file of the evaluation set
+    kenlm_model_file: Optional[str] = None  # The path of the KenLM binary model file
+    preds_output_folder: Optional[str] = None  # The optional folder where the predictions are stored
+    hyps_cache_file: Optional[str] = None  # The cache file for storing the logprobs of the model
+
+    # Parameters for inference
+    batch_size: int = 16  # The batch size
+    device: str = "cuda"  # The device to load the model onto to calculate log probabilities
+    use_amp: bool = False  # Whether to use AMP if available to calculate log probabilities
+
+    # Beam Search hyperparameters
+
+    # The decoding scheme to be used for evaluation.
+    # Can be one of ["greedy", "beamsearch", "beamsearch_ngram"]
+    decoding_mode: str = "beamsearch_ngram"
+
+    beam_width: List[int] = field(default_factory=lambda: [128])  # The width or list of the widths for the beam search decoding
+    beam_alpha: List[float] = field(default_factory=lambda: [1.0])  # The alpha parameter or list of the alphas for the beam search decoding
+    beam_beta: List[float] = field(default_factory=lambda: [0.0])  # The beta parameter or list of the betas for the beam search decoding
+
+    decoding_strategy: str = "beam"
+    decoding: ctc_beam_decoding.BeamCTCInferConfig = field(default_factory=lambda: ctc_beam_decoding.BeamCTCInferConfig(beam_size=128))
+
+    text_processing: Optional[TextProcessingConfig] = field(default_factory=lambda: TextProcessingConfig(
+        punctuation_marks = ".,?",
+        separate_punctuation = False,
+        do_lowercase = False,
+        rm_punctuation = False,
+    ))
+# fmt: on
+
+
+def apply_text_processing(
+    punctuation_capitalization: PunctuationCapitalization, cfg: EvalBeamSearchNGramConfig, text: List[str] | str
+) -> List[str] | str:
+    is_list = isinstance(text, list)
+    text_arr = text if is_list else [text]
+    if cfg.text_processing.do_lowercase:
+        text_arr = punctuation_capitalization.do_lowercase(text_arr)
+    if cfg.text_processing.rm_punctuation:
+        text_arr = punctuation_capitalization.rm_punctuation(text_arr)
+    if cfg.text_processing.separate_punctuation:
+        text_arr = punctuation_capitalization.separate_punctuation(text_arr)
+
+    return text_arr if is_list else text_arr[0]
+
+
+def beam_search_eval(
+    audio_filepaths,
+    model: nemo_asr.models.ASRModel,
+    cfg: EvalBeamSearchNGramConfig,
+    target_transcripts: List[str],
+    preds_output_file: str = None,
+    lm_path: str = None,
+    beam_alpha: float = 1.0,
+    beam_beta: float = 0.0,
+    beam_width: int = 128,
+    punctuation_capitalization: PunctuationCapitalization = None,
+):
+    level = logging.getEffectiveLevel()
+    logging.setLevel(logging.CRITICAL)
+    # Reset config
+    if isinstance(model, EncDecHybridRNNTCTCModel):
+        model.change_decoding_strategy(decoding_cfg=None, decoder_type="ctc")
+    else:
+        model.change_decoding_strategy(None)
+
+    # Override the beam search config with current search candidate configuration
+    cfg.decoding.beam_size = beam_width
+    cfg.decoding.ngram_lm_alpha = beam_alpha
+    cfg.decoding.beam_beta = beam_beta
+    cfg.decoding.return_best_hypothesis = False
+    cfg.decoding.ngram_lm_model = cfg.kenlm_model_file
+
+    # Update model's decoding strategy config
+    model.cfg.decoding.strategy = cfg.decoding_strategy
+    model.cfg.decoding.beam = cfg.decoding
+
+    # Update model's decoding strategy
+    if isinstance(model, EncDecHybridRNNTCTCModel):
+        model.change_decoding_strategy(model.cfg.decoding, decoder_type='ctc')
+    else:
+        model.change_decoding_strategy(model.cfg.decoding)
+    logging.setLevel(level)
+
+    all_hyps = model.transcribe(audio_filepaths, cfg.batch_size)
+
+    wer_dist_first = cer_dist_first = 0
+    wer_dist_best = cer_dist_best = 0
+    words_count = 0
+    chars_count = 0
+    if preds_output_file:
+        out_file = open(preds_output_file, 'w', encoding='utf_8', newline='\n')
+
+    for batch_idx, nbest_hyp in enumerate(all_hyps):
+        target = target_transcripts[batch_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 = float("inf")
+        for candidate_idx, candidate in enumerate(nbest_hyp):
+            pred_text = apply_text_processing(punctuation_capitalization, cfg, candidate.text)
+
+            pred_split_w = pred_text.split()
+            wer_dist = edit_distance(target_split_w, pred_split_w)['total']
+            pred_split_c = list(pred_text)
+            cer_dist = edit_distance(target_split_c, pred_split_c)['total']
+
+            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.score
+            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
+
+    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(
+        'Oracle WER/CER in candidates with perfect LM= {:.2%}/{:.2%}'.format(
+            wer_dist_best / words_count, cer_dist_best / chars_count
+        )
+    )
+
+    logging.info(f"=================================================================================")
+
+    return wer_dist_first / words_count, cer_dist_first / chars_count
+
+
+@hydra_runner(config_path=None, config_name='EvalBeamSearchNGramConfig', schema=EvalBeamSearchNGramConfig)
+def main(cfg: EvalBeamSearchNGramConfig):
+    if is_dataclass(cfg):
+        cfg = OmegaConf.structured(cfg)  # type: EvalBeamSearchNGramConfig
+
+    valid_decoding_modes = ["greedy", "beamsearch", "beamsearch_ngram"]
+    if cfg.decoding_mode not in valid_decoding_modes:
+        raise ValueError(
+            f"Given decoding_mode={cfg.decoding_mode} is invalid. Available options are :\n" f"{valid_decoding_modes}"
+        )
+
+    if cfg.nemo_model_file.endswith('.nemo'):
+        asr_model = nemo_asr.models.ASRModel.restore_from(cfg.nemo_model_file, map_location=torch.device(cfg.device))
+    else:
+        logging.warning(
+            "nemo_model_file does not end with .nemo, therefore trying to load a pretrained model with this name."
+        )
+        asr_model = nemo_asr.models.ASRModel.from_pretrained(
+            cfg.nemo_model_file, map_location=torch.device(cfg.device)
+        )
+
+    target_transcripts = []
+    manifest_dir = Path(cfg.input_manifest).parent
+    with open(cfg.input_manifest, 'r', encoding='utf_8') as manifest_file:
+        audio_file_paths = []
+        for line in tqdm(manifest_file, desc=f"Reading Manifest {cfg.input_manifest} ...", ncols=120):
+            data = json.loads(line)
+            audio_file = Path(data['audio_filepath'])
+            if not audio_file.is_file() and not audio_file.is_absolute():
+                audio_file = manifest_dir / audio_file
+            target_transcripts.append(data['text'])
+            audio_file_paths.append(str(audio_file.absolute()))
+
+    punctuation_capitalization = PunctuationCapitalization(cfg.text_processing.punctuation_marks)
+    target_transcripts = apply_text_processing(punctuation_capitalization, cfg, target_transcripts)
+
+    if cfg.hyps_cache_file and os.path.exists(cfg.hyps_cache_file):
+        logging.info(f"Found a cached file of hypotheses at '{cfg.hyps_cache_file}'.")
+        logging.info(f"Loading the cached file of hypotheses from '{cfg.hyps_cache_file}' ...")
+        with open(cfg.hyps_cache_file, 'rb') as probs_file:
+            all_hyps = msgpack.load(probs_file)
+
+        if len(all_hyps) != len(audio_file_paths):
+            raise ValueError(
+                f"The number of samples in the hypotheses file '{cfg.hyps_cache_file}' does not "
+                f"match the manifest file. You may need to delete the hypotheses cached file."
+            )
+    else:
+
+        with torch.amp.autocast(asr_model.device.type, enabled=cfg.use_amp):
+            with torch.no_grad():
+                if isinstance(asr_model, EncDecHybridRNNTCTCModel):
+                    asr_model.cur_decoder = 'ctc'
+                all_hyps = asr_model.transcribe(audio_file_paths, batch_size=cfg.batch_size)
+
+        if cfg.hyps_cache_file:
+            os.makedirs(os.path.split(cfg.hyps_cache_file)[0], exist_ok=True)
+            logging.info(f"Writing cached files of hypotheses at '{cfg.hyps_cache_file}'...")
+            with open(cfg.hyps_cache_file, 'wb') as f_dump:
+                msgpack.dump(all_hyps, f_dump)
+
+    wer_dist_greedy = 0
+    cer_dist_greedy = 0
+    words_count = 0
+    chars_count = 0
+    for batch_idx, hyp in enumerate(all_hyps):
+        pred_text = apply_text_processing(punctuation_capitalization, cfg, hyp.text)
+
+        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 = edit_distance(target_split_w, pred_split_w)['total']
+        cer_dist = edit_distance(target_split_c, pred_split_c)['total']
+
+        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))
+
+    asr_model = asr_model.to('cpu')
+
+    if cfg.decoding_mode == "beamsearch_ngram":
+        if not os.path.exists(cfg.kenlm_model_file):
+            raise FileNotFoundError(f"Could not find the KenLM model file '{cfg.kenlm_model_file}'.")
+        lm_path = cfg.kenlm_model_file
+    else:
+        lm_path = None
+
+    # 'greedy' decoding_mode would skip the beam search decoding
+    if cfg.decoding_mode in ["beamsearch_ngram", "beamsearch"]:
+        if cfg.beam_width is None or cfg.beam_alpha is None or cfg.beam_beta is None:
+            raise ValueError("beam_width, beam_alpha and beam_beta are needed to perform beam search decoding.")
+        params = {'beam_width': cfg.beam_width, 'beam_alpha': cfg.beam_alpha, 'beam_beta': cfg.beam_beta}
+        hp_grid = ParameterGrid(params)
+        hp_grid = list(hp_grid)
+
+        best_wer_beam_size, best_cer_beam_size = None, None
+        best_wer_alpha, best_cer_alpha = None, None
+        best_wer_beta, best_cer_beta = None, None
+        best_wer, best_cer = float("inf"), float("inf")
+
+        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 cfg.preds_output_folder and not os.path.exists(cfg.preds_output_folder):
+            os.mkdir(cfg.preds_output_folder)
+        for hp in hp_grid:
+            if cfg.preds_output_folder:
+                preds_output_file = os.path.join(
+                    cfg.preds_output_folder,
+                    f"preds_out_width{hp['beam_width']}_alpha{hp['beam_alpha']}_beta{hp['beam_beta']}.tsv",
+                )
+            else:
+                preds_output_file = None
+
+            candidate_wer, candidate_cer = beam_search_eval(
+                audio_file_paths,
+                asr_model,
+                cfg,
+                target_transcripts=target_transcripts,
+                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"],
+                punctuation_capitalization=punctuation_capitalization,
+            )
+
+            if candidate_cer < best_cer:
+                best_cer_beam_size, best_cer_alpha, best_cer_beta, best_cer = (
+                    hp["beam_width"],
+                    hp["beam_alpha"],
+                    hp["beam_beta"],
+                    candidate_cer,
+                )
+
+            if candidate_wer < best_wer:
+                best_wer_beam_size, best_wer_alpha, best_wer_beta, best_wer = (
+                    hp["beam_width"],
+                    hp["beam_alpha"],
+                    hp["beam_beta"],
+                    candidate_wer,
+                )
+
+        logging.info(
+            f'Best WER Candidate = {best_wer:.2%} :: Beam size = {best_wer_beam_size}, '
+            f'Beam alpha = {best_wer_alpha}, Beam beta = {best_wer_beta}'
+        )
+
+        logging.info(
+            f'Best CER Candidate = {best_cer:.2%} :: Beam size = {best_cer_beam_size}, '
+            f'Beam alpha = {best_cer_alpha}, Beam beta = {best_cer_beta}'
+        )
+        logging.info(f"=================================================================================")
+
+
+if __name__ == '__main__':
+    main()

scripts/asr_language_modeling/ngram_lm/eval_beamsearch_ngram_transducer.py

@@ -0,0 +1,444 @@
+# Copyright (c) 2022, 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 Transducer 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'.
+
+# Config Help
+
+To discover all arguments of the script, please run :
+python eval_beamsearch_ngram.py --help
+python eval_beamsearch_ngram.py --cfg job
+
+# USAGE
+
+python eval_beamsearch_ngram_transducer.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, separated with commas>] \
+           beam_alpha=[<list of the beam alphas, separated with commas>] \
+           preds_output_folder=<optional folder to store the predictions> \
+           probs_cache_file=null \
+           decoding_strategy=<greedy_batch or maes decoding>
+           maes_prefix_alpha=[<list of the maes prefix alphas, separated with commas>] \
+           maes_expansion_gamma=[<list of the maes expansion gammas, separated with commas>] \
+           hat_subtract_ilm=<in case of HAT model: subtract internal LM or not> \
+           hat_ilm_weight=[<in case of HAT model: list of the HAT internal LM weights, separated with commas>] \
+           ...
+
+
+# Grid Search for Hyper parameters
+
+For grid search, you can provide a list of arguments as follows -
+
+           beam_width=[4,8,16,....] \
+           beam_alpha=[-2.0,-1.0,...,1.0,2.0] \
+
+# 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
+import os
+import tempfile
+from dataclasses import dataclass, field, is_dataclass
+from pathlib import Path
+from typing import List, Optional
+
+import msgpack
+import numpy as np
+import torch
+from kaldialign import edit_distance
+from omegaconf import MISSING, OmegaConf
+from sklearn.model_selection import ParameterGrid
+from tqdm.auto import tqdm
+
+import nemo.collections.asr as nemo_asr
+from nemo.collections.asr.parts.submodules import rnnt_beam_decoding
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+
+# fmt: off
+
+
+@dataclass
+class EvalBeamSearchNGramConfig:
+    """
+    Evaluate an ASR model with beam search decoding and n-gram KenLM language model.
+    """
+    # # The path of the '.nemo' file of the ASR model or the name of a pretrained model (ngc / huggingface)
+    nemo_model_file: str = MISSING
+
+    # File paths
+    input_manifest: str = MISSING  # The manifest file of the evaluation set
+    kenlm_model_file: Optional[str] = None  # The path of the KenLM binary model file
+    preds_output_folder: Optional[str] = None  # The optional folder where the predictions are stored
+    probs_cache_file: Optional[str] = None  # The cache file for storing the logprobs of the model
+
+    # Parameters for inference
+    acoustic_batch_size: int = 128  # The batch size to calculate log probabilities
+    beam_batch_size: int = 128  # The batch size to be used for beam search decoding
+    device: str = "cuda"  # The device to load the model onto to calculate log probabilities
+    use_amp: bool = False  # Whether to use AMP if available to calculate log probabilities
+    num_workers: int = 1  # Number of workers for DataLoader
+
+    # The decoding scheme to be used for evaluation
+    decoding_strategy: str = "greedy_batch" # ["greedy_batch", "beam", "tsd", "alsd", "maes"]
+
+    # Beam Search hyperparameters
+    beam_width: List[int] = field(default_factory=lambda: [8])  # The width or list of the widths for the beam search decoding
+    beam_alpha: List[float] = field(default_factory=lambda: [0.2])  # The alpha parameter or list of the alphas for the beam search decoding
+
+    maes_prefix_alpha: List[int] = field(default_factory=lambda: [2])  # The maes_prefix_alpha or list of the maes_prefix_alpha for the maes decoding
+    maes_expansion_gamma: List[float] = field(default_factory=lambda: [2.3])  # The maes_expansion_gamma or list of the maes_expansion_gamma for the maes decoding
+
+    # HAT related parameters (only for internal lm subtraction)
+    hat_subtract_ilm: bool = False
+    hat_ilm_weight: List[float] = field(default_factory=lambda: [0.0])
+
+    decoding: rnnt_beam_decoding.BeamRNNTInferConfig = field(default_factory=lambda: rnnt_beam_decoding.BeamRNNTInferConfig(beam_size=128))
+
+
+# fmt: on
+
+
+def decoding_step(
+    model: nemo_asr.models.ASRModel,
+    cfg: EvalBeamSearchNGramConfig,
+    all_probs: List[torch.Tensor],
+    target_transcripts: List[str],
+    preds_output_file: str = None,
+    beam_batch_size: int = 128,
+    progress_bar: bool = True,
+):
+    level = logging.getEffectiveLevel()
+    logging.setLevel(logging.CRITICAL)
+    # Reset config
+    model.change_decoding_strategy(None)
+
+    cfg.decoding.hat_ilm_weight = cfg.decoding.hat_ilm_weight * cfg.hat_subtract_ilm
+    # Override the beam search config with current search candidate configuration
+    cfg.decoding.return_best_hypothesis = False
+    cfg.decoding.ngram_lm_model = cfg.kenlm_model_file
+    cfg.decoding.hat_subtract_ilm = cfg.hat_subtract_ilm
+
+    # Update model's decoding strategy config
+    model.cfg.decoding.strategy = cfg.decoding_strategy
+    model.cfg.decoding.beam = cfg.decoding
+
+    # Update model's decoding strategy
+    model.change_decoding_strategy(model.cfg.decoding)
+    logging.setLevel(level)
+
+    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', encoding='utf_8', newline='\n')
+
+    if progress_bar:
+        if cfg.decoding_strategy == "greedy_batch":
+            description = "Greedy_batch decoding.."
+        else:
+            description = f"{cfg.decoding_strategy} decoding with bw={cfg.decoding.beam_size}, ba={cfg.decoding.ngram_lm_alpha}, ma={cfg.decoding.maes_prefix_alpha}, mg={cfg.decoding.maes_expansion_gamma}, hat_ilmw={cfg.decoding.hat_ilm_weight}"
+        it = tqdm(range(int(np.ceil(len(all_probs) / beam_batch_size))), desc=description, ncols=120)
+    else:
+        it = range(int(np.ceil(len(all_probs) / beam_batch_size)))
+    for batch_idx in it:
+        # disabling type checking
+        probs_batch = all_probs[batch_idx * beam_batch_size : (batch_idx + 1) * beam_batch_size]
+        probs_lens = torch.tensor([prob.shape[-1] for prob in probs_batch])
+        with torch.no_grad():
+            packed_batch = torch.zeros(len(probs_batch), probs_batch[0].shape[0], max(probs_lens), device='cpu')
+
+            for prob_index in range(len(probs_batch)):
+                packed_batch[prob_index, :, : probs_lens[prob_index]] = torch.tensor(
+                    probs_batch[prob_index].unsqueeze(0), device=packed_batch.device, dtype=packed_batch.dtype
+                )
+            beams_batch = model.decoding.rnnt_decoder_predictions_tensor(
+                packed_batch,
+                probs_lens,
+                return_hypotheses=True,
+            )
+            best_hyp_batch = [dec_hyp[0] for dec_hyp in beams_batch]
+        if cfg.decoding_strategy == "greedy_batch":
+            beams_batch = [[x] for x in best_hyp_batch]
+
+        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):  # type: (int, rnnt_beam_decoding.rnnt_utils.Hypothesis)
+                pred_text = candidate.text
+                pred_split_w = pred_text.split()
+                wer_dist = edit_distance(target_split_w, pred_split_w)['total']
+                pred_split_c = list(pred_text)
+                cer_dist = edit_distance(target_split_c, pred_split_c)['total']
+
+                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.score
+                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 cfg.decoding_strategy == "greedy_batch":
+        return wer_dist_first / words_count, cer_dist_first / chars_count
+
+    if preds_output_file:
+        out_file.close()
+        logging.info(f"Stored the predictions of {cfg.decoding_strategy} decoding at '{preds_output_file}'.")
+
+    if cfg.decoding.ngram_lm_model:
+        logging.info(
+            f"WER/CER with {cfg.decoding_strategy} decoding and N-gram model = {wer_dist_first / words_count:.2%}/{cer_dist_first / chars_count:.2%}"
+        )
+    else:
+        logging.info(
+            f"WER/CER with {cfg.decoding_strategy} decoding = {wer_dist_first / words_count:.2%}/{cer_dist_first / chars_count:.2%}"
+        )
+    logging.info(
+        f"Oracle WER/CER in candidates with perfect LM= {wer_dist_best / words_count:.2%}/{cer_dist_best / chars_count:.2%}"
+    )
+    logging.info(f"=================================================================================")
+
+    return wer_dist_first / words_count, cer_dist_first / chars_count
+
+
+@hydra_runner(config_path=None, config_name='EvalBeamSearchNGramConfig', schema=EvalBeamSearchNGramConfig)
+def main(cfg: EvalBeamSearchNGramConfig):
+    if is_dataclass(cfg):
+        cfg = OmegaConf.structured(cfg)  # type: EvalBeamSearchNGramConfig
+
+    valid_decoding_strategis = ["greedy_batch", "beam", "tsd", "alsd", "maes"]
+    if cfg.decoding_strategy not in valid_decoding_strategis:
+        raise ValueError(
+            f"Given decoding_strategy={cfg.decoding_strategy} is invalid. Available options are :\n"
+            f"{valid_decoding_strategis}"
+        )
+
+    if cfg.nemo_model_file.endswith('.nemo'):
+        asr_model = nemo_asr.models.ASRModel.restore_from(cfg.nemo_model_file, map_location=torch.device(cfg.device))
+    else:
+        logging.warning(
+            "nemo_model_file does not end with .nemo, therefore trying to load a pretrained model with this name."
+        )
+        asr_model = nemo_asr.models.ASRModel.from_pretrained(
+            cfg.nemo_model_file, map_location=torch.device(cfg.device)
+        )
+
+    if cfg.kenlm_model_file:
+        if not os.path.exists(cfg.kenlm_model_file):
+            raise FileNotFoundError(f"Could not find the KenLM model file '{cfg.kenlm_model_file}'.")
+        if cfg.decoding_strategy != "maes":
+            raise ValueError(f"Decoding with kenlm model is supported only for maes decoding algorithm.")
+        lm_path = cfg.kenlm_model_file
+    else:
+        lm_path = None
+        cfg.beam_alpha = [0.0]
+    if cfg.hat_subtract_ilm:
+        assert lm_path, "kenlm must be set for hat internal lm subtraction"
+
+    if cfg.decoding_strategy != "maes":
+        cfg.maes_prefix_alpha, cfg.maes_expansion_gamma, cfg.hat_ilm_weight = [0], [0], [0]
+
+    target_transcripts = []
+    manifest_dir = Path(cfg.input_manifest).parent
+    with open(cfg.input_manifest, 'r', encoding='utf_8') as manifest_file:
+        audio_file_paths = []
+        for line in tqdm(manifest_file, desc=f"Reading Manifest {cfg.input_manifest} ...", ncols=120):
+            data = json.loads(line)
+            audio_file = Path(data['audio_filepath'])
+            if not audio_file.is_file() and not audio_file.is_absolute():
+                audio_file = manifest_dir / audio_file
+            target_transcripts.append(data['text'])
+            audio_file_paths.append(str(audio_file.absolute()))
+
+    if cfg.probs_cache_file and os.path.exists(cfg.probs_cache_file):
+        logging.info(f"Found a cached file of probabilities at '{cfg.probs_cache_file}'.")
+        logging.info(f"Loading the cached file of probabilities from '{cfg.probs_cache_file}' ...")
+        with open(cfg.probs_cache_file, 'rb') as probs_file:
+            all_probs = msgpack.load(probs_file)
+
+        if len(all_probs) != len(audio_file_paths):
+            raise ValueError(
+                f"The number of samples in the probabilities file '{cfg.probs_cache_file}' does not "
+                f"match the manifest file. You may need to delete the probabilities cached file."
+            )
+    else:
+
+        # manual calculation of encoder_embeddings
+        with torch.amp.autocast(asr_model.device.type, enabled=cfg.use_amp):
+            with torch.no_grad():
+                asr_model.eval()
+                asr_model.encoder.freeze()
+                device = next(asr_model.parameters()).device
+                all_probs = []
+                with tempfile.TemporaryDirectory() as tmpdir:
+                    with open(os.path.join(tmpdir, 'manifest.json'), 'w', encoding='utf-8') as fp:
+                        for audio_file in audio_file_paths:
+                            entry = {'audio_filepath': audio_file, 'duration': 100000, 'text': ''}
+                            fp.write(json.dumps(entry) + '\n')
+                    config = {
+                        'paths2audio_files': audio_file_paths,
+                        'batch_size': cfg.acoustic_batch_size,
+                        'temp_dir': tmpdir,
+                        'num_workers': cfg.num_workers,
+                        'channel_selector': None,
+                        'augmentor': None,
+                    }
+                    temporary_datalayer = asr_model._setup_transcribe_dataloader(config)
+                    for test_batch in tqdm(temporary_datalayer, desc="Transcribing", disable=True):
+                        encoded, encoded_len = asr_model.forward(
+                            input_signal=test_batch[0].to(device), input_signal_length=test_batch[1].to(device)
+                        )
+                        # dump encoder embeddings per file
+                        for idx in range(encoded.shape[0]):
+                            encoded_no_pad = encoded[idx, :, : encoded_len[idx]]
+                            all_probs.append(encoded_no_pad)
+
+        if cfg.probs_cache_file:
+            logging.info(f"Writing cached files of probabilities at '{cfg.probs_cache_file}'...")
+            with open(cfg.probs_cache_file, 'wb') as f_dump:
+                msgpack.dump(all_probs, f_dump)
+
+    if cfg.decoding_strategy == "greedy_batch":
+        asr_model = asr_model.to('cpu')
+        candidate_wer, candidate_cer = decoding_step(
+            asr_model,
+            cfg,
+            all_probs=all_probs,
+            target_transcripts=target_transcripts,
+            beam_batch_size=cfg.beam_batch_size,
+            progress_bar=True,
+        )
+        logging.info(f"Greedy batch WER/CER = {candidate_wer:.2%}/{candidate_cer:.2%}")
+
+    asr_model = asr_model.to('cpu')
+
+    # 'greedy_batch' decoding_strategy would skip the beam search decoding
+    if cfg.decoding_strategy in ["beam", "tsd", "alsd", "maes"]:
+        if cfg.beam_width is None or cfg.beam_alpha is None:
+            raise ValueError("beam_width and beam_alpha are needed to perform beam search decoding.")
+        params = {
+            'beam_width': cfg.beam_width,
+            'beam_alpha': cfg.beam_alpha,
+            'maes_prefix_alpha': cfg.maes_prefix_alpha,
+            'maes_expansion_gamma': cfg.maes_expansion_gamma,
+            'hat_ilm_weight': cfg.hat_ilm_weight,
+        }
+        hp_grid = ParameterGrid(params)
+        hp_grid = list(hp_grid)
+
+        best_wer_beam_size, best_cer_beam_size = None, None
+        best_wer_alpha, best_cer_alpha = None, None
+        best_wer, best_cer = 1e6, 1e6
+
+        logging.info(
+            f"==============================Starting the {cfg.decoding_strategy} decoding==============================="
+        )
+        logging.info(f"Grid search size: {len(hp_grid)}")
+        logging.info(f"It may take some time...")
+        logging.info(f"==============================================================================================")
+
+        if cfg.preds_output_folder and not os.path.exists(cfg.preds_output_folder):
+            os.mkdir(cfg.preds_output_folder)
+        for hp in hp_grid:
+            if cfg.preds_output_folder:
+                results_file = f"preds_out_{cfg.decoding_strategy}_bw{hp['beam_width']}"
+                if cfg.decoding_strategy == "maes":
+                    results_file = f"{results_file}_ma{hp['maes_prefix_alpha']}_mg{hp['maes_expansion_gamma']}"
+                    if cfg.kenlm_model_file:
+                        results_file = f"{results_file}_ba{hp['beam_alpha']}"
+                        if cfg.hat_subtract_ilm:
+                            results_file = f"{results_file}_hat_ilmw{hp['hat_ilm_weight']}"
+                preds_output_file = os.path.join(cfg.preds_output_folder, f"{results_file}.tsv")
+            else:
+                preds_output_file = None
+
+            cfg.decoding.beam_size = hp["beam_width"]
+            cfg.decoding.ngram_lm_alpha = hp["beam_alpha"]
+            cfg.decoding.maes_prefix_alpha = hp["maes_prefix_alpha"]
+            cfg.decoding.maes_expansion_gamma = hp["maes_expansion_gamma"]
+            cfg.decoding.hat_ilm_weight = hp["hat_ilm_weight"]
+
+            candidate_wer, candidate_cer = decoding_step(
+                asr_model,
+                cfg,
+                all_probs=all_probs,
+                target_transcripts=target_transcripts,
+                preds_output_file=preds_output_file,
+                beam_batch_size=cfg.beam_batch_size,
+                progress_bar=True,
+            )
+
+            if candidate_cer < best_cer:
+                best_cer_beam_size = hp["beam_width"]
+                best_cer_alpha = hp["beam_alpha"]
+                best_cer_ma = hp["maes_prefix_alpha"]
+                best_cer_mg = hp["maes_expansion_gamma"]
+                best_cer_hat_ilm_weight = hp["hat_ilm_weight"]
+                best_cer = candidate_cer
+
+            if candidate_wer < best_wer:
+                best_wer_beam_size = hp["beam_width"]
+                best_wer_alpha = hp["beam_alpha"]
+                best_wer_ma = hp["maes_prefix_alpha"]
+                best_wer_ga = hp["maes_expansion_gamma"]
+                best_wer_hat_ilm_weight = hp["hat_ilm_weight"]
+                best_wer = candidate_wer
+
+        wer_hat_parameter = ""
+        if cfg.hat_subtract_ilm:
+            wer_hat_parameter = f"HAT ilm weight = {best_wer_hat_ilm_weight}, "
+        logging.info(
+            f'Best WER Candidate = {best_wer:.2%} :: Beam size = {best_wer_beam_size}, '
+            f'Beam alpha = {best_wer_alpha}, {wer_hat_parameter}'
+            f'maes_prefix_alpha = {best_wer_ma}, maes_expansion_gamma = {best_wer_ga} '
+        )
+
+        cer_hat_parameter = ""
+        if cfg.hat_subtract_ilm:
+            cer_hat_parameter = f"HAT ilm weight = {best_cer_hat_ilm_weight}"
+        logging.info(
+            f'Best CER Candidate = {best_cer:.2%} :: Beam size = {best_cer_beam_size}, '
+            f'Beam alpha = {best_cer_alpha}, {cer_hat_parameter} '
+            f'maes_prefix_alpha = {best_cer_ma}, maes_expansion_gamma = {best_cer_mg}'
+        )
+        logging.info(f"=================================================================================")
+
+
+if __name__ == '__main__':
+    main()

scripts/asr_language_modeling/ngram_lm/eval_wfst_decoding_ctc.py

@@ -0,0 +1,424 @@
+# Copyright (c) 2025, 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 in ARPA format in
+# fusion with WFST decoders on top of a trained ASR model with CTC decoder.
+# NeMo's WFST decoders use WFST decoding graphs made from ARPA LMs
+# to find the best candidates. This script supports BPE level encodings only
+# and models which is detected automatically from the type of the model.
+# You may train the LM model with e.g. SRILM.
+
+# Config Help
+
+To discover all arguments of the script, please run :
+python eval_wfst_decoding_ctc.py --help
+python eval_wfst_decoding_ctc.py --cfg job
+
+# USAGE
+
+python eval_wfst_decoding_ctc.py nemo_model_file=<path to the .nemo file of the model> \
+           input_manifest=<path to the evaluation JSON manifest file> \
+           arpa_model_file=<path to the ARPA LM model> \
+           decoding_wfst_file=<path to the decoding WFST file> \
+           beam_width=[<list of the beam widths, separated with commas>] \
+           lm_weight=[<list of the LM weight multipliers, separated with commas>] \
+           decoding_mode=<decoding mode, affects output. Usually "nbest"> \
+           decoding_search_type=<a.k.a. decoding backend. Usually "riva"> \
+           open_vocabulary_decoding=<whether to use open vocabulary mode for WFST decoding> \
+           preds_output_folder=<optional folder to store the predictions> \
+           probs_cache_file=null
+           ...
+
+
+# Grid Search for Hyper parameters
+
+For grid search, you can provide a list of arguments as follows -
+
+           beam_width=[5.0,10.0,15.0,20.0] \
+           lm_weight=[0.1,0.5,0.6,0.7,0.8,0.9,1.0,1.1,1.2,1.3,1.4,1.5,2.0] \
+
+"""
+
+
+import contextlib
+import json
+import os
+from dataclasses import dataclass, field, is_dataclass
+from pathlib import Path
+from typing import List, Optional
+
+import msgpack
+import numpy as np
+import torch
+from kaldialign import edit_distance
+from omegaconf import MISSING, OmegaConf
+from sklearn.model_selection import ParameterGrid
+from tqdm.auto import tqdm
+
+import nemo.collections.asr as nemo_asr
+from nemo.collections.asr.models import EncDecHybridRNNTCTCModel
+from nemo.collections.asr.parts.submodules import ctc_beam_decoding
+from nemo.collections.asr.parts.utils.transcribe_utils import PunctuationCapitalization, TextProcessingConfig
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+
+# fmt: off
+
+
+@dataclass
+class EvalWFSTNGramConfig:
+    """
+    Evaluate an ASR model with WFST decoding and n-gram ARPA language model.
+    """
+    # # The path of the '.nemo' file of the ASR model or the name of a pretrained model (ngc / huggingface)
+    nemo_model_file: str = MISSING
+
+    # File paths
+    input_manifest: str = MISSING  # The manifest file of the evaluation set
+    arpa_model_file: Optional[str] = None  # The path of the ARPA model file
+    decoding_wfst_file: Optional[str] = None  # The path of the decoding WFST file
+    preds_output_folder: Optional[str] = None  # The optional folder where the predictions are stored
+    probs_cache_file: Optional[str] = None  # The cache file for storing the logprobs of the model
+
+    # Parameters for inference
+    acoustic_batch_size: int = 16  # The batch size to calculate log probabilities
+    beam_batch_size: int = 512  # The batch size to be used for beam search decoding
+    device: str = "cuda"  # The device to load the model onto to calculate log probabilities and run WFST decoding
+    use_amp: bool = False  # Whether to use AMP if available to calculate log probabilities
+
+    # WFST decoding hyperparameters
+
+    beam_width: List[float] = field(default_factory=lambda: [10])  # The width or list of the beam widths for the WFST decoding
+    lm_weight: List[float] = field(default_factory=lambda: [1.0])  # The language model weight parameter or list of parameters for the WFST decoding
+
+    open_vocabulary_decoding: bool = False  # Whether to use open vocabulary mode for WFST decoding
+    decoding_mode: str = "nbest"
+    decoding_search_type: str = "riva"
+    decoding: ctc_beam_decoding.WfstCTCInferConfig = field(
+        default_factory=lambda: ctc_beam_decoding.WfstCTCInferConfig(beam_size=1)
+    )
+
+    text_processing: Optional[TextProcessingConfig] = field(default_factory=lambda: TextProcessingConfig(
+        punctuation_marks = ".,?",
+        separate_punctuation = False,
+        do_lowercase = False,
+        rm_punctuation = False,
+    ))
+# fmt: on
+
+
+def beam_search_eval(
+    model: nemo_asr.models.ASRModel,
+    cfg: EvalWFSTNGramConfig,
+    all_probs: List[torch.Tensor],
+    target_transcripts: List[str],
+    preds_output_file: str = None,
+    lm_weight: float = 1.0,
+    beam_width: float = 10.0,
+    beam_batch_size: int = 512,
+    progress_bar: bool = True,
+    punctuation_capitalization: PunctuationCapitalization = None,
+):
+    level = logging.getEffectiveLevel()
+    logging.setLevel(logging.CRITICAL)
+    # Reset config
+    if isinstance(model, EncDecHybridRNNTCTCModel):
+        model.change_decoding_strategy(decoding_cfg=None, decoder_type="ctc")
+    else:
+        model.change_decoding_strategy(None)
+
+    # Override the beam search config with current search candidate configuration
+    cfg.decoding.beam_width = beam_width
+    cfg.decoding.lm_weight = lm_weight
+    cfg.decoding.open_vocabulary_decoding = cfg.open_vocabulary_decoding
+    cfg.decoding.return_best_hypothesis = False
+    cfg.decoding.arpa_lm_path = cfg.arpa_model_file
+    cfg.decoding.wfst_lm_path = cfg.decoding_wfst_file
+    cfg.decoding.device = cfg.device
+    cfg.decoding.decoding_mode = cfg.decoding_mode
+    cfg.decoding.search_type = cfg.decoding_search_type
+
+    # Update model's decoding strategy config
+    model.cfg.decoding.strategy = "wfst"
+    model.cfg.decoding.wfst = cfg.decoding
+
+    # Update model's decoding strategy
+    if isinstance(model, EncDecHybridRNNTCTCModel):
+        model.change_decoding_strategy(model.cfg.decoding, decoder_type='ctc')
+        decoding = model.ctc_decoding
+    else:
+        model.change_decoding_strategy(model.cfg.decoding)
+        decoding = model.decoding
+    logging.setLevel(level)
+
+    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', encoding='utf_8', newline='\n')
+
+    if progress_bar:
+        it = tqdm(
+            range(int(np.ceil(len(all_probs) / beam_batch_size))),
+            desc=f"Beam search decoding with width={beam_width}, lm_weight={lm_weight}",
+            ncols=120,
+        )
+    else:
+        it = range(int(np.ceil(len(all_probs) / beam_batch_size)))
+    for batch_idx in it:
+        # disabling type checking
+        probs_batch = all_probs[batch_idx * beam_batch_size : (batch_idx + 1) * beam_batch_size]
+        probs_lens = torch.tensor([prob.shape[0] for prob in probs_batch])
+        with torch.no_grad():
+            packed_batch = torch.zeros(len(probs_batch), max(probs_lens), probs_batch[0].shape[-1], device='cpu')
+
+            for prob_index in range(len(probs_batch)):
+                packed_batch[prob_index, : probs_lens[prob_index], :] = probs_batch[prob_index].to(
+                    device=packed_batch.device, dtype=packed_batch.dtype
+                )
+
+            _, beams_batch = decoding.ctc_decoder_predictions_tensor(
+                packed_batch,
+                decoder_lengths=probs_lens,
+                return_hypotheses=True,
+            )
+
+        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):  # type: (int, ctc_beam_decoding.rnnt_utils.Hypothesis)
+                pred_text = candidate.text
+                if cfg.text_processing.do_lowercase:
+                    pred_text = punctuation_capitalization.do_lowercase([pred_text])[0]
+                if cfg.text_processing.rm_punctuation:
+                    pred_text = punctuation_capitalization.rm_punctuation([pred_text])[0]
+                if cfg.text_processing.separate_punctuation:
+                    pred_text = punctuation_capitalization.separate_punctuation([pred_text])[0]
+                pred_split_w = pred_text.split()
+                wer_dist = edit_distance(target_split_w, pred_split_w)['total']
+                pred_split_c = list(pred_text)
+                cer_dist = edit_distance(target_split_c, pred_split_c)['total']
+
+                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.score
+                if preds_output_file:
+                    out_file.write(f'{pred_text}\t{score}\n')
+            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}'.")
+
+    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
+        )
+    )
+    logging.info(
+        'Oracle WER/CER in candidates with perfect LM= {:.2%}/{:.2%}'.format(
+            wer_dist_best / words_count, cer_dist_best / chars_count
+        )
+    )
+    logging.info(f"=================================================================================")
+
+    return wer_dist_first / words_count, cer_dist_first / chars_count
+
+
+@hydra_runner(config_path=None, config_name='EvalWFSTNGramConfig', schema=EvalWFSTNGramConfig)
+def main(cfg: EvalWFSTNGramConfig):
+    if is_dataclass(cfg):
+        cfg = OmegaConf.structured(cfg)  # type: EvalWFSTNGramConfig
+
+    if cfg.nemo_model_file.endswith('.nemo'):
+        asr_model = nemo_asr.models.ASRModel.restore_from(cfg.nemo_model_file, map_location=torch.device(cfg.device))
+    else:
+        logging.warning(
+            "nemo_model_file does not end with .nemo, therefore trying to load a pretrained model with this name."
+        )
+        asr_model = nemo_asr.models.ASRModel.from_pretrained(
+            cfg.nemo_model_file, map_location=torch.device(cfg.device)
+        )
+
+    target_transcripts = []
+    manifest_dir = Path(cfg.input_manifest).parent
+    with open(cfg.input_manifest, 'r', encoding='utf_8') as manifest_file:
+        audio_file_paths = []
+        for line in tqdm(manifest_file, desc=f"Reading Manifest {cfg.input_manifest} ...", ncols=120):
+            data = json.loads(line)
+            audio_file = Path(data['audio_filepath'])
+            if not audio_file.is_file() and not audio_file.is_absolute():
+                audio_file = manifest_dir / audio_file
+            target_transcripts.append(data['text'])
+            audio_file_paths.append(str(audio_file.absolute()))
+
+    punctuation_capitalization = PunctuationCapitalization(cfg.text_processing.punctuation_marks)
+    if cfg.text_processing.do_lowercase:
+        target_transcripts = punctuation_capitalization.do_lowercase(target_transcripts)
+    if cfg.text_processing.rm_punctuation:
+        target_transcripts = punctuation_capitalization.rm_punctuation(target_transcripts)
+    if cfg.text_processing.separate_punctuation:
+        target_transcripts = punctuation_capitalization.separate_punctuation(target_transcripts)
+
+    if cfg.probs_cache_file and os.path.exists(cfg.probs_cache_file):
+        logging.info(f"Found a cached file of probabilities at '{cfg.probs_cache_file}'.")
+        logging.info(f"Loading the cached file of probabilities from '{cfg.probs_cache_file}' ...")
+        with open(cfg.probs_cache_file, 'rb') as probs_file:
+            all_probs = msgpack.load(probs_file)
+
+        if len(all_probs) != len(audio_file_paths):
+            raise ValueError(
+                f"The number of samples in the probabilities file '{cfg.probs_cache_file}' does not "
+                f"match the manifest file. You may need to delete the probabilities cached file."
+            )
+    else:
+
+        with torch.amp.autocast(asr_model.device.type, enabled=cfg.use_amp):
+            with torch.no_grad():
+                if isinstance(asr_model, EncDecHybridRNNTCTCModel):
+                    asr_model.cur_decoder = 'ctc'
+                all_hyps = asr_model.transcribe(
+                    audio_file_paths, batch_size=cfg.acoustic_batch_size, return_hypotheses=True
+                )
+                all_logits = [h.y_sequence for h in all_hyps]
+
+        all_probs = all_logits
+        if cfg.probs_cache_file:
+            os.makedirs(os.path.split(cfg.probs_cache_file)[0], exist_ok=True)
+            logging.info(f"Writing cached files of probabilities at '{cfg.probs_cache_file}'...")
+            with open(cfg.probs_cache_file, 'wb') as f_dump:
+                msgpack.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 = preds.to(device='cpu').unsqueeze(0)
+        preds_lens = torch.tensor([preds_tensor.shape[1]], device='cpu')
+        if isinstance(asr_model, EncDecHybridRNNTCTCModel):
+            pred_text = asr_model.ctc_decoding.ctc_decoder_predictions_tensor(preds_tensor, preds_lens)[0]
+        else:
+            pred_text = asr_model.decoding.ctc_decoder_predictions_tensor(preds_tensor, preds_lens)[0]
+
+        if cfg.text_processing.do_lowercase:
+            pred_text = punctuation_capitalization.do_lowercase([pred_text])[0]
+        if cfg.text_processing.rm_punctuation:
+            pred_text = punctuation_capitalization.rm_punctuation([pred_text])[0]
+        if cfg.text_processing.separate_punctuation:
+            pred_text = punctuation_capitalization.separate_punctuation([pred_text])[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 = edit_distance(target_split_w, pred_split_w)['total']
+        cer_dist = edit_distance(target_split_c, pred_split_c)['total']
+
+        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))
+
+    asr_model = asr_model.to('cpu')
+
+    if (cfg.arpa_model_file is None or not os.path.exists(cfg.arpa_model_file)) and (
+        cfg.decoding_wfst_file is None or not os.path.exists(cfg.decoding_wfst_file)
+    ):
+        raise FileNotFoundError(
+            f"Could not find both the ARPA model file `{cfg.arpa_model_file}` "
+            f"and the decoding WFST file `{cfg.decoding_wfst_file}`."
+        )
+
+    if cfg.beam_width is None or cfg.lm_weight is None:
+        raise ValueError("beam_width and lm_weight are needed to perform WFST decoding.")
+    params = {'beam_width': cfg.beam_width, 'lm_weight': cfg.lm_weight}
+    hp_grid = ParameterGrid(params)
+    hp_grid = list(hp_grid)
+
+    best_wer_beam_width, best_cer_beam_width = None, None
+    best_wer_lm_weight, best_cer_lm_weight = None, None
+    best_wer, best_cer = 1e6, 1e6
+
+    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 cfg.preds_output_folder and not os.path.exists(cfg.preds_output_folder):
+        os.mkdir(cfg.preds_output_folder)
+    for hp in hp_grid:
+        if cfg.preds_output_folder:
+            preds_output_file = os.path.join(
+                cfg.preds_output_folder,
+                f"preds_out_beam_width{hp['beam_width']}_lm_weight{hp['lm_weight']}.tsv",
+            )
+        else:
+            preds_output_file = None
+
+        candidate_wer, candidate_cer = beam_search_eval(
+            asr_model,
+            cfg,
+            all_probs=all_probs,
+            target_transcripts=target_transcripts,
+            preds_output_file=preds_output_file,
+            beam_width=hp["beam_width"],
+            lm_weight=hp["lm_weight"],
+            beam_batch_size=cfg.beam_batch_size,
+            progress_bar=True,
+            punctuation_capitalization=punctuation_capitalization,
+        )
+
+        if candidate_cer < best_cer:
+            best_cer_beam_width = hp["beam_width"]
+            best_cer_lm_weight = hp["lm_weight"]
+            best_cer = candidate_cer
+
+        if candidate_wer < best_wer:
+            best_wer_beam_width = hp["beam_width"]
+            best_wer_lm_weight = hp["lm_weight"]
+            best_wer = candidate_wer
+
+    logging.info(
+        f'Best WER Candidate = {best_wer:.2%} :: Beam size = {best_wer_beam_width}, LM weight = {best_wer_lm_weight}'
+    )
+
+    logging.info(
+        f'Best CER Candidate = {best_cer:.2%} :: Beam size = {best_cer_beam_width}, LM weight = {best_cer_lm_weight}'
+    )
+    logging.info(f"=================================================================================")
+
+
+if __name__ == '__main__':
+    main()

scripts/asr_language_modeling/ngram_lm/install_beamsearch_decoders.sh

@@ -0,0 +1,69 @@
+#!/usr/bin/env bash
+# Copyright (c) 2022, 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.
+
+# Use this script to install KenLM, OpenSeq2Seq decoder, Flashlight decoder
+shopt -s expand_aliases
+
+NEMO_PATH=/workspace/nemo  # Path to NeMo folder: /workspace/nemo if you use NeMo/Dockerfile
+if [ "$#" -eq 1 ]; then
+  NEMO_PATH=$1
+fi
+KENLM_MAX_ORDER=10 # Maximum order of KenLM model, also specified in the setup_os2s_decoders.py
+
+if [ -d "$NEMO_PATH" ]; then
+  echo "The folder '$NEMO_PATH' exists."
+else
+  echo "Error: The folder '$NEMO_PATH' does not exist. Specify it as a first command line positional argument!"
+  exit 1
+fi
+cd $NEMO_PATH
+
+if [ $(id -u) -eq 0 ]; then
+  alias aptupdate='apt-get update'
+  alias b2install='./b2'
+else
+  alias aptupdate='sudo apt-get update'
+  alias b2install='sudo ./b2'
+fi
+
+aptupdate && apt-get upgrade -y && apt-get install -y swig liblzma-dev && rm -rf /var/lib/apt/lists/* # liblzma needed for flashlight decoder
+
+# install Boost package for KenLM
+wget https://archives.boost.io/release/1.80.0/source/boost_1_80_0.tar.bz2 && tar --bzip2 -xf $NEMO_PATH/boost_1_80_0.tar.bz2 && cd boost_1_80_0 && ./bootstrap.sh && b2install --layout=tagged link=static,shared threading=multi,single install -j4 && cd .. || echo FAILURE
+export BOOST_ROOT=$NEMO_PATH/boost_1_80_0
+
+git clone https://github.com/NVIDIA/OpenSeq2Seq
+cd OpenSeq2Seq
+git checkout ctc-decoders
+cd ..
+mv OpenSeq2Seq/decoders $NEMO_PATH/
+rm -rf OpenSeq2Seq
+cd $NEMO_PATH/decoders
+cp $NEMO_PATH/scripts/installers/setup_os2s_decoders.py ./setup.py
+./setup.sh
+
+# install KenLM
+cd $NEMO_PATH/decoders/kenlm/build && cmake -DKENLM_MAX_ORDER=$KENLM_MAX_ORDER .. && make -j2
+cd $NEMO_PATH/decoders/kenlm
+python setup.py install --max_order=$KENLM_MAX_ORDER
+export KENLM_LIB=$NEMO_PATH/decoders/kenlm/build/bin
+export KENLM_ROOT=$NEMO_PATH/decoders/kenlm
+cd ..
+
+# install Flashlight
+git clone https://github.com/flashlight/text && cd text
+python setup.py bdist_wheel
+pip install dist/*.whl
+cd ..

scripts/asr_language_modeling/ngram_lm/make_phone_lm.py

@@ -0,0 +1,944 @@
+#!/usr/bin/env python
+
+# Copyright (c) 2022, 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.
+
+# Copyright (c) 2016, Johns Hopkins University (Author: Daniel Povey).
+# 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 was copied from https://github.com/kaldi-asr/kaldi/blob/master/egs/wsj/s5/utils/lang/make_phone_lm.py
+# with minor python3 related changes.
+
+import argparse
+import math
+import sys
+from collections import defaultdict
+
+# note, this was originally based
+
+parser = argparse.ArgumentParser(
+    description="""
+This script creates a language model that's intended to be used in modeling
+phone sequences (either of sentences or of dictionary entries), although of
+course it will work for any type of data.  The easiest way
+to describe it is as a a Kneser-Ney language model (unmodified, with addition)
+with a fixed discounting constant equal to 1, except with no smoothing of the
+bigrams (and hence no unigram state).  This is (a) because we want to keep the
+graph after context expansion small, (b) because languages tend to have
+constraints on which phones can follow each other, and (c) in order to get valid
+sequences of word-position-dependent phones so that lattice-align-words can
+work.  It also includes have a special entropy-based pruning technique that
+backs off the statistics of pruned n-grams to lower-order states.
+This script reads lines from its standard input, each
+consisting of a sequence of integer symbol-ids (which should be > 0),
+representing the phone sequences of a sentence or dictionary entry.
+This script outputs a backoff language model in FST format""",
+    epilog="See also utils/lang/make_phone_bigram_lang.sh",
+)
+
+
+parser.add_argument(
+    "--phone-disambig-symbol",
+    type=int,
+    required=False,
+    help="Integer corresponding to an otherwise-unused "
+    "phone-level disambiguation symbol (e.g. #5).  This is "
+    "inserted at the beginning of the phone sequence and "
+    "whenever we back off.",
+)
+parser.add_argument(
+    "--ngram-order",
+    type=int,
+    default=4,
+    choices=[2, 3, 4, 5, 6, 7],
+    help="Order of n-gram to use (but see also --num-extra-states;" "the effective order after pruning may be less.",
+)
+parser.add_argument(
+    "--num-extra-ngrams",
+    type=int,
+    default=20000,
+    help="Target number of n-grams in addition to the n-grams in "
+    "the bigram LM states which can't be pruned away.  n-grams "
+    "will be pruned to reach this target.",
+)
+parser.add_argument(
+    "--no-backoff-ngram-order",
+    type=int,
+    default=2,
+    choices=[1, 2, 3, 4, 5],
+    help="This specifies the n-gram order at which (and below which) "
+    "no backoff or pruning should be done.  This is expected to normally "
+    "be bigram, but for testing purposes you may want to set it to "
+    "1.",
+)
+parser.add_argument(
+    "--print-as-arpa",
+    type=str,
+    default="false",
+    choices=["true", "false"],
+    help="If true, print LM in ARPA format (default is to print "
+    "as FST).  You must also set --no-backoff-ngram-order=1 or "
+    "this is not allowed.",
+)
+parser.add_argument("--verbose", type=int, default=0, choices=[0, 1, 2, 3, 4, 5], help="Verbose level")
+
+args = parser.parse_args()
+
+if args.verbose >= 1:
+    print(" ".join(sys.argv), file=sys.stderr)
+
+
+class CountsForHistory(object):
+    ## This class (which is more like a struct) stores the counts seen in a
+    ## particular history-state.  It is used inside class NgramCounts.
+    ## It really does the job of a dict from int to float, but it also
+    ## keeps track of the total count.
+    def __init__(self):
+        # The 'lambda: defaultdict(float)' is an anonymous function taking no
+        # arguments that returns a new defaultdict(float).
+        self.word_to_count = defaultdict(int)
+        self.total_count = 0
+
+    def Words(self):
+        return list(self.word_to_count.keys())
+
+    def __str__(self):
+        # e.g. returns ' total=12 3->4 4->6 -1->2'
+        return " total={0} {1}".format(
+            str(self.total_count),
+            " ".join(["{0} -> {1}".format(word, count) for word, count in self.word_to_count.items()]),
+        )
+
+    ## Adds a certain count (expected to be integer, but might be negative).  If
+    ## the resulting count for this word is zero, removes the dict entry from
+    ## word_to_count.
+    ## [note, though, that in some circumstances we 'add back' zero counts
+    ## where the presence of n-grams would be structurally required by the arpa,
+    ## specifically if a higher-order history state has a nonzero count,
+    ## we need to structurally have the count there in the states it backs
+    ## off to.
+    def AddCount(self, predicted_word, count):
+        self.total_count += count
+        assert self.total_count >= 0
+        old_count = self.word_to_count[predicted_word]
+        new_count = old_count + count
+        if new_count < 0:
+            print("predicted-word={0}, old-count={1}, count={2}".format(predicted_word, old_count, count))
+        assert new_count >= 0
+        if new_count == 0:
+            del self.word_to_count[predicted_word]
+        else:
+            self.word_to_count[predicted_word] = new_count
+
+
+class NgramCounts(object):
+    ## A note on data-structure.  Firstly, all words are represented as
+    ## integers.  We store n-gram counts as an array, indexed by (history-length
+    ## == n-gram order minus one) (note: python calls arrays "lists") of dicts
+    ## from histories to counts, where histories are arrays of integers and
+    ## "counts" are dicts from integer to float.  For instance, when
+    ## accumulating the 4-gram count for the '8' in the sequence '5 6 7 8', we'd
+    ## do as follows: self.counts[3][[5,6,7]][8] += 1.0 where the [3] indexes an
+    ## array, the [[5,6,7]] indexes a dict, and the [8] indexes a dict.
+    def __init__(self, ngram_order):
+        assert ngram_order >= 2
+        # Integerized counts will never contain negative numbers, so
+        # inside this program, we use -3 and -2 for the BOS and EOS symbols
+        # respectively.
+        # Note: it's actually important that the bos-symbol is the most negative;
+        # it helps ensure that we print the state with left-context <s> first
+        # when we print the FST, and this means that the start-state will have
+        # the correct value.
+        self.bos_symbol = -3
+        self.eos_symbol = -2
+        # backoff_symbol is kind of a pseudo-word, it's used in keeping track of
+        # the backoff counts in each state.
+        self.backoff_symbol = -1
+        self.total_num_words = 0  # count includes EOS but not BOS.
+        self.counts = []
+        for n in range(ngram_order):
+            self.counts.append(defaultdict(lambda: CountsForHistory()))
+
+    # adds a raw count (called while processing input data).
+    # Suppose we see the sequence '6 7 8 9' and ngram_order=4, 'history'
+    # would be (6,7,8) and 'predicted_word' would be 9; 'count' would be
+    # 1.
+    def AddCount(self, history, predicted_word, count):
+        self.counts[len(history)][history].AddCount(predicted_word, count)
+
+    # 'line' is a string containing a sequence of integer word-ids.
+    # This function adds the un-smoothed counts from this line of text.
+    def AddRawCountsFromLine(self, line):
+        try:
+            words = [self.bos_symbol] + [int(x) for x in line.split()] + [self.eos_symbol]
+        except Exception:
+            sys.exit("make_phone_lm.py: bad input line {0} (expected a sequence " "of integers)".format(line))
+
+        for n in range(1, len(words)):
+            predicted_word = words[n]
+            history_start = max(0, n + 1 - args.ngram_order)
+            history = tuple(words[history_start:n])
+            self.AddCount(history, predicted_word, 1)
+            self.total_num_words += 1
+
+    def AddRawCountsFromStandardInput(self):
+        lines_processed = 0
+        while True:
+            line = sys.stdin.readline()
+            if line == "":
+                break
+            self.AddRawCountsFromLine(line)
+            lines_processed += 1
+        if lines_processed == 0 or args.verbose > 0:
+            print(
+                "make_phone_lm.py: processed {0} lines of input".format(lines_processed),
+                file=sys.stderr,
+            )
+
+    # This backs off the counts by subtracting 1 and assigning the subtracted
+    # count to the backoff state.  It's like a special case of Kneser-Ney with D
+    # = 1.  The optimal D would likely be something like 0.9, but we plan to
+    # later do entropy-pruning, and the remaining small counts of 0.1 would
+    # essentially all get pruned away anyway, so we don't lose much by doing it
+    # like this.
+    def ApplyBackoff(self):
+        # note: in the normal case where args.no_backoff_ngram_order == 2 we
+        # don't do backoff for history-length = 1 (i.e. for bigrams)... this is
+        # a kind of special LM where we're not going to back off to unigram,
+        # there will be no unigram.
+        if args.verbose >= 1:
+            initial_num_ngrams = self.GetNumNgrams()
+        for n in reversed(list(range(args.no_backoff_ngram_order, args.ngram_order))):
+            this_order_counts = self.counts[n]
+            for hist, counts_for_hist in this_order_counts.items():
+                backoff_hist = hist[1:]
+                backoff_counts_for_hist = self.counts[n - 1][backoff_hist]
+                this_discount_total = 0
+                for word in counts_for_hist.Words():
+                    counts_for_hist.AddCount(word, -1)
+                    # You can interpret the following line as incrementing the
+                    # count-of-counts for the next-lower order.  Note, however,
+                    # that later when we remove n-grams, we'll also add their
+                    # counts to the next-lower-order history state, so the
+                    # resulting counts won't strictly speaking be
+                    # counts-of-counts.
+                    backoff_counts_for_hist.AddCount(word, 1)
+                    this_discount_total += 1
+                counts_for_hist.AddCount(self.backoff_symbol, this_discount_total)
+
+        if args.verbose >= 1:
+            # Note: because D == 1, we completely back off singletons.
+            print(
+                "make_phone_lm.py: ApplyBackoff() reduced the num-ngrams from "
+                "{0} to {1}".format(initial_num_ngrams, self.GetNumNgrams()),
+                file=sys.stderr,
+            )
+
+    # This function prints out to stderr the n-gram counts stored in this
+    # object; it's used for debugging.
+    def Print(self, info_string):
+        print(info_string, file=sys.stderr)
+        # these are useful for debug.
+        total = 0.0
+        total_excluding_backoff = 0.0
+        for this_order_counts in self.counts:
+            for hist, counts_for_hist in this_order_counts.items():
+                print(str(hist) + str(counts_for_hist), file=sys.stderr)
+                total += counts_for_hist.total_count
+                total_excluding_backoff += counts_for_hist.total_count
+                if self.backoff_symbol in counts_for_hist.word_to_count:
+                    total_excluding_backoff -= counts_for_hist.word_to_count[self.backoff_symbol]
+        print(
+            "total count = {0}, excluding backoff = {1}".format(total, total_excluding_backoff),
+            file=sys.stderr,
+        )
+
+    def GetHistToStateMap(self):
+        # This function, called from PrintAsFst, returns a map from
+        # history to integer FST-state.
+        hist_to_state = dict()
+        fst_state_counter = 0
+        for n in range(0, args.ngram_order):
+            for hist in self.counts[n].keys():
+                hist_to_state[hist] = fst_state_counter
+                fst_state_counter += 1
+        return hist_to_state
+
+    # Returns the probability of word 'word' in history-state 'hist'.
+    # If 'word' is self.backoff_symbol, returns the backoff prob
+    # of this history-state.
+    # Returns None if there is no such word in this history-state, or this
+    # history-state does not exist.
+    def GetProb(self, hist, word):
+        if len(hist) >= args.ngram_order or not hist in self.counts[len(hist)]:
+            return None
+        counts_for_hist = self.counts[len(hist)][hist]
+        total_count = float(counts_for_hist.total_count)
+        if not word in counts_for_hist.word_to_count:
+            print(
+                "make_phone_lm.py: no prob for {0} -> {1} " "[no such count]".format(hist, word),
+                file=sys.stderr,
+            )
+            return None
+        prob = float(counts_for_hist.word_to_count[word]) / total_count
+        if len(hist) > 0 and word != self.backoff_symbol and self.backoff_symbol in counts_for_hist.word_to_count:
+            prob_in_backoff = self.GetProb(hist[1:], word)
+            backoff_prob = float(counts_for_hist.word_to_count[self.backoff_symbol]) / total_count
+            try:
+                prob += backoff_prob * prob_in_backoff
+            except Exception:
+                sys.exit("problem, hist is {0}, word is {1}".format(hist, word))
+        return prob
+
+    def PruneEmptyStates(self):
+        # Removes history-states that have no counts.
+
+        # It's possible in principle for history-states to have no counts and
+        # yet they cannot be pruned away because a higher-order version of the
+        # state exists with nonzero counts, so we have to keep track of this.
+        protected_histories = set()
+
+        states_removed_per_hist_len = [0] * args.ngram_order
+
+        for n in reversed(list(range(args.no_backoff_ngram_order, args.ngram_order))):
+            num_states_removed = 0
+            for hist, counts_for_hist in self.counts[n].items():
+                l = len(counts_for_hist.word_to_count)
+                assert l > 0 and self.backoff_symbol in counts_for_hist.word_to_count
+                if l == 1 and not hist in protected_histories:  # only the backoff symbol has a count.
+                    del self.counts[n][hist]
+                    num_states_removed += 1
+                else:
+                    # if this state was not pruned away, then the state that
+                    # it backs off to may not be pruned away either.
+                    backoff_hist = hist[1:]
+                    protected_histories.add(backoff_hist)
+            states_removed_per_hist_len[n] = num_states_removed
+        if args.verbose >= 1:
+            print(
+                "make_phone_lm.py: in PruneEmptyStates(), num states removed for "
+                "each history-length was: " + str(states_removed_per_hist_len),
+                file=sys.stderr,
+            )
+
+    def EnsureStructurallyNeededNgramsExist(self):
+        # makes sure that if an n-gram like (6, 7, 8) -> 9 exists,
+        # then counts exist for (7, 8) -> 9 and (8,) -> 9.  It does so
+        # by adding zero counts where such counts were absent.
+        # [note: () -> 9 is guaranteed anyway by the backoff method, if
+        # we have a unigram state].
+        if args.verbose >= 1:
+            num_ngrams_initial = self.GetNumNgrams()
+        for n in reversed(list(range(args.no_backoff_ngram_order, args.ngram_order))):
+
+            for hist, counts_for_hist in self.counts[n].items():
+                # This loop ensures that if we have an n-gram like (6, 7, 8) -> 9,
+                # then, say, (7, 8) -> 9 and (8) -> 9 exist.
+                reduced_hist = hist
+                for m in reversed(list(range(args.no_backoff_ngram_order, n))):
+                    reduced_hist = reduced_hist[1:]  # shift an element off
+                    # the history.
+                    counts_for_backoff_hist = self.counts[m][reduced_hist]
+                    for word in counts_for_hist.word_to_count.keys():
+                        counts_for_backoff_hist.word_to_count[word] += 0
+                # This loop ensures that if we have an n-gram like (6, 7, 8) -> 9,
+                # then, say, (6, 7) -> 8 and (6) -> 7 exist.  This will be needed
+                # for FST representations of the ARPA LM.
+                reduced_hist = hist
+                for m in reversed(list(range(args.no_backoff_ngram_order, n))):
+                    this_word = reduced_hist[-1]
+                    reduced_hist = reduced_hist[:-1]  # pop an element off the
+                    # history
+                    counts_for_backoff_hist = self.counts[m][reduced_hist]
+                    counts_for_backoff_hist.word_to_count[this_word] += 0
+        if args.verbose >= 1:
+            print(
+                "make_phone_lm.py: in EnsureStructurallyNeededNgramsExist(), "
+                "added {0} n-grams".format(self.GetNumNgrams() - num_ngrams_initial),
+                file=sys.stderr,
+            )
+
+    # This function prints the estimated language model as an FST.
+    def PrintAsFst(self, word_disambig_symbol):
+        # n is the history-length (== order + 1).  We iterate over the
+        # history-length in the order 1, 0, 2, 3, and then iterate over the
+        # histories of each order in sorted order.  Putting order 1 first
+        # and sorting on the histories
+        # ensures that the bigram state with <s> as the left context comes first.
+        # (note: self.bos_symbol is the most negative symbol)
+
+        # History will map from history (as a tuple) to integer FST-state.
+        hist_to_state = self.GetHistToStateMap()
+
+        for n in [1, 0] + list(range(2, args.ngram_order)):
+            this_order_counts = self.counts[n]
+            # For order 1, make sure the keys are sorted.
+            keys = this_order_counts.keys() if n != 1 else sorted(this_order_counts.keys())
+            for hist in keys:
+                word_to_count = this_order_counts[hist].word_to_count
+                this_fst_state = hist_to_state[hist]
+
+                for word in word_to_count.keys():
+                    # work out this_cost.  Costs in OpenFst are negative logs.
+                    this_cost = -math.log(self.GetProb(hist, word))
+
+                    if word > 0:  # a real word.
+                        next_hist = hist + (word,)  # appending tuples
+                        while not next_hist in hist_to_state:
+                            next_hist = next_hist[1:]
+                        next_fst_state = hist_to_state[next_hist]
+                        print(this_fst_state, next_fst_state, word, word, this_cost)
+                    elif word == self.eos_symbol:
+                        # print final-prob for this state.
+                        print(this_fst_state, this_cost)
+                    else:
+                        assert word == self.backoff_symbol
+                        backoff_fst_state = hist_to_state[hist[1 : len(hist)]]
+                        print(
+                            this_fst_state,
+                            backoff_fst_state,
+                            word_disambig_symbol,
+                            0,
+                            this_cost,
+                        )
+
+    # This function returns a set of n-grams that cannot currently be pruned
+    # away, either because a higher-order form of the same n-gram already exists,
+    # or because the n-gram leads to an n-gram state that exists.
+    # [Note: as we prune, we remove any states that can be removed; see that
+    # PruneToIntermediateTarget() calls PruneEmptyStates().
+
+    def GetProtectedNgrams(self):
+        ans = set()
+        for n in range(args.no_backoff_ngram_order + 1, args.ngram_order):
+            for hist, counts_for_hist in self.counts[n].items():
+                # If we have an n-gram (6, 7, 8) -> 9, the following loop will
+                # add the backed-off n-grams (7, 8) -> 9 and (8) -> 9 to
+                # 'protected-ngrams'.
+                reduced_hist = hist
+                for _ in reversed(list(range(args.no_backoff_ngram_order, n))):
+                    reduced_hist = reduced_hist[1:]  # shift an element off
+                    # the history.
+
+                    for word in counts_for_hist.word_to_count.keys():
+                        if word != self.backoff_symbol:
+                            ans.add(reduced_hist + (word,))
+                # The following statement ensures that if we are in a
+                # history-state (6, 7, 8), then n-grams (6, 7, 8) and (6, 7) are
+                # protected.  This assures that the FST states are accessible.
+                reduced_hist = hist
+                for _ in reversed(list(range(args.no_backoff_ngram_order, n))):
+                    ans.add(reduced_hist)
+                    reduced_hist = reduced_hist[:-1]  # pop an element off the
+                    # history
+        return ans
+
+    def PruneNgram(self, hist, word):
+        counts_for_hist = self.counts[len(hist)][hist]
+        assert word != self.backoff_symbol and word in counts_for_hist.word_to_count
+        count = counts_for_hist.word_to_count[word]
+        del counts_for_hist.word_to_count[word]
+        counts_for_hist.word_to_count[self.backoff_symbol] += count
+        # the next call adds the count to the symbol 'word' in the backoff
+        # history-state, and also updates its 'total_count'.
+        self.counts[len(hist) - 1][hist[1:]].AddCount(word, count)
+
+    # The function PruningLogprobChange is the same as the same-named
+    # function in float-counts-prune.cc in pocolm.  Note, it doesn't access
+    # any class members.
+
+    # This function computes the log-likelihood change (<= 0) from backing off
+    # a particular symbol to the lower-order state.
+    # The value it returns can be interpreted as a lower bound the actual log-likelihood
+    # change.  By "the actual log-likelihood change" we mean of data generated by
+    # the model itself before making the change, then modeled with the changed model
+    # [and comparing the log-like with the log-like before changing the model].  That is,
+    # it's a K-L divergence, but with the caveat that we don't normalize by the
+    # overall count of the data, so it's a K-L divergence multiplied by the training-data
+    # count.
+
+    #  'count' is the count of the word (call it 'a') in this state.  It's an integer.
+    #  'discount' is the discount-count in this state (represented as the count
+    #         for the symbol self.backoff_symbol).  It's an integer.
+    #  [note: we don't care about the total-count in this state, it cancels out.]
+    #  'backoff_count' is the count of word 'a' in the lower-order state.
+    #                 [actually it is the augmented count, treating any
+    #                  extra probability from even-lower-order states as
+    #                  if it were a count].  It's a float.
+    #  'backoff_total' is the total count in the lower-order state.  It's a float.
+    def PruningLogprobChange(self, count, discount, backoff_count, backoff_total):
+        if count == 0:
+            return 0.0
+
+        assert discount > 0 and backoff_total >= backoff_count and backoff_total >= 0.99 * discount
+
+        # augmented_count is like 'count', but with the extra count for symbol
+        # 'a' due to backoff included.
+        augmented_count = count + discount * backoff_count / backoff_total
+
+        # We imagine a phantom symbol 'b' that represents all symbols other than
+        # 'a' appearing in this history-state that are accessed via backoff.  We
+        # treat these as being distinct symbols from the same symbol if accessed
+        # not-via-backoff.  (Treating same symbols as distinct gives an upper bound
+        # on the divergence).  We also treat them as distinct from the same symbols
+        # that are being accessed via backoff from other states.  b_count is the
+        # observed count of symbol 'b' in this state (the backed-off count is
+        # zero).  b_count is also the count of symbol 'b' in the backoff state.
+        # Note: b_count will not be negative because backoff_total >= backoff_count.
+        b_count = discount * ((backoff_total - backoff_count) / backoff_total)
+        assert b_count >= -0.001 * backoff_total
+
+        # We imagine a phantom symbol 'c' that represents all symbols other than
+        # 'a' and 'b' appearing in the backoff state, which got there from
+        # backing off other states (other than 'this' state).  Again, we imagine
+        # the symbols are distinct even though they may not be (i.e. that c and
+        # b represent disjoint sets of symbol, even though they might not really
+        # be disjoint), and this gives us an upper bound on the divergence.
+        c_count = backoff_total - backoff_count - b_count
+        assert c_count >= -0.001 * backoff_total
+
+        # a_other is the count of 'a' in the backoff state that comes from
+        # 'other sources', i.e. it was backed off from history-states other than
+        # the current history state.
+        a_other_count = backoff_count - discount * backoff_count / backoff_total
+        assert a_other_count >= -0.001 * backoff_count
+
+        # the following sub-expressions are the 'new' versions of certain
+        # quantities after we assign the total count 'count' to backoff.  it
+        # increases the backoff count in 'this' state, and also the total count
+        # in the backoff state, and the count of symbol 'a' in the backoff
+        # state.
+        new_backoff_count = backoff_count + count  # new count of symbol 'a' in
+        # backoff state
+        new_backoff_total = backoff_total + count  # new total count in
+        # backoff state.
+        new_discount = discount + count  # new discount-count in 'this' state.
+
+        # all the loglike changes below are of the form
+        # count-of-symbol * log(new prob / old prob)
+        # which can be more conveniently written (by canceling the denominators),
+        # count-of-symbol * log(new count / old count).
+
+        # this_a_change is the log-like change of symbol 'a' coming from 'this'
+        # state.  bear in mind that
+        # augmented_count = count + discount * backoff_count / backoff_total,
+        # and the 'count' term is zero in the numerator part of the log expression,
+        # because symbol 'a' is completely backed off in 'this' state.
+        this_a_change = augmented_count * math.log(
+            (new_discount * new_backoff_count / new_backoff_total) / augmented_count
+        )
+
+        # other_a_change is the log-like change of symbol 'a' coming from all
+        # other states than 'this'.  For speed reasons we don't examine the
+        # direct (non-backoff) counts of symbol 'a' in all other states than
+        # 'this' that back off to the backoff state-- it would be slower.
+        # Instead we just treat the direct part of the prob for symbol 'a' as a
+        # distinct symbol when it comes from those other states... as usual,
+        # doing so gives us an upper bound on the divergence.
+        other_a_change = a_other_count * math.log(
+            (new_backoff_count / new_backoff_total) / (backoff_count / backoff_total)
+        )
+
+        # b_change is the log-like change of phantom symbol 'b' coming from
+        # 'this' state (and note: it only comes from this state, that's how we
+        # defined it).
+        # note: the expression below could be more directly written as a
+        # ratio of pseudo-counts as follows, by converting the backoff probabilities
+        # into pseudo-counts in 'this' state:
+        #  b_count * logf((new_discount * b_count / new_backoff_total) /
+        #                 (discount * b_count / backoff_total),
+        # but we cancel b_count to give us the expression below.
+        b_change = b_count * math.log((new_discount / new_backoff_total) / (discount / backoff_total))
+
+        # c_change is the log-like change of phantom symbol 'c' coming from
+        # all other states that back off to the backoff sate (and all prob. mass of
+        # 'c' comes from those other states).  The expression below could be more
+        # directly written as a ratio of counts, as c_count * logf((c_count /
+        # new_backoff_total) / (c_count / backoff_total)), but we simplified it to
+        # the expression below.
+        c_change = c_count * math.log(backoff_total / new_backoff_total)
+
+        ans = this_a_change + other_a_change + b_change + c_change
+        # the answer should not be positive.
+        assert ans <= 0.0001 * (count + discount + backoff_count + backoff_total)
+        if args.verbose >= 4:
+            print(
+                "pruning-logprob-change for {0},{1},{2},{3} is {4}".format(
+                    count, discount, backoff_count, backoff_total, ans
+                ),
+                file=sys.stderr,
+            )
+        return ans
+
+    def GetLikeChangeFromPruningNgram(self, hist, word):
+        counts_for_hist = self.counts[len(hist)][hist]
+        counts_for_backoff_hist = self.counts[len(hist) - 1][hist[1:]]
+        assert word != self.backoff_symbol and word in counts_for_hist.word_to_count
+        count = counts_for_hist.word_to_count[word]
+        discount = counts_for_hist.word_to_count[self.backoff_symbol]
+        backoff_total = counts_for_backoff_hist.total_count
+        # backoff_count is a pseudo-count: it's like the count of 'word' in the
+        # backoff history-state, but adding something to account for further
+        # levels of backoff.
+        try:
+            backoff_count = self.GetProb(hist[1:], word) * backoff_total
+        except Exception:
+            print(
+                "problem getting backoff count: hist = {0}, word = {1}".format(hist, word),
+                file=sys.stderr,
+            )
+            sys.exit(1)
+
+        return self.PruningLogprobChange(float(count), float(discount), backoff_count, float(backoff_total))
+
+    # note: returns loglike change per word.
+    def PruneToIntermediateTarget(self, num_extra_ngrams):
+        protected_ngrams = self.GetProtectedNgrams()
+        initial_num_extra_ngrams = self.GetNumExtraNgrams()
+        num_ngrams_to_prune = initial_num_extra_ngrams - num_extra_ngrams
+        assert num_ngrams_to_prune > 0
+
+        num_candidates_per_order = [0] * args.ngram_order
+        num_pruned_per_order = [0] * args.ngram_order
+
+        # like_change_and_ngrams this will be a list of tuples consisting
+        # of the likelihood change as a float and then the words of the n-gram
+        # that we're considering pruning,
+        # e.g. (-0.164, 7, 8, 9)
+        # meaning that pruning the n-gram (7, 8) -> 9 leads to
+        # a likelihood change of -0.164.  We'll later sort this list
+        # so we can prune the n-grams that made the least-negative
+        # likelihood change.
+        like_change_and_ngrams = []
+        for n in range(args.no_backoff_ngram_order, args.ngram_order):
+            for hist, counts_for_hist in self.counts[n].items():
+                for word, count in counts_for_hist.word_to_count.items():
+                    if word != self.backoff_symbol:
+                        if not hist + (word,) in protected_ngrams:
+                            like_change = self.GetLikeChangeFromPruningNgram(hist, word)
+                            like_change_and_ngrams.append((like_change,) + hist + (word,))
+                            num_candidates_per_order[len(hist)] += 1
+
+        like_change_and_ngrams.sort(reverse=True)
+
+        if num_ngrams_to_prune > len(like_change_and_ngrams):
+            print(
+                "make_phone_lm.py: aimed to prune {0} n-grams but could only "
+                "prune {1}".format(num_ngrams_to_prune, len(like_change_and_ngrams)),
+                file=sys.stderr,
+            )
+            num_ngrams_to_prune = len(like_change_and_ngrams)
+
+        total_loglike_change = 0.0
+
+        for i in range(num_ngrams_to_prune):
+            total_loglike_change += like_change_and_ngrams[i][0]
+            hist = like_change_and_ngrams[i][1:-1]  # all but 1st and last elements
+            word = like_change_and_ngrams[i][-1]  # last element
+            num_pruned_per_order[len(hist)] += 1
+            self.PruneNgram(hist, word)
+
+        like_change_per_word = total_loglike_change / self.total_num_words
+
+        if args.verbose >= 1:
+            effective_threshold = (
+                like_change_and_ngrams[num_ngrams_to_prune - 1][0] if num_ngrams_to_prune >= 0 else 0.0
+            )
+            print(
+                "Pruned from {0} ngrams to {1}, with threshold {2}.  Candidates per order were {3}, "
+                "num-ngrams pruned per order were {4}.  Like-change per word was {5}".format(
+                    initial_num_extra_ngrams,
+                    initial_num_extra_ngrams - num_ngrams_to_prune,
+                    "%.4f" % effective_threshold,
+                    num_candidates_per_order,
+                    num_pruned_per_order,
+                    like_change_per_word,
+                ),
+                file=sys.stderr,
+            )
+
+        if args.verbose >= 3:
+            print(
+                "Pruning: like_change_and_ngrams is:\n"
+                + "\n".join([str(x) for x in like_change_and_ngrams[:num_ngrams_to_prune]])
+                + "\n-------- stop pruning here: ----------\n"
+                + "\n".join([str(x) for x in like_change_and_ngrams[num_ngrams_to_prune:]]),
+                file=sys.stderr,
+            )
+            self.Print(
+                "Counts after pruning to num-extra-ngrams={0}".format(initial_num_extra_ngrams - num_ngrams_to_prune)
+            )
+
+        self.PruneEmptyStates()
+        if args.verbose >= 3:
+            ngram_counts.Print("Counts after removing empty states [inside pruning algorithm]:")
+        return like_change_per_word
+
+    def PruneToFinalTarget(self, num_extra_ngrams):
+        # prunes to a specified num_extra_ngrams.  The 'extra_ngrams' refers to
+        # the count of n-grams of order higher than args.no_backoff_ngram_order.
+        # We construct a sequence of targets that gradually approaches
+        # this value.  Doing it iteratively like this is a good way
+        # to deal with the fact that sometimes we can't prune a certain
+        # n-gram before certain other n-grams are pruned (because
+        # they lead to a state that must be kept, or an n-gram exists
+        # that backs off to this n-gram).
+
+        current_num_extra_ngrams = self.GetNumExtraNgrams()
+
+        if num_extra_ngrams >= current_num_extra_ngrams:
+            print(
+                "make_phone_lm.py: not pruning since target num-extra-ngrams={0} is >= "
+                "current num-extra-ngrams={1}".format(num_extra_ngrams, current_num_extra_ngrams),
+                file=sys.stderr,
+            )
+            return
+
+        target_sequence = [num_extra_ngrams]
+        # two final iterations where the targets differ by factors of 1.1,
+        # preceded by two iterations where the targets differ by factors of 1.2.
+        for this_factor in [1.1, 1.2]:
+            for n in range(0, 2):
+                if int((target_sequence[-1] + 1) * this_factor) < current_num_extra_ngrams:
+                    target_sequence.append(int((target_sequence[-1] + 1) * this_factor))
+        # then change in factors of 1.3
+        while True:
+            this_factor = 1.3
+            if int((target_sequence[-1] + 1) * this_factor) < current_num_extra_ngrams:
+                target_sequence.append(int((target_sequence[-1] + 1) * this_factor))
+            else:
+                break
+
+        target_sequence = list(set(target_sequence))  # only keep unique targets.
+        target_sequence.sort(reverse=True)
+
+        print(
+            "make_phone_lm.py: current num-extra-ngrams={0}, pruning with "
+            "following sequence of targets: {1}".format(current_num_extra_ngrams, target_sequence),
+            file=sys.stderr,
+        )
+        total_like_change_per_word = 0.0
+        for target in target_sequence:
+            total_like_change_per_word += self.PruneToIntermediateTarget(target)
+
+        if args.verbose >= 1:
+            print(
+                "make_phone_lm.py: K-L divergence from pruning (upper bound) is " "%.4f" % total_like_change_per_word,
+                file=sys.stderr,
+            )
+
+    # returns the number of n-grams on top of those that can't be pruned away
+    # because their order is <= args.no_backoff_ngram_order.
+    def GetNumExtraNgrams(self):
+        ans = 0
+        for hist_len in range(args.no_backoff_ngram_order, args.ngram_order):
+            # note: hist_len + 1 is the actual order.
+            ans += self.GetNumNgrams(hist_len)
+        return ans
+
+    def GetNumNgrams(self, hist_len=None):
+        ans = 0
+        if hist_len is None:
+            for hist_len in range(args.ngram_order):
+                # note: hist_len + 1 is the actual order.
+                ans += self.GetNumNgrams(hist_len)
+            return ans
+        else:
+            for counts_for_hist in self.counts[hist_len].values():
+                ans += len(counts_for_hist.word_to_count)
+                if self.backoff_symbol in counts_for_hist.word_to_count:
+                    ans -= 1  # don't count the backoff symbol, it doesn't produce
+                    # its own n-gram line.
+            return ans
+
+    # this function, used in PrintAsArpa, converts an integer to
+    # a string by either printing it as a string, or for self.bos_symbol
+    # and self.eos_symbol, printing them as "<s>" and "</s>" respectively.
+    def IntToString(self, i):
+        if i == self.bos_symbol:
+            return "<s>"
+        elif i == self.eos_symbol:
+            return "</s>"
+        else:
+            assert i != self.backoff_symbol
+            return str(i)
+
+    def PrintAsArpa(self):
+        # Prints out the FST in ARPA format.
+        assert args.no_backoff_ngram_order == 1  # without unigrams we couldn't
+        # print as ARPA format.
+
+        print("\\data\\")
+        for hist_len in range(args.ngram_order):
+            # print the number of n-grams.  Add 1 for the 1-gram
+            # section because of <s>, we print -99 as the prob so we
+            # have a place to put the backoff prob.
+            print(
+                "ngram {0}={1}".format(
+                    hist_len + 1,
+                    self.GetNumNgrams(hist_len) + (1 if hist_len == 0 else 0),
+                )
+            )
+
+        print("")
+
+        for hist_len in range(args.ngram_order):
+            print("\\{0}-grams:".format(hist_len + 1))
+
+            # print fake n-gram for <s>, for its backoff prob.
+            if hist_len == 0:
+                backoff_prob = self.GetProb((self.bos_symbol,), self.backoff_symbol)
+                if backoff_prob != None:
+                    print("-99\t<s>\t{0}".format("%.5f" % math.log10(backoff_prob)))
+
+            for hist in self.counts[hist_len].keys():
+                for word in self.counts[hist_len][hist].word_to_count.keys():
+                    if word != self.backoff_symbol:
+                        prob = self.GetProb(hist, word)
+                        assert prob != None and prob > 0
+                        backoff_prob = self.GetProb((hist) + (word,), self.backoff_symbol)
+                        line = "{0}\t{1}".format(
+                            "%.5f" % math.log10(prob),
+                            " ".join(self.IntToString(x) for x in hist + (word,)),
+                        )
+                        if backoff_prob != None:
+                            line += "\t{0}".format("%.5f" % math.log10(backoff_prob))
+                        print(line)
+            print("")
+        print("\\end\\")
+
+
+ngram_counts = NgramCounts(args.ngram_order)
+ngram_counts.AddRawCountsFromStandardInput()
+
+if args.verbose >= 3:
+    ngram_counts.Print("Raw counts:")
+ngram_counts.ApplyBackoff()
+if args.verbose >= 3:
+    ngram_counts.Print("Counts after applying Kneser-Ney discounting:")
+ngram_counts.EnsureStructurallyNeededNgramsExist()
+if args.verbose >= 3:
+    ngram_counts.Print("Counts after adding structurally-needed n-grams (1st time):")
+ngram_counts.PruneEmptyStates()
+if args.verbose >= 3:
+    ngram_counts.Print("Counts after removing empty states:")
+ngram_counts.PruneToFinalTarget(args.num_extra_ngrams)
+
+ngram_counts.EnsureStructurallyNeededNgramsExist()
+if args.verbose >= 3:
+    ngram_counts.Print("Counts after adding structurally-needed n-grams (2nd time):")
+
+
+if args.print_as_arpa == "true":
+    ngram_counts.PrintAsArpa()
+else:
+    if args.phone_disambig_symbol is None:
+        sys.exit("make_phone_lm.py: --phone-disambig-symbol must be provided (unless " "you are writing as ARPA")
+    ngram_counts.PrintAsFst(args.phone_disambig_symbol)
+
+
+## Below are some little test commands that can be used to look at the detailed stats
+## for a kind of sanity check.
+# test comand:
+# (echo 6 7 8 4; echo 7 8 9; echo 7 8; echo 7 4; echo 8 4 ) | utils/lang/make_phone_lm.py --phone-disambig-symbol=400  --verbose=3
+#  (echo 6 7 8 4; echo 7 8 9; echo 7 8; echo 7 4; echo 8 4 ) | utils/lang/make_phone_lm.py --phone-disambig-symbol=400  --verbose=3 --num-extra-ngrams=0
+# (echo 6 7 8 4; echo 6 7 ) | utils/lang/make_phone_lm.py --print-as-arpa=true --no-backoff-ngram-order=1  --verbose=3
+
+
+## The following shows how we created some data suitable to do comparisons with
+## other language modeling toolkits.  Note: we're running in a configuration
+## where --no-backoff-ngram-order=1 (i.e. we have a unigram LM state) because
+## it's the only way to get perplexity calculations and to write an ARPA file.
+##
+# cd egs/tedlium/s5_r2
+# . ./path.sh
+# mkdir -p lm_test
+# ali-to-phones exp/tri3/final.mdl "ark:gunzip -c exp/tri3/ali.*.gz|" ark,t:-  | awk '{$1 = ""; print}' > lm_test/phone_seqs
+# wc lm_test/phone_seqs
+# 92464  8409563 27953288 lm_test/phone_seqs
+# head -n 20000 lm_test/phone_seqs > lm_test/train.txt
+# tail -n 1000 lm_test/phone_seqs > lm_test/test.txt
+
+## This shows make_phone_lm.py with the default number of extra-lm-states (20k)
+## You have to have SRILM on your path to ger perplexities [note: it should be on the
+## path if you installed it and you sourced the tedlium s5b path.sh, as above.]
+# utils/lang/make_phone_lm.py --print-as-arpa=true --no-backoff-ngram-order=1 --verbose=1 < lm_test/train.txt > lm_test/arpa_pr20k
+# ngram -order 4 -unk -lm lm_test/arpa_pr20k -ppl lm_test/test.txt
+# file lm_test/test.txt: 1000 sentences, 86489 words, 3 OOVs
+# 0 zeroprobs, logprob= -80130.1 ppl=*8.23985* ppl1= 8.44325
+# on training data: 0 zeroprobs, logprob= -1.6264e+06 ppl= 7.46947 ppl1= 7.63431
+
+## This shows make_phone_lm.py without any pruning (make --num-extra-ngrams very large).
+# utils/lang/make_phone_lm.py --print-as-arpa=true --num-extra-ngrams=1000000 --no-backoff-ngram-order=1 --verbose=1 < lm_test/train.txt > lm_test/arpa
+# ngram -order 4 -unk -lm lm_test/arpa -ppl lm_test/test.txt
+# file lm_test/test.txt: 1000 sentences, 86489 words, 3 OOVs
+# 0 zeroprobs, logprob= -74976 ppl=*7.19459* ppl1= 7.36064
+# on training data: 0 zeroprobs, logprob= -1.44198e+06 ppl= 5.94659 ppl1= 6.06279
+
+## This is SRILM without pruning (c.f. the 7.19 above, it's slightly better).
+# ngram-count -text lm_test/train.txt -order 4 -kndiscount2 -kndiscount3 -kndiscount4 -interpolate -lm lm_test/arpa_srilm
+# ngram -order 4 -unk -lm lm_test/arpa_srilm -ppl lm_test/test.txt
+# file lm_test/test.txt: 1000 sentences, 86489 words, 3 OOVs
+# 0 zeroprobs, logprob= -74742.2 ppl= *7.15044* ppl1= 7.31494
+
+
+## This is SRILM with a pruning beam tuned to get 20k n-grams above unigram
+##  (c.f. the 8.23 above, it's a lot worse).
+# ngram-count -text lm_test/train.txt -order 4 -kndiscount2 -kndiscount3 -kndiscount4 -interpolate -prune 1.65e-05 -lm lm_test/arpa_srilm.pr1.65e-5
+# the model has 20249 n-grams above unigram [c.f. our 20k]
+# ngram -order 4 -unk -lm lm_test/arpa_srilm.pr1.65e-5 -ppl lm_test/test.txt
+# file lm_test/test.txt: 1000 sentences, 86489 words, 3 OOVs
+# 0 zeroprobs, logprob= -86803.7 ppl=*9.82202* ppl1= 10.0849
+
+
+## This is pocolm..
+## Note: we have to hold out some of the training data as dev to
+## estimate the hyperparameters, but we'll fold it back in before
+## making the final LM. [--fold-dev-into=train]
+# mkdir -p lm_test/data/text
+# head -n 1000 lm_test/train.txt > lm_test/data/text/dev.txt
+# tail -n +1001 lm_test/train.txt > lm_test/data/text/train.txt
+## give it a 'large' num-words so it picks them all.
+# export PATH=$PATH:../../../tools/pocolm/scripts
+# train_lm.py --num-word=100000 --fold-dev-into=train lm_test/data/text 4 lm_test/data/lm_unpruned
+# get_data_prob.py lm_test/test.txt lm_test/data/lm_unpruned/100000_4.pocolm
+## compute-probs: average log-prob per word was -1.95956 (perplexity = *7.0962*) over 87489 words.
+## Note: we can compare this perplexity with 7.15 with SRILM and 7.19 with make_phone_lm.py.
+
+#   pruned_lm_dir=${lm_dir}/${num_word}_${order}_prune${threshold}.pocolm
+# prune_lm_dir.py --target-num-ngrams=20100 lm_test/data/lm_unpruned/100000_4.pocolm lm_test/data/lm_unpruned/100000_4_pr20k.pocolm
+# get_data_prob.py lm_test/test.txt lm_test/data/lm_unpruned/100000_4_pr20k.pocolm
+## compute-probs: average log-prob per word was -2.0409 (perplexity = 7.69757) over 87489 words.
+## note: the 7.69 can be compared with 9.82 from SRILM and 8.23 from pocolm.
+## format_arpa_lm.py lm_test/data/lm_unpruned/100000_4_pr20k.pocolm | head
+## .. it has 20488 n-grams above unigram.  More than 20k but not enough to explain the difference
+## .. in perplexity.
+
+## OK... if I reran after modifying prune_lm_dir.py to comment out the line
+## 'steps += 'EM EM'.split()' which adds the two EM stages per step, and got the
+## perplexity again, I got the following:
+## compute-probs: average log-prob per word was -2.09722 (perplexity = 8.14353) over 87489 words.
+## .. so it turns out the E-M is actually important.

scripts/asr_language_modeling/ngram_lm/ngram_merge.py

@@ -0,0 +1,485 @@
+# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES.  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 interpolate two arpa N-gram language models (LMs), 
+culculate perplexity of resulted LM, and make binary KenLM from it.
+
+Minimun usage example to interpolate two N-gram language models with weights:
+alpha * ngram_a + beta * ngram_b = 2 * ngram_a + 1 * ngram_b
+
+python3 ngram_merge.py  --kenlm_bin_path /workspace/nemo/decoders/kenlm/build/bin \
+                    --arpa_a /path/ngram_a.kenlm.tmp.arpa \
+                    --alpha 2 \
+                    --arpa_b /path/ngram_b.kenlm.tmp.arpa \
+                    --beta 1 \
+                    --out_path /path/out
+
+
+Merge two N-gram language models and calculate its perplexity with test_file.
+python3 ngram_merge.py  --kenlm_bin_path /workspace/nemo/decoders/kenlm/build/bin \
+                    --ngram_bin_path /workspace/nemo/decoders/ngram-1.3.14/src/bin \
+                    --arpa_a /path/ngram_a.kenlm.tmp.arpa \
+                    --alpha 0.5 \
+                    --arpa_b /path/ngram_b.kenlm.tmp.arpa \
+                    --beta 0.5 \
+                    --out_path /path/out \
+                    --nemo_model_file /path/to/model_tokenizer.nemo \
+                    --test_file /path/to/test_manifest.json \
+                    --force
+"""
+
+import argparse
+import os
+import shlex
+import subprocess
+import sys
+from typing import Tuple
+
+import torch
+
+import nemo.collections.asr as nemo_asr
+from nemo.collections.asr.modules.rnnt import RNNTDecoder
+from nemo.collections.asr.parts.submodules.ngram_lm import DEFAULT_TOKEN_OFFSET, kenlm_utils
+from nemo.utils import logging
+
+
+class NgramMerge:
+    def __init__(self, ngram_bin_path):
+        self.ngram_bin_path = ngram_bin_path
+
+    def ngrammerge(self, arpa_a: str, alpha: float, arpa_b: str, beta: float, arpa_c: str, force: bool) -> str:
+        """
+        Merge two ARPA n-gram language models using the ngrammerge command-line tool and output the result in ARPA format.
+
+        Args:
+            arpa_a (str): Path to the first input ARPA file.
+            alpha (float): Interpolation weight for the first model.
+            arpa_b (str): Path to the second input ARPA file.
+            beta (float): Interpolation weight for the second model.
+            arpa_c (str): Path to the output ARPA file.
+            force (bool): Whether to overwrite existing output files.
+
+        Returns:
+            str: Path to the output ARPA file in mod format.
+        """
+        mod_a = arpa_a + ".mod"
+        mod_b = arpa_b + ".mod"
+        mod_c = arpa_c + ".mod"
+        if os.path.isfile(mod_c) and not force:
+            logging.info("File " + mod_c + " exists. Skipping.")
+        else:
+            sh_args = [
+                os.path.join(self.ngram_bin_path, "ngrammerge"),
+                "--alpha=" + str(alpha),
+                "--beta=" + str(beta),
+                "--normalize",
+                # "--use_smoothing",
+                mod_a,
+                mod_b,
+                mod_c,
+            ]
+            logging.info(
+                "\n"
+                + str(
+                    subprocess.run(
+                        sh_args,
+                        capture_output=False,
+                        text=True,
+                        stdout=sys.stdout,
+                        stderr=sys.stderr,
+                    )
+                )
+                + "\n",
+            )
+        return mod_c
+
+    def arpa2mod(self, arpa_path: str, force: bool):
+        """
+        This function reads an ARPA n-gram model and converts it to a binary format. The binary model is saved to the same directory as the ARPA model with a ".mod" extension. If the binary model file already exists and force argument is False, then the function skips conversion and returns a message. Otherwise, it executes the command to create a binary model using the subprocess.run method.
+
+        Parameters:
+            arpa_path (string): The file path to the ARPA n-gram model.
+            force (bool): If True, the function will convert the ARPA model to binary even if the binary file already exists. If False and the binary file exists, the function will skip the conversion.
+        Returns:
+            If the binary model file already exists and force argument is False, returns a message indicating that the file exists and the conversion is skipped.
+            Otherwise, returns a subprocess.CompletedProcess object, which contains information about the executed command. The subprocess's output and error streams are redirected to stdout and stderr, respectively.
+        """
+        mod_path = arpa_path + ".mod"
+        if os.path.isfile(mod_path) and not force:
+            return "File " + mod_path + " exists. Skipping."
+        else:
+            sh_args = [
+                os.path.join(self.ngram_bin_path, "ngramread"),
+                "--ARPA",
+                arpa_path,
+                mod_path,
+            ]
+            return subprocess.run(
+                sh_args,
+                capture_output=False,
+                text=True,
+                stdout=sys.stdout,
+                stderr=sys.stderr,
+            )
+
+    def merge(
+        self, arpa_a: str, alpha: float, arpa_b: str, beta: float, out_path: str, force: bool
+    ) -> Tuple[str, str]:
+        """
+        Merges two ARPA language models using the ngrammerge tool.
+
+        Args:
+            arpa_a (str): Path to the first ARPA language model file.
+            alpha (float): Interpolation weight for the first model.
+            arpa_b (str): Path to the second ARPA language model file.
+            beta (float): Interpolation weight for the second model.
+            out_path (str): Path to the output directory for the merged ARPA model.
+            force (bool): Whether to force overwrite of existing files.
+
+        Returns:
+            Tuple[str, str]: A tuple containing the path to the merged binary language model file and the path to the
+            merged ARPA language model file.
+        """
+        logging.info("\n" + str(self.arpa2mod(arpa_a, force)) + "\n")
+
+        logging.info("\n" + str(self.arpa2mod(arpa_b, force)) + "\n")
+        arpa_c = os.path.join(
+            out_path,
+            f"{os.path.split(arpa_a)[1]}-{alpha}-{os.path.split(arpa_b)[1]}-{beta}.arpa",
+        )
+        mod_c = self.ngrammerge(arpa_a, alpha, arpa_b, beta, arpa_c, force)
+        return mod_c, arpa_c
+
+    def perplexity(self, ngram_mod: str, test_far: str) -> str:
+        """
+        Calculates perplexity of a given ngram model on a test file.
+
+        Args:
+            ngram_mod (str): The path to the ngram model file.
+            test_far (str): The path to the test file.
+
+        Returns:
+            str: A string representation of the perplexity calculated.
+
+        Raises:
+            AssertionError: If the subprocess to calculate perplexity returns a non-zero exit code.
+
+        Example:
+            >>> perplexity("/path/to/ngram_model", "/path/to/test_file")
+            'Perplexity: 123.45'
+        """
+        sh_args = [
+            os.path.join(self.ngram_bin_path, "ngramperplexity"),
+            "--v=1",
+            ngram_mod,
+            test_far,
+        ]
+        ps = subprocess.Popen(sh_args, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        stdout, stderr = ps.communicate()
+        exit_code = ps.wait()
+        command = " ".join(sh_args)
+        assert (
+            exit_code == 0
+        ), f"Exit_code must be 0.\n bash command: {command} \n stdout: {stdout} \n stderr: {stderr}"
+        perplexity_out = "\n".join(stdout.split("\n")[-6:-1])
+        return perplexity_out
+
+    def make_arpa(self, ngram_mod: str, ngram_arpa: str, force: bool):
+        """
+        Converts an ngram model in binary format to ARPA format.
+
+        Args:
+        - ngram_mod (str): The path to the ngram model in binary format.
+        - ngram_arpa (str): The desired path for the ARPA format output file.
+        - force (bool): If True, the ARPA format file will be generated even if it already exists.
+
+        Returns:
+        - Tuple[bytes, bytes]
+
+        Raises:
+        - AssertionError: If the shell command execution returns a non-zero exit code.
+        - FileNotFoundError: If the binary ngram model file does not exist.
+        """
+        if os.path.isfile(ngram_arpa) and not force:
+            logging.info("File " + ngram_arpa + " exists. Skipping.")
+            return None
+        else:
+            sh_args = [
+                os.path.join(self.ngram_bin_path, "ngramprint"),
+                "--ARPA",
+                ngram_mod,
+                ngram_arpa,
+            ]
+            return subprocess.run(
+                sh_args,
+                capture_output=False,
+                text=True,
+                stdout=sys.stdout,
+                stderr=sys.stderr,
+            )
+
+    def test_perplexity(self, mod_c: str, symbols: str, test_txt: str, nemo_model_file: str, tmp_path: str) -> str:
+        """
+        Tests the perplexity of a given ngram model on a test file.
+
+        Args:
+            mod_c (str): The path to the ngram model file.
+            symbols (str): The path to the symbol table file.
+            test_txt (str): The path to the test text file.
+            nemo_model_file (str): The path to the NeMo model file.
+            tmp_path (str): The path to the temporary directory where the test far file will be created.
+            force (bool): If True, overwrites any existing far file.
+
+        Returns:
+            str: A string representation of the perplexity calculated.
+
+        Example:
+            >>> test_perplexity("/path/to/ngram_model", "/path/to/symbol_table", "/path/to/test_file", "/path/to/tokenizer_model", "/path/to/tmp_dir", True)
+            'Perplexity: 123.45'
+        """
+
+        test_far = farcompile(symbols, test_txt, tmp_path, nemo_model_file)
+        res_p = self.perplexity(mod_c, test_far)
+        return res_p
+
+
+def farcompile(symbols: str, text_file: str, tmp_path: str, nemo_model_file: str) -> str:
+    """
+    Compiles a text file into a FAR file using the given symbol table or tokenizer.
+
+    Args:
+        symbols (str): The path to the symbol table file.
+        text_file (str): The path to the text file to compile.
+        tmp_path (str): The path to the temporary directory where the test far file will be created.
+        nemo_model_file (str): The path to the NeMo model file (.nemo).
+        force (bool): If True, overwrites any existing FAR file.
+
+    Returns:
+        test_far (str): The path to the resulting FAR file.
+
+    Example:
+        >>> farcompile("/path/to/symbol_table", "/path/to/text_file", "/path/to/far_file", "/path/to/tokenizer_model", "/path/to/nemo_model", True)
+    """
+    test_far = os.path.join(tmp_path, os.path.split(text_file)[1] + ".far")
+
+    sh_args = [
+        "farcompilestrings",
+        "--generate_keys=10",
+        "--fst_type=compact",
+        "--symbols=" + symbols,
+        "--keep_symbols",
+    ]
+
+    tokenizer, encoding_level, is_aggregate_tokenizer, _ = kenlm_utils.setup_tokenizer(nemo_model_file)
+
+    with open(test_far, "wb") as test_far_stdout:
+        ps = subprocess.Popen(
+            sh_args,
+            stdin=subprocess.PIPE,
+            stdout=test_far_stdout,
+            stderr=sys.stderr,
+        )
+
+        kenlm_utils.iter_files(
+            source_path=[text_file],
+            dest_path=ps.stdin,
+            tokenizer=tokenizer,
+            encoding_level=encoding_level,
+            is_aggregate_tokenizer=is_aggregate_tokenizer,
+            verbose=1,
+        )
+        stdout, stderr = ps.communicate()
+
+    exit_code = ps.returncode
+
+    command = f"{shlex.join(sh_args)} > {shlex.quote(test_far)}"
+    assert exit_code == 0, f"Exit_code must be 0.\n bash command: {command} \n stdout: {stdout} \n stderr: {stderr}"
+    return test_far
+
+
+def make_kenlm(kenlm_bin_path: str, ngram_arpa: str, force: bool):
+    """
+    Builds a language model from an ARPA format file using the KenLM toolkit.
+
+    Args:
+    - kenlm_bin_path (str): The path to the KenLM toolkit binary.
+    - ngram_arpa (str): The path to the ARPA format file.
+    - force (bool): If True, the KenLM language model will be generated even if it already exists.
+
+    Raises:
+    - AssertionError: If the shell command execution returns a non-zero exit code.
+    - FileNotFoundError: If the KenLM binary or ARPA format file does not exist.
+    """
+    ngram_kenlm = ngram_arpa + ".kenlm"
+    if os.path.isfile(ngram_kenlm) and not force:
+        logging.info("File " + ngram_kenlm + " exists. Skipping.")
+        return None
+    else:
+        sh_args = [os.path.join(kenlm_bin_path, "build_binary"), "trie", "-i", ngram_arpa, ngram_kenlm]
+        return subprocess.run(
+            sh_args,
+            capture_output=False,
+            text=True,
+            stdout=sys.stdout,
+            stderr=sys.stderr,
+        )
+
+
+def make_symbol_list(nemo_model_file, symbols, force):
+    """
+    Function: make_symbol_list
+
+    Create a symbol table for the input tokenizer model file.
+
+    Args:
+        nemo_model_file (str): Path to the NeMo model file.
+        symbols (str): Path to the file where symbol list will be saved.
+        force (bool): Flag to force creation of symbol list even if it already exists.
+
+    Returns:
+        None
+
+    Raises:
+        None
+    """
+    if os.path.isfile(symbols) and not force:
+        logging.info("File " + symbols + " exists. Skipping.")
+    else:
+        if nemo_model_file.endswith('.nemo'):
+            asr_model = nemo_asr.models.ASRModel.restore_from(nemo_model_file, map_location=torch.device('cpu'))
+        else:
+            logging.warning(
+                "nemo_model_file does not end with .nemo, therefore trying to load a pretrained model with this name."
+            )
+            asr_model = nemo_asr.models.ASRModel.from_pretrained(nemo_model_file, map_location=torch.device('cpu'))
+
+        if isinstance(asr_model.decoder, RNNTDecoder):
+            vocab_size = asr_model.decoder.blank_idx
+        else:
+            vocab_size = len(asr_model.decoder.vocabulary)
+
+        vocab = [chr(idx + DEFAULT_TOKEN_OFFSET) for idx in range(vocab_size)]
+        with open(symbols, "w", encoding="utf-8") as f:
+            for i, v in enumerate(vocab):
+                f.write(v + " " + str(i) + "\n")
+
+
+def main(
+    kenlm_bin_path: str,
+    ngram_bin_path: str,
+    arpa_a: str,
+    alpha: float,
+    arpa_b: str,
+    beta: float,
+    out_path: str,
+    test_file: str,
+    symbols: str,
+    nemo_model_file: str,
+    force: bool,
+) -> None:
+    """
+    Entry point function for merging ARPA format language models, testing perplexity, creating symbol list,
+    and making ARPA and Kenlm models.
+
+    Args:
+    - kenlm_bin_path (str): The path to the Kenlm binary.
+    - arpa_a (str): The path to the first ARPA format language model.
+    - alpha (float): The weight given to the first language model during merging.
+    - arpa_b (str): The path to the second ARPA format language model.
+    - beta (float): The weight given to the second language model during merging.
+    - out_path (str): The path where the output files will be saved.
+    - test_file (str): The path to the file on which perplexity needs to be calculated.
+    - symbols (str): The path to the file where symbol list for the tokenizer model will be saved.
+    - nemo_model_file (str): The path to the NeMo model file.
+    - force (bool): If True, overwrite existing files, otherwise skip the operations.
+
+    Returns:
+    - None
+    """
+    nm = NgramMerge(ngram_bin_path)
+    mod_c, arpa_c = nm.merge(arpa_a, alpha, arpa_b, beta, out_path, force)
+
+    if test_file and nemo_model_file:
+        if not symbols:
+            symbols = os.path.join(out_path, os.path.split(nemo_model_file)[1] + ".syms")
+            make_symbol_list(nemo_model_file, symbols, force)
+        for test_f in test_file.split(","):
+            test_p = nm.test_perplexity(mod_c, symbols, test_f, nemo_model_file, out_path)
+            logging.info("Perplexity summary " + test_f + " : " + test_p)
+
+    logging.info("Making ARPA and Kenlm model " + arpa_c)
+    out = nm.make_arpa(mod_c, arpa_c, force)
+    if out:
+        logging.info("\n" + str(out) + "\n")
+
+    out = make_kenlm(kenlm_bin_path, arpa_c, force)
+    if out:
+        logging.info("\n" + str(out) + "\n")
+
+
+def _parse_args():
+    parser = argparse.ArgumentParser(
+        description="Interpolate ARPA N-gram language models and make KenLM binary model to be used with beam search decoder of ASR models."
+    )
+    parser.add_argument(
+        "--kenlm_bin_path",
+        required=True,
+        type=str,
+        help="The path to the bin folder of KenLM library.",
+    )  # Use /workspace/nemo/decoders/kenlm/build/bin if installed it with scripts/asr_language_modeling/ngram_lm/install_beamsearch_decoders.sh
+    parser.add_argument(
+        "--ngram_bin_path",
+        required=True,
+        type=str,
+        help="The path to the bin folder of OpenGrm Ngram library.",
+    )  # Use /workspace/nemo/decoders/ngram-1.3.14/src/bin if installed it with scripts/installers/install_opengrm.sh
+    parser.add_argument("--arpa_a", required=True, type=str, help="Path to the arpa_a")
+    parser.add_argument("--alpha", required=True, type=float, help="Weight of arpa_a")
+    parser.add_argument("--arpa_b", required=True, type=str, help="Path to the arpa_b")
+    parser.add_argument("--beta", required=True, type=float, help="Weight of arpa_b")
+    parser.add_argument(
+        "--out_path",
+        required=True,
+        type=str,
+        help="Path to write tmp and resulted files.",
+    )
+    parser.add_argument(
+        "--test_file",
+        required=False,
+        type=str,
+        default=None,
+        help="Path to test file to count perplexity if provided.",
+    )
+    parser.add_argument(
+        "--symbols",
+        required=False,
+        type=str,
+        default=None,
+        help="Path to symbols (.syms) file . Could be calculated if it is not provided. Use as: --symbols /path/to/earnest.syms",
+    )
+    parser.add_argument(
+        "--nemo_model_file",
+        required=False,
+        type=str,
+        default=None,
+        help="The path to '.nemo' file of the ASR model, or name of a pretrained NeMo model",
+    )
+    parser.add_argument("--force", "-f", action="store_true", help="Whether to recompile and rewrite all files")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    main(**vars(_parse_args()))

scripts/asr_language_modeling/ngram_lm/train_kenlm.py

@@ -0,0 +1,196 @@
+# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES.  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_paths=<list of paths 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 logging
+import os
+import subprocess
+import sys
+from dataclasses import dataclass, field
+from glob import glob
+from typing import List
+
+from omegaconf import MISSING
+
+from nemo.collections.asr.parts.submodules.ngram_lm import NGramGPULanguageModel, kenlm_utils
+from nemo.core.config import hydra_runner
+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. 
+"""
+
+
+@dataclass
+class TrainKenlmConfig:
+    """
+    Train an N-gram language model with KenLM to be used with beam search decoder of ASR models.
+    """
+
+    train_paths: List[str] = (
+        MISSING  # List of training files or folders. Files can be a plain text file or ".json" manifest or ".json.gz". Example: [/path/to/manifest/file,/path/to/folder]
+    )
+
+    nemo_model_file: str = MISSING  # The path to '.nemo' file of the ASR model, or name of a pretrained NeMo model
+    kenlm_model_file: str = MISSING  # The path to store the KenLM binary model file
+    ngram_length: int = MISSING  # The order of N-gram LM
+    kenlm_bin_path: str = MISSING  # The path to the bin folder of KenLM.
+
+    preserve_arpa: bool = False  # Whether to preserve the intermediate ARPA file.
+    ngram_prune: List[int] = field(
+        default_factory=lambda: [0]
+    )  # List of digits to prune Ngram. Example: [0,0,1]. See Pruning section on the https://kheafield.com/code/kenlm/estimation
+    cache_path: str = ""  # Cache path to save tokenized files.
+    verbose: int = 1  # Verbose level, default is 1.
+    save_nemo: bool = False  # Save .nemo checkpoint to use with NGramGPULanguageModel
+    normalize_unk_nemo: bool = True  # Normalize the UNK token in the NGramGPULanguageModel model
+
+
+@hydra_runner(config_path=None, config_name='TrainKenlmConfig', schema=TrainKenlmConfig)
+def main(args: TrainKenlmConfig):
+    train_paths = kenlm_utils.get_train_list(args.train_paths)
+
+    if isinstance(args.ngram_prune, str):
+        args.ngram_prune = [args.ngram_prune]
+
+    tokenizer, encoding_level, is_aggregate_tokenizer, full_vocab_size = kenlm_utils.setup_tokenizer(
+        args.nemo_model_file
+    )
+
+    if encoding_level == "subword":
+        discount_arg = "--discount_fallback"  # --discount_fallback is needed for training KenLM for BPE-based models
+    else:
+        discount_arg = ""
+
+    arpa_file = f"{args.kenlm_model_file}.tmp.arpa"
+    """ LMPLZ ARGUMENT SETUP """
+    kenlm_args = [
+        os.path.join(args.kenlm_bin_path, 'lmplz'),
+        "-o",
+        str(args.ngram_length),
+        "--arpa",
+        arpa_file,
+        discount_arg,
+        "--prune",
+    ] + [str(n) for n in args.ngram_prune]
+
+    if args.cache_path:
+        if not os.path.exists(args.cache_path):
+            os.makedirs(args.cache_path, exist_ok=True)
+
+        """ DATASET SETUP """
+        encoded_train_files = []
+        for file_num, train_file in enumerate(train_paths):
+            logging.info(f"Encoding the train file '{train_file}' number {file_num+1} out of {len(train_paths)} ...")
+
+            cached_files = glob(os.path.join(args.cache_path, os.path.split(train_file)[1]) + "*")
+            encoded_train_file = os.path.join(args.cache_path, os.path.split(train_file)[1] + f"_{file_num}.tmp.txt")
+            if (
+                cached_files and cached_files[0] != encoded_train_file
+            ):  # cached_files exists but has another file name: f"_{file_num}.tmp.txt"
+                os.rename(cached_files[0], encoded_train_file)
+                logging.info("Rename", cached_files[0], "to", encoded_train_file)
+
+            encoded_train_files.append(encoded_train_file)
+
+        kenlm_utils.iter_files(
+            source_path=train_paths,
+            dest_path=encoded_train_files,
+            tokenizer=tokenizer,
+            encoding_level=encoding_level,
+            is_aggregate_tokenizer=is_aggregate_tokenizer,
+            verbose=args.verbose,
+        )
+
+        first_process_args = ["cat"] + encoded_train_files
+        first_process = subprocess.Popen(first_process_args, stdout=subprocess.PIPE, stderr=sys.stderr)
+
+        logging.info(f"Running lmplz command \n\n{' '.join(kenlm_args)}\n\n")
+        kenlm_p = subprocess.run(
+            kenlm_args,
+            stdin=first_process.stdout,
+            capture_output=False,
+            text=True,
+            stdout=sys.stdout,
+            stderr=sys.stderr,
+        )
+        first_process.wait()
+
+    else:
+        logging.info(f"Running lmplz command \n\n{' '.join(kenlm_args)}\n\n")
+        kenlm_p = subprocess.Popen(kenlm_args, stdout=sys.stdout, stdin=subprocess.PIPE, stderr=sys.stderr)
+        kenlm_utils.iter_files(
+            source_path=train_paths,
+            dest_path=kenlm_p.stdin,
+            tokenizer=tokenizer,
+            encoding_level=encoding_level,
+            is_aggregate_tokenizer=is_aggregate_tokenizer,
+            verbose=args.verbose,
+        )
+
+        kenlm_p.communicate()
+
+    if kenlm_p.returncode != 0:
+        raise RuntimeError("Training KenLM was not successful!")
+
+    """ BINARY BUILD """
+
+    kenlm_args = [
+        os.path.join(args.kenlm_bin_path, "build_binary"),
+        "trie",
+        arpa_file,
+        args.kenlm_model_file,
+    ]
+    logging.info(f"Running binary_build command \n\n{' '.join(kenlm_args)}\n\n")
+    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!")
+
+    if args.save_nemo:
+        if full_vocab_size is None:
+            raise NotImplementedError("Unknown vocab size, cannot convert the model")
+        nemo_model = NGramGPULanguageModel.from_arpa(
+            lm_path=arpa_file, vocab_size=full_vocab_size, normalize_unk=args.normalize_unk_nemo
+        )
+        nemo_model.save_to(f"{args.kenlm_model_file}.nemo")
+
+    if not args.preserve_arpa:
+        os.remove(arpa_file)
+        logging.info(f"Deleted the arpa file '{arpa_file}'.")
+
+
+if __name__ == '__main__':
+    main()