Introduce code from Rust monorepo. This is a mess

tts/nemo/collections/asr/models/__init__.py

@@ -0,0 +1,22 @@
+# 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.
+
+from nemo.collections.asr.models.asr_model import ASRModel
+from nemo.collections.asr.models.classification_models import EncDecClassificationModel
+from nemo.collections.asr.models.clustering_diarizer import ClusteringDiarizer
+from nemo.collections.asr.models.ctc_bpe_models import EncDecCTCModelBPE
+from nemo.collections.asr.models.ctc_models import EncDecCTCModel
+from nemo.collections.asr.models.label_models import EncDecSpeakerLabelModel, ExtractSpeakerEmbeddingsModel
+from nemo.collections.asr.models.rnnt_bpe_models import EncDecRNNTBPEModel
+from nemo.collections.asr.models.rnnt_models import EncDecRNNTModel

tts/nemo/collections/asr/models/asr_model.py

@@ -0,0 +1,91 @@
+# 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 os
+from abc import ABC, abstractmethod
+from typing import List, Optional, Union
+
+import torch
+from omegaconf import OmegaConf
+
+from nemo.core.classes import ModelPT
+from nemo.core.classes.exportable import Exportable
+from nemo.utils import model_utils
+
+__all__ = ['ASRModel']
+
+
+class ASRModel(ModelPT, ABC):
+    @abstractmethod
+    def transcribe(self, paths2audio_files: List[str], batch_size: int = 4) -> List[str]:
+        """
+        Takes paths to audio files and returns text transcription
+        Args:
+            paths2audio_files: paths to audio fragment to be transcribed
+
+        Returns:
+            transcription texts
+        """
+        pass
+
+    def multi_validation_epoch_end(self, outputs, dataloader_idx: int = 0):
+        val_loss_mean = torch.stack([x['val_loss'] for x in outputs]).mean()
+        wer_num = torch.stack([x['val_wer_num'] for x in outputs]).sum()
+        wer_denom = torch.stack([x['val_wer_denom'] for x in outputs]).sum()
+        tensorboard_logs = {'val_loss': val_loss_mean, 'val_wer': wer_num / wer_denom}
+        return {'val_loss': val_loss_mean, 'log': tensorboard_logs}
+
+    def multi_test_epoch_end(self, outputs, dataloader_idx: int = 0):
+        val_loss_mean = torch.stack([x['test_loss'] for x in outputs]).mean()
+        wer_num = torch.stack([x['test_wer_num'] for x in outputs]).sum()
+        wer_denom = torch.stack([x['test_wer_denom'] for x in outputs]).sum()
+        tensorboard_logs = {'test_loss': val_loss_mean, 'test_wer': wer_num / wer_denom}
+        return {'test_loss': val_loss_mean, 'log': tensorboard_logs}
+
+    @classmethod
+    def list_available_models(cls) -> 'List[PretrainedModelInfo]':
+        """
+        This method returns a list of pre-trained model which can be instantiated directly from NVIDIA's NGC cloud.
+        Returns:
+            List of available pre-trained models.
+        """
+        # recursively walk the subclasses to generate pretrained model info
+        list_of_models = model_utils.resolve_subclass_pretrained_model_info(cls)
+        return list_of_models
+
+
+class ExportableEncDecModel(Exportable):
+    """
+    Simple utiliy mix-in to export models that consist of encoder/decoder pair 
+    plus pre/post processor, but have to be exported as encoder/decoder pair only
+    (covers most ASR classes)
+    """
+
+    @property
+    def input_module(self):
+        return self.encoder
+
+    @property
+    def output_module(self):
+        return self.decoder
+
+    def forward_for_export(self, input, length=None):
+        encoder_output = self.input_module(input, length)
+        if isinstance(encoder_output, tuple):
+            return self.output_module(encoder_output[0])
+        else:
+            return self.output_module(encoder_output)
+
+    def _prepare_for_export(self, **kwargs):
+        self.input_module._prepare_for_export(**kwargs)
+        self.output_module._prepare_for_export(**kwargs)

tts/nemo/collections/asr/models/classification_models.py

@@ -0,0 +1,750 @@
+# 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 copy
+import json
+import os
+import tempfile
+from abc import abstractmethod
+from math import ceil
+from typing import Dict, List, Optional, Union
+
+import torch
+from omegaconf import DictConfig, ListConfig, OmegaConf
+from pytorch_lightning import Trainer
+from torchmetrics.regression import MeanAbsoluteError, MeanSquaredError
+
+from nemo.collections.asr.data import audio_to_label_dataset
+from nemo.collections.asr.models.asr_model import ASRModel, ExportableEncDecModel
+from nemo.collections.asr.parts.preprocessing.features import WaveformFeaturizer
+from nemo.collections.asr.parts.preprocessing.perturb import process_augmentations
+from nemo.collections.common.losses import CrossEntropyLoss, MSELoss
+from nemo.collections.common.metrics import TopKClassificationAccuracy
+from nemo.core.classes.common import PretrainedModelInfo, typecheck
+from nemo.core.neural_types import *
+from nemo.utils import logging, model_utils
+
+__all__ = ['EncDecClassificationModel', 'EncDecRegressionModel']
+
+
+class _EncDecBaseModel(ASRModel, ExportableEncDecModel):
+    """Encoder decoder Classification models."""
+
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+        # Get global rank and total number of GPU workers for IterableDataset partitioning, if applicable
+        # Global_rank and local_rank is set by LightningModule in Lightning 1.2.0
+        self.world_size = 1
+        if trainer is not None:
+            self.world_size = trainer.num_nodes * trainer.num_gpus
+
+        # Convert config to a DictConfig
+        cfg = model_utils.convert_model_config_to_dict_config(cfg)
+
+        # Convert config to support Hydra 1.0+ instantiation
+        cfg = model_utils.maybe_update_config_version(cfg)
+
+        self.is_regression_task = cfg.get('is_regression_task', False)
+        # Change labels if needed
+        self._update_decoder_config(cfg.labels, cfg.decoder)
+        super().__init__(cfg=cfg, trainer=trainer)
+
+        if hasattr(self._cfg, 'spec_augment') and self._cfg.spec_augment is not None:
+            self.spec_augmentation = ASRModel.from_config_dict(self._cfg.spec_augment)
+        else:
+            self.spec_augmentation = None
+        if hasattr(self._cfg, 'crop_or_pad_augment') and self._cfg.crop_or_pad_augment is not None:
+            self.crop_or_pad = ASRModel.from_config_dict(self._cfg.crop_or_pad_augment)
+        else:
+            self.crop_or_pad = None
+
+        self.preprocessor = self._setup_preprocessor()
+        self.encoder = self._setup_encoder()
+        self.decoder = self._setup_decoder()
+        self.loss = self._setup_loss()
+        self._setup_metrics()
+
+    @abstractmethod
+    def _setup_preprocessor(self):
+        """
+        Setup preprocessor for audio data
+        Returns: Preprocessor
+
+        """
+        pass
+
+    @abstractmethod
+    def _setup_encoder(self):
+        """
+        Setup encoder for the Encoder-Decoder network
+        Returns: Encoder
+        """
+        pass
+
+    @abstractmethod
+    def _setup_decoder(self):
+        """
+        Setup decoder for the Encoder-Decoder network
+        Returns: Decoder
+        """
+        pass
+
+    @abstractmethod
+    def _setup_loss(self):
+        """
+        Setup loss function for training
+        Returns: Loss function
+
+        """
+        pass
+
+    @abstractmethod
+    def _setup_metrics(self):
+        """
+        Setup metrics to be tracked in addition to loss
+        Returns: void
+
+        """
+        pass
+
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        if hasattr(self.preprocessor, '_sample_rate'):
+            audio_eltype = AudioSignal(freq=self.preprocessor._sample_rate)
+        else:
+            audio_eltype = AudioSignal()
+        return {
+            "input_signal": NeuralType(('B', 'T'), audio_eltype),
+            "input_signal_length": NeuralType(tuple('B'), LengthsType()),
+        }
+
+    @property
+    @abstractmethod
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        pass
+
+    def forward(self, input_signal, input_signal_length):
+        processed_signal, processed_signal_len = self.preprocessor(
+            input_signal=input_signal, length=input_signal_length,
+        )
+        # Crop or pad is always applied
+        if self.crop_or_pad is not None:
+            processed_signal, processed_signal_len = self.crop_or_pad(
+                input_signal=processed_signal, length=processed_signal_len
+            )
+        # Spec augment is not applied during evaluation/testing
+        if self.spec_augmentation is not None and self.training:
+            processed_signal = self.spec_augmentation(input_spec=processed_signal, length=processed_signal_len)
+        encoded, encoded_len = self.encoder(audio_signal=processed_signal, length=processed_signal_len)
+        logits = self.decoder(encoder_output=encoded)
+        return logits
+
+    def setup_training_data(self, train_data_config: Optional[Union[DictConfig, Dict]]):
+        if 'shuffle' not in train_data_config:
+            train_data_config['shuffle'] = True
+        # preserve config
+        self._update_dataset_config(dataset_name='train', config=train_data_config)
+
+        self._train_dl = self._setup_dataloader_from_config(config=DictConfig(train_data_config))
+
+        # Need to set this because if using an IterableDataset, the length of the dataloader is the total number
+        # of samples rather than the number of batches, and this messes up the tqdm progress bar.
+        # So we set the number of steps manually (to the correct number) to fix this.
+        if 'is_tarred' in train_data_config and train_data_config['is_tarred']:
+            # We also need to check if limit_train_batches is already set.
+            # If it's an int, we assume that the user has set it to something sane, i.e. <= # training batches,
+            # and don't change it. Otherwise, adjust batches accordingly if it's a float (including 1.0).
+            if isinstance(self._trainer.limit_train_batches, float):
+                self._trainer.limit_train_batches = int(
+                    self._trainer.limit_train_batches
+                    * ceil((len(self._train_dl.dataset) / self.world_size) / train_data_config['batch_size'])
+                )
+
+    def setup_validation_data(self, val_data_config: Optional[Union[DictConfig, Dict]]):
+        if 'shuffle' not in val_data_config:
+            val_data_config['shuffle'] = False
+
+        # preserve config
+        self._update_dataset_config(dataset_name='validation', config=val_data_config)
+
+        self._validation_dl = self._setup_dataloader_from_config(config=DictConfig(val_data_config))
+
+    def setup_test_data(self, test_data_config: Optional[Union[DictConfig, Dict]]):
+        if 'shuffle' not in test_data_config:
+            test_data_config['shuffle'] = False
+
+        # preserve config
+        self._update_dataset_config(dataset_name='test', config=test_data_config)
+
+        self._test_dl = self._setup_dataloader_from_config(config=DictConfig(test_data_config))
+
+    def test_dataloader(self):
+        if self._test_dl is not None:
+            return self._test_dl
+
+    def _setup_dataloader_from_config(self, config: DictConfig):
+
+        OmegaConf.set_struct(config, False)
+        config.is_regression_task = self.is_regression_task
+        OmegaConf.set_struct(config, True)
+
+        if 'augmentor' in config:
+            augmentor = process_augmentations(config['augmentor'])
+        else:
+            augmentor = None
+
+        featurizer = WaveformFeaturizer(
+            sample_rate=config['sample_rate'], int_values=config.get('int_values', False), augmentor=augmentor
+        )
+        shuffle = config['shuffle']
+
+        # Instantiate tarred dataset loader or normal dataset loader
+        if config.get('is_tarred', False):
+            if ('tarred_audio_filepaths' in config and config['tarred_audio_filepaths'] is None) or (
+                'manifest_filepath' in config and config['manifest_filepath'] is None
+            ):
+                logging.warning(
+                    "Could not load dataset as `manifest_filepath` is None or "
+                    f"`tarred_audio_filepaths` is None. Provided config : {config}"
+                )
+                return None
+
+            if 'vad_stream' in config and config['vad_stream']:
+                logging.warning("VAD inference does not support tarred dataset now")
+                return None
+
+            shuffle_n = config.get('shuffle_n', 4 * config['batch_size']) if shuffle else 0
+            dataset = audio_to_label_dataset.get_tarred_classification_label_dataset(
+                featurizer=featurizer,
+                config=OmegaConf.to_container(config),
+                shuffle_n=shuffle_n,
+                global_rank=self.global_rank,
+                world_size=self.world_size,
+            )
+            shuffle = False
+            batch_size = config['batch_size']
+            collate_func = dataset.collate_fn
+
+        else:
+            if 'manifest_filepath' in config and config['manifest_filepath'] is None:
+                logging.warning(f"Could not load dataset as `manifest_filepath` is None. Provided config : {config}")
+                return None
+
+            if 'vad_stream' in config and config['vad_stream']:
+                logging.info("Perform streaming frame-level VAD")
+                dataset = audio_to_label_dataset.get_speech_label_dataset(
+                    featurizer=featurizer, config=OmegaConf.to_container(config)
+                )
+                batch_size = 1
+                collate_func = dataset.vad_frame_seq_collate_fn
+            else:
+                dataset = audio_to_label_dataset.get_classification_label_dataset(
+                    featurizer=featurizer, config=OmegaConf.to_container(config)
+                )
+                batch_size = config['batch_size']
+                collate_func = dataset.collate_fn
+
+        return torch.utils.data.DataLoader(
+            dataset=dataset,
+            batch_size=batch_size,
+            collate_fn=collate_func,
+            drop_last=config.get('drop_last', False),
+            shuffle=shuffle,
+            num_workers=config.get('num_workers', 0),
+            pin_memory=config.get('pin_memory', False),
+        )
+
+    @torch.no_grad()
+    def transcribe(self, paths2audio_files: List[str], batch_size: int = 4, logprobs=False) -> List[str]:
+        """
+        Generate class labels for provided audio files. Use this method for debugging and prototyping.
+
+        Args:
+            paths2audio_files: (a list) of paths to audio files. \
+                Recommended length per file is approximately 1 second.
+            batch_size: (int) batch size to use during inference. \
+                Bigger will result in better throughput performance but would use more memory.
+            logprobs: (bool) pass True to get log probabilities instead of class labels.
+
+        Returns:
+
+            A list of transcriptions (or raw log probabilities if logprobs is True) in the same order as paths2audio_files
+        """
+        if paths2audio_files is None or len(paths2audio_files) == 0:
+            return []
+        # We will store transcriptions here
+        labels = []
+        # Model's mode and device
+        mode = self.training
+        device = next(self.parameters()).device
+        dither_value = self.preprocessor.featurizer.dither
+        pad_to_value = self.preprocessor.featurizer.pad_to
+
+        try:
+            self.preprocessor.featurizer.dither = 0.0
+            self.preprocessor.featurizer.pad_to = 0
+            # Switch model to evaluation mode
+            self.eval()
+            logging_level = logging.get_verbosity()
+            logging.set_verbosity(logging.WARNING)
+            # Work in tmp directory - will store manifest file there
+            with tempfile.TemporaryDirectory() as tmpdir:
+                with open(os.path.join(tmpdir, 'manifest.json'), 'w') as fp:
+                    for audio_file in paths2audio_files:
+                        label = 0.0 if self.is_regression_task else self.cfg.labels[0]
+                        entry = {'audio_filepath': audio_file, 'duration': 100000.0, 'label': label}
+                        fp.write(json.dumps(entry) + '\n')
+
+                config = {'paths2audio_files': paths2audio_files, 'batch_size': batch_size, 'temp_dir': tmpdir}
+
+                temporary_datalayer = self._setup_transcribe_dataloader(config)
+                for test_batch in temporary_datalayer:
+                    logits = self.forward(
+                        input_signal=test_batch[0].to(device), input_signal_length=test_batch[1].to(device)
+                    )
+                    if logprobs:
+                        # dump log probs per file
+                        for idx in range(logits.shape[0]):
+                            lg = logits[idx]
+                            labels.append(lg.cpu().numpy())
+                    else:
+                        labels_k = []
+                        top_ks = self._accuracy.top_k
+                        for top_k_i in top_ks:
+                            # replace top k value with current top k
+                            self._accuracy.top_k = top_k_i
+                            labels_k_i = self._accuracy.top_k_predicted_labels(logits)
+                            labels_k.append(labels_k_i)
+
+                        # convenience: if only one top_k, pop out the nested list
+                        if len(top_ks) == 1:
+                            labels_k = labels_k[0]
+
+                        labels += labels_k
+                        # reset top k to orignal value
+                        self._accuracy.top_k = top_ks
+                    del test_batch
+        finally:
+            # set mode back to its original value
+            self.train(mode=mode)
+            self.preprocessor.featurizer.dither = dither_value
+            self.preprocessor.featurizer.pad_to = pad_to_value
+            logging.set_verbosity(logging_level)
+        return labels
+
+    def _setup_transcribe_dataloader(self, config: Dict) -> 'torch.utils.data.DataLoader':
+        """
+        Setup function for a temporary data loader which wraps the provided audio file.
+
+        Args:
+            config: A python dictionary which contains the following keys:
+
+        Returns:
+            A pytorch DataLoader for the given audio file(s).
+        """
+        dl_config = {
+            'manifest_filepath': os.path.join(config['temp_dir'], 'manifest.json'),
+            'sample_rate': self.preprocessor._sample_rate,
+            'labels': self.cfg.labels,
+            'batch_size': min(config['batch_size'], len(config['paths2audio_files'])),
+            'trim_silence': False,
+            'shuffle': False,
+        }
+
+        temporary_datalayer = self._setup_dataloader_from_config(config=DictConfig(dl_config))
+        return temporary_datalayer
+
+    @abstractmethod
+    def _update_decoder_config(self, labels, cfg):
+        pass
+
+
+class EncDecClassificationModel(_EncDecBaseModel):
+    """Encoder decoder Classification models."""
+
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+
+        if cfg.get("is_regression_task", False):
+            raise ValueError(f"EndDecClassificationModel requires the flag is_regression_task to be set as false")
+
+        super().__init__(cfg=cfg, trainer=trainer)
+
+    def _setup_preprocessor(self):
+        return EncDecClassificationModel.from_config_dict(self._cfg.preprocessor)
+
+    def _setup_encoder(self):
+        return EncDecClassificationModel.from_config_dict(self._cfg.encoder)
+
+    def _setup_decoder(self):
+        return EncDecClassificationModel.from_config_dict(self._cfg.decoder)
+
+    def _setup_loss(self):
+        return CrossEntropyLoss()
+
+    def _setup_metrics(self):
+        self._accuracy = TopKClassificationAccuracy(dist_sync_on_step=True)
+
+    @classmethod
+    def list_available_models(cls) -> Optional[List[PretrainedModelInfo]]:
+        """
+        This method returns a list of pre-trained model which can be instantiated directly from NVIDIA's NGC cloud.
+
+        Returns:
+            List of available pre-trained models.
+        """
+        results = []
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="vad_telephony_marblenet",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:vad_telephony_marblenet",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/vad_telephony_marblenet/versions/1.0.0rc1/files/vad_telephony_marblenet.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="vad_marblenet",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:vad_marblenet",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/vad_marblenet/versions/1.0.0rc1/files/vad_marblenet.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="commandrecognition_en_matchboxnet3x1x64_v1",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:commandrecognition_en_matchboxnet3x1x64_v1",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/commandrecognition_en_matchboxnet3x1x64_v1/versions/1.0.0rc1/files/commandrecognition_en_matchboxnet3x1x64_v1.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="commandrecognition_en_matchboxnet3x2x64_v1",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:commandrecognition_en_matchboxnet3x2x64_v1",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/commandrecognition_en_matchboxnet3x2x64_v1/versions/1.0.0rc1/files/commandrecognition_en_matchboxnet3x2x64_v1.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="commandrecognition_en_matchboxnet3x1x64_v2",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:commandrecognition_en_matchboxnet3x1x64_v2",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/commandrecognition_en_matchboxnet3x1x64_v2/versions/1.0.0rc1/files/commandrecognition_en_matchboxnet3x1x64_v2.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="commandrecognition_en_matchboxnet3x2x64_v2",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:commandrecognition_en_matchboxnet3x2x64_v2",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/commandrecognition_en_matchboxnet3x2x64_v2/versions/1.0.0rc1/files/commandrecognition_en_matchboxnet3x2x64_v2.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="commandrecognition_en_matchboxnet3x1x64_v2_subset_task",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:commandrecognition_en_matchboxnet3x1x64_v2_subset_task",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/commandrecognition_en_matchboxnet3x1x64_v2_subset_task/versions/1.0.0rc1/files/commandrecognition_en_matchboxnet3x1x64_v2_subset_task.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="commandrecognition_en_matchboxnet3x2x64_v2_subset_task",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:commandrecognition_en_matchboxnet3x2x64_v2_subset_task",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/commandrecognition_en_matchboxnet3x2x64_v2_subset_task/versions/1.0.0rc1/files/commandrecognition_en_matchboxnet3x2x64_v2_subset_task.nemo",
+        )
+        results.append(model)
+        return results
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {"outputs": NeuralType(('B', 'D'), LogitsType())}
+
+    # PTL-specific methods
+    def training_step(self, batch, batch_nb):
+        self.training_step_end()
+        audio_signal, audio_signal_len, labels, labels_len = batch
+        logits = self.forward(input_signal=audio_signal, input_signal_length=audio_signal_len)
+        loss_value = self.loss(logits=logits, labels=labels)
+
+        self.log('train_loss', loss_value)
+        self.log('learning_rate', self._optimizer.param_groups[0]['lr'])
+
+        self._accuracy(logits=logits, labels=labels)
+        topk_scores = self._accuracy.compute()
+        self._accuracy.reset()
+
+        for top_k, score in zip(self._accuracy.top_k, topk_scores):
+            self.log('training_batch_accuracy_top@{}'.format(top_k), score)
+
+        return {
+            'loss': loss_value,
+        }
+
+    def validation_step(self, batch, batch_idx, dataloader_idx=0):
+        audio_signal, audio_signal_len, labels, labels_len = batch
+        logits = self.forward(input_signal=audio_signal, input_signal_length=audio_signal_len)
+        loss_value = self.loss(logits=logits, labels=labels)
+        acc = self._accuracy(logits=logits, labels=labels)
+        correct_counts, total_counts = self._accuracy.correct_counts_k, self._accuracy.total_counts_k
+        return {
+            'val_loss': loss_value,
+            'val_correct_counts': correct_counts,
+            'val_total_counts': total_counts,
+            'val_acc': acc,
+        }
+
+    def test_step(self, batch, batch_idx, dataloader_idx=0):
+        audio_signal, audio_signal_len, labels, labels_len = batch
+        logits = self.forward(input_signal=audio_signal, input_signal_length=audio_signal_len)
+        loss_value = self.loss(logits=logits, labels=labels)
+        acc = self._accuracy(logits=logits, labels=labels)
+        correct_counts, total_counts = self._accuracy.correct_counts_k, self._accuracy.total_counts_k
+        return {
+            'test_loss': loss_value,
+            'test_correct_counts': correct_counts,
+            'test_total_counts': total_counts,
+            'test_acc': acc,
+        }
+
+    def multi_validation_epoch_end(self, outputs, dataloader_idx: int = 0):
+        val_loss_mean = torch.stack([x['val_loss'] for x in outputs]).mean()
+        correct_counts = torch.stack([x['val_correct_counts'] for x in outputs]).sum(axis=0)
+        total_counts = torch.stack([x['val_total_counts'] for x in outputs]).sum(axis=0)
+
+        self._accuracy.correct_counts_k = correct_counts
+        self._accuracy.total_counts_k = total_counts
+        topk_scores = self._accuracy.compute()
+        self._accuracy.reset()
+
+        tensorboard_log = {'val_loss': val_loss_mean}
+        for top_k, score in zip(self._accuracy.top_k, topk_scores):
+            tensorboard_log['val_epoch_top@{}'.format(top_k)] = score
+
+        return {'log': tensorboard_log}
+
+    def multi_test_epoch_end(self, outputs, dataloader_idx: int = 0):
+        test_loss_mean = torch.stack([x['test_loss'] for x in outputs]).mean()
+        correct_counts = torch.stack([x['test_correct_counts'].unsqueeze(0) for x in outputs]).sum(axis=0)
+        total_counts = torch.stack([x['test_total_counts'].unsqueeze(0) for x in outputs]).sum(axis=0)
+
+        self._accuracy.correct_counts_k = correct_counts
+        self._accuracy.total_counts_k = total_counts
+        topk_scores = self._accuracy.compute()
+        self._accuracy.reset()
+
+        tensorboard_log = {'test_loss': test_loss_mean}
+        for top_k, score in zip(self._accuracy.top_k, topk_scores):
+            tensorboard_log['test_epoch_top@{}'.format(top_k)] = score
+
+        return {'log': tensorboard_log}
+
+    @typecheck()
+    def forward(self, input_signal, input_signal_length):
+        logits = super().forward(input_signal=input_signal, input_signal_length=input_signal_length)
+        return logits
+
+    def change_labels(self, new_labels: List[str]):
+        """
+        Changes labels used by the decoder model. Use this method when fine-tuning on from pre-trained model.
+        This method changes only decoder and leaves encoder and pre-processing modules unchanged. For example, you would
+        use it if you want to use pretrained encoder when fine-tuning on a data in another dataset.
+
+        If new_labels == self.decoder.vocabulary then nothing will be changed.
+
+        Args:
+
+            new_labels: list with new labels. Must contain at least 2 elements. Typically, \
+            this is set of labels for the dataset.
+
+        Returns: None
+
+        """
+        if new_labels is not None and not isinstance(new_labels, ListConfig):
+            new_labels = ListConfig(new_labels)
+
+        if self._cfg.labels == new_labels:
+            logging.warning(
+                f"Old labels ({self._cfg.labels}) and new labels ({new_labels}) match. Not changing anything"
+            )
+        else:
+            if new_labels is None or len(new_labels) == 0:
+                raise ValueError(f'New labels must be non-empty list of labels. But I got: {new_labels}')
+
+            # Update config
+            self._cfg.labels = new_labels
+
+            decoder_config = self.decoder.to_config_dict()
+            new_decoder_config = copy.deepcopy(decoder_config)
+            self._update_decoder_config(new_labels, new_decoder_config)
+            del self.decoder
+            self.decoder = EncDecClassificationModel.from_config_dict(new_decoder_config)
+
+            OmegaConf.set_struct(self._cfg.decoder, False)
+            self._cfg.decoder = new_decoder_config
+            OmegaConf.set_struct(self._cfg.decoder, True)
+
+            if 'train_ds' in self._cfg and self._cfg.train_ds is not None:
+                self._cfg.train_ds.labels = new_labels
+
+            if 'validation_ds' in self._cfg and self._cfg.validation_ds is not None:
+                self._cfg.validation_ds.labels = new_labels
+
+            if 'test_ds' in self._cfg and self._cfg.test_ds is not None:
+                self._cfg.test_ds.labels = new_labels
+
+            logging.info(f"Changed decoder output to {self.decoder.num_classes} labels.")
+
+    def _update_decoder_config(self, labels, cfg):
+        """
+        Update the number of classes in the decoder based on labels provided.
+
+        Args:
+            labels: The current labels of the model
+            cfg: The config of the decoder which will be updated.
+        """
+        OmegaConf.set_struct(cfg, False)
+
+        if 'params' in cfg:
+            cfg.params.num_classes = len(labels)
+        else:
+            cfg.num_classes = len(labels)
+
+        OmegaConf.set_struct(cfg, True)
+
+
+class EncDecRegressionModel(_EncDecBaseModel):
+    """Encoder decoder class for speech regression models.
+    Model class creates training, validation methods for setting up data
+    performing model forward pass.
+    """
+
+    @classmethod
+    def list_available_models(cls) -> List[PretrainedModelInfo]:
+        """
+        This method returns a list of pre-trained model which can be instantiated directly from NVIDIA's NGC cloud.
+        Returns:
+            List of available pre-trained models.
+        """
+        result = []
+
+        return result
+
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+        if not cfg.get('is_regression_task', False):
+            raise ValueError(f"EndDecRegressionModel requires the flag is_regression_task to be set as true")
+        super().__init__(cfg=cfg, trainer=trainer)
+
+    def _setup_preprocessor(self):
+        return EncDecRegressionModel.from_config_dict(self._cfg.preprocessor)
+
+    def _setup_encoder(self):
+        return EncDecRegressionModel.from_config_dict(self._cfg.encoder)
+
+    def _setup_decoder(self):
+        return EncDecRegressionModel.from_config_dict(self._cfg.decoder)
+
+    def _setup_loss(self):
+        return MSELoss()
+
+    def _setup_metrics(self):
+        self._mse = MeanSquaredError()
+        self._mae = MeanAbsoluteError()
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {"preds": NeuralType(tuple('B'), RegressionValuesType())}
+
+    @typecheck()
+    def forward(self, input_signal, input_signal_length):
+        logits = super().forward(input_signal=input_signal, input_signal_length=input_signal_length)
+        return logits.view(-1)
+
+    # PTL-specific methods
+    def training_step(self, batch, batch_idx):
+        self.training_step_end()
+        audio_signal, audio_signal_len, targets, targets_len = batch
+        logits = self.forward(input_signal=audio_signal, input_signal_length=audio_signal_len)
+        loss = self.loss(preds=logits, labels=targets)
+        train_mse = self._mse(preds=logits, target=targets)
+        train_mae = self._mae(preds=logits, target=targets)
+
+        tensorboard_logs = {
+            'train_loss': loss,
+            'train_mse': train_mse,
+            'train_mae': train_mae,
+            'learning_rate': self._optimizer.param_groups[0]['lr'],
+        }
+
+        self.log_dict(tensorboard_logs)
+        return {'loss': loss}
+
+    def validation_step(self, batch, batch_idx, dataloader_idx: int = 0):
+        audio_signal, audio_signal_len, targets, targets_len = batch
+        logits = self.forward(input_signal=audio_signal, input_signal_length=audio_signal_len)
+        loss_value = self.loss(preds=logits, labels=targets)
+        val_mse = self._mse(preds=logits, target=targets)
+        val_mae = self._mae(preds=logits, target=targets)
+
+        return {'val_loss': loss_value, 'val_mse': val_mse, 'val_mae': val_mae}
+
+    def test_step(self, batch, batch_idx, dataloader_idx: int = 0):
+        logs = self.validation_step(batch, batch_idx, dataloader_idx)
+
+        return {'test_loss': logs['val_loss'], 'test_mse': logs['test_mse'], 'test_mae': logs['val_mae']}
+
+    def multi_validation_epoch_end(self, outputs, dataloader_idx: int = 0):
+        val_loss_mean = torch.stack([x['val_loss'] for x in outputs]).mean()
+        val_mse = self._mse.compute()
+        self._mse.reset()
+        val_mae = self._mae.compute()
+        self._mae.reset()
+
+        tensorboard_logs = {'val_loss': val_loss_mean, 'val_mse': val_mse, 'val_mae': val_mae}
+
+        return {'val_loss': val_loss_mean, 'val_mse': val_mse, 'val_mae': val_mae, 'log': tensorboard_logs}
+
+    def multi_test_epoch_end(self, outputs, dataloader_idx: int = 0):
+        test_loss_mean = torch.stack([x['test_loss'] for x in outputs]).mean()
+        test_mse = self._mse.compute()
+        self._mse.reset()
+        test_mae = self._mae.compute()
+        self._mae.reset()
+
+        tensorboard_logs = {'test_loss': test_loss_mean, 'test_mse': test_mse, 'test_mae': test_mae}
+
+        return {'test_loss': test_loss_mean, 'test_mse': test_mse, 'test_mae': test_mae, 'log': tensorboard_logs}
+
+    @torch.no_grad()
+    def transcribe(self, paths2audio_files: List[str], batch_size: int = 4) -> List[float]:
+        """
+        Generate class labels for provided audio files. Use this method for debugging and prototyping.
+
+        Args:
+            paths2audio_files: (a list) of paths to audio files. \
+                Recommended length per file is approximately 1 second.
+            batch_size: (int) batch size to use during inference. \
+                Bigger will result in better throughput performance but would use more memory.
+
+        Returns:
+
+            A list of predictions in the same order as paths2audio_files
+        """
+        predictions = super().transcribe(paths2audio_files, batch_size, logprobs=True)
+        return [float(pred) for pred in predictions]
+
+    def _update_decoder_config(self, labels, cfg):
+
+        OmegaConf.set_struct(cfg, False)
+
+        if 'params' in cfg:
+            cfg.params.num_classes = 1
+        else:
+            cfg.num_classes = 1
+
+        OmegaConf.set_struct(cfg, True)

