chore: import upstream snapshot with attribution

scripts/pseudo_labeling/README.md

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

scripts/pseudo_labeling/update_inference_config.py

@@ -0,0 +1,137 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import argparse
+import glob
+import math
+import os
+from typing import List, Union
+
+
+from filelock import FileLock
+from omegaconf import ListConfig, OmegaConf
+
+
+def count_files_for_tarred_pseudo_labeling(manifest_filepath: Union[str, ListConfig]) -> int:
+    """
+    Counts the total number of entries across multiple manifest files.
+    Args:
+        manifest_filepath (Union[str, ListConfig]): The file path to the manifest files.
+    Returns:
+        int: The total number of entries across all matching manifest files.
+    """
+    # Convert ListConfig to string if needed
+    if isinstance(manifest_filepath, ListConfig):
+        manifest_filepath = manifest_filepath[0]  # Use the first element if it's a list or ListConfig
+    dir_path, filename = os.path.split(manifest_filepath)
+    prefix = filename.split('_', 1)[0]
+    number_of_files = 0
+    for full_path in glob.glob(os.path.join(dir_path, f"{prefix}_[0-9]*.json")):
+        with open(full_path, 'r') as f:
+            number_of_files += len(f.readlines())
+    return number_of_files
+
+
+def count_files_for_pseudo_labeling(manifest_filepath: Union[str, list, ListConfig]) -> int:
+    """
+    Counts the number of entries in a single manifest file .
+    Args:
+        manifest_filepath (Union[str, list, ListConfig]): The file path to the manifest file.
+    Returns:
+        int: The total number of entries (lines) in the manifest file.
+    """
+    # Convert ListConfig to string if needed
+    if isinstance(manifest_filepath, list) or isinstance(manifest_filepath, ListConfig):
+        manifest_filepath = manifest_filepath[0]
+    with open(manifest_filepath, 'r') as f:
+        number_of_files = len(f.readlines())
+    return number_of_files
+
+
+def export_limit_predict_batches(inference_configs: List[str], p_cache: float, num_gpus: int) -> None:
+    """
+    Updates inference configuration files to set `limit_predict_batches`.
+    This is done to force partial transcription of unlabeled dataset for dynamic update of PLs.
+
+    Args:
+        inference_configs (List[str]): A list of file paths to the inference configuration files.
+        p_cache (float): A scaling factor for the cache to adjust the number of batches.
+        num_gpus (int): The number of GPUs available for inference.
+
+    Returns:
+        None: The function modifies and saves the updated configuration files in-place.
+    """
+    for config_path in inference_configs:
+        config = OmegaConf.load(config_path)
+        tarred_audio_filepaths = config.predict_ds.get("tarred_audio_filepaths", None)
+        manifest_filepaths = config.predict_ds.manifest_filepath
+
+        if tarred_audio_filepaths:
+            number_of_files = count_files_for_tarred_pseudo_labeling(manifest_filepaths)
+        else:
+            number_of_files = count_files_for_pseudo_labeling(manifest_filepaths)
+
+        if hasattr(config.predict_ds, "batch_size"):
+            batch_size = config.predict_ds.batch_size
+            limit_predict_batches = math.ceil((number_of_files * p_cache) / (batch_size * num_gpus))
+            OmegaConf.update(config, "trainer.limit_predict_batches", limit_predict_batches)
+            OmegaConf.save(config, config_path)
+        elif hasattr(config.predict_ds, "batch_duration"):
+            batch_duration = config.predict_ds.batch_duration
+            average_audio_len = 10
+            limit_predict_batches = math.ceil(
+                (number_of_files * average_audio_len * p_cache) / (batch_duration * num_gpus)
+            )
+            OmegaConf.update(config, "trainer.limit_predict_batches", limit_predict_batches)
+            OmegaConf.save(config, config_path)
+        else:
+            batch_size = 32
+            limit_predict_batches = math.ceil((number_of_files * p_cache) / (batch_size * num_gpus))
+            OmegaConf.update(config, "trainer.limit_predict_batches", limit_predict_batches)
+            OmegaConf.save(config, config_path)
+
+
+def main():
+    rank = int(os.environ.get("RANK", 0))  # Default to 0 if not set
+
+    # Ensure only one process executes this block
+    parser = argparse.ArgumentParser(description="Export limit_predict_batches as environment variables.")
+    parser.add_argument(
+        "--inference_configs",
+        type=str,
+        nargs='+',  # Accepts one or more values as a list
+        required=True,
+        help="Paths to one or more inference config YAML files.",
+    )
+    parser.add_argument("--p_cache", type=float, required=True, help="Pseudo-label cache fraction.")
+    parser.add_argument("--num_gpus", type=int, required=True, help="Number of GPUs available.")
+
+    args = parser.parse_args()
+    lock_dir = os.path.dirname(args.inference_configs[0])
+    lock_file = lock_dir + "/my_script.lock"
+    # Code executed by all processes
+
+    #     # Code executed by a single process
+    with FileLock(lock_file):
+        if rank == 0:
+            export_limit_predict_batches(
+                inference_configs=args.inference_configs, p_cache=args.p_cache, num_gpus=args.num_gpus
+            )
+
+    # Remove the lock file after the FileLock context is exited
+    if os.path.exists(lock_file):
+        os.remove(lock_file)
+
+
+if __name__ == "__main__":
+    main()

