chore: import upstream snapshot with attribution

examples/asr/speech_pretraining/README.md

@@ -0,0 +1,30 @@
+# Speech Self-Supervised Learning
+
+This directory contains example scripts to self-supervised speech models. 
+
+There are two main types of supported self-supervised learning methods:
+- [Wav2vec-BERT](https://arxiv.org/abs/2108.06209): `speech_pre_training.py`
+- [NEST](https://arxiv.org/abs/2408.13106): `masked_token_pred_pretrain.py`
+    - For downstream tasks that use NEST as multi-layer feature extractor, please refer to `./downstream/speech_classification_mfa_train.py`
+    - For extracting multi-layer features from NEST, please refer to `<NEMO ROOT>/scripts/ssl/extract_features.py`
+    - For using NEST as weight initialization for downstream tasks, please refer to the usage of [maybe_init_from_pretrained_checkpoint](https://github.com/NVIDIA/NeMo/blob/main/nemo/core/classes/modelPT.py#L1242).
+
+
+For their corresponding usage, please refer to the example yaml config:
+- Wav2vec-BERT: `examples/asr/conf/ssl/fastconformer/fast-conformer.yaml`
+- NEST: `examples/asr/conf/ssl/nest/nest_fast-conformer.yaml`
+
+
+The dataset format follows that of ASR models, but no groundtruth transcriptions are needed. For example, the jsonl file specified in `manifest_filepath` should look like:
+```
+{"audio_filepath": "path/to/audio1.wav", "duration": 10.0, "text": ""}
+{"audio_filepath": "path/to/audio2.wav", "duration": 5.0, "text": ""}
+```
+
+Please refer to the [ASR dataset documentation](https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/asr/datasets.html#preparing-custom-asr-data%60) for more details.
+
+
+For most efficient data loading, please refer to 
+- [lhotse dataloading](https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/asr/datasets.html#lhotse-dataloading) 
+- [pre-compute bucket durations](https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/asr/datasets.html#pre-computing-bucket-duration-bins)
+- [optimizing GPU memory usage](https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/asr/datasets.html#pushing-gpu-utilization-to-the-limits-with-bucketing-and-oomptimizer)

examples/asr/speech_pretraining/downstream/speech_classification_mfa_train.py

@@ -0,0 +1,108 @@
+# 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 collections import OrderedDict
+
+import lightning.pytorch as pl
+import torch
+from omegaconf import OmegaConf
+
+from nemo.collections.asr.models import EncDecSpeakerLabelModel, SpeechEncDecSelfSupervisedModel
+from nemo.core.classes.common import typecheck
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+typecheck.set_typecheck_enabled(enabled=False)
+
+"""
+Example script for training a speech classification model with a self-supervised pre-trained encoder, and 
+use the SSL encoder for multi-layer feature extraction.
+
+# Example of training a speaker classification model with a self-supervised pre-trained encoder
+```sh
+python speech_classification_mfa_train.py \
+    # (Optional: --config-path=<path to dir of configs> --config-name=<name of config without .yaml>) \
+    ++init_from_nemo_model=<path to pre-trained SSL .nemo file> \
+    # or use ++init_from_pretrained_model=<model_name> \
+    model.train_ds.manifest_filepath=<path to train manifest> \
+    model.validation_ds.manifest_filepath=<path to val manifest> \
+    trainer.devices=-1 \
+    trainer.accelerator="gpu" \
+    strategy="ddp"  \
+    trainer.max_epochs=100 \
+    model.optim.name="adamw" \
+    model.optim.lr=0.001 \
+    model.optim.betas=[0.9,0.999] \
+    model.optim.weight_decay=0.0001 \
+    model.optim.sched.warmup_steps=2000
+    exp_manager.create_wandb_logger=True \
+    exp_manager.wandb_logger_kwargs.name="<Name of experiment>" \
+    exp_manager.wandb_logger_kwargs.project="<Namex of project>"
+```
+    
+"""
+
+
+def load_ssl_encoder(model, cfg):
+    if cfg.get("init_from_ptl_ckpt", None) is not None:
+        state_dict = torch.load(cfg.init_from_ptl_ckpt, map_location='cpu')['state_dict']
+        logging.info(f"Loading encoder from PyTorch Lightning checkpoint: {cfg.init_from_ptl_ckpt}")
+    elif cfg.get("init_from_nemo_model", None) is not None:
+        ssl_model = SpeechEncDecSelfSupervisedModel.restore_from(cfg.init_from_nemo_model, map_location='cpu')
+        state_dict = ssl_model.state_dict()
+        logging.info(f"Loading encoder from NeMo model: {cfg.init_from_nemo_model}")
+    elif cfg.get("init_from_pretrained_model", None) is not None:
+        ssl_model = SpeechEncDecSelfSupervisedModel.from_pretrained(cfg.init_from_pretrained_model, map_location='cpu')
+        state_dict = ssl_model.state_dict()
+        logging.info(f"Loading encoder from pretrained model: {cfg.init_from_pretrained_model}")
+    else:
+        logging.info("No model checkpoint or pretrained model specified for encoder initialization.")
+        return model
+
+    encoder_state_dict = OrderedDict()
+    for key, value in state_dict.items():
+        if key.startswith('encoder.'):
+            encoder_state_dict[f'preprocessor.feature_extractor.{key}'] = value
+
+    model.load_state_dict(encoder_state_dict, strict=False)
+    logging.info("Loaded ssl encoder state dict.")
+
+    return model
+
+
+@hydra_runner(config_path="../conf/ssl/nest/multi_layer_feat", config_name="nest_ecapa_tdnn_small")
+def main(cfg):
+
+    logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}')
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+
+    speaker_model = EncDecSpeakerLabelModel(cfg=cfg.model, trainer=trainer)
+
+    if cfg.model.preprocessor.get("encoder", None) is not None:
+        # multi-layer feature extractor
+        speaker_model = load_ssl_encoder(speaker_model, cfg)
+    else:
+        speaker_model.maybe_init_from_pretrained_checkpoint(cfg)
+
+    trainer.fit(speaker_model)
+
+    if hasattr(cfg.model, 'test_ds') and cfg.model.test_ds.manifest_filepath is not None:
+        if speaker_model.prepare_test(trainer):
+            trainer.test(speaker_model)
+
+
+if __name__ == '__main__':
+    main()

examples/asr/speech_pretraining/masked_token_pred_pretrain.py

@@ -0,0 +1,65 @@
+# 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 lightning.pytorch as pl
+from omegaconf import OmegaConf
+
+from nemo.collections.asr.models.ssl_models import EncDecDenoiseMaskedTokenPredModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+"""
+# Example of training a self-supervised denoising masksed token prediction model
+```sh
+python pretrain_masked_token_pred.py \
+    # (Optional: --config-path=<path to dir of configs> --config-name=<name of config without .yaml>) \
+    model.train_ds.manifest_filepath=<path to train manifest> \
+    model.train_ds.noise_manifest=<path to noise manifest> \
+    model.validation_ds.manifest_filepath=<path to val/test manifest> \
+    model.validation_ds.noise_manifest=<path to noise manifest> \
+    trainer.devices=-1 \
+    trainer.accelerator="gpu" \
+    strategy="ddp"  \
+    trainer.max_epochs=100 \
+    model.optim.name="adamw" \
+    model.optim.lr=0.001 \
+    model.optim.betas=[0.9,0.999] \
+    model.optim.weight_decay=0.0001 \
+    model.optim.sched.warmup_steps=2000
+    exp_manager.create_wandb_logger=True \
+    exp_manager.wandb_logger_kwargs.name="<Name of experiment>" \
+    exp_manager.wandb_logger_kwargs.project="<Namex of project>"
+```
+"""
+
+
+@hydra_runner(config_path="../conf/ssl/nest", config_name="nest_fast-conformer")
+def main(cfg):
+    logging.info(f"Hydra config: {OmegaConf.to_yaml(cfg)}")
+
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+    asr_model = EncDecDenoiseMaskedTokenPredModel(cfg=cfg.model, trainer=trainer)
+
+    # Initialize the weights of the model from another model, if provided via config
+    asr_model.maybe_init_from_pretrained_checkpoint(cfg)
+
+    trainer.fit(asr_model)
+
+
+if __name__ == "__main__":
+    main()

examples/asr/speech_pretraining/speech_pre_training.py

@@ -0,0 +1,68 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import lightning.pytorch as pl
+from omegaconf import OmegaConf
+
+from nemo.collections.asr.models.ssl_models import SpeechEncDecSelfSupervisedModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+from nemo.utils.trainer_utils import resolve_trainer_cfg
+
+"""
+# Example of unsupervised pre-training of a model
+```sh
+python speech_pre_training.py \
+    # (Optional: --config-path=<path to dir of configs> --config-name=<name of config without .yaml>) \
+    model.train_ds.manifest_filepath=<path to train manifest> \
+    model.validation_ds.manifest_filepath=<path to val/test manifest> \
+    trainer.devices=-1 \
+    trainer.accelerator="gpu" \
+    strategy="ddp"  \
+    trainer.max_epochs=100 \
+    model.optim.name="adamw" \
+    model.optim.lr=0.001 \
+    model.optim.betas=[0.9,0.999] \
+    model.optim.weight_decay=0.0001 \
+    model.optim.sched.warmup_steps=2000
+    exp_manager.create_wandb_logger=True \
+    exp_manager.wandb_logger_kwargs.name="<Name of experiment>" \
+    exp_manager.wandb_logger_kwargs.project="<Namex of project>"
+```
+
+For documentation on fine-tuning, please visit -
+https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/configs.html#fine-tuning-configurations
+When doing supervised fine-tuning from unsupervised pre-trained encoder, set flag init_strict to False
+
+"""
+
+
+@hydra_runner(config_path="../conf/ssl/fastconformer/", config_name="fast-conformer")
+def main(cfg):
+    logging.info(f"Hydra config: {OmegaConf.to_yaml(cfg)}")
+
+    trainer = pl.Trainer(**resolve_trainer_cfg(cfg.trainer))
+    exp_manager(trainer, cfg.get("exp_manager", None))
+    asr_model = SpeechEncDecSelfSupervisedModel(cfg=cfg.model, trainer=trainer)
+
+    # Initialize the weights of the model from another model, if provided via config
+    asr_model.maybe_init_from_pretrained_checkpoint(cfg)
+
+    trainer.fit(asr_model)
+
+
+if __name__ == "__main__":
+    main()