tts/nemo/collections/asr/models/clustering_diarizer.py

@@ -0,0 +1,439 @@
+# 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 json
+import os
+import pickle as pkl
+import shutil
+import tarfile
+import tempfile
+from collections import defaultdict
+from typing import List, Optional
+
+import torch
+from omegaconf import DictConfig, OmegaConf, open_dict
+from omegaconf.listconfig import ListConfig
+from pytorch_lightning.utilities import rank_zero_only
+from tqdm import tqdm
+
+from nemo.collections.asr.models.classification_models import EncDecClassificationModel
+from nemo.collections.asr.models.label_models import ExtractSpeakerEmbeddingsModel
+from nemo.collections.asr.parts.mixins.mixins import DiarizationMixin
+from nemo.collections.asr.parts.utils.speaker_utils import audio_rttm_map, perform_diarization, write_rttm2manifest
+from nemo.collections.asr.parts.utils.vad_utils import (
+    generate_overlap_vad_seq,
+    generate_vad_segment_table,
+    get_vad_stream_status,
+    prepare_manifest,
+)
+from nemo.core.classes import Model
+from nemo.utils import logging, model_utils
+from nemo.utils.exp_manager import NotFoundError
+
+try:
+    from torch.cuda.amp import autocast
+except ImportError:
+    from contextlib import contextmanager
+
+    @contextmanager
+    def autocast(enabled=None):
+        yield
+
+
+__all__ = ['ClusteringDiarizer']
+
+_MODEL_CONFIG_YAML = "model_config.yaml"
+_VAD_MODEL = "vad_model.nemo"
+_SPEAKER_MODEL = "speaker_model.nemo"
+
+
+def get_available_model_names(class_name):
+    available_models = class_name.list_available_models()
+    return list(map(lambda x: x.pretrained_model_name, available_models))
+
+
+class ClusteringDiarizer(Model, DiarizationMixin):
+    def __init__(self, cfg: DictConfig):
+        cfg = model_utils.convert_model_config_to_dict_config(cfg)
+        # Convert config to support Hydra 1.0+ instantiation
+        cfg = model_utils.maybe_update_config_version(cfg)
+        self._cfg = cfg
+        self._out_dir = self._cfg.diarizer.out_dir
+        if not os.path.exists(self._out_dir):
+            os.mkdir(self._out_dir)
+
+        # init vad model
+        self.has_vad_model = False
+        self.has_vad_model_to_save = False
+
+        self._speaker_manifest_path = self._cfg.diarizer.speaker_embeddings.oracle_vad_manifest
+        self.AUDIO_RTTM_MAP = None
+        self.paths2audio_files = self._cfg.diarizer.paths2audio_files
+
+        if self._cfg.diarizer.vad.model_path is not None:
+            self._init_vad_model()
+            self._vad_dir = os.path.join(self._out_dir, 'vad_outputs')
+            self._vad_out_file = os.path.join(self._vad_dir, "vad_out.json")
+            shutil.rmtree(self._vad_dir, ignore_errors=True)
+            os.makedirs(self._vad_dir)
+
+        # init speaker model
+        self._init_speaker_model()
+
+        if self._cfg.diarizer.get('num_speakers', None):
+            self._num_speakers = self._cfg.diarizer.num_speakers
+            logging.warning("in next release num_speakers will be changed to oracle_num_speakers")
+        else:
+            self._num_speakers = self._cfg.diarizer.oracle_num_speakers
+
+        if self._cfg.diarizer.get('max_num_speakers', None):
+            self.max_num_speakers = self._cfg.diarizer.max_num_speakers
+        else:
+            self.max_num_speakers = 8
+
+        self._device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    @classmethod
+    def list_available_models(cls):
+        pass
+
+    def _init_speaker_model(self):
+        model_path = self._cfg.diarizer.speaker_embeddings.model_path
+        if model_path is not None and model_path.endswith('.nemo'):
+            self._speaker_model = ExtractSpeakerEmbeddingsModel.restore_from(model_path)
+            logging.info("Speaker Model restored locally from {}".format(model_path))
+        else:
+            if model_path not in get_available_model_names(ExtractSpeakerEmbeddingsModel):
+                logging.warning(
+                    "requested {} model name not available in pretrained models, instead".format(model_path)
+                )
+                model_path = "speakerdiarization_speakernet"
+            logging.info("Loading pretrained {} model from NGC".format(model_path))
+            self._speaker_model = ExtractSpeakerEmbeddingsModel.from_pretrained(model_name=model_path)
+
+        self._speaker_dir = os.path.join(self._out_dir, 'speaker_outputs')
+
+    def set_vad_model(self, vad_config):
+        with open_dict(self._cfg):
+            self._cfg.diarizer.vad = vad_config
+        self._init_vad_model()
+
+    def _init_vad_model(self):
+        model_path = self._cfg.diarizer.vad.model_path
+        if model_path.endswith('.nemo'):
+            self._vad_model = EncDecClassificationModel.restore_from(model_path)
+            logging.info("VAD model loaded locally from {}".format(model_path))
+        else:
+            if model_path not in get_available_model_names(EncDecClassificationModel):
+                logging.warning(
+                    "requested {} model name not available in pretrained models, instead".format(model_path)
+                )
+                model_path = "vad_telephony_marblenet"
+            logging.info("Loading pretrained {} model from NGC".format(model_path))
+            self._vad_model = EncDecClassificationModel.from_pretrained(model_name=model_path)
+
+        self._vad_window_length_in_sec = self._cfg.diarizer.vad.window_length_in_sec
+        self._vad_shift_length_in_sec = self._cfg.diarizer.vad.shift_length_in_sec
+        self.has_vad_model_to_save = True
+        self.has_vad_model = True
+
+    def _setup_vad_test_data(self, manifest_vad_input):
+        vad_dl_config = {
+            'manifest_filepath': manifest_vad_input,
+            'sample_rate': self._cfg.sample_rate,
+            'vad_stream': True,
+            'labels': ['infer',],
+            'time_length': self._vad_window_length_in_sec,
+            'shift_length': self._vad_shift_length_in_sec,
+            'trim_silence': False,
+            'num_workers': self._cfg.num_workers,
+        }
+        self._vad_model.setup_test_data(test_data_config=vad_dl_config)
+
+    def _setup_spkr_test_data(self, manifest_file):
+        spk_dl_config = {
+            'manifest_filepath': manifest_file,
+            'sample_rate': self._cfg.sample_rate,
+            'batch_size': 1,
+            'time_length': self._cfg.diarizer.speaker_embeddings.window_length_in_sec,
+            'shift_length': self._cfg.diarizer.speaker_embeddings.shift_length_in_sec,
+            'trim_silence': False,
+            'embedding_dir': self._speaker_dir,
+            'labels': None,
+            'task': "diarization",
+            'num_workers': self._cfg.num_workers,
+        }
+        self._speaker_model.setup_test_data(spk_dl_config)
+
+    def _run_vad(self, manifest_file):
+        self._vad_model = self._vad_model.to(self._device)
+        self._vad_model.eval()
+
+        time_unit = int(self._vad_window_length_in_sec / self._vad_shift_length_in_sec)
+        trunc = int(time_unit / 2)
+        trunc_l = time_unit - trunc
+        all_len = 0
+        data = []
+        for line in open(manifest_file, 'r'):
+            file = os.path.basename(json.loads(line)['audio_filepath'])
+            data.append(os.path.splitext(file)[0])
+
+        status = get_vad_stream_status(data)
+        for i, test_batch in enumerate(tqdm(self._vad_model.test_dataloader())):
+            test_batch = [x.to(self._device) for x in test_batch]
+            with autocast():
+                log_probs = self._vad_model(input_signal=test_batch[0], input_signal_length=test_batch[1])
+                probs = torch.softmax(log_probs, dim=-1)
+                pred = probs[:, 1]
+                if status[i] == 'start':
+                    to_save = pred[:-trunc]
+                elif status[i] == 'next':
+                    to_save = pred[trunc:-trunc_l]
+                elif status[i] == 'end':
+                    to_save = pred[trunc_l:]
+                else:
+                    to_save = pred
+                all_len += len(to_save)
+                outpath = os.path.join(self._vad_dir, data[i] + ".frame")
+                with open(outpath, "a") as fout:
+                    for f in range(len(to_save)):
+                        fout.write('{0:0.4f}\n'.format(to_save[f]))
+            del test_batch
+            if status[i] == 'end' or status[i] == 'single':
+                all_len = 0
+
+        if not self._cfg.diarizer.vad.vad_decision_smoothing:
+            # Shift the window by 10ms to generate the frame and use the prediction of the window to represent the label for the frame;
+            self.vad_pred_dir = self._vad_dir
+
+        else:
+            # Generate predictions with overlapping input segments. Then a smoothing filter is applied to decide the label for a frame spanned by multiple segments.
+            # smoothing_method would be either in majority vote (median) or average (mean)
+            logging.info("Generating predictions with overlapping input segments")
+            smoothing_pred_dir = generate_overlap_vad_seq(
+                frame_pred_dir=self._vad_dir,
+                smoothing_method=self._cfg.diarizer.vad.smoothing_params.method,
+                overlap=self._cfg.diarizer.vad.smoothing_params.overlap,
+                seg_len=self._vad_window_length_in_sec,
+                shift_len=self._vad_shift_length_in_sec,
+                num_workers=self._cfg.num_workers,
+            )
+            self.vad_pred_dir = smoothing_pred_dir
+
+        logging.info("Converting frame level prediction to speech/no-speech segment in start and end times format.")
+        table_out_dir = generate_vad_segment_table(
+            vad_pred_dir=self.vad_pred_dir,
+            threshold=self._cfg.diarizer.vad.threshold,
+            shift_len=self._vad_shift_length_in_sec,
+            num_workers=self._cfg.num_workers,
+        )
+
+        vad_table_list = [os.path.join(table_out_dir, key + ".txt") for key in self.AUDIO_RTTM_MAP]
+        write_rttm2manifest(self._cfg.diarizer.paths2audio_files, vad_table_list, self._vad_out_file)
+        self._speaker_manifest_path = self._vad_out_file
+
+    def _extract_embeddings(self, manifest_file):
+        logging.info("Extracting embeddings for Diarization")
+        self._setup_spkr_test_data(manifest_file)
+        uniq_names = []
+        out_embeddings = defaultdict(list)
+        self._speaker_model = self._speaker_model.to(self._device)
+        self._speaker_model.eval()
+        with open(manifest_file, 'r') as manifest:
+            for line in manifest.readlines():
+                line = line.strip()
+                dic = json.loads(line)
+                uniq_names.append(dic['audio_filepath'].split('/')[-1].rsplit('.', 1)[0])
+
+        for i, test_batch in enumerate(tqdm(self._speaker_model.test_dataloader())):
+            test_batch = [x.to(self._device) for x in test_batch]
+            audio_signal, audio_signal_len, labels, slices = test_batch
+            with autocast():
+                _, embs = self._speaker_model.forward(input_signal=audio_signal, input_signal_length=audio_signal_len)
+                emb_shape = embs.shape[-1]
+                embs = embs.view(-1, emb_shape).type(torch.float32)
+                embs = embs.cpu().detach().numpy()
+                out_embeddings[uniq_names[i]].extend(embs)
+            del test_batch
+
+        embedding_dir = os.path.join(self._speaker_dir, 'embeddings')
+        if not os.path.exists(embedding_dir):
+            os.makedirs(embedding_dir, exist_ok=True)
+
+        prefix = manifest_file.split('/')[-1].rsplit('.', 1)[-2]
+
+        name = os.path.join(embedding_dir, prefix)
+        self._embeddings_file = name + '_embeddings.pkl'
+        pkl.dump(out_embeddings, open(self._embeddings_file, 'wb'))
+        logging.info("Saved embedding files to {}".format(embedding_dir))
+
+    def path2audio_files_to_manifest(self, paths2audio_files):
+        mfst_file = os.path.join(self._out_dir, 'manifest.json')
+        with open(mfst_file, 'w') as fp:
+            for audio_file in paths2audio_files:
+                audio_file = audio_file.strip()
+                entry = {'audio_filepath': audio_file, 'offset': 0.0, 'duration': None, 'text': '-', 'label': 'infer'}
+                fp.write(json.dumps(entry) + '\n')
+
+        return mfst_file
+
+    def diarize(self, paths2audio_files: List[str] = None, batch_size: int = 1):
+        """
+        """
+
+        if paths2audio_files:
+            self.paths2audio_files = paths2audio_files
+        else:
+            if self._cfg.diarizer.paths2audio_files is None:
+                raise ValueError("Pass path2audio files either through config or to diarize method")
+            else:
+                self.paths2audio_files = self._cfg.diarizer.paths2audio_files
+
+        if type(self.paths2audio_files) is str and os.path.isfile(self.paths2audio_files):
+            paths2audio_files = []
+            with open(self.paths2audio_files, 'r') as path2file:
+                for audiofile in path2file.readlines():
+                    audiofile = audiofile.strip()
+                    paths2audio_files.append(audiofile)
+
+        elif type(self.paths2audio_files) in [list, ListConfig]:
+            paths2audio_files = list(self.paths2audio_files)
+
+        else:
+            raise ValueError("paths2audio_files must be of type list or path to file containing audio files")
+
+        self.AUDIO_RTTM_MAP = audio_rttm_map(paths2audio_files, self._cfg.diarizer.path2groundtruth_rttm_files)
+
+        if self.has_vad_model:
+            logging.info("Performing VAD")
+            mfst_file = self.path2audio_files_to_manifest(paths2audio_files)
+            self._dont_auto_split = False
+            self._split_duration = 50
+            manifest_vad_input = mfst_file
+
+            if not self._dont_auto_split:
+                logging.info("Split long audio file to avoid CUDA memory issue")
+                logging.debug("Try smaller split_duration if you still have CUDA memory issue")
+                config = {
+                    'manifest_filepath': mfst_file,
+                    'time_length': self._vad_window_length_in_sec,
+                    'split_duration': self._split_duration,
+                    'num_workers': self._cfg.num_workers,
+                }
+                manifest_vad_input = prepare_manifest(config)
+            else:
+                logging.warning(
+                    "If you encounter CUDA memory issue, try splitting manifest entry by split_duration to avoid it."
+                )
+
+            self._setup_vad_test_data(manifest_vad_input)
+            self._run_vad(manifest_vad_input)
+        else:
+            if not os.path.exists(self._speaker_manifest_path):
+                raise NotFoundError("Oracle VAD based manifest file not found")
+
+        self._extract_embeddings(self._speaker_manifest_path)
+        out_rttm_dir = os.path.join(self._out_dir, 'pred_rttms')
+        os.makedirs(out_rttm_dir, exist_ok=True)
+
+        perform_diarization(
+            embeddings_file=self._embeddings_file,
+            reco2num=self._num_speakers,
+            manifest_path=self._speaker_manifest_path,
+            sample_rate=self._cfg.sample_rate,
+            window=self._cfg.diarizer.speaker_embeddings.window_length_in_sec,
+            shift=self._cfg.diarizer.speaker_embeddings.shift_length_in_sec,
+            audio_rttm_map=self.AUDIO_RTTM_MAP,
+            out_rttm_dir=out_rttm_dir,
+            max_num_speakers=self.max_num_speakers,
+        )
+
+    @staticmethod
+    def __make_nemo_file_from_folder(filename, source_dir):
+        with tarfile.open(filename, "w:gz") as tar:
+            tar.add(source_dir, arcname="./")
+
+    @rank_zero_only
+    def save_to(self, save_path: str):
+        """
+        Saves model instance (weights and configuration) into EFF archive or .
+         You can use "restore_from" method to fully restore instance from .nemo file.
+
+        .nemo file is an archive (tar.gz) with the following:
+            model_config.yaml - model configuration in .yaml format. You can deserialize this into cfg argument for model's constructor
+            model_wights.chpt - model checkpoint
+
+        Args:
+            save_path: Path to .nemo file where model instance should be saved
+        """
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            config_yaml = os.path.join(tmpdir, _MODEL_CONFIG_YAML)
+            spkr_model = os.path.join(tmpdir, _SPEAKER_MODEL)
+
+            self.to_config_file(path2yaml_file=config_yaml)
+            if self.has_vad_model:
+                vad_model = os.path.join(tmpdir, _VAD_MODEL)
+                self._vad_model.save_to(vad_model)
+            self._speaker_model.save_to(spkr_model)
+            self.__make_nemo_file_from_folder(filename=save_path, source_dir=tmpdir)
+
+    @staticmethod
+    def __unpack_nemo_file(path2file: str, out_folder: str) -> str:
+        if not os.path.exists(path2file):
+            raise FileNotFoundError(f"{path2file} does not exist")
+        tar = tarfile.open(path2file, "r:gz")
+        tar.extractall(path=out_folder)
+        tar.close()
+        return out_folder
+
+    @classmethod
+    def restore_from(
+        cls,
+        restore_path: str,
+        override_config_path: Optional[str] = None,
+        map_location: Optional[torch.device] = None,
+        strict: bool = False,
+    ):
+        # Get path where the command is executed - the artifacts will be "retrieved" there
+        # (original .nemo behavior)
+        cwd = os.getcwd()
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            try:
+                cls.__unpack_nemo_file(path2file=restore_path, out_folder=tmpdir)
+                os.chdir(tmpdir)
+                if override_config_path is None:
+                    config_yaml = os.path.join(tmpdir, _MODEL_CONFIG_YAML)
+                else:
+                    config_yaml = override_config_path
+                conf = OmegaConf.load(config_yaml)
+                if os.path.exists(os.path.join(tmpdir, _VAD_MODEL)):
+                    conf.diarizer.vad.model_path = os.path.join(tmpdir, _VAD_MODEL)
+                else:
+                    logging.info(
+                        f'Model {cls.__name__} does not contain a VAD model. A VAD model or manifest file with'
+                        f'speech segments need for diarization with this model'
+                    )
+
+                conf.diarizer.speaker_embeddings.model_path = os.path.join(tmpdir, _SPEAKER_MODEL)
+                conf.restore_map_location = map_location
+                OmegaConf.set_struct(conf, True)
+                instance = cls(cfg=conf)
+
+                logging.info(f'Model {cls.__name__} was successfully restored from {restore_path}.')
+            finally:
+                os.chdir(cwd)
+
+        return instance

tts/nemo/collections/asr/models/configs/__init__.py

@@ -0,0 +1,46 @@
+# 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.
+
+from nemo.collections.asr.models.configs.classification_models_config import (
+    EncDecClassificationConfig,
+    EncDecClassificationDatasetConfig,
+    EncDecClassificationModelConfig,
+)
+from nemo.collections.asr.models.configs.ctc_models_config import (
+    EncDecCTCConfig,
+    EncDecCTCDatasetConfig,
+    EncDecCTCModelConfig,
+)
+from nemo.collections.asr.models.configs.matchboxnet_config import (
+    EncDecClassificationModelConfigBuilder,
+    MatchboxNetModelConfig,
+    MatchboxNetVADModelConfig,
+)
+from nemo.collections.asr.models.configs.quartznet_config import (
+    EncDecCTCModelConfigBuilder,
+    JasperModelConfig,
+    QuartzNetModelConfig,
+)
+from nemo.collections.asr.modules.audio_preprocessing import (
+    AudioToMelSpectrogramPreprocessorConfig,
+    AudioToMFCCPreprocessorConfig,
+    CropOrPadSpectrogramAugmentationConfig,
+    SpectrogramAugmentationConfig,
+)
+from nemo.collections.asr.modules.conv_asr import (
+    ConvASRDecoderClassificationConfig,
+    ConvASRDecoderConfig,
+    ConvASREncoderConfig,
+    JasperEncoderConfig,
+)

