chore: import upstream snapshot with attribution

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

scripts/speech_recognition/convert_hf_dataset_to_nemo.py

@@ -0,0 +1,413 @@
+# 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.
+
+"""
+Python wrapper over HuggingFace Datasets to create preprocessed NeMo ASR Datasets.
+
+List of HuggingFace datasets : https://huggingface.co/datasets
+(Please filter by task: automatic-speech-recognition)
+
+# Setup
+After installation of huggingface datasets (pip install datasets), some datasets might require authentication
+- for example Mozilla Common Voice. You should go to the above link, register as a user and generate an API key.
+
+## Authenticated Setup Steps
+
+Website steps:
+- Visit https://huggingface.co/settings/profile
+- Visit "Access Tokens" on list of items.
+- Create new token - provide a name for the token and "read" access is sufficient.
+  - PRESERVE THAT TOKEN API KEY. You can copy that key for next step.
+- Visit the HuggingFace Dataset page for Mozilla Common Voice
+  - There should be a section that asks you for your approval.
+  - Make sure you are logged in and then read that agreement.
+  - If and only if you agree to the text, then accept the terms.
+
+Code steps:
+- Now on your machine, run `huggingface-cli login`
+- Paste your preserved HF TOKEN API KEY (from above).
+
+Now you should be logged in. When running the script, dont forget to set `use_auth_token=True` !
+
+# Usage
+The script supports two modes, but the offline mode is the preferred mechanism. The drawback of the offline mode
+is that it requires 3 copies of the dataset to exist simultanously -
+
+1) The .arrow files for HF cache
+2) The extracted dataset in HF cache
+3) The preprocessed audio files preserved in the output_dir provided in the script.
+
+Due to this, make sure your HDD is large enough to store the processed dataset !
+
+## Usage - Offline Mode
+
+python convert_hf_dataset_to_nemo.py \
+    output_dir=<Path to some storage drive that will hold preprocessed audio files> \
+    path=<`path` argument in HF datasets, cannot be null> \
+    name=<`name` argument in HF datasets, can be null> \
+    split=<`split` argument in HF datasets, can be null> \
+    use_auth_token=<Can be `True` or `False` depending on whether the dataset requires authentication>
+
+This will create an output directory of multiple sub-folders containing the preprocessed .wav files,
+along with a nemo compatible JSON manifest file.
+
+NOTE:
+    The JSON manifest itself is not preprocessed ! You should perform text normalization, and cleanup
+    inconsistent text by using NeMo Text Normalization tool and Speech Data Explorer toolkit !
+
+## Usage - Streaming Mode
+
+NOTE:
+    This mode is not well supported. It trades of speed for storage by only having one copy of the dataset in
+    output_dir, however the speed of processing is around 10x slower than offline mode. Some datasets (such as MCV)
+    fail to run entirely.
+
+    DO NOT USE if you have sufficient disk space.
+
+python convert_hf_dataset_to_nemo.py \
+    ... all the arguments from above \
+    streaming=True
+
+"""
+
+import json
+import os
+import traceback
+from dataclasses import dataclass, field, is_dataclass
+from typing import Optional
+
+import hydra
+import librosa
+import soundfile
+import tqdm
+from datasets import Audio, Dataset, IterableDataset, load_dataset
+from hydra.conf import HydraConf, RunDir
+from hydra.core.config_store import ConfigStore
+from omegaconf import OmegaConf
+
+
+@dataclass
+class HFDatasetConversionConfig:
+    # Nemo Dataset info
+    output_dir: str  # path to output directory where the files will be saved
+
+    # HF Dataset info
+    path: str  # HF dataset path
+    name: Optional[str] = None  # name of the dataset subset
+    split: Optional[str] = None  # split of the dataset subset
+    use_auth_token: bool = False  # whether authentication token should be passed or not (Required for MCV)
+
+    # NeMo dataset conversion
+    sampling_rate: int = 16000
+    streaming: bool = False  # Whether to use Streaming dataset API. [NOT RECOMMENDED]
+    num_proc: int = -1
+    ensure_ascii: bool = True  # When saving the JSON entry, whether to ensure ascii.
+
+    # Placeholders. Generated internally.
+    resolved_output_dir: str = ''
+    split_output_dir: Optional[str] = None
+
+    hydra: HydraConf = field(default_factory=lambda: HydraConf(run=RunDir(dir=".")))
+
+
+def prepare_output_dirs(cfg: HFDatasetConversionConfig):
+    """
+    Prepare output directories and subfolders as needed.
+    Also prepare the arguments of the config with these directories.
+    """
+    output_dir = os.path.abspath(cfg.output_dir)
+    output_dir = os.path.join(output_dir, cfg.path)
+
+    if cfg.name is not None:
+        output_dir = os.path.join(output_dir, cfg.name)
+
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir, exist_ok=True)
+
+    cfg.resolved_output_dir = output_dir
+    cfg.split_output_dir = None
+
+
+def infer_dataset_segments(batch):
+    """
+    Helper method to run in batch mode over a mapped Dataset.
+
+    Infers the path of the subdirectories for the dataset, removing {extracted/HASH}.
+
+    Returns:
+        A cleaned list of path segments
+    """
+    segments = []
+    segment, path = os.path.split(batch['audio']['path'])
+    segments.insert(0, path)
+    while segment not in ('', os.path.sep):
+        segment, path = os.path.split(segment)
+        segments.insert(0, path)
+
+    if 'extracted' in segments:
+        index_of_basedir = segments.index("extracted")
+        segments = segments[(index_of_basedir + 1 + 1) :]  # skip .../extracted/{hash}/
+
+    return segments
+
+
+def prepare_audio_filepath(audio_filepath):
+    """
+    Helper method to run in batch mode over a mapped Dataset.
+
+    Prepares the audio filepath and its subdirectories. Remaps the extension to .wav file.
+
+    Args:
+        audio_filepath: String path to the audio file.
+
+    Returns:
+        Cleaned filepath renamed to be a wav file.
+    """
+    audio_basefilepath = os.path.split(audio_filepath)[0]
+    if not os.path.exists(audio_basefilepath):
+        os.makedirs(audio_basefilepath, exist_ok=True)
+
+    # Remove temporary fmt file
+    if os.path.exists(audio_filepath):
+        os.remove(audio_filepath)
+
+    # replace any ext with .wav
+    audio_filepath, ext = os.path.splitext(audio_filepath)
+    audio_filepath = audio_filepath + '.wav'
+
+    # Remove previous run file
+    if os.path.exists(audio_filepath):
+        os.remove(audio_filepath)
+    return audio_filepath
+
+
+def build_map_dataset_to_nemo_func(cfg: HFDatasetConversionConfig, basedir):
+    """
+    Helper method to run in batch mode over a mapped Dataset.
+
+    Creates a function that can be passed to Dataset.map() containing the config and basedir.
+    Useful to map a HF dataset to NeMo compatible format in an efficient way for offline processing.
+
+    Returns:
+        A function pointer which can be used for Dataset.map()
+    """
+
+    def map_dataset_to_nemo(batch):
+        # Write audio file to correct path
+        if cfg.streaming:
+            batch['audio_filepath'] = batch['audio']['path'].split("::")[0].replace("zip://", "")
+        else:
+            segments = infer_dataset_segments(batch)
+            audio_filepath = os.path.join(*segments)
+            batch['audio_filepath'] = audio_filepath
+
+        batch['audio_filepath'] = os.path.abspath(os.path.join(basedir, batch['audio_filepath']))
+        audio_filepath = batch['audio_filepath']
+        audio_filepath = prepare_audio_filepath(audio_filepath)
+        batch['audio_filepath'] = audio_filepath  # update filepath with prepared path
+
+        soundfile.write(audio_filepath, batch['audio']['array'], samplerate=cfg.sampling_rate, format='wav')
+
+        batch['duration'] = librosa.get_duration(y=batch['audio']['array'], sr=batch['audio']['sampling_rate'])
+        return batch
+
+    return map_dataset_to_nemo
+
+
+def convert_offline_dataset_to_nemo(
+    dataset: Dataset,
+    cfg: HFDatasetConversionConfig,
+    basedir: str,
+    manifest_filepath: str,
+):
+    """
+    Converts a HF dataset to a audio-preprocessed Nemo dataset in Offline mode.
+    Also writes out a nemo compatible manifest file.
+
+    Args:
+        dataset: Iterable HF Dataset.
+        cfg: HFDatasetConvertionConfig.
+        basedir: Base output directory.
+        manifest_filepath: Filepath of manifest.
+    """
+    num_proc = cfg.num_proc
+    if num_proc < 0:
+        num_proc = max(1, os.cpu_count() // 2)
+
+    dataset = dataset.map(build_map_dataset_to_nemo_func(cfg, basedir), num_proc=num_proc)
+    ds_iter = iter(dataset)
+
+    with open(manifest_filepath, 'w') as manifest_f:
+        for idx, sample in enumerate(
+            tqdm.tqdm(
+                ds_iter, desc=f'Processing {cfg.path} (split : {cfg.split}):', total=len(dataset), unit=' samples'
+            )
+        ):
+            # remove large components from sample
+            del sample['audio']
+            if 'file' in sample:
+                del sample['file']
+            manifest_f.write(f"{json.dumps(sample, ensure_ascii=cfg.ensure_ascii)}\n")
+
+
+def convert_streaming_dataset_to_nemo(
+    dataset: IterableDataset, cfg: HFDatasetConversionConfig, basedir: str, manifest_filepath: str
+):
+    """
+    Converts a HF dataset to a audio-preprocessed Nemo dataset in Streaming mode.
+    Also writes out a nemo compatible manifest file.
+
+    Args:
+        dataset: Iterable HF Dataset.
+        cfg: HFDatasetConvertionConfig.
+        basedir: Base output directory.
+        manifest_filepath: Filepath of manifest.
+    """
+    # Disable until fix https://github.com/huggingface/datasets/pull/3556 is merged
+    # dataset = dataset.map(build_map_dataset_to_nemo_func(cfg, basedir))
+
+    ds_iter = iter(dataset)
+
+    with open(manifest_filepath, 'w') as manifest_f:
+        for idx, sample in enumerate(
+            tqdm.tqdm(ds_iter, desc=f'Processing {cfg.path} (split: {cfg.split}):', unit=' samples')
+        ):
+
+            audio_filepath = sample['audio']['path'].split("::")[0].replace("zip://", "")
+            audio_filepath = os.path.abspath(os.path.join(basedir, audio_filepath))
+            audio_filepath = prepare_audio_filepath(audio_filepath)
+
+            soundfile.write(audio_filepath, sample['audio']['array'], samplerate=cfg.sampling_rate, format='wav')
+
+            manifest_line = {
+                'audio_filepath': audio_filepath,
+                'text': sample['text'],
+                'duration': librosa.get_duration(sample['audio']['array'], sr=cfg.sampling_rate),
+            }
+
+            # remove large components from sample
+            del sample['audio']
+            del sample['text']
+            if 'file' in sample:
+                del sample['file']
+
+            manifest_line.update(sample)
+
+            manifest_f.write(f"{json.dumps(sample, ensure_ascii=cfg.ensure_ascii)}\n")
+
+
+def process_dataset(dataset: IterableDataset, cfg: HFDatasetConversionConfig):
+    """
+    Top level method that processes a given IterableDataset to Nemo compatible dataset.
+    It also writes out a nemo compatible manifest file.
+
+    Args:
+        dataset: HF Dataset.
+        cfg: HFDatasetConvertionConfig
+    """
+    dataset = dataset.cast_column("audio", Audio(cfg.sampling_rate, mono=True))
+
+    # for Common Voice, "sentence" is used instead of "text" to store the transcript.
+    if 'sentence' in dataset.features:
+        dataset = dataset.rename_column("sentence", "text")
+
+    if cfg.split_output_dir is None:
+        basedir = cfg.resolved_output_dir
+        manifest_filename = f"{cfg.path.replace('/', '_')}_manifest.json"
+    else:
+        basedir = cfg.split_output_dir
+        split = os.path.split(cfg.split_output_dir)[-1]
+        manifest_filename = f"{split}_{cfg.path.replace('/', '_')}_manifest.json"
+
+        if not os.path.exists(cfg.split_output_dir):
+            os.makedirs(cfg.split_output_dir, exist_ok=True)
+
+        cfg.split = split
+
+    manifest_filepath = os.path.abspath(os.path.join(basedir, manifest_filename))
+
+    if cfg.streaming:
+        convert_streaming_dataset_to_nemo(dataset, cfg, basedir=basedir, manifest_filepath=manifest_filepath)
+    else:
+        convert_offline_dataset_to_nemo(dataset, cfg, basedir=basedir, manifest_filepath=manifest_filepath)
+
+    print()
+    print("Dataset conversion finished !")
+
+
+@hydra.main(config_name='hfds_config', config_path=None)
+def main(cfg: HFDatasetConversionConfig):
+    # Convert dataclass to omegaconf
+    if is_dataclass(cfg):
+        cfg = OmegaConf.structured(cfg)
+
+    # Prepare output subdirs
+    prepare_output_dirs(cfg)
+
+    # Load dataset in offline/streaming mode
+    dataset = None
+    try:
+        dataset = load_dataset(
+            path=cfg.path,
+            name=cfg.name,
+            split=cfg.split,
+            cache_dir=None,
+            streaming=cfg.streaming,
+            token=cfg.use_auth_token,
+            trust_remote_code=True,
+        )
+
+    except Exception as e:
+        print(
+            "HuggingFace datasets failed due to some reason (stack trace below). \nFor certain datasets (eg: MCV), "
+            "it may be necessary to login to the huggingface-cli (via `huggingface-cli login`).\n"
+            "Once logged in, you need to set `use_auth_token=True` when calling this script.\n\n"
+            "Traceback error for reference :\n"
+        )
+        print(traceback.format_exc())
+        exit(1)
+
+    # Multiple datasets were provided at once, process them one by one into subdirs.
+    if isinstance(dataset, dict):
+        print()
+        print("Multiple splits found for dataset", cfg.path, ":", list(dataset.keys()))
+
+        keys = list(dataset.keys())
+        for key in keys:
+            ds_split = dataset[key]
+            print(f"Processing split {key} for dataset {cfg.path}")
+
+            cfg.split_output_dir = os.path.join(cfg.resolved_output_dir, key)
+            process_dataset(ds_split, cfg)
+
+            del dataset[key], ds_split
+
+        # reset the split output directory
+        cfg.split_output_dir = None
+
+    else:
+        # Single dataset was found, process into resolved directory.
+        print("Single split found for dataset", cfg.path, "| Split chosen =", cfg.split)
+
+        if cfg.split is not None:
+            cfg.split_output_dir = os.path.join(cfg.resolved_output_dir, cfg.split)
+
+        process_dataset(dataset, cfg)
+
+
+# Register the dataclass as a valid config
+ConfigStore.instance().store(name='hfds_config', node=HFDatasetConversionConfig)
+
+if __name__ == '__main__':
+    main()

scripts/speech_recognition/convert_torchaudio_nemo.py

@@ -0,0 +1,157 @@
+#!/usr/bin/env python3
+# Copyright (c) 2026, 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.
+
+"""
+Convert .nemo checkpoints that were trained with ``preprocessor.use_torchaudio=True``
+to the current format (non-torchaudio FilterbankFeatures).
+
+After torchaudio was removed as a dependency (PR #15211), models trained with the
+torchaudio-based preprocessor (FilterbankFeaturesTA) fail to load because the
+state dict keys no longer match:
+
+    Old (torchaudio):
+        preprocessor.featurizer._mel_spec_extractor.spectrogram.window
+        preprocessor.featurizer._mel_spec_extractor.mel_scale.fb
+
+    New (current):
+        preprocessor.featurizer.window
+        preprocessor.featurizer.fb
+
+This script renames those keys and also sets ``use_torchaudio: false`` in the model
+config so that the correct featurizer class is instantiated on load.
+
+Usage
+-----
+    python convert_torchaudio_nemo.py --nemo_file model.nemo --output_file model_converted.nemo
+"""
+
+import argparse
+import os
+import tarfile
+import tempfile
+
+import torch
+import yaml
+
+from nemo.utils.tar_utils import safe_extract
+
+
+MODEL_CONFIG_YAML = "model_config.yaml"
+MODEL_WEIGHTS_CKPT = "model_weights.ckpt"
+
+# Old torchaudio key suffix -> new key suffix
+KEY_MIGRATION = {
+    "featurizer._mel_spec_extractor.spectrogram.window": "featurizer.window",
+    "featurizer._mel_spec_extractor.mel_scale.fb": "featurizer.fb",
+}
+
+
+def migrate_state_dict(state_dict: dict) -> tuple[dict, list[tuple[str, str]]]:
+    """Rename torchaudio-era keys.  Returns (new_state_dict, list of (old, new) renames)."""
+    renames = []
+    for key in list(state_dict.keys()):
+        for old_suffix, new_suffix in KEY_MIGRATION.items():
+            if key.endswith(old_suffix):
+                new_key = key[: -len(old_suffix)] + new_suffix
+                if "featurizer.fb" in new_suffix:
+                    state_dict[new_key] = state_dict.pop(key).T.unsqueeze(0)
+                else:
+                    state_dict[new_key] = state_dict.pop(key)
+                renames.append((key, new_key))
+                break
+    return state_dict, renames
+
+
+def migrate_config(cfg: dict) -> bool:
+    """Set ``use_torchaudio: false`` in the preprocessor config.  Returns True if changed."""
+    preprocessor = cfg.get("preprocessor", {})
+    if preprocessor.get("use_torchaudio", False):
+        preprocessor["use_torchaudio"] = False
+        return True
+    return False
+
+
+def convert_nemo_file(nemo_path: str, output_path: str) -> None:
+    """Extract, migrate, and repack a .nemo archive."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # --- Unpack --------------------------------------------------------
+        # Older checkpoints may be gzipped; newer ones are plain tar.
+        try:
+            tar = tarfile.open(nemo_path, "r:")
+        except tarfile.ReadError:
+            tar = tarfile.open(nemo_path, "r:gz")
+        with tar:
+            safe_extract(tar, tmpdir)
+
+        # --- Migrate state dict --------------------------------------------
+        weights_path = os.path.join(tmpdir, MODEL_WEIGHTS_CKPT)
+        if not os.path.isfile(weights_path):
+            raise FileNotFoundError(
+                f"Could not find {MODEL_WEIGHTS_CKPT} inside the .nemo archive. "
+                "Are you sure this is a valid .nemo file?"
+            )
+
+        state_dict = torch.load(weights_path, map_location="cpu", weights_only=True)
+        state_dict, renames = migrate_state_dict(state_dict)
+        if not renames:
+            print("No torchaudio keys found in state dict — nothing to migrate.")
+            return
+
+        for old, new in renames:
+            print(f"  Renamed: {old}  ->  {new}")
+
+        torch.save(state_dict, weights_path)
+
+        # --- Migrate config ------------------------------------------------
+        config_path = os.path.join(tmpdir, MODEL_CONFIG_YAML)
+        if os.path.isfile(config_path):
+            with open(config_path) as f:
+                cfg = yaml.safe_load(f)
+            if migrate_config(cfg):
+                print("  Config:  set use_torchaudio=false")
+                with open(config_path, "w") as f:
+                    yaml.dump(cfg, f, default_flow_style=False)
+
+        # --- Repack --------------------------------------------------------
+        with tarfile.open(output_path, "w:") as tar:
+            tar.add(tmpdir, arcname=".")
+
+    print(f"\nConverted checkpoint saved to: {output_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Convert .nemo checkpoints from torchaudio preprocessor format to the current format.",
+    )
+    parser.add_argument(
+        "--nemo_file",
+        required=True,
+        help="Path to the source .nemo file.",
+    )
+    parser.add_argument(
+        "--output_file",
+        required=True,
+        help="Path to write the converted .nemo file.",
+    )
+    args = parser.parse_args()
+
+    if not os.path.isfile(args.nemo_file):
+        raise FileNotFoundError(f"File not found: {args.nemo_file}")
+
+    convert_nemo_file(args.nemo_file, args.output_file)
+
+
+if __name__ == "__main__":
+    main()

scripts/speech_recognition/create_dali_tarred_dataset_index.py

@@ -0,0 +1,95 @@
+# 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 glob
+import logging
+import os
+from dataclasses import dataclass
+
+import hydra
+from hydra.core.config_store import ConfigStore
+from joblib import Parallel, delayed
+from omegaconf import MISSING
+
+try:
+    from wds2idx import IndexCreator
+
+    INDEX_CREATOR_AVAILABLE = True
+except (ImportError, ModuleNotFoundError):
+    INDEX_CREATOR_AVAILABLE = False
+
+"""
+python create_dali_tarred_dataset_index.py \
+    tar_dir=<path to the directory which contains tarred dataset> \
+    workers=-1
+
+"""
+
+logging.basicConfig(level=logging.INFO)
+
+
+@dataclass
+class DALITarredIndexConfig:
+    tar_dir: str = MISSING  # Path to the existing dataset's manifest
+    workers: int = -1  # number of worker processes
+
+
+def process_index_path(tar_paths, index_dir):
+    """
+    Appends the folder `{index_dir}` to the filepath of all tarfiles.
+    Example:
+         /X/Y/Z/audio_0.tar -> /X/Y/Z/{index_dir}/audio_0.index
+    """
+    index_paths = []
+    for path in tar_paths:
+        basepath, filename = os.path.split(path)
+        path = filename.replace('.tar', '.index')
+        path = os.path.join(basepath, path)
+        base, name = os.path.split(path)
+        index_path = os.path.join(index_dir, name)
+        index_paths.append(index_path)
+
+    return index_paths
+
+
+def build_index(tarpath, indexfile):
+    with IndexCreator(tarpath, indexfile) as index:
+        index.create_index()
+
+
+@hydra.main(config_path=None, config_name='index_config', version_base="1.1")
+def main(cfg: DALITarredIndexConfig):
+    if not INDEX_CREATOR_AVAILABLE:
+        logging.error("`wds2idx` is not installed. Please install NVIDIA DALI >= 1.11")
+        exit(1)
+
+    tar_files = list(glob.glob(os.path.join(cfg.tar_dir, "*.tar")))
+
+    index_dir = os.path.join(cfg.tar_dir, "dali_index")
+    if not os.path.exists(index_dir):
+        os.makedirs(index_dir, exist_ok=True)
+
+    index_paths = process_index_path(tar_files, index_dir)
+
+    with Parallel(n_jobs=cfg.workers, verbose=len(tar_files)) as parallel:
+        _ = parallel(delayed(build_index)(tarpath, indexfile) for tarpath, indexfile in zip(tar_files, index_paths))
+
+    logging.info("Finished constructing index files !")
+
+
+ConfigStore.instance().store(name='index_config', node=DALITarredIndexConfig)
+
+
+if __name__ == '__main__':
+    main()

scripts/speech_recognition/estimate_data_weights.py

@@ -0,0 +1,172 @@
+# 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 contextlib import contextmanager
+from typing import Sequence
+
+import click
+import numpy as np
+from omegaconf import DictConfig, ListConfig, OmegaConf
+
+from nemo.collections.common.data.lhotse.cutset import get_parser_fn
+
+
+@click.command()
+@click.argument("input_cfgs", type=click.Path(exists=True, dir_okay=False), nargs=-1)
+@click.argument("output_cfg", type=click.Path())
+@click.option(
+    "-t",
+    "--temperature",
+    type=float,
+    default=None,
+    multiple=True,
+    help="Temperature for re-weighting datasets. 1 is a neutral value. "
+    "Lower temperature over-samples smaller datasets, and vice versa. "
+    "Can be specified multiple times to apply a different temperature to each group level in the YAML config.",
+)
+@click.option(
+    "-s",
+    "--strategy",
+    type=click.Choice(["num_hours", "num_examples"]),
+    default="num_hours",
+    help="Strategy for choosing weights for each dataset.",
+)
+def estimate_data_weights(input_cfgs: str, output_cfg: str, temperature: list[float], strategy: str):
+    """
+    Read a YAML specification of datasets from INPUT_CFGS, compute their weights, and save the result in OUTPUT_CFG.
+    The weight for each entry is determined by the number of hours in a given dataset.
+
+    If more than one config is provided as input, we will concatenate them and output a single merged config.
+
+    Optionally, apply temperature re-weighting to balance the datasets (specify TEMPERATURE lesser than 1).
+    """
+    data = ListConfig([])
+    for icfg in input_cfgs:
+        data.extend(OmegaConf.load(icfg))
+    temperature = parse_temperature(temperature)
+    validate(data)
+    count(data, weight_key=strategy)
+    aggregate_group_weights(data)
+    reweight(data, temperature=temperature)
+    OmegaConf.save(data, output_cfg)
+
+
+def validate(entry: DictConfig | ListConfig, _level: int = 0):
+    if isinstance(entry, ListConfig):
+        for subentry in entry:
+            validate(subentry, _level + 1)
+        return
+
+    assert "type" in entry, f"Invalid YAML data config at nesting level {_level}: missing key 'type' in entry={entry}"
+
+    if entry.type == "group":
+        for subentry in entry["input_cfg"]:
+            validate(subentry, _level + 1)
+
+
+def count(entry: DictConfig | ListConfig, weight_key: str) -> None:
+    if isinstance(entry, ListConfig):
+        for subentry in entry:
+            count(subentry, weight_key=weight_key)
+        return
+    if entry.type == "group":
+        for subentry in entry["input_cfg"]:
+            count(subentry, weight_key=weight_key)
+        return
+
+    with quick_iter_options(entry):
+        iterable, is_tarred = get_parser_fn(entry.type)(entry)
+        stats = {"num_hours": 0.0, "num_examples": 0}
+        for example in iterable:
+            if hasattr(example, "duration"):
+                stats["num_hours"] += example.duration
+            stats["num_examples"] += 1
+        stats["num_hours"] /= 3600.0
+
+    if weight_key == "num_hours" and stats[weight_key] == 0.0:
+        raise RuntimeError(
+            f"Cannot set weights based on 'num_hours': at least one dataset has examples without 'duration' property. "
+            f"Details: {entry=}"
+        )
+
+    entry["weight"] = stats[weight_key]
+
+
+def aggregate_group_weights(entry: DictConfig | ListConfig) -> None:
+    if isinstance(entry, ListConfig):
+        for subentry in entry:
+            aggregate_group_weights(subentry)
+        return
+
+    if entry.type != "group":
+        return
+
+    for subentry in entry["input_cfg"]:
+        if "weight" not in subentry:
+            aggregate_group_weights(subentry)
+
+    entry.weight = sum(subentry["weight"] for subentry in entry["input_cfg"])
+
+
+def reweight(entry: DictConfig | ListConfig, temperature: None | float | list[float]) -> None:
+    if not temperature or (isinstance(entry, DictConfig) and entry.type != "group"):
+        return
+
+    if isinstance(temperature, Sequence):
+        temperature, *next_temperatures = temperature
+    else:
+        next_temperatures = temperature
+
+    if isinstance(entry, ListConfig):
+        for subentry in entry:
+            reweight(subentry, temperature=next_temperatures)
+        new_weights = temperature_reweighting([se.weight for se in entry], temperature=temperature)
+        for se, nw in zip(entry, new_weights):
+            se.weight = nw
+        return
+
+    for subentry in entry["input_cfg"]:
+        reweight(subentry, temperature=next_temperatures)
+
+    new_weights = temperature_reweighting([se.weight for se in entry["input_cfg"]], temperature=temperature)
+    for se, nw in zip(entry["input_cfg"], new_weights):
+        se.weight = nw
+
+
+def temperature_reweighting(weights: list[float], temperature: float = 1.0):
+    """(w_i ^ alpha / sum(w_i ^ alpha))"""
+    weights = np.asarray(weights) ** temperature
+    return (weights / weights.sum()).tolist()
+
+
+@contextmanager
+def quick_iter_options(entry: DictConfig):
+    entry.metadata_only = True
+    entry.force_finite = True
+    yield entry
+    del entry["metadata_only"]
+    del entry["force_finite"]
+
+
+def parse_temperature(value: list[float]) -> float | list[float] | None:
+    match value:
+        case 0:
+            return None
+        case 1:
+            return value[0]
+        case _:
+            return value
+
+
+if __name__ == '__main__':
+    estimate_data_weights()

scripts/speech_recognition/estimate_duration_bins.py

@@ -0,0 +1,124 @@
+# 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 itertools import islice
+from pathlib import Path
+
+from lhotse.cut import Cut
+from lhotse.dataset.sampling.dynamic_bucketing import estimate_duration_buckets
+from omegaconf import OmegaConf
+
+from nemo.collections.common.data.lhotse.cutset import read_cutset_from_config
+from nemo.collections.common.data.lhotse.dataloader import LhotseDataLoadingConfig
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Estimate duration bins for Lhotse dynamic bucketing using a sample of the input dataset. "
+        "The dataset is read either from one or more manifest files and supports data weighting.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "input",
+        help='Data input. Options: '
+        '1) "path.json" - any single NeMo manifest; '
+        '2) "[[path1.json],[path2.json],...]" - any collection of NeMo manifests; '
+        '3) "[[path1.json,weight1],[path2.json,weight2],...]" - any collection of weighted NeMo manifests; '
+        '4) "input_cfg.yaml" - a new option supporting input configs, same as in model training \'input_cfg\' arg; '
+        '5) "path/to/shar_data" - a path to Lhotse Shar data directory; '
+        '6) "key=val" - in case none of the previous variants cover your case: "key" is the key you\'d use in NeMo training config with its corresponding value ',
+    )
+    parser.add_argument("-b", "--buckets", type=int, default=30, help="The desired number of buckets.")
+    parser.add_argument(
+        "-n",
+        "--num_examples",
+        type=int,
+        default=-1,
+        help="The number of examples (utterances) to estimate the bins. -1 means use all data "
+        "(be careful: it could be iterated over infinitely).",
+    )
+    parser.add_argument(
+        "-l",
+        "--min_duration",
+        type=float,
+        default=-float("inf"),
+        help="If specified, we'll filter out utterances shorter than this.",
+    )
+    parser.add_argument(
+        "-u",
+        "--max_duration",
+        type=float,
+        default=float("inf"),
+        help="If specified, we'll filter out utterances longer than this.",
+    )
+    parser.add_argument(
+        "-q", "--quiet", type=bool, default=False, help="When specified, only print the estimated duration bins."
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    if '=' in args.input:
+        inp_arg = args.input
+    elif args.input.endswith(".yaml"):
+        inp_arg = f"input_cfg={args.input}"
+    elif Path(args.input).is_dir():
+        inp_arg = f"shar_path={args.input}"
+    else:
+        inp_arg = f"manifest_filepath={args.input}"
+    config = OmegaConf.merge(
+        OmegaConf.structured(LhotseDataLoadingConfig),
+        OmegaConf.from_dotlist([inp_arg, "metadata_only=true"]),
+    )
+    cuts, _ = read_cutset_from_config(config)
+    min_dur, max_dur = args.min_duration, args.max_duration
+    nonaudio, discarded, tot = 0, 0, 0
+    observed_max_dur = 0
+
+    def duration_ok(cut) -> bool:
+        nonlocal nonaudio, discarded, tot, observed_max_dur
+        tot += 1
+        if not isinstance(cut, Cut):
+            nonaudio += 1
+            return False
+        if not (min_dur <= cut.duration <= max_dur):
+            discarded += 1
+            return False
+        observed_max_dur = max(cut.duration, observed_max_dur)
+        return True
+
+    cuts = cuts.filter(duration_ok)
+    if (N := args.num_examples) > 0:
+        cuts = islice(cuts, N)
+    duration_bins = estimate_duration_buckets(cuts, num_buckets=args.buckets)
+    duration_bins = f"[{','.join(str(round(b, ndigits=5)) for b in duration_bins)}]"
+    if args.quiet:
+        print(duration_bins)
+        return
+    if discarded:
+        ratio = discarded / tot
+        print(f"Note: we discarded {discarded}/{tot} ({ratio:.2%}) utterances due to min/max duration filtering.")
+    if nonaudio:
+        print(f"Note: we discarded {nonaudio} non-audio examples found during iteration.")
+    print(f"Used {tot - nonaudio - discarded} examples for the estimation.")
+    print("Use the following options in your config:")
+    print(f"\tnum_buckets={args.buckets}")
+    print(f"\tbucket_duration_bins={duration_bins}")
+    print(f"\tmax_duration={observed_max_dur}")
+
+
+if __name__ == "__main__":
+    main()

scripts/speech_recognition/estimate_duration_bins_2d.py

@@ -0,0 +1,415 @@
+# 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 ast
+import math
+import warnings
+from functools import partial
+from itertools import islice
+from pathlib import Path
+from typing import Callable, Iterable
+
+import numpy as np
+import pandas as pd
+from lhotse.cut import Cut
+from omegaconf import OmegaConf
+
+from nemo.collections.common.data import apply_prompt_format_fn
+from nemo.collections.common.data.lhotse.cutset import read_cutset_from_config
+from nemo.collections.common.data.lhotse.dataloader import LhotseDataLoadingConfig, tokenize
+from nemo.collections.common.data.lhotse.sampling import DurationFilter, FixedBucketBatchSizeConstraint2D
+from nemo.collections.common.prompts.formatter import PromptFormatter
+from nemo.collections.common.tokenizers import (
+    AggregateTokenizer,
+    CanaryTokenizer,
+    SentencePieceTokenizer,
+    TokenizerSpec,
+)
+from nemo.collections.common.tokenizers.aggregate_tokenizer import TokenizerWrapper
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Estimate duration bins for Lhotse dynamic bucketing using a sample of the input dataset. "
+        "The dataset is read either from one or more manifest files and supports data weighting. "
+        "Unlike estimate_duration_bins.py, this script prepares the setup for 2D bucketing. "
+        "This means that each main bucket for audio duration is sub-divided into sub-buckets "
+        "for the number of output tokens (supporting BPE and Aggregated tokenizers). "
+        "2D bucketing is especially useful for encoder-decoder models where input audio duration is often "
+        "not sufficient to stratify the sampling with an optimal GPU utilization.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "input",
+        help='Data input. Options: '
+        '1) "path.json" - any single NeMo manifest; '
+        '2) "[[path1.json],[path2.json],...]" - any collection of NeMo manifests; '
+        '3) "[[path1.json,weight1],[path2.json,weight2],...]" - any collection of weighted NeMo manifests; '
+        '4) "input_cfg.yaml" - a new option supporting input configs, same as in model training \'input_cfg\' arg; '
+        '5) "path/to/shar_data" - a path to Lhotse Shar data directory; '
+        '6) "key=val" - in case none of the previous variants cover your case: "key" is the key you\'d use in NeMo training config with its corresponding value ',
+    )
+    parser.add_argument(
+        "-t",
+        "--tokenizer",
+        nargs="+",
+        required=True,
+        help="Path to one or more SPE tokenizers. More than one means we'll use AggregateTokenizer and --langs argument must also be used. When provided, we'll estimate a 2D distribution for input and output sequence lengths.",
+    )
+    parser.add_argument(
+        "-a", "--langs", nargs="+", help="Language names for each of AggregateTokenizer sub-tokenizers."
+    )
+    parser.add_argument(
+        "-b",
+        "--buckets",
+        type=int,
+        default=30,
+        help="The desired number of buckets (dim0 => covers input sequence length / audio duration).",
+    )
+    parser.add_argument(
+        "-s",
+        "--sub-buckets",
+        type=int,
+        default=2,
+        help="The desired number of sub-buckets (dim1 => covers output sequence length / num_tokens).",
+    )
+    parser.add_argument("--text-field", default="text", help="The key in manifests to read transcripts from.")
+    parser.add_argument("--lang-field", default="lang", help="The key in manifests to read language from.")
+    parser.add_argument(
+        "-n",
+        "--num_examples",
+        type=int,
+        default=-1,
+        help="The number of examples (utterances) to estimate the bins. -1 means use all data "
+        "(be careful: it could be iterated over infinitely).",
+    )
+    parser.add_argument(
+        "-l",
+        "--min_duration",
+        type=float,
+        default=-float("inf"),
+        help="If specified, we'll filter out utterances shorter than this.",
+    )
+    parser.add_argument(
+        "-u",
+        "--max_duration",
+        type=float,
+        default=float("inf"),
+        help="If specified, we'll filter out utterances longer than this.",
+    )
+    parser.add_argument(
+        "--max_tps", type=float, default=None, help="Deprecated. TPS is automatically determined per bucket."
+    )
+    parser.add_argument(
+        "--token_outlier_threshold",
+        type=float,
+        default=4.0,
+        help="The lower this is, the more outliers in transcript token count will be filtered out. "
+        "By default allow token counts at 4 sigma away from distribution mean, computed separately for every bucket.",
+    )
+    parser.add_argument(
+        "-q", "--quiet", type=bool, default=False, help="When specified, only print the estimated duration bins."
+    )
+    parser.add_argument(
+        "-f",
+        "--prompt-format",
+        type=str,
+        help="When specified, we'll use a prompt formatter in addition to the tokenizer for the purpose of estimating token count bins. "
+        "This is useful for accurate 2D bucket estimation with models such as EncDecMultiTaskModel (Canary-1B), "
+        "or any model where the label sequence consists of a user prompt and a model's response.",
+    )
+    parser.add_argument(
+        "-p",
+        "--prompt",
+        type=str,
+        help="Prompt slots provided as a Python list of dicts. It is used together with --prompt-format option."
+        "For example, with Canary-1B you may use: [{'role':'user','slots':{'source_lang':'en','target_lang':'en','task':'asr','pnc':'yes'}]",
+    )
+    return parser.parse_args()
+
+
+def sort_two_arrays(A, B):
+    joint = np.rec.fromarrays([A, B])
+    joint.sort()
+    return joint.f0, joint.f1
+
+
+def estimate_duration_buckets(
+    cuts: Iterable[Cut],
+    num_buckets: int,
+    num_subbuckets: int,
+    max_tps: float,
+    max_duration: float,
+    token_outlier_threshold: float,
+    quiet: bool,
+) -> list[tuple[float, float]]:
+    """
+    This function is based on lhotse.dataset.sampling.dynamic_bucketing.estimate_duration_buckets.
+    It extends it to a 2D bucketing case.
+    """
+    assert num_buckets > 1
+
+    constraint = FixedBucketBatchSizeConstraint2D([(0.0, 0.0)], [0])
+
+    # Gather the duration and token count statistics for the dataset.
+    sizes = []
+    num_tokens = []
+    for c in cuts:
+        dur, toks = constraint.measure_length(c)
+        sizes.append(dur)
+        num_tokens.append(toks)
+    sizes = np.array(sizes, dtype=np.float32)
+    num_tokens = np.array(num_tokens, dtype=np.int32)
+    sizes, num_tokens = sort_two_arrays(sizes, num_tokens)
+
+    # We are building buckets with equal duration (empirically leads to more even bucket exhaustion over time).
+    # We need to determine how much duration to allocate per bucket.
+    size_per_bucket = sizes.sum() / num_buckets
+
+    if not quiet:
+        print("Duration distribution:")
+        print(pd.Series(sizes).describe(percentiles=[0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.99, 0.995, 0.999]))
+    if math.isinf(max_duration):
+        max_duration = round(sizes[-1], 3)  # Round to 3 decimal places to be consistent for the output format.
+
+    bins = []
+    tps_thresholds = []
+    bin_indexes = [0]
+    tot = 0.0
+
+    def _estimate_token_buckets(max_bucket_duration, start_idx, end_idx, corr_subbuckets=None):
+        # Since this is 2D bucketing, apply the same bin creation logic
+        # for the second dimension (i.e. token count) as for the first dimension (duration).
+        # That means we aim to have each bucket contain roughly the same number of tokens.
+        # Note that this estimation is biased towards more padding if you have
+        # a lot of zero-token examples (e.g. non-speech).
+        nonlocal bins
+
+        if not corr_subbuckets:
+            corr_subbuckets = num_subbuckets
+
+        # Start by discarding outlier examples as defined by token-per-second (TPS) attribute.
+        # We empirically determined high TPS examples to cause severe OOMs limiting batch sizes.
+        # We cap the TPS for each top-level bucket at 4 standard deviations of TPS.
+        # Examples exceeding that TPS value will be discarded during sampling at training time.
+        num_tokens_bucket_all = num_tokens[start_idx:end_idx]
+        sizes_bucket_all = sizes[start_idx:end_idx]
+        non_outlier_indexes = find_non_outliers_z_score(
+            num_tokens_bucket_all / sizes_bucket_all, threshold=token_outlier_threshold
+        )
+        num_tokens_bucket = num_tokens_bucket_all[non_outlier_indexes]
+        sizes_bucket = sizes_bucket_all[non_outlier_indexes]
+        max_tps_bucket = (num_tokens_bucket / sizes_bucket).max()
+        num_tokens_bucket, sizes_bucket = sort_two_arrays(num_tokens_bucket, sizes_bucket)
+        if not quiet:
+            outlier_tps = np.delete(num_tokens_bucket_all / sizes_bucket_all, non_outlier_indexes)
+            print(
+                f"[bucket <= {max_bucket_duration:.2f}s] [{num_tokens_bucket.min()} - {num_tokens_bucket.max()}] [approx-max-tps: {max_tps_bucket:.2f}] Discarded {end_idx - start_idx - len(num_tokens_bucket)} max token outliers",
+                end=" ",
+            )
+            if len(outlier_tps) > 0:
+                print(f"min-outlier: {outlier_tps.min():.2f}, max-outlier: {outlier_tps.max():.2f}).", end="")
+            print()
+
+        tokens_per_subbucket = num_tokens_bucket.sum() / corr_subbuckets
+        tot_toks = 0
+        # Iterate over token counts, and whenever we hit tokens_per_subbucket, create a new 2D bucket bin.
+        for num_toks, size in zip(num_tokens_bucket, sizes_bucket):
+            # Threshold hit: we are creating a new (max_duration, max_num_tokens) bin.
+            if tot_toks > tokens_per_subbucket:
+                bins.append((max_bucket_duration, num_toks))
+                tps_thresholds.append(max_tps_bucket)
+                tot_toks = 0
+            tot_toks += num_toks
+        bins.append((max_bucket_duration, num_toks))
+        tps_thresholds.append(max_tps_bucket)
+
+    duration_bins = []
+
+    # Iterate over data, and whenever we hit size_per_bucket, register it as a new duration bucket.
+    for binidx, size in enumerate(sizes):
+        if tot > size_per_bucket:
+            size = round(size, 3)  # Round to 3 decimal places to be consistent for the output format.
+            duration_bins.append(size)
+            bin_indexes.append(binidx)
+            tot = 0.0
+        tot += size
+
+    if not quiet:
+        print(f"Initial duration_bins={duration_bins}")
+
+    skipped_buckets = 1
+    start_idx = 0
+
+    # Iterate over newly created duration bins to handle cases where some bins have the same value —
+    # this usually happens when the data is skewed.
+    # If we detect such bins, we skip estimating token buckets for that particular bin.
+    # Instead, we keep track of how many bins got skipped because they had the same duration.
+    # Then, when we finally hit a bin with a different duration, we treat all those skipped bins as one "combined" bin.
+    # For that combined bin, we create more subbuckets — specifically, the number of skipped bins × `num_subbuckets` (set by the user).
+    #
+    # Example of durations bins created from skewed duration distribution: [5, 20, 30, 30, 30, 40]
+    # Here, we'd end up making token subbuckets for: [5, 20, 40]
+    # where [20, 40] bucket will have 4 times more subbuckets (as we combined 4 buckets into 1) than usual bucket in that settings.
+
+    for i, (duration_bin, binidx) in enumerate(zip(duration_bins, bin_indexes[1:])):
+        if (i != len(duration_bins) - 1 and duration_bins[i + 1] == duration_bin) or (
+            i == len(duration_bins) - 1 and max_duration == duration_bin
+        ):
+            skipped_buckets += 1
+            continue
+        _estimate_token_buckets(
+            max_bucket_duration=duration_bin,
+            start_idx=start_idx,
+            end_idx=binidx,
+            corr_subbuckets=num_subbuckets * skipped_buckets,
+        )
+        start_idx = binidx
+        skipped_buckets = 1
+
+    # Estimate an extra 2D bin set for global max duration.
+    # Also, if the last value in duration_bins is equal to max_duration,
+    # we need to make sure we properly handle any previously "skipped" buckets that ended at this max value.
+    _estimate_token_buckets(
+        max_bucket_duration=max_duration,
+        start_idx=start_idx,
+        end_idx=len(sizes),
+        corr_subbuckets=num_subbuckets * skipped_buckets,
+    )
+    return bins, tps_thresholds
+
+
+def find_non_outliers_z_score(data, threshold=4):
+    # Note: we don't apply abs() here because we only filter the upper end of the distribution.
+    # We don't mind low-token-counts for bucketing purposes.
+    z_scores = (data - np.mean(data)) / np.std(data)
+    return np.where(z_scores <= threshold)
+
+
+def load_tokenizer(paths: list[str], langs: list[str] = None, is_canary: bool = True) -> TokenizerSpec:
+    if len(paths) == 1:
+        tok = SentencePieceTokenizer(paths[0])
+    else:
+        assert langs is not None and len(paths) == len(
+            langs
+        ), f"Cannot create AggregateTokenizer; each tokenizer must have assigned a language via --langs option (we got --tokenizers={paths} and --langs={langs})"
+        if is_canary:
+            tokcls = CanaryTokenizer
+        else:
+            tokcls = AggregateTokenizer
+        tok = tokcls({lang: SentencePieceTokenizer(p) for lang, p in zip(langs, paths)})
+    return tok
+
+
+def apply_tokenizer(cut, tokenizer=None, prompt: PromptFormatter = None):
+    if prompt is not None:
+        encoded = apply_prompt_format_fn(cut, prompt)
+        cut.supervisions[0].tokens = encoded["input_ids"]
+
+    elif tokenizer is not None:
+        cut = tokenize(cut, TokenizerWrapper(tokenizer))
+
+    return cut
+
+
+class RejectionsCounter:
+    def __init__(self, predicate: Callable, message: str):
+        self.predicate = predicate
+        self.message = message
+        self.total = 0
+        self.rejected = 0
+
+    def __call__(self, example) -> bool:
+        ans = self.predicate(example)
+        self.total += 1
+        if not ans:
+            self.rejected += 1
+        return ans
+
+    def print_report(self) -> None:
+        if self.rejected:
+            print(f"{self.message} | Rejected {self.rejected}/{self.total} examples.")
+
+
+def main():
+    args = parse_args()
+
+    if not args.quiet:
+        pd.set_option('display.float_format', lambda x: '%.2f' % x)
+
+    if args.max_tps is not None:
+        warnings.warn(
+            "The option --max_tps has been deprecated in favor of "
+            "automatic TPS determination that's variable across buckets."
+        )
+
+    tokenizer = None
+    prompt = None
+    if args.tokenizer is not None:
+        tokenizer = load_tokenizer(
+            paths=args.tokenizer,
+            langs=args.langs,
+            is_canary=args.prompt_format is not None and 'canary' in args.prompt_format,
+        )
+        if args.prompt_format is not None:
+            prompt_defaults = None
+            if args.prompt is not None:
+                prompt_defaults = ast.literal_eval(args.prompt)
+            prompt = PromptFormatter.resolve(args.prompt_format)(tokenizer, defaults=prompt_defaults)
+
+    if '=' in args.input:
+        inp_arg = args.input
+    elif args.input.endswith(".yaml"):
+        inp_arg = f"input_cfg={args.input}"
+    elif Path(args.input).is_dir():
+        inp_arg = f"shar_path={args.input}"
+    else:
+        inp_arg = f"manifest_filepath={args.input}"
+    config = OmegaConf.merge(
+        OmegaConf.structured(LhotseDataLoadingConfig),
+        OmegaConf.from_dotlist(
+            [inp_arg, "metadata_only=true", f"text_field={args.text_field}", f"lang_field={args.lang_field}"]
+        ),
+    )
+    cuts, _ = read_cutset_from_config(config)
+    duration_filter = RejectionsCounter(DurationFilter(args.min_duration, args.max_duration), "Duration filtering")
+    cuts = cuts.filter(duration_filter)
+    cuts = cuts.map(partial(apply_tokenizer, tokenizer=tokenizer, prompt=prompt))
+    if (N := args.num_examples) > 0:
+        cuts = islice(cuts, N)
+
+    duration_bins, tps_thresholds = estimate_duration_buckets(
+        cuts,
+        num_buckets=args.buckets,
+        num_subbuckets=args.sub_buckets,
+        max_duration=args.max_duration,
+        max_tps=args.max_tps,
+        token_outlier_threshold=args.token_outlier_threshold,
+        quiet=args.quiet,
+    )
+    duration_bins = "[" + ','.join(f"[{b:.3f},{sb:d}]" for b, sb in duration_bins) + "]"
+    tps_thresholds = "[" + ",".join(f"{t:.2f}" for t in tps_thresholds) + "]"
+    if not args.quiet:
+        duration_filter.print_report()
+    print("Use the following options in your config:")
+    print(f"\tuse_bucketing=1")
+    print(f"\tnum_buckets={args.buckets}")
+    print(f"\tbucket_duration_bins={duration_bins}")
+    print(f"The max_tps setting below is optional, use it if your data has low quality long transcript outliers:")
+    print(f"\tmax_tps={tps_thresholds}")
+
+
+if __name__ == "__main__":
+    main()

scripts/speech_recognition/filter_tarred_audio_dataset.py

@@ -0,0 +1,156 @@
+# 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 functools import partial
+from io import BytesIO
+from pathlib import Path
+
+import click
+import lhotse
+import torch.utils.data
+from lhotse import CutSet, MonoCut
+from lhotse.audio.backend import LibsndfileBackend
+from lhotse.dataset import DynamicCutSampler, IterableDatasetWrapper
+from lhotse.shar import JsonlShardWriter, TarWriter
+from omegaconf import OmegaConf
+
+from nemo.collections.common.data.lhotse import read_cutset_from_config
+from nemo.collections.common.data.lhotse.dataloader import LhotseDataLoadingConfig
+
+
+@click.command()
+@click.argument("manifest_filepath")
+@click.argument("tarred_audio_filepaths")
+@click.argument("filtered_manifest_filepath")
+@click.argument("output_dir", type=click.Path())
+@click.option(
+    "-f",
+    "--output-format",
+    type=click.Choice(["lhotse_shar", "nemo_tarred"]),
+    default="lhotse_shar",
+    help="Which format should we use to save the filtered tarred data.",
+)
+@click.option("-s", "--shard-size", type=int, default=1000, help="Desired number of examples per output shard.")
+def filter_tarred(
+    manifest_filepath: str,
+    tarred_audio_filepaths: str,
+    filtered_manifest_filepath: str,
+    output_dir: str,
+    output_format: str,
+    shard_size: int,
+):
+    """
+    Given an existing tarred dataset and manifests that point to a subset of examples,
+    create a new tarred dataset corresponding to the subset.
+
+    This is useful if you want to "re-tar" an existing tarred dataset in order to efficiently
+    read some subset of it.
+    """
+    lhotse.set_dill_enabled(True)
+    all_cuts = read_cutset(manifest_filepath, tarred_audio_filepaths)
+    keep_cuts = {cut.id: cut for cut in read_cutset(filtered_manifest_filepath)}
+    filtered_cuts = bg_load(
+        all_cuts.filter(lambda c: c.id in keep_cuts).map(partial(attach_custom, cuts_with_custom=keep_cuts))
+    )
+    if not '://' in output_dir:  # we support object store writing too
+        Path(output_dir).mkdir(exist_ok=True, parents=True)
+    if output_format == "lhotse_shar":
+        filtered_cuts.to_shar(output_dir=output_dir, fields={"recording": "flac"}, shard_size=shard_size)
+    elif output_format == "nemo_tarred":
+        export_to_nemo_tarred(cuts=filtered_cuts, output_dir=output_dir, shard_size=shard_size)
+    else:
+        raise RuntimeError(f"Unsupported output format: '{output_format}'")
+
+
+def read_cutset(src: str, tar: str | None = None) -> CutSet:
+    inp_arg = ["force_finite=true"]
+    if tar is not None:
+        inp_arg += [f"manifest_filepath={src}", f"tarred_audio_filepaths={tar}"]
+    else:
+        inp_arg += ["metadata_only=true"]
+        if src.endswith(".yaml"):
+            inp_arg += [f"input_cfg={src}"]
+        elif Path(src).is_dir():
+            inp_arg += [f"shar_path={src}"]
+        else:
+            inp_arg += [f"manifest_filepath={src}"]
+    config = OmegaConf.merge(
+        OmegaConf.structured(LhotseDataLoadingConfig),
+        OmegaConf.from_dotlist(inp_arg),
+    )
+    cuts, _ = read_cutset_from_config(config)
+    return cuts
+
+
+def export_to_nemo_tarred(cuts: CutSet, output_dir: str, shard_size: int) -> None:
+    with (
+        TarWriter(pattern=f"{output_dir}/audio_%d.tar", shard_size=shard_size) as aw,
+        JsonlShardWriter(pattern=f"{output_dir}/manifest_%d.jsonl", shard_size=shard_size) as mw,
+    ):
+        for cut in cuts:
+            assert (
+                isinstance(cut, MonoCut) and len(cut.supervisions) == 1
+            ), f"Export to nemo_tarred format is possible only for mono cuts with a single supervision, but we got: {cut}"
+            # Prepare audio for writing.
+            audio_name = f"{cut.id}.flac"
+            audio = BytesIO()
+            LibsndfileBackend().save_audio(audio, cut.load_audio(), sampling_rate=cut.sampling_rate, format="flac")
+            audio.seek(0)
+            # Prepare manifest for writing.
+            ans = {"audio_filepath": audio_name, "duration": cut.duration}
+            if cut.supervisions[0].text:
+                ans["text"] = cut.supervisions[0].text
+            if cut.supervisions[0].language:
+                ans["lang"] = cut.supervisions[0].language
+            if cut.custom is not None:
+                # Ensure if we export anything custom, these are only simple built-in types compatible with JSON.
+                ans.update({k: v for k, v in cut.custom.items() if isinstance(v, (int, float, str, list, dict))})
+            # Set the right shard_id.
+            shard_id = max(0, mw.num_shards - 1)
+            if mw.num_items > 0 and mw.num_items % mw.shard_size == 0:
+                shard_id += 1
+            ans["shard_id"] = shard_id
+            # Write both items.
+            aw.write(audio_name, audio)
+            mw.write(ans)
+
+
+def attach_custom(cut, cuts_with_custom):
+    custom = cuts_with_custom[cut.id].custom
+    if custom is not None:
+        cut.custom.update(custom)
+    return cut
+
+
+class Identity(torch.utils.data.Dataset):
+    def __getitem__(self, x):
+        cut = x[0]
+        for k in ["dataloading_info", "shard_id"]:
+            cut.custom.pop(k, None)
+        return cut
+
+
+def bg_load(cuts: CutSet) -> CutSet:
+    return CutSet(
+        torch.utils.data.DataLoader(
+            IterableDatasetWrapper(Identity(), DynamicCutSampler(cuts, max_cuts=1)),
+            batch_size=None,
+            num_workers=1,
+            prefetch_factor=10,
+        )
+    )
+
+
+if __name__ == '__main__':
+    filter_tarred()

scripts/speech_recognition/oomptimizer.py

@@ -0,0 +1,548 @@
+#!/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 importlib
+import math
+import sys
+from numbers import Number
+
+import click
+import lightning.pytorch as pl
+import torch
+from lhotse import compute_num_samples
+from omegaconf import OmegaConf
+
+from nemo.collections.asr.models.asr_model import ASRModel
+from nemo.core.neural_types import AudioSignal, LabelsType, LengthsType, MaskType, NeuralType
+from nemo.utils import logging
+
+
+class ProfilingBatchGenerator:
+    """
+    ProfilingBatchGenerator is used to generate artificial mini-batches for model training
+    and tracking the progress of batch size optimization.
+
+    The high-level usage API is the following::
+
+        >>> gen = ProfilingBatchGenerator(schema)
+        ... finished = False
+        ... while not finished:
+        ...     batch = gen(input_seq_len, output_seq_len)
+        ...     try:
+        ...         training_step(model, batch)
+        ...         oom = False
+        ...     except torch.cuda.OutOfMemoryError:
+        ...         oom = True
+        ...     finished = gen.advance(oom)
+        ... solution = gen.max_batch_size  # The solution of the search problem.
+        ... gen.reset()  # Can re-use for other sequence lengths now.
+
+    The search terminates once the difference between max working batch size and min OOM batch size
+    divided by the latter is smaller than ``rel_gap_thresh`` that difference amounts to a single element.
+    For example, a max working batch size is 96 and min OOM batch size is 100 indicates a gap of 0.04,
+    which would terminate the search with threshold of 0.05.
+
+    In order to generate mini-batches compatible with a given model, the generator:
+
+    * accepts a ``schema`` argument in its constructor, and
+
+    * accepts input/output sequence lengths in each call to generate a mini-batch.
+
+    ``schema`` has the following structure::
+
+
+        >>> {
+        ...     "cls": tuple | MyBatchType,
+        ...     "inputs": [
+        ...         {
+        ...             "type": NeuralType(...) | Literal["dummy"],
+        ...             "seq_length": Literal["input", "output"],
+        ...             "vocab_size": int,  # optional, required only for LabelsType
+        ...             "name": str,  # optional, indicates kwarg
+        ...         },
+        ...         ...,
+        ...     ]
+        ... }
+
+    ``cls`` indicates how we should construct the mini-batch. Typically you can just use ``tuple`` for most
+    batch schemas. However, if the model expects a specific, e.g., dataclass, you can tell ``ProfilingBatchGenerator``
+    to use it. The mini-batch object will be constructed using the items in ``inputs``.
+
+    Each element of ``inputs`` specifies a NeMo NeuralType which needs to have a defined ``elements_type``.
+    The supported types are ``AudioSignal``, ``LengthsType`` and ``LabelsType``.
+    If "type" is not a NeuralType, we interpret that as a placeholder tensor that's not relevant but expected
+    by the model/batch constructor. In addition, ``"seq_length"`` key is used to determine whether we should apply
+    input or output sequence length to a given tensor.
+
+    Optional keys:
+
+    * ``vocab_size`` is required for ``LabelsType`` so that we can generate proper label values.
+
+    * ``name`` is required if objects of ``cls`` have to be constructed using keyword arguments.
+
+    A simple schema example for a model using audio/lengths tensor pair (unsupervised/self-supervised)::
+
+        >>> {
+        ...     "cls": tuple,
+        ...     "inputs": [
+        ...         {"type": NeuralType(("B", "T"), AudioSignal()), "seq_length": "input"},
+        ...         {"type": NeuralType(("B"), LengthsType()), "seq_length": "input"},
+        ...     ]
+        ... }
+
+    """
+
+    def __init__(
+        self,
+        schema: dict,
+        start_batch_size: int = 32,
+        rel_gap_thresh: float = 0.05,
+        device: str = "cuda",
+    ):
+        self.schema = schema
+        self.start_batch_size = start_batch_size
+        self.rel_gap_thresh = rel_gap_thresh
+        self.device = device
+        self.reset()
+
+    def __call__(self, input_seq_length: int, output_seq_length: int):
+        B = self._current
+        select_seq_length = {"input": input_seq_length, "output": output_seq_length}
+        batch = []
+        names = []
+        for item in self.schema["inputs"]:
+            nt = item["type"]
+            if isinstance(nt, str) and nt == "constant":
+                if isinstance(val := item["value"], str) and val == "batch":
+                    tnsr = torch.tensor([B], dtype=torch.long, device=self.device)
+                else:
+                    tnsr = torch.tensor([val], dtype=torch.long, device=self.device)
+            elif not isinstance(nt, NeuralType):  # placeholder
+                tnsr = torch.tensor([])
+            elif isinstance(nt.elements_type, AudioSignal):
+                seq_length = select_seq_length[item["seq_length"]]
+                tnsr = torch.randn(B, seq_length, dtype=torch.float32, device=self.device)
+            elif isinstance(nt.elements_type, LengthsType):
+                seq_length = select_seq_length[item["seq_length"]]
+                tnsr = torch.ones(B, dtype=torch.long, device=self.device) * seq_length
+            elif isinstance(nt.elements_type, LabelsType):
+                seq_length = select_seq_length[item["seq_length"]]
+                tnsr = torch.randint(0, item["vocab_size"], size=(B, seq_length), device=self.device)
+            elif isinstance(nt.elements_type, MaskType):
+                seq_length = select_seq_length[item["seq_length"]]
+                tnsr = torch.ones(B, seq_length, device=self.device)
+            else:
+                raise RuntimeError("Unexpected item in oomptimizer schema: {item}")
+            batch.append(tnsr)
+            names.append(item.get("name"))
+        args = [elem for name, elem in zip(names, batch) if name is None]
+        kwargs = {name: elem for name, elem in zip(names, batch) if name is not None}
+        if not kwargs and self.schema["cls"] == tuple:
+            return tuple(args)
+        return self.schema["cls"](*args, **kwargs)
+
+    @property
+    def max_batch_size(self) -> int | None:
+        """
+        Return the solution of the batch size search problem.
+        It will keep returning None until the search is done.
+        """
+        if (
+            self._max_ok is not None
+            and self._min_err is not None
+            and (self.current_rel_gap <= self.rel_gap_thresh or self._min_err - self._max_ok <= 1)
+        ):
+            return self._max_ok
+        return None
+
+    @property
+    def current_rel_gap(self) -> float | None:
+        """
+        Return the current gap between the largest batch that works and the smallest batch that triggers OOM.
+        The gap is defined as the batch size difference divided by the larger element.
+        E.g., if the best found batch size is 95 and the smallest that triggers OOM is 100, the gap is 0.05.
+        """
+        if self._min_err is None or self._max_ok is None:
+            return None
+        return (self._min_err - self._max_ok) / self._min_err
+
+    def reset(self):
+        """Reset the generator to prepare it for a new search."""
+        self._current = self.start_batch_size
+        self._max_ok = None  # max batch size that works
+        self._min_err = None  # min batch size that doesn't work
+
+    def advance(self, oom: bool) -> bool:
+        """
+        Adjusts the current batch size based on the outcome.
+        Returns a bool indicating whether the calibration is complete.
+        """
+        if self.max_batch_size is not None:
+            return True
+
+        if oom:
+            # Training step failed with OOM.
+            # Update the minimum known batch size that causes an error.
+            self._min_err = min(float("inf") if self._min_err is None else self._min_err, self._current)
+            # Training step failed on OOM
+            if self._max_ok is None:
+                # We haven't found a batch size that works yet, keep going 2x down.
+                self._current = round(self._current / 2)
+            else:
+                # Try the middle-point between the known extremes.
+                self._current = round((self._max_ok + self._min_err) / 2)
+        else:
+            # Training step successful.
+            # Update the maximum known batch size that works.
+            self._max_ok = max(-1 if self._max_ok is None else self._max_ok, self._current)
+            if self._min_err is None:
+                # We haven't found a batch size that causes an error yet, keep going 2x higher
+                self._current *= 2
+            else:
+                # Try the middle-point between the known extremes.
+                self._current = round((self._max_ok + self._min_err) / 2)
+
+        return False
+
+
+class FloatList(click.Option):
+    """Support passing bucket duration bins as [1.1,2.5,5.6,...]"""
+
+    name = "list[float]"
+
+    def type_cast_value(self, ctx, value):
+        if isinstance(value, list) and all(isinstance(v, float) for v in value):
+            return value
+        try:
+            import ast
+
+            ans = ast.literal_eval(value)
+            if isinstance(ans[0], list):
+                ans = [tuple(item) for item in ans]
+            return ans
+        except ValueError:
+            raise click.BadParameter(value)
+
+
+@click.command(context_settings={'show_default': True})
+@click.option(
+    "-n",
+    "--pretrained-name",
+    type=str,
+    default=None,
+    help="Name of a pretrained model to use, e.g. 'nvidia/canary-1b'.",
+)
+@click.option(
+    "-m",
+    "--module-name",
+    type=str,
+    default=None,
+    help="Full path to NeMo's module corresponding to CONFIG_PATH, e.g. 'nemo.collections.asr.models.EncDecMultiTaskModel'.",
+)
+@click.option(
+    "-c", "--config-path", type=str, default=None, help="Path to the training configuration file for MODULE_NAME."
+)
+@click.option("-o", "--optimizer-name", type=str, default="adamw", help="Name of optimizer to use.")
+@click.option(
+    "-b",
+    "--buckets",
+    cls=FloatList,
+    default=[5.0, 10.0, 15.0, 20.0, 25.0, 30.0],
+    help="List of upper-bound bucket bins (i.e. first bucket is [0.0 - item0), second bucket is [item0 - item1), etc.). "
+    "We also support a nested list for 2D bucketing, e.g. [[2.0, 10],[2.0,20],[4.5,15],[4.5,30],...], "
+    "where each item is a pair of (max_input_seq_len, max_output_seq_len) for a given bucket.",
+)
+@click.option(
+    "-t",
+    "--threshold",
+    type=float,
+    default=0.05,
+    help="Search stopping criterion in range [0, 1], lower is more precise. Interpret as the uncerainty gap, i.e. (min_oom_batch_size - max_ok_batch_size) / min_oom_batch_size.",
+)
+@click.option("-s", "--start-batch-size", type=int, default=32, help="Initial batch size to start the search from.")
+@click.option(
+    "-r",
+    "--ratio",
+    type=int,
+    default=12,  # conservative estimate towards longer transcripts
+    help="The output_sequence_length to input_sequence_length ratio for the purpose of determing the maximum output sequence lengths. "
+    "The interpretation depends on input and output modalities. Examples: for audio->text it's tokens per second. "
+    "For text->audio it's seconds per token. For audio->audio it's output seconds per input second. "
+    "For text->text it's output tokens per input token. "
+    "In general larger ratio means longer output sequences and increased memory consumption. "
+    "The default value is set adequately for automatic speech recognition. "
+    "This argument is ignored when 2D buckets are provided to --buckets option.",
+)
+@click.option(
+    "-f",
+    "--memory-fraction",
+    type=float,
+    default=0.9,
+    help="Limits the use of CUDA memory for this process to MEMORY_FRACTION of the total device memory. "
+    "By default we force 5% memory to be unused to account for non-training-loop related CUDA memory usage"
+    "in actual training scripts.",
+)
+@click.option(
+    "-d",
+    "--device",
+    default="cuda:0",
+    help="Device string to be passed to torch.device; due to MEMORY_FRACTION option, "
+    "it must specify the device index (e.g. cuda:0). "
+    "You can also leave the default index and select a specific GPU using env var CUDA_VISIBLE_DEVICES=<idx>",
+)
+@click.option(
+    "-y",
+    "--dtype",
+    default="bfloat16",
+    help="Float precision to use for computation (used together with autocast).",
+)
+@click.option(
+    "--ddp/--no-ddp",
+    type=bool,
+    default=True,
+    help="Whether we should simulate DDP GPU RAM usage. Stores an extra copy of the model in GPU memory. Enabled by default.",
+)
+def oomptimizer(
+    pretrained_name: str | None,
+    module_name: str | None,
+    config_path: str | None,
+    optimizer_name: str,
+    buckets: list[float],
+    threshold: float,
+    start_batch_size: int,
+    ratio: int,
+    memory_fraction: float,
+    device: str,
+    dtype: str,
+    ddp: bool,
+):
+    """
+    OOMptimizer finds the optimal batch sizes for training your model with bucketing dataloading.
+    It performs a search over batch sizes until it converges by measuring the GPU memory usage for
+    a model's training step and optimizer update.
+
+    \b
+    There are two main usage patterns: for using a pretrained model or an untrained model configuration.
+    The latter is more flexible but requires the user to provide two separate arguments. Examples:
+    * python oomptimizer.py --pretrained-name nvidia/canary-1b
+    * python oomptimizer.py --module-name nemo.collections.asr.models.EncDecMultiTaskModel \
+        --config-path examples/asr/conf/speech_multitask/fast-conformer_aed.yaml
+
+    Dynamic bucketing is notoriously difficult to tune as you risk running into CUDA OOM many steps into the training.
+    In order to simplify finding the optimal settings, OOMptimizer scans each bucket to find the maximum possible
+    batch size that doesn't trigger a CUDA OOM.
+
+    \b
+    The suggested workflow is the following:
+    1) Run scripts/speech_recognition/estimate_duration_bins.py to get the duration distribution of your data.
+        (consider running estimate_duration_bins_2d.py for models with a strong dependency on output sequence length
+        such as attention-encoder-decoder models).
+    2) Run OOMptimizer to find the optimal batch sizes for your specific model, optimizer, and GPU.
+    3) Use these optimal settings in your actual training script and enjoy optimal GPU utilization OOM-free.
+
+    In the unlikely event that OOMptimizer bucket batch sizes are still leading to OOMs,
+    please try a lower setting of the MEMORY_FRACTION option, e.g. 0.75 (75% of GPU memory).
+    This may be required in very complex setups where there are additional GPU RAM loads that can't be anticipated
+    through the combination of training_step and optimizer update.
+    """
+    if all(opt is None for opt in (pretrained_name, module_name, config_path)):
+        click.secho(
+            "You need to provide either PRETRAINED_NAME or the pair of MODULE_NAME and CONFIG_PATH.", fg="yellow"
+        )
+        sys.exit(1)
+    logging.setLevel(logging.CRITICAL)
+    torch.cuda.set_per_process_memory_fraction(memory_fraction, device)
+
+    trainer = pl.Trainer(barebones=True)
+    trainer.log_every_n_steps = 1000000
+    model_clones = []
+    for _ in range(2 if ddp else 1):
+        if pretrained_name is not None:
+            assert (
+                config_path is None and module_name is None
+            ), "--pretrained-name cannot be used together with --module-name/--config-path"
+            click.echo(f"Intializing ASR model from pretrained checkpoint {pretrained_name}.")
+            if pretrained_name.endswith('.nemo'):
+                model = ASRModel.restore_from(pretrained_name, trainer=trainer).to(device)
+            else:
+                model = ASRModel.from_pretrained(pretrained_name, trainer=trainer).to(device)
+        else:
+            assert config_path is not None, "--module-name requires --config-path to be specified as well."
+            assert module_name is not None, "--config-path requires --module-name to be specified as well."
+            cfg = OmegaConf.load(config_path)
+            namespace, name = module_name.rsplit('.', maxsplit=1)
+            model_cls = getattr(importlib.import_module(namespace), name)
+            model = model_cls(cfg=cfg.model, trainer=trainer).to(device)
+        model_clones.append(model)
+    model = model_clones[-1]
+
+    if not hasattr(model, "oomptimizer_schema"):
+        click.secho(
+            f"We read model of type {type(model)} which doesn't seem to support OOMptimizer "
+            f"(we could not find the property .oomptimizer_schema).",
+            fg="red",
+        )
+        sys.exit(1)
+
+    schema = model.oomptimizer_schema
+
+    click.echo("Setting up the optimizers.")
+    optimizer, _ = model.setup_optimization({"name": optimizer_name, "lr": 1e-7, "weight_decay": 0.0})
+
+    is_2d_bucketing = all(
+        isinstance(item, (list, tuple)) and len(item) == 2 and all(isinstance(v, Number) for v in item)
+        for item in buckets
+    )
+    # Determine modality for input and output.
+    modalities = [
+        (
+            "text"
+            if any(
+                isinstance(item["type"], NeuralType)
+                and isinstance(item["type"].elements_type, LabelsType)
+                and item["seq_length"] == direction
+                for item in schema["inputs"]
+                if item["type"] != "dummy"
+            )
+            else "audio"
+        )
+        for direction in ("input", "output")
+    ]
+
+    def get_max_seq_lens(buckets):
+
+        def _determine_lens_for_bucket(bin):
+            if is_2d_bucketing:
+                input_len, output_len = bin
+            else:
+                input_len = bin
+                output_len = math.ceil(ratio * input_len)
+            sampling_rate = getattr(
+                model, "sample_rate", 16000
+            )  # TODO: may need to extend schema for broader model coverage
+            match modalities:
+                case "audio", "audio":
+                    return (
+                        compute_num_samples(input_len, sampling_rate=sampling_rate),
+                        compute_num_samples(output_len, sampling_rate=sampling_rate),
+                    )
+                case "audio", "text":
+                    return (compute_num_samples(input_len, sampling_rate=sampling_rate), output_len)
+                case "text", "audio":
+                    return (
+                        input_len,
+                        compute_num_samples(output_len, sampling_rate=sampling_rate),
+                    )
+                case "text", "text":
+                    return input_len, output_len
+                case _:
+                    raise RuntimeError(f"Unexpected modality combination: {_}")
+
+        return [_determine_lens_for_bucket(bin) for bin in buckets]
+
+    click.echo("Starting profiling.")
+    max_seq_lens = get_max_seq_lens(buckets)
+    gen = ProfilingBatchGenerator(schema=schema, start_batch_size=start_batch_size, rel_gap_thresh=threshold)
+    profile = {}
+
+    # Iterate buckets from the largest to the smallest sequences. This usually ends up creating
+    # a tiny bit smaller batches, likely due to worse memory fragmentation.
+    with torch.autocast("cuda", getattr(torch, dtype)):
+        for bucket, (seq_len_in, seq_len_out) in reversed(list(zip(buckets, max_seq_lens))):
+            click.echo(f"The current sequence lengths are: input={seq_len_in} output={seq_len_out}.")
+            gen.reset()
+            batch_idx = 0
+
+            def step():
+                click.echo(
+                    f"\t[BEGIN step] [CUDA RAM CURRENT: {torch.cuda.memory_allocated() / (1024 * 1024):.1f}MB] [CUDA RAM MAX: {torch.cuda.max_memory_allocated() / (1024*1024):.1f}MB]"
+                )
+                batch = gen(seq_len_in, seq_len_out)
+                oom = False
+                try:
+                    click.echo(f"\tCurrent gap: {gen.current_rel_gap}... ", nl=False)
+                    optimizer.zero_grad()
+                    out = model.training_step(batch, batch_idx)
+                    out['loss'].sum().backward()
+                    optimizer.step()
+                except torch.cuda.OutOfMemoryError as e:
+                    click.secho(f"OOM!", fg="yellow")
+                    oom = True
+                except RuntimeError as e:
+                    if "cuFFT error: CUFFT_INTERNAL_ERROR" not in str(e):
+                        raise
+                    click.secho(f"OOM!", fg="yellow")
+                    oom = True
+                else:
+                    click.secho(f"OK!", fg="green")
+                finally:
+                    click.echo(
+                        f"\t[END step] [CUDA RAM CURRENT: {torch.cuda.memory_allocated() / (1024 * 1024):.1f}MB] [CUDA RAM MAX: {torch.cuda.max_memory_allocated() / (1024*1024):.1f}MB]"
+                    )
+                    del batch
+                    # Note: We could call empty_cache() to free up some more memory on the GPU,
+                    #       but we have found out empirically that this causes a mismatched condition
+                    #       between OOMptimizer and the actual training. During training, there is some
+                    #       degree of memory fragmentation and it's better to simulate that in OOMptimizer.
+                    # torch.cuda.memory.empty_cache()
+                    torch.cuda.reset_max_memory_allocated()
+                return oom
+
+            oom = step()
+            while not (finished := gen.advance(oom)):
+                click.echo("\t" + "=" * 80)
+                oom = step()
+
+            click.secho(
+                f"=> Optimal setting for bucket={bucket} (input={seq_len_in} output={seq_len_out}) is max_batch_size={gen.max_batch_size}",
+                fg="green",
+            )
+            profile[(bucket, seq_len_in, seq_len_out)] = gen.max_batch_size
+            gen.start_batch_size = gen.max_batch_size * 2
+
+    # Reverse the profile to be ascendingly sorted again.
+    profile = dict(reversed(list(profile.items())))
+
+    click.echo("The 1st stage profile is:")
+    for (bucket, seq_len_in, seq_len_out), bs in profile.items():
+        click.echo(f"Bucket={bucket} (input={seq_len_in} output={seq_len_out}) => max_batch_size={bs}")
+
+    if is_2d_bucketing:
+        # 2D bucketing doesn't support bucket merging.
+        final_profile = [["[" + ",".join(map(str, b)) + "]", bs] for (b, _, __), bs in profile.items()]
+    else:
+        click.echo("Bucket merging stage...")
+        final_profile = []
+        for idx, ((bucket, seq_len_in, seq_len_out), bs) in enumerate(profile.items()):
+            if idx == 0:
+                final_profile.append([bucket, bs])
+                continue
+            if bs == final_profile[-1][1]:
+                click.echo(f"Merging bucket {idx} with bucket {idx-1} due to identical batch sizes.")
+                final_profile[-1][0] = bucket
+                continue
+            final_profile.append([bucket, bs])
+
+    click.secho(f"The profile was created with the following settings:")
+    click.secho(f"* using {memory_fraction:.1%} of available GPU RAM.")
+    click.secho(f"* {'' if ddp else 'not '}simulating DDP memory overhead.")
+    click.secho(f"* using AMP with dtype={dtype}.")
+    click.secho("The final profile is:", bold=True)
+    click.secho("\tbucket_duration_bins=[" + ",".join(str(seqlen) for seqlen, bs in final_profile) + "]", bold=True)
+    click.secho("\tbucket_batch_size=[" + ",".join(str(bs) for seqlen, bs in final_profile) + "]", bold=True)
+
+
+if __name__ == "__main__":
+    oomptimizer()

scripts/speech_recognition/partial_conversion_to_tarred_audio_dataset.py

@@ -0,0 +1,199 @@
+# 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 json
+import os
+from dataclasses import dataclass, field
+from typing import Optional
+
+import hydra
+from convert_to_tarred_audio_dataset import ASRTarredDatasetBuilder, ASRTarredDatasetMetadata
+from hydra.core.config_store import ConfigStore
+from joblib import Parallel, delayed
+from omegaconf import MISSING
+from tqdm import tqdm
+
+"""
+# Partial Tarred Audio Dataset Creator
+
+## Overview
+
+This script facilitates the creation of tarred and sharded audio datasets from existing tarred manifests. It allows you to select specific shards from a manifest file and then tar them separately. 
+
+This is useful in several scenarios:
+- When you only need to process a specific subset of shards (e.g., for debugging or incremental dataset preparation).
+- When you want to parallelize shard creation across multiple SLURM jobs to accelerate the dataset generation process and overcome per-job time limits.
+
+## Prerequisites
+
+- Ensure that the `convert_to_tarred_audio_dataset` script is correctly configured and run with the `--only_manifests` flag to generate the necessary manifest files.
+- Make sure the paths to the manifest and metadata files are correct and accessible.
+
+## Usage
+
+### Script Execution
+
+To run the script, use the following command:
+
+python partial_convertion_to_tarred_audio_dataset.py \
+    # the path to the tarred manifest file that contains the entries for the shards you want to process. This option is mandatory.
+    --tarred_manifest_filepath=<path to the tarred manifest file > \
+    # any other optional argument
+    --output_dir=<output directory for tarred shards> \
+    --shards_to_tar=<shard IDs to be tarred> \
+    --num_workers=-1 \
+    --dataset_metadata_filepath=<dataset metadata YAML filepath>
+
+Example:
+python partial_convertion_to_tarred_audio_dataset.py \
+    tarred_manifest_filepath="path/to/manifest.json" \
+    shards_to_tar="0:3"
+"""
+
+
+def select_shards(manifest_filepath: str, shards_to_tar: str, slice_with_offset: bool = False):
+    """
+    Selects and returns a subset of shards from the tarred manifest file.
+
+    Args:
+        manifest_filepath (str): The path to the tarred manifest file.
+        shards_to_tar (str): A range or list of shard IDs to select, e.g., "0:5" or "0,1,2".
+        slice_with_offset (bool, optional): If True, slices entries based on audio offsets. Defaults to False.
+
+    Raises:
+        FileNotFoundError: If the manifest file does not exist.
+        KeyError: If `slice_with_offset` is enabled but required fields are missing in the manifest entries.
+
+    Returns:
+        Dict[int, List[Dict[str, any]]]: A dictionary where the keys are shard IDs and the values are lists of entries for those shards.
+    """
+    shard_ids = []
+    if shards_to_tar != "all":
+        if ":" not in shards_to_tar:
+            shard_ids = [int(shards_to_tar)]
+        else:
+            start_shard_idx, end_shard_idx = map(
+                lambda x: int(x.strip()) if x.strip() else None, shards_to_tar.split(":")
+            )
+            shard_ids = list(range(start_shard_idx, end_shard_idx))
+
+    entries_to_shard = {}
+    with open(manifest_filepath, 'r') as manifest:
+        for line in tqdm(manifest, desc="Selecting shards"):
+            entry = json.loads(line)
+            if shards_to_tar == "all" or entry['shard_id'] in shard_ids:
+                if entry['shard_id'] not in entries_to_shard:
+                    entries_to_shard[entry['shard_id']] = []
+
+                if slice_with_offset:
+                    if 'abs_audio_filepath' not in entry or 'source_audio_offset' not in entry:
+                        raise KeyError(
+                            f"`slice_with_offset` is enabled, but `abs_audio_filepath` and/or `source_audio_offset` are not found in the entry:\n{entry}."
+                        )
+                    entry['audio_filepath'] = entry.pop('abs_audio_filepath')
+                    entry['offset'] = entry.pop('source_audio_offset')
+
+                entries_to_shard[entry['shard_id']].append(entry)
+
+    return entries_to_shard
+
+
+@dataclass
+class PartialASRTarredDatasetConfig:
+    """
+    Configuration class for creating partial tarred audio dataset shards.
+
+    Attributes:
+        tarred_manifest_filepath (str): The path to the tarred manifest file.
+        output_dir (Optional[str]): Directory where the output tarred shards will be saved.
+        shards_to_tar (Optional[str]): A range or list of shard IDs to tar.
+        num_workers (int): Number of parallel workers to use for tar file creation.
+        dataset_metadata_filepath (Optional[str]): Path to the dataset metadata YAML file.
+        dataset_metadata (ASRTarredDatasetMetadata): Dataset metadata configuration.
+    """
+
+    tarred_manifest_filepath: str = MISSING
+    output_dir: Optional[str] = None
+    shards_to_tar: Optional[str] = "all"
+    num_workers: int = 1
+    dataset_metadata_filepath: Optional[str] = None
+    dataset_metadata: ASRTarredDatasetMetadata = field(default=ASRTarredDatasetMetadata)
+    slice_with_offset: bool = False
+
+
+def create_shards(cfg: PartialASRTarredDatasetConfig):
+    """
+    Creates tarred shards based on the provided configuration.
+
+    Args:
+        cfg (PartialASRTarredDatasetConfig): The configuration object containing paths, shard IDs, and metadata.
+
+    Raises:
+        ValueError: If the `tarred_manifest_filepath` is None.
+        FileNotFoundError: If the tarred manifest file or dataset metadata file does not exist.
+
+    Notes:
+        - Reads the tarred manifest file and selects the specified shards.
+        - Creates tarred shards in parallel using the `ASRTarredDatasetBuilder`.
+        - The `dataset_metadata_filepath` is inferred if not provided.
+    """
+    if cfg.tarred_manifest_filepath is None:
+        raise ValueError("The `tarred_manifest_filepath` cannot be `None`. Please check your configuration.")
+
+    if not os.path.exists(cfg.tarred_manifest_filepath):
+        raise FileNotFoundError(
+            f"The `tarred_manifest_filepath` was not found: {cfg.tarred_manifest_filepath}. Please verify that the filepath is correct."
+        )
+
+    if cfg.dataset_metadata_filepath is None:
+        cfg.dataset_metadata_filepath = os.path.join(os.path.dirname(cfg.tarred_manifest_filepath), "metadata.yaml")
+
+    if cfg.output_dir is None:
+        cfg.output_dir = os.path.dirname(cfg.tarred_manifest_filepath)
+
+    if not os.path.exists(cfg.dataset_metadata_filepath):
+        raise FileNotFoundError(
+            f"The `dataset_metadata_filepath` was not found: {cfg.dataset_metadata_filepath}. Please verify that the filepath is correct."
+        )
+    else:
+        cfg.dataset_metadata = ASRTarredDatasetMetadata.from_file(cfg.dataset_metadata_filepath)
+
+    entries_to_shard = select_shards(
+        cfg.tarred_manifest_filepath, cfg.shards_to_tar, cfg.dataset_metadata.dataset_config.slice_with_offset
+    )
+
+    builder = ASRTarredDatasetBuilder()
+    builder.configure(cfg.dataset_metadata.dataset_config)
+
+    with Parallel(n_jobs=cfg.num_workers, verbose=len(entries_to_shard)) as parallel:
+        # Call parallel tarfile construction
+        _ = parallel(
+            delayed(builder._create_shard)(
+                entries=entries_to_shard[shard_id],
+                target_dir=cfg.output_dir,
+                shard_id=shard_id,
+            )
+            for shard_id in entries_to_shard
+        )
+
+
+@hydra.main(config_path=None, config_name='partial_tar_config')
+def main(cfg: PartialASRTarredDatasetConfig):
+    create_shards(cfg)
+
+
+ConfigStore.instance().store(name='partial_tar_config', node=PartialASRTarredDatasetConfig)
+
+if __name__ == '__main__':
+    main()