chore: import upstream snapshot with attribution

scripts/asr_context_biasing/build_gpu_boosting_tree.py

@@ -0,0 +1,102 @@
+# Copyright (c) 2025, 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.
+
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+import torch
+from omegaconf import MISSING
+from torch.nn.utils.rnn import pad_sequence
+
+import nemo.collections.asr as nemo_asr
+from nemo.collections.asr.parts.context_biasing.boosting_graph_batched import (
+    BoostingTreeModelConfig,
+    GPUBoostingTreeModel,
+)
+from nemo.collections.common.tokenizers import AggregateTokenizer
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+
+
+@dataclass
+class BuildWordBoostingTreeConfig(BoostingTreeModelConfig):
+    """
+    Build GPU-accelerated phrase boosting tree (btree) to be used with greedy and beam search decoders of ASR models.
+    """
+
+    asr_pretrained_name: Optional[str] = None  # Name of a pretrained model
+    asr_model_path: Optional[str] = None  # The path to '.nemo' ASR checkpoint
+    save_to: str = MISSING  # The path to save the GPU-accelerated word boosting graph
+
+    # evaluation of obtained boosting tree with test_sentences (optional)
+    test_boosting_tree: bool = False  # Whether to test the GPU-accelerated word boosting tree after building it
+    test_sentences: List[str] = field(
+        default_factory=list
+    )  # The phrases to test boosting tree ["hello world","nvlink","nvlinz","omniverse cloud now","acupuncture"]
+
+
+@hydra_runner(config_path=None, config_name='BuildWordBoostingTreeConfig', schema=BuildWordBoostingTreeConfig)
+def main(cfg: BuildWordBoostingTreeConfig):
+
+    # 1. load asr model to obtain tokenizer
+    if cfg.asr_model_path is None and cfg.asr_pretrained_name is None:
+        raise ValueError("Either asr_model_path or asr_pretrained_name must be provided")
+    elif cfg.asr_model_path is not None:
+        asr_model = nemo_asr.models.ASRModel.restore_from(cfg.asr_model_path, map_location=torch.device('cpu'))
+    else:
+        asr_model = nemo_asr.models.ASRModel.from_pretrained(cfg.asr_pretrained_name)
+
+    is_aggregate_tokenizer = isinstance(asr_model.tokenizer, AggregateTokenizer)
+
+    # 2. Build GPU-accelerated word boosting tree from config
+    gpu_boosting_model = GPUBoostingTreeModel.from_config(cfg, tokenizer=asr_model.tokenizer)
+
+    # 3. save gpu boosting tree to nemo file
+    gpu_boosting_model.save_to(cfg.save_to)
+
+    # 4. test gpu boosting tree model
+    logging.info("testing gpu boosting tree model...")
+    if cfg.test_boosting_tree and cfg.test_sentences:
+        gpu_boosting_model_loaded = GPUBoostingTreeModel.from_nemo(
+            cfg.save_to, vocab_size=len(asr_model.tokenizer.vocab), use_triton=cfg.use_triton
+        )
+        device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+        gpu_boosting_model_loaded = gpu_boosting_model_loaded.cuda()
+
+        if not is_aggregate_tokenizer:
+            sentences_ids = [asr_model.tokenizer.text_to_ids(sentence) for sentence in cfg.test_sentences]
+            sentences_tokens = [asr_model.tokenizer.text_to_tokens(sentence) for sentence in cfg.test_sentences]
+        else:
+            sentences_ids = [
+                asr_model.tokenizer.text_to_ids(sentence, cfg.source_lang) for sentence in cfg.test_sentences
+            ]
+            sentences_tokens = []  # aggregate tokenizer does not support text_to_tokens
+
+        boosting_scores = gpu_boosting_model_loaded(
+            labels=pad_sequence([torch.LongTensor(sentence) for sentence in sentences_ids], batch_first=True).to(
+                device
+            ),
+            labels_lengths=torch.LongTensor([len(sentence) for sentence in sentences_ids]).to(device),
+            bos=False,
+            eos=False if not is_aggregate_tokenizer else True,
+        )
+
+        logging.info(f"[info]: boosting_scores: {boosting_scores}")
+        logging.info(f"[info]: test_sentences: {cfg.test_sentences}")
+        logging.info(f"[info]: test_sentences_tokens: {sentences_tokens}")
+        logging.info(f"[info]: test_sentences_ids: {sentences_ids}")
+
+
+if __name__ == '__main__':
+    main()

scripts/asr_context_biasing/compute_key_words_fscore.py

@@ -0,0 +1,50 @@
+# 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.
+
+import argparse
+from nemo.collections.asr.parts import context_biasing
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--input_manifest",
+        type=str,
+        required=True,
+        help="manifest with recognition results",
+    )
+    parser.add_argument("--key_words_file", type=str, required=True, help="file of key words for fscore calculation")
+    parser.add_argument(
+        "--ctcws-mode",
+        action='store_true',
+        help="whether to use ctcws mode to split the key words from transcriptions",
+    )
+
+    args = parser.parse_args()
+
+    key_words_list = []
+    with open(args.key_words_file, encoding='utf-8') as f:
+        for line in f.readlines():
+            if args.ctcws_mode:
+                item = line.strip().split("_")[0].lower()
+            else:
+                item = line.strip().lower()
+            if item not in key_words_list:
+                key_words_list.append(item)
+
+    context_biasing.compute_fscore(args.input_manifest, key_words_list, print_stats=True)
+
+
+if __name__ == '__main__':
+    main()

scripts/asr_context_biasing/eval_greedy_decoding_with_context_biasing.py

@@ -0,0 +1,511 @@
+# 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 evaluates CTC and Transducer (RNNT) models (only Hybrid Transducer-CTC in case of Transducer) in context biasing mode
+# by applying CTC-based Word Spotter (paper link) 
+
+# Config Help
+
+To discover all arguments of the script, please run :
+python eval_greedy_decoding_with_context_biasing.py --help
+python eval_greedy_decoding_with_context_biasing.py --cfg job
+
+# USAGE
+
+python eval_greedy_decoding_with_context_biasing.py \
+            nemo_model_file=<path to the .nemo file of the model> \
+            input_manifest=<path to the evaluation JSON manifest file \
+            preds_output_folder=<folder to store the predictions> \
+            decoder_type=<type of model decoder [ctc or rnnt]> \
+            acoustic_batch_size=<batch size to calculate log probabilities> \
+            apply_context_biasing=<True or False to apply context biasing> \
+            context_file=<path to the context biasing file with key words/phrases> \
+            beam_threshold=[<list of the beam thresholds, separated with commas>] \
+            context_score=[<list of the context scores, separated with commas>] \
+            ctc_ali_token_weight=[<list of the ctc alignment token weights, separated with commas>] \
+            ...
+
+# Description of context biasing graph:
+Context biasing file contains words/phrases with their spellings
+(one word/phrase per line, spellings are separated from word/phrase by underscore symbol):
+WORD1_SPELLING1
+WORD2_SPELLING1_SPELLING2
+...
+nvidia_nvidia
+gpu_gpu_g p u
+nvlink_nvlink_nv link
+...
+alternative spellings help to improve the recognition accuracy of abbreviations and complicated words,
+which are often recognized as separate words (gpu -> g p u, nvlink -> nv link, tensorrt -> tensor rt, and so on).
+
+
+# Grid Search for Hyper parameters
+
+For grid search, you can provide a list of arguments as follows -
+
+            beam_threshold=[4.0,5.0,6.0,....] \
+            context_score=[1.0,1.5,...,4.0,4.5] \
+            ctc_ali_token_weight=[0.1,0.2,...,0.7,0.8] \
+
+"""
+
+
+import contextlib
+import json
+import os
+import tempfile
+from dataclasses import dataclass, field, is_dataclass
+from pathlib import Path
+from typing import Dict, Optional
+
+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 EncDecCTCModelBPE, EncDecHybridRNNTCTCModel
+from nemo.collections.asr.parts import context_biasing
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+
+
+@dataclass
+class EvalContextBiasingConfig:
+    """
+    Evaluate CTC and Transducer (RNNT) ASR models in greedy decoding with context biasing.
+    """
+
+    # # 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
+    preds_output_folder: str = MISSING  # The folder where the predictions are stored
+
+    # 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
+    decoder_type: Optional[str] = None  # [ctc, rnnt] decoder type for asr model
+
+    # Context-Biasing params
+    apply_context_biasing: bool = False  # True in case of context biasing
+    context_file: str = MISSING  # text file with context biasing words and their spellings
+    spelling_separator: str = "_"  # separator between word and its spellings in context biasing file
+    beam_threshold: list[float] = field(default_factory=lambda: [5.0])  # beam pruning threshold for ctc-ws decoding
+    context_score: list[float] = field(default_factory=lambda: [3.0])  # per token weight for context biasing words
+    ctc_ali_token_weight: list[float] = field(
+        default_factory=lambda: [0.6]
+    )  # weight of CTC tokens to prevent false accept errors
+    print_cb_stats: bool = False  # print context biasing stats (mostly for debugging)
+
+    # Auxiliary parameters
+    sort_logits: bool = True  # do logits sorting before decoding - it reduces computation on puddings
+    softmax_temperature: float = 1.00
+    preserve_alignments: bool = False
+
+
+def decoding_step(
+    asr_model: nemo_asr.models.ASRModel,
+    cfg: EvalContextBiasingConfig,
+    encoder_outputs: list[torch.Tensor],
+    ctc_logprobs: list[np.ndarray],
+    target_transcripts: list[str],
+    audio_file_paths: list[str],
+    durations: list[str],
+    preds_output_manifest: str,
+    beam_batch_size: int = 128,
+    progress_bar: bool = True,
+    context_graph: context_biasing.ContextGraphCTC = None,
+    blank_idx: int = 0,
+    hp: Optional[Dict] = None,
+) -> tuple[float, float]:
+
+    # run CTC-based Word Spotter:
+    if cfg.apply_context_biasing:
+        ws_results = {}
+        for idx, logits in tqdm(
+            enumerate(ctc_logprobs), desc=f"Eval CTC-based Word Spotter...", ncols=120, total=len(ctc_logprobs)
+        ):
+            ws_results[audio_file_paths[idx]] = context_biasing.run_word_spotter(
+                logits,
+                context_graph,
+                asr_model,
+                blank_idx=blank_idx,
+                beam_threshold=hp['beam_threshold'],
+                cb_weight=hp['context_score'],
+                ctc_ali_token_weight=hp['ctc_ali_token_weight'],
+            )
+
+    level = logging.getEffectiveLevel()
+    logging.setLevel(logging.CRITICAL)
+    # reset config
+    asr_model.change_decoding_strategy(None)
+
+    # preserve alignment:
+    asr_model.cfg.decoding.preserve_alignments = cfg.preserve_alignments
+
+    # update model's decoding strategy config
+    if isinstance(asr_model, EncDecCTCModelBPE):
+        # in case of ctc
+        asr_model.cfg.decoding.strategy = "greedy"
+    else:
+        # in case of rnnt
+        asr_model.cfg.decoding.strategy = "greedy_batch"
+        # fast greedy batch decoding:
+        asr_model.cfg.decoding.greedy.loop_labels = True
+
+    # update model's decoding strategy
+    asr_model.change_decoding_strategy(asr_model.cfg.decoding)
+    logging.setLevel(level)
+
+    wer_dist_first = cer_dist_first = 0
+    words_count = chars_count = sample_idx = 0
+
+    out_manifest = open(preds_output_manifest, 'w', encoding='utf_8', newline='\n')
+
+    # ctc part for both EncDecCTCModelBPE and EncDecHybridRNNTCTCModel
+    if cfg.decoder_type == "ctc":
+        for batch_idx, probs in enumerate(ctc_logprobs):
+            preds = np.argmax(probs, axis=1)
+            if cfg.apply_context_biasing and ws_results[audio_file_paths[batch_idx]]:
+                # make new text by mearging alignment with ctc-ws predictions:
+                if cfg.print_cb_stats:
+                    logging.info("\n" + "********" * 10)
+                    logging.info(f"File name: {audio_file_paths[batch_idx]}")
+                pred_text, raw_text = context_biasing.merge_alignment_with_ws_hyps(
+                    preds,
+                    asr_model,
+                    ws_results[audio_file_paths[batch_idx]],
+                    decoder_type="ctc",
+                    blank_idx=blank_idx,
+                    print_stats=cfg.print_cb_stats,
+                )
+                if cfg.print_cb_stats:
+                    logging.info(f"raw text: {raw_text}")
+                    logging.info(f"hyp text: {pred_text}")
+                    logging.info(f"ref text: {target_transcripts[batch_idx]}")
+            else:
+                preds_tensor = torch.tensor(preds, device='cpu').unsqueeze(0)
+                if isinstance(asr_model, EncDecHybridRNNTCTCModel):
+                    hyp = asr_model.ctc_decoding.ctc_decoder_predictions_tensor(preds_tensor)[0]
+                else:
+                    hyp = asr_model.wer.decoding.ctc_decoder_predictions_tensor(preds_tensor)[0]
+                pred_text = 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_first += wer_dist
+            cer_dist_first += cer_dist
+            words_count += len(target_split_w)
+            chars_count += len(target_split_c)
+
+            if preds_output_manifest:
+                item = {
+                    'audio_filepath': audio_file_paths[batch_idx],
+                    'duration': durations[batch_idx],
+                    'text': target_transcripts[batch_idx],
+                    'pred_text': pred_text,
+                    'wer': f"{wer_dist/len(target_split_w):.4f}",
+                }
+                print(json.dumps(item), file=out_manifest)
+        out_manifest.close()
+
+        return wer_dist_first / words_count, cer_dist_first / chars_count
+
+    # rnnt part for EncDecHybridRNNTCTCModel
+    else:
+        if progress_bar:
+            description = "Greedy_batch decoding.."
+            it = tqdm(range(int(np.ceil(len(encoder_outputs) / beam_batch_size))), desc=description, ncols=120)
+        else:
+            it = range(int(np.ceil(len(encoder_outputs) / beam_batch_size)))
+        for batch_idx in it:
+            probs_batch = encoder_outputs[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
+                    )
+                best_hyp_batch = asr_model.decoding.rnnt_decoder_predictions_tensor(
+                    packed_batch,
+                    probs_lens,
+                    return_hypotheses=True,
+                )
+            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)
+                for candidate_idx, candidate in enumerate(beams):
+                    if cfg.apply_context_biasing and ws_results[audio_file_paths[sample_idx + beams_idx]]:
+                        # make new text by mearging alignment with ctc-ws predictions:
+                        if cfg.print_cb_stats:
+                            logging.info("\n" + "********" * 10)
+                            logging.info(f"File name: {audio_file_paths[batch_idx]}")
+                        pred_text, raw_text = context_biasing.merge_alignment_with_ws_hyps(
+                            candidate,
+                            asr_model,
+                            ws_results[audio_file_paths[sample_idx + beams_idx]],
+                            decoder_type="rnnt",
+                            blank_idx=blank_idx,
+                            print_stats=cfg.print_cb_stats,
+                        )
+                        if cfg.print_cb_stats:
+                            logging.info(f"raw text: {raw_text}")
+                            logging.info(f"hyp text: {pred_text}")
+                            logging.info(f"ref text: {target_transcripts[sample_idx + beams_idx]}")
+                    else:
+                        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']
+
+                    if candidate_idx == 0:
+                        # first candidate
+                        wer_dist_tosave = wer_dist
+                        wer_dist_first += wer_dist
+                        cer_dist_first += cer_dist
+
+                # write manifest with prediction results
+                alignment = []
+
+                if preds_output_manifest:
+                    item = {
+                        'audio_filepath': audio_file_paths[sample_idx + beams_idx],
+                        'duration': durations[sample_idx + beams_idx],
+                        'text': target_transcripts[sample_idx + beams_idx],
+                        'pred_text': pred_text,
+                        'wer': f"{wer_dist_tosave/len(target_split_w):.3f}",
+                        'alignment': f"{alignment}",
+                    }
+                    print(json.dumps(item), file=out_manifest)
+
+            sample_idx += len(probs_batch)
+        out_manifest.close()
+
+        return wer_dist_first / words_count, cer_dist_first / chars_count
+
+
+@hydra_runner(config_path=None, config_name='EvalContextBiasingConfig', schema=EvalContextBiasingConfig)
+def main(cfg: EvalContextBiasingConfig):
+    if is_dataclass(cfg):
+        cfg = OmegaConf.structured(cfg)
+
+    assert os.path.isfile(cfg.input_manifest), f"input_manifest {cfg.input_manifest} does not exist"
+    assert cfg.context_file, "context_file must be provided for f-score computation"
+    assert os.path.isfile(cfg.context_file), f"context_file {cfg.context_file} does not exist"
+    assert cfg.decoder_type in ["ctc", "rnnt"], "decoder_type must be ctc or rnnt"
+    assert cfg.preds_output_folder, "preds_output_folder must be provided"
+    assert os.path.isdir(cfg.preds_output_folder), f"preds_output_folder {cfg.preds_output_folder} does not exist"
+
+    # load nemo asr model
+    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 not isinstance(asr_model, (EncDecCTCModelBPE, EncDecHybridRNNTCTCModel)):
+        raise ValueError("ASR model must be CTC BPE or Hybrid Transducer-CTC")
+
+    # load nemo manifest
+    target_transcripts = []
+    durations = []
+    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'])
+            durations.append(data['duration'])
+            audio_file_paths.append(str(audio_file.absolute()))
+
+    # 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
+            encoder_outputs = []
+            ctc_logprobs = []
+            if isinstance(asr_model, EncDecCTCModelBPE):
+                # in case of EncDecCTCModelBPE
+                hyp_results = asr_model.transcribe(
+                    audio_file_paths, batch_size=cfg.acoustic_batch_size, return_hypotheses=True
+                )
+                ctc_logprobs = [hyp.alignments.cpu().numpy() for hyp in hyp_results]
+                blank_idx = asr_model.decoding.blank_id
+            else:
+                # in case of EncDecHybridRNNTCTCModel
+                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="Getting encoder and CTC decoder outputs...", disable=False
+                    ):
+                        encoded, encoded_len = asr_model.forward(
+                            input_signal=test_batch[0].to(device), input_signal_length=test_batch[1].to(device)
+                        )
+                        ctc_dec_outputs = asr_model.ctc_decoder(encoder_output=encoded).cpu()
+                        # dump encoder embeddings per file
+                        for idx in range(encoded.shape[0]):
+                            encoded_no_pad = encoded[idx, :, : encoded_len[idx]]
+                            ctc_dec_outputs_no_pad = ctc_dec_outputs[idx, : encoded_len[idx]]
+                            encoder_outputs.append(encoded_no_pad)
+                            ctc_logprobs.append(ctc_dec_outputs_no_pad.cpu().numpy())
+                    blank_idx = asr_model.decoder.blank_idx
+
+    # load context biasing words
+    context_transcripts = []
+    for line in open(cfg.context_file).readlines():
+        item = line.strip().lower().split(cfg.spelling_separator)
+        word = item[0]
+        word_tokenization = [asr_model.tokenizer.text_to_ids(x) for x in item[1:]]
+        context_transcripts.append([word, word_tokenization])
+    context_words = [item[0] for item in context_transcripts]
+    # build context graph:
+    if cfg.apply_context_biasing:
+        context_graph = context_biasing.ContextGraphCTC(blank_id=blank_idx)
+        context_graph.add_to_graph(context_transcripts)
+    else:
+        context_graph = None
+
+    # sort encoder_outputs according to length:
+    if cfg.decoder_type == "rnnt" and cfg.sort_logits:
+        encoder_outputs_with_indeces = sorted(enumerate(encoder_outputs), key=lambda x: x[1].size()[1], reverse=True)
+        encoder_outputs_sorted = []
+        target_transcripts_sorted = []
+        audio_file_paths_sorted = []
+        durations_sorted = []
+        ctc_logprobs_sorted = []
+        for pair in encoder_outputs_with_indeces:
+            encoder_outputs_sorted.append(pair[1])
+            target_transcripts_sorted.append(target_transcripts[pair[0]])
+            audio_file_paths_sorted.append(audio_file_paths[pair[0]])
+            durations_sorted.append(durations[pair[0]])
+            ctc_logprobs_sorted.append(ctc_logprobs[pair[0]])
+        encoder_outputs = encoder_outputs_sorted
+        target_transcripts = target_transcripts_sorted
+        audio_file_paths = audio_file_paths_sorted
+        durations = durations_sorted
+        ctc_logprobs = ctc_logprobs_sorted
+
+    # setup search parameters grid
+    params = {
+        'beam_threshold': cfg.beam_threshold,
+        'context_score': cfg.context_score,
+        'ctc_ali_token_weight': cfg.ctc_ali_token_weight,
+    }
+    hp_grid = ParameterGrid(params)
+    hp_grid = list(hp_grid)
+
+    logging.info(f"=========================Starting the decoding========================")
+    logging.info(f"Grid search size: {len(hp_grid)}")
+    logging.info(f"It may take some time...")
+    logging.info(f"======================================================================")
+
+    asr_model = asr_model.to('cpu')
+    best_wer = 1e6
+
+    # run decoding step for each hyper parameter set
+    for hp in hp_grid:
+        results_file = f"preds_out_manifest_bthr-{hp['beam_threshold']}_cs-{hp['context_score']}ctcw-{hp['ctc_ali_token_weight']}.json"
+        preds_output_manifest = os.path.join(cfg.preds_output_folder, results_file)
+        candidate_wer, candidate_cer = decoding_step(
+            asr_model,
+            cfg,
+            encoder_outputs=encoder_outputs,
+            target_transcripts=target_transcripts,
+            audio_file_paths=audio_file_paths,
+            durations=durations,
+            beam_batch_size=cfg.beam_batch_size,
+            progress_bar=True,
+            preds_output_manifest=preds_output_manifest,
+            context_graph=context_graph,
+            ctc_logprobs=ctc_logprobs,
+            blank_idx=blank_idx,
+            hp=hp,
+        )
+
+        # compute fscore
+        fscore_stats = context_biasing.compute_fscore(preds_output_manifest, context_words)
+
+        # find the best wer value
+        if candidate_wer < best_wer:
+            best_beam_threshold = hp["beam_threshold"]
+            best_context_score = hp["context_score"]
+            best_ctc_ali_token_weight = hp["ctc_ali_token_weight"]
+            best_wer = candidate_wer
+            best_fscore_stats = fscore_stats
+
+        logging.info(f"======================================================================")
+        logging.info(f"Greedy WER/CER = {candidate_wer:.2%}/{candidate_cer:.2%}")
+        logging.info(f"Precision/Recall/Fscore = {fscore_stats[0]:.4f}/{fscore_stats[1]:.4f}/{fscore_stats[2]:.4f}")
+        logging.info(
+            f"Params: b_thr = {hp['beam_threshold']}, cs = {hp['context_score']}, ctc_ali_weight = {hp['ctc_ali_token_weight']}"
+        )
+        logging.info(f"======================================================================")
+
+    if len(hp_grid) > 1:
+        logging.info(f"=========================Best Results=================================")
+        logging.info(f"Best WER = {best_wer:.2%}")
+        logging.info(
+            f"Best Precision/Recall/Fscore = {best_fscore_stats[0]:.4f}/{best_fscore_stats[1]:.4f}/{best_fscore_stats[2]:.4f}"
+        )
+        logging.info(
+            f"Best beam_threshold = {best_beam_threshold}, context_score = {best_context_score}, ctc_ali_token_weight = {best_ctc_ali_token_weight}"
+        )
+
+
+if __name__ == '__main__':
+    main()

scripts/asr_eou/add_eob_labels.py

@@ -0,0 +1,222 @@
+# 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.
+
+"""
+A script to add EOB labels to a manifest file.
+
+Example usage:
+
+```bash
+python add_eob_labels.py /path/to/manifest/dir
+```
+where output will be saved in the same directory with `-eob` suffix added to the filename.
+"""
+
+import argparse
+import json
+from pathlib import Path
+from string import punctuation
+
+from tqdm import tqdm
+
+parser = argparse.ArgumentParser(description="Add `is_backchannel` labels to manifest files.")
+parser.add_argument(
+    "input_manifest",
+    type=str,
+    help="Path to the input manifest file to be cleaned.",
+)
+parser.add_argument(
+    "-o",
+    "--output",
+    type=str,
+    default=None,
+    help="Path to the output manifest file after cleaning.",
+)
+parser.add_argument(
+    "-p",
+    "--pattern",
+    type=str,
+    default="*.json",
+    help="Pattern to match files in the input directory.",
+)
+
+
+def read_manifest(manifest_path):
+    manifest = []
+    with open(manifest_path, 'r') as f:
+        for line in f.readlines():
+            line = line.strip()
+            if line:
+                manifest.append(json.loads(line))
+    return manifest
+
+
+def write_manifest(manifest_path, manifest):
+    with open(manifest_path, 'w') as f:
+        for item in manifest:
+            f.write(json.dumps(item) + '\n')
+
+
+def clean_text(text):
+    text = text.translate(str.maketrans('', '', punctuation)).lower().strip()
+    valid_chars = "abcdefghijklmnopqrstuvwxyz'"
+    text = ''.join([c for c in text if c in valid_chars or c.isspace() or c == "'"])
+    return " ".join(text.split()).strip()
+
+
+backchannel_phrases = [
+    'absolutely',
+    'ah',
+    'all right',
+    'alright',
+    'but yeah',
+    'definitely',
+    'exactly',
+    'go ahead',
+    'good',
+    'great',
+    'great thanks',
+    'ha ha',
+    'hi',
+    'i know',
+    'i know right',
+    'i see',
+    'indeed',
+    'interesting',
+    'mhmm',
+    'mhmm mhmm',
+    'mhmm right',
+    'mhmm yeah',
+    'mhmm yes',
+    'nice',
+    'of course',
+    'oh',
+    'oh dear',
+    'oh man',
+    'oh okay',
+    'oh wow',
+    'oh yes',
+    'ok',
+    'ok thanks',
+    'okay',
+    'okay okay',
+    'okay thanks',
+    'perfect',
+    'really',
+    'right',
+    'right exactly',
+    'right right',
+    'right yeah',
+    'so yeah',
+    'sounds good',
+    'sure',
+    'thank you',
+    'thanks',
+    "that's awesome",
+    'thats right',
+    'thats true',
+    'true',
+    'uh-huh',
+    'uh-huh yeah',
+    'uhhuh',
+    'um-humm',
+    'well',
+    'what',
+    'wow',
+    'yeah',
+    'yeah i know',
+    'yeah i see',
+    'yeah mhmm',
+    'yeah okay',
+    'yeah right',
+    'yeah uh-huh',
+    'yeah yeah',
+    'yep',
+    'yes',
+    'yes please',
+    'yes yes',
+    'you know',
+    "you're right",
+]
+
+backchannel_phrases_nopc = [clean_text(phrase) for phrase in backchannel_phrases]
+
+
+def check_if_backchannel(text):
+    """
+    Check if the text is a backchannel phrase.
+    """
+    # Remove punctuation and convert to lowercase
+    text = clean_text(text)
+    # Check if the text is in the list of backchannel phrases
+    return text in backchannel_phrases_nopc
+
+
+def add_eob_labels(manifest_path):
+    """
+    Add EOB labels to a manifest file.
+    Args:
+        manifest_path: Path to the manifest file.
+
+    Returns:
+        manifest: List of dictionaries with the EOB label added.
+        num_eob: Number of EOB labels added.
+    """
+    num_eob = 0
+    manifest = read_manifest(manifest_path)
+    for i, item in enumerate(manifest):
+        text = item['text']
+        # Check if the text is a backchannel phrase
+        is_backchannel = check_if_backchannel(text)
+        # Add the EOB label to the text
+        if is_backchannel:
+            item['is_backchannel'] = True
+            num_eob += 1
+        else:
+            item['is_backchannel'] = False
+        manifest[i] = item
+    return manifest, num_eob
+
+
+def main():
+    args = parser.parse_args()
+    input_manifest = Path(args.input_manifest)
+
+    if input_manifest.is_dir():
+        manifest_list = list(input_manifest.glob(args.pattern))
+        if not manifest_list:
+            raise ValueError(f"No files found in {input_manifest} matching pattern `{args.pattern}`")
+    else:
+        manifest_list = [input_manifest]
+
+    if args.output is None:
+        output_dir = input_manifest if input_manifest.is_dir() else input_manifest.parent
+    else:
+        output_dir = Path(args.output)
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+    total_num_eob = 0
+    print(f"Processing {len(manifest_list)} manifest files...")
+    for manifest_path in tqdm(manifest_list, total=len(manifest_list)):
+        output_file = output_dir / f"{manifest_path.stem}-eob.json"
+        new_manifest, num_eob = add_eob_labels(manifest_path)
+        total_num_eob += num_eob
+        write_manifest(output_file, new_manifest)
+        print(f"Processed {manifest_path} and saved to {output_file}. Number of EOB labels added: {num_eob}")
+
+    print(f"Total number of EOB labels added: {total_num_eob}")
+
+
+if __name__ == "__main__":
+    main()

scripts/asr_eou/clean_manifest.py

@@ -0,0 +1,670 @@
+# 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.
+
+"""
+A rule-based text cleaning script for preparing text for ASR-EOU model training.
+
+Example usage:
+
+```bash
+python clean_manifest.py \
+    /path/to/manifest/dir \
+    -o /path/to/output/dir
+```
+
+"""
+
+import argparse
+import re
+import unicodedata
+from datetime import datetime
+from pathlib import Path
+from string import punctuation
+
+from num2words import num2words
+from whisper_normalizer.english import EnglishTextNormalizer
+
+from nemo.collections.asr.parts.utils.manifest_utils import read_manifest, write_manifest
+
+punctuations = punctuation.replace("'", "")
+
+text_normalizer = EnglishTextNormalizer()
+
+parser = argparse.ArgumentParser(description="Clean manifest file")
+parser.add_argument(
+    "input_manifest",
+    type=str,
+    help="Path to the input manifest file to be cleaned.",
+)
+parser.add_argument(
+    "-o",
+    "--output",
+    type=str,
+    default=None,
+    help="Path to the output manifest file after cleaning.",
+)
+parser.add_argument(
+    "-lower",
+    "--lowercase",
+    type=bool,
+    default=False,
+    help="Whether to convert the text to lowercase.",
+)
+parser.add_argument(
+    "-drop",
+    "--remove_punc",
+    type=bool,
+    default=False,
+    help="Whether to remove punctuation from the text.",
+)
+parser.add_argument(
+    "--normalize",
+    type=bool,
+    default=False,
+    help="Whether to normalize the text using Whisper text normalizer.",
+)
+parser.add_argument(
+    "-n2w",
+    "--replace_numbers",
+    type=bool,
+    default=True,
+    help="Whether to replace numbers with words.",
+)
+parser.add_argument(
+    "-p",
+    "--pattern",
+    type=str,
+    default="**/*.json",
+    help="Pattern to match files in the input directory.",
+)
+parser.add_argument(
+    "-t",
+    "--text_field",
+    type=str,
+    default="text",
+    help="Field in the manifest to clean. Default is 'text'.",
+)
+parser.add_argument(
+    "--auto_pc",
+    action="store_true",
+    help="If set, will add auto capitalization and punctuation at the end of the text.",
+)
+parser.add_argument(
+    "--format",
+    default="asr",
+    choices=["asr", "conv"],
+    help="Format of the manifest. Default is 'asr'.",
+)
+parser.add_argument(
+    "--keep_name",
+    action="store_true",
+    help="If set, will keep the original name of the manifest file.",
+)
+
+# Spoken representations
+
+MONTHS = [
+    "",
+    "January",
+    "February",
+    "March",
+    "April",
+    "May",
+    "June",
+    "July",
+    "August",
+    "September",
+    "October",
+    "November",
+    "December",
+]
+
+ORDINALS = {
+    1: "first",
+    2: "second",
+    3: "third",
+    4: "fourth",
+    5: "fifth",
+    6: "sixth",
+    7: "seventh",
+    8: "eighth",
+    9: "ninth",
+    10: "tenth",
+    11: "eleventh",
+    12: "twelfth",
+    13: "thirteenth",
+    14: "fourteenth",
+    15: "fifteenth",
+    16: "sixteenth",
+    17: "seventeenth",
+    18: "eighteenth",
+    19: "nineteenth",
+    20: "twentieth",
+    21: "twenty first",
+    22: "twenty second",
+    23: "twenty third",
+    24: "twenty fourth",
+    25: "twenty fifth",
+    26: "twenty sixth",
+    27: "twenty seventh",
+    28: "twenty eighth",
+    29: "twenty ninth",
+    30: "thirtieth",
+    31: "thirty first",
+}
+
+
+def speak_year(year: int) -> str:
+    if 2000 <= year <= 2099:
+        return f"twenty {speak_number(year % 100)}"
+    elif 1900 <= year <= 1999:
+        return f"nineteen {speak_number(year % 100)}"
+    else:
+        return str(year)
+
+
+def speak_number(n: int) -> str:
+    num_words = {
+        0: "zero",
+        1: "one",
+        2: "two",
+        3: "three",
+        4: "four",
+        5: "five",
+        6: "six",
+        7: "seven",
+        8: "eight",
+        9: "nine",
+        10: "ten",
+        11: "eleven",
+        12: "twelve",
+        13: "thirteen",
+        14: "fourteen",
+        15: "fifteen",
+        16: "sixteen",
+        17: "seventeen",
+        18: "eighteen",
+        19: "nineteen",
+        20: "twenty",
+        30: "thirty",
+        40: "forty",
+        50: "fifty",
+        60: "sixty",
+        70: "seventy",
+        80: "eighty",
+        90: "ninety",
+    }
+    if n <= 20:
+        return num_words[n]
+    elif n < 100:
+        tens, ones = divmod(n, 10)
+        return f"{num_words[tens * 10]} {num_words[ones]}" if ones else num_words[tens * 10]
+    else:
+        return str(n)
+
+
+def _parse_numeric_date(date_str: str, *, dayfirst: bool):
+    separator = "/" if "/" in date_str else "-"
+    parts = date_str.split(separator)
+
+    if len(parts) == 3 and len(parts[0]) == 4:
+        year, month, day = [int(part) for part in parts]
+    elif len(parts) == 3:
+        first, second, year = [int(part) for part in parts]
+        day, month = (first, second) if dayfirst else (second, first)
+    elif len(parts) == 2:
+        first, second = [int(part) for part in parts]
+        year = datetime.now().year
+        day, month = (first, second) if dayfirst else (second, first)
+    else:
+        raise ValueError(date_str)
+
+    if year < 100:
+        year += 2000 if year < 69 else 1900
+
+    return datetime(year, month, day)
+
+
+def parse_with_auto_dayfirst(date_str: str):
+    try:
+        parsed_us = _parse_numeric_date(date_str, dayfirst=False)
+    except Exception:
+        parsed_us = None
+
+    try:
+        parsed_eu = _parse_numeric_date(date_str, dayfirst=True)
+    except Exception:
+        parsed_eu = None
+
+    if parsed_us is None:
+        return parsed_eu
+    if parsed_eu is None:
+        return parsed_us
+
+    first = int(date_str.split("/" if "/" in date_str else "-")[0])
+    if first > 12 and len(str(first)) != 4:
+        return parsed_eu
+
+    # Default fallback assumes US style for ambiguous numeric dates.
+    return parsed_us
+
+
+def date_to_spoken_string(date_str: str) -> str:
+    parsed = parse_with_auto_dayfirst(date_str)
+    if not parsed:
+        return None
+
+    month = MONTHS[parsed.month]
+    day = ORDINALS[parsed.day]
+    spoken = f"{month} {day} {speak_year(parsed.year)}"
+
+    return spoken
+
+
+def replace_dates_in_text(text: str) -> str:
+    # Regex pattern to match common date formats like:
+    # 5/22, 05/22/2025, 22/05/2025, 2025-05-22
+    date_pattern = r'\b(?:\d{1,4}[-/])?\d{1,2}[-/]\d{1,4}\b'
+
+    def replace_match(match):
+        date_str = match.group(0)
+        spoken = date_to_spoken_string(date_str)
+        return spoken if spoken else date_str
+
+    return re.sub(date_pattern, replace_match, text)
+
+
+def convert_to_spoken(text: str) -> str:
+
+    text = replace_dates_in_text(text)  # Convert dates to spoken form
+
+    # Mapping of metric units to spoken forms
+    unit_map = {
+        "kg": "kilograms",
+        "g": "grams",
+        "mg": "milligrams",
+        "l": "liters",
+        "ml": "milliliters",
+        "cm": "centimeters",
+        "mm": "millimeters",
+        "m": "meters",
+        "km": "kilometers",
+        "°c": "degrees celsius",
+        "°f": "degrees fahrenheit",
+        "oz": "ounces",
+        "lb": "pounds",
+        "lbs": "pounds",
+    }
+
+    # Replace metric units like "12kg" or "5 ml"
+    def replace_metric(match):
+        number = match.group(1)
+        unit = match.group(2).lower()
+        spoken_unit = unit_map.get(unit, unit)
+        return f"{number} {spoken_unit}"
+
+    # Replace time like "5am" or "6PM"
+    def replace_ampm(match):
+        hour = match.group(1)
+        meridiem = match.group(2).lower()
+        return f"{hour} {'a m' if meridiem == 'am' else 'p m'}"
+
+    # Replace time like "1:30pm"
+    def replace_colon_time(match):
+        hour = match.group(1)
+        minute = match.group(2)
+        meridiem = match.group(3).lower()
+        return f"{hour} {minute} {'a m' if meridiem == 'am' else 'p m'}"
+
+    # Convert feet and inches like 5'11" to "5 feet 11 inches"
+    def replace_feet_inches(match):
+        feet = match.group(1)
+        inches = match.group(2)
+        return f"{feet} feet {inches} inches"
+
+    # Convert just feet (e.g., 6') to "6 feet"
+    def replace_feet_only(match):
+        feet = match.group(1)
+        return f"{feet} feet"
+
+    # Convert just inches (e.g., 10") to "10 inches"
+    def replace_inches_only(match):
+        inches = match.group(1)
+        return f"{inches} inches"
+
+    # Apply replacements
+    # First: time with colon (e.g., 1:30pm)
+    text = re.sub(r'\b(\d{1,2}):(\d{2})(am|pm)\b', replace_colon_time, text, flags=re.IGNORECASE)
+
+    # Then: basic am/pm (e.g., 5am)
+    text = re.sub(r'\b(\d{1,2})(am|pm)\b', replace_ampm, text, flags=re.IGNORECASE)
+
+    # Then: replace 1st, 2nd, 3rd, etc
+    text = text.replace("1st", "first")
+    text = text.replace("2nd", "second")
+    text = text.replace("3rd", "third")
+    text = text.replace("@", " at ")
+
+    # Finally: metric units
+    text = re.sub(
+        r'\b(\d+(?:\.\d+)?)\s?(kg|g|mg|l|ml|cm|mm|m|km|°c|°f|oz|lbs?|LB|LBS?)\b',
+        replace_metric,
+        text,
+        flags=re.IGNORECASE,
+    )
+    text = re.sub(r'\b(\d+)\'(\d+)"', replace_feet_inches, text)  # e.g., 5'11"
+    text = re.sub(r'\b(\d+)\'', replace_feet_only, text)  # e.g., 6'
+    text = re.sub(r'(\d+)"', replace_inches_only, text)  # e.g., 10"
+
+    return text
+
+
+def replace_numbers_with_words(text):
+    def convert_number(match):
+        num_str = match.group()
+        original = num_str
+
+        # Remove dollar sign
+        is_dollar = False
+        if num_str.startswith('$'):
+            is_dollar = True
+            num_str = num_str[1:]
+        elif num_str.endswith('$'):
+            is_dollar = True
+            num_str = num_str[:-1]
+
+        # Remove commas
+        num_str = num_str.replace(',', '')
+
+        try:
+            if '.' in num_str:
+                # Convert decimal number
+                integer_part, decimal_part = num_str.split('.')
+                words = num2words(int(integer_part)) + ' point ' + ' '.join(num2words(int(d)) for d in decimal_part)
+            else:
+                words = num2words(int(num_str))
+            if is_dollar:
+                words += ' dollars'
+            return words + " "
+        except Exception:
+            return original  # Return original if conversion fails
+
+    # Pattern matches: $3,000 or 3,000.45 or 1234
+    pattern = re.compile(r'\$?\d{1,3}(?:,\d{3})*(?:\.\d+)?|\$?\d+(?:\.\d+)?')
+    result = pattern.sub(convert_number, text)
+    result = result.replace("$", " dollars ")  # Handle dollar sign separately
+
+    def merge_th(text: str) -> str:
+        # merge th with the preceding digit
+        candidates = ["four th ", "five th ", "six th ", "seven th ", "eight th ", "nine th "]
+        for key in candidates:
+            if key in text:
+                if "five" in key:
+                    target = "fifth "
+                else:
+                    target = f"{key.split(' ')[0]}th "
+                text = text.replace(key, target)
+            elif text.endswith(key.strip()):
+                if "five" in key:
+                    target = "fifth"
+                else:
+                    target = f"{key.split(' ')[0]}th"
+                text = text.replace(key.strip(), target)
+        return text
+
+    result = merge_th(result)
+    result = " ".join(result.split())  # Remove extra spaces
+    return result
+
+
+def unicode_to_ascii(text: str) -> str:
+    """
+    Converts text with accented or special Latin characters (e.g., ó, ñ, ū, ō)
+    into their closest ASCII equivalents.
+    """
+    # Normalize the string to NFKD to separate base characters from diacritics
+    normalized = unicodedata.normalize('NFKD', text)
+
+    # Encode to ASCII bytes, ignoring characters that can't be converted
+    ascii_bytes = normalized.encode('ascii', 'ignore')
+
+    # Decode back to string
+    ascii_text = ascii_bytes.decode('ascii')
+
+    return ascii_text
+
+
+def drop_punctuations(text: str) -> str:
+    """
+    Clean the text by removing invalid characters and converting to lowercase.
+
+    :param text: Input text.
+    :return: Cleaned text.
+    """
+    valid_chars = "abcdefghijklmnopqrstuvwxyz'"
+    text = text.lower()
+    text = unicode_to_ascii(text)
+    text = text.replace(":", " ")
+    text = text.replace("-", " ")
+    text = ''.join([c for c in text if c in valid_chars or c.isspace()])
+    text = ' '.join(text.split()).strip()
+    return text
+
+
+def clean_label(_str: str) -> str:
+    """
+    Remove unauthorized characters in a string, lower it and remove unneeded spaces
+    """
+    # replace_with_space = [char for char in '/?*\",.:=?_{|}~¨«·»¡¿„…‧‹›≪≫!:;ː→']
+    replace_with_blank = [char for char in '`¨´‘’“”`ʻ‘’“"‘”']
+    replace_with_apos = [char for char in '‘’ʻ‘’‘'] + ["\u2019"]
+    _str = _str.strip()
+    for i in replace_with_blank:
+        _str = _str.replace(i, "")
+    for i in replace_with_apos:
+        _str = _str.replace(i, "'")
+
+    text = _str
+    text = text.replace("\u2103", "celsius")
+    text = text.replace("\u2109", "fahrenheit")
+    text = text.replace("\u00b0", "degrees")
+    text = text.replace("\u2019", "'")
+    text = text.replace("\\", ".")
+    text = text.replace("\n", " ")
+    text = text.replace("\r", " ")
+    text = text.replace("\t", " ")
+
+    ret = " ".join(text.split())
+    return ret
+
+
+def ends_with_punctuation(s: str) -> bool:
+    # Strip trailing whitespace
+    s = s.rstrip()
+
+    # consider this set to be punctuation that's acceptable to end a sentence with
+    puncturation_chars = [",", ".", ":", ";", "?", "!", "-", "—", "–", "…"]
+
+    # If string is empty after stripping, return False
+    if not s:
+        return False
+
+    # Get the last character
+    last_char = s[-1]
+
+    # Return True if the last character is punctuation, otherwise False
+    return last_char in puncturation_chars
+
+
+def add_period_if_needed(text: str) -> str:
+    """
+    Add a period at the end of the text if it does not already end with one.
+    """
+    if not ends_with_punctuation(text):
+        text += "."
+    return text.strip()
+
+
+def capitalize_self_i(text):
+    # Replace standalone lowercase "i" with "I"
+    # Handles "i", "i.", "i?", "i'll", "i'm", etc.
+    return re.sub(r'\b(i)(?=[\s.,!?;:\'\"-]|$)', r'I', text)
+
+
+def add_space_after_punctuation(text):
+    # Add a space after punctuation if it's not already followed by one or by the end of the string
+    return re.sub(r'([,\.?;:])(?=\S)', r'\1 ', text)
+
+
+def add_auto_capitalization(text):
+    if text.lower() != text:
+        # If the text is not all lowercase, we assume it has some capitalization
+        return text
+
+    # Remove space before punctuation (.,!?;:)
+    text = re.sub(r'\s+([.,!?;:])', r'\1', text)
+
+    # Capitalize the first letter of each sentence
+    def capitalize_sentences(match):
+        return match.group(1) + match.group(2).upper()
+
+    # Ensure first character is capitalized
+    text = text.strip()
+    if text:
+        text = text[0].upper() + text[1:]
+
+    text = capitalize_self_i(text)
+    text = add_space_after_punctuation(text)
+    # Capitalize after sentence-ending punctuation followed by space(s)
+    text = re.sub(r'([.!?]\s+)([a-z])', capitalize_sentences, text)
+    return text
+
+
+def unicode_to_ascii(text: str) -> str:
+    """
+    Converts text with accented or special Latin characters (e.g., ó, ñ, ū, ō)
+    into their closest ASCII equivalents.
+    """
+    # Normalize the string to NFKD to separate base characters from diacritics
+    normalized = unicodedata.normalize('NFKD', text)
+
+    # Encode to ASCII bytes, ignoring characters that can't be converted
+    ascii_bytes = normalized.encode('ascii', 'ignore')
+
+    # Decode back to string
+    ascii_text = ascii_bytes.decode('ascii')
+
+    return ascii_text
+
+
+def clean_text(text: str, args) -> str:
+    """
+    Clean the text based on the provided arguments.
+    """
+    text = unicode_to_ascii(text)
+    if args.normalize:
+        text = text_normalizer(text)
+    if args.replace_numbers:
+        text = convert_to_spoken(text)
+        text = replace_numbers_with_words(text)
+    if args.lowercase:
+        text = text.lower()
+    if args.remove_punc:
+        text = text.replace("-", " ")
+        text = text.replace("_", " ")
+        text = text.translate(str.maketrans("", "", punctuations))
+        text = drop_punctuations(text)
+    if args.auto_pc:
+        text = add_auto_capitalization(text)
+    return clean_label(text)
+
+
+def clean_asr_manifest(manifest, text_field, args):
+    for i, item in enumerate(manifest):
+        text = str(item[text_field])
+        manifest[i][f"origin_{text_field}"] = text
+        manifest[i][text_field] = clean_text(text, args)
+    return manifest
+
+
+def clean_conv_manifest(manifest, text_field, args):
+    new_manifest = []
+    for i, item in enumerate(manifest):
+        conversations = []
+        for turn in item["conversations"]:
+            conversations.append(
+                {
+                    "role": turn["role"],
+                    "value": clean_text(turn["value"], args),
+                    "type": turn.get("type", "text"),
+                }
+            )
+        item["conversations"] = conversations
+        new_manifest.append(item)
+    return manifest
+
+
+def main(args):
+    text_field = args.text_field
+    manifest_files = Path(args.input_manifest)
+    if manifest_files.is_dir():
+        manifest_files = list(manifest_files.glob(args.pattern))
+    elif manifest_files.is_file():
+        manifest_files = [manifest_files]
+    else:
+        raise ValueError(f"Invalid input manifest path: {args.input_manifest}")
+
+    for manifest_file in manifest_files:
+        print(f"Processing manifest file: {manifest_file}")
+        postfix = "-cleaned"
+        postfix += "_norm" if args.normalize else ""
+        postfix += "_n2w" if args.replace_numbers else ""
+        if args.lowercase and args.remove_punc:
+            postfix += "_noPC"
+        else:
+            postfix += "_lc" if args.lowercase else ""
+            postfix += "_np" if args.remove_punc else ""
+        postfix += "_aPC" if args.auto_pc else ""
+
+        output_manifest = manifest_file.with_name(f"{manifest_file.stem}{postfix}{manifest_file.suffix}")
+
+        if args.output:
+            if args.output.endswith(".json"):
+                if len(manifest_files) > 1:
+                    raise ValueError("Output path must be a directory when processing multiple manifest files.")
+                output_manifest = Path(args.output)
+            else:
+                output_dir = Path(args.output)
+                output_dir.mkdir(parents=True, exist_ok=True)
+                if args.keep_name:
+                    output_manifest = output_dir / manifest_file.name
+                else:
+                    output_manifest = output_dir / output_manifest.name
+
+        manifest = read_manifest(str(manifest_file))
+
+        if args.format == "asr":
+            manifest = clean_asr_manifest(manifest, text_field, args)
+        elif args.format == "conv":
+            manifest = clean_conv_manifest(manifest, text_field, args)
+        else:
+            raise ValueError(f"Unsupported manifest format: {args.format}")
+
+        write_manifest(str(output_manifest), manifest)
+        print(f"Cleaned manifest saved to {output_manifest}")
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    main(args)

scripts/asr_eou/conf/data.yaml

@@ -0,0 +1,46 @@
+
+output_dir: ???
+
+data:
+  pattern: "*.json"
+  manifest_filepath: ???
+  tarred_audio_filepaths: null
+  sample_rate: 16000
+  max_duration: 30 # you may need to update it for your dataset
+  min_duration: 0.1
+  batch_duration: 300  # you may disable batch_duration by setting it to `null`
+  batch_size: null 
+  shuffle: false
+  seed: 42
+  num_workers: 8
+  pin_memory: true
+  quadratic_duration: 30
+  num_buckets: 30
+  num_cuts_for_bins_estimate: 10000
+  bucket_buffer_size: 10000
+  shuffle_buffer_size: 10000
+
+  random_padding:
+      prob: 1.0
+      min_pad_duration: 0.0  # minimum duration of pre/post padding in seconds
+      max_pad_duration: 3.0  # maximum duration of pre/post padding in seconds
+      max_total_duration: 40.0  # maximum total duration of the padded audio in seconds
+      pad_distribution: 'constant'  # distribution of padding duration, 'uniform' or 'normal' or 'constant'
+      pre_pad_duration: 0.2
+      post_pad_duration: 3.0
+
+  augmentor:
+    white_noise:
+      prob: 0.0
+      min_level: -90
+      max_level: -40
+    gain:
+      prob: 0.0
+      min_gain_dbfs: -10.0
+      max_gain_dbfs: 10.0
+    noise:
+      prob: 1.0
+      manifest_path: ???
+      min_snr_db: 0
+      max_snr_db: 20
+      max_gain_db: 300.0

scripts/asr_eou/eval_eou_metrics.py

@@ -0,0 +1,176 @@
+# 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 calculates the EOU metrics using predictions and references in SegLST format.
+
+Example usage:
+
+The PREDICTION_ROOT and REFERENCE_ROOT directories should have the following structure:
+
+<PREDICTION_ROOT>:
+->dataset1/
+    -> sample1.json
+    -> sample2.json
+->dataset2/
+    -> sample1.json
+    -> sample2.json
+
+<REFERENCE_ROOT>:
+->dataset1/
+    -> sample1.json
+    -> sample2.json
+->dataset2/
+    -> sample1.json
+    -> sample2.json
+
+
+each sample.json should contain a list of dictionaries with the following fields:
+{
+    "session_id": str,
+    "start_time": float,  # start time in seconds
+    "end_time": float,  # end time in seconds
+    "words": str,  # transcription of the utterance
+    "audio_filepath": str,  # only in prediction
+    "eou_prob": float, # only in prediction, probability of EOU in range [0.1]
+    "eou_pred": bool, # only in prediction
+    "full_text": str, # only in prediction, which is the full transcription up to the end_time
+}
+    
+```bash
+python eval_eou_metrics.py \
+    --prediction $PREDICTION_ROOT \
+    --reference $REFERENCE_ROOT  \
+    --multiple
+```
+"""
+
+
+import argparse
+import json
+from pathlib import Path
+from typing import List
+
+from nemo.collections.asr.parts.utils.eou_utils import EOUResult, aggregate_eou_metrics, evaluate_eou
+
+parser = argparse.ArgumentParser(description="Evaluate end of utterance predictions against reference labels.")
+parser.add_argument(
+    "-p",
+    "--prediction",
+    type=str,
+    required=True,
+    help="Path to the directory containing the predictions.",
+)
+parser.add_argument(
+    "-r",
+    "--reference",
+    type=str,
+    required=True,
+    help="Path to the directory containing the groundtruth.",
+)
+parser.add_argument(
+    "--eob",
+    action="store_true",
+    help="Whether to evaluate end of backchannel predictions.",
+)
+parser.add_argument(
+    "--ignore_eob",
+    action="store_true",
+    help="Whether to ignore end of backchannel predictions.",
+)
+parser.add_argument(
+    "--multiple",
+    action="store_true",
+    help="Whether to evaluate multiple datasets.",
+)
+
+
+def load_segLST(directory: str, use_eob: bool = False, ignore_eob: bool = False) -> dict:
+    json_files = list(Path(directory).glob("*.json"))
+    segLST = {}
+    for json_file in json_files:
+        key = json_file.stem
+        with open(json_file, 'r') as f:
+            data = json.load(f)
+            assert isinstance(data, list), f"Data in {json_file} is not a list."
+            if not ignore_eob:
+                # get the data with the correct eob label
+                data = [x for x in data if (x.get("is_backchannel", False) == use_eob)]
+            segLST[key] = data
+    return segLST
+
+
+def evaluate_eou_predictions(
+    prediction_dir: str, reference_dir: str, use_eob: bool = False, ignore_eob: bool = False
+) -> List[EOUResult]:
+    prediction_segLST = load_segLST(prediction_dir, use_eob, ignore_eob)
+    reference_segLST = load_segLST(reference_dir, use_eob, ignore_eob)
+
+    eou_metrics = []
+    for key, reference in reference_segLST.items():
+        if key not in prediction_segLST:
+            raise ValueError(f"Key {key} in reference not found in predictions.")
+        prediction = prediction_segLST[key]
+        eou_result = evaluate_eou(
+            prediction=prediction, reference=reference, threshold=None, collar=0.0, do_sorting=True
+        )
+        eou_metrics.append(eou_result)
+
+    results = aggregate_eou_metrics(eou_metrics)
+
+    # add prefix to the keys of the results
+    prefix = Path(reference_dir).stem
+    prefix += "_eob" if use_eob else "_eou"
+    results = {f"{prefix}_{k}": v for k, v in results.items()}
+
+    return results
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+
+    prediction_dir = Path(args.prediction)
+    reference_dir = Path(args.reference)
+
+    if not prediction_dir.is_dir():
+        raise ValueError(f"Prediction directory {prediction_dir} does not exist or is not a directory.")
+    if not reference_dir.is_dir():
+        raise ValueError(f"Reference directory {reference_dir} does not exist or is not a directory.")
+
+    if args.multiple:
+        # get all subdirectories in the prediction and reference directories
+        prediction_dirs = sorted([x for x in prediction_dir.glob("*/") if x.is_dir()])
+        reference_dirs = sorted([x for x in reference_dir.glob("*/") if x.is_dir()])
+        if len(prediction_dirs) != len(reference_dirs):
+            raise ValueError(
+                f"Number of prediction directories {len(prediction_dirs)} must match number of reference directories {len(reference_dirs)}."
+            )
+    else:
+        prediction_dirs = [prediction_dir]
+        reference_dirs = [reference_dir]
+
+    for ref_dir, pred_dir in zip(reference_dirs, prediction_dirs):
+        if args.multiple and ref_dir.stem != pred_dir.stem:
+            raise ValueError(
+                f"Reference directory {ref_dir} and prediction directory {pred_dir} must have the same name."
+            )
+        results = evaluate_eou_predictions(
+            prediction_dir=str(pred_dir), reference_dir=str(ref_dir), use_eob=args.eob, ignore_eob=args.ignore_eob
+        )
+        # Print the results
+        print("==========================================")
+        print(f"Evaluation Results for: {pred_dir} against {ref_dir}")
+        for key, value in results.items():
+            print(f"{key}: {value:.4f}")
+        print("==========================================")

scripts/asr_eou/generate_noisy_eval_data.py

@@ -0,0 +1,224 @@
+# 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 is used to generate noisy evaluation data for ASR and end of utterance detection.
+
+Example usage with a single manifest input:
+python generate_noisy_eval_data.py \
+    --config-path conf/ \
+    --config-name data \
+    output_dir=/path/to/output \
+    data.manifest_filepath=/path/to/manifest.json \
+    data.seed=42 \
+    data.noise.manifest_path /path/to/noise_manifest.json
+
+
+Example usage with multiple manifests matching a pattern:
+python generate_noisy_eval_data.py \
+    --config-path conf/ \
+    --config-name data \
+    output_dir=/path/to/output/dir \
+    data.manifest_filepath=/path/to/manifest/dir/ \
+    data.pattern="*.json" \
+    data.seed=42 \
+    data.noise.manifest_path /path/to/noise_manifest.json
+
+"""
+
+from copy import deepcopy
+from pathlib import Path
+from shutil import rmtree
+
+import librosa
+import lightning.pytorch as pl
+import numpy as np
+import soundfile as sf
+import torch
+import yaml
+from lhotse.cut import MixedCut
+from omegaconf import DictConfig, ListConfig, OmegaConf, open_dict
+from tqdm import tqdm
+
+from nemo.collections.asr.data.audio_to_eou_label_lhotse import LhotseSpeechToTextBpeEOUDataset
+from nemo.collections.asr.parts.utils.manifest_utils import read_manifest, write_manifest
+from nemo.collections.common.data.lhotse import get_lhotse_dataloader_from_config
+from nemo.collections.common.parts.preprocessing import parsers
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+
+
+@hydra_runner(config_path="conf/", config_name="data")
+def main(cfg: DictConfig):
+    """
+    Generate noisy evaluation data for ASR and end of utterance detection.
+
+    Args:
+        cfg: DictConfig object containing the configuration.
+    """
+    # Seed everything for reproducibility
+    seed = cfg.data.get('seed', None)
+    if seed is None:
+        seed = np.random.randint(0, 2**32 - 1)
+        logging.info(f'No seed provided, using random seed: {seed}')
+    logging.info(f'Setting random seed to {seed}')
+    with open_dict(cfg):
+        cfg.data.seed = seed
+    logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}')
+    pl.seed_everything(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+    # Patch data config
+    with open_dict(cfg.data):
+        cfg.data.force_finite = True
+        cfg.data.force_map_dataset = True
+        cfg.data.shuffle = False
+        cfg.data.check_tokenizer = False  # No need to check tokenizer in LhotseSpeechToTextBpeEOUDataset
+
+    # Make output directory
+    output_dir = Path(cfg.output_dir)
+    if output_dir.exists() and cfg.get('overwrite', False):
+        logging.info(f'Removing existing output directory: {output_dir}')
+        rmtree(output_dir)
+    if not output_dir.exists():
+        logging.info(f'Creating output directory: {output_dir}')
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Dump the config to the output directory
+    config = OmegaConf.to_container(cfg, resolve=True)
+    with open(output_dir / 'config.yaml', 'w') as f:
+        yaml.dump(config, f)
+    logging.info(f'Config dumped to {output_dir / "config.yaml"}')
+
+    if isinstance(cfg.data.manifest_filepath, (list, ListConfig)):
+        manifest_list = [Path(x) for x in cfg.data.manifest_filepath]
+    else:
+        input_manifest_file = Path(cfg.data.manifest_filepath)
+        if input_manifest_file.is_dir():
+            pattern = cfg.data.get('pattern', '*.json')
+            manifest_list = list(input_manifest_file.glob(pattern))
+            if not manifest_list:
+                raise ValueError(f"No files found in {input_manifest_file} matching pattern `{pattern}`")
+        else:
+            manifest_list = [Path(x) for x in str(input_manifest_file).split(",")]
+
+    logging.info(f'Found {len(manifest_list)} manifest files to process...')
+
+    for i, manifest_file in enumerate(manifest_list):
+        logging.info(f'[{i+1}/{len(manifest_list)}] Processing {manifest_file}...')
+        data_cfg = deepcopy(cfg.data)
+        data_cfg.manifest_filepath = str(manifest_file)
+        process_manifest(data_cfg, output_dir)
+
+
+def process_manifest(data_cfg: DictConfig, output_dir: Path):
+    """
+    Process a manifest file and generate noisy evaluation data.
+
+    Args:
+        data_cfg: Configuration.
+        output_dir: Output directory.
+    """
+    # Load the input manifest
+    input_manifest = read_manifest(data_cfg.manifest_filepath)
+    logging.info(f'Found {len(input_manifest)} items in input manifest: {data_cfg.manifest_filepath}')
+    manifest_parent_dir = Path(data_cfg.manifest_filepath).parent
+    if Path(input_manifest[0]["audio_filepath"]).is_absolute():
+        output_audio_dir = output_dir / 'wav'
+        flatten_audio_path = True
+    else:
+        output_audio_dir = output_dir
+        flatten_audio_path = False
+
+    if "random_padding" in data_cfg and data_cfg.random_padding.pad_distribution == "constant":
+        is_constant_padding = True
+        pre_pad_dur = data_cfg.random_padding.pre_pad_duration
+    else:
+        is_constant_padding = False
+        pre_pad_dur = None
+
+    # Load the dataset
+    tokenizer = parsers.make_parser()  # dummy tokenizer
+    dataset = LhotseSpeechToTextBpeEOUDataset(cfg=data_cfg, tokenizer=tokenizer, return_cuts=True)
+
+    dataloader = get_lhotse_dataloader_from_config(
+        config=data_cfg,
+        global_rank=0,
+        world_size=1,
+        dataset=dataset,
+        tokenizer=tokenizer,
+    )
+
+    # Generate noisy evaluation data
+    manifest = []
+    for i, batch in enumerate(tqdm(dataloader, desc="Generating noisy evaluation data")):
+        audio_batch, audio_len_batch, cuts_batch = batch
+        audio_batch = audio_batch.cpu().numpy()
+        audio_len_batch = audio_len_batch.cpu().numpy()
+
+        for j in range(len(cuts_batch)):
+            cut = cuts_batch[j]
+            if isinstance(cut, MixedCut):
+                cut = cut.first_non_padding_cut
+
+            manifest_item = {}
+            for k, v in cut.custom.items():
+                if k == "dataloading_info":
+                    continue
+                manifest_item[k] = v
+            audio = audio_batch[j][: audio_len_batch[j]]
+            audio_file = cut.recording.sources[0].source
+
+            if flatten_audio_path:
+                output_audio_file = output_audio_dir / str(audio_file).replace('/', '_')[:255]  # type: Path
+            else:
+                output_audio_file = output_audio_dir / Path(audio_file).relative_to(manifest_parent_dir)  # type: Path
+
+            output_audio_file.parent.mkdir(parents=True, exist_ok=True)
+            sf.write(output_audio_file, audio, dataset.sample_rate)
+
+            manifest_item["audio_filepath"] = str(output_audio_file.relative_to(output_audio_dir))
+            manifest_item["offset"] = 0
+            manifest_item["duration"] = audio.shape[0] / dataset.sample_rate
+
+            if is_constant_padding:
+                # Adjust the sou_time and eou_time for constant padding
+                if 'sou_time' in manifest_item and 'eou_time' in manifest_item:
+                    if not isinstance(manifest_item['sou_time'], list):
+                        manifest_item['sou_time'] = manifest_item['sou_time'] + pre_pad_dur
+                        manifest_item['eou_time'] = manifest_item['eou_time'] + pre_pad_dur
+                    else:
+                        manifest_item['sou_time'] = [x + pre_pad_dur for x in manifest_item['sou_time']]
+                        manifest_item['eou_time'] = [x + pre_pad_dur for x in manifest_item['eou_time']]
+                else:
+                    # add sou_time and eou_time to the manifest item
+                    manifest_item['sou_time'] = pre_pad_dur
+                    manifest_item['eou_time'] = pre_pad_dur + librosa.get_duration(filename=audio_file)
+
+            manifest.append(manifest_item)
+
+    # Write the output manifest
+    output_manifest_file = output_dir / Path(data_cfg.manifest_filepath).name
+    write_manifest(output_manifest_file, manifest)
+    logging.info(f'Output manifest written to {output_manifest_file}')
+
+
+if __name__ == "__main__":
+    main()

scripts/asr_eou/tokenizers/add_special_tokens_to_sentencepiece.py

@@ -0,0 +1,214 @@
+# 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.
+
+import os
+
+os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
+import logging
+import sys
+import tempfile
+from argparse import ArgumentParser
+from pathlib import Path
+
+import sentencepiece as spm
+
+from nemo.collections.asr.data.audio_to_eou_label_lhotse import EOB_STRING, EOU_STRING
+from nemo.core.connectors.save_restore_connector import SaveRestoreConnector
+
+try:
+    import sentencepiece_model_pb2 as spt
+except (ImportError, ModuleNotFoundError):
+    raise Exception("Ensure that sentencepiece_model_pb2.py has been generated from the protoc compiler")
+
+
+SPECIAL_TOKENS = [EOU_STRING, EOB_STRING]
+
+"""Utility to add special tokens to existing sentencepiece models.
+
+Generate sentencepiece_model_pb2.py in the directory of this script before running
+To generate run `protoc --python_out=<path_to_NeMo>/scripts/asr_end_of_utterance/tokenizers sentencepiece_model.proto`
+inside the src folder in sentencepiece repo
+Refer: https://github.com/google/sentencepiece/issues/121
+
+Usage:
+python add_special_tokens_to_sentencepiece.py \
+    --input_file your_model.nemo \
+    --output_dir /path/to/new/tokenizer_dir/
+"""
+
+
+parser = ArgumentParser(description="Add special tokens to sentencepiece model")
+parser.add_argument(
+    "-i",
+    "--input_file",
+    type=str,
+    required=True,
+    help="Path to nemo model file, or sentencepiece model file",
+)
+parser.add_argument(
+    "-o",
+    "--output_dir",
+    type=str,
+    required=True,
+    help="Path to output directory for new tokenizer",
+)
+parser.add_argument(
+    "--tokens",
+    type=str,
+    nargs='+',
+    help="Special tokens to add to tokenizer",
+    default=SPECIAL_TOKENS,
+)
+parser.add_argument(
+    "--extract_only",
+    action="store_true",
+    help="Extract tokenizer without adding special tokens",
+)
+
+
+def extract_nemo_tokenizer(nemo_filepath: str, output_dir: Path) -> str:
+    """
+    Extract a tokenizer from a Nemo file.
+    Args:
+        nemo_filepath: Path to the Nemo file.
+        output_dir: Path to the output directory.
+
+    Returns:
+        tokenizer: Path to the tokenizer file.
+    """
+    SaveRestoreConnector._unpack_nemo_file(path2file=nemo_filepath, out_folder=output_dir)
+    tokenizer = None
+    for file in Path(output_dir).glob("**/*"):
+        if file.is_file() and file.name.endswith("tokenizer.model"):
+            tokenizer = file
+            break
+    if tokenizer is None:
+        raise ValueError(f"Tokenizer not found in {output_dir}: {os.listdir(output_dir)}")
+    return str(tokenizer.absolute())
+
+
+def edit_spt_model(input_file, output_dir, tokens, is_userdefined, extract_only=False):
+    """
+    Edit a sentencepiece model to add special tokens.
+    Args:
+        input_file: Path to the input sentencepiece model file.
+        output_dir: Path to the output directory.
+        tokens: List of special tokens to add.
+        is_userdefined: Whether the special tokens are user-defined.
+        extract_only: Whether to extract the tokenizer only.
+    """
+
+    if extract_only:
+        logging.info("Extracting tokenizer only, no special tokens will be added.")
+
+    output_dir = Path(output_dir)
+
+    if output_dir.exists():
+        logging.warning(f"Output directory {output_dir} already exists. Overwriting it.")
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    output_file = str(output_dir / "tokenizer.model")
+
+    token_type = 3
+    if is_userdefined:
+        token_type = 4
+
+    model = spt.ModelProto()
+    with open(input_file, 'rb') as f:
+        model.ParseFromString(f.read())
+
+    if not extract_only:
+        for token in tokens:
+            piece = model.SentencePiece(piece=token, score=0.0, type=token_type)
+            if piece in model.pieces:
+                logging.error(f"Special Token '{token}' already exists in the input model!")
+                sys.exit(1)
+            model.pieces.append(piece)
+
+    sp = spm.SentencePieceProcessor()
+    sp.LoadFromSerializedProto(model.SerializeToString())
+
+    if not extract_only:
+        try:
+            for token in tokens:
+                id = sp.piece_to_id(token)
+                logging.info(f"Created token '{token}' at ID {id}")
+            logging.info(f"New tokenizer vocab size: {sp.get_piece_size()}")
+        except Exception:
+            logging.error(
+                "Could not appropriately configure new tokenizer. Verify if the special tokens already exist."
+            )
+            sys.exit(1)
+
+    with open(output_file, 'wb') as outf:
+        outf.write(model.SerializeToString())
+    logging.info(f"Created new tokenizer at: {output_file}")
+
+    # Write the vocab to file
+    vocab_file = str(output_dir / "tokenizer.vocab")
+    with open(vocab_file, "w", encoding="utf-8") as f:
+        for i in range(sp.get_piece_size()):
+            piece = sp.id_to_piece(i)
+            score = sp.get_score(i)  # Optional: only available if using newer SentencePiece versions
+            f.write(f"{piece}\t{score}\n")  # Format follows the original vocab format
+    logging.info(f"Created new tokenizer vocab at: {vocab_file}")
+
+    special_tokens = ["<s>", "</s>", "<pad>", "<unk>"]
+    special_tokens.extend(tokens)
+    vocab_txt_file = str(output_dir / "vocab.txt")
+    with open(vocab_txt_file, "w", encoding="utf-8") as f:
+        for i in range(sp.get_piece_size()):
+            piece = sp.id_to_piece(i)
+            if piece in special_tokens:
+                # skip special tokens
+                continue
+            token = piece[1:] if piece.startswith("▁") else f"##{piece}"
+            if len(token) > 0:
+                f.write(f"{token}\n")  # Format follows the original vocab format
+    logging.info(f"Created new tokenizer vocab at: {vocab_txt_file}")
+
+
+def inject_special_tokens(input_file, output_dir, tokens, is_userdefined=True, extract_only=False):
+    """
+    Inject special tokens into a sentencepiece model.
+    NOTE: is_userdefined should be set to True in order for ASR model to work with the new special tokens properly.
+
+    Args:
+        input_file: Path to the input sentencepiece model file.
+        output_dir: Path to the output directory.
+        tokens: List of special tokens to add.
+        is_userdefined: Whether the special tokens are user-defined.
+        extract_only: Whether to extract the tokenizer only.
+    """
+
+    if not os.path.exists(input_file):
+        raise ValueError(f"Input file {input_file} does not exist")
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Check if input file is a Nemo file
+        if input_file.endswith(".nemo"):
+            input_file = extract_nemo_tokenizer(input_file, temp_dir)
+            logging.info(f"Extracted tokenizer from Nemo file: {input_file}")
+        else:
+            input_file = os.path.abspath(input_file)
+            logging.info(f"Using input file: {input_file}")
+
+        edit_spt_model(input_file, output_dir, tokens, is_userdefined, extract_only=extract_only)
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
+    args = parser.parse_args()
+    inject_special_tokens(args.input_file, args.output_dir, args.tokens, extract_only=args.extract_only)

scripts/asr_eou/transcribe_speech_sharded.py

@@ -0,0 +1,343 @@
+# 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.
+
+# Add the examples/asr directory to the Python path so that we can import the transcribe_speech.py file
+import sys
+from pathlib import Path
+
+nemo_root = Path(__file__).parent.parent.parent
+asr_examples_dir = nemo_root / "examples" / "asr"
+sys.path.insert(0, str(asr_examples_dir))
+
+from collections import defaultdict
+from copy import deepcopy
+from dataclasses import dataclass
+from math import ceil
+from pathlib import Path
+from typing import List
+
+from omegaconf import ListConfig
+from tqdm import tqdm
+from transcribe_speech import TranscriptionConfig as SingleTranscribeConfig  # type: ignore
+from transcribe_speech import main as single_transcribe_main  # type: ignore
+
+from nemo.collections.asr.parts.utils.manifest_utils import read_manifest, write_manifest
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+
+"""
+Transcribe audio manifests on distributed GPUs. Useful for transcription of moderate amounts of audio data.
+This script also supports splitting the manifest into chunks and merging the results back together.
+This script is a modified version of `transcribe_speech.py` that only takes manifest files as input.
+It is useful for transcribing a large amount of audio data that does not fit into a single job.
+
+# Arguments
+  model_path: path to .nemo ASR checkpoint
+  pretrained_name: name of pretrained ASR model (from NGC registry)
+  dataset_manifest: path to dataset JSON manifest file (in NeMo formats), can be a comma-separated list of manifest files
+                    or a directory containing manifest files
+  pattern: pattern to glob the manifest files if `dataset_manifest` is a directory
+  output_dir: directory to write the transcriptions
+
+  compute_langs: Bool to request language ID information (if the model supports it)
+  timestamps: Bool to request greedy time stamp information (if the model supports it) by default None 
+
+  (Optionally: You can limit the type of timestamp computations using below overrides)
+  ctc_decoding.ctc_timestamp_type="all"  # (default all, can be [all, char, word, segment])
+  rnnt_decoding.rnnt_timestamp_type="all"  # (default all, can be [all, char, word, segment])
+
+  output_filename: Output filename where the transcriptions will be written
+  batch_size: batch size during inference
+  presort_manifest: sorts the provided manifest by audio length for faster inference (default: True)
+
+  cuda: Optional int to enable or disable execution of model on certain CUDA device.
+  allow_mps: Bool to allow using MPS (Apple Silicon M-series GPU) device if available
+  amp: Bool to decide if Automatic Mixed Precision should be used during inference
+  audio_type: Str filetype of the audio. Supported = wav, flac, mp3
+
+  overwrite_transcripts: Bool which when set allows repeated transcriptions to overwrite previous results.
+
+  ctc_decoding: Decoding sub-config for CTC. Refer to documentation for specific values.
+  rnnt_decoding: Decoding sub-config for RNNT. Refer to documentation for specific values.
+
+  calculate_wer: Bool to decide whether to calculate wer/cer at end of this script
+  clean_groundtruth_text: Bool to clean groundtruth text
+  langid: Str used for convert_num_to_words during groundtruth cleaning
+  use_cer: Bool to use Character Error Rate (CER)  or Word Error Rate (WER)
+
+  calculate_rtfx: Bool to calculate the RTFx throughput to transcribe the input dataset.
+
+# Usage
+ASR model can be specified by either "model_path" or "pretrained_name".
+append_pred - optional. Allows you to add more than one prediction to an existing .json
+pred_name_postfix - optional. The name you want to be written for the current model
+Results are returned in a JSON manifest file.
+
+```bash
+CUDA_VISIBLE_DEVICES=1 python transcribe_speech_distributed.py \
+    model_path=<path to .nemo ASR checkpoint> \
+    dataset_manifest="<remove or path to manifest>" \
+    output_dir="<output directory>" \
+    output_filename="<remove or specify output filename>" \
+    clean_groundtruth_text=True \
+    langid='en' \
+    batch_size=32 \
+    timestamps=False \
+    compute_langs=False \
+    amp=True \
+    append_pred=False \
+    pred_name_postfix="<remove or use another model name for output filename>" \
+    split_size=10000 \
+    num_nodes=1 \
+    node_idx=0 \
+    num_gpus_per_node=1 \
+    gpu_idx=0
+```
+
+If you use Slurm, you can use this params to configure the script:
+```bash
+    gpu_idx=\$SLURM_LOCALID \
+    num_gpus_per_node=\$SLURM_GPUS_ON_NODE \
+    num_nodes=\$SLURM_JOB_NUM_NODES \
+    node_idx=\$SLURM_NODEID
+```
+
+"""
+
+
+@dataclass
+class TranscriptionConfig(SingleTranscribeConfig):
+    """
+    Transcription Configuration for audio to text transcription.
+    """
+
+    # General configs
+    pattern: str = "*.json"
+    output_dir: str = "transcribe_output/"
+
+    # Distributed config
+    num_nodes: int = 1  # total number of nodes
+    node_idx: int = 0  # index of the current node
+    num_gpus_per_node: int = 1  # number of GPUs per node
+    gpu_idx: int = 0  # index of the current GPU
+    bind_gpu_to_cuda: bool = (
+        False  # If False, the script will just do .cuda() on the model, otherwise it will do .to(f"cuda:{gpu_idx}")
+    )
+
+    # handle long manifest
+    split_size: int = -1  # -1 means no split, otherwise split the manifest into chunks of this size
+
+
+def get_unfinished_manifest(manifest_list: List[Path], output_dir: Path):
+    """
+    Get the manifest files that have not finished processing yet, including those that are partly processed.
+
+    Args:
+        manifest_list: List of manifest files to process.
+        output_dir: Directory to write the transcriptions.
+
+    Returns:
+        List of manifest files that have not finished processing yet.
+    """
+    unfinished = []
+    for manifest_file in manifest_list:
+        output_manifest_file = output_dir / manifest_file.name
+        if not output_manifest_file.exists():
+            unfinished.append(manifest_file)
+    return sorted(unfinished)
+
+
+def get_manifest_for_current_rank(
+    manifest_list: List[Path], gpu_id: int = 0, num_gpu: int = 1, node_idx: int = 0, num_node: int = 1
+):
+    """
+    Get the manifest files for the current rank.
+
+    Args:
+        manifest_list: List of manifest files to process.
+        gpu_id: ID of the current GPU.
+        num_gpu: Number of GPUs per node.
+        node_idx: Index of the current node.
+        num_node: Total number of nodes.
+
+    Returns:
+        List of manifest files for the current rank.
+    """
+    node_manifest_list = []
+    assert num_node > 0, f"num_node ({num_node}) must be greater than 0"
+    assert num_gpu > 0, f"num_gpu ({num_gpu}) must be greater than 0"
+    assert 0 <= gpu_id < num_gpu, f"gpu_id ({gpu_id}) must be in range [0, {num_gpu})"
+    assert 0 <= node_idx < num_node, f"node_idx ({node_idx}) must be in range [0, {num_node})"
+    for i, manifest_file in enumerate(manifest_list):
+        if (i + node_idx) % num_node == 0:
+            node_manifest_list.append(manifest_file)
+
+    gpu_manifest_list = []
+    for i, manifest_file in enumerate(node_manifest_list):
+        if (i + gpu_id) % num_gpu == 0:
+            gpu_manifest_list.append(manifest_file)
+    return gpu_manifest_list
+
+
+def maybe_split_manifest(manifest_list: List[Path], cfg: TranscriptionConfig) -> List[Path]:
+    """
+    Split the manifest files into chunks of the specified size.
+
+    Args:
+        manifest_list: List of manifest files to process.
+        cfg: Configuration.
+
+    Returns:
+        List of sharded manifest files.
+    """
+    if cfg.split_size is None or cfg.split_size <= 0:
+        return manifest_list
+
+    all_sharded_manifest_files = []
+    sharded_manifest_dir = Path(cfg.output_dir) / "sharded_manifest_todo"
+    sharded_manifest_dir.mkdir(parents=True, exist_ok=True)
+
+    sharded_manifest_done_dir = Path(cfg.output_dir) / "sharded_manifest_done"
+    sharded_manifest_done_dir.mkdir(parents=True, exist_ok=True)
+    cfg.output_dir = sharded_manifest_done_dir
+
+    logging.info(f"Splitting {len(manifest_list)} manifest files by every {cfg.split_size} samples.")
+    for manifest_file in tqdm(manifest_list, total=len(manifest_list), desc="Splitting manifest files"):
+        manifest = read_manifest(manifest_file)
+
+        num_chunks = ceil(len(manifest) / cfg.split_size)
+        for i in range(num_chunks):
+            chunk_manifest = manifest[i * cfg.split_size : (i + 1) * cfg.split_size]
+            sharded_manifest_file = sharded_manifest_dir / f"{manifest_file.stem}--tpart_{i}.json"
+            write_manifest(sharded_manifest_file, chunk_manifest)
+            all_sharded_manifest_files.append(sharded_manifest_file)
+
+    return all_sharded_manifest_files
+
+
+def maybe_merge_manifest(cfg: TranscriptionConfig):
+    """
+    Merge the sharded manifest files back into the original manifest files and write them to the output directory.
+
+    Args:
+        cfg: Configuration.
+
+    Returns:
+        None.
+    """
+    if cfg.split_size is None or cfg.split_size <= 0:
+        return
+
+    # only merge manifest on the first GPU of the first node
+    if not (cfg.gpu_idx == 0 and cfg.node_idx == 0):
+        return
+
+    sharded_manifest_dir = Path(cfg.output_dir)
+    sharded_manifests = list(sharded_manifest_dir.glob("*--tpart_*.json"))
+    if not sharded_manifests:
+        logging.info(f"No sharded manifest files found in {sharded_manifest_dir}")
+        return
+
+    logging.info(f"Merging {len(sharded_manifests)} sharded manifest files.")
+    manifest_dict = defaultdict(list)
+    for sharded_manifest in sharded_manifests:
+        data_name = sharded_manifest.stem.split("--tpart_")[0]
+        manifest_dict[data_name].append(sharded_manifest)
+
+    output_dir = Path(cfg.output_dir).parent
+    for data_name, sharded_manifest_list in tqdm(
+        manifest_dict.items(), total=len(manifest_dict), desc="Merging manifest files"
+    ):
+        merged_manifest = []
+        for sharded_manifest in sharded_manifest_list:
+            manifest = read_manifest(sharded_manifest)
+            merged_manifest.extend(manifest)
+        output_manifest = output_dir / f"{data_name}.json"
+        write_manifest(output_manifest, merged_manifest)
+    logging.info(f"Merged manifest files saved to {output_dir}")
+
+
+@hydra_runner(config_name="TranscriptionConfig", schema=TranscriptionConfig)
+def run_distributed_transcribe(cfg: TranscriptionConfig):
+    """
+    Run distributed transcription with the given configuration.
+    """
+    logging.info(f"Running distributed transcription with config: {cfg}")
+
+    if cfg.dataset_manifest is None:
+        raise ValueError("`dataset_manifest` is required")
+
+    # load the manifest
+    if isinstance(cfg.dataset_manifest, str) and "," in cfg.dataset_manifest:
+        manifest_list = cfg.dataset_manifest.split(",")
+    elif isinstance(cfg.dataset_manifest, (ListConfig, list)):
+        manifest_list = cfg.dataset_manifest
+    else:
+        input_manifest = Path(cfg.dataset_manifest)
+        if input_manifest.is_dir():
+            manifest_list = list(input_manifest.glob(cfg.pattern))
+        elif input_manifest.is_file():
+            manifest_list = [input_manifest]
+        else:
+            raise ValueError(f"Invalid manifest file or directory: {input_manifest}")
+
+    if not manifest_list:
+        raise ValueError(f"No manifest files found matching pattern: {cfg.pattern} in {input_manifest}")
+
+    manifest_list = maybe_split_manifest(manifest_list, cfg)
+    original_manifest_list = list(manifest_list)
+    logging.info(f"Found {len(manifest_list)} manifest files.")
+
+    output_dir = Path(cfg.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    unfinished_manifest = get_unfinished_manifest(manifest_list, output_dir=output_dir)
+    if not unfinished_manifest:
+        maybe_merge_manifest(cfg)
+        logging.info("All manifest files have been processed. Exiting.")
+        return
+    logging.info(f"Found {len(unfinished_manifest)} unfinished manifest files.")
+
+    manifest_list = get_manifest_for_current_rank(
+        unfinished_manifest,
+        gpu_id=cfg.gpu_idx,
+        num_gpu=cfg.num_gpus_per_node,
+        node_idx=cfg.node_idx,
+        num_node=cfg.num_nodes,
+    )
+    if not manifest_list:
+        logging.info(f"No manifest files found for GPU {cfg.gpu_idx} on node {cfg.node_idx}. Exiting.")
+        return
+
+    logging.info(f"Processing {len(manifest_list)} manifest files with GPU {cfg.gpu_idx} on node {cfg.node_idx}.")
+
+    cfg.cuda = cfg.gpu_idx if cfg.bind_gpu_to_cuda else None
+    for manifest_file in tqdm(manifest_list):
+        logging.info(f"Processing {manifest_file}...")
+        output_filename = output_dir / Path(manifest_file).name
+        curr_cfg = deepcopy(cfg)
+        curr_cfg.dataset_manifest = str(manifest_file)
+        curr_cfg.output_filename = str(output_filename)
+
+        single_transcribe_main(curr_cfg)
+
+    # check if all manifest files have been processed
+    unfinished_manifest = get_unfinished_manifest(original_manifest_list, output_dir=output_dir)
+    if not unfinished_manifest:
+        maybe_merge_manifest(cfg)
+        logging.info("All manifest files have been processed. Exiting.")
+        return
+
+
+if __name__ == '__main__':
+    run_distributed_transcribe()  # noqa pylint: disable=no-value-for-parameter

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()

scripts/audio_to_audio/convert_nemo_to_lhotse.py

@@ -0,0 +1,77 @@
+# 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.
+
+import argparse
+
+from nemo.collections.audio.data.audio_to_audio_lhotse import convert_manifest_nemo_to_lhotse
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Convert an audio-to-audio manifest from NeMo format to Lhotse format. "
+        "This step enables the use of Lhotse datasets for audio-to-audio processing. "
+    )
+    parser.add_argument("input", help='Path to the input NeMo manifest.')
+    parser.add_argument(
+        "output", help="Path where we'll write the output Lhotse manifest (supported extensions: .jsonl.gz and .jsonl)"
+    )
+    parser.add_argument(
+        "-i",
+        "--input_key",
+        default="audio_filepath",
+        help="Key of the input recording, mapped to Lhotse's 'Cut.recording'.",
+    )
+    parser.add_argument(
+        "-t",
+        "--target_key",
+        default="target_filepath",
+        help="Key of the target recording, mapped to Lhotse's 'Cut.target_recording'.",
+    )
+    parser.add_argument(
+        "-r",
+        "--reference_key",
+        default="reference_filepath",
+        help="Key of the reference recording, mapped to Lhotse's 'Cut.reference_recording'.",
+    )
+    parser.add_argument(
+        "-e",
+        "--embedding_key",
+        default="embedding_filepath",
+        help="Key of the embedding, mapped to Lhotse's 'Cut.embedding_vector'.",
+    )
+    parser.add_argument(
+        "-a",
+        "--force_absolute_paths",
+        action='store_true',
+        default=False,
+        help="Force absolute paths in the generated manifests.",
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    convert_manifest_nemo_to_lhotse(
+        input_manifest=args.input,
+        output_manifest=args.output,
+        input_key=args.input_key,
+        target_key=args.target_key,
+        reference_key=args.reference_key,
+        embedding_key=args.embedding_key,
+        force_absolute_paths=args.force_absolute_paths,
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/checkpoint_averaging/average_model_checkpoints.py

@@ -0,0 +1,173 @@
+# 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.
+
+# Copyright 2017 Johns Hopkins University (Shinji Watanabe)
+#
+# 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.
+
+"""
+# Changes to script
+Change the script to import the NeMo model class you would like to load a checkpoint for,
+then update the model constructor to use this model class. This can be found by the line:
+<<< Change model class here ! >>>
+By default, this script imports and creates the `EncDecCTCModelBPE` class but it can be
+changed to any NeMo Model.
+# Run the script
+## Saving a .nemo model file (loaded with ModelPT.restore_from(...))
+HYDRA_FULL_ERROR=1 python average_model_checkpoints.py \
+    --config-path="<path to config directory>" \
+    --config-name="<config name>" \
+    name=<name of the averaged checkpoint> \
+    +checkpoint_dir=<OPTIONAL: directory of checkpoint> \
+    +checkpoint_paths=\"[/path/to/ptl_1.ckpt,/path/to/ptl_2.ckpt,/path/to/ptl_3.ckpt,...]\"
+## Saving an averaged pytorch checkpoint (loaded with torch.load(...))
+HYDRA_FULL_ERROR=1 python average_model_checkpoints.py \
+    --config-path="<path to config directory>" \
+    --config-name="<config name>" \
+    name=<name of the averaged checkpoint> \
+     +checkpoint_dir=<OPTIONAL: directory of checkpoint> \
+    +checkpoint_paths=\"[/path/to/ptl_1.ckpt,/path/to/ptl_2.ckpt,/path/to/ptl_3.ckpt,...]\" \
+    +save_ckpt_only=true
+"""
+
+import os
+
+import lightning.pytorch as pl
+import torch
+from omegaconf import OmegaConf, open_dict
+
+# Change this import to the model you would like to average
+from nemo.collections.asr.models import EncDecCTCModelBPE
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+
+
+def process_config(cfg: OmegaConf):
+    """
+    Process config
+    """
+    if 'name' not in cfg or cfg.name is None:
+        raise ValueError("`cfg.name` must be provided to save a model checkpoint")
+
+    if 'checkpoint_paths' not in cfg or cfg.checkpoint_paths is None:
+        raise ValueError(
+            "`cfg.checkpoint_paths` must be provided as a list of one or more str paths to "
+            "pytorch lightning checkpoints"
+        )
+
+    save_ckpt_only = False
+
+    with open_dict(cfg):
+        name_prefix = cfg.name
+        checkpoint_paths = cfg.pop('checkpoint_paths')
+
+        if 'checkpoint_dir' in cfg:
+            checkpoint_dir = cfg.pop('checkpoint_dir')
+        else:
+            checkpoint_dir = None
+
+        if 'save_ckpt_only' in cfg:
+            save_ckpt_only = cfg.pop('save_ckpt_only')
+
+    if type(checkpoint_paths) not in (list, tuple):
+        checkpoint_paths = str(checkpoint_paths).replace("[", "").replace("]", "")
+        checkpoint_paths = checkpoint_paths.split(",")
+        checkpoint_paths = [ckpt_path.strip() for ckpt_path in checkpoint_paths]
+
+    if checkpoint_dir is not None:
+        checkpoint_paths = [os.path.join(checkpoint_dir, path) for path in checkpoint_paths]
+
+    return name_prefix, checkpoint_paths, save_ckpt_only
+
+
+@hydra_runner(config_path=None, config_name=None)
+def main(cfg):
+    """
+    Main function
+    """
+
+    logging.info("This script is deprecated and will be removed in the 25.01 release.")
+
+    name_prefix, checkpoint_paths, save_ckpt_only = process_config(cfg)
+
+    if not save_ckpt_only:
+        trainer = pl.Trainer(**cfg.trainer)
+
+        # <<< Change model class here ! >>>
+        # Model architecture which will contain the averaged checkpoints
+        # Change the model constructor to the one you would like (if needed)
+        model = EncDecCTCModelBPE(cfg=cfg.model, trainer=trainer)
+
+    """ < Checkpoint Averaging Logic > """
+    # load state dicts
+    n = len(checkpoint_paths)
+    avg_state = None
+
+    logging.info(f"Averaging {n} checkpoints ...")
+
+    for ix, path in enumerate(checkpoint_paths):
+        checkpoint = torch.load(path, map_location='cpu')
+
+        if 'state_dict' in checkpoint:
+            checkpoint = checkpoint['state_dict']
+
+        if ix == 0:
+            # Initial state
+            avg_state = checkpoint
+
+            logging.info(f"Initialized average state dict with checkpoint : {path}")
+        else:
+            # Accumulated state
+            for k in avg_state:
+                avg_state[k] = avg_state[k] + checkpoint[k]
+
+            logging.info(f"Updated average state dict with state from checkpoint : {path}")
+
+    for k in avg_state:
+        if str(avg_state[k].dtype).startswith("torch.int"):
+            # For int type, not averaged, but only accumulated.
+            # e.g. BatchNorm.num_batches_tracked
+            pass
+        else:
+            avg_state[k] = avg_state[k] / n
+
+    # Save model
+    if save_ckpt_only:
+        ckpt_name = name_prefix + '-averaged.ckpt'
+        torch.save(avg_state, ckpt_name)
+
+        logging.info(f"Averaged pytorch checkpoint saved as : {ckpt_name}")
+    else:
+        # Set model state
+        logging.info("Loading averaged state dict in provided model")
+        model.load_state_dict(avg_state, strict=True)
+
+        ckpt_name = name_prefix + '-averaged.nemo'
+        model.save_to(ckpt_name)
+
+        logging.info(f"Averaged model saved as : {ckpt_name}")
+
+
+if __name__ == '__main__':
+    main()

scripts/checkpoint_averaging/checkpoint_averaging.py

@@ -0,0 +1,158 @@
+# 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.
+"""
+Builds a .nemo file with average weights over multiple .ckpt files (assumes .ckpt files in same folder as .nemo file).
+Usage example for building *-averaged.nemo for a given .nemo file:
+NeMo/scripts/checkpoint_averaging/checkpoint_averaging.py my_model.nemo
+Usage example for building *-averaged.nemo files for all results in sub-directories under current path:
+find . -name '*.nemo' | grep -v -- "-averaged.nemo" | xargs NeMo/scripts/checkpoint_averaging/checkpoint_averaging.py
+NOTE: if yout get the following error `AttributeError: Can't get attribute '???' on <module '__main__' from '???'>`
+      use --import_fname_list <FILE> with all files that contains missing classes.
+"""
+
+import argparse
+import glob
+import importlib
+import os
+import sys
+
+import torch
+from tqdm.auto import tqdm
+
+from nemo.core import ModelPT
+from nemo.utils import logging, model_utils
+
+
+def main():
+    """
+    Main function
+    """
+
+    logging.info("This script is deprecated and will be removed in the 25.01 release.")
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        'model_fname_list',
+        metavar='NEMO_FILE_OR_FOLDER',
+        type=str,
+        nargs='+',
+        help='Input .nemo files (or folders who contains them) to parse',
+    )
+    parser.add_argument(
+        '--import_fname_list',
+        metavar='FILE',
+        type=str,
+        nargs='+',
+        default=[],
+        help='A list of Python file names to "from FILE import *"',
+    )
+    parser.add_argument(
+        '--class_path',
+        type=str,
+        default='',
+        help='A path to class "module.submodule.class" (if given)',
+    )
+    args = parser.parse_args()
+
+    logging.info(
+        f"\n\nIMPORTANT:\nIf you get the following error:\n\t"
+        "(AttributeError: Can't get attribute '???' on <module '__main__' from '???'>)\nuse:\n\t"
+        "--import_fname_list\nfor all files that contain missing classes.\n\n"
+    )
+
+    for fn in args.import_fname_list:
+        logging.info(f"Importing * from {fn}")
+        sys.path.insert(0, os.path.dirname(fn))
+        globals().update(importlib.import_module(os.path.splitext(os.path.basename(fn))[0]).__dict__)
+
+    device = torch.device("cpu")
+
+    # loop over all folders with .nemo files (or .nemo files)
+    for model_fname_i, model_fname in enumerate(args.model_fname_list):
+        if not model_fname.endswith(".nemo"):
+            # assume model_fname is a folder which contains a .nemo file
+            nemo_files = list(
+                filter(lambda fn: not fn.endswith("-averaged.nemo"), glob.glob(os.path.join(model_fname, "*.nemo")))
+            )
+            if len(nemo_files) != 1:
+                raise RuntimeError(f"Expected exactly one .nemo file but discovered {len(nemo_files)} .nemo files")
+
+            model_fname = nemo_files[0]
+
+        model_folder_path = os.path.dirname(model_fname)
+        fn, fe = os.path.splitext(model_fname)
+        avg_model_fname = f"{fn}-averaged{fe}"
+
+        logging.info(f"\n===> [{model_fname_i+1} / {len(args.model_fname_list)}] Parsing folder {model_folder_path}\n")
+
+        # restore model from .nemo file path
+        model_cfg = ModelPT.restore_from(restore_path=model_fname, return_config=True)
+        if args.class_path:
+            classpath = args.class_path
+        else:
+            classpath = model_cfg.target  # original class path
+        imported_class = model_utils.import_class_by_path(classpath)
+        logging.info(f"Loading model {model_fname}")
+        nemo_model = imported_class.restore_from(restore_path=model_fname, map_location=device)
+
+        # search for all checkpoints (ignore -last.ckpt)
+        checkpoint_paths = [
+            os.path.join(model_folder_path, x)
+            for x in os.listdir(model_folder_path)
+            if x.endswith('.ckpt') and not x.endswith('-last.ckpt')
+        ]
+        """ < Checkpoint Averaging Logic > """
+        # load state dicts
+        n = len(checkpoint_paths)
+        avg_state = None
+
+        logging.info(f"Averaging {n} checkpoints ...")
+
+        for ix, path in enumerate(tqdm(checkpoint_paths, total=n, desc='Averaging checkpoints')):
+            checkpoint = torch.load(path, map_location=device)
+
+            if 'state_dict' in checkpoint:
+                checkpoint = checkpoint['state_dict']
+            else:
+                raise RuntimeError(f"Checkpoint from {path} does not include a state_dict.")
+
+            if ix == 0:
+                # Initial state
+                avg_state = checkpoint
+
+                logging.info(f"Initialized average state dict with checkpoint:\n\t{path}")
+            else:
+                # Accumulated state
+                for k in avg_state:
+                    avg_state[k] = avg_state[k] + checkpoint[k]
+
+                logging.info(f"Updated average state dict with state from checkpoint:\n\t{path}")
+
+        for k in avg_state:
+            if str(avg_state[k].dtype).startswith("torch.int"):
+                # For int type, not averaged, but only accumulated.
+                # e.g. BatchNorm.num_batches_tracked
+                pass
+            else:
+                avg_state[k] = avg_state[k] / n
+
+        # restore merged weights into model
+        nemo_model.load_state_dict(avg_state, strict=True)
+        # Save model
+        logging.info(f"Saving average model to:\n\t{avg_model_fname}")
+        nemo_model.save_to(avg_model_fname)
+
+
+if __name__ == '__main__':
+    main()

scripts/dataset_processing/add_noise.py

@@ -0,0 +1,171 @@
+# 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.
+
+# USAGE: python add_noise.py --input_manifest=<manifest file of original "clean" dataset>
+#   --noise_manifest=<manifest file poinitng to noise data>
+#   --out_dir=<destination directory for noisy audio and manifests>
+#   --snrs=<list of snrs at which noise should be added to the audio>
+#   --seed=<seed for random number generator>
+#   --num_workers=<number of parallel workers>
+# To be able to reproduce the same noisy dataset, use a fixed seed and num_workers=1
+
+import argparse
+import copy
+import json
+import multiprocessing
+import os
+import random
+
+import numpy as np
+import soundfile as sf
+
+from nemo.collections.asr.parts.preprocessing.perturb import NoisePerturbation
+from nemo.collections.asr.parts.preprocessing.segment import AudioSegment
+
+rng = None
+att_factor = 0.8
+save_noise = False
+sample_rate = 16000
+
+
+def get_out_dir_name(out_dir, input_name, noise_name, snr):
+    return os.path.join(out_dir, input_name, noise_name + "_" + str(snr) + "db")
+
+
+def create_manifest(input_manifest, noise_manifest, snrs, out_path, save_noise):
+    os.makedirs(os.path.join(out_path, "manifests"), exist_ok=True)
+    for snr in snrs:
+        out_dir = get_out_dir_name(
+            out_path,
+            os.path.splitext(os.path.basename(input_manifest))[0],
+            os.path.splitext(os.path.basename(noise_manifest))[0],
+            snr,
+        )
+        out_mfst = os.path.join(
+            os.path.join(out_path, "manifests"),
+            os.path.splitext(os.path.basename(input_manifest))[0]
+            + "_"
+            + os.path.splitext(os.path.basename(noise_manifest))[0]
+            + "_"
+            + str(snr)
+            + "db"
+            + ".json",
+        )
+        with open(input_manifest, "r") as inf, open(out_mfst, "w") as outf:
+            for line in inf:
+                row = json.loads(line.strip())
+                row['audio_filepath'] = os.path.join(out_dir, os.path.basename(row['audio_filepath']))
+                if save_noise:
+                    file_ext = os.path.splitext(row['audio_filepath'])[1]
+                    noise_filename = os.path.basename(row['audio_filepath']).replace(file_ext, "_noise" + file_ext)
+                    row['noise_filepath'] = os.path.join(out_dir, noise_filename)
+                outf.write(json.dumps(row) + "\n")
+
+
+def process_row(row):
+    audio_file = row['audio_filepath']
+    global sample_rate
+    data_orig = AudioSegment.from_file(audio_file, target_sr=sample_rate, offset=0)
+    for snr in row['snrs']:
+        min_snr_db = snr
+        max_snr_db = snr
+        global att_factor
+        perturber = NoisePerturbation(
+            manifest_path=row['noise_manifest'], min_snr_db=min_snr_db, max_snr_db=max_snr_db, rng=rng
+        )
+        out_dir = get_out_dir_name(
+            row['out_dir'],
+            os.path.splitext(os.path.basename(row['input_manifest']))[0],
+            os.path.splitext(os.path.basename(row['noise_manifest']))[0],
+            snr,
+        )
+        os.makedirs(out_dir, exist_ok=True)
+        out_f = os.path.join(out_dir, os.path.basename(audio_file))
+        if os.path.exists(out_f):
+            continue
+        data = copy.deepcopy(data_orig)
+        perturber.perturb(data)
+
+        max_level = np.max(np.abs(data.samples))
+
+        norm_factor = att_factor / max_level
+        new_samples = norm_factor * data.samples
+        sf.write(out_f, new_samples.transpose(), sample_rate)
+
+        global save_noise
+        if save_noise:
+            noise_samples = new_samples - norm_factor * data_orig.samples
+            out_f_ext = os.path.splitext(out_f)[1]
+            out_f_noise = out_f.replace(out_f_ext, "_noise" + out_f_ext)
+            sf.write(out_f_noise, noise_samples.transpose(), sample_rate)
+
+
+def add_noise(infile, snrs, noise_manifest, out_dir, num_workers=1):
+    allrows = []
+
+    with open(infile, "r") as inf:
+        for line in inf:
+            row = json.loads(line.strip())
+            row['snrs'] = snrs
+            row['out_dir'] = out_dir
+            row['noise_manifest'] = noise_manifest
+            row['input_manifest'] = infile
+            allrows.append(row)
+    pool = multiprocessing.Pool(num_workers)
+    pool.map(process_row, allrows)
+    pool.close()
+    print('Done!')
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--input_manifest",
+        type=str,
+        required=True,
+        help="clean test set",
+    )
+    parser.add_argument("--noise_manifest", type=str, required=True, help="path to noise manifest file")
+    parser.add_argument("--out_dir", type=str, required=True, help="destination directory for audio and manifests")
+    parser.add_argument("--snrs", type=int, nargs="+", default=[0, 10, 20, 30])
+    parser.add_argument("--seed", type=int, default=None)
+    parser.add_argument("--num_workers", default=1, type=int)
+    parser.add_argument("--sample_rate", default=16000, type=int)
+    parser.add_argument(
+        "--attenuation_factor",
+        default=0.8,
+        type=float,
+        help="Attenuation factor applied on the normalized noise-added samples before writing to wave",
+    )
+    parser.add_argument(
+        "--save_noise", default=False, action="store_true", help="save the noise added to the input signal"
+    )
+
+    args = parser.parse_args()
+    global sample_rate
+    sample_rate = args.sample_rate
+    global att_factor
+    att_factor = args.attenuation_factor
+    global save_noise
+    save_noise = args.save_noise
+    global rng
+    rng = args.seed
+    num_workers = args.num_workers
+
+    add_noise(args.input_manifest, args.snrs, args.noise_manifest, args.out_dir, num_workers=num_workers)
+    create_manifest(args.input_manifest, args.noise_manifest, args.snrs, args.out_dir, args.save_noise)
+
+
+if __name__ == '__main__':
+    main()

scripts/dataset_processing/fisher_audio_to_wav.py

@@ -0,0 +1,122 @@
+# 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.
+
+# USAGE:
+# python fisher_audio_to_wav.py \
+#   --data_root=<FisherEnglishTrainingSpeech root> \
+#   --dest_root=<destination dir root>
+#
+# Converts all .sph audio files in the Fisher dataset to .wav.
+# Requires sph2pipe to be installed.
+import argparse
+import concurrent.futures
+import glob
+import logging
+import os
+import subprocess
+
+from tqdm import tqdm
+
+parser = argparse.ArgumentParser(description='Convert Fisher .sph to .wav')
+parser.add_argument(
+    "--data_root",
+    default=None,
+    type=str,
+    required=True,
+    help="The path to the root Fisher dataset folder.",
+)
+parser.add_argument(
+    "--dest_root",
+    default=None,
+    type=str,
+    required=True,
+    help="Path to the destination root directory.",
+)
+args = parser.parse_args()
+
+
+def __convert_audio(in_path, out_path):
+    """
+    Helper function that's called per thread, converts sph to wav.
+    Args:
+        in_path: source sph file to convert
+        out_path: destination for wav file
+    """
+    cmd = ["sph2pipe", "-f", "wav", "-p", in_path, out_path]
+    subprocess.run(cmd)
+
+
+def __process_set(data_root, dst_root):
+    """
+    Finds and converts all sph audio files in the given directory to wav.
+    Args:
+        data_folder: source directory with sph files to convert
+        dst_root: where wav files will be stored
+    """
+    sph_list = glob.glob(data_root)
+
+    if not os.path.exists(dst_root):
+        os.makedirs(dst_root)
+
+    # Set up and execute concurrent audio conversion
+    tp = concurrent.futures.ProcessPoolExecutor(max_workers=64)
+    futures = []
+
+    for sph_path in tqdm(sph_list, desc="Submitting sph futures", unit="file"):
+        audio_id, _ = os.path.splitext(os.path.basename(sph_path))
+        out_path = os.path.join(dst_root, "{}.wav".format(audio_id))
+        futures.append(tp.submit(__convert_audio, sph_path, out_path))
+
+    pbar = tqdm(total=len(sph_list), desc="Converting sph files", unit="file")
+    count = 0
+    for f in concurrent.futures.as_completed(futures):
+        count += 1
+        pbar.update()
+    tp.shutdown()
+    pbar.close()
+
+
+def main():
+    data_root = args.data_root
+    dest_root = args.dest_root
+
+    logging.info("\n\nConverting audio for Part 1")
+    __process_set(
+        os.path.join(
+            data_root,
+            "LDC2004S13-Part1",
+            "fisher_eng_tr_sp_d*",
+            "audio",
+            "*",
+            "*.sph",
+        ),
+        os.path.join(dest_root, "LDC2004S13-Part1", "audio_wav"),
+    )
+
+    logging.info("\n\nConverting audio for Part 2")
+    __process_set(
+        os.path.join(
+            data_root,
+            "LDC2005S13-Part2",
+            "fe_03_p2_sph*",
+            "audio",
+            "*",
+            "*.sph",
+        ),
+        os.path.join(dest_root, "LDC2005S13-Part2", "audio_wav"),
+    )
+
+
+if __name__ == '__main__':
+    main()

scripts/dataset_processing/g2p/convert_cmu_arpabet_to_ipa.py

@@ -0,0 +1,141 @@
+# Copyright (c) 2022, 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.
+
+import itertools
+import os
+from argparse import ArgumentParser
+from typing import Dict
+
+from syllabify import syllabify
+
+
+"""
+Usage:
+    cd NeMo/scripts && python dataset_processing/g2p/convert_cmu_arpabet_to_ipa.py
+"""
+
+
+def parse_args():
+    parser = ArgumentParser("ARPABET to IPA conversion sctipt")
+    parser.add_argument(
+        '--cmu_arpabet',
+        help="Path to CMU ARPABET dictionary file",
+        type=str,
+        default="tts_dataset_files/cmudict-0.7b_nv22.10",
+    )
+    parser.add_argument("--ipa_out", help="Path to save IPA version of the dictionary", type=str, required=True)
+    parser.add_argument(
+        "--mapping",
+        help="ARPABET to IPA phoneme mapping file",
+        type=str,
+        default="tts_dataset_files/cmudict-arpabet_to_ipa_nv22.10.tsv",
+    )
+    return parser.parse_args()
+
+
+def convert_arp_to_ipa(arp_to_ipa_dict: Dict[str, str], arp_input: str, remove_space: bool = False) -> str:
+    """
+    Converts ARPABET phoneme to IPA based on arp_to_ipa_dict mapping
+
+    Args:
+        arp_to_ipa_dict: ARPABET to IPA phonemes mapping
+        arp_input: ARPABET input
+        remove_space: set to TRUE to remove spaces between IPA phonemes
+
+    Returns:
+        input word in IPA form
+    """
+
+    primary_stress = "ˈ"
+    secondary_stress = "ˌ"
+    stress_dict = {"0": "", "1": primary_stress, "2": secondary_stress}
+
+    word_ipa = ""
+    phonemes = arp_input.split()
+
+    # split ARPABET phoneme input into syllables,
+    # e.g. syllabify(["HH", "AH0", "L", "OW1"]) -> [(['HH'], ['AH0'], []), (['L'], ['OW1'], [])]
+    syllables = syllabify(phonemes)
+
+    for syl_idx, syll in enumerate(syllables):
+        syll_stress = ""
+        syll_ipa = ""
+
+        # syll is a tuple of lists of phonemes, here we flatten it and get rid of empty entries,
+        # e.g. (['HH'], ['AH0'], []) -> ['HH', 'AH0']
+        syll = [x for x in itertools.chain.from_iterable(syll)]
+        for phon_idx, phon in enumerate(syll):
+            if phon[-1].isdigit():
+                syll_stress = phon[-1]
+                if syll_stress not in stress_dict:
+                    raise ValueError(f"{syll_stress} unknown")
+                syll_stress = stress_dict[syll_stress]
+
+            # some phonemes are followed by a digit that represents stress, e.g., `AH0`
+            if phon not in arp_to_ipa_dict and phon[-1].isdigit():
+                phon = phon[:-1]
+
+            if phon not in arp_to_ipa_dict:
+                raise ValueError(f"|{phon}| phoneme not found in |{arp_input}|")
+            else:
+                ipa_phone = arp_to_ipa_dict[phon]
+                syll_ipa += ipa_phone + " "
+
+        word_ipa += " " + syll_stress + syll_ipa.strip()
+
+    word_ipa = word_ipa.strip()
+    if remove_space:
+        word_ipa = word_ipa.replace(" ", "")
+    return word_ipa
+
+
+def _get_arpabet_to_ipa_mapping(arp_ipa_map_file: str) -> Dict[str, str]:
+    """
+    arp_ipa_map_file: Arpabet to IPA phonemes mapping
+    """
+    arp_to_ipa = {}
+    with open(arp_ipa_map_file, "r", encoding="utf-8") as f:
+        for line in f:
+            arp, ipa = line.strip().split("\t")
+            arp_to_ipa[arp] = ipa
+    return arp_to_ipa
+
+
+def convert_cmu_arpabet_to_ipa(arp_ipa_map_file: str, arp_dict_file: str, output_ipa_file: str):
+    """
+    Converts CMU ARPABET-based dictionary to IPA.
+
+    Args:
+        arp_ipa_map_file: ARPABET to IPA phoneme mapping file
+        arp_dict_file: path to ARPABET version of CMU dictionary
+        output_ipa_file: path to output IPA version of CMU dictionary
+    """
+    arp_to_ipa_dict = _get_arpabet_to_ipa_mapping(arp_ipa_map_file)
+    with open(arp_dict_file, "r", encoding="utf-8") as f_arp, open(output_ipa_file, "w", encoding="utf-8") as f_ipa:
+        for line in f_arp:
+            if line.startswith(";;;"):
+                f_ipa.write(line)
+            else:
+                # First, split the line at " #" if there are comments in the dictionary file following the mapping entries.
+                # Next, split at default "  " separator.
+                graphemes, phonemes = line.split(" #")[0].strip().split("  ")
+                ipa_form = convert_arp_to_ipa(arp_to_ipa_dict, phonemes, remove_space=True)
+                f_ipa.write(f"{graphemes}  {ipa_form}\n")
+
+    print(f"IPA version of {os.path.abspath(arp_dict_file)} saved in {os.path.abspath(output_ipa_file)}")
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    convert_cmu_arpabet_to_ipa(args.mapping, args.cmu_arpabet, args.ipa_out)

scripts/dataset_processing/g2p/syllabify.py

@@ -0,0 +1,288 @@
+# Copyright (c) 2022, 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.
+#
+# Copyright (c) 2012-2013 Kyle Gorman <gormanky@ohsu.edu>
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the
+# "Software"), to deal in the Software without restriction, including
+# without limitation the rights to use, copy, modify, merge, publish,
+# distribute, sublicense, and/or sell copies of the Software, and to
+# permit persons to whom the Software is furnished to do so, subject to
+# the following conditions:
+#
+# The above copyright notice and this permission notice shall be included
+# in all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
+# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
+# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
+# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
+# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+#
+# syllabify.py: prosodic parsing of ARPABET entries
+# source: https://github.com/kylebgorman/syllabify
+
+from itertools import chain
+
+## constants
+SLAX = {
+    "IH1",
+    "IH2",
+    "EH1",
+    "EH2",
+    "AE1",
+    "AE2",
+    "AH1",
+    "AH2",
+    "UH1",
+    "UH2",
+}
+VOWELS = {
+    "IY1",
+    "IY2",
+    "IY0",
+    "EY1",
+    "EY2",
+    "EY0",
+    "AA1",
+    "AA2",
+    "AA0",
+    "ER1",
+    "ER2",
+    "ER0",
+    "AW1",
+    "AW2",
+    "AW0",
+    "AO1",
+    "AO2",
+    "AO0",
+    "AY1",
+    "AY2",
+    "AY0",
+    "OW1",
+    "OW2",
+    "OW0",
+    "OY1",
+    "OY2",
+    "OY0",
+    "IH0",
+    "EH0",
+    "AE0",
+    "AH0",
+    "UH0",
+    "UW1",
+    "UW2",
+    "UW0",
+    "UW",
+    "IY",
+    "EY",
+    "AA",
+    "ER",
+    "AW",
+    "AO",
+    "AY",
+    "OW",
+    "OY",
+    "UH",
+    "IH",
+    "EH",
+    "AE",
+    "AH",
+    "UH",
+} | SLAX
+
+## licit medial onsets
+
+O2 = {
+    ("P", "R"),
+    ("T", "R"),
+    ("K", "R"),
+    ("B", "R"),
+    ("D", "R"),
+    ("G", "R"),
+    ("F", "R"),
+    ("TH", "R"),
+    ("P", "L"),
+    ("K", "L"),
+    ("B", "L"),
+    ("G", "L"),
+    ("F", "L"),
+    ("S", "L"),
+    ("K", "W"),
+    ("G", "W"),
+    ("S", "W"),
+    ("S", "P"),
+    ("S", "T"),
+    ("S", "K"),
+    ("HH", "Y"),  # "clerihew"
+    ("R", "W"),
+}
+O3 = {("S", "T", "R"), ("S", "K", "L"), ("T", "R", "W")}  # "octroi"
+
+# This does not represent anything like a complete list of onsets, but
+# merely those that need to be maximized in medial position.
+
+
+def syllabify(pron, alaska_rule=True):
+    """
+    Syllabifies a CMU dictionary (ARPABET) word string
+
+    # Alaska rule:
+    >>> pprint(syllabify('AH0 L AE1 S K AH0'.split())) # Alaska
+    '-AH0-.L-AE1-S.K-AH0-'
+    >>> pprint(syllabify('AH0 L AE1 S K AH0'.split(), 0)) # Alaska
+    '-AH0-.L-AE1-.S K-AH0-'
+
+    # huge medial onsets:
+    >>> pprint(syllabify('M IH1 N S T R AH0 L'.split())) # minstrel
+    'M-IH1-N.S T R-AH0-L'
+    >>> pprint(syllabify('AA1  K T R W AA0 R'.split())) # octroi
+    '-AA1-K.T R W-AA0-R'
+
+    # destressing
+    >>> pprint(destress(syllabify('M IH1 L AH0 T EH2 R IY0'.split())))
+    'M-IH-.L-AH-.T-EH-.R-IY-'
+
+    # normal treatment of 'j':
+    >>> pprint(syllabify('M EH1 N Y UW0'.split())) # menu
+    'M-EH1-N.Y-UW0-'
+    >>> pprint(syllabify('S P AE1 N Y AH0 L'.split())) # spaniel
+    'S P-AE1-N.Y-AH0-L'
+    >>> pprint(syllabify('K AE1 N Y AH0 N'.split())) # canyon
+    'K-AE1-N.Y-AH0-N'
+    >>> pprint(syllabify('M IH0 N Y UW2 EH1 T'.split())) # minuet
+    'M-IH0-N.Y-UW2-.-EH1-T'
+    >>> pprint(syllabify('JH UW1 N Y ER0'.split())) # junior
+    'JH-UW1-N.Y-ER0-'
+    >>> pprint(syllabify('K L EH R IH HH Y UW'.split())) # clerihew
+    'K L-EH-.R-IH-.HH Y-UW-'
+
+    # nuclear treatment of 'j'
+    >>> pprint(syllabify('R EH1 S K Y UW0'.split())) # rescue
+    'R-EH1-S.K-Y UW0-'
+    >>> pprint(syllabify('T R IH1 B Y UW0 T'.split())) # tribute
+    'T R-IH1-B.Y-UW0-T'
+    >>> pprint(syllabify('N EH1 B Y AH0 L AH0'.split())) # nebula
+    'N-EH1-B.Y-AH0-.L-AH0-'
+    >>> pprint(syllabify('S P AE1 CH UH0 L AH0'.split())) # spatula
+    'S P-AE1-.CH-UH0-.L-AH0-'
+    >>> pprint(syllabify('AH0 K Y UW1 M AH0 N'.split())) # acumen
+    '-AH0-K.Y-UW1-.M-AH0-N'
+    >>> pprint(syllabify('S AH1 K Y AH0 L IH0 N T'.split())) # succulent
+    'S-AH1-K.Y-AH0-.L-IH0-N T'
+    >>> pprint(syllabify('F AO1 R M Y AH0 L AH0'.split())) # formula
+    'F-AO1 R-M.Y-AH0-.L-AH0-'
+    >>> pprint(syllabify('V AE1 L Y UW0'.split())) # value
+    'V-AE1-L.Y-UW0-'
+
+    # everything else
+    >>> pprint(syllabify('N AO0 S T AE1 L JH IH0 K'.split())) # nostalgic
+    'N-AO0-.S T-AE1-L.JH-IH0-K'
+    >>> pprint(syllabify('CH ER1 CH M AH0 N'.split())) # churchmen
+    'CH-ER1-CH.M-AH0-N'
+    >>> pprint(syllabify('K AA1 M P AH0 N S EY2 T'.split())) # compensate
+    'K-AA1-M.P-AH0-N.S-EY2-T'
+    >>> pprint(syllabify('IH0 N S EH1 N S'.split())) # inCENSE
+    '-IH0-N.S-EH1-N S'
+    >>> pprint(syllabify('IH1 N S EH2 N S'.split())) # INcense
+    '-IH1-N.S-EH2-N S'
+    >>> pprint(syllabify('AH0 S EH1 N D'.split())) # ascend
+    '-AH0-.S-EH1-N D'
+    >>> pprint(syllabify('R OW1 T EY2 T'.split())) # rotate
+    'R-OW1-.T-EY2-T'
+    >>> pprint(syllabify('AA1 R T AH0 S T'.split())) # artist
+    '-AA1 R-.T-AH0-S T'
+    >>> pprint(syllabify('AE1 K T ER0'.split())) # actor
+    '-AE1-K.T-ER0-'
+    >>> pprint(syllabify('P L AE1 S T ER0'.split())) # plaster
+    'P L-AE1-S.T-ER0-'
+    >>> pprint(syllabify('B AH1 T ER0'.split())) # butter
+    'B-AH1-.T-ER0-'
+    >>> pprint(syllabify('K AE1 M AH0 L'.split())) # camel
+    'K-AE1-.M-AH0-L'
+    >>> pprint(syllabify('AH1 P ER0'.split())) # upper
+    '-AH1-.P-ER0-'
+    >>> pprint(syllabify('B AH0 L UW1 N'.split())) # balloon
+    'B-AH0-.L-UW1-N'
+    >>> pprint(syllabify('P R OW0 K L EY1 M'.split())) # proclaim
+    'P R-OW0-.K L-EY1-M'
+    >>> pprint(syllabify('IH0 N S EY1 N'.split())) # insane
+    '-IH0-N.S-EY1-N'
+    >>> pprint(syllabify('IH0 K S K L UW1 D'.split())) # exclude
+    '-IH0-K.S K L-UW1-D'
+    """
+    ## main pass
+    mypron = list(pron)
+    nuclei = []
+    onsets = []
+    i = -1
+    for j, seg in enumerate(mypron):
+        if seg in VOWELS:
+            nuclei.append([seg])
+            onsets.append(mypron[i + 1 : j])  # actually interludes, r.n.
+            i = j
+    codas = [mypron[i + 1 :]]
+    ## resolve disputes and compute coda
+    for i in range(1, len(onsets)):
+        coda = []
+        # boundary cases
+        if len(onsets[i]) > 1 and onsets[i][0] == "R":
+            nuclei[i - 1].append(onsets[i].pop(0))
+        if len(onsets[i]) > 2 and onsets[i][-1] == "Y":
+            nuclei[i].insert(0, onsets[i].pop())
+        if len(onsets[i]) > 1 and alaska_rule and nuclei[i - 1][-1] in SLAX and onsets[i][0] == "S":
+            coda.append(onsets[i].pop(0))
+        # onset maximization
+        depth = 1
+        if len(onsets[i]) > 1:
+            if tuple(onsets[i][-2:]) in O2:
+                depth = 3 if tuple(onsets[i][-3:]) in O3 else 2
+        for j in range(len(onsets[i]) - depth):
+            coda.append(onsets[i].pop(0))
+        # store coda
+        codas.insert(i - 1, coda)
+
+    ## verify that all segments are included in the ouput
+    output = list(zip(onsets, nuclei, codas))  # in Python3 zip is a generator
+    flat_output = list(chain.from_iterable(chain.from_iterable(output)))
+    if flat_output != mypron:
+        raise ValueError(f"could not syllabify {mypron}, got {flat_output}")
+    return output
+
+
+def pprint(syllab):
+    """
+    Pretty-print a syllabification
+    """
+    return ".".join("-".join(" ".join(p) for p in syl) for syl in syllab)
+
+
+def destress(syllab):
+    """
+    Generate a syllabification with nuclear stress information removed
+    """
+    syls = []
+    for onset, nucleus, coda in syllab:
+        nuke = [p[:-1] if p[-1] in {"0", "1", "2"} else p for p in nucleus]
+        syls.append((onset, nuke, coda))
+    return syls
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()

scripts/dataset_processing/get_aishell_data.py

@@ -0,0 +1,177 @@
+# 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.
+
+# USAGE: python get_aishell_data.py --data_root=<where to put data>
+
+import argparse
+import json
+import logging
+import os
+import subprocess
+import tarfile
+import urllib.request
+
+from tqdm import tqdm
+
+from nemo.utils.tar_utils import safe_extract
+
+parser = argparse.ArgumentParser(description="Aishell Data download")
+parser.add_argument("--data_root", required=True, default=None, type=str)
+args = parser.parse_args()
+
+URL = {"data_aishell": "http://www.openslr.org/resources/33/data_aishell.tgz"}
+
+
+def __retrieve_with_progress(source: str, filename: str):
+    """
+    Downloads source to destination
+    Displays progress bar
+    Args:
+        source: url of resource
+        destination: local filepath
+    Returns:
+    """
+    with open(filename, "wb") as f:
+        response = urllib.request.urlopen(source)
+        total = response.length
+
+        if total is None:
+            f.write(response.content)
+        else:
+            with tqdm(total=total, unit="B", unit_scale=True, unit_divisor=1024) as pbar:
+                for data in response:
+                    f.write(data)
+                    pbar.update(len(data))
+
+
+def __maybe_download_file(destination: str, source: str):
+    """
+    Downloads source to destination if it doesn't exist.
+    If exists, skips download
+    Args:
+        destination: local filepath
+        source: url of resource
+
+    Returns:
+
+    """
+    source = URL[source]
+    if not os.path.exists(destination):
+        logging.info("{0} does not exist. Downloading ...".format(destination))
+        __retrieve_with_progress(source, filename=destination + ".tmp")
+        os.rename(destination + ".tmp", destination)
+        logging.info("Downloaded {0}.".format(destination))
+    else:
+        logging.info("Destination {0} exists. Skipping.".format(destination))
+    return destination
+
+
+def __extract_all_files(filepath: str, data_root: str, data_dir: str):
+    if not os.path.exists(data_dir):
+        extract_file(filepath, data_root)
+        audio_dir = os.path.join(data_dir, "wav")
+        for subfolder, _, filelist in os.walk(audio_dir):
+            for ftar in filelist:
+                extract_file(os.path.join(subfolder, ftar), subfolder)
+    else:
+        logging.info("Skipping extracting. Data already there %s" % data_dir)
+
+
+def extract_file(filepath: str, data_dir: str):
+    try:
+        with tarfile.open(filepath) as tar:
+            safe_extract(tar, data_dir)
+    except Exception:
+        logging.info("Not extracting. Maybe already there?")
+
+
+def __process_data(data_folder: str, dst_folder: str):
+    """
+    To generate manifest
+    Args:
+        data_folder: source with wav files
+        dst_folder: where manifest files will be stored
+    Returns:
+
+    """
+
+    if not os.path.exists(dst_folder):
+        os.makedirs(dst_folder)
+
+    transcript_file = os.path.join(data_folder, "transcript", "aishell_transcript_v0.8.txt")
+    transcript_dict = {}
+    with open(transcript_file, "r", encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            audio_id, text = line.split(" ", 1)
+            # remove white space
+            text = text.replace(" ", "")
+            transcript_dict[audio_id] = text
+
+    data_types = ["train", "dev", "test"]
+    vocab_count = {}
+    for dt in data_types:
+        json_lines = []
+        audio_dir = os.path.join(data_folder, "wav", dt)
+        for sub_folder, _, file_list in os.walk(audio_dir):
+            for fname in file_list:
+                audio_path = os.path.join(sub_folder, fname)
+                audio_id = fname.strip(".wav")
+                if audio_id not in transcript_dict:
+                    continue
+                text = transcript_dict[audio_id]
+                for li in text:
+                    vocab_count[li] = vocab_count.get(li, 0) + 1
+                duration = subprocess.check_output(["soxi", "-D", audio_path])
+                duration = float(duration)
+                json_lines.append(
+                    json.dumps(
+                        {
+                            "audio_filepath": os.path.abspath(audio_path),
+                            "duration": duration,
+                            "text": text,
+                        },
+                        ensure_ascii=False,
+                    )
+                )
+
+        manifest_path = os.path.join(dst_folder, dt + ".json")
+        with open(manifest_path, "w", encoding="utf-8") as fout:
+            for line in json_lines:
+                fout.write(line + "\n")
+
+    vocab = sorted(vocab_count.items(), key=lambda k: k[1], reverse=True)
+    vocab_file = os.path.join(dst_folder, "vocab.txt")
+    with open(vocab_file, "w", encoding="utf-8") as f:
+        for v, c in vocab:
+            f.write(v + "\n")
+
+
+def main():
+    data_root = args.data_root
+    data_set = "data_aishell"
+    logging.info("\n\nWorking on: {0}".format(data_set))
+    file_path = os.path.join(data_root, data_set + ".tgz")
+    logging.info("Getting {0}".format(data_set))
+    __maybe_download_file(file_path, data_set)
+    logging.info("Extracting {0}".format(data_set))
+    data_folder = os.path.join(data_root, data_set)
+    __extract_all_files(file_path, data_root, data_folder)
+    logging.info("Processing {0}".format(data_set))
+    __process_data(data_folder, data_folder)
+    logging.info("Done!")
+
+
+if __name__ == "__main__":
+    main()

scripts/dataset_processing/get_commonvoice_data.py

@@ -0,0 +1,229 @@
+# 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.
+#
+
+# Copyright (c) 2020, SeanNaren.  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.
+#
+
+# To convert mp3 files to wav using sox, you must have installed sox with mp3 support
+# For example sudo apt-get install libsox-fmt-mp3
+import argparse
+import csv
+import json
+import logging
+import multiprocessing
+import os
+import sys
+import tarfile
+import urllib.request
+from multiprocessing.pool import ThreadPool
+from pathlib import Path
+from typing import List
+
+from tqdm import tqdm
+
+from nemo.utils.tar_utils import safe_extract
+
+parser = argparse.ArgumentParser(description='Downloads and processes Mozilla Common Voice dataset.')
+parser.add_argument("--data_root", default='CommonVoice_dataset/', type=str, help="Directory to store the dataset.")
+parser.add_argument('--manifest_dir', default='./', type=str, help='Output directory for manifests')
+parser.add_argument("--num_workers", default=multiprocessing.cpu_count(), type=int, help="Workers to process dataset.")
+parser.add_argument('--sample_rate', default=16000, type=int, help='Sample rate')
+parser.add_argument('--n_channels', default=1, type=int, help='Number of channels for output wav files')
+parser.add_argument("--log", dest="log", action="store_true", default=False)
+parser.add_argument("--cleanup", dest="cleanup", action="store_true", default=False)
+parser.add_argument(
+    '--files_to_process',
+    nargs='+',
+    default=['test.tsv', 'dev.tsv', 'train.tsv'],
+    type=str,
+    help='list of *.csv file names to process',
+)
+parser.add_argument(
+    '--version',
+    default='cv-corpus-5.1-2020-06-22',
+    type=str,
+    help='Version of the dataset (obtainable via https://commonvoice.mozilla.org/en/datasets',
+)
+parser.add_argument(
+    '--language',
+    default='en',
+    type=str,
+    help='Which language to download.(default english,'
+    'check https://commonvoice.mozilla.org/en/datasets for more language codes',
+)
+args = parser.parse_args()
+COMMON_VOICE_URL = (
+    f"https://voice-prod-bundler-ee1969a6ce8178826482b88e843c335139bd3fb4.s3.amazonaws.com/"
+    "{}/{}.tar.gz".format(args.version, args.language)
+)
+COMMON_VOICE_USER_AGENT = (
+    'Mozilla/5.0 (Windows NT 10.0; WOW64) ' 'AppleWebKit/537.36 (KHTML, like Gecko) Chrome/51.0.2704.103 Safari/537.36'
+)
+
+
+def _load_sox():
+    try:
+        import sox
+        from sox import Transformer
+    except ImportError:
+        raise ImportError(
+            "Optional dependency 'sox' is required by this script. Install it with: pip install sox"
+        ) from None
+
+    return sox, Transformer
+
+
+def download_commonvoice_archive(url: str, output_path: str):
+    request = urllib.request.Request(url, headers={'User-Agent': COMMON_VOICE_USER_AGENT})
+    with urllib.request.urlopen(request) as response, open(output_path, 'wb') as f:
+        while True:
+            chunk = response.read(1024 * 1024)
+            if not chunk:
+                break
+            f.write(chunk)
+
+
+def create_manifest(data: List[tuple], output_name: str, manifest_path: str):
+    output_file = Path(manifest_path) / output_name
+    output_file.parent.mkdir(exist_ok=True, parents=True)
+
+    with output_file.open(mode='w') as f:
+        for wav_path, duration, text in tqdm(data, total=len(data)):
+            if wav_path != '':
+                # skip invalid input files that could not be converted
+                f.write(
+                    json.dumps({'audio_filepath': os.path.abspath(wav_path), "duration": duration, 'text': text})
+                    + '\n'
+                )
+
+
+def process_files(csv_file, data_root, num_workers):
+    """Read *.csv file description, convert mp3 to wav, process text.
+        Save results to data_root.
+
+    Args:
+        csv_file: str, path to *.csv file with data description, usually start from 'cv-'
+        data_root: str, path to dir to save results; wav/ dir will be created
+    """
+    sox, Transformer = _load_sox()
+    wav_dir = os.path.join(data_root, 'wav/')
+    os.makedirs(wav_dir, exist_ok=True)
+    audio_clips_path = os.path.dirname(csv_file) + '/clips/'
+
+    def process(x):
+        file_path, text = x
+        file_name = os.path.splitext(os.path.basename(file_path))[0]
+        text = text.lower().strip()
+        audio_path = os.path.join(audio_clips_path, file_path)
+        if os.path.getsize(audio_path) == 0:
+            logging.warning(f'Skipping empty audio file {audio_path}')
+            return '', '', ''
+
+        output_wav_path = os.path.join(wav_dir, file_name + '.wav')
+
+        if not os.path.exists(output_wav_path):
+            tfm = Transformer()
+            tfm.rate(samplerate=args.sample_rate)
+            tfm.channels(n_channels=args.n_channels)
+            tfm.build(input_filepath=audio_path, output_filepath=output_wav_path)
+
+        duration = sox.file_info.duration(output_wav_path)
+        return output_wav_path, duration, text
+
+    logging.info('Converting mp3 to wav for {}.'.format(csv_file))
+    with open(csv_file) as csvfile:
+        reader = csv.DictReader(csvfile, delimiter='\t')
+        next(reader, None)  # skip the headers
+        data = []
+        for row in reader:
+            file_name = row['path']
+            # add the mp3 extension if the tsv entry does not have it
+            if not file_name.endswith('.mp3'):
+                file_name += '.mp3'
+            data.append((file_name, row['sentence']))
+        with ThreadPool(num_workers) as pool:
+            data = list(tqdm(pool.imap(process, data), total=len(data)))
+    return data
+
+
+def main():
+    if args.log:
+        logging.basicConfig(level=logging.INFO)
+
+    data_root = args.data_root
+    os.makedirs(data_root, exist_ok=True)
+
+    target_unpacked_dir = os.path.join(data_root, "CV_unpacked")
+
+    if os.path.exists(target_unpacked_dir):
+        logging.info('Find existing folder {}'.format(target_unpacked_dir))
+    else:
+        logging.info("Could not find Common Voice, Downloading corpus...")
+
+        # some dataset versions are packaged in different named files, so forcing
+        output_archive_filename = args.language + '.tar.gz'
+        output_archive_filename = os.path.join(data_root, output_archive_filename)
+
+        download_commonvoice_archive(COMMON_VOICE_URL, output_archive_filename)
+        filename = f"{args.language}.tar.gz"
+        target_file = os.path.join(data_root, os.path.basename(filename))
+
+        os.makedirs(target_unpacked_dir, exist_ok=True)
+        logging.info("Unpacking corpus to {} ...".format(target_unpacked_dir))
+        with tarfile.open(target_file) as tar:
+            safe_extract(tar, target_unpacked_dir)
+        if args.cleanup:
+            logging.info("removing tar archive to save space")
+            os.remove(target_file)
+
+    folder_path = os.path.join(target_unpacked_dir, args.version + f'/{args.language}/')
+    if not os.path.isdir(folder_path):
+        # try without language
+        folder_path = os.path.join(target_unpacked_dir, args.version)
+        if not os.path.isdir(folder_path):
+            # try without version
+            folder_path = target_unpacked_dir
+            if not os.path.isdir(folder_path):
+                logging.error(f'unable to locate unpacked files in {folder_path}')
+                sys.exit()
+
+    for csv_file in args.files_to_process:
+        data = process_files(
+            csv_file=os.path.join(folder_path, csv_file),
+            data_root=os.path.join(data_root, os.path.splitext(csv_file)[0]),
+            num_workers=args.num_workers,
+        )
+        logging.info('Creating manifests...')
+        create_manifest(
+            data=data,
+            output_name=f'commonvoice_{os.path.splitext(csv_file)[0]}_manifest.json',
+            manifest_path=args.manifest_dir,
+        )
+
+
+if __name__ == "__main__":
+    main()

scripts/dataset_processing/get_demand_data.py

@@ -0,0 +1,138 @@
+# 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.
+
+# USAGE: python get_demand_data.py --data_root=<where to put data>
+#        --data_set=<datasets_to_download>
+# where <datasets_to_download> can be: one or more of the 16 kHz noise profiles
+# listed at https://zenodo.org/record/1227121#.Ygb4avXMKJk ,
+# or ALL
+# You can put more than one data_set comma-separated:
+# --data_sets=DKITCHEN,DLIVING,NRIVER
+
+import argparse
+import glob
+import json
+import logging
+import os
+import shutil
+import subprocess
+import urllib.request
+
+parser = argparse.ArgumentParser(description='LibriSpeech Data download')
+parser.add_argument("--data_root", required=True, default=None, type=str)
+parser.add_argument("--data_sets", default="ALL", type=str)
+
+parser.add_argument('--log', dest='log', action='store_true', default=False)
+args = parser.parse_args()
+
+URLS = {
+    'DKITCHEN': ("https://zenodo.org/record/1227121/files/DKITCHEN_16k.zip"),
+    'DLIVING': ("https://zenodo.org/record/1227121/files/DLIVING_16k.zip"),
+    'DWASHING': ("https://zenodo.org/record/1227121/files/DWASHING_16k.zip"),
+    'NFIELD': ("https://zenodo.org/record/1227121/files/NFIELD_16k.zip"),
+    'NPARK': ("https://zenodo.org/record/1227121/files/NPARK_16k.zip"),
+    'NRIVER': ("https://zenodo.org/record/1227121/files/NRIVER_16k.zip"),
+    'OHALLWAY': ("https://zenodo.org/record/1227121/files/OHALLWAY_16k.zip"),
+    'OMEETING': ("https://zenodo.org/record/1227121/files/OMEETING_16k.zip"),
+    'OOFFICE': ("https://zenodo.org/record/1227121/files/OOFFICE_16k.zip"),
+    'PCAFETER': ("https://zenodo.org/record/1227121/files/PCAFETER_16k.zip"),
+    'PRESTO': ("https://zenodo.org/record/1227121/files/PRESTO_16k.zip"),
+    'PSTATION': ("https://zenodo.org/record/1227121/files/PSTATION_16k.zip"),
+    'SPSQUARE': ("https://zenodo.org/record/1227121/files/SPSQUARE_16k.zip"),
+    'STRAFFIC': ("https://zenodo.org/record/1227121/files/STRAFFIC_16k.zip"),
+    'TBUS': ("https://zenodo.org/record/1227121/files/TBUS_16k.zip"),
+    'TCAR': ("https://zenodo.org/record/1227121/files/TCAR_16k.zip"),
+    'TMETRO': ("https://zenodo.org/record/1227121/files/TMETRO_16k.zip"),
+}
+
+
+def __maybe_download_file(destination: str, source: str):
+    """
+    Downloads source to destination if it doesn't exist.
+    If exists, skips download
+    Args:
+        destination: local filepath
+        source: url of resource
+    Returns:
+    """
+    source = URLS[source]
+    if not os.path.exists(destination):
+        logging.info("{0} does not exist. Downloading ...".format(destination))
+        urllib.request.urlretrieve(source, filename=destination + '.tmp')
+        os.rename(destination + '.tmp', destination)
+        logging.info("Downloaded {0}.".format(destination))
+    else:
+        logging.info("Destination {0} exists. Skipping.".format(destination))
+    return destination
+
+
+def __extract_file(filepath: str, data_dir: str):
+    shutil.unpack_archive(filepath, data_dir)
+
+
+def __create_manifest(dst_folder: str):
+    """
+    Create manifests for the noise files
+    Args:
+        file_path: path to a source transcript  with flac sources
+        dst_folder: path where manifests will be created
+    Returns:
+
+        a list of metadata entries for processed files.
+    """
+    # Read directory
+    # Get all wav file names
+    # create line per wav file in manifest
+    noise_name = os.path.basename(dst_folder)
+    wav_files = glob.glob(dst_folder + "/*.wav")
+    wav_files.sort()
+    os.makedirs(os.path.join(os.path.dirname(dst_folder), "manifests"), exist_ok=True)
+    with open(os.path.join(os.path.dirname(dst_folder), "manifests", noise_name + ".json"), "w") as mfst_f:
+        for wav_f in wav_files:
+            dur = subprocess.check_output(["soxi", "-D", wav_f])
+            row = {"audio_filepath": wav_f, "text": "", "duration": float(dur)}
+            mfst_f.write(json.dumps(row) + "\n")
+
+
+def main():
+    data_root = args.data_root
+    data_sets = args.data_sets
+
+    if args.log:
+        print("here")
+        logging.basicConfig(level=logging.INFO)
+    if not os.path.exists(data_root):
+        os.makedirs(data_root)
+
+    if data_sets == "ALL":
+        data_sets = URLS.keys()
+    else:
+        data_sets = data_sets.split(',')
+
+    for data_set in data_sets:
+        if data_set not in URLS.keys():
+            raise ValueError(f"{data_sets} is not part of demand noise database")
+        logging.info("\n\nWorking on: {0}".format(data_set))
+        filepath = os.path.join(data_root, data_set + "_16k.zip")
+        logging.info("Getting {0}".format(data_set))
+        __maybe_download_file(filepath, data_set.upper())
+        logging.info("Extracting {0}".format(data_set))
+        __extract_file(filepath, data_root)
+        logging.info("Processing {0}".format(data_set))
+        __create_manifest(os.path.join(data_root, data_set))
+    logging.info('Done!')
+
+
+if __name__ == "__main__":
+    main()

scripts/dataset_processing/get_librispeech_data.py

@@ -0,0 +1,221 @@
+# 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.
+#
+# USAGE: python get_librispeech_data.py --data_root=<where to put data>
+#        --data_set=<datasets_to_download> --num_workers=<number of parallel workers>
+# where <datasets_to_download> can be: dev_clean, dev_other, test_clean,
+# test_other, train_clean_100, train_clean_360, train_other_500 or ALL
+# You can also put more than one data_set comma-separated:
+# --data_set=dev_clean,train_clean_100
+import argparse
+import fnmatch
+import functools
+import json
+import logging
+import multiprocessing
+import os
+import subprocess
+import tarfile
+import urllib.request
+
+from tqdm import tqdm
+
+from nemo.utils.tar_utils import safe_extract
+
+parser = argparse.ArgumentParser(description="LibriSpeech Data download")
+parser.add_argument("--data_root", required=True, default=None, type=str)
+parser.add_argument("--data_sets", default="dev_clean", type=str)
+parser.add_argument("--num_workers", default=4, type=int)
+parser.add_argument("--log", dest="log", action="store_true", default=False)
+args = parser.parse_args()
+
+URLS = {
+    "TRAIN_CLEAN_100": ("http://www.openslr.org/resources/12/train-clean-100.tar.gz"),
+    "TRAIN_CLEAN_360": ("http://www.openslr.org/resources/12/train-clean-360.tar.gz"),
+    "TRAIN_OTHER_500": ("http://www.openslr.org/resources/12/train-other-500.tar.gz"),
+    "DEV_CLEAN": "http://www.openslr.org/resources/12/dev-clean.tar.gz",
+    "DEV_OTHER": "http://www.openslr.org/resources/12/dev-other.tar.gz",
+    "TEST_CLEAN": "http://www.openslr.org/resources/12/test-clean.tar.gz",
+    "TEST_OTHER": "http://www.openslr.org/resources/12/test-other.tar.gz",
+    "DEV_CLEAN_2": "https://www.openslr.org/resources/31/dev-clean-2.tar.gz",
+    "TRAIN_CLEAN_5": "https://www.openslr.org/resources/31/train-clean-5.tar.gz",
+}
+
+
+def _load_sox_transformer():
+    try:
+        from sox import Transformer
+    except ImportError:
+        raise ImportError(
+            "Optional dependency 'sox' is required by this script. Install it with: pip install sox"
+        ) from None
+
+    return Transformer
+
+
+def __retrieve_with_progress(source: str, filename: str):
+    """
+    Downloads source to destination
+    Displays progress bar
+    Args:
+        source: url of resource
+        destination: local filepath
+    Returns:
+    """
+    with open(filename, "wb") as f:
+        response = urllib.request.urlopen(source)
+        total = response.length
+
+        if total is None:
+            f.write(response.content)
+        else:
+            with tqdm(total=total, unit="B", unit_scale=True, unit_divisor=1024) as pbar:
+                for data in response:
+                    f.write(data)
+                    pbar.update(len(data))
+
+
+def __maybe_download_file(destination: str, source: str):
+    """
+    Downloads source to destination if it doesn't exist.
+    If exists, skips download
+    Args:
+        destination: local filepath
+        source: url of resource
+    Returns:
+    """
+    source = URLS[source]
+    if not os.path.exists(destination):
+        logging.info("{0} does not exist. Downloading ...".format(destination))
+
+        __retrieve_with_progress(source, filename=destination + ".tmp")
+
+        os.rename(destination + ".tmp", destination)
+        logging.info("Downloaded {0}.".format(destination))
+    else:
+        logging.info("Destination {0} exists. Skipping.".format(destination))
+    return destination
+
+
+def __extract_file(filepath: str, data_dir: str):
+    try:
+        with tarfile.open(filepath) as tar:
+            safe_extract(tar, data_dir)
+    except Exception:
+        logging.info("Not extracting. Maybe already there?")
+
+
+def __process_transcript(file_path: str, dst_folder: str):
+    """
+    Converts flac files to wav from a given transcript, capturing the metadata.
+    Args:
+        file_path: path to a source transcript  with flac sources
+        dst_folder: path where wav files will be stored
+    Returns:
+        a list of metadata entries for processed files.
+    """
+    Transformer = _load_sox_transformer()
+    entries = []
+    root = os.path.dirname(file_path)
+    with open(file_path, encoding="utf-8") as fin:
+        for line in fin:
+            id, text = line[: line.index(" ")], line[line.index(" ") + 1 :]
+            transcript_text = text.lower().strip()
+
+            # Convert FLAC file to WAV
+            flac_file = os.path.join(root, id + ".flac")
+            wav_file = os.path.join(dst_folder, id + ".wav")
+            if not os.path.exists(wav_file):
+                Transformer().build(flac_file, wav_file)
+            # check duration
+            duration = subprocess.check_output(["soxi", "-D", wav_file])
+
+            entry = {}
+            entry["audio_filepath"] = os.path.abspath(wav_file)
+            entry["duration"] = float(duration)
+            entry["text"] = transcript_text
+            entries.append(entry)
+    return entries
+
+
+def __process_data(data_folder: str, dst_folder: str, manifest_file: str, num_workers: int):
+    """
+    Converts flac to wav and build manifests's json
+    Args:
+        data_folder: source with flac files
+        dst_folder: where wav files will be stored
+        manifest_file: where to store manifest
+        num_workers: number of parallel workers processing files
+    Returns:
+    """
+
+    if not os.path.exists(dst_folder):
+        os.makedirs(dst_folder)
+
+    files = []
+    entries = []
+
+    for root, dirnames, filenames in os.walk(data_folder):
+        for filename in fnmatch.filter(filenames, "*.trans.txt"):
+            files.append(os.path.join(root, filename))
+
+    with multiprocessing.Pool(num_workers) as p:
+        processing_func = functools.partial(__process_transcript, dst_folder=dst_folder)
+        results = p.imap(processing_func, files)
+        for result in tqdm(results, total=len(files)):
+            entries.extend(result)
+
+    with open(manifest_file, "w") as fout:
+        for m in entries:
+            fout.write(json.dumps(m) + "\n")
+
+
+def main():
+    data_root = args.data_root
+    data_sets = args.data_sets
+    num_workers = args.num_workers
+
+    if args.log:
+        logging.basicConfig(level=logging.INFO)
+
+    if data_sets == "ALL":
+        data_sets = "dev_clean,dev_other,train_clean_100,train_clean_360,train_other_500,test_clean,test_other"
+    if data_sets == "mini":
+        data_sets = "dev_clean_2,train_clean_5"
+    for data_set in data_sets.split(","):
+        logging.info("\n\nWorking on: {0}".format(data_set))
+        filepath = os.path.join(data_root, data_set + ".tar.gz")
+        logging.info("Getting {0}".format(data_set))
+        __maybe_download_file(filepath, data_set.upper())
+        logging.info("Extracting {0}".format(data_set))
+        __extract_file(filepath, data_root)
+        logging.info("Processing {0}".format(data_set))
+        __process_data(
+            os.path.join(
+                os.path.join(data_root, "LibriSpeech"),
+                data_set.replace("_", "-"),
+            ),
+            os.path.join(
+                os.path.join(data_root, "LibriSpeech"),
+                data_set.replace("_", "-"),
+            )
+            + "-processed",
+            os.path.join(data_root, data_set + ".json"),
+            num_workers=num_workers,
+        )
+    logging.info("Done!")
+
+
+if __name__ == "__main__":
+    main()

scripts/dataset_processing/get_openslr_rir_data.py

@@ -0,0 +1,165 @@
+# 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.
+
+# USAGE: python get_openslr_rir_data.py --data_root=<where to put data>
+# Data is downloaded from OpenSLR's "Room Impulse Response and Noise Database"
+# RIRs in multichannel files are separated into single channel files and
+# a json file that can be used as in input to NeMo is created
+
+import argparse
+import glob
+import json
+import logging
+import os
+import subprocess
+import urllib.request
+from shutil import copy, move
+from zipfile import ZipFile
+
+from tqdm import tqdm
+
+parser = argparse.ArgumentParser(description="OpenSLR RIR Data download and process")
+parser.add_argument("--data_root", required=True, default=None, type=str)
+args = parser.parse_args()
+
+URLS = {
+    "SLR28": ("http://www.openslr.org/resources/28/rirs_noises.zip"),
+}
+
+
+def __retrieve_with_progress(source: str, filename: str):
+    """
+    Downloads source to destination
+    Displays progress bar
+    Args:
+        source: url of resource
+        destination: local filepath
+    Returns:
+    """
+    with open(filename, "wb") as f:
+        response = urllib.request.urlopen(source)
+        total = response.length
+
+        if total is None:
+            f.write(response.content)
+        else:
+            with tqdm(total=total, unit="B", unit_scale=True, unit_divisor=1024) as pbar:
+                for data in response:
+                    f.write(data)
+                    pbar.update(len(data))
+
+
+def __maybe_download_file(destination: str, source: str):
+    """
+    Downloads source to destination if it doesn't exist.
+    If exists, skips download
+    Args:
+        destination: local filepath
+        source: url of resource
+    Returns:
+    """
+    source = URLS[source]
+    if not os.path.exists(destination):
+        logging.info("{0} does not exist. Downloading ...".format(destination))
+        __retrieve_with_progress(source, filename=destination + ".tmp")
+        os.rename(destination + ".tmp", destination)
+        logging.info("Downloaded {0}.".format(destination))
+    else:
+        logging.info("Destination {0} exists. Skipping.".format(destination))
+    return destination
+
+
+def __extract_file(filepath: str, data_dir: str):
+    try:
+        with ZipFile(filepath, "r") as zipObj:
+            zipObj.extractall(data_dir)
+    except Exception:
+        logging.info("Not extracting. Maybe already there?")
+
+
+def __process_data(data_folder: str, dst_folder: str, manifest_file: str):
+    """
+    Converts flac to wav and build manifests's json
+    Args:
+        data_folder: source with flac files
+        dst_folder: where wav files will be stored
+        manifest_file: where to store manifest
+    Returns:
+    """
+    if not os.path.exists(dst_folder):
+        os.makedirs(dst_folder)
+
+    real_rir_list = os.path.join(data_folder, "RIRS_NOISES", "real_rirs_isotropic_noises", "rir_list")
+    rirfiles = []
+    with open(real_rir_list, "r") as rir_f:
+        for line in rir_f:
+            rirfiles.append(os.path.join(data_folder, line.rstrip().split(" ")[4]))
+
+    real_rir_folder = os.path.join(dst_folder, "real_rirs")
+    if not os.path.exists(real_rir_folder):
+        os.makedirs(real_rir_folder)
+    # split multi-channel rir files to single channel
+    for rir_f in rirfiles:
+        n_chans = int(subprocess.check_output(["soxi", "-c", rir_f]))
+        if n_chans == 1:
+            copy(rir_f, real_rir_folder)
+        else:
+            for chan in range(1, n_chans + 1):
+                chan_file_name = os.path.join(
+                    real_rir_folder,
+                    os.path.splitext(os.path.basename(rir_f))[0] + "-" + str(chan) + ".wav",
+                )
+                _ = subprocess.check_output(["sox", rir_f, chan_file_name, "remix", str(chan)])
+
+    # move simulated rirs to processed
+    if not os.path.exists(os.path.join(dst_folder, "simulated_rirs")):
+        move(os.path.join(data_folder, "RIRS_NOISES", "simulated_rirs"), dst_folder)
+
+    os.chdir(dst_folder)
+    all_rirs = glob.glob("**/*.wav", recursive=True)
+    with open(manifest_file, "w") as man_f:
+        entry = {}
+        for rir in all_rirs:
+            rir_file = os.path.join(dst_folder, rir)
+            duration = subprocess.check_output(["soxi", "-D", rir_file])
+            entry["audio_filepath"] = rir_file
+            entry["duration"] = float(duration)
+            entry["offset"] = 0
+            entry["text"] = "_"
+            man_f.write(json.dumps(entry) + "\n")
+
+    print("Done!")
+
+
+def main():
+    data_root = os.path.abspath(args.data_root)
+    data_set = "slr28"
+    logging.getLogger().setLevel(logging.INFO)
+    logging.info("\n\nWorking on: {0}".format(data_set))
+    filepath = os.path.join(data_root, data_set + ".zip")
+    logging.info("Getting {0}".format(data_set))
+    __maybe_download_file(filepath, data_set.upper())
+    logging.info("Extracting {0}".format(data_set))
+    __extract_file(filepath, data_root)
+    logging.info("Processing {0}".format(data_set))
+    __process_data(
+        data_root,
+        os.path.join(os.path.join(data_root, "processed")),
+        os.path.join(os.path.join(data_root, "processed", "rir.json")),
+    )
+    logging.info("Done!")
+
+
+if __name__ == "__main__":
+    main()

scripts/dataset_processing/kaldi2json.py

@@ -0,0 +1,135 @@
+# 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.
+
+import argparse
+import json
+import os
+
+import pandas as pd
+
+from nemo.utils import logging
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Convert kaldi data folder to manifest.json")
+    parser.add_argument(
+        "--data_dir",
+        required=True,
+        type=str,
+        help="data in kaldi format",
+    )
+    parser.add_argument(
+        "--manifest",
+        required=True,
+        type=str,
+        help="path to store the manifest file",
+    )
+    parser.add_argument(
+        "--with_aux_data",
+        default=False,
+        action="store_true",
+        help="whether to include auxiliary data in the manifest",
+    )
+    args = parser.parse_args()
+
+    kaldi_folder = args.data_dir
+    required_data = {
+        "audio_filepath": os.path.join(kaldi_folder, "wav.scp"),
+        "duration": os.path.join(kaldi_folder, "segments"),
+        "text": os.path.join(kaldi_folder, "text"),
+    }
+    aux_data = {
+        "speaker": os.path.join(kaldi_folder, "utt2spk"),
+        "gender": os.path.join(kaldi_folder, "utt2gender"),
+    }
+    output_names = list(required_data.keys())
+
+    # check if required files exist
+    for name, file in required_data.items():
+        if not os.path.exists(file):
+            raise ValueError(f"{os.path.basename(file)} is not in {kaldi_folder}.")
+
+    # read wav.scp
+    wavscp = pd.read_csv(required_data["audio_filepath"], sep=" ", header=None)
+    if wavscp.shape[1] > 2:
+        logging.warning(
+            f"""More than two columns in 'wav.scp': {wavscp.shape[1]}.
+            Maybe it contains pipes? Pipe processing can be slow at runtime."""
+        )
+        wavscp = pd.read_csv(
+            required_data["audio_filepath"],
+            sep="^([^ ]+) ",
+            engine="python",
+            header=None,
+            usecols=[1, 2],
+            names=["wav_label", "audio_filepath"],
+        )
+    else:
+        wavscp = wavscp.rename(columns={0: "wav_label", 1: "audio_filepath"})
+
+    # read text
+    text = pd.read_csv(
+        required_data["text"],
+        sep="^([^ ]+) ",
+        engine="python",
+        header=None,
+        usecols=[1, 2],
+        names=["label", "text"],
+    )
+
+    # read segments
+    segments = pd.read_csv(
+        required_data["duration"],
+        sep=" ",
+        header=None,
+        names=["label", "wav_label", "offset", "end"],
+    )
+    # add offset if needed
+    if len(segments.offset) > len(segments.offset[segments.offset == 0.0]):
+        logging.info("Adding offset field.")
+        output_names.insert(2, "offset")
+    segments["duration"] = (segments.end - segments.offset).round(decimals=3)
+
+    # merge data
+    wav_segments_text = pd.merge(
+        pd.merge(segments, wavscp, how="inner", on="wav_label"),
+        text,
+        how="inner",
+        on="label",
+    )
+
+    if args.with_aux_data:
+        # check if auxiliary data is present
+        for name, aux_file in aux_data.items():
+            if os.path.exists(aux_file):
+                logging.info(f"Adding info from '{os.path.basename(aux_file)}'.")
+                wav_segments_text = pd.merge(
+                    wav_segments_text,
+                    pd.read_csv(aux_file, sep=" ", header=None, names=["label", name]),
+                    how="left",
+                    on="label",
+                )
+                output_names.append(name)
+            else:
+                logging.info(f"'{os.path.basename(aux_file)}' does not exist. Skipping ...")
+
+    # write data to .json
+    entries = wav_segments_text[output_names].to_dict(orient="records")
+    with open(args.manifest, "w", encoding="utf-8") as fout:
+        for m in entries:
+            fout.write(json.dumps(m, ensure_ascii=False) + "\n")
+
+
+if __name__ == "__main__":
+    main()

scripts/dataset_processing/process_aishell2_data.py

@@ -0,0 +1,109 @@
+# 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.
+
+# USAGE: python process_aishell2_data.py
+#                   --audio_folder=<source data>
+#                   --dest_folder=<where to store the results>
+import argparse
+import json
+import os
+import subprocess
+
+parser = argparse.ArgumentParser(description="Processing Aishell2 Data")
+parser.add_argument("--audio_folder", default=None, type=str, required=True, help="Audio (wav) data directory.")
+parser.add_argument("--dest_folder", default=None, type=str, required=True, help="Destination directory.")
+args = parser.parse_args()
+
+
+def __process_data(data_folder: str, dst_folder: str):
+    """
+    To generate manifest
+    Args:
+        data_folder: source with wav files
+        dst_folder: where manifest files will be stored
+    Returns:
+    """
+    if not os.path.exists(dst_folder):
+        os.makedirs(dst_folder)
+    data_type = ['dev', 'test', 'train']
+    for data in data_type:
+        dst_file = os.path.join(dst_folder, data + ".json")
+        uttrances = []
+        wav_dir = os.path.join(data_folder, "wav", data)
+        transcript_file = os.path.join(data_folder, "transcript", data, "trans.txt")
+        trans_text = {}
+        with open(transcript_file, "r", encoding='utf-8') as f:
+            for line in f:
+                line = line.strip().split()
+                utterance_id, text = line[0], " ".join(line[1:])
+                trans_text[utterance_id] = text.upper()
+        session_list = os.listdir(wav_dir)
+        for sessions in session_list:
+            cur_dir = os.path.join(wav_dir, sessions)
+            for wavs in os.listdir(cur_dir):
+                audio_id = wavs.strip(".wav")
+                audio_filepath = os.path.abspath(os.path.join(cur_dir, wavs))
+                duration = subprocess.check_output(["soxi", "-D", audio_filepath])
+                duration = float(duration)
+                text = trans_text[audio_id]
+                uttrances.append(
+                    json.dumps(
+                        {"audio_filepath": audio_filepath, "duration": duration, "text": text}, ensure_ascii=False
+                    )
+                )
+        with open(dst_file, "w") as f:
+            for line in uttrances:
+                f.write(line + "\n")
+
+
+def __get_vocab(data_folder: str, des_dir: str):
+    """
+    To generate the vocabulary file
+    Args:
+        data_folder: source with the transcript file
+        dst_folder: where the file will be stored
+    Returns:
+    """
+    if not os.path.exists(des_dir):
+        os.makedirs(des_dir)
+    trans_file = os.path.join(data_folder, "transcript", "train", "trans.txt")
+    vocab_dict = {}
+    with open(trans_file, "r", encoding='utf-8') as f:
+        for line in f:
+            line = line.strip().split()
+            text = " ".join(line[1:])
+            for i in text.upper():
+                if i in vocab_dict:
+                    vocab_dict[i] += 1
+                else:
+                    vocab_dict[i] = 1
+    vocab_dict = sorted(vocab_dict.items(), key=lambda k: k[1], reverse=True)
+    vocab = os.path.join(des_dir, "vocab.txt")
+    vocab = open(vocab, "w", encoding='utf-8')
+    for k in vocab_dict:
+        vocab.write(k[0] + "\n")
+    vocab.close()
+
+
+def main():
+    source_data = args.audio_folder
+    des_dir = args.dest_folder
+    print("begin to process data...")
+    __process_data(source_data, des_dir)
+    __get_vocab(source_data, des_dir)
+    print("finish all!")
+
+
+if __name__ == "__main__":
+    main()

scripts/dataset_processing/process_an4_data.py

@@ -0,0 +1,89 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import argparse
+import glob
+import json
+import logging
+import os
+import subprocess
+
+import librosa
+
+parser = argparse.ArgumentParser(description="AN4 dataset download and processing")
+parser.add_argument("--data_root", required=True, default=None, type=str)
+args = parser.parse_args()
+
+
+def build_manifest(data_root, transcripts_path, manifest_path, wav_path):
+    with open(transcripts_path, 'r') as fin:
+        with open(manifest_path, 'w') as fout:
+            for line in fin:
+                # Lines look like this:
+                # <s> transcript </s> (fileID)
+                transcript = line[: line.find('(') - 1].lower()
+                transcript = transcript.replace('<s>', '').replace('</s>', '')
+                transcript = transcript.strip()
+
+                file_id = line[line.find('(') + 1 : -2]  # e.g. "cen4-fash-b"
+                audio_path = os.path.join(
+                    data_root,
+                    wav_path,
+                    file_id[file_id.find('-') + 1 : file_id.rfind('-')],
+                    file_id + '.wav',
+                )
+
+                duration = librosa.core.get_duration(filename=audio_path)
+
+                # Write the metadata to the manifest
+                metadata = {
+                    "audio_filepath": audio_path,
+                    "duration": duration,
+                    "text": transcript,
+                }
+                json.dump(metadata, fout)
+                fout.write('\n')
+
+
+def main():
+    data_root = os.path.abspath(args.data_root)
+
+    # Convert from .sph to .wav
+    logging.info("Converting audio files to .wav...")
+    sph_list = glob.glob(os.path.join(data_root, 'an4/**/*.sph'), recursive=True)
+    for sph_path in sph_list:
+        wav_path = sph_path[:-4] + '.wav'
+        cmd = ['sox', sph_path, wav_path]
+        subprocess.run(cmd)
+    logging.info("Finished conversion.")
+
+    # Build manifests
+    logging.info("Building training manifest...")
+    train_transcripts = os.path.join(data_root, 'an4/etc/an4_train.transcription')
+    train_manifest = os.path.join(data_root, 'an4/train_manifest.json')
+    train_wavs = os.path.join(data_root, 'an4/wav/an4_clstk')
+    build_manifest(data_root, train_transcripts, train_manifest, train_wavs)
+    logging.info("Training manifests created.")
+
+    logging.info("Building test manifest...")
+    test_transcripts = os.path.join(data_root, 'an4/etc/an4_test.transcription')
+    test_manifest = os.path.join(data_root, 'an4/test_manifest.json')
+    test_wavs = os.path.join(data_root, 'an4/wav/an4test_clstk')
+    build_manifest(data_root, test_transcripts, test_manifest, test_wavs)
+    logging.info("Test manifest created.")
+
+    logging.info("Done with AN4 processing!")
+
+
+if __name__ == '__main__':
+    main()

scripts/dataset_processing/process_fisher_data.py

@@ -0,0 +1,434 @@
+# 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.
+
+# USAGE:
+# python process_fisher_data.py \
+#   --audio_root=<audio (.wav) directory>
+#   --transcript_root=<LDC Fisher dataset directory> \
+#   --dest_root=<destination directory> \
+#   --data_sets=LDC2004S13-Part1,LDC2005S13-Part2 \
+#   --remove_noises
+#
+# Matches Fisher dataset transcripts to the corresponding audio file (.wav),
+# and slices them into min_slice_duration segments with one speaker.
+# Also performs some other processing on transcripts.
+#
+# Heavily derived from Patter's Fisher processing script.
+
+import argparse
+import glob
+import json
+import os
+import re
+from math import ceil, floor
+
+import numpy as np
+import scipy.io.wavfile as wavfile
+from tqdm import tqdm
+
+parser = argparse.ArgumentParser(description="Fisher Data Processing")
+parser.add_argument(
+    "--audio_root",
+    default=None,
+    type=str,
+    required=True,
+    help="The path to the root of the audio (wav) data folder.",
+)
+parser.add_argument(
+    "--transcript_root",
+    default=None,
+    type=str,
+    required=True,
+    help="The path to the root of the transcript data folder.",
+)
+parser.add_argument(
+    "--dest_root",
+    default=None,
+    type=str,
+    required=True,
+    help="Path to the destination root directory.",
+)
+
+# Optional arguments
+parser.add_argument(
+    "--min_slice_duration",
+    default=10.0,
+    type=float,
+    help="Minimum audio slice duration after processing.",
+)
+parser.add_argument(
+    "--keep_low_conf",
+    action="store_true",
+    help="Keep all utterances with low confidence transcripts",
+)
+parser.add_argument(
+    "--remove_noises",
+    action="store_true",
+    help="Removes transcripted noises such as [laughter].",
+)
+parser.add_argument(
+    "--noises_to_emoji",
+    action="store_true",
+    help="Converts transcripts for noises to an emoji character.",
+)
+args = parser.parse_args()
+
+# Total number of files before segmenting, and train/val/test splits
+NUM_FILES = 5850 + 5849
+TRAIN_END_IDX = int(NUM_FILES * 0.8)
+VAL_END_IDX = int(NUM_FILES * 0.9)
+
+# Known transcription errors and their fixes (from Mozilla)
+TRANSCRIPT_BUGS = {
+    "fe_03_00265-B-3353-3381": "correct",
+    "fe_03_00991-B-52739-52829": "that's one of those",
+    "fe_03_10282-A-34442-34484.wav": "they don't want",
+    "fe_03_10677-B-10104-10641": "uh my mine yeah the german shepherd "
+    + "pitbull mix he snores almost as loud "
+    + "as i do",
+    "fe_03_00027-B-39380-39405": None,
+    "fe_03_11487-B-3109-23406": None,
+    "fe_03_01326-A-30742-30793": None,
+}
+
+TRANSCRIPT_NUMBERS = {
+    "401k": "four o one k",
+    "f16": "f sixteen",
+    "m16": "m sixteen",
+    "ak47": "a k forty seven",
+    "v8": "v eight",
+    "y2k": "y two k",
+    "mp3": "m p three",
+    "vh1": "v h one",
+    "90210": "nine o two one o",
+    "espn2": "e s p n two",
+    "u2": "u two",
+    "dc3s": "d c threes",
+    "book 2": "book two",
+    "s2b": "s two b",
+    "3d": "three d",
+}
+
+TAG_MAP = {
+    "[laughter]": "🤣",
+    "[laugh]": "🤣",
+    "[noise]": "😕",
+    "[sigh]": "😕",
+    "[cough]": "😕",
+    "[mn]": "😕",
+    "[breath]": "😕",
+    "[lipsmack]": "😕",
+    "[[skip]]": "",
+    "[pause]": "",
+    "[sneeze]": "😕",
+}
+
+
+def __write_sample(dest, file_id, count, file_count, sample_rate, audio, duration, transcript):
+    """
+    Writes one slice to the given target directory.
+    Args:
+        dest: the destination directory
+        file_id: name of the transcript/audio file for this block
+        count: the count of segments in the file so far
+        file_count: the total number of filse processed so far
+        sample rate: sample rate of the audio data
+        audio: audio data of the current sample
+        duration: audio duration of the current sample
+        transcript: transcript of the current sample
+    """
+    partition = __partition_name(file_count)
+    audio_path = os.path.join(dest, partition, f"{file_id}_{count:03}.wav")
+
+    # Write audio
+    wavfile.write(audio_path, sample_rate, audio)
+
+    # Write transcript info
+    transcript = {
+        "audio_filepath": audio_path,
+        "duration": duration,
+        "text": transcript,
+    }
+
+    # Append to manifest
+    manifest_path = os.path.join(dest, f"manifest_{partition}.json")
+    with open(manifest_path, 'a') as f:
+        json.dump(transcript, f)
+        f.write('\n')
+
+
+def __normalize(utt):
+    replace_table = str.maketrans(dict.fromkeys('()*;:"!&{},.-?'))
+    utt = (
+        utt.lower()
+        .replace('[uh]', 'uh')
+        .replace('[um]', 'um')
+        .replace('<noise>', '[noise]')
+        .replace('<spoken_noise>', '[vocalized-noise]')
+        .replace('.period', 'period')
+        .replace('.dot', 'dot')
+        .replace('-hyphen', 'hyphen')
+        .replace('._', ' ')
+        .translate(replace_table)
+    )
+    utt = re.sub(r"'([a-z]+)'", r'\1', utt)  # Unquote quoted words
+    return utt
+
+
+def __process_utterance(file_id, trans_path, line, keep_low_conf, rem_noises, emojify):
+    """
+    Processes one utterance (one line of a transcript).
+    Args:
+        file_id: the ID of the transcript file
+        trans_path: transcript path
+        line: one line in the transcript file
+        keep_low_conf: whether to keep low confidence lines
+        rem_noises: whether to remove noise symbols
+        emojify: whether to convert noise symbols to emoji, lower precedence
+    """
+    # Check for lines to skip (comments, empty, low confidence)
+    if line.startswith('#') or not line.strip() or (not keep_low_conf and '((' in line):
+        return None, None, None, None
+
+    # Data and sanity checks
+    line = line.split()
+
+    t_start, t_end = float(line[0]), float(line[1])
+    if (t_start < 0) or (t_end < t_start):
+        print(f"Invalid time: {t_start} to {t_end} in {trans_path}")
+        return None, None, None, None
+
+    channel = line[2]
+    idx = 0 if line[2] == 'A:' else 1
+
+    if channel not in ('A:', 'B:'):
+        print(f"Could not read channel info ({channel}) in {trans_path}")
+        return None, None, None, None
+
+    # Replacements as necessary
+    line_id = '-'.join([file_id, channel[0], str(t_start * 10), str(t_end * 10)])
+
+    content = TRANSCRIPT_BUGS.get(line_id, ' '.join(line[3:]))
+
+    if content is None:
+        return None, None, None, None
+
+    for tag, newtag in TRANSCRIPT_NUMBERS.items():
+        content = content.replace(tag, newtag)
+
+    content = __normalize(content)
+
+    if rem_noises:
+        for k, _ in TAG_MAP.items():
+            content = content.replace(k, '')
+    elif emojify:
+        for k, v in TAG_MAP.items():
+            content = content.replace(k, v)
+
+    return t_start, t_end, idx, content
+
+
+def __process_one_file(
+    trans_path,
+    sample_rate,
+    audio_data,
+    file_id,
+    dst_root,
+    min_slice_duration,
+    file_count,
+    keep_low_conf,
+    rem_noises,
+    emojify,
+):
+    """
+    Creates one block of audio slices and their corresponding transcripts.
+    Args:
+        trans_path: filepath to transcript
+        sample_rate: sample rate of the audio
+        audio_data: numpy array of shape [samples, channels]
+        file_id: identifying label, e.g. 'fe_03_01102'
+        dst_root: path to destination directory
+        min_slice_duration: min number of seconds for an audio slice
+        file_count: total number of files processed so far
+        keep_low_conf: keep utterances with low-confidence transcripts
+        rem_noises: remove noise symbols
+        emojify: convert noise symbols into emoji characters
+    """
+    count = 0
+
+    with open(trans_path, encoding="utf-8") as fin:
+        fin.readline()  # Comment w/ corresponding sph filename
+        fin.readline()  # Comment about transcriber
+
+        transcript_buffers = ['', '']  # [A buffer, B buffer]
+        audio_buffers = [[], []]
+        buffer_durations = [0.0, 0.0]
+
+        for line in fin:
+            t_start, t_end, idx, content = __process_utterance(
+                file_id, trans_path, line, keep_low_conf, rem_noises, emojify
+            )
+
+            if content is None or not content:
+                continue
+
+            duration = t_end - t_start
+
+            # Append utterance to buffer
+            transcript_buffers[idx] += content
+            audio_buffers[idx].append(
+                audio_data[
+                    floor(t_start * sample_rate) : ceil(t_end * sample_rate),
+                    idx,
+                ]
+            )
+            buffer_durations[idx] += duration
+
+            if buffer_durations[idx] < min_slice_duration:
+                transcript_buffers[idx] += ' '
+            else:
+                # Write out segment and transcript
+                count += 1
+                __write_sample(
+                    dst_root,
+                    file_id,
+                    count,
+                    file_count,
+                    sample_rate,
+                    np.concatenate(audio_buffers[idx], axis=0),
+                    buffer_durations[idx],
+                    transcript_buffers[idx],
+                )
+
+                # Clear buffers
+                transcript_buffers[idx] = ''
+                audio_buffers[idx] = []
+                buffer_durations[idx] = 0.0
+
+            # Note: We drop any shorter "scraps" at the end of the file, if
+            #   they end up shorter than min_slice_duration.
+
+
+def __partition_name(file_count):
+    if file_count >= VAL_END_IDX:
+        return "test"
+    elif file_count >= TRAIN_END_IDX:
+        return "val"
+    else:
+        return "train"
+
+
+def __process_data(
+    audio_root,
+    transcript_root,
+    dst_root,
+    min_slice_duration,
+    file_count,
+    keep_low_conf,
+    rem_noises,
+    emojify,
+):
+    """
+    Converts Fisher wav files to numpy arrays, segments audio and transcripts.
+    Args:
+        audio_root: source directory with the wav files
+        transcript_root: source directory with the transcript files
+            (can be the same as audio_root)
+        dst_root: where the processed and segmented files will be stored
+        min_slice_duration: minimum number of seconds for a slice of output
+        file_count: total number of files processed so far
+        keep_low_conf: whether or not to keep low confidence transcriptions
+        rem_noises: whether to remove noise symbols
+        emojify: whether to convert noise symbols to emoji, lower precedence
+    Assumes:
+        1. There is exactly one transcripts directory in data_folder
+        2. Audio files are all: <audio_root>/audio-wav/fe_03_xxxxx.wav
+    """
+    transcript_list = glob.glob(os.path.join(transcript_root, "fe_03_p*_tran*", "data", "trans", "*", "*.txt"))
+    print("Found {} transcripts.".format(len(transcript_list)))
+
+    count = file_count
+
+    # Grab audio file associated with each transcript, and slice
+    for trans_path in tqdm(transcript_list, desc="Matching and segmenting"):
+        file_id, _ = os.path.splitext(os.path.basename(trans_path))
+        audio_path = os.path.join(audio_root, "audio_wav", file_id + ".wav")
+
+        sample_rate, audio_data = wavfile.read(audio_path)
+
+        # Create a set of segments (a block) for each file
+        __process_one_file(
+            trans_path,
+            sample_rate,
+            audio_data,
+            file_id,
+            dst_root,
+            min_slice_duration,
+            count,
+            keep_low_conf,
+            rem_noises,
+            emojify,
+        )
+        count += 1
+
+    return count
+
+
+def main():
+    # Arguments to the script
+    audio_root = args.audio_root
+    transcript_root = args.transcript_root
+    dest_root = args.dest_root
+
+    min_slice_duration = args.min_slice_duration
+    keep_low_conf = args.keep_low_conf
+    rem_noises = args.remove_noises
+    emojify = args.noises_to_emoji
+
+    print(f"Expected number of files to segment: {NUM_FILES}")
+    print("With a 80/10/10 split:")
+    print(f"Number of training files: {TRAIN_END_IDX}")
+    print(f"Number of validation files: {VAL_END_IDX - TRAIN_END_IDX}")
+    print(f"Number of test files: {NUM_FILES - VAL_END_IDX}")
+
+    if not os.path.exists(os.path.join(dest_root, 'train/')):
+        os.makedirs(os.path.join(dest_root, 'train/'))
+        os.makedirs(os.path.join(dest_root, 'val/'))
+        os.makedirs(os.path.join(dest_root, 'test/'))
+    else:
+        # Wipe manifest contents first
+        open(os.path.join(dest_root, "manifest_train.json"), 'w').close()
+        open(os.path.join(dest_root, "manifest_val.json"), 'w').close()
+        open(os.path.join(dest_root, "manifest_test.json"), 'w').close()
+
+    file_count = 0
+
+    for data_set in ['LDC2004S13-Part1', 'LDC2005S13-Part2']:
+        print(f"\n\nWorking on dataset: {data_set}")
+        file_count = __process_data(
+            os.path.join(audio_root, data_set),
+            os.path.join(transcript_root, data_set),
+            dest_root,
+            min_slice_duration,
+            file_count,
+            keep_low_conf,
+            rem_noises,
+            emojify,
+        )
+
+        print(f"Total file count so far: {file_count}")
+
+
+if __name__ == "__main__":
+    main()

scripts/dataset_processing/process_hub5_data.py

@@ -0,0 +1,285 @@
+# 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 is heavily derived from the Patter HUB5 processing script written
+# by Ryan Leary
+
+import argparse
+import glob
+import json
+import os
+import re
+import subprocess
+import sys
+from collections import namedtuple
+from math import ceil, floor
+from operator import attrgetter
+
+import numpy as np
+import scipy.io.wavfile as wavfile
+from tqdm import tqdm
+
+parser = argparse.ArgumentParser(description="Prepare HUB5 data for training/eval")
+parser.add_argument(
+    "--data_root",
+    default=None,
+    type=str,
+    required=True,
+    help="The path to the root LDC HUB5 dataset directory.",
+)
+parser.add_argument(
+    "--dest_root",
+    default=None,
+    type=str,
+    required=True,
+    help="Path to the destination root directory for processed files.",
+)
+
+# Optional arguments
+parser.add_argument(
+    "--min_slice_duration",
+    default=10.0,
+    type=float,
+    help="Minimum audio slice duration after processing.",
+)
+
+args = parser.parse_args()
+
+StmUtterance = namedtuple(
+    'StmUtterance',
+    [
+        'filename',
+        'channel',
+        'speaker_id',
+        'begin',
+        'end',
+        'label',
+        'transcript',
+    ],
+)
+STM_LINE_FMT = re.compile(r"^(\w+)\s+(\w+)\s+(\w+)\s+([0-9.]+)\s+([0-9.]+)\s+(<.*>)?\s+(.+)$")
+
+# Transcription errors and their fixes
+TRANSCRIPT_BUGS = {"en_4622-B-12079-12187": "KIND OF WEIRD BUT"}
+
+
+def get_utt_id(segment):
+    """
+    Gives utterance IDs in a form like: en_4156-a-36558-37113
+    """
+    return "{}-{}-{}-{}".format(
+        segment.filename,
+        segment.channel,
+        int(segment.begin * 100),
+        int(segment.end * 100),
+    )
+
+
+def convert_utterances(sph_path, wav_path):
+    """
+    Converts a sphere audio file to wav.
+    """
+    cmd = ["sph2pipe", "-f", "wav", "-p", sph_path, wav_path]
+    subprocess.run(cmd)
+
+
+def create_wavs(data_root, dest_root):
+    """
+    Converts the English sph files to wav using sph2pipe.
+    """
+    sph_root = os.path.join(data_root, "hub5e_00", "english")
+    sph_list = glob.glob(os.path.join(sph_root, "*.sph"))
+
+    # Iterate over each sphere file and conver to wav
+    for sph_path in tqdm(sph_list, desc="Converting to wav", unit="file"):
+        sph_name, _ = os.path.splitext(os.path.basename(sph_path))
+        wav_path = os.path.join(dest_root, 'full_audio_wav', sph_name + ".wav")
+        cmd = ["sph2pipe", "-f", "wav", "-p", sph_path, wav_path]
+        subprocess.run(cmd)
+
+
+def process_transcripts(dataset_root):
+    """
+    Reads in transcripts for each audio segment and processes them.
+    """
+    stm_path = os.path.join(
+        dataset_root,
+        "2000_hub5_eng_eval_tr",
+        "reference",
+        "hub5e00.english.000405.stm",
+    )
+    results = []
+    chars = set()
+
+    with open(stm_path, "r") as fh:
+        for line in fh:
+            # lines with ';;' are comments
+            if line.startswith(";;"):
+                continue
+
+            if "IGNORE_TIME_SEGMENT_" in line:
+                continue
+            line = line.replace("<B_ASIDE>", "").replace("<E_ASIDE>", "")
+            line = line.replace("(%HESITATION)", "UH")
+            line = line.replace("-", "")
+            line = line.replace("(%UH)", "UH")
+            line = line.replace("(%AH)", "UH")
+            line = line.replace("(", "").replace(")", "")
+
+            line = line.lower()
+
+            m = STM_LINE_FMT.search(line.strip())
+            utt = StmUtterance(*m.groups())
+
+            # Convert begin/end times to float
+            utt = utt._replace(begin=float(utt.begin))
+            utt = utt._replace(end=float(utt.end))
+
+            # Check for utterance in dict of transcript mistakes
+            transcript_update = TRANSCRIPT_BUGS.get(get_utt_id(utt))
+            if transcript_update is not None:
+                utt = utt._replace(transcript=transcript_update)
+
+            results.append(utt)
+            chars.update(list(utt.transcript))
+    return results, chars
+
+
+def write_one_segment(dest_root, speaker_id, count, audio, sr, duration, transcript):
+    """
+    Writes out one segment of audio, and writes its corresponding transcript
+    in the manifest.
+
+    Args:
+        dest_root: the path to the output directory root
+        speaker_id: ID of the speaker, used in file naming
+        count: number of segments from this speaker so far
+        audio: the segment's audio data
+        sr: sample rate of the audio
+        duration: duration of the audio
+        transcript: the corresponding transcript
+    """
+    audio_path = os.path.join(dest_root, "audio", f"{speaker_id}_{count:03}.wav")
+
+    manifest_path = os.path.join(dest_root, "manifest_hub5.json")
+
+    # Write audio
+    wavfile.write(audio_path, sr, audio)
+
+    # Write transcript
+    transcript = {
+        "audio_filepath": audio_path,
+        "duration": duration,
+        "text": transcript,
+    }
+    with open(manifest_path, 'a') as f:
+        json.dump(transcript, f)
+        f.write('\n')
+
+
+def segment_audio(info_list, dest_root, min_slice_duration):
+    """
+    Combines audio into >= min_slice_duration segments of the same speaker,
+    and writes the combined transcripts into a manifest.
+
+    Args:
+        info_list: list of StmUtterance objects with transcript information.
+        dest_root: path to output destination
+        min_slice_duration: min number of seconds per output audio slice
+    """
+    info_list = sorted(info_list, key=attrgetter('speaker_id', 'begin'))
+
+    prev_id = None  # For checking audio concatenation
+    id_count = 0
+
+    sample_rate, audio_data = None, None
+    transcript_buffer = ''
+    audio_buffer = []
+    buffer_duration = 0.0
+
+    # Iterate through utterances to build segments
+    for info in info_list:
+        if info.speaker_id != prev_id:
+            # Scrap the remainder in the buffers and start next segment
+            prev_id = info.speaker_id
+            id_count = 0
+
+            sample_rate, audio_data = wavfile.read(os.path.join(dest_root, 'full_audio_wav', info.filename + '.wav'))
+            transcript_buffer = ''
+            audio_buffer = []
+            buffer_duration = 0.0
+
+        # Append utterance info to buffers
+        transcript_buffer += info.transcript
+        channel = 0 if info.channel.lower() == 'a' else 1
+        audio_buffer.append(
+            audio_data[
+                floor(info.begin * sample_rate) : ceil(info.end * sample_rate),
+                channel,
+            ]
+        )
+        buffer_duration += info.end - info.begin
+
+        if buffer_duration < min_slice_duration:
+            transcript_buffer += ' '
+        else:
+            # Write out segment and transcript
+            id_count += 1
+            write_one_segment(
+                dest_root,
+                info.speaker_id,
+                id_count,
+                np.concatenate(audio_buffer, axis=0),
+                sample_rate,
+                buffer_duration,
+                transcript_buffer,
+            )
+
+            transcript_buffer = ''
+            audio_buffer = []
+            buffer_duration = 0.0
+
+
+def main():
+    data_root = args.data_root
+    dest_root = args.dest_root
+
+    min_slice_duration = args.min_slice_duration
+
+    if not os.path.exists(os.path.join(dest_root, 'full_audio_wav')):
+        os.makedirs(os.path.join(dest_root, 'full_audio_wav'))
+    if not os.path.exists(os.path.join(dest_root, 'audio')):
+        os.makedirs(os.path.join(dest_root, 'audio'))
+
+    # Create/wipe manifest contents
+    open(os.path.join(dest_root, "manifest_hub5.json"), 'w').close()
+
+    # Convert full audio files from .sph to .wav
+    create_wavs(data_root, dest_root)
+
+    # Get each audio transcript from transcript file
+    info_list, chars = process_transcripts(data_root)
+
+    print("Writing out vocab file", file=sys.stderr)
+    with open(os.path.join(dest_root, "vocab.txt"), 'w') as fh:
+        for x in sorted(list(chars)):
+            fh.write(x + "\n")
+
+    # Segment the audio data
+    print("Segmenting audio and writing manifest")
+    segment_audio(info_list, dest_root, min_slice_duration)
+
+
+if __name__ == '__main__':
+    main()

scripts/dataset_processing/process_speech_commands_data.py

@@ -0,0 +1,521 @@
+# 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.
+
+"""
+Usage:
+
+python process_speech_commands_data.py \
+    --data_root=<absolute path to where the data should be stored> \
+    --data_version=<either 1 or 2, indicating version of the dataset> \
+    --class_split=<either "all" or "sub", indicates whether all 30/35 classes should be used, or the 10+2 split should be used> \
+    --num_processes=<number of processes to use for data preprocessing> \
+    --rebalance \
+    --log
+"""
+
+import argparse
+import glob
+import json
+import logging
+import os
+import re
+import tarfile
+import urllib.request
+from collections import defaultdict
+from functools import partial
+from multiprocessing import Pool
+from typing import Dict, List, Set, Tuple
+
+import librosa
+import numpy as np
+import soundfile
+from tqdm import tqdm
+
+from nemo.utils.tar_utils import safe_extract
+
+URL_v1 = 'http://download.tensorflow.org/data/speech_commands_v0.01.tar.gz'
+URL_v2 = 'http://download.tensorflow.org/data/speech_commands_v0.02.tar.gz'
+
+
+def __maybe_download_file(destination: str, source: str) -> str:
+    """
+    Downloads source to destination if it doesn't exist.
+    If exists, skips download
+
+    Args:
+        destination: local filepath
+        source: url of resource
+
+    Returns:
+        Local filepath of the downloaded file
+    """
+    if not os.path.exists(destination):
+        logging.info(f'{destination} does not exist. Downloading ...')
+        urllib.request.urlretrieve(source, filename=destination + '.tmp')
+        os.rename(destination + '.tmp', destination)
+        logging.info(f'Downloaded {destination}.')
+    else:
+        logging.info(f'Destination {destination} exists. Skipping.')
+    return destination
+
+
+def __extract_all_files(filepath: str, data_dir: str):
+    if not os.path.exists(data_dir):
+        extract_file(filepath, data_dir)
+    else:
+        logging.info(f'Skipping extracting. Data already there {data_dir}')
+
+
+def extract_file(filepath: str, data_dir: str):
+    try:
+        with tarfile.open(filepath) as tar:
+            safe_extract(tar, data_dir)
+    except Exception:
+        logging.info('Not extracting. Maybe already there?')
+
+
+def __get_mp_chunksize(dataset_size: int, num_processes: int) -> int:
+    """
+    Returns the number of chunks to split the dataset into for multiprocessing.
+
+    Args:
+        dataset_size: size of the dataset
+        num_processes: number of processes to use for multiprocessing
+
+        Returns:
+            Number of chunks to split the dataset into for multiprocessing
+    """
+    chunksize = dataset_size // num_processes
+    return chunksize if chunksize > 0 else 1
+
+
+def __construct_filepaths(
+    all_files: List[str],
+    valset_uids: Set[str],
+    testset_uids: Set[str],
+    class_split: str,
+    class_subset: List[str],
+    pattern: str,
+) -> Tuple[Dict[str, int], Dict[str, List[tuple]], List[tuple], List[tuple], List[tuple], List[tuple], List[tuple]]:
+    """
+    Prepares the filepaths for the dataset.
+
+    Args:
+        all_files: list of all files in the dataset
+        valset_uids: set of uids of files in the validation set
+        testset_uids: set of uids of files in the test set
+        class_split: whether to use all classes as distinct labels, or to use
+            10 classes subset and rest of the classes as noise or background
+        class_subset: list of classes to consider if `class_split` is set to `sub`
+        pattern: regex pattern to match the file names in the dataset
+    """
+
+    label_count = defaultdict(int)
+    label_filepaths = defaultdict(list)
+    unknown_val_filepaths = []
+    unknown_test_filepaths = []
+
+    train, val, test = [], [], []
+    for entry in all_files:
+        r = re.match(pattern, entry)
+        if r:
+            label, uid = r.group(2), r.group(3)
+
+            if label == '_background_noise_' or label == 'silence':
+                continue
+
+        if class_split == 'sub' and label not in class_subset:
+            label = 'unknown'
+
+            if uid in valset_uids:
+                unknown_val_filepaths.append((label, entry))
+            elif uid in testset_uids:
+                unknown_test_filepaths.append((label, entry))
+
+        if uid not in valset_uids and uid not in testset_uids:
+            label_count[label] += 1
+            label_filepaths[label].append((label, entry))
+
+        if label == 'unknown':
+            continue
+
+        if uid in valset_uids:
+            val.append((label, entry))
+        elif uid in testset_uids:
+            test.append((label, entry))
+        else:
+            train.append((label, entry))
+
+    return {
+        'label_count': label_count,
+        'label_filepaths': label_filepaths,
+        'unknown_val_filepaths': unknown_val_filepaths,
+        'unknown_test_filepaths': unknown_test_filepaths,
+        'train': train,
+        'val': val,
+        'test': test,
+    }
+
+
+def __construct_silence_set(
+    rng: np.random.RandomState, sampling_rate: int, silence_stride: int, data_folder: str, background_noise: str
+) -> List[str]:
+    """
+    Creates silence files given a background noise.
+
+    Args:
+        rng: Random state for random number generator
+        sampling_rate: sampling rate of the audio
+        silence_stride: stride for creating silence files
+        data_folder: folder containing the silence directory
+        background_noise: filepath of the background noise
+
+    Returns:
+        List of filepaths of silence files
+    """
+    silence_files = []
+    if '.wav' in background_noise:
+        y, sr = librosa.load(background_noise, sr=sampling_rate)
+
+        for i in range(0, len(y) - sampling_rate, silence_stride):
+            file_path = f'silence/{os.path.basename(background_noise)[:-4]}_{i}.wav'
+            y_slice = y[i : i + sampling_rate] * rng.uniform(0.0, 1.0)
+            out_file_path = os.path.join(data_folder, file_path)
+            soundfile.write(out_file_path, y_slice, sr)
+
+            silence_files.append(('silence', out_file_path))
+
+    return silence_files
+
+
+def __rebalance_files(max_count: int, label_filepath: str) -> Tuple[str, List[str], int]:
+    """
+    Rebalance the number of samples for a class.
+
+    Args:
+        max_count: maximum number of samples for a class
+        label_filepath: list of filepaths for a class
+
+    Returns:
+        Rebalanced list of filepaths along with the label name and the number of samples
+    """
+    command, samples = label_filepath
+    filepaths = [sample[1] for sample in samples]
+
+    rng = np.random.RandomState(0)
+    filepaths = np.asarray(filepaths)
+    num_samples = len(filepaths)
+
+    if num_samples < max_count:
+        difference = max_count - num_samples
+        duplication_ids = rng.choice(num_samples, difference, replace=True)
+        filepaths = np.append(filepaths, filepaths[duplication_ids], axis=0)
+
+    return command, filepaths, num_samples
+
+
+def __prepare_metadata(skip_duration, sample: Tuple[str, str]) -> dict:
+    """
+    Creates the manifest entry for a file.
+
+    Args:
+        skip_duration: Whether to skip the computation of duration
+        sample: Tuple of label and filepath
+
+    Returns:
+        Manifest entry of the file
+    """
+    label, audio_path = sample
+    return json.dumps(
+        {
+            'audio_filepath': audio_path,
+            'duration': 0.0 if skip_duration else librosa.core.get_duration(filename=audio_path),
+            'command': label,
+        }
+    )
+
+
+def __process_data(
+    data_folder: str,
+    dst_folder: str,
+    num_processes: int = 1,
+    rebalance: bool = False,
+    class_split: str = 'all',
+    skip_duration: bool = False,
+):
+    """
+    Processes the data and generates the manifests.
+
+    Args:
+        data_folder: source with wav files and validation / test lists
+        dst_folder: where manifest files will be stored
+        num_processes: number of processes
+        rebalance: rebalance the classes to have same number of samples
+        class_split: whether to use all classes as distinct labels, or to use
+            10 classes subset and rest of the classes as noise or background
+        skip_duration: Bool whether to skip duration computation. Use this only for
+            colab notebooks where knowing duration is not necessary for demonstration
+    """
+
+    os.makedirs(dst_folder, exist_ok=True)
+
+    # Used for 10 classes + silence + unknown class setup - Only used when class_split is 'sub'
+    class_subset = ['yes', 'no', 'up', 'down', 'left', 'right', 'on', 'off', 'stop', 'go']
+
+    pattern = re.compile(r'(.+\/)?(\w+)\/([^_]+)_.+wav')
+    all_files = glob.glob(os.path.join(data_folder, '*/*wav'))
+
+    # Get files in the validation set
+    valset_uids = set()
+    with open(os.path.join(data_folder, 'validation_list.txt')) as fin:
+        for line in fin:
+            r = re.match(pattern, line)
+            if r:
+                valset_uids.add(r.group(3))
+
+    # Get files in the test set
+    testset_uids = set()
+    with open(os.path.join(data_folder, 'testing_list.txt')) as fin:
+        for line in fin:
+            r = re.match(pattern, line)
+            if r:
+                testset_uids.add(r.group(3))
+
+    logging.info('Validation and test set lists extracted')
+
+    filepath_info = __construct_filepaths(all_files, valset_uids, testset_uids, class_split, class_subset, pattern)
+    label_count = filepath_info['label_count']
+    label_filepaths = filepath_info['label_filepaths']
+    unknown_val_filepaths = filepath_info['unknown_val_filepaths']
+    unknown_test_filepaths = filepath_info['unknown_test_filepaths']
+    train = filepath_info['train']
+    val = filepath_info['val']
+    test = filepath_info['test']
+
+    logging.info('Prepared filepaths for dataset')
+
+    pool = Pool(num_processes)
+
+    # Add silence and unknown class label samples
+    if class_split == 'sub':
+        logging.info('Perforiming 10+2 class subsplit')
+
+        silence_path = os.path.join(data_folder, 'silence')
+        os.makedirs(silence_path, exist_ok=True)
+
+        silence_stride = 1000  # 0.0625 second stride
+        sampling_rate = 16000
+        folder = os.path.join(data_folder, '_background_noise_')
+
+        silence_files = []
+        rng = np.random.RandomState(0)
+
+        background_noise_files = [os.path.join(folder, x) for x in os.listdir(folder)]
+        silence_set_fn = partial(__construct_silence_set, rng, sampling_rate, silence_stride, data_folder)
+        for silence_flist in tqdm(
+            pool.imap(
+                silence_set_fn, background_noise_files, __get_mp_chunksize(len(background_noise_files), num_processes)
+            ),
+            total=len(background_noise_files),
+            desc='Constructing silence set',
+        ):
+            silence_files.extend(silence_flist)
+
+        rng = np.random.RandomState(0)
+        rng.shuffle(silence_files)
+        logging.info(f'Constructed silence set of {len(silence_files)}')
+
+        # Create the splits
+        rng = np.random.RandomState(0)
+        silence_split = 0.1
+        unknown_split = 0.1
+
+        # train split
+        num_total_samples = sum([label_count[cls] for cls in class_subset])
+        num_silence_samples = int(np.ceil(silence_split * num_total_samples))
+
+        # initialize sample
+        label_count['silence'] = 0
+        label_filepaths['silence'] = []
+
+        for silence_id in range(num_silence_samples):
+            label_count['silence'] += 1
+            label_filepaths['silence'].append(silence_files[silence_id])
+
+        train.extend(label_filepaths['silence'])
+
+        # Update train unknown set
+        unknown_train_samples = label_filepaths['unknown']
+
+        rng.shuffle(unknown_train_samples)
+        unknown_size = int(np.ceil(unknown_split * num_total_samples))
+
+        label_count['unknown'] = unknown_size
+        label_filepaths['unknown'] = unknown_train_samples[:unknown_size]
+
+        train.extend(label_filepaths['unknown'])
+
+        logging.info('Train set prepared')
+
+        # val set silence
+        num_val_samples = len(val)
+        num_silence_samples = int(np.ceil(silence_split * num_val_samples))
+
+        val_idx = label_count['silence'] + 1
+        for silence_id in range(num_silence_samples):
+            val.append(silence_files[val_idx + silence_id])
+
+        # Update val unknown set
+        rng.shuffle(unknown_val_filepaths)
+        unknown_size = int(np.ceil(unknown_split * num_val_samples))
+
+        val.extend(unknown_val_filepaths[:unknown_size])
+
+        logging.info('Validation set prepared')
+
+        # test set silence
+        num_test_samples = len(test)
+        num_silence_samples = int(np.ceil(silence_split * num_test_samples))
+
+        test_idx = val_idx + num_silence_samples + 1
+        for silence_id in range(num_silence_samples):
+            test.append(silence_files[test_idx + silence_id])
+
+        # Update test unknown set
+        rng.shuffle(unknown_test_filepaths)
+        unknown_size = int(np.ceil(unknown_split * num_test_samples))
+
+        test.extend(unknown_test_filepaths[:unknown_size])
+
+        logging.info('Test set prepared')
+
+    max_command = None
+    max_count = -1
+    for command, count in label_count.items():
+        if command == 'unknown':
+            continue
+
+        if count > max_count:
+            max_count = count
+            max_command = command
+
+    if rebalance:
+        logging.info(f'Command with maximum number of samples = {max_command} with {max_count} samples')
+        logging.info(f'Rebalancing dataset by duplicating classes with less than {max_count} samples...')
+
+        rebalance_fn = partial(__rebalance_files, max_count)
+        for command, filepaths, num_samples in tqdm(
+            pool.imap(rebalance_fn, label_filepaths.items(), __get_mp_chunksize(len(label_filepaths), num_processes)),
+            total=len(label_filepaths),
+            desc='Rebalancing dataset',
+        ):
+            if num_samples < max_count:
+                logging.info(f'Extended class label {command} from {num_samples} samples to {len(filepaths)} samples')
+                label_filepaths[command] = [(command, filepath) for filepath in filepaths]
+
+        del train
+        train = []
+        for label, samples in label_filepaths.items():
+            train.extend(samples)
+
+    manifests = [
+        ('train_manifest.json', train),
+        ('validation_manifest.json', val),
+        ('test_manifest.json', test),
+    ]
+
+    metadata_fn = partial(__prepare_metadata, skip_duration)
+    for manifest_filename, dataset in manifests:
+        num_files = len(dataset)
+
+        logging.info(f'Preparing manifest : {manifest_filename} with #{num_files} files')
+
+        manifest = [
+            metadata
+            for metadata in tqdm(
+                pool.imap(metadata_fn, dataset, __get_mp_chunksize(len(dataset), num_processes)),
+                total=num_files,
+                desc=f'Preparing {manifest_filename}',
+            )
+        ]
+
+        with open(os.path.join(dst_folder, manifest_filename), 'w') as fout:
+            for metadata in manifest:
+                fout.write(metadata + '\n')
+
+        logging.info(f'Finished construction of manifest. Path: {os.path.join(dst_folder, manifest_filename)}')
+
+    pool.close()
+
+    if skip_duration:
+        logging.info(
+            f'\n<<NOTE>> Duration computation was skipped for demonstration purposes on Colaboratory.\n'
+            f'In order to replicate paper results and properly perform data augmentation, \n'
+            f'please recompute the manifest file without the `--skip_duration` flag !\n'
+        )
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Google Speech Commands Data download and preprocessing')
+    parser.add_argument('--data_root', required=True, help='Root directory for storing data')
+    parser.add_argument(
+        '--data_version',
+        required=True,
+        default=1,
+        type=int,
+        choices=[1, 2],
+        help='Version of the speech commands dataset to download',
+    )
+    parser.add_argument(
+        '--class_split', default='all', choices=['all', 'sub'], help='Whether to consider all classes or only a subset'
+    )
+    parser.add_argument('--num_processes', default=1, type=int, help='Number of processes')
+    parser.add_argument('--rebalance', action='store_true', help='Rebalance the number of samples in each class')
+    parser.add_argument('--skip_duration', action='store_true', help='Skip computing duration of audio files')
+    parser.add_argument('--log', action='store_true', help='Generate logs')
+    args = parser.parse_args()
+
+    if args.log:
+        logging.basicConfig(level=logging.DEBUG)
+
+    data_root = args.data_root
+    data_set = f'google_speech_recognition_v{args.data_version}'
+    data_folder = os.path.join(data_root, data_set)
+
+    logging.info(f'Working on: {data_set}')
+
+    URL = URL_v1 if args.data_version == 1 else URL_v2
+
+    # Download and extract
+    if not os.path.exists(data_folder):
+        file_path = os.path.join(data_root, data_set + '.tar.bz2')
+        logging.info(f'Getting {data_set}')
+        __maybe_download_file(file_path, URL)
+        logging.info(f'Extracting {data_set}')
+        __extract_all_files(file_path, data_folder)
+
+    logging.info(f'Processing {data_set}')
+    __process_data(
+        data_folder,
+        data_folder,
+        num_processes=args.num_processes,
+        rebalance=args.rebalance,
+        class_split=args.class_split,
+        skip_duration=args.skip_duration,
+    )
+    logging.info('Done!')
+
+
+if __name__ == '__main__':
+    main()

scripts/dataset_processing/process_vad_data.py

@@ -0,0 +1,503 @@
+# 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.
+"""
+Usage:
+
+python process_vad_data.py \
+    --out_dir=<output path to where the generated manifest should be stored> \
+    --speech_data_root=<path where the speech data are stored> \
+    --background_data_root=<path where the background data are stored> \
+    --rebalance_method=<'under' or 'over' or 'fixed'> \ 
+    --log
+    (Optional --demo (for demonstration in tutorial). If you want to use your own background noise data, make sure to delete --demo)
+"""
+import argparse
+import glob
+import json
+import logging
+import os
+import tarfile
+import urllib.request
+
+import librosa
+import numpy as np
+import soundfile as sf
+from sklearn.model_selection import train_test_split
+
+from nemo.utils.tar_utils import safe_extract
+
+sr = 16000
+
+# google speech command v2
+URL = "http://download.tensorflow.org/data/speech_commands_v0.02.tar.gz"
+
+
+def __maybe_download_file(destination: str, source: str):
+    """
+    Downloads source to destination if it doesn't exist.
+    If exists, skips download
+    Args:
+        destination: local filepath
+        source: url of resource
+
+    Returns:
+
+    """
+    if not os.path.exists(destination):
+        logging.info(f"{destination} does not exist. Downloading ...")
+        urllib.request.urlretrieve(source, filename=destination + '.tmp')
+        os.rename(destination + '.tmp', destination)
+        logging.info(f"Downloaded {destination}.")
+    else:
+        logging.info(f"Destination {destination} exists. Skipping.")
+    return destination
+
+
+def extract_file(filepath: str, data_dir: str):
+    try:
+        with tarfile.open(filepath) as tar:
+            safe_extract(tar, data_dir)
+    except Exception:
+        logging.info('Not extracting. Maybe already there?')
+
+
+def __extract_all_files(filepath: str, data_root: str, data_dir: str):
+    if not os.path.exists(data_dir):
+        extract_file(filepath, data_dir)
+    else:
+        logging.info(f'Skipping extracting. Data already there {data_dir}')
+
+
+def split_train_val_test(data_dir, file_type, test_size=0.1, val_size=0.1, demo=False):
+    X = []
+    if file_type == "speech":
+        for o in os.listdir(data_dir):
+            if os.path.isdir(os.path.join(data_dir, o)) and o.split("/")[-1] != "_background_noise_":
+                X.extend(glob.glob(os.path.join(data_dir, o) + '/*.wav'))
+
+        if demo:
+            logging.info(
+                f"For Demonstration, we use {int(len(X)/100)}/{len(X)} speech data. Make sure to remove --demo flag when you actually train your model!"
+            )
+            X = np.random.choice(X, int(len(X) / 100), replace=False)
+
+    else:
+        for o in os.listdir(data_dir):
+            if os.path.isdir(os.path.join(data_dir, o)):
+                X.extend(glob.glob(os.path.join(data_dir, o) + '/*.wav'))
+            else:  # for using "_background_noise_" from google speech commands as background data
+                if o.endswith(".wav"):
+                    X.append(os.path.join(data_dir, o))
+
+    X_train, X_test = train_test_split(X, test_size=test_size, random_state=1)
+    val_size_tmp = val_size / (1 - test_size)
+    X_train, X_val = train_test_split(X_train, test_size=val_size_tmp, random_state=1)
+
+    with open(os.path.join(data_dir, file_type + "_training_list.txt"), "w") as outfile:
+        outfile.write("\n".join(X_train))
+    with open(os.path.join(data_dir, file_type + "_testing_list.txt"), "w") as outfile:
+        outfile.write("\n".join(X_test))
+    with open(os.path.join(data_dir, file_type + "_validation_list.txt"), "w") as outfile:
+        outfile.write("\n".join(X_val))
+
+    logging.info(f'Overall: {len(X)}, Train: {len(X_train)}, Validatoin: {len(X_val)}, Test: {len(X_test)}')
+    logging.info(f"Finish spliting train, val and test for {file_type}. Write to files!")
+
+
+def process_google_speech_train(data_dir):
+    X = []
+    for o in os.listdir(data_dir):
+        if os.path.isdir(os.path.join(data_dir, o)) and o.split("/")[-1] != "_background_noise_":
+            X.extend(glob.glob(os.path.join(data_dir, o) + '/*.wav'))
+
+    short_files = [i.split(data_dir)[1] for i in files]
+
+    with open(os.path.join(data_dir, 'testing_list.txt'), 'r') as allfile:
+        testing_list = allfile.read().splitlines()
+
+    with open(os.path.join(data_dir, 'validation_list.txt'), 'r') as allfile:
+        validation_list = allfile.read().splitlines()
+
+    exist_set = set(testing_list).copy()
+    exist_set.update(set(validation_list))
+
+    training_list = [i for i in short_files if i not in exist_set]
+
+    with open(os.path.join(data_dir, "training_list.txt"), "w") as outfile:
+        outfile.write("\n".join(training_list))
+
+    logging.info(
+        f'Overall: {len(files)}, Train: {len(training_list)}, Validatoin: {len(validation_list)}, Test: {len(testing_list)}'
+    )
+
+
+def write_manifest(
+    out_dir,
+    files,
+    prefix,
+    manifest_name,
+    start=0.0,
+    end=None,
+    duration_stride=1.0,
+    duration_max=None,
+    duration_limit=100.0,
+    filter_long=False,
+):
+    """
+    Given a list of files, segment each file and write them to manifest with restrictions.
+    Args:
+        out_dir: directory of generated manifest
+        files: list of files to be processed
+        prefix: label of samples
+        manifest_name: name of generated manifest
+        start: beginning of audio of generating segment
+        end: end of audio of generating segment
+        duration_stride: stride for segmenting audio samples
+        duration_max: duration for each segment
+        duration_limit: duration threshold for filtering out long audio samples
+        filter_long: boolean to determine whether to filter out long audio samples
+    Returns:
+    """
+    seg_num = 0
+    skip_num = 0
+    if duration_max is None:
+        duration_max = 1e9
+
+    if not os.path.exists(out_dir):
+        logging.info(f'Outdir {out_dir} does not exist. Creat directory.')
+        os.mkdir(out_dir)
+
+    output_path = os.path.join(out_dir, manifest_name + '.json')
+    with open(output_path, 'w') as fout:
+        for file in files:
+            label = prefix
+
+            try:
+                x, _sr = librosa.load(file, sr=sr)
+                duration = librosa.get_duration(y=x, sr=sr)
+
+            except Exception:
+                continue
+
+            if filter_long and duration > duration_limit:
+                skip_num += 1
+                continue
+
+            offsets = []
+            durations = []
+
+            if duration > duration_max:
+                current_offset = start
+
+                while current_offset < duration:
+                    if end is not None and current_offset > end:
+                        break
+
+                    difference = duration - current_offset
+
+                    if difference < duration_max:
+                        break
+
+                    offsets.append(current_offset)
+                    durations.append(duration_max)
+
+                    current_offset += duration_stride
+
+            else:
+                # Duration is not long enough! Skip
+                skip_num += 1
+
+            for duration, offset in zip(durations, offsets):
+                metadata = {
+                    'audio_filepath': file,
+                    'duration': duration,
+                    'label': label,
+                    'text': '_',  # for compatibility with ASRAudioText
+                    'offset': offset,
+                }
+                json.dump(metadata, fout)
+                fout.write('\n')
+                fout.flush()
+                seg_num += 1
+    return skip_num, seg_num, output_path
+
+
+def load_list_write_manifest(
+    data_dir,
+    out_dir,
+    filename,
+    prefix,
+    start,
+    end,
+    duration_stride=1.0,
+    duration_max=1.0,
+    duration_limit=100.0,
+    filter_long=True,
+):
+
+    filename = prefix + '_' + filename
+    file_path = os.path.join(data_dir, filename)
+
+    with open(file_path, 'r') as allfile:
+        files = allfile.read().splitlines()
+
+    manifest_name = filename.split('_list.txt')[0] + '_manifest'
+    skip_num, seg_num, output_path = write_manifest(
+        out_dir,
+        files,
+        prefix,
+        manifest_name,
+        start,
+        end,
+        duration_stride,
+        duration_max,
+        duration_limit,
+        filter_long=True,
+    )
+    return skip_num, seg_num, output_path
+
+
+def rebalance_json(data_dir, data_json, num, prefix):
+    data = []
+    seg = 0
+    with open(data_json, 'r') as f:
+        for line in f:
+            data.append(json.loads(line))
+
+    filename = data_json.split('/')[-1]
+    fout_path = os.path.join(data_dir, prefix + "_" + filename)
+
+    if len(data) >= num:
+        selected_sample = np.random.choice(data, num, replace=False)
+    else:
+        selected_sample = np.random.choice(data, num, replace=True)
+
+    with open(fout_path, 'a') as fout:
+        for i in selected_sample:
+            seg += 1
+            json.dump(i, fout)
+            fout.write('\n')
+            fout.flush()
+
+    logging.info(f'Get {seg}/{num} to  {fout_path} from {data_json}')
+    return fout_path
+
+
+def generate_variety_noise(data_dir, filename, prefix):
+
+    curr_dir = data_dir.split("_background_noise_")[0]
+    silence_path = os.path.join(curr_dir, "_background_noise_more")
+
+    if not os.path.exists(silence_path):
+        os.mkdir(silence_path)
+
+    silence_stride = 1000  # stride = 1/16 seconds
+    sampling_rate = 16000
+
+    silence_files = []
+    rng = np.random.RandomState(0)
+
+    filename = prefix + '_' + filename
+    file_path = os.path.join(data_dir, filename)
+
+    with open(file_path, 'r') as allfile:
+        files = allfile.read().splitlines()
+
+    for file in files:
+        y, sr = librosa.load(path=file, sr=sampling_rate)
+
+        for i in range(
+            0, len(y) - sampling_rate, silence_stride * 100
+        ):  # stride * 100 to generate less samples for demo
+            file_name = "{}_{}.wav".format(file.split("/")[-1], i)
+            y_slice = y[i : i + sampling_rate]
+            magnitude = rng.uniform(0.0, 1.0)
+            y_slice *= magnitude
+            out_file_path = os.path.join(silence_path, file_name)
+            sf.write(out_file_path, y_slice, sr)
+
+            silence_files.append(out_file_path)
+
+    new_list_file = os.path.join(silence_path, filename)
+    with open(new_list_file, "w") as outfile:
+        outfile.write("\n".join(silence_files))
+
+    logging.info(f"Generate {len(out_file_path)} background files for {file_path}. => {new_list_file} !")
+    return len(silence_files)
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Speech and backgound data download and preprocess')
+    parser.add_argument("--out_dir", required=False, default='./manifest/', type=str)
+    parser.add_argument("--speech_data_root", required=True, default=None, type=str)
+    parser.add_argument("--background_data_root", required=True, default=None, type=str)
+    parser.add_argument('--test_size', required=False, default=0.1, type=float)
+    parser.add_argument('--val_size', required=False, default=0.1, type=float)
+    parser.add_argument('--window_length_in_sec', required=False, default=0.63, type=float)
+    parser.add_argument('--log', required=False, action='store_true')
+    parser.add_argument('--rebalance_method', required=False, default=None, type=str)
+    parser.add_argument('--demo', required=False, action='store_true')
+    parser.set_defaults(log=False, generate=False)
+    args = parser.parse_args()
+
+    if not args.rebalance_method:
+        rebalance = False
+    else:
+        if args.rebalance_method != 'over' and args.rebalance_method != 'under' and args.rebalance_method != 'fixed':
+            raise NameError("Please select a valid sampling method: over/under/fixed.")
+        else:
+            rebalance = True
+
+    if args.log:
+        logging.basicConfig(level=logging.DEBUG)
+
+    # Download speech data
+    speech_data_root = args.speech_data_root
+    data_set = "google_speech_recognition_v2"
+    speech_data_folder = os.path.join(speech_data_root, data_set)
+
+    background_data_folder = args.background_data_root
+    logging.info(f"Working on: {data_set}")
+
+    # Download and extract speech data
+    if not os.path.exists(speech_data_folder):
+        file_path = os.path.join(speech_data_root, data_set + ".tar.bz2")
+        logging.info(f"Getting {data_set}")
+        __maybe_download_file(file_path, URL)
+        logging.info(f"Extracting {data_set}")
+        __extract_all_files(file_path, speech_data_root, speech_data_folder)
+
+    logging.info(f"Split speech data!")
+    # dataset provide testing.txt and validation.txt feel free to split data using that with process_google_speech_train
+    split_train_val_test(speech_data_folder, "speech", args.test_size, args.val_size, args.demo)
+
+    logging.info(f"Split background data!")
+    split_train_val_test(background_data_folder, "background", args.test_size, args.val_size)
+
+    out_dir = args.out_dir
+
+    # Process Speech manifest
+    logging.info(f"=== Write speech data to manifest!")
+    skip_num_val, speech_seg_num_val, speech_val = load_list_write_manifest(
+        speech_data_folder,
+        out_dir,
+        'validation_list.txt',
+        'speech',
+        0.2,
+        0.8,
+        args.window_length_in_sec,
+        args.window_length_in_sec,
+    )
+    skip_num_test, speech_seg_num_test, speech_test = load_list_write_manifest(
+        speech_data_folder, out_dir, 'testing_list.txt', 'speech', 0.2, 0.8, 0.01, args.window_length_in_sec
+    )
+    skip_num_train, speech_seg_num_train, speech_train = load_list_write_manifest(
+        speech_data_folder,
+        out_dir,
+        'training_list.txt',
+        'speech',
+        0.2,
+        0.8,
+        args.window_length_in_sec,
+        args.window_length_in_sec,
+    )
+
+    logging.info(f'Val: Skip {skip_num_val} samples. Get {speech_seg_num_val} segments! => {speech_val} ')
+    logging.info(f'Test: Skip {skip_num_test} samples. Get {speech_seg_num_test} segments! => {speech_test}')
+    logging.info(f'Train: Skip {skip_num_train} samples. Get {speech_seg_num_train} segments!=> {speech_train}')
+
+    # Process background manifest
+    # if we select to generate more background noise data
+    if args.demo:
+        logging.info("Start generating more background noise data")
+        generate_variety_noise(background_data_folder, 'validation_list.txt', 'background')
+        generate_variety_noise(background_data_folder, 'training_list.txt', 'background')
+        generate_variety_noise(background_data_folder, 'testing_list.txt', 'background')
+        background_data_folder = os.path.join(
+            background_data_folder.split("_background_noise_")[0], "_background_noise_more"
+        )
+
+    logging.info(f"=== Write background data to manifest!")
+    skip_num_val, background_seg_num_val, background_val = load_list_write_manifest(
+        background_data_folder, out_dir, 'validation_list.txt', 'background', 0, None, 0.15, args.window_length_in_sec
+    )
+    skip_num_test, background_seg_num_test, background_test = load_list_write_manifest(
+        background_data_folder, out_dir, 'testing_list.txt', 'background', 0, None, 0.01, args.window_length_in_sec
+    )
+    skip_num_train, background_seg_num_train, background_train = load_list_write_manifest(
+        background_data_folder, out_dir, 'training_list.txt', 'background', 0, None, 0.15, args.window_length_in_sec
+    )
+
+    logging.info(f'Val: Skip {skip_num_val} samples. Get {background_seg_num_val} segments! => {background_val}')
+    logging.info(f'Test: Skip {skip_num_test} samples. Get {background_seg_num_test} segments! => {background_test}')
+    logging.info(
+        f'Train: Skip {skip_num_train} samples. Get {background_seg_num_train} segments! => {background_train}'
+    )
+    min_val, max_val = min(speech_seg_num_val, background_seg_num_val), max(speech_seg_num_val, background_seg_num_val)
+    min_test, max_test = (
+        min(speech_seg_num_test, background_seg_num_test),
+        max(speech_seg_num_test, background_seg_num_test),
+    )
+    min_train, max_train = (
+        min(speech_seg_num_train, background_seg_num_train),
+        max(speech_seg_num_train, background_seg_num_train),
+    )
+
+    logging.info('Finish generating manifest!')
+
+    if rebalance:
+        # Random Oversampling: Randomly duplicate examples in the minority class.
+        # Random Undersampling: Randomly delete examples in the majority class.
+        if args.rebalance_method == 'under':
+            logging.info(f"Rebalancing number of samples in classes using {args.rebalance_method} sampling.")
+            logging.info(f'Val: {min_val} Test: {min_test} Train: {min_train}!')
+
+            rebalance_json(out_dir, background_val, min_val, 'balanced')
+            rebalance_json(out_dir, background_test, min_test, 'balanced')
+            rebalance_json(out_dir, background_train, min_train, 'balanced')
+
+            rebalance_json(out_dir, speech_val, min_val, 'balanced')
+            rebalance_json(out_dir, speech_test, min_test, 'balanced')
+            rebalance_json(out_dir, speech_train, min_train, 'balanced')
+
+        if args.rebalance_method == 'over':
+            logging.info(f"Rebalancing number of samples in classes using {args.rebalance_method} sampling.")
+            logging.info(f'Val: {max_val} Test: {max_test} Train: {max_train}!')
+
+            rebalance_json(out_dir, background_val, max_val, 'balanced')
+            rebalance_json(out_dir, background_test, max_test, 'balanced')
+            rebalance_json(out_dir, background_train, max_train, 'balanced')
+
+            rebalance_json(out_dir, speech_val, max_val, 'balanced')
+            rebalance_json(out_dir, speech_test, max_test, 'balanced')
+            rebalance_json(out_dir, speech_train, max_train, 'balanced')
+
+        if args.rebalance_method == 'fixed':
+            fixed_test, fixed_val, fixed_train = 200, 100, 500
+            logging.info(f"Rebalancing number of samples in classes using {args.rebalance_method} sampling.")
+            logging.info(f'Val: {fixed_val} Test: {fixed_test} Train: {fixed_train}!')
+
+            rebalance_json(out_dir, background_val, fixed_val, 'balanced')
+            rebalance_json(out_dir, background_test, fixed_test, 'balanced')
+            rebalance_json(out_dir, background_train, fixed_train, 'balanced')
+
+            rebalance_json(out_dir, speech_val, fixed_val, 'balanced')
+            rebalance_json(out_dir, speech_test, fixed_test, 'balanced')
+            rebalance_json(out_dir, speech_train, fixed_train, 'balanced')
+    else:
+        logging.info("Don't rebalance number of samples in classes.")
+
+
+if __name__ == '__main__':
+    main()

scripts/dataset_processing/speaker_tasks/README.md

@@ -0,0 +1,5 @@
+# Speaker Tasks Dataset Scripts
+
+In this folder are scripts to download speaker tasks (mainly for diarization) datasets. These scripts will return NeMo format manifest files to use with Diarization.
+
+We also have scripts for CallHome and DIHARD3, however the data has to be downloaded separately. If you require the scripts please leave an issue.

scripts/dataset_processing/speaker_tasks/get_aishell_diarization_data.py

@@ -0,0 +1,118 @@
+# 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.
+
+# downloads the training/eval set for AISHELL Diarization.
+# the training dataset is around 170GiB, to skip pass the --skip_train flag.
+
+import argparse
+import glob
+import logging
+import os
+import tarfile
+import urllib.request
+from pathlib import Path
+
+from nemo.collections.asr.parts.utils.manifest_utils import create_manifest
+from nemo.utils.tar_utils import safe_extract
+
+train_url = "https://www.openslr.org/resources/111/train_{}.tar.gz"
+train_datasets = ["S", "M", "L"]
+
+eval_url = "https://www.openslr.org/resources/111/test.tar.gz"
+
+
+def _load_sox_transformer():
+    try:
+        from sox import Transformer
+    except ImportError:
+        raise ImportError(
+            "Optional dependency 'sox' is required by this script. Install it with: pip install sox"
+        ) from None
+
+    return Transformer
+
+
+def extract_file(filepath: str, data_dir: str):
+    try:
+        with tarfile.open(filepath) as tar:
+            safe_extract(tar, data_dir)
+    except Exception:
+        logging.info("Not extracting. Maybe already there?")
+
+
+def __process_data(dataset_url: str, dataset_path: Path, manifest_output_path: Path):
+    os.makedirs(dataset_path, exist_ok=True)
+    tar_file_path = os.path.join(dataset_path, os.path.basename(dataset_url))
+    if not os.path.exists(tar_file_path):
+        urllib.request.urlretrieve(dataset_url, filename=tar_file_path)
+    extract_file(tar_file_path, str(dataset_path))
+    wav_path = dataset_path / 'converted_wav/'
+    extracted_dir = Path(tar_file_path).stem.replace('.tar', '')
+    flac_path = dataset_path / (extracted_dir + '/wav/')
+    __process_flac_audio(flac_path, wav_path)
+
+    audio_files = [os.path.join(os.path.abspath(wav_path), file) for file in os.listdir(str(wav_path))]
+    rttm_files = glob.glob(str(dataset_path / (extracted_dir + '/TextGrid/*.rttm')))
+    rttm_files = [os.path.abspath(file) for file in rttm_files]
+
+    audio_list = dataset_path / 'audio_files.txt'
+    rttm_list = dataset_path / 'rttm_files.txt'
+    with open(audio_list, 'w') as f:
+        f.write('\n'.join(audio_files))
+    with open(rttm_list, 'w') as f:
+        f.write('\n'.join(rttm_files))
+    create_manifest(
+        str(audio_list),
+        manifest_output_path,
+        rttm_path=str(rttm_list),
+    )
+
+
+def __process_flac_audio(flac_path, wav_path):
+    Transformer = _load_sox_transformer()
+    os.makedirs(wav_path, exist_ok=True)
+    flac_files = os.listdir(flac_path)
+    for flac_file in flac_files:
+        # Convert FLAC file to WAV
+        id = Path(flac_file).stem
+        wav_file = os.path.join(wav_path, id + ".wav")
+        if not os.path.exists(wav_file):
+            Transformer().build(os.path.join(flac_path, flac_file), wav_file)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Aishell Data download")
+    parser.add_argument("--data_root", default='./', type=str)
+    parser.add_argument("--output_manifest_path", default='aishell_diar_manifest.json', type=str)
+    parser.add_argument("--skip_train", help="skip downloading the training dataset", action="store_true")
+    args = parser.parse_args()
+    data_root = Path(args.data_root)
+    data_root.mkdir(exist_ok=True, parents=True)
+
+    if not args.skip_train:
+        for tag in train_datasets:
+            dataset_url = train_url.format(tag)
+            dataset_path = data_root / f'{tag}/'
+            manifest_output_path = data_root / f'train_{tag}_manifest.json'
+            __process_data(
+                dataset_url=dataset_url, dataset_path=dataset_path, manifest_output_path=manifest_output_path
+            )
+    # create test dataset
+    dataset_path = data_root / f'eval/'
+    manifest_output_path = data_root / f'eval_manifest.json'
+    __process_data(dataset_url=eval_url, dataset_path=dataset_path, manifest_output_path=manifest_output_path)
+
+
+if __name__ == "__main__":
+    main()

scripts/dataset_processing/speaker_tasks/get_ami_data.py

@@ -0,0 +1,102 @@
+# 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.
+#
+# Download the AMI test dataset used to evaluate Speaker Diarization
+# More information here: https://groups.inf.ed.ac.uk/ami/corpus/
+# USAGE: python get_ami_data.py
+import argparse
+import os
+import subprocess
+
+from nemo.collections.asr.parts.utils.manifest_utils import create_manifest
+
+rttm_url = "https://raw.githubusercontent.com/BUTSpeechFIT/AMI-diarization-setup/main/only_words/rttms/{}/{}.rttm"
+uem_url = "https://raw.githubusercontent.com/BUTSpeechFIT/AMI-diarization-setup/main/uems/{}/{}.uem"
+list_url = "https://raw.githubusercontent.com/BUTSpeechFIT/AMI-diarization-setup/main/lists/{}.meetings.txt"
+
+
+audio_types = ['Mix-Headset', 'Array1-01']
+
+# these two IDs in the train set are missing download links for Array1-01.
+# We exclude them as a result.
+not_found_ids = ['IS1007d', 'IS1003b']
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Download the AMI Corpus Dataset for Speaker Diarization")
+    parser.add_argument(
+        "--test_manifest_filepath",
+        help="path to output test manifest file",
+        type=str,
+        default='AMI_test_manifest.json',
+    )
+    parser.add_argument(
+        "--dev_manifest_filepath",
+        help="path to output dev manifest file",
+        type=str,
+        default='AMI_dev_manifest.json',
+    )
+    parser.add_argument(
+        "--train_manifest_filepath",
+        help="path to output train manifest file",
+        type=str,
+        default='AMI_train_manifest.json',
+    )
+    parser.add_argument("--data_root", help="path to output data directory", type=str, default="ami_dataset")
+    args = parser.parse_args()
+
+    data_path = os.path.abspath(args.data_root)
+    os.makedirs(data_path, exist_ok=True)
+
+    for manifest_path, split in (
+        (args.test_manifest_filepath, 'test'),
+        (args.dev_manifest_filepath, 'dev'),
+        (args.train_manifest_filepath, 'train'),
+    ):
+        split_path = os.path.join(data_path, split)
+        audio_path = os.path.join(split_path, "audio")
+        os.makedirs(split_path, exist_ok=True)
+        rttm_path = os.path.join(split_path, "rttm")
+        uem_path = os.path.join(split_path, "uem")
+
+        subprocess.run(["wget", "-P", split_path, list_url.format(split)])
+        with open(os.path.join(split_path, f"{split}.meetings.txt")) as f:
+            ids = f.read().strip().split('\n')
+        for id in [file_id for file_id in ids if file_id not in not_found_ids]:
+            for audio_type in audio_types:
+                audio_type_path = os.path.join(audio_path, audio_type)
+                os.makedirs(audio_type_path, exist_ok=True)
+                audio_download = (
+                    f"https://groups.inf.ed.ac.uk/ami/AMICorpusMirror//amicorpus/{id}/audio/" f"{id}.{audio_type}.wav"
+                )
+                subprocess.run(["wget", "-P", audio_type_path, audio_download])
+            rttm_download = rttm_url.format(split, id)
+            subprocess.run(["wget", "-P", rttm_path, rttm_download])
+            uem_download = uem_url.format(split, id)
+            subprocess.run(["wget", "-P", uem_path, uem_download])
+
+        rttm_files_path = os.path.join(split_path, 'rttm_files.txt')
+        with open(rttm_files_path, 'w') as f:
+            f.write('\n'.join(os.path.join(rttm_path, p) for p in os.listdir(rttm_path)))
+        uem_files_path = os.path.join(split_path, 'uem_files.txt')
+        with open(uem_files_path, 'w') as f:
+            f.write('\n'.join(os.path.join(uem_path, p) for p in os.listdir(uem_path)))
+        for audio_type in audio_types:
+            audio_type_path = os.path.join(audio_path, audio_type)
+            audio_files_path = os.path.join(split_path, f'audio_files_{audio_type}.txt')
+            with open(audio_files_path, 'w') as f:
+                f.write('\n'.join(os.path.join(audio_type_path, p) for p in os.listdir(audio_type_path)))
+            audio_type_manifest_path = manifest_path.replace('.json', f'.{audio_type}.json')
+            create_manifest(
+                audio_files_path, audio_type_manifest_path, rttm_path=rttm_files_path, uem_path=uem_files_path
+            )

scripts/dataset_processing/speaker_tasks/get_hi-mia_data.py

@@ -0,0 +1,208 @@
+# 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.
+
+# USAGE: python get_hi-mia_data.py --data_root=<where to put data>
+
+import argparse
+import json
+import logging as _logging
+import os
+import tarfile
+import urllib.request
+from glob import glob
+
+import librosa as l
+from sklearn.model_selection import StratifiedShuffleSplit
+from tqdm import tqdm
+
+from nemo.utils.tar_utils import safe_extract
+
+parser = argparse.ArgumentParser(description="HI-MIA Data download")
+parser.add_argument("--data_root", required=True, default=None, type=str)
+parser.add_argument("--log_level", default=20, type=int)
+args = parser.parse_args()
+logging = _logging.getLogger(__name__)
+logging.addHandler(_logging.StreamHandler())
+logging.setLevel(args.log_level)
+
+URL = {
+    "dev": "http://www.openslr.org/resources/85/dev.tar.gz",
+    "test": "http://www.openslr.org/resources/85/test.tar.gz",
+    "train": "http://www.openslr.org/resources/85/train.tar.gz",
+}
+
+
+def __retrieve_with_progress(source: str, filename: str):
+    """
+    Downloads source to destination
+    Displays progress bar
+    Args:
+        source: url of resource
+        destination: local filepath
+    Returns:
+    """
+    with open(filename, "wb") as f:
+        response = urllib.request.urlopen(source)
+        total = response.length
+
+        if total is None:
+            f.write(response.content)
+        else:
+            with tqdm(total=total, unit="B", unit_scale=True, unit_divisor=1024) as pbar:
+                for data in response:
+                    f.write(data)
+                    pbar.update(len(data))
+
+
+def __maybe_download_file(destination: str, source: str):
+    """
+    Downloads source to destination if it doesn't exist.
+    If exists, skips download
+    Args:
+        destination: local filepath
+        source: url of resource
+
+    Returns:
+
+    """
+    source = URL[source]
+    if not os.path.exists(destination) and not os.path.exists(os.path.splitext(destination)[0]):
+        logging.info("{0} does not exist. Downloading ...".format(destination))
+        __retrieve_with_progress(source, filename=destination + ".tmp")
+        os.rename(destination + ".tmp", destination)
+        logging.info("Downloaded {0}.".format(destination))
+    elif os.path.exists(destination):
+        logging.info("Destination {0} exists. Skipping.".format(destination))
+    elif os.path.exists(os.path.splitext(destination)[0]):
+        logging.warning(
+            "Assuming extracted folder %s contains the extracted files from %s. Will not download.",
+            os.path.basename(destination),
+            destination,
+        )
+    return destination
+
+
+def __extract_all_files(filepath: str, data_root: str, data_dir: str):
+    if not os.path.exists(data_dir):
+        extract_file(filepath, data_root)
+        audio_dir = os.path.join(data_dir, "wav")
+        for subfolder, _, filelist in os.walk(audio_dir):
+            for ftar in filelist:
+                extract_file(os.path.join(subfolder, ftar), subfolder)
+    else:
+        logging.info("Skipping extracting. Data already there %s" % data_dir)
+
+
+def extract_file(filepath: str, data_dir: str):
+    try:
+        with tarfile.open(filepath, encoding='utf-8') as tar:
+            safe_extract(tar, data_dir)
+    except Exception:
+        logging.info("Not extracting. Maybe already there?")
+
+
+def __remove_tarred_files(filepath: str, data_dir: str):
+    if os.path.exists(data_dir) and os.path.isfile(filepath):
+        logging.info("Deleting %s" % filepath)
+        os.remove(filepath)
+
+
+def write_file(name, lines, idx):
+    with open(name, "w") as fout:
+        for i in idx:
+            dic = lines[i]
+            json.dump(dic, fout)
+            fout.write("\n")
+    logging.info("wrote %s", name)
+
+
+def __process_data(data_folder: str, data_set: str):
+    """
+    To generate manifest
+    Args:
+        data_folder: source with wav files
+    Returns:
+
+    """
+    fullpath = os.path.abspath(data_folder)
+    filelist = glob(fullpath + "/**/*.wav", recursive=True)
+    out = os.path.join(fullpath, data_set + "_all.json")
+    utt2spk = os.path.join(fullpath, "utt2spk")
+    utt2spk_file = open(utt2spk, "w")
+    id = -2  # speaker id
+
+    if os.path.exists(out):
+        logging.warning(
+            "%s already exists and is assumed to be processed. If not, please delete %s and rerun this script",
+            out,
+            out,
+        )
+        return
+
+    speakers = []
+    lines = []
+    with open(out, "w") as outfile:
+        for line in tqdm(filelist):
+            line = line.strip()
+            y, sr = l.load(line, sr=None)
+            if sr != 16000:
+                y, sr = l.load(line, sr=16000)
+                l.output.write_wav(line, y, sr)
+            dur = l.get_duration(y=y, sr=sr)
+            if data_set == "test":
+                speaker = line.split("/")[-1].split(".")[0].split("_")[0]
+            else:
+                speaker = line.split("/")[id]
+            speaker = list(speaker)
+            speaker = "".join(speaker)
+            speakers.append(speaker)
+            meta = {"audio_filepath": line, "duration": float(dur), "label": speaker}
+            lines.append(meta)
+            json.dump(meta, outfile)
+            outfile.write("\n")
+            utt2spk_file.write(line.split("/")[-1] + "\t" + speaker + "\n")
+
+    utt2spk_file.close()
+
+    if data_set != "test":
+        sss = StratifiedShuffleSplit(n_splits=1, test_size=0.1, random_state=42)
+        for train_idx, test_idx in sss.split(speakers, speakers):
+            print(len(train_idx))
+
+        out = os.path.join(fullpath, "train.json")
+        write_file(out, lines, train_idx)
+        out = os.path.join(fullpath, "dev.json")
+        write_file(out, lines, test_idx)
+
+
+def main():
+    data_root = args.data_root
+    for data_set in URL.keys():
+
+        # data_set = 'data_aishell'
+        logging.info("\n\nWorking on: {0}".format(data_set))
+        file_path = os.path.join(data_root, data_set + ".tgz")
+        logging.info("Getting {0}".format(data_set))
+        __maybe_download_file(file_path, data_set)
+        logging.info("Extracting {0}".format(data_set))
+        data_folder = os.path.join(data_root, data_set)
+        __extract_all_files(file_path, data_root, data_folder)
+        __remove_tarred_files(file_path, data_folder)
+        logging.info("Processing {0}".format(data_set))
+        __process_data(data_folder, data_set)
+        logging.info("Done!")
+
+
+if __name__ == "__main__":
+    main()

scripts/dataset_processing/speaker_tasks/get_voxconverse.py

@@ -0,0 +1,95 @@
+# 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.
+
+# downloads the training/eval set for VoxConverse.
+
+import argparse
+import logging
+import os
+import urllib.request
+import zipfile
+from pathlib import Path
+
+from nemo.collections.asr.parts.utils.manifest_utils import create_manifest
+
+dev_url = "https://www.robots.ox.ac.uk/~vgg/data/voxconverse/data/voxconverse_dev_wav.zip"
+test_url = "https://www.robots.ox.ac.uk/~vgg/data/voxconverse/data/voxconverse_test_wav.zip"
+rttm_annotations_url = "https://github.com/joonson/voxconverse/archive/refs/heads/master.zip"
+
+
+def extract_file(filepath: Path, data_dir: Path):
+    try:
+        with zipfile.ZipFile(str(filepath), 'r') as zip_ref:
+            zip_ref.extractall(str(data_dir))
+    except Exception:
+        logging.info("Not extracting. Maybe already there?")
+
+
+def download_file(url: str, destination: Path) -> Path:
+    urllib.request.urlretrieve(url, filename=str(destination))
+    return destination
+
+
+def _generate_manifest(data_root: Path, audio_path: Path, rttm_path: Path, manifest_output_path: Path):
+    audio_list = str(data_root / 'audio_file.txt')
+    rttm_list = str(data_root / 'rttm_file.txt')
+    with open(audio_list, 'w') as f:
+        f.write('\n'.join([str(os.path.join(rttm_path, x)) for x in os.listdir(audio_path)]))
+    with open(rttm_list, 'w') as f:
+        f.write('\n'.join([str(os.path.join(rttm_path, x)) for x in os.listdir(rttm_path)]))
+    create_manifest(
+        audio_list,
+        str(manifest_output_path),
+        rttm_path=rttm_list,
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser(description="VoxConverse Data download")
+    parser.add_argument("--data_root", default='./', type=str)
+    args = parser.parse_args()
+    data_root = Path(args.data_root)
+    data_root.mkdir(exist_ok=True, parents=True)
+
+    test_path = data_root / os.path.basename(test_url)
+    dev_path = data_root / os.path.basename(dev_url)
+    rttm_path = data_root / os.path.basename(rttm_annotations_url)
+
+    if not os.path.exists(test_path):
+        test_path = download_file(test_url, test_path)
+    if not os.path.exists(dev_path):
+        dev_path = download_file(dev_url, dev_path)
+    if not os.path.exists(rttm_path):
+        rttm_path = download_file(rttm_annotations_url, rttm_path)
+
+    extract_file(test_path, data_root / 'test/')
+    extract_file(dev_path, data_root / 'dev/')
+    extract_file(rttm_path, data_root)
+
+    _generate_manifest(
+        data_root=data_root,
+        audio_path=os.path.abspath(data_root / 'test/voxconverse_test_wav/'),
+        rttm_path=os.path.abspath(data_root / 'voxconverse-master/test/'),
+        manifest_output_path=data_root / 'test_manifest.json',
+    )
+    _generate_manifest(
+        data_root=data_root,
+        audio_path=os.path.abspath(data_root / 'dev/audio/'),
+        rttm_path=os.path.abspath(data_root / 'voxconverse-master/dev/'),
+        manifest_output_path=data_root / 'dev_manifest.json',
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/dataset_processing/spoken_wikipedia/preprocess.py

@@ -0,0 +1,204 @@
+# Copyright (c) 2022, 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 can be used to preprocess Spoken Wikipedia corpus before running ctc-segmentation.
+The input folder consists of subfolders with following stricture
+  ├── <Name of Wikipedia article>
+  │   ├── aligned.swc
+  │   ├── audiometa.txt
+  │   ├── audio.ogg
+  │   ├── info.json
+  │   ├── wiki.html
+  │   ├── wiki.txt
+  │   └── wiki.xml
+
+
+## The destination folder will contain look enumerated .ogg and .txt files like this:
+  ├── audio
+  |   ├── 1.ogg
+  |   ├── 2.ogg
+  |   ...
+  └── text
+      ├── 1.txt     
+      ├── 2.txt     
+      ...
+"""
+
+import argparse
+import os
+import re
+import shutil
+import subprocess
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--input_folder", required=True, type=str, help="Input folder in which each subfolder contains an article"
+)
+parser.add_argument(
+    "--destination_folder", required=True, type=str, help="Destination folder with audio and text subfolder"
+)
+args = parser.parse_args()
+
+
+def replace_diacritics(text):
+    text = re.sub(r"[éèëēêęěė]", "e", text)
+    text = re.sub(r"[ãâāáäăâàąåạả]", "a", text)
+    text = re.sub(r"[úūüùưûů]", "u", text)
+    text = re.sub(r"[ôōóöõòő]", "o", text)
+    text = re.sub(r"[ćçč]", "c", text)
+    text = re.sub(r"[ïīíîıì]", "i", text)
+    text = re.sub(r"[ñńňņ]", "n", text)
+    text = re.sub(r"[țť]", "t", text)
+    text = re.sub(r"[łľ]", "l", text)
+    text = re.sub(r"[żžź]", "z", text)
+    text = re.sub(r"[ğ]", "g", text)
+    text = re.sub(r"[ř]", "r", text)
+    text = re.sub(r"[ý]", "y", text)
+    text = re.sub(r"[æ]", "ae", text)
+    text = re.sub(r"[œ]", "oe", text)
+    text = re.sub(r"[șşšś]", "s", text)
+    return text
+
+
+def get_audio(name, n):
+    """
+    Copies .ogg file. If there are several .ogg files, concatenates them.
+
+    Args:
+        name - name of folder within Spoken Wikipedia
+        n - integer that will serve as output file name, e.g. if n=1, file 1.ogg will be created
+    """
+    audio_path = os.path.join(args.input_folder, name, "audio.ogg")
+    if not os.path.exists(audio_path):
+        ##  Some folders have multiple .ogg files, so we need to first combine them into one file. Example:
+        ##  |── Universe
+        ##  │   ├── aligned.swc
+        ##  │   ├── audio1.ogg
+        ##  │   ├── audio2.ogg
+        ##  │   ├── audio3.ogg
+        ##  │   ├── audio4.ogg
+        ##  │   ├── audiometa.txt
+        ##  │   ├── info.json
+        ##  │   ├── wiki.html
+        ##  │   ├── wiki.txt
+        ##  │   └── wiki.xml
+
+        multiple_ogg_files = []
+        for i in range(1, 5):
+            path = os.path.join(args.input_folder, name, "audio" + str(i) + ".ogg")
+            if os.path.exists(path):
+                multiple_ogg_files.append(path)
+            else:
+                break
+        if len(multiple_ogg_files) == 0:
+            return
+        elif len(multiple_ogg_files) == 1:
+            shutil.copy(multiple_ogg_files[0], audio_path)
+        else:
+            tmp_file_name = "ffmeg_inputs.txt"
+            print("tmp_file_name=", tmp_file_name)
+            with open(tmp_file_name, "w", encoding="utf-8") as tmp_file:
+                for path in multiple_ogg_files:
+                    tmp_file.write("file '" + path + "'\n")
+            ffmpeg_cmd = ["ffmpeg", "-f", "concat", "-i", tmp_file_name, "-c", "copy", audio_path]
+            print("ffmpeg command:", ffmpeg_cmd)
+            subprocess.run(ffmpeg_cmd, check=True)
+
+    output_audio_path = args.destination_folder + "/audio/" + str(n) + ".ogg"
+    shutil.copy(audio_path, output_audio_path)
+
+
+def get_text(name, n):
+    """
+    Cleans wiki.txt.
+
+    Args:
+        name - name of folder within Spoken Wikipedia
+        n - integer that will serve as output file name, e.g. if n=1, file 1.txt will be created
+    """
+
+    # Then we need to clean the text
+    out_text = open(args.destination_folder + "/text/" + str(n) + ".txt", "w", encoding="utf-8")
+    with open(args.input_folder + "/" + name + "/wiki.txt", "r", encoding="utf-8") as f:
+        for line in f:
+            do_break = False
+            line2 = line.strip()
+            ref_parts = line2.split("<ref")
+            for idx, s in enumerate(ref_parts):
+                if idx != 0:
+                    s = "<ref" + s
+                if s.startswith("[[Image") and s.endswith("]]"):
+                    continue
+                if s.startswith("[[File") and s.endswith("]]"):
+                    continue
+                if s.startswith(":"):
+                    continue
+                if s.startswith("{|") or s.startswith("|}") or s.startswith("|") or s.startswith("!"):
+                    continue
+                if s.startswith("{{") and (s.endswith("}}") or "}}" not in s):
+                    continue
+                if s.startswith("{{pp-move"):
+                    continue
+                s = re.sub(r"\[\[Image\:[^\]]+\]\]", r"", s)
+                s = re.sub(r"\[\[File\:[^\]]+\]\]", r"", s)
+                s = re.sub(r"\[http[^\]]+\]", r"", s)
+                s = re.sub(r"<math>[^<>]+</math>", r"", s)
+                s = re.sub(r"<!\-\-.+\-\->", r"", s)  # <!--DashBot-->    can be inside <ref>
+                s = re.sub(r"<ref>.+</ref>", r"", s)
+                s = re.sub(r"<ref .+</ref>", r"", s)
+                s = re.sub(r"<ref[^<>]+/>", r"", s)
+                s = re.sub(r"<[^ <>]+>", r"", s)  # <sub>, <sup>, </u>
+                if (
+                    re.match(r"== *Notes *==", s)
+                    or re.match(r"== *References *==", s)
+                    or re.match(r"== *External links *==", s)
+                    or re.match(r"== *See also *==", s)
+                ):
+                    do_break = True
+                    break
+                s = re.sub(r"{{convert\|(\d+)\|(\w+)\|[^}]+}}", r"\g<1> \g<2>", s)  # {{convert|7600|lb|kg}}
+                s = re.sub(r"{{cquote\|", r"", s)
+                s = re.sub(r"{{[^{}]+}}", r"", s)
+                s = s.replace("{{", "").replace("}}", "")
+                s = re.sub(r"(lang[^()]+)", r"", s)  # (lang-bn...)
+                s = re.sub(r"==+", r"", s)
+                s = re.sub(r"''+", r" ", s)  # remove multiple quotes
+                s = re.sub(r" '", r" ", s)  # remove quote at the beginning
+                s = re.sub(r"' ", r" ", s)  # remove quote at the end
+                s = re.sub(r"[…\*]", r" ", s)
+                s = re.sub(r"\\u....", r" ", s)  # remove unicode
+                s = re.sub(r"&[^ ;&]+;", r"", s)  # &nbsp; &mdash;
+
+                s = replace_diacritics(s)
+
+                s = re.sub(r"\[\[[^\]]+\|([^\]]+)\]\]", r"\g<1>", s)  # if several variants, take the last one
+                s = re.sub(r"\[\[([^\]]+)\]\]", r"\g<1>", s)
+
+                out_text.write(s + "\n")
+            if do_break:
+                break
+    out_text.close()
+
+
+if __name__ == "__main__":
+    n = 0
+    for name in os.listdir(args.input_folder):
+        n += 1
+        if not os.path.exists(args.input_folder + "/" + name + "/wiki.txt"):
+            print("wiki.txt does not exist in " + name)
+            continue
+        get_audio(name, n)
+        get_text(name, n)

scripts/dataset_processing/spoken_wikipedia/run.sh

@@ -0,0 +1,135 @@
+#!/bin/bash
+
+# Copyright (c) 2022, 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.
+
+## Download the Spoken Wikipedia corpus for English
+## Note, that there are some other languages available
+## @InProceedings{KHN16.518,
+##  author = {Arne K{\"o}hn and Florian Stegen and Timo Baumann},
+##  title = {Mining the Spoken Wikipedia for Speech Data and Beyond},
+##  booktitle = {Proceedings of the Tenth International Conference on Language Resources and Evaluation (LREC 2016)},
+##  year = {2016},
+##  month = {may},
+##  date = {23-28},
+##  location = {Portorož, Slovenia},
+##  editor = {Nicoletta Calzolari (Conference Chair) and Khalid Choukri and Thierry Declerck and Marko Grobelnik and Bente Maegaard and Joseph Mariani and Asuncion Moreno and Jan Odijk and Stelios Piperidis},
+##  publisher = {European Language Resources Association (ELRA)},
+##  address = {Paris, France},
+##  isbn = {978-2-9517408-9-1},
+##  islrn = {684-927-624-257-3/},
+##  language = {english}
+## }
+
+wget https://corpora.uni-hamburg.de/hzsk/de/islandora/object/file:swc-2.0_en-with-audio/datastream/TAR/en-with-audio.tar .
+tar -xvf en-with-audio.tar
+
+##  We get a folder English with 1339 subfolders, each subfolder corresponds to a Wikipedia article. Example:
+##  ├── Universal_suffrage
+##  │   ├── aligned.swc
+##  │   ├── audiometa.txt
+##  │   ├── audio.ogg
+##  │   ├── info.json
+##  │   ├── wiki.html
+##  │   ├── wiki.txt
+##  │   └── wiki.xml
+
+##  We will use two files: audio.ogg and wiki.txt
+
+## Some folders have multiple .ogg files, this will be handled during preprocess.py. Example:
+##  |── Universe
+##  │   ├── aligned.swc
+##  │   ├── audio1.ogg
+##  │   ├── audio2.ogg
+##  │   ├── audio3.ogg
+##  │   ├── audio4.ogg
+##  │   ├── audiometa.txt
+##  │   ├── info.json
+##  │   ├── wiki.html
+##  │   ├── wiki.txt
+##  │   └── wiki.xml
+
+## Some rare folders are incomplete, these will be skipped during preprocessing.
+
+## Rename some folders with special symbols because they cause problems to ffmpeg when concatening multiple .ogg files
+mv "english/The_Hitchhiker%27s_Guide_to_the_Galaxy" "english/The_Hitchhikers_guide_to_the_Galaxy"
+mv "english/SummerSlam_(2003)" "english/SummerSlam_2003"
+mv "english/Over_the_Edge_(1999)" "english/Over_the_Edge_1999"
+mv "english/Lost_(TV_series)" "english/Lost_TV_series"
+mv "english/S._A._Andr%c3%a9e%27s_Arctic_Balloon_Expedition_of_1897" "english/S_A_Andres_Arctic_Balloon_Expedition_of_1897"
+
+## path to NeMo repository, e.g. /home/user/NeMo
+NEMO_PATH=
+
+INPUT_DIR="english"
+OUTPUT_DIR=${INPUT_DIR}_result
+
+rm -rf $OUTPUT_DIR
+rm -rf ${INPUT_DIR}_prepared
+mkdir ${INPUT_DIR}_prepared
+mkdir ${INPUT_DIR}_prepared/audio
+mkdir ${INPUT_DIR}_prepared/text
+python ${NEMO_PATH}/scripts/dataset_processing/spoken_wikipedia/preprocess.py --input_folder ${INPUT_DIR} --destination_folder ${INPUT_DIR}_prepared
+
+## Now we have ${INPUT_DIR}_prepared folder with the following structure:
+##  ├── audio
+##  |   ├── 1.ogg
+##  |   ├── 2.ogg
+##  |   ...
+##  └── text
+##      ├── 1.txt     
+##      ├── 2.txt     
+##      ...
+
+MODEL_FOR_SEGMENTATION="stt_en_fastconformer_ctc_large"
+MODEL_FOR_RECOGNITION="stt_en_conformer_ctc_large"
+## We set this threshold as very permissive, later we will use other metrics for filtering
+THRESHOLD=-10
+
+${NEMO_PATH}/tools/ctc_segmentation/run_segmentation.sh \
+--SCRIPTS_DIR=${NEMO_PATH}/tools/ctc_segmentation/scripts \
+--MODEL_NAME_OR_PATH=${MODEL_FOR_SEGMENTATION} \
+--DATA_DIR=${INPUT_DIR}_prepared \
+--OUTPUT_DIR=${OUTPUT_DIR} \
+--MIN_SCORE=${THRESHOLD}
+
+# Thresholds for filtering
+CER_THRESHOLD=20
+WER_THRESHOLD=30
+CER_EDGE_THRESHOLD=30
+LEN_DIFF_RATIO_THRESHOLD=0.15
+EDGE_LEN=25
+BATCH_SIZE=1
+
+${NEMO_PATH}/tools/ctc_segmentation/run_filter.sh \
+--SCRIPTS_DIR=${NEMO_PATH}/tools/ctc_segmentation/scripts \
+--MODEL_NAME_OR_PATH=${MODEL_FOR_RECOGNITION} \
+--BATCH_SIZE=${BATCH_SIZE} \
+--MANIFEST=$OUTPUT_DIR/manifests/manifest.json \
+--INPUT_AUDIO_DIR=${INPUT_DIR}_prepared/audio/ \
+--EDGE_LEN=${EDGE_LEN} \
+--CER_THRESHOLD=${CER_THRESHOLD} \
+--WER_THRESHOLD=${WER_THRESHOLD} \
+--CER_EDGE_THRESHOLD=${CER_EDGE_THRESHOLD} \
+--LEN_DIFF_RATIO_THRESHOLD=${LEN_DIFF_RATIO_THRESHOLD}
+
+python ${NEMO_PATH}/examples/asr/speech_to_text_eval.py \
+  dataset_manifest=${OUTPUT_DIR}/manifests/manifest_transcribed_metrics_filtered.json \
+  use_cer=True \
+  only_score_manifest=True
+
+python ${NEMO_PATH}/examples/asr/speech_to_text_eval.py \
+  dataset_manifest=${OUTPUT_DIR}/manifests/manifest_transcribed_metrics_filtered.json \
+  use_cer=False \
+  only_score_manifest=True

scripts/dataset_processing/tts/aishell3/ds_conf/ds_for_fastpitch_align.yaml

@@ -0,0 +1,49 @@
+name: "ds_for_fastpitch_align"
+
+manifest_filepath: "train_manifest.json"
+sup_data_path: "sup_data"
+sup_data_types: [ "align_prior_matrix", "pitch", "speaker_id"]
+phoneme_dict_path: "scripts/tts_dataset_files/zh/24finals/pinyin_dict_nv_22.10.txt"
+
+dataset:
+  _target_: nemo.collections.tts.data.dataset.TTSDataset
+  manifest_filepath: ${manifest_filepath}
+  sample_rate: 22050
+  sup_data_path: ${sup_data_path}
+  sup_data_types: ${sup_data_types}
+  n_fft: 1024
+  win_length: 1024
+  hop_length: 256
+  window: "hann"
+  n_mels: 80
+  lowfreq: 0
+  highfreq: null
+  max_duration: null
+  min_duration: 0.1
+  ignore_file: null
+  trim: true
+  trim_top_db: 50
+  trim_frame_length: 1024
+  trim_hop_length: 256
+  pitch_fmin: 65.40639132514966
+  pitch_fmax: 2093.004522404789
+
+  text_normalizer:
+    _target_: nemo_text_processing.text_normalization.normalize.Normalizer
+    lang: zh
+    input_case: cased
+
+  text_normalizer_call_kwargs:
+    verbose: false
+    punct_pre_process: true
+    punct_post_process: true
+
+  text_tokenizer:
+    _target_: nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.ChinesePhonemesTokenizer
+    punct: true
+    apostrophe: true
+    pad_with_space: true
+    g2p:
+      _target_: nemo.collections.tts.g2p.models.zh_cn_pinyin.ChineseG2p
+      phoneme_dict: ${phoneme_dict_path}
+      word_segmenter: jieba # Only jieba is supported now.

scripts/dataset_processing/tts/aishell3/get_data.py

@@ -0,0 +1,176 @@
+# 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.
+
+# Disclaimer:
+#   Each user is responsible for checking the content of datasets and the applicable licenses and determining if suitable for the intended use.
+
+import argparse
+import json
+import os
+import random
+import subprocess
+import tarfile
+import urllib.request
+from pathlib import Path
+
+import numpy as np
+from nemo_text_processing.text_normalization.normalize import Normalizer
+from opencc import OpenCC
+
+from nemo.utils.tar_utils import safe_extract
+
+URL = "https://www.openslr.org/resources/93/data_aishell3.tgz"
+
+
+def get_args():
+    parser = argparse.ArgumentParser(
+        description='Prepare SF_bilingual dataset and create manifests with predefined split'
+    )
+
+    parser.add_argument(
+        "--data-root",
+        type=Path,
+        help="where the dataset will reside",
+        default="./DataChinese/sf_bilingual_speech_zh_en_vv1/SF_bilingual/",
+    )
+    parser.add_argument(
+        "--manifests-path", type=Path, help="where the resulting manifests files will reside", default="./"
+    )
+    parser.add_argument("--val-size", default=0.01, type=float, help="eval set split")
+    parser.add_argument("--test-size", default=0.01, type=float, help="test set split")
+    parser.add_argument(
+        "--seed-for-ds-split",
+        default=100,
+        type=float,
+        help="Seed for deterministic split of train/dev/test, NVIDIA's default is 100",
+    )
+
+    args = parser.parse_args()
+    return args
+
+
+def __maybe_download_file(source_url, destination_path):
+    if not destination_path.exists():
+        tmp_file_path = destination_path.with_suffix('.tmp')
+        urllib.request.urlretrieve(source_url, filename=str(tmp_file_path))
+        tmp_file_path.rename(destination_path)
+
+
+def __extract_file(filepath, data_dir):
+    try:
+        with tarfile.open(filepath) as tar:
+            safe_extract(tar, str(data_dir))
+    except Exception:
+        print(f"Error while extracting {filepath}. Already extracted?")
+
+
+def __process_transcript(file_path: str):
+    # Create directory for processed wav files
+    Path(file_path / "processed").mkdir(parents=True, exist_ok=True)
+    # Create zh-TW to zh-simplify converter
+    cc = OpenCC('t2s')
+    # Create normalizer
+    text_normalizer = Normalizer(
+        lang="zh",
+        input_case="cased",
+        overwrite_cache=True,
+        cache_dir=str(file_path / "cache_dir"),
+    )
+    text_normalizer_call_kwargs = {"punct_pre_process": True, "punct_post_process": True}
+    normalizer_call = lambda x: text_normalizer.normalize(x, **text_normalizer_call_kwargs)
+    entries = []
+    SPEAKER_LEN = 7
+
+    candidates = []
+    speakers = set()
+    with open(file_path / "train" / "content.txt", encoding="utf-8") as fin:
+        for line in fin:
+            content = line.split()
+            wav_name, text = content[0], "".join(content[1::2]) + "。"
+            wav_name = wav_name.replace(u'\ufeff', '')
+            speaker = wav_name[:SPEAKER_LEN]
+            speakers.add(speaker)
+            wav_file = file_path / "train" / "wav" / speaker / wav_name
+            assert os.path.exists(wav_file), f"{wav_file} not found!"
+            duration = subprocess.check_output(["soxi", "-D", str(wav_file)])
+            if float(duration) <= 3.0:  # filter out wav files shorter than 3 seconds
+                continue
+            processed_file = file_path / "processed" / wav_name
+            # convert wav to mono 22050HZ, 16 bit (as SFSpeech dataset)
+            subprocess.run(["sox", str(wav_file), "-r", "22050", "-c", "1", "-b", "16", str(processed_file)])
+            candidates.append((processed_file, duration, text, speaker))
+
+    # remapping the speakder to speaker_id (start from 1)
+    remapping = {}
+    for index, speaker in enumerate(sorted(speakers)):
+        remapping[speaker] = index + 1
+
+    for processed_file, duration, text, speaker in candidates:
+        simplified_text = cc.convert(text)
+        normalized_text = normalizer_call(simplified_text)
+        entry = {
+            'audio_filepath': os.path.abspath(processed_file),
+            'duration': float(duration),
+            'text': text,
+            'normalized_text': normalized_text,
+            'speaker_raw': speaker,
+            'speaker': remapping[speaker],
+        }
+
+        entries.append(entry)
+
+    return entries
+
+
+def __process_data(dataset_path, val_size, test_size, seed_for_ds_split, manifests_dir):
+    entries = __process_transcript(dataset_path)
+
+    random.Random(seed_for_ds_split).shuffle(entries)
+
+    train_size = 1.0 - val_size - test_size
+    train_entries, validate_entries, test_entries = np.split(
+        entries, [int(len(entries) * train_size), int(len(entries) * (train_size + val_size))]
+    )
+
+    assert len(train_entries) > 0, "Not enough data for train, val and test"
+
+    def save(p, data):
+        with open(p, 'w') as f:
+            for d in data:
+                f.write(json.dumps(d) + '\n')
+
+    save(manifests_dir / "train_manifest.json", train_entries)
+    save(manifests_dir / "val_manifest.json", validate_entries)
+    save(manifests_dir / "test_manifest.json", test_entries)
+
+
+def main():
+    args = get_args()
+
+    tarred_data_path = args.data_root / "data_aishell3.tgz"
+
+    __maybe_download_file(URL, tarred_data_path)
+    __extract_file(str(tarred_data_path), str(args.data_root))
+
+    __process_data(
+        args.data_root,
+        args.val_size,
+        args.test_size,
+        args.seed_for_ds_split,
+        args.manifests_path,
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/dataset_processing/tts/compute_feature_stats.py

@@ -0,0 +1,230 @@
+# 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 is to compute global and speaker-level feature statistics for a given TTS training manifest.
+
+This script should be run after compute_features.py as it loads the precomputed feature data.
+
+$ python <nemo_root_path>/scripts/dataset_processing/tts/compute_feature_stats.py \
+    --feature_config_path=<nemo_root_path>/examples/tts/conf/features/feature_22050.yaml
+    --manifest_path=<data_root_path>/manifest1.json \
+    --manifest_path=<data_root_path>/manifest2.json \
+    --audio_dir=<data_root_path>/audio1 \
+    --audio_dir=<data_root_path>/audio2 \
+    --feature_dir=<data_root_path>/features1 \
+    --feature_dir=<data_root_path>/features2 \
+    --stats_path=<data_root_path>/feature_stats.json
+
+The output dictionary will contain the feature statistics for every speaker, as well as a "default" entry
+with the global statistics.
+
+For example:
+
+{
+    "default": {
+        "pitch_mean": 100.0,
+        "pitch_std": 50.0,
+        "energy_mean": 7.5,
+        "energy_std": 4.5
+    },
+    "speaker1": {
+        "pitch_mean": 105.0,
+        "pitch_std": 45.0,
+        "energy_mean": 7.0,
+        "energy_std": 5.0
+    },
+    "speaker2": {
+        "pitch_mean": 110.0,
+        "pitch_std": 30.0,
+        "energy_mean": 5.0,
+        "energy_std": 2.5
+    }
+}
+
+"""
+
+import argparse
+import json
+from collections import defaultdict
+from pathlib import Path
+from typing import List, Tuple
+
+import torch
+from omegaconf import OmegaConf
+from tqdm import tqdm
+
+from nemo.collections.asr.parts.utils.manifest_utils import read_manifest
+from nemo.core.classes.common import safe_instantiate
+
+
+def get_args():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        description="Compute TTS feature statistics.",
+    )
+    parser.add_argument(
+        "--feature_config_path",
+        required=True,
+        type=Path,
+        help="Path to feature config file.",
+    )
+    parser.add_argument(
+        "--manifest_path",
+        required=True,
+        type=Path,
+        action="append",
+        help="Path(s) to training manifest.",
+    )
+    parser.add_argument(
+        "--audio_dir",
+        required=True,
+        type=Path,
+        action="append",
+        help="Path(s) to base directory with audio data.",
+    )
+    parser.add_argument(
+        "--feature_dir",
+        required=True,
+        type=Path,
+        action="append",
+        help="Path(s) to directory where feature data was stored.",
+    )
+    parser.add_argument(
+        "--feature_names",
+        default="pitch,energy",
+        type=str,
+        help="Comma separated list of features to process.",
+    )
+    parser.add_argument(
+        "--mask_field",
+        default="voiced_mask",
+        type=str,
+        help="If provided, stat computation will ignore non-masked frames.",
+    )
+    parser.add_argument(
+        "--stats_path",
+        default=Path("feature_stats.json"),
+        type=Path,
+        help="Path to output JSON file with dataset feature statistics.",
+    )
+    parser.add_argument(
+        "--overwrite",
+        action=argparse.BooleanOptionalAction,
+        help="Whether to overwrite the output stats file if it exists.",
+    )
+
+    args = parser.parse_args()
+    return args
+
+
+def _compute_stats(values: List[torch.Tensor]) -> Tuple[float, float]:
+    values_tensor = torch.cat(values, dim=0)
+    mean = values_tensor.mean().item()
+    std = values_tensor.std(dim=0).item()
+    return mean, std
+
+
+def main():
+    args = get_args()
+
+    feature_config_path = args.feature_config_path
+    manifest_paths = args.manifest_path
+    audio_dirs = args.audio_dir
+    feature_dirs = args.feature_dir
+    feature_name_str = args.feature_names
+    mask_field = args.mask_field
+    stats_path = args.stats_path
+    overwrite = args.overwrite
+
+    if not (len(manifest_paths) == len(audio_dirs) == len(feature_dirs)):
+        raise ValueError(
+            f"Need same number of manifest, audio_dir, and feature_dir. Received: "
+            f"{len(manifest_paths)}, "
+            f"{len(audio_dirs)}, "
+            f"{len(feature_dirs)}"
+        )
+
+    for manifest_path, audio_dir, feature_dir in zip(manifest_paths, audio_dirs, feature_dirs):
+        if not manifest_path.exists():
+            raise ValueError(f"Manifest {manifest_path} does not exist.")
+
+        if not audio_dir.exists():
+            raise ValueError(f"Audio directory {audio_dir} does not exist.")
+
+        if not feature_dir.exists():
+            raise ValueError(
+                f"Feature directory {feature_dir} does not exist. "
+                f"Please check that the path is correct and that you ran compute_features.py"
+            )
+
+    if stats_path.exists():
+        if overwrite:
+            print(f"Will overwrite existing stats path: {stats_path}")
+        else:
+            raise ValueError(f"Stats path already exists: {stats_path}")
+
+    feature_config = OmegaConf.load(feature_config_path)
+    feature_config = safe_instantiate(feature_config)
+    featurizer_dict = feature_config.featurizers
+
+    print(f"Found featurizers for {list(featurizer_dict.keys())}.")
+    featurizers = featurizer_dict.values()
+
+    feature_names = feature_name_str.split(",")
+    # For each feature, we have a dictionary mapping speaker IDs to a list containing all features
+    # for that speaker
+    feature_stats = {name: defaultdict(list) for name in feature_names}
+
+    for manifest_path, audio_dir, feature_dir in zip(manifest_paths, audio_dirs, feature_dirs):
+        entries = read_manifest(manifest_path)
+
+        for entry in tqdm(entries):
+            speaker = entry["speaker"]
+
+            entry_dict = {}
+            for featurizer in featurizers:
+                feature_dict = featurizer.load(manifest_entry=entry, audio_dir=audio_dir, feature_dir=feature_dir)
+                entry_dict.update(feature_dict)
+
+            if mask_field:
+                mask = entry_dict[mask_field]
+            else:
+                mask = None
+
+            for feature_name in feature_names:
+                values = entry_dict[feature_name]
+                if mask is not None:
+                    values = values[mask]
+
+                feature_stat_dict = feature_stats[feature_name]
+                feature_stat_dict["default"].append(values)
+                feature_stat_dict[speaker].append(values)
+
+    stat_dict = defaultdict(dict)
+    for feature_name in feature_names:
+        mean_key = f"{feature_name}_mean"
+        std_key = f"{feature_name}_std"
+        feature_stat_dict = feature_stats[feature_name]
+        for speaker_id, values in feature_stat_dict.items():
+            speaker_mean, speaker_std = _compute_stats(values)
+            stat_dict[speaker_id][mean_key] = speaker_mean
+            stat_dict[speaker_id][std_key] = speaker_std
+
+    with open(stats_path, 'w', encoding="utf-8") as stats_f:
+        json.dump(stat_dict, stats_f, indent=4)
+
+
+if __name__ == "__main__":
+    main()

scripts/dataset_processing/tts/compute_features.py

@@ -0,0 +1,130 @@
+# 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 computes features for TTS models prior to training, such as pitch and energy.
+The resulting features will be stored in the provided 'feature_dir'.
+
+$ python <nemo_root_path>/scripts/dataset_processing/tts/compute_features.py \
+    --feature_config_path=<nemo_root_path>/examples/tts/conf/features/feature_22050.yaml \
+    --manifest_path=<data_root_path>/manifest.json \
+    --audio_dir=<data_root_path>/audio \
+    --feature_dir=<data_root_path>/features \
+    --overwrite \
+    --num_workers=1
+"""
+
+import argparse
+from pathlib import Path
+
+from joblib import Parallel, delayed
+from omegaconf import OmegaConf
+from tqdm import tqdm
+
+from nemo.collections.asr.parts.utils.manifest_utils import read_manifest
+from nemo.core.classes.common import safe_instantiate
+
+
+def get_args():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        description="Compute TTS features.",
+    )
+    parser.add_argument(
+        "--feature_config_path",
+        required=True,
+        type=Path,
+        help="Path to feature config file.",
+    )
+    parser.add_argument(
+        "--manifest_path",
+        required=True,
+        type=Path,
+        help="Path to training manifest.",
+    )
+    parser.add_argument(
+        "--audio_dir",
+        required=True,
+        type=Path,
+        help="Path to base directory with audio data.",
+    )
+    parser.add_argument(
+        "--feature_dir",
+        required=True,
+        type=Path,
+        help="Path to directory where feature data will be stored.",
+    )
+    parser.add_argument(
+        "--dedupe_files",
+        action=argparse.BooleanOptionalAction,
+        help="If given, will only process the first manifest entry found for each audio file.",
+    )
+    parser.add_argument(
+        "--overwrite",
+        action=argparse.BooleanOptionalAction,
+        help="Whether to overwrite existing feature files.",
+    )
+    parser.add_argument(
+        "--num_workers", default=1, type=int, help="Number of parallel threads to use. If -1 all CPUs are used."
+    )
+
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+    feature_config_path = args.feature_config_path
+    manifest_path = args.manifest_path
+    audio_dir = args.audio_dir
+    feature_dir = args.feature_dir
+    dedupe_files = args.dedupe_files
+    overwrite = args.overwrite
+    num_workers = args.num_workers
+
+    if not manifest_path.exists():
+        raise ValueError(f"Manifest {manifest_path} does not exist.")
+
+    if not audio_dir.exists():
+        raise ValueError(f"Audio directory {audio_dir} does not exist.")
+
+    feature_config = OmegaConf.load(feature_config_path)
+    feature_config = safe_instantiate(feature_config)
+    featurizers = feature_config.featurizers
+
+    entries = read_manifest(manifest_path)
+
+    if dedupe_files:
+        final_entries = []
+        audio_filepath_set = set()
+        for entry in entries:
+            audio_filepath = entry["audio_filepath"]
+            if audio_filepath in audio_filepath_set:
+                continue
+            final_entries.append(entry)
+            audio_filepath_set.add(audio_filepath)
+        entries = final_entries
+
+    for feature_name, featurizer in featurizers.items():
+        print(f"Computing: {feature_name}")
+        Parallel(n_jobs=num_workers)(
+            delayed(featurizer.save)(
+                manifest_entry=entry, audio_dir=audio_dir, feature_dir=feature_dir, overwrite=overwrite
+            )
+            for entry in tqdm(entries)
+        )
+
+
+if __name__ == "__main__":
+    main()

scripts/dataset_processing/tts/compute_speaker_stats.py

@@ -0,0 +1,144 @@
+# Copyright (c) 2022, 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 is to compute speaker-level statistics, such as pitch mean & standard deviation, for a given
+TTS training manifest.
+
+This script should be run after extract_sup_data.py as it uses the precomputed supplemental features.
+
+$ python <nemo_root_path>/scripts/dataset_processing/tts/compute_speaker_stats.py \
+    --manifest_path=<data_root_path>/fastpitch_manifest.json \
+    --sup_data_path=<data_root_path>/sup_data \
+    --pitch_stats_path=<data_root_path>/pitch_stats.json
+"""
+
+import argparse
+import json
+import os
+from collections import defaultdict
+from pathlib import Path
+from typing import List, Tuple
+
+import torch
+from tqdm import tqdm
+
+from nemo.collections.asr.parts.utils.manifest_utils import read_manifest
+from nemo.collections.tts.parts.utils.tts_dataset_utils import get_base_dir
+from nemo.collections.tts.torch.tts_data_types import Pitch
+from nemo.utils import logging
+
+
+def get_args():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        description="Compute speaker level pitch statistics.",
+    )
+    parser.add_argument(
+        "--manifest_path",
+        required=True,
+        type=Path,
+        help="Path to training manifest.",
+    )
+    parser.add_argument(
+        "--sup_data_path",
+        default=Path("sup_data"),
+        type=Path,
+        help="Path to base directory with supplementary data.",
+    )
+    parser.add_argument(
+        "--pitch_stats_path",
+        default=Path("pitch_stats.json"),
+        type=Path,
+        help="Path to output JSON file with speaker pitch statistics.",
+    )
+    args = parser.parse_args()
+    return args
+
+
+def _compute_stats(values: List[torch.Tensor]) -> Tuple[float, float]:
+    values_tensor = torch.cat(values, dim=0)
+    mean = values_tensor.mean().item()
+    std = values_tensor.std(dim=0).item()
+    return mean, std
+
+
+def _get_sup_data_filepath(manifest_entry: dict, audio_dir: Path, sup_data_dir: Path) -> Path:
+    """
+    Get the absolute path of a supplementary data type for the input manifest entry.
+
+    Example: audio_filepath "<audio_dir>/speaker1/audio1.wav" becomes "<sup_data_dir>/speaker1_audio1.pt"
+
+    Args:
+        manifest_entry: Manifest entry dictionary.
+        audio_dir: base directory where audio is stored.
+        sup_data_dir: base directory where supplementary data is stored.
+
+    Returns:
+        Path to the supplementary data file.
+    """
+    audio_path = Path(manifest_entry["audio_filepath"])
+    rel_audio_path = audio_path.relative_to(audio_dir)
+    rel_sup_data_path = rel_audio_path.with_suffix(".pt")
+    sup_data_filename = str(rel_sup_data_path).replace(os.sep, "_")
+    sup_data_filepath = sup_data_dir / sup_data_filename
+    return sup_data_filepath
+
+
+def main():
+    args = get_args()
+    manifest_path = args.manifest_path
+    sup_data_path = args.sup_data_path
+    pitch_stats_path = args.pitch_stats_path
+
+    pitch_data_path = Path(os.path.join(sup_data_path, Pitch.name))
+    if not os.path.exists(pitch_data_path):
+        raise ValueError(
+            f"Pitch directory {pitch_data_path} does not exist. Make sure 'sup_data_path' is correct "
+            f"and that you have computed the pitch using extract_sup_data.py"
+        )
+
+    entries = read_manifest(manifest_path)
+
+    audio_paths = [entry["audio_filepath"] for entry in entries]
+    base_dir = get_base_dir(audio_paths)
+
+    global_pitch_values = []
+    speaker_pitch_values = defaultdict(list)
+    for entry in tqdm(entries):
+        pitch_path = _get_sup_data_filepath(manifest_entry=entry, audio_dir=base_dir, sup_data_dir=pitch_data_path)
+        if not os.path.exists(pitch_path):
+            logging.warning(f"Unable to find pitch file for {entry}")
+            continue
+
+        pitch = torch.load(pitch_path)
+        # Filter out non-speech frames
+        pitch = pitch[pitch != 0]
+        global_pitch_values.append(pitch)
+        if "speaker" in entry:
+            speaker_id = entry["speaker"]
+            speaker_pitch_values[speaker_id].append(pitch)
+
+    global_pitch_mean, global_pitch_std = _compute_stats(global_pitch_values)
+    pitch_stats = {"default": {"pitch_mean": global_pitch_mean, "pitch_std": global_pitch_std}}
+    for speaker_id, pitch_values in speaker_pitch_values.items():
+        pitch_mean, pitch_std = _compute_stats(pitch_values)
+        pitch_stats[speaker_id] = {"pitch_mean": pitch_mean, "pitch_std": pitch_std}
+
+    with open(pitch_stats_path, 'w', encoding="utf-8") as stats_f:
+        json.dump(pitch_stats, stats_f, indent=4)
+
+
+if __name__ == "__main__":
+    main()

scripts/dataset_processing/tts/create_speaker_map.py

@@ -0,0 +1,104 @@
+# 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 takes a list of TTS manifests and creates a JSON mapping the input speaker names to
+unique indices for multi-speaker TTS training.
+
+To ensure that speaker names are unique across datasets, it is recommended that you prepend the speaker
+names in your manifest with the name of the dataset.
+
+$ python <nemo_root_path>/scripts/dataset_processing/tts/create_speaker_map.py \
+    --manifest_path=manifest1.json \
+    --manifest_path=manifest2.json \
+    --speaker_map_path=speakers.json
+
+Example output:
+
+{
+    "vctk_p225": 0,
+    "vctk_p226": 1,
+    "vctk_p227": 2,
+    ...
+}
+
+"""
+
+import argparse
+import json
+from pathlib import Path
+
+from nemo.collections.asr.parts.utils.manifest_utils import read_manifest
+
+
+def get_args():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        description="Create mapping from speaker names to numerical speaker indices.",
+    )
+    parser.add_argument(
+        "--manifest_path",
+        required=True,
+        type=Path,
+        action="append",
+        help="Path to training manifest(s).",
+    )
+    parser.add_argument(
+        "--speaker_map_path",
+        required=True,
+        type=Path,
+        help="Path for output speaker index JSON",
+    )
+    parser.add_argument(
+        "--overwrite",
+        action=argparse.BooleanOptionalAction,
+        help="Whether to overwrite the output speaker file if it exists.",
+    )
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+    manifest_paths = args.manifest_path
+    speaker_map_path = args.speaker_map_path
+    overwrite = args.overwrite
+
+    for manifest_path in manifest_paths:
+        if not manifest_path.exists():
+            raise ValueError(f"Manifest {manifest_path} does not exist.")
+
+    if speaker_map_path.exists():
+        if overwrite:
+            print(f"Will overwrite existing speaker path: {speaker_map_path}")
+        else:
+            raise ValueError(f"Speaker path already exists: {speaker_map_path}")
+
+    speaker_set = set()
+    for manifest_path in manifest_paths:
+        entries = read_manifest(manifest_path)
+        for entry in entries:
+            speaker = str(entry["speaker"])
+            speaker_set.add(speaker)
+
+    speaker_list = list(speaker_set)
+    speaker_list.sort()
+    speaker_index_map = {speaker_list[i]: i for i in range(len(speaker_list))}
+
+    with open(speaker_map_path, 'w', encoding="utf-8") as stats_f:
+        json.dump(speaker_index_map, stats_f, indent=4)
+
+
+if __name__ == "__main__":
+    main()

scripts/dataset_processing/tts/extract_sup_data.py

@@ -0,0 +1,62 @@
+# 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.
+
+
+import torch
+from tqdm import tqdm
+
+from nemo.core.classes.common import safe_instantiate
+from nemo.core.config import hydra_runner
+
+
+def get_pitch_stats(pitch_list):
+    pitch_tensor = torch.cat(pitch_list)
+    pitch_mean, pitch_std = pitch_tensor.mean().item(), pitch_tensor.std().item()
+    pitch_min, pitch_max = pitch_tensor.min().item(), pitch_tensor.max().item()
+    print(f"PITCH_MEAN={pitch_mean}, PITCH_STD={pitch_std}")
+    print(f"PITCH_MIN={pitch_min}, PITCH_MAX={pitch_max}")
+
+
+def preprocess_ds_for_fastpitch_align(dataloader):
+    pitch_list = []
+    for batch in tqdm(dataloader, total=len(dataloader)):
+        audios, audio_lengths, tokens, tokens_lengths, align_prior_matrices, pitches, pitches_lengths, *_ = batch
+        pitch = pitches.squeeze(0)
+        pitch_list.append(pitch[pitch != 0])
+
+    get_pitch_stats(pitch_list)
+
+
+CFG_NAME2FUNC = {
+    "ds_for_fastpitch_align": preprocess_ds_for_fastpitch_align,
+    "ds_for_mixer_tts": preprocess_ds_for_fastpitch_align,
+}
+
+
+@hydra_runner(config_path='ljspeech/ds_conf', config_name='ds_for_fastpitch_align')
+def main(cfg):
+    dataset = safe_instantiate(cfg.dataset)
+    dataloader = torch.utils.data.DataLoader(
+        dataset=dataset,
+        batch_size=1,
+        collate_fn=dataset._collate_fn,
+        num_workers=cfg.get("dataloader_params", {}).get("num_workers", 4),
+    )
+
+    print(f"Processing {cfg.manifest_filepath}:")
+    CFG_NAME2FUNC[cfg.name](dataloader)
+
+
+if __name__ == '__main__':
+    main()  # noqa pylint: disable=no-value-for-parameter

scripts/dataset_processing/tts/generate_mels.py

@@ -0,0 +1,185 @@
+# 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 is to generate mel spectrograms from a Fastpitch model checkpoint. Please see general usage below. It runs
+on GPUs by default, but you can add `--num-workers 5 --cpu` as an option to run on CPUs.
+
+$ python scripts/dataset_processing/tts/generate_mels.py \
+    --fastpitch-model-ckpt ./models/fastpitch/multi_spk/FastPitch--val_loss\=1.4473-epoch\=209.ckpt \
+    --input-json-manifests /home/xueyang/HUI-Audio-Corpus-German-clean/test_manifest_text_normed_phonemes.json
+    --output-json-manifest-root /home/xueyang/experiments/multi_spk_tts_de
+"""
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+import torch
+from joblib import Parallel, delayed
+from tqdm import tqdm
+
+from nemo.collections.tts.models import FastPitchModel
+from nemo.collections.tts.parts.utils.tts_dataset_utils import (
+    BetaBinomialInterpolator,
+    beta_binomial_prior_distribution,
+)
+from nemo.utils import logging
+
+
+def get_args():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        description="Generate mel spectrograms with pretrained FastPitch model, and create manifests for finetuning Hifigan.",
+    )
+    parser.add_argument(
+        "--fastpitch-model-ckpt",
+        required=True,
+        type=Path,
+        help="Specify a full path of a fastpitch model checkpoint with the suffix of either .ckpt or .nemo.",
+    )
+    parser.add_argument(
+        "--input-json-manifests",
+        nargs="+",
+        required=True,
+        type=Path,
+        help="Specify a full path of a JSON manifest. You could add multiple manifests.",
+    )
+    parser.add_argument(
+        "--output-json-manifest-root",
+        required=True,
+        type=Path,
+        help="Specify a full path of output root that would contain new manifests.",
+    )
+    parser.add_argument(
+        "--num-workers",
+        default=-1,
+        type=int,
+        help="Specify the max number of concurrently Python workers processes. "
+        "If -1 all CPUs are used. If 1 no parallel computing is used.",
+    )
+    parser.add_argument("--cpu", action='store_true', default=False, help="Generate mel spectrograms using CPUs.")
+    args = parser.parse_args()
+    return args
+
+
+def __load_wav(audio_file):
+    with sf.SoundFile(audio_file, 'r') as f:
+        samples = f.read(dtype='float32')
+    return samples.transpose()
+
+
+def __generate_mels(entry, spec_model, device, use_beta_binomial_interpolator, mel_root):
+    # Generate a spectrograms (we need to use ground truth alignment for correct matching between audio and mels)
+    audio = __load_wav(entry["audio_filepath"])
+    audio = torch.from_numpy(audio).unsqueeze(0).to(device)
+    audio_len = torch.tensor(audio.shape[1], dtype=torch.long, device=device).unsqueeze(0)
+
+    if spec_model.fastpitch.speaker_emb is not None and "speaker" in entry:
+        speaker = torch.tensor([entry['speaker']]).to(device)
+    else:
+        speaker = None
+
+    with torch.no_grad():
+        if "normalized_text" in entry:
+            text = spec_model.parse(entry["normalized_text"], normalize=False)
+        else:
+            text = spec_model.parse(entry['text'])
+
+        text_len = torch.tensor(text.shape[-1], dtype=torch.long, device=device).unsqueeze(0)
+        spect, spect_len = spec_model.preprocessor(input_signal=audio, length=audio_len)
+
+        # Generate attention prior and spectrogram inputs for HiFi-GAN
+        if use_beta_binomial_interpolator:
+            beta_binomial_interpolator = BetaBinomialInterpolator()
+            attn_prior = (
+                torch.from_numpy(beta_binomial_interpolator(spect_len.item(), text_len.item()))
+                .unsqueeze(0)
+                .to(text.device)
+            )
+        else:
+            attn_prior = (
+                torch.from_numpy(beta_binomial_prior_distribution(text_len.item(), spect_len.item()))
+                .unsqueeze(0)
+                .to(text.device)
+            )
+
+        spectrogram = spec_model.forward(
+            text=text,
+            input_lens=text_len,
+            spec=spect,
+            mel_lens=spect_len,
+            attn_prior=attn_prior,
+            speaker=speaker,
+        )[0]
+
+        save_path = mel_root / f"{Path(entry['audio_filepath']).stem}.npy"
+        np.save(save_path, spectrogram[0].to('cpu').numpy())
+        entry["mel_filepath"] = str(save_path)
+
+    return entry
+
+
+def main():
+    args = get_args()
+    ckpt_path = args.fastpitch_model_ckpt
+    input_manifest_filepaths = args.input_json_manifests
+    output_json_manifest_root = args.output_json_manifest_root
+
+    mel_root = output_json_manifest_root / "mels"
+    mel_root.mkdir(exist_ok=True, parents=True)
+
+    # load pretrained FastPitch model checkpoint
+    suffix = ckpt_path.suffix
+    if suffix == ".nemo":
+        spec_model = FastPitchModel.restore_from(ckpt_path).eval()
+    elif suffix == ".ckpt":
+        spec_model = FastPitchModel.load_from_checkpoint(ckpt_path).eval()
+    else:
+        raise ValueError(f"Unsupported suffix: {suffix}")
+    if not args.cpu:
+        spec_model.cuda()
+    device = spec_model.device
+
+    use_beta_binomial_interpolator = spec_model.cfg.train_ds.dataset.get("use_beta_binomial_interpolator", False)
+
+    for manifest in input_manifest_filepaths:
+        logging.info(f"Processing {manifest}.")
+        entries = []
+        with open(manifest, "r") as fjson:
+            for line in fjson:
+                entries.append(json.loads(line.strip()))
+
+        if device == "cpu":
+            new_entries = Parallel(n_jobs=args.num_workers)(
+                delayed(__generate_mels)(entry, spec_model, device, use_beta_binomial_interpolator, mel_root)
+                for entry in entries
+            )
+        else:
+            new_entries = []
+            for entry in tqdm(entries):
+                new_entry = __generate_mels(entry, spec_model, device, use_beta_binomial_interpolator, mel_root)
+                new_entries.append(new_entry)
+
+        mel_manifest_path = output_json_manifest_root / f"{manifest.stem}_mel{manifest.suffix}"
+        with open(mel_manifest_path, "w") as fmel:
+            for entry in new_entries:
+                fmel.write(json.dumps(entry) + "\n")
+        logging.info(f"Processing {manifest} is complete --> {mel_manifest_path}")
+
+
+if __name__ == "__main__":
+    main()

scripts/dataset_processing/tts/hifitts/get_data.py

@@ -0,0 +1,123 @@
+# 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.
+
+import argparse
+import glob
+import json
+import re
+import tarfile
+import urllib.request
+from pathlib import Path
+
+from tqdm import tqdm
+
+from nemo.utils.tar_utils import safe_extract
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description='Download HiFiTTS and create manifests with predefined split')
+    parser.add_argument(
+        "--data-root",
+        required=True,
+        type=Path,
+        help='Directory into which to download and extract dataset. \{data-root\}/hi_fi_tts_v0 will be created.',
+    )
+    parser.add_argument(
+        '--split',
+        type=str,
+        default='all',
+        help='Choose to generate manifest for all or one of (train, test, split), note that this will still download the full dataset.',
+    )
+
+    args = parser.parse_args()
+    return args
+
+
+URL = "https://us.openslr.org/resources/109/hi_fi_tts_v0.tar.gz"
+
+
+def __maybe_download_file(source_url, destination_path):
+    if not destination_path.exists():
+        tmp_file_path = destination_path.with_suffix('.tmp')
+        urllib.request.urlretrieve(source_url, filename=str(tmp_file_path))
+        tmp_file_path.rename(destination_path)
+
+
+def __extract_file(filepath, data_dir):
+    try:
+        with tarfile.open(filepath) as tar:
+            safe_extract(tar, str(data_dir))
+    except Exception:
+        print(f"Error while extracting {filepath}. Already extracted?")
+
+
+def __process_data(data_root, filelists):
+    # Create manifests (based on predefined NVIDIA's split)
+    for split in tqdm(filelists):
+        manifest_target = data_root / f"{split}_manifest.json"
+        print(f"Creating manifest for {split}.")
+
+        entries = []
+        for manifest_src in glob.glob(str(data_root / f"*_{split}.json")):
+            try:
+                search_res = re.search('.*\/([0-9]+)_manifest_([a-z]+)_.*.json', manifest_src)
+                speaker_id = search_res.group(1)
+                audio_quality = search_res.group(2)
+            except Exception:
+                print(f"Failed to find speaker id or audio quality for {manifest_src}, check formatting.")
+                continue
+
+            with open(manifest_src, 'r') as f_in:
+                for input_json_entry in f_in:
+                    data = json.loads(input_json_entry)
+
+                    # Make sure corresponding wavfile exists
+                    wav_path = data_root / data['audio_filepath']
+                    assert wav_path.exists(), f"{wav_path} does not exist!"
+
+                    entry = {
+                        'audio_filepath': data['audio_filepath'],
+                        'duration': data['duration'],
+                        'text': data['text'],
+                        'normalized_text': data['text_normalized'],
+                        'speaker': int(speaker_id),
+                        # Audio_quality is either clean or other.
+                        # The clean set includes recordings with high sound-to-noise ratio and wide bandwidth.
+                        # The books with noticeable noise or narrow bandwidth are included in the other subset.
+                        # Note: some speaker_id's have both clean and other audio quality.
+                        'audio_quality': audio_quality,
+                    }
+                    entries.append(entry)
+
+        with open(manifest_target, 'w') as f_out:
+            for m in entries:
+                f_out.write(json.dumps(m) + '\n')
+
+
+def main():
+    args = get_args()
+
+    split = ['train', 'dev', 'test'] if args.split == 'all' else list(args.split)
+
+    tarred_data_path = args.data_root / "hi_fi_tts_v0.tar.gz"
+
+    __maybe_download_file(URL, tarred_data_path)
+    __extract_file(str(tarred_data_path), str(args.data_root))
+
+    data_root = args.data_root / "hi_fi_tts_v0"
+    __process_data(data_root, split)
+
+
+if __name__ == '__main__':
+    main()

scripts/dataset_processing/tts/hui_acg/ds_conf/ds_for_fastpitch_align.yaml

@@ -0,0 +1,45 @@
+name: "ds_for_fastpitch_align"
+
+manifest_filepath: ???
+sup_data_path: ???
+sup_data_types: [ "align_prior_matrix", "pitch" ]
+
+dataset:
+  _target_: nemo.collections.tts.data.dataset.TTSDataset
+  manifest_filepath: ${manifest_filepath}
+  sample_rate: 44100
+  sup_data_path: ${sup_data_path}
+  sup_data_types: ${sup_data_types}
+  n_fft: 2048
+  win_length: 2048
+  hop_length: 512
+  window: "hann"
+  n_mels: 80
+  lowfreq: 0
+  highfreq: null
+  max_duration: 15
+  min_duration: 0.1
+  ignore_file: null
+  trim: false
+  pitch_fmin: 65.40639132514966
+  pitch_fmax: 2093.004522404789
+  use_beta_binomial_interpolator: false
+
+  text_normalizer:
+    _target_: nemo_text_processing.text_normalization.normalize.Normalizer
+    lang: de
+    input_case: cased
+
+  text_normalizer_call_kwargs:
+    verbose: false
+    punct_pre_process: true
+    punct_post_process: true
+
+  text_tokenizer:
+    _target_: nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.GermanPhonemesTokenizer
+    punct: true
+    apostrophe: true
+    pad_with_space: true
+
+dataloader_params:
+  num_workers: 12

scripts/dataset_processing/tts/hui_acg/get_data.py

@@ -0,0 +1,334 @@
+# 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.
+
+import argparse
+import json
+import random
+import shutil
+import urllib.request
+from pathlib import Path
+
+import pandas as pd
+from joblib import Parallel, delayed
+from tqdm import tqdm
+
+try:
+    from nemo_text_processing.text_normalization.normalize import Normalizer
+except (ImportError, ModuleNotFoundError):
+    raise ModuleNotFoundError(
+        "The package `nemo_text_processing` was not installed in this environment. Please refer to"
+        " https://github.com/NVIDIA/NeMo-text-processing and install this package before using "
+        "this script"
+    )
+
+from nemo.utils import logging
+
+# full corpus.
+URLS_FULL = {
+    "Bernd_Ungerer": "https://opendata.iisys.de/opendata/Datasets/HUI-Audio-Corpus-German/dataset_full/Bernd_Ungerer.zip",
+    "Eva_K": "https://opendata.iisys.de/opendata/Datasets/HUI-Audio-Corpus-German/dataset_full/Eva_K.zip",
+    "Friedrich": "https://opendata.iisys.de/opendata/Datasets/HUI-Audio-Corpus-German/dataset_full/Friedrich.zip",
+    "Hokuspokus": "https://opendata.iisys.de/opendata/Datasets/HUI-Audio-Corpus-German/dataset_full/Hokuspokus.zip",
+    "Karlsson": "https://opendata.iisys.de/opendata/Datasets/HUI-Audio-Corpus-German/dataset_full/Karlsson.zip",
+    "others": "https://opendata.iisys.de/opendata/Datasets/HUI-Audio-Corpus-German/dataset_full/others.zip",
+}
+URL_STATS_FULL = "https://opendata.iisys.de/opendata/Datasets/HUI-Audio-Corpus-German/datasetStatistic.zip"
+
+# the clean subset of the full corpus.
+URLS_CLEAN = {
+    "Bernd_Ungerer": "https://opendata.iisys.de/opendata/Datasets/HUI-Audio-Corpus-German/dataset_clean/Bernd_Ungerer_Clean.zip",
+    "Eva_K": "https://opendata.iisys.de/opendata/Datasets/HUI-Audio-Corpus-German/dataset_clean/Eva_K_Clean.zip",
+    "Friedrich": "https://opendata.iisys.de/opendata/Datasets/HUI-Audio-Corpus-German/dataset_clean/Friedrich_Clean.zip",
+    "Hokuspokus": "https://opendata.iisys.de/opendata/Datasets/HUI-Audio-Corpus-German/dataset_clean/Hokuspokus_Clean.zip",
+    "Karlsson": "https://opendata.iisys.de/opendata/Datasets/HUI-Audio-Corpus-German/dataset_clean/Karlsson_Clean.zip",
+    "others": "https://opendata.iisys.de/opendata/Datasets/HUI-Audio-Corpus-German/dataset_clean/others_Clean.zip",
+}
+URL_STATS_CLEAN = "https://opendata.iisys.de/opendata/Datasets/HUI-Audio-Corpus-German/datasetStatisticClean.zip"
+
+
+def get_args():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        description="Download HUI-Audio-Corpus-German and create manifests with predefined split. "
+        "Please check details about the corpus in https://github.com/iisys-hof/HUI-Audio-Corpus-German.",
+    )
+    parser.add_argument("--data-root", required=True, type=Path, help="where the resulting dataset will reside.")
+    parser.add_argument("--manifests-root", required=True, type=Path, help="where the manifests files will reside.")
+    parser.add_argument("--set-type", default="clean", choices=["full", "clean"], type=str)
+    parser.add_argument("--min-duration", default=0.1, type=float)
+    parser.add_argument("--max-duration", default=15, type=float)
+    parser.add_argument(
+        "--num-workers",
+        default=-1,
+        type=int,
+        help="Specify the max number of concurrently Python workers processes. "
+        "If -1 all CPUs are used. If 1 no parallel computing is used.",
+    )
+    parser.add_argument(
+        "--normalize-text",
+        default=False,
+        action='store_true',
+        help="Normalize original text and add a new entry 'normalized_text' to .json file if True.",
+    )
+    parser.add_argument(
+        "--val-num-utts-per-speaker",
+        default=1,
+        type=int,
+        help="Specify the number of utterances for each speaker in val split. All speakers are covered.",
+    )
+    parser.add_argument(
+        "--test-num-utts-per-speaker",
+        default=1,
+        type=int,
+        help="Specify the number of utterances for each speaker in test split. All speakers are covered.",
+    )
+    parser.add_argument(
+        "--seed-for-ds-split",
+        default=100,
+        type=float,
+        help="Seed for deterministic split of train/dev/test, NVIDIA's default is 100",
+    )
+    args = parser.parse_args()
+    return args
+
+
+def __maybe_download_file(source_url, destination_path):
+    if not destination_path.exists():
+        logging.info(f"Downloading data: {source_url} --> {destination_path}")
+        tmp_file_path = destination_path.with_suffix(".tmp")
+        urllib.request.urlretrieve(source_url, filename=tmp_file_path)
+        tmp_file_path.rename(destination_path)
+    else:
+        logging.info(f"Skipped downloading data because it exists: {destination_path}")
+
+
+def __extract_file(filepath, data_dir):
+    logging.info(f"Unzipping data: {filepath} --> {data_dir}")
+    shutil.unpack_archive(filepath, data_dir)
+    logging.info(f"Unzipping data is complete: {filepath}.")
+
+
+def __save_json(json_file, dict_list):
+    logging.info(f"Saving JSON split to {json_file}.")
+    with open(json_file, "w") as f:
+        for d in dict_list:
+            f.write(json.dumps(d) + "\n")
+
+
+def __process_data(
+    dataset_path,
+    stat_path_root,
+    speaker_id,
+    min_duration,
+    max_duration,
+    val_size,
+    test_size,
+    seed_for_ds_split,
+):
+    logging.info(f"Preparing JSON split for speaker {speaker_id}.")
+    # parse statistic.txt
+    stat_path = stat_path_root / "statistic.txt"
+    with open(stat_path, 'r') as fstat:
+        lines = fstat.readlines()
+        num_utts = int(lines[4].strip().split()[-1])
+        hours = round(float(lines[9].strip().split()[-1]) / 3600.0, 2)
+
+    # parse overview.csv to generate JSON splits.
+    overview_path = stat_path_root / "overview.csv"
+    entries = []
+    with open(overview_path, 'r') as foverview:
+        # Let's skip the header
+        foverview.readline()
+        for line in tqdm(foverview):
+            file_stem, duration, *_, text = line.strip().split("|")
+            duration = float(duration)
+
+            # file_stem -> dir_name (e.g. maerchen_01_f000051 -> maerchen)
+            dir_name = "_".join(file_stem.split("_")[:-2])
+            audio_path = dataset_path / dir_name / "wavs" / f"{file_stem}.wav"
+
+            if min_duration <= duration <= max_duration:
+                entry = {
+                    "audio_filepath": str(audio_path),
+                    "duration": duration,
+                    "text": text,
+                    "speaker": speaker_id,
+                }
+                entries.append(entry)
+
+    random.Random(seed_for_ds_split).shuffle(entries)
+    train_size = len(entries) - val_size - test_size
+    if train_size <= 0:
+        logging.warning(f"Skipped speaker {speaker_id}. Not enough data for train, val and test.")
+        train, val, test, is_skipped = [], [], [], True
+    else:
+        logging.info(f"Preparing JSON split for speaker {speaker_id} is complete.")
+        train, val, test, is_skipped = (
+            entries[:train_size],
+            entries[train_size : train_size + val_size],
+            entries[train_size + val_size :],
+            False,
+        )
+
+    return {
+        "train": train,
+        "val": val,
+        "test": test,
+        "is_skipped": is_skipped,
+        "hours": hours,
+        "num_utts": num_utts,
+    }
+
+
+def __text_normalization(json_file, num_workers=-1):
+    text_normalizer_call_kwargs = {
+        "punct_pre_process": True,
+        "punct_post_process": True,
+    }
+    text_normalizer = Normalizer(
+        lang="de",
+        input_case="cased",
+        overwrite_cache=True,
+        cache_dir=str(json_file.parent / "cache_dir"),
+    )
+
+    def normalizer_call(x):
+        return text_normalizer.normalize(x, **text_normalizer_call_kwargs)
+
+    def add_normalized_text(line_dict):
+        normalized_text = normalizer_call(line_dict["text"])
+        line_dict.update({"normalized_text": normalized_text})
+        return line_dict
+
+    logging.info(f"Normalizing text for {json_file}.")
+    with open(json_file, 'r', encoding='utf-8') as fjson:
+        lines = fjson.readlines()
+        # Note: you need to verify which backend works well on your cluster.
+        # backend="loky" is fine on multi-core Ubuntu OS; backend="threading" on Slurm.
+        dict_list = Parallel(n_jobs=num_workers)(
+            delayed(add_normalized_text)(json.loads(line)) for line in tqdm(lines)
+        )
+
+    json_file_text_normed = json_file.parent / f"{json_file.stem}_text_normed{json_file.suffix}"
+    with open(json_file_text_normed, 'w', encoding="utf-8") as fjson_norm:
+        for dct in dict_list:
+            fjson_norm.write(json.dumps(dct) + "\n")
+    logging.info(f"Normalizing text is complete: {json_file} --> {json_file_text_normed}")
+
+
+def main():
+    args = get_args()
+    data_root = args.data_root
+    manifests_root = args.manifests_root
+    set_type = args.set_type
+
+    dataset_root = data_root / f"HUI-Audio-Corpus-German-{set_type}"
+    dataset_root.mkdir(parents=True, exist_ok=True)
+
+    if set_type == "full":
+        data_source = URLS_FULL
+        stats_source = URL_STATS_FULL
+    elif set_type == "clean":
+        data_source = URLS_CLEAN
+        stats_source = URL_STATS_CLEAN
+    else:
+        raise ValueError(f"Unknown {set_type}. Please choose either clean or full.")
+
+    # download and unzip dataset stats
+    zipped_stats_path = dataset_root / Path(stats_source).name
+    __maybe_download_file(stats_source, zipped_stats_path)
+    __extract_file(zipped_stats_path, dataset_root)
+
+    # download datasets
+    # Note: you need to verify which backend works well on your cluster.
+    # backend="loky" is fine on multi-core Ubuntu OS; backend="threading" on Slurm.
+    Parallel(n_jobs=args.num_workers)(
+        delayed(__maybe_download_file)(data_url, dataset_root / Path(data_url).name)
+        for _, data_url in data_source.items()
+    )
+
+    # unzip datasets
+    # Note: you need to verify which backend works well on your cluster.
+    # backend="loky" is fine on multi-core Ubuntu OS; backend="threading" on Slurm.
+    Parallel(n_jobs=args.num_workers)(
+        delayed(__extract_file)(dataset_root / Path(data_url).name, dataset_root)
+        for _, data_url in data_source.items()
+    )
+
+    # generate json files for train/val/test splits
+    stats_path_root = dataset_root / Path(stats_source).stem / "speacker"
+    entries_train, entries_val, entries_test = [], [], []
+    speaker_entries = []
+    num_speakers = 0
+    for child in stats_path_root.iterdir():
+        if child.is_dir():
+            speaker = child.name
+            num_speakers += 1
+            speaker_stats_root = stats_path_root / speaker
+            speaker_data_path = dataset_root / speaker
+
+            logging.info(f"Processing Speaker: {speaker}")
+            results = __process_data(
+                speaker_data_path,
+                speaker_stats_root,
+                num_speakers,
+                args.min_duration,
+                args.max_duration,
+                args.val_num_utts_per_speaker,
+                args.test_num_utts_per_speaker,
+                args.seed_for_ds_split,
+            )
+
+            entries_train.extend(results["train"])
+            entries_val.extend(results["val"])
+            entries_test.extend(results["test"])
+
+            speaker_entry = {
+                "speaker_name": speaker,
+                "speaker_id": num_speakers,
+                "hours": results["hours"],
+                "num_utts": results["num_utts"],
+                "is_skipped": results["is_skipped"],
+            }
+            speaker_entries.append(speaker_entry)
+
+    # shuffle in place across multiple speakers
+    random.Random(args.seed_for_ds_split).shuffle(entries_train)
+    random.Random(args.seed_for_ds_split).shuffle(entries_val)
+    random.Random(args.seed_for_ds_split).shuffle(entries_test)
+
+    # save speaker stats.
+    df = pd.DataFrame.from_records(speaker_entries)
+    df.sort_values(by="hours", ascending=False, inplace=True)
+    spk2id_file_path = manifests_root / "spk2id.csv"
+    df.to_csv(spk2id_file_path, index=False)
+    logging.info(f"Saving Speaker to ID mapping to {spk2id_file_path}.")
+
+    # save json splits.
+    train_json = manifests_root / "train_manifest.json"
+    val_json = manifests_root / "val_manifest.json"
+    test_json = manifests_root / "test_manifest.json"
+    __save_json(train_json, entries_train)
+    __save_json(val_json, entries_val)
+    __save_json(test_json, entries_test)
+
+    # normalize text if requested. New json file, train_manifest_text_normed.json, will be generated.
+    if args.normalize_text:
+        __text_normalization(train_json, args.num_workers)
+        __text_normalization(val_json, args.num_workers)
+        __text_normalization(test_json, args.num_workers)
+
+
+if __name__ == "__main__":
+    main()

scripts/dataset_processing/tts/libritts/get_data.py

@@ -0,0 +1,134 @@
+# 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.
+#
+# USAGE: python get_data.py --data-root=<where to put data> --data-set=<datasets_to_download> --num-workers=<number of parallel workers>
+# where <datasets_to_download> can be: dev_clean, dev_other, test_clean,
+# test_other, train_clean_100, train_clean_360, train_other_500 or ALL
+# You can also put more than one data_set comma-separated:
+# --data-set=dev_clean,train_clean_100
+import argparse
+import fnmatch
+import functools
+import json
+import multiprocessing
+import os
+import subprocess
+import tarfile
+import urllib.request
+from pathlib import Path
+
+from tqdm import tqdm
+
+from nemo.utils.tar_utils import safe_extract
+
+parser = argparse.ArgumentParser(description='Download LibriTTS and create manifests')
+parser.add_argument("--data-root", required=True, type=Path)
+parser.add_argument("--data-sets", default="dev_clean", type=str)
+parser.add_argument("--num-workers", default=4, type=int)
+args = parser.parse_args()
+
+URLS = {
+    'TRAIN_CLEAN_100': "https://www.openslr.org/resources/60/train-clean-100.tar.gz",
+    'TRAIN_CLEAN_360': "https://www.openslr.org/resources/60/train-clean-360.tar.gz",
+    'TRAIN_OTHER_500': "https://www.openslr.org/resources/60/train-other-500.tar.gz",
+    'DEV_CLEAN': "https://www.openslr.org/resources/60/dev-clean.tar.gz",
+    'DEV_OTHER': "https://www.openslr.org/resources/60/dev-other.tar.gz",
+    'TEST_CLEAN': "https://www.openslr.org/resources/60/test-clean.tar.gz",
+    'TEST_OTHER': "https://www.openslr.org/resources/60/test-other.tar.gz",
+}
+
+
+def __maybe_download_file(source_url, destination_path):
+    if not destination_path.exists():
+        tmp_file_path = destination_path.with_suffix('.tmp')
+        urllib.request.urlretrieve(source_url, filename=str(tmp_file_path))
+        tmp_file_path.rename(destination_path)
+
+
+def __extract_file(filepath, data_dir):
+    try:
+        with tarfile.open(filepath) as tar:
+            safe_extract(tar, str(data_dir))
+    except Exception:
+        print(f"Error while extracting {filepath}. Already extracted?")
+
+
+def __process_transcript(file_path: str):
+    entries = []
+    with open(file_path, encoding="utf-8") as fin:
+        text = fin.readlines()[0].strip()
+
+        # TODO(oktai15): add normalized text via Normalizer/NormalizerWithAudio
+        wav_file = file_path.replace(".normalized.txt", ".wav")
+        speaker_id = file_path.split('/')[-3]
+        assert os.path.exists(wav_file), f"{wav_file} not found!"
+        duration = subprocess.check_output(["soxi", "-D", wav_file])
+        entry = {
+            'audio_filepath': os.path.abspath(wav_file),
+            'duration': float(duration),
+            'text': text,
+            'speaker': int(speaker_id),
+        }
+
+        entries.append(entry)
+
+    return entries
+
+
+def __process_data(data_folder, manifest_file, num_workers):
+    files = []
+    entries = []
+
+    for root, dirnames, filenames in os.walk(data_folder):
+        # we will use normalized text provided by the original dataset
+        for filename in fnmatch.filter(filenames, '*.normalized.txt'):
+            files.append(os.path.join(root, filename))
+
+    with multiprocessing.Pool(num_workers) as p:
+        processing_func = functools.partial(__process_transcript)
+        results = p.imap(processing_func, files)
+        for result in tqdm(results, total=len(files)):
+            entries.extend(result)
+
+    with open(manifest_file, 'w') as fout:
+        for m in entries:
+            fout.write(json.dumps(m) + '\n')
+
+
+def main():
+    data_root = args.data_root
+    data_sets = args.data_sets
+    num_workers = args.num_workers
+
+    if data_sets == "ALL":
+        data_sets = "dev_clean,dev_other,train_clean_100,train_clean_360,train_other_500,test_clean,test_other"
+    if data_sets == "mini":
+        data_sets = "dev_clean,train_clean_100"
+    for data_set in data_sets.split(','):
+        filepath = data_root / f"{data_set}.tar.gz"
+        print(f"Downloading data for {data_set}...")
+        __maybe_download_file(URLS[data_set.upper()], filepath)
+        print("Extracting...")
+        __extract_file(str(filepath), str(data_root))
+
+        print("Processing and building manifest.")
+        __process_data(
+            str(data_root / "LibriTTS" / data_set.replace("_", "-")),
+            str(data_root / "LibriTTS" / f"{data_set}.json"),
+            num_workers=num_workers,
+        )
+
+
+if __name__ == "__main__":
+    main()

scripts/dataset_processing/tts/ljspeech/ds_conf/ds_for_fastpitch_align.yaml

@@ -0,0 +1,49 @@
+name: "ds_for_fastpitch_align"
+
+manifest_filepath: "train_manifest.json"
+sup_data_path: "sup_data"
+sup_data_types: [ "align_prior_matrix", "pitch" ]
+phoneme_dict_path: "scripts/tts_dataset_files/cmudict-0.7b_nv22.10"
+heteronyms_path: "scripts/tts_dataset_files/heteronyms-052722"
+
+dataset:
+  _target_: nemo.collections.tts.data.dataset.TTSDataset
+  manifest_filepath: ${manifest_filepath}
+  sample_rate: 22050
+  sup_data_path: ${sup_data_path}
+  sup_data_types: ${sup_data_types}
+  n_fft: 1024
+  win_length: 1024
+  hop_length: 256
+  window: "hann"
+  n_mels: 80
+  lowfreq: 0
+  highfreq: 8000
+  max_duration: null
+  min_duration: 0.1
+  ignore_file: null
+  trim: false
+  pitch_fmin: 65.40639132514966
+  pitch_fmax: 2093.004522404789
+
+  text_normalizer:
+    _target_: nemo_text_processing.text_normalization.normalize.Normalizer
+    lang: en
+    input_case: cased
+
+  text_normalizer_call_kwargs:
+    verbose: false
+    punct_pre_process: true
+    punct_post_process: true
+
+  text_tokenizer:
+    _target_: nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.EnglishPhonemesTokenizer
+    punct: true
+    stresses: true
+    chars: true
+    apostrophe: true
+    pad_with_space: true
+    g2p:
+      _target_: nemo.collections.tts.g2p.models.en_us_arpabet.EnglishG2p
+      phoneme_dict: ${phoneme_dict_path}
+      heteronyms: ${heteronyms_path}

scripts/dataset_processing/tts/ljspeech/ds_conf/ds_for_mixer_tts.yaml

@@ -0,0 +1,49 @@
+name: "ds_for_mixer_tts"
+
+manifest_filepath: "train_manifest.json"
+sup_data_path: "sup_data"
+sup_data_types: [ "align_prior_matrix", "pitch" ]
+phoneme_dict_path: "scripts/tts_dataset_files/cmudict-0.7b_nv22.10"
+heteronyms_path: "scripts/tts_dataset_files/heteronyms-052722"
+
+dataset:
+  _target_: nemo.collections.tts.data.dataset.TTSDataset
+  manifest_filepath: ${manifest_filepath}
+  sample_rate: 22050
+  sup_data_path: ${sup_data_path}
+  sup_data_types: ${sup_data_types}
+  n_fft: 1024
+  win_length: 1024
+  hop_length: 256
+  window: "hann"
+  n_mels: 80
+  lowfreq: 0
+  highfreq: 8000
+  max_duration: null
+  min_duration: 0.1
+  ignore_file: null
+  trim: false
+  pitch_fmin: 65.40639132514966
+  pitch_fmax: 2093.004522404789
+
+  text_normalizer:
+    _target_: nemo_text_processing.text_normalization.normalize.Normalizer
+    lang: en
+    input_case: cased
+
+  text_normalizer_call_kwargs:
+    verbose: false
+    punct_pre_process: true
+    punct_post_process: true
+
+  text_tokenizer:
+    _target_: nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.EnglishPhonemesTokenizer
+    punct: true
+    stresses: true
+    chars: true
+    apostrophe: true
+    pad_with_space: true
+    g2p:
+      _target_: nemo.collections.tts.g2p.models.en_us_arpabet.EnglishG2p
+      phoneme_dict: ${phoneme_dict_path}
+      heteronyms: ${heteronyms_path}

scripts/dataset_processing/tts/ljspeech/get_data.py

@@ -0,0 +1,134 @@
+# 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.
+
+import argparse
+import json
+import tarfile
+import urllib.request
+from pathlib import Path
+
+from tqdm import tqdm
+
+from nemo.utils.tar_utils import safe_extract
+
+try:
+    from nemo_text_processing.text_normalization.normalize import Normalizer
+except (ImportError, ModuleNotFoundError):
+    raise ModuleNotFoundError(
+        "The package `nemo_text_processing` was not installed in this environment. Please refer to"
+        " https://github.com/NVIDIA/NeMo-text-processing and install this package before using "
+        "this script"
+    )
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description='Download LJSpeech and create manifests with predefined split')
+    parser.add_argument("--data-root", required=True, type=Path)
+
+    args = parser.parse_args()
+    return args
+
+
+URL = "https://data.keithito.com/data/speech/LJSpeech-1.1.tar.bz2"
+FILELIST_BASE = 'https://raw.githubusercontent.com/NVIDIA/tacotron2/master/filelists'
+
+
+def _load_sox():
+    try:
+        import sox
+    except ImportError:
+        raise ImportError(
+            "Optional dependency 'sox' is required by this script. Install it with: pip install sox"
+        ) from None
+
+    return sox
+
+
+def __maybe_download_file(source_url, destination_path):
+    if not destination_path.exists():
+        tmp_file_path = destination_path.with_suffix('.tmp')
+        urllib.request.urlretrieve(source_url, filename=str(tmp_file_path))
+        tmp_file_path.rename(destination_path)
+
+
+def __extract_file(filepath, data_dir):
+    try:
+        with tarfile.open(filepath) as tar:
+            safe_extract(tar, str(data_dir))
+    except Exception:
+        print(f"Error while extracting {filepath}. Already extracted?")
+
+
+def __process_data(data_root):
+    sox = _load_sox()
+    text_normalizer = Normalizer(
+        lang="en",
+        input_case="cased",
+        overwrite_cache=True,
+        cache_dir=data_root / "cache_dir",
+    )
+    text_normalizer_call_kwargs = {"punct_pre_process": True, "punct_post_process": True}
+    normalizer_call = lambda x: text_normalizer.normalize(x, **text_normalizer_call_kwargs)
+
+    # Create manifests (based on predefined NVIDIA's split)
+    filelists = ['train', 'val', 'test']
+    for split in tqdm(filelists):
+        # Download file list if necessary
+        filelist_path = data_root / f"ljs_audio_text_{split}_filelist.txt"
+
+        if not filelist_path.exists():
+            urllib.request.urlretrieve(
+                f"{FILELIST_BASE}/ljs_audio_text_{split}_filelist.txt",
+                filename=str(filelist_path),
+            )
+
+        manifest_target = data_root / f"{split}_manifest.json"
+        with open(manifest_target, 'w') as f_out:
+            with open(filelist_path, 'r') as filelist:
+                print(f"\nCreating {manifest_target}...")
+                for line in tqdm(filelist):
+                    basename = line[6:16]
+
+                    text = line[21:].strip()
+                    norm_text = normalizer_call(text)
+
+                    # Make sure corresponding wavfile exists
+                    wav_path = data_root / 'wavs' / f"{basename}.wav"
+                    assert wav_path.exists(), f"{wav_path} does not exist!"
+
+                    entry = {
+                        'audio_filepath': str(wav_path),
+                        'duration': sox.file_info.duration(wav_path),
+                        'text': text,
+                        'normalized_text': norm_text,
+                    }
+
+                    f_out.write(json.dumps(entry) + '\n')
+
+
+def main():
+    args = get_args()
+
+    tarred_data_path = args.data_root / "LJSpeech-1.1.tar.bz2"
+
+    __maybe_download_file(URL, tarred_data_path)
+    __extract_file(str(tarred_data_path), str(args.data_root))
+
+    data_root = args.data_root / "LJSpeech-1.1"
+
+    __process_data(data_root)
+
+
+if __name__ == '__main__':
+    main()

scripts/dataset_processing/tts/ljspeech/lj_speech.tsv

@@ -0,0 +1,21 @@
+Mr.	mister
+Mrs.	misses
+Dr.	doctor
+Drs.	doctors
+Co.	company
+Lt.	lieutenant
+Sgt.	sergeant
+St.	saint
+Jr.	junior
+Maj.	major
+Hon.	honorable
+Gov.	governor
+Capt.	captain
+Esq.	esquire
+Gen.	general
+Ltd.	limited
+Rev.	reverend
+Col.	colonel
+Mt.	mount
+Ft.	fort
+etc.	et cetera

scripts/dataset_processing/tts/preprocess_audio.py

@@ -0,0 +1,280 @@
+# Copyright (c) 2022, 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 is used to preprocess audio before TTS model training.
+
+It can be configured to do several processing steps such as silence trimming, volume normalization,
+and duration filtering.
+
+These can be done separately through multiple executions of the script, or all at once to avoid saving
+too many copies of the same audio.
+
+Most of these can also be done by the TTS data loader at training time, but doing them ahead of time
+lets us implement more complex processing, validate the correctness of the output, and save on compute time.
+
+$ python <nemo_root_path>/scripts/dataset_processing/tts/preprocess_audio.py \
+    --input_manifest="<data_root_path>/manifest.json" \
+    --output_manifest="<data_root_path>/manifest_processed.json" \
+    --input_audio_dir="<data_root_path>/audio" \
+    --output_audio_dir="<data_root_path>/audio_processed" \
+    --num_workers=1 \
+    --trim_config_path="<nemo_root_path>/examples/tts/conf/trim/energy.yaml" \
+    --output_sample_rate=22050 \
+    --output_format=flac \
+    --volume_level=0.95 \
+    --min_duration=0.5 \
+    --max_duration=20.0 \
+    --filter_file="filtered.txt"
+"""
+
+import argparse
+import os
+from pathlib import Path
+from typing import Tuple
+
+import librosa
+import soundfile as sf
+from joblib import Parallel, delayed
+from omegaconf import OmegaConf
+from tqdm import tqdm
+
+from nemo.collections.asr.parts.utils.manifest_utils import read_manifest, write_manifest
+from nemo.collections.tts.parts.preprocessing.audio_trimming import AudioTrimmer
+from nemo.collections.tts.parts.utils.tts_dataset_utils import get_abs_rel_paths, normalize_volume
+from nemo.core.classes.common import safe_instantiate
+from nemo.utils import logging
+
+
+def get_args():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        description="Compute speaker level pitch statistics.",
+    )
+    parser.add_argument(
+        "--input_manifest",
+        required=True,
+        type=Path,
+        help="Path to input training manifest.",
+    )
+    parser.add_argument(
+        "--input_audio_dir",
+        required=True,
+        type=Path,
+        help="Path to base directory with audio files.",
+    )
+    parser.add_argument(
+        "--output_manifest",
+        required=True,
+        type=Path,
+        help="Path to output training manifest with processed audio.",
+    )
+    parser.add_argument(
+        "--output_audio_dir",
+        required=True,
+        type=Path,
+        help="Path to output directory for audio files.",
+    )
+    parser.add_argument(
+        "--overwrite_audio",
+        action=argparse.BooleanOptionalAction,
+        help="Whether to reprocess and overwrite existing audio files in output_audio_dir.",
+    )
+    parser.add_argument(
+        "--overwrite_manifest",
+        action=argparse.BooleanOptionalAction,
+        help="Whether to overwrite the output manifest file if it exists.",
+    )
+    parser.add_argument(
+        "--num_workers", default=1, type=int, help="Number of parallel threads to use. If -1 all CPUs are used."
+    )
+    parser.add_argument(
+        "--trim_config_path",
+        required=False,
+        type=Path,
+        help="Path to config file for nemo.collections.tts.data.audio_trimming.AudioTrimmer",
+    )
+    parser.add_argument(
+        "--max_entries", default=0, type=int, help="If provided, maximum number of entries in the manifest to process."
+    )
+    parser.add_argument(
+        "--output_sample_rate", default=0, type=int, help="If provided, rate to resample the audio to."
+    )
+    parser.add_argument(
+        "--output_format",
+        default="wav",
+        type=str,
+        help="If provided, format output audio will be saved as. If not provided, will keep original format.",
+    )
+    parser.add_argument(
+        "--volume_level", default=0.0, type=float, help="If provided, peak volume to normalize audio to."
+    )
+    parser.add_argument(
+        "--min_duration", default=0.0, type=float, help="If provided, filter out utterances shorter than min_duration."
+    )
+    parser.add_argument(
+        "--max_duration", default=0.0, type=float, help="If provided, filter out utterances longer than max_duration."
+    )
+    parser.add_argument(
+        "--filter_file",
+        required=False,
+        type=Path,
+        help="If provided, output filter_file will contain list of " "utterances filtered out.",
+    )
+    args = parser.parse_args()
+    return args
+
+
+def _process_entry(
+    entry: dict,
+    input_audio_dir: Path,
+    output_audio_dir: Path,
+    overwrite_audio: bool,
+    audio_trimmer: AudioTrimmer,
+    output_sample_rate: int,
+    output_format: str,
+    volume_level: float,
+) -> Tuple[dict, float, float]:
+    audio_filepath = Path(entry["audio_filepath"])
+
+    audio_path, audio_path_rel = get_abs_rel_paths(input_path=audio_filepath, base_path=input_audio_dir)
+
+    if not output_format:
+        output_format = audio_path.suffix
+
+    output_path = output_audio_dir / audio_path_rel
+    output_path = output_path.with_suffix(output_format)
+    output_path.parent.mkdir(exist_ok=True, parents=True)
+
+    if output_path.exists() and not overwrite_audio:
+        original_duration = librosa.get_duration(path=audio_path)
+        output_duration = librosa.get_duration(path=output_path)
+    else:
+        audio, sample_rate = librosa.load(audio_path, sr=None)
+        original_duration = librosa.get_duration(y=audio, sr=sample_rate)
+        if audio_trimmer is not None:
+            audio, start_i, end_i = audio_trimmer.trim_audio(
+                audio=audio, sample_rate=int(sample_rate), audio_id=str(audio_path)
+            )
+
+        if output_sample_rate:
+            audio = librosa.resample(y=audio, orig_sr=sample_rate, target_sr=output_sample_rate)
+            sample_rate = output_sample_rate
+
+        if volume_level:
+            audio = normalize_volume(audio, volume_level=volume_level)
+
+        if audio.size > 0:
+            sf.write(file=output_path, data=audio, samplerate=sample_rate)
+            output_duration = librosa.get_duration(y=audio, sr=sample_rate)
+        else:
+            output_duration = 0.0
+
+    entry["duration"] = round(output_duration, 2)
+
+    if os.path.isabs(audio_filepath):
+        entry["audio_filepath"] = str(output_path)
+    else:
+        output_filepath = audio_path_rel.with_suffix(output_format)
+        entry["audio_filepath"] = str(output_filepath)
+
+    return entry, original_duration, output_duration
+
+
+def main():
+    args = get_args()
+
+    input_manifest_path = args.input_manifest
+    output_manifest_path = args.output_manifest
+    input_audio_dir = args.input_audio_dir
+    output_audio_dir = args.output_audio_dir
+    overwrite_audio = args.overwrite_audio
+    overwrite_manifest = args.overwrite_manifest
+    num_workers = args.num_workers
+    max_entries = args.max_entries
+    output_sample_rate = args.output_sample_rate
+    output_format = args.output_format
+    volume_level = args.volume_level
+    min_duration = args.min_duration
+    max_duration = args.max_duration
+    filter_file = args.filter_file
+
+    if output_manifest_path.exists():
+        if overwrite_manifest:
+            print(f"Will overwrite existing manifest path: {output_manifest_path}")
+        else:
+            raise ValueError(f"Manifest path already exists: {output_manifest_path}")
+
+    if args.trim_config_path:
+        audio_trimmer_config = OmegaConf.load(args.trim_config_path)
+        audio_trimmer = safe_instantiate(audio_trimmer_config)
+    else:
+        audio_trimmer = None
+
+    if output_format:
+        if output_format.upper() not in sf.available_formats():
+            raise ValueError(f"Unsupported output audio format: {output_format}")
+        output_format = f".{output_format}"
+
+    output_audio_dir.mkdir(exist_ok=True, parents=True)
+
+    entries = read_manifest(input_manifest_path)
+    if max_entries:
+        entries = entries[:max_entries]
+
+    # 'threading' backend is required when parallelizing torch models.
+    job_outputs = Parallel(n_jobs=num_workers, backend='threading')(
+        delayed(_process_entry)(
+            entry=entry,
+            input_audio_dir=input_audio_dir,
+            output_audio_dir=output_audio_dir,
+            overwrite_audio=overwrite_audio,
+            audio_trimmer=audio_trimmer,
+            output_sample_rate=output_sample_rate,
+            output_format=output_format,
+            volume_level=volume_level,
+        )
+        for entry in tqdm(entries)
+    )
+
+    output_entries = []
+    filtered_entries = []
+    original_durations = 0.0
+    output_durations = 0.0
+    for output_entry, original_duration, output_duration in job_outputs:
+        original_durations += original_duration
+
+        if (
+            output_duration == 0.0
+            or (min_duration and output_duration < min_duration)
+            or (max_duration and output_duration > max_duration)
+        ):
+            if output_duration != original_duration:
+                output_entry["original_duration"] = original_duration
+            filtered_entries.append(output_entry)
+            continue
+
+        output_durations += output_duration
+        output_entries.append(output_entry)
+
+    write_manifest(output_path=output_manifest_path, target_manifest=output_entries, ensure_ascii=False)
+    if filter_file:
+        write_manifest(output_path=str(filter_file), target_manifest=filtered_entries, ensure_ascii=False)
+
+    logging.info(f"Duration of original audio: {original_durations / 3600:.2f} hours")
+    logging.info(f"Duration of processed audio: {output_durations / 3600:.2f} hours")
+
+
+if __name__ == "__main__":
+    main()

scripts/dataset_processing/tts/preprocess_text.py

@@ -0,0 +1,183 @@
+# 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 is used to preprocess text before TTS model training. This is needed mainly for text normalization,
+which is slow to rerun during training.
+
+The output manifest will be the same as the input manifest but with final text stored in the 'normalized_text' field.
+
+$ python <nemo_root_path>/scripts/dataset_processing/tts/preprocess_text.py \
+    --input_manifest="<data_root_path>/manifest.json" \
+    --output_manifest="<data_root_path>/manifest_processed.json" \
+    --normalizer_config_path="<nemo_root_path>/examples/tts/conf/text/normalizer_en.yaml" \
+    --lower_case \
+    --num_workers=4 \
+    --joblib_batch_size=16
+"""
+
+import argparse
+from pathlib import Path
+
+from joblib import Parallel, delayed
+from omegaconf import OmegaConf
+from tqdm import tqdm
+
+from nemo.core.classes.common import safe_instantiate
+
+try:
+    from nemo_text_processing.text_normalization.normalize import Normalizer
+except (ImportError, ModuleNotFoundError):
+    raise ModuleNotFoundError(
+        "The package `nemo_text_processing` was not installed in this environment. Please refer to"
+        " https://github.com/NVIDIA/NeMo-text-processing and install this package before using "
+        "this script"
+    )
+
+from nemo.collections.asr.parts.utils.manifest_utils import read_manifest, write_manifest
+
+
+def get_args():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        description="Process and normalize text data.",
+    )
+    parser.add_argument(
+        "--input_manifest",
+        required=True,
+        type=Path,
+        help="Path to input training manifest.",
+    )
+    parser.add_argument(
+        "--output_manifest",
+        required=True,
+        type=Path,
+        help="Path to output training manifest with processed text.",
+    )
+    parser.add_argument(
+        "--overwrite",
+        action=argparse.BooleanOptionalAction,
+        help="Whether to overwrite the output manifest file if it exists.",
+    )
+    parser.add_argument(
+        "--text_key",
+        default="text",
+        type=str,
+        help="Input text field to normalize.",
+    )
+    parser.add_argument(
+        "--normalized_text_key",
+        default="normalized_text",
+        type=str,
+        help="Output field to save normalized text to.",
+    )
+    parser.add_argument(
+        "--lower_case",
+        action=argparse.BooleanOptionalAction,
+        help="Whether to convert the final text to lower case.",
+    )
+    parser.add_argument(
+        "--normalizer_config_path",
+        required=False,
+        type=Path,
+        help="Path to config file for nemo_text_processing.text_normalization.normalize.Normalizer.",
+    )
+    parser.add_argument(
+        "--num_workers", default=1, type=int, help="Number of parallel threads to use. If -1 all CPUs are used."
+    )
+    parser.add_argument(
+        "--joblib_batch_size", type=int, help="Batch size for joblib workers. Defaults to 'auto' if not provided."
+    )
+    parser.add_argument(
+        "--max_entries", default=0, type=int, help="If provided, maximum number of entries in the manifest to process."
+    )
+
+    args = parser.parse_args()
+    return args
+
+
+def _process_entry(
+    entry: dict,
+    normalizer: Normalizer,
+    text_key: str,
+    normalized_text_key: str,
+    lower_case: bool,
+    lower_case_norm: bool,
+) -> dict:
+    text = entry[text_key]
+
+    if normalizer is not None:
+        if lower_case_norm:
+            text = text.lower()
+        text = normalizer.normalize(text, punct_pre_process=True, punct_post_process=True)
+
+    if lower_case:
+        text = text.lower()
+
+    entry[normalized_text_key] = text
+
+    return entry
+
+
+def main():
+    args = get_args()
+
+    input_manifest_path = args.input_manifest
+    output_manifest_path = args.output_manifest
+    text_key = args.text_key
+    normalized_text_key = args.normalized_text_key
+    lower_case = args.lower_case
+    num_workers = args.num_workers
+    batch_size = args.joblib_batch_size
+    max_entries = args.max_entries
+    overwrite = args.overwrite
+
+    if output_manifest_path.exists():
+        if overwrite:
+            print(f"Will overwrite existing manifest path: {output_manifest_path}")
+        else:
+            raise ValueError(f"Manifest path already exists: {output_manifest_path}")
+
+    if args.normalizer_config_path:
+        normalizer_config = OmegaConf.load(args.normalizer_config_path)
+        normalizer = safe_instantiate(normalizer_config)
+        lower_case_norm = normalizer.input_case == "lower_cased"
+    else:
+        normalizer = None
+        lower_case_norm = False
+
+    entries = read_manifest(input_manifest_path)
+    if max_entries:
+        entries = entries[:max_entries]
+
+    if not batch_size:
+        batch_size = 'auto'
+
+    output_entries = Parallel(n_jobs=num_workers, batch_size=batch_size)(
+        delayed(_process_entry)(
+            entry=entry,
+            normalizer=normalizer,
+            text_key=text_key,
+            normalized_text_key=normalized_text_key,
+            lower_case=lower_case,
+            lower_case_norm=lower_case_norm,
+        )
+        for entry in tqdm(entries)
+    )
+
+    write_manifest(output_path=output_manifest_path, target_manifest=output_entries, ensure_ascii=False)
+
+
+if __name__ == "__main__":
+    main()

scripts/dataset_processing/tts/resynthesize_dataset.py

@@ -0,0 +1,252 @@
+# 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.
+
+"""
+This script is a helper for resynthesizing TTS dataset using a pretrained text-to-spectrogram model.
+Goal of resynthesis (as opposed to text-to-speech) is to use the largest amount of ground-truth features from existing speech data.
+For example, for resynthesis we want to have the same pitch and durations instead of ones predicted by the model.
+The results are to be used for some other task: vocoder finetuning, spectrogram enhancer training, etc.
+
+Let's say we have the following toy dataset:
+/dataset/manifest.json
+/dataset/1/foo.wav
+/dataset/2/bar.wav
+/dataset/sup_data/pitch/1_foo.pt
+/dataset/sup_data/pitch/2_bar.pt
+
+manifest.json has two entries for "/dataset/1/foo.wav" and "/dataset/2/bar.wav"
+(sup_data folder contains pitch files precomputed during training a FastPitch model on this dataset.)
+(If you lost your sup_data - don't worry, we use TTSDataset class so they would be created on-the-fly)
+
+Our script call is
+$ python scripts/dataset_processing/tts/resynthesize_dataset.py \
+    --model-path ./models/fastpitch/multi_spk/FastPitch--val_loss\=1.4473-epoch\=209.ckpt \
+    --input-json-manifest "/dataset/manifest.json" \
+    --input-sup-data-path "/dataset/sup_data/" \
+    --output-folder "/output/" \
+    --device "cuda:0" \
+    --batch-size 1 \
+    --num-workers 1
+
+Then we get output dataset with following directory structure:
+/output/manifest_mel.json
+/output/mels/foo.npy
+/output/mels/foo_gt.npy
+/output/mels/bar.npy
+/output/mels/bar_gt.npy
+
+/output/manifest_mel.json has the same entries as /dataset/manifest.json but with new fields for spectrograms.
+"mel_filepath" is path to the resynthesized spectrogram .npy, "mel_gt_filepath" is path to ground-truth spectrogram .npy
+
+The output structure is similar to generate_mels.py script for compatibility reasons.
+"""
+
+import argparse
+import itertools
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, Iterable, Iterator, List
+
+import numpy as np
+import torch
+from omegaconf import DictConfig, OmegaConf
+from tqdm import tqdm
+
+from nemo.collections.asr.parts.utils.manifest_utils import read_manifest, write_manifest
+from nemo.collections.tts.models import FastPitchModel
+from nemo.collections.tts.models.base import SpectrogramGenerator
+from nemo.collections.tts.parts.utils.helpers import process_batch, to_device_recursive
+
+
+def chunks(iterable: Iterable, size: int) -> Iterator[List]:
+    # chunks([1, 2, 3, 4, 5], size=2) -> [[1, 2], [3, 4], [5]]
+    # assumes iterable does not have any `None`s
+    args = [iter(iterable)] * size
+    for chunk in itertools.zip_longest(*args, fillvalue=None):
+        chunk = list(item for item in chunk if item is not None)
+        if chunk:
+            yield chunk
+
+
+def load_model(path: Path, device: torch.device) -> SpectrogramGenerator:
+    model = None
+    if path.suffix == ".nemo":
+        model = SpectrogramGenerator.restore_from(path, map_location=device)
+    elif path.suffix == ".ckpt":
+        model = SpectrogramGenerator.load_from_checkpoint(path, map_location=device)
+    else:
+        raise ValueError(f"Unknown checkpoint type {path.suffix} ({path})")
+
+    return model.eval().to(device)
+
+
+@dataclass
+class TTSDatasetResynthesizer:
+    """
+    Reuses internals of a SpectrogramGenerator to resynthesize dataset using ground truth features.
+    Default setup is FastPitch with learned alignment.
+    If your use case requires different setup, you can either contribute to this script or subclass this class.
+    """
+
+    model: SpectrogramGenerator
+    device: torch.device
+
+    @torch.no_grad()
+    def resynthesize_batch(self, batch: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Resynthesizes a single batch.
+        Takes a dict with main data and sup data.
+        Outputs a dict with model outputs.
+        """
+        if not isinstance(self.model, FastPitchModel):
+            raise NotImplementedError(
+                "This script supports only FastPitch. Please implement resynthesizing routine for your desired model."
+            )
+
+        batch = to_device_recursive(batch, self.device)
+
+        mels, mel_lens = self.model.preprocessor(input_signal=batch["audio"], length=batch["audio_lens"])
+
+        reference_audio = batch.get("reference_audio", None)
+        reference_audio_len = batch.get("reference_audio_lens", None)
+        reference_spec, reference_spec_len = None, None
+        if reference_audio is not None:
+            reference_spec, reference_spec_len = self.model.preprocessor(
+                input_signal=reference_audio, length=reference_audio_len
+            )
+
+        outputs_tuple = self.model.forward(
+            text=batch["text"],
+            durs=None,
+            pitch=batch["pitch"],
+            speaker=batch.get("speaker"),
+            pace=1.0,
+            spec=mels,
+            attn_prior=batch.get("attn_prior"),
+            mel_lens=mel_lens,
+            input_lens=batch["text_lens"],
+            reference_spec=reference_spec,
+            reference_spec_lens=reference_spec_len,
+        )
+        names = self.model.fastpitch.output_types.keys()
+        return {"spec": mels, "mel_lens": mel_lens, **dict(zip(names, outputs_tuple))}
+
+    def resynthesized_batches(self) -> Iterator[Dict[str, Any]]:
+        """
+        Returns a generator of resynthesized batches.
+        Each returned batch is a dict containing main data, sup data, and model output
+        """
+        self.model.setup_training_data(self.model._cfg["train_ds"])
+
+        for batch_tuple in iter(self.model._train_dl):
+            batch = process_batch(batch_tuple, sup_data_types_set=self.model._train_dl.dataset.sup_data_types)
+            yield self.resynthesize_batch(batch)
+
+
+def prepare_paired_mel_spectrograms(
+    model_path: Path,
+    input_json_manifest: Path,
+    input_sup_data_path: Path,
+    output_folder: Path,
+    device: torch.device,
+    batch_size: int,
+    num_workers: int,
+):
+    model = load_model(model_path, device)
+
+    dataset_config_overrides = {
+        "dataset": {
+            "manifest_filepath": str(input_json_manifest.absolute()),
+            "sup_data_path": str(input_sup_data_path.absolute()),
+        },
+        "dataloader_params": {"batch_size": batch_size, "num_workers": num_workers, "shuffle": False},
+    }
+    model._cfg.train_ds = OmegaConf.merge(model._cfg.train_ds, DictConfig(dataset_config_overrides))
+    resynthesizer = TTSDatasetResynthesizer(model, device)
+
+    input_manifest = read_manifest(input_json_manifest)
+
+    output_manifest = []
+    output_json_manifest = output_folder / f"{input_json_manifest.stem}_mel{input_json_manifest.suffix}"
+    output_mels_folder = output_folder / "mels"
+    output_mels_folder.mkdir(exist_ok=True, parents=True)
+    for batch, batch_manifest in tqdm(
+        zip(resynthesizer.resynthesized_batches(), chunks(input_manifest, size=batch_size)), desc="Batch #"
+    ):
+        pred_mels = batch["spect"].cpu()  # key from fastpitch.output_types
+        true_mels = batch["spec"].cpu()  # key from code above
+        mel_lens = batch["mel_lens"].cpu().flatten()  # key from code above
+
+        for i, (manifest_entry, length) in enumerate(zip(batch_manifest, mel_lens.tolist())):
+            print(manifest_entry["audio_filepath"])
+            filename = Path(manifest_entry["audio_filepath"]).stem
+
+            # note that lengths match
+            pred_mel = pred_mels[i, :, :length].clone().numpy()
+            true_mel = true_mels[i, :, :length].clone().numpy()
+
+            pred_mel_path = output_mels_folder / f"{filename}.npy"
+            true_mel_path = output_mels_folder / f"{filename}_gt.npy"
+
+            np.save(pred_mel_path, pred_mel)
+            np.save(true_mel_path, true_mel)
+
+            new_manifest_entry = {
+                **manifest_entry,
+                "mel_filepath": str(pred_mel_path),
+                "mel_gt_filepath": str(true_mel_path),
+            }
+            output_manifest.append(new_manifest_entry)
+
+    write_manifest(output_json_manifest, output_manifest, ensure_ascii=False)
+
+
+def argument_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        description="Resynthesize TTS dataset using a pretrained text-to-spectrogram model",
+    )
+    parser.add_argument(
+        "--model-path",
+        required=True,
+        type=Path,
+        help="Path to a checkpoint (either .nemo or .ckpt)",
+    )
+    parser.add_argument(
+        "--input-json-manifest",
+        required=True,
+        type=Path,
+        help="Path to the input JSON manifest",
+    )
+    parser.add_argument(
+        "--input-sup-data-path",
+        required=True,
+        type=Path,
+        help="sup_data_path for the JSON manifest",
+    )
+    parser.add_argument(
+        "--output-folder",
+        required=True,
+        type=Path,
+        help="Path to the output folder. Will contain updated manifest and mels/ folder with spectrograms in .npy files",
+    )
+    parser.add_argument("--device", required=True, type=torch.device, help="Device ('cpu', 'cuda:0', ...)")
+    parser.add_argument("--batch-size", required=True, type=int, help="Batch size in the DataLoader")
+    parser.add_argument("--num-workers", required=True, type=int, help="Num workers in the DataLoader")
+    return parser
+
+
+if __name__ == "__main__":
+    arguments = argument_parser().parse_args()
+    prepare_paired_mel_spectrograms(**vars(arguments))

scripts/dataset_processing/tts/sfbilingual/ds_conf/ds_for_fastpitch_align.yaml

@@ -0,0 +1,49 @@
+name: "ds_for_fastpitch_align"
+
+manifest_filepath: "train_manifest.json"
+sup_data_path: "sup_data"
+sup_data_types: [ "align_prior_matrix", "pitch" ]
+phoneme_dict_path: "scripts/tts_dataset_files/zh/24finals/pinyin_dict_nv_22.10.txt"
+
+dataset:
+  _target_: nemo.collections.tts.data.dataset.TTSDataset
+  manifest_filepath: ${manifest_filepath}
+  sample_rate: 22050
+  sup_data_path: ${sup_data_path}
+  sup_data_types: ${sup_data_types}
+  n_fft: 1024
+  win_length: 1024
+  hop_length: 256
+  window: "hann"
+  n_mels: 80
+  lowfreq: 0
+  highfreq: null
+  max_duration: null
+  min_duration: 0.1
+  ignore_file: null
+  trim: true
+  trim_top_db: 50
+  trim_frame_length: 1024
+  trim_hop_length: 256
+  pitch_fmin: 65.40639132514966
+  pitch_fmax: 2093.004522404789
+
+  text_normalizer:
+    _target_: nemo_text_processing.text_normalization.normalize.Normalizer
+    lang: zh
+    input_case: cased
+
+  text_normalizer_call_kwargs:
+    verbose: false
+    punct_pre_process: true
+    punct_post_process: true
+
+  text_tokenizer:
+    _target_: nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.ChinesePhonemesTokenizer
+    punct: true
+    apostrophe: true
+    pad_with_space: true
+    g2p:
+      _target_: nemo.collections.tts.g2p.models.zh_cn_pinyin.ChineseG2p
+      phoneme_dict: ${phoneme_dict_path}
+      word_segmenter: jieba # Only jieba is supported now.

scripts/dataset_processing/tts/sfbilingual/get_data.py

@@ -0,0 +1,137 @@
+# 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.
+
+import argparse
+import json
+import os
+import random
+import subprocess
+from pathlib import Path
+
+import numpy as np
+from opencc import OpenCC
+
+try:
+    from nemo_text_processing.text_normalization.normalize import Normalizer
+except (ImportError, ModuleNotFoundError):
+    raise ModuleNotFoundError(
+        "The package `nemo_text_processing` was not installed in this environment. Please refer to"
+        " https://github.com/NVIDIA/NeMo-text-processing and install this package before using "
+        "this script"
+    )
+
+
+def get_args():
+    parser = argparse.ArgumentParser(
+        description='Prepare SF_bilingual dataset and create manifests with predefined split'
+    )
+
+    parser.add_argument(
+        "--data-root",
+        type=Path,
+        help="where the dataset will reside",
+        default="./DataChinese/sf_bilingual_speech_zh_en_vv1/SF_bilingual/",
+    )
+    parser.add_argument(
+        "--manifests-path", type=Path, help="where the resulting manifests files will reside", default="./"
+    )
+    parser.add_argument("--val-size", default=0.01, type=float, help="eval set split")
+    parser.add_argument("--test-size", default=0.01, type=float, help="test set split")
+    parser.add_argument(
+        "--seed-for-ds-split",
+        default=100,
+        type=float,
+        help="Seed for deterministic split of train/dev/test, NVIDIA's default is 100",
+    )
+
+    args = parser.parse_args()
+    return args
+
+
+def __process_transcript(file_path: str):
+    # Create zh-TW to zh-simplify converter
+    cc = OpenCC('t2s')
+    # Create normalizer
+    text_normalizer = Normalizer(
+        lang="zh",
+        input_case="cased",
+        overwrite_cache=True,
+        cache_dir=str(file_path / "cache_dir"),
+    )
+    text_normalizer_call_kwargs = {"punct_pre_process": True, "punct_post_process": True}
+    normalizer_call = lambda x: text_normalizer.normalize(x, **text_normalizer_call_kwargs)
+    entries = []
+    i = 0
+    with open(file_path / "text_SF.txt", encoding="utf-8") as fin:
+        for line in fin:
+            content = line.split()
+            wav_name, text = content[0], "".join(content[1:])
+            wav_name = wav_name.replace(u'\ufeff', '')
+            # WAR: change DL to SF, e.g. real wave file com_SF_ce2727.wav, wav name in text_SF
+            # com_DL_ce2727. It would be fixed through the dataset in the future.
+            wav_name = wav_name.replace('DL', 'SF')
+            wav_file = file_path / "wavs" / (wav_name + ".wav")
+            assert os.path.exists(wav_file), f"{wav_file} not found!"
+            duration = subprocess.check_output(["soxi", "-D", str(wav_file)])
+            simplified_text = cc.convert(text)
+            normalized_text = normalizer_call(simplified_text)
+            entry = {
+                'audio_filepath': os.path.abspath(wav_file),
+                'duration': float(duration),
+                'text': text,
+                'normalized_text': normalized_text,
+            }
+
+            i += 1
+            entries.append(entry)
+    return entries
+
+
+def __process_data(dataset_path, val_size, test_size, seed_for_ds_split, manifests_dir):
+    entries = __process_transcript(dataset_path)
+
+    random.Random(seed_for_ds_split).shuffle(entries)
+
+    train_size = 1.0 - val_size - test_size
+    train_entries, validate_entries, test_entries = np.split(
+        entries, [int(len(entries) * train_size), int(len(entries) * (train_size + val_size))]
+    )
+
+    assert len(train_entries) > 0, "Not enough data for train, val and test"
+
+    def save(p, data):
+        with open(p, 'w') as f:
+            for d in data:
+                f.write(json.dumps(d) + '\n')
+
+    save(manifests_dir / "train_manifest.json", train_entries)
+    save(manifests_dir / "val_manifest.json", validate_entries)
+    save(manifests_dir / "test_manifest.json", test_entries)
+
+
+def main():
+    args = get_args()
+    dataset_root = args.data_root
+    dataset_root.mkdir(parents=True, exist_ok=True)
+    __process_data(
+        dataset_root,
+        args.val_size,
+        args.test_size,
+        args.seed_for_ds_split,
+        args.manifests_path,
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/dataset_processing/tts/thorsten_neutral/ds_conf/ds_for_fastpitch_align.yaml

@@ -0,0 +1,44 @@
+name: "ds_for_fastpitch_align"
+
+manifest_filepath: "train_manifest.json"
+sup_data_path: "sup_data"
+sup_data_types: [ "align_prior_matrix", "pitch" ]
+
+dataset:
+  _target_: nemo.collections.tts.data.dataset.TTSDataset
+  manifest_filepath: ${manifest_filepath}
+  sample_rate: 22050
+  sup_data_path: ${sup_data_path}
+  sup_data_types: ${sup_data_types}
+  n_fft: 1024
+  win_length: 1024
+  hop_length: 256
+  window: "hann"
+  n_mels: 80
+  lowfreq: 0
+  highfreq: null
+  max_duration: null
+  min_duration: 0.1
+  ignore_file: null
+  trim: true
+  trim_top_db: 50
+  trim_frame_length: ${dataset.win_length}
+  trim_hop_length: ${dataset.hop_length}
+  pitch_fmin: 65.40639132514966
+  pitch_fmax: 2093.004522404789
+
+  text_normalizer:
+    _target_: nemo_text_processing.text_normalization.normalize.Normalizer
+    lang: de
+    input_case: cased
+
+  text_normalizer_call_kwargs:
+    verbose: false
+    punct_pre_process: true
+    punct_post_process: true
+
+  text_tokenizer:
+    _target_: nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.GermanCharsTokenizer
+    punct: true
+    apostrophe: true
+    pad_with_space: true

scripts/dataset_processing/tts/thorsten_neutral/get_data.py

@@ -0,0 +1,274 @@
+# 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.
+
+"""
+This script is used to generate JSON manifests for mel-generator model training. The usage is below.
+
+$ python scripts/dataset_processing/tts/thorsten_neutral/get_data.py \
+    --data-root ~/experiments/thorsten_neutral \
+    --manifests-root ~/experiments/thorsten_neutral \
+    --data-version "22_10" \
+    --min-duration 0.1 \
+    --normalize-text
+"""
+
+import argparse
+import json
+import random
+import shutil
+import subprocess
+import urllib.request
+from pathlib import Path
+
+from joblib import Parallel, delayed
+from tqdm import tqdm
+
+try:
+    from nemo_text_processing.text_normalization.normalize import Normalizer
+except (ImportError, ModuleNotFoundError):
+    raise ModuleNotFoundError(
+        "The package `nemo_text_processing` was not installed in this environment. Please refer to"
+        " https://github.com/NVIDIA/NeMo-text-processing and install this package before using "
+        "this script"
+    )
+
+from nemo.utils import logging
+
+# Thorsten Müller published two neural voice datasets, 21.02 and 22.10.
+THORSTEN_NEUTRAL = {
+    "21_02": {
+        "url": "https://zenodo.org/record/5525342/files/thorsten-neutral_v03.tgz?download=1",
+        "dir_name": "thorsten-de_v03",
+        "metadata": ["metadata.csv"],
+    },
+    "22_10": {
+        "url": "https://zenodo.org/record/7265581/files/ThorstenVoice-Dataset_2022.10.zip?download=1",
+        "dir_name": "ThorstenVoice-Dataset_2022.10",
+        "metadata": ["metadata_train.csv", "metadata_dev.csv", "metadata_test.csv"],
+    },
+}
+
+
+def get_args():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        description="Download Thorsten Müller's neutral voice dataset and create manifests with predefined split. "
+        "Thorsten Müller published two neural voice datasets, 21.02 and 22.10, where 22.10 provides better "
+        "audio quality. Please choose one of the two for your TTS models. Details about the dataset are "
+        "in https://github.com/thorstenMueller/Thorsten-Voice.",
+    )
+    parser.add_argument("--data-root", required=True, type=Path, help="where the resulting dataset will reside.")
+    parser.add_argument("--manifests-root", required=True, type=Path, help="where the manifests files will reside.")
+    parser.add_argument("--data-version", default="22_10", choices=["21_02", "22_10"], type=str)
+    parser.add_argument("--min-duration", default=0.1, type=float)
+    parser.add_argument("--max-duration", default=float('inf'), type=float)
+    parser.add_argument("--val-size", default=100, type=int)
+    parser.add_argument("--test-size", default=100, type=int)
+    parser.add_argument(
+        "--num-workers",
+        default=-1,
+        type=int,
+        help="Specify the max number of concurrent Python worker processes. "
+        "If -1 all CPUs are used. If 1 no parallel computing is used.",
+    )
+    parser.add_argument(
+        "--normalize-text",
+        default=False,
+        action='store_true',
+        help="Normalize original text and add a new entry 'normalized_text' to .json file if True.",
+    )
+    parser.add_argument(
+        "--seed-for-ds-split",
+        default=100,
+        type=float,
+        help="Seed for deterministic split of train/dev/test, NVIDIA's default is 100.",
+    )
+    args = parser.parse_args()
+    return args
+
+
+def __maybe_download_file(source_url, destination_path):
+    if not destination_path.exists():
+        logging.info(f"Downloading data: {source_url} --> {destination_path}")
+        tmp_file_path = destination_path.with_suffix(".tmp")
+        urllib.request.urlretrieve(source_url, filename=tmp_file_path)
+        tmp_file_path.rename(destination_path)
+    else:
+        logging.info(f"Skipped downloading data because it exists: {destination_path}")
+
+
+def __extract_file(filepath, data_dir):
+    logging.info(f"Unzipping data: {filepath} --> {data_dir}")
+    shutil.unpack_archive(filepath, data_dir)
+    logging.info(f"Unzipping data is complete: {filepath}.")
+
+
+def __save_json(json_file, dict_list):
+    logging.info(f"Saving JSON split to {json_file}.")
+    with open(json_file, "w") as f:
+        for d in dict_list:
+            f.write(json.dumps(d) + "\n")
+
+
+def __text_normalization(json_file, num_workers=-1):
+    text_normalizer_call_kwargs = {
+        "punct_pre_process": True,
+        "punct_post_process": True,
+    }
+    text_normalizer = Normalizer(
+        lang="de",
+        input_case="cased",
+        overwrite_cache=True,
+        cache_dir=str(json_file.parent / "cache_dir"),
+    )
+
+    def normalizer_call(x):
+        return text_normalizer.normalize(x, **text_normalizer_call_kwargs)
+
+    def add_normalized_text(line_dict):
+        normalized_text = normalizer_call(line_dict["text"])
+        line_dict.update({"normalized_text": normalized_text})
+        return line_dict
+
+    logging.info(f"Normalizing text for {json_file}.")
+    with open(json_file, 'r', encoding='utf-8') as fjson:
+        lines = fjson.readlines()
+        # Note: you need to verify which backend works well on your cluster.
+        # backend="loky" is fine on multi-core Ubuntu OS; backend="threading" on Slurm.
+        dict_list = Parallel(n_jobs=num_workers)(
+            delayed(add_normalized_text)(json.loads(line)) for line in tqdm(lines)
+        )
+
+    json_file_text_normed = json_file.parent / f"{json_file.stem}_text_normed{json_file.suffix}"
+    with open(json_file_text_normed, 'w', encoding="utf-8") as fjson_norm:
+        for dct in dict_list:
+            fjson_norm.write(json.dumps(dct) + "\n")
+    logging.info(f"Normalizing text is complete: {json_file} --> {json_file_text_normed}")
+
+
+def __process_data(
+    unzipped_dataset_path, metadata, min_duration, max_duration, val_size, test_size, seed_for_ds_split
+):
+    logging.info("Preparing JSON train/val/test splits.")
+
+    entries = list()
+    not_found_wavs = list()
+    wrong_duration_wavs = list()
+
+    for metadata_fname in metadata:
+        meta_file = unzipped_dataset_path / metadata_fname
+        with open(meta_file, 'r') as fmeta:
+            for line in tqdm(fmeta):
+                items = line.strip().split('|')
+                wav_file_stem, text = items[0], items[1]
+                wav_file = unzipped_dataset_path / "wavs" / f"{wav_file_stem}.wav"
+
+                # skip audios if they do not exist.
+                if not wav_file.exists():
+                    not_found_wavs.append(wav_file)
+                    logging.warning(f"Skipping {wav_file}: it is not found.")
+                    continue
+
+                # skip audios if their duration is out of range.
+                duration = subprocess.check_output(["soxi", "-D", str(wav_file)])
+                duration = float(duration)
+                if min_duration <= duration <= max_duration:
+                    entry = {
+                        'audio_filepath': str(wav_file),
+                        'duration': duration,
+                        'text': text,
+                    }
+                    entries.append(entry)
+                elif duration < min_duration:
+                    wrong_duration_wavs.append(wav_file)
+                    logging.warning(f"Skipping {wav_file}: it is too short, less than {min_duration} seconds.")
+                    continue
+                else:
+                    wrong_duration_wavs.append(wav_file)
+                    logging.warning(f"Skipping {wav_file}: it is too long, greater than {max_duration} seconds.")
+                    continue
+
+    random.Random(seed_for_ds_split).shuffle(entries)
+    train_size = len(entries) - val_size - test_size
+    if train_size <= 0:
+        raise ValueError("Not enough data for the train split.")
+
+    logging.info("Preparing JSON train/val/test splits is complete.")
+    train, val, test = (
+        entries[:train_size],
+        entries[train_size : train_size + val_size],
+        entries[train_size + val_size :],
+    )
+
+    return train, val, test, not_found_wavs, wrong_duration_wavs
+
+
+def main():
+    args = get_args()
+    data_root = args.data_root
+    manifests_root = args.manifests_root
+    data_version = args.data_version
+
+    dataset_root = data_root / f"ThorstenVoice-Dataset-{data_version}"
+    dataset_root.mkdir(parents=True, exist_ok=True)
+
+    # download and extract dataset
+    dataset_url = THORSTEN_NEUTRAL[data_version]["url"]
+    zipped_dataset_path = dataset_root / Path(dataset_url).name.split("?")[0]
+    __maybe_download_file(dataset_url, zipped_dataset_path)
+    __extract_file(zipped_dataset_path, dataset_root)
+
+    # generate train/dev/test splits
+    unzipped_dataset_path = dataset_root / THORSTEN_NEUTRAL[data_version]["dir_name"]
+    entries_train, entries_val, entries_test, not_found_wavs, wrong_duration_wavs = __process_data(
+        unzipped_dataset_path=unzipped_dataset_path,
+        metadata=THORSTEN_NEUTRAL[data_version]["metadata"],
+        min_duration=args.min_duration,
+        max_duration=args.max_duration,
+        val_size=args.val_size,
+        test_size=args.test_size,
+        seed_for_ds_split=args.seed_for_ds_split,
+    )
+
+    # save json splits.
+    train_json = manifests_root / "train_manifest.json"
+    val_json = manifests_root / "val_manifest.json"
+    test_json = manifests_root / "test_manifest.json"
+    __save_json(train_json, entries_train)
+    __save_json(val_json, entries_val)
+    __save_json(test_json, entries_test)
+
+    # save skipped audios that are not found into a file.
+    if len(not_found_wavs) > 0:
+        skipped_not_found_file = manifests_root / "skipped_not_found_wavs.list"
+        with open(skipped_not_found_file, "w") as f_notfound:
+            for line in not_found_wavs:
+                f_notfound.write(f"{line}\n")
+
+    # save skipped audios that are too short or too long into a file.
+    if len(wrong_duration_wavs) > 0:
+        skipped_wrong_duration_file = manifests_root / "skipped_wrong_duration_wavs.list"
+        with open(skipped_wrong_duration_file, "w") as f_wrong_dur:
+            for line in wrong_duration_wavs:
+                f_wrong_dur.write(f"{line}\n")
+
+    # normalize text if requested. New json file, train_manifest_text_normed.json, will be generated.
+    if args.normalize_text:
+        __text_normalization(train_json, args.num_workers)
+        __text_normalization(val_json, args.num_workers)
+        __text_normalization(test_json, args.num_workers)
+
+
+if __name__ == "__main__":
+    main()

scripts/installers/Dockerfile.ngramtools

@@ -0,0 +1,30 @@
+#!/usr/bin/env bash
+# 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 script to install KenLM, OpenSeq2Seq decoder, Flashlight decoder, OpenGRM Ngram tool to contaner 
+
+# How to use? Build it from NeMo root folder:
+# 1. git clone https://github.com/NVIDIA/NeMo.git && cd NeMo
+# 2. DOCKER_BUILDKIT=1 docker build -t nemo:23.03.1 -f ./scripts/installers/Dockerfile.ngramtools .
+
+from nvcr.io/nvidia/nemo:23.03
+
+WORKDIR /workspace/nemo
+
+COPY scripts/. /workspace/nemo/scripts/
+
+RUN /bin/bash scripts/asr_language_modeling/ngram_lm/install_beamsearch_decoders.sh
+
+RUN /bin/bash scripts/installers/install_opengrm.sh

scripts/installers/Dockerfile.speech_translation_vllm

@@ -0,0 +1,48 @@
+# 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 Dockerfile installs the necessary dependencies for streaming speech translation with vLLM.
+
+ARG BASE_IMAGE=pytorch/pytorch:2.9.0-cuda12.8-cudnn9-runtime
+
+FROM ${BASE_IMAGE}
+
+ARG GIT_URL="https://github.com/NVIDIA-NeMo/NeMo.git"
+ARG CHECKOUT="main"
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+# Install dependencies
+RUN apt-get update && \
+  apt-get upgrade -y && \
+  apt-get install -y \
+  libtool libltdl-dev \
+  libsndfile1 sox \
+  libfreetype6 \
+  libsox-fmt-all \
+  swig \
+  ffmpeg \
+  libavdevice-dev \
+  git \
+  build-essential \
+  cmake && \
+  rm -rf /var/lib/apt/lists/*
+
+# Install NeMo
+RUN git clone --recursive ${GIT_URL} nemo_dir && cd nemo_dir && \
+  git checkout $CHECKOUT && \
+  pip install ".[all]"
+
+# Install vLLM for streaming speech translation
+RUN pip install vllm==0.12.0

scripts/installers/install_ais_cli_latest.sh

@@ -0,0 +1,30 @@
+#!/bin/bash
+
+# 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.
+
+echo "Install latest AIS CLI"
+AIS_CLI_URL=https://github.com/NVIDIA/aistore/releases/latest/download/ais-linux-amd64.tar.gz
+
+echo "Download AIS CLI from ${AIS_CLI_URL}"
+curl -LO ${AIS_CLI_URL}
+
+echo "Extract"
+tar -xzvf ais-linux-amd64.tar.gz
+
+echo "Move to /usr/local/bin/"
+mv ./ais /usr/local/bin/. 
+
+echo "Cleanup"
+rm ais-linux-amd64.tar.gz

scripts/installers/install_graphviz.sh

@@ -0,0 +1,51 @@
+#!/bin/bash
+
+# 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.
+
+DOCKER=false
+MAYBE_SUDO=""
+
+GRAPHVIZ_REPO=https://gitlab.com/graphviz/graphviz.git
+GRAPHVIZ_LATEST_RELEASE=9.0.0
+GRAPHVIZ_PY_REPO=https://github.com/xflr6/graphviz
+GRAPHVIZ_PY_LATEST_RELEASE=448d1a0  # Temporary fix until 0.20.2
+
+if [[ $* == *--docker* ]]; then
+    echo "Docker installation"
+    DOCKER=true
+else
+    echo "Local installation"
+    if [[ $(sudo -n -v 2) ]]; then
+        MAYBE_SUDO="sudo"
+    else
+        echo "No sudo detected"
+    fi
+fi
+
+{
+    $MAYBE_SUDO apt-get update
+    $MAYBE_SUDO apt-get remove -y graphviz
+    pip uninstall -y graphviz
+    if [[ $DOCKER == false ]]; then
+        $MAYBE_SUDO apt-get install -y libtool libltdl-dev automake autoconf bison flex tcl \
+            ghostscript libgd-dev fontconfig libcairo2-dev libpango1.0-dev libgts-dev
+    fi
+    git clone ${GRAPHVIZ_REPO} -b ${GRAPHVIZ_LATEST_RELEASE} && cd graphviz
+    ./autogen.sh && ./configure --disable-python --disable-perl
+    $MAYBE_SUDO make -j && $MAYBE_SUDO make install
+    cd .. && $MAYBE_SUDO rm -rf graphviz
+    pip install -v "git+${GRAPHVIZ_PY_REPO}@${GRAPHVIZ_PY_LATEST_RELEASE}#egg=graphviz"
+} || { echo "graphviz installed with errors! Please check installation manually."; exit 1; }
+echo "graphviz (re-) installed successfully!"

scripts/installers/install_k2.sh

@@ -0,0 +1,29 @@
+#!/bin/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.
+
+K2_REPO=https://github.com/k2-fsa/k2
+LATEST_RELEASE=0a09f67 # fix for PyTorch 2.6.0
+# uncomment the following line after the next k2 version is released (>1.24.4)
+#LATEST_RELEASE=$(git -c 'versionsort.suffix=-' \
+#    ls-remote --exit-code --refs --sort='version:refname' --tags ${K2_REPO} '*.*' \
+#    | tail --lines=1 \
+#    | cut -d '/' -f 3)
+# "cut --delimiter '/' --fields 3" doesn't work on macOS, use "-d ... -f ..." instead
+
+pip install wheel setuptools cmake
+K2_MAKE_ARGS="-j" pip install -v "git+${K2_REPO}@${LATEST_RELEASE}#egg=k2" || { echo "k2 could not be installed!"; exit 1; }
+python3 -m k2.version > /dev/null || { echo "k2 installed with errors! Please check installation manually."; exit 1; }
+echo "k2 installed successfully!"

scripts/installers/install_opengrm.sh

@@ -0,0 +1,32 @@
+#!/bin/bash
+
+# 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 install OpenFST and Ngram tools from OpenGRM library
+# Optionally, you can specify a path where to install it as a first positional argument: scripts/installers/install_opengrm.sh /path/to/install/openfst .
+# Alternatively, in the Linux Debian you can use: sudo apt install libngram-tools
+
+DECODERS_PATH=/workspace/nemo/decoders  # Path to decoders folder: /workspace/nemo/decoders if you use NeMo/Dockerfile
+if [ "$#" -eq 1 ]
+then
+  DECODERS_PATH=$1
+fi
+cd $DECODERS_PATH
+
+# Install OpenGrm OpenFST
+wget https://www.openfst.org/twiki/pub/FST/FstDownload/openfst-1.8.2.tar.gz && tar xvzf openfst-1.8.2.tar.gz && cd openfst-1.8.2 && ./configure --enable-grm && make -j4 && make -j4 install && cd ..
+
+# Install OpenGrm Ngram
+OPENFSTPREFIX=$DECODERS_PATH/openfst-1.8.2/src &&  wget https://www.opengrm.org/twiki/pub/GRM/NGramDownload/ngram-1.3.14.tar.gz && tar xvzf ngram-1.3.14.tar.gz && cd ngram-1.3.14 && LDFLAGS="-L${OPENFSTPREFIX}/lib" CXXFLAGS="-I${OPENFSTPREFIX}/include" ./configure --prefix ${OPENFSTPREFIX} && make -j4 && make -j4 install && cd ..

scripts/installers/install_riva_decoder.sh

@@ -0,0 +1,17 @@
+#!/bin/bash
+
+# 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.
+
+pip install kaldifst kaldilm riva-asrlib-decoder

scripts/installers/setup_os2s_decoders.py

@@ -0,0 +1,144 @@
+# 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.
+
+"""Script to build and install decoder package.
+
+It is used by scripts/asr_language_modeling/ngram_lm/install_beamsearch_decoders.sh to install
+KenLM and OpenSeq2Seq decoder.
+
+You can set the order of KenLM model by changing -DKENLM_MAX_ORDER=10 argument.
+"""
+from __future__ import absolute_import, division, print_function
+
+import argparse
+import distutils.ccompiler
+import glob
+import multiprocessing.pool
+import os
+import platform
+import sys
+
+from setuptools import Extension, distutils, setup
+
+parser = argparse.ArgumentParser(description=__doc__)
+parser.add_argument(
+    "--num_processes", default=1, type=int, help="Number of cpu processes to build package. (default: %(default)d)"
+)
+args = parser.parse_known_args()
+
+# reconstruct sys.argv to pass to setup below
+sys.argv = [sys.argv[0]] + args[1]
+
+
+# monkey-patch for parallel compilation
+# See: https://stackoverflow.com/a/13176803
+def parallelCCompile(
+    self,
+    sources,
+    output_dir=None,
+    macros=None,
+    include_dirs=None,
+    debug=0,
+    extra_preargs=None,
+    extra_postargs=None,
+    depends=None,
+):
+    # those lines are copied from distutils.ccompiler.CCompiler directly
+    macros, objects, extra_postargs, pp_opts, build = self._setup_compile(
+        output_dir, macros, include_dirs, sources, depends, extra_postargs
+    )
+    cc_args = self._get_cc_args(pp_opts, debug, extra_preargs)
+
+    # parallel code
+    def _single_compile(obj):
+        try:
+            src, ext = build[obj]
+        except KeyError:
+            return
+        self._compile(obj, src, ext, cc_args, extra_postargs, pp_opts)
+
+    # convert to list, imap is evaluated on-demand
+    thread_pool = multiprocessing.pool.ThreadPool(args[0].num_processes)
+    list(thread_pool.imap(_single_compile, objects))
+    return objects
+
+
+def compile_test(header, library):
+    dummy_path = os.path.join(os.path.dirname(__file__), "dummy")
+    command = (
+        "bash -c \"g++ -include "
+        + header
+        + " -l"
+        + library
+        + " -x c++ - <<<'int main() {}' -o "
+        + dummy_path
+        + " >/dev/null 2>/dev/null && rm "
+        + dummy_path
+        + " 2>/dev/null\""
+    )
+    return os.system(command) == 0
+
+
+# hack compile to support parallel compiling
+distutils.ccompiler.CCompiler.compile = parallelCCompile
+
+FILES = glob.glob('kenlm/util/*.cc') + glob.glob('kenlm/lm/*.cc') + glob.glob('kenlm/util/double-conversion/*.cc')
+
+FILES += glob.glob('openfst-1.6.3/src/lib/*.cc')
+
+FILES = [fn for fn in FILES if not (fn.endswith('main.cc') or fn.endswith('test.cc') or fn.endswith('unittest.cc'))]
+
+LIBS = ['stdc++']
+if platform.system() != 'Darwin':
+    LIBS.append('rt')
+
+ARGS = ['-O3', '-DNDEBUG', '-DKENLM_MAX_ORDER=10', '-std=c++11']
+
+if compile_test('zlib.h', 'z'):
+    ARGS.append('-DHAVE_ZLIB')
+    LIBS.append('z')
+
+if compile_test('bzlib.h', 'bz2'):
+    ARGS.append('-DHAVE_BZLIB')
+    LIBS.append('bz2')
+
+if compile_test('lzma.h', 'lzma'):
+    ARGS.append('-DHAVE_XZLIB')
+    LIBS.append('lzma')
+
+os.system('swig -python -c++ ./decoders.i')
+
+decoders_module = [
+    Extension(
+        name='_swig_decoders',
+        sources=FILES + glob.glob('*.cxx') + glob.glob('*.cpp'),
+        language='c++',
+        include_dirs=[
+            '.',
+            'kenlm',
+            'openfst-1.6.3/src/include',
+            'ThreadPool',
+        ],
+        libraries=LIBS,
+        extra_compile_args=ARGS,
+    )
+]
+
+setup(
+    name='ctc_decoders',
+    version='1.1',
+    description="""CTC decoders""",
+    ext_modules=decoders_module,
+    py_modules=['ctc_decoders', 'swig_decoders'],
+)

scripts/magpietts/README_magpie_po.md

@@ -0,0 +1,248 @@
+### Offline Preference Alignment (DPO/RPO)
+
+Code: `nemo/collections/tts/models/magpietts_preference_optimization.py`
+
+Preference Alignment (DPO/RPO) involves the following steps
+1) Create a list of text-context pairs for which we will generate preference data.
+2) For each text-context pair generate multiple audios from a base TTS checkpoint and calculate metrics (CER/SSIM) for each generation.
+3) Create chosen-rejected pairs from the generated audio.
+4) Finetune the base TTS checkpoint on the chosen-rejected pairs.
+
+#### 1. Create text-context pairs
+We pair a list of challenging texts with context audios from our speech datasets. We add a similar number of regular transcripts our datasets such as LibriTTS paired with random context audios. We also include examples with text contexts. There are other options for generating text-context pairs.
+
+```
+python scripts/magpietts/dpo/create_text_contextpairs.py \
+    --challenging_texts /Data/DPOPairsInputData/challenging_texts_nemollm.txt \
+    --regular_texts_for_audiocontext /Data/DPOPairsInputData/regular_texts_for_audiocontext.txt \
+    --regular_texts_for_textcontext /Data/DPOPairsInputData/regular_texts_for_textcontext.txt \
+    --audio_contexts /Data/DPOPairsInputData/audio_context_list.json \
+    --text_contexts /Data/DPOPairsInputData/text_context_list.txt \
+    --output_manifest /Data/DPOPairsInputData/text_context_pairs_v2.json \
+    --nsamples_perpair 6 ;
+```
+Each pair is repeated `nsamples_perpair` times which specifies how many samples we want to generate for each pair. The output manifest serves as the input for the next step.
+
+We can also explore other options for these text-context pairs as well depending on the task.
+
+#### 2. Generate audios for each text-context pair
+
+Next, we can generate audios from a base TTS checkpoint using the following command. We pass the `audio_dir` as "/" since our text context pairs contains absolute paths. Model config arguments should be modified accordingly to match the base checkpoint architecture. We can run the below command on cluster to generate audios across multiple nodes. This command saves the generated audios along with the metrics for each generation in the `exp_dir`. Each generated audio file is accompanied with a `.json` file that has the CER/SSIM metrics.
+
+
+```
+python examples/tts/magpietts.py \
+--config-name=magpietts_po_inference \
+mode=test \
+batch_size=64 \
++init_from_ptl_ckpt=<PATH_TO_MAGPIE_CKPT> \
+exp_manager.exp_dir=<PO_EXP_DIR> \
++test_ds_meta.textcontextpairs.manifest_path=<OUTPUT_MANIFEST_FROM_STEP_1> \
++test_ds_meta.textcontextpairs.audio_dir="/" \
++test_ds_meta.textcontextpairs.feature_dir="/" \
+model.model_type="decoder_context_tts"  # Change this as needed \
+model.context_duration_min=5.0 \
+model.context_duration_max=5.0 \
+model.use_text_conditioning_encoder=true \
+model.codecmodel_path=<PATH_TO_CODEC_MODEL> \
+model.alignment_loss_scale=0.002 \
+model.prior_scaling_factor=null \
+model.load_cached_codes_if_available=false
+```
+#### 3. Create chosen-rejected pairs from the generations
+
+Next, we go through the generated audio directory and create chosen-rejected pairs.
+
+```
+python scripts/magpietts/dpo/create_preference_pairs.py \
+--input_manifest <OUTPUT_MANIFEST_FROM_STEP_1> \
+--generated_audio_dir <PO_EXP_DIR>/MagpieTTS-PO-Infer/version_0/audio \
+--group_size 6 \
+--cer_threshold 0.01 \
+--val_size 256 ;
+```
+
+`cer_threshold=0.01` means that filter out pairs in which the chosen CER > 0.01.
+
+This command should save train and val manifests for DPO finetuning in the base directory of the generated_audio_dir, that is, `<PO_EXP_DIR>/MagpieTTS-PO-Infer/version_0/manifests/`
+
+#### 4. DPO Finetuning Command
+
+Finally, we perform DPO finetuning using the following command:
+
+```
+python examples/tts/magpietts.py \
+batch_size=4 \
++init_from_ptl_ckpt=<PATH_TO_MAGPIE_CKPT> \
++mode="dpo_train" \
+max_epochs=10 \
+exp_manager.exp_dir=<PO_EXP_DIR> \
+exp_manager.checkpoint_callback_params.always_save_nemo=false \
+model.train_ds.datasets._target_="nemo.collections.tts.data.text_to_speech_dataset.MagpieTTSDatasetDPO" \
+model.validation_ds.datasets._target_="nemo.collections.tts.data.text_to_speech_dataset.MagpieTTSDatasetDPO" \
++train_ds_meta.dpopreftrain.manifest_path="<PO_EXP_DIR>/MagpieTTS-PO-Infer/version_0/manifests/" \
++train_ds_meta.dpopreftrain.audio_dir="/" \
++train_ds_meta.dpopreftrain.feature_dir="/" \
++val_ds_meta.dpoprefval.manifest_path="<PO_EXP_DIR>/MagpieTTS-PO-Infer/version_0/manifests/dpo_val_manifest.json" \
++val_ds_meta.dpoprefval.audio_dir="/" \
++val_ds_meta.dpoprefval.feature_dir="/" \
++model.dpo_beta=0.01 \
++model.dpo_sft_loss_weight=0.0 \
+model.model_type="decoder_context_tts"  # Change this as needed \
+model.context_duration_min=5.0 \
+model.context_duration_max=5.0 \
+model.use_text_conditioning_encoder=true \
+model.codecmodel_path=<PATH_TO_CODEC_MODEL> \
+model.alignment_loss_scale=0.001 \
+model.prior_scaling_factor=null \
+trainer.val_check_interval=200 \
+trainer.log_every_n_steps=10 \
+model.optim.lr=2e-7 \
+~model.optim.sched
+```
+
+Note the following overrides in the above command:
+
+```
++mode="dpo_train" \
+model.train_ds.datasets._target_="nemo.collections.tts.data.text_to_speech_dataset.MagpieTTSDatasetDPO" \
+model.validation_ds.datasets._target_="nemo.collections.tts.data.text_to_speech_dataset.MagpieTTSDatasetDPO" \
+```
+
+Again, our manifest contain absolute paths so we specify `audio_dir="/"` .
+
+### Online Preference Optimization (GRPO)
+
+For online preference optmization, process is much simpler.
+
+1) Create a list of text-context pairs for which we will generate preference data (just one pair for a text-context not repeated).
+We'll use the same process as above, just set `nsamples_perpair 1` in the command.
+```
+python scripts/magpietts/dpo/create_text_contextpairs.py \
+    --challenging_texts /Data/DPOPairsInputData/challenging_texts_nemollm.txt \
+    --regular_texts_for_audiocontext /Data/DPOPairsInputData/regular_texts_for_audiocontext.txt \
+    --regular_texts_for_textcontext /Data/DPOPairsInputData/regular_texts_for_textcontext.txt \
+    --audio_contexts /Data/DPOPairsInputData/audio_context_list.json \
+    --text_contexts /Data/DPOPairsInputData/text_context_list.txt \
+    --output_manifest /Data/DPOPairsInputData/text_context_pairs_v2.json \
+    --nsamples_perpair 1 ;
+```
+
+2. Train using GRPO
+
+To train with GRPO, we use a similar training command as the base model training with a few modifications.
+
+1. We start from a pretrained checkpoint supplied using `+init_from_ptl_ckpt`
+2. We add `+mode="onlinepo_train"` to specify preference optimization based training.
+3. Use a small batch size (bs=2) since we generate `num_generations_per_item` samples per item in the batch and the effective batch size becomes `bs*num_generations_per_item`
+4. The manifest should contain absolute audio paths and the `audio_dir` is specified as "/" in the `train_ds_meta` command.
+5. Use the same model specific overrides as the base model (eg. x-attn heads, is_causal, num_layers, local transformer etc).
+6. Set dropout probs to 0 for all modules - This is especially important if we are not using reference free mode. KL divergence loss becomes very spiky and unstable. Set prob to 0 by `model.decoder.p_dropout=0.0`.
+7. Dont use attention prior or CTC loss during GRPO.
+8. Add the following GRPO specific arguments in the training command.
+
+```
++model.grpo_beta=0.0 \ # Coeffecient for KL loss (if not using reference free mode)
++model.num_generations_per_item=12 \ # 12 samples generated for each item and we compute reward for each
++model.reference_free=true \ # Reference free means we dont use KL loss term. Only optimize for rewards
++model.inference_cfg_prob=0.0 \ # fraction of generations generated using CFG. Can set > 0.0 if we want to optimize for both CFG and non CFG modes of generation
++model.use_local_transformer_prob=0.5 \ # fraction of generations generated using Local Transformer. Set it between 0.0 and 1.0 to improve both LT outputs and non LT outputs for models with an LT
++model.inference_cfg_scale=2.5 \ # CFG scale for samples generated using CFG
++model.cer_reward_weight=0.33 \ # weightage of CER reward in the overall reward
++model.ssim_reward_weight=0.33 \ # weightage of SSIM reward in the overall reward
++model.pesq_reward_weight=0.33 \ # weightage of PESQ reward in the overall reward
++model.use_pesq=true \ # set this is true is using pesq reward
++model.reward_asr_model="whisper" \ # Use whisper only for multilingual settings, dont specify for English
+model.cfg_unconditional_prob=0.0 \ # Set this to 0, we dont want want to drop out unconditional input
++model.inference_topk=2016 \ # Top-K - Not yet sure if we should use topk=80 or not. top_k 2016 just disable top_k in a way.
++model.inference_temperature=0.8 \ # Slightly higher temperature for more variety of generations in preference optimization
++model.use_kv_cache_during_online_po=true \ # Use KV caching while generating samples for GRPO
++model.loss_type="grpo" \ # can be grpo or dr_grpo. grpo works better in my experiments.
++model.scale_rewards=true \ # Whether to divide advantages by std deviation or not (set true for GRPO and false for DR_GRPO)
++model.max_decoder_steps=430 \ # Max steps for generation
+```
+
+9. We also want to validate more frequently during GRPO since each step takes longer. So we add the following args.
+```
+~trainer.check_val_every_n_epoch \
++trainer.val_check_interval=50 \
+```
+
+10. We use a lower learning rate and save the best checkpoints based on lowest CER on our validation set using:
+```
+model.optim.lr=1e-7 \
+~model.optim.sched \
+exp_manager.checkpoint_callback_params.monitor="val_cer_gt" \
+exp_manager.checkpoint_callback_params.mode="min" \
+```
+
+11. Specify precision and gradient clipping as necessary
+```
+trainer.precision=32 \
++trainer.gradient_clip_val=2.5 \
+```
+
+
+Below is a sample training command for multilingual GRPO:
+
+```
+python examples/tts/magpietts.py \
+--config-name=magpietts_multilingual_v1  #TODO(blisc) after updating yamls\
+batch_size=2 \
++init_from_ptl_ckpt=<PATH_TO_MAGPIE_CKPT> \
++mode="onlinepo_train" \
++model.text_tokenizers.chartokenizer._target_=AutoTokenizer  # Change this as needed \
++model.text_tokenizers.chartokenizer.pretrained_model="google/byt5-small" # Change this as needed \
+max_epochs=20 \
+exp_manager.exp_dir=<GRPO_EXP_DIR> \
++exp_manager.version=0 \
+exp_manager.checkpoint_callback_params.always_save_nemo=false \
++train_ds_meta.dpopreftrain.manifest_path=<TRAIN_MANIFEST_FROM_STEP_1> \
++train_ds_meta.dpopreftrain.audio_dir="/" \
++train_ds_meta.dpopreftrain.feature_dir="/" \
++train_ds_meta.dpopreftrain.tokenizer_names="[chartokenizer]"  #Change this as needed \
++val_ds_meta.dpoprefval.manifest_path=<VAL_MANIFEST_FROM_STEP_1> \
++val_ds_meta.dpoprefval.audio_dir="/" \
++val_ds_meta.dpoprefval.feature_dir="/" \
++val_ds_meta.dpoprefval.tokenizer_names="[chartokenizer]" #Change this as needed \
++model.grpo_beta=0.0 \
++model.num_generations_per_item=12 \
++model.reference_free=true \
++model.inference_cfg_prob=0.0 \
++model.inference_cfg_scale=2.5 \
++model.cer_reward_weight=0.33 \
++model.ssim_reward_weight=0.33 \
++model.pesq_reward_weight=0.33 \
++model.use_pesq=true \
++model.reward_asr_model="whisper" \
+model.cfg_unconditional_prob=0.0 \
++model.inference_topk=2016 \
++model.inference_temperature=0.8 \
++model.use_kv_cache_during_online_po=true \
++model.loss_type="grpo" \
++model.scale_rewards=true \
++model.max_decoder_steps=430 \
+model.model_type="decoder_context_tts" #Change this as needed \
+model.context_duration_min=5.0 \
+model.context_duration_max=5.0 \
+model.decoder.p_dropout=0.0 \
+model.encoder.p_dropout=0.0 \
+model.local_transformer_type="autoregressive" #Change this as needed \
+model.local_transformer_n_layers=1 #Change this as needed \
+model.local_transformer_n_heads=1 #Change this as needed \
+model.local_transformer_hidden_dim=256 #Change this as needed \
+model.use_text_conditioning_encoder=true #Change this as needed \
+model.codecmodel_path=<PATH_TO_CODEC_MODEL> \
+model.alignment_loss_scale=0.0 \
+model.prior_scaling_factor=null \
+~trainer.check_val_every_n_epoch \
++trainer.val_check_interval=50 \
+trainer.log_every_n_steps=10 \
+model.optim.lr=1e-7 \
+~model.optim.sched \
+exp_manager.checkpoint_callback_params.monitor="val_cer_gt" \
+exp_manager.checkpoint_callback_params.mode="min" \
+trainer.precision=32 \
++trainer.gradient_clip_val=2.5 \
+```
+

scripts/magpietts/create_lhotse_shar_from_nemo_manifest.py

@@ -0,0 +1,561 @@
+# Copyright (c) 2025, 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 requires the following updates to lhotse: add `shard_offset` in lhotse's writers.
+$ pip install git+https://github.com/lhotse-speech/lhotse.git@883c24b5f6cdc4bbc73e89186e99f7907262b59c
+
+Example of manifest:
+    {
+        "audio_filepath": "train-clean-360/4098/11547/4098_11547_000032_000000.wav",
+        "text": "\"Isn't it?\" queried Theo.",
+        "speaker": "| Language:en Dataset:LibriTTS Speaker:4098 |",
+        "chapter_id": "11547",
+        "utter_id": "000032_000000",
+        "duration": 1.9700416666666667,
+        "normalized_text": "\"Isn't it?\" queried Theo.",
+        "context_speaker_similarity": 0.7800518870353699,
+        "context_audio_filepath": "train-clean-360/4098/11547/4098_11547_000031_000001.wav",
+        "context_audio_duration": 9.45
+    }
+
+Example usage:
+    python scripts/magpietts/create_lhotse_shar_from_nemo_manifest.py \
+        --manifest-path ${MANIFEST} \
+        --audio-base-dir ${AUDIO_BASE_DIR} \
+        --output-dir ${OUTPUT_DIR} \
+        --num-jobs ${NUM_JOBS} \
+        --processing-chunk-size ${CHUNK_SIZE} \
+        --audio-format ${AUDIO_FORMAT} \
+        --log-level ${LOG_LEVEL} \
+        --shuffle \
+        --shuffle-seed 42 \
+    2>&1 | tee ./log/create_lhotse_shar_from_nemo_manifest.stdout
+
+Expected output:
+    $ tree ${OUTPUT_DIR}
+    ${OUTPUT_DIR}/
+        cuts/
+            cuts.000000.jsonl.gz
+            cuts.000001.jsonl.gz
+            ...
+        target_audio/
+            recording.000000.tar
+            recording.000001.tar
+            ...
+        context_audio/
+            recording.000000.tar
+            recording.000001.tar
+            ...
+"""
+
+import argparse
+import itertools
+import logging
+import math
+import os
+import random
+import re
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from functools import partial
+from pathlib import Path
+from typing import Any, Dict, Tuple
+
+from lhotse import AudioSource, MonoCut, Recording, SupervisionSegment, compute_num_samples, fastcopy
+from lhotse.serialization import load_jsonl
+from lhotse.shar.writers import AudioTarWriter, JsonlShardWriter
+from tqdm import tqdm
+
+NEMO_KEYS_NO_NEED_TO_LOG_IN_CUSTOM_FIELDS_FOR_SUPERVISION = [
+    "audio_filepath",
+    "context_audio_filepath",
+    "text",
+    "offset",
+    "duration",
+    "speaker",
+]
+
+
+def to_shar_placeholder(recording: Recording, cut: MonoCut) -> Recording:
+    """this function is borrowed from lhotse.shar.writers.to_shar_placeholder. The only change for Recording instance is to update the id with cut.id."""
+    return fastcopy(
+        recording,
+        id=cut.id,
+        # Creates a single AudioSource out of multiple ones.
+        sources=[AudioSource(type="shar", channels=recording.channel_ids, source="")],
+        # Removes the transform metadata because they were already executed.
+        transforms=None,
+        duration=cut.duration,
+        num_samples=compute_num_samples(cut.duration, recording.sampling_rate),
+    )
+
+
+def check_speaker_format(item: str):
+    """Enforce speaker format like '| Language:en Dataset:HiFiTTS Speaker:9136_other |'"""
+    pattern = r"\| Language:\w+ Dataset:[\w\d\W]+ Speaker:[\w\d\W]+ \|"
+    if not isinstance(item, str):
+        return False
+    return bool(re.match(pattern, item))
+
+
+def get_recording_id(relative_path: str) -> str:
+    """Generate a recording ID from the relative audio path."""
+    return "rec-" + relative_path.rsplit(".", 1)[0].replace("/", "-")
+
+
+def process_manifest_entry(entry: Dict[str, Any], audio_base_dir: Path) -> Tuple[MonoCut, MonoCut] | None:
+    """
+    Processes a single entry from the NeMo manifest to create Lhotse objects.
+
+    Returns:
+        tuple: (target_cut, context_cut) or None if an error occurs.
+    """
+    try:
+        # Required fields
+        target_audio_path_relative = entry.get("audio_filepath")
+        context_audio_path_relative = entry.get("context_audio_filepath")
+        target_audio_duration = entry.get("duration")
+        context_audio_duration = entry.get("context_audio_duration")
+        text = entry.get("text")
+        # observed cases when text is empty while normalized_text is not.
+        if not text or not text.strip():
+            text = entry.get("normalized_text")
+        speaker = entry.get("speaker")
+
+        # Check required fields
+        if not all(
+            [
+                target_audio_path_relative,
+                context_audio_path_relative,
+                target_audio_duration,
+                context_audio_duration,
+                text,
+                speaker,
+            ]
+        ):
+            logging.warning(f"Skipping entry due to missing fields: {entry}")
+            return None
+
+        # Check speaker format
+        if not check_speaker_format(speaker):
+            logging.warning(f"Skipping entry due to incorrect speaker format: {entry}")
+            return None
+
+        target_audio_filepath = audio_base_dir / target_audio_path_relative
+        context_audio_filepath = audio_base_dir / context_audio_path_relative
+
+        if not target_audio_filepath.is_file():
+            logging.warning(
+                f"Skipping entry due to missing target audio file: {target_audio_filepath} from entry: {entry}"
+            )
+            return None
+        if not context_audio_filepath.is_file():
+            logging.warning(
+                f"Skipping entry due to missing context audio file: {context_audio_filepath} from entry: {entry}"
+            )
+            return None
+
+        # Create IDs
+        target_recording_id = get_recording_id(target_audio_path_relative)
+        context_recording_id = get_recording_id(context_audio_path_relative)
+
+        # Create Recordings
+        # TODO: if input is FLAC, then we should set AudioSegment.from_file(int_values=True). Does this applies to lhotse?
+        target_recording = Recording.from_file(target_audio_filepath, recording_id=target_recording_id)
+        context_recording = Recording.from_file(context_audio_filepath, recording_id=context_recording_id)
+
+        # Custom fields exist in manifests, so better to keep them for future usage.
+        custom_fields = {
+            key: val
+            for key, val in entry.items()
+            if key not in NEMO_KEYS_NO_NEED_TO_LOG_IN_CUSTOM_FIELDS_FOR_SUPERVISION
+        }
+        custom_fields["context_recording_id"] = context_recording_id
+
+        # Extract language from speaker string
+        lang_match = re.search(r"Language:(\w+)", speaker)
+        language = lang_match.group(1) if lang_match else None
+
+        # offset in seconds
+        target_offset_in_seconds = entry.get("offset", 0.0)
+        context_offset_in_seconds = entry.get("context_audio_offset", 0.0)
+
+        # Create Supervision for target cut. We constrain one supervision per cut for now.
+        supervision = SupervisionSegment(
+            id=f"sup-{target_recording_id}",
+            recording_id=target_recording_id,
+            start=target_offset_in_seconds,
+            duration=target_audio_duration,  # duration from manifest
+            channel=0,  # only support mono audio for now
+            text=text,
+            language=language,
+            speaker=speaker,
+            custom=custom_fields,
+        )
+
+        # Create target cut
+        target_cut_id = f"cut-{target_recording_id}-{target_offset_in_seconds:.2f}-{target_audio_duration:.2f}"
+        target_cut = MonoCut(
+            id=target_cut_id,
+            start=target_offset_in_seconds,
+            duration=target_audio_duration,
+            channel=0,  # only support mono audio for now
+            recording=target_recording,
+            supervisions=[supervision],
+        )
+        if not math.isclose(target_cut.duration, target_audio_duration, abs_tol=0.1):
+            logging.warning(
+                f"Manifest duration ({target_audio_duration}) differs significantly from cut duration ({target_cut.duration}) for {target_recording_id}. Using cut duration."
+            )
+            target_cut.supervisions[0].duration = target_cut.duration
+
+        # Create context cut. This cut is only used to load segmented audio and would not be stored in the final manifest.
+        context_cut_id = (
+            f"context_cut-{context_recording_id}-{context_offset_in_seconds:.2f}-{context_audio_duration:.2f}"
+        )
+        if context_cut_id.split("-", 1)[1] == target_cut_id.split("-", 1)[1]:
+            logging.warning(f"Context cut has the same recording segment as target cut. Skipping entry: {entry}")
+            return None
+
+        context_cut = MonoCut(
+            id=context_cut_id,
+            start=context_offset_in_seconds,
+            duration=context_audio_duration,
+            channel=0,  # only support mono audio for now
+            recording=context_recording,
+        )
+        return target_cut, context_cut
+
+    except Exception as e:
+        logging.error(f"Skipping entry due to error during metadata processing: {entry}: {e}", exc_info=True)
+        return None
+
+
+def shuffle_jsonl_file(input_path: Path, seed: int = None) -> Path:
+    """
+    Shuffle lines in a JSONL file and write to a shuffled copy.
+
+    Args:
+        input_path: Path to the original JSONL file
+        seed: Random seed for reproducible shuffling
+
+    Returns:
+        Path to the shuffled file
+    """
+    if seed is not None:
+        random.seed(seed)
+
+    logging.info(f"Reading and shuffling manifest entries from {input_path}")
+
+    # Read all lines into memory
+    with open(input_path, 'r', encoding='utf-8') as f:
+        lines = f.readlines()
+
+    logging.info(f"Loaded {len(lines)} entries, now shuffling...")
+
+    # Shuffle the lines
+    random.shuffle(lines)
+
+    # Create output path with "_shuffled" suffix
+    shuffled_path = input_path.parent / f"{input_path.stem}_shuffled{input_path.suffix}"
+
+    # Write shuffled content
+    with open(shuffled_path, 'w', encoding='utf-8') as f:
+        f.writelines(lines)
+
+    logging.info(f"Shuffled manifest written to: {shuffled_path}")
+    return shuffled_path
+
+
+def chunked_iterator(iterable, chunk_size):
+    """Yield successive chunks from iterable."""
+    _it = iter(iterable)
+    while _chunk := tuple(itertools.islice(_it, chunk_size)):
+        yield _chunk
+
+
+def process_and_write_chunk(
+    manifest_chunk_with_idx: Tuple[int, Tuple[Dict[str, Any], ...]],
+    audio_base_dir: Path,
+    output_dir: Path,
+    audio_format: str,
+) -> Dict[str, int]:
+    """
+    Processes a chunk of manifest entries, loads audio, and writes corresponding
+    single shard files for cuts, target audio, and context audio.
+    Designed to be run in a parallel worker process.
+    Loads and writes audio iteratively to save memory.
+
+    Returns a dict containing processing stats like 'processed', 'initial_errors', 'audio_load_errors'.
+    """
+    chunk_idx, manifest_chunk = manifest_chunk_with_idx
+    worker_pid = os.getpid()
+    logging.debug(f"[Worker {worker_pid}, Chunk {chunk_idx}] Starting processing {len(manifest_chunk)} entries.")
+
+    # --- 1. Process manifest entries to get Cut objects ---
+    chunk_metadata = []
+    initial_errors = 0
+    for entry in manifest_chunk:
+        result = process_manifest_entry(entry, audio_base_dir=audio_base_dir)
+        if result is not None:
+            chunk_metadata.append(result)
+        else:
+            initial_errors += 1
+
+    if not chunk_metadata:
+        logging.warning(
+            f"[Worker {worker_pid}, Chunk {chunk_idx}] No valid entries after initial processing. Skipping."
+        )
+        return {"processed": 0, "initial_errors": initial_errors, "audio_load_errors": 0, "write_errors": 0}
+
+    logging.debug(
+        f"[Worker {worker_pid}, Chunk {chunk_idx}] Collected {len(chunk_metadata)} cut pairs after initial processing."
+    )
+
+    # --- 2. Initialize writers and perform iterative load-and-write ---
+    cuts_dir = output_dir / "cuts"
+    target_recordings_dir = output_dir / "target_audio"
+    context_recordings_dir = output_dir / "context_audio"
+
+    cuts_pattern = str(cuts_dir / "cuts.%06d.jsonl.gz")
+    target_rec_pattern = str(target_recordings_dir / "recording.%06d.tar")
+    context_rec_pattern = str(context_recordings_dir / "recording.%06d.tar")
+
+    chunk_processed_count = 0
+    chunk_audio_load_errors = 0  # Errors during audio loading phase for this chunk
+    chunk_write_errors = 0  # Errors during write phase for this chunk
+
+    logging.debug(
+        f"[Worker {worker_pid}, Chunk {chunk_idx}] Initializing writers with offset {chunk_idx} and processing {len(chunk_metadata)} pairs iteratively..."
+    )
+    try:
+        # Specify shard_size with len(chunk_metadata) and shard_offset with chunk_idx, ensuring each chunk is written to a separate shard file.
+        shard_size_for_worker = len(chunk_metadata)
+        with (
+            JsonlShardWriter(
+                pattern=cuts_pattern, shard_size=shard_size_for_worker, shard_offset=chunk_idx
+            ) as cut_writer,
+            AudioTarWriter(
+                pattern=target_rec_pattern,
+                shard_size=shard_size_for_worker,
+                format=audio_format,
+                shard_offset=chunk_idx,
+            ) as target_rec_writer,
+            AudioTarWriter(
+                pattern=context_rec_pattern,
+                shard_size=shard_size_for_worker,
+                format=audio_format,
+                shard_offset=chunk_idx,
+            ) as context_rec_writer,
+        ):
+            # Iterate directly over chunk_metadata
+            for target_cut, context_cut in chunk_metadata:
+                # 1. load target/context audio given the audio offset
+                try:
+                    target_audio = target_cut.load_audio()
+                    context_audio = context_cut.load_audio()
+                except Exception as e:
+                    logging.error(
+                        f"[Worker {worker_pid}, Chunk {chunk_idx}] Error loading target/context audio for cut {target_cut}: {e}",
+                        exc_info=True,
+                    )
+                    chunk_audio_load_errors += 1
+                    continue
+
+                # 2. Write target audio and context audio
+                try:
+                    target_rec_writer.write(
+                        key=target_cut.id,
+                        value=target_audio,
+                        sampling_rate=target_cut.sampling_rate,
+                        manifest=to_shar_placeholder(
+                            target_cut.recording, target_cut
+                        ),  # update manifest.id with target_cut.id that has the audio offset and duration
+                    )
+                    context_rec_writer.write(
+                        key=target_cut.id,  # use target cut id as key for context audio to ensure reference
+                        value=context_audio,
+                        sampling_rate=context_cut.sampling_rate,
+                        manifest=to_shar_placeholder(
+                            context_cut.recording, context_cut
+                        ),  # update manifest.id with context_cut.id that has the audio offset and duration
+                    )
+                except Exception as e:
+                    logging.error(
+                        f"[Worker {worker_pid}, Chunk {chunk_idx}] Error writing target/context audio for target cut {target_cut}: {e}",
+                        exc_info=True,
+                    )
+                    chunk_write_errors += 1
+                    continue
+
+                # 3. write cut metadata
+                try:
+                    cut_writer.write(target_cut)
+                except Exception as e:
+                    logging.error(
+                        f"[Worker {worker_pid}, Chunk {chunk_idx}] Error writing cut metadata for cut {target_cut}: {e}",
+                        exc_info=True,
+                    )
+                    chunk_write_errors += 1
+                    continue
+
+                chunk_processed_count += 1
+
+    except Exception as e:
+        logging.error(
+            f"[Worker {worker_pid}, Chunk {chunk_idx}] CRITICAL error during writer initialization: {e}", exc_info=True
+        )
+        chunk_write_errors = len(chunk_metadata)
+        chunk_processed_count = 0
+
+    # This part is only reached if the main try block completes without critical errors
+    logging.debug(
+        f"[Worker {worker_pid}, Chunk {chunk_idx}] Finished chunk. Processed: {chunk_processed_count}, Audio Load Errors: {chunk_audio_load_errors}, Write Errors: {chunk_write_errors}"
+    )
+
+    return {
+        "processed": chunk_processed_count,
+        "initial_errors": initial_errors,  # Errors from initial metadata processing
+        "audio_load_errors": chunk_audio_load_errors,
+        "write_errors": chunk_write_errors,
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        description="Convert NeMo manifest to sharded Lhotse JSONL/TARs using parallel workers per chunk.",
+    )
+    parser.add_argument("--manifest-path", required=True, type=Path, help="Path to the input NeMo JSON manifest file.")
+    parser.add_argument(
+        "--audio-base-dir", required=True, type=Path, help="Base directory where audio files are located."
+    )
+    parser.add_argument("--output-dir", required=True, type=Path, help="Base directory to save the sharded outputs.")
+    parser.add_argument(
+        "--num-jobs",
+        type=int,
+        default=max(1, os.cpu_count() // 2),
+        help="Number of parallel worker processes (each processing a whole chunk/shard).",
+    )
+    parser.add_argument(
+        "--processing-chunk-size",
+        type=int,
+        default=4000,
+        help="Number of manifest entries per chunk (effectively the items per output shard file).",
+    )
+    parser.add_argument(
+        "--audio-format", type=str, default="flac", help="Audio format for TAR writers (e.g., flac, wav, opus)."
+    )
+    parser.add_argument(
+        "--log-level",
+        type=str,
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
+        help="Set the logging level for the main process and workers.",
+    )
+    parser.add_argument(
+        "--shuffle",
+        action="store_true",
+        help="Shuffle the manifest entries before processing.",
+    )
+    parser.add_argument(
+        "--shuffle-seed",
+        type=int,
+        default=None,
+        help="Random seed for reproducible shuffling (only used if --shuffle is enabled).",
+    )
+
+    args = parser.parse_args()
+
+    # Configure logging based on argument
+    log_level = getattr(logging, args.log_level.upper(), logging.INFO)
+    log_format = '%(asctime)s - PID:%(process)d - %(levelname)s - %(message)s'
+    logging.basicConfig(level=log_level, format=log_format)
+
+    # Ensure output directories exist
+    cuts_dir = args.output_dir / "cuts"
+    target_recordings_dir = args.output_dir / "target_audio"
+    context_recordings_dir = args.output_dir / "context_audio"
+    cuts_dir.mkdir(parents=True, exist_ok=True)
+    target_recordings_dir.mkdir(parents=True, exist_ok=True)
+    context_recordings_dir.mkdir(parents=True, exist_ok=True)
+
+    # Handle shuffling if requested
+    if args.shuffle:
+        logging.info(f"Shuffling manifest entries from: {args.manifest_path}")
+        shuffled_manifest_path = shuffle_jsonl_file(args.manifest_path, seed=args.shuffle_seed)
+        manifest_iterable = load_jsonl(shuffled_manifest_path)
+        logging.info(f"Using shuffled manifest for processing: {shuffled_manifest_path}")
+    else:
+        logging.info(f"Reading NeMo manifest lazily from: {args.manifest_path}")
+        manifest_iterable = load_jsonl(args.manifest_path)
+
+    logging.info(
+        f"Processing manifest in chunks of {args.processing_chunk_size} using {args.num_jobs} parallel workers..."
+    )
+
+    total_processed_count = 0
+    total_initial_errors = 0
+    total_audio_load_errors = 0
+    total_write_errors = 0
+    num_chunks = 0
+
+    worker_func = partial(
+        process_and_write_chunk,
+        audio_base_dir=args.audio_base_dir,
+        output_dir=args.output_dir,
+        audio_format=args.audio_format,
+    )
+
+    with ProcessPoolExecutor(max_workers=args.num_jobs) as executor:
+        # Enumerate chunks to pass index to worker. Each index is the same as the shard_offset.
+        chunk_iterator = enumerate(chunked_iterator(manifest_iterable, args.processing_chunk_size))
+        futures = {
+            executor.submit(worker_func, chunk_with_idx): chunk_with_idx[0] for chunk_with_idx in chunk_iterator
+        }
+        num_chunks = len(futures)
+
+        logging.info(f"Submitted {num_chunks} chunks to workers.")
+
+        for future in tqdm(as_completed(futures), total=num_chunks, desc="Processing Chunks"):
+            chunk_idx = futures[future]
+            try:
+                result = future.result()
+                total_processed_count += result["processed"]
+                total_initial_errors += result["initial_errors"]
+                total_audio_load_errors += result["audio_load_errors"]
+                total_write_errors += result["write_errors"]
+                logging.debug(f"Chunk {chunk_idx} finished with result: {result}")
+            except Exception as e:
+                logging.error(f"Chunk {chunk_idx} failed with exception: {e}", exc_info=True)
+                # Increment error count based on chunk size. Difficult to know precisely. Assume all failed.
+                total_initial_errors += args.processing_chunk_size
+
+    logging.info("=" * 30 + " Processing Summary " + "=" * 30)
+    logging.info(f"Total chunks processed: {num_chunks}")
+    logging.info(f"Successfully processed and wrote data for approximately {total_processed_count} entries.")
+    total_errors = total_initial_errors + total_audio_load_errors + total_write_errors
+    if total_errors > 0:
+        logging.warning(f"Encountered errors/skips in {total_errors} potential entries:")
+        logging.warning(f"  - Initial processing errors/skips: {total_initial_errors}")
+        logging.warning(f"  - Audio loading errors/skips (affecting writes): {total_audio_load_errors}")
+        logging.warning(f"  - Writing errors: {total_write_errors}")
+        logging.warning("Check logs above (use DEBUG level for more details) for specific entry issues.")
+    else:
+        logging.info("No significant errors reported.")
+    logging.info("Manifest creation finished.")
+
+
+if __name__ == "__main__":
+    main()

scripts/magpietts/dpo/create_preference_pairs.py

@@ -0,0 +1,349 @@
+# Copyright (c) 2025, 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.
+import argparse
+import copy
+import json
+import math
+import os
+import random
+
+from tqdm import tqdm
+
+
+def main():
+    """
+    This script creates chosen-rejected pairs for DPO/RPO.
+    We match the manifest records with the generated audio files and metrics.
+    The script then creates a new manifest with chosen-rejected pairs.
+    which is used for training and validation manifest for DPO training.
+
+    Arguments:
+        --input_manifest: Path to the input JSON manifest file containing text/context records.
+        --generated_audio_dir: Directory containing generated audio files and associated metadata.
+        --group_size: Number of records per group used for ranking.
+        --cer_threshold: CER threshold for chosen records. Only records with CER <= threshold are retained.
+        --val_size: Number of validation samples to retain.
+    """
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input_manifest", type=str)
+    parser.add_argument(
+        "--generated_audio_dir",
+        type=str,
+    )
+    parser.add_argument("--group_size", type=int, default=4)
+    parser.add_argument("--cer_threshold", type=float, default=0.02)
+    parser.add_argument(
+        "--min_length_threshold",
+        type=float,
+        default=1.5,
+        help="Minimum length permitted. Set this shorter to allow very short sentences (which can be useful for DPO tuning.",
+    )
+    parser.add_argument("--val_size", type=int, default=64)
+    args = parser.parse_args()
+
+    records = read_manifest(args.input_manifest)
+    audio_files, codec_files, metric_files = find_audio_files(args.generated_audio_dir)
+    assert len(records) <= len(
+        audio_files
+    ), "Mismatch between number of records and number of generated audio files {} vs {}".format(
+        len(records), len(audio_files)
+    )
+
+    for idx, record in tqdm(enumerate(records)):
+        if idx % 100 == 0:
+            print("At idx: ", idx, len(records))
+        record['audio_filepath'] = audio_files[idx]
+        record['target_audio_codes_path'] = codec_files[idx]
+        with open(metric_files[idx], 'r') as f:
+            metrics = json.load(f)
+            record['duration'] = metrics['duration']
+            record['cer_gts'] = metrics['cer_gt']
+            record['wer_gts'] = metrics['wer_gt']
+            record['pred_context_similarity'] = metrics['spk_similarity']
+            record['pred_transcript'] = metrics['pred_transcript']
+            record['gt_transcript'] = metrics['gt_transcript']
+
+    out_manifest_dir = args.generated_audio_dir.replace("/audios", "/manifests")
+    if not os.path.exists(out_manifest_dir):
+        os.makedirs(out_manifest_dir)
+
+    out_manifest = os.path.join(out_manifest_dir, "manifest_with_metrics.json")
+    write_manifest(out_manifest, records)
+
+    group_size = args.group_size
+    val_size = args.val_size
+
+    for num_chosen_per_group in [1, 2]:
+        all_best_records, all_worst_records = create_chosen_rejected_records(records, group_size, num_chosen_per_group)
+        print("Len all_best_records: ", len(all_best_records))
+        print("Len all_worst_records: ", len(all_worst_records))
+        best_records, worst_records = filter_best_and_worst_records(
+            all_best_records, all_worst_records, args.cer_threshold, args.min_length_threshold
+        )
+        print("Len filtered best_records: ", len(best_records))
+        print("Len filtered worst_records: ", len(worst_records))
+        worst_records = normalize_rejected_rewards(worst_records)
+        paired_records = [
+            (best_record, worst_record) for best_record, worst_record in zip(best_records, worst_records)
+        ]
+        random.shuffle(paired_records)
+
+        final_records = []
+        for best_record, worst_record in paired_records:
+            assert best_record['reward'] == 1
+            assert worst_record['reward'] < 1
+            final_records.append(best_record)
+            final_records.append(worst_record)
+
+        final_records_val = final_records[:val_size]
+        final_records_train = final_records[val_size:]
+
+        train_manifest = os.path.join(
+            out_manifest_dir, "dpo_train_manifest_numchosen_per_group_{}.json".format(num_chosen_per_group)
+        )
+        val_manifest = os.path.join(
+            out_manifest_dir, "dpo_val_manifest_numchosen_per_group_{}.json".format(num_chosen_per_group)
+        )
+
+        write_manifest(train_manifest, final_records_train)
+        write_manifest(val_manifest, final_records_val)
+
+
+def read_manifest(manifest_path):
+    with open(manifest_path, 'r') as f:
+        lines = f.readlines()
+        records = []
+        for line in lines:
+            records.append(json.loads(line.strip()))
+    return records
+
+
+def write_manifest(fp, records):
+    with open(fp, "w") as f:
+        for record in records:
+            f.write(json.dumps(record) + "\n")
+    print("Wrote {} records to: {}".format(len(records), fp))
+
+
+def find_audio_files(directory):
+    audio_files = []
+    unique_ranks = {}
+    for root, dirs, files in os.walk(directory):
+        for file in files:
+            if file.endswith(".wav"):
+                rank_num = int(file.split("Rank")[1].split("_")[0])
+                unique_ranks[rank_num] = True
+                audio_num = int(file.split(".wav")[0].split("_")[-1])
+                audio_files.append((rank_num, audio_num, os.path.join(root, file)))
+    ranked_audio_files = []
+    for af in audio_files:
+        rank, num, path = af
+        audio_num = num * len(unique_ranks) + rank
+        ranked_audio_files.append((audio_num, path))
+    ranked_audio_files = sorted(ranked_audio_files, key=lambda x: x[0])
+    ranked_audio_files = [x[1] for x in ranked_audio_files]
+    ranked_codec_files = [f.replace(".wav", "_codes.pt") for f in ranked_audio_files]
+    metric_files = [f.replace(".wav", "_metrics.json") for f in ranked_audio_files]
+
+    return ranked_audio_files, ranked_codec_files, metric_files
+
+
+def pareto_rank(items):
+    """
+    Given a list of (cer, ssim, item_idx), return the list of items
+    sorted by their Pareto rank (rank 1 is best). Items in the same
+    rank are sorted by ascending cer.
+
+    :param items: List of tuples (cer, ssim, item_idx).
+    :return: A list of tuples (rank, cer, ssim, item_idx), sorted first by rank,
+             then by ascending cer within the same rank.
+    """
+
+    # A helper function to check if item A is dominated by item B
+    # A: (cerA, ssimA), B: (cerB, ssimB)
+    def is_dominated(A, B):
+        assert len(A) == 2
+        assert len(B) == 2
+        return (B[0] <= A[0]) and (B[1] >= A[1]) and (B != A)
+        # Equivalently, check at least one strict inequality:
+        # (B[0] < A[0]) or (B[1] > A[1])
+        # can be factored into the condition:
+        # (B[0] <= A[0]) and (B[1] >= A[1]) and (B != A)
+
+    # Make a working copy so we can remove items
+    remaining = items[:]
+
+    ranked_items = []  # Will hold tuples of (rank, cer, ssim, item_idx)
+    current_rank = 1
+
+    while remaining:
+        # Find all non-dominated items in the current set 'remaining'
+        non_dominated = []
+        for i in range(len(remaining)):
+            dominated = False
+            for j in range(len(remaining)):
+                if i != j:
+                    if is_dominated(remaining[i][:2], remaining[j][:2]):
+                        dominated = True
+                        break
+            if not dominated:
+                non_dominated.append(remaining[i])
+
+        # Assign current_rank to all non-dominated items
+        # and remove them from remaining
+        for nd in non_dominated:
+            ranked_items.append((current_rank, nd[0], nd[1], nd[2]))
+            remaining.remove(nd)
+
+        current_rank += 1
+
+    # Now sort the ranked items by (rank asc, cer asc, ssim desc)
+    ranked_items.sort(key=lambda x: (x[0], x[1], -x[2]))
+
+    return ranked_items
+
+
+def standard_normal_cdf(z):
+    """
+    Compute the standard normal cumulative distribution function (CDF) for a given z-score.
+    """
+    return 0.5 * (1 + math.erf(z / math.sqrt(2)))
+
+
+def normalize_rejected_rewards(worst_records):
+    cer_deltas = [record['cer_delta'] for record in worst_records]
+    sim_deltas = [record['sim_delta'] for record in worst_records]
+    cer_mean = sum(cer_deltas) / len(cer_deltas)
+    cer_std = math.sqrt(sum([(d - cer_mean) ** 2 for d in cer_deltas]) / len(cer_deltas))
+    sim_mean = sum(sim_deltas) / len(sim_deltas)
+    sim_std = math.sqrt(sum([(d - sim_mean) ** 2 for d in sim_deltas]) / len(sim_deltas))
+
+    for record in worst_records:
+        cer_z_score = (record['cer_delta'] - cer_mean) / cer_std
+        sim_z_score = (record['sim_delta'] - sim_mean) / sim_std
+        record['reward'] = 1.0 - (standard_normal_cdf(cer_z_score) + standard_normal_cdf(sim_z_score))  # Range -1 to 1
+
+    return worst_records
+
+
+def create_chosen_rejected_records(records_orig, group_size=6, num_chosen_per_group=1):
+    records = copy.deepcopy(records_orig)
+    assert len(records) % group_size == 0
+    num_groups = len(records) // group_size
+    best_records = []
+    worst_records = []
+    num_skipped = 0
+
+    if num_chosen_per_group == 1:
+        chosen_group_indices = [0]
+        rejected_group_indices = [group_size - 1]
+    elif num_chosen_per_group == 2:
+        chosen_group_indices = [0, 1]
+        rejected_group_indices = [group_size - 1, group_size - 2]
+    else:
+        raise ValueError("num_chosen_per_group must be 1 or 2")
+
+    for gidx in range(num_groups):
+        gsi = gidx * group_size
+        gei = (gidx + 1) * group_size
+        group = records[gsi:gei]
+
+        cer_sim_indices = []
+        skip_group = False
+        for sidx, record in enumerate(group):
+            if record['pred_transcript'] == "<INVALID>":
+                print(f"Skipping group starting at index {gsi} due to invalid entries.")
+                num_skipped += len(group)
+                skip_group = True
+                break
+            cer_sim_indices.append((record['cer_gts'], record['pred_context_similarity'], sidx))
+        if skip_group:
+            continue
+        cer_sim_indices_orig = copy.deepcopy(cer_sim_indices)
+        cer_sim_indices = pareto_rank(cer_sim_indices)
+
+        for cgi in chosen_group_indices:
+            for rji in rejected_group_indices:
+                best_record = group[cer_sim_indices[cgi][3]]
+                worst_record = group[cer_sim_indices[rji][3]]
+                best_record['reward'] = 1
+                reward_delta = (worst_record['cer_gts'] - best_record['cer_gts']) + (
+                    best_record['pred_context_similarity'] - worst_record['pred_context_similarity']
+                )
+                if (
+                    reward_delta <= 0
+                    or worst_record['cer_gts'] < best_record['cer_gts']
+                    or worst_record['pred_context_similarity'] > best_record['pred_context_similarity']
+                ):
+                    print(
+                        "Warning reward_delta is not positive",
+                        reward_delta,
+                        best_record['cer_gts'],
+                        worst_record['cer_gts'],
+                        best_record['pred_context_similarity'],
+                        worst_record['pred_context_similarity'],
+                    )
+                    print(cer_sim_indices_orig)
+                    print(cer_sim_indices)
+                else:
+                    # Never add pairs in which rejected has better CER than chosen or better context similarity
+                    reward_delta = max(0.001, reward_delta)
+                    worst_record['reward'] = 1.0 - reward_delta
+                    worst_record['cer_delta'] = worst_record['cer_gts'] - best_record['cer_gts']
+                    worst_record['sim_delta'] = (
+                        best_record['pred_context_similarity'] - worst_record['pred_context_similarity']
+                    )
+                    best_records.append(best_record)
+                    worst_records.append(worst_record)
+
+    print(f"Skipped {num_skipped} records due to invalid entries.")
+    return best_records, worst_records
+
+
+def filter_best_and_worst_records(best_records, worst_records, cer_threshold=0.02, min_length_threshold=1.5):
+    ridx = 0
+    filtered_best_records = []
+    filtered_worst_records = []
+    best_cer_avg = 0.0
+    worst_cer_avg = 0.0
+    skipped_records = 0
+    while ridx < len(best_records):
+        # print(ridx, len(best_records))
+        best_record = best_records[ridx]
+        if best_record['cer_gts'] < cer_threshold:
+            worst_record = worst_records[ridx]
+            if (worst_record['duration'] > 19.0 or best_record['duration'] > 19.0) or (
+                worst_record['duration'] < min_length_threshold or best_record['duration'] < min_length_threshold
+            ):
+                skipped_records += 1
+                ridx += 1
+                continue
+            assert best_record['cer_gts'] <= worst_record['cer_gts']
+            if worst_record['cer_gts'] == best_record['cer_gts']:
+                assert worst_record['pred_context_similarity'] <= best_record['pred_context_similarity']
+
+            filtered_best_records.append(best_record)
+            filtered_worst_records.append(worst_record)
+            best_cer_avg += best_record['cer_gts']
+            worst_cer_avg += worst_record['cer_gts']
+        ridx += 1
+
+    best_cer_avg /= len(filtered_best_records)
+    worst_cer_avg /= len(filtered_worst_records)
+    print(f"Best CER avg: {best_cer_avg}, Worst CER avg: {worst_cer_avg}")
+    return filtered_best_records, filtered_worst_records
+
+
+if __name__ == "__main__":
+    main()

scripts/magpietts/dpo/create_text_contextpairs.py

@@ -0,0 +1,180 @@
+# Copyright (c) 2025, 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.
+import argparse
+import json
+import random
+
+
+def write_manifest(fp, records):
+    """
+    Writes a list of records to a JSON file, where each record is written as a new line.
+
+    Args:
+        fp (str): File path where the records should be written.
+        records (list): List of records (dictionaries) to write.
+    """
+    with open(fp, "w") as f:
+        for record in records:
+            f.write(json.dumps(record) + "\n")
+    print("Wrote {} records to: {}".format(len(records), fp))
+
+
+def main():
+    """
+    Processes text and audio context data to create text-context pairs.
+    The resulting dataset is saved as a JSON manifest file.
+
+    Example usage:
+    python scripts/magpietts/dpo/create_text_contextpairs.py \
+    --challenging_texts /Data/DPOPairsInputDatav2/challenging_with_short.txt \
+    --regular_texts_for_audiocontext /Data/DPOPairsInputDatav2/regular_texts_for_audiocontext.txt \
+    --regular_texts_for_textcontext /Data/DPOPairsInputDatav2/regular_texts_for_textcontext.txt \
+    --audio_contexts /Data/DPOPairsInputDatav2/audio_context_list.json \
+    --text_contexts /Data/DPOPairsInputDatav2/text_context_list_with_audio.txt \
+    --output_manifest /Data/DPOPairsInputDatav2/grpo_train_with_short.json \
+    --n_audio_contexts_per_challenging_text 2 \
+    --n_text_contexts_per_challenging_text 2 \
+    --n_audio_contexts_per_regular_text 1 \
+    --n_text_contexts_per_regular_text 1 \
+    --nsamples_perpair 1 ;
+    """
+    parser = argparse.ArgumentParser(description='Create text-context pairs for DPO')
+    parser.add_argument("--challenging_texts", type=str, help="Text file containing challenging texts")
+    parser.add_argument("--regular_texts_for_audiocontext", type=str, help="Text file containing regular texts")
+    parser.add_argument("--regular_texts_for_textcontext", type=str, help="Text file containing regular texts")
+    parser.add_argument(
+        "--audio_contexts", type=str, help="Manifest containing audio contexts"
+    )  # This manifest should contain 'context_audio_filepath', 'context_audio_duration' and (optionally) 'context_audio_codes_path'
+    parser.add_argument("--text_contexts", type=str, help="Text file containing text contexts")
+    parser.add_argument("--n_audio_contexts_per_challenging_text", type=int, default=10)
+    parser.add_argument("--n_audio_contexts_per_regular_text", type=int, default=1)
+    parser.add_argument("--n_text_contexts_per_challenging_text", type=int, default=3)
+    parser.add_argument("--n_text_contexts_per_regular_text", type=int, default=1)
+    parser.add_argument("--nsamples_perpair", type=int, default=6)
+    parser.add_argument("--output_manifest", type=str)
+    args = parser.parse_args()
+
+    with open(args.challenging_texts, 'r') as f:
+        challenging_texts = f.readlines()
+        challenging_texts = [text.strip() for text in challenging_texts if text.strip() != '']
+
+    with open(args.regular_texts_for_audiocontext, 'r') as f:
+        regular_texts_for_audiocontext = f.readlines()
+        regular_texts_for_audiocontext = [
+            text.strip() for text in regular_texts_for_audiocontext if text.strip() != ''
+        ]
+
+    with open(args.regular_texts_for_textcontext, 'r') as f:
+        regular_texts_for_textcontext = f.readlines()
+        regular_texts_for_textcontext = [text.strip() for text in regular_texts_for_textcontext if text.strip() != '']
+
+    with open(args.audio_contexts, 'r') as f:
+        audio_contexts = f.readlines()
+        audio_contexts = [json.loads(context.strip()) for context in audio_contexts if context.strip() != '']
+
+    with open(args.text_contexts, 'r') as f:
+        text_contexts = f.readlines()
+        text_contexts = [text for text in text_contexts if text.strip() != '']
+
+    all_records = []
+    for challenging_text in challenging_texts:
+        for _ in range(args.n_audio_contexts_per_challenging_text):
+            audio_context = random.choice(audio_contexts)
+            record = create_audio_context_record(challenging_text, audio_context, 'challenging')
+            all_records.append(record)
+
+        for _ in range(args.n_text_contexts_per_challenging_text):
+            text_context = random.choice(text_contexts)
+            record = create_text_context_record(challenging_text, text_context, 'challenging')
+            all_records.append(record)
+
+    for regular_text in regular_texts_for_audiocontext:
+        for _ in range(args.n_audio_contexts_per_regular_text):
+            audio_context = random.choice(audio_contexts)
+            record = create_audio_context_record(regular_text, audio_context, 'regular')
+            all_records.append(record)
+
+    for regular_text in regular_texts_for_textcontext:
+        for _ in range(args.n_text_contexts_per_regular_text):
+            text_context = random.choice(text_contexts)
+            record = create_text_context_record(regular_text, text_context, 'regular')
+            all_records.append(record)
+
+    random.shuffle(all_records)
+    repeated_records = []
+    for record in all_records:
+        for _ in range(args.nsamples_perpair):
+            repeated_records.append(record)
+
+    write_manifest(args.output_manifest, repeated_records)
+    write_manifest(
+        args.output_manifest.replace(".json", "_tinysubset.json"), repeated_records[: 100 * args.nsamples_perpair]
+    )
+
+
+def create_audio_context_record(text, audio_context, record_type):
+    """
+    Creates a record for a text-context pair with audio context.
+
+    Args:
+        text (str): The main text content.
+        audio_context (dict): Dictionary containing audio context information.
+        record_type (str): Type of record ('challenging' or 'regular').
+
+    Returns:
+        dict: A dictionary representing the audio context record.
+    """
+    record = {
+        'text': text,
+        'duration': 6.0,  # Does not matter, avoids filtering out in DPO,
+        'audio_filepath': audio_context['context_audio_filepath'],
+        'context_audio_filepath': audio_context['context_audio_filepath'],
+        'context_audio_duration': audio_context['context_audio_duration'],
+        'record_type': record_type,  # challenging or regular
+    }
+    if 'context_audio_codes_path' in audio_context:
+        record['context_audio_codes_path'] = audio_context['context_audio_codes_path']
+        record['target_audio_codes_path'] = audio_context['context_audio_codes_path']
+
+    return record
+
+
+def create_text_context_record(text, text_context, record_type):
+    """
+    Creates a record for a text-context pair with text context.
+
+    Args:
+        text (str): The main text content.
+        text_context (str): The associated text context.
+        record_type (str): Type of record ('challenging' or 'regular').
+
+    Returns:
+        dict: A dictionary representing the text context record.
+    """
+    if text_context.endswith("\n"):
+        text_context = text_context[:-1]
+    record = {
+        'text': text,
+        'duration': 6.0,  # Does not matter, avoids filtering out in DPO,
+        'audio_filepath': text_context.split(",")[1],
+        'context_text': text_context.split(",")[0],
+        'record_type': record_type,  # challenging or regular
+    }
+    if text_context.split(",")[-1].endswith(".pt"):
+        record['target_audio_codes_path'] = text_context.split(",")[-1]
+    return record
+
+
+if __name__ == '__main__':
+    main()

scripts/magpietts/extend_lhotse_shards_with_audio_codes.py

@@ -0,0 +1,905 @@
+# Copyright (c) 2025, 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 extends the Lhotse shards with audio codec codes.
+
+Example of input shards:
+    $ tree ${CUTS_DIR}
+    ${CUTS_DIR}/
+        cuts.000000.jsonl.gz
+        cuts.000001.jsonl.gz
+        ...
+
+    $ tree ${TARGET_AUDIO_DIR}
+    ${TARGET_AUDIO_DIR}/
+        recording.000000.tar
+        recording.000001.tar
+        ...
+
+    $ tree ${CONTEXT_AUDIO_DIR}
+    ${CONTEXT_AUDIO_DIR}/
+        recording.000000.tar
+        recording.000001.tar
+        ...
+
+Example usage:
+    export WANDB_API_KEY=${WANDB}
+    python -u ${CODE_DIR}/scripts/magpietts/extend_lhotse_shards_with_audio_codes.py \
+        --cuts-dir ${CUTS_DIR} \
+        --target-audio-dir ${TARGET_AUDIO_DIR} \
+        --context-audio-dir ${CONTEXT_AUDIO_DIR} \
+        --output-dir ${RESULTS} \
+        --codec-model-name ${CODEC_MODEL_NAME} \
+        --codec-model-path ${CODEC_MODEL_PATH} \
+        --codec-frame-rate ${CODEC_FRAME_RATE} \
+        --devices ${DEVICES} \
+        --num-nodes ${NUM_NODES} \
+        --batch-size ${BATCH_SIZE} \
+        --buffer-size ${BUFFER_SIZE} \
+        --wandb-entity ${WANDB_ENTITY} \
+        --wandb-project ${WANDB_PROJECT} \
+        --wandb-name ${WANDB_NAME} \
+        --log-level "DEBUG" \
+    2>&1 | tee ${LOG}/${WANDB_NAME}.stdout
+
+Expected output:
+    $ tree ${RESULTS}
+    ${RESULTS}/
+        21fpsCausalDecoder/
+            target_codes/
+                codes.000000.tar
+                codes.000001.tar
+                ...
+            context_codes/
+                codes.000000.tar
+                codes.000001.tar
+                ...
+"""
+
+import argparse
+import glob
+import logging
+import os
+import re
+import threading
+from collections import defaultdict
+from concurrent.futures import Future, ThreadPoolExecutor
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Union
+
+import lightning.pytorch as pl
+import torch
+import wandb
+from lhotse import CutSet
+from lhotse.array import Array, TemporalArray
+from lhotse.dataset import IterableDatasetWrapper, SimpleCutSampler
+from lhotse.shar.writers.array import ArrayTarWriter
+from lightning.pytorch import Trainer
+from lightning.pytorch.callbacks import BasePredictionWriter
+from lightning.pytorch.loggers import WandbLogger
+from lightning.pytorch.strategies import DDPStrategy
+from torch.utils.data import DataLoader, Dataset
+from tqdm import tqdm
+
+from nemo.collections.tts.models import AudioCodecModel
+
+
+def compute_effective_audio_length(original_audio_tensor: torch.Tensor, samples_per_frame: int) -> int:
+    """Computes the effective length of an audio tensor, padded to be a multiple of samples_per_frame."""
+    original_len = original_audio_tensor.shape[0]
+    effective_len = original_len
+    if samples_per_frame > 0:
+        effective_len = ((original_len + samples_per_frame - 1) // samples_per_frame) * samples_per_frame
+    return effective_len
+
+
+def collate_audio_vectors(
+    audio_list: List[torch.Tensor], audio_lens_list: List[int], padding_value: Union[float, int]
+) -> torch.Tensor:
+    """
+    Collate a list of audio vectors into a single tensor, handling padding for variable lengths.
+    Returns a padded tensor.
+    """
+    assert all(len(t.shape) == 1 for t in audio_list), "Expected only 1-D input tensors."
+    assert len(audio_list) == len(audio_lens_list), "Expected the same number of audio vectors and lengths."
+
+    # Create a padded tensor with the maximum audio length from audio_lens_list, where its max length could be longer than
+    # max length of `audio_list``. For example, `audio_lens_list` could be a multiple of the codec model samples per frame.
+    result = audio_list[0].new_ones(len(audio_lens_list), max(audio_lens_list)) * padding_value
+    for i, t in enumerate(audio_list):
+        result[i, : t.shape[0]] = t
+    return result
+
+
+class AudioPairLhotseDataset(Dataset):
+    """
+    A Lhotse Dataset that processes a batch of MonoCuts (received as a CutSet)
+    containing target and context audio.
+    Designed to be used with a Lhotse sampler yielding CutSet batches.
+    Handles loading audio and collating the batch within __getitem__.
+    """
+
+    def __init__(self, target_sample_rate: int, codec_model_samples_per_frame: int):
+        self.target_sample_rate = target_sample_rate
+        self.codec_model_samples_per_frame = codec_model_samples_per_frame
+
+    def __getitem__(self, cuts: CutSet) -> Optional[Dict[str, Any]]:
+        original_target_audios_list = []
+        effective_target_lengths_list = []
+        original_context_audios_list = []
+        effective_context_lengths_list = []
+        target_cut_ids_list = []
+        shard_indices_list = []
+
+        for cut in cuts:
+            if not cut.has_custom("shard_origin"):
+                err_msg = f"Cut {cut} is missing required key 'shard_origin'."
+                logging.error(err_msg)
+                raise ValueError(err_msg)
+            if not cut.has_custom("context_recording"):
+                err_msg = f"Cut {cut} is missing required key 'context_recording'."
+                logging.error(err_msg)
+                raise ValueError(err_msg)
+
+            # Parse shard index from the custom field, handling potential errors
+            origin_path = cut.custom["shard_origin"]
+            match = re.search(r"cuts\.(\d+)\.jsonl\.gz$", origin_path)
+            if match is None:
+                raise ValueError(f"Could not parse shard index from shard_origin: {origin_path}")
+            shard_idx_origin = int(match.group(1))
+
+            # audio shape: (num_channels (1), num_samples) -> (num_samples)
+            # resample to target sample rate
+            target_audio = torch.from_numpy(cut.recording.resample(self.target_sample_rate).load_audio().squeeze(0))
+            context_audio = torch.from_numpy(
+                cut.context_recording.resample(self.target_sample_rate).load_audio().squeeze(0)
+            )
+            original_target_audios_list.append(target_audio)
+            original_context_audios_list.append(context_audio)
+
+            eff_target_len = compute_effective_audio_length(target_audio, self.codec_model_samples_per_frame)
+            effective_target_lengths_list.append(eff_target_len)
+
+            eff_context_len = compute_effective_audio_length(context_audio, self.codec_model_samples_per_frame)
+            effective_context_lengths_list.append(eff_context_len)
+
+            target_cut_ids_list.append(cut.id)
+            shard_indices_list.append(shard_idx_origin)
+
+        # Ensure lists are not empty before calling collate_audio_vectors.
+        if not original_target_audios_list:
+            err_msg = "AudioPairLhotseDataset.__getitem__ processed an empty CutSet or failed to load any audio data, resulting in an empty audio list."
+            logging.error(err_msg)
+            raise ValueError(err_msg)
+
+        target_audio_padded_batch = collate_audio_vectors(
+            original_target_audios_list, effective_target_lengths_list, padding_value=0.0
+        )
+        context_audio_padded_batch = collate_audio_vectors(
+            original_context_audios_list, effective_context_lengths_list, padding_value=0.0
+        )
+
+        # TODO: is it really necessary to convert lengths to torch.int64? currently applying torch.int32.
+        target_audio_lens_collated = torch.IntTensor(effective_target_lengths_list)
+        context_audio_lens_collated = torch.IntTensor(effective_context_lengths_list)
+
+        return {
+            "target_audios": target_audio_padded_batch,
+            "target_audio_lens": target_audio_lens_collated,
+            "context_audios": context_audio_padded_batch,
+            "context_audio_lens": context_audio_lens_collated,
+            "target_cut_id": target_cut_ids_list,
+            "shard_idx_origin": shard_indices_list,
+        }
+
+
+class CodecExtractor(pl.LightningModule):
+    """
+    LightningModule to extract codec codes. Manages DataLoader creation and
+    distribution via predict_dataloader hook.
+    """
+
+    def __init__(
+        self,
+        model_path: str,
+        cuts_dir: str,
+        target_audio_dir: str,
+        context_audio_dir: str,
+        batch_size: int,
+    ):
+        super().__init__()
+        self.model_path = model_path
+        self.cuts_dir = Path(cuts_dir)
+        self.target_audio_dir = Path(target_audio_dir)
+        self.context_audio_dir = Path(context_audio_dir)
+        self.batch_size = batch_size
+
+        logging.info(f"Initializing `AudioPairLhotseDataset` with model path: {self.model_path}")
+        # load the model. mapping to cpu is to avoid GPU mem spikes when initializing the model
+        self.codec_model = AudioCodecModel.restore_from(restore_path=self.model_path, map_location='cpu', strict=False)
+        self.codec_model.eval()
+        logging.info("Codec model loaded.")
+
+        # Placeholder for the rank-specific list of dataloaders
+        self._rank_dataloaders: Optional[List[DataLoader]] = None
+
+    def predict_dataloader(self) -> List[DataLoader]:
+        """
+        Creates and returns the list of DataLoaders assigned to the current rank.
+        Caches the result to avoid redundant creation.
+
+        This function is called by the Trainer to get the dataloaders for the current rank. This happens after
+        intializing `model.predict()` but before any actual prediction steps (ie. calls to `model.predict_step()`) are executed.
+        """
+        # Return cached dataloaders if already created for this rank
+        if self._rank_dataloaders is not None:
+            return self._rank_dataloaders
+
+        # Determine rank and world size
+        try:
+            # Prefer trainer attributes if available
+            current_global_rank = self.global_rank
+            world_size = self.trainer.world_size
+        except AttributeError:
+            # Fallback to torch.distributed if trainer attributes aren't set yet
+            current_global_rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
+            world_size = torch.distributed.get_world_size() if torch.distributed.is_initialized() else 1
+
+        logging.info(f"[Rank {current_global_rank}/{world_size}] Creating assigned subset of dataloaders...")
+
+        # Find all shard files globally
+        cuts_shard_pattern = str(self.cuts_dir / "cuts.*.jsonl.gz")
+        all_cuts_shard_paths = sorted(glob.glob(cuts_shard_pattern))
+
+        if not all_cuts_shard_paths:
+            msg = f"[Rank {current_global_rank}/{world_size}] No input cut shards found matching pattern: {cuts_shard_pattern}. Cannot proceed."
+            logging.error(msg)
+            raise FileNotFoundError(msg)
+
+        num_total_shards = len(all_cuts_shard_paths)
+
+        # Verify shard indices are contiguous and start from 0 based on filenames (globally)
+        first_idx_str = re.search(r"cuts\.(\d+)\.jsonl\.gz$", all_cuts_shard_paths[0]).group(1)
+        last_idx_str = re.search(r"cuts\.(\d+)\.jsonl\.gz$", all_cuts_shard_paths[-1]).group(1)
+        first_idx = int(first_idx_str)
+        last_idx = int(last_idx_str)
+        expected_last_idx = num_total_shards - 1
+        if first_idx != 0:
+            raise ValueError(f"Expected first shard index to be 0, but found {first_idx} in {all_cuts_shard_paths[0]}")
+        if last_idx != expected_last_idx:
+            raise ValueError(
+                f"Expected last shard index to be {expected_last_idx}, but found {last_idx} in {all_cuts_shard_paths[-1]}"
+            )
+        logging.info(
+            f"[Rank {current_global_rank}/{world_size}] Verified {num_total_shards} total shard files globally, with indices from {first_idx} to {last_idx}."
+        )
+
+        # Calculate the slice of original shard indices assigned to this rank
+        is_distributed = world_size > 1
+        assigned_shard_indices_for_rank = []
+
+        if num_total_shards > 0:
+            if not is_distributed:
+                assigned_shard_indices_for_rank = list(range(num_total_shards))
+                logging.info(
+                    f"[Rank {current_global_rank}/{world_size}] Non-distributed mode. Will process all {num_total_shards} shards."
+                )
+            else:
+                num_per_rank_base = num_total_shards // world_size
+                num_with_extra = num_total_shards % world_size
+
+                if current_global_rank < num_with_extra:
+                    start_shard_offset = current_global_rank * (num_per_rank_base + 1)
+                    num_shards_for_rank = num_per_rank_base + 1
+                else:
+                    # Offset by the shards handled by ranks with an extra one
+                    start_shard_offset = num_with_extra + current_global_rank * num_per_rank_base
+                    num_shards_for_rank = num_per_rank_base
+
+                end_shard_offset = start_shard_offset + num_shards_for_rank
+                assigned_shard_indices_for_rank = list(range(start_shard_offset, end_shard_offset))
+
+                logging.info(
+                    f"[Rank {current_global_rank}/{world_size}] Assigned original shard indices "
+                    f"{start_shard_offset} through {end_shard_offset -1} "
+                    f"({len(assigned_shard_indices_for_rank)} shards)"
+                )
+
+        if not assigned_shard_indices_for_rank:
+            logging.info(
+                f"[Rank {current_global_rank}/{world_size}] No shards assigned to this rank. Returning empty dataloader list. This usually happens when the number of shards is less than the number of ranks."
+            )
+            self._rank_dataloaders = []
+            return []
+
+        # Create DataLoaders only for the shards assigned to this rank
+        dataloaders_for_rank = []
+        for original_shard_idx in tqdm(
+            assigned_shard_indices_for_rank,
+            total=len(assigned_shard_indices_for_rank),
+            desc=f">>> [Rank {current_global_rank}/{world_size}] Creating DataLoaders for its assigned shards",
+        ):
+            logging.debug(f"[Rank {current_global_rank}] Processing original shard {original_shard_idx}...")
+            fields = {
+                "cuts": [str(self.cuts_dir / f"cuts.{original_shard_idx:06d}.jsonl.gz")],
+                "recording": [str(self.target_audio_dir / f"recording.{original_shard_idx:06d}.tar")],
+                "context_recording": [str(self.context_audio_dir / f"recording.{original_shard_idx:06d}.tar")],
+            }
+            # Verify if all files exist
+            if not all(Path(shard_filepaths[0]).is_file() for shard_filepaths in fields.values()):
+                err_msg = f"[Rank {current_global_rank}/{world_size}] Missing one or more files for shard {original_shard_idx}. Files: {fields}"
+                logging.error(err_msg)
+                raise FileNotFoundError(err_msg)
+
+            try:
+                logging.debug(
+                    f"[Rank {current_global_rank}] Loading CutSet for original shard {original_shard_idx}..."
+                )
+                shard_cutset = CutSet.from_shar(fields=fields)
+                logging.debug(f"[Rank {current_global_rank}] Loaded CutSet for original shard {original_shard_idx}.")
+            except Exception as e:
+                logging.critical(
+                    f"[Rank {current_global_rank}/{world_size}] CRITICAL ERROR: Failed to load CutSet from shar for original shard index {original_shard_idx}. \
+                    Files attempted: {fields}. \
+                    Error: {e}",
+                    exc_info=True,
+                )
+                raise
+
+            logging.debug(f"[Rank {current_global_rank}] Creating Sampler for original shard {original_shard_idx}...")
+            # Explicitly set rank=0, world_size=1 to ensure sampler iterates the whole shard_cutset
+            sampler = SimpleCutSampler(
+                shard_cutset, max_cuts=self.batch_size, shuffle=False, drop_last=False, rank=0, world_size=1
+            )
+            logging.debug(f"[Rank {current_global_rank}] Creating Dataset for original shard {original_shard_idx}...")
+            shard_dataset = AudioPairLhotseDataset(
+                target_sample_rate=self.codec_model.sample_rate,
+                codec_model_samples_per_frame=self.codec_model.samples_per_frame,
+            )
+            logging.debug(f"[Rank {current_global_rank}] Wrapping Dataset for original shard {original_shard_idx}...")
+            iterable_dataset = IterableDatasetWrapper(
+                dataset=shard_dataset,
+                sampler=sampler,
+            )
+            logging.debug(
+                f"[Rank {current_global_rank}] Creating DataLoader for original shard {original_shard_idx}..."
+            )
+            dl = DataLoader(
+                dataset=iterable_dataset,
+                batch_size=None,
+                num_workers=1,  # Keep num_workers=1 for `IterableDatasetWrapper + SimpleCutSampler` to avoid duplicate batches.
+                pin_memory=True,
+            )
+            logging.debug(
+                f"[Rank {current_global_rank}] Appending DataLoader for original shard {original_shard_idx}..."
+            )
+            dataloaders_for_rank.append(dl)
+            logging.debug(f"[Rank {current_global_rank}] Finished processing original shard {original_shard_idx}.")
+
+        logging.info(
+            f"[Rank {current_global_rank}/{world_size}] Created {len(dataloaders_for_rank)} DataLoaders for this rank."
+        )
+        # Cache the created dataloaders for this rank
+        self._rank_dataloaders = dataloaders_for_rank
+        return self._rank_dataloaders
+
+    def forward(
+        self,
+        target_audios: torch.Tensor,
+        target_audio_lens: torch.Tensor,
+        context_audios: torch.Tensor,
+        context_audio_lens: torch.Tensor,
+    ) -> Optional[Dict[str, torch.Tensor]]:
+        try:
+            target_audios = target_audios.to(self.device)
+            target_audio_lens = target_audio_lens.to(self.device)
+            context_audios = context_audios.to(self.device)
+            context_audio_lens = context_audio_lens.to(self.device)
+            with torch.inference_mode():
+                target_tokens, target_audios_encoded_len = self.codec_model.encode(
+                    audio=target_audios, audio_len=target_audio_lens
+                )
+                context_tokens, context_audios_encoded_len = self.codec_model.encode(
+                    audio=context_audios, audio_len=context_audio_lens
+                )
+            return {
+                "target_codes": target_tokens.to(dtype=torch.uint16, device="cpu"),
+                "target_codes_lengths": target_audios_encoded_len.to(device="cpu"),
+                "context_codes": context_tokens.to(dtype=torch.uint16, device="cpu"),
+                "context_codes_lengths": context_audios_encoded_len.to(device="cpu"),
+            }
+        except Exception as e:
+            logging.error(
+                f"[Rank {self.global_rank}/{self.world_size}] Error during batched codec encoding: {e}", exc_info=True
+            )
+            raise e
+
+    def predict_step(
+        self, batch: Dict[str, Any], batch_idx: int, dataloader_idx: int = 0
+    ) -> Optional[List[Dict[str, Any]]]:
+        codes_dict = self(
+            target_audios=batch["target_audios"],
+            target_audio_lens=batch["target_audio_lens"],
+            context_audios=batch["context_audios"],
+            context_audio_lens=batch["context_audio_lens"],
+        )
+
+        target_codes_batch = codes_dict["target_codes"]
+        target_codes_lens = codes_dict["target_codes_lengths"]
+        context_codes_batch = codes_dict["context_codes"]
+        context_codes_lens = codes_dict["context_codes_lengths"]
+
+        target_cut_ids = batch["target_cut_id"]
+        shard_indices_in_batch = batch["shard_idx_origin"]
+
+        # The shard_indices list should ideally contain the *same* original index
+        # for all items in a batch, because each DataLoader loads from only one shard.
+        results = []
+        batch_size = batch["target_audios"].shape[0]
+        original_shard_idx = shard_indices_in_batch[0]
+        if not all(idx == original_shard_idx for idx in shard_indices_in_batch):
+            raise ValueError(
+                f"Inconsistent shard indices within batch! Batch Index: {batch_idx}, Dataloader Index: {dataloader_idx}. Indices: {shard_indices_in_batch}."
+            )
+
+        if len(target_cut_ids) != batch_size or target_codes_batch.shape[0] != batch_size:
+            raise ValueError(
+                f"Batch size mismatch after inference! Input IDs: {len(target_cut_ids)}, "
+                f"Input Audio Batch: {batch_size}, Output Codes Batch: {target_codes_batch.shape[0]}. "
+                f"Batch Index: {batch_idx}, Dataloader Index: {dataloader_idx}"
+            )
+
+        for target_cut_id, target_codes, context_codes, target_codes_len, context_codes_len in zip(
+            target_cut_ids, target_codes_batch, context_codes_batch, target_codes_lens, context_codes_lens
+        ):
+            results.append(
+                {
+                    "target_cut_id": target_cut_id,
+                    "shard_idx": original_shard_idx,
+                    "target_codes": target_codes[:, :target_codes_len],
+                    "context_codes": context_codes[:, :context_codes_len],
+                }
+            )
+
+        return results
+
+
+class CodecPredictionWriter(BasePredictionWriter):
+    """
+    Writes codec predictions (target and context codes) to ArrayTarWriter shards asynchronously.
+    Uses a ThreadPoolExecutor with a single worker to serialize writes and closing operations per shard,
+    allowing potential overlap between prediction computation and I/O while closing writers early.
+    """
+
+    def __init__(
+        self,
+        output_dir: str,
+        codec_model_name: str,
+        codec_frame_rate: float,
+    ):
+        super().__init__(write_interval="batch")
+        self.output_dir_base = Path(output_dir)
+        self.codec_model_name = codec_model_name
+        self.codec_frame_rate = codec_frame_rate
+        self.rank: int = -1
+        self.world_size: int = -1
+        self.target_writers: Dict[int, ArrayTarWriter] = {}
+        self.context_writers: Dict[int, ArrayTarWriter] = {}
+        self.target_codes_dir: Optional[Path] = None
+        self.context_codes_dir: Optional[Path] = None
+
+        # Attributes for asynchronous writing and closing
+        self.writer_lock: Optional[threading.Lock] = None
+        self.bg_worker_thread: Optional[ThreadPoolExecutor] = None
+        self.futures_per_shard: Optional[Dict[int, List[Future]]] = None
+        self.closer_futures: Optional[List[Future]] = None  # Futures for the _wait_and_close_worker tasks
+        self.last_processed_shard_idx: int = -1
+
+    def setup(self, trainer: Trainer, pl_module: pl.LightningModule, stage: Optional[str] = None) -> None:
+        self.rank = trainer.global_rank
+        self.world_size = trainer.world_size
+        logging.info(
+            f"[Rank {self.rank}/{self.world_size}] Setting up CodecPredictionWriter for async writing with early close."
+        )
+
+        # Initialize async components
+        self.writer_lock = threading.Lock()
+        # Single worker ensures sequential execution of writes AND closes
+        self.bg_worker_thread = ThreadPoolExecutor(max_workers=1, thread_name_prefix=f'CodecWriterRank{self.rank}')
+        self.futures_per_shard = defaultdict(list)
+        self.closer_futures = []
+        self.last_processed_shard_idx = -1
+
+        # Create directories
+        self.target_codes_dir = self.output_dir_base / self.codec_model_name / "target_codes"
+        self.context_codes_dir = self.output_dir_base / self.codec_model_name / "context_codes"
+        if self.rank == 0:
+            self.target_codes_dir.mkdir(parents=True, exist_ok=True)
+            self.context_codes_dir.mkdir(parents=True, exist_ok=True)
+        if trainer.world_size > 1:
+            torch.distributed.barrier()
+        logging.info(f"[Rank {self.rank}/{self.world_size}] Setup complete. Writers will be created on demand.")
+
+    def _get_or_create_writer(
+        self, writer_dict: Dict[int, ArrayTarWriter], shard_idx: int, base_dir: Path
+    ) -> ArrayTarWriter:
+        # Lock needed as this might be called from main thread while closer task modifies dicts
+        with self.writer_lock:
+            if shard_idx not in writer_dict:
+                output_filename = str(base_dir / f"codes.{shard_idx:06d}.tar")
+                logging.debug(
+                    f"[Rank {self.rank}/{self.world_size}] Creating writer for shard {shard_idx} (Thread-safe check): {output_filename}"
+                )
+                try:
+                    writer = ArrayTarWriter(pattern=output_filename, shard_size=None, compression="numpy")
+                    writer.__enter__()
+                    writer_dict[shard_idx] = writer
+                    logging.info(f"[Rank {self.rank}/{self.world_size}] Created writer for shard {shard_idx}")
+                except Exception as e:
+                    msg = f"[Rank {self.rank}/{self.world_size}] Failed to create writer for shard {shard_idx} (file: {output_filename}): {e}"
+                    logging.error(msg, exc_info=True)
+                    raise ValueError(msg)
+
+            # Return writer even if it might be closed soon by a background task
+            # The background task handles the actual closing.
+            return writer_dict[shard_idx]
+
+    def _write_worker(
+        self,
+        target_cut_id: str,
+        shard_idx: int,
+        target_codes: torch.Tensor,
+        context_codes: torch.Tensor,
+        target_writer: ArrayTarWriter,
+        context_writer: ArrayTarWriter,
+    ):
+        """Worker function executed by the background thread to write a single item."""
+        # Assuming target_writer and context_writer are valid when this task starts
+        try:
+            target_codes_array_manifest = TemporalArray(
+                array=Array(storage_type="shar", storage_path="", storage_key="", shape=list(target_codes.shape)),
+                temporal_dim=-1,
+                frame_shift=1 / self.codec_frame_rate,
+                start=0,
+            )
+            context_codes_array_manifest = TemporalArray(
+                array=Array(storage_type="shar", storage_path="", storage_key="", shape=list(context_codes.shape)),
+                temporal_dim=-1,
+                frame_shift=1 / self.codec_frame_rate,
+                start=0,
+            )
+            target_writer.write(key=target_cut_id, value=target_codes.numpy(), manifest=target_codes_array_manifest)
+            context_writer.write(key=target_cut_id, value=context_codes.numpy(), manifest=context_codes_array_manifest)
+            logging.debug(f"[Worker Rank {self.rank}] Wrote item {target_cut_id} for shard {shard_idx}")
+        except Exception as e:
+            msg = f"[Worker Rank {self.rank}] CRITICAL I/O ERROR writing item {target_cut_id} for shard {shard_idx}: {e}. Writer might be closed prematurely?"
+            logging.error(msg, exc_info=True)
+            raise ValueError(msg)
+
+    def _wait_and_close_worker(self, shard_idx_to_close: int):
+        """Waits for all write tasks of a shard, then closes and removes its writers."""
+        logging.info(f"[Worker Rank {self.rank}] Starting closure process for shard {shard_idx_to_close}")
+        # 1. Retrieve and remove the list of write futures for this shard
+        # Do this early to prevent new futures being added for this closing shard?
+        # No, write_on_batch_end logic prevents submission for old shards.
+        write_futures = self.futures_per_shard.pop(shard_idx_to_close, [])
+
+        # 2. Wait for all write operations for this shard to complete
+        logging.info(
+            f"[Worker Rank {self.rank}] Waiting for {len(write_futures)} write tasks for shard {shard_idx_to_close}..."
+        )
+        processed_write_futures = 0
+        if write_futures:
+            for f in write_futures:
+                try:
+                    f.result()  # Wait for completion
+                    processed_write_futures += 1
+                except Exception as e:
+                    # Write worker already logged this, but log context here
+                    logging.error(
+                        f"[Worker Rank {self.rank}] Exception during write future.result() for shard {shard_idx_to_close}: {e}",
+                        exc_info=False,
+                    )
+            logging.info(
+                f"[Worker Rank {self.rank}] Completed {processed_write_futures}/{len(write_futures)} write tasks for shard {shard_idx_to_close}."
+            )
+        else:
+            logging.warning(
+                f"[Worker Rank {self.rank}] No write futures found to wait for shard {shard_idx_to_close} during close."
+            )
+
+        # 3. Safely remove and close the writers
+        writers_closed_count = 0
+        with self.writer_lock:  # Protect access to the writer dictionaries
+            target_writer = self.target_writers.pop(shard_idx_to_close, None)
+            context_writer = self.context_writers.pop(shard_idx_to_close, None)
+
+        if target_writer:
+            try:
+                target_writer.close()
+                logging.info(f"[Worker Rank {self.rank}] Closed target writer for shard {shard_idx_to_close}.")
+                writers_closed_count += 1
+            except Exception as e:
+                logging.error(
+                    f"[Worker Rank {self.rank}] Error closing target writer for shard {shard_idx_to_close}: {e}",
+                    exc_info=True,
+                )
+        else:
+            logging.warning(
+                f"[Worker Rank {self.rank}] Target writer for shard {shard_idx_to_close} not found during close."
+            )
+
+        if context_writer:
+            try:
+                context_writer.close()
+                logging.info(f"[Worker Rank {self.rank}] Closed context writer for shard {shard_idx_to_close}.")
+                writers_closed_count += 1
+            except Exception as e:
+                logging.error(
+                    f"[Worker Rank {self.rank}] Error closing context writer for shard {shard_idx_to_close}: {e}",
+                    exc_info=True,
+                )
+        else:
+            logging.warning(
+                f"[Worker Rank {self.rank}] Context writer for shard {shard_idx_to_close} not found during close."
+            )
+
+        logging.info(
+            f"[Worker Rank {self.rank}] Finished closure process for shard {shard_idx_to_close}. Closed {writers_closed_count} writers."
+        )
+
+    def write_on_batch_end(
+        self,
+        trainer: Trainer,
+        pl_module: pl.LightningModule,
+        predictions: Optional[List[Dict[str, Any]]],
+        batch_indices: Optional[List[int]],
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int,
+    ) -> None:
+
+        if not predictions:
+            err_msg = f"[Rank {self.rank}/{self.world_size}] Received empty predictions list for batch_idx {batch_idx}, dataloader_idx {dataloader_idx}. Skipping."
+            logging.error(err_msg)
+            raise ValueError(err_msg)
+
+        current_shard_idx = predictions[0]["shard_idx"]
+        if not all(p["shard_idx"] == current_shard_idx for p in predictions):
+            raise ValueError(
+                f"[Rank {self.rank}] Inconsistent shard indices within batch! Batch Index: {batch_idx}, Dataloader Index: {dataloader_idx}."
+            )
+
+        # Check for shard change and submit closer task for the previous shard
+        if current_shard_idx != self.last_processed_shard_idx and self.last_processed_shard_idx != -1:
+            logging.info(
+                f"[Rank {self.rank}] Shard index changed from {self.last_processed_shard_idx} to {current_shard_idx}. "
+                f"Submitting closure task for shard {self.last_processed_shard_idx}."
+            )
+            try:
+                closer_future = self.bg_worker_thread.submit(
+                    self._wait_and_close_worker, self.last_processed_shard_idx
+                )
+                self.closer_futures.append(closer_future)
+            except Exception as e:
+                msg = f"[Rank {self.rank}] Failed to submit closer task for shard {self.last_processed_shard_idx}: {e}"
+                logging.error(msg, exc_info=True)
+                raise ValueError(msg)
+
+        self.last_processed_shard_idx = current_shard_idx
+
+        # Submit write tasks for each item in the current batch
+        for prediction in predictions:
+            try:
+                target_cut_id = prediction["target_cut_id"]
+                shard_idx = prediction["shard_idx"]
+                target_codes = prediction["target_codes"]
+                context_codes = prediction["context_codes"]
+
+                # This needs the lock because the closer task might be removing entries concurrently
+                target_writer = self._get_or_create_writer(self.target_writers, shard_idx, self.target_codes_dir)
+                context_writer = self._get_or_create_writer(self.context_writers, shard_idx, self.context_codes_dir)
+
+                # Submit the writing task
+                write_future = self.bg_worker_thread.submit(
+                    self._write_worker,
+                    target_cut_id,
+                    shard_idx,
+                    target_codes,
+                    context_codes,
+                    target_writer,
+                    context_writer,
+                )
+                self.futures_per_shard[shard_idx].append(write_future)
+                logging.debug(f"[Rank {self.rank}] Submitted write task for item {target_cut_id}, shard {shard_idx}")
+
+            except Exception as e:
+                msg = f"[Rank {self.rank}] Error processing prediction item {prediction.get('target_cut_id', 'UNKNOWN')} from batch {batch_idx}: {e}"
+                logging.error(msg, exc_info=True)
+                raise ValueError(msg)
+
+    def teardown(self, trainer: Trainer, pl_module: pl.LightningModule, stage: Optional[str] = None) -> None:
+        logging.info(
+            f"[Rank {self.rank}/{self.world_size}] Tearing down CodecPredictionWriter. Handling final shard and waiting for closers..."
+        )
+
+        # 1. Submit closer task for the very last processed shard (if any)
+        final_shard_processed = self.last_processed_shard_idx
+        if final_shard_processed != -1 and final_shard_processed in self.futures_per_shard:
+            logging.info(
+                f"[Rank {self.rank}] Submitting final closure task for last processed shard {final_shard_processed}."
+            )
+            try:
+                closer_future = self.bg_worker_thread.submit(self._wait_and_close_worker, final_shard_processed)
+                self.closer_futures.append(closer_future)
+            except Exception as e:
+                msg = f"[Rank {self.rank}] Failed to submit final closer task for shard {final_shard_processed}: {e}"
+                logging.error(msg, exc_info=True)
+                raise ValueError(msg)
+
+        # 2. Wait for all closer tasks to complete
+        num_closer_futures = len(self.closer_futures)
+        logging.info(
+            f"[Rank {self.rank}/{self.world_size}] Waiting for {num_closer_futures} background closer tasks to complete."
+        )
+        processed_closer_futures = 0
+        if self.closer_futures:
+            for future in tqdm(
+                self.closer_futures,
+                total=num_closer_futures,
+                desc=f"[Rank {self.rank}/{self.world_size}] Finalizing Shard Closures",
+                leave=False,
+            ):
+                try:
+                    future.result()  # Wait and check for exceptions from the closer worker
+                    processed_closer_futures += 1
+                except Exception as e:
+                    msg = f"[Rank {self.rank}/{self.world_size}] Exception caught during closer future.result(): {e}"
+                    logging.error(msg, exc_info=True)
+                    raise ValueError(msg)
+
+            logging.info(
+                f"[Rank {self.rank}/{self.world_size}] Completed {processed_closer_futures}/{num_closer_futures} closer tasks."
+            )
+        else:
+            logging.info(f"[Rank {self.rank}/{self.world_size}] No closer tasks were submitted.")
+
+        # 3. Shutdown the executor gracefully (all tasks should be done now)
+        if self.bg_worker_thread:
+            logging.info(f"[Rank {self.rank}/{self.world_size}] Shutting down background worker thread.")
+            self.bg_worker_thread.shutdown(wait=True)
+            self.bg_worker_thread = None
+
+        # 4. Final sanity checks and cleanup
+        remaining_writers = len(self.target_writers) + len(self.context_writers)
+        if remaining_writers > 0:
+            msg = f"[Rank {self.rank}/{self.world_size}] {remaining_writers} writers remain after teardown! This should not happen. Keys: Target {list(self.target_writers.keys())}, Context {list(self.context_writers.keys())}"
+            logging.error(msg)
+            raise ValueError(msg)
+
+        remaining_futures = sum(len(futs) for futs in self.futures_per_shard.values())
+        if remaining_futures > 0:
+            msg = f"[Rank {self.rank}/{self.world_size}] {remaining_futures} write futures remain after teardown! This should not happen. Shards: {list(self.futures_per_shard.keys())}"
+            logging.error(msg)
+            raise ValueError(msg)
+
+        self.target_writers.clear()
+        self.context_writers.clear()
+        self.futures_per_shard.clear()
+        self.closer_futures.clear()
+
+        logging.info(f"[Rank {self.rank}/{self.world_size}] Teardown complete.")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--cuts-dir", type=str, required=True, help="Directory containing input cuts/cuts.*.jsonl.gz shards."
+    )
+    parser.add_argument(
+        "--target-audio-dir", type=str, required=True, help="Directory containing target_audio/recording.*.tar shards."
+    )
+    parser.add_argument(
+        "--context-audio-dir",
+        type=str,
+        required=True,
+        help="Directory containing context_audio/recording.*.tar shards.",
+    )
+    parser.add_argument("--output-dir", type=str, required=True, help="Base directory to save the output code shards.")
+    parser.add_argument(
+        "--codec-model-name",
+        type=str,
+        default="21fpsCausalDecoder",
+        help="Name for codec model (used in output path).",
+    )
+    parser.add_argument(
+        "--codec-model-path", type=str, required=True, help="Path to the NeMo codec model (.nemo file)."
+    )
+    parser.add_argument("--codec-frame-rate", type=float, default=21.5, help="Frame rate for codec model.")
+    parser.add_argument("--devices", type=int, default=-1, help="Number of GPUs per node (-1 for all).")
+    parser.add_argument("--num-nodes", type=int, default=1, help="Number of nodes for distributed processing.")
+    parser.add_argument("--batch-size", type=int, default=32, help="Batch size PER GPU for codec inference.")
+    parser.add_argument(
+        "--buffer-size", type=int, default=256, help="Number of items to buffer before writing to TAR files."
+    )
+    parser.add_argument("--wandb-entity", type=str, default=None, help="Wandb entity.")
+    parser.add_argument("--wandb-project", type=str, default="lhotse_codes_extraction", help="Wandb project.")
+    parser.add_argument("--wandb-name", type=str, default=None, help="Wandb run name.")
+    parser.add_argument(
+        "--log-level",
+        type=str,
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
+        help="Set the logging level.",
+    )
+    args = parser.parse_args()
+
+    log_level_val = getattr(logging, args.log_level.upper(), logging.INFO)
+    log_format = '%(asctime)s - PID:%(process)d - %(levelname)s - %(message)s'
+    logging.basicConfig(level=log_level_val, format=log_format)
+
+    codec_extractor = CodecExtractor(
+        model_path=args.codec_model_path,
+        cuts_dir=args.cuts_dir,
+        target_audio_dir=args.target_audio_dir,
+        context_audio_dir=args.context_audio_dir,
+        batch_size=args.batch_size,
+    )
+
+    pred_writer = CodecPredictionWriter(
+        output_dir=args.output_dir,
+        codec_model_name=args.codec_model_name,
+        codec_frame_rate=args.codec_frame_rate,
+    )
+
+    wandb_logger = None
+    if args.wandb_entity and args.wandb_project:
+        wandb_logger = WandbLogger(
+            project=args.wandb_project,
+            entity=args.wandb_entity,
+            name=args.wandb_name or f"extract_codes_{args.codec_model_name}_{os.path.basename(args.cuts_dir)}",
+            log_model=False,
+        )
+        logging.info(f"Wandb logging enabled to {args.wandb_entity}/{args.wandb_project}")
+
+    strategy = DDPStrategy(find_unused_parameters=False) if torch.cuda.is_available() and args.devices != 1 else "auto"
+    trainer = Trainer(
+        devices=args.devices if torch.cuda.is_available() else 1,
+        num_nodes=args.num_nodes,
+        accelerator="gpu" if torch.cuda.is_available() else "cpu",
+        strategy=strategy,
+        logger=wandb_logger,
+        callbacks=[pred_writer],
+        use_distributed_sampler=False,  # we should disable replacing or wrapping Lhostse CutSampler with a `DistributedSamplerWrapper` since Lhotse's sampler already handles distributed sampling.
+    )
+
+    logging.info(f"Starting prediction with {trainer.world_size} ranks.")
+    trainer.predict(codec_extractor, return_predictions=False)
+    logging.info("Prediction finished.")
+
+    if trainer.is_global_zero and wandb_logger:
+        wandb.finish()
+        logging.info("Wandb run finished.")
+
+
+if __name__ == "__main__":
+    import torch.multiprocessing
+
+    try:
+        torch.multiprocessing.set_start_method('spawn')
+    except RuntimeError:
+        # This exception occurs if the start method has already been set. We can safely ignore it.
+        pass
+    main()

scripts/pseudo_labeling/README.md

@@ -0,0 +1,27 @@
+# 🧠 TopIPL: Iterative Pseudo-Labeling for ASR
+
+TopIPL is an **iterative pseudo-labeling algorithm** for training speech recognition models using both labeled and unlabeled data. It integrates seamlessly into the NeMo ASR pipeline and enables **self-training** across epochs with minimal manual intervention.
+
+## 🚀 Key Features
+
+- ⚙️ Supports **semi-supervised ASR training** with dynamic iterative pseudo-label refinement.
+- 🧪 Designed for large-scale training using both labeled and unlabeled speech data.
+- 🔁 Automatically writes pseudo-labels and updates training configs between iterations.
+
+## 📦 Required Components
+
+TopIPL relies on the following components:
+
+- **[`SDPNeMoRunIPLProcessor`]**  
+  Commands for running IPL are generated and submitted using SDP processors and NeMo-Run.  
+  See instructions for usage [here](https://github.com/NVIDIA/NeMo-speech-data-processor/blob/main/sdp/processors/ipl/README.md).
+
+- **Training Callback: `IPLEpochStopperCallback`**  
+  Add this to your training config under `exp_manager` to **stop training at the end of each epoch**, enabling pseudo-label update:
+
+```yaml
+exp_manager:
+  create_ipl_epoch_stopper_callback: True
+  ipl_epoch_stopper_callback_params:
+    stop_every_n_epochs: n # Stop training after every n epochs (default: 1)
+

scripts/pseudo_labeling/update_inference_config.py

@@ -0,0 +1,137 @@
+# 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.
+import argparse
+import glob
+import math
+import os
+from typing import List, Union
+
+
+from filelock import FileLock
+from omegaconf import ListConfig, OmegaConf
+
+
+def count_files_for_tarred_pseudo_labeling(manifest_filepath: Union[str, ListConfig]) -> int:
+    """
+    Counts the total number of entries across multiple manifest files.
+    Args:
+        manifest_filepath (Union[str, ListConfig]): The file path to the manifest files.
+    Returns:
+        int: The total number of entries across all matching manifest files.
+    """
+    # Convert ListConfig to string if needed
+    if isinstance(manifest_filepath, ListConfig):
+        manifest_filepath = manifest_filepath[0]  # Use the first element if it's a list or ListConfig
+    dir_path, filename = os.path.split(manifest_filepath)
+    prefix = filename.split('_', 1)[0]
+    number_of_files = 0
+    for full_path in glob.glob(os.path.join(dir_path, f"{prefix}_[0-9]*.json")):
+        with open(full_path, 'r') as f:
+            number_of_files += len(f.readlines())
+    return number_of_files
+
+
+def count_files_for_pseudo_labeling(manifest_filepath: Union[str, list, ListConfig]) -> int:
+    """
+    Counts the number of entries in a single manifest file .
+    Args:
+        manifest_filepath (Union[str, list, ListConfig]): The file path to the manifest file.
+    Returns:
+        int: The total number of entries (lines) in the manifest file.
+    """
+    # Convert ListConfig to string if needed
+    if isinstance(manifest_filepath, list) or isinstance(manifest_filepath, ListConfig):
+        manifest_filepath = manifest_filepath[0]
+    with open(manifest_filepath, 'r') as f:
+        number_of_files = len(f.readlines())
+    return number_of_files
+
+
+def export_limit_predict_batches(inference_configs: List[str], p_cache: float, num_gpus: int) -> None:
+    """
+    Updates inference configuration files to set `limit_predict_batches`.
+    This is done to force partial transcription of unlabeled dataset for dynamic update of PLs.
+
+    Args:
+        inference_configs (List[str]): A list of file paths to the inference configuration files.
+        p_cache (float): A scaling factor for the cache to adjust the number of batches.
+        num_gpus (int): The number of GPUs available for inference.
+
+    Returns:
+        None: The function modifies and saves the updated configuration files in-place.
+    """
+    for config_path in inference_configs:
+        config = OmegaConf.load(config_path)
+        tarred_audio_filepaths = config.predict_ds.get("tarred_audio_filepaths", None)
+        manifest_filepaths = config.predict_ds.manifest_filepath
+
+        if tarred_audio_filepaths:
+            number_of_files = count_files_for_tarred_pseudo_labeling(manifest_filepaths)
+        else:
+            number_of_files = count_files_for_pseudo_labeling(manifest_filepaths)
+
+        if hasattr(config.predict_ds, "batch_size"):
+            batch_size = config.predict_ds.batch_size
+            limit_predict_batches = math.ceil((number_of_files * p_cache) / (batch_size * num_gpus))
+            OmegaConf.update(config, "trainer.limit_predict_batches", limit_predict_batches)
+            OmegaConf.save(config, config_path)
+        elif hasattr(config.predict_ds, "batch_duration"):
+            batch_duration = config.predict_ds.batch_duration
+            average_audio_len = 10
+            limit_predict_batches = math.ceil(
+                (number_of_files * average_audio_len * p_cache) / (batch_duration * num_gpus)
+            )
+            OmegaConf.update(config, "trainer.limit_predict_batches", limit_predict_batches)
+            OmegaConf.save(config, config_path)
+        else:
+            batch_size = 32
+            limit_predict_batches = math.ceil((number_of_files * p_cache) / (batch_size * num_gpus))
+            OmegaConf.update(config, "trainer.limit_predict_batches", limit_predict_batches)
+            OmegaConf.save(config, config_path)
+
+
+def main():
+    rank = int(os.environ.get("RANK", 0))  # Default to 0 if not set
+
+    # Ensure only one process executes this block
+    parser = argparse.ArgumentParser(description="Export limit_predict_batches as environment variables.")
+    parser.add_argument(
+        "--inference_configs",
+        type=str,
+        nargs='+',  # Accepts one or more values as a list
+        required=True,
+        help="Paths to one or more inference config YAML files.",
+    )
+    parser.add_argument("--p_cache", type=float, required=True, help="Pseudo-label cache fraction.")
+    parser.add_argument("--num_gpus", type=int, required=True, help="Number of GPUs available.")
+
+    args = parser.parse_args()
+    lock_dir = os.path.dirname(args.inference_configs[0])
+    lock_file = lock_dir + "/my_script.lock"
+    # Code executed by all processes
+
+    #     # Code executed by a single process
+    with FileLock(lock_file):
+        if rank == 0:
+            export_limit_predict_batches(
+                inference_configs=args.inference_configs, p_cache=args.p_cache, num_gpus=args.num_gpus
+            )
+
+    # Remove the lock file after the FileLock context is exited
+    if os.path.exists(lock_file):
+        os.remove(lock_file)
+
+
+if __name__ == "__main__":
+    main()

scripts/pseudo_labeling/write_transcribed_files.py

@@ -0,0 +1,260 @@
+# 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.
+import argparse
+import glob
+import json
+import os
+from typing import List
+
+
+from filelock import FileLock
+
+from nemo.utils import logging
+
+
+def create_transcribed_shard_manifests(prediction_filepaths: List[str]) -> List[str]:
+    """
+    Creates transcribed shard manifest files by processing predictions and organizing them by shard ID.
+
+    This function reads a `predictions_all.json` file from each given directory, organizes the data by
+    shard IDs, and writes the entries to separate shard manifest files. For each shard, the `pred_text`
+    field is updated as the main transcription (`text`), and the original transcription (`text`) is
+    stored as `orig_text`.
+
+    Args:
+        prediction_filepaths (List[str]): A list of file paths to directories containing
+            `predictions_all.json` files with prediction data, including shard IDs.
+
+    Returns:
+        List[str]: A list of file paths to the combined manifest files (`transcribed_manifest__OP_0..CL_.json`)
+        created for each directory.
+    """
+    all_manifest_filepaths = []
+    for prediction_filepath in prediction_filepaths:
+        max_shard_id = 0
+        shard_data = {}
+        full_path = os.path.join(prediction_filepath, "predictions_all.json")
+        with open(full_path, 'r') as f:
+            for line in f.readlines():
+                data_entry = json.loads(line)
+                shard_id = data_entry.get("shard_id")
+                if max_shard_id < shard_id:
+                    max_shard_id = shard_id
+                if shard_id not in shard_data:
+                    shard_data[shard_id] = []
+                shard_data[shard_id].append(data_entry)
+        for shard_id, entries in shard_data.items():
+            output_filename = os.path.join(prediction_filepath, f"transcribed_manifest_{shard_id}.json")
+            with open(output_filename, 'w') as f:
+                for data_entry in entries:
+                    if data_entry['audio_filepath'].endswith(".wav"):
+                        if 'text' in data_entry:
+                            data_entry['orig_text'] = data_entry.pop('text')
+                        data_entry['text'] = data_entry.pop('pred_text')
+                        json.dump(data_entry, f, ensure_ascii=False)
+                        f.write("\n")
+        shard_manifest_filepath = os.path.join(
+            prediction_filepath, f"transcribed_manifest__OP_0..{max_shard_id}_CL_.json"
+        )
+        all_manifest_filepaths.append(shard_manifest_filepath)
+    return all_manifest_filepaths
+
+
+def create_transcribed_manifests(prediction_filepaths: List[str]) -> List[str]:
+    """
+    Creates updated transcribed manifest files by processing predictions.
+
+    This function reads prediction files (`predictions_all.json`) from the provided directories,
+    updates the transcription data by renaming the `pred_text` field to `text`, and stores the
+    original `text` field as `orig_text`. The updated data is written to new transcribed manifest
+    files (`transcribed_manifest.json`) in each directory.
+
+    Args:
+        prediction_filepaths (List[str]): A list of file paths to directories containing
+            prediction files (`predictions_all.json`).
+
+    Returns:
+        List[str]: A list of file paths to the newly created transcribed manifest files
+        (`transcribed_manifest.json`).
+    """
+    all_manifest_filepaths = []
+    for prediction_filepath in prediction_filepaths:
+        prediction_name = os.path.join(prediction_filepath, "predictions_all.json")
+        transcripted_name = os.path.join(prediction_filepath, f"transcribed_manifest.json")
+
+        # Open and read the original predictions_all.json file
+        with open(transcripted_name, 'w', encoding='utf-8') as f:
+            with open(prediction_name, 'r', encoding='utf-8') as pred_f:
+
+                for line in pred_f.readlines():
+                    data_entry = json.loads(line)
+                    if 'text' in data_entry:
+                        data_entry['orig_text'] = data_entry.pop('text')
+                    data_entry['text'] = data_entry.pop('pred_text')
+                    json.dump(data_entry, f, ensure_ascii=False)
+                    f.write("\n")
+            # Append the path of the new manifest file to the list
+            all_manifest_filepaths.append(transcripted_name)
+
+    return all_manifest_filepaths
+
+
+def write_sampled_shard_transcriptions(manifest_filepaths: List[str]) -> List[List[str]]:
+    """
+    Updates transcriptions by merging predicted shard data and transcribed manifest data.
+    This function processes prediction and transcribed manifest files, merges them
+    by matching the shard_id and audio file paths. For each shard, the corresponding
+    data entries are written to a new file.
+    Args:
+        manifest_filepaths (List[str]): A list of file paths to directories containing
+            prediction and transcribed manifest files.
+    Returns:
+        List[List[str]]: A list of lists containing the file paths to the generated
+            transcribed shard manifest files.
+    """
+    all_manifest_filepaths = []
+
+    # Process each prediction directory
+    for prediction_filepath in manifest_filepaths:
+        predicted_shard_data = {}
+        # Collect entries from prediction files based on shard id
+        prediction_path = os.path.join(prediction_filepath, "predictions_all.json")
+        with open(prediction_path, 'r') as f:
+            for line in f:
+                data_entry = json.loads(line)
+                shard_id = data_entry.get("shard_id")
+                audio_filepath = data_entry['audio_filepath']
+                predicted_shard_data.setdefault(shard_id, {})[audio_filepath] = data_entry
+    max_shard_id = 0
+    for full_path in glob.glob(os.path.join(prediction_filepath, f"transcribed_manifest_[0-9]*.json")):
+        all_data_entries = []
+        with open(full_path, 'r') as f:
+            for line in f:
+                data_entry = json.loads(line)
+                shard_id = data_entry.get("shard_id")
+                max_shard_id = max(max_shard_id, shard_id)
+                all_data_entries.append(data_entry)
+        # Write the merged data to a new manifest file keeping new transcriptions
+        output_filename = os.path.join(prediction_filepath, f"transcribed_manifest_{shard_id}.json")
+        with open(output_filename, 'w') as f:
+            for data_entry in all_data_entries:
+                audio_filepath = data_entry['audio_filepath']
+                # Escape duplicated audio files that end with *dup
+                if audio_filepath.endswith(".wav"):
+                    if shard_id in predicted_shard_data and audio_filepath in predicted_shard_data[shard_id]:
+                        predicted_data_entry = predicted_shard_data[shard_id][audio_filepath]
+                        if 'text' in predicted_data_entry:
+                            predicted_data_entry['orig_text'] = predicted_data_entry.pop('text')
+                        if "pred_text" in predicted_data_entry:
+                            predicted_data_entry['text'] = predicted_data_entry.pop('pred_text')
+                        json.dump(predicted_data_entry, f, ensure_ascii=False)
+                    else:
+                        json.dump(data_entry, f, ensure_ascii=False)
+                    f.write("\n")
+
+    shard_manifest_filepath = os.path.join(prediction_filepath, f"transcribed_manifest__OP_0..{max_shard_id}_CL_.json")
+    all_manifest_filepaths.append([shard_manifest_filepath])
+
+    return all_manifest_filepaths
+
+
+def write_sampled_transcriptions(manifest_filepaths: List[str]) -> List[str]:
+    """
+    Updates transcriptions by merging predicted data with transcribed manifest data.
+
+    This function processes prediction and transcribed manifest files within given directories.
+    It matches audio file paths to update transcriptions with predictions, ensuring each audio file
+    is properly transcribed. The updated data is written to the transcribed manifest file.
+
+    Args:
+        manifest_filepaths (List[str]): A list of file paths to directories containing
+            the prediction file (`predictions_all.json`) and the transcribed manifest file
+            (`transcribed_manifest.json`).
+
+    Returns:
+        List[str]: A list of file paths to the updated transcribed manifest files.
+    """
+    all_manifest_filepaths = []
+    for prediction_filepath in manifest_filepaths:
+        predicted_data = {}
+        prediction_path = os.path.join(prediction_filepath, "predictions_all.json")
+        with open(prediction_path, 'r') as f:
+            for line in f:
+                data_entry = json.loads(line)
+                path = data_entry['audio_filepath']
+                predicted_data[path] = data_entry
+
+        full_path = os.path.join(prediction_filepath, f"transcribed_manifest.json")
+        all_data_entries = []
+        with open(full_path, 'r') as f:
+            for line in f:
+                data_entry = json.loads(line)
+                all_data_entries.append(data_entry)
+
+        output_filename = os.path.join(prediction_filepath, f"transcribed_manifest.json")
+        with open(output_filename, 'w') as f:
+            for data_entry in all_data_entries:
+                audio_filepath = data_entry['audio_filepath']
+                if audio_filepath.endswith(".wav"):
+                    if audio_filepath in predicted_data:
+                        predicted_data_entry = predicted_data[audio_filepath]
+                        if 'text' in predicted_data_entry:
+                            predicted_data_entry['orig_text'] = predicted_data_entry.pop('text')
+                        predicted_data_entry['text'] = predicted_data_entry.pop('pred_text')
+                        json.dump(predicted_data_entry, f, ensure_ascii=False)
+                        f.write("\n")
+                    else:
+                        json.dump(data_entry, f, ensure_ascii=False)
+                        f.write("\n")
+        all_manifest_filepaths.append(output_filename)
+    return all_manifest_filepaths
+
+
+if __name__ == "__main__":
+    rank = int(os.environ.get("RANK", 0))  # Default to 0 if not set
+
+    parser = argparse.ArgumentParser(description="Script to create or write transcriptions")
+    parser.add_argument("--is_tarred", action="store_true", help="If true, processes tarred manifests")
+    parser.add_argument("--full_pass", action="store_true", help="If true, processes full pass manifests")
+    parser.add_argument(
+        "--prediction_filepaths",
+        type=str,
+        nargs='+',  # Accepts one or more values as a list
+        required=True,
+        help="Paths to one or more inference config YAML files.",
+    )
+
+    args = parser.parse_args()
+
+    lock_dir = os.path.dirname(args.prediction_filepaths[0])
+    lock_file = lock_dir + "/my_script.lock"
+
+    with FileLock(lock_file):
+        if rank == 0:
+            if args.is_tarred:
+                result = (
+                    write_sampled_shard_transcriptions(args.prediction_filepaths)
+                    if not args.full_pass
+                    else create_transcribed_shard_manifests(args.prediction_filepaths)
+                )
+            else:
+                result = (
+                    write_sampled_transcriptions(args.prediction_filepaths)
+                    if not args.full_pass
+                    else create_transcribed_manifests(args.prediction_filepaths)
+                )
+
+    # Remove the lock file after the FileLock context is exited
+    if os.path.exists(lock_file):
+        os.remove(lock_file)

scripts/speaker_tasks/create_alignment_manifest.py

@@ -0,0 +1,388 @@
+# 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.
+
+import argparse
+import os
+import shutil
+from pathlib import Path
+from typing import List
+
+from nemo.collections.asr.parts.utils.manifest_utils import get_ctm_line, read_manifest, write_ctm, write_manifest
+from nemo.utils import logging
+
+
+def get_seg_info_from_ctm_line(
+    ctm_list: List[str],
+    output_precision: int,
+    speaker_index: int = 7,
+    start_time_index: int = 2,
+    duration_index: int = 3,
+):
+    """
+    Get time stamp information and speaker labels from CTM lines.
+    This is following CTM format appeared in `Rich Transcription Meeting Eval Plan: RT09` document.
+
+    CTM Format:
+        <SOURCE>< <CHANNEL> <BEG-TIME> <DURATION> <TOKEN> <CONF> <TYPE> <SPEAKER>
+
+    Args:
+        ctm_list (list): List containing CTM items. e.g.: ['sw02001-A', '1', '0.000', '0.200', 'hello', '0.98', 'lex', 'speaker3']
+        output_precision (int): Precision for CTM outputs in integer.
+
+    Returns:
+        start (float): Start time of the segment.
+        end (float): End time of the segment.
+        speaker_id (str): Speaker ID of the segment.
+    """
+    speaker_id = ctm_list[speaker_index]
+    start = float(ctm_list[start_time_index])
+    end = float(ctm_list[start_time_index]) + float(ctm_list[duration_index])
+    start = round(start, output_precision)
+    end = round(end, output_precision)
+    if type(speaker_id) == str:
+        speaker_id = speaker_id.strip()
+    return start, end, speaker_id
+
+
+def get_unaligned_files(unaligned_path: str) -> List[str]:
+    """
+    Get files without alignments in order to filter them out (as they cannot be used for data simulation).
+    In the unaligned file, each line contains the file name and the reason for the unalignment, if necessary to specify.
+
+    Example: unaligned.txt
+
+    <utterance_id> <comment>
+    1272-128104-0000 (no such file)
+    2289-152257-0025 (no such file)
+    2289-152257-0026 (mapping failed)
+    ...
+
+    Args:
+        unaligned_path (str): Path to the file containing unaligned examples
+
+    Returns:
+        skip_files (list): Unaligned file names to skip
+    """
+    skip_files = []
+    with open(unaligned_path, 'r', encoding='utf-8') as f:
+        for line in f.readlines():
+            line = line.strip()
+            if not line:
+                continue
+            unaligned_file = line.split()[0]
+            skip_files.append(unaligned_file)
+    return skip_files
+
+
+def get_new_ctm_lines_from_alignments(session_name, speaker_id, wordlist, alignments, output_precision=3) -> List[str]:
+    """
+    Create new CTM entry (to write to output ctm file)
+
+    Args:
+        session_name (str): Current session name.
+        speaker_id (int): LibriSpeech speaker ID for the current entry.
+        wordlist (list): List of words
+        alignments (list): List of alignments
+        output_precision (int): Precision for CTM outputs
+    Returns:
+        arr (list): List of ctm entries, each entry is a tuple of (start_time, text)
+    """
+    arr = []
+    for i in range(len(wordlist)):
+        word = wordlist[i]
+        if word != "":
+            # note that using the current alignments the first word is always empty, so there is no error from indexing the array with i-1
+            align1 = float(round(alignments[i - 1], output_precision))
+            align2 = float(
+                round(
+                    alignments[i] - alignments[i - 1],
+                    output_precision,
+                )
+            )
+            text = get_ctm_line(
+                source=session_name,
+                channel=speaker_id,
+                start_time=align1,
+                duration=align2,
+                token=word,
+                conf=None,
+                type_of_token='lex',
+                speaker=speaker_id,
+                output_precision=output_precision,
+            )
+            arr.append((align1, text))
+    return arr
+
+
+def load_librispeech_alignment(alignment_filepath: str) -> dict:
+    """
+    Load alignment data for librispeech
+
+    Args:
+        alignment_filepath (str): Path to the file containing alignments
+    Returns:
+        alignments (dict[tuple]): A dictionary containing file index and alignments
+    """
+    alignments = {}
+    with open(alignment_filepath, "r") as fin:
+        for line in fin.readlines():
+            line = line.strip()
+            if not line:
+                continue
+            file_id, words, timestamps = line.split()
+            alignments[file_id] = (words, timestamps)
+    return alignments
+
+
+def create_librispeech_ctm_alignments(
+    input_manifest_filepath, base_alignment_path, ctm_output_directory, libri_dataset_split
+):
+    """
+    Create new CTM alignments using input LibriSpeech word alignments.
+
+    Args:
+        input_manifest_filepath (str): Path to the input LibriSpeech manifest file
+        base_alignment_path (str): Path to the base directory containing the LibriSpeech word alignments
+        ctm_source_dir (str): Directory to write the CTM files to
+        libri_dataset_split (str): Which split of the LibriSpeech dataset is being used
+    """
+    manifest = read_manifest(input_manifest_filepath)
+    unaligned_path = os.path.join(base_alignment_path, "unaligned.txt")
+
+    if os.path.exists(unaligned_path):
+        unaligned_file_ids = set(get_unaligned_files(unaligned_path))
+    else:
+        unaligned_file_ids = set()
+
+    libri_dataset_split = libri_dataset_split.replace("_", "-")
+
+    # delete output directory if it exists or throw warning
+    if os.path.isdir(ctm_output_directory):
+        logging.info(f"Removing existing output directory: {ctm_output_directory}")
+        shutil.rmtree(ctm_output_directory)
+    if not os.path.exists(ctm_output_directory):
+        logging.info(f"Creating output directory: {ctm_output_directory}")
+        os.mkdir(ctm_output_directory)
+
+    if len(manifest) == 0:
+        raise Exception(f"Input manifest is empty: {input_manifest_filepath}")
+
+    for entry in manifest:
+        audio_file = entry['audio_filepath']
+        file_id = Path(audio_file).stem
+
+        if file_id in unaligned_file_ids:
+            continue
+
+        speaker_id = file_id.split('-')[0]
+        book_id = file_id.split('-')[1]
+        book_dir = os.path.join(base_alignment_path, "LibriSpeech", libri_dataset_split, speaker_id, book_id)
+        alignment_filepath = os.path.join(book_dir, f"{speaker_id}-{book_id}.alignment.txt")
+
+        alignment_data = load_librispeech_alignment(alignment_filepath)
+        if file_id not in alignment_data:
+            logging.warning(f"Cannot find alignment data for {audio_file} in {alignment_filepath}")
+            continue
+
+        words, end_times = alignment_data[file_id]
+        words = words.replace('\"', '').lower().split(',')
+        end_times = [float(e) for e in end_times.replace('\"', '').split(',')]
+        ctm_list = get_new_ctm_lines_from_alignments(file_id, speaker_id, words, end_times)
+        write_ctm(os.path.join(ctm_output_directory, file_id + '.ctm'), ctm_list)
+
+
+def create_manifest_with_alignments(
+    input_manifest_filepath,
+    ctm_source_dir,
+    output_manifest_filepath,
+    data_format_style,
+    silence_dur_threshold=0.1,
+    output_precision=3,
+):
+    """
+    Create new manifest file with word alignments using CTM files
+
+    Args:
+        input_manifest_filepath (str): Path to the input manifest file
+        ctm_source_dir (str): Directory to read the CTM files from
+        output_manifest_filepath (str): Path to the output manifest file containing word alignments
+        precision (int): How many decimal places to keep in the manifest file
+    """
+    manifest = read_manifest(input_manifest_filepath)
+
+    target_manifest = []
+    src_i = 0
+    tgt_i = 0
+    while src_i < len(manifest):
+        f = manifest[src_i]
+        fn = f['audio_filepath'].split('/')[-1]
+        filename = fn.split('.')[0]  # assuming that there is only one period in the input filenames
+        if "voxceleb" in data_format_style:
+            fn_split = f['audio_filepath'].split('/')
+            filename = fn_split[-3] + '-' + fn_split[-2] + '-' + fn_split[-1].split('.')[0]
+            ctm_filepath = os.path.join(ctm_source_dir, filename + '.ctm')
+        else:
+            ctm_filepath = os.path.join(ctm_source_dir, filename + '.ctm')
+
+        if not os.path.isfile(ctm_filepath):
+            logging.info(f"Skipping {filename}.wav as there is no corresponding CTM file")
+            src_i += 1
+            continue
+
+        with open(ctm_filepath, 'r') as ctm_file:
+            lines = ctm_file.readlines()
+
+        # One-word samples should be filtered out.
+        if len(lines) <= 1:
+            src_i += 1
+            continue
+
+        words = []
+        end_times = []
+        i = 0
+        prev_end = 0
+        for i in range(len(lines)):
+            ctm = lines[i].split(' ')
+            start, end, speaker_id = get_seg_info_from_ctm_line(ctm_list=ctm, output_precision=output_precision)
+            interval = start - prev_end
+
+            if (i == 0 and interval > 0) or (i > 0 and interval > silence_dur_threshold):
+                words.append("")
+                end_times.append(start)
+            elif i > 0:
+                end_times[-1] = start
+
+            words.append(ctm[4])
+            end_times.append(end)
+
+            i += 1
+            prev_end = end
+
+        # append last end
+        if f['duration'] > prev_end:
+            words.append("")
+            end_times.append(f['duration'])
+
+        # build target manifest entry
+        target_manifest.append(
+            {
+                'audio_filepath': f['audio_filepath'],
+                'duration': f['duration'],
+                'text': f['text'],
+                'words': words,
+                'alignments': end_times,
+                'speaker_id': speaker_id,
+            }
+        )
+
+        src_i += 1
+        tgt_i += 1
+
+    logging.info(f"Writing output manifest file to {output_manifest_filepath}")
+    write_manifest(output_manifest_filepath, target_manifest)
+
+
+def main():
+    """
+    Create a combined manifest file including word alignments and speaker IDs
+    """
+    input_manifest_filepath = args.input_manifest_filepath
+    base_alignment_path = args.base_alignment_path
+    output_manifest_filepath = args.output_manifest_filepath
+    ctm_output_directory = args.ctm_output_directory
+    libri_dataset_split = args.libri_dataset_split
+    use_ctm_alignment_source = args.use_ctm_alignment_source
+    output_precision = args.output_precision
+
+    # Case 1: args.base_alignment_path is containing the ctm files
+    if use_ctm_alignment_source:
+        ctm_source_dir = args.base_alignment_path
+    # Case 2: args.base_alignment_path is containing *.lab style alignments for the dataset
+    else:
+        create_librispeech_ctm_alignments(
+            input_manifest_filepath, base_alignment_path, ctm_output_directory, libri_dataset_split
+        )
+        ctm_source_dir = ctm_output_directory
+
+    create_manifest_with_alignments(
+        input_manifest_filepath,
+        ctm_source_dir,
+        output_manifest_filepath,
+        data_format_style=args.data_format_style,
+        silence_dur_threshold=args.silence_dur_threshold,
+        output_precision=output_precision,
+    )
+
+
+if __name__ == "__main__":
+    """
+    This script creates a manifest file to be used for generating synthetic
+    multispeaker audio sessions. The script takes in the default manifest file
+    for a LibriSpeech dataset and corresponding word alignments and produces
+    a combined manifest file that contains word alignments and speaker IDs
+    per example. It can also be used to produce a manifest file for a different
+    dataset if alignments are passed in CTM files.
+
+    The alignments are obtained from: https://github.com/CorentinJ/librispeech-alignments
+
+    Args:
+        input_manifest_filepath (str): Path to input manifest file
+        base_alignment_path (str): Path to the base directory for the LibriSpeech alignment dataset
+                                   (specifically to the LibriSpeech-Alignments directory containing
+                                   both the LibriSpeech folder as well as the unaligned.txt file)
+                                   or to a directory containing the requisite CTM files
+        output_manifest_filepath (str): Path to output manifest file
+        ctm_output_directory (str): Path to output CTM directory (only used for LibriSpeech)
+        libri_dataset_split (str): Which dataset split to create a combined manifest file for
+        use_ctm_alignment_source (bool): If true, base_alignment_path points to a directory containing ctm files
+    """
+    parser = argparse.ArgumentParser(description="LibriSpeech Alignment Manifest Creator")
+    parser.add_argument("--input_manifest_filepath", help="path to input manifest file", type=str, required=True)
+    parser.add_argument("--base_alignment_path", help="path to alignments (LibriSpeech)", type=str, required=False)
+    parser.add_argument("--output_manifest_filepath", help="path to output manifest file", type=str, required=True)
+    parser.add_argument(
+        "--ctm_output_directory",
+        help="path to output ctm directory for LibriSpeech (or to input CTM directory)",
+        type=str,
+        required=True,
+    )
+    parser.add_argument(
+        "--libri_dataset_split",
+        help="which test/dev/training set to create a manifest for (only used for LibriSpeech)",
+        type=str,
+        required=False,
+        default="",
+    )
+    parser.add_argument(
+        "--use_ctm_alignment_source",
+        help="if true, base_alignment_path points to a directory containing ctm files",
+        action='store_true',
+        required=False,
+    )
+    parser.add_argument(
+        "--data_format_style",
+        help="Use specific format for speaker IDs and utterance IDs. e.g. 'voxceleb', 'librispeech', 'swbd'",
+        default="",
+        type=str,
+        required=False,
+    )
+    parser.add_argument(
+        "--output_precision", help="precision for output alignments", type=int, required=False, default=3
+    )
+    parser.add_argument(
+        "--silence_dur_threshold", help="threshold for inserting silence", type=float, required=False, default=0.1
+    )
+    args = parser.parse_args()
+
+    main()

scripts/speaker_tasks/create_diarization_train_dataset.py

@@ -0,0 +1,205 @@
+# 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 creates a manifest file for diarization training. If you specify `pairwise_rttm_output_folder`, the script generates
+a two-speaker subset of the original RTTM files. For example, an RTTM file with 4 speakers will obtain 6 different pairs and
+6 RTTM files with two speakers in each RTTM file.
+
+Args:
+   --input_manifest_path: input json file name
+   --output_manifest_path: output manifest_file name
+   --pairwise_rttm_output_folder: Save two-speaker pair RTTM files
+   --window: Window length for segmentation
+   --shift: Shift length for segmentation
+   --decimals: Rounding decimals
+"""
+
+import argparse
+import copy
+import itertools
+import os
+import random
+
+from tqdm import tqdm
+
+from nemo.collections.asr.parts.utils.manifest_utils import (
+    get_input_manifest_dict,
+    get_subsegment_dict,
+    rreplace,
+    write_truncated_subsegments,
+)
+from nemo.collections.asr.parts.utils.speaker_utils import (
+    audio_rttm_map,
+    rttm_to_labels,
+    segments_manifest_to_subsegments_manifest,
+    write_rttm2manifest,
+)
+from nemo.utils import logging
+
+random.seed(42)
+
+
+def labels_to_rttmfile(labels, uniq_id, filename, out_rttm_dir):
+    """
+    Write rttm file with uniq_id name in out_rttm_dir with time_stamps in labels
+    """
+    filename = os.path.join(out_rttm_dir, filename + '.rttm')
+    with open(filename, 'w') as f:
+        for line in labels:
+            line = line.strip()
+            start, end, speaker = line.split()
+            duration = float(end) - float(start)
+            start = float(start)
+            log = 'SPEAKER {} 1   {:.3f}   {:.3f} <NA> <NA> {} <NA> <NA>\n'.format(uniq_id, start, duration, speaker)
+            f.write(log)
+
+    return filename
+
+
+def split_into_pairwise_rttm(audio_rttm_map, input_manifest_path, output_dir):
+    """
+    Create pairwise RTTM files and save it to `output_dir`. This function picks two speakers from the original RTTM files
+    then saves the two-speaker subset of RTTM to `output_dir`.
+
+    Args:
+        audio_rttm_map (dict):
+            A dictionary with keys of uniq id, which is being used to map audio files and corresponding rttm files
+        input_manifest_path (str):
+            Path of the input manifest file.
+        output_dir (str):
+            Path to the directory where the new RTTM files are saved.
+    """
+    input_manifest_dict = get_input_manifest_dict(input_manifest_path)
+    rttmlist = []
+    rttm_split_manifest_dict = {}
+    split_audio_rttm_map = {}
+    logging.info("Creating split RTTM files.")
+    for uniq_id, line in tqdm(input_manifest_dict.items(), total=len(input_manifest_dict)):
+        audiopath = line['audio_filepath']
+        num_speakers = line['num_speakers']
+        rttm_filepath = line['rttm_filepath']
+
+        rttm = rttm_to_labels(rttm_filepath)
+        speakers = []
+        j = 0
+        while len(speakers) < num_speakers:
+            if rttm[j].split(' ')[2] not in speakers:
+                speakers.append(rttm[j].split(' ')[2])
+            j += 1
+        base_fn = audiopath.split('/')[-1].replace('.wav', '')
+        for pair in itertools.combinations(speakers, 2):
+            i, target_rttm = 0, []
+            while i < len(rttm):
+                entry = rttm[i]
+                sp_id = entry.split(' ')[2]
+                if sp_id in pair:
+                    target_rttm.append(entry)
+                i += 1
+
+            pair_string = f".{pair[0]}_{pair[1]}"
+            uniq_id_pair = uniq_id + pair_string
+            filename = base_fn + pair_string
+            labels_to_rttmfile(target_rttm, base_fn, filename, output_dir)
+            rttm_path = output_dir + filename + ".rttm"
+            rttmlist.append(rttm_path)
+            line_mod = copy.deepcopy(line)
+            line_mod['rttm_filepath'] = rttm_path
+            meta = copy.deepcopy(audio_rttm_map[uniq_id])
+            meta['rttm_filepath'] = rttm_path
+            rttm_split_manifest_dict[uniq_id_pair] = line_mod
+            split_audio_rttm_map[uniq_id_pair] = meta
+
+    return rttm_split_manifest_dict, split_audio_rttm_map
+
+
+def main(input_manifest_path, output_manifest_path, pairwise_rttm_output_folder, window, shift, step_count, decimals):
+
+    if '.json' not in input_manifest_path:
+        raise ValueError("input_manifest_path file should be .json file format")
+    if output_manifest_path and '.json' not in output_manifest_path:
+        raise ValueError("output_manifest_path file should be .json file format")
+    elif not output_manifest_path:
+        output_manifest_path = rreplace(input_manifest_path, '.json', f'.{step_count}seg.json')
+
+    if pairwise_rttm_output_folder is not None:
+        if not pairwise_rttm_output_folder.endswith('/'):
+            pairwise_rttm_output_folder = f"{pairwise_rttm_output_folder}/"
+        org_audio_rttm_map = audio_rttm_map(input_manifest_path)
+        input_manifest_dict, AUDIO_RTTM_MAP = split_into_pairwise_rttm(
+            audio_rttm_map=org_audio_rttm_map,
+            input_manifest_path=input_manifest_path,
+            output_dir=pairwise_rttm_output_folder,
+        )
+    else:
+        input_manifest_dict = get_input_manifest_dict(input_manifest_path)
+        AUDIO_RTTM_MAP = audio_rttm_map(input_manifest_path)
+
+    segment_manifest_path = rreplace(input_manifest_path, '.json', '_seg.json')
+    subsegment_manifest_path = rreplace(input_manifest_path, '.json', '_subseg.json')
+
+    # todo: do we need to expose this?
+    min_subsegment_duration = 0.05
+    step_count = int(step_count)
+
+    segments_manifest_file = write_rttm2manifest(AUDIO_RTTM_MAP, segment_manifest_path, decimals)
+    subsegments_manifest_file = subsegment_manifest_path
+
+    logging.info("Creating subsegments.")
+    segments_manifest_to_subsegments_manifest(
+        segments_manifest_file=segments_manifest_file,
+        subsegments_manifest_file=subsegments_manifest_file,
+        window=window,
+        shift=shift,
+        min_subsegment_duration=min_subsegment_duration,
+        include_uniq_id=True,
+    )
+    subsegments_dict = get_subsegment_dict(subsegments_manifest_file, window, shift, decimals)
+    write_truncated_subsegments(input_manifest_dict, subsegments_dict, output_manifest_path, step_count, decimals)
+    os.remove(segment_manifest_path)
+    os.remove(subsegment_manifest_path)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input_manifest_path", help="input json file name", type=str, required=True)
+    parser.add_argument(
+        "--output_manifest_path", help="output manifest_file name", type=str, default=None, required=False
+    )
+    parser.add_argument(
+        "--pairwise_rttm_output_folder",
+        help="Save two-speaker pair RTTM files",
+        type=str,
+        default=None,
+        required=False,
+    )
+    parser.add_argument("--window", help="Window length for segmentation", type=float, required=True)
+    parser.add_argument("--shift", help="Shift length for segmentation", type=float, required=True)
+    parser.add_argument("--decimals", help="Rounding decimals", type=int, default=3, required=False)
+    parser.add_argument(
+        "--step_count",
+        help="Number of the unit segments you want to create per utterance",
+        required=True,
+    )
+    args = parser.parse_args()
+
+    main(
+        args.input_manifest_path,
+        args.output_manifest_path,
+        args.pairwise_rttm_output_folder,
+        args.window,
+        args.shift,
+        args.step_count,
+        args.decimals,
+    )

scripts/speaker_tasks/create_synth_vad_manifest.py

@@ -0,0 +1,108 @@
+# 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.
+
+import argparse
+import multiprocessing as mp
+from itertools import repeat
+from pathlib import Path
+
+import librosa
+from tqdm import tqdm
+
+from nemo.collections.asr.parts.utils.manifest_utils import read_manifest, write_manifest
+from nemo.collections.asr.parts.utils.vad_utils import get_frame_labels, load_speech_segments_from_rttm
+
+"""
+This script generates a manifest file for synthetic data generated using the NeMo multispeaker speech data simulator.
+The audio created from the simulator can be used to train a VAD model using the manifest file contains the following fields:
+The manifest file contains the following fields:
+    
+    audio_filepath (str): Path to audio file.
+    offset (float): Offset in seconds for the start of the audio file.
+    duration (float): Duration in seconds for the audio file.
+    text (str): Transcription of the audio file.
+    label (list): List of frame labels for the audio file.
+    orig_sample_rate (int): Original sample rate of the audio file.
+    vad_frame_unit_secs (float): Duration in seconds for each frame label.
+
+Usage:
+    python build_synthetic_vad_manifest.py \
+        --input_dir /path/to/synthetic/data \
+        --frame_length 0.04 \
+        --output_file /path/to/output/manifest.json
+"""
+
+
+def generate_manifest_entry(inputs):
+    """
+    Generates a manifest entry for a single audio file.
+    This function is parallelized using multiprocessing.Pool.
+
+    Args:
+        inputs (tuple): Tuple containing audio file path and frame length in seconds.
+            inputs[0]:
+                audio_filepath (str): Path to audio file.
+            inputs[1]:
+                vad_frame_unit_secs (float): Duration in seconds for each frame label.
+
+    Returns:
+        entry (dict): Dictionary containing manifest entry.
+    """
+    audio_filepath, vad_frame_unit_secs = inputs
+    audio_filepath = Path(audio_filepath)
+    y, sr = librosa.load(str(audio_filepath))
+    dur = librosa.get_duration(y=y, sr=sr)
+
+    manifest_path = audio_filepath.parent / Path(f"{audio_filepath.stem}.json")
+    audio_manifest = read_manifest(manifest_path)
+    text = " ".join([x["text"] for x in audio_manifest])
+
+    rttm_path = audio_filepath.parent / Path(f"{audio_filepath.stem}.rttm")
+    segments = load_speech_segments_from_rttm(rttm_path)
+    labels = get_frame_labels(segments, vad_frame_unit_secs, 0.0, dur)
+
+    entry = {
+        "audio_filepath": str(audio_filepath.absolute()),
+        "offset": 0.0,
+        "duration": dur,
+        "text": text,
+        "label": labels,
+        "orig_sample_rate": sr,
+        "vad_frame_unit_secs": vad_frame_unit_secs,
+    }
+    return entry
+
+
+def main(args):
+    wav_list = list(Path(args.input_dir).glob("*.wav"))
+    print(f"Found {len(wav_list)} in directory: {args.input_dir}")
+
+    inputs = zip(wav_list, repeat(args.frame_length))
+    with mp.Pool(processes=mp.cpu_count()) as pool:
+        manifest_data = list(tqdm(pool.imap(generate_manifest_entry, inputs), total=len(wav_list)))
+
+    write_manifest(args.output_file, manifest_data)
+    print(f"Manifest saved to: {args.output_file}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("input_dir", default=None, help="Path to directory containing synthetic data")
+    parser.add_argument(
+        "-l", "--frame_length", default=0.04, type=float, help="Duration in seconds for each frame label"
+    )
+    parser.add_argument("-o", "--output_file", default=None, help="Path to output manifest file")
+
+    args = parser.parse_args()
+    main(args)

scripts/speaker_tasks/eval_diar_with_asr.py

@@ -0,0 +1,242 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import argparse
+import json
+import os
+
+from nemo.collections.asr.metrics.der import evaluate_der
+from nemo.collections.asr.parts.utils.diarization_utils import OfflineDiarWithASR
+from nemo.collections.asr.parts.utils.manifest_utils import read_file
+from nemo.collections.asr.parts.utils.speaker_utils import (
+    get_uniqname_from_filepath,
+    labels_to_supervisions,
+    rttm_to_labels,
+)
+
+
+"""
+Evaluation script for diarization with ASR.
+Calculates Diarization Error Rate (DER) with RTTM files and WER and cpWER with CTM files.
+In the output ctm_eval.csv file in the output folder,
+session-level DER, WER, cpWER and speaker counting accuracies are evaluated.
+
+- Evaluation mode
+
+diar_eval_mode == "full":
+    DIHARD challenge style evaluation, the most strict way of evaluating diarization
+    (collar, ignore_overlap) = (0.0, False)
+diar_eval_mode == "fair":
+    Evaluation setup used in VoxSRC challenge
+    (collar, ignore_overlap) = (0.25, False)
+diar_eval_mode == "forgiving":
+    Traditional evaluation setup
+    (collar, ignore_overlap) = (0.25, True)
+diar_eval_mode == "all":
+    Compute all three modes (default)
+
+
+Use CTM files to calculate WER and cpWER
+```
+python eval_diar_with_asr.py \
+ --hyp_rttm_list="/path/to/hypothesis_rttm_filepaths.list" \
+ --ref_rttm_list="/path/to/reference_rttm_filepaths.list" \
+ --hyp_ctm_list="/path/to/hypothesis_ctm_filepaths.list" \
+ --ref_ctm_list="/path/to/reference_ctm_filepaths.list" \
+ --root_path="/path/to/output/directory"
+```
+
+Use .json files to calculate WER and cpWER
+```
+python eval_diar_with_asr.py \
+ --hyp_rttm_list="/path/to/hypothesis_rttm_filepaths.list" \
+ --ref_rttm_list="/path/to/reference_rttm_filepaths.list" \
+ --hyp_json_list="/path/to/hypothesis_json_filepaths.list" \
+ --ref_ctm_list="/path/to/reference_ctm_filepaths.list" \
+ --root_path="/path/to/output/directory"
+```
+
+Only use RTTMs to calculate DER
+```
+python eval_diar_with_asr.py \
+ --hyp_rttm_list="/path/to/hypothesis_rttm_filepaths.list" \
+ --ref_rttm_list="/path/to/reference_rttm_filepaths.list" \
+ --root_path="/path/to/output/directory"
+```
+
+"""
+
+
+def get_supervisions_from_rttms(rttm_file_path_list):
+    """Generate diarization annotation objects from a list of RTTM files.
+
+    Each entry in the returned list is ``[uniq_id, list[SupervisionSegment]]``.
+    """
+    annotation_obj_list = []
+    for rttm_file in rttm_file_path_list:
+        rttm_file = rttm_file.strip()
+        if rttm_file is not None and os.path.exists(rttm_file):
+            uniq_id = get_uniqname_from_filepath(rttm_file)
+            ref_labels = rttm_to_labels(rttm_file)
+            reference = labels_to_supervisions(ref_labels, uniq_name=uniq_id)
+            annotation_obj_list.append([uniq_id, reference])
+    return annotation_obj_list
+
+
+def make_meta_dict(hyp_rttm_list, ref_rttm_list):
+    """Create a temporary `audio_rttm_map_dict` for evaluation"""
+    meta_dict = {}
+    for k, rttm_file in enumerate(ref_rttm_list):
+        uniq_id = get_uniqname_from_filepath(rttm_file)
+        meta_dict[uniq_id] = {"rttm_filepath": rttm_file.strip()}
+        if hyp_rttm_list is not None:
+            hyp_rttm_file = hyp_rttm_list[k]
+            meta_dict[uniq_id].update({"hyp_rttm_filepath": hyp_rttm_file.strip()})
+    return meta_dict
+
+
+def make_trans_info_dict(hyp_json_list_path):
+    """Create `trans_info_dict` from the `.json` files"""
+    trans_info_dict = {}
+    for json_file in hyp_json_list_path:
+        json_file = json_file.strip()
+        with open(json_file) as jsf:
+            json_data = json.load(jsf)
+        uniq_id = get_uniqname_from_filepath(json_file)
+        trans_info_dict[uniq_id] = json_data
+    return trans_info_dict
+
+
+def read_file_path(list_path):
+    """Read file path and strip to remove line change symbol"""
+    return sorted([x.strip() for x in read_file(list_path)])
+
+
+def main(
+    hyp_rttm_list_path: str,
+    ref_rttm_list_path: str,
+    hyp_ctm_list_path: str,
+    ref_ctm_list_path: str,
+    hyp_json_list_path: str,
+    diar_eval_mode: str = "all",
+    root_path: str = "./",
+):
+
+    # Read filepath list files
+    hyp_rttm_list = read_file_path(hyp_rttm_list_path) if hyp_rttm_list_path else None
+    ref_rttm_list = read_file_path(ref_rttm_list_path) if ref_rttm_list_path else None
+    hyp_ctm_list = read_file_path(hyp_ctm_list_path) if hyp_ctm_list_path else None
+    ref_ctm_list = read_file_path(ref_ctm_list_path) if ref_ctm_list_path else None
+    hyp_json_list = read_file_path(hyp_json_list_path) if hyp_json_list_path else None
+
+    audio_rttm_map_dict = make_meta_dict(hyp_rttm_list, ref_rttm_list)
+
+    trans_info_dict = make_trans_info_dict(hyp_json_list) if hyp_json_list else None
+
+    all_hypothesis = get_supervisions_from_rttms(hyp_rttm_list)
+    all_reference = get_supervisions_from_rttms(ref_rttm_list)
+
+    diar_score = evaluate_der(
+        audio_rttm_map_dict=audio_rttm_map_dict,
+        all_reference=all_reference,
+        all_hypothesis=all_hypothesis,
+        diar_eval_mode=diar_eval_mode,
+    )
+
+    # Get session-level diarization error rate and speaker counting error
+    der_results = OfflineDiarWithASR.gather_eval_results(
+        diar_score=diar_score,
+        audio_rttm_map_dict=audio_rttm_map_dict,
+        trans_info_dict=trans_info_dict,
+        root_path=root_path,
+    )
+
+    if ref_ctm_list is not None:
+        # Calculate WER and cpWER if reference CTM files exist
+        if hyp_ctm_list is not None:
+            wer_results = OfflineDiarWithASR.evaluate(
+                audio_file_list=hyp_rttm_list,
+                hyp_trans_info_dict=None,
+                hyp_ctm_file_list=hyp_ctm_list,
+                ref_ctm_file_list=ref_ctm_list,
+            )
+        elif hyp_json_list is not None:
+            wer_results = OfflineDiarWithASR.evaluate(
+                audio_file_list=hyp_rttm_list,
+                hyp_trans_info_dict=trans_info_dict,
+                hyp_ctm_file_list=None,
+                ref_ctm_file_list=ref_ctm_list,
+            )
+        else:
+            raise ValueError("Hypothesis information is not provided in the correct format.")
+    else:
+        wer_results = {}
+
+    # Print average DER, WER and cpWER
+    OfflineDiarWithASR.print_errors(der_results=der_results, wer_results=wer_results)
+
+    # Save detailed session-level evaluation results in `root_path`.
+    OfflineDiarWithASR.write_session_level_result_in_csv(
+        der_results=der_results,
+        wer_results=wer_results,
+        root_path=root_path,
+        csv_columns=OfflineDiarWithASR.get_csv_columns(),
+    )
+    return None
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--hyp_rttm_list", help="path to the filelist of hypothesis RTTM files", type=str, required=True, default=None
+    )
+    parser.add_argument(
+        "--ref_rttm_list", help="path to the filelist of reference RTTM files", type=str, required=True, default=None
+    )
+    parser.add_argument(
+        "--hyp_ctm_list", help="path to the filelist of hypothesis CTM files", type=str, required=False, default=None
+    )
+    parser.add_argument(
+        "--ref_ctm_list", help="path to the filelist of reference CTM files", type=str, required=False, default=None
+    )
+    parser.add_argument(
+        "--hyp_json_list",
+        help="(Optional) path to the filelist of hypothesis JSON files",
+        type=str,
+        required=False,
+        default=None,
+    )
+    parser.add_argument(
+        "--diar_eval_mode",
+        help='evaluation mode: "all", "full", "fair", "forgiving"',
+        type=str,
+        required=False,
+        default="all",
+    )
+    parser.add_argument(
+        "--root_path", help='directory for saving result files', type=str, required=False, default="./"
+    )
+
+    args = parser.parse_args()
+
+    main(
+        args.hyp_rttm_list,
+        args.ref_rttm_list,
+        args.hyp_ctm_list,
+        args.ref_ctm_list,
+        args.hyp_json_list,
+        args.diar_eval_mode,
+        args.root_path,
+    )

scripts/speaker_tasks/filelist_to_manifest.py

@@ -0,0 +1,259 @@
+# 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 converts a filelist file where each line contains
+<absolute path of wav file> to a manifest json file.
+Optionally post processes the manifest file to create dev and train split for speaker embedding
+training, also optionally segment an audio file in to segments of random DURATIONS and create those
+wav files in CWD.
+
+Args:
+--filelist: path to file containing list of audio files
+--manifest(optional): if you already have manifest file, but would like to process it for creating
+    segments and splitting then use manifest ignoring filelist
+--id: index of speaker label in filename present in filelist file that is separated by '/'
+--out: output manifest file name
+--split: if you would want to split the  manifest file for training purposes
+    you may not need this for test set. output file names is <out>_<train/dev>.json, defaults to False
+--create_segments: if you would want to segment each manifest line to segments of [1,2,3,4] sec or less
+    you may not need this for test set, defaults to False
+--min_spkrs_count: min number of samples per speaker to consider and ignore otherwise, defaults to 0 (all speakers)
+"""
+
+import argparse
+import json
+import os
+import random
+
+import librosa as l
+import numpy as np
+import soundfile as sf
+from sklearn.model_selection import StratifiedShuffleSplit
+from tqdm.contrib.concurrent import process_map
+
+from nemo.collections.asr.parts.utils.manifest_utils import read_manifest
+
+random.seed(42)
+
+DURATIONS = sorted([3], reverse=True)
+MIN_ENERGY = 0.01
+CWD = os.getcwd()
+
+
+def _load_sox():
+    try:
+        import sox
+    except ImportError:
+        raise ImportError(
+            "Optional dependency 'sox' is required by this script. Install it with: pip install sox"
+        ) from None
+
+    return sox
+
+
+def filter_manifest_line(manifest_line):
+    split_manifest = []
+    audio_path = manifest_line['audio_filepath']
+    start = manifest_line.get('offset', 0)
+    dur = manifest_line['duration']
+    label = manifest_line['label']
+    endname = os.path.splitext(audio_path.split(label, 1)[-1])[0]
+    to_path = os.path.join(CWD, 'segments', label)
+    to_path = os.path.join(to_path, endname[1:])
+    os.makedirs(os.path.dirname(to_path), exist_ok=True)
+
+    if dur >= min(DURATIONS):
+        signal, sr = sf.read(audio_path)
+        remaining_dur = dur - start
+
+        segments = DURATIONS.copy()
+        mode = int(remaining_dur // sum(DURATIONS))
+        rem = remaining_dur % sum(DURATIONS)
+        segments = mode * segments
+
+        for val in DURATIONS:
+            if rem >= val:
+                segments.append(val)
+                rem = rem - val
+
+        for temp_dur in segments:
+            segment_audio = signal[int(start * sr) : int(start * sr + temp_dur * sr)]
+            if l.feature.rms(y=segment_audio).mean() > MIN_ENERGY:
+                final_string = '_' + str(start) + '_' + str(temp_dur)
+                final_string = final_string.replace('.', '-')
+                to_file = to_path + final_string + '.wav'
+
+                c_start = int(float(start * sr))
+                c_end = c_start + int(float(temp_dur * sr))
+                segment = signal[c_start:c_end]
+                sf.write(to_file, segment, sr)
+
+                meta = manifest_line.copy()
+                meta['audio_filepath'] = to_file
+                meta['offset'] = 0
+                meta['duration'] = temp_dur
+                split_manifest.append(meta)
+
+            start = start + temp_dur
+
+    return split_manifest
+
+
+def count_and_consider_only(speakers, lines, min_count=10):
+    """
+    consider speakers only if samples per speaker is at least min_count
+    """
+    uniq_speakers, indices, counts = np.unique(speakers, return_index=True, return_counts=True)
+    print("speaker count before filtering minimum number of speaker counts: ", len(uniq_speakers))
+    required_speakers = {}
+    for idx, count in enumerate(counts):
+        if count >= min_count:
+            required_speakers[uniq_speakers[idx]] = count
+
+    print("speaker count after filtering minimum number of speaker counts: ", len(required_speakers))
+    required_lines = []
+    speakers_only = []
+    for idx, speaker in enumerate(speakers):
+        if speaker in required_speakers:
+            required_lines.append(lines[idx])
+            speakers_only.append(speaker)
+
+    return speakers_only, required_lines
+
+
+def write_file(name, lines, idx):
+    with open(name, 'w', encoding='utf-8') as fout:
+        for i in idx:
+            dic = lines[i]
+            json.dump(dic, fout)
+            fout.write('\n')
+    print("wrote", name)
+
+
+def read_file(filelist, id=-1):
+    json_lines = []
+    with open(filelist, 'r') as fo:
+        lines = fo.readlines()
+        lines = sorted(lines)
+        for line in lines:
+            line = line.strip()
+            speaker = line.split('/')[id]
+            speaker = list(speaker)
+            speaker = ''.join(speaker)
+            meta = {"audio_filepath": line, "offset": 0, "duration": None, "label": speaker}
+            json_lines.append(meta)
+    return json_lines
+
+
+def get_duration(json_line):
+    dur = json_line['duration']
+    if dur is None:
+        sox = _load_sox()
+        wav_path = json_line['audio_filepath']
+        json_line['duration'] = sox.file_info.duration(wav_path)
+    return json_line
+
+
+def get_labels(lines):
+    labels = []
+    for line in lines:
+        label = line['label']
+        labels.append(label)
+    return labels
+
+
+def main(filelist, manifest, id, out, split=False, create_segments=False, min_count=10):
+    if os.path.exists(out):
+        os.remove(out)
+    if filelist:
+        lines = read_file(filelist=filelist, id=id)
+        lines = process_map(get_duration, lines, chunksize=100)
+        out_file = os.path.splitext(filelist)[0] + '_manifest.json'
+        write_file(out_file, lines, range(len(lines)))
+    else:
+        lines = read_manifest(manifest)
+
+    lines = process_map(get_duration, lines, chunksize=100)
+
+    if create_segments:
+        print(f"creating and writing segments to {CWD}")
+        lines = process_map(filter_manifest_line, lines, chunksize=100)
+        temp = []
+        for line in lines:
+            temp.extend(line)
+        del lines
+        lines = temp
+
+    speakers = [x['label'] for x in lines]
+
+    if min_count:
+        speakers, lines = count_and_consider_only(speakers, lines, abs(min_count))
+
+    write_file(out, lines, range(len(lines)))
+    path = os.path.dirname(out)
+    if split:
+        speakers = [x['label'] for x in lines]
+        sss = StratifiedShuffleSplit(n_splits=1, test_size=0.1, random_state=42)
+        for train_idx, test_idx in sss.split(speakers, speakers):
+            print("number of train samples after split: ", len(train_idx))
+
+        out = os.path.join(path, 'train.json')
+        write_file(out, lines, train_idx)
+        out = os.path.join(path, 'dev.json')
+        write_file(out, lines, test_idx)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--filelist", help="path to filelist file", type=str, required=False, default=None)
+    parser.add_argument("--manifest", help="manifest file name", type=str, required=False, default=None)
+    parser.add_argument(
+        "--id",
+        help="field num seperated by '/' to be considered as speaker label from filelist file, can be ignored if manifest file is already provided with labels",
+        type=int,
+        required=False,
+        default=None,
+    )
+    parser.add_argument("--out", help="manifest_file name", type=str, required=True)
+    parser.add_argument(
+        "--split",
+        help="bool if you would want to split the manifest file for training purposes",
+        required=False,
+        action='store_true',
+    )
+    parser.add_argument(
+        "--create_segments",
+        help="bool if you would want to segment each manifest line to segments of 4 sec or less",
+        required=False,
+        action='store_true',
+    )
+    parser.add_argument(
+        "--min_spkrs_count",
+        default=0,
+        type=int,
+        help="min number of samples per speaker to consider and ignore otherwise",
+    )
+
+    args = parser.parse_args()
+
+    main(
+        args.filelist,
+        args.manifest,
+        args.id,
+        args.out,
+        args.split,
+        args.create_segments,
+        args.min_spkrs_count,
+    )

scripts/speaker_tasks/multispeaker_data_analysis.py

@@ -0,0 +1,301 @@
+# 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.
+
+import argparse
+import multiprocessing
+import shutil
+from collections import OrderedDict
+from pathlib import Path
+from pprint import pprint
+from typing import Dict
+
+import numpy as np
+from scipy.stats import expon
+from tqdm import tqdm
+
+from nemo.collections.asr.parts.utils.vad_utils import (
+    get_nonspeech_segments,
+    load_speech_overlap_segments_from_rttm,
+    plot_sample_from_rttm,
+)
+from nemo.utils.dependency import import_optional_dependency
+
+"""
+This script analyzes multi-speaker speech dataset and generates statistics.
+The input directory </path/to/rttm_and_wav_directory> is required to contain the following files:
+    - rttm files (*.rttm)
+    - wav files (*.wav)
+
+Usage:
+    python <NEMO_ROOT>/scripts/speaker_tasks/multispeaker_data_analysis.py \
+        </path/to/rttm_and_wav_directory> \
+        --session_dur 20 \
+        --silence_mean 0.2 \
+        --silence_var 100 \
+        --overlap_mean 0.15 \
+        --overlap_var 50 \
+        --num_workers 8 \
+        --num_samples 10 \
+        --output_dir <path/to/output_directory>
+"""
+
+
+def process_sample(sess_dict: Dict) -> Dict:
+    """
+    Process each synthetic sample
+
+    Args:
+        sess_dict (dict): dictionary containing the following keys
+            rttm_file (str): path to the rttm file
+            session_dur (float): duration of the session (specified by argument)
+            precise (bool): whether to measure the precise duration of the session using sox
+
+    Returns:
+        results (dict): dictionary containing the following keys
+            session_dur (float): duration of the session
+            silence_len_list (list): list of silence durations of each silence occurrence
+            silence_dur (float): total silence duration in a session
+            silence_ratio (float): ratio of silence duration to session duration
+            overlap_len_list (list): list of overlap durations of each overlap occurrence
+            overlap_dur (float): total overlap duration
+            overlap_ratio (float): ratio of overlap duration to speech (non-silence) duration
+    """
+
+    rttm_file = sess_dict["rttm_file"]
+    session_dur = sess_dict["session_dur"]
+    precise = sess_dict["precise"]
+    if precise or session_dur is None:
+        sox = import_optional_dependency("sox")
+        wav_file = rttm_file.parent / Path(rttm_file.stem + ".wav")
+        session_dur = sox.file_info.duration(str(wav_file))
+
+    speech_seg, overlap_seg = load_speech_overlap_segments_from_rttm(rttm_file)
+    speech_dur = sum([sess_dict[1] - sess_dict[0] for sess_dict in speech_seg])
+
+    silence_seg = get_nonspeech_segments(speech_seg, session_dur)
+    silence_len_list = [sess_dict[1] - sess_dict[0] for sess_dict in silence_seg]
+    silence_dur = max(0, session_dur - speech_dur)
+    silence_ratio = silence_dur / session_dur
+
+    overlap_len_list = [sess_dict[1] - sess_dict[0] for sess_dict in overlap_seg]
+    overlap_dur = sum(overlap_len_list) if len(overlap_len_list) else 0
+    overlap_ratio = overlap_dur / speech_dur
+
+    results = {
+        "session_dur": session_dur,
+        "silence_len_list": silence_len_list,
+        "silence_dur": silence_dur,
+        "silence_ratio": silence_ratio,
+        "overlap_len_list": overlap_len_list,
+        "overlap_dur": overlap_dur,
+        "overlap_ratio": overlap_ratio,
+    }
+
+    return results
+
+
+def run_multispeaker_data_analysis(
+    input_dir,
+    session_dur=None,
+    silence_mean=None,
+    silence_var=None,
+    overlap_mean=None,
+    overlap_var=None,
+    precise=False,
+    save_path=None,
+    num_workers=1,
+) -> Dict:
+    rttm_list = list(Path(input_dir).glob("*.rttm"))
+    """
+    Analyze the multispeaker data and plot the distribution of silence and overlap durations.
+
+    Args:
+        input_dir (str): path to the directory containing the rttm files
+        session_dur (float): duration of the session (specified by argument)
+        silence_mean (float): mean of the silence duration distribution
+        silence_var (float): variance of the silence duration distribution
+        overlap_mean (float): mean of the overlap duration distribution
+        overlap_var (float): variance of the overlap duration distribution
+        precise (bool): whether to measure the precise duration of the session using sox
+        save_path (str): path to save the plots
+
+    Returns:
+        stats (dict): dictionary containing the statistics of the analyzed data
+    """
+    import matplotlib.pyplot as plt
+
+    sns = import_optional_dependency("seaborn")
+
+    print(f"Found {len(rttm_list)} files to be processed")
+    if len(rttm_list) == 0:
+        raise ValueError(f"No rttm files found in {input_dir}")
+
+    silence_duration = 0.0
+    total_duration = 0.0
+    overlap_duration = 0.0
+
+    silence_ratio_all = []
+    overlap_ratio_all = []
+    silence_length_all = []
+    overlap_length_all = []
+
+    queue = []
+    for rttm_file in tqdm(rttm_list):
+        queue.append(
+            {
+                "rttm_file": rttm_file,
+                "session_dur": session_dur,
+                "precise": precise,
+            }
+        )
+
+    if num_workers <= 1:
+        results = [process_sample(sess_dict) for sess_dict in tqdm(queue)]
+    else:
+        with multiprocessing.Pool(processes=num_workers) as p:
+            results = list(
+                tqdm(
+                    p.imap(process_sample, queue),
+                    total=len(queue),
+                    desc='Processing',
+                    leave=True,
+                )
+            )
+
+    for item in results:
+        total_duration += item["session_dur"]
+        silence_duration += item["silence_dur"]
+        overlap_duration += item["overlap_dur"]
+
+        silence_length_all += item["silence_len_list"]
+        overlap_length_all += item["overlap_len_list"]
+
+        silence_ratio_all.append(item["silence_ratio"])
+        overlap_ratio_all.append(item["overlap_ratio"])
+
+    actual_silence_mean = silence_duration / total_duration
+    actual_silence_var = np.var(silence_ratio_all)
+    actual_overlap_mean = overlap_duration / (total_duration - silence_duration)
+    actual_overlap_var = np.var(overlap_ratio_all)
+
+    stats = OrderedDict()
+    stats["total duration (hours)"] = f"{total_duration / 3600:.2f}"
+    stats["number of sessions"] = len(rttm_list)
+    stats["average session duration (seconds)"] = f"{total_duration / len(rttm_list):.2f}"
+    stats["actual silence ratio mean/var"] = f"{actual_silence_mean:.4f}/{actual_silence_var:.4f}"
+    stats["actual overlap ratio mean/var"] = f"{actual_overlap_mean:.4f}/{actual_overlap_var:.4f}"
+    stats["expected silence ratio mean/var"] = f"{silence_mean}/{silence_var}"
+    stats["expected overlap ratio mean/var"] = f"{overlap_mean}/{overlap_var}"
+    stats["save_path"] = save_path
+
+    print("-----------------------------------------------")
+    print("                    Results                    ")
+    print("-----------------------------------------------")
+    for k, v in stats.items():
+        print(k, ": ", v)
+    print("-----------------------------------------------")
+
+    fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(14, 14))
+    fig.suptitle(
+        f"Average session={total_duration/len(rttm_list):.2f} seconds, num sessions={len(rttm_list)}, total={total_duration/3600:.2f} hours"
+    )
+    sns.histplot(silence_ratio_all, ax=ax1)
+    ax1.set_xlabel("Silence ratio in a session")
+    ax1.set_title(
+        f"Target silence mean={silence_mean}, var={silence_var}. \nActual silence ratio={actual_silence_mean:.4f}, var={actual_silence_var:.4f}"
+    )
+
+    _, scale = expon.fit(silence_length_all, floc=0)
+    sns.histplot(silence_length_all, ax=ax2)
+    ax2.set_xlabel("Per-silence length in seconds")
+    ax2.set_title(f"Per-silence length histogram, \nfitted exponential distribution with mean={scale:.4f}")
+
+    sns.histplot(overlap_ratio_all, ax=ax3)
+    ax3.set_title(
+        f"Target overlap mean={overlap_mean}, var={overlap_var}. \nActual ratio={actual_overlap_mean:.4f}, var={actual_overlap_var:.4f}"
+    )
+    ax3.set_xlabel("Overlap ratio in a session")
+    _, scale2 = expon.fit(overlap_length_all, floc=0)
+    sns.histplot(overlap_length_all, ax=ax4)
+    ax4.set_title(f"Per overlap length histogram, \nfitted exponential distribution with mean={scale2:.4f}")
+    ax4.set_xlabel("Duration in seconds")
+
+    if save_path:
+        fig.savefig(save_path)
+        print(f"Figure saved at: {save_path}")
+
+    return stats
+
+
+def visualize_multispeaker_data(input_dir: str, output_dir: str, num_samples: int = 10) -> None:
+    """
+    Visualize a set of randomly sampled data in the input directory
+
+    Args:
+        input_dir (str): Path to the input directory
+        output_dir (str): Path to the output directory
+        num_samples (int): Number of samples to visualize
+    """
+    rttm_list = list(Path(input_dir).glob("*.rttm"))
+    idx_list = np.random.permutation(len(rttm_list))[:num_samples]
+    print(f"Visualizing {num_samples} random samples")
+    for idx in idx_list:
+        rttm_file = rttm_list[idx]
+        audio_file = rttm_file.parent / Path(rttm_file.stem + ".wav")
+        output_file = Path(output_dir) / Path(rttm_file.stem + ".png")
+        plot_sample_from_rttm(audio_file=audio_file, rttm_file=rttm_file, save_path=str(output_file), show=False)
+    print(f"Sample plots saved at: {output_dir}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("input_dir", default="", help="Input directory")
+    parser.add_argument("-sd", "--session_dur", default=None, type=float, help="Duration per session in seconds")
+    parser.add_argument("-sm", "--silence_mean", default=None, type=float, help="Expected silence ratio mean")
+    parser.add_argument("-sv", "--silence_var", default=None, type=float, help="Expected silence ratio variance")
+    parser.add_argument("-om", "--overlap_mean", default=None, type=float, help="Expected overlap ratio mean")
+    parser.add_argument("-ov", "--overlap_var", default=None, type=float, help="Expected overlap ratio variance")
+    parser.add_argument("-w", "--num_workers", default=1, type=int, help="Number of CPU workers to use")
+    parser.add_argument("-s", "--num_samples", default=10, type=int, help="Number of random samples to plot")
+    parser.add_argument("-o", "--output_dir", default="analysis/", type=str, help="Directory for saving output figure")
+    parser.add_argument(
+        "--precise", action="store_true", help="Set to get precise duration, with significant time cost"
+    )
+    args = parser.parse_args()
+
+    print("Running with params:")
+    pprint(vars(args))
+
+    output_dir = Path(args.output_dir)
+    if output_dir.exists():
+        print(f"Removing existing output directory: {args.output_dir}")
+        shutil.rmtree(str(output_dir))
+    output_dir.mkdir(parents=True)
+
+    run_multispeaker_data_analysis(
+        input_dir=args.input_dir,
+        session_dur=args.session_dur,
+        silence_mean=args.silence_mean,
+        silence_var=args.silence_var,
+        overlap_mean=args.overlap_mean,
+        overlap_var=args.overlap_var,
+        precise=args.precise,
+        save_path=str(Path(args.output_dir, "statistics.png")),
+        num_workers=args.num_workers,
+    )
+
+    visualize_multispeaker_data(input_dir=args.input_dir, output_dir=args.output_dir, num_samples=args.num_samples)
+
+    print("The multispeaker data analysis has been completed.")
+    print(f"Please check the output directory: \n{args.output_dir}")

scripts/speaker_tasks/pathfiles_to_diarize_manifest.py

@@ -0,0 +1,69 @@
+# 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.
+
+import argparse
+import random
+
+from nemo.collections.asr.parts.utils.manifest_utils import create_manifest
+
+random.seed(42)
+
+"""
+This script creates manifest file for speaker diarization inference purposes.
+Useful to get manifest when you have list of audio files and optionally rttm and uem files for evaluation
+
+Note: make sure basename for each file is unique and rttm files also has the corresponding base name for mapping
+"""
+
+
+def main(
+    wav_path, text_path=None, rttm_path=None, uem_path=None, ctm_path=None, manifest_filepath=None, add_duration=False
+):
+    create_manifest(
+        wav_path,
+        manifest_filepath,
+        text_path=text_path,
+        rttm_path=rttm_path,
+        uem_path=uem_path,
+        ctm_path=ctm_path,
+        add_duration=add_duration,
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--paths2audio_files", help="path to text file containing list of audio files", type=str, required=True
+    )
+    parser.add_argument("--paths2txt_files", help="path to text file containing list of transcription files", type=str)
+    parser.add_argument("--paths2rttm_files", help="path to text file containing list of rttm files", type=str)
+    parser.add_argument("--paths2uem_files", help="path to uem files", type=str)
+    parser.add_argument("--paths2ctm_files", help="path to ctm files", type=str)
+    parser.add_argument("--manifest_filepath", help="path to output manifest file", type=str, required=True)
+    parser.add_argument(
+        "--add_duration",
+        help="add duration of audio files to output manifest files.",
+        action='store_true',
+    )
+    args = parser.parse_args()
+
+    main(
+        args.paths2audio_files,
+        args.paths2txt_files,
+        args.paths2rttm_files,
+        args.paths2uem_files,
+        args.paths2ctm_files,
+        args.manifest_filepath,
+        args.add_duration,
+    )

scripts/speech_recognition/canary/build_canary_1b_special_tokenizer.py

@@ -0,0 +1,53 @@
+#!/usr/bin/env python
+# 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.
+from pathlib import Path
+
+import click
+
+from nemo.collections.common.tokenizers import CanaryTokenizer
+
+
+@click.command()
+@click.argument("output_dir", type=click.Path())
+def main(output_dir: str) -> None:
+    """
+    Builds the special tokens tokenizer for NVIDIA Canary-1B model.
+    It's intended to be used with CanaryTokenizer (a specialized AggregateTokenizer)
+    under name ``spl_tokens``.
+    """
+    CanaryTokenizer.build_special_tokenizer(
+        [
+            "<|endoftext|>",
+            "<|startoftranscript|>",
+            "<|transcribe|>",
+            "<|translate|>",
+            "<|nopnc|>",
+            "<|pnc|>",
+            "<|nospeech|>",
+        ]
+        + [
+            "<|en|>",
+            "<|es|>",
+            "<|de|>",
+            "<|fr|>",
+        ]
+        + [f"<|spltoken{i}|>" for i in range(16)],
+        model_dir=output_dir,
+        force_rebuild=True,
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/speech_recognition/canary/build_canary_2_special_tokenizer.py

@@ -0,0 +1,278 @@
+#!/usr/bin/env python
+# 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.
+import math
+
+import click
+
+from nemo.collections.common.tokenizers import CanaryTokenizer
+
+
+@click.command()
+@click.argument("output_dir", type=click.Path())
+def main(output_dir: str) -> None:
+    """
+    Builds the special tokens tokenizer for NVIDIA Canary-2.0 model.
+    It's intended to be used with CanaryTokenizer (a specialized AggregateTokenizer)
+    under name ``spl_tokens``.
+    """
+
+    tokens = (
+        [
+            # Generic special tokens
+            "<|endoftext|>",
+            "<|startoftranscript|>",
+            "<|nopnc|>",
+            "<|pnc|>",
+            "<|nospeech|>",
+            "<|startofcontext|>",
+            "<|itn|>",
+            "<|noitn|>",
+            "<|timestamp|>",
+            "<|notimestamp|>",
+            "<|diarize|>",
+            "<|nodiarize|>",
+            "<|spkchange|>",
+            "<|audioseparator|>",
+            "<|emo:undefined|>",
+            "<|emo:neutral|>",
+            "<|emo:happy|>",
+            "<|emo:sad|>",
+            "<|emo:angry|>",
+        ]
+        # Language special tokens
+        + [
+            "<|unklang|>",
+        ]
+        + ISO_LANGS
+        # Timestamp frame special tokens
+        + [f"<|{i}|>" for i in range(900)]
+        # Speaker indicator special tokens
+        + [f"<|spk{i}|>" for i in range(16)]
+    )
+
+    num_tokens = len(tokens) + 3  # count "<pad>", "<unk>", "_" too
+    print(f"We have {num_tokens} special tokens.")
+    final_num_tokens = next_multiple_of_64(num_tokens)
+    num_extra_tokens = final_num_tokens - num_tokens
+    print(f"Adding extra {num_extra_tokens} unused special tokens for a total vocab size of {final_num_tokens}")
+
+    tokens += [
+        # Timestamp related special tokens
+        f"<|spltoken{i}|>"
+        for i in range(num_extra_tokens)
+    ]
+
+    tokenizer = CanaryTokenizer.build_special_tokenizer(
+        tokens=tokens,
+        model_dir=output_dir,
+        force_rebuild=True,
+    )
+
+    assert tokenizer.vocab_size == 1152, tokenizer.vocab_size
+
+
+def next_multiple_of_64(n):
+    return ((n + 63) // 64) * 64
+
+
+ISO_LANGS = [
+    "<|aa|>",
+    "<|ab|>",
+    "<|af|>",
+    "<|ak|>",
+    "<|sq|>",
+    "<|am|>",
+    "<|ar|>",
+    "<|an|>",
+    "<|hy|>",
+    "<|as|>",
+    "<|av|>",
+    "<|ae|>",
+    "<|ay|>",
+    "<|az|>",
+    "<|bm|>",
+    "<|ba|>",
+    "<|eu|>",
+    "<|be|>",
+    "<|bn|>",
+    "<|bi|>",
+    "<|bs|>",
+    "<|br|>",
+    "<|bg|>",
+    "<|my|>",
+    "<|ca|>",
+    "<|ch|>",
+    "<|ce|>",
+    "<|ny|>",
+    "<|zh|>",
+    "<|cu|>",
+    "<|cv|>",
+    "<|kw|>",
+    "<|co|>",
+    "<|cr|>",
+    "<|hr|>",
+    "<|cs|>",
+    "<|da|>",
+    "<|dv|>",
+    "<|nl|>",
+    "<|dz|>",
+    "<|en|>",
+    "<|eo|>",
+    "<|et|>",
+    "<|ee|>",
+    "<|fo|>",
+    "<|fj|>",
+    "<|fi|>",
+    "<|fr|>",
+    "<|fy|>",
+    "<|ff|>",
+    "<|gd|>",
+    "<|gl|>",
+    "<|lg|>",
+    "<|ka|>",
+    "<|de|>",
+    "<|el|>",
+    "<|kl|>",
+    "<|gn|>",
+    "<|gu|>",
+    "<|ht|>",
+    "<|ha|>",
+    "<|he|>",
+    "<|hz|>",
+    "<|hi|>",
+    "<|ho|>",
+    "<|hu|>",
+    "<|is|>",
+    "<|io|>",
+    "<|ig|>",
+    "<|id|>",
+    "<|ia|>",
+    "<|ie|>",
+    "<|iu|>",
+    "<|ik|>",
+    "<|ga|>",
+    "<|it|>",
+    "<|ja|>",
+    "<|jv|>",
+    "<|kn|>",
+    "<|kr|>",
+    "<|ks|>",
+    "<|kk|>",
+    "<|km|>",
+    "<|ki|>",
+    "<|rw|>",
+    "<|ky|>",
+    "<|kv|>",
+    "<|kg|>",
+    "<|ko|>",
+    "<|kj|>",
+    "<|ku|>",
+    "<|lo|>",
+    "<|la|>",
+    "<|lv|>",
+    "<|li|>",
+    "<|ln|>",
+    "<|lt|>",
+    "<|lu|>",
+    "<|lb|>",
+    "<|mk|>",
+    "<|mg|>",
+    "<|ms|>",
+    "<|ml|>",
+    "<|mt|>",
+    "<|gv|>",
+    "<|mi|>",
+    "<|mr|>",
+    "<|mh|>",
+    "<|mn|>",
+    "<|na|>",
+    "<|nv|>",
+    "<|nd|>",
+    "<|nr|>",
+    "<|ng|>",
+    "<|ne|>",
+    "<|no|>",
+    "<|nb|>",
+    "<|nn|>",
+    "<|oc|>",
+    "<|oj|>",
+    "<|or|>",
+    "<|om|>",
+    "<|os|>",
+    "<|pi|>",
+    "<|ps|>",
+    "<|fa|>",
+    "<|pl|>",
+    "<|pt|>",
+    "<|pa|>",
+    "<|qu|>",
+    "<|ro|>",
+    "<|rm|>",
+    "<|rn|>",
+    "<|ru|>",
+    "<|se|>",
+    "<|sm|>",
+    "<|sg|>",
+    "<|sa|>",
+    "<|sc|>",
+    "<|sr|>",
+    "<|sn|>",
+    "<|sd|>",
+    "<|si|>",
+    "<|sk|>",
+    "<|sl|>",
+    "<|so|>",
+    "<|st|>",
+    "<|es|>",
+    "<|su|>",
+    "<|sw|>",
+    "<|ss|>",
+    "<|sv|>",
+    "<|tl|>",
+    "<|ty|>",
+    "<|tg|>",
+    "<|ta|>",
+    "<|tt|>",
+    "<|te|>",
+    "<|th|>",
+    "<|bo|>",
+    "<|ti|>",
+    "<|to|>",
+    "<|ts|>",
+    "<|tn|>",
+    "<|tr|>",
+    "<|tk|>",
+    "<|tw|>",
+    "<|ug|>",
+    "<|uk|>",
+    "<|ur|>",
+    "<|uz|>",
+    "<|ve|>",
+    "<|vi|>",
+    "<|vo|>",
+    "<|wa|>",
+    "<|cy|>",
+    "<|wo|>",
+    "<|xh|>",
+    "<|ii|>",
+    "<|yi|>",
+    "<|yo|>",
+    "<|za|>",
+    "<|zu|>",
+]
+
+
+if __name__ == "__main__":
+    main()

scripts/speech_recognition/code_switching/README.md

@@ -0,0 +1,17 @@
+# Scripts for creation of synthetic code-switched data from monolingual sources
+Follow the 2 steps listed below in order - 
+
+
+1. Create the (intermediate) manifest file using `code_switching_manifest_creation.py`.  It's usage is as follows:
+
+    `python code_switching_manifest_creation.py --manifest_language1 <absolute path of Language 1's manifest file> --manifest_language2 <absolute path of Language 2's manifest file> --manifest_save_path <absolute path to save the created manifest> --id_language1 <language code for language 1 (e.g. en)> --id_language2 <language code for language 2 (e.g. es)> --max_sample_duration_sec <maximum duration of generated sample in seconds> --min_sample_duration_sec <maximum duration of generated sample in seconds> --dataset_size_required_hrs <size of generated synthetic dataset required in hrs>`
+
+    Estimated runtime for dataset_size_required_hrs=10,000 is ~2 mins
+
+2. Create the synthetic audio data and the corresponding manifest file using `code_switching_audio_data_creation.py` It's usage is as follows: 
+
+    `python code_switching_audio_data_creation.py --manifest_path <absolute path to intermediate CS manifest generated in step 1> --audio_save_folder_path <absolute path to directory where you want to save the synthesized audios> --manifest_save_path <absolute path to save the created manifest> --audio_normalized_amplitude <scaled normalized amplitude desired> --cs_data_sampling_rate <desired sampling rate for generated audios> --sample_beginning_pause_msec <pause to be added to the beginning of the generated sample in milli seconds> --sample_joining_pause_msec <pause to be added between segments while joining, in milli seconds> --sample_end_pause_msec <pause to be added to the end of the generated sample in milli seconds> --is_lid_manifest <boolean to create manifest in the multi-sample lid format for the text field, true by default> --workers <number of worker processes>`
+
+    Example of the multi-sample LID format: ```[{“str”:“esta muestra ” “lang”:”es”},{“str”:“was generated synthetically”: “lang”:”en”}]```
+    
+    Estimated runtime for generating a 10,000 hour corpus is ~40 hrs with a single worker

scripts/speech_recognition/code_switching/code_switching_audio_data_creation.py

@@ -0,0 +1,288 @@
+# 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.
+
+import argparse
+import json
+import logging
+import os
+
+import librosa
+import numpy as np
+from joblib import Parallel, delayed
+from scipy.io import wavfile
+from tqdm import tqdm
+from nemo.collections.asr.parts.utils.manifest_utils import read_manifest
+
+parser = argparse.ArgumentParser(description='Create synthetic code-switching data audio data from monolingual data')
+parser.add_argument("--manifest_path", default=None, type=str, help='Path to CS indermediate manifest', required=True)
+parser.add_argument(
+    "--audio_save_folder_path",
+    default=None,
+    type=str,
+    help='Path to directory where created synthetic set would be saved',
+    required=True,
+)
+parser.add_argument(
+    "--manifest_save_path", default=None, type=str, help='Path to save the created manifest', required=True
+)
+parser.add_argument(
+    "--audio_normalized_amplitude", default=15000, type=int, help='Normalized amplitdue of audio samples'
+)
+parser.add_argument(
+    "--cs_data_sampling_rate",
+    default=16000,
+    type=int,
+    help='Desired sampling rate for the audios in the generated dataset',
+)
+parser.add_argument(
+    "--sample_beginning_pause_msec",
+    default=20,
+    type=int,
+    help='Pause to be added at the beginning of the sample (msec)',
+)
+parser.add_argument(
+    "--sample_joining_pause_msec",
+    default=100,
+    type=int,
+    help='Pause to be added between different phrases of the sample (msec)',
+)
+parser.add_argument(
+    "--sample_end_pause_msec", default=20, type=int, help='Pause to be added at the end of the sample (msec)'
+)
+parser.add_argument(
+    "--is_lid_manifest",
+    default=True,
+    type=bool,
+    help='If true, generate manifest in the multi-sample lid format, else the standard manifest format',
+)
+parser.add_argument("--workers", default=1, type=int, help='Number of worker processes')
+
+args = parser.parse_args()
+
+
+def split_list(input_list: list, num_splits: int):
+    """
+    Args:
+        input_list: the input list to split
+        num_splits: number of splits required
+
+    Returns:
+        iterator of split lists
+
+    """
+    k, m = divmod(len(input_list), num_splits)
+    return (input_list[i * k + min(i, m) : (i + 1) * k + min(i + 1, m)] for i in range(num_splits))
+
+
+def combine_manifests(manifest_save_path: str, num_split: int):
+    """
+    Args:
+        manifest_save_path: absolute path to save the combined manifest
+        num_splits: number of splits of manifest
+
+    Returns:
+        num_samples_combined: the total number of samples in the generated dataset
+    """
+    num_samples_combined = 0
+    base_directory = os.path.dirname(manifest_save_path)
+
+    with open(manifest_save_path, 'w') as outfile:
+        for i in range(num_split):
+            split_manifest_path = base_directory + '/temp_' + str(i) + '.json'
+            data_split = read_manifest(split_manifest_path)
+
+            for elem in data_split:
+                s = json.dumps(elem)
+                outfile.write(s + '\n')
+                num_samples_combined += 1
+
+            # removing the intermediate file
+            os.remove(split_manifest_path)
+
+    return num_samples_combined
+
+
+def create_cs_data(
+    intermediate_cs_manifest_list: list,
+    audio_save_folder: str,
+    manfest_save_path: str,
+    audio_amplitude_normalization: int,
+    pause_beg_msec: int,
+    pause_join_msec: int,
+    pause_end_msec: int,
+    cs_data_sampling_rate: int,
+    is_lid_manifest: bool,
+):
+    """
+    Args:
+        intermediate_cs_manifest_list: the intermediate cs manifest obtained from code_switching_manifest_creation.py as a list
+        audio_save_folder: Absolute path to save the generated audio samples
+        manfest_save_path: Absolute path to save the corresponding manifest
+        audio_amplitude_normalization: The amplitude to scale to after normalization
+        pause_beg_msec: Pause to be added at the beginning of the sample (msec)
+        pause_join_msec: Pause to be added between different phrases of the sample (msec)
+        pause_end_msec: Pause to be added at the end of the sample (msec)
+        cs_data_sampling_rate: Desired sampling rate of the generated samples
+        is_lid_manifest: If true, generate manifest in the multi-sample lid format, else the standard manifest format
+
+    Returns:
+
+    """
+
+    fs = cs_data_sampling_rate
+    incorrect_sample_flag = 0
+
+    with open(manfest_save_path, 'w') as outfile:
+        for data in tqdm(intermediate_cs_manifest_list):
+
+            combined_audio = []
+
+            staring_pause = np.zeros(int(pause_beg_msec * fs / 1000))
+            combined_audio += list(staring_pause)
+
+            text_entry_list = []
+            for index in range(len(data['lang_ids'])):
+
+                phrase_entry = {}
+                # dictionary to store the phrase information which will be added to the complete sentence
+
+                data_sample, fs_sample = librosa.load(data['paths'][index], sr=fs)
+                # Alternative-  fs_sample, data_sample = wavfile.read(data['paths'][index])
+
+                if fs_sample != fs:
+                    logging.error('Sampling rate error inside create_cs_data function')
+                    exit
+
+                # Remove leading and trailing zeros
+                data_sample = np.trim_zeros(data_sample)
+
+                # take care of empty arrays: rare
+                if data_sample.size == 0:
+                    incorrect_sample_flag = 1
+                    continue
+
+                # normalizing data
+                data_sample_norm = (
+                    data_sample
+                    / np.maximum(np.abs(data_sample.max()), np.abs(data_sample.min()))
+                    * audio_amplitude_normalization
+                )
+
+                combined_audio += list(data_sample_norm)
+
+                phrase_entry['str'] = data['texts'][index]
+                phrase_entry['lang'] = data['lang_ids'][index]
+
+                text_entry_list.append(phrase_entry)
+
+                # adding small pause between semgments
+                if index != (len(data['lang_ids']) - 1):
+                    pause = np.zeros(int(pause_join_msec * fs / 1000))
+                    combined_audio += list(pause)
+
+            if incorrect_sample_flag == 1:
+                incorrect_sample_flag = 0
+                continue
+
+            ending_pause = np.zeros(int(pause_end_msec * fs / 1000))
+            combined_audio += list(ending_pause)
+
+            sample_id = data['uid']
+            audio_file_path = audio_save_folder + '/' + str(sample_id) + ".wav"
+
+            # saving audio
+            wavfile.write(audio_file_path, fs, np.array(combined_audio).astype(np.int16))
+            # Alternative-  librosa.output.write_wav(audio_file_path, combined_audio, fs)
+
+            metadata_json = {}
+            metadata_json['audio_filepath'] = audio_file_path
+            metadata_json['duration'] = float(len(combined_audio) / fs)
+            if is_lid_manifest:
+                metadata_json['text'] = text_entry_list
+            else:
+                metadata_json['text'] = ' '.join(data['texts'])
+
+            metadata_json['language_ids'] = data['lang_ids']
+            metadata_json['original_texts'] = data['texts']
+            metadata_json['original_paths'] = data['paths']
+            metadata_json['original_durations'] = data['durations']
+
+            s = json.dumps(metadata_json)
+            outfile.write(s + '\n')
+
+
+def main():
+
+    cs_intermediate_manifest_path = args.manifest_path
+    audio_save_folder = args.audio_save_folder_path
+    manifest_save_path = args.manifest_save_path
+    audio_amplitude_normalization = args.audio_normalized_amplitude
+    pause_beg_msec = args.sample_beginning_pause_msec
+    pause_join_msec = args.sample_joining_pause_msec
+    pause_end_msec = args.sample_end_pause_msec
+    cs_data_sampling_rate = args.cs_data_sampling_rate
+    is_lid_manifest = args.is_lid_manifest
+    num_process = args.workers
+
+    # Sanity Checks
+    if (cs_intermediate_manifest_path is None) or (not os.path.exists(cs_intermediate_manifest_path)):
+        logging.error('Please provide correct CS manifest (obtained from code_switching_manifest_creation.py)')
+        exit
+
+    if (audio_save_folder is None) or (not os.path.exists(audio_save_folder)):
+        logging.error('audio_save_folder_path is incorrect or does not exist')
+        exit
+
+    if manifest_save_path is None:
+        logging.error('Please provide valid manifest_save_path')
+        exit
+
+    # Reading data
+    logging.info('Reading manifests')
+    intermediate_cs_manifest = read_manifest(cs_intermediate_manifest_path)
+
+    # Spliting the data
+    data_split = split_list(intermediate_cs_manifest, num_process)
+
+    # Creating Audio data
+    logging.info('Creating synthetic audio data')
+    base_directory = os.path.dirname(manifest_save_path)
+
+    Parallel(n_jobs=num_process)(
+        delayed(create_cs_data)(
+            split_manifest,
+            audio_save_folder,
+            base_directory + '/temp_' + str(idx) + '.json',
+            audio_amplitude_normalization,
+            pause_beg_msec,
+            pause_join_msec,
+            pause_end_msec,
+            cs_data_sampling_rate,
+            is_lid_manifest,
+        )
+        for idx, split_manifest in enumerate(data_split)
+    )
+
+    # Combining manifests
+    num_samples_combined = combine_manifests(manifest_save_path, num_process)
+
+    print("Synthetic CS audio data saved at :", audio_save_folder)
+    print("Synthetic CS manifest saved at :", manifest_save_path)
+    print("Total number of samples in the generated dataset :", str(num_samples_combined))
+
+    logging.info('Done!')
+
+
+if __name__ == "__main__":
+    main()

scripts/speech_recognition/code_switching/code_switching_manifest_creation.py

@@ -0,0 +1,177 @@
+# 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.
+
+import argparse
+import logging
+import os
+import random
+from nemo.collections.asr.parts.utils.manifest_utils import read_manifest, write_manifest
+
+# Checks -
+# (Recommendation) Please normalize the text for each language (avoid numbers, special characters, punctuation)
+# Please ensure that the audio_filepaths are absolute locations
+
+
+parser = argparse.ArgumentParser(description='Create synthetic code-switching data manifest from monolingual data')
+
+parser.add_argument("--manifest_language1", default=None, type=str, help='Manifest file for language 1', required=True)
+parser.add_argument("--manifest_language2", default=None, type=str, help='Manifest file for language 2', required=True)
+parser.add_argument(
+    "--manifest_save_path", default=None, type=str, help='Path to save created CS indermediate manifest', required=True
+)
+parser.add_argument(
+    "--id_language1", default=None, type=str, help='Identifier for language 1, eg: en, es, hi', required=True
+)
+parser.add_argument(
+    "--id_language2", default=None, type=str, help='Identifier for language 2, eg: en, es, hi', required=True
+)
+parser.add_argument("--max_sample_duration_sec", default=19, type=int, help='Maximum duration of sample (sec)')
+parser.add_argument("--min_sample_duration_sec", default=16, type=int, help='Minimum duration of sample (sec)')
+parser.add_argument("--dataset_size_required_hrs", default=1, type=int, help='Duration of dataset required (hrs)')
+
+args = parser.parse_args()
+
+
+def create_cs_manifest(
+    data_lang_0: list,
+    data_lang_1: list,
+    lid_lang_0: str,
+    lid_lang_1: str,
+    max_sample_duration_sec: int,
+    min_sample_duration_sec: int,
+    data_requirement_hrs: int,
+):
+    """
+    Args:
+        data_lang_0: Manifest entries from first langauge
+        data_lang_1: Manifest entries from second langauge
+        lid_lang_0: Language ID marker for first langauge
+        lid_lang_1: Language ID marker for second langauge
+        max_sample_duration_sec: Maximum permissible duration of generated CS sample in sec
+        min_sample_duration_sec: Minimum permissible duration of generated CS sample in sec
+        data_requirement_hrs: Required size of generated corpus
+
+    Returns:
+        Created synthetic CS manifest as list
+
+    """
+
+    total_duration = 0
+    constructed_data = []
+    sample_id = 0
+
+    num_samples_lang0 = len(data_lang_0)
+    num_samples_lang1 = len(data_lang_1)
+
+    while total_duration < (data_requirement_hrs * 3600):
+
+        created_sample_duration_sec = 0
+        created_sample_dict = {}
+        created_sample_dict['lang_ids'] = []
+        created_sample_dict['texts'] = []
+        created_sample_dict['paths'] = []
+        created_sample_dict['durations'] = []
+
+        while created_sample_duration_sec < min_sample_duration_sec:
+
+            lang_selection = random.randint(0, 1)
+
+            if lang_selection == 0:
+                index = random.randint(0, num_samples_lang0 - 1)
+                sample = data_lang_0[index]
+                lang_id = lid_lang_0
+            else:
+                index = random.randint(0, num_samples_lang1 - 1)
+                sample = data_lang_1[index]
+                lang_id = lid_lang_1
+
+            if (created_sample_duration_sec + sample['duration']) > max_sample_duration_sec:
+                continue
+            else:
+                created_sample_duration_sec += sample['duration']
+                created_sample_dict['lang_ids'].append(lang_id)
+                created_sample_dict['texts'].append(sample['text'])
+                created_sample_dict['paths'].append(sample['audio_filepath'])
+                created_sample_dict['durations'].append(sample['duration'])
+
+        created_sample_dict['total_duration'] = created_sample_duration_sec
+
+        # adding a uid which will be used to save the generated audio file later
+        created_sample_dict['uid'] = sample_id
+        sample_id += 1
+
+        constructed_data.append(created_sample_dict)
+        total_duration += created_sample_duration_sec
+
+    return constructed_data
+
+
+def main():
+
+    manifest0 = args.manifest_language1
+    manifest1 = args.manifest_language2
+    lid0 = args.id_language1
+    lid1 = args.id_language2
+    min_sample_duration = args.min_sample_duration_sec
+    max_sample_duration = args.max_sample_duration_sec
+    dataset_requirement = args.dataset_size_required_hrs
+    manifest_save_path = args.manifest_save_path
+
+    # Sanity Checks
+    if (manifest0 is None) or (not os.path.exists(manifest0)):
+        logging.error('Manifest for language 1 is incorrect')
+        exit
+
+    if (manifest1 is None) or (not os.path.exists(manifest1)):
+        logging.error('Manifest for language 2 is incorrect')
+        exit
+
+    if lid0 is None:
+        logging.error('Please provide correct language code for language 1')
+        exit
+
+    if lid1 is None:
+        logging.error('Please provide correct language code for language 2')
+        exit
+
+    if manifest_save_path is None:
+        logging.error('Please provide correct manifest save path')
+        exit
+
+    if min_sample_duration >= max_sample_duration:
+        logging.error('Please ensure max_sample_duration > min_sample_duration')
+        exit
+
+    # Reading data
+    logging.info('Reading manifests')
+    data_language0 = read_manifest(manifest0)
+    data_language1 = read_manifest(manifest1)
+
+    # Creating the CS data Manifest
+    logging.info('Creating CS manifest')
+    constructed_data = create_cs_manifest(
+        data_language0, data_language1, lid0, lid1, max_sample_duration, min_sample_duration, dataset_requirement
+    )
+
+    # Saving Manifest
+    logging.info('saving manifest')
+    write_manifest(manifest_save_path, constructed_data)
+
+    print("Synthetic CS manifest saved at :", manifest_save_path)
+
+    logging.info('Done!')
+
+
+if __name__ == "__main__":
+    main()

scripts/speech_recognition/confidence/benchmark_asr_confidence.py

@@ -0,0 +1,308 @@
+# 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.
+
+import json
+import os
+from dataclasses import dataclass, field, is_dataclass
+from pathlib import Path
+from typing import Optional
+
+import lightning.pytorch as pl
+import torch
+from omegaconf import MISSING, OmegaConf
+from sklearn.model_selection import ParameterGrid
+
+from nemo.collections.asr.models import ASRModel, EncDecRNNTModel
+from nemo.collections.asr.parts.submodules.ctc_decoding import CTCDecodingConfig
+from nemo.collections.asr.parts.submodules.rnnt_decoding import RNNTDecodingConfig
+from nemo.collections.asr.parts.utils.asr_confidence_benchmarking_utils import (
+    apply_confidence_parameters,
+    run_confidence_benchmark,
+)
+from nemo.collections.asr.parts.utils.asr_confidence_utils import ConfidenceConfig
+from nemo.core.config import hydra_runner
+from nemo.utils import logging, model_utils
+
+"""
+Get confidence metrics and curve plots for a given model, dataset, and confidence parameters.
+
+# Arguments
+  model_path: Path to .nemo ASR checkpoint
+  pretrained_name: Name of pretrained ASR model (from NGC registry)
+  dataset_manifest: Path to dataset JSON manifest file (in NeMo format)
+  output_dir: Output directory to store a report and curve plot directories
+  
+  batch_size: batch size during inference
+  num_workers: number of workers during inference
+  
+  cuda: Optional int to enable or disable execution of model on certain CUDA device
+  amp: Bool to decide if Automatic Mixed Precision should be used during inference
+  audio_type: Str filetype of the audio. Supported = wav, flac, mp3
+  
+  target_level: Word- or token-level confidence. Supported = word, token, auto (for computing both word and token)
+  confidence_cfg: Config with confidence parameters
+  grid_params: Dictionary with lists of parameters to iteratively benchmark on
+
+# Usage
+ASR model can be specified by either "model_path" or "pretrained_name".
+Data for transcription are defined with "dataset_manifest".
+Results are returned as a benchmark report and curve plots.
+
+python benchmark_asr_confidence.py \
+    model_path=null \
+    pretrained_name=null \
+    dataset_manifest="" \
+    output_dir="" \
+    batch_size=64 \
+    num_workers=8 \
+    cuda=0 \
+    amp=True \
+    target_level="word" \
+    confidence_cfg.exclude_blank=False \
+    'grid_params="{\"aggregation\": [\"min\", \"prod\"], \"alpha\": [0.33, 0.5]}"'
+"""
+
+
+def get_experiment_params(cfg):
+    """Get experiment parameters from a confidence config and generate the experiment name.
+
+    Returns:
+        List of experiment parameters.
+        String with the experiment name.
+    """
+    blank = "no_blank" if cfg.exclude_blank else "blank"
+    duration = "duration" if cfg.tdt_include_duration else "no_duration"
+    aggregation = cfg.aggregation
+    method_name = cfg.method_cfg.name
+    alpha = cfg.method_cfg.alpha
+    if method_name == "entropy":
+        entropy_type = cfg.method_cfg.entropy_type
+        entropy_norm = cfg.method_cfg.entropy_norm
+        experiment_param_list = [
+            aggregation,
+            str(cfg.exclude_blank),
+            str(cfg.tdt_include_duration),
+            method_name,
+            entropy_type,
+            entropy_norm,
+            str(alpha),
+        ]
+        experiment_str = "-".join([aggregation, blank, duration, method_name, entropy_type, entropy_norm, str(alpha)])
+    else:
+        experiment_param_list = [
+            aggregation,
+            str(cfg.exclude_blank),
+            str(cfg.tdt_include_duration),
+            method_name,
+            "-",
+            "-",
+            str(alpha),
+        ]
+        experiment_str = "-".join([aggregation, blank, duration, method_name, str(alpha)])
+    return experiment_param_list, experiment_str
+
+
+@dataclass
+class ConfidenceBenchmarkingConfig:
+    # Required configs
+    model_path: Optional[str] = None  # Path to a .nemo file
+    pretrained_name: Optional[str] = None  # Name of a pretrained model
+    dataset_manifest: str = MISSING
+    output_dir: str = MISSING
+
+    # General configs
+    batch_size: int = 32
+    num_workers: int = 4
+
+    # Set `cuda` to int to define CUDA device. If 'None', will look for CUDA
+    # device anyway, and do inference on CPU only if CUDA device is not found.
+    # If `cuda` is a negative number, inference will be on CPU only.
+    cuda: Optional[int] = None
+    amp: bool = False
+    audio_type: str = "wav"
+
+    # Confidence configs
+    target_level: str = "auto"  # Choices: "word", "token", "auto" (for both word- and token-level confidence)
+    confidence_cfg: ConfidenceConfig = field(
+        default_factory=lambda: ConfidenceConfig(preserve_word_confidence=True, preserve_token_confidence=True)
+    )
+    grid_params: Optional[str] = None  # a dictionary with lists of parameters to iteratively benchmark on
+
+
+@hydra_runner(config_name="ConfidenceBenchmarkingConfig", schema=ConfidenceBenchmarkingConfig)
+def main(cfg: ConfidenceBenchmarkingConfig):
+    torch.set_grad_enabled(False)
+
+    logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}')
+
+    if is_dataclass(cfg):
+        cfg = OmegaConf.structured(cfg)
+
+    if cfg.model_path is None and cfg.pretrained_name is None:
+        raise ValueError("Both cfg.model_path and cfg.pretrained_name cannot be None!")
+
+    # setup GPU
+    if cfg.cuda is None:
+        if torch.cuda.is_available():
+            device = [0]  # use 0th CUDA device
+            accelerator = 'gpu'
+        else:
+            device = 1
+            accelerator = 'cpu'
+    else:
+        device = [cfg.cuda]
+        accelerator = 'gpu'
+
+    map_location = torch.device('cuda:{}'.format(device[0]) if accelerator == 'gpu' else 'cpu')
+
+    # setup model
+    if cfg.model_path is not None:
+        # restore model from .nemo file path
+        model_cfg = ASRModel.restore_from(restore_path=cfg.model_path, return_config=True)
+        classpath = model_cfg.target  # original class path
+        imported_class = model_utils.import_class_by_path(classpath)  # type: ASRModel
+        logging.info(f"Restoring model : {imported_class.__name__}")
+        asr_model = imported_class.restore_from(
+            restore_path=cfg.model_path, map_location=map_location
+        )  # type: ASRModel
+    else:
+        # restore model by name
+        asr_model = ASRModel.from_pretrained(
+            model_name=cfg.pretrained_name, map_location=map_location
+        )  # type: ASRModel
+
+    trainer = pl.Trainer(devices=device, accelerator=accelerator)
+    asr_model.set_trainer(trainer)
+    asr_model = asr_model.eval()
+
+    # Check if ctc or rnnt model
+    is_rnnt = isinstance(asr_model, EncDecRNNTModel)
+
+    # Check that the model has the `change_decoding_strategy` method
+    if not hasattr(asr_model, 'change_decoding_strategy'):
+        raise RuntimeError("The asr_model you are using must have the `change_decoding_strategy` method.")
+
+    # get filenames and reference texts from manifest
+    filepaths = []
+    reference_texts = []
+    if os.stat(cfg.dataset_manifest).st_size == 0:
+        logging.error(f"The input dataset_manifest {cfg.dataset_manifest} is empty. Exiting!")
+        return None
+    manifest_dir = Path(cfg.dataset_manifest).parent
+    with open(cfg.dataset_manifest, 'r') as f:
+        for line in f:
+            item = json.loads(line)
+            audio_file = Path(item['audio_filepath'])
+            if not audio_file.is_file() and not audio_file.is_absolute():
+                audio_file = manifest_dir / audio_file
+            filepaths.append(str(audio_file.absolute()))
+            reference_texts.append(item['text'])
+
+    # do grid-based benchmarking if grid_params is provided, otherwise a regular one
+    work_dir = Path(cfg.output_dir)
+    os.makedirs(work_dir, exist_ok=True)
+    report_legend = (
+        ",".join(
+            [
+                "model_type",
+                "aggregation",
+                "blank",
+                "duration",
+                "method_name",
+                "entropy_type",
+                "entropy_norm",
+                "alpha",
+                "target_level",
+                "auc_roc",
+                "auc_pr",
+                "auc_nt",
+                "nce",
+                "ece",
+                "auc_yc",
+                "std_yc",
+                "max_yc",
+            ]
+        )
+        + "\n"
+    )
+    model_typename = "RNNT" if is_rnnt else "CTC"
+    report_file = work_dir / Path("report.csv")
+    if cfg.grid_params:
+        asr_model.change_decoding_strategy(
+            RNNTDecodingConfig(fused_batch_size=-1, strategy="greedy_batch", confidence_cfg=cfg.confidence_cfg)
+            if is_rnnt
+            else CTCDecodingConfig(confidence_cfg=cfg.confidence_cfg)
+        )
+        params = json.loads(cfg.grid_params)
+        hp_grid = ParameterGrid(params)
+        hp_grid = list(hp_grid)
+
+        logging.info(f"==============================Running a benchmarking with grid search=========================")
+        logging.info(f"Grid search size: {len(hp_grid)}")
+        logging.info(f"Results will be written to:\nreport file `{report_file}`\nand plot directories near the file")
+        logging.info(f"==============================================================================================")
+
+        with open(report_file, "tw", encoding="utf-8") as f:
+            f.write(report_legend)
+            f.flush()
+            for i, hp in enumerate(hp_grid):
+                logging.info(f"Run # {i + 1}, grid: `{hp}`")
+                asr_model.change_decoding_strategy(apply_confidence_parameters(asr_model.cfg.decoding, hp))
+                param_list, experiment_name = get_experiment_params(asr_model.cfg.decoding.confidence_cfg)
+                plot_dir = work_dir / Path(experiment_name)
+                results = run_confidence_benchmark(
+                    asr_model,
+                    cfg.target_level,
+                    filepaths,
+                    reference_texts,
+                    cfg.batch_size,
+                    cfg.num_workers,
+                    plot_dir,
+                    cfg.amp,
+                )
+                for level, result in results.items():
+                    f.write(f"{model_typename},{','.join(param_list)},{level},{','.join([str(r) for r in result])}\n")
+                    f.flush()
+    else:
+        asr_model.change_decoding_strategy(
+            RNNTDecodingConfig(fused_batch_size=-1, strategy="greedy_batch", confidence_cfg=cfg.confidence_cfg)
+            if is_rnnt
+            else CTCDecodingConfig(confidence_cfg=cfg.confidence_cfg)
+        )
+        param_list, experiment_name = get_experiment_params(asr_model.cfg.decoding.confidence_cfg)
+        plot_dir = work_dir / Path(experiment_name)
+
+        logging.info(f"==============================Running a single benchmarking===================================")
+        logging.info(f"Results will be written to:\nreport file `{report_file}`\nand plot directory `{plot_dir}`")
+
+        with open(report_file, "tw", encoding="utf-8") as f:
+            f.write(report_legend)
+            f.flush()
+            results = run_confidence_benchmark(
+                asr_model,
+                cfg.batch_size,
+                cfg.num_workers,
+                cfg.target_level,
+                filepaths,
+                reference_texts,
+                plot_dir,
+                cfg.amp,
+            )
+            for level, result in results.items():
+                f.write(f"{model_typename},{','.join(param_list)},{level},{','.join([str(r) for r in result])}\n")
+    logging.info(f"===========================================Done===============================================")
+
+
+if __name__ == '__main__':
+    main()