tts/nemo/collections/asr/models/configs/classification_models_config.py

@@ -0,0 +1,95 @@
+# 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.
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+from omegaconf import MISSING
+
+import nemo.core.classes.dataset
+from nemo.collections.asr.modules.audio_preprocessing import (
+    AudioToMFCCPreprocessorConfig,
+    CropOrPadSpectrogramAugmentationConfig,
+    SpectrogramAugmentationConfig,
+)
+from nemo.collections.asr.modules.conv_asr import ConvASRDecoderClassificationConfig, ConvASREncoderConfig
+from nemo.core.config import modelPT as model_cfg
+
+
+@dataclass
+class EncDecClassificationDatasetConfig(nemo.core.classes.dataset.DatasetConfig):
+    manifest_filepath: Optional[str] = None
+    sample_rate: int = MISSING
+    labels: List[str] = MISSING
+    trim_silence: bool = False
+
+    # Tarred dataset support
+    is_tarred: bool = False
+    tarred_audio_filepaths: Optional[str] = None
+    tarred_shard_strategy: str = "scatter"
+    shuffle_n: int = 0
+
+    # Optional
+    int_values: Optional[int] = None
+    augmentor: Optional[Dict[str, Any]] = None
+    max_duration: Optional[float] = None
+    min_duration: Optional[float] = None
+
+    # VAD Optional
+    vad_stream: Optional[bool] = None
+    time_length: float = 0.31
+    shift_length: float = 0.01
+    normalize_audio: bool = False
+    is_regression_task: bool = False
+
+
+@dataclass
+class EncDecClassificationConfig(model_cfg.ModelConfig):
+    # Model global arguments
+    sample_rate: int = 16000
+    repeat: int = 1
+    dropout: float = 0.0
+    separable: bool = True
+    kernel_size_factor: float = 1.0
+    labels: List[str] = MISSING
+    timesteps: int = MISSING
+
+    # Dataset configs
+    train_ds: EncDecClassificationDatasetConfig = EncDecClassificationDatasetConfig(
+        manifest_filepath=None, shuffle=True, trim_silence=False
+    )
+    validation_ds: EncDecClassificationDatasetConfig = EncDecClassificationDatasetConfig(
+        manifest_filepath=None, shuffle=False
+    )
+    test_ds: EncDecClassificationDatasetConfig = EncDecClassificationDatasetConfig(
+        manifest_filepath=None, shuffle=False
+    )
+
+    # Optimizer / Scheduler config
+    optim: Optional[model_cfg.OptimConfig] = model_cfg.OptimConfig(sched=model_cfg.SchedConfig())
+
+    # Model component configs
+    preprocessor: AudioToMFCCPreprocessorConfig = AudioToMFCCPreprocessorConfig()
+    spec_augment: Optional[SpectrogramAugmentationConfig] = SpectrogramAugmentationConfig()
+    crop_or_pad_augment: Optional[CropOrPadSpectrogramAugmentationConfig] = CropOrPadSpectrogramAugmentationConfig(
+        audio_length=timesteps
+    )
+
+    encoder: ConvASREncoderConfig = ConvASREncoderConfig()
+    decoder: ConvASRDecoderClassificationConfig = ConvASRDecoderClassificationConfig()
+
+
+@dataclass
+class EncDecClassificationModelConfig(model_cfg.NemoConfig):
+    model: EncDecClassificationConfig = EncDecClassificationConfig()

tts/nemo/collections/asr/models/configs/ctc_models_config.py

@@ -0,0 +1,85 @@
+# 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.
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+from omegaconf import MISSING
+
+import nemo.core.classes.dataset
+from nemo.collections.asr.modules.audio_preprocessing import (
+    AudioToMelSpectrogramPreprocessorConfig,
+    SpectrogramAugmentationConfig,
+)
+from nemo.collections.asr.modules.conv_asr import ConvASRDecoderConfig, ConvASREncoderConfig
+from nemo.core.config import modelPT as model_cfg
+
+
+@dataclass
+class EncDecCTCDatasetConfig(nemo.core.classes.dataset.DatasetConfig):
+    manifest_filepath: Optional[str] = None
+    sample_rate: int = MISSING
+    labels: List[str] = MISSING
+    trim_silence: bool = False
+
+    # Tarred dataset support
+    is_tarred: bool = False
+    tarred_audio_filepaths: Optional[str] = None
+    tarred_shard_strategy: str = "scatter"
+    shuffle_n: int = 0
+
+    # Optional
+    int_values: Optional[int] = None
+    augmentor: Optional[Dict[str, Any]] = None
+    max_duration: Optional[float] = None
+    min_duration: Optional[float] = None
+    max_utts: int = 0
+    blank_index: int = -1
+    unk_index: int = -1
+    normalize: bool = False
+    trim: bool = True
+    parser: Optional[str] = 'en'
+    eos_id: Optional[int] = None
+    bos_id: Optional[int] = None
+    pad_id: int = 0
+    use_start_end_token: bool = False
+
+
+@dataclass
+class EncDecCTCConfig(model_cfg.ModelConfig):
+    # Model global arguments
+    sample_rate: int = 16000
+    repeat: int = 1
+    dropout: float = 0.0
+    separable: bool = False
+    labels: List[str] = MISSING
+
+    # Dataset configs
+    train_ds: EncDecCTCDatasetConfig = EncDecCTCDatasetConfig(manifest_filepath=None, shuffle=True)
+    validation_ds: EncDecCTCDatasetConfig = EncDecCTCDatasetConfig(manifest_filepath=None, shuffle=False)
+    test_ds: EncDecCTCDatasetConfig = EncDecCTCDatasetConfig(manifest_filepath=None, shuffle=False)
+
+    # Optimizer / Scheduler config
+    optim: Optional[model_cfg.OptimConfig] = model_cfg.OptimConfig(sched=model_cfg.SchedConfig())
+
+    # Model component configs
+    preprocessor: AudioToMelSpectrogramPreprocessorConfig = AudioToMelSpectrogramPreprocessorConfig()
+    spec_augment: Optional[SpectrogramAugmentationConfig] = SpectrogramAugmentationConfig()
+    encoder: ConvASREncoderConfig = ConvASREncoderConfig()
+    decoder: ConvASRDecoderConfig = ConvASRDecoderConfig()
+
+
+@dataclass
+class EncDecCTCModelConfig(model_cfg.NemoConfig):
+    model: EncDecCTCConfig = EncDecCTCConfig()

tts/nemo/collections/asr/models/configs/matchboxnet_config.py

@@ -0,0 +1,253 @@
+# 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.
+
+from dataclasses import dataclass, field
+from typing import Any, Callable, List, Optional
+
+from omegaconf import MISSING
+
+from nemo.collections.asr.models.configs import classification_models_config as clf_cfg
+from nemo.collections.asr.modules.audio_preprocessing import (
+    AudioToMFCCPreprocessorConfig,
+    CropOrPadSpectrogramAugmentationConfig,
+    SpectrogramAugmentationConfig,
+)
+from nemo.collections.asr.modules.conv_asr import (
+    ConvASRDecoderClassificationConfig,
+    ConvASREncoderConfig,
+    JasperEncoderConfig,
+)
+from nemo.core.config import modelPT as model_cfg
+
+
+# fmt: off
+def matchboxnet_3x1x64():
+    config = [
+        JasperEncoderConfig(filters=128, repeat=1, kernel=[11], stride=[1], dilation=[1], dropout=0.0,
+                            residual=False, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=64, repeat=1, kernel=[13], stride=[1], dilation=[1], dropout=0.0,
+                            residual=True, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=64, repeat=1, kernel=[15], stride=[1], dilation=[1], dropout=0.0,
+                            residual=True, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=64, repeat=1, kernel=[17], stride=[1], dilation=[1], dropout=0.0,
+                            residual=True, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=128, repeat=1, kernel=[29], stride=[1], dilation=[2], dropout=0.0,
+                            residual=False, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=128, repeat=1, kernel=[1], stride=[1], dilation=[1], dropout=0.0,
+                            residual=False, groups=1, separable=False, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False)
+    ]
+    return config
+
+
+def matchboxnet_3x1x64_vad():
+    config = [
+        JasperEncoderConfig(filters=128, repeat=1, kernel=[11], stride=[1], dilation=[1], dropout=0.0,
+                            residual=False, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=64, repeat=1, kernel=[13], stride=[1], dilation=[1], dropout=0.0,
+                            residual=True, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=64, repeat=1, kernel=[15], stride=[1], dilation=[1], dropout=0.0,
+                            residual=True, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=64, repeat=1, kernel=[17], stride=[1], dilation=[1], dropout=0.0,
+                            residual=True, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=128, repeat=1, kernel=[29], stride=[1], dilation=[2], dropout=0.0,
+                            residual=False, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=128, repeat=1, kernel=[1], stride=[1], dilation=[1], dropout=0.0,
+                            residual=False, groups=1, separable=False, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False)
+    ]
+    return config
+
+
+# fmt: on
+
+
+@dataclass
+class MatchboxNetModelConfig(clf_cfg.EncDecClassificationConfig):
+    # Model global arguments
+    sample_rate: int = 16000
+    repeat: int = 1
+    dropout: float = 0.0
+    separable: bool = True
+    kernel_size_factor: float = 1.0
+    timesteps: int = 128
+    labels: List[str] = MISSING
+
+    # Dataset configs
+    train_ds: clf_cfg.EncDecClassificationDatasetConfig = clf_cfg.EncDecClassificationDatasetConfig(
+        manifest_filepath=None, shuffle=True, trim_silence=False
+    )
+    validation_ds: clf_cfg.EncDecClassificationDatasetConfig = clf_cfg.EncDecClassificationDatasetConfig(
+        manifest_filepath=None, shuffle=False
+    )
+    test_ds: clf_cfg.EncDecClassificationDatasetConfig = clf_cfg.EncDecClassificationDatasetConfig(
+        manifest_filepath=None, shuffle=False
+    )
+
+    # Optimizer / Scheduler config
+    optim: Optional[model_cfg.OptimConfig] = model_cfg.OptimConfig(sched=model_cfg.SchedConfig())
+
+    # Model general component configs
+    preprocessor: AudioToMFCCPreprocessorConfig = AudioToMFCCPreprocessorConfig(window_size=0.025)
+    spec_augment: Optional[SpectrogramAugmentationConfig] = SpectrogramAugmentationConfig(
+        freq_masks=2, time_masks=2, freq_width=15, time_width=25, rect_masks=5, rect_time=25, rect_freq=15
+    )
+    crop_or_pad_augment: Optional[CropOrPadSpectrogramAugmentationConfig] = CropOrPadSpectrogramAugmentationConfig(
+        audio_length=128
+    )
+
+    encoder: ConvASREncoderConfig = ConvASREncoderConfig(activation="relu")
+    decoder: ConvASRDecoderClassificationConfig = ConvASRDecoderClassificationConfig()
+
+
+@dataclass
+class MatchboxNetVADModelConfig(MatchboxNetModelConfig):
+    timesteps: int = 64
+    labels: List[str] = field(default_factory=lambda: ['background', 'speech'])
+
+    crop_or_pad_augment: Optional[CropOrPadSpectrogramAugmentationConfig] = None
+
+
+class EncDecClassificationModelConfigBuilder(model_cfg.ModelConfigBuilder):
+    VALID_CONFIGS = ['matchboxnet_3x1x64', 'matchboxnet_3x1x64_vad']
+
+    def __init__(self, name: str = 'matchboxnet_3x1x64', encoder_cfg_func: Optional[Callable[[], List[Any]]] = None):
+        if name not in EncDecClassificationModelConfigBuilder.VALID_CONFIGS:
+            raise ValueError("`name` must be one of : \n" f"{EncDecClassificationModelConfigBuilder.VALID_CONFIGS}")
+
+        self.name = name
+
+        if 'matchboxnet_3x1x64_vad' in name:
+            if encoder_cfg_func is None:
+                encoder_cfg_func = matchboxnet_3x1x64_vad
+
+            model_cfg = MatchboxNetVADModelConfig(
+                repeat=1,
+                separable=True,
+                encoder=ConvASREncoderConfig(jasper=encoder_cfg_func(), activation="relu"),
+                decoder=ConvASRDecoderClassificationConfig(),
+            )
+
+        elif 'matchboxnet_3x1x64' in name:
+            if encoder_cfg_func is None:
+                encoder_cfg_func = matchboxnet_3x1x64
+
+            model_cfg = MatchboxNetModelConfig(
+                repeat=1,
+                separable=False,
+                spec_augment=SpectrogramAugmentationConfig(rect_masks=5, rect_freq=50, rect_time=120),
+                encoder=ConvASREncoderConfig(jasper=encoder_cfg_func(), activation="relu"),
+                decoder=ConvASRDecoderClassificationConfig(),
+            )
+
+        else:
+            raise ValueError(f"Invalid config name submitted to {self.__class__.__name__}")
+
+        super(EncDecClassificationModelConfigBuilder, self).__init__(model_cfg)
+        self.model_cfg: clf_cfg.EncDecClassificationConfig = model_cfg  # enable type hinting
+
+    def set_labels(self, labels: List[str]):
+        self.model_cfg.labels = labels
+
+    def set_separable(self, separable: bool):
+        self.model_cfg.separable = separable
+
+    def set_repeat(self, repeat: int):
+        self.model_cfg.repeat = repeat
+
+    def set_sample_rate(self, sample_rate: int):
+        self.model_cfg.sample_rate = sample_rate
+
+    def set_dropout(self, dropout: float = 0.0):
+        self.model_cfg.dropout = dropout
+
+    def set_timesteps(self, timesteps: int):
+        self.model_cfg.timesteps = timesteps
+
+    def set_is_regression_task(self, is_regression_task: bool):
+        self.model_cfg.is_regression_task = is_regression_task
+
+    # Note: Autocomplete for users wont work without these overrides
+    # But practically it is not needed since python will infer at runtime
+
+    # def set_train_ds(self, cfg: Optional[clf_cfg.EncDecClassificationDatasetConfig] = None):
+    #     super().set_train_ds(cfg)
+    #
+    # def set_validation_ds(self, cfg: Optional[clf_cfg.EncDecClassificationDatasetConfig] = None):
+    #     super().set_validation_ds(cfg)
+    #
+    # def set_test_ds(self, cfg: Optional[clf_cfg.EncDecClassificationDatasetConfig] = None):
+    #     super().set_test_ds(cfg)
+
+    def _finalize_cfg(self):
+        # propagate labels
+        self.model_cfg.train_ds.labels = self.model_cfg.labels
+        self.model_cfg.validation_ds.labels = self.model_cfg.labels
+        self.model_cfg.test_ds.labels = self.model_cfg.labels
+        self.model_cfg.decoder.vocabulary = self.model_cfg.labels
+
+        # propagate num classes
+        self.model_cfg.decoder.num_classes = len(self.model_cfg.labels)
+
+        # propagate sample rate
+        self.model_cfg.sample_rate = self.model_cfg.sample_rate
+        self.model_cfg.preprocessor.sample_rate = self.model_cfg.sample_rate
+        self.model_cfg.train_ds.sample_rate = self.model_cfg.sample_rate
+        self.model_cfg.validation_ds.sample_rate = self.model_cfg.sample_rate
+        self.model_cfg.test_ds.sample_rate = self.model_cfg.sample_rate
+
+        # propagate filters
+        self.model_cfg.encoder.feat_in = self.model_cfg.preprocessor.features
+        self.model_cfg.decoder.feat_in = self.model_cfg.encoder.jasper[-1].filters
+
+        # propagate timeteps
+        if self.model_cfg.crop_or_pad_augment is not None:
+            self.model_cfg.crop_or_pad_augment.audio_length = self.model_cfg.timesteps
+
+        # propagate separable
+        for layer in self.model_cfg.encoder.jasper[:-1]:  # type: JasperEncoderConfig
+            layer.separable = self.model_cfg.separable
+
+        # propagate repeat
+        for layer in self.model_cfg.encoder.jasper[1:-2]:  # type: JasperEncoderConfig
+            layer.repeat = self.model_cfg.repeat
+
+        # propagate dropout
+        for layer in self.model_cfg.encoder.jasper:  # type: JasperEncoderConfig
+            layer.dropout = self.model_cfg.dropout
+
+    def build(self) -> clf_cfg.EncDecClassificationConfig:
+        return super().build()

tts/nemo/collections/asr/models/configs/quartznet_config.py

@@ -0,0 +1,308 @@
+# 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.
+
+from dataclasses import dataclass
+from typing import Any, Callable, List, Optional
+
+from omegaconf import MISSING
+
+from nemo.collections.asr.models.configs import ctc_models_config as ctc_cfg
+from nemo.collections.asr.modules.audio_preprocessing import (
+    AudioToMelSpectrogramPreprocessorConfig,
+    SpectrogramAugmentationConfig,
+)
+from nemo.collections.asr.modules.conv_asr import ConvASRDecoderConfig, ConvASREncoderConfig, JasperEncoderConfig
+from nemo.core.config import modelPT as model_cfg
+
+
+# fmt: off
+def qn_15x5():
+    config = [
+        JasperEncoderConfig(filters=256, repeat=1, kernel=[33], stride=[2], dilation=[1], dropout=0.0,
+                            residual=False, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=256, repeat=5, kernel=[33], stride=[1], dilation=[1], dropout=0.0,
+                            residual=True, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=256, repeat=5, kernel=[33], stride=[1], dilation=[1], dropout=0.0,
+                            residual=True, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=256, repeat=5, kernel=[33], stride=[1], dilation=[1], dropout=0.0,
+                            residual=True, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=256, repeat=5, kernel=[39], stride=[1], dilation=[1], dropout=0.0,
+                            residual=True, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=256, repeat=5, kernel=[39], stride=[1], dilation=[1], dropout=0.0,
+                            residual=True, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=256, repeat=5, kernel=[39], stride=[1], dilation=[1], dropout=0.0,
+                            residual=True, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=512, repeat=5, kernel=[51], stride=[1], dilation=[1], dropout=0.0,
+                            residual=True, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=512, repeat=5, kernel=[51], stride=[1], dilation=[1], dropout=0.0,
+                            residual=True, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=512, repeat=5, kernel=[51], stride=[1], dilation=[1], dropout=0.0,
+                            residual=True, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=512, repeat=5, kernel=[63], stride=[1], dilation=[1], dropout=0.0,
+                            residual=True, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=512, repeat=5, kernel=[63], stride=[1], dilation=[1], dropout=0.0,
+                            residual=True, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=512, repeat=5, kernel=[63], stride=[1], dilation=[1], dropout=0.0,
+                            residual=True, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=512, repeat=5, kernel=[75], stride=[1], dilation=[1], dropout=0.0,
+                            residual=True, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=512, repeat=5, kernel=[75], stride=[1], dilation=[1], dropout=0.0,
+                            residual=True, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=512, repeat=5, kernel=[75], stride=[1], dilation=[1], dropout=0.0,
+                            residual=True, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=512, repeat=1, kernel=[87], stride=[1], dilation=[2], dropout=0.0,
+                            residual=False, groups=1, separable=True, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=1024, repeat=1, kernel=[1], stride=[1], dilation=[1], dropout=0.0,
+                            residual=False, groups=1, separable=False, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False)
+    ]
+    return config
+
+
+def jasper_10x5_dr():
+    config = [
+        JasperEncoderConfig(filters=256, repeat=1, kernel=[11], stride=[2], dilation=[1], dropout=0.2,
+                            residual=False, groups=1, separable=False, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=256, repeat=5, kernel=[11], stride=[1], dilation=[1], dropout=0.2,
+                            residual=True, groups=1, separable=False, heads=-1, residual_mode='add',
+                            residual_dense=True, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=256, repeat=5, kernel=[11], stride=[1], dilation=[1], dropout=0.2,
+                            residual=True, groups=1, separable=False, heads=-1, residual_mode='add',
+                            residual_dense=True, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=384, repeat=5, kernel=[13], stride=[1], dilation=[1], dropout=0.2,
+                            residual=True, groups=1, separable=False, heads=-1, residual_mode='add',
+                            residual_dense=True, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=384, repeat=5, kernel=[13], stride=[1], dilation=[1], dropout=0.2,
+                            residual=True, groups=1, separable=False, heads=-1, residual_mode='add',
+                            residual_dense=True, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=512, repeat=5, kernel=[17], stride=[1], dilation=[1], dropout=0.2,
+                            residual=True, groups=1, separable=False, heads=-1, residual_mode='add',
+                            residual_dense=True, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=512, repeat=5, kernel=[17], stride=[1], dilation=[1], dropout=0.2,
+                            residual=True, groups=1, separable=False, heads=-1, residual_mode='add',
+                            residual_dense=True, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=640, repeat=5, kernel=[21], stride=[1], dilation=[1], dropout=0.3,
+                            residual=True, groups=1, separable=False, heads=-1, residual_mode='add',
+                            residual_dense=True, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=640, repeat=5, kernel=[21], stride=[1], dilation=[1], dropout=0.3,
+                            residual=True, groups=1, separable=False, heads=-1, residual_mode='add',
+                            residual_dense=True, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=768, repeat=5, kernel=[25], stride=[1], dilation=[1], dropout=0.3,
+                            residual=True, groups=1, separable=False, heads=-1, residual_mode='add',
+                            residual_dense=True, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=768, repeat=5, kernel=[25], stride=[1], dilation=[1], dropout=0.3,
+                            residual=True, groups=1, separable=False, heads=-1, residual_mode='add',
+                            residual_dense=True, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=896, repeat=1, kernel=[29], stride=[1], dilation=[2], dropout=0.4,
+                            residual=False, groups=1, separable=False, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False),
+        JasperEncoderConfig(filters=1024, repeat=1, kernel=[1], stride=[1], dilation=[1], dropout=0.4,
+                            residual=False, groups=1, separable=False, heads=-1, residual_mode='add',
+                            residual_dense=False, se=False, se_reduction_ratio=8, se_context_size=-1,
+                            se_interpolation_mode='nearest', kernel_size_factor=1.0, stride_last=False)
+    ]
+    return config
+# fmt: on
+
+
+@dataclass
+class JasperModelConfig(ctc_cfg.EncDecCTCConfig):
+    # Model global arguments
+    sample_rate: int = 16000
+    repeat: int = 1
+    dropout: float = 0.0
+    separable: bool = False
+    labels: List[str] = MISSING
+
+    # Dataset configs
+    train_ds: ctc_cfg.EncDecCTCDatasetConfig = ctc_cfg.EncDecCTCDatasetConfig(
+        manifest_filepath=None, shuffle=True, trim_silence=True
+    )
+    validation_ds: ctc_cfg.EncDecCTCDatasetConfig = ctc_cfg.EncDecCTCDatasetConfig(
+        manifest_filepath=None, shuffle=False
+    )
+    test_ds: ctc_cfg.EncDecCTCDatasetConfig = ctc_cfg.EncDecCTCDatasetConfig(manifest_filepath=None, shuffle=False)
+
+    # Optimizer / Scheduler config
+    optim: Optional[model_cfg.OptimConfig] = model_cfg.OptimConfig(sched=model_cfg.SchedConfig())
+
+    # Model general component configs
+    preprocessor: AudioToMelSpectrogramPreprocessorConfig = AudioToMelSpectrogramPreprocessorConfig()
+    spec_augment: Optional[SpectrogramAugmentationConfig] = SpectrogramAugmentationConfig()
+    encoder: ConvASREncoderConfig = ConvASREncoderConfig(activation="relu")
+    decoder: ConvASRDecoderConfig = ConvASRDecoderConfig()
+
+
+@dataclass
+class QuartzNetModelConfig(JasperModelConfig):
+    separable: bool = True
+
+
+class EncDecCTCModelConfigBuilder(model_cfg.ModelConfigBuilder):
+    VALID_CONFIGS = ['quartznet_15x5', 'quartznet_15x5_zh', 'jasper_10x5dr']
+
+    def __init__(self, name: str = 'quartznet_15x5', encoder_cfg_func: Optional[Callable[[], List[Any]]] = None):
+        if name not in EncDecCTCModelConfigBuilder.VALID_CONFIGS:
+            raise ValueError("`name` must be one of : \n" f"{EncDecCTCModelConfigBuilder.VALID_CONFIGS}")
+
+        self.name = name
+
+        if 'quartznet_15x5' in name:
+            if encoder_cfg_func is None:
+                encoder_cfg_func = qn_15x5
+
+            model_cfg = QuartzNetModelConfig(
+                repeat=5,
+                separable=True,
+                spec_augment=SpectrogramAugmentationConfig(rect_masks=5, rect_freq=50, rect_time=120),
+                encoder=ConvASREncoderConfig(jasper=encoder_cfg_func(), activation="relu"),
+                decoder=ConvASRDecoderConfig(),
+            )
+
+        elif 'jasper_10x5' in name:
+            if encoder_cfg_func is None:
+                encoder_cfg_func = jasper_10x5_dr
+
+            model_cfg = JasperModelConfig(
+                repeat=5,
+                separable=False,
+                spec_augment=SpectrogramAugmentationConfig(rect_masks=5, rect_freq=50, rect_time=120),
+                encoder=ConvASREncoderConfig(jasper=encoder_cfg_func(), activation="relu"),
+                decoder=ConvASRDecoderConfig(),
+            )
+
+        else:
+            raise ValueError(f"Invalid config name submitted to {self.__class__.__name__}")
+
+        super(EncDecCTCModelConfigBuilder, self).__init__(model_cfg)
+        self.model_cfg: ctc_cfg.EncDecCTCConfig = model_cfg  # enable type hinting
+
+        if 'zh' in name:
+            self.set_dataset_normalize(normalize=False)
+
+    def set_labels(self, labels: List[str]):
+        self.model_cfg.labels = labels
+
+    def set_separable(self, separable: bool):
+        self.model_cfg.separable = separable
+
+    def set_repeat(self, repeat: int):
+        self.model_cfg.repeat = repeat
+
+    def set_sample_rate(self, sample_rate: int):
+        self.model_cfg.sample_rate = sample_rate
+
+    def set_dropout(self, dropout: float = 0.0):
+        self.model_cfg.dropout = dropout
+
+    def set_dataset_normalize(self, normalize: bool):
+        self.model_cfg.train_ds.normalize = normalize
+        self.model_cfg.validation_ds.normalize = normalize
+        self.model_cfg.test_ds.normalize = normalize
+
+    # Note: Autocomplete for users wont work without these overrides
+    # But practically it is not needed since python will infer at runtime
+
+    # def set_train_ds(self, cfg: Optional[ctc_cfg.EncDecCTCDatasetConfig] = None):
+    #     super().set_train_ds(cfg)
+    #
+    # def set_validation_ds(self, cfg: Optional[ctc_cfg.EncDecCTCDatasetConfig] = None):
+    #     super().set_validation_ds(cfg)
+    #
+    # def set_test_ds(self, cfg: Optional[ctc_cfg.EncDecCTCDatasetConfig] = None):
+    #     super().set_test_ds(cfg)
+
+    def _finalize_cfg(self):
+        # propagate labels
+        self.model_cfg.train_ds.labels = self.model_cfg.labels
+        self.model_cfg.validation_ds.labels = self.model_cfg.labels
+        self.model_cfg.test_ds.labels = self.model_cfg.labels
+        self.model_cfg.decoder.vocabulary = self.model_cfg.labels
+
+        # propagate num classes
+        self.model_cfg.decoder.num_classes = len(self.model_cfg.labels)
+
+        # propagate sample rate
+        self.model_cfg.sample_rate = self.model_cfg.sample_rate
+        self.model_cfg.preprocessor.sample_rate = self.model_cfg.sample_rate
+        self.model_cfg.train_ds.sample_rate = self.model_cfg.sample_rate
+        self.model_cfg.validation_ds.sample_rate = self.model_cfg.sample_rate
+        self.model_cfg.test_ds.sample_rate = self.model_cfg.sample_rate
+
+        # propagate filters
+        self.model_cfg.encoder.feat_in = self.model_cfg.preprocessor.features
+        self.model_cfg.decoder.feat_in = self.model_cfg.encoder.jasper[-1].filters
+
+        # propagate separable
+        for layer in self.model_cfg.encoder.jasper[:-1]:  # type: JasperEncoderConfig
+            layer.separable = self.model_cfg.separable
+
+        # propagate repeat
+        for layer in self.model_cfg.encoder.jasper[1:-2]:  # type: JasperEncoderConfig
+            layer.repeat = self.model_cfg.repeat
+
+        # propagate dropout
+        for layer in self.model_cfg.encoder.jasper:  # type: JasperEncoderConfig
+            layer.dropout = self.model_cfg.dropout
+
+    def build(self) -> ctc_cfg.EncDecCTCConfig:
+        return super().build()