scripts/pseudo_labeling/write_transcribed_files.py

@@ -0,0 +1,260 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import argparse
+import glob
+import json
+import os
+from typing import List
+
+
+from filelock import FileLock
+
+from nemo.utils import logging
+
+
+def create_transcribed_shard_manifests(prediction_filepaths: List[str]) -> List[str]:
+    """
+    Creates transcribed shard manifest files by processing predictions and organizing them by shard ID.
+
+    This function reads a `predictions_all.json` file from each given directory, organizes the data by
+    shard IDs, and writes the entries to separate shard manifest files. For each shard, the `pred_text`
+    field is updated as the main transcription (`text`), and the original transcription (`text`) is
+    stored as `orig_text`.
+
+    Args:
+        prediction_filepaths (List[str]): A list of file paths to directories containing
+            `predictions_all.json` files with prediction data, including shard IDs.
+
+    Returns:
+        List[str]: A list of file paths to the combined manifest files (`transcribed_manifest__OP_0..CL_.json`)
+        created for each directory.
+    """
+    all_manifest_filepaths = []
+    for prediction_filepath in prediction_filepaths:
+        max_shard_id = 0
+        shard_data = {}
+        full_path = os.path.join(prediction_filepath, "predictions_all.json")
+        with open(full_path, 'r') as f:
+            for line in f.readlines():
+                data_entry = json.loads(line)
+                shard_id = data_entry.get("shard_id")
+                if max_shard_id < shard_id:
+                    max_shard_id = shard_id
+                if shard_id not in shard_data:
+                    shard_data[shard_id] = []
+                shard_data[shard_id].append(data_entry)
+        for shard_id, entries in shard_data.items():
+            output_filename = os.path.join(prediction_filepath, f"transcribed_manifest_{shard_id}.json")
+            with open(output_filename, 'w') as f:
+                for data_entry in entries:
+                    if data_entry['audio_filepath'].endswith(".wav"):
+                        if 'text' in data_entry:
+                            data_entry['orig_text'] = data_entry.pop('text')
+                        data_entry['text'] = data_entry.pop('pred_text')
+                        json.dump(data_entry, f, ensure_ascii=False)
+                        f.write("\n")
+        shard_manifest_filepath = os.path.join(
+            prediction_filepath, f"transcribed_manifest__OP_0..{max_shard_id}_CL_.json"
+        )
+        all_manifest_filepaths.append(shard_manifest_filepath)
+    return all_manifest_filepaths
+
+
+def create_transcribed_manifests(prediction_filepaths: List[str]) -> List[str]:
+    """
+    Creates updated transcribed manifest files by processing predictions.
+
+    This function reads prediction files (`predictions_all.json`) from the provided directories,
+    updates the transcription data by renaming the `pred_text` field to `text`, and stores the
+    original `text` field as `orig_text`. The updated data is written to new transcribed manifest
+    files (`transcribed_manifest.json`) in each directory.
+
+    Args:
+        prediction_filepaths (List[str]): A list of file paths to directories containing
+            prediction files (`predictions_all.json`).
+
+    Returns:
+        List[str]: A list of file paths to the newly created transcribed manifest files
+        (`transcribed_manifest.json`).
+    """
+    all_manifest_filepaths = []
+    for prediction_filepath in prediction_filepaths:
+        prediction_name = os.path.join(prediction_filepath, "predictions_all.json")
+        transcripted_name = os.path.join(prediction_filepath, f"transcribed_manifest.json")
+
+        # Open and read the original predictions_all.json file
+        with open(transcripted_name, 'w', encoding='utf-8') as f:
+            with open(prediction_name, 'r', encoding='utf-8') as pred_f:
+
+                for line in pred_f.readlines():
+                    data_entry = json.loads(line)
+                    if 'text' in data_entry:
+                        data_entry['orig_text'] = data_entry.pop('text')
+                    data_entry['text'] = data_entry.pop('pred_text')
+                    json.dump(data_entry, f, ensure_ascii=False)
+                    f.write("\n")
+            # Append the path of the new manifest file to the list
+            all_manifest_filepaths.append(transcripted_name)
+
+    return all_manifest_filepaths
+
+
+def write_sampled_shard_transcriptions(manifest_filepaths: List[str]) -> List[List[str]]:
+    """
+    Updates transcriptions by merging predicted shard data and transcribed manifest data.
+    This function processes prediction and transcribed manifest files, merges them
+    by matching the shard_id and audio file paths. For each shard, the corresponding
+    data entries are written to a new file.
+    Args:
+        manifest_filepaths (List[str]): A list of file paths to directories containing
+            prediction and transcribed manifest files.
+    Returns:
+        List[List[str]]: A list of lists containing the file paths to the generated
+            transcribed shard manifest files.
+    """
+    all_manifest_filepaths = []
+
+    # Process each prediction directory
+    for prediction_filepath in manifest_filepaths:
+        predicted_shard_data = {}
+        # Collect entries from prediction files based on shard id
+        prediction_path = os.path.join(prediction_filepath, "predictions_all.json")
+        with open(prediction_path, 'r') as f:
+            for line in f:
+                data_entry = json.loads(line)
+                shard_id = data_entry.get("shard_id")
+                audio_filepath = data_entry['audio_filepath']
+                predicted_shard_data.setdefault(shard_id, {})[audio_filepath] = data_entry
+    max_shard_id = 0
+    for full_path in glob.glob(os.path.join(prediction_filepath, f"transcribed_manifest_[0-9]*.json")):
+        all_data_entries = []
+        with open(full_path, 'r') as f:
+            for line in f:
+                data_entry = json.loads(line)
+                shard_id = data_entry.get("shard_id")
+                max_shard_id = max(max_shard_id, shard_id)
+                all_data_entries.append(data_entry)
+        # Write the merged data to a new manifest file keeping new transcriptions
+        output_filename = os.path.join(prediction_filepath, f"transcribed_manifest_{shard_id}.json")
+        with open(output_filename, 'w') as f:
+            for data_entry in all_data_entries:
+                audio_filepath = data_entry['audio_filepath']
+                # Escape duplicated audio files that end with *dup
+                if audio_filepath.endswith(".wav"):
+                    if shard_id in predicted_shard_data and audio_filepath in predicted_shard_data[shard_id]:
+                        predicted_data_entry = predicted_shard_data[shard_id][audio_filepath]
+                        if 'text' in predicted_data_entry:
+                            predicted_data_entry['orig_text'] = predicted_data_entry.pop('text')
+                        if "pred_text" in predicted_data_entry:
+                            predicted_data_entry['text'] = predicted_data_entry.pop('pred_text')
+                        json.dump(predicted_data_entry, f, ensure_ascii=False)
+                    else:
+                        json.dump(data_entry, f, ensure_ascii=False)
+                    f.write("\n")
+
+    shard_manifest_filepath = os.path.join(prediction_filepath, f"transcribed_manifest__OP_0..{max_shard_id}_CL_.json")
+    all_manifest_filepaths.append([shard_manifest_filepath])
+
+    return all_manifest_filepaths
+
+
+def write_sampled_transcriptions(manifest_filepaths: List[str]) -> List[str]:
+    """
+    Updates transcriptions by merging predicted data with transcribed manifest data.
+
+    This function processes prediction and transcribed manifest files within given directories.
+    It matches audio file paths to update transcriptions with predictions, ensuring each audio file
+    is properly transcribed. The updated data is written to the transcribed manifest file.
+
+    Args:
+        manifest_filepaths (List[str]): A list of file paths to directories containing
+            the prediction file (`predictions_all.json`) and the transcribed manifest file
+            (`transcribed_manifest.json`).
+
+    Returns:
+        List[str]: A list of file paths to the updated transcribed manifest files.
+    """
+    all_manifest_filepaths = []
+    for prediction_filepath in manifest_filepaths:
+        predicted_data = {}
+        prediction_path = os.path.join(prediction_filepath, "predictions_all.json")
+        with open(prediction_path, 'r') as f:
+            for line in f:
+                data_entry = json.loads(line)
+                path = data_entry['audio_filepath']
+                predicted_data[path] = data_entry
+
+        full_path = os.path.join(prediction_filepath, f"transcribed_manifest.json")
+        all_data_entries = []
+        with open(full_path, 'r') as f:
+            for line in f:
+                data_entry = json.loads(line)
+                all_data_entries.append(data_entry)
+
+        output_filename = os.path.join(prediction_filepath, f"transcribed_manifest.json")
+        with open(output_filename, 'w') as f:
+            for data_entry in all_data_entries:
+                audio_filepath = data_entry['audio_filepath']
+                if audio_filepath.endswith(".wav"):
+                    if audio_filepath in predicted_data:
+                        predicted_data_entry = predicted_data[audio_filepath]
+                        if 'text' in predicted_data_entry:
+                            predicted_data_entry['orig_text'] = predicted_data_entry.pop('text')
+                        predicted_data_entry['text'] = predicted_data_entry.pop('pred_text')
+                        json.dump(predicted_data_entry, f, ensure_ascii=False)
+                        f.write("\n")
+                    else:
+                        json.dump(data_entry, f, ensure_ascii=False)
+                        f.write("\n")
+        all_manifest_filepaths.append(output_filename)
+    return all_manifest_filepaths
+
+
+if __name__ == "__main__":
+    rank = int(os.environ.get("RANK", 0))  # Default to 0 if not set
+
+    parser = argparse.ArgumentParser(description="Script to create or write transcriptions")
+    parser.add_argument("--is_tarred", action="store_true", help="If true, processes tarred manifests")
+    parser.add_argument("--full_pass", action="store_true", help="If true, processes full pass manifests")
+    parser.add_argument(
+        "--prediction_filepaths",
+        type=str,
+        nargs='+',  # Accepts one or more values as a list
+        required=True,
+        help="Paths to one or more inference config YAML files.",
+    )
+
+    args = parser.parse_args()
+
+    lock_dir = os.path.dirname(args.prediction_filepaths[0])
+    lock_file = lock_dir + "/my_script.lock"
+
+    with FileLock(lock_file):
+        if rank == 0:
+            if args.is_tarred:
+                result = (
+                    write_sampled_shard_transcriptions(args.prediction_filepaths)
+                    if not args.full_pass
+                    else create_transcribed_shard_manifests(args.prediction_filepaths)
+                )
+            else:
+                result = (
+                    write_sampled_transcriptions(args.prediction_filepaths)
+                    if not args.full_pass
+                    else create_transcribed_manifests(args.prediction_filepaths)
+                )
+
+    # Remove the lock file after the FileLock context is exited
+    if os.path.exists(lock_file):
+        os.remove(lock_file)