tts/nemo/collections/asr/models/ctc_bpe_models.py

@@ -0,0 +1,304 @@
+# 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 copy
+import os
+from typing import Dict, Optional
+
+import torch
+from omegaconf import DictConfig, ListConfig, OmegaConf, open_dict
+
+from nemo.collections.asr.data import audio_to_text_dataset
+from nemo.collections.asr.losses.ctc import CTCLoss
+from nemo.collections.asr.metrics.wer_bpe import WERBPE
+from nemo.collections.asr.models.ctc_models import EncDecCTCModel
+from nemo.collections.asr.parts.mixins import ASRBPEMixin
+from nemo.collections.asr.parts.preprocessing.perturb import process_augmentations
+from nemo.core.classes.common import PretrainedModelInfo
+from nemo.utils import logging, model_utils
+
+__all__ = ['EncDecCTCModelBPE']
+
+
+class EncDecCTCModelBPE(EncDecCTCModel, ASRBPEMixin):
+    """Encoder decoder CTC-based models with Byte Pair Encoding."""
+
+    @classmethod
+    def list_available_models(cls) -> Optional[PretrainedModelInfo]:
+        """
+        This method returns a list of pre-trained model which can be instantiated directly from NVIDIA's NGC cloud.
+
+        Returns:
+            List of available pre-trained models.
+        """
+        results = []
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="stt_en_citrinet_256",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_citrinet_256",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_en_citrinet_256/versions/1.0.0rc1/files/stt_en_citrinet_256.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="stt_en_citrinet_512",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_citrinet_512",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_en_citrinet_512/versions/1.0.0rc1/files/stt_en_citrinet_512.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="stt_en_citrinet_1024",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_citrinet_1024",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_en_citrinet_1024/versions/1.0.0rc1/files/stt_en_citrinet_1024.nemo",
+        )
+
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="stt_es_citrinet_512",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_es_citrinet_512",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_es_citrinet_512/versions/1.0.0/files/stt_es_citrinet_512.nemo",
+        )
+
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="stt_en_conformer_ctc_small",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_small",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_en_conformer_ctc_small/versions/1.0.0/files/stt_en_conformer_ctc_small.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="stt_en_conformer_ctc_medium",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_medium",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_en_conformer_ctc_medium/versions/1.0.0/files/stt_en_conformer_ctc_medium.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="stt_en_conformer_ctc_large",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_large",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_en_conformer_ctc_large/versions/1.0.0/files/stt_en_conformer_ctc_large.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="stt_en_conformer_ctc_small_ls",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_small_ls",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_en_conformer_ctc_small_ls/versions/1.0.0/files/stt_en_conformer_ctc_small_ls.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="stt_en_conformer_ctc_medium_ls",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_medium_ls",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_en_conformer_ctc_medium_ls/versions/1.0.0/files/stt_en_conformer_ctc_medium_ls.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="stt_en_conformer_ctc_large_ls",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_large_ls",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_en_conformer_ctc_large_ls/versions/1.0.0/files/stt_en_conformer_ctc_large_ls.nemo",
+        )
+        results.append(model)
+
+        return results
+
+    def __init__(self, cfg: DictConfig, trainer=None):
+        # Convert to Hydra 1.0 compatible DictConfig
+        cfg = model_utils.convert_model_config_to_dict_config(cfg)
+        cfg = model_utils.maybe_update_config_version(cfg)
+
+        if 'tokenizer' not in cfg:
+            raise ValueError("`cfg` must have `tokenizer` config to create a tokenizer !")
+
+        # Setup the tokenizer
+        self._setup_tokenizer(cfg.tokenizer)
+
+        # Initialize a dummy vocabulary
+        vocabulary = self.tokenizer.tokenizer.get_vocab()
+
+        # Set the new vocabulary
+        with open_dict(cfg):
+            cfg.decoder.vocabulary = ListConfig(list(vocabulary.keys()))
+
+        # Override number of classes if placeholder provided
+        num_classes = cfg.decoder["num_classes"]
+
+        if num_classes < 1:
+            logging.info(
+                "\nReplacing placeholder number of classes ({}) with actual number of classes - {}".format(
+                    num_classes, len(vocabulary)
+                )
+            )
+            cfg.decoder["num_classes"] = len(vocabulary)
+
+        super().__init__(cfg=cfg, trainer=trainer)
+
+        # Setup metric objects
+        self._wer = WERBPE(
+            tokenizer=self.tokenizer,
+            batch_dim_index=0,
+            use_cer=self._cfg.get('use_cer', False),
+            ctc_decode=True,
+            dist_sync_on_step=True,
+            log_prediction=self._cfg.get("log_prediction", False),
+        )
+
+    def _setup_dataloader_from_config(self, config: Optional[Dict]):
+        if 'augmentor' in config:
+            augmentor = process_augmentations(config['augmentor'])
+        else:
+            augmentor = None
+
+        shuffle = config['shuffle']
+
+        # Instantiate tarred dataset loader or normal dataset loader
+        if config.get('is_tarred', False):
+            if ('tarred_audio_filepaths' in config and config['tarred_audio_filepaths'] is None) or (
+                'manifest_filepath' in config and config['manifest_filepath'] is None
+            ):
+                logging.warning(
+                    "Could not load dataset as `manifest_filepath` was None or "
+                    f"`tarred_audio_filepaths` is None. Provided config : {config}"
+                )
+                return None
+
+            shuffle_n = config.get('shuffle_n', 4 * config['batch_size']) if shuffle else 0
+            dataset = audio_to_text_dataset.get_tarred_bpe_dataset(
+                config=config,
+                tokenizer=self.tokenizer,
+                shuffle_n=shuffle_n,
+                global_rank=self.global_rank,
+                world_size=self.world_size,
+                augmentor=augmentor,
+            )
+            shuffle = False
+        else:
+            if 'manifest_filepath' in config and config['manifest_filepath'] is None:
+                logging.warning(f"Could not load dataset as `manifest_filepath` was None. Provided config : {config}")
+                return None
+
+            dataset = audio_to_text_dataset.get_bpe_dataset(
+                config=config, tokenizer=self.tokenizer, augmentor=augmentor
+            )
+
+        return torch.utils.data.DataLoader(
+            dataset=dataset,
+            batch_size=config['batch_size'],
+            collate_fn=dataset.collate_fn,
+            drop_last=config.get('drop_last', False),
+            shuffle=shuffle,
+            num_workers=config.get('num_workers', 0),
+            pin_memory=config.get('pin_memory', False),
+        )
+
+    def _setup_transcribe_dataloader(self, config: Dict) -> 'torch.utils.data.DataLoader':
+        """
+        Setup function for a temporary data loader which wraps the provided audio file.
+
+        Args:
+            config: A python dictionary which contains the following keys:
+            paths2audio_files: (a list) of paths to audio files. The files should be relatively short fragments. \
+                Recommended length per file is between 5 and 25 seconds.
+            batch_size: (int) batch size to use during inference. \
+                Bigger will result in better throughput performance but would use more memory.
+            temp_dir: (str) A temporary directory where the audio manifest is temporarily
+                stored.
+
+        Returns:
+            A pytorch DataLoader for the given audio file(s).
+        """
+        dl_config = {
+            'manifest_filepath': os.path.join(config['temp_dir'], 'manifest.json'),
+            'sample_rate': self.preprocessor._sample_rate,
+            'batch_size': min(config['batch_size'], len(config['paths2audio_files'])),
+            'shuffle': False,
+        }
+
+        temporary_datalayer = self._setup_dataloader_from_config(config=DictConfig(dl_config))
+        return temporary_datalayer
+
+    def change_vocabulary(self, new_tokenizer_dir: str, new_tokenizer_type: str):
+        """
+        Changes vocabulary of the tokenizer used during CTC decoding process.
+        Use this method when fine-tuning on from pre-trained model.
+        This method changes only decoder and leaves encoder and pre-processing modules unchanged. For example, you would
+        use it if you want to use pretrained encoder when fine-tuning on a data in another language, or when you'd need
+        model to learn capitalization, punctuation and/or special characters.
+
+        Args:
+            new_tokenizer_dir: Path to the new tokenizer directory.
+            new_tokenizer_type: Either `bpe` or `wpe`. `bpe` is used for SentencePiece tokenizers,
+                whereas `wpe` is used for `BertTokenizer`.
+
+        Returns: None
+
+        """
+        if not os.path.isdir(new_tokenizer_dir):
+            raise NotADirectoryError(
+                f'New tokenizer dir must be non-empty path to a directory. But I got: {new_tokenizer_dir}'
+            )
+
+        if new_tokenizer_type.lower() not in ('bpe', 'wpe'):
+            raise ValueError(f'New tokenizer type must be either `bpe` or `wpe`')
+
+        tokenizer_cfg = OmegaConf.create({'dir': new_tokenizer_dir, 'type': new_tokenizer_type})
+
+        # Setup the tokenizer
+        self._setup_tokenizer(tokenizer_cfg)
+
+        # Initialize a dummy vocabulary
+        vocabulary = self.tokenizer.tokenizer.get_vocab()
+
+        # Set the new vocabulary
+        decoder_config = copy.deepcopy(self.decoder.to_config_dict())
+        decoder_config.vocabulary = ListConfig(list(vocabulary.keys()))
+
+        decoder_num_classes = decoder_config['num_classes']
+
+        # Override number of classes if placeholder provided
+        logging.info(
+            "\nReplacing old number of classes ({}) with new number of classes - {}".format(
+                decoder_num_classes, len(vocabulary)
+            )
+        )
+
+        decoder_config['num_classes'] = len(vocabulary)
+
+        del self.decoder
+        self.decoder = EncDecCTCModelBPE.from_config_dict(decoder_config)
+        del self.loss
+        self.loss = CTCLoss(
+            num_classes=self.decoder.num_classes_with_blank - 1,
+            zero_infinity=True,
+            reduction=self._cfg.get("ctc_reduction", "mean_batch"),
+        )
+        self._wer = WERBPE(
+            tokenizer=self.tokenizer,
+            batch_dim_index=0,
+            use_cer=self._cfg.get('use_cer', False),
+            ctc_decode=True,
+            log_prediction=self._cfg.get("log_prediction", False),
+        )
+
+        # Update config
+        OmegaConf.set_struct(self._cfg.decoder, False)
+        self._cfg.decoder = decoder_config
+        OmegaConf.set_struct(self._cfg.decoder, True)
+
+        logging.info(f"Changed tokenizer to {self.decoder.vocabulary} vocabulary.")

tts/nemo/collections/asr/models/ctc_models.py

@@ -0,0 +1,649 @@
+# 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 copy
+import json
+import os
+import tempfile
+from math import ceil
+from typing import Dict, List, Optional, Union
+
+import torch
+from omegaconf import DictConfig, OmegaConf, open_dict
+from pytorch_lightning import Trainer
+from tqdm.auto import tqdm
+
+from nemo.collections.asr.data import audio_to_text_dataset
+from nemo.collections.asr.data.audio_to_text_dali import DALIOutputs
+from nemo.collections.asr.losses.ctc import CTCLoss
+from nemo.collections.asr.metrics.wer import WER
+from nemo.collections.asr.models.asr_model import ASRModel, ExportableEncDecModel
+from nemo.collections.asr.parts.mixins import ASRModuleMixin
+from nemo.collections.asr.parts.preprocessing.perturb import process_augmentations
+from nemo.core.classes.common import PretrainedModelInfo, typecheck
+from nemo.core.neural_types import AudioSignal, LabelsType, LengthsType, LogprobsType, NeuralType, SpectrogramType
+from nemo.utils import logging
+
+__all__ = ['EncDecCTCModel']
+
+
+class EncDecCTCModel(ASRModel, ExportableEncDecModel, ASRModuleMixin):
+    """Base class for encoder decoder CTC-based models."""
+
+    @classmethod
+    def list_available_models(cls) -> Optional[PretrainedModelInfo]:
+        """
+        This method returns a list of pre-trained model which can be instantiated directly from NVIDIA's NGC cloud.
+
+        Returns:
+            List of available pre-trained models.
+        """
+        results = []
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="QuartzNet15x5Base-En",
+            description="QuartzNet15x5 model trained on six datasets: LibriSpeech, Mozilla Common Voice (validated clips from en_1488h_2019-12-10), WSJ, Fisher, Switchboard, and NSC Singapore English. It was trained with Apex/Amp optimization level O1 for 600 epochs. The model achieves a WER of 3.79% on LibriSpeech dev-clean, and a WER of 10.05% on dev-other. Please visit https://ngc.nvidia.com/catalog/models/nvidia:nemospeechmodels for further details.",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemospeechmodels/versions/1.0.0a5/files/QuartzNet15x5Base-En.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="stt_en_quartznet15x5",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_quartznet15x5",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_en_quartznet15x5/versions/1.0.0rc1/files/stt_en_quartznet15x5.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="stt_zh_quartznet15x5",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_zh_quartznet15x5",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_zh_quartznet15x5/versions/1.0.0rc1/files/stt_zh_quartznet15x5.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="stt_en_jasper10x5dr",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_jasper10x5dr",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_en_jasper10x5dr/versions/1.0.0rc1/files/stt_en_jasper10x5dr.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="stt_ca_quartznet15x5",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_ca_quartznet15x5",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_ca_quartznet15x5/versions/1.0.0rc1/files/stt_ca_quartznet15x5.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="stt_it_quartznet15x5",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_it_quartznet15x5",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_it_quartznet15x5/versions/1.0.0rc1/files/stt_it_quartznet15x5.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="stt_fr_quartznet15x5",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_fr_quartznet15x5",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_fr_quartznet15x5/versions/1.0.0rc1/files/stt_fr_quartznet15x5.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="stt_es_quartznet15x5",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_es_quartznet15x5",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_es_quartznet15x5/versions/1.0.0rc1/files/stt_es_quartznet15x5.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="stt_de_quartznet15x5",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_de_quartznet15x5",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_de_quartznet15x5/versions/1.0.0rc1/files/stt_de_quartznet15x5.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="stt_pl_quartznet15x5",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_pl_quartznet15x5",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_pl_quartznet15x5/versions/1.0.0rc1/files/stt_pl_quartznet15x5.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="stt_ru_quartznet15x5",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_ru_quartznet15x5",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_ru_quartznet15x5/versions/1.0.0rc1/files/stt_ru_quartznet15x5.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="stt_zh_citrinet_512",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_zh_citrinet_512",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_zh_citrinet_512/versions/1.0.0rc1/files/stt_zh_citrinet_512.nemo",
+        )
+        results.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="asr_talknet_aligner",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:asr_talknet_aligner",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/asr_talknet_aligner/versions/1.0.0rc1/files/qn5x5_libri_tts_phonemes.nemo",
+        )
+        results.append(model)
+
+        return results
+
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+        # Get global rank and total number of GPU workers for IterableDataset partitioning, if applicable
+        # Global_rank and local_rank is set by LightningModule in Lightning 1.2.0
+        self.world_size = 1
+        if trainer is not None:
+            self.world_size = trainer.num_nodes * trainer.num_gpus
+
+        super().__init__(cfg=cfg, trainer=trainer)
+        self.preprocessor = EncDecCTCModel.from_config_dict(self._cfg.preprocessor)
+        self.encoder = EncDecCTCModel.from_config_dict(self._cfg.encoder)
+
+        with open_dict(self._cfg):
+            if "feat_in" not in self._cfg.decoder or (
+                not self._cfg.decoder.feat_in and hasattr(self.encoder, '_feat_out')
+            ):
+                self._cfg.decoder.feat_in = self.encoder._feat_out
+            if "feat_in" not in self._cfg.decoder or not self._cfg.decoder.feat_in:
+                raise ValueError("param feat_in of the decoder's config is not set!")
+
+        self.decoder = EncDecCTCModel.from_config_dict(self._cfg.decoder)
+
+        self.loss = CTCLoss(
+            num_classes=self.decoder.num_classes_with_blank - 1,
+            zero_infinity=True,
+            reduction=self._cfg.get("ctc_reduction", "mean_batch"),
+        )
+
+        if hasattr(self._cfg, 'spec_augment') and self._cfg.spec_augment is not None:
+            self.spec_augmentation = EncDecCTCModel.from_config_dict(self._cfg.spec_augment)
+        else:
+            self.spec_augmentation = None
+
+        # Setup metric objects
+        self._wer = WER(
+            vocabulary=self.decoder.vocabulary,
+            batch_dim_index=0,
+            use_cer=self._cfg.get('use_cer', False),
+            ctc_decode=True,
+            dist_sync_on_step=True,
+            log_prediction=self._cfg.get("log_prediction", False),
+        )
+
+    @torch.no_grad()
+    def transcribe(
+        self,
+        paths2audio_files: List[str],
+        batch_size: int = 4,
+        logprobs: bool = False,
+        return_hypotheses: bool = False,
+    ) -> List[str]:
+        """
+        Uses greedy decoding to transcribe audio files. Use this method for debugging and prototyping.
+
+        Args:
+            paths2audio_files: (a list) of paths to audio files. \
+                Recommended length per file is between 5 and 25 seconds. \
+                But it is possible to pass a few hours long file if enough GPU memory is available.
+            batch_size: (int) batch size to use during inference.
+                Bigger will result in better throughput performance but would use more memory.
+            logprobs: (bool) pass True to get log probabilities instead of transcripts.
+            return_hypotheses: (bool) Either return hypotheses or text
+                With hypotheses can do some postprocessing like getting timestamp or rescoring
+
+        Returns:
+            A list of transcriptions (or raw log probabilities if logprobs is True) in the same order as paths2audio_files
+        """
+        if paths2audio_files is None or len(paths2audio_files) == 0:
+            return {}
+
+        if return_hypotheses and logprobs:
+            raise ValueError(
+                "Either `return_hypotheses` or `logprobs` can be True at any given time."
+                "Returned hypotheses will contain the logprobs."
+            )
+
+        # We will store transcriptions here
+        hypotheses = []
+        # Model's mode and device
+        mode = self.training
+        device = next(self.parameters()).device
+        dither_value = self.preprocessor.featurizer.dither
+        pad_to_value = self.preprocessor.featurizer.pad_to
+
+        try:
+            self.preprocessor.featurizer.dither = 0.0
+            self.preprocessor.featurizer.pad_to = 0
+            # Switch model to evaluation mode
+            self.eval()
+            # Freeze the encoder and decoder modules
+            self.encoder.freeze()
+            self.decoder.freeze()
+            logging_level = logging.get_verbosity()
+            logging.set_verbosity(logging.WARNING)
+            # Work in tmp directory - will store manifest file there
+            with tempfile.TemporaryDirectory() as tmpdir:
+                with open(os.path.join(tmpdir, 'manifest.json'), 'w') as fp:
+                    for audio_file in paths2audio_files:
+                        entry = {'audio_filepath': audio_file, 'duration': 100000, 'text': 'nothing'}
+                        fp.write(json.dumps(entry) + '\n')
+
+                config = {'paths2audio_files': paths2audio_files, 'batch_size': batch_size, 'temp_dir': tmpdir}
+
+                temporary_datalayer = self._setup_transcribe_dataloader(config)
+                for test_batch in tqdm(temporary_datalayer, desc="Transcribing"):
+                    logits, logits_len, greedy_predictions = self.forward(
+                        input_signal=test_batch[0].to(device), input_signal_length=test_batch[1].to(device)
+                    )
+                    if logprobs:
+                        # dump log probs per file
+                        for idx in range(logits.shape[0]):
+                            lg = logits[idx][: logits_len[idx]]
+                            hypotheses.append(lg.cpu().numpy())
+                    else:
+                        current_hypotheses = self._wer.ctc_decoder_predictions_tensor(
+                            greedy_predictions, predictions_len=logits_len, return_hypotheses=return_hypotheses,
+                        )
+
+                        if return_hypotheses:
+                            # dump log probs per file
+                            for idx in range(logits.shape[0]):
+                                current_hypotheses[idx].y_sequence = logits[idx][: logits_len[idx]]
+
+                        hypotheses += current_hypotheses
+
+                    del greedy_predictions
+                    del logits
+                    del test_batch
+        finally:
+            # set mode back to its original value
+            self.train(mode=mode)
+            self.preprocessor.featurizer.dither = dither_value
+            self.preprocessor.featurizer.pad_to = pad_to_value
+            if mode is True:
+                self.encoder.unfreeze()
+                self.decoder.unfreeze()
+            logging.set_verbosity(logging_level)
+        return hypotheses
+
+    def change_vocabulary(self, new_vocabulary: List[str]):
+        """
+        Changes vocabulary used during CTC decoding process. Use this method when fine-tuning on from pre-trained model.
+        This method changes only decoder and leaves encoder and pre-processing modules unchanged. For example, you would
+        use it if you want to use pretrained encoder when fine-tuning on a data in another language, or when you'd need
+        model to learn capitalization, punctuation and/or special characters.
+
+        If new_vocabulary == self.decoder.vocabulary then nothing will be changed.
+
+        Args:
+
+            new_vocabulary: list with new vocabulary. Must contain at least 2 elements. Typically, \
+            this is target alphabet.
+
+        Returns: None
+
+        """
+        if self.decoder.vocabulary == new_vocabulary:
+            logging.warning(f"Old {self.decoder.vocabulary} and new {new_vocabulary} match. Not changing anything.")
+        else:
+            if new_vocabulary is None or len(new_vocabulary) == 0:
+                raise ValueError(f'New vocabulary must be non-empty list of chars. But I got: {new_vocabulary}')
+            decoder_config = self.decoder.to_config_dict()
+            new_decoder_config = copy.deepcopy(decoder_config)
+            new_decoder_config['vocabulary'] = new_vocabulary
+            new_decoder_config['num_classes'] = len(new_vocabulary)
+
+            del self.decoder
+            self.decoder = EncDecCTCModel.from_config_dict(new_decoder_config)
+            del self.loss
+            self.loss = CTCLoss(
+                num_classes=self.decoder.num_classes_with_blank - 1,
+                zero_infinity=True,
+                reduction=self._cfg.get("ctc_reduction", "mean_batch"),
+            )
+            self._wer = WER(
+                vocabulary=self.decoder.vocabulary,
+                batch_dim_index=0,
+                use_cer=self._cfg.get('use_cer', False),
+                ctc_decode=True,
+                dist_sync_on_step=True,
+                log_prediction=self._cfg.get("log_prediction", False),
+            )
+
+            # Update config
+            OmegaConf.set_struct(self._cfg.decoder, False)
+            self._cfg.decoder = new_decoder_config
+            OmegaConf.set_struct(self._cfg.decoder, True)
+
+            ds_keys = ['train_ds', 'validation_ds', 'test_ds']
+            for key in ds_keys:
+                if key in self.cfg:
+                    with open_dict(self.cfg[key]):
+                        self.cfg[key]['labels'] = OmegaConf.create(new_vocabulary)
+
+            logging.info(f"Changed decoder to output to {self.decoder.vocabulary} vocabulary.")
+
+    def _setup_dataloader_from_config(self, config: Optional[Dict]):
+        if 'augmentor' in config:
+            augmentor = process_augmentations(config['augmentor'])
+        else:
+            augmentor = None
+
+        # Automatically inject args from model config to dataloader config
+        audio_to_text_dataset.inject_dataloader_value_from_model_config(self.cfg, config, key='sample_rate')
+        audio_to_text_dataset.inject_dataloader_value_from_model_config(self.cfg, config, key='labels')
+
+        shuffle = config['shuffle']
+        device = 'gpu' if torch.cuda.is_available() else 'cpu'
+        if config.get('use_dali', False):
+            device_id = self.local_rank if device == 'gpu' else None
+            dataset = audio_to_text_dataset.get_dali_char_dataset(
+                config=config,
+                shuffle=shuffle,
+                device_id=device_id,
+                global_rank=self.global_rank,
+                world_size=self.world_size,
+                preprocessor_cfg=self._cfg.preprocessor,
+            )
+            return dataset
+
+        # Instantiate tarred dataset loader or normal dataset loader
+        if config.get('is_tarred', False):
+            if ('tarred_audio_filepaths' in config and config['tarred_audio_filepaths'] is None) or (
+                'manifest_filepath' in config and config['manifest_filepath'] is None
+            ):
+                logging.warning(
+                    "Could not load dataset as `manifest_filepath` was None or "
+                    f"`tarred_audio_filepaths` is None. Provided config : {config}"
+                )
+                return None
+
+            shuffle_n = config.get('shuffle_n', 4 * config['batch_size']) if shuffle else 0
+            dataset = audio_to_text_dataset.get_tarred_char_dataset(
+                config=config,
+                shuffle_n=shuffle_n,
+                global_rank=self.global_rank,
+                world_size=self.world_size,
+                augmentor=augmentor,
+            )
+            shuffle = False
+        else:
+            if 'manifest_filepath' in config and config['manifest_filepath'] is None:
+                logging.warning(f"Could not load dataset as `manifest_filepath` was None. Provided config : {config}")
+                return None
+
+            dataset = audio_to_text_dataset.get_char_dataset(config=config, augmentor=augmentor)
+
+        return torch.utils.data.DataLoader(
+            dataset=dataset,
+            batch_size=config['batch_size'],
+            collate_fn=dataset.collate_fn,
+            drop_last=config.get('drop_last', False),
+            shuffle=shuffle,
+            num_workers=config.get('num_workers', 0),
+            pin_memory=config.get('pin_memory', False),
+        )
+
+    def setup_training_data(self, train_data_config: Optional[Union[DictConfig, Dict]]):
+        """
+        Sets up the training data loader via a Dict-like object.
+
+        Args:
+            train_data_config: A config that contains the information regarding construction
+                of an ASR Training dataset.
+
+        Supported Datasets:
+            -   :class:`~nemo.collections.asr.data.audio_to_text.AudioToCharDataset`
+            -   :class:`~nemo.collections.asr.data.audio_to_text.AudioToBPEDataset`
+            -   :class:`~nemo.collections.asr.data.audio_to_text.TarredAudioToCharDataset`
+            -   :class:`~nemo.collections.asr.data.audio_to_text.TarredAudioToBPEDataset`
+            -   :class:`~nemo.collections.asr.data.audio_to_text_dali.AudioToCharDALIDataset`
+        """
+        if 'shuffle' not in train_data_config:
+            train_data_config['shuffle'] = True
+
+        # preserve config
+        self._update_dataset_config(dataset_name='train', config=train_data_config)
+
+        self._train_dl = self._setup_dataloader_from_config(config=train_data_config)
+
+        # Need to set this because if using an IterableDataset, the length of the dataloader is the total number
+        # of samples rather than the number of batches, and this messes up the tqdm progress bar.
+        # So we set the number of steps manually (to the correct number) to fix this.
+        if 'is_tarred' in train_data_config and train_data_config['is_tarred']:
+            # We also need to check if limit_train_batches is already set.
+            # If it's an int, we assume that the user has set it to something sane, i.e. <= # training batches,
+            # and don't change it. Otherwise, adjust batches accordingly if it's a float (including 1.0).
+            if isinstance(self._trainer.limit_train_batches, float):
+                self._trainer.limit_train_batches = int(
+                    self._trainer.limit_train_batches
+                    * ceil((len(self._train_dl.dataset) / self.world_size) / train_data_config['batch_size'])
+                )
+
+    def setup_validation_data(self, val_data_config: Optional[Union[DictConfig, Dict]]):
+        """
+        Sets up the validation data loader via a Dict-like object.
+
+        Args:
+            val_data_config: A config that contains the information regarding construction
+                of an ASR Training dataset.
+
+        Supported Datasets:
+            -   :class:`~nemo.collections.asr.data.audio_to_text.AudioToCharDataset`
+            -   :class:`~nemo.collections.asr.data.audio_to_text.AudioToBPEDataset`
+            -   :class:`~nemo.collections.asr.data.audio_to_text.TarredAudioToCharDataset`
+            -   :class:`~nemo.collections.asr.data.audio_to_text.TarredAudioToBPEDataset`
+            -   :class:`~nemo.collections.asr.data.audio_to_text_dali.AudioToCharDALIDataset`
+        """
+        if 'shuffle' not in val_data_config:
+            val_data_config['shuffle'] = False
+
+        # preserve config
+        self._update_dataset_config(dataset_name='validation', config=val_data_config)
+
+        self._validation_dl = self._setup_dataloader_from_config(config=val_data_config)
+
+    def setup_test_data(self, test_data_config: Optional[Union[DictConfig, Dict]]):
+        """
+        Sets up the test data loader via a Dict-like object.
+
+        Args:
+            test_data_config: A config that contains the information regarding construction
+                of an ASR Training dataset.
+
+        Supported Datasets:
+            -   :class:`~nemo.collections.asr.data.audio_to_text.AudioToCharDataset`
+            -   :class:`~nemo.collections.asr.data.audio_to_text.AudioToBPEDataset`
+            -   :class:`~nemo.collections.asr.data.audio_to_text.TarredAudioToCharDataset`
+            -   :class:`~nemo.collections.asr.data.audio_to_text.TarredAudioToBPEDataset`
+            -   :class:`~nemo.collections.asr.data.audio_to_text_dali.AudioToCharDALIDataset`
+        """
+        if 'shuffle' not in test_data_config:
+            test_data_config['shuffle'] = False
+
+        # preserve config
+        self._update_dataset_config(dataset_name='test', config=test_data_config)
+
+        self._test_dl = self._setup_dataloader_from_config(config=test_data_config)
+
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        if hasattr(self.preprocessor, '_sample_rate'):
+            input_signal_eltype = AudioSignal(freq=self.preprocessor._sample_rate)
+        else:
+            input_signal_eltype = AudioSignal()
+        return {
+            "input_signal": NeuralType(('B', 'T'), input_signal_eltype, optional=True),
+            "input_signal_length": NeuralType(tuple('B'), LengthsType(), optional=True),
+            "processed_signal": NeuralType(('B', 'D', 'T'), SpectrogramType(), optional=True),
+            "processed_signal_length": NeuralType(tuple('B'), LengthsType(), optional=True),
+        }
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {
+            "outputs": NeuralType(('B', 'T', 'D'), LogprobsType()),
+            "encoded_lengths": NeuralType(tuple('B'), LengthsType()),
+            "greedy_predictions": NeuralType(('B', 'T'), LabelsType()),
+        }
+
+    @typecheck()
+    def forward(
+        self, input_signal=None, input_signal_length=None, processed_signal=None, processed_signal_length=None
+    ):
+        """
+        Forward pass of the model.
+
+        Args:
+            input_signal: Tensor that represents a batch of raw audio signals,
+                of shape [B, T]. T here represents timesteps, with 1 second of audio represented as
+                `self.sample_rate` number of floating point values.
+            input_signal_length: Vector of length B, that contains the individual lengths of the audio
+                sequences.
+            processed_signal: Tensor that represents a batch of processed audio signals,
+                of shape (B, D, T) that has undergone processing via some DALI preprocessor.
+            processed_signal_length: Vector of length B, that contains the individual lengths of the
+                processed audio sequences.
+
+        Returns:
+            A tuple of 3 elements -
+            1) The log probabilities tensor of shape [B, T, D].
+            2) The lengths of the acoustic sequence after propagation through the encoder, of shape [B].
+            3) The greedy token predictions of the model of shape [B, T] (via argmax)
+        """
+        has_input_signal = input_signal is not None and input_signal_length is not None
+        has_processed_signal = processed_signal is not None and processed_signal_length is not None
+        if (has_input_signal ^ has_processed_signal) == False:
+            raise ValueError(
+                f"{self} Arguments ``input_signal`` and ``input_signal_length`` are mutually exclusive "
+                " with ``processed_signal`` and ``processed_signal_len`` arguments."
+            )
+
+        if not has_processed_signal:
+            processed_signal, processed_signal_length = self.preprocessor(
+                input_signal=input_signal, length=input_signal_length,
+            )
+
+        if self.spec_augmentation is not None and self.training:
+            processed_signal = self.spec_augmentation(input_spec=processed_signal, length=processed_signal_length)
+
+        encoded, encoded_len = self.encoder(audio_signal=processed_signal, length=processed_signal_length)
+        log_probs = self.decoder(encoder_output=encoded)
+        greedy_predictions = log_probs.argmax(dim=-1, keepdim=False)
+
+        return log_probs, encoded_len, greedy_predictions
+
+    # PTL-specific methods
+    def training_step(self, batch, batch_nb):
+        signal, signal_len, transcript, transcript_len = batch
+        if isinstance(batch, DALIOutputs) and batch.has_processed_signal:
+            log_probs, encoded_len, predictions = self.forward(
+                processed_signal=signal, processed_signal_length=signal_len
+            )
+        else:
+            log_probs, encoded_len, predictions = self.forward(input_signal=signal, input_signal_length=signal_len)
+
+        loss_value = self.loss(
+            log_probs=log_probs, targets=transcript, input_lengths=encoded_len, target_lengths=transcript_len
+        )
+
+        tensorboard_logs = {'train_loss': loss_value, 'learning_rate': self._optimizer.param_groups[0]['lr']}
+
+        if hasattr(self, '_trainer') and self._trainer is not None:
+            log_every_n_steps = self._trainer.log_every_n_steps
+        else:
+            log_every_n_steps = 1
+
+        if (batch_nb + 1) % log_every_n_steps == 0:
+            self._wer.update(
+                predictions=predictions,
+                targets=transcript,
+                target_lengths=transcript_len,
+                predictions_lengths=encoded_len,
+            )
+            wer, _, _ = self._wer.compute()
+            self._wer.reset()
+            tensorboard_logs.update({'training_batch_wer': wer})
+
+        return {'loss': loss_value, 'log': tensorboard_logs}
+
+    def validation_step(self, batch, batch_idx, dataloader_idx=0):
+        signal, signal_len, transcript, transcript_len = batch
+        if isinstance(batch, DALIOutputs) and batch.has_processed_signal:
+            log_probs, encoded_len, predictions = self.forward(
+                processed_signal=signal, processed_signal_length=signal_len
+            )
+        else:
+            log_probs, encoded_len, predictions = self.forward(input_signal=signal, input_signal_length=signal_len)
+
+        loss_value = self.loss(
+            log_probs=log_probs, targets=transcript, input_lengths=encoded_len, target_lengths=transcript_len
+        )
+        self._wer.update(
+            predictions=predictions, targets=transcript, target_lengths=transcript_len, predictions_lengths=encoded_len
+        )
+        wer, wer_num, wer_denom = self._wer.compute()
+        self._wer.reset()
+        return {
+            'val_loss': loss_value,
+            'val_wer_num': wer_num,
+            'val_wer_denom': wer_denom,
+            'val_wer': wer,
+        }
+
+    def test_step(self, batch, batch_idx, dataloader_idx=0):
+        logs = self.validation_step(batch, batch_idx, dataloader_idx=dataloader_idx)
+        test_logs = {
+            'test_loss': logs['val_loss'],
+            'test_wer_num': logs['val_wer_num'],
+            'test_wer_denom': logs['val_wer_denom'],
+            'test_wer': logs['val_wer'],
+        }
+        return test_logs
+
+    def test_dataloader(self):
+        if self._test_dl is not None:
+            return self._test_dl
+
+    def _setup_transcribe_dataloader(self, config: Dict) -> 'torch.utils.data.DataLoader':
+        """
+        Setup function for a temporary data loader which wraps the provided audio file.
+
+        Args:
+            config: A python dictionary which contains the following keys:
+            paths2audio_files: (a list) of paths to audio files. The files should be relatively short fragments. \
+                Recommended length per file is between 5 and 25 seconds.
+            batch_size: (int) batch size to use during inference. \
+                Bigger will result in better throughput performance but would use more memory.
+            temp_dir: (str) A temporary directory where the audio manifest is temporarily
+                stored.
+
+        Returns:
+            A pytorch DataLoader for the given audio file(s).
+        """
+        dl_config = {
+            'manifest_filepath': os.path.join(config['temp_dir'], 'manifest.json'),
+            'sample_rate': self.preprocessor._sample_rate,
+            'labels': self.decoder.vocabulary,
+            'batch_size': min(config['batch_size'], len(config['paths2audio_files'])),
+            'trim_silence': True,
+            'shuffle': False,
+        }
+
+        temporary_datalayer = self._setup_dataloader_from_config(config=DictConfig(dl_config))
+        return temporary_datalayer

tts/nemo/collections/asr/models/label_models.py

@@ -0,0 +1,409 @@
+# 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 copy
+import json
+import os
+import pickle as pkl
+from typing import Dict, List, Optional, Union
+
+import librosa
+import torch
+from omegaconf import DictConfig
+from omegaconf.omegaconf import open_dict
+from pytorch_lightning import Trainer
+
+from nemo.collections.asr.data.audio_to_label import AudioToSpeechLabelDataset
+from nemo.collections.asr.losses.angularloss import AngularSoftmaxLoss
+from nemo.collections.asr.models.asr_model import ExportableEncDecModel
+from nemo.collections.asr.parts.preprocessing.features import WaveformFeaturizer
+from nemo.collections.asr.parts.preprocessing.perturb import process_augmentations
+from nemo.collections.common.losses import CrossEntropyLoss as CELoss
+from nemo.collections.common.metrics import TopKClassificationAccuracy
+from nemo.core.classes import ModelPT
+from nemo.core.classes.common import PretrainedModelInfo, typecheck
+from nemo.core.neural_types import *
+from nemo.utils import logging
+
+__all__ = ['EncDecSpeakerLabelModel', 'ExtractSpeakerEmbeddingsModel']
+
+
+class EncDecSpeakerLabelModel(ModelPT, ExportableEncDecModel):
+    """Encoder decoder class for speaker label models.
+    Model class creates training, validation methods for setting up data
+    performing model forward pass.
+    Expects config dict for
+    * preprocessor
+    * Jasper/Quartznet Encoder
+    * Speaker Decoder
+    """
+
+    @classmethod
+    def list_available_models(cls) -> List[PretrainedModelInfo]:
+        """
+        This method returns a list of pre-trained model which can be instantiated directly from NVIDIA's NGC cloud.
+        Returns:
+            List of available pre-trained models.
+        """
+        result = []
+        model = PretrainedModelInfo(
+            pretrained_model_name="speakerrecognition_speakernet",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/speakerrecognition_speakernet/versions/1.0.0rc1/files/speakerrecognition_speakernet.nemo",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:speakerrecognition_speakernet",
+        )
+        result.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="speakerverification_speakernet",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/speakerverification_speakernet/versions/1.0.0rc1/files/speakerverification_speakernet.nemo",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:speakerverification_speakernet",
+        )
+        result.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="speakerdiarization_speakernet",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/speakerdiarization_speakernet/versions/1.0.0rc1/files/speakerdiarization_speakernet.nemo",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:speakerdiarization_speakernet",
+        )
+        result.append(model)
+
+        return result
+
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+        super().__init__(cfg=cfg, trainer=trainer)
+        self.preprocessor = EncDecSpeakerLabelModel.from_config_dict(cfg.preprocessor)
+        self.encoder = EncDecSpeakerLabelModel.from_config_dict(cfg.encoder)
+        self.decoder = EncDecSpeakerLabelModel.from_config_dict(cfg.decoder)
+        if 'angular' in cfg.decoder and cfg.decoder['angular']:
+            logging.info("Training with Angular Softmax Loss")
+            scale = cfg.loss.scale
+            margin = cfg.loss.margin
+            self.loss = AngularSoftmaxLoss(scale=scale, margin=margin)
+        else:
+            logging.info("Training with Softmax-CrossEntropy loss")
+            self.loss = CELoss()
+        self.task = None
+        self._accuracy = TopKClassificationAccuracy(top_k=[1])
+
+    def __setup_dataloader_from_config(self, config: Optional[Dict]):
+        if 'augmentor' in config:
+            augmentor = process_augmentations(config['augmentor'])
+        else:
+            augmentor = None
+
+        featurizer = WaveformFeaturizer(
+            sample_rate=config['sample_rate'], int_values=config.get('int_values', False), augmentor=augmentor
+        )
+        self.dataset = AudioToSpeechLabelDataset(
+            manifest_filepath=config['manifest_filepath'],
+            labels=config['labels'],
+            featurizer=featurizer,
+            max_duration=config.get('max_duration', None),
+            min_duration=config.get('min_duration', None),
+            trim=False,
+            time_length=config.get('time_length', 8),
+            shift_length=config.get('shift_length', 0.75),
+        )
+
+        if self.task == 'diarization':
+            logging.info("Setting up diarization parameters")
+            _collate_func = self.dataset.sliced_seq_collate_fn
+            batch_size = 1
+            shuffle = False
+        else:
+            logging.info("Setting up identification parameters")
+            _collate_func = self.dataset.fixed_seq_collate_fn
+            batch_size = config['batch_size']
+            shuffle = config.get('shuffle', False)
+
+        return torch.utils.data.DataLoader(
+            dataset=self.dataset,
+            batch_size=batch_size,
+            collate_fn=_collate_func,
+            drop_last=config.get('drop_last', False),
+            shuffle=shuffle,
+            num_workers=config.get('num_workers', 0),
+            pin_memory=config.get('pin_memory', False),
+        )
+
+    def setup_training_data(self, train_data_layer_config: Optional[Union[DictConfig, Dict]]):
+        if 'shuffle' not in train_data_layer_config:
+            train_data_layer_config['shuffle'] = True
+        self.task = 'identification'
+        self._train_dl = self.__setup_dataloader_from_config(config=train_data_layer_config)
+
+    def setup_validation_data(self, val_data_layer_config: Optional[Union[DictConfig, Dict]]):
+        val_data_layer_config['labels'] = self.dataset.labels
+        self.task = 'identification'
+        self._validation_dl = self.__setup_dataloader_from_config(config=val_data_layer_config)
+
+    def setup_test_data(self, test_data_layer_params: Optional[Union[DictConfig, Dict]]):
+        if hasattr(self, 'dataset'):
+            test_data_layer_params['labels'] = self.dataset.labels
+
+        if 'task' in test_data_layer_params and test_data_layer_params['task']:
+            self.task = test_data_layer_params['task'].lower()
+            self.time_length = test_data_layer_params.get('time_length', 1.5)
+            self.shift_length = test_data_layer_params.get('shift_length', 0.75)
+        else:
+            self.task = 'identification'
+
+        self.embedding_dir = test_data_layer_params.get('embedding_dir', './')
+        self._test_dl = self.__setup_dataloader_from_config(config=test_data_layer_params)
+        self.test_manifest = test_data_layer_params.get('manifest_filepath', None)
+
+    def test_dataloader(self):
+        if self._test_dl is not None:
+            return self._test_dl
+
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        if hasattr(self.preprocessor, '_sample_rate'):
+            audio_eltype = AudioSignal(freq=self.preprocessor._sample_rate)
+        else:
+            audio_eltype = AudioSignal()
+        return {
+            "input_signal": NeuralType(('B', 'T'), audio_eltype),
+            "input_signal_length": NeuralType(tuple('B'), LengthsType()),
+        }
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {
+            "logits": NeuralType(('B', 'D'), LogitsType()),
+            "embs": NeuralType(('B', 'D'), AcousticEncodedRepresentation()),
+        }
+
+    @typecheck()
+    def forward(self, input_signal, input_signal_length):
+        processed_signal, processed_signal_len = self.preprocessor(
+            input_signal=input_signal, length=input_signal_length,
+        )
+
+        encoded, _ = self.encoder(audio_signal=processed_signal, length=processed_signal_len)
+        logits, embs = self.decoder(encoder_output=encoded)
+        return logits, embs
+
+    # PTL-specific methods
+    def training_step(self, batch, batch_idx):
+        audio_signal, audio_signal_len, labels, _ = batch
+        logits, _ = self.forward(input_signal=audio_signal, input_signal_length=audio_signal_len)
+        loss = self.loss(logits=logits, labels=labels)
+
+        self.log('loss', loss)
+        self.log('learning_rate', self._optimizer.param_groups[0]['lr'])
+
+        self._accuracy(logits=logits, labels=labels)
+        top_k = self._accuracy.compute()
+        self._accuracy.reset()
+        for i, top_i in enumerate(top_k):
+            self.log(f'training_batch_accuracy_top@{i}', top_i)
+
+        return {'loss': loss}
+
+    def validation_step(self, batch, batch_idx, dataloader_idx: int = 0):
+        audio_signal, audio_signal_len, labels, _ = batch
+        logits, _ = self.forward(input_signal=audio_signal, input_signal_length=audio_signal_len)
+        loss_value = self.loss(logits=logits, labels=labels)
+        acc_top_k = self._accuracy(logits=logits, labels=labels)
+        correct_counts, total_counts = self._accuracy.correct_counts_k, self._accuracy.total_counts_k
+
+        return {
+            'val_loss': loss_value,
+            'val_correct_counts': correct_counts,
+            'val_total_counts': total_counts,
+            'val_acc_top_k': acc_top_k,
+        }
+
+    def multi_validation_epoch_end(self, outputs, dataloader_idx: int = 0):
+        val_loss_mean = torch.stack([x['val_loss'] for x in outputs]).mean()
+        correct_counts = torch.stack([x['val_correct_counts'] for x in outputs]).sum(axis=0)
+        total_counts = torch.stack([x['val_total_counts'] for x in outputs]).sum(axis=0)
+
+        self._accuracy.correct_counts_k = correct_counts
+        self._accuracy.total_counts_k = total_counts
+        topk_scores = self._accuracy.compute()
+        self._accuracy.reset()
+
+        logging.info("val_loss: {:.3f}".format(val_loss_mean))
+        self.log('val_loss', val_loss_mean)
+        for top_k, score in zip(self._accuracy.top_k, topk_scores):
+            self.log('val_epoch_accuracy_top@{}'.format(top_k), score)
+
+        return {
+            'val_loss': val_loss_mean,
+            'val_acc_top_k': topk_scores,
+        }
+
+    def test_step(self, batch, batch_idx, dataloader_idx: int = 0):
+        audio_signal, audio_signal_len, labels, _ = batch
+        logits, _ = self.forward(input_signal=audio_signal, input_signal_length=audio_signal_len)
+        loss_value = self.loss(logits=logits, labels=labels)
+        acc_top_k = self._accuracy(logits=logits, labels=labels)
+        correct_counts, total_counts = self._accuracy.correct_counts_k, self._accuracy.total_counts_k
+
+        return {
+            'test_loss': loss_value,
+            'test_correct_counts': correct_counts,
+            'test_total_counts': total_counts,
+            'test_acc_top_k': acc_top_k,
+        }
+
+    def multi_test_epoch_end(self, outputs, dataloader_idx: int = 0):
+        test_loss_mean = torch.stack([x['test_loss'] for x in outputs]).mean()
+        correct_counts = torch.stack([x['test_correct_counts'] for x in outputs]).sum(axis=0)
+        total_counts = torch.stack([x['test_total_counts'] for x in outputs]).sum(axis=0)
+
+        self._accuracy.correct_counts_k = correct_counts
+        self._accuracy.total_counts_k = total_counts
+        topk_scores = self._accuracy.compute()
+        self._accuracy.reset()
+
+        logging.info("test_loss: {:.3f}".format(test_loss_mean))
+        self.log('test_loss', test_loss_mean)
+        for top_k, score in zip(self._accuracy.top_k, topk_scores):
+            self.log('test_epoch_accuracy_top@{}'.format(top_k), score)
+
+        return {
+            'test_loss': test_loss_mean,
+            'test_acc_top_k': topk_scores,
+        }
+
+    def setup_finetune_model(self, model_config: DictConfig):
+        """
+        setup_finetune_model method sets up training data, validation data and test data with new
+        provided config, this checks for the previous labels set up during training from scratch, if None,
+        it sets up labels for provided finetune data from manifest files
+        Args:
+        model_config: cfg which has train_ds, optional validation_ds, optional test_ds and
+        mandatory encoder and decoder model params
+        make sure you set num_classes correctly for finetune data
+        Returns: None
+        """
+        if hasattr(self, 'dataset'):
+            scratch_labels = self.dataset.labels
+        else:
+            scratch_labels = None
+
+        logging.info("Setting up data loaders with manifests provided from model_config")
+
+        if 'train_ds' in model_config and model_config.train_ds is not None:
+            self.setup_training_data(model_config.train_ds)
+        else:
+            raise KeyError("train_ds is not found in model_config but you need it for fine tuning")
+
+        if self.dataset.labels is None or len(self.dataset.labels) == 0:
+            raise ValueError(f'New labels must be non-empty list of labels. But I got: {self.dataset.labels}')
+
+        if 'validation_ds' in model_config and model_config.validation_ds is not None:
+            self.setup_multiple_validation_data(model_config.validation_ds)
+
+        if 'test_ds' in model_config and model_config.test_ds is not None:
+            self.setup_multiple_test_data(model_config.test_ds)
+
+        if scratch_labels == self.dataset.labels:  # checking for new finetune dataset labels
+            logging.warning(
+                "Trained dataset labels are same as finetune dataset labels -- continuing change of decoder parameters"
+            )
+        elif scratch_labels is None:
+            logging.warning(
+                "Either you provided a dummy manifest file during training from scratch or you restored from a pretrained nemo file"
+            )
+
+        decoder_config = model_config.decoder
+        new_decoder_config = copy.deepcopy(decoder_config)
+        if new_decoder_config['num_classes'] != len(self.dataset.labels):
+            raise ValueError(
+                "number of classes provided {} is not same as number of different labels in finetuning data: {}".format(
+                    new_decoder_config['num_classes'], len(self.dataset.labels)
+                )
+            )
+
+        del self.decoder
+        self.decoder = EncDecSpeakerLabelModel.from_config_dict(new_decoder_config)
+
+        with open_dict(self._cfg.decoder):
+            self._cfg.decoder = new_decoder_config
+
+        logging.info(f"Changed decoder output to # {self.decoder._num_classes} classes.")
+
+    @torch.no_grad()
+    def get_embedding(self, path2audio_file):
+        audio, sr = librosa.load(path2audio_file, sr=None)
+        target_sr = self._cfg.train_ds.get('sample_rate', 16000)
+        if sr != target_sr:
+            audio = librosa.core.resample(audio, sr, target_sr)
+        audio_length = audio.shape[0]
+        device = self.device
+        audio_signal, audio_signal_len = (
+            torch.tensor([audio], device=device),
+            torch.tensor([audio_length], device=device),
+        )
+        mode = self.training
+        self.freeze()
+
+        _, embs = self.forward(input_signal=audio_signal, input_signal_length=audio_signal_len)
+
+        self.train(mode=mode)
+        if mode is True:
+            self.unfreeze()
+        del audio_signal, audio_signal_len
+        return embs
+
+
+class ExtractSpeakerEmbeddingsModel(EncDecSpeakerLabelModel):
+    """
+    This Model class facilitates extraction of speaker embeddings from a pretrained model.
+    Respective embedding file is saved in self.embedding dir passed through cfg
+    """
+
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+        super().__init__(cfg=cfg, trainer=trainer)
+
+    def test_step(self, batch, batch_ix):
+        audio_signal, audio_signal_len, labels, slices = batch
+        _, embs = self.forward(input_signal=audio_signal, input_signal_length=audio_signal_len)
+        return {'embs': embs, 'labels': labels, 'slices': slices}
+
+    def test_epoch_end(self, outputs):
+        embs = torch.cat([x['embs'] for x in outputs])
+        slices = torch.cat([x['slices'] for x in outputs])
+        emb_shape = embs.shape[-1]
+        embs = embs.view(-1, emb_shape).cpu().numpy()
+        out_embeddings = {}
+        start_idx = 0
+        with open(self.test_manifest, 'r') as manifest:
+            for idx, line in enumerate(manifest.readlines()):
+                line = line.strip()
+                dic = json.loads(line)
+                structure = dic['audio_filepath'].split('/')[-3:]
+                uniq_name = '@'.join(structure)
+                if uniq_name in out_embeddings:
+                    raise KeyError("Embeddings for label {} already present in emb dictionary".format(uniq_name))
+                num_slices = slices[idx]
+                end_idx = start_idx + num_slices
+                out_embeddings[uniq_name] = embs[start_idx:end_idx].mean(axis=0)
+                start_idx = end_idx
+
+        embedding_dir = os.path.join(self.embedding_dir, 'embeddings')
+        if not os.path.exists(embedding_dir):
+            os.mkdir(embedding_dir)
+
+        prefix = self.test_manifest.split('/')[-1].split('.')[-2]
+
+        name = os.path.join(embedding_dir, prefix)
+        pkl.dump(out_embeddings, open(name + '_embeddings.pkl', 'wb'))
+        logging.info("Saved embedding files to {}".format(embedding_dir))
+
+        return {}

tts/nemo/collections/asr/models/rnnt_bpe_models.py

@@ -0,0 +1,290 @@
+# 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 copy
+import os
+from typing import Dict, Optional
+
+import torch
+from omegaconf import DictConfig, ListConfig, OmegaConf, open_dict
+from pytorch_lightning import Trainer
+
+from nemo.collections.asr.data import audio_to_text_dataset
+from nemo.collections.asr.losses.rnnt import RNNTLoss
+from nemo.collections.asr.metrics.rnnt_wer_bpe import RNNTBPEWER, RNNTBPEDecoding
+from nemo.collections.asr.models.rnnt_models import EncDecRNNTModel
+from nemo.collections.asr.parts.mixins import ASRBPEMixin
+from nemo.collections.asr.parts.preprocessing.perturb import process_augmentations
+from nemo.core.classes.common import PretrainedModelInfo
+from nemo.utils import logging, model_utils
+
+
+class EncDecRNNTBPEModel(EncDecRNNTModel, ASRBPEMixin):
+    """Base class for encoder decoder RNNT-based models with subword tokenization."""
+
+    @classmethod
+    def list_available_models(cls) -> Optional[PretrainedModelInfo]:
+        """
+        This method returns a list of pre-trained model which can be instantiated directly from NVIDIA's NGC cloud.
+
+        Returns:
+            List of available pre-trained models.
+        """
+        result = []
+        return result
+
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+        # Convert to Hydra 1.0 compatible DictConfig
+        cfg = model_utils.convert_model_config_to_dict_config(cfg)
+        cfg = model_utils.maybe_update_config_version(cfg)
+
+        # Tokenizer is necessary for this model
+        if 'tokenizer' not in cfg:
+            raise ValueError("`cfg` must have `tokenizer` config to create a tokenizer !")
+
+        if not isinstance(cfg, DictConfig):
+            cfg = OmegaConf.create(cfg)
+
+        # Setup the tokenizer
+        self._setup_tokenizer(cfg.tokenizer)
+
+        # Initialize a dummy vocabulary
+        vocabulary = self.tokenizer.tokenizer.get_vocab()
+
+        # Set the new vocabulary
+        with open_dict(cfg):
+            cfg.labels = ListConfig(list(vocabulary))
+
+        with open_dict(cfg.decoder):
+            cfg.decoder.vocab_size = len(vocabulary)
+
+        with open_dict(cfg.joint):
+            cfg.joint.num_classes = len(vocabulary)
+            cfg.joint.vocabulary = ListConfig(list(vocabulary))
+            cfg.joint.jointnet.encoder_hidden = cfg.model_defaults.enc_hidden
+            cfg.joint.jointnet.pred_hidden = cfg.model_defaults.pred_hidden
+
+        super().__init__(cfg=cfg, trainer=trainer)
+
+        # Setup decoding object
+        self.decoding = RNNTBPEDecoding(
+            decoding_cfg=self.cfg.decoding, decoder=self.decoder, joint=self.joint, tokenizer=self.tokenizer,
+        )
+
+        # Setup wer object
+        self.wer = RNNTBPEWER(
+            decoding=self.decoding,
+            batch_dim_index=0,
+            use_cer=self._cfg.get('use_cer', False),
+            log_prediction=self._cfg.get('log_prediction', True),
+            dist_sync_on_step=True,
+        )
+
+        # Setup fused Joint step if flag is set
+        if self.joint.fuse_loss_wer:
+            self.joint.set_loss(self.loss)
+            self.joint.set_wer(self.wer)
+
+    def change_vocabulary(
+        self, new_tokenizer_dir: str, new_tokenizer_type: str, decoding_cfg: Optional[DictConfig] = None
+    ):
+        """
+        Changes vocabulary used during RNNT decoding process. Use this method when fine-tuning on from pre-trained model.
+        This method changes only decoder and leaves encoder and pre-processing modules unchanged. For example, you would
+        use it if you want to use pretrained encoder when fine-tuning on data in another language, or when you'd need
+        model to learn capitalization, punctuation and/or special characters.
+
+        Args:
+            new_tokenizer_dir: Directory path to tokenizer.
+            new_tokenizer_type: Type of tokenizer. Can be either `bpe` or `wpe`.
+            decoding_cfg: A config for the decoder, which is optional. If the decoding type
+                needs to be changed (from say Greedy to Beam decoding etc), the config can be passed here.
+
+        Returns: None
+
+        """
+        if not os.path.isdir(new_tokenizer_dir):
+            raise NotADirectoryError(
+                f'New tokenizer dir must be non-empty path to a directory. But I got: {new_tokenizer_dir}'
+            )
+
+        if new_tokenizer_type.lower() not in ('bpe', 'wpe'):
+            raise ValueError(f'New tokenizer type must be either `bpe` or `wpe`')
+
+        tokenizer_cfg = OmegaConf.create({'dir': new_tokenizer_dir, 'type': new_tokenizer_type})
+
+        # Setup the tokenizer
+        self._setup_tokenizer(tokenizer_cfg)
+
+        # Initialize a dummy vocabulary
+        vocabulary = self.tokenizer.tokenizer.get_vocab()
+
+        joint_config = self.joint.to_config_dict()
+        new_joint_config = copy.deepcopy(joint_config)
+        new_joint_config['vocabulary'] = ListConfig(list(vocabulary.keys()))
+        new_joint_config['num_classes'] = len(vocabulary)
+        del self.joint
+        self.joint = EncDecRNNTBPEModel.from_config_dict(new_joint_config)
+
+        decoder_config = self.decoder.to_config_dict()
+        new_decoder_config = copy.deepcopy(decoder_config)
+        new_decoder_config.vocab_size = len(vocabulary)
+        del self.decoder
+        self.decoder = EncDecRNNTBPEModel.from_config_dict(new_decoder_config)
+
+        del self.loss
+        self.loss = RNNTLoss(num_classes=self.joint.num_classes_with_blank - 1)
+
+        if decoding_cfg is None:
+            # Assume same decoding config as before
+            decoding_cfg = self.cfg.decoding
+
+        self.decoding = RNNTBPEDecoding(
+            decoding_cfg=decoding_cfg, decoder=self.decoder, joint=self.joint, tokenizer=self.tokenizer,
+        )
+
+        self.wer = RNNTBPEWER(
+            decoding=self.decoding,
+            batch_dim_index=self.wer.batch_dim_index,
+            use_cer=self.wer.use_cer,
+            log_prediction=self.wer.log_prediction,
+            dist_sync_on_step=True,
+        )
+
+        # Setup fused Joint step
+        if self.joint.fuse_loss_wer:
+            self.joint.set_loss(self.loss)
+            self.joint.set_wer(self.wer)
+
+        # Update config
+        with open_dict(self.cfg.joint):
+            self.cfg.joint = new_joint_config
+
+        with open_dict(self.cfg.decoder):
+            self.cfg.decoder = new_decoder_config
+
+        with open_dict(self.cfg.decoding):
+            self.cfg.decoding = decoding_cfg
+
+        logging.info(f"Changed decoder to output to {self.joint.vocabulary} vocabulary.")
+
+    def change_decoding_strategy(self, decoding_cfg: DictConfig):
+        """
+        Changes decoding strategy used during RNNT decoding process.
+
+        Args:
+            decoding_cfg: A config for the decoder, which is optional. If the decoding type
+                needs to be changed (from say Greedy to Beam decoding etc), the config can be passed here.
+        """
+        if decoding_cfg is None:
+            # Assume same decoding config as before
+            logging.info("No `decoding_cfg` passed when changing decoding strategy, using internal config")
+            decoding_cfg = self.cfg.decoding
+
+        self.decoding = RNNTBPEDecoding(
+            decoding_cfg=decoding_cfg, decoder=self.decoder, joint=self.joint, tokenizer=self.tokenizer,
+        )
+
+        self.wer = RNNTBPEWER(
+            decoding=self.decoding,
+            batch_dim_index=self.wer.batch_dim_index,
+            use_cer=self.wer.use_cer,
+            log_prediction=self.wer.log_prediction,
+            dist_sync_on_step=True,
+        )
+
+        # Setup fused Joint step
+        if self.joint.fuse_loss_wer:
+            self.joint.set_loss(self.loss)
+            self.joint.set_wer(self.wer)
+
+        # Update config
+        with open_dict(self.cfg.decoding):
+            self.cfg.decoding = decoding_cfg
+
+        logging.info(f"Changed decoding strategy to \n{OmegaConf.to_yaml(self.cfg.decoding)}")
+
+    def _setup_dataloader_from_config(self, config: Optional[Dict]):
+        if 'augmentor' in config:
+            augmentor = process_augmentations(config['augmentor'])
+        else:
+            augmentor = None
+
+        shuffle = config['shuffle']
+
+        # Instantiate tarred dataset loader or normal dataset loader
+        if config.get('is_tarred', False):
+            if ('tarred_audio_filepaths' in config and config['tarred_audio_filepaths'] is None) or (
+                'manifest_filepath' in config and config['manifest_filepath'] is None
+            ):
+                logging.warning(
+                    "Could not load dataset as `manifest_filepath` was None or "
+                    f"`tarred_audio_filepaths` is None. Provided config : {config}"
+                )
+                return None
+
+            shuffle_n = config.get('shuffle_n', 4 * config['batch_size']) if shuffle else 0
+            dataset = audio_to_text_dataset.get_tarred_bpe_dataset(
+                config=config,
+                tokenizer=self.tokenizer,
+                shuffle_n=shuffle_n,
+                global_rank=self.global_rank,
+                world_size=self.world_size,
+                augmentor=augmentor,
+            )
+            shuffle = False
+        else:
+            if 'manifest_filepath' in config and config['manifest_filepath'] is None:
+                logging.warning(f"Could not load dataset as `manifest_filepath` was None. Provided config : {config}")
+                return None
+
+            dataset = audio_to_text_dataset.get_bpe_dataset(
+                config=config, tokenizer=self.tokenizer, augmentor=augmentor
+            )
+
+        return torch.utils.data.DataLoader(
+            dataset=dataset,
+            batch_size=config['batch_size'],
+            collate_fn=dataset.collate_fn,
+            drop_last=config.get('drop_last', False),
+            shuffle=shuffle,
+            num_workers=config.get('num_workers', 0),
+            pin_memory=config.get('pin_memory', False),
+        )
+
+    def _setup_transcribe_dataloader(self, config: Dict) -> 'torch.utils.data.DataLoader':
+        """
+        Setup function for a temporary data loader which wraps the provided audio file.
+
+        Args:
+            config: A python dictionary which contains the following keys:
+            paths2audio_files: (a list) of paths to audio files. The files should be relatively short fragments. \
+                Recommended length per file is between 5 and 25 seconds.
+            batch_size: (int) batch size to use during inference. \
+                Bigger will result in better throughput performance but would use more memory.
+            temp_dir: (str) A temporary directory where the audio manifest is temporarily
+                stored.
+
+        Returns:
+            A pytorch DataLoader for the given audio file(s).
+        """
+        dl_config = {
+            'manifest_filepath': os.path.join(config['temp_dir'], 'manifest.json'),
+            'sample_rate': self.preprocessor._sample_rate,
+            'batch_size': min(config['batch_size'], len(config['paths2audio_files'])),
+            'shuffle': False,
+        }
+
+        temporary_datalayer = self._setup_dataloader_from_config(config=DictConfig(dl_config))
+        return temporary_datalayer

tts/nemo/collections/asr/models/rnnt_models.py

@@ -0,0 +1,835 @@
+# 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 copy
+import json
+import os
+import tempfile
+from math import ceil
+from typing import Dict, List, Optional, Union
+
+import torch
+from omegaconf import DictConfig, OmegaConf, open_dict
+from pytorch_lightning import Trainer
+from tqdm.auto import tqdm
+
+from nemo.collections.asr.data import audio_to_text_dataset
+from nemo.collections.asr.data.audio_to_text_dali import DALIOutputs
+from nemo.collections.asr.losses.rnnt import RNNTLoss, resolve_rnnt_default_loss_name
+from nemo.collections.asr.metrics.rnnt_wer import RNNTWER, RNNTDecoding
+from nemo.collections.asr.models.asr_model import ASRModel
+from nemo.collections.asr.parts.mixins import ASRModuleMixin
+from nemo.collections.asr.parts.preprocessing.perturb import process_augmentations
+from nemo.core.classes.common import PretrainedModelInfo, typecheck
+from nemo.core.neural_types import AcousticEncodedRepresentation, AudioSignal, LengthsType, NeuralType, SpectrogramType
+from nemo.utils import logging
+
+
+class EncDecRNNTModel(ASRModel, ASRModuleMixin):
+    """Base class for encoder decoder RNNT-based models."""
+
+    @classmethod
+    def list_available_models(cls) -> Optional[PretrainedModelInfo]:
+        """
+        This method returns a list of pre-trained model which can be instantiated directly from NVIDIA's NGC cloud.
+
+        Returns:
+            List of available pre-trained models.
+        """
+        result = []
+        return result
+
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+        # Get global rank and total number of GPU workers for IterableDataset partitioning, if applicable
+        # Global_rank and local_rank is set by LightningModule in Lightning 1.2.0
+        self.world_size = 1
+        if trainer is not None:
+            self.world_size = trainer.num_nodes * trainer.num_gpus
+
+        super().__init__(cfg=cfg, trainer=trainer)
+
+        # Initialize components
+        self.preprocessor = EncDecRNNTModel.from_config_dict(self.cfg.preprocessor)
+        self.encoder = EncDecRNNTModel.from_config_dict(self.cfg.encoder)
+
+        # Update config values required by components dynamically
+        with open_dict(self.cfg.decoder):
+            self.cfg.decoder.vocab_size = len(self.cfg.labels)
+
+        with open_dict(self.cfg.joint):
+            self.cfg.joint.num_classes = len(self.cfg.labels)
+            self.cfg.joint.vocabulary = self.cfg.labels
+            self.cfg.joint.jointnet.encoder_hidden = self.cfg.model_defaults.enc_hidden
+            self.cfg.joint.jointnet.pred_hidden = self.cfg.model_defaults.pred_hidden
+
+        self.decoder = EncDecRNNTModel.from_config_dict(self.cfg.decoder)
+        self.joint = EncDecRNNTModel.from_config_dict(self.cfg.joint)
+
+        # Setup RNNT Loss
+        loss_name, loss_kwargs = self.extract_rnnt_loss_cfg(self.cfg.get("loss", None))
+
+        self.loss = RNNTLoss(
+            num_classes=self.joint.num_classes_with_blank - 1, loss_name=loss_name, loss_kwargs=loss_kwargs
+        )
+
+        if hasattr(self.cfg, 'spec_augment') and self._cfg.spec_augment is not None:
+            self.spec_augmentation = EncDecRNNTModel.from_config_dict(self.cfg.spec_augment)
+        else:
+            self.spec_augmentation = None
+
+        # Setup decoding objects
+        self.decoding = RNNTDecoding(
+            decoding_cfg=self.cfg.decoding, decoder=self.decoder, joint=self.joint, vocabulary=self.joint.vocabulary,
+        )
+        # Setup WER calculation
+        self.wer = RNNTWER(
+            decoding=self.decoding,
+            batch_dim_index=0,
+            use_cer=self._cfg.get('use_cer', False),
+            log_prediction=self._cfg.get('log_prediction', True),
+            dist_sync_on_step=True,
+        )
+
+        # Whether to compute loss during evaluation
+        if 'compute_eval_loss' in self.cfg:
+            self.compute_eval_loss = self.cfg.compute_eval_loss
+        else:
+            self.compute_eval_loss = True
+
+        # Setup fused Joint step if flag is set
+        if self.joint.fuse_loss_wer:
+            self.joint.set_loss(self.loss)
+            self.joint.set_wer(self.wer)
+
+        self.setup_optim_normalization()
+
+    def setup_optim_normalization(self):
+        """
+        Helper method to setup normalization of certain parts of the model prior to the optimization step.
+
+        Supported pre-optimization normalizations are as follows:
+
+        .. code-block:: yaml
+
+            # Variation Noise injection
+            model:
+                variational_noise:
+                    std: 0.0
+                    start_step: 0
+
+            # Joint - Length normalization
+            model:
+                normalize_joint_txu: false
+
+            # Encoder Network - gradient normalization
+            model:
+                normalize_encoder_norm: false
+
+            # Decoder / Prediction Network - gradient normalization
+            model:
+                normalize_decoder_norm: false
+
+            # Joint - gradient normalization
+            model:
+                normalize_joint_norm: false
+        """
+        # setting up the variational noise for the decoder
+        if hasattr(self.cfg, 'variational_noise'):
+            self._optim_variational_noise_std = self.cfg['variational_noise'].get('std', 0)
+            self._optim_variational_noise_start = self.cfg['variational_noise'].get('start_step', 0)
+        else:
+            self._optim_variational_noise_std = 0
+            self._optim_variational_noise_start = 0
+
+        # Setup normalized gradients for model joint by T x U scaling factor (joint length normalization)
+        self._optim_normalize_joint_txu = self.cfg.get('normalize_joint_txu', False)
+        self._optim_normalize_txu = None
+
+        # Setup normalized encoder norm for model
+        self._optim_normalize_encoder_norm = self.cfg.get('normalize_encoder_norm', False)
+
+        # Setup normalized decoder norm for model
+        self._optim_normalize_decoder_norm = self.cfg.get('normalize_decoder_norm', False)
+
+        # Setup normalized joint norm for model
+        self._optim_normalize_joint_norm = self.cfg.get('normalize_joint_norm', False)
+
+    def extract_rnnt_loss_cfg(self, cfg: Optional[DictConfig]):
+        """
+        Helper method to extract the rnnt loss name, and potentially its kwargs
+        to be passed.
+
+        Args:
+            cfg: Should contain `loss_name` as a string which is resolved to a RNNT loss name.
+                If the default should be used, then `default` can be used.
+                Optionally, one can pass additional kwargs to the loss function. The subdict
+                should have a keyname as follows : `{loss_name}_kwargs`.
+
+                Note that whichever loss_name is selected, that corresponding kwargs will be
+                selected. For the "default" case, the "{resolved_default}_kwargs" will be used.
+
+        Examples:
+            .. code-block:: yaml
+
+                loss_name: "default"
+                warprnnt_numba_kwargs:
+                    kwargs2: some_other_val
+
+        Returns:
+            A tuple, the resolved loss name as well as its kwargs (if found).
+        """
+        if cfg is None:
+            cfg = DictConfig({})
+
+        loss_name = cfg.get("loss_name", "default")
+
+        if loss_name == "default":
+            loss_name = resolve_rnnt_default_loss_name()
+
+        loss_kwargs = cfg.get(f"{loss_name}_kwargs", None)
+
+        logging.info(f"Using RNNT Loss : {loss_name}\n" f"Loss {loss_name}_kwargs: {loss_kwargs}")
+
+        return loss_name, loss_kwargs
+
+    @torch.no_grad()
+    def transcribe(
+        self, paths2audio_files: List[str], batch_size: int = 4, return_hypotheses: bool = False
+    ) -> List[str]:
+        """
+        Uses greedy decoding to transcribe audio files. Use this method for debugging and prototyping.
+
+        Args:
+
+            paths2audio_files: (a list) of paths to audio files. \
+        Recommended length per file is between 5 and 25 seconds. \
+        But it is possible to pass a few hours long file if enough GPU memory is available.
+            batch_size: (int) batch size to use during inference. \
+        Bigger will result in better throughput performance but would use more memory.
+            return_hypotheses: (bool) Either return hypotheses or text
+        With hypotheses can do some postprocessing like getting timestamp or rescoring
+        Returns:
+
+            A list of transcriptions in the same order as paths2audio_files
+        """
+        if paths2audio_files is None or len(paths2audio_files) == 0:
+            return {}
+        # We will store transcriptions here
+        hypotheses = []
+        # Model's mode and device
+        mode = self.training
+        device = next(self.parameters()).device
+        try:
+            # Switch model to evaluation mode
+            self.eval()
+            # Freeze the encoder and decoder modules
+            self.encoder.freeze()
+            self.decoder.freeze()
+            self.joint.freeze()
+            logging_level = logging.get_verbosity()
+            logging.set_verbosity(logging.WARNING)
+            # Work in tmp directory - will store manifest file there
+            with tempfile.TemporaryDirectory() as tmpdir:
+                with open(os.path.join(tmpdir, 'manifest.json'), 'w') as fp:
+                    for audio_file in paths2audio_files:
+                        entry = {'audio_filepath': audio_file, 'duration': 100000, 'text': 'nothing'}
+                        fp.write(json.dumps(entry) + '\n')
+
+                config = {'paths2audio_files': paths2audio_files, 'batch_size': batch_size, 'temp_dir': tmpdir}
+
+                temporary_datalayer = self._setup_transcribe_dataloader(config)
+                for test_batch in tqdm(temporary_datalayer, desc="Transcribing"):
+                    encoded, encoded_len = self.forward(
+                        input_signal=test_batch[0].to(device), input_signal_length=test_batch[1].to(device)
+                    )
+                    hypotheses += self.decoding.rnnt_decoder_predictions_tensor(
+                        encoded, encoded_len, return_hypotheses=return_hypotheses
+                    )
+                    del encoded
+                    del test_batch
+        finally:
+            # set mode back to its original value
+            self.train(mode=mode)
+            logging.set_verbosity(logging_level)
+            if mode is True:
+                self.encoder.unfreeze()
+                self.decoder.unfreeze()
+                self.joint.unfreeze()
+        return hypotheses
+
+    def change_vocabulary(self, new_vocabulary: List[str], decoding_cfg: Optional[DictConfig] = None):
+        """
+        Changes vocabulary used during RNNT decoding process. Use this method when fine-tuning a pre-trained model.
+        This method changes only decoder and leaves encoder and pre-processing modules unchanged. For example, you would
+        use it if you want to use pretrained encoder when fine-tuning on data in another language, or when you'd need
+        model to learn capitalization, punctuation and/or special characters.
+
+        Args:
+            new_vocabulary: list with new vocabulary. Must contain at least 2 elements. Typically, \
+                this is target alphabet.
+            decoding_cfg: A config for the decoder, which is optional. If the decoding type
+                needs to be changed (from say Greedy to Beam decoding etc), the config can be passed here.
+
+        Returns: None
+
+        """
+        if self.joint.vocabulary == new_vocabulary:
+            logging.warning(f"Old {self.joint.vocabulary} and new {new_vocabulary} match. Not changing anything.")
+        else:
+            if new_vocabulary is None or len(new_vocabulary) == 0:
+                raise ValueError(f'New vocabulary must be non-empty list of chars. But I got: {new_vocabulary}')
+
+            joint_config = self.joint.to_config_dict()
+            new_joint_config = copy.deepcopy(joint_config)
+            new_joint_config['vocabulary'] = new_vocabulary
+            new_joint_config['num_classes'] = len(new_vocabulary)
+            del self.joint
+            self.joint = EncDecRNNTModel.from_config_dict(new_joint_config)
+
+            decoder_config = self.decoder.to_config_dict()
+            new_decoder_config = copy.deepcopy(decoder_config)
+            new_decoder_config.vocab_size = len(new_vocabulary)
+            del self.decoder
+            self.decoder = EncDecRNNTModel.from_config_dict(new_decoder_config)
+
+            del self.loss
+            loss_name, loss_kwargs = self.extract_rnnt_loss_cfg(self.cfg.get('loss', None))
+            self.loss = RNNTLoss(
+                num_classes=self.joint.num_classes_with_blank - 1, loss_name=loss_name, loss_kwargs=loss_kwargs
+            )
+
+            if decoding_cfg is None:
+                # Assume same decoding config as before
+                decoding_cfg = self.cfg.decoding
+
+            self.decoding = RNNTDecoding(
+                decoding_cfg=decoding_cfg, decoder=self.decoder, joint=self.joint, vocabulary=self.joint.vocabulary,
+            )
+
+            self.wer = RNNTWER(
+                decoding=self.decoding,
+                batch_dim_index=self.wer.batch_dim_index,
+                use_cer=self.wer.use_cer,
+                log_prediction=self.wer.log_prediction,
+                dist_sync_on_step=True,
+            )
+
+            # Setup fused Joint step
+            if self.joint.fuse_loss_wer:
+                self.joint.set_loss(self.loss)
+                self.joint.set_wer(self.wer)
+
+            # Update config
+            with open_dict(self.cfg.joint):
+                self.cfg.joint = new_joint_config
+
+            with open_dict(self.cfg.decoder):
+                self.cfg.decoder = new_decoder_config
+
+            with open_dict(self.cfg.decoding):
+                self.cfg.decoding = decoding_cfg
+
+            ds_keys = ['train_ds', 'validation_ds', 'test_ds']
+            for key in ds_keys:
+                if key in self.cfg:
+                    with open_dict(self.cfg[key]):
+                        self.cfg[key]['labels'] = OmegaConf.create(new_vocabulary)
+
+            logging.info(f"Changed decoder to output to {self.joint.vocabulary} vocabulary.")
+
+    def change_decoding_strategy(self, decoding_cfg: DictConfig):
+        """
+        Changes decoding strategy used during RNNT decoding process.
+
+        Args:
+            decoding_cfg: A config for the decoder, which is optional. If the decoding type
+                needs to be changed (from say Greedy to Beam decoding etc), the config can be passed here.
+        """
+        if decoding_cfg is None:
+            # Assume same decoding config as before
+            logging.info("No `decoding_cfg` passed when changing decoding strategy, using internal config")
+            decoding_cfg = self.cfg.decoding
+
+        self.decoding = RNNTDecoding(
+            decoding_cfg=decoding_cfg, decoder=self.decoder, joint=self.joint, vocabulary=self.joint.vocabulary,
+        )
+
+        self.wer = RNNTWER(
+            decoding=self.decoding,
+            batch_dim_index=self.wer.batch_dim_index,
+            use_cer=self.wer.use_cer,
+            log_prediction=self.wer.log_prediction,
+            dist_sync_on_step=True,
+        )
+
+        # Setup fused Joint step
+        if self.joint.fuse_loss_wer:
+            self.joint.set_loss(self.loss)
+            self.joint.set_wer(self.wer)
+
+        # Update config
+        with open_dict(self.cfg.decoding):
+            self.cfg.decoding = decoding_cfg
+
+        logging.info(f"Changed decoding strategy to \n{OmegaConf.to_yaml(self.cfg.decoding)}")
+
+    def _setup_dataloader_from_config(self, config: Optional[Dict]):
+        if 'augmentor' in config:
+            augmentor = process_augmentations(config['augmentor'])
+        else:
+            augmentor = None
+
+        # Automatically inject args from model config to dataloader config
+        audio_to_text_dataset.inject_dataloader_value_from_model_config(self.cfg, config, key='sample_rate')
+        audio_to_text_dataset.inject_dataloader_value_from_model_config(self.cfg, config, key='labels')
+
+        shuffle = config['shuffle']
+        device = 'gpu' if torch.cuda.is_available() else 'cpu'
+        if config.get('use_dali', False):
+            device_id = self.local_rank if device == 'gpu' else None
+            dataset = audio_to_text_dataset.get_dali_char_dataset(
+                config=config,
+                shuffle=shuffle,
+                device_id=device_id,
+                global_rank=self.global_rank,
+                world_size=self.world_size,
+                preprocessor_cfg=self._cfg.preprocessor,
+            )
+            return dataset
+
+        # Instantiate tarred dataset loader or normal dataset loader
+        if config.get('is_tarred', False):
+            if ('tarred_audio_filepaths' in config and config['tarred_audio_filepaths'] is None) or (
+                'manifest_filepath' in config and config['manifest_filepath'] is None
+            ):
+                logging.warning(
+                    "Could not load dataset as `manifest_filepath` was None or "
+                    f"`tarred_audio_filepaths` is None. Provided config : {config}"
+                )
+                return None
+
+            shuffle_n = config.get('shuffle_n', 4 * config['batch_size']) if shuffle else 0
+            dataset = audio_to_text_dataset.get_tarred_char_dataset(
+                config=config,
+                shuffle_n=shuffle_n,
+                global_rank=self.global_rank,
+                world_size=self.world_size,
+                augmentor=augmentor,
+            )
+            shuffle = False
+        else:
+            if 'manifest_filepath' in config and config['manifest_filepath'] is None:
+                logging.warning(f"Could not load dataset as `manifest_filepath` was None. Provided config : {config}")
+                return None
+
+            dataset = audio_to_text_dataset.get_char_dataset(config=config, augmentor=augmentor)
+
+        return torch.utils.data.DataLoader(
+            dataset=dataset,
+            batch_size=config['batch_size'],
+            collate_fn=dataset.collate_fn,
+            drop_last=config.get('drop_last', False),
+            shuffle=shuffle,
+            num_workers=config.get('num_workers', 0),
+            pin_memory=config.get('pin_memory', False),
+        )
+
+    def setup_training_data(self, train_data_config: Optional[Union[DictConfig, Dict]]):
+        """
+        Sets up the training data loader via a Dict-like object.
+
+        Args:
+            train_data_config: A config that contains the information regarding construction
+                of an ASR Training dataset.
+
+        Supported Datasets:
+            -   :class:`~nemo.collections.asr.data.audio_to_text.AudioToCharDataset`
+            -   :class:`~nemo.collections.asr.data.audio_to_text.AudioToBPEDataset`
+            -   :class:`~nemo.collections.asr.data.audio_to_text.TarredAudioToCharDataset`
+            -   :class:`~nemo.collections.asr.data.audio_to_text.TarredAudioToBPEDataset`
+            -   :class:`~nemo.collections.asr.data.audio_to_text_dali.AudioToCharDALIDataset`
+        """
+        if 'shuffle' not in train_data_config:
+            train_data_config['shuffle'] = True
+
+        # preserve config
+        self._update_dataset_config(dataset_name='train', config=train_data_config)
+
+        self._train_dl = self._setup_dataloader_from_config(config=train_data_config)
+
+        # Need to set this because if using an IterableDataset, the length of the dataloader is the total number
+        # of samples rather than the number of batches, and this messes up the tqdm progress bar.
+        # So we set the number of steps manually (to the correct number) to fix this.
+        if 'is_tarred' in train_data_config and train_data_config['is_tarred']:
+            # We also need to check if limit_train_batches is already set.
+            # If it's an int, we assume that the user has set it to something sane, i.e. <= # training batches,
+            # and don't change it. Otherwise, adjust batches accordingly if it's a float (including 1.0).
+            if isinstance(self._trainer.limit_train_batches, float):
+                self._trainer.limit_train_batches = int(
+                    self._trainer.limit_train_batches
+                    * ceil((len(self._train_dl.dataset) / self.world_size) / train_data_config['batch_size'])
+                )
+
+    def setup_validation_data(self, val_data_config: Optional[Union[DictConfig, Dict]]):
+        """
+        Sets up the validation data loader via a Dict-like object.
+
+        Args:
+            val_data_config: A config that contains the information regarding construction
+                of an ASR Training dataset.
+
+        Supported Datasets:
+            -   :class:`~nemo.collections.asr.data.audio_to_text.AudioToCharDataset`
+            -   :class:`~nemo.collections.asr.data.audio_to_text.AudioToBPEDataset`
+            -   :class:`~nemo.collections.asr.data.audio_to_text.TarredAudioToCharDataset`
+            -   :class:`~nemo.collections.asr.data.audio_to_text.TarredAudioToBPEDataset`
+            -   :class:`~nemo.collections.asr.data.audio_to_text_dali.AudioToCharDALIDataset`
+        """
+        if 'shuffle' not in val_data_config:
+            val_data_config['shuffle'] = False
+
+        # preserve config
+        self._update_dataset_config(dataset_name='validation', config=val_data_config)
+
+        self._validation_dl = self._setup_dataloader_from_config(config=val_data_config)
+
+    def setup_test_data(self, test_data_config: Optional[Union[DictConfig, Dict]]):
+        """
+        Sets up the test data loader via a Dict-like object.
+
+        Args:
+            test_data_config: A config that contains the information regarding construction
+                of an ASR Training dataset.
+
+        Supported Datasets:
+            -   :class:`~nemo.collections.asr.data.audio_to_text.AudioToCharDataset`
+            -   :class:`~nemo.collections.asr.data.audio_to_text.AudioToBPEDataset`
+            -   :class:`~nemo.collections.asr.data.audio_to_text.TarredAudioToCharDataset`
+            -   :class:`~nemo.collections.asr.data.audio_to_text.TarredAudioToBPEDataset`
+            -   :class:`~nemo.collections.asr.data.audio_to_text_dali.AudioToCharDALIDataset`
+        """
+        if 'shuffle' not in test_data_config:
+            test_data_config['shuffle'] = False
+
+        # preserve config
+        self._update_dataset_config(dataset_name='test', config=test_data_config)
+
+        self._test_dl = self._setup_dataloader_from_config(config=test_data_config)
+
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        if hasattr(self.preprocessor, '_sample_rate'):
+            input_signal_eltype = AudioSignal(freq=self.preprocessor._sample_rate)
+        else:
+            input_signal_eltype = AudioSignal()
+
+        return {
+            "input_signal": NeuralType(('B', 'T'), input_signal_eltype, optional=True),
+            "input_signal_length": NeuralType(tuple('B'), LengthsType(), optional=True),
+            "processed_signal": NeuralType(('B', 'D', 'T'), SpectrogramType(), optional=True),
+            "processed_signal_length": NeuralType(tuple('B'), LengthsType(), optional=True),
+        }
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {
+            "outputs": NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()),
+            "encoded_lengths": NeuralType(tuple('B'), LengthsType()),
+        }
+
+    @typecheck()
+    def forward(
+        self, input_signal=None, input_signal_length=None, processed_signal=None, processed_signal_length=None
+    ):
+        """
+        Forward pass of the model. Note that for RNNT Models, the forward pass of the model is a 3 step process,
+        and this method only performs the first step - forward of the acoustic model.
+
+        Please refer to the `training_step` in order to see the full `forward` step for training - which
+        performs the forward of the acoustic model, the prediction network and then the joint network.
+        Finally, it computes the loss and possibly compute the detokenized text via the `decoding` step.
+
+        Please refer to the `validation_step` in order to see the full `forward` step for inference - which
+        performs the forward of the acoustic model, the prediction network and then the joint network.
+        Finally, it computes the decoded tokens via the `decoding` step and possibly compute the batch metrics.
+
+        Args:
+            input_signal: Tensor that represents a batch of raw audio signals,
+                of shape [B, T]. T here represents timesteps, with 1 second of audio represented as
+                `self.sample_rate` number of floating point values.
+            input_signal_length: Vector of length B, that contains the individual lengths of the audio
+                sequences.
+            processed_signal: Tensor that represents a batch of processed audio signals,
+                of shape (B, D, T) that has undergone processing via some DALI preprocessor.
+            processed_signal_length: Vector of length B, that contains the individual lengths of the
+                processed audio sequences.
+
+        Returns:
+            A tuple of 2 elements -
+            1) The log probabilities tensor of shape [B, T, D].
+            2) The lengths of the acoustic sequence after propagation through the encoder, of shape [B].
+        """
+        has_input_signal = input_signal is not None and input_signal_length is not None
+        has_processed_signal = processed_signal is not None and processed_signal_length is not None
+        if (has_input_signal ^ has_processed_signal) is False:
+            raise ValueError(
+                f"{self} Arguments ``input_signal`` and ``input_signal_length`` are mutually exclusive "
+                " with ``processed_signal`` and ``processed_signal_len`` arguments."
+            )
+
+        if not has_processed_signal:
+            processed_signal, processed_signal_length = self.preprocessor(
+                input_signal=input_signal, length=input_signal_length,
+            )
+
+        # Spec augment is not applied during evaluation/testing
+        if self.spec_augmentation is not None and self.training:
+            processed_signal = self.spec_augmentation(input_spec=processed_signal, length=processed_signal_length)
+
+        encoded, encoded_len = self.encoder(audio_signal=processed_signal, length=processed_signal_length)
+        return encoded, encoded_len
+
+    # PTL-specific methods
+    def training_step(self, batch, batch_nb):
+        signal, signal_len, transcript, transcript_len = batch
+
+        # forward() only performs encoder forward
+        if isinstance(batch, DALIOutputs) and batch.has_processed_signal:
+            encoded, encoded_len = self.forward(processed_signal=signal, processed_signal_length=signal_len)
+        else:
+            encoded, encoded_len = self.forward(input_signal=signal, input_signal_length=signal_len)
+        del signal
+
+        # During training, loss must be computed, so decoder forward is necessary
+        decoder, target_length = self.decoder(targets=transcript, target_length=transcript_len)
+
+        if hasattr(self, '_trainer') and self._trainer is not None:
+            log_every_n_steps = self._trainer.log_every_n_steps
+            sample_id = self._trainer.global_step
+
+        else:
+            log_every_n_steps = 1
+            sample_id = batch_nb
+
+        # If experimental fused Joint-Loss-WER is not used
+        if not self.joint.fuse_loss_wer:
+            # Compute full joint and loss
+            joint = self.joint(encoder_outputs=encoded, decoder_outputs=decoder)
+            loss_value = self.loss(
+                log_probs=joint, targets=transcript, input_lengths=encoded_len, target_lengths=target_length
+            )
+
+            tensorboard_logs = {'train_loss': loss_value, 'learning_rate': self._optimizer.param_groups[0]['lr']}
+
+            if (sample_id + 1) % log_every_n_steps == 0:
+                self.wer.update(encoded, encoded_len, transcript, transcript_len)
+                _, scores, words = self.wer.compute()
+                self.wer.reset()
+                tensorboard_logs.update({'training_batch_wer': scores.float() / words})
+
+        else:
+            # If experimental fused Joint-Loss-WER is used
+            if (sample_id + 1) % log_every_n_steps == 0:
+                compute_wer = True
+            else:
+                compute_wer = False
+
+            # Fused joint step
+            loss_value, wer, _, _ = self.joint(
+                encoder_outputs=encoded,
+                decoder_outputs=decoder,
+                encoder_lengths=encoded_len,
+                transcripts=transcript,
+                transcript_lengths=transcript_len,
+                compute_wer=compute_wer,
+            )
+
+            tensorboard_logs = {'train_loss': loss_value, 'learning_rate': self._optimizer.param_groups[0]['lr']}
+
+            if compute_wer:
+                tensorboard_logs.update({'training_batch_wer': wer})
+
+        # Log items
+        self.log_dict(tensorboard_logs)
+
+        # Preserve batch acoustic model T and language model U parameters if normalizing
+        if self._optim_normalize_joint_txu:
+            self._optim_normalize_txu = [encoded_len.max(), transcript_len.max()]
+
+        return {'loss': loss_value}
+
+    def validation_step(self, batch, batch_idx, dataloader_idx=0):
+        signal, signal_len, transcript, transcript_len = batch
+
+        # forward() only performs encoder forward
+        if isinstance(batch, DALIOutputs) and batch.has_processed_signal:
+            encoded, encoded_len = self.forward(processed_signal=signal, processed_signal_length=signal_len)
+        else:
+            encoded, encoded_len = self.forward(input_signal=signal, input_signal_length=signal_len)
+        del signal
+
+        tensorboard_logs = {}
+
+        # If experimental fused Joint-Loss-WER is not used
+        if not self.joint.fuse_loss_wer:
+            if self.compute_eval_loss:
+                decoder, target_length = self.decoder(targets=transcript, target_length=transcript_len)
+                joint = self.joint(encoder_outputs=encoded, decoder_outputs=decoder)
+
+                loss_value = self.loss(
+                    log_probs=joint, targets=transcript, input_lengths=encoded_len, target_lengths=target_length
+                )
+
+                tensorboard_logs['val_loss'] = loss_value
+
+            self.wer.update(encoded, encoded_len, transcript, transcript_len)
+            wer, wer_num, wer_denom = self.wer.compute()
+            self.wer.reset()
+
+            tensorboard_logs['val_wer_num'] = wer_num
+            tensorboard_logs['val_wer_denom'] = wer_denom
+            tensorboard_logs['val_wer'] = wer
+
+        else:
+            # If experimental fused Joint-Loss-WER is used
+            compute_wer = True
+
+            if self.compute_eval_loss:
+                decoded, target_len = self.decoder(targets=transcript, target_length=transcript_len)
+            else:
+                decoded = None
+                target_len = transcript_len
+
+            # Fused joint step
+            loss_value, wer, wer_num, wer_denom = self.joint(
+                encoder_outputs=encoded,
+                decoder_outputs=decoded,
+                encoder_lengths=encoded_len,
+                transcripts=transcript,
+                transcript_lengths=target_len,
+                compute_wer=compute_wer,
+            )
+
+            if loss_value is not None:
+                tensorboard_logs['val_loss'] = loss_value
+
+            tensorboard_logs['val_wer_num'] = wer_num
+            tensorboard_logs['val_wer_denom'] = wer_denom
+            tensorboard_logs['val_wer'] = wer
+
+        return tensorboard_logs
+
+    def test_step(self, batch, batch_idx, dataloader_idx=0):
+        logs = self.validation_step(batch, batch_idx, dataloader_idx=dataloader_idx)
+        test_logs = {
+            'test_wer_num': logs['val_wer_num'],
+            'test_wer_denom': logs['val_wer_denom'],
+            # 'test_wer': logs['val_wer'],
+        }
+        if 'val_loss' in logs:
+            test_logs['test_loss'] = logs['val_loss']
+        return test_logs
+
+    def multi_validation_epoch_end(self, outputs, dataloader_idx: int = 0):
+        if self.compute_eval_loss:
+            val_loss_mean = torch.stack([x['val_loss'] for x in outputs]).mean()
+            val_loss_log = {'val_loss': val_loss_mean}
+        else:
+            val_loss_log = {}
+        wer_num = torch.stack([x['val_wer_num'] for x in outputs]).sum()
+        wer_denom = torch.stack([x['val_wer_denom'] for x in outputs]).sum()
+        tensorboard_logs = {**val_loss_log, 'val_wer': wer_num.float() / wer_denom}
+        return {**val_loss_log, 'log': tensorboard_logs}
+
+    def multi_test_epoch_end(self, outputs, dataloader_idx: int = 0):
+        if self.compute_eval_loss:
+            test_loss_mean = torch.stack([x['test_loss'] for x in outputs]).mean()
+            test_loss_log = {'test_loss': test_loss_mean}
+        else:
+            test_loss_log = {}
+        wer_num = torch.stack([x['test_wer_num'] for x in outputs]).sum()
+        wer_denom = torch.stack([x['test_wer_denom'] for x in outputs]).sum()
+        tensorboard_logs = {**test_loss_log, 'test_wer': wer_num.float() / wer_denom}
+        return {**test_loss_log, 'log': tensorboard_logs}
+
+    def _setup_transcribe_dataloader(self, config: Dict) -> 'torch.utils.data.DataLoader':
+        """
+        Setup function for a temporary data loader which wraps the provided audio file.
+
+        Args:
+            config: A python dictionary which contains the following keys:
+            paths2audio_files: (a list) of paths to audio files. The files should be relatively short fragments. \
+                Recommended length per file is between 5 and 25 seconds.
+            batch_size: (int) batch size to use during inference. \
+                Bigger will result in better throughput performance but would use more memory.
+            temp_dir: (str) A temporary directory where the audio manifest is temporarily
+                stored.
+
+        Returns:
+            A pytorch DataLoader for the given audio file(s).
+        """
+        dl_config = {
+            'manifest_filepath': os.path.join(config['temp_dir'], 'manifest.json'),
+            'sample_rate': self.preprocessor._sample_rate,
+            'labels': self.joint.vocabulary,
+            'batch_size': min(config['batch_size'], len(config['paths2audio_files'])),
+            'trim_silence': True,
+            'shuffle': False,
+        }
+
+        temporary_datalayer = self._setup_dataloader_from_config(config=DictConfig(dl_config))
+        return temporary_datalayer
+
+    def on_after_backward(self):
+        super().on_after_backward()
+        if self._optim_variational_noise_std > 0 and self.global_step >= self._optim_variational_noise_start:
+            for param_name, param in self.decoder.named_parameters():
+                if param.grad is not None:
+                    noise = torch.normal(
+                        mean=0.0,
+                        std=self._optim_variational_noise_std,
+                        size=param.size(),
+                        device=param.device,
+                        dtype=param.dtype,
+                    )
+                    param.grad.data.add_(noise)
+
+        if self._optim_normalize_joint_txu:
+            T, U = self._optim_normalize_txu
+            if T is not None and U is not None:
+                for param_name, param in self.encoder.named_parameters():
+                    if param.grad is not None:
+                        param.grad.data.div_(U)
+
+                for param_name, param in self.decoder.named_parameters():
+                    if param.grad is not None:
+                        param.grad.data.div_(T)
+
+        if self._optim_normalize_encoder_norm:
+            for param_name, param in self.encoder.named_parameters():
+                if param.grad is not None:
+                    norm = param.grad.norm()
+                    param.grad.data.div_(norm)
+
+        if self._optim_normalize_decoder_norm:
+            for param_name, param in self.decoder.named_parameters():
+                if param.grad is not None:
+                    norm = param.grad.norm()
+                    param.grad.data.div_(norm)
+
+        if self._optim_normalize_joint_norm:
+            for param_name, param in self.joint.named_parameters():
+                if param.grad is not None:
+                    norm = param.grad.norm()
+                    param.grad.data.div_(norm)

tts/nemo/collections/asr/models/wav2vec/__init__.py

@@ -0,0 +1,14 @@
+# 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.
+from nemo.collections.asr.models.wav2vec.wav2vec_model import Wav2VecEncoderModel

tts/nemo/collections/asr/models/wav2vec/wav2vec_config.py

@@ -0,0 +1,180 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from dataclasses import dataclass, field
+from enum import Enum, auto
+from typing import Any, Dict, List, Optional
+
+
+class Wav2VecActivationType(Enum):
+    relu = 'relu'
+    gelu = 'gelu'
+
+
+class Wav2VecMaskType(Enum):
+    """
+    Used to select configuration to compute mask lengths
+        static = fixed size
+        uniform = sample from uniform distribution [mask_other, mask_length*2]
+        normal = sample from normal distribution with mean mask_length and stdev mask_other. mask is min 1 element
+        poisson = sample from possion distribution with lambda = mask length
+    """
+
+    static = auto()
+    uniform = auto()
+    normal = auto()
+    poisson = auto()
+
+
+class Wav2VecConvExtractorMode(Enum):
+    """
+    Mode for feature extractor. default has a single group norm with d groups in the first conv block,
+    whereas layer_norm has layer norms in every block.
+    """
+
+    default = auto()
+    layer_norm = auto()
+
+
+@dataclass
+class ConvConfig:
+    conv_pos: int = field(default=128, metadata={'help': 'Number of filters for convolutional positional embeddings'})
+    conv_pos_groups: int = field(
+        default=16, metadata={'help': 'Number of groups for convolutional positional embeddings'}
+    )
+
+
+@dataclass
+class Wav2VecTransformerEncoderConfig:
+    encoder_layers: int = field(default=12, metadata={'help': 'Number of encoder layers in transformer model'})
+    encoder_layerdrop: float = field(default=0.05, metadata={'help': 'Probability of dropping transformer layers'})
+    embedding_dim: int = field(default=768, metadata={'help': 'Encoder embedding dim'})
+    ffn_embedding_dim: int = field(default=3072, metadata={'help': 'Encoder embedding dim for feed forward'})
+    num_attention_heads: int = field(default=8, metadata={'help': 'Number of encoder attention heads'})
+    dropout: float = field(default=0.1, metadata={'help': 'Dropout probability for transformer encoder'})
+    activation_fn: Wav2VecActivationType = field(
+        default=Wav2VecActivationType.gelu, metadata={'help': 'Activation for transformer'}
+    )
+    layer_norm_first: bool = field(default=False, metadata={'help': 'Apply layer norm first within the transformer'})
+
+
+@dataclass
+class Wav2VecTransformerConfig:
+    dropout: float = field(default=0.1, metadata={'help': 'Dropout probability for the transformer'})
+    conv: ConvConfig = ConvConfig()
+    encoder: Wav2VecTransformerEncoderConfig = Wav2VecTransformerEncoderConfig()
+
+
+@dataclass
+class QuantizerConfig:
+    quantize_targets: bool = field(default=True, metadata={'help': 'Use quantized targets'})
+    quantize_input: bool = field(default=False, metadata={'help': 'Use quantized inputs'})
+    same_quantizer: bool = field(default=False, metadata={'help': 'Use the same quantizer for inputs and targets'})
+    latent_vars: int = field(
+        default=320, metadata={'help': 'Number of latent variables in each group of the codebook'}
+    )
+    latent_groups: int = field(default=2, metadata={'help': 'Number of groups within the codebook'})
+    latent_dim: int = field(
+        default=0,
+        metadata={
+            'help': 'If greater than 0, use dim for latent variables, else infered by final_dim / latent_groups'
+        },
+    )
+    latent_temp: tuple = field(
+        default=(2, 0.5, 0.999995), metadata={'help': 'Quantize temperature (start, stop, decay factor)'}
+    )
+
+
+@dataclass
+class ConvFeatureEncoderConfig:
+    extractor_mode: Wav2VecConvExtractorMode = field(default=Wav2VecConvExtractorMode.default)
+    conv_bias: bool = field(default=False, metadata={'help': 'Include bias in convolution feature extractor model'})
+    conv_feature_layers: List = field(
+        default_factory=lambda: [(512, 10, 5)] + [(512, 3, 2)] * 4 + [(512, 2, 2)] + [(512, 2, 2)],
+        metadata={'help': 'convolutional feature extraction layers [(dim, kernel_size, stride), ...'},
+    )
+
+
+@dataclass
+class LossConfig:
+    prob_ppl_weight: float = field(default=0.1, metadata={'help': 'Weight applied to quantized prob perplexity loss'})
+    feature_loss_weight: float = field(default=0, metadata={'help': 'Weight applied to feature L2 Norm'})
+
+
+@dataclass
+class Wav2VecMaskingConfig:
+    mask_prob: float = field(default=0.65, metadata={'help': 'Probability of replacing token with mask'})
+    mask_type: Wav2VecMaskType = field(default=Wav2VecMaskType.static,)
+    mask_other: int = field(
+        default=0,
+        metadata={'help': 'Secondary mask used for complex distributions (see help in compute_mask_indices)'},
+    )
+    mask_length: int = field(default=10, metadata={'help': 'Length of mask when masking time steps'})
+    no_mask_overlap: bool = field(default=False, metadata={'help': 'Whether to allow masks to overlap'})
+    mask_min_space: int = field(
+        default=1, metadata={'help': 'Minimum space beetween spans (if no overlap is enabled)'}
+    )
+    mask_channel_prob: float = field(default=0, metadata={'help': 'Probability of replacing a feature with 0'})
+    mask_channel_type: Wav2VecMaskType = field(default=Wav2VecMaskType.static,)
+    mask_channel_other: int = field(
+        default=0,
+        metadata={
+            'help': 'Secondary mask argument (used for more complex distributions (see help in compute_mask_indices)'
+        },
+    )
+    mask_channel_length: int = field(default=10, metadata={'help': 'Length of masks for features (channels)'})
+    no_mask_channel_overlap: bool = field(
+        default=False, metadata={'help': 'Whether to allow channel masks to overlap'}
+    )
+    mask_channel_min_space: int = field(
+        default=1, metadata={'help': 'Minimum space between spans (if no overlap is enabled)'}
+    )
+
+
+@dataclass
+class Wav2VecEncoderModelConfig:
+    loss: LossConfig = LossConfig()
+    quantizer: QuantizerConfig = QuantizerConfig()
+    conv_feature_encoder: ConvFeatureEncoderConfig = ConvFeatureEncoderConfig()
+    transformer_encoder: Wav2VecTransformerConfig = Wav2VecTransformerConfig()
+    masking: Wav2VecMaskingConfig = Wav2VecMaskingConfig()
+
+    dropout_input: float = field(default=0.1, metadata={'help': 'Dropout applied to input raw features'})
+    dropout_features: float = field(
+        default=0.1, metadata={'help': 'Dropout applied to the features generator by convolutions'}
+    )
+    final_dim: int = field(default=0, metadata={'help': 'Project final representations and targets to this dimension'})
+    n_negatives: int = field(
+        default=100, metadata={'help': 'Number of negatives to sample from the same audio sample'}
+    )
+    cross_sample_negatives: int = field(
+        default=0, metadata={'help': 'Number of negatives to sample from any sample in the batch'}
+    )
+    codebook_negatives: int = field(default=0, metadata={'help': 'Number of negative examples in codebook'})
+    negatives_from_everywhere: bool = field(
+        default=False, metadata={'help': 'Sample negatives from everywhere, not just masked states'}
+    )
+    logit_temp: float = field(default=0.1, metadata={'help': 'Temperature to divide logits by'})
+    target_glu: bool = field(default=False, metadata={'help': 'Adds project and applies GLU to targets'})
+    feature_grad_mult: float = field(default=0.1, metadata={'help': 'Multiply extracted feature gradients'})
+
+    train_ds: Optional[Dict[Any, Any]] = None
+    validation_ds: Optional[Dict[Any, Any]] = None
+    test_ds: Optional[Dict[Any, Any]] = None
+    optim: Optional[Dict[Any, Any]] = None

tts/nemo/collections/asr/models/wav2vec/wav2vec_model.py

@@ -0,0 +1,504 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+import logging
+from math import ceil
+from typing import Dict, Optional, Union
+
+import torch
+from omegaconf import DictConfig, OmegaConf
+from pytorch_lightning import Trainer
+from torch import nn
+
+from nemo.collections.asr.data import audio_to_text_dataset
+from nemo.collections.asr.losses.wav2vecloss import Wav2VecLoss
+from nemo.collections.asr.models.wav2vec.wav2vec_config import Wav2VecEncoderModelConfig
+from nemo.collections.asr.modules.wav2vec_modules import GumbelVectorQuantizer, compute_mask_indices
+from nemo.collections.asr.parts.preprocessing.perturb import process_augmentations
+from nemo.collections.asr.parts.submodules.wav2vec import ConvFeatureEncoder, GradMultiply, Wav2VecTransformerEncoder
+from nemo.core import ModelPT
+from nemo.core.classes.common import PretrainedModelInfo, typecheck
+from nemo.core.neural_types import AudioSignal, EncodedRepresentation, LossType, MaskType, NeuralType
+from nemo.core.neural_types.elements import BoolType, FloatType
+
+
+def buffered_arange(max):
+    if not hasattr(buffered_arange, "buf"):
+        buffered_arange.buf = torch.LongTensor()
+    if max > buffered_arange.buf.numel():
+        buffered_arange.buf.resize_(max)
+        torch.arange(max, out=buffered_arange.buf)
+    return buffered_arange.buf[:max]
+
+
+class Wav2VecEncoderModel(ModelPT):
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+        # Get global rank and total number of GPU workers for IterableDataset partitioning, if applicable
+        self.global_rank = 0
+        self.world_size = 0
+        self.local_rank = 0
+        if trainer is not None:
+            self.global_rank = (trainer.node_rank * trainer.num_gpus) + trainer.local_rank
+            self.world_size = trainer.num_nodes * trainer.num_gpus
+            self.local_rank = trainer.local_rank
+
+        super().__init__(cfg=cfg, trainer=trainer)
+
+        schema = OmegaConf.structured(Wav2VecEncoderModelConfig)
+        if isinstance(cfg, dict):
+            cfg = OmegaConf.create(cfg)
+        elif not isinstance(cfg, DictConfig):
+            raise ValueError(f"cfg was type: {type(cfg)}. Expected either a dict or a DictConfig")
+
+        cfg = OmegaConf.create(OmegaConf.to_container(cfg, resolve=True))
+        cfg = OmegaConf.merge(schema, cfg)
+
+        feature_enc_layers = cfg.conv_feature_encoder.conv_feature_layers
+        self.embed = feature_enc_layers[-1][0]  # Select last conv output layer dimension
+
+        self.feature_extractor = ConvFeatureEncoder(
+            conv_layers=feature_enc_layers,
+            mode=cfg.conv_feature_encoder.extractor_mode,
+            conv_bias=cfg.conv_feature_encoder.conv_bias,
+        )
+
+        encoder_embed_dim = cfg.transformer_encoder.encoder.embedding_dim
+        self.post_extract_proj = (
+            nn.Linear(self.embed, encoder_embed_dim)
+            if self.embed != encoder_embed_dim and not cfg.quantizer.quantize_input
+            else None
+        )
+
+        self.mask_cfg = cfg.masking
+
+        self.dropout_input = nn.Dropout(cfg.dropout_input)
+        self.dropout_features = nn.Dropout(cfg.dropout_features)
+
+        self.feature_grad_mult = cfg.feature_grad_mult
+
+        self.quantizer = None
+        self.input_quantizer = None
+
+        self.n_negatives = cfg.n_negatives
+        self.cross_sample_negatives = cfg.cross_sample_negatives
+        self.codebook_negatives = cfg.codebook_negatives
+        self.negatives_from_everywhere = cfg.negatives_from_everywhere
+
+        final_dim = cfg.final_dim if cfg.final_dim > 0 else encoder_embed_dim
+        self.final_dim = final_dim
+        if cfg.quantizer.quantize_targets:
+            vq_dim = cfg.quantizer.latent_dim if cfg.quantizer.latent_dim > 0 else final_dim
+            self.quantizer = GumbelVectorQuantizer(
+                dim=self.embed,
+                num_vars=cfg.quantizer.latent_vars,
+                temp=cfg.quantizer.latent_temp,
+                groups=cfg.quantizer.latent_groups,
+                combine_groups=False,
+                vq_dim=vq_dim,
+                time_first=True,
+            )
+            self.project_q = nn.Linear(vq_dim, final_dim)
+        else:
+            self.project_q = nn.Linear(self.embed, final_dim)
+
+        if cfg.quantizer.quantize_input:
+            if cfg.quantizer.same_quantizer and self.quantizer is not None:
+                vq_dim = final_dim
+                self.input_quantizer = self.quantizer
+            else:
+                vq_dim = cfg.quantizer.latent_dim if cfg.quantizer.latent_dim > 0 else encoder_embed_dim
+                self.input_quantizer = GumbelVectorQuantizer(
+                    dim=self.embed,
+                    num_vars=cfg.quantizer.latent_vars,
+                    temp=cfg.quantizer.latent_temp,
+                    groups=cfg.quantizer.latent_groups,
+                    combine_groups=False,
+                    vq_dim=vq_dim,
+                    time_first=True,
+                )
+            self.project_inp = nn.Linear(vq_dim, encoder_embed_dim)
+
+        self.mask_emb = nn.Parameter(torch.FloatTensor(encoder_embed_dim).uniform_())
+
+        self.encoder = Wav2VecTransformerEncoder(cfg.transformer_encoder)
+        self.layer_norm = nn.LayerNorm(self.embed)
+
+        self.target_glu = None
+        if cfg.target_glu:
+            self.target_glu = nn.Sequential(nn.Linear(final_dim, final_dim * 2), nn.GLU())
+
+        self.final_proj = nn.Linear(encoder_embed_dim, final_dim)
+        self.loss = Wav2VecLoss(
+            feature_loss_weight=cfg.loss.feature_loss_weight,
+            prob_ppl_weight=cfg.loss.prob_ppl_weight,
+            logit_temp=cfg.logit_temp,
+        )
+
+    def training_step(self, batch, batch_idx):
+        loss, feature_loss, prob_ppl_loss = self._step(batch)
+
+        self.log('learning_rate', self._optimizer.param_groups[0]['lr'])
+        self.log('loss', loss, on_epoch=True)
+        self.log('feature_loss', feature_loss, on_epoch=True, sync_dist=True)
+        self.log('prob_ppl_loss', prob_ppl_loss, on_epoch=True, sync_dist=True)
+        return {'loss': loss}
+
+    def validation_step(self, batch, batch_idx, dataloader_idx=0):
+        loss, feature_loss, prob_ppl_loss = self._step(batch)
+        self.log('val_loss', loss, prog_bar=True, on_epoch=True, sync_dist=True)
+
+    def test_step(self, batch, batch_idx, dataloader_idx=0):
+        loss, feature_loss, prob_ppl_loss = self._step(batch)
+        self.log('test_loss', loss, prog_bar=True, on_epoch=True, sync_dist=True)
+
+    def _step(self, batch):
+        audio_signal, audio_lengths, _, _ = batch
+
+        padding_mask = self._create_padding_mask(audio_lengths)
+
+        self._update_quantizer_temp()
+        logits, targets, sampled_negatives, _, features_penalty, prob_ppl_loss, _ = self(
+            source=audio_signal, padding_mask=padding_mask
+        )
+        loss, feature_loss, prob_ppl_loss = self.loss(
+            logits=logits,
+            targets=targets,
+            negatives=sampled_negatives,
+            prob_ppl_loss=prob_ppl_loss,
+            feature_loss=features_penalty,
+        )
+        return loss, feature_loss, prob_ppl_loss
+
+    @classmethod
+    def list_available_models(cls) -> Optional[PretrainedModelInfo]:
+        return None
+
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {
+            "source": NeuralType(('B', 'T'), AudioSignal()),
+            "padding_mask": NeuralType(('B', 'T'), MaskType(), optional=True),
+            "mask": NeuralType(elements_type=BoolType(), optional=True),
+            "features_only": NeuralType(elements_type=BoolType(), optional=True),
+        }
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {
+            "logits": NeuralType(('B', 'T', 'D'), EncodedRepresentation()),
+            "targets": NeuralType(('B', 'T', 'D'), EncodedRepresentation(), optional=True),
+            "sampled_negatives": NeuralType(('N', 'B', 'T', 'D'), EncodedRepresentation(), optional=True),
+            "padding_mask": NeuralType(('B', 'T'), MaskType(), optional=True),
+            "features_penalty": NeuralType(elements_type=LossType(), optional=True),
+            "prob_ppl_loss": NeuralType(elements_type=LossType(), optional=True),
+            "cur_codebook_temp": NeuralType(elements_type=FloatType(), optional=True),
+        }
+
+    @typecheck()
+    def forward(self, source, padding_mask=None, mask=True, features_only=False) -> tuple:
+        prob_ppl_loss, cur_codebook_temp = None, None
+
+        if self.feature_grad_mult > 0:
+            features = self.feature_extractor(source)
+            if self.feature_grad_mult != 1.0:
+                features = GradMultiply.apply(features, self.feature_grad_mult)
+        else:
+            with torch.no_grad():
+                features = self.feature_extractor(source)
+
+        features_penalty = features.float().pow(2).mean()  # L2 Norm on features
+
+        features = features.transpose(1, 2)
+        features = self.layer_norm(features)
+        unmasked_features = features.clone()
+
+        if padding_mask is not None:
+            extra = padding_mask.size(1) % features.size(1)
+            if extra > 0:
+                padding_mask = padding_mask[:, :-extra]
+            padding_mask = padding_mask.view(padding_mask.size(0), features.size(1), -1)
+            padding_mask = padding_mask.all(-1)
+
+        if self.post_extract_proj is not None:
+            features = self.post_extract_proj(features)
+
+        features = self.dropout_input(features)
+        unmasked_features = self.dropout_features(unmasked_features)
+
+        if self.input_quantizer:
+            features, prob_ppl_loss, cur_codebook_temp = self.input_quantizer(features)
+            features = self.project_inp(features)
+        if mask:
+            logits, mask_indices = self.apply_mask(features, padding_mask)
+            if mask_indices is not None:
+                targets = unmasked_features[mask_indices].view(
+                    unmasked_features.size(0), -1, unmasked_features.size(-1)
+                )
+            else:
+                targets = unmasked_features
+        else:
+            logits = features
+            targets = unmasked_features
+            mask_indices = None
+
+        logits = self.encoder(logits, padding_mask=padding_mask)
+
+        if features_only:
+            return logits, padding_mask
+
+        if self.quantizer:
+            targets, prob_ppl_loss, cur_codebook_temp = self.quantizer(targets)
+            targets = self.project_q(targets)
+
+            if self.negatives_from_everywhere:
+                neg_cands, *_ = self.quantizer(unmasked_features)
+                sampled_negatives, _ = self.sample_negatives(neg_cands, targets.size(1))
+                sampled_negatives = self.project_q(sampled_negatives)
+            else:
+                sampled_negatives, _ = self.sample_negatives(targets, targets.size(1))
+
+            if self.codebook_negatives > 0:
+                cb_negs = self.quantizer.sample_from_codebook(
+                    targets.size(0) * targets.size(1), self.codebook_negatives
+                )
+                cb_negs = cb_negs.view(
+                    self.codebook_negatives, targets.size(0), targets.size(1), -1
+                )  # order doesnt matter
+                cb_negs = self.project_q(cb_negs)
+                sampled_negatives = torch.cat([sampled_negatives, cb_negs], dim=0)
+        else:
+            targets = self.project_q(targets)
+
+            if self.negatives_from_everywhere:
+                sampled_negatives, _ = self.sample_negatives(unmasked_features, targets.size(1))
+                sampled_negatives = self.project_q(sampled_negatives)
+            else:
+                sampled_negatives, _ = self.sample_negatives(targets, targets.size(1))
+
+        logits = logits[mask_indices].view(logits.size(0), -1, logits.size(-1))
+
+        if self.target_glu:
+            targets = self.target_glu(targets)
+            sampled_negatives = self.target_glu(sampled_negatives)
+
+        logits = self.final_proj(logits)
+
+        return logits, targets, sampled_negatives, padding_mask, features_penalty, prob_ppl_loss, cur_codebook_temp
+
+    def extract_features(self, source, audio_lengths, mask=False):
+        padding_mask = self._create_padding_mask(audio_lengths)
+        return self(source=source, padding_mask=padding_mask, mask=mask, features_only=True)
+
+    def remove_pretraining_modules(self):
+        self.quantizer = None
+        self.project_q = None
+        self.target_glu = None
+        self.final_proj = None
+
+    def _update_quantizer_temp(self):
+        if self.quantizer:
+            self.quantizer.set_num_updates(self.trainer.global_step)
+        if self.input_quantizer:
+            self.input_quantizer.set_num_updates(self.trainer.global_step)
+
+    def apply_mask(self, x, padding_mask):
+        B, T, C = x.shape
+        if self.mask_cfg.mask_prob > 0:
+            mask_indices = compute_mask_indices(
+                (B, T),
+                padding_mask,
+                self.mask_cfg.mask_prob,
+                self.mask_cfg.mask_length,
+                self.mask_cfg.mask_type,
+                self.mask_cfg.mask_other,
+                min_masks=2,
+                no_overlap=self.mask_cfg.no_mask_overlap,
+                min_space=self.mask_cfg.mask_min_space,
+            )
+            mask_indices = torch.from_numpy(mask_indices).to(x.device)
+            mask_emb = self.mask_emb.type_as(x)
+            x[mask_indices] = mask_emb
+        else:
+            mask_indices = None
+
+        if self.mask_cfg.mask_channel_prob > 0:
+            mask_channel_indices = compute_mask_indices(
+                (B, C),
+                None,
+                self.mask_cfg.mask_channel_prob,
+                self.mask_cfg.mask_channel_length,
+                self.mask_cfg.mask_channel_type,
+                self.mask_cfg.mask_channel_other,
+                no_overlap=self.mask_cfg.no_mask_channel_overlap,
+                min_space=self.mask_cfg.mask_channel_min_space,
+            )
+            mask_channel_indices = torch.from_numpy(mask_channel_indices).to(x.device).unsqueeze(1).expand(-1, T, -1)
+            x[mask_channel_indices] = 0
+
+        return x, mask_indices
+
+    def sample_negatives(self, y, num):
+
+        if self.n_negatives == 0 and self.cross_sample_negatives == 0:
+            return y.new(0)
+
+        bsz, tsz, fsz = y.shape
+        y = y.view(-1, fsz)  # BTC => (BxT)C
+
+        cross_high = tsz * bsz
+        high = tsz
+        with torch.no_grad():
+            assert high > 1, f"{bsz, tsz, fsz}"
+
+            if self.n_negatives > 0:
+                tszs = buffered_arange(num).unsqueeze(-1).expand(-1, self.n_negatives).flatten()
+
+                neg_idxs = torch.randint(low=0, high=high - 1, size=(bsz, self.n_negatives * num))
+                neg_idxs[neg_idxs >= tszs] += 1
+
+            if self.cross_sample_negatives > 0:
+                tszs = buffered_arange(num).unsqueeze(-1).expand(-1, self.cross_sample_negatives).flatten()
+
+                cross_neg_idxs = torch.randint(
+                    low=0, high=cross_high - 1, size=(bsz, self.cross_sample_negatives * num),
+                )
+                cross_neg_idxs[cross_neg_idxs >= tszs] += 1
+
+        if self.n_negatives > 0:
+            for i in range(1, bsz):
+                neg_idxs[i] += i * high
+        else:
+            neg_idxs = cross_neg_idxs
+
+        if self.cross_sample_negatives > 0 and self.n_negatives > 0:
+            neg_idxs = torch.cat([neg_idxs, cross_neg_idxs], dim=1)
+
+        negs = y[neg_idxs.view(-1)]
+        negs = negs.view(bsz, num, self.n_negatives + self.cross_sample_negatives, fsz).permute(
+            2, 0, 1, 3
+        )  # to NxBxTxC
+        return negs, neg_idxs
+
+    def _create_padding_mask(self, audio_lengths):
+        # Broadcast to vectorize creating the padding mask
+        max_len = max(audio_lengths)
+        padding_mask = torch.arange(max_len, device=self.device)
+        padding_mask = padding_mask.expand(len(audio_lengths), max_len) < audio_lengths.unsqueeze(1)
+        # Negate to false where no padding
+        padding_mask = ~padding_mask
+        return padding_mask
+
+    def setup_training_data(self, train_data_config: Optional[Union[DictConfig, Dict]]):
+        if 'shuffle' not in train_data_config:
+            train_data_config['shuffle'] = True
+
+        # preserve config
+        self._update_dataset_config(dataset_name='train', config=train_data_config)
+
+        self._train_dl = self._setup_dataloader_from_config(config=train_data_config)
+
+        # Need to set this because if using an IterableDataset, the length of the dataloader is the total number
+        # of samples rather than the number of batches, and this messes up the tqdm progress bar.
+        # So we set the number of steps manually (to the correct number) to fix this.
+        if 'is_tarred' in train_data_config and train_data_config['is_tarred']:
+            # We also need to check if limit_train_batches is already set.
+            # If it's an int, we assume that the user has set it to something sane, i.e. <= # training batches,
+            # and don't change it. Otherwise, adjust batches accordingly if it's a float (including 1.0).
+            if isinstance(self._trainer.limit_train_batches, float):
+                self._trainer.limit_train_batches = int(
+                    self._trainer.limit_train_batches
+                    * ceil((len(self._train_dl.dataset) / self.world_size) / train_data_config['batch_size'])
+                )
+
+    def setup_validation_data(self, val_data_config: Optional[Union[DictConfig, Dict]]):
+        if 'shuffle' not in val_data_config:
+            val_data_config['shuffle'] = False
+
+        # preserve config
+        self._update_dataset_config(dataset_name='validation', config=val_data_config)
+
+        self._validation_dl = self._setup_dataloader_from_config(config=val_data_config)
+
+    def setup_test_data(self, test_data_config: Optional[Union[DictConfig, Dict]]):
+        if 'shuffle' not in test_data_config:
+            test_data_config['shuffle'] = False
+
+        # preserve config
+        self._update_dataset_config(dataset_name='test', config=test_data_config)
+
+        self._test_dl = self._setup_dataloader_from_config(config=test_data_config)
+
+    def _setup_dataloader_from_config(self, config: Optional[Dict]):
+
+        if 'augmentor' in config:
+            augmentor = process_augmentations(config['augmentor'])
+        else:
+            augmentor = None
+
+        shuffle = config['shuffle']
+        device = 'gpu' if torch.cuda.is_available() else 'cpu'
+        if config.get('use_dali', False):
+            device_id = self.local_rank if device == 'gpu' else None
+            dataset = audio_to_text_dataset.get_dali_char_dataset(
+                config=config,
+                shuffle=shuffle,
+                device_id=device_id,
+                global_rank=self.global_rank,
+                world_size=self.world_size,
+                preprocessor_cfg=self._cfg.preprocessor,
+            )
+            return dataset
+
+        # Instantiate tarred dataset loader or normal dataset loader
+        if config.get('is_tarred', False):
+            if ('tarred_audio_filepaths' in config and config['tarred_audio_filepaths'] is None) or (
+                'manifest_filepath' in config and config['manifest_filepath'] is None
+            ):
+                logging.warning(
+                    "Could not load dataset as `manifest_filepath` was None or "
+                    f"`tarred_audio_filepaths` is None. Provided config : {config}"
+                )
+                return None
+
+            shuffle_n = config.get('shuffle_n', 4 * config['batch_size']) if shuffle else 0
+            dataset = audio_to_text_dataset.get_tarred_char_dataset(
+                config=config,
+                shuffle_n=shuffle_n,
+                global_rank=self.global_rank,
+                world_size=self.world_size,
+                augmentor=augmentor,
+            )
+            shuffle = False
+        else:
+            if 'manifest_filepath' in config and config['manifest_filepath'] is None:
+                logging.warning(f"Could not load dataset as `manifest_filepath` was None. Provided config : {config}")
+                return None
+
+            dataset = audio_to_text_dataset.get_char_dataset(config=config, augmentor=augmentor)
+
+        return torch.utils.data.DataLoader(
+            dataset=dataset,
+            batch_size=config['batch_size'],
+            collate_fn=dataset.collate_fn,
+            drop_last=config.get('drop_last', False),
+            shuffle=shuffle,
+            num_workers=config.get('num_workers', 0),
+            pin_memory=config.get('pin_memory', False),
+        )