Introduce code from Rust monorepo. This is a mess

tts/nemo/README.md

@@ -0,0 +1,9 @@
+NeMo (**Ne**ural **Mo**dules) is a toolkit for creating AI applications built around **neural modules**, conceptual blocks of neural networks that take *typed* inputs and produce *typed* outputs.
+
+**NeMo Core** provides common APIs all modules and models have to implement.
+
+**NeMo Collections**
+
+* ASR - collection of modules and models for building speech recognition networks
+* TTS - collection of modules and models for building speech synthesis networks
+* NLP - collection of modules and models for building NLP networks

tts/nemo/__init__.py

@@ -0,0 +1,28 @@
+# 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.package_info import (
+    __contact_emails__,
+    __contact_names__,
+    __description__,
+    __download_url__,
+    __homepage__,
+    __keywords__,
+    __license__,
+    __package_name__,
+    __repository_url__,
+    __shortversion__,
+    __version__,
+)

tts/nemo/collections/__init__.py

@@ -0,0 +1,13 @@
+# 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.

tts/nemo/collections/asr/__init__.py

@@ -0,0 +1,25 @@
+# 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 import data, losses, models, modules
+from nemo.package_info import __version__
+
+# Set collection version equal to NeMo version.
+__version = __version__
+
+# Authorship.
+__author__ = "NVIDIA Corporation"
+
+# Set collection name.
+__description__ = "Automatic Speech Recognition collection"

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

@@ -0,0 +1,13 @@
+# 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.

tts/nemo/collections/asr/data/audio_to_label.py

@@ -0,0 +1,921 @@
+# 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 io
+import math
+import os
+from typing import Dict, List, Optional, Union
+
+import braceexpand
+import torch
+import webdataset as wd
+
+from nemo.collections.asr.parts.preprocessing.segment import available_formats as valid_sf_formats
+from nemo.collections.common.parts.preprocessing import collections
+from nemo.core.classes import Dataset, IterableDataset
+from nemo.core.neural_types import AudioSignal, LabelsType, LengthsType, NeuralType, RegressionValuesType
+from nemo.utils import logging
+
+# List of valid file formats (prioritized by order of importance)
+VALID_FILE_FORMATS = ';'.join(['wav', 'mp3', 'flac'] + [fmt.lower() for fmt in valid_sf_formats.keys()])
+
+
+def repeat_signal(signal, sig_len, required_length):
+    """repeat signal to make short signal to have required_length
+    Args:
+        signal (FloatTensor): input signal
+        sig_len (LongTensor): length of input signal
+        required_length(float) : length of generated signal
+    Returns:
+        signal (FloatTensor): generated signal of required_length by repeating itself.
+    """
+    repeat = required_length // sig_len
+    rem = required_length % sig_len
+    sub = signal[-rem:] if rem > 0 else torch.tensor([])
+    rep_sig = torch.cat(repeat * [signal])
+    signal = torch.cat((rep_sig, sub))
+    return signal
+
+
+def normalize(signal):
+    """normalize signal
+    Args:
+        signal(FloatTensor): signal to be normalized.
+    """
+    signal_minusmean = signal - signal.mean()
+    return signal_minusmean / signal_minusmean.abs().max()
+
+
+def count_occurence(manifest_file_id):
+    """Count number of wav files in Dict manifest_file_id. Use for _TarredAudioToLabelDataset.
+    Args:
+        manifest_file_id (Dict): Dict of files and their corresponding id. {'A-sub0' : 1, ..., 'S-sub10':100}
+    Returns:
+        count (Dict): Dict of wav files {'A' : 2, ..., 'S':10}
+    """
+    count = dict()
+    for i in manifest_file_id:
+        audio_filename = i.split("-sub")[0]
+        count[audio_filename] = count.get(audio_filename, 0) + 1
+    return count
+
+
+def _speech_collate_fn(batch, pad_id):
+    """collate batch of audio sig, audio len, tokens, tokens len
+    Args:
+        batch (Optional[FloatTensor], Optional[LongTensor], LongTensor,
+               LongTensor):  A tuple of tuples of signal, signal lengths,
+               encoded tokens, and encoded tokens length.  This collate func
+               assumes the signals are 1d torch tensors (i.e. mono audio).
+    """
+    _, audio_lengths, _, tokens_lengths = zip(*batch)
+    max_audio_len = 0
+    has_audio = audio_lengths[0] is not None
+    if has_audio:
+        max_audio_len = max(audio_lengths).item()
+    max_tokens_len = max(tokens_lengths).item()
+
+    audio_signal, tokens = [], []
+    for sig, sig_len, tokens_i, tokens_i_len in batch:
+        if has_audio:
+            sig_len = sig_len.item()
+            if sig_len < max_audio_len:
+                pad = (0, max_audio_len - sig_len)
+                sig = torch.nn.functional.pad(sig, pad)
+            audio_signal.append(sig)
+        tokens_i_len = tokens_i_len.item()
+        if tokens_i_len < max_tokens_len:
+            pad = (0, max_tokens_len - tokens_i_len)
+            tokens_i = torch.nn.functional.pad(tokens_i, pad, value=pad_id)
+        tokens.append(tokens_i)
+
+    if has_audio:
+        audio_signal = torch.stack(audio_signal)
+        audio_lengths = torch.stack(audio_lengths)
+    else:
+        audio_signal, audio_lengths = None, None
+    tokens = torch.stack(tokens)
+    tokens_lengths = torch.stack(tokens_lengths)
+
+    return audio_signal, audio_lengths, tokens, tokens_lengths
+
+
+def _fixed_seq_collate_fn(self, batch):
+    """collate batch of audio sig, audio len, tokens, tokens len
+        Args:
+            batch (Optional[FloatTensor], Optional[LongTensor], LongTensor,
+                LongTensor):  A tuple of tuples of signal, signal lengths,
+                encoded tokens, and encoded tokens length.  This collate func
+                assumes the signals are 1d torch tensors (i.e. mono audio).
+        """
+    fixed_length = self.featurizer.sample_rate * self.time_length
+    _, audio_lengths, _, tokens_lengths = zip(*batch)
+
+    has_audio = audio_lengths[0] is not None
+    fixed_length = int(min(fixed_length, max(audio_lengths)))
+
+    audio_signal, tokens, new_audio_lengths = [], [], []
+    for sig, sig_len, tokens_i, _ in batch:
+        if has_audio:
+            sig_len = sig_len.item()
+            chunck_len = sig_len - fixed_length
+
+            if chunck_len < 0:
+                repeat = fixed_length // sig_len
+                rem = fixed_length % sig_len
+                sub = sig[-rem:] if rem > 0 else torch.tensor([])
+                rep_sig = torch.cat(repeat * [sig])
+                signal = torch.cat((rep_sig, sub))
+                new_audio_lengths.append(torch.tensor(fixed_length))
+            else:
+                start_idx = torch.randint(0, chunck_len, (1,)) if chunck_len else torch.tensor(0)
+                end_idx = start_idx + fixed_length
+                signal = sig[start_idx:end_idx]
+                new_audio_lengths.append(torch.tensor(fixed_length))
+
+            audio_signal.append(signal)
+        tokens.append(tokens_i)
+
+    if has_audio:
+        audio_signal = torch.stack(audio_signal)
+        audio_lengths = torch.stack(new_audio_lengths)
+    else:
+        audio_signal, audio_lengths = None, None
+    tokens = torch.stack(tokens)
+    tokens_lengths = torch.stack(tokens_lengths)
+
+    return audio_signal, audio_lengths, tokens, tokens_lengths
+
+
+def _sliced_seq_collate_fn(self, batch):
+    """collate batch of audio sig, audio len, tokens, tokens len
+    Args:
+        batch (Optional[FloatTensor], Optional[LongTensor], LongTensor,
+            LongTensor):  A tuple of tuples of signal, signal lengths,
+            encoded tokens, and encoded tokens length.  This collate func
+            assumes the signals are 1d torch tensors (i.e. mono audio).
+    """
+    slice_length = self.featurizer.sample_rate * self.time_length
+    _, audio_lengths, _, tokens_lengths = zip(*batch)
+    shift = self.shift_length * self.featurizer.sample_rate
+    has_audio = audio_lengths[0] is not None
+
+    audio_signal, num_slices, tokens, audio_lengths = [], [], [], []
+    for sig, sig_len, tokens_i, _ in batch:
+        if has_audio:
+            sig_len = sig_len.item()
+            dur = sig_len / self.featurizer.sample_rate
+            base = math.ceil((dur - self.time_length) / self.shift_length)
+            slice_length = int(slice_length)
+            shift = int(shift)
+            slices = 1 if base < 0 else base + 1
+            for slice_id in range(slices):
+                start_idx = slice_id * shift
+                end_idx = start_idx + slice_length
+                signal = sig[start_idx:end_idx]
+                if len(signal) < slice_length:
+                    signal = repeat_signal(signal, len(signal), slice_length)
+                audio_signal.append(signal)
+
+            num_slices.append(slices)
+            tokens.extend([tokens_i] * slices)
+            audio_lengths.extend([slice_length] * slices)
+
+    if has_audio:
+        audio_signal = torch.stack(audio_signal)
+        audio_lengths = torch.tensor(audio_lengths)
+    else:
+        audio_signal, audio_lengths = None, None
+    tokens = torch.stack(tokens)
+    tokens_lengths = torch.tensor(num_slices)  # each embedding length
+
+    return audio_signal, audio_lengths, tokens, tokens_lengths
+
+
+def _vad_frame_seq_collate_fn(self, batch):
+    """collate batch of audio sig, audio len, tokens, tokens len
+    Args:
+        batch (Optional[FloatTensor], Optional[LongTensor], LongTensor,
+            LongTensor):  A tuple of tuples of signal, signal lengths,
+            encoded tokens, and encoded tokens length.  This collate func
+            assumes the signals are 1d torch tensors (i.e. mono audio).
+            batch size equals to 1.
+    """
+    slice_length = int(self.featurizer.sample_rate * self.time_length)
+    _, audio_lengths, _, tokens_lengths = zip(*batch)
+    slice_length = min(slice_length, max(audio_lengths))
+    shift = int(self.featurizer.sample_rate * self.shift_length)
+    has_audio = audio_lengths[0] is not None
+
+    audio_signal, num_slices, tokens, audio_lengths = [], [], [], []
+
+    append_len_start = slice_length // 2
+    append_len_end = slice_length - slice_length // 2
+    for sig, sig_len, tokens_i, _ in batch:
+        if self.normalize_audio:
+            sig = normalize(sig)
+        start = torch.zeros(append_len_start)
+        end = torch.zeros(append_len_end)
+        sig = torch.cat((start, sig, end))
+        sig_len += slice_length
+
+        if has_audio:
+            slices = (sig_len - slice_length) // shift
+            for slice_id in range(slices):
+                start_idx = slice_id * shift
+                end_idx = start_idx + slice_length
+                signal = sig[start_idx:end_idx]
+                audio_signal.append(signal)
+
+            num_slices.append(slices)
+            tokens.extend([tokens_i] * slices)
+            audio_lengths.extend([slice_length] * slices)
+
+    if has_audio:
+        audio_signal = torch.stack(audio_signal)
+        audio_lengths = torch.tensor(audio_lengths)
+    else:
+        audio_signal, audio_lengths = None, None
+
+    tokens = torch.stack(tokens)
+    tokens_lengths = torch.tensor(num_slices)
+    return audio_signal, audio_lengths, tokens, tokens_lengths
+
+
+class _AudioLabelDataset(Dataset):
+    """
+    Dataset that loads tensors via a json file containing paths to audio files,
+    labels, and durations and offsets(in seconds). Each new line is a
+    different sample. Example below:
+    and their target labels. JSON files should be of the following format::
+        {"audio_filepath": "/path/to/audio_wav_0.wav", "duration": time_in_sec_0, "label": \
+target_label_0, "offset": offset_in_sec_0}
+        ...
+        {"audio_filepath": "/path/to/audio_wav_n.wav", "duration": time_in_sec_n, "label": \
+target_label_n, "offset": offset_in_sec_n}
+    Args:
+        manifest_filepath (str): Dataset parameter. Path to JSON containing data.
+        labels (list): Dataset parameter. List of target classes that can be output by the speaker recognition model.
+        featurizer
+        min_duration (float): Dataset parameter. All training files which have a duration less than min_duration
+            are dropped. Note: Duration is read from the manifest JSON.
+            Defaults to 0.1.
+        max_duration (float): Dataset parameter.
+            All training files which have a duration more than max_duration
+            are dropped. Note: Duration is read from the manifest JSON.
+            Defaults to None.
+        trim (bool): Whether to use trim silence from beginning and end of audio signal using librosa.effects.trim().
+            Defaults to False.
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Returns definitions of module output ports.
+        """
+
+        output_types = {
+            'audio_signal': NeuralType(
+                ('B', 'T'),
+                AudioSignal(freq=self._sample_rate)
+                if self is not None and hasattr(self, '_sample_rate')
+                else AudioSignal(),
+            ),
+            'a_sig_length': NeuralType(tuple('B'), LengthsType()),
+        }
+
+        if self.is_regression_task:
+            output_types.update(
+                {
+                    'targets': NeuralType(tuple('B'), RegressionValuesType()),
+                    'targets_length': NeuralType(tuple('B'), LengthsType()),
+                }
+            )
+        else:
+
+            output_types.update(
+                {'label': NeuralType(tuple('B'), LabelsType()), 'label_length': NeuralType(tuple('B'), LengthsType()),}
+            )
+
+        return output_types
+
+    def __init__(
+        self,
+        *,
+        manifest_filepath: str,
+        labels: List[str],
+        featurizer,
+        min_duration: Optional[float] = 0.1,
+        max_duration: Optional[float] = None,
+        trim: bool = False,
+        is_regression_task: bool = False,
+    ):
+        super().__init__()
+        self.collection = collections.ASRSpeechLabel(
+            manifests_files=manifest_filepath.split(','),
+            min_duration=min_duration,
+            max_duration=max_duration,
+            is_regression_task=is_regression_task,
+        )
+
+        self.featurizer = featurizer
+        self.trim = trim
+        self.is_regression_task = is_regression_task
+
+        if not is_regression_task:
+            self.labels = labels if labels else self.collection.uniq_labels
+            self.num_classes = len(self.labels) if self.labels is not None else 1
+            self.label2id, self.id2label = {}, {}
+            for label_id, label in enumerate(self.labels):
+                self.label2id[label] = label_id
+                self.id2label[label_id] = label
+
+            for idx in range(len(self.labels[:5])):
+                logging.debug(" label id {} and its mapped label {}".format(idx, self.id2label[idx]))
+
+        else:
+            self.labels = []
+            self.num_classes = 1
+
+    def __len__(self):
+        return len(self.collection)
+
+    def __getitem__(self, index):
+        sample = self.collection[index]
+
+        offset = sample.offset
+
+        if offset is None:
+            offset = 0
+
+        features = self.featurizer.process(sample.audio_file, offset=offset, duration=sample.duration, trim=self.trim)
+        f, fl = features, torch.tensor(features.shape[0]).long()
+
+        if not self.is_regression_task:
+            t = torch.tensor(self.label2id[sample.label]).long()
+        else:
+            t = torch.tensor(sample.label).float()
+
+        tl = torch.tensor(1).long()  # For compatibility with collate_fn used later
+
+        return f, fl, t, tl
+
+
+# Ported from https://github.com/NVIDIA/OpenSeq2Seq/blob/master/open_seq2seq/data/speech2text/speech_commands.py
+class AudioToClassificationLabelDataset(_AudioLabelDataset):
+    """
+    Dataset that loads tensors via a json file containing paths to audio
+    files, command class, and durations (in seconds). Each new line is a
+    different sample. Example below:
+    {"audio_filepath": "/path/to/audio_wav_0.wav", "duration": time_in_sec_0, "label": \
+        target_label_0, "offset": offset_in_sec_0}
+    ...
+    {"audio_filepath": "/path/to/audio_wav_n.wav", "duration": time_in_sec_n, "label": \
+        target_label_n, "offset": offset_in_sec_n}
+    Args:
+        manifest_filepath: Path to manifest json as described above. Can
+            be comma-separated paths.
+        labels (Optional[list]): String containing all the possible labels to map to
+            if None then automatically picks from ASRSpeechLabel collection.
+        featurizer: Initialized featurizer class that converts paths of
+            audio to feature tensors
+        max_duration: If audio exceeds this length, do not include in dataset
+        min_duration: If audio is less than this length, do not include
+            in dataset
+        trim: Boolean flag whether to trim the audio
+    """
+
+    # self.labels = labels if labels else self.collection.uniq_labels
+    # self.num_commands = len(self.labels)
+
+    def _collate_fn(self, batch):
+        return _speech_collate_fn(batch, pad_id=0)
+
+
+class AudioToSpeechLabelDataset(_AudioLabelDataset):
+    """
+    Dataset that loads tensors via a json file containing paths to audio
+    files, command class, and durations (in seconds). Each new line is a
+    different sample. Example below:
+    {"audio_filepath": "/path/to/audio_wav_0.wav", "duration": time_in_sec_0, "label": \
+        target_label_0, "offset": offset_in_sec_0}
+    ...
+    {"audio_filepath": "/path/to/audio_wav_n.wav", "duration": time_in_sec_n, "label": \
+        target_label_n, "offset": offset_in_sec_n}
+    Args:
+        manifest_filepath (str): Path to manifest json as described above. Can
+            be comma-separated paths.
+        labels (Optional[list]): String containing all the possible labels to map to
+            if None then automatically picks from ASRSpeechLabel collection.
+        min_duration (float): Dataset parameter.
+            All training files which have a duration less than min_duration
+            are dropped. Note: Duration is read from the manifest JSON.
+            Defaults to 0.1.
+        max_duration (float): Dataset parameter.
+            All training files which have a duration more than max_duration
+            are dropped. Note: Duration is read from the manifest JSON.
+            Defaults to None.
+        trim (bool): Whether to use trim silence from beginning and end
+            of audio signal using librosa.effects.trim().
+            Defaults to False.
+        time_length (float): time length of slice (in seconds)
+            Use this for speaker recognition and VAD tasks.
+        shift_length (float): amount of shift of window for generating the frame for VAD task in a batch
+            Use this for VAD task during inference.
+        normalize_audio (bool): Whether to normalize audio signal.
+            Defaults to False.
+        is_regression_task (bool): Whether the dataset is for a regression task instead of classification
+    """
+
+    def __init__(
+        self,
+        *,
+        manifest_filepath: str,
+        labels: List[str],
+        featurizer,
+        min_duration: Optional[float] = 0.1,
+        max_duration: Optional[float] = None,
+        trim: bool = False,
+        time_length: Optional[float] = 8,
+        shift_length: Optional[float] = 1,
+        normalize_audio: bool = False,
+        is_regression_task: bool = False,
+    ):
+        logging.info("Time length considered for collate func is {}".format(time_length))
+        logging.info("Shift length considered for collate func is {}".format(shift_length))
+        self.time_length = time_length
+        self.shift_length = shift_length
+        self.normalize_audio = normalize_audio
+
+        super().__init__(
+            manifest_filepath=manifest_filepath,
+            labels=labels,
+            featurizer=featurizer,
+            min_duration=min_duration,
+            max_duration=max_duration,
+            trim=trim,
+            is_regression_task=is_regression_task,
+        )
+
+    def fixed_seq_collate_fn(self, batch):
+        return _fixed_seq_collate_fn(self, batch)
+
+    def sliced_seq_collate_fn(self, batch):
+        return _sliced_seq_collate_fn(self, batch)
+
+    def vad_frame_seq_collate_fn(self, batch):
+        return _vad_frame_seq_collate_fn(self, batch)
+
+
+class _TarredAudioLabelDataset(IterableDataset):
+    """
+    A similar Dataset to the AudioLabelDataSet, but which loads tarred audio files.
+
+    Accepts a single comma-separated JSON manifest file (in the same style as for the AudioToSpeechLabelDataset),
+    as well as the path(s) to the tarball(s) containing the wav files. Each line of the manifest should
+    contain the information for one audio file, including at least the label and name of the audio
+    file within the tarball.
+
+    Valid formats for the audio_tar_filepaths argument include:
+    (1) a single string that can be brace-expanded, e.g. 'path/to/audio.tar' or 'path/to/audio_{1..100}.tar.gz', or
+    (2) a list of file paths that will not be brace-expanded, e.g. ['audio_1.tar', 'audio_2.tar', ...].
+
+    Note: For brace expansion in (1), there may be cases where `{x..y}` syntax cannot be used due to shell interference.
+    This occurs most commonly inside SLURM scripts. Therefore we provide a few equivalent replacements.
+    Supported opening braces - { <=> (, [, < and the special tag _OP_.
+    Supported closing braces - } <=> ), ], > and the special tag _CL_.
+    For SLURM based tasks, we suggest the use of the special tags for ease of use.
+
+    See the documentation for more information about accepted data and input formats.
+
+    If using multiple processes the number of shards should be divisible by the number of workers to ensure an
+    even split among workers. If it is not divisible, logging will give a warning but training will proceed.
+    In addition, if using mutiprocessing, each shard MUST HAVE THE SAME NUMBER OF ENTRIES after filtering
+    is applied. We currently do not check for this, but your program may hang if the shards are uneven!
+
+    Notice that a few arguments are different from the AudioLabelDataSet; for example, shuffle (bool) has been
+    replaced by shuffle_n (int).
+
+    Additionally, please note that the len() of this DataLayer is assumed to be the length of the manifest
+    after filtering. An incorrect manifest length may lead to some DataLoader issues down the line.
+
+    Args:
+        audio_tar_filepaths: Either a list of audio tarball filepaths, or a
+            string (can be brace-expandable).
+        manifest_filepath (str): Path to the manifest.
+        labels (list): Dataset parameter.
+            List of target classes that can be output by the speaker recognition model.
+        featurizer
+        shuffle_n (int): How many samples to look ahead and load to be shuffled.
+            See WebDataset documentation for more details.
+            Defaults to 0.
+        min_duration (float): Dataset parameter.
+            All training files which have a duration less than min_duration
+            are dropped. Note: Duration is read from the manifest JSON.
+            Defaults to 0.1.
+        max_duration (float): Dataset parameter.
+            All training files which have a duration more than max_duration
+            are dropped. Note: Duration is read from the manifest JSON.
+            Defaults to None.
+        trim(bool): Whether to use trim silence from beginning and end
+            of audio signal using librosa.effects.trim().
+            Defaults to False.
+        time_length (float): time length of slice (in seconds) # Pass this only for speaker recognition and VAD task
+        shift_length (float): amount of shift of window for generating the frame for VAD task. in a batch # Pass this only for VAD task during inference.
+        normalize_audio (bool): Whether to normalize audio signal. Defaults to False.
+        shard_strategy (str): Tarred dataset shard distribution strategy chosen as a str value during ddp.
+            -   `scatter`: The default shard strategy applied by WebDataset, where each node gets
+                a unique set of shards, which are permanently pre-allocated and never changed at runtime.
+            -   `replicate`: Optional shard strategy, where each node gets all of the set of shards
+                available in the tarred dataset, which are permanently pre-allocated and never changed at runtime.
+                The benefit of replication is that it allows each node to sample data points from the entire
+                dataset independently of other nodes, and reduces dependence on value of `shuffle_n`.
+
+                Note: Replicated strategy allows every node to sample the entire set of available tarfiles,
+                and therefore more than one node may sample the same tarfile, and even sample the same
+                data points! As such, there is no assured guarantee that all samples in the dataset will be
+                sampled at least once during 1 epoch.
+        global_rank (int): Worker rank, used for partitioning shards. Defaults to 0.
+        world_size (int): Total number of processes, used for partitioning shards. Defaults to 0.
+        is_regression_task (bool): Whether it is a regression task. Defualts to False.
+    """
+
+    def __init__(
+        self,
+        *,
+        audio_tar_filepaths: Union[str, List[str]],
+        manifest_filepath: str,
+        labels: List[str],
+        featurizer,
+        shuffle_n: int = 0,
+        min_duration: Optional[float] = 0.1,
+        max_duration: Optional[float] = None,
+        trim: bool = False,
+        shard_strategy: str = "scatter",
+        global_rank: int = 0,
+        world_size: int = 0,
+        is_regression_task: bool = False,
+    ):
+        self.collection = collections.ASRSpeechLabel(
+            manifests_files=manifest_filepath.split(','),
+            min_duration=min_duration,
+            max_duration=max_duration,
+            index_by_file_id=True,  # Must set this so the manifest lines can be indexed by file ID
+        )
+
+        self.file_occurence = count_occurence(self.collection.mapping)
+
+        self.featurizer = featurizer
+        self.trim = trim
+
+        self.labels = labels if labels else self.collection.uniq_labels
+        self.num_classes = len(self.labels)
+
+        self.label2id, self.id2label = {}, {}
+        for label_id, label in enumerate(self.labels):
+            self.label2id[label] = label_id
+            self.id2label[label_id] = label
+
+        for idx in range(len(self.labels[:5])):
+            logging.debug(" label id {} and its mapped label {}".format(idx, self.id2label[idx]))
+
+        valid_shard_strategies = ['scatter', 'replicate']
+        if shard_strategy not in valid_shard_strategies:
+            raise ValueError(f"`shard_strategy` must be one of {valid_shard_strategies}")
+
+        if isinstance(audio_tar_filepaths, str):
+            # Replace '(' and '[' with '{'
+            brace_keys_open = ['(', '[', '<', '_OP_']
+            for bkey in brace_keys_open:
+                if bkey in audio_tar_filepaths:
+                    audio_tar_filepaths = audio_tar_filepaths.replace(bkey, "{")
+
+            # Replace ')' and ']' with '}'
+            brace_keys_close = [')', ']', '>', '_CL_']
+            for bkey in brace_keys_close:
+                if bkey in audio_tar_filepaths:
+                    audio_tar_filepaths = audio_tar_filepaths.replace(bkey, "}")
+
+        # Check for distributed and partition shards accordingly
+        if world_size > 1:
+            if isinstance(audio_tar_filepaths, str):
+                # Brace expand
+                audio_tar_filepaths = list(braceexpand.braceexpand(audio_tar_filepaths))
+
+            if shard_strategy == 'scatter':
+                logging.info("All tarred dataset shards will be scattered evenly across all nodes.")
+
+                if len(audio_tar_filepaths) % world_size != 0:
+                    logging.warning(
+                        f"Number of shards in tarred dataset ({len(audio_tar_filepaths)}) is not divisible "
+                        f"by number of distributed workers ({world_size})."
+                    )
+
+                begin_idx = (len(audio_tar_filepaths) // world_size) * global_rank
+                end_idx = begin_idx + (len(audio_tar_filepaths) // world_size)
+                audio_tar_filepaths = audio_tar_filepaths[begin_idx:end_idx]
+                logging.info(
+                    "Partitioning tarred dataset: process (%d) taking shards [%d, %d)", global_rank, begin_idx, end_idx
+                )
+
+            elif shard_strategy == 'replicate':
+                logging.info("All tarred dataset shards will be replicated across all nodes.")
+
+            else:
+                raise ValueError(f"Invalid shard strategy ! Allowed values are : {valid_shard_strategies}")
+
+        # Put together WebDataset
+        self._dataset = wd.WebDataset(urls=audio_tar_filepaths, nodesplitter=None)
+
+        if shuffle_n > 0:
+            self._dataset = self._dataset.shuffle(shuffle_n)
+        else:
+            logging.info("WebDataset will not shuffle files within the tar files.")
+
+        self._dataset = (
+            self._dataset.rename(audio=VALID_FILE_FORMATS, key='__key__')
+            .to_tuple('audio', 'key')
+            .pipe(self._filter)
+            .map(f=self._build_sample)
+        )
+
+    def _filter(self, iterator):
+        """This function is used to remove samples that have been filtered out by ASRSpeechLabel already.
+        Otherwise, we would get a KeyError as _build_sample attempts to find the manifest entry for a sample
+        that was filtered out (e.g. for duration).
+        Note that if using multi-GPU training, filtering may lead to an imbalance in samples in each shard,
+        which may make your code hang as one process will finish before the other.
+        """
+
+        class TarredAudioFilter:
+            def __init__(self, collection, file_occurence):
+                self.iterator = iterator
+                self.collection = collection
+                self.file_occurence = file_occurence
+                self._iterable = None
+
+            def __iter__(self):
+                self._iterable = self._internal_generator()
+                return self
+
+            def __next__(self):
+                try:
+                    values = next(self._iterable)
+                except StopIteration:
+                    # reset generator
+                    self._iterable = self._internal_generator()
+                    values = next(self._iterable)
+
+                return values
+
+            def _internal_generator(self):
+                """
+                WebDataset requires an Iterator, but we require an iterable that yields 1-or-more
+                values per value inside self.iterator.
+
+                Therefore wrap the iterator with a generator function that will yield 1-or-more
+                values per sample in the iterator.
+                """
+                for _, tup in enumerate(self.iterator):
+                    audio_bytes, audio_filename = tup
+
+                    file_id, _ = os.path.splitext(os.path.basename(audio_filename))
+                    if audio_filename in self.file_occurence:
+                        for j in range(0, self.file_occurence[file_id]):
+                            if j == 0:
+                                audio_filename = file_id
+                            else:
+                                audio_filename = file_id + "-sub" + str(j)
+                            yield audio_bytes, audio_filename
+
+        return TarredAudioFilter(self.collection, self.file_occurence)
+
+    def _build_sample(self, tup):
+        """Builds the training sample by combining the data from the WebDataset with the manifest info.
+        """
+        audio_bytes, audio_filename = tup
+        # Grab manifest entry from self.collection
+        file_id, _ = os.path.splitext(os.path.basename(audio_filename))
+
+        manifest_idx = self.collection.mapping[file_id]
+        manifest_entry = self.collection[manifest_idx]
+
+        offset = manifest_entry.offset
+        if offset is None:
+            offset = 0
+
+        # Convert audio bytes to IO stream for processing (for SoundFile to read)
+        audio_filestream = io.BytesIO(audio_bytes)
+        features = self.featurizer.process(
+            audio_filestream, offset=offset, duration=manifest_entry.duration, trim=self.trim,
+        )
+
+        audio_filestream.close()
+
+        # Audio features
+        f, fl = features, torch.tensor(features.shape[0]).long()
+
+        t = self.label2id[manifest_entry.label]
+        tl = 1  # For compatibility with collate_fn used later
+
+        return f, fl, torch.tensor(t).long(), torch.tensor(tl).long()
+
+    def __iter__(self):
+        return self._dataset.__iter__()
+
+    def __len__(self):
+        return len(self.collection)
+
+
+class TarredAudioToClassificationLabelDataset(_TarredAudioLabelDataset):
+    """
+    A similar Dataset to the AudioToClassificationLabelDataset, but which loads tarred audio files.
+
+    Accepts a single comma-separated JSON manifest file (in the same style as for the AudioToClassificationLabelDataset),
+    as well as the path(s) to the tarball(s) containing the wav files. Each line of the manifest should
+    contain the information for one audio file, including at least the transcript and name of the audio
+    file within the tarball.
+
+    Valid formats for the audio_tar_filepaths argument include:
+    (1) a single string that can be brace-expanded, e.g. 'path/to/audio.tar' or 'path/to/audio_{1..100}.tar.gz', or
+    (2) a list of file paths that will not be brace-expanded, e.g. ['audio_1.tar', 'audio_2.tar', ...].
+
+    See the WebDataset documentation for more information about accepted data and input formats.
+
+    If using multiple processes the number of shards should be divisible by the number of workers to ensure an
+    even split among workers. If it is not divisible, logging will give a warning but training will proceed.
+    In addition, if using mutiprocessing, each shard MUST HAVE THE SAME NUMBER OF ENTRIES after filtering
+    is applied. We currently do not check for this, but your program may hang if the shards are uneven!
+
+    Notice that a few arguments are different from the AudioToBPEDataset; for example, shuffle (bool) has been
+    replaced by shuffle_n (int).
+
+    Additionally, please note that the len() of this DataLayer is assumed to be the length of the manifest
+    after filtering. An incorrect manifest length may lead to some DataLoader issues down the line.
+
+    Args:
+        audio_tar_filepaths: Either a list of audio tarball filepaths, or a
+            string (can be brace-expandable).
+        manifest_filepath (str): Path to the manifest.
+        labels (list): Dataset parameter.
+            List of target classes that can be output by the speaker recognition model.
+        featurizer
+        shuffle_n (int): How many samples to look ahead and load to be shuffled.
+            See WebDataset documentation for more details.
+            Defaults to 0.
+        min_duration (float): Dataset parameter.
+            All training files which have a duration less than min_duration
+            are dropped. Note: Duration is read from the manifest JSON.
+            Defaults to 0.1.
+        max_duration (float): Dataset parameter.
+            All training files which have a duration more than max_duration
+            are dropped. Note: Duration is read from the manifest JSON.
+            Defaults to None.
+        trim(bool): Whether to use trim silence from beginning and end
+            of audio signal using librosa.effects.trim().
+            Defaults to False.
+        shard_strategy (str): Tarred dataset shard distribution strategy chosen as a str value during ddp.
+            -   `scatter`: The default shard strategy applied by WebDataset, where each node gets
+                a unique set of shards, which are permanently pre-allocated and never changed at runtime.
+            -   `replicate`: Optional shard strategy, where each node gets all of the set of shards
+                available in the tarred dataset, which are permanently pre-allocated and never changed at runtime.
+                The benefit of replication is that it allows each node to sample data points from the entire
+                dataset independently of other nodes, and reduces dependence on value of `shuffle_n`.
+
+                Note: Replicated strategy allows every node to sample the entire set of available tarfiles,
+                and therefore more than one node may sample the same tarfile, and even sample the same
+                data points! As such, there is no assured guarantee that all samples in the dataset will be
+                sampled at least once during 1 epoch.
+        global_rank (int): Worker rank, used for partitioning shards. Defaults to 0.
+        world_size (int): Total number of processes, used for partitioning shards. Defaults to 0.
+        is_regression_task (bool): Whether it is a regression task. Defualts to False.
+    """
+
+    # self.labels = labels if labels else self.collection.uniq_labels
+    # self.num_commands = len(self.labels)
+
+    def _collate_fn(self, batch):
+        return _speech_collate_fn(batch, pad_id=0)
+
+
+class TarredAudioToSpeechLabelDataset(_TarredAudioLabelDataset):
+    """
+    A similar Dataset to the AudioToSpeechLabelDataset, but which loads tarred audio files.
+
+    Accepts a single comma-separated JSON manifest file (in the same style as for the AudioToSpeechLabelDataset),
+    as well as the path(s) to the tarball(s) containing the wav files. Each line of the manifest should
+    contain the information for one audio file, including at least the transcript and name of the audio
+    file within the tarball.
+
+    Valid formats for the audio_tar_filepaths argument include:
+    (1) a single string that can be brace-expanded, e.g. 'path/to/audio.tar' or 'path/to/audio_{1..100}.tar.gz', or
+    (2) a list of file paths that will not be brace-expanded, e.g. ['audio_1.tar', 'audio_2.tar', ...].
+
+    See the WebDataset documentation for more information about accepted data and input formats.
+
+    If using multiple processes the number of shards should be divisible by the number of workers to ensure an
+    even split among workers. If it is not divisible, logging will give a warning but training will proceed.
+    In addition, if using mutiprocessing, each shard MUST HAVE THE SAME NUMBER OF ENTRIES after filtering
+    is applied. We currently do not check for this, but your program may hang if the shards are uneven!
+
+    Notice that a few arguments are different from the AudioToBPEDataset; for example, shuffle (bool) has been
+    replaced by shuffle_n (int).
+
+    Additionally, please note that the len() of this DataLayer is assumed to be the length of the manifest
+    after filtering. An incorrect manifest length may lead to some DataLoader issues down the line.
+
+    Args:
+        audio_tar_filepaths: Either a list of audio tarball filepaths, or a
+            string (can be brace-expandable).
+        manifest_filepath (str): Path to the manifest.
+        labels (list): Dataset parameter.
+            List of target classes that can be output by the speaker recognition model.
+        featurizer
+        shuffle_n (int): How many samples to look ahead and load to be shuffled.
+            See WebDataset documentation for more details.
+            Defaults to 0.
+        min_duration (float): Dataset parameter.
+            All training files which have a duration less than min_duration
+            are dropped. Note: Duration is read from the manifest JSON.
+            Defaults to 0.1.
+        max_duration (float): Dataset parameter.
+            All training files which have a duration more than max_duration
+            are dropped. Note: Duration is read from the manifest JSON.
+            Defaults to None.
+        trim(bool): Whether to use trim silence from beginning and end
+            of audio signal using librosa.effects.trim().
+            Defaults to False.
+        time_length (float): time length of slice (in seconds) # Pass this only for speaker recognition and VAD task
+        shift_length (float): amount of shift of window for generating the frame for VAD task. in a batch # Pass this only for VAD task during inference.
+        normalize_audio (bool): Whether to normalize audio signal. Defaults to False.
+        shard_strategy (str): Tarred dataset shard distribution strategy chosen as a str value during ddp.
+            -   `scatter`: The default shard strategy applied by WebDataset, where each node gets
+                a unique set of shards, which are permanently pre-allocated and never changed at runtime.
+            -   `replicate`: Optional shard strategy, where each node gets all of the set of shards
+                available in the tarred dataset, which are permanently pre-allocated and never changed at runtime.
+                The benefit of replication is that it allows each node to sample data points from the entire
+                dataset independently of other nodes, and reduces dependence on value of `shuffle_n`.
+
+                Note: Replicated strategy allows every node to sample the entire set of available tarfiles,
+                and therefore more than one node may sample the same tarfile, and even sample the same
+                data points! As such, there is no assured guarantee that all samples in the dataset will be
+                sampled at least once during 1 epoch.
+        global_rank (int): Worker rank, used for partitioning shards. Defaults to 0.
+        world_size (int): Total number of processes, used for partitioning shards. Defaults to 0.
+    """
+
+    def __init__(
+        self,
+        *,
+        audio_tar_filepaths: Union[str, List[str]],
+        manifest_filepath: str,
+        labels: List[str],
+        featurizer,
+        shuffle_n: int = 0,
+        min_duration: Optional[float] = 0.1,
+        max_duration: Optional[float] = None,
+        trim: bool = False,
+        time_length: Optional[float] = 8,
+        shift_length: Optional[float] = 1,
+        normalize_audio: bool = False,
+        shard_strategy: str = "scatter",
+        global_rank: int = 0,
+        world_size: int = 0,
+    ):
+        logging.info("Time length considered for collate func is {}".format(time_length))
+        logging.info("Shift length considered for collate func is {}".format(shift_length))
+        self.time_length = time_length
+        self.shift_length = shift_length
+        self.normalize_audio = normalize_audio
+
+        super().__init__(
+            audio_tar_filepaths=audio_tar_filepaths,
+            manifest_filepath=manifest_filepath,
+            labels=labels,
+            featurizer=featurizer,
+            shuffle_n=shuffle_n,
+            min_duration=min_duration,
+            max_duration=max_duration,
+            trim=trim,
+            shard_strategy=shard_strategy,
+            global_rank=global_rank,
+            world_size=world_size,
+        )
+
+    def fixed_seq_collate_fn(self, batch):
+        return _fixed_seq_collate_fn(self, batch)
+
+    def sliced_seq_collate_fn(self, batch):
+        return _sliced_seq_collate_fn(self, batch)
+
+    def vad_frame_seq_collate_fn(self, batch):
+        return _vad_frame_seq_collate_fn(self, batch)

tts/nemo/collections/asr/data/audio_to_label_dataset.py

@@ -0,0 +1,129 @@
+# 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.data import audio_to_label
+
+
+def get_classification_label_dataset(featurizer, config: dict) -> audio_to_label.AudioToClassificationLabelDataset:
+    """
+    Instantiates a Classification AudioLabelDataset.
+
+    Args:
+        config: Config of the AudioToClassificationLabelDataset.
+
+    Returns:
+        An instance of AudioToClassificationLabelDataset.
+    """
+    dataset = audio_to_label.AudioToClassificationLabelDataset(
+        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=config.get('trim_silence', False),
+        is_regression_task=config.get('is_regression_task', False),
+    )
+    return dataset
+
+
+def get_speech_label_dataset(featurizer, config: dict) -> audio_to_label.AudioToSpeechLabelDataset:
+    """
+    Instantiates a Speech Label (e.g. VAD, speaker recognition) AudioLabelDataset.
+
+    Args:
+        config: Config of the AudioToSpeechLabelDataSet.
+
+    Returns:
+        An instance of AudioToSpeechLabelDataset.
+    """
+    dataset = audio_to_label.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=config.get('trim_silence', False),
+        time_length=config.get('time_length', 0.31),
+        shift_length=config.get('shift_length', 0.01),
+        normalize_audio=config.get('normalize_audio', False),
+    )
+    return dataset
+
+
+def get_tarred_classification_label_dataset(
+    featurizer, config: dict, shuffle_n: int, global_rank: int, world_size: int
+) -> audio_to_label.TarredAudioToClassificationLabelDataset:
+    """
+    Instantiates a Classification TarredAudioLabelDataset.
+
+    Args:
+        config: Config of the TarredAudioToClassificationLabelDataset.
+        shuffle_n: How many samples to look ahead and load to be shuffled.
+            See WebDataset documentation for more details.
+        global_rank: Global rank of this device.
+        world_size: Global world size in the training method.
+
+    Returns:
+        An instance of TarredAudioToClassificationLabelDataset.
+    """
+    dataset = audio_to_label.TarredAudioToClassificationLabelDataset(
+        audio_tar_filepaths=config['tarred_audio_filepaths'],
+        manifest_filepath=config['manifest_filepath'],
+        labels=config['labels'],
+        featurizer=featurizer,
+        shuffle_n=shuffle_n,
+        max_duration=config.get('max_duration', None),
+        min_duration=config.get('min_duration', None),
+        trim=config.get('trim_silence', False),
+        shard_strategy=config.get('tarred_shard_strategy', 'scatter'),
+        global_rank=global_rank,
+        world_size=world_size,
+        is_regression_task=config.get('is_regression_task', False),
+    )
+    return dataset
+
+
+def get_tarred_speech_label_dataset(
+    featurizer, config: dict, shuffle_n: int, global_rank: int, world_size: int,
+) -> audio_to_label.TarredAudioToSpeechLabelDataset:
+    """
+    InInstantiates a Speech Label (e.g. VAD, speaker recognition) TarredAudioLabelDataset.
+
+    Args:
+        config: Config of the TarredAudioToSpeechLabelDataset.
+        shuffle_n: How many samples to look ahead and load to be shuffled.
+            See WebDataset documentation for more details.
+        global_rank: Global rank of this device.
+        world_size: Global world size in the training method.
+
+    Returns:
+        An instance of TarredAudioToSpeechLabelDataset.
+    """
+    dataset = audio_to_label.TarredAudioToSpeechLabelDataset(
+        audio_tar_filepaths=config['tarred_audio_filepaths'],
+        manifest_filepath=config['manifest_filepath'],
+        labels=config['labels'],
+        featurizer=featurizer,
+        shuffle_n=shuffle_n,
+        max_duration=config.get('max_duration', None),
+        min_duration=config.get('min_duration', None),
+        trim=config.get('trim_silence', False),
+        time_length=config.get('time_length', 0.31),
+        shift_length=config.get('shift_length', 0.01),
+        normalize_audio=config.get('normalize_audio', False),
+        shard_strategy=config.get('tarred_shard_strategy', 'scatter'),
+        global_rank=global_rank,
+        world_size=world_size,
+    )
+    return dataset

tts/nemo/collections/asr/data/audio_to_text_dali.py

@@ -0,0 +1,444 @@
+# 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 math
+from collections.abc import Iterator
+from typing import Callable, List, Optional, Union
+
+import torch
+from omegaconf import DictConfig
+
+from nemo.collections.common.parts.preprocessing import parsers
+from nemo.utils.decorators import experimental
+
+try:
+    import nvidia.dali as dali
+    from nvidia.dali.pipeline import Pipeline
+    from nvidia.dali.plugin.pytorch import DALIGenericIterator as DALIPytorchIterator
+    from nvidia.dali.plugin.pytorch import LastBatchPolicy as LastBatchPolicy
+
+    HAVE_DALI = True
+except (ImportError, ModuleNotFoundError):
+    HAVE_DALI = False
+
+__all__ = [
+    'AudioToCharDALIDataset',
+]
+
+
+class DALIOutputs(object):
+    def __init__(self, out_dict):
+        self._has_processed_signal = 'processed_signal' in out_dict and 'processed_signal_len' in out_dict
+        if not self._has_processed_signal:
+            assert 'audio' in out_dict and 'audio_len' in out_dict
+        assert 'transcript' in out_dict and 'transcript_len' in out_dict
+        if self._has_processed_signal:
+            self._outs = (
+                out_dict['processed_signal'],
+                out_dict['processed_signal_len'].reshape(-1),
+                out_dict['transcript'],
+                out_dict['transcript_len'].reshape(-1),
+            )
+        else:
+            self._outs = (
+                out_dict['audio'],
+                out_dict['audio_len'].reshape(-1),
+                out_dict['transcript'],
+                out_dict['transcript_len'].reshape(-1),
+            )
+
+    @property
+    def has_processed_signal(self):
+        return self._has_processed_signal
+
+    def __getitem__(self, key):
+        return self._outs[key]
+
+    def __len__(self):
+        return len(self._outs)
+
+
+@experimental
+class AudioToCharDALIDataset(Iterator):
+    """
+    NVIDIA DALI pipeline that loads tensors via one or more manifest files where each line containing a sample descriptor in JSON,
+    including audio files, transcripts, and durations (in seconds).
+    Here's an example:
+    {"audio_filepath": "/path/to/audio.wav", "text_filepath": "/path/to/audio.txt", "duration": 23.147}
+    ...
+    {"audio_filepath": "/path/to/audio.wav", "text": "the transcription", "offset": 301.75, "duration": 0.82, "utt":
+    "utterance_id", "ctm_utt": "en_4156", "side": "A"}
+    Args:
+        manifest_filepath: Path to manifest file with the format described above. Can be comma-separated paths.
+        labels: String containing all the possible characters to map to.
+        sample_rate (int): Sample rate to resample loaded audio to.
+        batch_size (int): Number of samples in a batch.
+        num_threads (int): Number of CPU processing threads to be created by the DALI pipeline.
+        max_duration (float): Determines the maximum allowed duration, in seconds, of the loaded audio files.
+        min_duration (float): Determines the minimum allowed duration, in seconds, of the loaded audio files.
+        blank_index (int): blank character index, default = -1
+        unk_index (int): unk_character index, default = -1
+        normalize (bool): whether to normalize transcript text (default): True
+        bos_id (int): Id of beginning of sequence symbol to append if not None
+        eos_id (int): Id of end of sequence symbol to append if not None
+        trim (bool): If True, it will extract the nonsilent region of the loaded audio signal.
+        shuffle (bool): If set to True, the dataset will shuffled after loading.
+        drop_last (bool): If set to True, the last batch will be dropped if incomplete. This will be the case when the shard size is not divisible by the batch size.
+                          If set to False and the size of dataset is not divisible by the batch size, then the last batch will be smaller.
+        device (str): Determines the device type to be used for preprocessing. Allowed values are: 'cpu', 'gpu'.
+        device_id (int): Index of the GPU to be used (local_rank). Only applicable when device == 'gpu'. Defaults to 0.
+        global_rank (int): Worker rank, used for partitioning shards. Defaults to 0.
+        world_size (int): Total number of processes, used for partitioning shards. Defaults to 1.
+        preprocessor_cfg (DictConfig): Preprocessor configuration. Supports AudioToMelSpectrogramPreprocessor and AudioToMFCCPreprocessor.
+    """
+
+    def __init__(
+        self,
+        manifest_filepath: str,
+        device: str,
+        batch_size: int,
+        labels: Union[str, List[str]],
+        sample_rate: int = 16000,
+        num_threads: int = 4,
+        max_duration: float = 0.0,
+        min_duration: float = 0.0,
+        blank_index: int = -1,
+        unk_index: int = -1,
+        normalize: bool = True,
+        bos_id: Optional[int] = None,
+        eos_id: Optional[int] = None,
+        trim: bool = False,
+        shuffle: bool = True,
+        drop_last: bool = False,
+        parser: Union[str, Callable] = 'en',
+        device_id: int = 0,
+        global_rank: int = 0,
+        world_size: int = 1,
+        preprocessor_cfg: DictConfig = None,
+    ):
+        self.drop_last = drop_last  # used by lr_scheduler
+        if not HAVE_DALI:
+            raise ModuleNotFoundError(
+                f"{self} requires NVIDIA DALI to be installed. "
+                f"See: https://docs.nvidia.com/deeplearning/dali/user-guide/docs/installation.html#id1"
+            )
+
+        if device not in ('cpu', 'gpu'):
+            raise ValueError(
+                f"{self} received an unexpected device argument {device}. Supported values are: 'cpu', 'gpu'"
+            )
+
+        self.batch_size = batch_size  # Used by NeMo
+
+        self.device = device
+        self.device_id = device_id
+
+        if world_size > 1:
+            self.shard_id = global_rank
+            self.num_shards = world_size
+        else:
+            self.shard_id = None
+            self.num_shards = None
+
+        self.labels = labels
+        if self.labels is None or len(self.labels) == 0:
+            raise ValueError(f"{self} expects non empty labels list")
+
+        self.parser = parsers.make_parser(
+            labels=labels, name=parser, unk_id=unk_index, blank_id=blank_index, do_normalize=normalize,
+        )
+
+        self.eos_id = eos_id
+        self.bos_id = bos_id
+        self.sample_rate = sample_rate
+
+        self.pipe = Pipeline(
+            batch_size=batch_size,
+            num_threads=num_threads,
+            device_id=self.device_id,
+            exec_async=True,
+            exec_pipelined=True,
+        )
+
+        has_preprocessor = preprocessor_cfg is not None
+        if has_preprocessor:
+            if preprocessor_cfg._target_ == "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor":
+                feature_type = "mel_spectrogram"
+            elif preprocessor_cfg._target_ == "nemo.collections.asr.modules.AudioToMFCCPreprocessor":
+                feature_type = "mfcc"
+            else:
+                raise ValueError(
+                    f"{self} received an unexpected preprocessor configuration: {preprocessor_cfg._target_}."
+                    f" Supported preprocessors are: AudioToMelSpectrogramPreprocessor, AudioToMFCCPreprocessor"
+                )
+
+            # Default values taken from AudioToMelSpectrogramPreprocessor
+            params = preprocessor_cfg
+            self.dither = params['dither'] if 'dither' in params else 0.0
+            self.preemph = params['preemph'] if 'preemph' in params else 0.97
+            self.window_size_sec = params['window_size'] if 'window_size' in params else 0.02
+            self.window_stride_sec = params['window_stride'] if 'window_stride' in params else 0.01
+            self.sample_rate = params['sample_rate'] if 'sample_rate' in params else sample_rate
+            self.window_size = int(self.window_size_sec * self.sample_rate)
+            self.window_stride = int(self.window_size_sec * self.sample_rate)
+
+            normalize = params['normalize'] if 'normalize' in params else 'per_feature'
+            if normalize == 'per_feature':  # Each freq channel independently
+                self.normalization_axes = (1,)
+            elif normalize == 'all_features':
+                self.normalization_axes = (0, 1)
+            else:
+                raise ValueError(
+                    f"{self} received {normalize} for the normalize parameter."
+                    f" It must be either 'per_feature' or 'all_features'."
+                )
+
+            self.window = None
+            window_name = params['window'] if 'window' in params else None
+            torch_windows = {
+                'hamming': torch.hamming_window,
+                'blackman': torch.blackman_window,
+                'bartlett': torch.bartlett_window,
+            }
+            if window_name is None or window_name == 'hann':
+                self.window = None  # Hann is DALI's default
+            elif window_name == 'ones':
+                self.window = torch.ones(self.window_size)
+            else:
+                try:
+                    window_fn = torch_windows.get(window_name, None)
+                    self.window = window_fn(self.window_size, periodic=False)
+                except:
+                    raise ValueError(
+                        f"{self} received {window_name} for the window parameter."
+                        f" It must be one of: ('hann', 'ones', 'hamming', 'blackman', 'bartlett', None)."
+                        f" None is equivalent to 'hann'."
+                    )
+
+            self.n_fft = params['n_fft'] if 'n_fft' in params else None  # None means default
+            self.n_mels = params['n_mels'] if 'n_mels' in params else 64
+            self.n_mfcc = params['n_mfcc'] if 'n_mfcc' in params else 64
+
+            features = params['features'] if 'features' in params else 0
+            if features > 0:
+                if feature_type == 'mel_spectrogram':
+                    self.n_mels = features
+                elif feature_type == 'mfcc':
+                    self.n_mfcc = features
+
+            # TODO Implement frame splicing
+            if 'frame_splicing' in params:
+                assert params['frame_splicing'] == 1, "Frame splicing is not implemented"
+
+            self.freq_low = params['lowfreq'] if 'lowfreq' in params else 0.0
+            self.freq_high = params['highfreq'] if 'highfreq' in params else self.sample_rate / 2.0
+            self.log_features = params['log'] if 'log' in params else True
+
+            # We want to avoid taking the log of zero
+            # There are two options: either adding or clamping to a small value
+
+            self.log_zero_guard_type = params['log_zero_guard_type'] if 'log_zero_guard_type' in params else 'add'
+            if self.log_zero_guard_type not in ["add", "clamp"]:
+                raise ValueError(
+                    f"{self} received {self.log_zero_guard_type} for the "
+                    f"log_zero_guard_type parameter. It must be either 'add' or "
+                    f"'clamp'."
+                )
+
+            self.log_zero_guard_value = params['log_zero_guard_value'] if 'log_zero_guard_value' in params else 1e-05
+            if isinstance(self.log_zero_guard_value, str):
+                if self.log_zero_guard_value == "tiny":
+                    self.log_zero_guard_value = torch.finfo(torch.float32).tiny
+                elif self.log_zero_guard_value == "eps":
+                    self.log_zero_guard_value = torch.finfo(torch.float32).eps
+                else:
+                    raise ValueError(
+                        f"{self} received {self.log_zero_guard_value} for the log_zero_guard_type parameter."
+                        f"It must be either a number, 'tiny', or 'eps'"
+                    )
+
+            self.mag_power = params['mag_power'] if 'mag_power' in params else 2
+            if self.mag_power != 1.0 and self.mag_power != 2.0:
+                raise ValueError(
+                    f"{self} received {self.mag_power} for the mag_power parameter." f" It must be either 1.0 or 2.0."
+                )
+
+            self.pad_to = params['pad_to'] if 'pad_to' in params else 16
+            self.pad_value = params['pad_value'] if 'pad_value' in params else 0.0
+
+        with self.pipe:
+            audio, transcript = dali.fn.nemo_asr_reader(
+                name="Reader",
+                manifest_filepaths=manifest_filepath.split(','),
+                dtype=dali.types.FLOAT,
+                downmix=True,
+                sample_rate=float(self.sample_rate),
+                min_duration=min_duration,
+                max_duration=max_duration,
+                read_sample_rate=False,
+                read_text=True,
+                random_shuffle=shuffle,
+                shard_id=self.shard_id,
+                num_shards=self.num_shards,
+                pad_last_batch=False,
+            )
+
+            transcript_len = dali.fn.shapes(dali.fn.reshape(transcript, shape=[-1]))
+            transcript = dali.fn.pad(transcript)
+
+            # Extract nonsilent region, if necessary
+            if trim:
+                # Need to extract non-silent region before moving to the GPU
+                roi_start, roi_len = dali.fn.nonsilent_region(audio, cutoff_db=-60)
+                audio = audio.gpu() if self.device == 'gpu' else audio
+                audio = dali.fn.slice(
+                    audio, roi_start, roi_len, normalized_anchor=False, normalized_shape=False, axes=[0]
+                )
+            else:
+                audio = audio.gpu() if self.device == 'gpu' else audio
+
+            if not has_preprocessor:
+                # No preprocessing, the output is the audio signal
+                audio_len = dali.fn.shapes(dali.fn.reshape(audio, shape=[-1]))
+                audio = dali.fn.pad(audio)
+                self.pipe.set_outputs(audio, audio_len, transcript, transcript_len)
+            else:
+                # Additive gaussian noise (dither)
+                if self.dither > 0.0:
+                    gaussian_noise = dali.fn.normal_distribution(device=self.device)
+                    audio = audio + self.dither * gaussian_noise
+
+                # Preemphasis filter
+                if self.preemph > 0.0:
+                    audio = dali.fn.preemphasis_filter(audio, preemph_coeff=self.preemph)
+
+                # Power spectrogram
+                spec = dali.fn.spectrogram(
+                    audio, nfft=self.n_fft, window_length=self.window_size, window_step=self.window_stride
+                )
+
+                if feature_type == 'mel_spectrogram' or feature_type == 'mfcc':
+                    # Spectrogram to Mel Spectrogram
+                    spec = dali.fn.mel_filter_bank(
+                        spec,
+                        sample_rate=self.sample_rate,
+                        nfilter=self.n_mels,
+                        normalize=True,
+                        freq_low=self.freq_low,
+                        freq_high=self.freq_high,
+                    )
+                    # Mel Spectrogram to MFCC
+                    if feature_type == 'mfcc':
+                        spec = dali.fn.mfcc(spec, n_mfcc=self.n_mfcc)
+
+                # Logarithm
+                if self.log_zero_guard_type == 'add':
+                    spec = spec + self.log_zero_guard_value
+
+                spec = dali.fn.to_decibels(
+                    spec, multiplier=math.log(10), reference=1.0, cutoff_db=math.log(self.log_zero_guard_value)
+                )
+
+                # Normalization
+                spec = dali.fn.normalize(spec, axes=self.normalization_axes)
+
+                # Extracting the length of the spectrogram
+                spec_len = dali.fn.slice(dali.fn.shapes(spec), 1, 1, axes=(0,))
+
+                # Pads feature dimension to be a multiple of `pad_to` and the temporal dimension to be as big as the largest sample (shape -1)
+                spec = dali.fn.pad(spec, fill_value=self.pad_value, axes=(0, 1), align=(self.pad_to, 1), shape=(1, -1))
+                self.pipe.set_outputs(spec, spec_len, transcript, transcript_len)
+        # Building DALI pipeline
+        self.pipe.build()
+
+        if has_preprocessor:
+            output_names = ['processed_signal', 'processed_signal_len', 'transcript_raw', 'transcript_raw_len']
+        else:
+            output_names = ['audio', 'audio_len', 'transcript_raw', 'transcript_raw_len']
+
+        last_batch_policy = LastBatchPolicy.DROP if drop_last else LastBatchPolicy.PARTIAL
+        self._iter = DALIPytorchIterator(
+            [self.pipe],
+            output_map=output_names,
+            reader_name="Reader",
+            last_batch_policy=last_batch_policy,
+            dynamic_shape=True,
+            auto_reset=True,
+        )
+
+        # TODO come up with a better solution
+        class DummyDataset:
+            def __init__(self, parent):
+                self.parent = parent
+
+            def __len__(self):
+                return self.parent.size
+
+        self.dataset = DummyDataset(self)  # Used by NeMo
+
+    def reset(self):
+        self._iter.reset()
+
+    def __iter__(self):
+        return self
+
+    def next(self):
+        return self.__next__()
+
+    @property
+    def size(self):
+        return self._iter.size
+
+    def __len__(self):
+        return len(self._iter)
+
+    def __next__(self):
+        outputs = self._iter.next()
+        assert len(outputs) == 1
+        dali_out = outputs[0]
+        text_raw_len = dali_out['transcript_raw_len'].numpy()
+        text_raw = dali_out['transcript_raw'].numpy()
+
+        out = {}
+        out_names = ['processed_signal', 'processed_signal_len', 'audio', 'audio_len']
+        for out_name in out_names:
+            if out_name in dali_out:
+                out[out_name] = dali_out[out_name].detach().clone()
+
+        text_tokens = []
+        text_tokens_len = []
+        max_len = 0
+        batch_size = text_raw.shape[0]
+        for i, text in enumerate(text_raw):
+            n = text_raw_len[i][0]
+            tbytes = str(text[:n].tobytes(), encoding='utf8')
+            ttokens = self.parser(tbytes)
+            if self.bos_id is not None:
+                ttokens = [self.bos_id] + ttokens
+            if self.eos_id is not None:
+                ttokens = ttokens + [self.eos_id]
+            ttokens_len = len(ttokens)
+            text_tokens_len.append(ttokens_len)
+            text_tokens.append(ttokens)
+            if ttokens_len > max_len:
+                max_len = ttokens_len
+
+        transcript_out = torch.zeros(batch_size, max_len, dtype=torch.long)
+        for i, n in enumerate(text_tokens_len):
+            transcript_out[i, :n] = torch.tensor(text_tokens[i], dtype=torch.long)
+        transcript_len_out = torch.tensor(text_tokens_len, dtype=torch.long)
+
+        out['transcript'] = transcript_out
+        out['transcript_len'] = transcript_len_out
+        return DALIOutputs(out)

tts/nemo/collections/asr/data/audio_to_text_dataset.py

@@ -0,0 +1,247 @@
+# 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 typing import Optional
+
+import torch
+from omegaconf import DictConfig, open_dict
+
+from nemo.collections.asr.data import audio_to_text, audio_to_text_dali
+from nemo.utils import logging
+
+
+def inject_dataloader_value_from_model_config(model_cfg: dict, dataloader_cfg: dict, key: str):
+    """
+    Extracts the label set provided at the top level of the model, and propagates it to the dataloader
+    config.
+
+    Args:
+        model_cfg: A DictConfig representing the model's config.
+        dataloader_cfg: A DictConfig representing the individual data loader
+        key: A str value representing a key in the model_cfg whose value will be propagated to the
+            dataloader config.
+    """
+    if key not in model_cfg:
+        logging.info(
+            f"Model level config does not container `{key}`, please explicitly provide `{key}` to the dataloaders."
+        )
+        return
+
+    # If key exists in the data loader config (either set explicitly or as a placeholder (via None))
+    if key in dataloader_cfg:
+        # Dataloader `labels` is provided and is non-null
+        if dataloader_cfg[key] is not None and model_cfg[key] != dataloader_cfg[key]:
+            # Model level `labels` dont match Dataloader level `labels`
+            logging.warning(
+                f'`{key}` is explicitly provided to the data loader, and is different from '
+                f'the `{key}` provided at the model level config.\n'
+                f'If this is incorrect, please set the dataloader\'s `{key}` to None.'
+            )
+
+        else:
+            # Dataloader `key` is None or values match
+            # Propagate from model level `key` (even if they match)
+            with open_dict(dataloader_cfg):
+                dataloader_cfg[key] = model_cfg[key]
+
+    else:
+        # If key key doesnt even exist in dataloader_cfg, inject it explicitly
+        with open_dict(dataloader_cfg):
+            dataloader_cfg[key] = model_cfg[key]
+
+
+def get_char_dataset(config: dict, augmentor: Optional['AudioAugmentor'] = None) -> audio_to_text.AudioToCharDataset:
+    """
+    Instantiates a Character Encoding based AudioToCharDataset.
+
+    Args:
+        config: Config of the AudioToCharDataset.
+        augmentor: Optional AudioAugmentor object for augmentations on audio data.
+
+    Returns:
+        An instance of AudioToCharDataset.
+    """
+    dataset = audio_to_text.AudioToCharDataset(
+        manifest_filepath=config['manifest_filepath'],
+        labels=config['labels'],
+        sample_rate=config['sample_rate'],
+        int_values=config.get('int_values', False),
+        augmentor=augmentor,
+        max_duration=config.get('max_duration', None),
+        min_duration=config.get('min_duration', None),
+        max_utts=config.get('max_utts', 0),
+        blank_index=config.get('blank_index', -1),
+        unk_index=config.get('unk_index', -1),
+        normalize=config.get('normalize_transcripts', False),
+        trim=config.get('trim_silence', False),
+        parser=config.get('parser', 'en'),
+    )
+    return dataset
+
+
+def get_bpe_dataset(
+    config: dict, tokenizer: 'TokenizerSpec', augmentor: Optional['AudioAugmentor'] = None
+) -> audio_to_text.AudioToBPEDataset:
+    """
+    Instantiates a Byte Pair Encoding / Word Piece Encoding based AudioToBPEDataset.
+
+    Args:
+        config: Config of the AudioToBPEDataset.
+        tokenizer: An instance of a TokenizerSpec object.
+        augmentor: Optional AudioAugmentor object for augmentations on audio data.
+
+    Returns:
+        An instance of AudioToBPEDataset.
+    """
+    dataset = audio_to_text.AudioToBPEDataset(
+        manifest_filepath=config['manifest_filepath'],
+        tokenizer=tokenizer,
+        sample_rate=config['sample_rate'],
+        int_values=config.get('int_values', False),
+        augmentor=augmentor,
+        max_duration=config.get('max_duration', None),
+        min_duration=config.get('min_duration', None),
+        max_utts=config.get('max_utts', 0),
+        trim=config.get('trim_silence', False),
+        use_start_end_token=config.get('use_start_end_token', True),
+    )
+    return dataset
+
+
+def get_tarred_char_dataset(
+    config: dict, shuffle_n: int, global_rank: int, world_size: int, augmentor: Optional['AudioAugmentor'] = None
+) -> audio_to_text.TarredAudioToCharDataset:
+    """
+    Instantiates a Character Encoding based TarredAudioToCharDataset.
+
+    Args:
+        config: Config of the TarredAudioToCharDataset.
+        shuffle_n: How many samples to look ahead and load to be shuffled.
+            See WebDataset documentation for more details.
+        global_rank: Global rank of this device.
+        world_size: Global world size in the training method.
+        augmentor: Optional AudioAugmentor object for augmentations on audio data.
+
+    Returns:
+        An instance of TarredAudioToCharDataset.
+    """
+    dataset = audio_to_text.TarredAudioToCharDataset(
+        audio_tar_filepaths=config['tarred_audio_filepaths'],
+        manifest_filepath=config['manifest_filepath'],
+        labels=config['labels'],
+        sample_rate=config['sample_rate'],
+        int_values=config.get('int_values', False),
+        augmentor=augmentor,
+        shuffle_n=shuffle_n,
+        max_duration=config.get('max_duration', None),
+        min_duration=config.get('min_duration', None),
+        max_utts=config.get('max_utts', 0),
+        blank_index=config.get('blank_index', -1),
+        unk_index=config.get('unk_index', -1),
+        normalize=config.get('normalize_transcripts', False),
+        trim=config.get('trim_silence', False),
+        parser=config.get('parser', 'en'),
+        shard_strategy=config.get('tarred_shard_strategy', 'scatter'),
+        global_rank=global_rank,
+        world_size=world_size,
+    )
+    return dataset
+
+
+def get_tarred_bpe_dataset(
+    config: dict,
+    tokenizer: 'TokenizerSpec',
+    shuffle_n: int,
+    global_rank: int,
+    world_size: int,
+    augmentor: Optional['AudioAugmentor'] = None,
+) -> audio_to_text.TarredAudioToBPEDataset:
+    """
+    Instantiates a Byte Pair Encoding / Word Piece Encoding based TarredAudioToBPEDataset.
+
+    Args:
+        config: Config of the TarredAudioToBPEDataset.
+        tokenizer: An instance of a TokenizerSpec object.
+        shuffle_n: How many samples to look ahead and load to be shuffled.
+            See WebDataset documentation for more details.
+        global_rank: Global rank of this device.
+        world_size: Global world size in the training method.
+        augmentor: Optional AudioAugmentor object for augmentations on audio data.
+
+    Returns:
+        An instance of TarredAudioToBPEDataset.
+    """
+    dataset = audio_to_text.TarredAudioToBPEDataset(
+        audio_tar_filepaths=config['tarred_audio_filepaths'],
+        manifest_filepath=config['manifest_filepath'],
+        tokenizer=tokenizer,
+        sample_rate=config['sample_rate'],
+        int_values=config.get('int_values', False),
+        augmentor=augmentor,
+        shuffle_n=shuffle_n,
+        max_duration=config.get('max_duration', None),
+        min_duration=config.get('min_duration', None),
+        max_utts=config.get('max_utts', 0),
+        trim=config.get('trim_silence', False),
+        use_start_end_token=config.get('use_start_end_token', True),
+        shard_strategy=config.get('tarred_shard_strategy', 'scatter'),
+        global_rank=global_rank,
+        world_size=world_size,
+    )
+    return dataset
+
+
+def get_dali_char_dataset(
+    config: dict,
+    shuffle: bool,
+    device_id: int,
+    global_rank: int,
+    world_size: int,
+    preprocessor_cfg: Optional[DictConfig] = None,
+) -> audio_to_text_dali.AudioToCharDALIDataset:
+    """
+    Instantiates a Character Encoding based AudioToCharDALIDataset.
+
+    Args:
+        config: Config of the AudioToCharDALIDataset.
+        shuffle: Bool flag whether to shuffle the dataset.
+        device_id: Index of the GPU to be used (local_rank). Only applicable when device == 'gpu'. Defaults to 0.
+        global_rank: Global rank of this device.
+        world_size: Global world size in the training method.
+        augmentor: Optional AudioAugmentor object for augmentations on audio data.
+
+    Returns:
+        An instance of AudioToCharDALIDataset.
+    """
+    device = 'gpu' if torch.cuda.is_available() else 'cpu'
+    dataset = audio_to_text_dali.AudioToCharDALIDataset(
+        manifest_filepath=config['manifest_filepath'],
+        device=device,
+        batch_size=config['batch_size'],
+        labels=config['labels'],
+        sample_rate=config['sample_rate'],
+        max_duration=config.get('max_duration', None),
+        min_duration=config.get('min_duration', None),
+        blank_index=config.get('blank_index', -1),
+        unk_index=config.get('unk_index', -1),
+        normalize=config.get('normalize_transcripts', False),
+        trim=config.get('trim_silence', False),
+        parser=config.get('parser', 'en'),
+        shuffle=shuffle,
+        device_id=device_id,
+        global_rank=global_rank,
+        world_size=world_size,
+        preprocessor_cfg=preprocessor_cfg,
+    )
+    return dataset

tts/nemo/collections/asr/data/feature_to_label.py

@@ -0,0 +1,139 @@
+# 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 typing import Dict, List, Optional
+
+import torch
+
+from nemo.collections.common.parts.preprocessing import collections
+from nemo.core.classes import Dataset
+from nemo.core.neural_types import AcousticEncodedRepresentation, LabelsType, LengthsType, NeuralType
+from nemo.utils import logging
+
+
+def _feature_collate_fn(batch):
+    """collate batch of feat sig, feat len, tokens, tokens len
+    Args:
+        batch (FloatTensor, LongTensor, LongTensor, LongTensor):  A tuple of tuples of feature, feature lengths,
+               encoded tokens, and encoded tokens length. 
+    """
+    _, feat_lengths, _, tokens_lengths = zip(*batch)
+    feat_signal, tokens = [], []
+    for sig, sig_len, tokens_i, tokens_i_len in batch:
+        feat_signal.append(sig)
+        tokens.append(tokens_i)
+
+    feat_signal = torch.stack(feat_signal)
+    feat_lengths = torch.stack(feat_lengths)
+
+    tokens = torch.stack(tokens)
+    tokens_lengths = torch.stack(tokens_lengths)
+
+    return feat_signal, feat_lengths, tokens, tokens_lengths
+
+
+class _FeatureSeqSpeakerLabelDataset(Dataset):
+    """
+    Dataset that loads tensors via a json file containing paths to feature files, sequences of labels. 
+    Each new line is a different sample. Example below:
+    and their target labels. JSON files should be of the following format:
+        {"feature_filepath": "/path/to/feature_0.p", "seq_label": speakerA speakerB SpeakerA ....} \
+        ...
+        {"feature_filepath": "/path/to/feature_n.p", "seq_label": target_seq_label_n} 
+    target_seq_label_n is the string of sequence of speaker label, separated by space.
+
+    Args:
+        manifest_filepath (str): Dataset parameter. Path to JSON containing data.
+        labels (Optional[list]): Dataset parameter. List of unique labels collected from all samples.
+        feature_loader : Dataset parameter. Feature loader to load (external) feature.       
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Returns definitions of module output ports.
+        """
+        # TODO output type for external features
+        output_types = {
+            'external_feat': NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()),
+            'feat_length': NeuralType(tuple('B'), LengthsType()),
+        }
+
+        if self.is_speaker_emb:
+            output_types.update(
+                {
+                    'embs': NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()),
+                    'embs_length': NeuralType(tuple('B'), LengthsType()),
+                    'label': NeuralType(('B', 'T'), LabelsType()),
+                    'label_length': NeuralType(tuple('B'), LengthsType()),
+                }
+            )
+        else:
+            output_types.update(
+                {'label': NeuralType(('B', 'T'), LabelsType()), 'label_length': NeuralType(tuple('B'), LengthsType()),}
+            )
+
+        return output_types
+
+    def __init__(
+        self, *, manifest_filepath: str, labels: List[str], feature_loader, is_speaker_emb: bool = False,
+    ):
+        super().__init__()
+        self.collection = collections.ASRFeatureSequenceLabel(manifests_files=manifest_filepath.split(','),)
+
+        self.feature_loader = feature_loader
+        self.labels = labels if labels else self.collection.uniq_labels
+        self.is_speaker_emb = is_speaker_emb
+
+        self.label2id, self.id2label = {}, {}
+        for label_id, label in enumerate(self.labels):
+            self.label2id[label] = label_id
+            self.id2label[label_id] = label
+
+        for idx in range(len(self.labels[:5])):
+            logging.debug(" label id {} and its mapped label {}".format(idx, self.id2label[idx]))
+
+    def __len__(self):
+        return len(self.collection)
+
+    def __getitem__(self, index):
+        sample = self.collection[index]
+
+        features = self.feature_loader.process(sample.feature_file)
+        f, fl = features, torch.tensor(features.shape[0]).long()
+
+        t = torch.tensor(sample.seq_label).float()
+        tl = torch.tensor(len(sample.seq_label)).long()
+
+        return f, fl, t, tl
+
+
+class FeatureToSeqSpeakerLabelDataset(_FeatureSeqSpeakerLabelDataset):
+    """
+    Dataset that loads tensors via a json file containing paths to feature
+    files and sequence of speakers. Each new line is a
+    different sample. Example below:
+    {"feature_filepath": "/path/to/feature_0.p", "seq_label": speakerA speakerB SpeakerA ....} \
+    ...
+    {"feature_filepath": "/path/to/feature_n.p", "seq_label": target_seq_label_n} 
+    target_seq_label_n is the string of sequence of speaker label, separated by space.
+
+    Args:
+        manifest_filepath (str): Path to manifest json as described above. Canbe comma-separated paths.
+        labels (Optional[list]): String containing all the possible labels to map to
+            if None then automatically picks from ASRFeatureSequenceLabel collection.
+        feature_loader, Feature load to loader (external) feature.
+    
+    """
+
+    def _collate_fn(self, batch):
+        return _feature_collate_fn(batch)

tts/nemo/collections/asr/data/feature_to_label_dataset.py

@@ -0,0 +1,31 @@
+# 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.data import feature_to_label
+
+
+def get_feature_seq_speakerlabel_dataset(
+    feature_loader, config: dict
+) -> feature_to_label.FeatureToSeqSpeakerLabelDataset:
+    """
+    Instantiates a FeatureSeqSpeakerLabelDataset.
+    Args:
+        config: Config of the FeatureToSeqSpeakerLabelDataset.
+
+    Returns:
+        An instance of FeatureToSeqSpeakerLabelDataset.
+    """
+    dataset = feature_to_label.FeatureToSeqSpeakerLabelDataset(
+        manifest_filepath=config['manifest_filepath'], labels=config['labels'], feature_loader=feature_loader,
+    )
+    return dataset

tts/nemo/collections/asr/data/vocabs.py

@@ -0,0 +1,352 @@
+# 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 abc
+import itertools
+import re
+import string
+import unicodedata
+from builtins import str as unicode
+from typing import List
+
+import nltk
+
+from nemo.collections.common.parts.preprocessing import parsers
+
+try:
+    import g2p_en  # noqa
+
+    _g2p = g2p_en.G2p()
+    _g2p.variables = None
+
+    HAVE_G2P = True
+except (FileNotFoundError, LookupError):
+    HAVE_G2P = False
+
+_words_re = re.compile("([a-z\-]+'[a-z\-]+|[a-z\-]+)|([^a-z{}]+)")
+
+
+def _text_preprocessing(text):
+    text = unicode(text)
+    text = ''.join(char for char in unicodedata.normalize('NFD', text) if unicodedata.category(char) != 'Mn')
+    text = text.lower()
+    text = re.sub("[^ a-z'\".,?!()\[\]:;\-]", "", text)
+    return text
+
+
+def _word_tokenize(text):
+    words = _words_re.findall(text)
+    words = [re.sub(r'\s(\d)', r'\1', word[1].upper()) if word[0] == '' else word[0] for word in words]
+    return words
+
+
+class G2p:
+    def __init__(
+        self,
+        phoneme_dict_path=None,
+        use_seq2seq_for_oov=False,
+        ignore_ambiguous_words=True,
+        text_preprocessing_func=_text_preprocessing,
+        word_tokenize_func=_word_tokenize,
+    ):
+        # Download NLTK datasets if this class is to be instantiated
+        try:
+            nltk.data.find('taggers/averaged_perceptron_tagger.zip')
+        except LookupError:
+            nltk.download('averaged_perceptron_tagger', quiet=True)
+        try:
+            nltk.data.find('corpora/cmudict.zip')
+        except LookupError:
+            nltk.download('cmudict', quiet=True)
+
+        self.homograph2features = _g2p.homograph2features
+        self.g2p_dict = self._construct_grapheme2phoneme_dict(phoneme_dict_path)
+        self.use_seq2seq_for_oov = use_seq2seq_for_oov
+        self.ignore_ambiguous_words = ignore_ambiguous_words
+
+        self.text_preprocessing_func = text_preprocessing_func
+        self.word_tokenize_func = word_tokenize_func
+
+    @staticmethod
+    def _construct_grapheme2phoneme_dict(phoneme_dict_path=None, encoding='latin-1'):
+        if phoneme_dict_path is None:
+            from nltk.corpus import cmudict
+
+            return cmudict.dict()
+
+        _alt_re = re.compile(r'\([0-9]+\)')
+        g2p_dict = {}
+        with open(phoneme_dict_path, encoding=encoding) as file:
+            for line in file:
+                if len(line) and ('A' <= line[0] <= 'Z' or line[0] == "'"):
+                    parts = line.split('  ')
+                    word = re.sub(_alt_re, '', parts[0])
+                    word = word.lower()
+
+                    pronunciation = parts[1].strip().split(" ")
+                    if word in g2p_dict:
+                        g2p_dict[word].append(pronunciation)
+                    else:
+                        g2p_dict[word] = [pronunciation]
+        return g2p_dict
+
+    def handle_ambiguous(self, word):
+        if not self.ignore_ambiguous_words or len(self.g2p_dict[word]) == 1:
+            return True
+        return False
+
+    def __call__(self, text):
+        text = self.text_preprocessing_func(text)
+        words = self.word_tokenize_func(text)
+        words_and_pos_tags = nltk.pos_tag(words)
+
+        prons = []
+        for word, pos in words_and_pos_tags:
+            word_by_hyphen = word.split("-")
+
+            # punctuation
+            if re.search("[a-zA-Z]", word) is None:
+                pron = list(word)
+            # homograph
+            elif word in self.homograph2features:
+                pron1, pron2, pos1 = self.homograph2features[word]
+                if pos.startswith(pos1):
+                    pron = pron1
+                else:
+                    pron = pron2
+            # `'s` suffix
+            elif (
+                len(word) > 2
+                and word.endswith("'s")
+                and (word not in self.g2p_dict)
+                and (word[:-2] in self.g2p_dict)
+                and self.handle_ambiguous(word[:-2])
+            ):
+                pron = self.g2p_dict[word[:-2]][0] + ["Z"]
+            # `s` suffix
+            elif (
+                len(word) > 1
+                and word.endswith("s")
+                and (word not in self.g2p_dict)
+                and (word[:-1] in self.g2p_dict)
+                and self.handle_ambiguous(word[:-1])
+            ):
+                pron = self.g2p_dict[word[:-1]][0] + ["Z"]
+            # g2p dict
+            elif word in self.g2p_dict and self.handle_ambiguous(word):
+                pron = self.g2p_dict[word][0]
+            # word with hyphens
+            elif len(word_by_hyphen) > 1 and all(
+                [sub_word in self.g2p_dict and self.handle_ambiguous(sub_word) for sub_word in word_by_hyphen]
+            ):
+                pron = []
+                for sub_word in word_by_hyphen:
+                    pron.extend(self.g2p_dict[sub_word][0])
+                    pron.extend(["-"])
+                pron.pop()
+            else:
+                if self.use_seq2seq_for_oov:
+                    # run gru-based seq2seq model from _g2p
+                    pron = _g2p.predict(word)
+                else:
+                    pron = word
+
+            prons.extend(pron)
+
+        return prons
+
+
+class Base(abc.ABC):
+    """Vocabulary for turning str text to list of int tokens."""
+
+    # fmt: off
+    PUNCT = (  # Derived from LJSpeech and "/" additionally
+        ',', '.', '!', '?', '-',
+        ':', ';', '/', '"', '(',
+        ')', '[', ']', '{', '}',
+    )
+    # fmt: on
+    PAD, BLANK, OOV = '<pad>', '<blank>', '<oov>'
+
+    def __init__(self, labels, *, pad=PAD, blank=BLANK, oov=OOV, sep='', add_blank_at="last_but_one"):
+        super().__init__()
+
+        labels = list(labels)
+        self.pad, labels = len(labels), labels + [pad]  # Padding
+
+        if add_blank_at is not None:
+            self.blank, labels = len(labels), labels + [blank]  # Reserved for blank from QN
+        else:
+            # use add_blank_at=None only for ASR where blank is added automatically
+            self.blank = -1
+
+        self.oov, labels = len(labels), labels + [oov]  # Out Of Vocabulary
+
+        if add_blank_at == "last":
+            labels[-1], labels[-2] = labels[-2], labels[-1]
+            self.oov, self.blank = self.blank, self.oov
+
+        self.labels = labels
+        self.sep = sep
+
+        self._util_ids = {self.pad, self.blank, self.oov}
+        self._label2id = {l: i for i, l in enumerate(labels)}
+        self._id2label = labels
+
+    def __call__(self, text: str) -> List[int]:
+        return self.encode(text)
+
+    @abc.abstractmethod
+    def encode(self, text: str) -> List[int]:
+        """Turns str text into int tokens."""
+        pass
+
+    def decode(self, tokens: List[int]) -> str:
+        """Turns ints tokens into str text."""
+        return self.sep.join(self._id2label[t] for t in tokens if t not in self._util_ids)
+
+
+class Chars(Base):
+    """Chars vocabulary."""
+
+    def __init__(
+        self, punct=True, spaces=False, apostrophe=True, add_blank_at="last_but_one",
+    ):
+        labels = []
+        self.space, labels = len(labels), labels + [' ']  # Space
+        labels.extend(string.ascii_lowercase)
+        if apostrophe:
+            labels.append("'")  # Apostrophe for saving "don't" and "Joe's"
+
+        if punct:
+            labels.extend(self.PUNCT)
+
+        super().__init__(labels, add_blank_at=add_blank_at)
+
+        self.punct = punct
+        self.spaces = spaces
+
+        self._parser = parsers.ENCharParser(labels)
+
+    def encode(self, text):
+        """See base class."""
+        text = self._parser._normalize(text)  # noqa
+
+        if self.spaces:
+            for p in set(text) & set(self.PUNCT):
+                text = text.replace(p, f' {p} ')
+            text = text.strip().replace('  ', ' ')
+
+        return self._parser._tokenize(text)  # noqa
+
+
+class Phonemes(Base):
+    """Phonemes vocabulary."""
+
+    # fmt: off
+    VOWELS = (
+        'AA', 'AE', 'AH', 'AO', 'AW',
+        'AY', 'EH', 'ER', 'EY', 'IH',
+        'IY', 'OW', 'OY', 'UH', 'UW',
+    )
+    CONSONANTS = (
+        'B', 'CH', 'D', 'DH', 'F', 'G',
+        'HH', 'JH', 'K', 'L', 'M', 'N',
+        'NG', 'P', 'R', 'S', 'SH', 'T',
+        'TH', 'V', 'W', 'Y', 'Z', 'ZH',
+    )
+    # fmt: on
+
+    def __init__(
+        self,
+        punct=True,
+        stresses=False,
+        spaces=True,
+        chars=False,
+        *,
+        space=' ',
+        silence=None,
+        apostrophe=True,
+        oov=Base.OOV,
+        sep='|',  # To be able to distinguish between 2/3 letters codes.
+        add_blank_at="last_but_one",
+        pad_with_space=False,
+        improved_version_g2p=False,
+        phoneme_dict_path=None,
+    ):
+        labels = []
+        self.space, labels = len(labels), labels + [space]  # Space
+
+        if silence is not None:
+            self.silence, labels = len(labels), labels + [silence]  # Silence
+
+        labels.extend(self.CONSONANTS)
+        vowels = list(self.VOWELS)
+
+        if stresses:
+            vowels = [f'{p}{s}' for p, s in itertools.product(vowels, (0, 1, 2))]
+        labels.extend(vowels)
+
+        if chars:
+            labels.extend(string.ascii_lowercase)
+
+        if apostrophe:
+            labels.append("'")  # Apostrophe
+
+        if punct:
+            labels.extend(self.PUNCT)
+
+        super().__init__(labels, oov=oov, sep=sep, add_blank_at=add_blank_at)
+
+        self.punct = punct
+        self.stresses = stresses
+        self.spaces = spaces
+        self.pad_with_space = pad_with_space
+
+        if improved_version_g2p:
+            self.g2p = G2p(phoneme_dict_path)
+        else:
+            self.g2p = _g2p
+
+    def encode(self, text):
+        """See base class."""
+        ps, space, labels = [], self.labels[self.space], set(self.labels)
+
+        for p in self.g2p(text):  # noqa
+            # Remove stress
+            if p.isalnum() and len(p) == 3 and not self.stresses:
+                p = p[:2]
+
+            # Add space if last one isn't one
+            if p == space and len(ps) > 0 and ps[-1] != space:
+                ps.append(p)
+
+            # Add next phoneme
+            if (p.isalnum() or p == "'") and p in labels:
+                ps.append(p)
+
+            # Add punct and remove space if needed
+            if (p in self.PUNCT) and self.punct:
+                if not self.spaces and len(ps) > 0 and ps[-1] == space:
+                    ps.pop()
+                ps.append(p)
+
+        # Remove trailing spaces
+        while ps[-1] == space:
+            ps.pop()
+
+        if self.pad_with_space:
+            ps = [space] + ps + [space]
+
+        return [self._label2id[p] for p in ps]

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

@@ -0,0 +1,16 @@
+# 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.losses.angularloss import AngularSoftmaxLoss
+from nemo.collections.asr.losses.ctc import CTCLoss

tts/nemo/collections/asr/losses/angularloss.py

@@ -0,0 +1,68 @@
+# ! /usr/bin/python
+# 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 torch
+
+from nemo.core.classes import Loss, Typing, typecheck
+from nemo.core.neural_types import LabelsType, LogitsType, LossType, NeuralType
+
+__all__ = ['AngularSoftmaxLoss']
+
+
+class AngularSoftmaxLoss(Loss, Typing):
+    """
+    Computes ArcFace Angular softmax angle loss
+    reference: https://openaccess.thecvf.com/content_CVPR_2019/papers/Deng_ArcFace_Additive_Angular_Margin_Loss_for_Deep_Face_Recognition_CVPR_2019_paper.pdf
+    args:
+    scale: scale value for cosine angle
+    margin: margin value added to cosine angle 
+    """
+
+    @property
+    def input_types(self):
+        """Input types definitions for AnguarLoss.
+        """
+        return {
+            "logits": NeuralType(('B', 'D'), LogitsType()),
+            "labels": NeuralType(('B',), LabelsType()),
+        }
+
+    @property
+    def output_types(self):
+        """Output types definitions for AngularLoss.
+        loss:
+            NeuralType(None)
+        """
+        return {"loss": NeuralType(elements_type=LossType())}
+
+    def __init__(self, scale=20.0, margin=1.35):
+        super().__init__()
+
+        self.eps = 1e-7
+        self.scale = scale
+        self.margin = margin
+
+    @typecheck()
+    def forward(self, logits, labels):
+        numerator = self.scale * torch.cos(
+            torch.acos(torch.clamp(torch.diagonal(logits.transpose(0, 1)[labels]), -1.0 + self.eps, 1 - self.eps))
+            + self.margin
+        )
+        excl = torch.cat(
+            [torch.cat((logits[i, :y], logits[i, y + 1 :])).unsqueeze(0) for i, y in enumerate(labels)], dim=0
+        )
+        denominator = torch.exp(numerator) + torch.sum(torch.exp(self.scale * excl), dim=1)
+        L = numerator - torch.log(denominator)
+        return -torch.mean(L)

tts/nemo/collections/asr/losses/ctc.py

@@ -0,0 +1,70 @@
+# ! /usr/bin/python
+# 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 torch
+from torch import nn
+
+from nemo.core.classes import Serialization, Typing, typecheck
+from nemo.core.neural_types import LabelsType, LengthsType, LogprobsType, LossType, NeuralType
+
+__all__ = ['CTCLoss']
+
+
+class CTCLoss(nn.CTCLoss, Serialization, Typing):
+    @property
+    def input_types(self):
+        """Input types definitions for CTCLoss.
+        """
+        return {
+            "log_probs": NeuralType(('B', 'T', 'D'), LogprobsType()),
+            "targets": NeuralType(('B', 'T'), LabelsType()),
+            "input_lengths": NeuralType(tuple('B'), LengthsType()),
+            "target_lengths": NeuralType(tuple('B'), LengthsType()),
+        }
+
+    @property
+    def output_types(self):
+        """Output types definitions for CTCLoss.
+        loss:
+            NeuralType(None)
+        """
+        return {"loss": NeuralType(elements_type=LossType())}
+
+    def __init__(self, num_classes, zero_infinity=False, reduction='mean_batch'):
+        self._blank = num_classes
+        # Don't forget to properly call base constructor
+        if reduction == 'mean_batch':
+            ctc_reduction = 'none'
+            self._apply_batch_mean = True
+        elif reduction in ['sum', 'mean', 'none']:
+            ctc_reduction = reduction
+            self._apply_batch_mean = False
+        super().__init__(blank=self._blank, reduction=ctc_reduction, zero_infinity=zero_infinity)
+
+    @typecheck()
+    def forward(self, log_probs, targets, input_lengths, target_lengths):
+        # override forward implementation
+        # custom logic, if necessary
+        input_lengths = input_lengths.long()
+        target_lengths = target_lengths.long()
+        targets = targets.long()
+        # here we transpose because we expect [B, T, D] while PyTorch assumes [T, B, D]
+        log_probs = log_probs.transpose(1, 0)
+        loss = super().forward(
+            log_probs=log_probs, targets=targets, input_lengths=input_lengths, target_lengths=target_lengths
+        )
+        if self._apply_batch_mean:
+            loss = torch.mean(loss)
+        return loss

tts/nemo/collections/asr/losses/rnnt.py

@@ -0,0 +1,288 @@
+# ! /usr/bin/python
+# 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 2018-2019, Mingkun Huang
+#
+# 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 operator
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+
+from nemo.core.classes import Loss, typecheck
+from nemo.core.neural_types import LabelsType, LengthsType, LogprobsType, LossType, NeuralType
+from nemo.core.utils.numba_utils import NUMBA_INSTALLATION_MESSAGE
+from nemo.utils import logging, model_utils
+
+try:
+    import warprnnt_pytorch as warprnnt
+
+    WARP_RNNT_AVAILABLE = True
+except (ImportError, ModuleNotFoundError):
+    WARP_RNNT_AVAILABLE = False
+
+try:
+    from nemo.collections.asr.parts.numba.rnnt_loss import RNNTLossNumba
+
+    NUMBA_RNNT_AVAILABLE = True
+except (ImportError, ModuleNotFoundError):
+    NUMBA_RNNT_AVAILABLE = False
+
+
+WARP_RNNT_INSTALLATION_MESSAGE = (
+    "Could not import `warprnnt_pytorch`.\n"
+    "Please visit https://github.com/HawkAaron/warp-transducer "
+    "and follow the steps in the readme to build and install the "
+    "pytorch bindings for RNNT Loss, or use the provided docker "
+    "container that supports RNN-T loss."
+)
+
+
+@dataclass
+class RNNTLossConfig:
+    loss_name: str
+    lib_name: str
+    is_available: bool = False
+    installation_msg: str = ""
+    min_version: Optional[str] = None
+
+
+# Resolved list of available RNNT losses
+RNNT_LOSS_RESOLVER = {
+    "warprnnt": RNNTLossConfig(
+        loss_name="warprnnt",
+        lib_name="warprnnt_pytorch",
+        is_available=WARP_RNNT_AVAILABLE,
+        installation_msg=WARP_RNNT_INSTALLATION_MESSAGE,
+    ),
+    "warprnnt_numba": RNNTLossConfig(
+        loss_name="warprnnt_numba",
+        lib_name="numba",
+        min_version='0.53.0',
+        is_available=NUMBA_RNNT_AVAILABLE,
+        installation_msg=NUMBA_INSTALLATION_MESSAGE,
+    ),
+}
+
+RNNT_LOSS_RESOLVER['default'] = RNNT_LOSS_RESOLVER['warprnnt_numba']
+
+
+def _warn_unused_additional_kwargs(loss_name, kwargs):
+    if len(kwargs) > 0:
+        logging.warning(
+            f"Loss function `{loss_name}` was provided with following additional kwargs,\n"
+            f"however they were ignored as it is unused.\n"
+            f"{kwargs}"
+        )
+
+
+def resolve_rnnt_default_loss_name() -> str:
+    return RNNT_LOSS_RESOLVER['default'].loss_name
+
+
+def resolve_rnnt_loss(loss_name: str, blank_idx: int, loss_kwargs: dict = None) -> torch.nn.Module:
+    loss_function_names = list(RNNT_LOSS_RESOLVER.keys())
+
+    if loss_name not in loss_function_names:
+        raise ValueError(
+            f"Provided `loss_name` {loss_name} not in list of available RNNT losses \n" f"{loss_function_names}"
+        )
+
+    all_available_losses = {name: config for name, config in RNNT_LOSS_RESOLVER.items() if config.is_available}
+
+    loss_config = RNNT_LOSS_RESOLVER[loss_name]  # type: RNNTLossConfig
+
+    # Re-raise import error with installation message
+    if not loss_config.is_available:
+        msg = (
+            f"Installed RNNT losses are : {list(all_available_losses.keys())}.\n"
+            f"****************************************************************\n"
+            f"To install the selected loss function, please follow the steps below:\n"
+            f"{loss_config.installation_msg}"
+        )
+        raise ImportError(msg)
+
+    # Library version check
+    if loss_config.min_version is not None:
+        ver_matched, msg = model_utils.check_lib_version(
+            loss_config.lib_name, checked_version=loss_config.min_version, operator=operator.ge
+        )
+
+        if ver_matched is False:
+            msg = (
+                f"{msg}\n"
+                f"****************************************************************\n"
+                f"To update the selected loss function, please follow the steps below:\n"
+                f"{loss_config.installation_msg}"
+            )
+            raise RuntimeError(msg)
+
+    # Resolve loss functions sequentially
+    loss_kwargs = {} if loss_kwargs is None else loss_kwargs
+
+    # Get actual loss name for `default`
+    if loss_name == 'default':
+        loss_name = loss_config.loss_name
+
+    """
+    Resolve RNNT loss functions
+    """
+    if loss_name == 'warprnnt':
+        loss_func = warprnnt.RNNTLoss(blank=blank_idx, reduction='none')
+        _warn_unused_additional_kwargs(loss_name, loss_kwargs)
+
+    elif loss_name == 'warprnnt_numba':
+        loss_func = RNNTLossNumba(blank=blank_idx, reduction='none')
+        _warn_unused_additional_kwargs(loss_name, loss_kwargs)
+
+    else:
+        raise ValueError(
+            f"Invalid value of `loss_name`: {loss_name}. Allowed loss names are :" f"{loss_function_names}"
+        )
+
+    return loss_func
+
+
+class RNNTLoss(Loss):
+    @property
+    def input_types(self):
+        """Input types definitions for CTCLoss.
+        """
+        return {
+            "log_probs": NeuralType(('B', 'T', 'D', 'D'), LogprobsType()),
+            "targets": NeuralType(('B', 'T'), LabelsType()),
+            "input_lengths": NeuralType(tuple('B'), LengthsType()),
+            "target_lengths": NeuralType(tuple('B'), LengthsType()),
+        }
+
+    @property
+    def output_types(self):
+        """Output types definitions for CTCLoss.
+        loss:
+            NeuralType(None)
+        """
+        return {"loss": NeuralType(elements_type=LossType())}
+
+    def __init__(self, num_classes, reduction: str = 'mean_batch', loss_name: str = "default", loss_kwargs=None):
+        """
+        RNN-T Loss function based on https://github.com/HawkAaron/warp-transducer.
+        Optionally, can utilize a numba implementation of the same loss without having to compile the loss,
+        albiet there is a small speed penalty for JIT numba compile.
+
+        Note:
+            Requires the pytorch bindings to be installed prior to calling this class.
+
+        Warning:
+            In the case that GPU memory is exhausted in order to compute RNNTLoss, it might cause
+            a core dump at the cuda level with the following error message.
+
+            ```
+                ...
+                costs = costs.to(acts.device)
+            RuntimeError: CUDA error: an illegal memory access was encountered
+            terminate called after throwing an instance of 'c10::Error'
+            ```
+
+            Please kill all remaining python processes after this point, and use a smaller batch size
+            for train, validation and test sets so that CUDA memory is not exhausted.
+
+        Args:
+            num_classes: Number of target classes for the joint network to predict.
+                (Excluding the RNN-T blank token).
+
+            reduction: Type of reduction to perform on loss. Possibly values are `mean`, `sum` or None.
+                None will return a torch vector comprising the individual loss values of the batch.
+
+            loss_name: String that is resolved into an RNNT loss function. Available list of losses
+                is ininitialized in `RNNT_LOSS_RESOLVER` dictionary.
+
+            loss_kwargs: Optional Dict of (str, value) pairs that are passed to the instantiated loss
+                function.
+        """
+        super(RNNTLoss, self).__init__()
+
+        if reduction not in [None, 'mean', 'sum', 'mean_batch']:
+            raise ValueError('`reduction` must be one of [mean, sum, mean_batch]')
+
+        self._blank = num_classes
+        self.reduction = reduction
+        self._loss = resolve_rnnt_loss(loss_name, blank_idx=self._blank, loss_kwargs=loss_kwargs)
+
+    @typecheck()
+    def forward(self, log_probs, targets, input_lengths, target_lengths):
+        # Cast to int 32
+        targets = targets.int()
+        input_lengths = input_lengths.int()
+        target_lengths = target_lengths.int()
+
+        max_logit_len = input_lengths.max()
+        max_targets_len = target_lengths.max()
+
+        # Force cast joint to float32
+        # TODO: Remove once Numba supports FP16
+        if log_probs.dtype != torch.float32:
+            logits_orig = log_probs
+            log_probs = log_probs.float()
+            del logits_orig  # save memory *before* computing the loss
+
+        # Ensure that shape mismatch does not occur due to padding
+        # Due to padding and subsequent downsampling, it may be possible that
+        # max sequence length computed does not match the actual max sequence length
+        # of the log_probs tensor, therefore we increment the input_lengths by the difference.
+        # This difference is generally small.
+        if log_probs.shape[1] != max_logit_len:
+            log_probs = log_probs.narrow(dim=1, start=0, length=max_logit_len).contiguous()
+
+        # Reduce transcript length to correct alignment if additional padding was applied.
+        # Transcript: [B, L] -> [B, L']; If L' < L
+        if targets.shape[1] != max_targets_len:
+            targets = targets.narrow(dim=1, start=0, length=max_targets_len)
+
+        # Loss reduction can be dynamic, so set it prior to call
+        if self.reduction != 'mean_batch':
+            self._loss.reduction = self.reduction
+
+        # Compute RNNT loss
+        loss = self._loss(acts=log_probs, labels=targets, act_lens=input_lengths, label_lens=target_lengths)
+
+        # Loss reduction can be dynamic, so reset it after call
+        if self.reduction != 'mean_batch':
+            self._loss.reduction = 'none'
+
+        # Loss reduction only for mean_batch mode
+        if self.reduction == 'mean_batch':
+            loss = torch.mean(loss)
+
+        # del new variables that may have been created
+        del (
+            log_probs,
+            targets,
+            input_lengths,
+            target_lengths,
+        )
+
+        return loss

tts/nemo/collections/asr/losses/wav2vecloss.py

@@ -0,0 +1,118 @@
+# 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.
+
+
+import torch
+import torch.nn.functional as F
+
+from nemo.core import Loss, typecheck
+from nemo.core.neural_types import EncodedRepresentation, LossType, NeuralType
+
+
+class Wav2VecLoss(Loss):
+    @property
+    def input_types(self):
+        """Input types definitions for Wav2VecLoss.
+        """
+        return {
+            "logits": NeuralType(('B', 'T', 'D'), EncodedRepresentation()),
+            "targets": NeuralType(('B', 'T', 'D'), EncodedRepresentation()),
+            "negatives": NeuralType(('N', 'B', 'T', 'D'), EncodedRepresentation()),
+            "prob_ppl_loss": NeuralType(elements_type=LossType()),
+            "feature_loss": NeuralType(elements_type=LossType()),
+        }
+
+    @property
+    def output_types(self):
+        """Output types definitions for Wav2VecLoss.
+        """
+        return {
+            "loss": NeuralType(elements_type=LossType()),
+            "feature_loss": NeuralType(elements_type=LossType()),
+            "prob_ppl_loss": NeuralType(elements_type=LossType()),
+        }
+
+    def __init__(self, feature_loss_weight: float, prob_ppl_weight: float, logit_temp: float, reduce: bool = True):
+        """
+        Compute the contrastive loss with respect to the model outputs and sampled negatives from quantizer codebooks.
+        Args:
+            feature_loss_weight: Feature penalty weight (L2 Norm)
+            prob_ppl_weight: Perplexity Loss with respect to probabilities during quantization
+            logit_temp: Temperature normalization applied in loss.
+            reduce: Reduce loss via sum reduction (Default true)
+        """
+        super().__init__()
+        self.feature_loss_weight = feature_loss_weight
+        self.prob_ppl_weight = prob_ppl_weight
+        self.logit_temp = logit_temp
+        self.reduce = reduce
+
+    @typecheck()
+    def forward(
+        self,
+        logits: torch.tensor,
+        targets: torch.tensor,
+        negatives: torch.tensor,
+        prob_ppl_loss: torch.tensor,
+        feature_loss: torch.tensor,
+    ) -> [torch.tensor, torch.tensor, torch.tensor]:
+        """
+        Args:
+            logits: Model activations
+            targets: The true target quantized representations
+            negatives: Sampled negatives from the quantizer codebooks. Sampled from all other timesteps.
+            feature_loss: Feature penalty (L2 Norm)
+            prob_ppl_loss: Perplexity Loss with respect to probs in quantization
+        Returns:
+            output loss values, feature loss, prob_ppl loss (after scaling).
+        """
+
+        # Calculate similarity between logits and all targets, returning FxBxT
+        similarity_scores = self._calculate_similarity(logits, negatives, targets)
+
+        # Create targets of size B*T
+        similarity_targets = logits.new_zeros(similarity_scores.size(1) * similarity_scores.size(2), dtype=torch.long)
+
+        # Transpose similarity scores to (T*B)xF for loss
+        similarity_scores = similarity_scores.transpose(0, 2)
+        similarity_scores = similarity_scores.reshape(-1, similarity_scores.size(-1))
+
+        loss = F.cross_entropy(similarity_scores, similarity_targets, reduction="sum" if self.reduce else "none")
+
+        sample_size = similarity_targets.numel()
+
+        if self.prob_ppl_weight != 0:
+            prob_ppl_loss = self.prob_ppl_weight * prob_ppl_loss * sample_size
+            loss += prob_ppl_loss
+
+        if self.feature_loss_weight != 0:
+            feature_loss = self.feature_loss_weight * feature_loss * sample_size
+            loss += feature_loss
+
+        return loss, feature_loss, prob_ppl_loss
+
+    def _calculate_similarity(self, logits, negatives, targets):
+        neg_is_pos = (targets == negatives).all(-1)
+        targets = targets.unsqueeze(0)
+        targets = torch.cat([targets, negatives], dim=0)
+        logits = torch.cosine_similarity(logits.float(), targets.float(), dim=-1).type_as(logits)
+        logits /= self.logit_temp
+        if neg_is_pos.any():
+            logits[1:][neg_is_pos] = float("-inf")
+        return logits

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

@@ -0,0 +1,13 @@
+# 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.

tts/nemo/collections/asr/metrics/rnnt_wer.py

@@ -0,0 +1,478 @@
+# 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 abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import List, Optional, Union
+
+import editdistance
+import torch
+from torchmetrics import Metric
+
+from nemo.collections.asr.parts.submodules import rnnt_beam_decoding as beam_decode
+from nemo.collections.asr.parts.submodules import rnnt_greedy_decoding as greedy_decode
+from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis, NBestHypotheses
+from nemo.utils import logging
+
+__all__ = ['RNNTDecoding', 'RNNTWER']
+
+
+class AbstractRNNTDecoding(ABC):
+    """
+    Used for performing RNN-T auto-regressive decoding of the Decoder+Joint network given the encoder state.
+
+    Args:
+        decoding_cfg: A dict-like object which contains the following key-value pairs.
+            strategy: str value which represents the type of decoding that can occur.
+                Possible values are :
+                -   greedy, greedy_batch (for greedy decoding).
+                -   beam, tsd, alsd (for beam search decoding).
+
+            compute_hypothesis_token_set: A bool flag, which determines whether to compute a list of decoded
+                tokens as well as the decoded string. Default is False in order to avoid double decoding
+                unless required.
+
+            preserve_alignments: Bool flag which preserves the history of logprobs generated during
+                decoding (sample / batched). When set to true, the Hypothesis will contain
+                the non-null value for `logprobs` in it. Here, `logprobs` is a List of torch.Tensors.
+
+                In order to obtain this hypothesis, please utilize `rnnt_decoder_predictions_tensor` function
+                with the `return_hypotheses` flag set to True.
+
+                The length of the list corresponds to the Acoustic Length (T).
+                Each value in the list (Ti) is a torch.Tensor (U), representing 1 or more targets from a vocabulary.
+                U is the number of target tokens for the current timestep Ti.
+
+            The config may further contain the following sub-dictionaries:
+            "greedy":
+                max_symbols: int, describing the maximum number of target tokens to decode per
+                    timestep during greedy decoding. Setting to larger values allows longer sentences
+                    to be decoded, at the cost of increased execution time.
+
+            "beam":
+                beam_size: int, defining the beam size for beam search. Must be >= 1.
+                    If beam_size == 1, will perform cached greedy search. This might be slightly different
+                    results compared to the greedy search above.
+
+                score_norm: optional bool, whether to normalize the returned beam score in the hypotheses.
+                    Set to True by default.
+
+                return_best_hypothesis: optional bool, whether to return just the best hypothesis or all of the
+                    hypotheses after beam search has concluded. This flag is set by default.
+
+                tsd_max_sym_exp: optional int, determines number of symmetric expansions of the target symbols
+                    per timestep of the acoustic model. Larger values will allow longer sentences to be decoded,
+                    at increased cost to execution time.
+
+                alsd_max_target_len: optional int or float, determines the potential maximum target sequence length.
+                    If an integer is provided, it can decode sequences of that particular maximum length.
+                    If a float is provided, it can decode sequences of int(alsd_max_target_len * seq_len),
+                    where seq_len is the length of the acoustic model output (T).
+
+                    NOTE:
+                        If a float is provided, it can be greater than 1!
+                        By default, a float of 2.0 is used so that a target sequence can be at most twice
+                        as long as the acoustic model output length T.
+
+        decoder: The Decoder/Prediction network module.
+        joint: The Joint network module.
+        blank_id: The id of the RNNT blank token.
+    """
+
+    def __init__(self, decoding_cfg, decoder, joint, blank_id: int):
+        super(AbstractRNNTDecoding, self).__init__()
+        self.cfg = decoding_cfg
+        self.blank_id = blank_id
+        self.compute_hypothesis_token_set = self.cfg.get("compute_hypothesis_token_set", False)
+        self.preserve_alignments = self.cfg.get('preserve_alignments', False)
+
+        possible_strategies = ['greedy', 'greedy_batch', 'beam', 'tsd', 'alsd']
+        if self.cfg.strategy not in possible_strategies:
+            raise ValueError(f"Decoding strategy must be one of {possible_strategies}")
+
+        if self.cfg.strategy == 'greedy':
+            self.decoding = greedy_decode.GreedyRNNTInfer(
+                decoder_model=decoder,
+                joint_model=joint,
+                blank_index=self.blank_id,
+                max_symbols_per_step=self.cfg.greedy.get('max_symbols', None),
+                preserve_alignments=self.preserve_alignments,
+            )
+
+        elif self.cfg.strategy == 'greedy_batch':
+            self.decoding = greedy_decode.GreedyBatchedRNNTInfer(
+                decoder_model=decoder,
+                joint_model=joint,
+                blank_index=self.blank_id,
+                max_symbols_per_step=self.cfg.greedy.get('max_symbols', None),
+                preserve_alignments=self.preserve_alignments,
+            )
+
+        elif self.cfg.strategy == 'beam':
+
+            self.decoding = beam_decode.BeamRNNTInfer(
+                decoder_model=decoder,
+                joint_model=joint,
+                beam_size=self.cfg.beam.beam_size,
+                return_best_hypothesis=decoding_cfg.beam.get('return_best_hypothesis', True),
+                search_type='default',
+                score_norm=self.cfg.beam.get('score_norm', True),
+                preserve_alignments=self.preserve_alignments,
+            )
+
+        elif self.cfg.strategy == 'tsd':
+
+            self.decoding = beam_decode.BeamRNNTInfer(
+                decoder_model=decoder,
+                joint_model=joint,
+                beam_size=self.cfg.beam.beam_size,
+                return_best_hypothesis=decoding_cfg.beam.get('return_best_hypothesis', True),
+                search_type='tsd',
+                score_norm=self.cfg.beam.get('score_norm', True),
+                tsd_max_sym_exp_per_step=self.cfg.beam.get('tsd_max_sym_exp', 50),
+                preserve_alignments=self.preserve_alignments,
+            )
+
+        elif self.cfg.strategy == 'alsd':
+
+            self.decoding = beam_decode.BeamRNNTInfer(
+                decoder_model=decoder,
+                joint_model=joint,
+                beam_size=self.cfg.beam.beam_size,
+                return_best_hypothesis=decoding_cfg.beam.get('return_best_hypothesis', True),
+                search_type='alsd',
+                score_norm=self.cfg.beam.get('score_norm', True),
+                alsd_max_target_len=self.cfg.beam.get('alsd_max_target_len', 2),
+                preserve_alignments=self.preserve_alignments,
+            )
+
+    def rnnt_decoder_predictions_tensor(
+        self, encoder_output: torch.Tensor, encoded_lengths: torch.Tensor, return_hypotheses: bool = False
+    ) -> (List[str], Optional[List[List[str]]], Optional[Union[Hypothesis, NBestHypotheses]]):
+        """
+        Decode an encoder output by autoregressive decoding of the Decoder+Joint networks.
+
+        Args:
+            encoder_output: torch.Tensor of shape [B, D, T].
+            encoded_lengths: torch.Tensor containing lengths of the padded encoder outputs. Shape [B].
+            return_hypotheses: bool. If set to True it will return list of Hypothesis or NBestHypotheses
+
+        Returns:
+            If `return_best_hypothesis` is set:
+                A tuple (hypotheses, None):
+                hypotheses - list of Hypothesis (best hypothesis per sample).
+                    Look at rnnt_utils.Hypothesis for more information.
+
+            If `return_best_hypothesis` is not set:
+                A tuple(hypotheses, all_hypotheses)
+                hypotheses - list of Hypothesis (best hypothesis per sample).
+                    Look at rnnt_utils.Hypothesis for more information.
+                all_hypotheses - list of NBestHypotheses. Each NBestHypotheses further contains a sorted
+                    list of all the hypotheses of the model per sample.
+                    Look at rnnt_utils.NBestHypotheses for more information.
+        """
+        # Compute hypotheses
+        with torch.no_grad():
+            hypotheses_list = self.decoding(
+                encoder_output=encoder_output, encoded_lengths=encoded_lengths
+            )  # type: [List[Hypothesis]]
+
+            # extract the hypotheses
+            hypotheses_list = hypotheses_list[0]  # type: List[Hypothesis]
+
+        prediction_list = hypotheses_list
+
+        if isinstance(prediction_list[0], NBestHypotheses):
+            hypotheses = []
+            all_hypotheses = []
+            for nbest_hyp in prediction_list:  # type: NBestHypotheses
+                n_hyps = nbest_hyp.n_best_hypotheses  # Extract all hypotheses for this sample
+                decoded_hyps = self.decode_hypothesis(n_hyps)  # type: List[str]
+                hypotheses.append(decoded_hyps[0])  # best hypothesis
+                all_hypotheses.append(decoded_hyps)
+            if return_hypotheses:
+                return hypotheses, all_hypotheses
+            best_hyp_text = [h.text for h in hypotheses]
+            all_hyp_text = [h.text for hh in all_hypotheses for h in hh]
+            return best_hyp_text, all_hyp_text
+        else:
+            hypotheses = self.decode_hypothesis(prediction_list)  # type: List[str]
+            if return_hypotheses:
+                return hypotheses, None
+            best_hyp_text = [h.text for h in hypotheses]
+            return best_hyp_text, None
+
+    def decode_hypothesis(self, hypotheses_list: List[Hypothesis]) -> List[Union[Hypothesis, NBestHypotheses]]:
+        """
+        Decode a list of hypotheses into a list of strings.
+
+        Args:
+            hypotheses_list: List of Hypothesis.
+
+        Returns:
+            A list of strings.
+        """
+        for ind in range(len(hypotheses_list)):
+            # Extract the integer encoded hypothesis
+            prediction = hypotheses_list[ind].y_sequence
+
+            if type(prediction) != list:
+                prediction = prediction.tolist()
+
+            # RNN-T sample level is already preprocessed by implicit CTC decoding
+            # Simply remove any blank tokens
+            prediction = [p for p in prediction if p != self.blank_id]
+
+            # De-tokenize the integer tokens
+            hypothesis = self.decode_tokens_to_str(prediction)
+            hypotheses_list[ind].text = hypothesis
+
+            if self.compute_hypothesis_token_set:
+                hypotheses_list[ind].tokens = self.decode_ids_to_tokens(prediction)
+        return hypotheses_list
+
+    @abstractmethod
+    def decode_tokens_to_str(self, tokens: List[int]) -> str:
+        """
+        Implemented by subclass in order to decoder a token id list into a string.
+
+        Args:
+            tokens: List of int representing the token ids.
+
+        Returns:
+            A decoded string.
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def decode_ids_to_tokens(self, tokens: List[int]) -> List[str]:
+        """
+        Implemented by subclass in order to decode a token id list into a token list.
+        A token list is the string representation of each token id.
+
+        Args:
+            tokens: List of int representing the token ids.
+
+        Returns:
+            A list of decoded tokens.
+        """
+        raise NotImplementedError()
+
+
+class RNNTDecoding(AbstractRNNTDecoding):
+    """
+    Used for performing RNN-T auto-regressive decoding of the Decoder+Joint network given the encoder state.
+
+    Args:
+        decoding_cfg: A dict-like object which contains the following key-value pairs.
+            strategy: str value which represents the type of decoding that can occur.
+                Possible values are :
+                -   greedy, greedy_batch (for greedy decoding).
+                -   beam, tsd, alsd (for beam search decoding).
+
+            compute_hypothesis_token_set: A bool flag, which determines whether to compute a list of decoded
+                tokens as well as the decoded string. Default is False in order to avoid double decoding
+                unless required.
+
+            preserve_alignments: Bool flag which preserves the history of logprobs generated during
+                decoding (sample / batched). When set to true, the Hypothesis will contain
+                the non-null value for `logprobs` in it. Here, `logprobs` is a List of torch.Tensors.
+
+                In order to obtain this hypothesis, please utilize `rnnt_decoder_predictions_tensor` function
+                with the `return_hypotheses` flag set to True.
+
+                The length of the list corresponds to the Acoustic Length (T).
+                Each value in the list (Ti) is a torch.Tensor (U), representing 1 or more targets from a vocabulary.
+                U is the number of target tokens for the current timestep Ti.
+
+            The config may further contain the following sub-dictionaries:
+            "greedy":
+                max_symbols: int, describing the maximum number of target tokens to decode per
+                    timestep during greedy decoding. Setting to larger values allows longer sentences
+                    to be decoded, at the cost of increased execution time.
+
+            "beam":
+                beam_size: int, defining the beam size for beam search. Must be >= 1.
+                    If beam_size == 1, will perform cached greedy search. This might be slightly different
+                    results compared to the greedy search above.
+
+                score_norm: optional bool, whether to normalize the returned beam score in the hypotheses.
+                    Set to True by default.
+
+                return_best_hypothesis: optional bool, whether to return just the best hypothesis or all of the
+                    hypotheses after beam search has concluded. This flag is set by default.
+
+                tsd_max_sym_exp: optional int, determines number of symmetric expansions of the target symbols
+                    per timestep of the acoustic model. Larger values will allow longer sentences to be decoded,
+                    at increased cost to execution time.
+
+                alsd_max_target_len: optional int or float, determines the potential maximum target sequence length.
+                    If an integer is provided, it can decode sequences of that particular maximum length.
+                    If a float is provided, it can decode sequences of int(alsd_max_target_len * seq_len),
+                    where seq_len is the length of the acoustic model output (T).
+
+                    NOTE:
+                        If a float is provided, it can be greater than 1!
+                        By default, a float of 2.0 is used so that a target sequence can be at most twice
+                        as long as the acoustic model output length T.
+
+        decoder: The Decoder/Prediction network module.
+        joint: The Joint network module.
+        vocabulary: The vocabulary (excluding the RNNT blank token) which will be used for decoding.
+    """
+
+    def __init__(
+        self, decoding_cfg, decoder, joint, vocabulary,
+    ):
+        blank_id = len(vocabulary)
+        self.labels_map = dict([(i, vocabulary[i]) for i in range(len(vocabulary))])
+
+        super(RNNTDecoding, self).__init__(decoding_cfg=decoding_cfg, decoder=decoder, joint=joint, blank_id=blank_id)
+
+    def decode_tokens_to_str(self, tokens: List[int]) -> str:
+        """
+        Implemented by subclass in order to decoder a token list into a string.
+
+        Args:
+            tokens: List of int representing the token ids.
+
+        Returns:
+            A decoded string.
+        """
+        hypothesis = ''.join(self.decode_ids_to_tokens(tokens))
+        return hypothesis
+
+    def decode_ids_to_tokens(self, tokens: List[int]) -> List[str]:
+        """
+        Implemented by subclass in order to decode a token id list into a token list.
+        A token list is the string representation of each token id.
+
+        Args:
+            tokens: List of int representing the token ids.
+
+        Returns:
+            A list of decoded tokens.
+        """
+        token_list = [self.labels_map[c] for c in tokens if c != self.blank_id]
+        return token_list
+
+
+class RNNTWER(Metric):
+    """
+    This metric computes numerator and denominator for Overall Word Error Rate (WER) between prediction and reference texts.
+    When doing distributed training/evaluation the result of res=WER(predictions, targets, target_lengths) calls
+    will be all-reduced between all workers using SUM operations.
+    Here contains two numbers res=[wer_numerator, wer_denominator]. WER=wer_numerator/wer_denominator.
+
+    If used with PytorchLightning LightningModule, include wer_numerator and wer_denominators inside validation_step results.
+    Then aggregate (sum) then at the end of validation epoch to correctly compute validation WER.
+
+    Example:
+        def validation_step(self, batch, batch_idx):
+            ...
+            wer_num, wer_denom = self.__wer(predictions, transcript, transcript_len)
+            return {'val_loss': loss_value, 'val_wer_num': wer_num, 'val_wer_denom': wer_denom}
+
+        def validation_epoch_end(self, outputs):
+            ...
+            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 = {'validation_loss': val_loss_mean, 'validation_avg_wer': wer_num / wer_denom}
+            return {'val_loss': val_loss_mean, 'log': tensorboard_logs}
+
+    Args:
+        decoding: RNNTDecoding object that will perform autoregressive decoding of the RNNT model.
+        batch_dim_index: Index of the batch dimension.
+        use_cer: Whether to use Character Error Rate isntead of Word Error Rate.
+        log_prediction: Whether to log a single decoded sample per call.
+
+    Returns:
+        res: a torch.Tensor object with two elements: [wer_numerator, wer_denominator]. To correctly compute average
+        text word error rate, compute wer=wer_numerator/wer_denominator
+    """
+
+    def __init__(
+        self, decoding: RNNTDecoding, batch_dim_index=0, use_cer=False, log_prediction=True, dist_sync_on_step=False
+    ):
+        super(RNNTWER, self).__init__(dist_sync_on_step=dist_sync_on_step, compute_on_step=False)
+        self.decoding = decoding
+        self.batch_dim_index = batch_dim_index
+        self.use_cer = use_cer
+        self.log_prediction = log_prediction
+        self.blank_id = self.decoding.blank_id
+        self.labels_map = self.decoding.labels_map
+
+        self.add_state("scores", default=torch.tensor(0), dist_reduce_fx='sum', persistent=False)
+        self.add_state("words", default=torch.tensor(0), dist_reduce_fx='sum', persistent=False)
+
+    def update(
+        self,
+        encoder_output: torch.Tensor,
+        encoded_lengths: torch.Tensor,
+        targets: torch.Tensor,
+        target_lengths: torch.Tensor,
+    ) -> torch.Tensor:
+        words = 0.0
+        scores = 0.0
+        references = []
+        with torch.no_grad():
+            # prediction_cpu_tensor = tensors[0].long().cpu()
+            targets_cpu_tensor = targets.long().cpu()
+            tgt_lenths_cpu_tensor = target_lengths.long().cpu()
+
+            # iterate over batch
+            for ind in range(targets_cpu_tensor.shape[self.batch_dim_index]):
+                tgt_len = tgt_lenths_cpu_tensor[ind].item()
+                target = targets_cpu_tensor[ind][:tgt_len].numpy().tolist()
+
+                reference = self.decoding.decode_tokens_to_str(target)
+                references.append(reference)
+
+            hypotheses, _ = self.decoding.rnnt_decoder_predictions_tensor(encoder_output, encoded_lengths)
+
+        if self.log_prediction:
+            logging.info(f"\n")
+            logging.info(f"reference :{references[0]}")
+            logging.info(f"predicted :{hypotheses[0]}")
+
+        for h, r in zip(hypotheses, references):
+            if self.use_cer:
+                h_list = list(h)
+                r_list = list(r)
+            else:
+                h_list = h.split()
+                r_list = r.split()
+            words += len(r_list)
+            # Compute Levenshtein's distance
+            scores += editdistance.eval(h_list, r_list)
+
+        self.scores += torch.tensor(scores, device=self.scores.device, dtype=self.scores.dtype)
+        self.words += torch.tensor(words, device=self.words.device, dtype=self.words.dtype)
+        # return torch.tensor([scores, words]).to(predictions.device)
+
+    def compute(self):
+        wer = self.scores.float() / self.words
+        return wer, self.scores.detach(), self.words.detach()
+
+
+@dataclass
+class RNNTDecodingConfig:
+    strategy: str = "greedy_batch"
+    compute_hypothesis_token_set: bool = False
+
+    # greedy decoding config
+    greedy: greedy_decode.GreedyRNNTInferConfig = greedy_decode.GreedyRNNTInferConfig()
+
+    # beam decoding config
+    beam: beam_decode.BeamRNNTInferConfig = beam_decode.BeamRNNTInferConfig(beam_size=4)

tts/nemo/collections/asr/metrics/rnnt_wer_bpe.py

@@ -0,0 +1,242 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from typing import List
+
+import editdistance
+import torch
+from torchmetrics import Metric
+
+from nemo.collections.asr.metrics.rnnt_wer import AbstractRNNTDecoding
+from nemo.collections.asr.parts.submodules import rnnt_beam_decoding as beam_decode
+from nemo.collections.asr.parts.submodules import rnnt_greedy_decoding as greedy_decode
+from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
+from nemo.utils import logging
+
+__all__ = ['RNNTBPEDecoding', 'RNNTBPEWER']
+
+
+class RNNTBPEDecoding(AbstractRNNTDecoding):
+    """
+    Used for performing RNN-T auto-regressive decoding of the Decoder+Joint network given the encoder state.
+
+    Args:
+        decoding_cfg: A dict-like object which contains the following key-value pairs.
+            strategy: str value which represents the type of decoding that can occur.
+                Possible values are :
+                -   greedy, greedy_batch (for greedy decoding).
+                -   beam, tsd, alsd (for beam search decoding).
+
+            compute_hypothesis_token_set: A bool flag, which determines whether to compute a list of decoded
+                tokens as well as the decoded string. Default is False in order to avoid double decoding
+                unless required.
+
+            preserve_alignments: Bool flag which preserves the history of logprobs generated during
+                decoding (sample / batched). When set to true, the Hypothesis will contain
+                the non-null value for `logprobs` in it. Here, `logprobs` is a List of torch.Tensors.
+
+                In order to obtain this hypothesis, please utilize `rnnt_decoder_predictions_tensor` function
+                with the `return_hypotheses` flag set to True.
+
+                The length of the list corresponds to the Acoustic Length (T).
+                Each value in the list (Ti) is a torch.Tensor (U), representing 1 or more targets from a vocabulary.
+                U is the number of target tokens for the current timestep Ti.
+
+            The config may further contain the following sub-dictionaries:
+            "greedy":
+                max_symbols: int, describing the maximum number of target tokens to decode per
+                    timestep during greedy decoding. Setting to larger values allows longer sentences
+                    to be decoded, at the cost of increased execution time.
+
+            "beam":
+                beam_size: int, defining the beam size for beam search. Must be >= 1.
+                    If beam_size == 1, will perform cached greedy search. This might be slightly different
+                    results compared to the greedy search above.
+
+                score_norm: optional bool, whether to normalize the returned beam score in the hypotheses.
+                    Set to True by default.
+
+                return_best_hypothesis: optional bool, whether to return just the best hypothesis or all of the
+                    hypotheses after beam search has concluded.
+
+                tsd_max_sym_exp: optional int, determines number of symmetric expansions of the target symbols
+                    per timestep of the acoustic model. Larger values will allow longer sentences to be decoded,
+                    at increased cost to execution time.
+
+                alsd_max_target_len: optional int or float, determines the potential maximum target sequence length.
+                    If an integer is provided, it can decode sequences of that particular maximum length.
+                    If a float is provided, it can decode sequences of int(alsd_max_target_len * seq_len),
+                    where seq_len is the length of the acoustic model output (T).
+
+                    NOTE:
+                        If a float is provided, it can be greater than 1!
+                        By default, a float of 2.0 is used so that a target sequence can be at most twice
+                        as long as the acoustic model output length T.
+
+        decoder: The Decoder/Prediction network module.
+        joint: The Joint network module.
+        tokenizer: The tokenizer which will be used for decoding.
+    """
+
+    def __init__(self, decoding_cfg, decoder, joint, tokenizer: TokenizerSpec):
+        blank_id = tokenizer.tokenizer.vocab_size
+        self.tokenizer = tokenizer
+
+        super(RNNTBPEDecoding, self).__init__(
+            decoding_cfg=decoding_cfg, decoder=decoder, joint=joint, blank_id=blank_id
+        )
+
+    def decode_tokens_to_str(self, tokens: List[int]) -> str:
+        """
+        Implemented by subclass in order to decoder a token list into a string.
+
+        Args:
+            tokens: List of int representing the token ids.
+
+        Returns:
+            A decoded string.
+        """
+        hypothesis = self.tokenizer.ids_to_text(tokens)
+        return hypothesis
+
+    def decode_ids_to_tokens(self, tokens: List[int]) -> List[str]:
+        """
+        Implemented by subclass in order to decode a token id list into a token list.
+        A token list is the string representation of each token id.
+
+        Args:
+            tokens: List of int representing the token ids.
+
+        Returns:
+            A list of decoded tokens.
+        """
+        token_list = self.tokenizer.ids_to_tokens(tokens)
+        return token_list
+
+
+class RNNTBPEWER(Metric):
+    """
+    This metric computes numerator and denominator for Overall Word Error Rate (WER) between prediction and reference texts.
+    When doing distributed training/evaluation the result of res=WER(predictions, targets, target_lengths) calls
+    will be all-reduced between all workers using SUM operations.
+    Here contains two numbers res=[wer_numerator, wer_denominator]. WER=wer_numerator/wer_denominator.
+
+    If used with PytorchLightning LightningModule, include wer_numerator and wer_denominators inside validation_step results.
+    Then aggregate (sum) then at the end of validation epoch to correctly compute validation WER.
+
+    Example:
+        def validation_step(self, batch, batch_idx):
+            ...
+            wer_num, wer_denom = self.__wer(predictions, transcript, transcript_len)
+            return {'val_loss': loss_value, 'val_wer_num': wer_num, 'val_wer_denom': wer_denom}
+
+        def validation_epoch_end(self, outputs):
+            ...
+            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 = {'validation_loss': val_loss_mean, 'validation_avg_wer': wer_num / wer_denom}
+            return {'val_loss': val_loss_mean, 'log': tensorboard_logs}
+
+    Args:
+        decoding: RNNTBPEDecoding object that will perform autoregressive decoding of the RNNT model.
+        batch_dim_index: Index of the batch dimension.
+        use_cer: Whether to use Character Error Rate isntead of Word Error Rate.
+        log_prediction: Whether to log a single decoded sample per call.
+
+    Returns:
+        res: a torch.Tensor object with two elements: [wer_numerator, wer_denominator]. To correctly compute average
+        text word error rate, compute wer=wer_numerator/wer_denominator
+    """
+
+    def __init__(
+        self,
+        decoding: RNNTBPEDecoding,
+        batch_dim_index=0,
+        use_cer: bool = False,
+        log_prediction: bool = True,
+        dist_sync_on_step=False,
+    ):
+        super(RNNTBPEWER, self).__init__(dist_sync_on_step=dist_sync_on_step, compute_on_step=False)
+        self.decoding = decoding
+        self.batch_dim_index = batch_dim_index
+        self.use_cer = use_cer
+        self.log_prediction = log_prediction
+        self.blank_id = self.decoding.blank_id
+        self.tokenizer = self.decoding.tokenizer
+
+        self.add_state("scores", default=torch.tensor(0), dist_reduce_fx='sum', persistent=False)
+        self.add_state("words", default=torch.tensor(0), dist_reduce_fx='sum', persistent=False)
+
+    def update(
+        self,
+        encoder_output: torch.Tensor,
+        encoded_lengths: torch.Tensor,
+        targets: torch.Tensor,
+        target_lengths: torch.Tensor,
+    ) -> torch.Tensor:
+        words = 0.0
+        scores = 0.0
+        references = []
+        with torch.no_grad():
+            # prediction_cpu_tensor = tensors[0].long().cpu()
+            targets_cpu_tensor = targets.long().cpu()
+            tgt_lenths_cpu_tensor = target_lengths.long().cpu()
+
+            # iterate over batch
+            for ind in range(targets_cpu_tensor.shape[self.batch_dim_index]):
+                tgt_len = tgt_lenths_cpu_tensor[ind].item()
+                target = targets_cpu_tensor[ind][:tgt_len].numpy().tolist()
+                reference = self.decoding.decode_tokens_to_str(target)
+                references.append(reference)
+
+            hypotheses, _ = self.decoding.rnnt_decoder_predictions_tensor(encoder_output, encoded_lengths)
+
+        if self.log_prediction:
+            logging.info(f"\n")
+            logging.info(f"reference :{references[0]}")
+            logging.info(f"predicted :{hypotheses[0]}")
+
+        for h, r in zip(hypotheses, references):
+            if self.use_cer:
+                h_list = list(h)
+                r_list = list(r)
+            else:
+                h_list = h.split()
+                r_list = r.split()
+            words += len(r_list)
+            # Compute Levenshtein's distance
+            scores += editdistance.eval(h_list, r_list)
+
+        del hypotheses
+
+        self.scores += torch.tensor(scores, device=self.scores.device, dtype=self.scores.dtype)
+        self.words += torch.tensor(words, device=self.words.device, dtype=self.words.dtype)
+        # return torch.tensor([scores, words]).to(predictions.device)
+
+    def compute(self):
+        wer = self.scores.float() / self.words
+        return wer, self.scores.detach(), self.words.detach()
+
+
+@dataclass
+class RNNTBPEDecodingConfig:
+    strategy: str = "greedy_batch"
+    compute_hypothesis_token_set: bool = False
+
+    # greedy decoding config
+    greedy: greedy_decode.GreedyRNNTInferConfig = greedy_decode.GreedyRNNTInferConfig()
+
+    # beam decoding config
+    beam: beam_decode.BeamRNNTInferConfig = beam_decode.BeamRNNTInferConfig(beam_size=4)

tts/nemo/collections/asr/metrics/wer.py

@@ -0,0 +1,247 @@
+# 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 typing import List
+
+import editdistance
+import torch
+from torchmetrics import Metric
+
+from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
+from nemo.utils import logging
+
+__all__ = ['word_error_rate', 'WER']
+
+
+def word_error_rate(hypotheses: List[str], references: List[str], use_cer=False) -> float:
+    """
+    Computes Average Word Error rate between two texts represented as
+    corresponding lists of string. Hypotheses and references must have same
+    length.
+    Args:
+      hypotheses: list of hypotheses
+      references: list of references
+      use_cer: bool, set True to enable cer
+    Returns:
+      (float) average word error rate
+    """
+    scores = 0
+    words = 0
+    if len(hypotheses) != len(references):
+        raise ValueError(
+            "In word error rate calculation, hypotheses and reference"
+            " lists must have the same number of elements. But I got:"
+            "{0} and {1} correspondingly".format(len(hypotheses), len(references))
+        )
+    for h, r in zip(hypotheses, references):
+        if use_cer:
+            h_list = list(h)
+            r_list = list(r)
+        else:
+            h_list = h.split()
+            r_list = r.split()
+        words += len(r_list)
+        scores += editdistance.eval(h_list, r_list)
+    if words != 0:
+        wer = 1.0 * scores / words
+    else:
+        wer = float('inf')
+    return wer
+
+
+class WER(Metric):
+    """
+    This metric computes numerator and denominator for Overall Word Error Rate (WER) between prediction and reference texts.
+    When doing distributed training/evaluation the result of res=WER(predictions, targets, target_lengths) calls
+    will be all-reduced between all workers using SUM operations.
+    Here contains two numbers res=[wer_numerator, wer_denominator]. WER=wer_numerator/wer_denominator.
+
+    If used with PytorchLightning LightningModule, include wer_numerator and wer_denominators inside validation_step results.
+    Then aggregate (sum) then at the end of validation epoch to correctly compute validation WER.
+
+    Example:
+        def validation_step(self, batch, batch_idx):
+            ...
+            wer_num, wer_denom = self.__wer(predictions, transcript, transcript_len)
+            return {'val_loss': loss_value, 'val_wer_num': wer_num, 'val_wer_denom': wer_denom}
+
+        def validation_epoch_end(self, outputs):
+            ...
+            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 = {'validation_loss': val_loss_mean, 'validation_avg_wer': wer_num / wer_denom}
+            return {'val_loss': val_loss_mean, 'log': tensorboard_logs}
+
+    Args:
+        vocabulary: List of strings that describes the vocabulary of the dataset.
+        batch_dim_index: Index of the batch dimension.
+        use_cer: Whether to use Character Error Rate instead of Word Error Rate.
+        ctc_decode: Whether to use CTC decoding or not. Currently, must be set.
+        log_prediction: Whether to log a single decoded sample per call.
+
+    Returns:
+        res: a torch.Tensor object with two elements: [wer_numerator, wer_denominator]. To correctly compute average
+        text word error rate, compute wer=wer_numerator/wer_denominator
+    """
+
+    def __init__(
+        self,
+        vocabulary,
+        batch_dim_index=0,
+        use_cer=False,
+        ctc_decode=True,
+        log_prediction=True,
+        dist_sync_on_step=False,
+    ):
+        super().__init__(dist_sync_on_step=dist_sync_on_step, compute_on_step=False)
+        self.batch_dim_index = batch_dim_index
+        self.blank_id = len(vocabulary)
+        self.labels_map = dict([(i, vocabulary[i]) for i in range(len(vocabulary))])
+        self.use_cer = use_cer
+        self.ctc_decode = ctc_decode
+        self.log_prediction = log_prediction
+
+        self.add_state("scores", default=torch.tensor(0), dist_reduce_fx='sum', persistent=False)
+        self.add_state("words", default=torch.tensor(0), dist_reduce_fx='sum', persistent=False)
+
+    def ctc_decoder_predictions_tensor(
+        self, predictions: torch.Tensor, predictions_len: torch.Tensor = None, return_hypotheses: bool = False,
+    ) -> List[str]:
+        """
+        Decodes a sequence of labels to words
+
+        Args:
+            predictions: A torch.Tensor of shape [Batch, Time] of integer indices that correspond
+                to the index of some character in the label set.
+            predictions_len: Optional tensor of length `Batch` which contains the integer lengths
+                of the sequence in the padded `predictions` tensor.
+            return_hypotheses: Bool flag whether to return just the decoding predictions of the model
+                or a Hypothesis object that holds information such as the decoded `text`,
+                the `alignment` of emited by the CTC Model, and the `length` of the sequence (if available).
+                May also contain the log-probabilities of the decoder (if this method is called via
+                transcribe())
+
+        Returns:
+            Either a list of str which represent the CTC decoded strings per sample,
+            or a list of Hypothesis objects containing additional information.
+        """
+        hypotheses = []
+        # Drop predictions to CPU
+        prediction_cpu_tensor = predictions.long().cpu()
+        # iterate over batch
+        for ind in range(prediction_cpu_tensor.shape[self.batch_dim_index]):
+            prediction = prediction_cpu_tensor[ind].detach().numpy().tolist()
+            if predictions_len is not None:
+                prediction = prediction[: predictions_len[ind]]
+            # CTC decoding procedure
+            decoded_prediction = []
+            previous = self.blank_id
+            for p in prediction:
+                if (p != previous or previous == self.blank_id) and p != self.blank_id:
+                    decoded_prediction.append(p)
+                previous = p
+
+            text = self.decode_tokens_to_str(decoded_prediction)
+
+            if not return_hypotheses:
+                hypothesis = text
+            else:
+                hypothesis = Hypothesis(
+                    y_sequence=None,
+                    score=-1.0,
+                    text=text,
+                    alignments=prediction,
+                    length=predictions_len[ind] if predictions_len is not None else 0,
+                )
+
+            hypotheses.append(hypothesis)
+        return hypotheses
+
+    def decode_tokens_to_str(self, tokens: List[int]) -> str:
+        """
+        Implemented in order to decoder a token list into a string.
+
+        Args:
+            tokens: List of int representing the token ids.
+
+        Returns:
+            A decoded string.
+        """
+        hypothesis = ''.join(self.decode_ids_to_tokens(tokens))
+        return hypothesis
+
+    def decode_ids_to_tokens(self, tokens: List[int]) -> List[str]:
+        """
+        Implemented in order to decode a token id list into a token list.
+        A token list is the string representation of each token id.
+
+        Args:
+            tokens: List of int representing the token ids.
+
+        Returns:
+            A list of decoded tokens.
+        """
+        token_list = [self.labels_map[c] for c in tokens if c != self.blank_id]
+        return token_list
+
+    def update(
+        self,
+        predictions: torch.Tensor,
+        targets: torch.Tensor,
+        target_lengths: torch.Tensor,
+        predictions_lengths: torch.Tensor = None,
+    ) -> torch.Tensor:
+        words = 0.0
+        scores = 0.0
+        references = []
+        with torch.no_grad():
+            # prediction_cpu_tensor = tensors[0].long().cpu()
+            targets_cpu_tensor = targets.long().cpu()
+            tgt_lenths_cpu_tensor = target_lengths.long().cpu()
+
+            # iterate over batch
+            for ind in range(targets_cpu_tensor.shape[self.batch_dim_index]):
+                tgt_len = tgt_lenths_cpu_tensor[ind].item()
+                target = targets_cpu_tensor[ind][:tgt_len].numpy().tolist()
+                reference = self.decode_tokens_to_str(target)
+                references.append(reference)
+            if self.ctc_decode:
+                hypotheses = self.ctc_decoder_predictions_tensor(predictions, predictions_lengths)
+            else:
+                raise NotImplementedError("Implement me if you need non-CTC decode on predictions")
+
+        if self.log_prediction:
+            logging.info(f"\n")
+            logging.info(f"reference:{references[0]}")
+            logging.info(f"predicted:{hypotheses[0]}")
+
+        for h, r in zip(hypotheses, references):
+            if self.use_cer:
+                h_list = list(h)
+                r_list = list(r)
+            else:
+                h_list = h.split()
+                r_list = r.split()
+            words += len(r_list)
+            # Compute Levenstein's distance
+            scores += editdistance.eval(h_list, r_list)
+
+        self.scores = torch.tensor(scores, device=self.scores.device, dtype=self.scores.dtype)
+        self.words = torch.tensor(words, device=self.words.device, dtype=self.words.dtype)
+        # return torch.tensor([scores, words]).to(predictions.device)
+
+    def compute(self):
+        scores = self.scores.detach().float()
+        words = self.words.detach().float()
+        return scores / words, scores, words

tts/nemo/collections/asr/metrics/wer_bpe.py

@@ -0,0 +1,210 @@
+# 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 typing import List
+
+import editdistance
+import torch
+from torchmetrics import Metric
+
+from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
+from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
+from nemo.utils import logging
+
+
+class WERBPE(Metric):
+    """
+    This metric computes numerator and denominator for Overall Word Error Rate for BPE tokens (WER-BPE) between prediction and reference texts.
+    When doing distributed training/evaluation the result of res=WERBPE(predictions, targets, target_lengths) calls
+    will be all-reduced between all workers using SUM operations.
+    Here contains two numbers res=[wer_numerator, wer_denominator]. WERBPE=wer_numerator/wer_denominator.
+
+    If used with PytorchLightning LightningModule, include wer_numerator and wer_denominators inside validation_step results.
+    Then aggregate (sum) then at the end of validation epoch to correctly compute validation WER.
+
+    Example:
+       def validation_step(self, batch, batch_idx):
+           ...
+           wer_num, wer_denom = self.__wer(predictions, transcript, transcript_len)
+           return {'val_loss': loss_value, 'val_wer_num': wer_num, 'val_wer_denom': wer_denom}
+
+       def validation_epoch_end(self, outputs):
+           ...
+           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 = {'validation_loss': val_loss_mean, 'validation_avg_wer': wer_num / wer_denom}
+           return {'val_loss': val_loss_mean, 'log': tensorboard_logs}
+
+    Args:
+       vocabulary: NeMo tokenizer object, which inherits from TokenizerSpec.
+       batch_dim_index: Index of the batch dimension.
+       use_cer: Whether to compute word-error-rate or character-error-rate.
+       ctc_decode: Whether to perform CTC decode.
+       log_prediction: Whether to log a single decoded sample per call.
+
+    Returns:
+       res: a torch.Tensor object with two elements: [wer_numerator, wer_denominators]. To correctly compute average
+       text word error rate, compute wer=wer_numerator/wer_denominators
+    """
+
+    def __init__(
+        self,
+        tokenizer: TokenizerSpec,
+        batch_dim_index=0,
+        use_cer=False,
+        ctc_decode=True,
+        log_prediction=True,
+        dist_sync_on_step=False,
+    ):
+        super().__init__(dist_sync_on_step=dist_sync_on_step, compute_on_step=False)
+        self.tokenizer = tokenizer
+        self.batch_dim_index = batch_dim_index
+        self.blank_id = tokenizer.tokenizer.vocab_size
+        self.use_cer = use_cer
+        self.ctc_decode = ctc_decode
+        self.log_prediction = log_prediction
+
+        self.add_state("scores", default=torch.tensor(0), dist_reduce_fx='sum', persistent=False)
+        self.add_state("words", default=torch.tensor(0), dist_reduce_fx='sum', persistent=False)
+
+    def ctc_decoder_predictions_tensor(
+        self, predictions: torch.Tensor, predictions_len: torch.Tensor = None, return_hypotheses: bool = False
+    ) -> List[str]:
+        """
+        Decodes a sequence of labels to words
+
+        Args:
+            predictions: A torch.Tensor of shape [Batch, Time] of integer indices that correspond
+                to the index of some character in the vocabulary of the tokenizer.
+            predictions_len: Optional tensor of length `Batch` which contains the integer lengths
+                of the sequence in the padded `predictions` tensor.
+            return_hypotheses: Bool flag whether to return just the decoding predictions of the model
+                or a Hypothesis object that holds information such as the decoded `text`,
+                the `alignment` of emited by the CTC Model, and the `length` of the sequence (if available).
+                May also contain the log-probabilities of the decoder (if this method is called via
+                transcribe()) inside `y_sequence`, otherwise it is set None as it is a duplicate of
+                `alignments`.
+
+        Returns:
+            Either a list of str which represent the CTC decoded strings per sample,
+            or a list of Hypothesis objects containing additional information.
+        """
+        hypotheses = []
+        # Drop predictions to CPU
+        prediction_cpu_tensor = predictions.long().cpu()
+        # iterate over batch
+        for ind in range(prediction_cpu_tensor.shape[self.batch_dim_index]):
+            prediction = prediction_cpu_tensor[ind].detach().numpy().tolist()
+            if predictions_len is not None:
+                prediction = prediction[: predictions_len[ind]]
+            # CTC decoding procedure
+            decoded_prediction = []
+            previous = self.blank_id
+            for p in prediction:
+                if (p != previous or previous == self.blank_id) and p != self.blank_id:
+                    decoded_prediction.append(p)
+                previous = p
+
+            text = self.decode_tokens_to_str(decoded_prediction)
+
+            if not return_hypotheses:
+                hypothesis = text
+            else:
+                hypothesis = Hypothesis(
+                    y_sequence=None,  # logprob info added by transcribe method
+                    score=-1.0,
+                    text=text,
+                    alignments=prediction,
+                    length=predictions_len[ind] if predictions_len is not None else 0,
+                )
+            hypotheses.append(hypothesis)
+        return hypotheses
+
+    def decode_tokens_to_str(self, tokens: List[int]) -> str:
+        """
+        Implemented in order to decoder a token list into a string.
+
+        Args:
+            tokens: List of int representing the token ids.
+
+        Returns:
+            A decoded string.
+        """
+        hypothesis = self.tokenizer.ids_to_text(tokens)
+        return hypothesis
+
+    def decode_ids_to_tokens(self, tokens: List[int]) -> List[str]:
+        """
+        Implemented in order to decode a token id list into a token list.
+        A token list is the string representation of each token id.
+
+        Args:
+            tokens: List of int representing the token ids.
+
+        Returns:
+            A list of decoded tokens.
+        """
+        token_list = self.tokenizer.ids_to_tokens(tokens)
+        return token_list
+
+    def update(
+        self,
+        predictions: torch.Tensor,
+        targets: torch.Tensor,
+        target_lengths: torch.Tensor,
+        predictions_lengths: torch.Tensor = None,
+    ):
+        words = 0.0
+        scores = 0.0
+        references = []
+        with torch.no_grad():
+            # prediction_cpu_tensor = tensors[0].long().cpu()
+            targets_cpu_tensor = targets.long().cpu()
+            tgt_lenths_cpu_tensor = target_lengths.long().cpu()
+
+            # iterate over batch
+            for ind in range(targets_cpu_tensor.shape[self.batch_dim_index]):
+                tgt_len = tgt_lenths_cpu_tensor[ind].item()
+                target = targets_cpu_tensor[ind][:tgt_len].numpy().tolist()
+                reference = self.decode_tokens_to_str(target)
+                references.append(reference)
+            if self.ctc_decode:
+                hypotheses = self.ctc_decoder_predictions_tensor(predictions, predictions_lengths)
+            else:
+                raise NotImplementedError("Implement me if you need non-CTC decode on predictions")
+
+        if self.log_prediction:
+            logging.info(f"\n")
+            logging.info(f"reference:{references[0]}")
+            logging.info(f"predicted:{hypotheses[0]}")
+
+        for h, r in zip(hypotheses, references):
+            if self.use_cer:
+                h_list = list(h)
+                r_list = list(r)
+            else:
+                h_list = h.split()
+                r_list = r.split()
+            words += len(r_list)
+            # Compute Levenstein's distance
+            scores += editdistance.eval(h_list, r_list)
+
+        self.scores = torch.tensor(scores, device=self.scores.device, dtype=self.scores.dtype)
+        self.words = torch.tensor(words, device=self.words.device, dtype=self.words.dtype)
+        # return torch.tensor([scores, words]).to(predictions.device)
+
+    def compute(self):
+        scores = self.scores.detach().float()
+        words = self.words.detach().float()
+        return scores / words, scores, words

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),
+        )

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

@@ -0,0 +1,30 @@
+# 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.modules.audio_preprocessing import (
+    AudioToMelSpectrogramPreprocessor,
+    AudioToMFCCPreprocessor,
+    CropOrPadSpectrogramAugmentation,
+    SpectrogramAugmentation,
+)
+from nemo.collections.asr.modules.beam_search_decoder import BeamSearchDecoderWithLM
+from nemo.collections.asr.modules.conformer_encoder import ConformerEncoder
+from nemo.collections.asr.modules.conv_asr import (
+    ConvASRDecoder,
+    ConvASRDecoderClassification,
+    ConvASREncoder,
+    SpeakerDecoder,
+)
+from nemo.collections.asr.modules.lstm_decoder import LSTMDecoder
+from nemo.collections.asr.modules.rnnt import RNNTDecoder, RNNTJoint

tts/nemo/collections/asr/modules/audio_preprocessing.py

@@ -0,0 +1,643 @@
+# 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 math
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import torch
+from packaging import version
+
+from nemo.collections.asr.parts.numba.spec_augment import SpecAugmentNumba, spec_augment_launch_heuristics
+from nemo.collections.asr.parts.preprocessing.features import FilterbankFeatures
+from nemo.collections.asr.parts.submodules.spectr_augment import SpecAugment, SpecCutout
+from nemo.core.classes import NeuralModule, typecheck
+from nemo.core.neural_types import (
+    AudioSignal,
+    LengthsType,
+    MelSpectrogramType,
+    MFCCSpectrogramType,
+    NeuralType,
+    SpectrogramType,
+)
+from nemo.core.utils import numba_utils
+from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
+from nemo.utils import logging
+
+try:
+    import torchaudio
+    import torchaudio.functional
+    import torchaudio.transforms
+
+    TORCHAUDIO_VERSION = version.parse(torchaudio.__version__)
+    TORCHAUDIO_VERSION_MIN = version.parse('0.5')
+
+    HAVE_TORCHAUDIO = True
+except ModuleNotFoundError:
+    HAVE_TORCHAUDIO = False
+
+__all__ = [
+    'AudioToMelSpectrogramPreprocessor',
+    'AudioToMFCCPreprocessor',
+    'SpectrogramAugmentation',
+    'CropOrPadSpectrogramAugmentation',
+]
+
+
+class AudioPreprocessor(NeuralModule, ABC):
+    """
+        An interface for Neural Modules that performs audio pre-processing,
+        transforming the wav files to features.
+    """
+
+    def __init__(self, win_length, hop_length):
+        super().__init__()
+
+        self.win_length = win_length
+        self.hop_length = hop_length
+
+        self.torch_windows = {
+            'hann': torch.hann_window,
+            'hamming': torch.hamming_window,
+            'blackman': torch.blackman_window,
+            'bartlett': torch.bartlett_window,
+            'ones': torch.ones,
+            None: torch.ones,
+        }
+
+    @typecheck()
+    @torch.no_grad()
+    def forward(self, input_signal, length):
+        processed_signal, processed_length = self.get_features(input_signal, length)
+
+        return processed_signal, processed_length
+
+    @abstractmethod
+    def get_features(self, input_signal, length):
+        # Called by forward(). Subclasses should implement this.
+        pass
+
+
+class AudioToMelSpectrogramPreprocessor(AudioPreprocessor):
+    """Featurizer module that converts wavs to mel spectrograms.
+        We don't use torchaudio's implementation here because the original
+        implementation is not the same, so for the sake of backwards-compatibility
+        this will use the old FilterbankFeatures for now.
+        Args:
+            sample_rate (int): Sample rate of the input audio data.
+                Defaults to 16000
+            window_size (float): Size of window for fft in seconds
+                Defaults to 0.02
+            window_stride (float): Stride of window for fft in seconds
+                Defaults to 0.01
+            n_window_size (int): Size of window for fft in samples
+                Defaults to None. Use one of window_size or n_window_size.
+            n_window_stride (int): Stride of window for fft in samples
+                Defaults to None. Use one of window_stride or n_window_stride.
+            window (str): Windowing function for fft. can be one of ['hann',
+                'hamming', 'blackman', 'bartlett']
+                Defaults to "hann"
+            normalize (str): Can be one of ['per_feature', 'all_features']; all
+                other options disable feature normalization. 'all_features'
+                normalizes the entire spectrogram to be mean 0 with std 1.
+                'pre_features' normalizes per channel / freq instead.
+                Defaults to "per_feature"
+            n_fft (int): Length of FT window. If None, it uses the smallest power
+                of 2 that is larger than n_window_size.
+                Defaults to None
+            preemph (float): Amount of pre emphasis to add to audio. Can be
+                disabled by passing None.
+                Defaults to 0.97
+            features (int): Number of mel spectrogram freq bins to output.
+                Defaults to 64
+            lowfreq (int): Lower bound on mel basis in Hz.
+                Defaults to 0
+            highfreq  (int): Lower bound on mel basis in Hz.
+                Defaults to None
+            log (bool): Log features.
+                Defaults to True
+            log_zero_guard_type(str): Need to avoid taking the log of zero. There
+                are two options: "add" or "clamp".
+                Defaults to "add".
+            log_zero_guard_value(float, or str): Add or clamp requires the number
+                to add with or clamp to. log_zero_guard_value can either be a float
+                or "tiny" or "eps". torch.finfo is used if "tiny" or "eps" is
+                passed.
+                Defaults to 2**-24.
+            dither (float): Amount of white-noise dithering.
+                Defaults to 1e-5
+            pad_to (int): Ensures that the output size of the time dimension is
+                a multiple of pad_to.
+                Defaults to 16
+            frame_splicing (int): Defaults to 1
+            exact_pad (bool): If True, sets stft center to False and adds padding, such that num_frames = audio_length
+                // hop_length. Defaults to False.
+            stft_exact_pad (bool): If True, uses pytorch_stft and convolutions with
+                padding such that num_frames = num_samples / hop_length. If False,
+                stft_conv will be used to determine how stft will be performed.
+                Defaults to False. TODO:This feature is deprecated and will be removed in 1.1.0
+            stft_conv (bool): If True, uses pytorch_stft and convolutions. If
+                False, uses torch.stft. TODO:This feature is deprecated and will be removed in 1.1.0
+                Defaults to False
+            pad_value (float): The value that shorter mels are padded with.
+                Defaults to 0
+            mag_power (float): The power that the linear spectrogram is raised to
+                prior to multiplication with mel basis.
+                Defaults to 2 for a power spec
+        """
+
+    def save_to(self, save_path: str):
+        pass
+
+    @classmethod
+    def restore_from(cls, restore_path: str):
+        pass
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input ports.
+        """
+        return {
+            "input_signal": NeuralType(('B', 'T'), AudioSignal(freq=self._sample_rate)),
+            "length": NeuralType(tuple('B'), LengthsType()),
+        }
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output ports.
+        processed_signal:
+            0: AxisType(BatchTag)
+            1: AxisType(MelSpectrogramSignalTag)
+            2: AxisType(ProcessedTimeTag)
+        processed_length:
+            0: AxisType(BatchTag)
+        """
+        return {
+            "processed_signal": NeuralType(('B', 'D', 'T'), MelSpectrogramType()),
+            "processed_length": NeuralType(tuple('B'), LengthsType()),
+        }
+
+    def __init__(
+        self,
+        sample_rate=16000,
+        window_size=0.02,
+        window_stride=0.01,
+        n_window_size=None,
+        n_window_stride=None,
+        window="hann",
+        normalize="per_feature",
+        n_fft=None,
+        preemph=0.97,
+        features=64,
+        lowfreq=0,
+        highfreq=None,
+        log=True,
+        log_zero_guard_type="add",
+        log_zero_guard_value=2 ** -24,
+        dither=1e-5,
+        pad_to=16,
+        frame_splicing=1,
+        exact_pad=False,
+        stft_exact_pad=False,
+        stft_conv=False,
+        pad_value=0,
+        mag_power=2.0,
+    ):
+        super().__init__(n_window_size, n_window_stride)
+
+        self._sample_rate = sample_rate
+        if window_size and n_window_size:
+            raise ValueError(f"{self} received both window_size and " f"n_window_size. Only one should be specified.")
+        if window_stride and n_window_stride:
+            raise ValueError(
+                f"{self} received both window_stride and " f"n_window_stride. Only one should be specified."
+            )
+        if window_size:
+            n_window_size = int(window_size * self._sample_rate)
+        if window_stride:
+            n_window_stride = int(window_stride * self._sample_rate)
+
+        self.featurizer = FilterbankFeatures(
+            sample_rate=self._sample_rate,
+            n_window_size=n_window_size,
+            n_window_stride=n_window_stride,
+            window=window,
+            normalize=normalize,
+            n_fft=n_fft,
+            preemph=preemph,
+            nfilt=features,
+            lowfreq=lowfreq,
+            highfreq=highfreq,
+            log=log,
+            log_zero_guard_type=log_zero_guard_type,
+            log_zero_guard_value=log_zero_guard_value,
+            dither=dither,
+            pad_to=pad_to,
+            frame_splicing=frame_splicing,
+            exact_pad=exact_pad,
+            stft_exact_pad=stft_exact_pad,
+            stft_conv=stft_conv,
+            pad_value=pad_value,
+            mag_power=mag_power,
+        )
+
+    def get_features(self, input_signal, length):
+        return self.featurizer(input_signal, length)
+
+    @property
+    def filter_banks(self):
+        return self.featurizer.filter_banks
+
+
+class AudioToMFCCPreprocessor(AudioPreprocessor):
+    """Preprocessor that converts wavs to MFCCs.
+    Uses torchaudio.transforms.MFCC.
+    Args:
+        sample_rate: The sample rate of the audio.
+            Defaults to 16000.
+        window_size: Size of window for fft in seconds. Used to calculate the
+            win_length arg for mel spectrogram.
+            Defaults to 0.02
+        window_stride: Stride of window for fft in seconds. Used to caculate
+            the hop_length arg for mel spect.
+            Defaults to 0.01
+        n_window_size: Size of window for fft in samples
+            Defaults to None. Use one of window_size or n_window_size.
+        n_window_stride: Stride of window for fft in samples
+            Defaults to None. Use one of window_stride or n_window_stride.
+        window: Windowing function for fft. can be one of ['hann',
+            'hamming', 'blackman', 'bartlett', 'none', 'null'].
+            Defaults to 'hann'
+        n_fft: Length of FT window. If None, it uses the smallest power of 2
+            that is larger than n_window_size.
+            Defaults to None
+        lowfreq (int): Lower bound on mel basis in Hz.
+            Defaults to 0
+        highfreq  (int): Lower bound on mel basis in Hz.
+            Defaults to None
+        n_mels: Number of mel filterbanks.
+            Defaults to 64
+        n_mfcc: Number of coefficients to retain
+            Defaults to 64
+        dct_type: Type of discrete cosine transform to use
+        norm: Type of norm to use
+        log: Whether to use log-mel spectrograms instead of db-scaled.
+            Defaults to True.
+    """
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input ports.
+        """
+        return {
+            "input_signal": NeuralType(('B', 'T'), AudioSignal(freq=self._sample_rate)),
+            "length": NeuralType(tuple('B'), LengthsType()),
+        }
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output ports.
+        """
+        return {
+            "processed_signal": NeuralType(('B', 'D', 'T'), MFCCSpectrogramType()),
+            "processed_length": NeuralType(tuple('B'), LengthsType()),
+        }
+
+    def save_to(self, save_path: str):
+        pass
+
+    @classmethod
+    def restore_from(cls, restore_path: str):
+        pass
+
+    def __init__(
+        self,
+        sample_rate=16000,
+        window_size=0.02,
+        window_stride=0.01,
+        n_window_size=None,
+        n_window_stride=None,
+        window='hann',
+        n_fft=None,
+        lowfreq=0.0,
+        highfreq=None,
+        n_mels=64,
+        n_mfcc=64,
+        dct_type=2,
+        norm='ortho',
+        log=True,
+    ):
+        self._sample_rate = sample_rate
+        if not HAVE_TORCHAUDIO:
+            logging.error('Could not import torchaudio. Some features might not work.')
+
+            raise ModuleNotFoundError(
+                "torchaudio is not installed but is necessary for "
+                "AudioToMFCCPreprocessor. We recommend you try "
+                "building it from source for the PyTorch version you have."
+            )
+        if window_size and n_window_size:
+            raise ValueError(f"{self} received both window_size and " f"n_window_size. Only one should be specified.")
+        if window_stride and n_window_stride:
+            raise ValueError(
+                f"{self} received both window_stride and " f"n_window_stride. Only one should be specified."
+            )
+        # Get win_length (n_window_size) and hop_length (n_window_stride)
+        if window_size:
+            n_window_size = int(window_size * self._sample_rate)
+        if window_stride:
+            n_window_stride = int(window_stride * self._sample_rate)
+
+        super().__init__(n_window_size, n_window_stride)
+
+        mel_kwargs = {}
+
+        mel_kwargs['f_min'] = lowfreq
+        mel_kwargs['f_max'] = highfreq
+        mel_kwargs['n_mels'] = n_mels
+
+        mel_kwargs['n_fft'] = n_fft or 2 ** math.ceil(math.log2(n_window_size))
+
+        mel_kwargs['win_length'] = n_window_size
+        mel_kwargs['hop_length'] = n_window_stride
+
+        # Set window_fn. None defaults to torch.ones.
+        window_fn = self.torch_windows.get(window, None)
+        if window_fn is None:
+            raise ValueError(
+                f"Window argument for AudioProcessor is invalid: {window}."
+                f"For no window function, use 'ones' or None."
+            )
+        mel_kwargs['window_fn'] = window_fn
+
+        # Use torchaudio's implementation of MFCCs as featurizer
+        self.featurizer = torchaudio.transforms.MFCC(
+            sample_rate=self._sample_rate,
+            n_mfcc=n_mfcc,
+            dct_type=dct_type,
+            norm=norm,
+            log_mels=log,
+            melkwargs=mel_kwargs,
+        )
+
+    def get_features(self, input_signal, length):
+        features = self.featurizer(input_signal)
+        seq_len = torch.ceil(length.to(torch.float32) / self.hop_length).to(dtype=torch.long)
+        return features, seq_len
+
+
+class SpectrogramAugmentation(NeuralModule):
+    """
+    Performs time and freq cuts in one of two ways.
+    SpecAugment zeroes out vertical and horizontal sections as described in
+    SpecAugment (https://arxiv.org/abs/1904.08779). Arguments for use with
+    SpecAugment are `freq_masks`, `time_masks`, `freq_width`, and `time_width`.
+    SpecCutout zeroes out rectangulars as described in Cutout
+    (https://arxiv.org/abs/1708.04552). Arguments for use with Cutout are
+    `rect_masks`, `rect_freq`, and `rect_time`.
+    Args:
+        freq_masks (int): how many frequency segments should be cut.
+            Defaults to 0.
+        time_masks (int): how many time segments should be cut
+            Defaults to 0.
+        freq_width (int): maximum number of frequencies to be cut in one
+            segment.
+            Defaults to 10.
+        time_width (int): maximum number of time steps to be cut in one
+            segment
+            Defaults to 10.
+        rect_masks (int): how many rectangular masks should be cut
+            Defaults to 0.
+        rect_freq (int): maximum size of cut rectangles along the frequency
+            dimension
+            Defaults to 5.
+        rect_time (int): maximum size of cut rectangles along the time
+            dimension
+            Defaults to 25.
+    """
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input types
+        """
+        return {
+            "input_spec": NeuralType(('B', 'D', 'T'), SpectrogramType()),
+            "length": NeuralType(tuple('B'), LengthsType(), optional=True),
+        }
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output types
+        """
+        return {"augmented_spec": NeuralType(('B', 'D', 'T'), SpectrogramType())}
+
+    def __init__(
+        self,
+        freq_masks=0,
+        time_masks=0,
+        freq_width=10,
+        time_width=10,
+        rect_masks=0,
+        rect_time=5,
+        rect_freq=20,
+        rng=None,
+        mask_value=0.0,
+        use_numba_spec_augment: bool = True,
+    ):
+        super().__init__()
+
+        if rect_masks > 0:
+            self.spec_cutout = SpecCutout(rect_masks=rect_masks, rect_time=rect_time, rect_freq=rect_freq, rng=rng,)
+            # self.spec_cutout.to(self._device)
+        else:
+            self.spec_cutout = lambda input_spec: input_spec
+
+        if freq_masks + time_masks > 0:
+            self.spec_augment = SpecAugment(
+                freq_masks=freq_masks,
+                time_masks=time_masks,
+                freq_width=freq_width,
+                time_width=time_width,
+                rng=rng,
+                mask_value=mask_value,
+            )
+        else:
+            self.spec_augment = lambda input_spec: input_spec
+
+        # Check if numba is supported, and use a Numba kernel if it is
+        if use_numba_spec_augment and numba_utils.numba_cuda_is_supported(__NUMBA_MINIMUM_VERSION__):
+            self.spec_augment_numba = SpecAugmentNumba(
+                freq_masks=freq_masks,
+                time_masks=time_masks,
+                freq_width=freq_width,
+                time_width=time_width,
+                rng=rng,
+                mask_value=mask_value,
+            )
+        else:
+            self.spec_augment_numba = None
+
+    @typecheck()
+    def forward(self, input_spec, length=None):
+        augmented_spec = self.spec_cutout(input_spec=input_spec)
+
+        # To run the Numba kernel, correct numba version is required as well as
+        # tensor must be on GPU and length must be provided
+        if self.spec_augment_numba is not None and spec_augment_launch_heuristics(augmented_spec, length):
+            augmented_spec = self.spec_augment_numba(input_spec=augmented_spec, length=length)
+        else:
+            augmented_spec = self.spec_augment(input_spec=augmented_spec)
+        return augmented_spec
+
+
+class CropOrPadSpectrogramAugmentation(NeuralModule):
+    """
+    Pad or Crop the incoming Spectrogram to a certain shape.
+    Args:
+        audio_length (int): the final number of timesteps that is required.
+            The signal will be either padded or cropped temporally to this
+            size.
+    """
+
+    def __init__(self, audio_length):
+        super(CropOrPadSpectrogramAugmentation, self).__init__()
+        self.audio_length = audio_length
+
+    @typecheck()
+    @torch.no_grad()
+    def forward(self, input_signal, length):
+        image = input_signal
+        num_images = image.shape[0]
+
+        audio_length = self.audio_length
+        image_len = image.shape[-1]
+
+        # Crop long signal
+        if image_len > audio_length:  # randomly slice
+            cutout_images = []
+            offset = torch.randint(low=0, high=image_len - audio_length + 1, size=[num_images])
+
+            for idx, offset in enumerate(offset):
+                cutout_images.append(image[idx : idx + 1, :, offset : offset + audio_length])
+
+            image = torch.cat(cutout_images, dim=0)
+            del cutout_images
+
+        else:  # symmetrically pad short signal with zeros
+            pad_left = (audio_length - image_len) // 2
+            pad_right = (audio_length - image_len) // 2
+
+            if (audio_length - image_len) % 2 == 1:
+                pad_right += 1
+
+            image = torch.nn.functional.pad(image, [pad_left, pad_right], mode="constant", value=0)
+
+        # Replace dynamic length sequences with static number of timesteps
+        length = (length * 0) + audio_length
+
+        return image, length
+
+    @property
+    def input_types(self):
+        """Returns definitions of module output ports.
+        """
+        return {
+            "input_signal": NeuralType(('B', 'D', 'T'), SpectrogramType()),
+            "length": NeuralType(tuple('B'), LengthsType()),
+        }
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output ports.
+        """
+        return {
+            "processed_signal": NeuralType(('B', 'D', 'T'), SpectrogramType()),
+            "processed_length": NeuralType(tuple('B'), LengthsType()),
+        }
+
+    def save_to(self, save_path: str):
+        pass
+
+    @classmethod
+    def restore_from(cls, restore_path: str):
+        pass
+
+
+@dataclass
+class AudioToMelSpectrogramPreprocessorConfig:
+    _target_: str = "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor"
+    sample_rate: int = 16000
+    window_size: float = 0.02
+    window_stride: float = 0.01
+    n_window_size: Optional[int] = None
+    n_window_stride: Optional[int] = None
+    window: str = "hann"
+    normalize: str = "per_feature"
+    n_fft: Optional[int] = None
+    preemph: float = 0.97
+    features: int = 64
+    lowfreq: int = 0
+    highfreq: Optional[int] = None
+    log: bool = True
+    log_zero_guard_type: str = "add"
+    log_zero_guard_value: float = 2 ** -24
+    dither: float = 1e-5
+    pad_to: int = 16
+    frame_splicing: int = 1
+    exact_pad: bool = False
+    stft_exact_pad: bool = False
+    stft_conv: bool = False
+    pad_value: int = 0
+    mag_power: float = 2.0
+
+
+@dataclass
+class AudioToMFCCPreprocessorConfig:
+    _target_: str = 'nemo.collections.asr.modules.AudioToMFCCPreprocessor'
+    sample_rate: int = 16000
+    window_size: float = 0.02
+    window_stride: float = 0.01
+    n_window_size: Optional[int] = None
+    n_window_stride: Optional[int] = None
+    window: str = 'hann'
+    n_fft: Optional[int] = None
+    lowfreq: Optional[float] = 0.0
+    highfreq: Optional[float] = None
+    n_mels: int = 64
+    n_mfcc: int = 64
+    dct_type: int = 2
+    norm: str = 'ortho'
+    log: bool = True
+
+
+@dataclass
+class SpectrogramAugmentationConfig:
+    _target_: str = "nemo.collections.asr.modules.SpectrogramAugmentation"
+    freq_masks: int = 0
+    time_masks: int = 0
+    freq_width: int = 0
+    time_width: Optional[Any] = 0
+    rect_masks: int = 0
+    rect_time: int = 0
+    rect_freq: int = 0
+    mask_value: float = 0
+    rng: Optional[Any] = None  # random.Random() type
+    use_numba_spec_augment: bool = True
+
+
+@dataclass
+class CropOrPadSpectrogramAugmentationConfig:
+    audio_length: int
+    _target_: str = "nemo.collections.asr.modules.CropOrPadSpectrogramAugmentation"

tts/nemo/collections/asr/modules/beam_search_decoder.py

@@ -0,0 +1,103 @@
+# 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 torch
+
+from nemo.core.classes import NeuralModule, typecheck
+from nemo.core.neural_types import LengthsType, LogprobsType, NeuralType, PredictionsType
+
+
+class BeamSearchDecoderWithLM(NeuralModule):
+    """Neural Module that does CTC beam search with a N-gram language model.
+    It takes a batch of log_probabilities. Note the bigger the batch, the
+    better as processing is parallelized. Outputs a list of size batch_size.
+    Each element in the list is a list of size beam_search, and each element
+    in that list is a tuple of (final_log_prob, hyp_string).
+    Args:
+        vocab (list): List of characters that can be output by the ASR model. For English, this is the 28 character set
+            {a-z '}. The CTC blank symbol is automatically added.
+        beam_width (int): Size of beams to keep and expand upon. Larger beams result in more accurate but slower
+            predictions
+        alpha (float): The amount of importance to place on the N-gram language model. Larger alpha means more
+            importance on the LM and less importance on the acoustic model.
+        beta (float): A penalty term given to longer word sequences. Larger beta will result in shorter sequences.
+        lm_path (str): Path to N-gram language model
+        num_cpus (int): Number of CPUs to use
+        cutoff_prob (float): Cutoff probability in vocabulary pruning, default 1.0, no pruning
+        cutoff_top_n (int): Cutoff number in pruning, only top cutoff_top_n characters with highest probs in
+            vocabulary will be used in beam search, default 40.
+        input_tensor (bool): Set to True if you intend to pass PyTorch Tensors, set to False if you intend to pass
+            NumPy arrays.
+    """
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input ports.
+        """
+        return {
+            "log_probs": NeuralType(('B', 'T', 'D'), LogprobsType()),
+            "log_probs_length": NeuralType(tuple('B'), LengthsType()),
+        }
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output ports.
+        """
+        return {"predictions": NeuralType(('B', 'T'), PredictionsType())}
+
+    def __init__(
+        self, vocab, beam_width, alpha, beta, lm_path, num_cpus, cutoff_prob=1.0, cutoff_top_n=40, input_tensor=False
+    ):
+
+        try:
+            from ctc_decoders import Scorer, ctc_beam_search_decoder_batch
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError(
+                "BeamSearchDecoderWithLM requires the installation of ctc_decoders "
+                "from scripts/asr_language_modeling/ngram_lm/install_beamsearch_decoders.sh"
+            )
+
+        super().__init__()
+
+        if lm_path is not None:
+            self.scorer = Scorer(alpha, beta, model_path=lm_path, vocabulary=vocab)
+        else:
+            self.scorer = None
+        self.beam_search_func = ctc_beam_search_decoder_batch
+        self.vocab = vocab
+        self.beam_width = beam_width
+        self.num_cpus = num_cpus
+        self.cutoff_prob = cutoff_prob
+        self.cutoff_top_n = cutoff_top_n
+        self.input_tensor = input_tensor
+
+    @typecheck(ignore_collections=True)
+    @torch.no_grad()
+    def forward(self, log_probs, log_probs_length):
+        probs_list = log_probs
+        if self.input_tensor:
+            probs = torch.exp(log_probs)
+            probs_list = []
+            for i, prob in enumerate(probs):
+                probs_list.append(prob[: log_probs_length[i], :])
+        res = self.beam_search_func(
+            probs_list,
+            self.vocab,
+            beam_size=self.beam_width,
+            num_processes=self.num_cpus,
+            ext_scoring_func=self.scorer,
+            cutoff_prob=self.cutoff_prob,
+            cutoff_top_n=self.cutoff_top_n,
+        )
+        return res

tts/nemo/collections/asr/modules/conformer_encoder.py

@@ -0,0 +1,255 @@
+# 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 math
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+
+from nemo.collections.asr.parts.submodules.conformer_modules import ConformerLayer
+from nemo.collections.asr.parts.submodules.multi_head_attention import PositionalEncoding, RelPositionalEncoding
+from nemo.collections.asr.parts.submodules.subsampling import ConvSubsampling
+from nemo.core.classes.common import typecheck
+from nemo.core.classes.exportable import Exportable
+from nemo.core.classes.module import NeuralModule
+from nemo.core.neural_types import AcousticEncodedRepresentation, LengthsType, NeuralType, SpectrogramType
+
+__all__ = ['ConformerEncoder']
+
+
+class ConformerEncoder(NeuralModule, Exportable):
+    """
+    The encoder for ASR model of Conformer.
+    Based on this paper:
+    'Conformer: Convolution-augmented Transformer for Speech Recognition' by Anmol Gulati et al.
+    https://arxiv.org/abs/2005.08100
+
+    Args:
+        feat_in (int): the size of feature channels
+        n_layers (int): number of layers of ConformerBlock
+        d_model (int): the hidden size of the model
+        feat_out (int): the size of the output features
+            Defaults to -1 (means feat_out is d_model)
+        subsampling (str): the method of subsampling, choices=['vggnet', 'striding']
+            Defaults to striding.
+        subsampling_factor (int): the subsampling factor which should be power of 2
+            Defaults to 4.
+        subsampling_conv_channels (int): the size of the convolutions in the subsampling module
+            Defaults to -1 which would set it to d_model.
+        ff_expansion_factor (int): the expansion factor in feed forward layers
+            Defaults to 4.
+        self_attention_model (str): type of the attention layer and positional encoding
+            'rel_pos': relative positional embedding and Transformer-XL
+            'abs_pos': absolute positional embedding and Transformer
+            default is rel_pos.
+        pos_emb_max_len (int): the maximum length of positional embeddings
+            Defaulst to 5000
+        n_heads (int): number of heads in multi-headed attention layers
+            Defaults to 4.
+        xscaling (bool): enables scaling the inputs to the multi-headed attention layers by sqrt(d_model)
+            Defaults to True.
+        untie_biases (bool): whether to not share (untie) the bias weights between layers of Transformer-XL
+            Defaults to True.
+        conv_kernel_size (int): the size of the convolutions in the convolutional modules
+            Defaults to 31.
+        dropout (float): the dropout rate used in all layers except the attention layers
+            Defaults to 0.1.
+        dropout_emb (float): the dropout rate used for the positional embeddings
+            Defaults to 0.1.
+        dropout_att (float): the dropout rate used for the attention layer
+            Defaults to 0.0.
+    """
+
+    def input_example(self):
+        """
+        Generates input examples for tracing etc.
+        Returns:
+            A tuple of input examples.
+        """
+        input_example = torch.randn(16, self._feat_in, 256).to(next(self.parameters()).device)
+        input_example_length = torch.randint(0, 256, (16,)).to(next(self.parameters()).device)
+        return tuple([input_example, input_example_length])
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input ports.
+        """
+        return OrderedDict(
+            {
+                "audio_signal": NeuralType(('B', 'D', 'T'), SpectrogramType()),
+                "length": NeuralType(tuple('B'), LengthsType()),
+            }
+        )
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output ports.
+        """
+        return OrderedDict(
+            {
+                "outputs": NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()),
+                "encoded_lengths": NeuralType(tuple('B'), LengthsType()),
+            }
+        )
+
+    def __init__(
+        self,
+        feat_in,
+        n_layers,
+        d_model,
+        feat_out=-1,
+        subsampling='striding',
+        subsampling_factor=4,
+        subsampling_conv_channels=-1,
+        ff_expansion_factor=4,
+        self_attention_model='rel_pos',
+        n_heads=4,
+        att_context_size=None,
+        xscaling=True,
+        untie_biases=True,
+        pos_emb_max_len=5000,
+        conv_kernel_size=31,
+        dropout=0.1,
+        dropout_emb=0.1,
+        dropout_att=0.0,
+    ):
+        super().__init__()
+
+        d_ff = d_model * ff_expansion_factor
+        self.d_model = d_model
+        self._feat_in = feat_in
+        self.scale = math.sqrt(self.d_model)
+        if att_context_size:
+            self.att_context_size = att_context_size
+        else:
+            self.att_context_size = [-1, -1]
+
+        if xscaling:
+            self.xscale = math.sqrt(d_model)
+        else:
+            self.xscale = None
+
+        if subsampling_conv_channels == -1:
+            subsampling_conv_channels = d_model
+        if subsampling and subsampling_factor > 1:
+            self.pre_encode = ConvSubsampling(
+                subsampling=subsampling,
+                subsampling_factor=subsampling_factor,
+                feat_in=feat_in,
+                feat_out=d_model,
+                conv_channels=subsampling_conv_channels,
+                activation=nn.ReLU(),
+            )
+            self._feat_out = d_model
+        else:
+            self._feat_out = d_model
+            self.pre_encode = nn.Linear(feat_in, d_model)
+
+        if not untie_biases and self_attention_model == "rel_pos":
+            d_head = d_model // n_heads
+            pos_bias_u = nn.Parameter(torch.Tensor(n_heads, d_head))
+            pos_bias_v = nn.Parameter(torch.Tensor(n_heads, d_head))
+            nn.init.zeros_(pos_bias_u)
+            nn.init.zeros_(pos_bias_v)
+        else:
+            pos_bias_u = None
+            pos_bias_v = None
+
+        if self_attention_model == "rel_pos":
+            self.pos_enc = RelPositionalEncoding(
+                d_model=d_model,
+                dropout_rate=dropout,
+                max_len=pos_emb_max_len,
+                xscale=self.xscale,
+                dropout_rate_emb=dropout_emb,
+            )
+        elif self_attention_model == "abs_pos":
+            pos_bias_u = None
+            pos_bias_v = None
+            self.pos_enc = PositionalEncoding(
+                d_model=d_model, dropout_rate=dropout, max_len=pos_emb_max_len, xscale=self.xscale
+            )
+        else:
+            raise ValueError(f"Not valid self_attention_model: '{self_attention_model}'!")
+
+        self.layers = nn.ModuleList()
+        for i in range(n_layers):
+            layer = ConformerLayer(
+                d_model=d_model,
+                d_ff=d_ff,
+                self_attention_model=self_attention_model,
+                n_heads=n_heads,
+                conv_kernel_size=conv_kernel_size,
+                dropout=dropout,
+                dropout_att=dropout_att,
+                pos_bias_u=pos_bias_u,
+                pos_bias_v=pos_bias_v,
+            )
+            self.layers.append(layer)
+
+        if feat_out > 0 and feat_out != self.output_dim:
+            self.out_proj = nn.Linear(self.feat_out, feat_out)
+            self._feat_out = feat_out
+        else:
+            self.out_proj = None
+            self._feat_out = d_model
+
+    @typecheck()
+    def forward(self, audio_signal, length=None):
+        if length is None:
+            length = torch.tensor(audio_signal.size(-1)).repeat(audio_signal.size(0)).to(audio_signal)
+        audio_signal = torch.transpose(audio_signal, 1, 2)
+
+        if isinstance(self.pre_encode, ConvSubsampling):
+            audio_signal, length = self.pre_encode(audio_signal, length)
+        else:
+            audio_signal = self.pre_encode(audio_signal)
+
+        audio_signal, pos_emb = self.pos_enc(audio_signal)
+        bs, xmax, idim = audio_signal.size()
+
+        # Create the self-attention and padding masks
+        pad_mask = self.make_pad_mask(length, max_time=xmax, device=audio_signal.device)
+        att_mask = pad_mask.unsqueeze(1).repeat([1, xmax, 1])
+        att_mask = att_mask & att_mask.transpose(1, 2)
+        if self.att_context_size[0] >= 0:
+            att_mask = att_mask.triu(diagonal=-self.att_context_size[0])
+        if self.att_context_size[1] >= 0:
+            att_mask = att_mask.tril(diagonal=self.att_context_size[1])
+        att_mask = ~att_mask
+        pad_mask = ~pad_mask
+
+        for lth, layer in enumerate(self.layers):
+            audio_signal = layer(x=audio_signal, att_mask=att_mask, pos_emb=pos_emb, pad_mask=pad_mask)
+
+        if self.out_proj is not None:
+            audio_signal = self.out_proj(audio_signal)
+
+        audio_signal = torch.transpose(audio_signal, 1, 2)
+        return audio_signal, length
+
+    @staticmethod
+    def make_pad_mask(seq_lens, max_time, device=None):
+        """Make masking for padding."""
+        bs = seq_lens.size(0)
+        seq_range = torch.arange(0, max_time, dtype=torch.int32)
+        seq_range_expand = seq_range.unsqueeze(0).expand(bs, max_time)
+        seq_lens = seq_lens.type(seq_range_expand.dtype).to(seq_range_expand.device)
+        seq_length_expand = seq_lens.unsqueeze(-1)
+        mask = seq_range_expand < seq_length_expand
+
+        if device:
+            mask = mask.to(device)
+        return mask

tts/nemo/collections/asr/modules/conv_asr.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.
+from collections import OrderedDict
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from omegaconf import MISSING, ListConfig, OmegaConf
+
+from nemo.collections.asr.parts.submodules.jasper import (
+    JasperBlock,
+    MaskedConv1d,
+    StatsPoolLayer,
+    init_weights,
+    jasper_activations,
+)
+from nemo.core.classes.common import typecheck
+from nemo.core.classes.exportable import Exportable
+from nemo.core.classes.module import NeuralModule
+from nemo.core.neural_types import (
+    AcousticEncodedRepresentation,
+    LengthsType,
+    LogitsType,
+    LogprobsType,
+    NeuralType,
+    SpectrogramType,
+)
+from nemo.utils import logging
+
+__all__ = ['ConvASRDecoder', 'ConvASREncoder', 'ConvASRDecoderClassification']
+
+
+class ConvASREncoder(NeuralModule, Exportable):
+    """
+    Convolutional encoder for ASR models. With this class you can implement JasperNet and QuartzNet models.
+
+    Based on these papers:
+        https://arxiv.org/pdf/1904.03288.pdf
+        https://arxiv.org/pdf/1910.10261.pdf
+    """
+
+    def _prepare_for_export(self, **kwargs):
+        m_count = 0
+        for m in self.modules():
+            if isinstance(m, MaskedConv1d):
+                m.use_mask = False
+                m_count += 1
+        Exportable._prepare_for_export(self, **kwargs)
+        logging.warning(f"Turned off {m_count} masked convolutions")
+
+    def input_example(self):
+        """
+        Generates input examples for tracing etc.
+        Returns:
+            A tuple of input examples.
+        """
+        input_example = torch.randn(16, self._feat_in, 256).to(next(self.parameters()).device)
+        return tuple([input_example])
+
+    @property
+    def disabled_deployment_input_names(self):
+        """Implement this method to return a set of input names disabled for export"""
+        return set(["length"])
+
+    @property
+    def disabled_deployment_output_names(self):
+        """Implement this method to return a set of output names disabled for export"""
+        return set(["encoded_lengths"])
+
+    def save_to(self, save_path: str):
+        pass
+
+    @classmethod
+    def restore_from(cls, restore_path: str):
+        pass
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input ports.
+        """
+        return OrderedDict(
+            {
+                "audio_signal": NeuralType(('B', 'D', 'T'), SpectrogramType()),
+                "length": NeuralType(tuple('B'), LengthsType()),
+            }
+        )
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output ports.
+        """
+        return OrderedDict(
+            {
+                "outputs": NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()),
+                "encoded_lengths": NeuralType(tuple('B'), LengthsType()),
+            }
+        )
+
+    def __init__(
+        self,
+        jasper,
+        activation: str,
+        feat_in: int,
+        normalization_mode: str = "batch",
+        residual_mode: str = "add",
+        norm_groups: int = -1,
+        conv_mask: bool = True,
+        frame_splicing: int = 1,
+        init_mode: Optional[str] = 'xavier_uniform',
+        quantize: bool = False,
+    ):
+        super().__init__()
+        if isinstance(jasper, ListConfig):
+            jasper = OmegaConf.to_container(jasper)
+
+        activation = jasper_activations[activation]()
+
+        # If the activation can be executed in place, do so.
+        if hasattr(activation, 'inplace'):
+            activation.inplace = True
+
+        feat_in = feat_in * frame_splicing
+
+        self._feat_in = feat_in
+
+        residual_panes = []
+        encoder_layers = []
+        self.dense_residual = False
+        for lcfg in jasper:
+            dense_res = []
+            if lcfg.get('residual_dense', False):
+                residual_panes.append(feat_in)
+                dense_res = residual_panes
+                self.dense_residual = True
+            groups = lcfg.get('groups', 1)
+            separable = lcfg.get('separable', False)
+            heads = lcfg.get('heads', -1)
+            residual_mode = lcfg.get('residual_mode', residual_mode)
+            se = lcfg.get('se', False)
+            se_reduction_ratio = lcfg.get('se_reduction_ratio', 8)
+            se_context_window = lcfg.get('se_context_size', -1)
+            se_interpolation_mode = lcfg.get('se_interpolation_mode', 'nearest')
+            kernel_size_factor = lcfg.get('kernel_size_factor', 1.0)
+            stride_last = lcfg.get('stride_last', False)
+            future_context = lcfg.get('future_context', -1)
+            encoder_layers.append(
+                JasperBlock(
+                    feat_in,
+                    lcfg['filters'],
+                    repeat=lcfg['repeat'],
+                    kernel_size=lcfg['kernel'],
+                    stride=lcfg['stride'],
+                    dilation=lcfg['dilation'],
+                    dropout=lcfg['dropout'],
+                    residual=lcfg['residual'],
+                    groups=groups,
+                    separable=separable,
+                    heads=heads,
+                    residual_mode=residual_mode,
+                    normalization=normalization_mode,
+                    norm_groups=norm_groups,
+                    activation=activation,
+                    residual_panes=dense_res,
+                    conv_mask=conv_mask,
+                    se=se,
+                    se_reduction_ratio=se_reduction_ratio,
+                    se_context_window=se_context_window,
+                    se_interpolation_mode=se_interpolation_mode,
+                    kernel_size_factor=kernel_size_factor,
+                    stride_last=stride_last,
+                    future_context=future_context,
+                    quantize=quantize,
+                )
+            )
+            feat_in = lcfg['filters']
+
+        self._feat_out = feat_in
+
+        self.encoder = torch.nn.Sequential(*encoder_layers)
+        self.apply(lambda x: init_weights(x, mode=init_mode))
+
+    @typecheck()
+    def forward(self, audio_signal, length=None):
+        s_input, length = self.encoder(([audio_signal], length))
+        if length is None:
+            return s_input[-1]
+
+        return s_input[-1], length
+
+
+class ConvASRDecoder(NeuralModule, Exportable):
+    """Simple ASR Decoder for use with CTC-based models such as JasperNet and QuartzNet
+
+     Based on these papers:
+        https://arxiv.org/pdf/1904.03288.pdf
+        https://arxiv.org/pdf/1910.10261.pdf
+        https://arxiv.org/pdf/2005.04290.pdf
+    """
+
+    def save_to(self, save_path: str):
+        pass
+
+    @classmethod
+    def restore_from(cls, restore_path: str):
+        pass
+
+    @property
+    def input_types(self):
+        return OrderedDict({"encoder_output": NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation())})
+
+    @property
+    def output_types(self):
+        return OrderedDict({"logprobs": NeuralType(('B', 'T', 'D'), LogprobsType())})
+
+    def __init__(self, feat_in, num_classes, init_mode="xavier_uniform", vocabulary=None):
+        super().__init__()
+
+        if vocabulary is not None:
+            if num_classes != len(vocabulary):
+                raise ValueError(
+                    f"If vocabulary is specified, it's length should be equal to the num_classes. Instead got: num_classes={num_classes} and len(vocabulary)={len(vocabulary)}"
+                )
+            self.__vocabulary = vocabulary
+        self._feat_in = feat_in
+        # Add 1 for blank char
+        self._num_classes = num_classes + 1
+
+        self.decoder_layers = torch.nn.Sequential(
+            torch.nn.Conv1d(self._feat_in, self._num_classes, kernel_size=1, bias=True)
+        )
+        self.apply(lambda x: init_weights(x, mode=init_mode))
+
+    @typecheck()
+    def forward(self, encoder_output):
+        return torch.nn.functional.log_softmax(self.decoder_layers(encoder_output).transpose(1, 2), dim=-1)
+
+    def input_example(self):
+        """
+        Generates input examples for tracing etc.
+        Returns:
+            A tuple of input examples.
+        """
+        bs = 8
+        seq = 64
+        input_example = torch.randn(bs, self._feat_in, seq).to(next(self.parameters()).device)
+        return tuple([input_example])
+
+    def _prepare_for_export(self, **kwargs):
+        m_count = 0
+        for m in self.modules():
+            if type(m).__name__ == "MaskedConv1d":
+                m.use_mask = False
+                m_count += 1
+        if m_count > 0:
+            logging.warning(f"Turned off {m_count} masked convolutions")
+        Exportable._prepare_for_export(self, **kwargs)
+
+    @property
+    def vocabulary(self):
+        return self.__vocabulary
+
+    @property
+    def num_classes_with_blank(self):
+        return self._num_classes
+
+
+class ConvASRDecoderClassification(NeuralModule, Exportable):
+    """Simple ASR Decoder for use with classification models such as JasperNet and QuartzNet
+
+     Based on these papers:
+        https://arxiv.org/pdf/2005.04290.pdf
+    """
+
+    def input_example(self):
+        """
+        Generates input examples for tracing etc.
+        Returns:
+            A tuple of input examples.
+        """
+        input_example = torch.randn(16, self._feat_in, 128).to(next(self.parameters()).device)
+        return tuple([input_example])
+
+    @property
+    def input_types(self):
+        return OrderedDict({"encoder_output": NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation())})
+
+    @property
+    def output_types(self):
+        return OrderedDict({"logits": NeuralType(('B', 'D'), LogitsType())})
+
+    def __init__(
+        self,
+        feat_in: int,
+        num_classes: int,
+        init_mode: Optional[str] = "xavier_uniform",
+        return_logits: bool = True,
+        pooling_type='avg',
+    ):
+        super().__init__()
+
+        self._feat_in = feat_in
+        self._return_logits = return_logits
+        self._num_classes = num_classes
+
+        if pooling_type == 'avg':
+            self.pooling = torch.nn.AdaptiveAvgPool1d(1)
+        elif pooling_type == 'max':
+            self.pooling = torch.nn.AdaptiveMaxPool1d(1)
+        else:
+            raise ValueError('Pooling type chosen is not valid. Must be either `avg` or `max`')
+
+        self.decoder_layers = torch.nn.Sequential(torch.nn.Linear(self._feat_in, self._num_classes, bias=True))
+        self.apply(lambda x: init_weights(x, mode=init_mode))
+
+    @typecheck()
+    def forward(self, encoder_output):
+        batch, in_channels, timesteps = encoder_output.size()
+
+        encoder_output = self.pooling(encoder_output).view(batch, in_channels)  # [B, C]
+        logits = self.decoder_layers(encoder_output)  # [B, num_classes]
+
+        if self._return_logits:
+            return logits
+
+        return torch.nn.functional.softmax(logits, dim=-1)
+
+    @property
+    def num_classes(self):
+        return self._num_classes
+
+
+class SpeakerDecoder(NeuralModule, Exportable):
+    """
+    Speaker Decoder creates the final neural layers that maps from the outputs
+    of Jasper Encoder to the embedding layer followed by speaker based softmax loss.
+    Args:
+        feat_in (int): Number of channels being input to this module
+        num_classes (int): Number of unique speakers in dataset
+        emb_sizes (list) : shapes of intermediate embedding layers (we consider speaker embbeddings from 1st of this layers)
+                Defaults to [1024,1024]
+        pool_mode (str) : Pooling stratergy type. options are 'gram','xvector','superVector'.
+                Defaults to 'xvector'
+        init_mode (str): Describes how neural network parameters are
+            initialized. Options are ['xavier_uniform', 'xavier_normal',
+            'kaiming_uniform','kaiming_normal'].
+            Defaults to "xavier_uniform".
+    """
+
+    def input_example(self):
+        """
+        Generates input examples for tracing etc.
+        Returns:
+            A tuple of input examples.
+        """
+        input_example = torch.randn(16, self.input_feat_in, 256).to(next(self.parameters()).device)
+        return tuple([input_example])
+
+    @property
+    def input_types(self):
+        return OrderedDict({"encoder_output": NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation())})
+
+    @property
+    def output_types(self):
+        return OrderedDict(
+            {
+                "logits": NeuralType(('B', 'D'), LogitsType()),
+                "embs": NeuralType(('B', 'D'), AcousticEncodedRepresentation()),
+            }
+        )
+
+    def __init__(
+        self, feat_in, num_classes, emb_sizes=None, pool_mode='xvector', angular=False, init_mode="xavier_uniform",
+    ):
+        super().__init__()
+        self.angular = angular
+        self.emb_id = 2
+        if self.angular:
+            bias = False
+        else:
+            bias = True
+
+        if type(emb_sizes) is str:
+            emb_sizes = emb_sizes.split(',')
+        elif type(emb_sizes) is int:
+            emb_sizes = [emb_sizes]
+        else:
+            emb_sizes = [512, 512]
+
+        self.input_feat_in = feat_in
+        self._num_classes = num_classes
+        self._pooling = StatsPoolLayer(feat_in=feat_in, pool_mode=pool_mode)
+        self._feat_in = self._pooling.feat_in
+
+        shapes = [self._feat_in]
+        for size in emb_sizes:
+            shapes.append(int(size))
+
+        emb_layers = []
+        for shape_in, shape_out in zip(shapes[:-1], shapes[1:]):
+            layer = self.affineLayer(shape_in, shape_out, learn_mean=False)
+            emb_layers.append(layer)
+
+        self.emb_layers = nn.ModuleList(emb_layers)
+
+        self.final = nn.Linear(shapes[-1], self._num_classes, bias=bias)
+
+        self.apply(lambda x: init_weights(x, mode=init_mode))
+
+    def affineLayer(self, inp_shape, out_shape, learn_mean=True):
+        layer = nn.Sequential(
+            nn.Linear(inp_shape, out_shape),
+            nn.BatchNorm1d(out_shape, affine=learn_mean, track_running_stats=True),
+            nn.ReLU(),
+        )
+
+        return layer
+
+    @typecheck()
+    def forward(self, encoder_output):
+        pool = self._pooling(encoder_output)
+        embs = []
+
+        for layer in self.emb_layers:
+            pool, emb = layer(pool), layer[: self.emb_id](pool)
+            embs.append(emb)
+
+        if self.angular:
+            for W in self.final.parameters():
+                W = F.normalize(W, p=2, dim=1)
+            pool = F.normalize(pool, p=2, dim=1)
+
+        out = self.final(pool)
+
+        return out, embs[-1]
+
+
+@dataclass
+class JasperEncoderConfig:
+    filters: int = MISSING
+    repeat: int = MISSING
+    kernel: List[int] = MISSING
+    stride: List[int] = MISSING
+    dilation: List[int] = MISSING
+    dropout: float = MISSING
+    residual: bool = MISSING
+
+    # Optional arguments
+    groups: int = 1
+    separable: bool = False
+    heads: int = -1
+    residual_mode: str = "add"
+    residual_dense: bool = False
+    se: bool = False
+    se_reduction_ratio: int = 8
+    se_context_size: int = -1
+    se_interpolation_mode: str = 'nearest'
+    kernel_size_factor: float = 1.0
+    stride_last: bool = False
+
+
+@dataclass
+class ConvASREncoderConfig:
+    _target_: str = 'nemo.collections.asr.modules.ConvASREncoder'
+    jasper: Optional[JasperEncoderConfig] = field(default_factory=list)
+    activation: str = MISSING
+    feat_in: int = MISSING
+    normalization_mode: str = "batch"
+    residual_mode: str = "add"
+    norm_groups: int = -1
+    conv_mask: bool = True
+    frame_splicing: int = 1
+    init_mode: Optional[str] = "xavier_uniform"
+
+
+@dataclass
+class ConvASRDecoderConfig:
+    _target_: str = 'nemo.collections.asr.modules.ConvASRDecoder'
+    feat_in: int = MISSING
+    num_classes: int = MISSING
+    init_mode: Optional[str] = "xavier_uniform"
+    vocabulary: Optional[List[str]] = field(default_factory=list)
+
+
+@dataclass
+class ConvASRDecoderClassificationConfig:
+    _target_: str = 'nemo.collections.asr.modules.ConvASRDecoderClassification'
+    feat_in: int = MISSING
+    num_classes: int = MISSING
+    init_mode: Optional[str] = "xavier_uniform"
+    return_logits: bool = True
+    pooling_type: str = 'avg'

tts/nemo/collections/asr/modules/lstm_decoder.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 collections import OrderedDict
+
+import torch
+import torch.nn as nn
+
+from nemo.core.classes.common import typecheck
+from nemo.core.classes.exportable import Exportable
+from nemo.core.classes.module import NeuralModule
+from nemo.core.neural_types import AcousticEncodedRepresentation, LogprobsType, NeuralType
+
+__all__ = ['LSTMDecoder']
+
+
+class LSTMDecoder(NeuralModule, Exportable):
+    """
+    Simple LSTM Decoder for ASR models
+    Args:
+        feat_in (int): size of the input features
+        num_classes (int): the size of the vocabulary
+        lstm_hidden_size (int): hidden size of the LSTM layers
+        vocabulary (vocab): The vocabulary
+        bidirectional (bool): default is False. Whether LSTMs are bidirectional or not
+        num_layers (int): default is 1. Number of LSTM layers stacked
+    """
+
+    @property
+    def input_types(self):
+        return OrderedDict({"encoder_output": NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation())})
+
+    @property
+    def output_types(self):
+        return OrderedDict({"logprobs": NeuralType(('B', 'T', 'D'), LogprobsType())})
+
+    def __init__(self, feat_in, num_classes, lstm_hidden_size, vocabulary=None, bidirectional=False, num_layers=1):
+        super().__init__()
+
+        if vocabulary is not None:
+            if num_classes != len(vocabulary):
+                raise ValueError(
+                    f"If vocabulary is specified, it's length should be equal to the num_classes. "
+                    f"Instead got: num_classes={num_classes} and len(vocabulary)={len(vocabulary)}"
+                )
+            self.__vocabulary = vocabulary
+        self._feat_in = feat_in
+        # Add 1 for blank char
+        self._num_classes = num_classes + 1
+
+        self.lstm_layer = nn.LSTM(
+            input_size=feat_in,
+            hidden_size=lstm_hidden_size,
+            num_layers=num_layers,
+            batch_first=True,
+            bidirectional=bidirectional,
+        )
+        self.linear_layer = torch.nn.Linear(in_features=lstm_hidden_size, out_features=self._num_classes)
+
+    @typecheck()
+    def forward(self, encoder_output):
+        output = encoder_output.transpose(1, 2)
+        output, _ = self.lstm_layer(output)
+        output = self.linear_layer(output)
+        return torch.nn.functional.log_softmax(output, dim=-1)
+
+    def input_example(self):
+        """
+        Generates input examples for tracing etc.
+        Returns:
+            A tuple of input examples.
+        """
+        bs = 8
+        seq = 64
+        input_example = torch.randn(bs, self._feat_in, seq).to(next(self.parameters()).device)
+        return tuple([input_example])
+
+    @property
+    def vocabulary(self):
+        return self.__vocabulary
+
+    @property
+    def num_classes_with_blank(self):
+        return self._num_classes

tts/nemo/collections/asr/modules/rnnt.py

@@ -0,0 +1,958 @@
+# 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 2017 Johns Hopkins University (Shinji Watanabe)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+
+from nemo.collections.asr.modules import rnnt_abstract
+from nemo.collections.asr.parts.utils import rnnt_utils
+from nemo.collections.common.parts import rnn
+from nemo.core.classes import typecheck
+from nemo.core.neural_types import (
+    AcousticEncodedRepresentation,
+    ElementType,
+    EmbeddedTextType,
+    LabelsType,
+    LengthsType,
+    LogprobsType,
+    LossType,
+    NeuralType,
+)
+from nemo.utils import logging
+
+
+class RNNTDecoder(rnnt_abstract.AbstractRNNTDecoder):
+    """A Recurrent Neural Network Transducer Decoder / Prediction Network (RNN-T Prediction Network).
+    An RNN-T Decoder/Prediction network, comprised of a stateful LSTM model.
+
+    Args:
+        prednet: A dict-like object which contains the following key-value pairs.
+            pred_hidden: int specifying the hidden dimension of the prediction net.
+            pred_rnn_layers: int specifying the number of rnn layers.
+
+            Optionally, it may also contain the following:
+            forget_gate_bias: float, set by default to 1.0, which constructs a forget gate
+                initialized to 1.0.
+                Reference:
+                [An Empirical Exploration of Recurrent Network Architectures](http://proceedings.mlr.press/v37/jozefowicz15.pdf)
+            t_max: int value, set to None by default. If an int is specified, performs Chrono Initialization
+                of the LSTM network, based on the maximum number of timesteps `t_max` expected during the course
+                of training.
+                Reference:
+                [Can recurrent neural networks warp time?](https://openreview.net/forum?id=SJcKhk-Ab)
+            weights_init_scale: Float scale of the weights after initialization. Setting to lower than one
+                sometimes helps reduce variance between runs.
+            hidden_hidden_bias_scale: Float scale for the hidden-to-hidden bias scale. Set to 0.0 for
+                the default behaviour.
+            dropout: float, set to 0.0 by default. Optional dropout applied at the end of the final LSTM RNN layer.
+
+        vocab_size: int, specifying the vocabulary size of the embedding layer of the Prediction network,
+            excluding the RNNT blank token.
+
+        normalization_mode: Can be either None, 'batch' or 'layer'. By default, is set to None.
+            Defines the type of normalization applied to the RNN layer.
+
+        random_state_sampling: bool, set to False by default. When set, provides normal-distribution
+            sampled state tensors instead of zero tensors during training.
+            Reference:
+            [Recognizing long-form speech using streaming end-to-end models](https://arxiv.org/abs/1910.11455)
+
+        blank_as_pad: bool, set to True by default. When set, will add a token to the Embedding layer of this
+            prediction network, and will treat this token as a pad token. In essence, the RNNT pad token will
+            be treated as a pad token, and the embedding layer will return a zero tensor for this token.
+
+            It is set by default as it enables various batch optimizations required for batched beam search.
+            Therefore, it is not recommended to disable this flag.
+    """
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input ports.
+        """
+        return {
+            "targets": NeuralType(('B', 'T'), LabelsType()),
+            "target_length": NeuralType(tuple('B'), LengthsType()),
+            "states": NeuralType(('D', 'B', 'D'), ElementType(), optional=True),
+        }
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output ports.
+        """
+        return {
+            "outputs": NeuralType(('B', 'D', 'T'), EmbeddedTextType()),
+            "encoded_lengths": NeuralType(tuple('B'), LengthsType()),
+        }
+
+    def __init__(
+        self,
+        prednet: Dict[str, Any],
+        vocab_size: int,
+        normalization_mode: Optional[str] = None,
+        random_state_sampling: bool = False,
+        blank_as_pad: bool = True,
+    ):
+        # Required arguments
+        self.pred_hidden = prednet['pred_hidden']
+        self.pred_rnn_layers = prednet["pred_rnn_layers"]
+        self.blank_idx = vocab_size
+
+        # Initialize the model (blank token increases vocab size by 1)
+        super().__init__(vocab_size=vocab_size, blank_idx=self.blank_idx, blank_as_pad=blank_as_pad)
+
+        # Optional arguments
+        forget_gate_bias = prednet.get('forget_gate_bias', 1.0)
+        t_max = prednet.get('t_max', None)
+        weights_init_scale = prednet.get('weights_init_scale', 1.0)
+        hidden_hidden_bias_scale = prednet.get('hidden_hidden_bias_scale', 0.0)
+        dropout = prednet.get('dropout', 0.0)
+        self.random_state_sampling = random_state_sampling
+
+        self.prediction = self._predict(
+            vocab_size=vocab_size,  # add 1 for blank symbol
+            pred_n_hidden=self.pred_hidden,
+            pred_rnn_layers=self.pred_rnn_layers,
+            forget_gate_bias=forget_gate_bias,
+            t_max=t_max,
+            norm=normalization_mode,
+            weights_init_scale=weights_init_scale,
+            hidden_hidden_bias_scale=hidden_hidden_bias_scale,
+            dropout=dropout,
+        )
+
+    @typecheck()
+    def forward(self, targets, target_length, states=None):
+        # y: (B, U)
+        y = rnn.label_collate(targets)
+
+        # state maintenance is unnecessary during training forward call
+        # to get state, use .predict() method.
+        g, _ = self.predict(y, state=states, add_sos=True)  # (B, U, D)
+        g = g.transpose(1, 2)  # (B, D, U)
+
+        return g, target_length
+
+    def predict(
+        self,
+        y: Optional[torch.Tensor] = None,
+        state: Optional[List[torch.Tensor]] = None,
+        add_sos: bool = True,
+        batch_size: Optional[int] = None,
+    ) -> (torch.Tensor, List[torch.Tensor]):
+        """
+        Stateful prediction of scores and state for a (possibly null) tokenset.
+        This method takes various cases into consideration :
+        - No token, no state - used for priming the RNN
+        - No token, state provided - used for blank token scoring
+        - Given token, states - used for scores + new states
+
+        Here:
+        B - batch size
+        U - label length
+        H - Hidden dimension size of RNN
+        L - Number of RNN layers
+
+        Args:
+            y: Optional torch tensor of shape [B, U] of dtype long which will be passed to the Embedding.
+                If None, creates a zero tensor of shape [B, 1, H] which mimics output of pad-token on Embedding.
+
+            state: An optional list of states for the RNN. Eg: For LSTM, it is the state list length is 2.
+                Each state must be a tensor of shape [L, B, H].
+                If None, and during training mode and `random_state_sampling` is set, will sample a
+                normal distribution tensor of the above shape. Otherwise, None will be passed to the RNN.
+
+            add_sos: bool flag, whether a zero vector describing a "start of signal" token should be
+                prepended to the above "y" tensor. When set, output size is (B, U + 1, H).
+
+            batch_size: An optional int, specifying the batch size of the `y` tensor.
+                Can be infered if `y` and `state` is None. But if both are None, then batch_size cannot be None.
+
+        Returns:
+            A tuple  (g, hid) such that -
+
+            If add_sos is False:
+                g: (B, U, H)
+                hid: (h, c) where h is the final sequence hidden state and c is the final cell state:
+                    h (tensor), shape (L, B, H)
+                    c (tensor), shape (L, B, H)
+
+            If add_sos is True:
+                g: (B, U + 1, H)
+                hid: (h, c) where h is the final sequence hidden state and c is the final cell state:
+                    h (tensor), shape (L, B, H)
+                    c (tensor), shape (L, B, H)
+
+        """
+        # Get device and dtype of current module
+        _p = next(self.parameters())
+        device = _p.device
+        dtype = _p.dtype
+
+        # If y is not None, it is of shape [B, U] with dtype long.
+        if y is not None:
+            if y.device != device:
+                y = y.to(device)
+
+            # (B, U) -> (B, U, H)
+            y = self.prediction["embed"](y)
+        else:
+            # Y is not provided, assume zero tensor with shape [B, 1, H] is required
+            # Emulates output of embedding of pad token.
+            if batch_size is None:
+                B = 1 if state is None else state[0].size(1)
+            else:
+                B = batch_size
+
+            y = torch.zeros((B, 1, self.pred_hidden), device=device, dtype=dtype)
+
+        # Prepend blank "start of sequence" symbol (zero tensor)
+        if add_sos:
+            B, U, H = y.shape
+            start = torch.zeros((B, 1, H), device=y.device, dtype=y.dtype)
+            y = torch.cat([start, y], dim=1).contiguous()  # (B, U + 1, H)
+        else:
+            start = None  # makes del call later easier
+
+        # If in training mode, and random_state_sampling is set,
+        # initialize state to random normal distribution tensor.
+        if state is None:
+            if self.random_state_sampling and self.training:
+                state = self.initialize_state(y)
+
+        # Forward step through RNN
+        y = y.transpose(0, 1)  # (U + 1, B, H)
+        g, hid = self.prediction["dec_rnn"](y, state)
+        g = g.transpose(0, 1)  # (B, U + 1, H)
+
+        del y, start, state
+        return g, hid
+
+    def _predict(
+        self,
+        vocab_size,
+        pred_n_hidden,
+        pred_rnn_layers,
+        forget_gate_bias,
+        t_max,
+        norm,
+        weights_init_scale,
+        hidden_hidden_bias_scale,
+        dropout,
+    ):
+        """
+        Prepare the trainable parameters of the Prediction Network.
+
+        Args:
+            vocab_size: Vocab size (excluding the blank token).
+            pred_n_hidden: Hidden size of the RNNs.
+            pred_rnn_layers: Number of RNN layers.
+            forget_gate_bias: Whether to perform unit forget gate bias.
+            t_max: Whether to perform Chrono LSTM init.
+            norm: Type of normalization to perform in RNN.
+            weights_init_scale: Float scale of the weights after initialization. Setting to lower than one
+                sometimes helps reduce variance between runs.
+            hidden_hidden_bias_scale: Float scale for the hidden-to-hidden bias scale. Set to 0.0 for
+                the default behaviour.
+            dropout: Whether to apply dropout to RNN.
+        """
+        if self.blank_as_pad:
+            embed = torch.nn.Embedding(vocab_size + 1, pred_n_hidden, padding_idx=self.blank_idx)
+        else:
+            embed = torch.nn.Embedding(vocab_size, pred_n_hidden)
+
+        layers = torch.nn.ModuleDict(
+            {
+                "embed": embed,
+                "dec_rnn": rnn.rnn(
+                    input_size=pred_n_hidden,
+                    hidden_size=pred_n_hidden,
+                    num_layers=pred_rnn_layers,
+                    norm=norm,
+                    forget_gate_bias=forget_gate_bias,
+                    t_max=t_max,
+                    dropout=dropout,
+                    weights_init_scale=weights_init_scale,
+                    hidden_hidden_bias_scale=hidden_hidden_bias_scale,
+                ),
+            }
+        )
+        return layers
+
+    def initialize_state(self, y: torch.Tensor) -> List[torch.Tensor]:
+        """
+        Initialize the state of the RNN layers, with same dtype and device as input `y`.
+
+        Args:
+            y: A torch.Tensor whose device the generated states will be placed on.
+
+        Returns:
+            List of torch.Tensor, each of shape [L, B, H], where
+                L = Number of RNN layers
+                B = Batch size
+                H = Hidden size of RNN.
+        """
+        batch = y.size(0)
+        if self.random_state_sampling and self.training:
+            state = [
+                torch.randn(self.pred_rnn_layers, batch, self.pred_hidden, dtype=y.dtype, device=y.device),
+                torch.randn(self.pred_rnn_layers, batch, self.pred_hidden, dtype=y.dtype, device=y.device),
+            ]
+
+        else:
+            state = [
+                torch.zeros(self.pred_rnn_layers, batch, self.pred_hidden, dtype=y.dtype, device=y.device),
+                torch.zeros(self.pred_rnn_layers, batch, self.pred_hidden, dtype=y.dtype, device=y.device),
+            ]
+        return state
+
+    def score_hypothesis(
+        self, hypothesis: rnnt_utils.Hypothesis, cache: Dict[Tuple[int], Any]
+    ) -> (torch.Tensor, List[torch.Tensor], torch.Tensor):
+        """
+        Similar to the predict() method, instead this method scores a Hypothesis during beam search.
+        Hypothesis is a dataclass representing one hypothesis in a Beam Search.
+
+        Args:
+            hypothesis: Refer to rnnt_utils.Hypothesis.
+            cache: Dict which contains a cache to avoid duplicate computations.
+
+        Returns:
+            Returns a tuple (y, states, lm_token) such that:
+            y is a torch.Tensor of shape [1, 1, H] representing the score of the last token in the Hypothesis.
+            state is a list of RNN states, each of shape [L, 1, H].
+            lm_token is the final integer token of the hypothesis.
+        """
+        if hypothesis.dec_state is not None:
+            device = hypothesis.dec_state[0].device
+        else:
+            _p = next(self.parameters())
+            device = _p.device
+
+        # parse "blank" tokens in hypothesis
+        if len(hypothesis.y_sequence) > 0 and hypothesis.y_sequence[-1] == self.blank_idx:
+            blank_state = True
+        else:
+            blank_state = False
+
+        # Convert last token of hypothesis to torch.Tensor
+        target = torch.full([1, 1], fill_value=hypothesis.y_sequence[-1], device=device, dtype=torch.long)
+        lm_token = target[:, -1]  # [1]
+
+        # Convert current hypothesis into a tuple to preserve in cache
+        sequence = tuple(hypothesis.y_sequence)
+
+        if sequence in cache:
+            y, new_state = cache[sequence]
+        else:
+            # Obtain score for target token and new states
+            if blank_state:
+                y, new_state = self.predict(None, state=None, add_sos=False, batch_size=1)  # [1, 1, H]
+
+            else:
+                y, new_state = self.predict(
+                    target, state=hypothesis.dec_state, add_sos=False, batch_size=1
+                )  # [1, 1, H]
+
+            y = y[:, -1:, :]  # Extract just last state : [1, 1, H]
+            cache[sequence] = (y, new_state)
+
+        return y, new_state, lm_token
+
+    def batch_score_hypothesis(
+        self, hypotheses: List[rnnt_utils.Hypothesis], cache: Dict[Tuple[int], Any], batch_states: List[torch.Tensor]
+    ) -> (torch.Tensor, List[torch.Tensor], torch.Tensor):
+        """
+        Used for batched beam search algorithms. Similar to score_hypothesis method.
+
+        Args:
+            hypothesis: List of Hypotheses. Refer to rnnt_utils.Hypothesis.
+            cache: Dict which contains a cache to avoid duplicate computations.
+            batch_states: List of torch.Tensor which represent the states of the RNN for this batch.
+                Each state is of shape [L, B, H]
+
+        Returns:
+            Returns a tuple (b_y, b_states, lm_tokens) such that:
+            b_y is a torch.Tensor of shape [B, 1, H] representing the scores of the last tokens in the Hypotheses.
+            b_state is a list of list of RNN states, each of shape [L, B, H].
+                Represented as B x List[states].
+            lm_token is a list of the final integer tokens of the hypotheses in the batch.
+        """
+        final_batch = len(hypotheses)
+
+        if final_batch == 0:
+            raise ValueError("No hypotheses was provided for the batch!")
+
+        _p = next(self.parameters())
+        device = _p.device
+        dtype = _p.dtype
+
+        tokens = []
+        process = []
+        done = [None for _ in range(final_batch)]
+
+        # For each hypothesis, cache the last token of the sequence and the current states
+        for i, hyp in enumerate(hypotheses):
+            sequence = tuple(hyp.y_sequence)
+
+            if sequence in cache:
+                done[i] = cache[sequence]
+            else:
+                tokens.append(hyp.y_sequence[-1])
+                process.append((sequence, hyp.dec_state))
+
+        if process:
+            batch = len(process)
+
+            # convert list of tokens to torch.Tensor, then reshape.
+            tokens = torch.tensor(tokens, device=device, dtype=torch.long).view(batch, -1)
+            dec_states = self.initialize_state(tokens.to(dtype=dtype))  # [L, B, H]
+            dec_states = self.batch_initialize_states(dec_states, [d_state for seq, d_state in process])
+
+            y, dec_states = self.predict(
+                tokens, state=dec_states, add_sos=False, batch_size=batch
+            )  # [B, 1, H], List([L, 1, H])
+
+        # Update done states and cache shared by entire batch.
+        j = 0
+        for i in range(final_batch):
+            if done[i] is None:
+                # Select sample's state from the batch state list
+                new_state = self.batch_select_state(dec_states, j)
+
+                # Cache [1, H] scores of the current y_j, and its corresponding state
+                done[i] = (y[j], new_state)
+                cache[process[j][0]] = (y[j], new_state)
+
+                j += 1
+
+        # Set the incoming batch states with the new states obtained from `done`.
+        batch_states = self.batch_initialize_states(batch_states, [d_state for y_j, d_state in done])
+
+        # Create batch of all output scores
+        # List[1, 1, H] -> [B, 1, H]
+        batch_y = torch.stack([y_j for y_j, d_state in done])
+
+        # Extract the last tokens from all hypotheses and convert to a tensor
+        lm_tokens = torch.tensor([h.y_sequence[-1] for h in hypotheses], device=device, dtype=torch.long).view(
+            final_batch
+        )
+
+        return batch_y, batch_states, lm_tokens
+
+    def batch_initialize_states(self, batch_states: List[torch.Tensor], decoder_states: List[List[torch.Tensor]]):
+        """
+        Create batch of decoder states.
+
+       Args:
+           batch_states (list): batch of decoder states
+              ([L x (B, H)], [L x (B, H)])
+
+           decoder_states (list of list): list of decoder states
+               [B x ([L x (1, H)], [L x (1, H)])]
+
+       Returns:
+           batch_states (tuple): batch of decoder states
+               ([L x (B, H)], [L x (B, H)])
+       """
+        # LSTM has 2 states
+        for layer in range(self.pred_rnn_layers):
+            for state_id in range(len(batch_states)):
+                batch_states[state_id][layer] = torch.stack([s[state_id][layer] for s in decoder_states])
+
+        return batch_states
+
+    def batch_select_state(self, batch_states: List[torch.Tensor], idx: int) -> List[List[torch.Tensor]]:
+        """Get decoder state from batch of states, for given id.
+
+        Args:
+            batch_states (list): batch of decoder states
+                ([L x (B, H)], [L x (B, H)])
+
+            idx (int): index to extract state from batch of states
+
+        Returns:
+            (tuple): decoder states for given id
+                ([L x (1, H)], [L x (1, H)])
+        """
+        state_list = []
+        for state_id in range(len(batch_states)):
+            states = [batch_states[state_id][layer][idx] for layer in range(self.pred_rnn_layers)]
+            state_list.append(states)
+
+        return state_list
+
+
+class RNNTJoint(rnnt_abstract.AbstractRNNTJoint):
+    """A Recurrent Neural Network Transducer Joint Network (RNN-T Joint Network).
+    An RNN-T Joint network, comprised of a feedforward model.
+
+    Args:
+        jointnet: A dict-like object which contains the following key-value pairs.
+            encoder_hidden: int specifying the hidden dimension of the encoder net.
+            pred_hidden: int specifying the hidden dimension of the prediction net.
+            joint_hidden: int specifying the hidden dimension of the joint net
+            activation: Activation function used in the joint step. Can be one of
+                ['relu', 'tanh', 'sigmoid'].
+
+            Optionally, it may also contain the following:
+            dropout: float, set to 0.0 by default. Optional dropout applied at the end of the joint net.
+
+        num_classes: int, specifying the vocabulary size that the joint network must predict,
+            excluding the RNNT blank token.
+
+        vocabulary: Optional list of strings/tokens that comprise the vocabulary of the joint network.
+            Unused and kept only for easy access for character based encoding RNNT models.
+
+        log_softmax: Optional bool, set to None by default. If set as None, will compute the log_softmax()
+            based on the value provided.
+
+        preserve_memory: Optional bool, set to False by default. If the model crashes due to the memory
+            intensive joint step, one might try this flag to empty the tensor cache in pytorch.
+
+            Warning: This will make the forward-backward pass much slower than normal.
+            It also might not fix the OOM if the GPU simply does not have enough memory to compute the joint.
+
+        experimental_fuse_loss_wer: Optional bool, set to False by default.
+            NOTE: This is an experimental feature that attempts to trade of compute time for memory preservation.
+            There may be undetermined effects to convergence behaviour.
+
+            Fuses the joint forward, loss forward and
+            wer forward steps. In doing so, it trades of speed for memory conservation by creating sub-batches
+            of the provided batch of inputs, and performs Joint forward, loss forward and wer forward (optional),
+            all on sub-batches, then collates results to be exactly equal to results from the entire batch.
+
+            When this flag is set, prior to calling forward, the fields `loss` and `wer` (either one) *must*
+            be set using the `RNNTJoint.set_loss()` or `RNNTJoint.set_wer()` methods.
+
+            Further, when this flag is set, the following argument `fused_batch_size` *must* be provided
+            as a non negative integer. This value refers to the size of the sub-batch.
+
+            When the flag is set, the input and output signature of `forward()` of this method changes.
+            Input - in addition to `encoder_outputs` (mandatory argument), the following arguments can be provided.
+                - decoder_outputs (optional). Required if loss computation is required.
+                - encoder_lengths (required)
+                - transcripts (optional). Required for wer calculation.
+                - transcript_lengths (optional). Required for wer calculation.
+                - compute_wer (bool, default false). Whether to compute WER or not for the fused batch.
+
+            Output - instead of the usual `joint` log prob tensor, the following results can be returned.
+                - loss (optional). Returned if decoder_outputs, transcripts and transript_lengths are not None.
+                - wer_numerator + wer_denominator (optional). Returned if transcripts, transcripts_lengths are provided
+                    and compute_wer is set.
+
+        fused_batch_size: Optional int, required if `fuse_loss_wer` flag is set. Determines the size of the
+            sub-batches. Should be any value below the actual batch size per GPU.
+    """
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input ports.
+        """
+        return {
+            "encoder_outputs": NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()),
+            "decoder_outputs": NeuralType(('B', 'D', 'T'), EmbeddedTextType()),
+            "encoder_lengths": NeuralType(tuple('B'), LengthsType(), optional=True),
+            "transcripts": NeuralType(('B', 'T'), LabelsType(), optional=True),
+            "transcript_lengths": NeuralType(tuple('B'), LengthsType(), optional=True),
+            "compute_wer": NeuralType(optional=True),
+        }
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output ports.
+        """
+        if not self._fuse_loss_wer:
+            return {
+                "outputs": NeuralType(('B', 'T', 'D', 'D'), LogprobsType()),
+            }
+
+        else:
+            return {
+                "loss": NeuralType(elements_type=LossType(), optional=True),
+                "wer": NeuralType(elements_type=ElementType(), optional=True),
+                "wer_numer": NeuralType(elements_type=ElementType(), optional=True),
+                "wer_denom": NeuralType(elements_type=ElementType(), optional=True),
+            }
+
+    def __init__(
+        self,
+        jointnet: Dict[str, Any],
+        num_classes: int,
+        vocabulary: Optional[List] = None,
+        log_softmax: Optional[bool] = None,
+        preserve_memory: bool = False,
+        experimental_fuse_loss_wer: bool = False,
+        fused_batch_size: Optional[int] = None,
+    ):
+        super().__init__()
+
+        self.vocabulary = vocabulary
+
+        self._vocab_size = num_classes
+        self._num_classes = num_classes + 1  # add 1 for blank symbol
+
+        self._fuse_loss_wer = experimental_fuse_loss_wer
+        self._fused_batch_size = fused_batch_size
+
+        if experimental_fuse_loss_wer and (fused_batch_size is None):
+            raise ValueError("If `fuse_loss_wer` is set, then `fused_batch_size` cannot be None!")
+
+        if experimental_fuse_loss_wer:
+            logging.warning(
+                "\nFused joint step is an experimental technique. Please be aware that it "
+                "may have unintended side effects!\n"
+            )
+
+        self._loss = None
+        self._wer = None
+
+        # Log softmax should be applied explicitly only for CPU
+        self.log_softmax = log_softmax
+        self.preserve_memory = preserve_memory
+
+        if preserve_memory:
+            logging.warning(
+                "`preserve_memory` was set for the Joint Model. Please be aware this will severely impact "
+                "the forward-backward step time. It also might not solve OOM issues if the GPU simply "
+                "does not have enough memory to compute the joint."
+            )
+
+        # Required arguments
+        self.encoder_hidden = jointnet['encoder_hidden']
+        self.pred_hidden = jointnet['pred_hidden']
+        self.joint_hidden = jointnet['joint_hidden']
+        self.activation = jointnet['activation']
+
+        # Optional arguments
+        dropout = jointnet.get('dropout', 0.0)
+
+        self.pred, self.enc, self.joint_net = self._joint_net(
+            num_classes=self._num_classes,  # add 1 for blank symbol
+            pred_n_hidden=self.pred_hidden,
+            enc_n_hidden=self.encoder_hidden,
+            joint_n_hidden=self.joint_hidden,
+            activation=self.activation,
+            dropout=dropout,
+        )
+
+    @typecheck()
+    def forward(
+        self,
+        encoder_outputs: torch.Tensor,
+        decoder_outputs: Optional[torch.Tensor],
+        encoder_lengths: Optional[torch.Tensor] = None,
+        transcripts: Optional[torch.Tensor] = None,
+        transcript_lengths: Optional[torch.Tensor] = None,
+        compute_wer: bool = False,
+    ) -> Union[torch.Tensor, List[Optional[torch.Tensor]]]:
+        # encoder = (B, D, T)
+        # decoder = (B, D, U) if passed, else None
+        encoder_outputs = encoder_outputs.transpose(1, 2)  # (B, T, D)
+
+        if decoder_outputs is not None:
+            decoder_outputs = decoder_outputs.transpose(1, 2)  # (B, U, D)
+
+        if not self._fuse_loss_wer:
+            if decoder_outputs is None:
+                raise ValueError(
+                    "decoder_outputs passed is None, and `fuse_loss_wer` is not set. "
+                    "decoder_outputs can only be None for fused step!"
+                )
+
+            out = self.joint(encoder_outputs, decoder_outputs)  # [B, T, U, V + 1]
+            return out
+
+        else:
+            # At least the loss module must be supplied during fused joint
+            if self._loss is None or self._wer is None:
+                raise ValueError("`fuse_loss_wer` flag is set, but `loss` and `wer` modules were not provided! ")
+
+            # If fused joint step is required, fused batch size is required as well
+            if self._fused_batch_size is None:
+                raise ValueError("If `experimental_fuse_loss_wer` is set, then `fused_batch_size` cannot be None!")
+
+            # When using fused joint step, both encoder and transcript lengths must be provided
+            if (encoder_lengths is None) or (transcript_lengths is None):
+                raise ValueError(
+                    "`experimental_fuse_loss_wer` is set, therefore encoder and target lengths "
+                    "must be provided as well!"
+                )
+
+            losses = []
+            wer_numer_list = []
+            wer_denom_list = []
+            batch_size = int(encoder_outputs.size(0))  # actual batch size
+
+            # Iterate over batch using fused_batch_size steps
+            for batch_idx in range(0, batch_size, self._fused_batch_size):
+                begin = batch_idx
+                end = min(begin + self._fused_batch_size, batch_size)
+
+                # Extract the sub batch inputs
+                # sub_enc = encoder_outputs[begin:end, ...]
+                # sub_transcripts = transcripts[begin:end, ...]
+                sub_enc = encoder_outputs.narrow(dim=0, start=begin, length=end - begin)
+                sub_transcripts = transcripts.narrow(dim=0, start=begin, length=end - begin)
+
+                sub_enc_lens = encoder_lengths[begin:end]
+                sub_transcript_lens = transcript_lengths[begin:end]
+
+                # Sub transcripts does not need the full padding of the entire batch
+                # Therefore reduce the decoder time steps to match
+                max_sub_enc_length = sub_enc_lens.max()
+                max_sub_transcript_length = sub_transcript_lens.max()
+
+                if decoder_outputs is not None:
+                    # Reduce encoder length to preserve computation
+                    # Encoder: [sub-batch, T, D] -> [sub-batch, T', D]; T' < T
+                    if sub_enc.shape[1] != max_sub_enc_length:
+                        sub_enc = sub_enc.narrow(dim=1, start=0, length=max_sub_enc_length)
+
+                    # sub_dec = decoder_outputs[begin:end, ...]  # [sub-batch, U, D]
+                    sub_dec = decoder_outputs.narrow(dim=0, start=begin, length=end - begin)  # [sub-batch, U, D]
+
+                    # Reduce decoder length to preserve computation
+                    # Decoder: [sub-batch, U, D] -> [sub-batch, U', D]; U' < U
+                    if sub_dec.shape[1] != max_sub_transcript_length + 1:
+                        sub_dec = sub_dec.narrow(dim=1, start=0, length=max_sub_transcript_length + 1)
+
+                    # Perform joint => [sub-batch, T', U', V + 1]
+                    sub_joint = self.joint(sub_enc, sub_dec)
+
+                    del sub_dec
+
+                    # Reduce transcript length to correct alignment
+                    # Transcript: [sub-batch, L] -> [sub-batch, L']; L' <= L
+                    if sub_transcripts.shape[1] != max_sub_transcript_length:
+                        sub_transcripts = sub_transcripts.narrow(dim=1, start=0, length=max_sub_transcript_length)
+
+                    # Compute sub batch loss
+                    # preserve loss reduction type
+                    loss_reduction = self.loss.reduction
+
+                    # override loss reduction to sum
+                    self.loss.reduction = None
+
+                    # compute and preserve loss
+                    loss_batch = self.loss(
+                        log_probs=sub_joint,
+                        targets=sub_transcripts,
+                        input_lengths=sub_enc_lens,
+                        target_lengths=sub_transcript_lens,
+                    )
+                    losses.append(loss_batch)
+
+                    # reset loss reduction type
+                    self.loss.reduction = loss_reduction
+
+                else:
+                    losses = None
+
+                # Compute WER for sub batch
+                if compute_wer:
+                    sub_enc = sub_enc.transpose(1, 2)  # [B, T, D] -> [B, D, T]
+                    sub_enc = sub_enc.detach()
+                    sub_transcripts = sub_transcripts.detach()
+
+                    original_log_prediction = self.wer.log_prediction
+                    if batch_idx == 0:
+                        self.wer.log_prediction = True
+                    else:
+                        self.wer.log_prediction = False
+
+                    # Compute the wer (with logging for just 1st sub-batch)
+                    self.wer.update(sub_enc, sub_enc_lens, sub_transcripts, sub_transcript_lens)
+                    wer, wer_num, wer_denom = self.wer.compute()
+                    self.wer.reset()
+
+                    wer_numer_list.append(wer_num)
+                    wer_denom_list.append(wer_denom)
+
+                    # Reset logging default
+                    self.wer.log_prediction = original_log_prediction
+
+                else:
+                    wer = None
+
+                del sub_enc, sub_transcripts, sub_enc_lens, sub_transcript_lens
+
+            # Collect sub batch loss results
+            if losses is not None:
+                losses = torch.cat(losses, 0)
+                losses = losses.mean()  # global batch size average
+            else:
+                losses = None
+
+            # Collect sub batch wer results
+            if compute_wer:
+                wer_num = torch.tensor(wer_numer_list, dtype=torch.long)
+                wer_denom = torch.tensor(wer_denom_list, dtype=torch.long)
+
+                wer_num = wer_num.sum()  # global sum of correct words/chars
+                wer_denom = wer_denom.sum()  # global sum of all words/chars
+            else:
+                wer_num = None
+                wer_denom = None
+
+            return losses, wer, wer_num, wer_denom
+
+    def joint(self, f: torch.Tensor, g: torch.Tensor) -> torch.Tensor:
+        """
+        Compute the joint step of the network.
+
+        Here,
+        B = Batch size
+        T = Acoustic model timesteps
+        U = Target sequence length
+        H1, H2 = Hidden dimensions of the Encoder / Decoder respectively
+        H = Hidden dimension of the Joint hidden step.
+        V = Vocabulary size of the Decoder (excluding the RNNT blank token).
+
+        NOTE:
+            The implementation of this model is slightly modified from the original paper.
+            The original paper proposes the following steps :
+            (enc, dec) -> Expand + Concat + Sum [B, T, U, H1+H2] -> Forward through joint hidden [B, T, U, H] -- *1
+            *1 -> Forward through joint final [B, T, U, V + 1].
+
+            We instead split the joint hidden into joint_hidden_enc and joint_hidden_dec and act as follows:
+            enc -> Forward through joint_hidden_enc -> Expand [B, T, 1, H] -- *1
+            dec -> Forward through joint_hidden_dec -> Expand [B, 1, U, H] -- *2
+            (*1, *2) -> Sum [B, T, U, H] -> Forward through joint final [B, T, U, V + 1].
+
+        Args:
+            f: Output of the Encoder model. A torch.Tensor of shape [B, T, H1]
+            g: Output of the Decoder model. A torch.Tensor of shape [B, U, H2]
+
+        Returns:
+            Logits / log softmaxed tensor of shape (B, T, U, V + 1).
+        """
+        # f = [B, T, H1]
+        f = self.enc(f)
+        f.unsqueeze_(dim=2)  # (B, T, 1, H)
+
+        # g = [B, U, H2]
+        g = self.pred(g)
+        g.unsqueeze_(dim=1)  # (B, 1, U, H)
+
+        inp = f + g  # [B, T, U, H]
+
+        del f, g
+
+        res = self.joint_net(inp)  # [B, T, U, V + 1]
+
+        del inp
+
+        if self.preserve_memory:
+            torch.cuda.empty_cache()
+
+        # If log_softmax is automatic
+        if self.log_softmax is None:
+            if not res.is_cuda:  # Use log softmax only if on CPU
+                res = res.log_softmax(dim=-1)
+        else:
+            if self.log_softmax:
+                res = res.log_softmax(dim=-1)
+
+        return res
+
+    def _joint_net(self, num_classes, pred_n_hidden, enc_n_hidden, joint_n_hidden, activation, dropout):
+        """
+        Prepare the trainable modules of the Joint Network
+
+        Args:
+            num_classes: Number of output classes (vocab size) excluding the RNNT blank token.
+            pred_n_hidden: Hidden size of the prediction network.
+            enc_n_hidden: Hidden size of the encoder network.
+            joint_n_hidden: Hidden size of the joint network.
+            activation: Activation of the joint. Can be one of [relu, tanh, sigmoid]
+            dropout: Dropout value to apply to joint.
+        """
+        pred = torch.nn.Linear(pred_n_hidden, joint_n_hidden)
+        enc = torch.nn.Linear(enc_n_hidden, joint_n_hidden)
+
+        if activation not in ['relu', 'sigmoid', 'tanh']:
+            raise ValueError("Unsupported activation for joint step - please pass one of " "[relu, sigmoid, tanh]")
+
+        activation = activation.lower()
+
+        if activation == 'relu':
+            activation = torch.nn.ReLU(inplace=True)
+        elif activation == 'sigmoid':
+            activation = torch.nn.Sigmoid()
+        elif activation == 'tanh':
+            activation = torch.nn.Tanh()
+
+        layers = (
+            [activation]
+            + ([torch.nn.Dropout(p=dropout)] if dropout else [])
+            + [torch.nn.Linear(joint_n_hidden, num_classes)]
+        )
+        return pred, enc, torch.nn.Sequential(*layers)
+
+    @property
+    def num_classes_with_blank(self):
+        return self._num_classes
+
+    @property
+    def loss(self):
+        return self._loss
+
+    def set_loss(self, loss):
+        if not self._fuse_loss_wer:
+            raise ValueError("Attempting to set loss module even though `fuse_loss_wer` is not set!")
+
+        self._loss = loss
+
+    @property
+    def wer(self):
+        return self._wer
+
+    def set_wer(self, wer):
+        if not self._fuse_loss_wer:
+            raise ValueError("Attempting to set WER module even though `fuse_loss_wer` is not set!")
+
+        self._wer = wer
+
+    @property
+    def fuse_loss_wer(self):
+        return self._fuse_loss_wer
+
+    def set_fuse_loss_wer(self, fuse_loss_wer):
+        self._fuse_loss_wer = fuse_loss_wer
+
+        if self._fuse_loss_wer is False:
+            self._loss = None
+            self._wer = None
+
+    @property
+    def fused_batch_size(self):
+        return self._fuse_loss_wer
+
+    def set_fused_batch_size(self, fused_batch_size):
+        self._fused_batch_size = fused_batch_size

tts/nemo/collections/asr/modules/rnnt_abstract.py

@@ -0,0 +1,236 @@
+# 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 abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+
+from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
+from nemo.core import NeuralModule
+
+
+class AbstractRNNTJoint(NeuralModule, ABC):
+    """
+    An abstract RNNT Joint framework, which can possibly integrate with GreedyRNNTInfer and BeamRNNTInfer classes.
+    Represents the abstract RNNT Joint network, which accepts the acoustic model and prediction network
+    embeddings in order to compute the joint of the two prior to decoding the output sequence.
+    """
+
+    @abstractmethod
+    def joint(self, f: torch.Tensor, g: torch.Tensor) -> torch.Tensor:
+        """
+        Compute the joint step of the network.
+
+        Here,
+        B = Batch size
+        T = Acoustic model timesteps
+        U = Target sequence length
+        H1, H2 = Hidden dimensions of the Encoder / Decoder respectively
+        H = Hidden dimension of the Joint hidden step.
+        V = Vocabulary size of the Decoder (excluding the RNNT blank token).
+
+        NOTE:
+            The implementation of this model is slightly modified from the original paper.
+            The original paper proposes the following steps :
+            (enc, dec) -> Expand + Concat + Sum [B, T, U, H1+H2] -> Forward through joint hidden [B, T, U, H] -- *1
+            *1 -> Forward through joint final [B, T, U, V + 1].
+
+            We instead split the joint hidden into joint_hidden_enc and joint_hidden_dec and act as follows:
+            enc -> Forward through joint_hidden_enc -> Expand [B, T, 1, H] -- *1
+            dec -> Forward through joint_hidden_dec -> Expand [B, 1, U, H] -- *2
+            (*1, *2) -> Sum [B, T, U, H] -> Forward through joint final [B, T, U, V + 1].
+
+        Args:
+            f: Output of the Encoder model. A torch.Tensor of shape [B, T, H1]
+            g: Output of the Decoder model. A torch.Tensor of shape [B, U, H2]
+
+        Returns:
+            Logits / log softmaxed tensor of shape (B, T, U, V + 1).
+        """
+        raise NotImplementedError()
+
+    @property
+    def num_classes_with_blank(self):
+        raise NotImplementedError()
+
+
+class AbstractRNNTDecoder(NeuralModule, ABC):
+    """
+    An abstract RNNT Decoder framework, which can possibly integrate with GreedyRNNTInfer and BeamRNNTInfer classes.
+    Represents the abstract RNNT Prediction/Decoder stateful network, which performs autoregressive decoding
+    in order to construct the output sequence.
+
+    Args:
+        vocab_size: Size of the vocabulary, excluding the RNNT blank token.
+        blank_idx: Index of the blank token. Can be 0 or size(vocabulary).
+        blank_as_pad: Bool flag, whether to allocate an additional token in the Embedding layer
+            of this module in order to treat all RNNT `blank` tokens as pad tokens, thereby letting
+            the Embedding layer batch tokens more efficiently.
+
+            It is mandatory to use this for certain Beam RNNT Infer methods - such as TSD, ALSD.
+            It is also more efficient to use greedy batch decoding with this flag.
+    """
+
+    def __init__(self, vocab_size, blank_idx, blank_as_pad):
+        super().__init__()
+
+        self.vocab_size = vocab_size
+        self.blank_idx = blank_idx  # first or last index of vocabulary
+        self.blank_as_pad = blank_as_pad
+
+        if blank_idx not in [0, vocab_size]:
+            raise ValueError("`blank_idx` must be either 0 or the final token of the vocabulary")
+
+    @abstractmethod
+    def predict(
+        self,
+        y: Optional[torch.Tensor] = None,
+        state: Optional[torch.Tensor] = None,
+        add_sos: bool = False,
+        batch_size: Optional[int] = None,
+    ) -> (torch.Tensor, List[torch.Tensor]):
+        """
+        Stateful prediction of scores and state for a (possibly null) tokenset.
+        This method takes various cases into consideration :
+        - No token, no state - used for priming the RNN
+        - No token, state provided - used for blank token scoring
+        - Given token, states - used for scores + new states
+
+        Here:
+        B - batch size
+        U - label length
+        H - Hidden dimension size of RNN
+        L - Number of RNN layers
+
+        Args:
+            y: Optional torch tensor of shape [B, U] of dtype long which will be passed to the Embedding.
+                If None, creates a zero tensor of shape [B, 1, H] which mimics output of pad-token on Embedding.
+
+            state: An optional list of states for the RNN. Eg: For LSTM, it is the state list length is 2.
+                Each state must be a tensor of shape [L, B, H].
+                If None, and during training mode and `random_state_sampling` is set, will sample a
+                normal distribution tensor of the above shape. Otherwise, None will be passed to the RNN.
+
+            add_sos: bool flag, whether a zero vector describing a "start of signal" token should be
+                prepended to the above "y" tensor. When set, output size is (B, U + 1, H).
+
+            batch_size: An optional int, specifying the batch size of the `y` tensor.
+                Can be infered if `y` and `state` is None. But if both are None, then batch_size cannot be None.
+
+        Returns:
+            A tuple  (g, hid) such that -
+
+            If add_sos is False:
+                g: (B, U, H)
+                hid: (h, c) where h is the final sequence hidden state and c is the final cell state:
+                    h (tensor), shape (L, B, H)
+                    c (tensor), shape (L, B, H)
+
+            If add_sos is True:
+                g: (B, U + 1, H)
+                hid: (h, c) where h is the final sequence hidden state and c is the final cell state:
+                    h (tensor), shape (L, B, H)
+                    c (tensor), shape (L, B, H)
+
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def initialize_state(self, y: torch.Tensor) -> List[torch.Tensor]:
+        """
+        Initialize the state of the RNN layers, with same dtype and device as input `y`.
+
+        Args:
+            y: A torch.Tensor whose device the generated states will be placed on.
+
+        Returns:
+            List of torch.Tensor, each of shape [L, B, H], where
+                L = Number of RNN layers
+                B = Batch size
+                H = Hidden size of RNN.
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def score_hypothesis(
+        self, hypothesis: Hypothesis, cache: Dict[Tuple[int], Any]
+    ) -> (torch.Tensor, List[torch.Tensor], torch.Tensor):
+        """
+        Similar to the predict() method, instead this method scores a Hypothesis during beam search.
+        Hypothesis is a dataclass representing one hypothesis in a Beam Search.
+
+        Args:
+            hypothesis: Refer to rnnt_utils.Hypothesis.
+            cache: Dict which contains a cache to avoid duplicate computations.
+
+        Returns:
+            Returns a tuple (y, states, lm_token) such that:
+            y is a torch.Tensor of shape [1, 1, H] representing the score of the last token in the Hypothesis.
+            state is a list of RNN states, each of shape [L, 1, H].
+            lm_token is the final integer token of the hypothesis.
+        """
+        raise NotImplementedError()
+
+    def batch_score_hypothesis(
+        self, hypotheses: List[Hypothesis], cache: Dict[Tuple[int], Any], batch_states: List[torch.Tensor]
+    ) -> (torch.Tensor, List[torch.Tensor], torch.Tensor):
+        """
+        Used for batched beam search algorithms. Similar to score_hypothesis method.
+
+        Args:
+            hypothesis: List of Hypotheses. Refer to rnnt_utils.Hypothesis.
+            cache: Dict which contains a cache to avoid duplicate computations.
+            batch_states: List of torch.Tensor which represent the states of the RNN for this batch.
+                Each state is of shape [L, B, H]
+
+        Returns:
+            Returns a tuple (b_y, b_states, lm_tokens) such that:
+            b_y is a torch.Tensor of shape [B, 1, H] representing the scores of the last tokens in the Hypotheses.
+            b_state is a list of list of RNN states, each of shape [L, B, H].
+                Represented as B x List[states].
+            lm_token is a list of the final integer tokens of the hypotheses in the batch.
+        """
+        raise NotImplementedError()
+
+    def batch_initialize_states(self, batch_states: List[torch.Tensor], decoder_states: List[List[torch.Tensor]]):
+        """
+        Create batch of decoder states.
+
+       Args:
+           batch_states (list): batch of decoder states
+              ([L x (B, H)], [L x (B, H)])
+
+           decoder_states (list of list): list of decoder states
+               [B x ([L x (1, H)], [L x (1, H)])]
+
+       Returns:
+           batch_states (tuple): batch of decoder states
+               ([L x (B, H)], [L x (B, H)])
+       """
+        raise NotImplementedError()
+
+    def batch_select_state(self, batch_states: List[torch.Tensor], idx: int) -> List[List[torch.Tensor]]:
+        """Get decoder state from batch of states, for given id.
+
+        Args:
+            batch_states (list): batch of decoder states
+                ([L x (B, H)], [L x (B, H)])
+
+            idx (int): index to extract state from batch of states
+
+        Returns:
+            (tuple): decoder states for given id
+                ([L x (1, H)], [L x (1, H)])
+        """
+        raise NotImplementedError()

tts/nemo/collections/asr/modules/wav2vec_modules.py

@@ -0,0 +1,306 @@
+# 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 typing import Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from nemo.collections.asr.models.wav2vec.wav2vec_config import Wav2VecMaskType
+from nemo.core import NeuralModule
+from nemo.core.neural_types import EncodedRepresentation, LossType, NeuralType
+
+
+class GumbelVectorQuantizer(NeuralModule):
+    def __init__(
+        self,
+        dim,
+        num_vars,
+        temp,
+        groups,
+        combine_groups,
+        vq_dim,
+        time_first,
+        activation=nn.GELU(),
+        weight_proj_depth=1,
+        weight_proj_factor=1,
+    ):
+        """Vector quantization using gumbel softmax
+
+        Args:
+            dim: input dimension (channels)
+            num_vars: number of quantized vectors per group
+            temp: temperature for training. this should be a tuple of 3 elements: (start, stop, decay factor)
+            groups: number of groups for vector quantization
+            combine_groups: whether to use the vectors for all groups
+            vq_dim: dimensionality of the resulting quantized vector
+            time_first: if true, expect input in BxTxC format, otherwise in BxCxT
+            activation: what activation to use (should be a module). this is only used if weight_proj_depth is > 1
+            weight_proj_depth: number of layers (with activation in between) to project input before computing logits
+            weight_proj_factor: this is used only if weight_proj_depth is > 1. scales the inner dimensionality of
+                                projections by this factor
+        """
+        super().__init__()
+
+        self.groups = groups
+        self.combine_groups = combine_groups
+        self.input_dim = dim
+        self.num_vars = num_vars
+        self.time_first = time_first
+
+        assert vq_dim % groups == 0, f"dim {vq_dim} must be divisible by groups {groups} for concatenation"
+
+        var_dim = vq_dim // groups
+        num_groups = groups if not combine_groups else 1
+
+        self.vars = nn.Parameter(torch.FloatTensor(1, num_groups * num_vars, var_dim))
+        nn.init.uniform_(self.vars)
+
+        if weight_proj_depth > 1:
+
+            def block(input_dim, output_dim):
+                return nn.Sequential(nn.Linear(input_dim, output_dim), activation)
+
+            inner_dim = self.input_dim * weight_proj_factor
+            self.weight_proj = nn.Sequential(
+                *[block(self.input_dim if i == 0 else inner_dim, inner_dim) for i in range(weight_proj_depth - 1)],
+                nn.Linear(inner_dim, groups * num_vars),
+            )
+        else:
+            self.weight_proj = nn.Linear(self.input_dim, groups * num_vars)
+            nn.init.normal_(self.weight_proj.weight, mean=0, std=1)
+            nn.init.zeros_(self.weight_proj.bias)
+
+        assert len(temp) == 3, "Quantize temperature should be a tuple of 3 elements: (start, stop, decay factor)"
+
+        self.max_temp, self.min_temp, self.temp_decay = temp
+        self.curr_temp = self.max_temp
+        self.codebook_indices = None
+
+    def set_num_updates(self, num_updates):
+        self.curr_temp = max(self.max_temp * self.temp_decay ** num_updates, self.min_temp)
+
+    def get_codebook_indices(self):
+        if self.codebook_indices is None:
+            from itertools import product
+
+            p = [range(self.num_vars)] * self.groups
+            inds = list(product(*p))
+            self.codebook_indices = torch.tensor(inds, dtype=torch.long, device=self.vars.device).flatten()
+
+            if not self.combine_groups:
+                self.codebook_indices = self.codebook_indices.view(self.num_vars ** self.groups, -1)
+                for b in range(1, self.groups):
+                    self.codebook_indices[:, b] += self.num_vars * b
+                self.codebook_indices = self.codebook_indices.flatten()
+        return self.codebook_indices
+
+    def sample_from_codebook(self, b, n):
+        indices = self.get_codebook_indices()
+        indices = indices.view(-1, self.groups)
+        cb_size = indices.size(0)
+        assert n < cb_size, f"sample size {n} is greater than size of codebook {cb_size}"
+        sample_idx = torch.randint(low=0, high=cb_size, size=(b * n,))
+        indices = indices[sample_idx]
+
+        z = self.vars.squeeze(0).index_select(0, indices.flatten()).view(b, n, -1)
+        return z
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input ports.
+        """
+        if self.time_first:
+            return {"x": NeuralType(('B', 'T', 'D'), EncodedRepresentation())}
+        return {"x": NeuralType(('B', 'D', 'T'), EncodedRepresentation())}
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output ports.
+        """
+        if self.time_first:
+            return {
+                "x": NeuralType(('B', 'T', 'D'), EncodedRepresentation()),
+                "quantize_prob_ppl": NeuralType(elements_type=LossType()),
+            }
+        return {
+            "x": NeuralType(('B', 'D', 'T'), EncodedRepresentation()),
+            "quantize_prob_ppl": NeuralType(elements_type=LossType()),
+        }
+
+    def forward(self, x):
+
+        if not self.time_first:
+            x = x.transpose(1, 2)
+
+        bsz, tsz, fsz = x.shape
+        x = x.reshape(-1, fsz)
+        x = self.weight_proj(x)
+        x = x.view(bsz * tsz * self.groups, -1)
+
+        _, k = x.max(-1)
+        hard_x = x.new_zeros(*x.shape).scatter_(-1, k.view(-1, 1), 1.0).view(bsz * tsz, self.groups, -1)
+
+        # Calculate quantize prob perplexity
+        num_vars = self.num_vars * self.groups
+        avg_probs = torch.softmax(x.view(bsz * tsz, self.groups, -1).float(), dim=-1).mean(dim=0)
+        quantize_prob_ppl = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-7), dim=-1)).sum()
+        quantize_prob_ppl = (num_vars - quantize_prob_ppl) / num_vars
+
+        if self.training:
+            x = F.gumbel_softmax(x.float(), tau=self.curr_temp, hard=True).type_as(x)
+        else:
+            x = hard_x
+
+        x = x.view(bsz * tsz, -1)
+
+        vars = self.vars
+        if self.combine_groups:
+            vars = vars.repeat(1, self.groups, 1)
+
+        x = x.unsqueeze(-1) * vars
+        x = x.view(bsz * tsz, self.groups, self.num_vars, -1)
+        x = x.sum(-2)
+        x = x.view(bsz, tsz, -1)
+
+        cur_codebook_temp = self.curr_temp
+
+        if not self.time_first:
+            x = x.transpose(1, 2)  # BTC -> BCT
+
+        return x, quantize_prob_ppl, cur_codebook_temp
+
+
+def compute_mask_indices(
+    shape: Tuple[int, int],
+    padding_mask: Optional[torch.Tensor],
+    mask_prob: float,
+    mask_length: int,
+    mask_type: Wav2VecMaskType = Wav2VecMaskType.static,
+    mask_other: float = 0.0,
+    min_masks: int = 0,
+    no_overlap: bool = False,
+    min_space: int = 0,
+) -> np.ndarray:
+    """
+    Computes random mask spans for a given shape
+    Args:
+        shape: the the shape for which to compute masks.
+            should be of size 2 where first element is batch size and 2nd is timesteps
+        padding_mask: optional padding mask of the same size as shape, which will prevent masking padded elements
+        mask_prob: probability for each token to be chosen as start of the span to be masked. this will be multiplied by
+            number of timesteps divided by length of mask span to mask approximately this percentage of all elements.
+            however due to overlaps, the actual number will be smaller (unless no_overlap is True)
+        mask_type: how 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
+        min_masks: minimum number of masked spans
+        no_overlap: if false, will switch to an alternative recursive algorithm that prevents spans from overlapping
+        min_space: only used if no_overlap is True, this is how many elements to keep unmasked between spans
+    """
+
+    bsz, all_sz = shape
+    mask = np.full((bsz, all_sz), False)
+
+    all_num_mask = int(
+        # add a random number for probabilistic rounding
+        mask_prob * all_sz / float(mask_length)
+        + np.random.rand()
+    )
+
+    all_num_mask = max(min_masks, all_num_mask)
+
+    mask_idcs = []
+    for i in range(bsz):
+        if padding_mask is not None:
+            sz = all_sz - padding_mask[i].long().sum().item()
+            num_mask = int(
+                # add a random number for probabilistic rounding
+                mask_prob * sz / float(mask_length)
+                + np.random.rand()
+            )
+            num_mask = max(min_masks, num_mask)
+        else:
+            sz = all_sz
+            num_mask = all_num_mask
+
+        if mask_type.value is Wav2VecMaskType.static.value:
+            lengths = np.full(num_mask, mask_length)
+        elif mask_type.value is Wav2VecMaskType.uniform:
+            lengths = np.random.randint(mask_other, mask_length * 2 + 1, size=num_mask)
+        elif mask_type.value is Wav2VecMaskType.normal.value:
+            lengths = np.random.normal(mask_length, mask_other, size=num_mask)
+            lengths = [max(1, int(round(x))) for x in lengths]
+        elif mask_type.value is Wav2VecMaskType.poisson.value:
+            lengths = np.random.poisson(mask_length, size=num_mask)
+            lengths = [int(round(x)) for x in lengths]
+        else:
+            raise Exception("unknown mask selection " + str(mask_type))
+
+        if sum(lengths) == 0:
+            lengths[0] = min(mask_length, sz - 1)
+
+        if no_overlap:
+            mask_idc = []
+
+            def arrange(s, e, length, keep_length):
+                span_start = np.random.randint(s, e - length)
+                mask_idc.extend(span_start + i for i in range(length))
+
+                new_parts = []
+                if span_start - s - min_space >= keep_length:
+                    new_parts.append((s, span_start - min_space + 1))
+                if e - span_start - keep_length - min_space > keep_length:
+                    new_parts.append((span_start + length + min_space, e))
+                return new_parts
+
+            parts = [(0, sz)]
+            min_length = min(lengths)
+            for length in sorted(lengths, reverse=True):
+                lens = np.fromiter((e - s if e - s >= length + min_space else 0 for s, e in parts), np.int)
+                l_sum = np.sum(lens)
+                if l_sum == 0:
+                    break
+                probs = lens / np.sum(lens)
+                c = np.random.choice(len(parts), p=probs)
+                s, e = parts.pop(c)
+                parts.extend(arrange(s, e, length, min_length))
+            mask_idc = np.asarray(mask_idc)
+        else:
+            min_len = min(lengths)
+            if sz - min_len <= num_mask:
+                min_len = sz - num_mask - 1
+
+            mask_idc = np.random.choice(sz - min_len, num_mask, replace=False)
+
+            mask_idc = np.asarray([mask_idc[j] + offset for j in range(len(mask_idc)) for offset in range(lengths[j])])
+
+        mask_idcs.append(np.unique(mask_idc[mask_idc < sz]))
+
+    min_len = min([len(m) for m in mask_idcs])
+    for i, mask_idc in enumerate(mask_idcs):
+        if len(mask_idc) > min_len:
+            mask_idc = np.random.choice(mask_idc, min_len, replace=False)
+        mask[i, mask_idc] = True
+
+    return mask

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

@@ -0,0 +1,13 @@
+# 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.

tts/nemo/collections/asr/parts/features.py

@@ -0,0 +1,39 @@
+# 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) 2018 Ryan Leary
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+# This file contains code artifacts adapted from https://github.com/ryanleary/patter
+
+"""
+ALIAS FILE for backward compatibility
+"""
+from nemo.collections.asr.parts.preprocessing.features import *

tts/nemo/collections/asr/parts/mixins/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nemo.collections.asr.parts.mixins.mixins import ASRBPEMixin, ASRModuleMixin, DiarizationMixin

tts/nemo/collections/asr/parts/mixins/mixins.py

@@ -0,0 +1,193 @@
+# 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
+
+from omegaconf import DictConfig, OmegaConf, open_dict
+
+from nemo.collections.asr.parts.utils import asr_module_utils
+from nemo.collections.common import tokenizers
+from nemo.utils import logging
+
+
+class ASRBPEMixin(ABC):
+    """ ASR BPE Mixin class that sets up a Tokenizer via a config
+
+    This mixin class adds the method `_setup_tokenizer(...)`, which can be used by ASR models
+    which depend on subword tokenization.
+
+    The setup_tokenizer method adds the following parameters to the class -
+        -   tokenizer_cfg: The resolved config supplied to the tokenizer (with `dir` and `type` arguments).
+        -   tokenizer_dir: The directory path to the tokenizer vocabulary + additional metadata.
+        -   tokenizer_type: The type of the tokenizer. Currently supports `bpe` and `wpe`.
+        -   vocab_path: Resolved path to the vocabulary text file.
+
+    In addition to these variables, the method will also instantiate and preserve a tokenizer
+    (subclass of TokenizerSpec) if successful, and assign it to self.tokenizer.
+    """
+
+    def _setup_tokenizer(self, tokenizer_cfg: DictConfig):
+        # Prevent tokenizer parallelism (unless user has explicitly set it)
+        if 'TOKENIZERS_PARALLELISM' not in os.environ:
+            os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+
+        self.tokenizer_cfg = OmegaConf.to_container(tokenizer_cfg, resolve=True)  # type: dict
+        self.tokenizer_dir = self.tokenizer_cfg.pop('dir')  # Remove tokenizer directory
+        self.tokenizer_type = self.tokenizer_cfg.pop('type').lower()  # Remove tokenizer_type
+
+        self.hf_tokenizer_kwargs = self.tokenizer_cfg.pop("hf_kwargs", {})  # Remove HF tokenizer kwargs
+
+        # Preserve config
+        if hasattr(self, 'cfg') and 'tokenizer' in self.cfg:
+            self.cfg.tokenizer.dir = self.tokenizer_dir
+            self.cfg.tokenizer.type = self.tokenizer_type
+
+            if 'hf_kwargs' in tokenizer_cfg:
+                with open_dict(self.cfg.tokenizer):
+                    self.cfg.tokenizer.hf_kwargs = tokenizer_cfg.get('hf_kwargs')
+
+        if self.tokenizer_type not in ['bpe', 'wpe']:
+            raise ValueError(
+                "`tokenizer.type` must be either `bpe` for SentencePiece tokenizer or "
+                "`wpe` for BERT based tokenizer"
+            )
+
+        if self.tokenizer_type == 'bpe':
+            # This is a BPE Tokenizer
+            if 'model_path' in self.tokenizer_cfg:
+                model_path = self.tokenizer_cfg.get('model_path')
+            else:
+                model_path = os.path.join(self.tokenizer_dir, 'tokenizer.model')
+            model_path = self.register_artifact('tokenizer.model_path', model_path)
+            self.model_path = model_path
+
+            if 'special_tokens' in self.tokenizer_cfg:
+                special_tokens = self.tokenizer_cfg['special_tokens']
+            else:
+                special_tokens = None
+
+            # Update special tokens
+            self.tokenizer = tokenizers.SentencePieceTokenizer(
+                model_path=model_path, special_tokens=special_tokens, legacy=True
+            )
+
+            if 'vocab_path' in self.tokenizer_cfg:
+                vocab_path = self.tokenizer_cfg.get('vocab_path')
+            else:
+                vocab_path = os.path.join(self.tokenizer_dir, 'vocab.txt')
+            vocab_path = self.register_artifact('tokenizer.vocab_path', vocab_path)
+            self.vocab_path = vocab_path
+
+            try:
+                if 'spe_tokenizer_vocab' in self.tokenizer_cfg:
+                    spe_vocab_path = self.tokenizer_cfg.get('spe_tokenizer_vocab')
+                else:
+                    spe_vocab_path = os.path.join(self.tokenizer_dir, 'tokenizer.vocab')
+                spe_vocab_path = self.register_artifact('tokenizer.spe_tokenizer_vocab', spe_vocab_path)
+                self.spe_vocab_path = spe_vocab_path
+            except FileNotFoundError:
+                # fallback case for older checkpoints that did not preserve the tokenizer.vocab
+                self.spe_vocab_path = None
+
+            vocabulary = {'<unk>': 0}
+            with open(vocab_path) as f:
+                for i, piece in enumerate(f):
+                    piece = piece.replace('\n', '')
+                    vocabulary[piece] = i + 1
+
+            # wrapper method to get vocabulary conveniently
+            def get_vocab():
+                return vocabulary
+
+            # attach utility values to the tokenizer wrapper
+            self.tokenizer.tokenizer.vocab_size = len(vocabulary)
+            self.tokenizer.tokenizer.get_vocab = get_vocab
+            self.tokenizer.tokenizer.all_special_tokens = self.tokenizer.special_token_to_id
+
+        else:
+            # This is a WPE Tokenizer
+            # If path from previous registration exists, remove it
+            if 'vocab_path' in self.tokenizer_cfg:
+                vocab_path = self.tokenizer_cfg.get('vocab_path')
+            else:
+                vocab_path = os.path.join(self.tokenizer_dir, 'vocab.txt')
+            vocab_path = self.register_artifact('tokenizer.vocab_path', vocab_path)
+            self.vocab_path = vocab_path
+
+            # If path from previous registration exists, remove it
+            if 'vocab_path' in self.tokenizer_cfg:
+                self.tokenizer_cfg.pop('vocab_path')
+
+            self.tokenizer = tokenizers.AutoTokenizer(
+                pretrained_model_name='bert-base-cased',
+                vocab_file=self.vocab_path,
+                mask_token=self.hf_tokenizer_kwargs.get('mask_token', None),
+                bos_token=self.hf_tokenizer_kwargs.get('bos_token', None),
+                eos_token=self.hf_tokenizer_kwargs.get('eos_token', None),
+                pad_token=self.hf_tokenizer_kwargs.get('pad_token', None),
+                sep_token=self.hf_tokenizer_kwargs.get('sep_token', None),
+                cls_token=self.hf_tokenizer_kwargs.get('cls_token', None),
+                unk_token=self.hf_tokenizer_kwargs.get('unk_token', None),
+                use_fast=self.hf_tokenizer_kwargs.get('use_fast', False),
+            )
+
+        logging.info(
+            "Tokenizer {} initialized with {} tokens".format(
+                self.tokenizer.__class__.__name__, self.tokenizer.vocab_size
+            )
+        )
+
+
+class ASRModuleMixin(ABC):
+    """
+    ASRModuleMixin is a mixin class added to ASR models in order to add methods that are specific
+    to a particular instantiation of a module inside of an ASRModel.
+
+    Each method should first check that the module is present within the subclass, and support additional
+    functionality if the corresponding module is present.
+    """
+
+    def change_conv_asr_se_context_window(self, context_window: int, update_config: bool = True):
+        """
+        Update the context window of the SqueezeExcitation module if the provided model contains an
+        `encoder` which is an instance of `ConvASREncoder`.
+
+        Args:
+            context_window:  An integer representing the number of input timeframes that will be used
+                to compute the context. Each timeframe corresponds to a single window stride of the
+                STFT features.
+
+                Say the window_stride = 0.01s, then a context window of 128 represents 128 * 0.01 s
+                of context to compute the Squeeze step.
+            update_config: Whether to update the config or not with the new context window.
+        """
+        asr_module_utils.change_conv_asr_se_context_window(
+            self, context_window=context_window, update_config=update_config
+        )
+
+
+class DiarizationMixin(ABC):
+    @abstractmethod
+    def diarize(self, paths2audio_files: List[str], batch_size: int = 1) -> List[str]:
+        """
+        Takes paths to audio files and returns speaker labels
+        Args:
+            paths2audio_files: paths to audio fragment to be transcribed
+
+        Returns:
+            Speaker labels
+        """
+        pass

tts/nemo/collections/asr/parts/numba/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nemo.collections.asr.parts.numba.rnnt_loss.rnnt_pytorch import RNNTLossNumba

tts/nemo/collections/asr/parts/numba/rnnt_loss/__init__.py

@@ -0,0 +1,16 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nemo.collections.asr.parts.numba.rnnt_loss.rnnt import rnnt_loss_cpu, rnnt_loss_gpu
+from nemo.collections.asr.parts.numba.rnnt_loss.rnnt_pytorch import RNNTLossNumba

tts/nemo/collections/asr/parts/numba/rnnt_loss/rnnt.py

@@ -0,0 +1,225 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Copyright 2018-2019, Mingkun Huang
+#
+# 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 multiprocessing
+
+import torch
+from numba import cuda
+
+from nemo.collections.asr.parts.numba.rnnt_loss.utils import global_constants, rnnt_helper
+from nemo.collections.asr.parts.numba.rnnt_loss.utils.cpu_utils import cpu_rnnt
+from nemo.collections.asr.parts.numba.rnnt_loss.utils.cuda_utils import gpu_rnnt
+
+
+def rnnt_loss_cpu(
+    acts: torch.Tensor,
+    labels: torch.Tensor,
+    input_lengths: torch.Tensor,
+    label_lengths: torch.Tensor,
+    costs: torch.Tensor,
+    grads: torch.Tensor,
+    blank_label: int,
+    num_threads: int,
+):
+    """
+    Wrapper method for accessing CPU RNNT loss.
+
+    CPU implementation ported from [HawkAaron/warp-transducer](https://github.com/HawkAaron/warp-transducer).
+
+    Args:
+        acts: Activation tensor of shape [B, T, U, V+1].
+        labels: Ground truth labels of shape [B, U].
+        input_lengths: Lengths of the acoustic sequence as a vector of ints [B].
+        label_lengths: Lengths of the target sequence as a vector of ints [B].
+        costs: Zero vector of length [B] in which costs will be set.
+        grads: Zero tensor of shape [B, T, U, V+1] where the gradient will be set.
+        blank_label: Index of the blank token in the vocabulary.
+        num_threads: Number of threads for OpenMP.
+    """
+    # aliases
+    log_probs = acts
+    flat_labels = labels
+
+    minibatch_size = log_probs.shape[0]
+    maxT = log_probs.shape[1]
+    maxU = log_probs.shape[2]
+    alphabet_size = log_probs.shape[3]
+
+    if num_threads < 0:
+        num_threads = multiprocessing.cpu_count()
+
+    num_threads = max(1, num_threads)  # have to use at least 1 thread
+
+    gpu_size, status = rnnt_helper.get_workspace_size(maxT, maxU, minibatch_size, gpu=False)
+    if status != global_constants.RNNTStatus.RNNT_STATUS_SUCCESS:
+        raise RuntimeError("Invalid parameter passed when calculating working space memory")
+
+    cpu_workspace = torch.zeros(gpu_size, device=log_probs.device, dtype=log_probs.dtype, requires_grad=False)
+
+    ### VIEW TENSORS AS VECTORS FOR POINTER INDEXING ###
+    log_probs, acts_shape = rnnt_helper.flatten_tensor(log_probs)
+    flat_labels, labels_shape = rnnt_helper.flatten_tensor(flat_labels)
+
+    wrapper = cpu_rnnt.CPURNNT(
+        minibatch=minibatch_size,
+        maxT=maxT,
+        maxU=maxU,
+        alphabet_size=alphabet_size,
+        workspace=cpu_workspace,
+        blank=blank_label,
+        num_threads=num_threads,
+        batch_first=True,
+    )
+
+    if grads is None:
+        status = wrapper.score_forward(
+            log_probs=log_probs.data,
+            costs=costs,
+            flat_labels=flat_labels.data,
+            label_lengths=label_lengths.data,
+            input_lengths=input_lengths.data,
+        )
+
+        if status != global_constants.RNNTStatus.RNNT_STATUS_SUCCESS:
+            raise RuntimeError("Could not calculate forward scores")
+
+    else:
+        ### FLATTEN GRAD TENSOR ###
+        grads, grads_shape = rnnt_helper.flatten_tensor(grads)
+
+        status = wrapper.cost_and_grad(
+            log_probs=log_probs.data,
+            grads=grads.data,
+            costs=costs,
+            flat_labels=flat_labels.data,
+            label_lengths=label_lengths.data,
+            input_lengths=input_lengths.data,
+        )
+
+        if status != global_constants.RNNTStatus.RNNT_STATUS_SUCCESS:
+            raise RuntimeError("Could not calculate forward scores")
+
+    del cpu_workspace, wrapper
+    return True
+
+
+def rnnt_loss_gpu(
+    acts: torch.Tensor,
+    labels: torch.Tensor,
+    input_lengths: torch.Tensor,
+    label_lengths: torch.Tensor,
+    costs: torch.Tensor,
+    grads: torch.Tensor,
+    blank_label: int,
+    num_threads: int,
+):
+    """
+    Wrapper method for accessing GPU RNNT loss.
+
+    CUDA implementation ported from [HawkAaron/warp-transducer](https://github.com/HawkAaron/warp-transducer).
+
+    Args:
+        acts: Activation tensor of shape [B, T, U, V+1].
+        labels: Ground truth labels of shape [B, U].
+        input_lengths: Lengths of the acoustic sequence as a vector of ints [B].
+        label_lengths: Lengths of the target sequence as a vector of ints [B].
+        costs: Zero vector of length [B] in which costs will be set.
+        grads: Zero tensor of shape [B, T, U, V+1] where the gradient will be set.
+        blank_label: Index of the blank token in the vocabulary.
+        num_threads: Number of threads for OpenMP.
+    """
+    minibatch_size = acts.shape[0]
+    maxT = acts.shape[1]
+    maxU = acts.shape[2]
+    alphabet_size = acts.shape[3]
+
+    if hasattr(cuda, 'external_stream'):
+        stream = cuda.external_stream(torch.cuda.current_stream(acts.device).cuda_stream)
+    else:
+        stream = cuda.default_stream()
+
+    if num_threads < 0:
+        num_threads = multiprocessing.cpu_count()
+
+    num_threads = max(1, num_threads)  # have to use at least 1 thread
+
+    gpu_size, status = rnnt_helper.get_workspace_size(maxT, maxU, minibatch_size, gpu=True)
+    if status != global_constants.RNNTStatus.RNNT_STATUS_SUCCESS:
+        raise RuntimeError("Invalid parameter passed when calculating working space memory")
+
+    # Select GPU index
+    cuda.select_device(acts.device.index)
+    gpu_workspace = torch.zeros(gpu_size, device=acts.device, dtype=acts.dtype, requires_grad=False)
+
+    ### VIEW TENSORS AS VECTORS FOR POINTER INDEXING ###
+    acts, acts_shape = rnnt_helper.flatten_tensor(acts)
+
+    ### REPRESENT THE CUDA ARRAY INTERFACE OF COSTS VECTOR ###
+    costs_repr = cuda.as_cuda_array(costs, sync=False)  # NO COPY OF DATA, JUST CHANGE REPRESENTATION
+
+    wrapper = gpu_rnnt.GPURNNT(
+        minibatch=minibatch_size,
+        maxT=maxT,
+        maxU=maxU,
+        alphabet_size=alphabet_size,
+        workspace=gpu_workspace,
+        blank=blank_label,
+        num_threads=num_threads,
+        stream=stream,
+    )
+
+    if grads is None:
+        status = wrapper.score_forward(
+            acts=acts.data,
+            costs=costs_repr,
+            pad_labels=labels.data,
+            label_lengths=label_lengths.data,
+            input_lengths=input_lengths.data,
+        )
+
+        if status != global_constants.RNNTStatus.RNNT_STATUS_SUCCESS:
+            raise RuntimeError("Could not calculate forward scores")
+
+    else:
+        ### FLATTEN GRAD TENSOR ###
+        grads, grads_shape = rnnt_helper.flatten_tensor(grads)
+
+        status = wrapper.cost_and_grad(
+            acts=acts.data,
+            grads=grads.data,
+            costs=costs_repr,
+            pad_labels=labels.data,
+            label_lengths=label_lengths.data,
+            input_lengths=input_lengths.data,
+        )
+
+        if status != global_constants.RNNTStatus.RNNT_STATUS_SUCCESS:
+            raise RuntimeError("Could not calculate forward scores")
+
+    del gpu_workspace, wrapper
+    return True

tts/nemo/collections/asr/parts/numba/rnnt_loss/rnnt_numpy.py

@@ -0,0 +1,213 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Copyright 2018-2019, Mingkun Huang
+#
+# 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 numpy as np
+import torch
+from torch.autograd import Function, Variable
+from torch.nn import Module
+
+
+def _assert_no_grad(tensor):
+    assert not tensor.requires_grad, (
+        "gradients only computed for log_probs - please " "mark other tensors as not requiring gradients"
+    )
+
+
+def forward_pass(log_probs, labels, blank):
+    """
+    Computes probability of the forward variable alpha.
+
+    Args:
+        log_probs: Tensor of shape [T, U, V+1]
+        labels: Labels of shape [B, U]
+        blank: Index of the blank token.
+
+    Returns:
+        A tuple of the forward variable probabilities - alpha of shape [T, U]
+        and the log likelihood of this forward step.
+    """
+    T, U, _ = log_probs.shape
+    alphas = np.zeros((T, U), dtype='f')
+
+    for t in range(1, T):
+        alphas[t, 0] = alphas[t - 1, 0] + log_probs[t - 1, 0, blank]
+
+    for u in range(1, U):
+        alphas[0, u] = alphas[0, u - 1] + log_probs[0, u - 1, labels[u - 1]]
+    for t in range(1, T):
+        for u in range(1, U):
+            no_emit = alphas[t - 1, u] + log_probs[t - 1, u, blank]
+            emit = alphas[t, u - 1] + log_probs[t, u - 1, labels[u - 1]]
+            alphas[t, u] = np.logaddexp(emit, no_emit)
+
+    loglike = alphas[T - 1, U - 1] + log_probs[T - 1, U - 1, blank]
+    return alphas, loglike
+
+
+def backward_pass(log_probs, labels, blank):
+    """
+    Computes probability of the backward variable beta.
+
+    Args:
+        log_probs: Tensor of shape [T, U, V+1]
+        labels: Labels of shape [B, U]
+        blank: Index of the blank token.
+
+    Returns:
+        A tuple of the backward variable probabilities - beta of shape [T, U]
+        and the log likelihood of this backward step.
+    """
+    T, U, _ = log_probs.shape
+    betas = np.zeros((T, U), dtype='f')
+    betas[T - 1, U - 1] = log_probs[T - 1, U - 1, blank]
+
+    for t in reversed(range(T - 1)):
+        betas[t, U - 1] = betas[t + 1, U - 1] + log_probs[t, U - 1, blank]
+
+    for u in reversed(range(U - 1)):
+        betas[T - 1, u] = betas[T - 1, u + 1] + log_probs[T - 1, u, labels[u]]
+
+    for t in reversed(range(T - 1)):
+        for u in reversed(range(U - 1)):
+            no_emit = betas[t + 1, u] + log_probs[t, u, blank]
+            emit = betas[t, u + 1] + log_probs[t, u, labels[u]]
+            betas[t, u] = np.logaddexp(emit, no_emit)
+
+    return betas, betas[0, 0]
+
+
+def compute_gradient(log_probs, alphas, betas, labels, blank):
+    """
+    Computes the gradients of the log_probs with respect to the log probability of this step occuring.
+
+    Args:
+    Args:
+        log_probs: Tensor of shape [T, U, V+1]
+        alphas: Tensor of shape [T, U] which represents the forward variable.
+        betas: Tensor of shape [T, U] which represents the backward variable.
+        labels: Labels of shape [B, U]
+        blank: Index of the blank token.
+
+    Returns:
+        Gradients of shape [T, U, V+1] with respect to the forward log probability
+    """
+    T, U, _ = log_probs.shape
+    grads = np.full(log_probs.shape, -float("inf"))
+    log_like = betas[0, 0]
+
+    # // grad to last blank transition
+    grads[T - 1, U - 1, blank] = alphas[T - 1, U - 1]
+
+    grads[: T - 1, :, blank] = alphas[: T - 1, :] + betas[1:, :]
+    for u, l in enumerate(labels):
+        grads[:, u, l] = alphas[:, u] + betas[:, u + 1]
+
+    grads = -np.exp(grads + log_probs - log_like)
+    return grads
+
+
+def transduce(log_probs, labels, blank=0):
+    """
+    Args:
+        log_probs: 3D array with shape
+              [input len, output len + 1, vocab size]
+        labels: 1D array with shape [output time steps]
+    Returns:
+        float: The negative log-likelihood
+        3D array: Gradients with respect to the
+                    unnormalized input actications
+    """
+    alphas, ll_forward = forward_pass(log_probs, labels, blank)
+    betas, ll_backward = backward_pass(log_probs, labels, blank)
+    grads = compute_gradient(log_probs, alphas, betas, labels, blank)
+    return -ll_forward, grads
+
+
+def transduce_batch(log_probs, labels, flen, glen, blank=0):
+    """
+    Compute the transducer loss of the batch.
+
+    Args:
+        log_probs: [B, T, U, V+1]. Activation matrix normalized with log-softmax.
+        labels: [B, U+1] - ground truth labels with <SOS> padded as blank token in the beginning.
+        flen: Length vector of the acoustic sequence.
+        glen: Length vector of the target sequence.
+        blank: Id of the blank token.
+
+    Returns:
+        Batch of transducer forward log probabilities (loss) and the gradients of the activation matrix.
+    """
+    grads = np.zeros_like(log_probs)
+    costs = []
+    for b in range(log_probs.shape[0]):
+        t = int(flen[b])
+        u = int(glen[b]) + 1
+        ll, g = transduce(log_probs[b, :t, :u, :], labels[b, : u - 1], blank)
+        grads[b, :t, :u, :] = g
+        costs.append(ll)
+    return costs, grads
+
+
+class _RNNT(Function):
+    @staticmethod
+    def forward(ctx, acts, labels, act_lens, label_lens, blank):
+        costs, grads = transduce_batch(
+            acts.detach().cpu().numpy(), labels.cpu().numpy(), act_lens.cpu().numpy(), label_lens.cpu().numpy(), blank,
+        )
+
+        costs = torch.FloatTensor([sum(costs)])
+        grads = torch.Tensor(grads).to(acts)
+
+        ctx.grads = Variable(grads)
+        return costs
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return ctx.grads, None, None, None, None
+
+
+class RNNTLoss(Module):
+    """
+    Parameters:
+        `blank_label` (int): default 0 - label index of blank token
+    """
+
+    def __init__(self, blank: int = 0):
+        super(RNNTLoss, self).__init__()
+        self.blank = blank
+        self.rnnt = _RNNT.apply
+
+    def forward(self, acts, labels, act_lens, label_lens):
+        assert len(labels.size()) == 2
+        _assert_no_grad(labels)
+        _assert_no_grad(act_lens)
+        _assert_no_grad(label_lens)
+
+        acts = torch.nn.functional.log_softmax(acts, -1)
+        return self.rnnt(acts, labels, act_lens, label_lens, self.blank)

tts/nemo/collections/asr/parts/numba/rnnt_loss/rnnt_pytorch.py

@@ -0,0 +1,186 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Copyright 2018-2019, Mingkun Huang
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+from torch.autograd import Function
+from torch.nn import Module
+
+from nemo.collections.asr.parts.numba.rnnt_loss import rnnt
+
+__all__ = ['rnnt_loss', 'RNNTLossNumba']
+
+
+class _RNNTNumba(Function):
+    @staticmethod
+    def forward(ctx, acts, labels, act_lens, label_lens, blank, reduction):
+        """
+        log_probs: Tensor of (batch x seqLength x labelLength x outputDim) containing output from network
+        labels: 2 dimensional Tensor containing all the targets of the batch with zero padded
+        act_lens: Tensor of size (batch) containing size of each output sequence from the network
+        label_lens: Tensor of (batch) containing label length of each example
+        """
+        is_cuda = acts.is_cuda
+
+        certify_inputs(acts, labels, act_lens, label_lens)
+
+        loss_func = rnnt.rnnt_loss_gpu if is_cuda else rnnt.rnnt_loss_cpu
+        grads = torch.zeros_like(acts) if acts.requires_grad else None
+        minibatch_size = acts.size(0)
+        costs = torch.zeros(minibatch_size, device=acts.device, dtype=acts.dtype)
+
+        loss_func(
+            acts,
+            labels=labels,
+            input_lengths=act_lens,
+            label_lengths=label_lens,
+            costs=costs,
+            grads=grads,
+            blank_label=blank,
+            num_threads=0,
+        )
+
+        if reduction in ['sum', 'mean']:
+            costs = costs.sum().unsqueeze_(-1)
+            if reduction == 'mean':
+                costs /= minibatch_size
+
+                if grads is not None:
+                    grads /= minibatch_size
+
+        # costs = costs.to(log_probs.device)
+        ctx.grads = grads
+
+        return costs
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        if grad_output is not None and ctx.grads is not None:
+            grad_output = grad_output.view(-1, 1, 1, 1).to(ctx.grads)
+            return ctx.grads.mul_(grad_output), None, None, None, None, None
+
+
+def rnnt_loss(acts, labels, act_lens, label_lens, blank=0, reduction='mean'):
+    """RNN Transducer Loss
+    Args:
+        acts: Tensor of (batch x seqLength x labelLength x outputDim) containing output from network
+        labels: 2 dimensional Tensor containing all the targets of the batch with zero padded
+        act_lens: Tensor of size (batch) containing size of each output sequence from the network
+        label_lens: Tensor of (batch) containing label length of each example
+        blank (int, optional): blank label. Default: 0.
+        reduction (string, optional): Specifies the reduction to apply to the output:
+            'none' | 'mean' | 'sum'. 'none': no reduction will be applied,
+            'mean': the output losses will be divided by the target lengths and
+            then the mean over the batch is taken. Default: 'mean'
+    """
+    if not acts.is_cuda:
+        acts = torch.nn.functional.log_softmax(acts, -1)
+
+    return _RNNTNumba.apply(acts, labels, act_lens, label_lens, blank, reduction)
+
+
+class RNNTLossNumba(Module):
+    """
+    Parameters:
+        blank (int, optional): blank label. Default: 0.
+        reduction (string, optional): Specifies the reduction to apply to the output:
+            'none' | 'mean' | 'sum'. 'none': no reduction will be applied,
+            'mean': the output losses will be divided by the target lengths and
+            then the mean over the batch is taken. Default: 'mean'
+    """
+
+    def __init__(self, blank=0, reduction='mean'):
+        super(RNNTLossNumba, self).__init__()
+        self.blank = blank
+        self.reduction = reduction
+        self.loss = _RNNTNumba.apply
+
+    def forward(self, acts, labels, act_lens, label_lens):
+        """
+        log_probs: Tensor of (batch x seqLength x labelLength x outputDim) containing output from network
+        labels: 2 dimensional Tensor containing all the targets of the batch with zero padded
+        act_lens: Tensor of size (batch) containing size of each output sequence from the network
+        label_lens: Tensor of (batch) containing label length of each example
+        """
+        if not acts.is_cuda:
+            # NOTE manually done log_softmax for CPU version,
+            # log_softmax is computed within GPU version.
+            acts = torch.nn.functional.log_softmax(acts, -1)
+
+        return self.loss(acts, labels, act_lens, label_lens, self.blank, self.reduction)
+
+
+def check_type(var, t, name):
+    if var.dtype is not t:
+        raise TypeError("{} must be {}".format(name, t))
+
+
+def check_contiguous(var, name):
+    if not var.is_contiguous():
+        raise ValueError("{} must be contiguous".format(name))
+
+
+def check_dim(var, dim, name):
+    if len(var.shape) != dim:
+        raise ValueError("{} must be {}D".format(name, dim))
+
+
+def certify_inputs(log_probs, labels, lengths, label_lengths):
+    # check_type(log_probs, torch.float32, "log_probs")
+    check_type(labels, torch.int32, "labels")
+    check_type(label_lengths, torch.int32, "label_lengths")
+    check_type(lengths, torch.int32, "lengths")
+    check_contiguous(log_probs, "log_probs")
+    check_contiguous(labels, "labels")
+    check_contiguous(label_lengths, "label_lengths")
+    check_contiguous(lengths, "lengths")
+
+    if lengths.shape[0] != log_probs.shape[0]:
+        raise ValueError(
+            f"Must have a length per example. "
+            f"Given lengths dim: {lengths.shape[0]}, "
+            f"Log probs dim : {log_probs.shape[0]}"
+        )
+    if label_lengths.shape[0] != log_probs.shape[0]:
+        raise ValueError(
+            "Must have a label length per example. "
+            f"Given label lengths dim : {label_lengths.shape[0]}, "
+            f"Log probs dim : {log_probs.shape[0]}"
+        )
+
+    check_dim(log_probs, 4, "log_probs")
+    check_dim(labels, 2, "labels")
+    check_dim(lengths, 1, "lenghts")
+    check_dim(label_lengths, 1, "label_lenghts")
+    max_T = torch.max(lengths)
+    max_U = torch.max(label_lengths)
+    T, U = log_probs.shape[1:3]
+    if T != max_T:
+        raise ValueError(f"Input length mismatch! Given T: {T}, Expected max T from input lengths: {max_T}")
+    if U != max_U + 1:
+        raise ValueError(f"Output length mismatch! Given U: {U}, Expected max U from target lengths: {max_U} + 1")

tts/nemo/collections/asr/parts/numba/rnnt_loss/utils/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tts/nemo/collections/asr/parts/numba/rnnt_loss/utils/cpu_utils/__init__.py

@@ -0,0 +1,27 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Copyright 2018-2019, Mingkun Huang
+#
+# 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.

tts/nemo/collections/asr/parts/numba/rnnt_loss/utils/cpu_utils/cpu_rnnt.py

@@ -0,0 +1,378 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Copyright 2018-2019, Mingkun Huang
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+import multiprocessing
+from typing import Optional
+
+import numba
+import torch
+
+from nemo.collections.asr.parts.numba.rnnt_loss.utils import global_constants
+
+
+def log_sum_exp(a: torch.Tensor, b: torch.Tensor):
+    """
+    Logsumexp with safety checks for infs.
+    """
+    if torch.isinf(a):
+        return b
+
+    if torch.isinf(b):
+        return a
+
+    if a > b:
+        return math.log1p(math.exp(b - a)) + a
+    else:
+        return math.log1p(math.exp(a - b)) + b
+
+
+class CpuRNNT_index:
+    def __init__(self, U: int, maxU: int, minibatch: int, alphabet_size: int, batch_first: bool):
+        """
+        A placeholder Index computation class that emits the resolved index in a flattened tensor,
+        mimicing pointer indexing in CUDA kernels on the CPU.
+
+        Args:
+            U: Length of the current target sample (without padding).
+            maxU: Max Length of the padded target samples.
+            minibatch: Minibatch index
+            alphabet_size: Size of the vocabulary including RNNT blank - V+1.
+            batch_first: Bool flag determining if batch index is first or third.
+        """
+        super(CpuRNNT_index, self).__init__()
+        self.U = U
+        self.maxU = maxU
+        self.minibatch = minibatch
+        self.alphabet_size = alphabet_size
+        self.batch_first = batch_first
+
+    def __call__(self, t: int, u: int, v: Optional[int] = None):
+        # if indexing all the values of the vocabulary, then only t, u are provided
+        if v is None:
+            return t * self.U + u
+        else:
+            # otherwise, t, u, v are provided to index particular value in the vocabulary.
+            if self.batch_first:
+                return (t * self.maxU + u) * self.alphabet_size + v
+            else:
+                return (t * self.maxU + u) * self.minibatch * self.alphabet_size + v
+
+
+class CpuRNNT_metadata:
+    def __init__(
+        self,
+        T: int,
+        U: int,
+        workspace: torch.Tensor,
+        bytes_used: int,
+        blank: int,
+        labels: torch.Tensor,
+        log_probs: torch.Tensor,
+        idx: CpuRNNT_index,
+    ):
+        """
+        Metadata for CPU based RNNT loss calculation. Holds the working space memory.
+
+        Args:
+            T: Length of the acoustic sequence (without padding).
+            U: Length of the target sequence (without padding).
+            workspace: Working space memory for the CPU.
+            bytes_used: Number of bytes currently used for indexing the working space memory. Generally 0.
+            blank: Index of the blank token in the vocabulary.
+            labels: Ground truth padded labels matrix of shape [B, U]
+            log_probs: Log probs / activation matrix of flattented shape [B, T, U, V+1]
+            idx:
+        """
+        super(CpuRNNT_metadata, self).__init__()
+
+        self.alphas = workspace[bytes_used : bytes_used + T * U]
+        bytes_used += T * U
+
+        self.betas = workspace[bytes_used : bytes_used + T * U]
+        bytes_used += T * U
+
+        self.log_probs2 = workspace[bytes_used : bytes_used + T * U * 2]  # // only store blank & label
+        bytes_used += T * U * 2
+
+        self.bytes_used = bytes_used
+
+        self.setup_probs(T, U, labels, blank, log_probs, idx)
+
+    def setup_probs(
+        self, T: int, U: int, labels: torch.Tensor, blank: int, log_probs: torch.Tensor, idx: CpuRNNT_index
+    ):
+        # initialize the log probs memory for blank and label token.
+        for t in range(T):
+            for u in range(U):
+                offset = (t * U + u) * 2  # mult with 2 is for selecting either blank or label token. Odd idx is blank.
+                self.log_probs2[offset] = log_probs[idx(t, u, blank)]
+                # // labels do not have first blank
+                if u < U - 1:
+                    self.log_probs2[offset + 1] = log_probs[idx(t, u, labels[u])]
+
+
+class CPURNNT:
+    def __init__(
+        self,
+        minibatch: int,
+        maxT: int,
+        maxU: int,
+        alphabet_size: int,
+        workspace: torch.Tensor,
+        blank: int,
+        num_threads: int,
+        batch_first: bool,
+    ):
+        """
+        Helper class to compute the Transducer Loss on CPU.
+
+        Args:
+            minibatch: Size of the minibatch b.
+            maxT: The maximum possible acoustic sequence length. Represents T in the logprobs tensor.
+            maxU: The maximum possible target sequence length. Represents U in the logprobs tensor.
+            alphabet_size: The vocabulary dimension V+1 (inclusive of RNNT blank).
+            workspace: An allocated chunk of memory that will be sliced off and reshaped into required
+                blocks used as working memory.
+            blank: Index of the RNNT blank token in the vocabulary. Generally the first or last token in the vocab.
+            num_threads: Number of OMP threads to launch.
+            batch_first: Bool that decides if batch dimension is first or third.
+        """
+        self.minibatch_ = minibatch
+        self.maxT_ = maxT
+        self.maxU_ = maxU
+        self.alphabet_size_ = alphabet_size
+        self.workspace = workspace  # a flat vector of floatX numbers that represents allocated memory slices
+        self.blank_ = blank
+        self.num_threads_ = num_threads
+        self.batch_first = batch_first
+
+        if num_threads > 0:
+            numba.set_num_threads(min(multiprocessing.cpu_count(), num_threads))
+        else:
+            self.num_threads_ = numba.get_num_threads()
+
+    def cost_and_grad_kernel(
+        self,
+        log_probs: torch.Tensor,
+        grad: torch.Tensor,
+        labels: torch.Tensor,
+        mb: int,
+        T: int,
+        U: int,
+        bytes_used: int,
+    ):
+        idx = CpuRNNT_index(U, self.maxU_, self.minibatch_, self.alphabet_size_, self.batch_first)
+        rnntm = CpuRNNT_metadata(T, U, self.workspace, bytes_used, self.blank_, labels, log_probs, idx)
+
+        if self.batch_first:
+            # zero grads
+            grad *= 0.0
+
+        llForward = self.compute_alphas(rnntm.log_probs2, T, U, rnntm.alphas)
+        llBackward = self.compute_betas_and_grads(
+            grad, rnntm.log_probs2, T, U, rnntm.alphas, rnntm.betas, labels, llForward
+        )
+
+        diff = (llForward - llBackward).abs()
+        if diff > 0.1:
+            print(f"WARNING: Forward backward likelihood mismatch : {diff}")
+
+        return -llForward
+
+    def compute_alphas(self, log_probs: torch.Tensor, T: int, U: int, alphas: torch.Tensor):
+        """
+        Compute the probability of the forward variable alpha.
+
+        Args:
+            log_probs: Flattened tensor [B, T, U, V+1]
+            T: Length of the acoustic sequence T (not padded).
+            U: Length of the target sequence U (not padded).
+            alphas: Working space memory for alpha of shape [B, T, U].
+
+        Returns:
+            Loglikelihood of the forward variable alpha.
+        """
+        idx = CpuRNNT_index(U, self.maxU_, self.minibatch_, self.alphabet_size_, self.batch_first)
+
+        alphas[0] = 0
+        for t in range(T):
+            for u in range(U):
+                if u == 0 and t > 0:
+                    alphas[idx(t, 0)] = alphas[idx(t - 1, 0)] + log_probs[idx(t - 1, 0) * 2]
+
+                if t == 0 and u > 0:
+                    alphas[idx(0, u)] = alphas[idx(0, u - 1)] + log_probs[idx(0, u - 1) * 2 + 1]
+
+                if t > 0 and u > 0:
+                    no_emit = alphas[idx(t - 1, u)] + log_probs[idx(t - 1, u) * 2]
+                    emit = alphas[idx(t, u - 1)] + log_probs[idx(t, u - 1) * 2 + 1]
+                    alphas[idx(t, u)] = log_sum_exp(emit, no_emit)
+
+        loglike = alphas[idx(T - 1, U - 1)] + log_probs[idx(T - 1, U - 1) * 2]
+        return loglike
+
+    def compute_betas_and_grads(
+        self,
+        grad: torch.Tensor,
+        log_probs: torch.Tensor,
+        T: int,
+        U: int,
+        alphas: torch.Tensor,
+        betas: torch.Tensor,
+        labels: torch.Tensor,
+        logll: torch.Tensor,
+    ):
+        """
+        Compute backward variable beta as well as gradients of the activation matrix wrt loglikelihood
+        of forward variable.
+
+        Args:
+            grad: Working space memory of flattened shape [B, T, U, V+1]
+            log_probs: Activatio tensor of flattented shape [B, T, U, V+1]
+            T: Length of the acoustic sequence T (not padded).
+            U: Length of the target sequence U (not padded).
+            alphas: Working space memory for alpha of shape [B, T, U].
+            betas: Working space memory for alpha of shape [B, T, U].
+            labels: Ground truth label of shape [B, U]
+            logll: Loglikelihood of the forward variable.
+
+        Returns:
+            Loglikelihood of the forward variable and inplace updates the grad tensor.
+        """
+        idx = CpuRNNT_index(U, self.maxU_, self.minibatch_, self.alphabet_size_, self.batch_first)
+        betas[idx(T - 1, U - 1)] = log_probs[idx(T - 1, U - 1) * 2]
+
+        for t in range(T - 1, -1, -1):
+            for u in range(U - 1, -1, -1):
+                if (u == U - 1) and (t < T - 1):
+                    betas[idx(t, U - 1)] = betas[idx(t + 1, U - 1)] + log_probs[idx(t, U - 1) * 2]
+
+                if (t == T - 1) and (u < U - 1):
+                    betas[idx(T - 1, u)] = betas[idx(T - 1, u + 1)] + log_probs[idx(T - 1, u) * 2 + 1]
+
+                if (t < T - 1) and (u < U - 1):
+                    no_emit = betas[idx(t + 1, u)] + log_probs[idx(t, u) * 2]
+                    emit = betas[idx(t, u + 1)] + log_probs[idx(t, u) * 2 + 1]
+                    betas[idx(t, u)] = log_sum_exp(emit, no_emit)
+
+        loglike = betas[0]
+        # // Gradients w.r.t. log probabilities
+        for t in range(T):
+            for u in range(U):
+                if t < T - 1:
+                    g = alphas[idx(t, u)] + betas[idx(t + 1, u)]
+                    grad[idx(t, u, self.blank_)] = -torch.exp(log_probs[idx(t, u) * 2] + g - loglike)
+
+                if u < U - 1:
+                    g = alphas[idx(t, u)] + betas[idx(t, u + 1)]
+                    grad[idx(t, u, labels[u])] = -torch.exp(log_probs[idx(t, u) * 2 + 1] + g - loglike)
+
+        # // gradient to the last blank transition
+        grad[idx(T - 1, U - 1, self.blank_)] = -torch.exp(
+            log_probs[idx(T - 1, U - 1) * 2] + alphas[idx(T - 1, U - 1)] - loglike
+        )
+
+        return loglike
+
+    def cost_and_grad(
+        self,
+        log_probs: torch.Tensor,
+        grads: torch.Tensor,
+        costs: torch.Tensor,
+        flat_labels: torch.Tensor,
+        label_lengths: torch.Tensor,
+        input_lengths: torch.Tensor,
+    ) -> global_constants.RNNTStatus:
+        # // per minibatch memory
+        per_minibatch_bytes = 0
+
+        # // alphas & betas
+        per_minibatch_bytes += self.maxT_ * self.maxU_ * 2
+
+        # // blank & label log probability cache
+        per_minibatch_bytes += self.maxT_ * self.maxU_ * 2
+
+        for mb in range(self.minibatch_):
+            T = input_lengths[mb]  # // Length of utterance (time)
+            U = label_lengths[mb] + 1  # // Number of labels in transcription
+            batch_size = self.alphabet_size_
+            if self.batch_first:
+                batch_size = self.maxT_ * self.maxU_ * self.alphabet_size_
+
+            costs[mb] = self.cost_and_grad_kernel(
+                log_probs[(mb * batch_size) :],
+                grads[(mb * batch_size) :],
+                flat_labels[(mb * (self.maxU_ - 1)) :],
+                mb,
+                T,
+                U,
+                mb * per_minibatch_bytes,
+            )
+
+        return global_constants.RNNTStatus.RNNT_STATUS_SUCCESS
+
+    def score_forward(
+        self,
+        log_probs: torch.Tensor,
+        costs: torch.Tensor,
+        flat_labels: torch.Tensor,
+        label_lengths: torch.Tensor,
+        input_lengths: torch.Tensor,
+    ):
+        # // per minibatch memory
+        per_minibatch_bytes = 0
+
+        # // alphas & betas
+        per_minibatch_bytes += self.maxT_ * self.maxU_ * 2
+
+        # // blank & label log probability cache
+        per_minibatch_bytes += self.maxT_ * self.maxU_ * 2
+
+        for mb in range(self.minibatch_):
+            T = input_lengths[mb]  # // Length of utterance (time)
+            U = label_lengths[mb] + 1  # // Number of labels in transcription
+            batch_size = self.alphabet_size_
+            if self.batch_first:
+                batch_size = self.maxT_ * self.maxU_ * self.alphabet_size_
+
+            idx = CpuRNNT_index(U, self.maxU_, self.minibatch_, self.alphabet_size_, self.batch_first)
+            rnntm = CpuRNNT_metadata(
+                T,
+                U,
+                self.workspace,
+                mb * per_minibatch_bytes,
+                self.blank_,
+                flat_labels[(mb * (self.maxU_ - 1)) :],
+                log_probs[(mb * batch_size) :],
+                idx,
+            )
+
+            costs[mb] = -self.compute_alphas(rnntm.log_probs2, T, U, rnntm.alphas)
+
+        return global_constants.RNNTStatus.RNNT_STATUS_SUCCESS

tts/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/__init__.py

@@ -0,0 +1,27 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Copyright 2018-2019, Mingkun Huang
+#
+# 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.

tts/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/gpu_rnnt.py

@@ -0,0 +1,254 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Copyright 2018-2019, Mingkun Huang
+#
+# 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 multiprocessing
+from typing import Optional
+
+import numba
+import torch
+from numba import cuda
+
+from nemo.collections.asr.parts.numba.rnnt_loss.utils import global_constants
+from nemo.collections.asr.parts.numba.rnnt_loss.utils.cuda_utils import gpu_rnnt_kernel, reduce
+
+
+class GPURNNT:
+    def __init__(
+        self, minibatch: int, maxT: int, maxU: int, alphabet_size: int, workspace, blank: int, num_threads: int, stream
+    ):
+        """
+        Helper class to launch the CUDA Kernels to compute the Transducer Loss.
+
+        Args:
+            minibatch: Int representing the batch size.
+            maxT: The maximum possible acoustic sequence length. Represents T in the logprobs tensor.
+            maxU: The maximum possible target sequence length. Represents U in the logprobs tensor.
+            alphabet_size: The vocabulary dimension V+1 (inclusive of RNNT blank).
+            workspace: An allocated chunk of memory that will be sliced off and reshaped into required
+                blocks used as working memory.
+            blank: Index of the RNNT blank token in the vocabulary. Generally the first or last token in the vocab.
+            num_threads: Number of OMP threads to launch.
+            stream: Numba Cuda Stream.
+        """
+        self.minibatch_ = minibatch
+        self.maxT_ = maxT
+        self.maxU_ = maxU
+        self.alphabet_size_ = alphabet_size
+        self.gpu_workspace = cuda.as_cuda_array(
+            workspace
+        )  # a flat vector of floatX numbers that represents allocated memory slices
+        self.blank_ = blank
+        self.num_threads_ = num_threads
+        self.stream_ = stream  # type: cuda.cudadrv.driver.Stream
+
+        if num_threads > 0:
+            numba.set_num_threads(min(multiprocessing.cpu_count(), num_threads))
+        else:
+            self.num_threads_ = numba.get_num_threads()
+
+    def log_softmax(self, acts: torch.Tensor, denom: torch.Tensor):
+        """
+        Computes the log softmax denominator of the input activation tensor
+        and stores the result in denom.
+
+        Args:
+            acts: Activation tensor of shape [B, T, U, V+1]. The input must be represented as a flat tensor
+                of shape [B * T * U * (V+1)] to allow pointer indexing.
+            denom: A zero tensor of same shape as acts.
+
+        Updates:
+            This kernel inplace updates the `denom` tensor
+        """
+        # // trans_acts + pred_acts -> log_softmax denominator
+        reduce.reduce_max(
+            acts,
+            denom,
+            rows=self.alphabet_size_,
+            cols=self.minibatch_ * self.maxT_ * self.maxU_,
+            minus=False,
+            stream=self.stream_,
+        )
+
+        reduce.reduce_exp(
+            acts,
+            denom,
+            rows=self.alphabet_size_,
+            cols=self.minibatch_ * self.maxT_ * self.maxU_,
+            minus=True,
+            stream=self.stream_,
+        )
+
+    def compute_cost_and_score(
+        self,
+        acts: torch.Tensor,
+        grads: Optional[torch.Tensor],
+        costs: torch.Tensor,
+        labels: torch.Tensor,
+        label_lengths: torch.Tensor,
+        input_lengths: torch.Tensor,
+    ) -> global_constants.RNNTStatus:
+        """
+        Compute both the loss and the gradients.
+
+        Args:
+            acts: A flattened tensor of shape [B, T, U, V+1] representing the activation matrix.
+            grad: A flattented zero tensor of same shape as acts.
+            costs: A zero vector of length B which will be updated inplace with the log probability costs.
+            flat_labels: A flattened matrix of labels of shape [B, U]
+            label_lengths: A vector of length B that contains the original lengths of the acoustic sequence.
+            input_lengths: A vector of length B that contains the original lengths of the target sequence.
+
+        Updates:
+            This will launch kernels that will update inline the following variables:
+            -   grads: Gradients of the activation matrix wrt the costs vector.
+            -   costs: Negative log likelihood of the forward variable.
+
+        Returns:
+            An enum that either represents a successful RNNT operation or failure.
+        """
+        training = grads is not None
+        used_offset = 0
+
+        # // denom
+        denom = self.gpu_workspace[used_offset : used_offset + self.maxT_ * self.maxU_ * self.minibatch_]
+        used_offset += self.maxT_ * self.maxU_ * self.minibatch_
+
+        # // alphas & betas
+        alphas = self.gpu_workspace[used_offset : used_offset + self.maxT_ * self.maxU_ * self.minibatch_]
+        used_offset += self.maxT_ * self.maxU_ * self.minibatch_
+        betas = self.gpu_workspace[used_offset : used_offset + self.maxT_ * self.maxU_ * self.minibatch_]
+        used_offset += self.maxT_ * self.maxU_ * self.minibatch_
+
+        # // logllh
+        llForward = self.gpu_workspace[used_offset : used_offset + self.minibatch_]
+        used_offset += self.minibatch_
+        llBackward = self.gpu_workspace[used_offset : used_offset + self.minibatch_]
+        used_offset += self.minibatch_
+
+        if training:
+            grads *= 0.0  # zero grads
+
+        ######## START EXECUTION ########
+        self.log_softmax(acts, denom)
+
+        # Compute alphas
+        gpu_rnnt_kernel.compute_alphas_kernel[self.minibatch_, self.maxU_, self.stream_, 0](
+            acts,
+            denom,
+            alphas,
+            llForward,
+            input_lengths,
+            label_lengths,
+            labels,
+            self.minibatch_,
+            self.maxT_,
+            self.maxU_,
+            self.alphabet_size_,
+            self.blank_,
+        )
+
+        if training:
+            # Compute betas
+            gpu_rnnt_kernel.compute_betas_kernel[self.minibatch_, self.maxU_, self.stream_, 0](
+                acts,
+                denom,
+                betas,
+                llBackward,
+                input_lengths,
+                label_lengths,
+                labels,
+                self.minibatch_,
+                self.maxT_,
+                self.maxU_,
+                self.alphabet_size_,
+                self.blank_,
+            )
+
+            # Compute gradient
+            grad_blocks_per_grid = self.minibatch_ * self.maxT_ * self.maxU_
+            grad_threads_per_block = gpu_rnnt_kernel.GPU_RNNT_THREAD_SIZE
+            gpu_rnnt_kernel.compute_grad_kernel[grad_blocks_per_grid, grad_threads_per_block, self.stream_, 0](
+                grads,
+                acts,
+                denom,
+                alphas,
+                betas,
+                llForward,
+                input_lengths,
+                label_lengths,
+                labels,
+                self.minibatch_,
+                self.maxT_,
+                self.maxU_,
+                self.alphabet_size_,
+                self.blank_,
+            )
+
+        # // cost
+        costs.copy_to_device(llForward, stream=self.stream_)
+        self.stream_.synchronize()
+
+        # compute negative log likelihood.
+        for mb in range(self.minibatch_):
+            costs[mb] = -costs[mb]
+
+        return global_constants.RNNTStatus.RNNT_STATUS_SUCCESS
+
+    def cost_and_grad(
+        self,
+        acts: torch.Tensor,
+        grads: torch.Tensor,
+        costs: torch.Tensor,
+        pad_labels: torch.Tensor,
+        label_lengths: torch.Tensor,
+        input_lengths: torch.Tensor,
+    ):
+        if (
+            acts is None
+            or grads is None
+            or costs is None
+            or pad_labels is None
+            or label_lengths is None
+            or input_lengths is None
+        ):
+            return global_constants.RNNTStatus.RNNT_STATUS_INVALID_VALUE
+
+        return self.compute_cost_and_score(acts, grads, costs, pad_labels, label_lengths, input_lengths)
+
+    def score_forward(
+        self,
+        acts: torch.Tensor,
+        costs: torch.Tensor,
+        pad_labels: torch.Tensor,
+        label_lengths: torch.Tensor,
+        input_lengths: torch.Tensor,
+    ):
+        if acts is None or costs is None or pad_labels is None or label_lengths is None or input_lengths is None:
+            return global_constants.RNNTStatus.RNNT_STATUS_INVALID_VALUE
+
+        return self.compute_cost_and_score(acts, None, costs, pad_labels, label_lengths, input_lengths)

tts/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/gpu_rnnt_kernel.py

@@ -0,0 +1,365 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Copyright 2018-2019, Mingkun Huang
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+
+import torch
+from numba import cuda
+
+from nemo.collections.asr.parts.numba.rnnt_loss.utils import rnnt_helper
+
+GPU_RNNT_THREAD_SIZE = 128
+
+
+@cuda.jit(device=True, inline=True)
+def logp(
+    denom: torch.Tensor, acts: torch.Tensor, maxT: int, maxU: int, alphabet_size: int, mb: int, t: int, u: int, v: int
+):
+    """
+    Compute the sum of log probability from the activation tensor and its denominator.
+
+    Args:
+        denom: Tensor of shape [B, T, U] flattened. Represents the denominator of the logprobs activation tensor
+            across entire vocabulary.
+        acts: Tensor of shape [B, T, U, V+1] flattened. Represents the logprobs activation tensor.
+        maxT: The maximum possible acoustic sequence length. Represents T in the logprobs tensor.
+        maxU: The maximum possible target sequence length. Represents U in the logprobs tensor.
+        alphabet_size: The vocabulary dimension V+1 (inclusive of RNNT blank).
+        mb: Batch indexer.
+        t: Acoustic sequence timestep indexer.
+        u: Target sequence timestep indexer.
+        v: Vocabulary token indexer.
+
+    Returns:
+        The sum of logprobs[mb, t, u, v] + denom[mb, t, u]
+    """
+    col = (mb * maxT + t) * maxU + u
+    return denom[col] + acts[col * alphabet_size + v]
+
+
+@cuda.jit()
+def compute_alphas_kernel(
+    acts: torch.Tensor,
+    denom: torch.Tensor,
+    alphas: torch.Tensor,
+    llForward: torch.Tensor,
+    xlen: torch.Tensor,
+    ylen: torch.Tensor,
+    mlabels: torch.Tensor,  # [B]
+    minibatch: int,
+    maxT: int,
+    maxU: int,
+    alphabet_size: int,
+    blank_: int,
+):
+    """
+    Compute alpha (forward variable) probabilities over the transduction step.
+
+    Args:
+        acts: Tensor of shape [B, T, U, V+1] flattened. Represents the logprobs activation tensor.
+        denom: Tensor of shape [B, T, U] flattened. Represents the denominator of the logprobs activation tensor
+            across entire vocabulary.
+        alphas: Zero tensor of shape [B, T, U]. Will be updated inside the kernel with the forward variable
+            probabilities.
+        llForward: Zero tensor of shape [B]. Represents the log-likelihood of the forward pass.
+            Returned as the forward pass loss that is reduced by the optimizer.
+        xlen: Vector of length B which contains the actual acoustic sequence lengths in the padded
+            activation tensor.
+        ylen: Vector of length B which contains the actual target sequence lengths in the padded
+            activation tensor.
+        mlabels: Matrix of shape [B, U+1] (+1 here is due to <SOS> token - usually the RNNT blank).
+            The matrix contains the padded target transcription that must be predicted.
+        minibatch: Int representing the batch size.
+        maxT: The maximum possible acoustic sequence length. Represents T in the logprobs tensor.
+        maxU: The maximum possible target sequence length. Represents U in the logprobs tensor.
+        alphabet_size: The vocabulary dimension V+1 (inclusive of RNNT blank).
+        blank_: Index of the RNNT blank token in the vocabulary. Generally the first or last token in the vocab.
+
+    Updates:
+        Kernel inplace updates the following inputs:
+        -   alphas: forward variable scores.
+        -   llForward: log-likelihood of forward variable.
+    """
+    # // launch B blocks, each block has U threads
+    b = cuda.blockIdx.x  # // batch id
+    u = cuda.threadIdx.x  # label id, u
+    T = xlen[b]  # select AM length of current sample
+    U = ylen[b] + 1  # select target length of current sample, +1 for the blank token
+
+    labels: torch.Tensor = mlabels[b]  # mb label start point, equivalent to mlabels + b * (maxU - 1)
+    offset = b * maxT * maxU  # pointer indexing offset
+
+    # alphas += offset # pointer offset, ignored since we explicitly add offset
+
+    # Initilize alpha[b, t=0, u=0] for all b in B
+    if u == 0:
+        alphas[offset] = 0
+
+    # sync until all alphas are initialized
+    cuda.syncthreads()
+
+    # Ordinary alpha calculations, broadcast across B=b and U=u
+    # Look up forward variable calculation from rnnt_numpy.forward_pass()
+    for n in range(1, T + U - 1):
+        t = n - u
+
+        if u == 0:
+            # for t in range(1, T) step to initialize alphas[b, t, 0]
+            if t > 0 and t < T:
+                alphas[offset + t * maxU + u] = alphas[offset + (t - 1) * maxU + u] + logp(
+                    denom, acts, maxT, maxU, alphabet_size, b, t - 1, 0, blank_
+                )
+        elif u < U:
+            # for u in range(1, U) step to initialize alphas[b, 0, u]
+            if t == 0:
+                alphas[offset + u] = alphas[offset + u - 1] + logp(
+                    denom, acts, maxT, maxU, alphabet_size, b, 0, u - 1, labels[u - 1]
+                )
+
+            # for t in range(1, T) for u in range(1, U) step to compute alphas[b, t, u]
+            elif t > 0 and t < T:
+                no_emit = alphas[offset + (t - 1) * maxU + u] + logp(
+                    denom, acts, maxT, maxU, alphabet_size, b, t - 1, u, blank_
+                )
+                emit = alphas[offset + t * maxU + u - 1] + logp(
+                    denom, acts, maxT, maxU, alphabet_size, b, t, u - 1, labels[u - 1]
+                )
+
+                alphas[offset + t * maxU + u] = rnnt_helper.log_sum_exp(emit, no_emit)
+
+        # sync across all B=b and U=u
+        cuda.syncthreads()
+
+    # After final sync, alphas[b, T-1, U - 1] + logprobs[b, T-1, U-1, blank] + denom[b, T-1, U-1] gives
+    # log-likelihood of forward pass.
+    if u == 0:
+        loglike = alphas[offset + (T - 1) * maxU + U - 1] + logp(
+            denom, acts, maxT, maxU, alphabet_size, b, T - 1, U - 1, blank_
+        )
+        llForward[b] = loglike
+
+
+@cuda.jit()
+def compute_betas_kernel(
+    acts: torch.Tensor,
+    denom: torch.Tensor,
+    betas: torch.Tensor,
+    llBackward: torch.Tensor,
+    xlen: torch.Tensor,
+    ylen: torch.Tensor,
+    mlabels: torch.Tensor,  # [B, U]
+    minibatch: int,
+    maxT: int,
+    maxU: int,
+    alphabet_size: int,
+    blank_: int,
+):
+    """
+    Compute beta (backward variable) probabilities over the transduction step.
+
+    Args:
+        acts: Tensor of shape [B, T, U, V+1] flattened. Represents the logprobs activation tensor.
+        denom: Tensor of shape [B, T, U] flattened. Represents the denominator of the logprobs activation tensor
+            across entire vocabulary.
+        betas: Zero tensor of shape [B, T, U]. Will be updated inside the kernel with the backward variable
+            probabilities.
+        llBackward: Zero tensor of shape [B]. Represents the log-likelihood of the backward pass.
+            Returned as the backward pass loss that is reduced by the optimizer.
+        xlen: Vector of length B which contains the actual acoustic sequence lengths in the padded
+            activation tensor.
+        ylen: Vector of length B which contains the actual target sequence lengths in the padded
+            activation tensor.
+        mlabels: Matrix of shape [B, U+1] (+1 here is due to <SOS> token - usually the RNNT blank).
+            The matrix contains the padded target transcription that must be predicted.
+        minibatch: Int representing the batch size.
+        maxT: The maximum possible acoustic sequence length. Represents T in the logprobs tensor.
+        maxU: The maximum possible target sequence length. Represents U in the logprobs tensor.
+        alphabet_size: The vocabulary dimension V+1 (inclusive of RNNT blank).
+        blank_: Index of the RNNT blank token in the vocabulary. Generally the first or last token in the vocab.
+
+    Updates:
+        Kernel inplace updates the following inputs:
+        -   betas: backward variable scores.
+        -   llBackward: log-likelihood of backward variable.
+    """
+    # // launch B blocks, each block has U threads
+    b = cuda.blockIdx.x  # // batch id
+    u = cuda.threadIdx.x  # label id, u
+    T = xlen[b]  # select AM length of current sample
+    U = ylen[b] + 1  # select target length of current sample, +1 for the blank token
+
+    labels: torch.Tensor = mlabels[b]  # mb label start point, equivalent to mlabels + b * (maxU - 1)
+    offset = b * maxT * maxU  # pointer indexing offset
+
+    # betas += offset # pointer offset, ignored since we explicitly add offset
+
+    # Initilize beta[b, t=T-1, u=U-1] for all b in B with log_probs[b, t=T-1, u=U-1, blank]
+    if u == 0:
+        betas[offset + (T - 1) * maxU + U - 1] = logp(denom, acts, maxT, maxU, alphabet_size, b, T - 1, U - 1, blank_)
+
+    # sync until all betas are initialized
+    cuda.syncthreads()
+
+    # Ordinary beta calculations, broadcast across B=b and U=u
+    # Look up backward variable calculation from rnnt_numpy.backward_pass()
+    for n in range(T + U - 2, -1, -1):
+        t = n - u
+
+        if u == (U - 1):
+            # for t in reversed(range(T - 1)) step to initialize betas[b, t, U-1]
+            if t >= 0 and t < (T - 1):
+                betas[offset + t * maxU + U - 1] = betas[offset + (t + 1) * maxU + U - 1] + logp(
+                    denom, acts, maxT, maxU, alphabet_size, b, t, U - 1, blank_
+                )
+        elif u < U:
+            if t == T - 1:
+                # for u in reversed(range(U - 1)) step to initialize betas[b, T-1, u]
+                betas[offset + (T - 1) * maxU + u] = betas[offset + (T - 1) * maxU + u + 1] + logp(
+                    denom, acts, maxT, maxU, alphabet_size, b, T - 1, u, labels[u]
+                )
+            elif (t >= 0) and (t < T - 1):
+                # for t in reversed(range(T - 1)) for u in reversed(range(U - 1)) step to compute betas[b, t, u]
+                no_emit = betas[offset + (t + 1) * maxU + u] + logp(
+                    denom, acts, maxT, maxU, alphabet_size, b, t, u, blank_
+                )
+                emit = betas[offset + t * maxU + u + 1] + logp(
+                    denom, acts, maxT, maxU, alphabet_size, b, t, u, labels[u]
+                )
+                betas[offset + t * maxU + u] = rnnt_helper.log_sum_exp(emit, no_emit)
+
+        # sync across all B=b and U=u
+        cuda.syncthreads()
+
+    # After final sync, betas[b, 0, 0] gives
+    # log-likelihood of backward pass.
+    if u == 0:
+        llBackward[b] = betas[offset]
+
+
+@cuda.jit()
+def compute_grad_kernel(
+    grads: torch.Tensor,
+    acts: torch.Tensor,
+    denom: torch.Tensor,
+    alphas: torch.Tensor,
+    betas: torch.Tensor,
+    logll: torch.Tensor,
+    xlen: torch.Tensor,
+    ylen: torch.Tensor,
+    mlabels: torch.Tensor,  # [B, U]
+    minibatch: int,
+    maxT: int,
+    maxU: int,
+    alphabet_size: int,
+    blank_: int,
+):
+    """
+    Compute gradients over the transduction step.
+
+    Args:
+        grads: Zero Tensor of shape [B, T, U, V+1]. Is updated by this kernel to contain the gradients
+            of this batch of samples.
+        acts: Tensor of shape [B, T, U, V+1] flattened. Represents the logprobs activation tensor.
+        denom: Tensor of shape [B, T, U] flattened. Represents the denominator of the logprobs activation tensor
+            across entire vocabulary.
+        alphas: Alpha variable, contains forward probabilities. A tensor of shape [B, T, U].
+        betas: Beta varoable, contains backward probabilities. A tensor of shape [B, T, U].
+        logll: Log-likelihood of the forward variable, represented as a vector of shape [B].
+            Represents the log-likelihood of the forward pass.
+        xlen: Vector of length B which contains the actual acoustic sequence lengths in the padded
+            activation tensor.
+        ylen: Vector of length B which contains the actual target sequence lengths in the padded
+            activation tensor.
+        mlabels: Matrix of shape [B, U+1] (+1 here is due to <SOS> token - usually the RNNT blank).
+            The matrix contains the padded target transcription that must be predicted.
+        minibatch: Int representing the batch size.
+        maxT: The maximum possible acoustic sequence length. Represents T in the logprobs tensor.
+        maxU: The maximum possible target sequence length. Represents U in the logprobs tensor.
+        alphabet_size: The vocabulary dimension V+1 (inclusive of RNNT blank).
+        blank_: Index of the RNNT blank token in the vocabulary. Generally the first or last token in the vocab.
+
+    Updates:
+        Kernel inplace updates the following inputs:
+        -   grads: Gradients with respect to the log likelihood (logll).
+    """
+    # Kernel call:
+    # blocks_per_grid = minibatch (b) * maxT (t) * maxU (u)
+    # threads_per_block = constant buffer size of parallel threads (v :: Constant)
+    tid = cuda.threadIdx.x  # represents v, taking steps of some constant size
+    idx = tid  # index of v < V+1; in steps of constant buffer size
+    col = cuda.blockIdx.x  # represents a fused index of b * t * u
+
+    # Decompose original indices from fused `col`
+    u = col % maxU  # (b * t * u) % u = u
+    bt = (col - u) // maxU  # (b * t * u - u) // U = b * t
+    t = bt % maxT  # (b * t) % t = t
+    mb = (bt - t) // maxT  # (b * t - t) // T = b
+
+    # constants
+    T = xlen[mb]  # select AM length of current sample
+    U = ylen[mb] + 1  # select target length of current sample, +1 for the blank token
+    labels: torch.Tensor = mlabels[mb]  # labels = mlabels + mb * (maxU - 1);
+
+    # Buffered gradient calculations, broadcast across B=b, T=t and U=u, looped over V with some constant stride.
+    # Look up gradient calculation from rnnt_numpy.compute_gradient()
+    if t < T and u < U:
+        # For cuda kernels, maximum number of threads per block is limited to some value.
+        # However, it may be the case that vocabulary size is larger than this limit
+        # To work around this, an arbitrary thread buffer size is chosen such that,
+        # 1) each element within the thread pool operates independently of the other
+        # 2) An inner while loop moves the index of each buffer element by the size of the buffer itself,
+        #    such that all elements of the vocabulary size are covered in (V + 1 // thread_buffer) number of steps.
+        # As such, each thread will perform the while loop at least (V + 1 // thread_buffer) number of times
+        while idx < alphabet_size:
+            # remember, `col` represents the tri-index [b, t, u]
+            # therefore; logpk = denom[b, t, u] + acts[b, t, u, v]
+            logpk = denom[col] + acts[col * alphabet_size + idx]
+            # initialize the grad of the sample acts[b, t, u, v]
+            grad = math.exp(alphas[col] + betas[col] + logpk - logll[mb])
+
+            # // grad to last blank transition
+            # grad[b, T-1, U-1, v=blank] -= exp(alphas[b, t, u) + logpk - logll[b])
+            if (idx == blank_) and (t == T - 1) and (u == U - 1):
+                grad -= math.exp(alphas[col] + logpk - logll[mb])
+
+            # grad of blank across t < T;
+            # grad[b, t<T-1, u, v=blank] -= exp(alphas[b, t, u] + logpk - logll[b] betas[b, t + 1, u])
+            if (idx == blank_) and (t < T - 1):
+                grad -= math.exp(alphas[col] + logpk - logll[mb] + betas[col + maxU])
+
+            # grad of correct token across u < U;
+            # grad[b, t, u<U-1, v=label[u]] -= exp(alphas[b, t, u] + logpk - logll[b] + betas[b, t, u+1])
+            if (u < U - 1) and (idx == labels[u]):
+                grad -= math.exp(alphas[col] + logpk - logll[mb] + betas[col + 1])
+
+            # update grads[b, t, u, v] = grad
+            grads[col * alphabet_size + idx] = grad
+
+            # update internal index through the thread_buffer;
+            # until idx < V + 1, such that entire vocabulary has been updated.
+            idx += GPU_RNNT_THREAD_SIZE

tts/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/reduce.py

@@ -0,0 +1,362 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Copyright 2018-2019, Mingkun Huang
+#
+# 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 enum
+import math
+
+import torch
+from numba import cuda
+
+from nemo.collections.asr.parts.numba.rnnt_loss.utils import global_constants, rnnt_helper
+
+warp_size = global_constants.warp_size()
+dtype = global_constants.dtype()
+
+CTA_REDUCE_SIZE = 128
+
+
+class I_Op(enum.Enum):
+    """
+    Represents an operation that is performed on the input tensor
+    """
+
+    EXPONENTIAL = 0
+    IDENTITY = 1
+
+
+class R_Op(enum.Enum):
+    """
+    Represents a reduction operation performed on the input tensor
+    """
+
+    ADD = 0
+    MAXIMUM = 1
+
+
+@cuda.jit(device=True)
+def CTAReduce(tid: int, x, storage, count: int, R_opid: int):
+    """
+    CUDA Warp reduction kernel.
+
+    It is a device kernel to be called by other kernels.
+
+    The data will be read from the right segement recursively, and reduced (ROP) onto the left half.
+    Operation continues while warp size is larger than a given offset.
+    Beyond this offset, warp reduction is performed via `shfl_down_sync`, which halves the reduction
+    space and sums the two halves at each call.
+
+    Note:
+        Efficient warp occurs at input shapes of 2 ^ K.
+
+    References:
+        - Warp Primitives [https://developer.nvidia.com/blog/using-cuda-warp-level-primitives/]
+
+    Args:
+        tid: CUDA thread index
+        x: activation. Single float.
+        storage: shared memory of size CTA_REDUCE_SIZE used for reduction in parallel threads.
+        count: equivalent to num_rows, which is equivalent to alphabet_size (V+1)
+        R_opid: Operator ID for reduction. See R_Op for more information.
+    """
+    storage[tid] = x
+
+    cuda.syncthreads()
+
+    # Fold the data in half with each pass
+    offset = CTA_REDUCE_SIZE // 2
+    while offset >= warp_size:
+        if (tid + offset) < count and tid < offset:
+            # Read from the right half and store to the left half.
+            if R_opid == 0:
+                x = rnnt_helper.add(x, storage[offset + tid])
+            else:
+                x = rnnt_helper.maximum(x, storage[offset + tid])
+
+            storage[tid] = x
+
+        cuda.syncthreads()
+        offset = offset // 2
+
+    offset = warp_size // 2
+    while offset > 0:
+        # warp reduction and sync
+        shuff = cuda.shfl_down_sync(0xFFFFFFFF, x, offset)
+
+        if (tid + offset < count) and (tid < offset):
+            if R_opid == 0:
+                x = rnnt_helper.add(x, shuff)
+            else:
+                x = rnnt_helper.maximum(x, shuff)
+
+        offset = offset // 2
+
+    return x
+
+
+@cuda.jit()
+def _reduce_rows(I_opid: int, R_opid: int, acts, output, num_rows: int):
+    """
+    CUDA Warp reduction kernel which reduces via the R_Op.Maximum
+
+    Reduces the input data such that I_Op = Identity and R_op = Maximum.
+    The result is stored in the blockIdx, and is stored as an identity op.
+
+    Note:
+        Efficient warp occurs at input shapes of 2 ^ K.
+
+    References:
+        - Warp Primitives [https://developer.nvidia.com/blog/using-cuda-warp-level-primitives/]
+
+    Args:
+        I_opid: Operator ID for input. See I_Op for more information. For this kernel,
+            the Identity op is chosen in general, and therefore the input is reduced in place
+            without scaling.
+        R_opid: Operator ID for reduction. See R_Op for more information.
+            For this kernel, generally Maximum op is chosen. It reduces the kernel via max.
+        acts: Flatened activation matrix of shape [B * T * U * (V+1)].
+        output: Flatened output matrix of shape [B * T * U * (V+1)]. Data will be overwritten.
+        num_rows: Vocabulary size (including blank token) - V+1.
+    """
+    tid = cuda.threadIdx.x
+    idx = tid
+    col = cuda.blockIdx.x
+
+    # allocate shared thread memory
+    storage = cuda.shared.array(shape=(CTA_REDUCE_SIZE,), dtype=acts.dtype)
+
+    max = output[col]
+
+    # // Each block works on a column
+    if idx < num_rows:
+        curr = acts[col * num_rows + idx] - max
+        if I_opid == 0:
+            curr = rnnt_helper.exponential(curr)
+        else:
+            curr = rnnt_helper.identity(curr)
+
+    idx += CTA_REDUCE_SIZE
+
+    while idx < num_rows:
+        activation_ = acts[col * num_rows + idx] - max
+
+        if I_opid == 0 and R_opid == 0:
+            curr = rnnt_helper.add(curr, rnnt_helper.exponential(activation_))
+        elif I_opid == 0 and R_opid == 1:
+            curr = rnnt_helper.maximum(curr, rnnt_helper.exponential(activation_))
+        elif I_opid == 1 and R_opid == 0:
+            curr = rnnt_helper.add(curr, rnnt_helper.identity(activation_))
+        else:
+            curr = rnnt_helper.maximum(curr, rnnt_helper.identity(activation_))
+
+        idx += CTA_REDUCE_SIZE
+
+    # // Sum thread-totals over the CTA.
+    curr = CTAReduce(tid, curr, storage, num_rows, R_opid)
+
+    # // Store result in out (inplace, I_op: identity)
+    if tid == 0:
+        output[col] = curr
+
+
+@cuda.jit()
+def _reduce_minus(I_opid: int, R_opid: int, acts, output, num_rows: int):
+    """
+    CUDA Warp reduction kernel which reduces via the R_Op.Add
+
+    Reduces the input data such that I_Op = Exponential and R_op = Add.
+    The result is stored in the blockIdx, and is stored as an exp op.
+
+    Note:
+        Efficient warp occurs at input shapes of 2 ^ K.
+
+    References:
+        - Warp Primitives [https://developer.nvidia.com/blog/using-cuda-warp-level-primitives/]
+
+    Args:
+        I_opid: Operator ID for input. See I_Op for more information. For this kernel,
+            the Exponential op is chosen in general, and therefore the input is reduced in place
+            with scaling.
+        R_opid: Operator ID for reduction. See R_Op for more information.
+            For this kernel, generally Add op is chosen. It reduces the kernel via summation.
+        acts: Flatened activation matrix of shape [B * T * U * (V+1)].
+        output: Flatened output matrix of shape [B * T * U * (V+1)]. Data will be overwritten.
+        num_rows: Vocabulary size (including blank token) - V+1.
+    """
+    tid = cuda.threadIdx.x
+    idx = tid
+    col = cuda.blockIdx.x
+
+    # allocate shared thread memory
+    storage = cuda.shared.array(shape=(CTA_REDUCE_SIZE,), dtype=acts.dtype)
+
+    max = output[col]
+
+    # // Each block works on a column
+    if idx < num_rows:
+        curr = acts[col * num_rows + idx] - max
+        if I_opid == 0:
+            curr = rnnt_helper.exponential(curr)
+        else:
+            curr = rnnt_helper.identity(curr)
+
+    idx += CTA_REDUCE_SIZE
+
+    while idx < num_rows:
+        activation_ = acts[col * num_rows + idx] - max
+
+        if I_opid == 0 and R_opid == 0:
+            curr = rnnt_helper.add(curr, rnnt_helper.exponential(activation_))
+        elif I_opid == 0 and R_opid == 1:
+            curr = rnnt_helper.maximum(curr, rnnt_helper.exponential(activation_))
+        elif I_opid == 1 and R_opid == 0:
+            curr = rnnt_helper.add(curr, rnnt_helper.identity(activation_))
+        else:
+            curr = rnnt_helper.maximum(curr, rnnt_helper.identity(activation_))
+
+        idx += CTA_REDUCE_SIZE
+
+    # // Sum thread-totals over the CTA.
+    curr = CTAReduce(tid, curr, storage, num_rows, R_opid)
+
+    # // Store result in out (inplace, I_op: exponential)
+    if tid == 0:
+        output[col] = -max - math.log(curr)
+
+
+def ReduceHelper(
+    I_opid: int,
+    R_opid: int,
+    acts: torch.Tensor,
+    output: torch.Tensor,
+    num_rows: int,
+    num_cols: int,
+    minus: bool,
+    stream,
+):
+    """
+    CUDA Warp reduction kernel helper which reduces via the R_Op.Add and writes
+    the result to `output` according to I_op id.
+
+    The result is stored in the blockIdx.
+
+    Note:
+        Efficient warp occurs at input shapes of 2 ^ K.
+
+    References:
+        - Warp Primitives [https://developer.nvidia.com/blog/using-cuda-warp-level-primitives/]
+
+    Args:
+        I_opid: Operator ID for input. See I_Op for more information.
+        R_opid: Operator ID for reduction. See R_Op for more information.
+        acts: Flatened activation matrix of shape [B * T * U * (V+1)].
+        output: Flatened output matrix of shape [B * T * U * (V+1)]. Data will be overwritten.
+        num_rows: Vocabulary size (including blank token) - V+1.
+            Represents the number of threads per block.
+        num_cols: Flattened shape of activation matrix, without vocabulary dimension (B * T * U).
+            Represents number of blocks per grid.
+        minus: Bool flag whether to add or subtract as reduction.
+            If minus is set; calls _reduce_minus, else calls _reduce_rows kernel.
+        stream: CUDA Stream.
+    """
+    if minus:
+        grid_size = num_cols
+        # call kernel
+        _reduce_minus[grid_size, CTA_REDUCE_SIZE, stream, 0](I_opid, R_opid, acts, output, num_rows)
+
+    else:
+        grid_size = num_cols
+        # call kernel
+        _reduce_rows[grid_size, CTA_REDUCE_SIZE, stream, 0](I_opid, R_opid, acts, output, num_rows)
+
+    return True
+
+
+def reduce_exp(acts: torch.Tensor, denom, rows: int, cols: int, minus: bool, stream):
+    """
+    Helper method to call the Warp Reduction Kernel to perform `exp` reduction.
+
+    Note:
+        Efficient warp occurs at input shapes of 2 ^ K.
+
+    References:
+        - Warp Primitives [https://developer.nvidia.com/blog/using-cuda-warp-level-primitives/]
+
+    Args:
+        acts: Flatened activation matrix of shape [B * T * U * (V+1)].
+        output: Flatened output matrix of shape [B * T * U * (V+1)]. Data will be overwritten.
+        rows: Vocabulary size (including blank token) - V+1.
+            Represents the number of threads per block.
+        cols: Flattened shape of activation matrix, without vocabulary dimension (B * T * U).
+            Represents number of blocks per grid.
+        minus: Bool flag whether to add or subtract as reduction.
+            If minus is set; calls _reduce_minus, else calls _reduce_rows kernel.
+        stream: CUDA Stream.
+    """
+    return ReduceHelper(
+        I_opid=I_Op.EXPONENTIAL.value,
+        R_opid=R_Op.ADD.value,
+        acts=acts,
+        output=denom,
+        num_rows=rows,
+        num_cols=cols,
+        minus=minus,
+        stream=stream,
+    )
+
+
+def reduce_max(acts: torch.Tensor, denom, rows: int, cols: int, minus: bool, stream):
+    """
+    Helper method to call the Warp Reduction Kernel to perform `max` reduction.
+
+    Note:
+        Efficient warp occurs at input shapes of 2 ^ K.
+
+    References:
+        - Warp Primitives [https://developer.nvidia.com/blog/using-cuda-warp-level-primitives/]
+
+    Args:
+        acts: Flatened activation matrix of shape [B * T * U * (V+1)].
+        output: Flatened output matrix of shape [B * T * U * (V+1)]. Data will be overwritten.
+        rows: Vocabulary size (including blank token) - V+1.
+            Represents the number of threads per block.
+        cols: Flattened shape of activation matrix, without vocabulary dimension (B * T * U).
+            Represents number of blocks per grid.
+        minus: Bool flag whether to add or subtract as reduction.
+            If minus is set; calls _reduce_minus, else calls _reduce_rows kernel.
+        stream: CUDA Stream.
+    """
+    return ReduceHelper(
+        I_opid=I_Op.IDENTITY.value,
+        R_opid=R_Op.MAXIMUM.value,
+        acts=acts,
+        output=denom,
+        num_rows=rows,
+        num_cols=cols,
+        minus=minus,
+        stream=stream,
+    )

tts/nemo/collections/asr/parts/numba/rnnt_loss/utils/global_constants.py

@@ -0,0 +1,68 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Copyright 2018-2019, Mingkun Huang
+#
+# 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 enum
+
+import numpy as np
+from numba import float32
+
+# Internal globals
+_THREADS_PER_BLOCK = 32
+_WARP_SIZE = 32
+_DTYPE = float32
+
+# Constants
+FP32_INF = np.inf
+FP32_NEG_INF = -np.inf
+THRESHOLD = 1e-1
+
+"""
+Getters
+"""
+
+
+def threads_per_block():
+    global _THREADS_PER_BLOCK
+    return _THREADS_PER_BLOCK
+
+
+def warp_size():
+    global _WARP_SIZE
+    return _WARP_SIZE
+
+
+def dtype():
+    global _DTYPE
+    return _DTYPE
+
+
+# RNNT STATUS
+class RNNTStatus(enum.Enum):
+    RNNT_STATUS_SUCCESS = 0
+    RNNT_STATUS_INVALID_VALUE = 1

tts/nemo/collections/asr/parts/numba/rnnt_loss/utils/rnnt_helper.py

@@ -0,0 +1,129 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Copyright 2018-2019, Mingkun Huang
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import math
+from typing import Optional
+
+import torch
+from numba import cuda
+
+from nemo.collections.asr.parts.numba.rnnt_loss.utils import global_constants
+
+threshold = global_constants.THRESHOLD
+
+
+@cuda.jit(device=True, inline=True)
+def log_sum_exp(a: float, b: float):
+    if a == global_constants.FP32_NEG_INF:
+        return b
+
+    if b == global_constants.FP32_NEG_INF:
+        return a
+
+    if a > b:
+        return math.log1p(math.exp(b - a)) + a
+    else:
+        return math.log1p(math.exp(a - b)) + b
+
+
+@cuda.jit(device=True, inline=True)
+def div_up(x: int, y: int):
+    return (x + y - 1) // y
+
+
+@cuda.jit(device=True)
+def maximum(x, y):
+    if x < y:
+        return y
+    else:
+        return x
+
+
+@cuda.jit(device=True)
+def add(x, y):
+    return x + y
+
+
+@cuda.jit(device=True)
+def identity(x):
+    return x
+
+
+@cuda.jit(device=True)
+def negate(x):
+    return -x
+
+
+@cuda.jit(device=True)
+def exponential(x):
+    return math.exp(x)
+
+
+@cuda.jit(device=True)
+def log_plus(p1: float, p2: float):
+    if p1 == global_constants.FP32_NEG_INF:
+        return p2
+
+    if p2 == global_constants.FP32_NEG_INF:
+        return p1
+
+    result = math.log1p(math.exp(-math.fabs(p1 - p2))) + maximum(p1, p2)
+    return result
+
+
+def get_workspace_size(
+    maxT: int, maxU: int, minibatch: int, gpu: bool
+) -> (Optional[int], global_constants.RNNTStatus):
+
+    if minibatch <= 0 or maxT <= 0 or maxU <= 0:
+        return (None, global_constants.RNNTStatus.RNNT_STATUS_INVALID_VALUE)
+
+    # per minibatch memory
+    per_minibatch_size = 0
+
+    # alphas & betas
+    per_minibatch_size += maxT * maxU * 2
+
+    if not gpu:
+        # // blank & label log probability cache
+        per_minibatch_size += maxT * maxU * 2
+    else:
+        # // softmax denominator
+        per_minibatch_size += maxT * maxU
+        # // forward - backward loglikelihood
+        per_minibatch_size += 2
+
+    size = per_minibatch_size * minibatch
+    return (size, global_constants.RNNTStatus.RNNT_STATUS_SUCCESS)
+
+
+def flatten_tensor(x: torch.Tensor):
+    original_shape = x.shape
+    x = x.view([-1])
+    return x, original_shape

tts/nemo/collections/asr/parts/numba/spec_augment/__init__.py

@@ -0,0 +1,18 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nemo.collections.asr.parts.numba.spec_augment.spec_aug_numba import (
+    SpecAugmentNumba,
+    spec_augment_launch_heuristics,
+)

tts/nemo/collections/asr/parts/numba/spec_augment/spec_aug_numba.py

@@ -0,0 +1,290 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import torch.nn as nn
+from numba import cuda
+
+from nemo.core.classes import Typing, typecheck
+from nemo.core.neural_types import LengthsType, NeuralType, SpectrogramType
+from nemo.utils import logging
+
+MAX_THREAD_BUFFER = 512
+
+
+@cuda.jit()
+def spec_augment_kernel(
+    x: torch.Tensor,
+    x_len: torch.Tensor,
+    freq_starts: torch.Tensor,
+    freq_widths: torch.Tensor,
+    time_starts: torch.Tensor,
+    time_widths: torch.Tensor,
+    mask_value: float,
+):
+    """
+    Numba CUDA kernel to perform SpecAugment in-place on the GPU.
+    Parallelize over freq and time axis, parallel threads over batch.
+    Sequential over masks (adaptive in time).
+
+    Args:
+        x: Pytorch tensor of shape [B, F, T] with the acoustic features.
+        x_len: Pytorch tensor of shape [B] with the lengths of the padded sequence.
+        freq_starts: Pytorch tensor of shape [B, M_f] with the start indices of freq masks.
+        freq_widths: Pytorch tensor of shape [B, M_f] with the width of freq masks.
+        time_starts: Pytorch tensor of shape [B, M_t] with the start indices of time masks.
+        time_widths: Pytorch tensor of shape [B, M_t] with the width of time masks.
+        mask_value: Float value that will be used as mask value.
+    """
+    f = cuda.blockIdx.x  # indexes the Freq dim
+    t = cuda.blockIdx.y  # indexes the Time dim
+    tid = cuda.threadIdx.x  # index of the current mask
+    threads_per_block = cuda.blockDim.x
+
+    # Compute the number of masks over freq axis
+    len_f = freq_starts.shape[1]
+    # For all samples in the batch, apply the freq mask
+    for bidx in range(0, x.shape[0], threads_per_block):
+        # Resolve the index of the batch (case where more masks than MAX_THREAD_BUFFER)
+        bm_idx = bidx * threads_per_block + tid
+
+        # For `len_f` number of freq masks that must be applied
+        for fidx in range(0, len_f):
+            # If resolved freq mask index < total number of freq masks
+            if fidx < len_f:
+                # Access the start index and width of this freq mask
+                f_start = freq_starts[bm_idx, fidx]
+                f_width = freq_widths[bm_idx, fidx]
+
+                # If block idx `f` >= start and < (start + width) of this freq mask
+                if f >= f_start and f < (f_start + f_width):
+                    x[bm_idx, f, t] = mask_value
+
+    # Compute the number of masks over time axis
+    len_t = time_starts.shape[1]
+    # For all samples in the batch, apply the time mask
+    for b_idx in range(0, x.shape[0], threads_per_block):
+        # Resolve the index of the batch (case where more masks than MAX_THREAD_BUFFER)
+        bm_idx = bidx * threads_per_block + tid
+
+        # For `len_t` number of freq masks that must be applied
+        for tidx in range(0, len_t):
+            # If resolved time mask index < total number of time masks
+            if tidx < len_t:
+                # Access the start index and width of this time mask
+                t_start = time_starts[bm_idx, tidx]
+                t_width = time_widths[bm_idx, tidx]
+
+                # If block idx `t` >= start and < (start + width) of this time mask
+                if t >= t_start and t < (t_start + t_width):
+                    # Current block idx `t` < current seq length x_len[b]
+                    # This ensure that we mask only upto the length of that sample
+                    # Everything after that index is padded value so unnecessary to mask
+                    if t < x_len[bm_idx]:
+                        x[bm_idx, f, t] = mask_value
+
+
+def spec_augment_launch_heuristics(x: torch.Tensor, length: torch.Tensor):
+    """
+    Heuristics to determins whether pytorch implementation or numba implementation is selected.
+    Assumes numba cuda is supported.
+
+    Args:
+        x: Torch tensor of shape [B, F, T]
+        length: Optional, Torch of tensor of shape [B] - containing lengths of the tensor.
+
+    Returns:
+        True if numba kernel should be selected, else False
+    """
+    if not x.is_cuda:
+        return False
+
+    if length is None:
+        return False
+
+    if x.shape[0] < 8:
+        return False
+
+    return True
+
+
+def launch_spec_augment_kernel(
+    x: torch.Tensor,
+    x_len: torch.Tensor,
+    freq_starts: torch.Tensor,
+    freq_lengths: torch.Tensor,
+    time_starts: torch.Tensor,
+    time_lengths: torch.Tensor,
+    freq_masks: int,
+    time_masks: int,
+    mask_value: float,
+):
+    """
+    Helper method to launch the SpecAugment kernel
+
+    Args:
+        x: Pytorch tensor of shape [B, F, T] with the acoustic features.
+        x_len: Pytorch tensor of shape [B] with the lengths of the padded sequence.
+        freq_starts: Pytorch tensor of shape [B, M_f] with the start indices of freq masks.
+        freq_widths: Pytorch tensor of shape [B, M_f] with the width of freq masks.
+        time_starts: Pytorch tensor of shape [B, M_t] with the start indices of time masks.
+        time_widths: Pytorch tensor of shape [B, M_t] with the width of time masks.
+        freq_masks: Int value that determines the number of time masks.
+        time_masks: Int value that determines the number of freq masks.
+        mask_value: Float value that will be used as mask value.
+
+    Returns:
+        The spec augmented tensor 'x'
+    """
+    # Setup CUDA stream
+    sh = x.shape
+    stream = cuda.external_stream(torch.cuda.current_stream(x.device).cuda_stream)
+
+    if time_masks > 0 or freq_masks > 0:
+        # Parallelize over freq and time axis, parallel threads over batch
+        # Sequential over masks (adaptive in time).
+        blocks_per_grid = [sh[1], sh[2]]
+        # threads_per_block = min(MAX_THREAD_BUFFER, max(freq_masks, time_masks))
+        threads_per_block = min(MAX_THREAD_BUFFER, x.shape[0])
+
+        # Numba does not support fp16, force cast to fp32 temporarily at the expense of memory
+        original_dtype = x.dtype
+        cast_x = False
+        if x.dtype == torch.float16:
+            x = x.float()
+            cast_x = True
+
+        # Launch CUDA kernel
+        spec_augment_kernel[blocks_per_grid, threads_per_block, stream, 0](
+            x, x_len, freq_starts, freq_lengths, time_starts, time_lengths, mask_value
+        )
+        torch.cuda.synchronize()
+
+        # Recast back to original dtype if earlier cast was performed
+        if cast_x:
+            x = x.to(dtype=original_dtype)
+
+    return x
+
+
+class SpecAugmentNumba(nn.Module, Typing):
+    """
+    Zeroes out(cuts) random continuous horisontal or
+    vertical segments of the spectrogram as described in
+    SpecAugment (https://arxiv.org/abs/1904.08779).
+
+    Utilizes a Numba CUDA kernel to perform inplace edit of the input without loops.
+    Parallelize over freq and time axis, parallel threads over batch.
+    Sequential over masks (adaptive in time).
+
+    Args:
+        freq_masks - how many frequency segments should be cut
+        time_masks - how many time segments should be cut
+        freq_width - maximum number of frequencies to be cut in one segment
+        time_width - maximum number of time steps to be cut in one segment.
+            Can be a positive integer or a float value in the range [0, 1].
+            If positive integer value, defines maximum number of time steps
+            to be cut in one segment.
+            If a float value, defines maximum percentage of timesteps that
+            are cut adaptively.
+        rng: Ignored.
+    """
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input types
+        """
+        return {
+            "input_spec": NeuralType(('B', 'D', 'T'), SpectrogramType()),
+            "length": NeuralType(tuple('B'), LengthsType()),
+        }
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output types
+        """
+        return {"augmented_spec": NeuralType(('B', 'D', 'T'), SpectrogramType())}
+
+    def __init__(
+        self, freq_masks=0, time_masks=0, freq_width=10, time_width=0.1, rng=None, mask_value=0.0,
+    ):
+        super().__init__()
+        # Message to mention that numba specaugment kernel will be available
+        # if input device is CUDA and lengths are provided
+        logging.debug("Numba SpecAugment kernel is available")
+
+        self.freq_masks = freq_masks
+        self.time_masks = time_masks
+
+        self.freq_width = freq_width
+        self.time_width = time_width
+
+        self.mask_value = mask_value
+
+        # Unused
+        self.rng = rng
+        if self.rng is not None:
+            logging.warning("`rng` was supplied to SpecAugmentNumba, but it is not used.")
+
+        if isinstance(time_width, int):
+            self.adaptive_temporal_width = False
+        else:
+            if time_width > 1.0 or time_width < 0.0:
+                raise ValueError('If `time_width` is a float value, must be in range [0, 1]')
+
+            self.adaptive_temporal_width = True
+
+    @typecheck()
+    @torch.no_grad()
+    def forward(self, input_spec, length):
+        sh = input_spec.shape
+        bs = sh[0]
+
+        if self.adaptive_temporal_width:
+            time_width = max(1, int(sh[2] * self.time_width))
+        else:
+            time_width = self.time_width
+
+        # Construct the freq and time masks as well as start positions
+        if self.freq_masks > 0:
+            freq_starts = torch.randint(
+                0, sh[1] - self.freq_width + 1, size=[bs, self.freq_masks], device=input_spec.device
+            )
+            freq_lengths = torch.randint(0, self.freq_width + 1, size=[bs, self.freq_masks], device=input_spec.device)
+        else:
+            freq_starts = torch.zeros([bs, 1], dtype=torch.int64, device=input_spec.device)
+            freq_lengths = torch.zeros([bs, 1], dtype=torch.int64, device=input_spec.device)
+
+        if self.time_masks > 0:
+            time_starts = torch.randint(
+                0, sh[2] - time_width + 1, size=[bs, self.time_masks], device=input_spec.device
+            )
+            time_lengths = torch.randint(0, time_width + 1, size=[bs, self.time_masks], device=input_spec.device)
+        else:
+            time_starts = torch.zeros([bs, 1], dtype=torch.int64, device=input_spec.device)
+            time_lengths = torch.zeros([bs, 1], dtype=torch.int64, device=input_spec.device)
+
+        x = launch_spec_augment_kernel(
+            input_spec,
+            length,
+            freq_starts=freq_starts,
+            freq_lengths=freq_lengths,
+            time_starts=time_starts,
+            time_lengths=time_lengths,
+            freq_masks=self.freq_masks,
+            time_masks=self.time_masks,
+            mask_value=self.mask_value,
+        )
+
+        return x

tts/nemo/collections/asr/parts/preprocessing/__init__.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nemo.collections.asr.parts.preprocessing.feature_loader import ExternalFeatureLoader
+from nemo.collections.asr.parts.preprocessing.features import (
+    STFT,
+    FeaturizerFactory,
+    FilterbankFeatures,
+    STFTExactPad,
+    STFTPatch,
+    WaveformFeaturizer,
+)
+from nemo.collections.asr.parts.preprocessing.perturb import (
+    AudioAugmentor,
+    AugmentationDataset,
+    GainPerturbation,
+    ImpulsePerturbation,
+    NoisePerturbation,
+    Perturbation,
+    RirAndNoisePerturbation,
+    ShiftPerturbation,
+    SpeedPerturbation,
+    TimeStretchPerturbation,
+    TranscodePerturbation,
+    WhiteNoisePerturbation,
+    perturbation_types,
+    process_augmentations,
+    register_perturbation,
+)
+from nemo.collections.asr.parts.preprocessing.segment import AudioSegment

tts/nemo/collections/asr/parts/preprocessing/feature_loader.py

@@ -0,0 +1,69 @@
+# 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 numpy as np
+import torch
+
+
+class ExternalFeatureLoader(object):
+    """Feature loader that load external features store in certain format. 
+    Currently support pickle, npy and npz format.
+    """
+
+    def __init__(self, file_path, sample_rate=16000, augmentor=None):
+        """Load feature samples from file_path.
+        file_path (str) is the path of the file that stores feature/sample.
+        """
+        self._file_path = file_path
+        self.augmentor = augmentor
+
+    def load_feature_from_file(self, file_path):
+        """Load samples from file_path and convert it to be of type float32
+        file_path (str) is the path of the file that stores feature/sample.
+        """
+        # load pickle/npy/npz file
+        try:
+            samples = np.load(file_path, allow_pickle=True)
+            return self._convert_samples_to_float32(samples)
+        except:
+            raise TypeError(
+                "Error open feature files. Probably not supported the format yet. Support pkl, npz, and npy format now. "
+            )
+        # TODO load other type of files such as kaldi io ark
+
+    @staticmethod
+    def _convert_samples_to_float32(samples):
+        """Convert sample type to float32.
+        Integers will be scaled to [-1, 1] in float32.
+        """
+        float32_samples = samples.astype('float32')
+        if samples.dtype in np.sctypes['int']:
+            bits = np.iinfo(samples.dtype).bits
+            float32_samples *= 1.0 / 2 ** (bits - 1)
+        elif samples.dtype in np.sctypes['float']:
+            pass
+        else:
+            raise TypeError("Unsupported sample type: %s." % samples.dtype)
+        return float32_samples
+
+    def process(self, file_path, offset=0, duration=0, trim=False, orig_sr=None):
+        feature = self.load_feature_from_file(file_path)
+        return self.process_segment(feature)
+
+    def process_segment(self, feature_segment):
+        if self.augmentor:
+            # augmentor for external features. Here possible augmentor for external embedding feature is Diaconis Augmentation and might be implemented later
+            self.augmentor.perturb(feature_segment)
+            return torch.tensor(feature_segment, dtype=torch.float)
+
+        return torch.tensor(feature_segment, dtype=torch.float)

tts/nemo/collections/asr/parts/preprocessing/features.py

@@ -0,0 +1,432 @@
+# 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) 2018 Ryan Leary
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+# This file contains code artifacts adapted from https://github.com/ryanleary/patter
+import math
+
+import librosa
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from librosa.util import tiny
+from torch.autograd import Variable
+from torch_stft import STFT
+
+from nemo.collections.asr.parts.preprocessing.perturb import AudioAugmentor
+from nemo.collections.asr.parts.preprocessing.segment import AudioSegment
+from nemo.collections.common.parts.patch_utils import stft_patch
+from nemo.utils import logging
+
+CONSTANT = 1e-5
+
+
+def normalize_batch(x, seq_len, normalize_type):
+    if normalize_type == "per_feature":
+        x_mean = torch.zeros((seq_len.shape[0], x.shape[1]), dtype=x.dtype, device=x.device)
+        x_std = torch.zeros((seq_len.shape[0], x.shape[1]), dtype=x.dtype, device=x.device)
+        for i in range(x.shape[0]):
+            if x[i, :, : seq_len[i]].shape[1] == 1:
+                raise ValueError(
+                    "normalize_batch with `per_feature` normalize_type received a tensor of length 1. This will result "
+                    "in torch.std() returning nan"
+                )
+            x_mean[i, :] = x[i, :, : seq_len[i]].mean(dim=1)
+            x_std[i, :] = x[i, :, : seq_len[i]].std(dim=1)
+        # make sure x_std is not zero
+        x_std += CONSTANT
+        return (x - x_mean.unsqueeze(2)) / x_std.unsqueeze(2)
+    elif normalize_type == "all_features":
+        x_mean = torch.zeros(seq_len.shape, dtype=x.dtype, device=x.device)
+        x_std = torch.zeros(seq_len.shape, dtype=x.dtype, device=x.device)
+        for i in range(x.shape[0]):
+            x_mean[i] = x[i, :, : seq_len[i].item()].mean()
+            x_std[i] = x[i, :, : seq_len[i].item()].std()
+        # make sure x_std is not zero
+        x_std += CONSTANT
+        return (x - x_mean.view(-1, 1, 1)) / x_std.view(-1, 1, 1)
+    elif "fixed_mean" in normalize_type and "fixed_std" in normalize_type:
+        x_mean = torch.tensor(normalize_type["fixed_mean"], device=x.device)
+        x_std = torch.tensor(normalize_type["fixed_std"], device=x.device)
+        return (x - x_mean.view(x.shape[0], x.shape[1]).unsqueeze(2)) / x_std.view(x.shape[0], x.shape[1]).unsqueeze(2)
+    else:
+        return x
+
+
+def splice_frames(x, frame_splicing):
+    """ Stacks frames together across feature dim
+
+    input is batch_size, feature_dim, num_frames
+    output is batch_size, feature_dim*frame_splicing, num_frames
+
+    """
+    seq = [x]
+    for n in range(1, frame_splicing):
+        seq.append(torch.cat([x[:, :, :n], x[:, :, n:]], dim=2))
+    return torch.cat(seq, dim=1)
+
+
+class WaveformFeaturizer(object):
+    def __init__(self, sample_rate=16000, int_values=False, augmentor=None):
+        self.augmentor = augmentor if augmentor is not None else AudioAugmentor()
+        self.sample_rate = sample_rate
+        self.int_values = int_values
+
+    def max_augmentation_length(self, length):
+        return self.augmentor.max_augmentation_length(length)
+
+    def process(self, file_path, offset=0, duration=0, trim=False, orig_sr=None):
+        audio = AudioSegment.from_file(
+            file_path,
+            target_sr=self.sample_rate,
+            int_values=self.int_values,
+            offset=offset,
+            duration=duration,
+            trim=trim,
+            orig_sr=orig_sr,
+        )
+        return self.process_segment(audio)
+
+    def process_segment(self, audio_segment):
+        self.augmentor.perturb(audio_segment)
+        return torch.tensor(audio_segment.samples, dtype=torch.float)
+
+    @classmethod
+    def from_config(cls, input_config, perturbation_configs=None):
+        if perturbation_configs is not None:
+            aa = AudioAugmentor.from_config(perturbation_configs)
+        else:
+            aa = None
+
+        sample_rate = input_config.get("sample_rate", 16000)
+        int_values = input_config.get("int_values", False)
+
+        return cls(sample_rate=sample_rate, int_values=int_values, augmentor=aa)
+
+
+class FeaturizerFactory(object):
+    def __init__(self):
+        pass
+
+    @classmethod
+    def from_config(cls, input_cfg, perturbation_configs=None):
+        return WaveformFeaturizer.from_config(input_cfg, perturbation_configs=perturbation_configs)
+
+
+# Create helper class to patch forward func for use with AMP
+class STFTPatch(STFT):
+    def forward(self, input_data):
+        return super().transform(input_data)[0]
+
+
+# Create helper class for STFT that yields num_frames = num_samples // hop_length
+class STFTExactPad(STFTPatch):
+    """adapted from Prem Seetharaman's https://github.com/pseeth/pytorch-stft"""
+
+    def __init__(self, *params, **kw_params):
+        super().__init__(*params, **kw_params)
+        self.pad_amount = (self.filter_length - self.hop_length) // 2
+
+    def inverse(self, magnitude, phase):
+        recombine_magnitude_phase = torch.cat([magnitude * torch.cos(phase), magnitude * torch.sin(phase)], dim=1)
+
+        inverse_transform = F.conv_transpose1d(
+            recombine_magnitude_phase,
+            Variable(self.inverse_basis, requires_grad=False),
+            stride=self.hop_length,
+            padding=0,
+        )
+
+        if self.window is not None:
+            window_sum = librosa.filters.window_sumsquare(
+                self.window,
+                magnitude.size(-1),
+                hop_length=self.hop_length,
+                win_length=self.win_length,
+                n_fft=self.filter_length,
+                dtype=np.float32,
+            )
+            # remove modulation effects
+            approx_nonzero_indices = torch.from_numpy(np.where(window_sum > tiny(window_sum))[0])
+            window_sum = torch.autograd.Variable(torch.from_numpy(window_sum), requires_grad=False).to(
+                magnitude.device
+            )
+            inverse_transform[..., approx_nonzero_indices] /= window_sum[approx_nonzero_indices]
+
+            # scale by hop ratio
+            inverse_transform *= self.filter_length / self.hop_length
+
+        inverse_transform = inverse_transform[..., self.pad_amount :]
+        inverse_transform = inverse_transform[..., : -self.pad_amount :]
+        inverse_transform = inverse_transform.squeeze(1)
+
+        return inverse_transform
+
+
+class FilterbankFeatures(nn.Module):
+    """Featurizer that converts wavs to Mel Spectrograms.
+    See AudioToMelSpectrogramPreprocessor for args.
+    """
+
+    def __init__(
+        self,
+        sample_rate=16000,
+        n_window_size=320,
+        n_window_stride=160,
+        window="hann",
+        normalize="per_feature",
+        n_fft=None,
+        preemph=0.97,
+        nfilt=64,
+        lowfreq=0,
+        highfreq=None,
+        log=True,
+        log_zero_guard_type="add",
+        log_zero_guard_value=2 ** -24,
+        dither=CONSTANT,
+        pad_to=16,
+        max_duration=16.7,
+        frame_splicing=1,
+        exact_pad=False,
+        stft_exact_pad=False,  # TODO: Remove this in 1.1.0
+        stft_conv=False,  # TODO: Remove this in 1.1.0
+        pad_value=0,
+        mag_power=2.0,
+        use_grads=False,
+    ):
+        super().__init__()
+        if stft_conv or stft_exact_pad:
+            logging.warning(
+                "Using torch_stft is deprecated and will be removed in 1.1.0. Please set stft_conv and stft_exact_pad "
+                "to False for FilterbankFeatures and AudioToMelSpectrogramPreprocessor. Please set exact_pad to True "
+                "as needed."
+            )
+        if (exact_pad or stft_exact_pad) and n_window_stride % 2 == 1:
+            raise NotImplementedError(
+                f"{self} received exact_pad == True, but hop_size was odd. If audio_length % hop_size == 0. Then the "
+                "returned spectrogram would not be of length audio_length // hop_size. Please use an even hop_size."
+            )
+        self.log_zero_guard_value = log_zero_guard_value
+        if (
+            n_window_size is None
+            or n_window_stride is None
+            or not isinstance(n_window_size, int)
+            or not isinstance(n_window_stride, int)
+            or n_window_size <= 0
+            or n_window_stride <= 0
+        ):
+            raise ValueError(
+                f"{self} got an invalid value for either n_window_size or "
+                f"n_window_stride. Both must be positive ints."
+            )
+        logging.info(f"PADDING: {pad_to}")
+
+        self.win_length = n_window_size
+        self.hop_length = n_window_stride
+        self.n_fft = n_fft or 2 ** math.ceil(math.log2(self.win_length))
+        self.stft_pad_amount = (self.n_fft - self.hop_length) // 2 if exact_pad else None
+        self.stft_exact_pad = stft_exact_pad
+        self.stft_conv = stft_conv
+
+        if stft_conv:
+            logging.info("STFT using conv")
+            if stft_exact_pad:
+                logging.info("STFT using exact pad")
+                self.stft = STFTExactPad(self.n_fft, self.hop_length, self.win_length, window)
+            else:
+                self.stft = STFTPatch(self.n_fft, self.hop_length, self.win_length, window)
+        else:
+            logging.info("STFT using torch")
+            if exact_pad:
+                logging.info("STFT using exact pad")
+            torch_windows = {
+                'hann': torch.hann_window,
+                'hamming': torch.hamming_window,
+                'blackman': torch.blackman_window,
+                'bartlett': torch.bartlett_window,
+                'none': None,
+            }
+            window_fn = torch_windows.get(window, None)
+            window_tensor = window_fn(self.win_length, periodic=False) if window_fn else None
+            self.register_buffer("window", window_tensor)
+            self.stft = lambda x: stft_patch(
+                x,
+                n_fft=self.n_fft,
+                hop_length=self.hop_length,
+                win_length=self.win_length,
+                center=False if exact_pad else True,
+                window=self.window.to(dtype=torch.float),
+                return_complex=False,
+            )
+
+        self.normalize = normalize
+        self.log = log
+        self.dither = dither
+        self.frame_splicing = frame_splicing
+        self.nfilt = nfilt
+        self.preemph = preemph
+        self.pad_to = pad_to
+        highfreq = highfreq or sample_rate / 2
+
+        filterbanks = torch.tensor(
+            librosa.filters.mel(sample_rate, self.n_fft, n_mels=nfilt, fmin=lowfreq, fmax=highfreq), dtype=torch.float
+        ).unsqueeze(0)
+        self.register_buffer("fb", filterbanks)
+
+        # Calculate maximum sequence length
+        max_length = self.get_seq_len(torch.tensor(max_duration * sample_rate, dtype=torch.float))
+        max_pad = pad_to - (max_length % pad_to) if pad_to > 0 else 0
+        self.max_length = max_length + max_pad
+        self.pad_value = pad_value
+        self.mag_power = mag_power
+
+        # We want to avoid taking the log of zero
+        # There are two options: either adding or clamping to a small value
+        if log_zero_guard_type not in ["add", "clamp"]:
+            raise ValueError(
+                f"{self} received {log_zero_guard_type} for the "
+                f"log_zero_guard_type parameter. It must be either 'add' or "
+                f"'clamp'."
+            )
+
+        self.use_grads = use_grads
+        if not use_grads:
+            self.forward = torch.no_grad()(self.forward)
+
+        # log_zero_guard_value is the the small we want to use, we support
+        # an actual number, or "tiny", or "eps"
+        self.log_zero_guard_type = log_zero_guard_type
+        logging.debug(f"sr: {sample_rate}")
+        logging.debug(f"n_fft: {self.n_fft}")
+        logging.debug(f"win_length: {self.win_length}")
+        logging.debug(f"hop_length: {self.hop_length}")
+        logging.debug(f"n_mels: {nfilt}")
+        logging.debug(f"fmin: {lowfreq}")
+        logging.debug(f"fmax: {highfreq}")
+        logging.debug(f"using grads: {use_grads}")
+
+    def log_zero_guard_value_fn(self, x):
+        if isinstance(self.log_zero_guard_value, str):
+            if self.log_zero_guard_value == "tiny":
+                return torch.finfo(x.dtype).tiny
+            elif self.log_zero_guard_value == "eps":
+                return torch.finfo(x.dtype).eps
+            else:
+                raise ValueError(
+                    f"{self} received {self.log_zero_guard_value} for the "
+                    f"log_zero_guard_type parameter. It must be either a "
+                    f"number, 'tiny', or 'eps'"
+                )
+        else:
+            return self.log_zero_guard_value
+
+    def get_seq_len(self, seq_len):
+        if isinstance(self.stft, STFT):
+            pad_amount = self.stft.pad_amount * 2
+        else:
+            # Assuming that center is True is stft_pad_amount = 0
+            pad_amount = self.stft_pad_amount * 2 if self.stft_pad_amount is not None else self.n_fft // 2 * 2
+        seq_len = torch.floor((seq_len + pad_amount - self.n_fft) / self.hop_length) + 1
+        return seq_len.to(dtype=torch.long)
+
+    @property
+    def filter_banks(self):
+        return self.fb
+
+    def forward(self, x, seq_len):
+        seq_len = self.get_seq_len(seq_len.float())
+
+        if self.stft_pad_amount is not None:
+            x = torch.nn.functional.pad(
+                x.unsqueeze(1), (self.stft_pad_amount, self.stft_pad_amount), "reflect"
+            ).squeeze(1)
+
+        # dither (only in training mode for eval determinism)
+        if self.training and self.dither > 0:
+            x += self.dither * torch.randn_like(x)
+
+        # do preemphasis
+        if self.preemph is not None:
+            x = torch.cat((x[:, 0].unsqueeze(1), x[:, 1:] - self.preemph * x[:, :-1]), dim=1)
+
+        # disable autocast to get full range of stft values
+        with torch.cuda.amp.autocast(enabled=False):
+            x = self.stft(x)
+
+        # torch returns real, imag; so convert to magnitude
+        if not self.stft_conv:
+            # guard is needed for sqrt if grads are passed through
+            guard = 0 if not self.use_grads else CONSTANT
+            if x.dtype in [torch.cfloat, torch.cdouble]:
+                x = torch.view_as_real(x)
+            x = torch.sqrt(x.pow(2).sum(-1) + guard)
+
+        # get power spectrum
+        if self.mag_power != 1.0:
+            x = x.pow(self.mag_power)
+
+        # dot with filterbank energies
+        x = torch.matmul(self.fb.to(x.dtype), x)
+
+        # log features if required
+        if self.log:
+            if self.log_zero_guard_type == "add":
+                x = torch.log(x + self.log_zero_guard_value_fn(x))
+            elif self.log_zero_guard_type == "clamp":
+                x = torch.log(torch.clamp(x, min=self.log_zero_guard_value_fn(x)))
+            else:
+                raise ValueError("log_zero_guard_type was not understood")
+
+        # frame splicing if required
+        if self.frame_splicing > 1:
+            x = splice_frames(x, self.frame_splicing)
+
+        # normalize if required
+        if self.normalize:
+            x = normalize_batch(x, seq_len, normalize_type=self.normalize)
+
+        # mask to zero any values beyond seq_len in batch, pad to multiple of `pad_to` (for efficiency)
+        max_len = x.size(-1)
+        mask = torch.arange(max_len).to(x.device)
+        mask = mask.expand(x.size(0), max_len) >= seq_len.unsqueeze(1)
+        x = x.masked_fill(mask.unsqueeze(1).type(torch.bool).to(device=x.device), self.pad_value)
+        del mask
+        pad_to = self.pad_to
+        if pad_to == "max":
+            x = nn.functional.pad(x, (0, self.max_length - x.size(-1)), value=self.pad_value)
+        elif pad_to > 0:
+            pad_amt = x.size(-1) % pad_to
+            if pad_amt != 0:
+                x = nn.functional.pad(x, (0, pad_to - pad_amt), value=self.pad_value)
+
+        return x, seq_len

tts/nemo/collections/asr/parts/preprocessing/perturb.py

@@ -0,0 +1,895 @@
+# 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) 2018 Ryan Leary
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+# This file contains code artifacts adapted from https://github.com/ryanleary/patter
+import copy
+import io
+import os
+import random
+import subprocess
+from tempfile import NamedTemporaryFile
+from typing import List, Optional, Union
+
+import librosa
+import numpy as np
+import soundfile as sf
+import webdataset as wd
+from omegaconf import DictConfig, OmegaConf
+from scipy import signal
+from torch.utils.data import IterableDataset
+
+from nemo.collections.asr.parts.preprocessing.segment import AudioSegment
+from nemo.collections.common.parts.preprocessing import collections, parsers
+from nemo.utils import logging
+
+try:
+    from nemo.collections.asr.parts.utils import numba_utils
+
+    HAVE_NUMBA = True
+except (ImportError, ModuleNotFoundError):
+    HAVE_NUMBA = False
+
+
+def read_one_audiosegment(manifest, target_sr, rng, tarred_audio=False, audio_dataset=None):
+
+    if tarred_audio:
+        if audio_dataset is None:
+            raise TypeError("Expected augmentation dataset but got None")
+        audio_file, file_id = next(audio_dataset)
+        manifest_idx = manifest.mapping[file_id]
+        manifest_entry = manifest[manifest_idx]
+
+        offset = 0 if manifest_entry.offset is None else manifest_entry.offset
+        duration = 0 if manifest_entry.duration is None else manifest_entry.duration
+
+    else:
+        audio_record = rng.sample(manifest.data, 1)[0]
+        audio_file = audio_record.audio_file
+        offset = 0 if audio_record.offset is None else audio_record.offset
+        duration = 0 if audio_record.duration is None else audio_record.duration
+
+    return AudioSegment.from_file(audio_file, target_sr=target_sr, offset=offset, duration=duration)
+
+
+class Perturbation(object):
+    def max_augmentation_length(self, length):
+        return length
+
+    def perturb(self, data):
+        raise NotImplementedError
+
+
+class SpeedPerturbation(Perturbation):
+    """
+    Performs Speed Augmentation by re-sampling the data to a different sampling rate,
+    which does not preserve pitch.
+
+    Note: This is a very slow operation for online augmentation. If space allows,
+    it is preferable to pre-compute and save the files to augment the dataset.
+
+    Args:
+        sr: Original sampling rate.
+        resample_type: Type of resampling operation that will be performed.
+            For better speed using `resampy`'s fast resampling method, use `resample_type='kaiser_fast'`.
+            For high-quality resampling, set `resample_type='kaiser_best'`.
+            To use `scipy.signal.resample`, set `resample_type='fft'` or `resample_type='scipy'`
+        min_speed_rate: Minimum sampling rate modifier.
+        max_speed_rate: Maximum sampling rate modifier.
+        num_rates: Number of discrete rates to allow. Can be a positive or negative
+            integer.
+            If a positive integer greater than 0 is provided, the range of
+            speed rates will be discretized into `num_rates` values.
+            If a negative integer or 0 is provided, the full range of speed rates
+            will be sampled uniformly.
+            Note: If a positive integer is provided and the resultant discretized
+            range of rates contains the value '1.0', then those samples with rate=1.0,
+            will not be augmented at all and simply skipped. This is to unnecessary
+            augmentation and increase computation time. Effective augmentation chance
+            in such a case is = `prob * (num_rates - 1 / num_rates) * 100`% chance
+            where `prob` is the global probability of a sample being augmented.
+        rng: Random seed number.
+    """
+
+    def __init__(self, sr, resample_type, min_speed_rate=0.9, max_speed_rate=1.1, num_rates=5, rng=None):
+
+        min_rate = min(min_speed_rate, max_speed_rate)
+        if min_rate < 0.0:
+            raise ValueError("Minimum sampling rate modifier must be > 0.")
+
+        if resample_type not in ('kaiser_best', 'kaiser_fast', 'fft', 'scipy'):
+            raise ValueError("Supported `resample_type` values are ('kaiser_best', 'kaiser_fast', 'fft', 'scipy')")
+
+        self._sr = sr
+        self._min_rate = min_speed_rate
+        self._max_rate = max_speed_rate
+        self._num_rates = num_rates
+        if num_rates > 0:
+            self._rates = np.linspace(self._min_rate, self._max_rate, self._num_rates, endpoint=True)
+        self._res_type = resample_type
+        self._rng = random.Random() if rng is None else rng
+
+    def max_augmentation_length(self, length):
+        return length * self._max_rate
+
+    def perturb(self, data):
+        # Select speed rate either from choice or random sample
+        if self._num_rates < 0:
+            speed_rate = self._rng.uniform(self._min_rate, self._max_rate)
+        else:
+            speed_rate = self._rng.choice(self._rates)
+
+        # Skip perturbation in case of identity speed rate
+        if speed_rate == 1.0:
+            return
+
+        new_sr = int(self._sr * speed_rate)
+        data._samples = librosa.core.resample(data._samples, self._sr, new_sr, res_type=self._res_type)
+
+
+class TimeStretchPerturbation(Perturbation):
+    """
+    Time-stretch an audio series by a fixed rate while preserving pitch, based on [1, 2].
+
+    Note:
+    This is a simplified implementation, intended primarily for reference and pedagogical purposes.
+    It makes no attempt to handle transients, and is likely to produce audible artifacts.
+
+    Reference
+    [1] [Ellis, D. P. W. “A phase vocoder in Matlab.” Columbia University, 2002.]
+    (http://www.ee.columbia.edu/~dpwe/resources/matlab/pvoc/)
+    [2] [librosa.effects.time_stretch]
+    (https://librosa.github.io/librosa/generated/librosa.effects.time_stretch.html)
+
+    Args:
+        min_speed_rate: Minimum sampling rate modifier.
+        max_speed_rate: Maximum sampling rate modifier.
+        num_rates: Number of discrete rates to allow. Can be a positive or negative
+            integer.
+            If a positive integer greater than 0 is provided, the range of
+            speed rates will be discretized into `num_rates` values.
+            If a negative integer or 0 is provided, the full range of speed rates
+            will be sampled uniformly.
+            Note: If a positive integer is provided and the resultant discretized
+            range of rates contains the value '1.0', then those samples with rate=1.0,
+            will not be augmented at all and simply skipped. This is to avoid unnecessary
+            augmentation and increase computation time. Effective augmentation chance
+            in such a case is = `prob * (num_rates - 1 / num_rates) * 100`% chance
+            where `prob` is the global probability of a sample being augmented.
+        n_fft: Number of fft filters to be computed.
+        rng: Random seed number.
+    """
+
+    def __init__(self, min_speed_rate=0.9, max_speed_rate=1.1, num_rates=5, n_fft=512, rng=None):
+
+        min_rate = min(min_speed_rate, max_speed_rate)
+        if min_rate < 0.0:
+            raise ValueError("Minimum sampling rate modifier must be > 0.")
+
+        self._min_rate = min_speed_rate
+        self._max_rate = max_speed_rate
+        self._num_rates = num_rates
+        if num_rates > 0:
+            self._rates = np.linspace(self._min_rate, self._max_rate, self._num_rates, endpoint=True)
+        self._rng = random.Random() if rng is None else rng
+
+        # Pre-compute constants
+        self._n_fft = int(n_fft)
+        self._hop_length = int(n_fft // 2)
+
+        # Pre-allocate buffers
+        self._phi_advance_fast = np.linspace(0, np.pi * self._hop_length, self._hop_length + 1)
+        self._scale_buffer_fast = np.empty(self._hop_length + 1, dtype=np.float32)
+
+        self._phi_advance_slow = np.linspace(0, np.pi * self._n_fft, self._n_fft + 1)
+        self._scale_buffer_slow = np.empty(self._n_fft + 1, dtype=np.float32)
+
+    def max_augmentation_length(self, length):
+        return length * self._max_rate
+
+    def perturb(self, data):
+        # Select speed rate either from choice or random sample
+        if self._num_rates < 0:
+            speed_rate = self._rng.uniform(self._min_rate, self._max_rate)
+        else:
+            speed_rate = self._rng.choice(self._rates)
+
+        # Skip perturbation in case of identity speed rate
+        if speed_rate == 1.0:
+            return
+
+        # Increase `n_fft` based on task (speed up or slow down audio)
+        # This greatly reduces upper bound of maximum time taken
+        # to compute slowed down audio segments.
+        if speed_rate >= 1.0:  # Speed up audio
+            fft_multiplier = 1
+            phi_advance = self._phi_advance_fast
+            scale_buffer = self._scale_buffer_fast
+
+        else:  # Slow down audio
+            fft_multiplier = 2
+            phi_advance = self._phi_advance_slow
+            scale_buffer = self._scale_buffer_slow
+
+        n_fft = int(self._n_fft * fft_multiplier)
+        hop_length = int(self._hop_length * fft_multiplier)
+
+        # Perform short-term Fourier transform (STFT)
+        stft = librosa.core.stft(data._samples, n_fft=n_fft, hop_length=hop_length)
+
+        # Stretch by phase vocoding
+        if HAVE_NUMBA:
+            stft_stretch = numba_utils.phase_vocoder(stft, speed_rate, phi_advance, scale_buffer)
+
+        else:
+            stft_stretch = librosa.core.phase_vocoder(stft, speed_rate, hop_length)
+
+        # Predict the length of y_stretch
+        len_stretch = int(round(len(data._samples) / speed_rate))
+
+        # Invert the STFT
+        y_stretch = librosa.core.istft(
+            stft_stretch, dtype=data._samples.dtype, hop_length=hop_length, length=len_stretch
+        )
+
+        data._samples = y_stretch
+
+
+class GainPerturbation(Perturbation):
+    """
+    Applies random gain to the audio.
+
+    Args:
+        min_gain_dbfs (float): Min gain level in dB
+        max_gain_dbfs (float): Max gain level in dB
+        rng: Random number generator
+    """
+
+    def __init__(self, min_gain_dbfs=-10, max_gain_dbfs=10, rng=None):
+        self._min_gain_dbfs = min_gain_dbfs
+        self._max_gain_dbfs = max_gain_dbfs
+        self._rng = random.Random() if rng is None else rng
+
+    def perturb(self, data):
+        gain = self._rng.uniform(self._min_gain_dbfs, self._max_gain_dbfs)
+        # logging.debug("gain: %d", gain)
+        data._samples = data._samples * (10.0 ** (gain / 20.0))
+
+
+class ImpulsePerturbation(Perturbation):
+    """
+    Convolves audio with a Room Impulse Response.
+
+    Args:
+        manifest_path (list): Manifest file for RIRs
+        audio_tar_filepaths (list): Tar files, if RIR audio files are tarred
+        shuffle_n (int): Shuffle parameter for shuffling buffered files from the tar files
+        shift_impulse (bool): Shift impulse response to adjust for delay at the beginning
+    """
+
+    def __init__(self, manifest_path=None, rng=None, audio_tar_filepaths=None, shuffle_n=128, shift_impulse=False):
+        self._manifest = collections.ASRAudioText(manifest_path, parser=parsers.make_parser([]), index_by_file_id=True)
+        self._audiodataset = None
+        self._tarred_audio = False
+        self._shift_impulse = shift_impulse
+        self._data_iterator = None
+
+        if audio_tar_filepaths:
+            self._tarred_audio = True
+            self._audiodataset = AugmentationDataset(manifest_path, audio_tar_filepaths, shuffle_n)
+            self._data_iterator = iter(self._audiodataset)
+
+        self._rng = random.Random() if rng is None else rng
+
+    def perturb(self, data):
+        impulse = read_one_audiosegment(
+            self._manifest,
+            data.sample_rate,
+            self._rng,
+            tarred_audio=self._tarred_audio,
+            audio_dataset=self._data_iterator,
+        )
+        if not self._shift_impulse:
+            impulse_norm = (impulse.samples - min(impulse.samples)) / (max(impulse.samples) - min(impulse.samples))
+            data._samples = signal.fftconvolve(data._samples, impulse_norm, "same")
+        else:
+            # Find peak and shift peak to left
+            impulse_norm = (impulse.samples - min(impulse.samples)) / (max(impulse.samples) - min(impulse.samples))
+            max_ind = np.argmax(np.abs(impulse_norm))
+
+            impulse_resp = impulse_norm[max_ind:]
+            delay_after = len(impulse_resp)
+            data._samples = signal.fftconvolve(data._samples, impulse_resp, "full")[:-delay_after]
+
+
+class ShiftPerturbation(Perturbation):
+    """
+    Perturbs audio by shifting the audio in time by a random amount between min_shift_ms and max_shift_ms.
+    The final length of the audio is kept unaltered by padding the audio with zeros.
+
+
+    Args:
+        min_shift_ms (float): Minimum time in milliseconds by which audio will be shifted
+        max_shift_ms (float): Maximum time in milliseconds by which audio will be shifted
+        rng: Random number generator
+    """
+
+    def __init__(self, min_shift_ms=-5.0, max_shift_ms=5.0, rng=None):
+        self._min_shift_ms = min_shift_ms
+        self._max_shift_ms = max_shift_ms
+        self._rng = random.Random() if rng is None else rng
+
+    def perturb(self, data):
+        shift_ms = self._rng.uniform(self._min_shift_ms, self._max_shift_ms)
+        if abs(shift_ms) / 1000 > data.duration:
+            # TODO: do something smarter than just ignore this condition
+            return
+        shift_samples = int(shift_ms * data.sample_rate // 1000)
+        # logging.debug("shift: %s", shift_samples)
+        if shift_samples < 0:
+            data._samples[-shift_samples:] = data._samples[:shift_samples]
+            data._samples[:-shift_samples] = 0
+        elif shift_samples > 0:
+            data._samples[:-shift_samples] = data._samples[shift_samples:]
+            data._samples[-shift_samples:] = 0
+
+
+class NoisePerturbation(Perturbation):
+    """
+    Perturbation that adds noise to input audio.
+
+    Args:
+        manifest_path (str): Manifest file with paths to noise files
+        min_snr_db (float): Minimum SNR of audio after noise is added
+        max_snr_db (float): Maximum SNR of audio after noise is added
+        max_gain_db (float): Maximum gain that can be applied on the noise sample
+        audio_tar_filepaths (list) : Tar files, if noise audio files are tarred
+        shuffle_n (int): Shuffle parameter for shuffling buffered files from the tar files
+        orig_sr (int): Original sampling rate of the noise files
+        rng: Random number generator
+    """
+
+    def __init__(
+        self,
+        manifest_path=None,
+        min_snr_db=10,
+        max_snr_db=50,
+        max_gain_db=300.0,
+        rng=None,
+        audio_tar_filepaths=None,
+        shuffle_n=100,
+        orig_sr=16000,
+    ):
+        self._manifest = collections.ASRAudioText(manifest_path, parser=parsers.make_parser([]), index_by_file_id=True)
+        self._audiodataset = None
+        self._tarred_audio = False
+        self._orig_sr = orig_sr
+        self._data_iterator = None
+
+        if audio_tar_filepaths:
+            self._tarred_audio = True
+            self._audiodataset = AugmentationDataset(manifest_path, audio_tar_filepaths, shuffle_n)
+            self._data_iterator = iter(self._audiodataset)
+
+        self._rng = random.Random() if rng is None else rng
+        self._min_snr_db = min_snr_db
+        self._max_snr_db = max_snr_db
+        self._max_gain_db = max_gain_db
+
+    @property
+    def orig_sr(self):
+        return self._orig_sr
+
+    def get_one_noise_sample(self, target_sr):
+        return read_one_audiosegment(
+            self._manifest, target_sr, self._rng, tarred_audio=self._tarred_audio, audio_dataset=self._data_iterator
+        )
+
+    def perturb(self, data):
+        noise = read_one_audiosegment(
+            self._manifest,
+            data.sample_rate,
+            self._rng,
+            tarred_audio=self._tarred_audio,
+            audio_dataset=self._data_iterator,
+        )
+        self.perturb_with_input_noise(data, noise)
+
+    def perturb_with_input_noise(self, data, noise, data_rms=None):
+        snr_db = self._rng.uniform(self._min_snr_db, self._max_snr_db)
+        if data_rms is None:
+            data_rms = data.rms_db
+        noise_gain_db = min(data_rms - noise.rms_db - snr_db, self._max_gain_db)
+        # logging.debug("noise: %s %s %s", snr_db, noise_gain_db, noise_record.audio_file)
+
+        # calculate noise segment to use
+        start_time = self._rng.uniform(0.0, noise.duration - data.duration)
+        if noise.duration > (start_time + data.duration):
+            noise.subsegment(start_time=start_time, end_time=start_time + data.duration)
+
+        # adjust gain for snr purposes and superimpose
+        noise.gain_db(noise_gain_db)
+
+        if noise._samples.shape[0] < data._samples.shape[0]:
+            noise_idx = self._rng.randint(0, data._samples.shape[0] - noise._samples.shape[0])
+            data._samples[noise_idx : noise_idx + noise._samples.shape[0]] += noise._samples
+
+        else:
+            data._samples += noise._samples
+
+    def perturb_with_foreground_noise(
+        self, data, noise, data_rms=None, max_noise_dur=2, max_additions=1,
+    ):
+        snr_db = self._rng.uniform(self._min_snr_db, self._max_snr_db)
+        if not data_rms:
+            data_rms = data.rms_db
+
+        noise_gain_db = min(data_rms - noise.rms_db - snr_db, self._max_gain_db)
+        n_additions = self._rng.randint(1, max_additions)
+
+        for i in range(n_additions):
+            noise_dur = self._rng.uniform(0.0, max_noise_dur)
+            start_time = self._rng.uniform(0.0, noise.duration)
+            start_sample = int(round(start_time * noise.sample_rate))
+            end_sample = int(round(min(noise.duration, (start_time + noise_dur)) * noise.sample_rate))
+            noise_samples = np.copy(noise._samples[start_sample:end_sample])
+            # adjust gain for snr purposes and superimpose
+            noise_samples *= 10.0 ** (noise_gain_db / 20.0)
+
+            if noise_samples.shape[0] > data._samples.shape[0]:
+                noise_samples = noise_samples[0 : data._samples.shape[0]]
+
+            noise_idx = self._rng.randint(0, data._samples.shape[0] - noise_samples.shape[0])
+            data._samples[noise_idx : noise_idx + noise_samples.shape[0]] += noise_samples
+
+
+class WhiteNoisePerturbation(Perturbation):
+    """
+    Perturbation that adds white noise to an audio file in the training dataset.
+
+    Args:
+        min_level (int): Minimum level in dB at which white noise should be added
+        max_level (int): Maximum level in dB at which white noise should be added
+        rng: Random number generator
+    """
+
+    def __init__(self, min_level=-90, max_level=-46, rng=None):
+        self.min_level = int(min_level)
+        self.max_level = int(max_level)
+        self._rng = np.random.RandomState() if rng is None else rng
+
+    def perturb(self, data):
+        noise_level_db = self._rng.randint(self.min_level, self.max_level, dtype='int32')
+        noise_signal = self._rng.randn(data._samples.shape[0]) * (10.0 ** (noise_level_db / 20.0))
+        data._samples += noise_signal
+
+
+class RirAndNoisePerturbation(Perturbation):
+    """
+        RIR augmentation with additive foreground and background noise.
+        In this implementation audio data is augmented by first convolving the audio with a Room Impulse Response
+        and then adding foreground noise and background noise at various SNRs. RIR, foreground and background noises
+        should either be supplied with a manifest file or as tarred audio files (faster).
+
+        Different sets of noise audio files based on the original sampling rate of the noise. This is useful while
+        training a mixed sample rate model. For example, when training a mixed model with 8 kHz and 16 kHz audio with a
+        target sampling rate of 16 kHz, one would want to augment 8 kHz data with 8 kHz noise rather than 16 kHz noise.
+
+        Args:
+            rir_manifest_path: Manifest file for RIRs
+            rir_tar_filepaths: Tar files, if RIR audio files are tarred
+            rir_prob: Probability of applying a RIR
+            noise_manifest_paths: Foreground noise manifest path
+            min_snr_db: Min SNR for foreground noise
+            max_snr_db: Max SNR for background noise,
+            noise_tar_filepaths: Tar files, if noise files are tarred
+            apply_noise_rir: Whether to convolve foreground noise with a a random RIR
+            orig_sample_rate: Original sampling rate of foreground noise audio
+            max_additions: Max number of times foreground noise is added to an utterance,
+            max_duration: Max duration of foreground noise
+            bg_noise_manifest_paths: Background noise manifest path
+            bg_min_snr_db: Min SNR for background noise
+            bg_max_snr_db: Max SNR for background noise
+            bg_noise_tar_filepaths: Tar files, if noise files are tarred
+            bg_orig_sample_rate: Original sampling rate of background noise audio
+
+    """
+
+    def __init__(
+        self,
+        rir_manifest_path=None,
+        rir_prob=0.5,
+        noise_manifest_paths=None,
+        min_snr_db=0,
+        max_snr_db=50,
+        rir_tar_filepaths=None,
+        rir_shuffle_n=100,
+        noise_tar_filepaths=None,
+        apply_noise_rir=False,
+        orig_sample_rate=None,
+        max_additions=5,
+        max_duration=2.0,
+        bg_noise_manifest_paths=None,
+        bg_min_snr_db=10,
+        bg_max_snr_db=50,
+        bg_noise_tar_filepaths=None,
+        bg_orig_sample_rate=None,
+    ):
+
+        logging.info("Called Rir aug init")
+        self._rir_prob = rir_prob
+        self._rng = random.Random()
+        self._rir_perturber = ImpulsePerturbation(
+            manifest_path=rir_manifest_path,
+            audio_tar_filepaths=rir_tar_filepaths,
+            shuffle_n=rir_shuffle_n,
+            shift_impulse=True,
+        )
+        self._fg_noise_perturbers = {}
+        self._bg_noise_perturbers = {}
+        if noise_manifest_paths:
+            for i in range(len(noise_manifest_paths)):
+                if orig_sample_rate is None:
+                    orig_sr = 16000
+                else:
+                    orig_sr = orig_sample_rate[i]
+                self._fg_noise_perturbers[orig_sr] = NoisePerturbation(
+                    manifest_path=noise_manifest_paths[i],
+                    min_snr_db=min_snr_db[i],
+                    max_snr_db=max_snr_db[i],
+                    audio_tar_filepaths=noise_tar_filepaths[i],
+                    orig_sr=orig_sr,
+                )
+        self._max_additions = max_additions
+        self._max_duration = max_duration
+        if bg_noise_manifest_paths:
+            for i in range(len(bg_noise_manifest_paths)):
+                if bg_orig_sample_rate is None:
+                    orig_sr = 16000
+                else:
+                    orig_sr = bg_orig_sample_rate[i]
+                self._bg_noise_perturbers[orig_sr] = NoisePerturbation(
+                    manifest_path=bg_noise_manifest_paths[i],
+                    min_snr_db=bg_min_snr_db[i],
+                    max_snr_db=bg_max_snr_db[i],
+                    audio_tar_filepaths=bg_noise_tar_filepaths[i],
+                    orig_sr=orig_sr,
+                )
+
+        self._apply_noise_rir = apply_noise_rir
+
+    def perturb(self, data):
+        prob = self._rng.uniform(0.0, 1.0)
+
+        if prob < self._rir_prob:
+            self._rir_perturber.perturb(data)
+
+        orig_sr = data.orig_sr
+        if orig_sr not in self._fg_noise_perturbers:
+            orig_sr = max(self._fg_noise_perturbers.keys())
+        fg_perturber = self._fg_noise_perturbers[orig_sr]
+
+        orig_sr = data.orig_sr
+        if orig_sr not in self._bg_noise_perturbers:
+            orig_sr = max(self._bg_noise_perturbers.keys())
+        bg_perturber = self._bg_noise_perturbers[orig_sr]
+
+        data_rms = data.rms_db
+        noise = fg_perturber.get_one_noise_sample(data.sample_rate)
+        if self._apply_noise_rir:
+            self._rir_perturber.perturb(noise)
+        fg_perturber.perturb_with_foreground_noise(
+            data, noise, data_rms=data_rms, max_noise_dur=self._max_duration, max_additions=self._max_additions
+        )
+        noise = bg_perturber.get_one_noise_sample(data.sample_rate)
+        bg_perturber.perturb_with_input_noise(data, noise, data_rms=data_rms)
+
+
+class TranscodePerturbation(Perturbation):
+    """
+        Audio codec augmentation. This implementation uses sox to transcode audio with low rate audio codecs,
+        so users need to make sure that the installed sox version supports the codecs used here (G711 and amr-nb).
+
+        Args:
+            rng: Random number generator
+    """
+
+    def __init__(self, rng=None):
+        self._rng = np.random.RandomState() if rng is None else rng
+        self._codecs = ["g711", "amr-nb"]
+
+    def perturb(self, data):
+        att_factor = 0.8
+        max_level = np.max(np.abs(data._samples))
+        norm_factor = att_factor / max_level
+        norm_samples = norm_factor * data._samples
+        orig_f = NamedTemporaryFile(suffix=".wav")
+        sf.write(orig_f.name, norm_samples.transpose(), 16000)
+
+        codec_ind = random.randint(0, len(self._codecs) - 1)
+        if self._codecs[codec_ind] == "amr-nb":
+            transcoded_f = NamedTemporaryFile(suffix="_amr.wav")
+            rates = list(range(0, 8))
+            rate = rates[random.randint(0, len(rates) - 1)]
+            _ = subprocess.check_output(
+                f"sox {orig_f.name} -V0 -C {rate} -t amr-nb - | sox -t amr-nb - -V0 -b 16 -r 16000 {transcoded_f.name}",
+                shell=True,
+            )
+        elif self._codecs[codec_ind] == "g711":
+            transcoded_f = NamedTemporaryFile(suffix="_g711.wav")
+            _ = subprocess.check_output(
+                f"sox {orig_f.name} -V0  -r 8000 -c 1 -e a-law {transcoded_f.name}", shell=True
+            )
+
+        new_data = AudioSegment.from_file(transcoded_f.name, target_sr=16000)
+        data._samples = new_data._samples[0 : data._samples.shape[0]]
+        return
+
+
+perturbation_types = {
+    "speed": SpeedPerturbation,
+    "time_stretch": TimeStretchPerturbation,
+    "gain": GainPerturbation,
+    "impulse": ImpulsePerturbation,
+    "shift": ShiftPerturbation,
+    "noise": NoisePerturbation,
+    "white_noise": WhiteNoisePerturbation,
+    "rir_noise_aug": RirAndNoisePerturbation,
+    "transcode_aug": TranscodePerturbation,
+}
+
+
+def register_perturbation(name: str, perturbation: Perturbation):
+    if name in perturbation_types.keys():
+        raise KeyError(
+            f"Perturbation with the name {name} exists. " f"Type of perturbation : {perturbation_types[name]}."
+        )
+
+    perturbation_types[name] = perturbation
+
+
+class AudioAugmentor(object):
+    def __init__(self, perturbations=None, rng=None):
+        self._rng = random.Random() if rng is None else rng
+        self._pipeline = perturbations if perturbations is not None else []
+
+    def perturb(self, segment):
+        for (prob, p) in self._pipeline:
+            if self._rng.random() < prob:
+                p.perturb(segment)
+        return
+
+    def max_augmentation_length(self, length):
+        newlen = length
+        for (prob, p) in self._pipeline:
+            newlen = p.max_augmentation_length(newlen)
+        return newlen
+
+    @classmethod
+    def from_config(cls, config):
+        ptbs = []
+        for p in config:
+            if p['aug_type'] not in perturbation_types:
+                logging.warning("%s perturbation not known. Skipping.", p['aug_type'])
+                continue
+            perturbation = perturbation_types[p['aug_type']]
+            ptbs.append((p['prob'], perturbation(**p['cfg'])))
+        return cls(perturbations=ptbs)
+
+
+def process_augmentations(augmenter) -> Optional[AudioAugmentor]:
+    """Process list of online data augmentations.
+    Accepts either an AudioAugmentor object with pre-defined augmentations,
+    or a dictionary that points to augmentations that have been defined.
+    If a dictionary is passed, must follow the below structure:
+    Dict[str, Dict[str, Any]]: Which refers to a dictionary of string
+    names for augmentations, defined in `asr/parts/perturb.py`.
+    The inner dictionary may contain key-value arguments of the specific
+    augmentation, along with an essential key `prob`. `prob` declares the
+    probability of the augmentation being applied, and must be a float
+    value in the range [0, 1].
+    # Example in YAML config file
+    Augmentations are generally applied only during training, so we can add
+    these augmentations to our yaml config file, and modify the behaviour
+    for training and evaluation.
+    ```yaml
+    AudioToSpeechLabelDataLayer:
+        ...  # Parameters shared between train and evaluation time
+        train:
+            augmentor:
+                shift:
+                    prob: 0.5
+                    min_shift_ms: -5.0
+                    max_shift_ms: 5.0
+                white_noise:
+                    prob: 1.0
+                    min_level: -90
+                    max_level: -46
+                ...
+        eval:
+            ...
+    ```
+    Then in the training script,
+    ```python
+    import copy
+    from ruamel.yaml import YAML
+    yaml = YAML(typ="safe")
+    with open(model_config) as f:
+        params = yaml.load(f)
+    # Train Config for Data Loader
+    train_dl_params = copy.deepcopy(params["AudioToTextDataLayer"])
+    train_dl_params.update(params["AudioToTextDataLayer"]["train"])
+    del train_dl_params["train"]
+    del train_dl_params["eval"]
+    data_layer_train = nemo_asr.AudioToTextDataLayer(
+        ...,
+        **train_dl_params,
+    )
+    # Evaluation Config for Data Loader
+    eval_dl_params = copy.deepcopy(params["AudioToTextDataLayer"])
+    eval_dl_params.update(params["AudioToTextDataLayer"]["eval"])
+    del eval_dl_params["train"]
+    del eval_dl_params["eval"]
+    data_layer_eval = nemo_asr.AudioToTextDataLayer(
+        ...,
+        **eval_dl_params,
+    )
+    ```
+    # Registering your own Augmentations
+    To register custom augmentations to obtain the above convenience of
+    the declaring the augmentations in YAML, you can put additional keys in
+    `perturbation_types` dictionary as follows.
+    ```python
+    from nemo.collections.asr.parts import perturb
+    # Define your own perturbation here
+    class CustomPerturbation(perturb.Perturbation):
+        ...
+    perturb.register_perturbation(name_of_perturbation, CustomPerturbation)
+    ```
+    Args:
+        augmenter: AudioAugmentor object or
+            dictionary of str -> kwargs (dict) which is parsed and used
+            to initialize an AudioAugmentor.
+            Note: It is crucial that each individual augmentation has
+            a keyword `prob`, that defines a float probability in the
+            the range [0, 1] of this augmentation being applied.
+            If this keyword is not present, then the augmentation is
+            disabled and a warning is logged.
+    Returns: AudioAugmentor object
+    """
+    if augmenter is None:
+        return None
+
+    if isinstance(augmenter, AudioAugmentor):
+        return augmenter
+
+    if not type(augmenter) in {dict, DictConfig}:
+        raise ValueError("Cannot parse augmenter. Must be a dict or an AudioAugmentor object ")
+
+    if isinstance(augmenter, DictConfig):
+        augmenter = OmegaConf.to_container(augmenter, resolve=True)
+
+    augmenter = copy.deepcopy(augmenter)
+
+    augmentations = []
+    for augment_name, augment_kwargs in augmenter.items():
+        prob = augment_kwargs.get('prob', None)
+
+        if prob is None:
+            raise KeyError(
+                f'Augmentation "{augment_name}" will not be applied as '
+                f'keyword argument "prob" was not defined for this augmentation.'
+            )
+
+        else:
+            _ = augment_kwargs.pop('prob')
+
+            if prob < 0.0 or prob > 1.0:
+                raise ValueError("`prob` must be a float value between 0 and 1.")
+
+            try:
+                augmentation = perturbation_types[augment_name](**augment_kwargs)
+                augmentations.append([prob, augmentation])
+            except KeyError:
+                raise KeyError(f"Invalid perturbation name. Allowed values : {perturbation_types.keys()}")
+
+    augmenter = AudioAugmentor(perturbations=augmentations)
+    return augmenter
+
+
+class AugmentationDataset(IterableDataset):
+    """
+        A class that loads tarred audio files and cycles over the files in the dataset.
+
+        Accepts a single comma-separated JSON manifest file (in the same style as for the AudioToCharDataset/AudioToBPEDataset),
+        as well as the path(s) to the tarball(s) containing the wav files. Each line of the manifest should
+        contain the information for one audio file, including at least the transcript and name of the audio
+        file within the tarball.
+
+        Valid formats for the audio_tar_filepaths argument include:
+        (1) a single string that can be brace-expanded, e.g. 'path/to/audio.tar' or 'path/to/audio_{1..100}.tar.gz', or
+        (2) a list of file paths that will not be brace-expanded, e.g. ['audio_1.tar', 'audio_2.tar', ...].
+
+        Note: For brace expansion in (1), there may be cases where `{x..y}` syntax cannot be used due to shell interference.
+        This occurs most commonly inside SLURM scripts. Therefore we provide a few equivalent replacements.
+        Supported opening braces - { <=> (, [, < and the special tag _OP_.
+        Supported closing braces - } <=> ), ], > and the special tag _CL_.
+        For SLURM based tasks, we suggest the use of the special tags for ease of use.
+
+        See the WebDataset documentation for more information about accepted data and input formats.
+    """
+
+    def __init__(self, manifest_path: str, tar_filepaths: Union[str, List[str]], shuffle_n: int = 128):
+        self._manifest = collections.ASRAudioText(manifest_path, parser=parsers.make_parser([]), index_by_file_id=True)
+
+        if isinstance(tar_filepaths, str):
+            # Replace '(' and '[' with '{'
+            brace_keys_open = ['(', '[', '<', '_OP_']
+            for bkey in brace_keys_open:
+                if bkey in tar_filepaths:
+                    tar_filepaths = tar_filepaths.replace(bkey, "{")
+
+            # Replace ')' and ']' with '}'
+            brace_keys_close = [')', ']', '>', '_CL_']
+            for bkey in brace_keys_close:
+                if bkey in tar_filepaths:
+                    tar_filepaths = tar_filepaths.replace(bkey, "}")
+
+        self.audio_dataset = wd.WebDataset(urls=tar_filepaths, nodesplitter=None)
+
+        if shuffle_n > 0:
+            self.audio_dataset = self.audio_dataset.shuffle(shuffle_n)
+        else:
+            logging.info("WebDataset will not shuffle files within the tar files.")
+
+        self.audio_dataset = self.audio_dataset.rename(audio='wav', key='__key__').to_tuple('audio', 'key')
+        self.audio_iter = iter(self.audio_dataset)
+
+    def __len__(self):
+        return len(self._manifest)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        while True:
+            try:
+                audio_bytes, audio_filename = next(self.audio_iter)
+
+            except StopIteration:
+                self.audio_iter = iter(self.audio_dataset)
+                audio_bytes, audio_filename = next(self.audio_iter)
+            file_id, _ = os.path.splitext(os.path.basename(audio_filename))
+
+            # Convert audio bytes to IO stream for processing (for SoundFile to read)
+            audio_file = io.BytesIO(audio_bytes)
+            return audio_file, file_id

tts/nemo/collections/asr/parts/preprocessing/segment.py

@@ -0,0 +1,270 @@
+# 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) 2018 Ryan Leary
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+# This file contains code artifacts adapted from https://github.com/ryanleary/patter
+
+import os
+import random
+
+import librosa
+import numpy as np
+import soundfile as sf
+from kaldiio.matio import read_kaldi
+from kaldiio.utils import open_like_kaldi
+from pydub import AudioSegment as Audio
+from pydub.exceptions import CouldntDecodeError
+
+from nemo.utils import logging
+
+available_formats = sf.available_formats()
+sf_supported_formats = ["." + i.lower() for i in available_formats.keys()]
+
+
+class AudioSegment(object):
+    """Monaural audio segment abstraction.
+    :param samples: Audio samples [num_samples x num_channels].
+    :type samples: ndarray.float32
+    :param sample_rate: Audio sample rate.
+    :type sample_rate: int
+    :raises TypeError: If the sample data type is not float or int.
+    """
+
+    def __init__(self, samples, sample_rate, target_sr=None, trim=False, trim_db=60, orig_sr=None):
+        """Create audio segment from samples.
+        Samples are convert float32 internally, with int scaled to [-1, 1].
+        """
+        samples = self._convert_samples_to_float32(samples)
+        if target_sr is not None and target_sr != sample_rate:
+            samples = librosa.core.resample(samples, sample_rate, target_sr)
+            sample_rate = target_sr
+        if trim:
+            samples, _ = librosa.effects.trim(samples, trim_db)
+        self._samples = samples
+        self._sample_rate = sample_rate
+        if self._samples.ndim >= 2:
+            self._samples = np.mean(self._samples, 1)
+
+        self._orig_sr = orig_sr if orig_sr is not None else sample_rate
+
+    def __eq__(self, other):
+        """Return whether two objects are equal."""
+        if type(other) is not type(self):
+            return False
+        if self._sample_rate != other._sample_rate:
+            return False
+        if self._samples.shape != other._samples.shape:
+            return False
+        if np.any(self.samples != other._samples):
+            return False
+        return True
+
+    def __ne__(self, other):
+        """Return whether two objects are unequal."""
+        return not self.__eq__(other)
+
+    def __str__(self):
+        """Return human-readable representation of segment."""
+        return "%s: num_samples=%d, sample_rate=%d, duration=%.2fsec, rms=%.2fdB" % (
+            type(self),
+            self.num_samples,
+            self.sample_rate,
+            self.duration,
+            self.rms_db,
+        )
+
+    @staticmethod
+    def _convert_samples_to_float32(samples):
+        """Convert sample type to float32.
+        Audio sample type is usually integer or float-point.
+        Integers will be scaled to [-1, 1] in float32.
+        """
+        float32_samples = samples.astype('float32')
+        if samples.dtype in np.sctypes['int']:
+            bits = np.iinfo(samples.dtype).bits
+            float32_samples *= 1.0 / 2 ** (bits - 1)
+        elif samples.dtype in np.sctypes['float']:
+            pass
+        else:
+            raise TypeError("Unsupported sample type: %s." % samples.dtype)
+        return float32_samples
+
+    @classmethod
+    def from_file(
+        cls, audio_file, target_sr=None, int_values=False, offset=0, duration=0, trim=False, orig_sr=None,
+    ):
+        """
+        Load a file supported by librosa and return as an AudioSegment.
+        :param audio_file: path of file to load
+        :param target_sr: the desired sample rate
+        :param int_values: if true, load samples as 32-bit integers
+        :param offset: offset in seconds when loading audio
+        :param duration: duration in seconds when loading audio
+        :return: numpy array of samples
+        """
+        samples = None
+        if not isinstance(audio_file, str) or os.path.splitext(audio_file)[-1] in sf_supported_formats:
+            try:
+                with sf.SoundFile(audio_file, 'r') as f:
+                    dtype = 'int32' if int_values else 'float32'
+                    sample_rate = f.samplerate
+                    if offset > 0:
+                        f.seek(int(offset * sample_rate))
+                    if duration > 0:
+                        samples = f.read(int(duration * sample_rate), dtype=dtype)
+                    else:
+                        samples = f.read(dtype=dtype)
+                samples = samples.transpose()
+            except RuntimeError as e:
+                logging.error(
+                    f"Loading {audio_file} via SoundFile raised RuntimeError: `{e}`. "
+                    f"NeMo will fallback to loading via pydub."
+                )
+        elif isinstance(audio_file, str) and audio_file.strip()[-1] == "|":
+            f = open_like_kaldi(audio_file, "rb")
+            sample_rate, samples = read_kaldi(f)
+            if offset > 0:
+                samples = samples[int(offset * sample_rate) :]
+            if duration > 0:
+                samples = samples[: int(duration * sample_rate)]
+            if not int_values:
+                abs_max_value = np.abs(samples).max()
+                samples = np.array(samples, dtype=np.float) / abs_max_value
+
+        if samples is None:
+            try:
+                samples = Audio.from_file(audio_file)
+                sample_rate = samples.frame_rate
+                if offset > 0:
+                    # pydub does things in milliseconds
+                    seconds = offset * 1000
+                    samples = samples[int(seconds * sample_rate) :]
+                if duration > 0:
+                    seconds = duration * 1000
+                    samples = samples[: int(seconds)]
+                samples = np.array(samples.get_array_of_samples())
+            except CouldntDecodeError as e:
+                logging.error(f"Loading {audio_file} via pydub raised CouldntDecodeError: `{e}`.")
+
+        return cls(samples, sample_rate, target_sr=target_sr, trim=trim, orig_sr=orig_sr)
+
+    @classmethod
+    def segment_from_file(cls, audio_file, target_sr=None, n_segments=0, trim=False, orig_sr=None):
+        """Grabs n_segments number of samples from audio_file randomly from the
+        file as opposed to at a specified offset.
+
+        Note that audio_file can be either the file path, or a file-like object.
+        """
+        try:
+            with sf.SoundFile(audio_file, 'r') as f:
+                sample_rate = f.samplerate
+                if n_segments > 0 and len(f) > n_segments:
+                    max_audio_start = len(f) - n_segments
+                    audio_start = random.randint(0, max_audio_start)
+                    f.seek(audio_start)
+                    samples = f.read(n_segments, dtype='float32')
+                else:
+                    samples = f.read(dtype='float32')
+            samples = samples.transpose()
+        except RuntimeError as e:
+            logging.error(f"Loading {audio_file} via SoundFile raised RuntimeError: `{e}`.")
+
+        samples = samples.transpose()
+        return cls(samples, sample_rate, target_sr=target_sr, trim=trim, orig_sr=orig_sr)
+
+    @property
+    def samples(self):
+        return self._samples.copy()
+
+    @property
+    def sample_rate(self):
+        return self._sample_rate
+
+    @property
+    def num_samples(self):
+        return self._samples.shape[0]
+
+    @property
+    def duration(self):
+        return self._samples.shape[0] / float(self._sample_rate)
+
+    @property
+    def rms_db(self):
+        mean_square = np.mean(self._samples ** 2)
+        return 10 * np.log10(mean_square)
+
+    @property
+    def orig_sr(self):
+        return self._orig_sr
+
+    def gain_db(self, gain):
+        self._samples *= 10.0 ** (gain / 20.0)
+
+    def pad(self, pad_size, symmetric=False):
+        """Add zero padding to the sample. The pad size is given in number
+        of samples.
+        If symmetric=True, `pad_size` will be added to both sides. If false,
+        `pad_size`
+        zeros will be added only to the end.
+        """
+        self._samples = np.pad(self._samples, (pad_size if symmetric else 0, pad_size), mode='constant',)
+
+    def subsegment(self, start_time=None, end_time=None):
+        """Cut the AudioSegment between given boundaries.
+        Note that this is an in-place transformation.
+        :param start_time: Beginning of subsegment in seconds.
+        :type start_time: float
+        :param end_time: End of subsegment in seconds.
+        :type end_time: float
+        :raise ValueError: If start_time or end_time is incorrectly set,
+        e.g. out
+                           of bounds in time.
+        """
+        start_time = 0.0 if start_time is None else start_time
+        end_time = self.duration if end_time is None else end_time
+        if start_time < 0.0:
+            start_time = self.duration + start_time
+        if end_time < 0.0:
+            end_time = self.duration + end_time
+        if start_time < 0.0:
+            raise ValueError("The slice start position (%f s) is out of bounds." % start_time)
+        if end_time < 0.0:
+            raise ValueError("The slice end position (%f s) is out of bounds." % end_time)
+        if start_time > end_time:
+            raise ValueError(
+                "The slice start position (%f s) is later than the end position (%f s)." % (start_time, end_time)
+            )
+        if end_time > self.duration:
+            raise ValueError("The slice end position (%f s) is out of bounds (> %f s)" % (end_time, self.duration))
+        start_sample = int(round(start_time * self._sample_rate))
+        end_sample = int(round(end_time * self._sample_rate))
+        self._samples = self._samples[start_sample:end_sample]

tts/nemo/collections/asr/parts/submodules/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tts/nemo/collections/asr/parts/submodules/conformer_modules.py

@@ -0,0 +1,190 @@
+# 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 torch
+from torch import nn as nn
+from torch.nn import LayerNorm
+
+from nemo.collections.asr.parts.submodules.multi_head_attention import (
+    MultiHeadAttention,
+    RelPositionMultiHeadAttention,
+)
+from nemo.collections.asr.parts.utils.activations import Swish
+
+__all__ = ['ConformerConvolution', 'ConformerFeedForward', 'ConformerLayer']
+
+
+class ConformerLayer(torch.nn.Module):
+    """A single block of the Conformer encoder.
+
+    Args:
+        d_model (int): input dimension of MultiheadAttentionMechanism and PositionwiseFeedForward
+        d_ff (int): hidden dimension of PositionwiseFeedForward
+        n_heads (int): number of heads for multi-head attention
+        conv_kernel_size (int): kernel size for depthwise convolution in convolution module
+        dropout (float): dropout probabilities for linear layers
+        dropout_att (float): dropout probabilities for attention distributions
+    """
+
+    def __init__(
+        self,
+        d_model,
+        d_ff,
+        self_attention_model='rel_pos',
+        n_heads=4,
+        conv_kernel_size=31,
+        dropout=0.1,
+        dropout_att=0.1,
+        pos_bias_u=None,
+        pos_bias_v=None,
+    ):
+        super(ConformerLayer, self).__init__()
+
+        self.self_attention_model = self_attention_model
+        self.n_heads = n_heads
+        self.fc_factor = 0.5
+
+        # first feed forward module
+        self.norm_feed_forward1 = LayerNorm(d_model)
+        self.feed_forward1 = ConformerFeedForward(d_model=d_model, d_ff=d_ff, dropout=dropout)
+
+        # convolution module
+        self.norm_conv = LayerNorm(d_model)
+        self.conv = ConformerConvolution(d_model=d_model, kernel_size=conv_kernel_size)
+
+        # multi-headed self-attention module
+        self.norm_self_att = LayerNorm(d_model)
+        if self_attention_model == 'rel_pos':
+            self.self_attn = RelPositionMultiHeadAttention(
+                n_head=n_heads, n_feat=d_model, dropout_rate=dropout_att, pos_bias_u=pos_bias_u, pos_bias_v=pos_bias_v
+            )
+        elif self_attention_model == 'abs_pos':
+            self.self_attn = MultiHeadAttention(n_head=n_heads, n_feat=d_model, dropout_rate=dropout_att)
+        else:
+            raise ValueError(
+                f"'{self_attention_model}' is not not a valid value for 'self_attention_model', "
+                f"valid values can be from ['rel_pos', 'abs_pos']"
+            )
+
+        # second feed forward module
+        self.norm_feed_forward2 = LayerNorm(d_model)
+        self.feed_forward2 = ConformerFeedForward(d_model=d_model, d_ff=d_ff, dropout=dropout)
+
+        self.dropout = nn.Dropout(dropout)
+        self.norm_out = LayerNorm(d_model)
+
+    def forward(self, x, att_mask=None, pos_emb=None, pad_mask=None):
+        """
+        Args:
+            x (torch.Tensor): input signals (B, T, d_model)
+            att_mask (torch.Tensor): attention masks(B, T, T)
+            pos_emb (torch.Tensor): (L, 1, d_model)
+            pad_mask (torch.tensor): padding mask
+        Returns:
+            x (torch.Tensor): (B, T, d_model)
+        """
+        residual = x
+        x = self.norm_feed_forward1(x)
+        x = self.feed_forward1(x)
+        x = self.fc_factor * self.dropout(x) + residual
+
+        residual = x
+        x = self.norm_self_att(x)
+        if self.self_attention_model == 'rel_pos':
+            x = self.self_attn(query=x, key=x, value=x, mask=att_mask, pos_emb=pos_emb)
+        elif self.self_attention_model == 'abs_pos':
+            x = self.self_attn(query=x, key=x, value=x, mask=att_mask)
+        else:
+            x = None
+        x = self.dropout(x) + residual
+
+        residual = x
+        x = self.norm_conv(x)
+        x = self.conv(x, pad_mask)
+        x = self.dropout(x) + residual
+
+        residual = x
+        x = self.norm_feed_forward2(x)
+        x = self.feed_forward2(x)
+        x = self.fc_factor * self.dropout(x) + residual
+
+        x = self.norm_out(x)
+        return x
+
+
+class ConformerConvolution(nn.Module):
+    """The convolution module for the Conformer model.
+    Args:
+        d_model (int): hidden dimension
+        kernel_size (int): kernel size for depthwise convolution
+    """
+
+    def __init__(self, d_model, kernel_size):
+        super(ConformerConvolution, self).__init__()
+        assert (kernel_size - 1) % 2 == 0
+        self.d_model = d_model
+
+        self.pointwise_conv1 = nn.Conv1d(
+            in_channels=d_model, out_channels=d_model * 2, kernel_size=1, stride=1, padding=0, bias=True
+        )
+        self.depthwise_conv = nn.Conv1d(
+            in_channels=d_model,
+            out_channels=d_model,
+            kernel_size=kernel_size,
+            stride=1,
+            padding=(kernel_size - 1) // 2,
+            groups=d_model,
+            bias=True,
+        )
+        self.batch_norm = nn.BatchNorm1d(d_model)
+
+        self.activation = Swish()
+        self.pointwise_conv2 = nn.Conv1d(
+            in_channels=d_model, out_channels=d_model, kernel_size=1, stride=1, padding=0, bias=True
+        )
+
+    def forward(self, x, pad_mask=None):
+        x = x.transpose(1, 2)
+        x = self.pointwise_conv1(x)
+        x = nn.functional.glu(x, dim=1)
+
+        if pad_mask is not None:
+            x.masked_fill_(pad_mask.unsqueeze(1), 0.0)
+
+        x = self.depthwise_conv(x)
+        x = self.batch_norm(x)
+        x = self.activation(x)
+        x = self.pointwise_conv2(x)
+        x = x.transpose(1, 2)
+        return x
+
+
+class ConformerFeedForward(nn.Module):
+    """
+    feed-forward module of Conformer model.
+    """
+
+    def __init__(self, d_model, d_ff, dropout, activation=Swish()):
+        super(ConformerFeedForward, self).__init__()
+        self.linear1 = nn.Linear(d_model, d_ff)
+        self.activation = activation
+        self.dropout = nn.Dropout(p=dropout)
+        self.linear2 = nn.Linear(d_ff, d_model)
+
+    def forward(self, x):
+        x = self.linear1(x)
+        x = self.activation(x)
+        x = self.dropout(x)
+        x = self.linear2(x)
+        return x

tts/nemo/collections/asr/parts/submodules/jasper.py

@@ -0,0 +1,960 @@
+# 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 math
+from typing import Callable, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.nn.init import _calculate_correct_fan
+from torch.nn.modules.utils import _single
+
+from nemo.collections.asr.parts.utils.activations import Swish
+from nemo.utils import logging
+
+try:
+    from pytorch_quantization import calib
+    from pytorch_quantization import nn as quant_nn
+    from pytorch_quantization import quant_modules
+    from pytorch_quantization.tensor_quant import QuantDescriptor
+
+    PYTORCH_QUANTIZATION_AVAILABLE = True
+except ImportError:
+    PYTORCH_QUANTIZATION_AVAILABLE = False
+
+jasper_activations = {"hardtanh": nn.Hardtanh, "relu": nn.ReLU, "selu": nn.SELU, "swish": Swish, "silu": nn.SiLU}
+
+
+def tds_uniform_(tensor, mode='fan_in'):
+    """
+    Uniform Initialization from the paper [Sequence-to-Sequence Speech Recognition with Time-Depth Separable Convolutions](https://www.isca-speech.org/archive/Interspeech_2019/pdfs/2460.pdf)
+    Normalized to -
+
+    .. math::
+        \text{bound} = \text{2} \times \sqrt{\frac{1}{\text{fan\_mode}}}
+
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        mode: either ``'fan_in'`` (default) or ``'fan_out'``. Choosing ``'fan_in'``
+            preserves the magnitude of the variance of the weights in the
+            forward pass. Choosing ``'fan_out'`` preserves the magnitudes in the
+            backwards pass.
+    """
+    fan = _calculate_correct_fan(tensor, mode)
+    gain = 2.0  # sqrt(4.0) = 2
+    std = gain / math.sqrt(fan)  # sqrt(4.0 / fan_in)
+    bound = std  # Calculate uniform bounds from standard deviation
+    with torch.no_grad():
+        return tensor.uniform_(-bound, bound)
+
+
+def tds_normal_(tensor, mode='fan_in'):
+    """
+    Normal Initialization from the paper [Sequence-to-Sequence Speech Recognition with Time-Depth Separable Convolutions](https://www.isca-speech.org/archive/Interspeech_2019/pdfs/2460.pdf)
+    Normalized to -
+
+    .. math::
+        \text{bound} = \text{2} \times \sqrt{\frac{1}{\text{fan\_mode}}}
+
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        mode: either ``'fan_in'`` (default) or ``'fan_out'``. Choosing ``'fan_in'``
+            preserves the magnitude of the variance of the weights in the
+            forward pass. Choosing ``'fan_out'`` preserves the magnitudes in the
+            backwards pass.
+    """
+    fan = _calculate_correct_fan(tensor, mode)
+    gain = 2.0
+    std = gain / math.sqrt(fan)  # sqrt(4.0 / fan_in)
+    bound = std  # Calculate uniform bounds from standard deviation
+    with torch.no_grad():
+        return tensor.normal_(0.0, bound)
+
+
+def init_weights(m, mode: Optional[str] = 'xavier_uniform'):
+    if isinstance(m, MaskedConv1d):
+        init_weights(m.conv, mode)
+    if isinstance(m, (nn.Conv1d, nn.Linear)):
+        if mode is not None:
+            if mode == 'xavier_uniform':
+                nn.init.xavier_uniform_(m.weight, gain=1.0)
+            elif mode == 'xavier_normal':
+                nn.init.xavier_normal_(m.weight, gain=1.0)
+            elif mode == 'kaiming_uniform':
+                nn.init.kaiming_uniform_(m.weight, nonlinearity="relu")
+            elif mode == 'kaiming_normal':
+                nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
+            elif mode == 'tds_uniform':
+                tds_uniform_(m.weight)
+            elif mode == 'tds_normal':
+                tds_normal_(m.weight)
+            else:
+                raise ValueError("Unknown Initialization mode: {0}".format(mode))
+    elif isinstance(m, nn.BatchNorm1d):
+        if m.track_running_stats:
+            m.running_mean.zero_()
+            m.running_var.fill_(1)
+            m.num_batches_tracked.zero_()
+        if m.affine:
+            nn.init.ones_(m.weight)
+            nn.init.zeros_(m.bias)
+
+
+def compute_new_kernel_size(kernel_size, kernel_width):
+    new_kernel_size = max(int(kernel_size * kernel_width), 1)
+    # If kernel is even shape, round up to make it odd
+    if new_kernel_size % 2 == 0:
+        new_kernel_size += 1
+    return new_kernel_size
+
+
+def get_same_padding(kernel_size, stride, dilation) -> int:
+    if stride > 1 and dilation > 1:
+        raise ValueError("Only stride OR dilation may be greater than 1")
+    return (dilation * (kernel_size - 1)) // 2
+
+
+def get_asymtric_padding(kernel_size, stride, dilation, future_context):
+    if stride > 1 and dilation > 1:
+        raise ValueError("Only stride OR dilation may be greater than 1")
+
+    left_context = kernel_size - 1 - future_context
+    right_context = future_context
+
+    symmetric_padding = get_same_padding(kernel_size, stride, dilation)
+
+    if kernel_size <= future_context:
+        # kernel size is smaller than future context, equivalent to using entire context of kernel
+        # simply return symmetric padding for this scenario
+        logging.warning(
+            f"Future context window is larger than the kernel size!\n"
+            f"Left context = {left_context} | Right context = greater than {right_context} | "
+            f"Kernel size = {kernel_size}\n"
+            f"Switching to symmetric padding (left context = right context = {symmetric_padding})"
+        )
+        return symmetric_padding
+
+    if left_context < symmetric_padding:
+        logging.warning(
+            f"Future context window is larger than half the kernel size!\n"
+            f"Conv layer therefore uses more future information than past to compute its output!\n"
+            f"Left context = {left_context} | Right context = {right_context} | "
+            f"Kernel size = {kernel_size}"
+        )
+
+    if dilation > 1:
+        left_context = dilation * kernel_size - 1 - dilation * future_context
+        right_context = dilation * future_context
+        return (left_context, right_context)
+
+    return (left_context, right_context)
+
+
+class StatsPoolLayer(nn.Module):
+    def __init__(self, feat_in, pool_mode='xvector'):
+        super().__init__()
+        self.feat_in = 0
+        if pool_mode == 'gram':
+            gram = True
+            super_vector = False
+        elif pool_mode == 'superVector':
+            gram = True
+            super_vector = True
+        else:
+            gram = False
+            super_vector = False
+
+        if gram:
+            self.feat_in += feat_in ** 2
+        else:
+            self.feat_in += 2 * feat_in
+
+        if super_vector and gram:
+            self.feat_in += 2 * feat_in
+
+        self.gram = gram
+        self.super = super_vector
+
+    def forward(self, encoder_output):
+
+        mean = encoder_output.mean(dim=-1)  # Time Axis
+        std = encoder_output.std(dim=-1)
+
+        pooled = torch.cat([mean, std], dim=-1)
+
+        if self.gram:
+            time_len = encoder_output.shape[-1]
+            # encoder_output = encoder_output
+            cov = encoder_output.bmm(encoder_output.transpose(2, 1))  # cov matrix
+            cov = cov.view(cov.shape[0], -1) / time_len
+
+        if self.gram and not self.super:
+            return cov
+
+        if self.super and self.gram:
+            pooled = torch.cat([pooled, cov], dim=-1)
+
+        return pooled
+
+
+class MaskedConv1d(nn.Module):
+    __constants__ = ["use_conv_mask", "real_out_channels", "heads"]
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        padding=0,
+        dilation=1,
+        groups=1,
+        heads=-1,
+        bias=False,
+        use_mask=True,
+        quantize=False,
+    ):
+        super(MaskedConv1d, self).__init__()
+
+        if not (heads == -1 or groups == in_channels):
+            raise ValueError("Only use heads for depthwise convolutions")
+
+        self.real_out_channels = out_channels
+        if heads != -1:
+            in_channels = heads
+            out_channels = heads
+            groups = heads
+
+        # preserve original padding
+        self._padding = padding
+
+        # if padding is a tuple/list, it is considered as asymmetric padding
+        if type(padding) in (tuple, list):
+            self.pad_layer = nn.ConstantPad1d(padding, value=0.0)
+            # reset padding for conv since pad_layer will handle this
+            padding = 0
+        else:
+            self.pad_layer = None
+
+        if PYTORCH_QUANTIZATION_AVAILABLE and quantize:
+            self.conv = quant_nn.QuantConv1d(
+                in_channels,
+                out_channels,
+                kernel_size,
+                stride=stride,
+                padding=padding,
+                dilation=dilation,
+                groups=groups,
+                bias=bias,
+            )
+        elif not PYTORCH_QUANTIZATION_AVAILABLE and quantize:
+            raise ImportError(
+                "pytorch-quantization is not installed. Install from "
+                "https://github.com/NVIDIA/TensorRT/tree/master/tools/pytorch-quantization."
+            )
+        else:
+            self.conv = nn.Conv1d(
+                in_channels,
+                out_channels,
+                kernel_size,
+                stride=stride,
+                padding=padding,
+                dilation=dilation,
+                groups=groups,
+                bias=bias,
+            )
+        self.use_mask = use_mask
+        self.heads = heads
+
+        # Calculations for "same" padding cache
+        self.same_padding = (self.conv.stride[0] == 1) and (
+            2 * self.conv.padding[0] == self.conv.dilation[0] * (self.conv.kernel_size[0] - 1)
+        )
+        if self.pad_layer is None:
+            self.same_padding_asymmetric = False
+        else:
+            self.same_padding_asymmetric = (self.conv.stride[0] == 1) and (
+                sum(self._padding) == self.conv.dilation[0] * (self.conv.kernel_size[0] - 1)
+            )
+
+        # `self.lens` caches consecutive integers from 0 to `self.max_len` that are used to compute the mask for a
+        # batch. Recomputed to bigger size as needed. Stored on a device of the latest batch lens.
+        if self.use_mask:
+            self.max_len = 0
+            self.lens = None
+
+    def get_seq_len(self, lens):
+        if self.same_padding or self.same_padding_asymmetric:
+            return lens
+
+        if self.pad_layer is None:
+            return (
+                lens + 2 * self.conv.padding[0] - self.conv.dilation[0] * (self.conv.kernel_size[0] - 1) - 1
+            ) // self.conv.stride[0] + 1
+        else:
+            return (
+                lens + sum(self._padding) - self.conv.dilation[0] * (self.conv.kernel_size[0] - 1) - 1
+            ) // self.conv.stride[0] + 1
+
+    def forward(self, x, lens):
+        if self.use_mask:
+            max_len = x.size(2)
+            if max_len > self.max_len:
+                self.lens = torch.arange(max_len)
+                self.max_len = max_len
+
+            self.lens = self.lens.to(lens.device)
+            mask = self.lens[:max_len].unsqueeze(0) < lens.unsqueeze(1)
+            x = x * mask.unsqueeze(1).to(device=x.device)
+            lens = self.get_seq_len(lens)
+
+        # asymmtric pad if necessary
+        if self.pad_layer is not None:
+            x = self.pad_layer(x)
+
+        sh = x.shape
+        if self.heads != -1:
+            x = x.view(-1, self.heads, sh[-1])
+
+        out = self.conv(x)
+
+        if self.heads != -1:
+            out = out.view(sh[0], self.real_out_channels, -1)
+
+        return out, lens
+
+
+class GroupShuffle(nn.Module):
+    def __init__(self, groups, channels):
+        super(GroupShuffle, self).__init__()
+
+        self.groups = groups
+        self.channels_per_group = channels // groups
+
+    def forward(self, x):
+        sh = x.shape
+
+        x = x.view(-1, self.groups, self.channels_per_group, sh[-1])
+
+        x = torch.transpose(x, 1, 2).contiguous()
+
+        x = x.view(-1, self.groups * self.channels_per_group, sh[-1])
+
+        return x
+
+
+class SqueezeExcite(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        reduction_ratio: int,
+        context_window: int = -1,
+        interpolation_mode: str = 'nearest',
+        activation: Optional[Callable] = None,
+        quantize: bool = False,
+    ):
+        """
+        Squeeze-and-Excitation sub-module.
+
+        Args:
+            channels: Input number of channels.
+            reduction_ratio: Reduction ratio for "squeeze" layer.
+            context_window: Integer number of timesteps that the context
+                should be computed over, using stride 1 average pooling.
+                If value < 1, then global context is computed.
+            interpolation_mode: Interpolation mode of timestep dimension.
+                Used only if context window is > 1.
+                The modes available for resizing are: `nearest`, `linear` (3D-only),
+                `bilinear`, `area`
+            activation: Intermediate activation function used. Must be a
+                callable activation function.
+        """
+        super(SqueezeExcite, self).__init__()
+        self.interpolation_mode = interpolation_mode
+        self._quantize = quantize
+
+        self.pool = None  # prepare a placeholder which will be updated
+        self.change_context_window(context_window=context_window)
+
+        if activation is None:
+            activation = nn.ReLU(inplace=True)
+
+        if PYTORCH_QUANTIZATION_AVAILABLE and quantize:
+            self.fc = nn.Sequential(
+                quant_nn.QuantLinear(channels, channels // reduction_ratio, bias=False),
+                activation,
+                quant_nn.QuantLinear(channels // reduction_ratio, channels, bias=False),
+            )
+        elif not PYTORCH_QUANTIZATION_AVAILABLE and quantize:
+            raise ImportError(
+                "pytorch-quantization is not installed. Install from "
+                "https://github.com/NVIDIA/TensorRT/tree/master/tools/pytorch-quantization."
+            )
+        else:
+            self.fc = nn.Sequential(
+                nn.Linear(channels, channels // reduction_ratio, bias=False),
+                activation,
+                nn.Linear(channels // reduction_ratio, channels, bias=False),
+            )
+        self.gap = nn.AdaptiveAvgPool1d(1)
+
+    def forward(self, x):
+        # The use of negative indices on the transpose allow for expanded SqueezeExcite
+        batch, channels, timesteps = x.size()[:3]
+        # Computes in float32 to avoid instabilities during training with AMP.
+        with torch.cuda.amp.autocast(enabled=False):
+            x = x.float()
+            if timesteps < self.context_window:
+                y = self.gap(x)
+            else:
+                y = self.pool(x)  # [B, C, T - context_window + 1]
+            y = y.transpose(1, -1)  # [B, T - context_window + 1, C]
+            y = self.fc(y)  # [B, T - context_window + 1, C]
+            y = y.transpose(1, -1)  # [B, C, T - context_window + 1]
+
+        if self.context_window > 0:
+            y = torch.nn.functional.interpolate(y, size=timesteps, mode=self.interpolation_mode)
+
+        y = torch.sigmoid(y)
+        return x * y
+
+    def change_context_window(self, context_window: int):
+        """
+        Update the context window of the SqueezeExcitation module, in-place if possible.
+
+        Will update the pooling layer to either nn.AdaptiveAvgPool1d() (for global SE) or nn.AvgPool1d()
+        (for limited context SE).
+
+        If only the context window is changing but still a limited SE context block - then
+        the earlier instance of nn.AvgPool1d() will be updated.
+
+        Args:
+            context_window: An integer representing the number of input timeframes that will be used
+                to compute the context. Each timeframe corresponds to a single window stride of the
+                STFT features.
+
+                Say the window_stride = 0.01s, then a context window of 128 represents 128 * 0.01 s
+                of context to compute the Squeeze step.
+        """
+        if hasattr(self, 'context_window'):
+            logging.info(f"Changing Squeeze-Excitation context window from {self.context_window} to {context_window}")
+
+        self.context_window = int(context_window)
+
+        if self.context_window <= 0:
+            if PYTORCH_QUANTIZATION_AVAILABLE and self._quantize:
+                if not isinstance(self.pool, quant_nn.QuantAdaptiveAvgPool1d(1)):
+                    self.pool = quant_nn.QuantAdaptiveAvgPool1d(1)  # context window = T
+
+            elif not PYTORCH_QUANTIZATION_AVAILABLE and self._quantize:
+                raise ImportError(
+                    "pytorch-quantization is not installed. Install from "
+                    "https://github.com/NVIDIA/TensorRT/tree/master/tools/pytorch-quantization."
+                )
+
+            else:
+                if not isinstance(self.pool, nn.AdaptiveAvgPool1d):
+                    self.pool = nn.AdaptiveAvgPool1d(1)  # context window = T
+        else:
+            if PYTORCH_QUANTIZATION_AVAILABLE and self._quantize:
+                if not isinstance(self.pool, quant_nn.QuantAvgPool1d):
+                    self.pool = quant_nn.QuantAvgPool1d(self.context_window, stride=1)
+
+            elif not PYTORCH_QUANTIZATION_AVAILABLE and self._quantize:
+                raise ImportError(
+                    "pytorch-quantization is not installed. Install from "
+                    "https://github.com/NVIDIA/TensorRT/tree/master/tools/pytorch-quantization."
+                )
+
+            else:
+                if not isinstance(self.pool, nn.AvgPool1d):
+                    self.pool = nn.AvgPool1d(self.context_window, stride=1)
+                else:
+                    # update the context window
+                    self.pool.kernel_size = _single(self.context_window)
+
+
+class JasperBlock(nn.Module):
+    """
+    Constructs a single "Jasper" block. With modified parameters, also constructs other blocks for models
+    such as `QuartzNet` and `Citrinet`.
+
+    - For `Jasper`    : `separable` flag should be False
+    - For `QuartzNet` : `separable` flag should be True
+    - For `Citrinet`  : `separable` flag and `se` flag should be True
+
+    Note that above are general distinctions, each model has intricate differences that expand over
+    multiple such blocks.
+
+    For further information about the differences between models which use JasperBlock, please review
+    the configs for ASR models found in the ASR examples directory.
+
+    Args:
+        inplanes: Number of input channels.
+        planes: Number of output channels.
+        repeat: Number of repeated sub-blocks (R) for this block.
+        kernel_size: Convolution kernel size across all repeated sub-blocks.
+        kernel_size_factor: Floating point scale value that is multiplied with kernel size,
+            then rounded down to nearest odd integer to compose the kernel size. Defaults to 1.0.
+        stride: Stride of the convolutional layers.
+        dilation: Integer which defined dilation factor of kernel. Note that when dilation > 1, stride must
+            be equal to 1.
+        padding: String representing type of padding. Currently only supports "same" padding,
+            which symmetrically pads the input tensor with zeros.
+        dropout: Floating point value, determins percentage of output that is zeroed out.
+        activation: String representing activation functions. Valid activation functions are :
+            {"hardtanh": nn.Hardtanh, "relu": nn.ReLU, "selu": nn.SELU, "swish": Swish}.
+            Defaults to "relu".
+        residual: Bool that determined whether a residual branch should be added or not.
+            All residual branches are constructed using a pointwise convolution kernel, that may or may not
+            perform strided convolution depending on the parameter `residual_mode`.
+        groups: Number of groups for Grouped Convolutions. Defaults to 1.
+        separable: Bool flag that describes whether Time-Channel depthwise separable convolution should be
+            constructed, or ordinary convolution should be constructed.
+        heads: Number of "heads" for the masked convolution. Defaults to -1, which disables it.
+        normalization: String that represents type of normalization performed. Can be one of
+            "batch", "group", "instance" or "layer" to compute BatchNorm1D, GroupNorm1D, InstanceNorm or
+            LayerNorm (which are special cases of GroupNorm1D).
+        norm_groups: Number of groups used for GroupNorm (if `normalization` == "group").
+        residual_mode: String argument which describes whether the residual branch should be simply
+            added ("add") or should first stride, then add ("stride_add"). Required when performing stride on
+            parallel branch as well as utilizing residual add.
+        residual_panes: Number of residual panes, used for Jasper-DR models. Please refer to the paper.
+        conv_mask: Bool flag which determines whether to utilize masked convolutions or not. In general,
+            it should be set to True.
+        se: Bool flag that determines whether Squeeze-and-Excitation layer should be used.
+        se_reduction_ratio: Integer value, which determines to what extend the hidden dimension of the SE
+            intermediate step should be reduced. Larger values reduce number of parameters, but also limit
+            the effectiveness of SE layers.
+        se_context_window: Integer value determining the number of timesteps that should be utilized in order
+            to compute the averaged context window. Defaults to -1, which means it uses global context - such
+            that all timesteps are averaged. If any positive integer is used, it will utilize limited context
+            window of that size.
+        se_interpolation_mode: String used for interpolation mode of timestep dimension for SE blocks.
+            Used only if context window is > 1.
+            The modes available for resizing are: `nearest`, `linear` (3D-only),
+            `bilinear`, `area`.
+        stride_last: Bool flag that determines whether all repeated blocks should stride at once,
+            (stride of S^R when this flag is False) or just the last repeated block should stride
+            (stride of S when this flag is True).
+        future_context: Int value that determins how many "right" / "future" context frames will be utilized
+            when calculating the output of the conv kernel. All calculations are done for odd kernel sizes only.
+
+            By default, this is -1, which is recomputed as the symmetric padding case.
+
+            When future_context >= 0, will compute the asymmetric padding as follows :
+            (left context, right context) = [K - 1 - future_context, future_context]
+
+            Determining an exact formula to limit future context is dependent on global layout of the model.
+            As such, we provide both "local" and "global" guidelines below.
+
+            Local context limit (should always be enforced)
+            - future context should be <= half the kernel size for any given layer
+            - future context > kernel size defaults to symmetric kernel
+            - future context of layer = number of future frames * width of each frame (dependent on stride)
+
+            Global context limit (should be carefully considered)
+            - future context should be layed out in an ever reducing pattern. Initial layers should restrict
+            future context less than later layers, since shallow depth (and reduced stride) means each frame uses
+            less amounts of future context.
+            - Beyond a certain point, future context should remain static for a given stride level. This is
+            the upper bound of the amount of future context that can be provided to the model on a global scale.
+            - future context is calculated (roughly) as - (2 ^ stride) * (K // 2) number of future frames.
+            This resultant value should be bound to some global maximum number of future seconds of audio (in ms).
+
+            Note: In the special case where K < future_context, it is assumed that the kernel is too small to limit
+            its future context, so symmetric padding is used instead.
+
+            Note: There is no explicit limitation on the amount of future context used, as long as
+            K > future_context constraint is maintained. This might lead to cases where future_context is
+            more than half the actual kernel size K! In such cases, the conv layer is utilizing more of the future
+            context than its current and past context to compute the output. While this is possible to do,
+            it is not recommended and the layer will raise a warning to notify the user of such cases.
+            It is advised to simply use symmetric padding for such cases.
+
+            Example:
+            Say we have a model that performs 8x stride and receives spectrogram frames with stride of 0.01s.
+            Say we wish to upper bound future context to 80 ms.
+
+            Layer ID, Kernel Size, Stride, Future Context, Global Context
+            0, K=5,  S=1, FC=8, GC= 2 * (2^0) = 2 * 0.01 ms  (special case, K < FC so use symmetric pad)
+            1, K=7,  S=1, FC=3, GC= 3 * (2^0) = 3 * 0.01 ms  (note that symmetric pad here uses 3 FC frames!)
+            2, K=11, S=2, FC=4, GC= 4 * (2^1) = 8 * 0.01 ms  (note that symmetric pad here uses 5 FC frames!)
+            3, K=15, S=1, FC=4, GC= 4 * (2^1) = 8 * 0.01 ms  (note that symmetric pad here uses 7 FC frames!)
+            4, K=21, S=2, FC=2, GC= 2 * (2^2) = 8 * 0.01 ms  (note that symmetric pad here uses 10 FC frames!)
+            5, K=25, S=2, FC=1, GC= 1 * (2^3) = 8 * 0.01 ms  (note that symmetric pad here uses 14 FC frames!)
+            6, K=29, S=1, FC=1, GC= 1 * (2^3) = 8 * 0.01 ms ...
+        quantize: Bool flag whether to quantize the Convolutional blocks.
+    """
+
+    __constants__ = ["conv_mask", "separable", "residual_mode", "res", "mconv"]
+
+    def __init__(
+        self,
+        inplanes,
+        planes,
+        repeat=3,
+        kernel_size=11,
+        kernel_size_factor=1,
+        stride=1,
+        dilation=1,
+        padding='same',
+        dropout=0.2,
+        activation=None,
+        residual=True,
+        groups=1,
+        separable=False,
+        heads=-1,
+        normalization="batch",
+        norm_groups=1,
+        residual_mode='add',
+        residual_panes=[],
+        conv_mask=False,
+        se=False,
+        se_reduction_ratio=16,
+        se_context_window=-1,
+        se_interpolation_mode='nearest',
+        stride_last=False,
+        future_context: int = -1,
+        quantize=False,
+    ):
+        super(JasperBlock, self).__init__()
+
+        if padding != "same":
+            raise ValueError("currently only 'same' padding is supported")
+
+        kernel_size_factor = float(kernel_size_factor)
+        if type(kernel_size) in (list, tuple):
+            kernel_size = [compute_new_kernel_size(k, kernel_size_factor) for k in kernel_size]
+        else:
+            kernel_size = compute_new_kernel_size(kernel_size, kernel_size_factor)
+
+        if future_context < 0:
+            padding_val = get_same_padding(kernel_size[0], stride[0], dilation[0])
+        else:
+            padding_val = get_asymtric_padding(kernel_size[0], stride[0], dilation[0], future_context)
+
+        self.conv_mask = conv_mask
+        self.separable = separable
+        self.residual_mode = residual_mode
+        self.se = se
+        self.quantize = quantize
+
+        inplanes_loop = inplanes
+        conv = nn.ModuleList()
+
+        for _ in range(repeat - 1):
+            # Stride last means only the last convolution in block will have stride
+            if stride_last:
+                stride_val = [1]
+            else:
+                stride_val = stride
+
+            conv.extend(
+                self._get_conv_bn_layer(
+                    inplanes_loop,
+                    planes,
+                    kernel_size=kernel_size,
+                    stride=stride_val,
+                    dilation=dilation,
+                    padding=padding_val,
+                    groups=groups,
+                    heads=heads,
+                    separable=separable,
+                    normalization=normalization,
+                    norm_groups=norm_groups,
+                    quantize=quantize,
+                )
+            )
+
+            conv.extend(self._get_act_dropout_layer(drop_prob=dropout, activation=activation))
+
+            inplanes_loop = planes
+
+        conv.extend(
+            self._get_conv_bn_layer(
+                inplanes_loop,
+                planes,
+                kernel_size=kernel_size,
+                stride=stride,
+                dilation=dilation,
+                padding=padding_val,
+                groups=groups,
+                heads=heads,
+                separable=separable,
+                normalization=normalization,
+                norm_groups=norm_groups,
+                quantize=quantize,
+            )
+        )
+
+        if se:
+            conv.append(
+                SqueezeExcite(
+                    planes,
+                    reduction_ratio=se_reduction_ratio,
+                    context_window=se_context_window,
+                    interpolation_mode=se_interpolation_mode,
+                    activation=activation,
+                    quantize=quantize,
+                )
+            )
+
+        self.mconv = conv
+
+        res_panes = residual_panes.copy()
+        self.dense_residual = residual
+
+        if residual:
+            res_list = nn.ModuleList()
+
+            if residual_mode == 'stride_add':
+                stride_val = stride
+            else:
+                stride_val = [1]
+
+            if len(residual_panes) == 0:
+                res_panes = [inplanes]
+                self.dense_residual = False
+            for ip in res_panes:
+                res = nn.ModuleList(
+                    self._get_conv_bn_layer(
+                        ip,
+                        planes,
+                        kernel_size=1,
+                        normalization=normalization,
+                        norm_groups=norm_groups,
+                        stride=stride_val,
+                        quantize=quantize,
+                    )
+                )
+
+                res_list.append(res)
+
+            self.res = res_list
+            if PYTORCH_QUANTIZATION_AVAILABLE and self.quantize:
+                self.residual_quantizer = quant_nn.TensorQuantizer(quant_nn.QuantConv2d.default_quant_desc_input)
+            elif not PYTORCH_QUANTIZATION_AVAILABLE and quantize:
+                raise ImportError(
+                    "pytorch-quantization is not installed. Install from "
+                    "https://github.com/NVIDIA/TensorRT/tree/master/tools/pytorch-quantization."
+                )
+        else:
+            self.res = None
+
+        self.mout = nn.Sequential(*self._get_act_dropout_layer(drop_prob=dropout, activation=activation))
+
+    def _get_conv(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size=11,
+        stride=1,
+        dilation=1,
+        padding=0,
+        bias=False,
+        groups=1,
+        heads=-1,
+        separable=False,
+        quantize=False,
+    ):
+        use_mask = self.conv_mask
+        if use_mask:
+            return MaskedConv1d(
+                in_channels,
+                out_channels,
+                kernel_size,
+                stride=stride,
+                dilation=dilation,
+                padding=padding,
+                bias=bias,
+                groups=groups,
+                heads=heads,
+                use_mask=use_mask,
+                quantize=quantize,
+            )
+        else:
+            if PYTORCH_QUANTIZATION_AVAILABLE and quantize:
+                return quant_nn.QuantConv1d(
+                    in_channels,
+                    out_channels,
+                    kernel_size,
+                    stride=stride,
+                    dilation=dilation,
+                    padding=padding,
+                    bias=bias,
+                    groups=groups,
+                )
+            elif not PYTORCH_QUANTIZATION_AVAILABLE and quantize:
+                raise ImportError(
+                    "pytorch-quantization is not installed. Install from "
+                    "https://github.com/NVIDIA/TensorRT/tree/master/tools/pytorch-quantization."
+                )
+            else:
+                return nn.Conv1d(
+                    in_channels,
+                    out_channels,
+                    kernel_size,
+                    stride=stride,
+                    dilation=dilation,
+                    padding=padding,
+                    bias=bias,
+                    groups=groups,
+                )
+
+    def _get_conv_bn_layer(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size=11,
+        stride=1,
+        dilation=1,
+        padding=0,
+        bias=False,
+        groups=1,
+        heads=-1,
+        separable=False,
+        normalization="batch",
+        norm_groups=1,
+        quantize=False,
+    ):
+        if norm_groups == -1:
+            norm_groups = out_channels
+
+        if separable:
+            layers = [
+                self._get_conv(
+                    in_channels,
+                    in_channels,
+                    kernel_size,
+                    stride=stride,
+                    dilation=dilation,
+                    padding=padding,
+                    bias=bias,
+                    groups=in_channels,
+                    heads=heads,
+                    quantize=quantize,
+                ),
+                self._get_conv(
+                    in_channels,
+                    out_channels,
+                    kernel_size=1,
+                    stride=1,
+                    dilation=1,
+                    padding=0,
+                    bias=bias,
+                    groups=groups,
+                    quantize=quantize,
+                ),
+            ]
+        else:
+            layers = [
+                self._get_conv(
+                    in_channels,
+                    out_channels,
+                    kernel_size,
+                    stride=stride,
+                    dilation=dilation,
+                    padding=padding,
+                    bias=bias,
+                    groups=groups,
+                    quantize=quantize,
+                )
+            ]
+
+        if normalization == "group":
+            layers.append(nn.GroupNorm(num_groups=norm_groups, num_channels=out_channels))
+        elif normalization == "instance":
+            layers.append(nn.GroupNorm(num_groups=out_channels, num_channels=out_channels))
+        elif normalization == "layer":
+            layers.append(nn.GroupNorm(num_groups=1, num_channels=out_channels))
+        elif normalization == "batch":
+            layers.append(nn.BatchNorm1d(out_channels, eps=1e-3, momentum=0.1))
+        else:
+            raise ValueError(
+                f"Normalization method ({normalization}) does not match" f" one of [batch, layer, group, instance]."
+            )
+
+        if groups > 1:
+            layers.append(GroupShuffle(groups, out_channels))
+        return layers
+
+    def _get_act_dropout_layer(self, drop_prob=0.2, activation=None):
+        if activation is None:
+            activation = nn.Hardtanh(min_val=0.0, max_val=20.0)
+        layers = [activation, nn.Dropout(p=drop_prob)]
+        return layers
+
+    def forward(self, input_: Tuple[List[Tensor], Optional[Tensor]]):
+        """
+        Forward pass of the module.
+
+        Args:
+            input_: The input is a tuple of two values - the preprocessed audio signal as well as the lengths
+                of the audio signal. The audio signal is padded to the shape [B, D, T] and the lengths are
+                a torch vector of length B.
+
+        Returns:
+            The output of the block after processing the input through `repeat` number of sub-blocks,
+            as well as the lengths of the encoded audio after padding/striding.
+        """
+        # type: (Tuple[List[Tensor], Optional[Tensor]]) -> Tuple[List[Tensor], Optional[Tensor]] # nopep8
+        lens_orig = None
+        xs = input_[0]
+        if len(input_) == 2:
+            xs, lens_orig = input_
+
+        # compute forward convolutions
+        out = xs[-1]
+
+        lens = lens_orig
+        for i, l in enumerate(self.mconv):
+            # if we're doing masked convolutions, we need to pass in and
+            # possibly update the sequence lengths
+            # if (i % 4) == 0 and self.conv_mask:
+            if isinstance(l, MaskedConv1d):
+                out, lens = l(out, lens)
+            else:
+                out = l(out)
+
+        # compute the residuals
+        if self.res is not None:
+            for i, layer in enumerate(self.res):
+                res_out = xs[i]
+                for j, res_layer in enumerate(layer):
+                    if isinstance(res_layer, MaskedConv1d):
+                        res_out, _ = res_layer(res_out, lens_orig)
+                    else:
+                        res_out = res_layer(res_out)
+
+                if self.residual_mode == 'add' or self.residual_mode == 'stride_add':
+                    if PYTORCH_QUANTIZATION_AVAILABLE and self.quantize:
+                        out = self.residual_quantizer(out) + res_out
+                    elif not PYTORCH_QUANTIZATION_AVAILABLE and self.quantize:
+                        raise ImportError(
+                            "pytorch-quantization is not installed. Install from "
+                            "https://github.com/NVIDIA/TensorRT/tree/master/tools/pytorch-quantization."
+                        )
+                    else:
+                        out = out + res_out
+                else:
+                    out = torch.max(out, res_out)
+
+        # compute the output
+        out = self.mout(out)
+        if self.res is not None and self.dense_residual:
+            return xs + [out], lens
+
+        return [out], lens

tts/nemo/collections/asr/parts/submodules/multi_head_attention.py

@@ -0,0 +1,340 @@
+# 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 2017 Johns Hopkins University (Shinji Watanabe)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+Part of this code is adopted from https://github.com/espnet/espnet
+"""
+
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+__all__ = [
+    'RelPositionMultiHeadAttention',
+    'RelPositionalEncoding',
+    'PositionalEncoding',
+]
+
+
+class MultiHeadAttention(nn.Module):
+    """Multi-Head Attention layer of Transformer.
+    Args:
+        n_head (int): number of heads
+        n_feat (int): size of the features
+        dropout_rate (float): dropout rate
+    """
+
+    def __init__(self, n_head, n_feat, dropout_rate):
+        """Construct an MultiHeadedAttention object."""
+        super(MultiHeadAttention, self).__init__()
+        assert n_feat % n_head == 0
+        # We assume d_v always equals d_k
+        self.d_k = n_feat // n_head
+        self.h = n_head
+        self.linear_q = nn.Linear(n_feat, n_feat)
+        self.linear_k = nn.Linear(n_feat, n_feat)
+        self.linear_v = nn.Linear(n_feat, n_feat)
+        self.linear_out = nn.Linear(n_feat, n_feat)
+        self.attn = None
+        self.dropout = nn.Dropout(p=dropout_rate)
+
+    def forward_qkv(self, query, key, value):
+        """Transforms query, key and value.
+        Args:
+            query (torch.Tensor): (batch, time1, size)
+            key (torch.Tensor): (batch, time2, size)
+            value (torch.Tensor): (batch, time2, size)
+        returns:
+            q (torch.Tensor): (batch, head, time1, size)
+            k (torch.Tensor): (batch, head, time2, size)
+            v (torch.Tensor): (batch, head, time2, size)
+        """
+        n_batch = query.size(0)
+        q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k)
+        k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k)
+        v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k)
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+
+        return q, k, v
+
+    def forward_attention(self, value, scores, mask):
+        """Compute attention context vector.
+        Args:
+            value (torch.Tensor): (batch, time2, size)
+            scores(torch.Tensor): (batch, time1, time2)
+            mask(torch.Tensor): (batch, time1, time2)
+        returns:
+            value (torch.Tensor): transformed `value` (batch, time2, d_model) weighted by the attention scores
+        """
+        n_batch = value.size(0)
+        if mask is not None:
+            mask = mask.unsqueeze(1)  # (batch, 1, time1, time2)
+            if scores.dtype == torch.float16:
+                dtype = np.float16
+            else:
+                dtype = np.float32
+            min_value = np.finfo(dtype).min
+            scores = scores.masked_fill(mask, min_value)
+            self.attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0)  # (batch, head, time1, time2)
+        else:
+            self.attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)
+
+        p_attn = self.dropout(self.attn)
+        x = torch.matmul(p_attn, value)  # (batch, head, time1, d_k)
+        x = x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)  # (batch, time1, d_model)
+
+        return self.linear_out(x)  # (batch, time1, d_model)
+
+    def forward(self, query, key, value, mask, pos_emb=None):
+        """Compute 'Scaled Dot Product Attention'.
+        Args:
+            query (torch.Tensor): (batch, time1, size)
+            key (torch.Tensor): (batch, time2, size)
+            value(torch.Tensor): (batch, time2, size)
+            mask (torch.Tensor): (batch, time1, time2)
+        returns:
+            output (torch.Tensor): transformed `value` (batch, time1, d_model) weighted by the query dot key attention
+        """
+        q, k, v = self.forward_qkv(query, key, value)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
+        return self.forward_attention(v, scores, mask)
+
+
+class RelPositionMultiHeadAttention(MultiHeadAttention):
+    """Multi-Head Attention layer of Transformer-XL with support of relative positional encoding.
+    Paper: https://arxiv.org/abs/1901.02860
+    Args:
+        n_head (int): number of heads
+        n_feat (int): size of the features
+        dropout_rate (float): dropout rate
+    """
+
+    def __init__(self, n_head, n_feat, dropout_rate, pos_bias_u, pos_bias_v):
+        """Construct an RelPositionMultiHeadedAttention object."""
+        super().__init__(n_head, n_feat, dropout_rate)
+        # linear transformation for positional encoding
+        self.linear_pos = nn.Linear(n_feat, n_feat, bias=False)
+        # these two learnable biases are used in matrix c and matrix d
+        # as described in https://arxiv.org/abs/1901.02860 Section 3.3
+        if pos_bias_u is None or pos_bias_v is None:
+            self.pos_bias_u = nn.Parameter(torch.FloatTensor(self.h, self.d_k))
+            self.pos_bias_v = nn.Parameter(torch.FloatTensor(self.h, self.d_k))
+            # nn.init.normal_(self.pos_bias_u, 0.0, 0.02)
+            # nn.init.normal_(self.pos_bias_v, 0.0, 0.02)
+            nn.init.zeros_(self.pos_bias_u)
+            nn.init.zeros_(self.pos_bias_v)
+        else:
+            self.pos_bias_u = pos_bias_u
+            self.pos_bias_v = pos_bias_v
+
+    def rel_shift(self, x):
+        """Compute relative positional encoding.
+        Args:
+            x (torch.Tensor): (batch, nheads, time, 2*time-1)
+        """
+        x_size = x.size()  # (b, h, t1, t2)
+        qlen = x_size[-2]
+        pos_len = x_size[-1]
+
+        # need to add a column of zeros on the left side of last two dimensions to perform the relative shifting
+        x = torch.nn.functional.pad(x, pad=(1, 0, 0, 0))  # (b, h, t1, t2+1)
+        x = x.view(*x.size()[:-2], pos_len + 1, qlen)  # (b, h, t2+1, t1)
+        # need to drop the first row
+        x = x[:, :, 1:].view(*x_size)  # (b, h, t1, t2)
+        return x
+
+    def forward(self, query, key, value, mask, pos_emb):
+        """Compute 'Scaled Dot Product Attention' with rel. positional encoding.
+        Args:
+            query (torch.Tensor): (batch, time1, size)
+            key (torch.Tensor): (batch, time2, size)
+            value(torch.Tensor): (batch, time2, size)
+            mask (torch.Tensor): (batch, time1, time2)
+            pos_emb (torch.Tensor) : (batch, time1, size)
+        Returns:
+            output (torch.Tensor): transformed `value` (batch, time1, d_model) weighted by the query dot key attention
+        """
+        q, k, v = self.forward_qkv(query, key, value)
+        q = q.transpose(1, 2)  # (batch, time1, head, d_k)
+
+        n_batch_pos = pos_emb.size(0)
+        p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k)
+        p = p.transpose(1, 2)  # (batch, head, time1, d_k)
+
+        # (batch, head, time1, d_k)
+        q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2)
+        # (batch, head, time1, d_k)
+        q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2)
+
+        # compute attention score
+        # first compute matrix a and matrix c
+        # as described in https://arxiv.org/abs/1901.02860 Section 3.3
+        # (batch, head, time1, time2)
+        matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))
+
+        # compute matrix b and matrix d
+        # (batch, head, time1, time2)
+        matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))
+        matrix_bd = self.rel_shift(matrix_bd)
+        # drops extra elements in the matrix_bd if there are any left from rel_shift() to match the matrix_ac's size
+        matrix_bd = matrix_bd[:, :, :, : matrix_ac.size(-1)]
+
+        scores = (matrix_ac + matrix_bd) / math.sqrt(self.d_k)  # (batch, head, time1, time2)
+
+        return self.forward_attention(v, scores, mask)
+
+
+class PositionalEncoding(torch.nn.Module):
+    """Fixed sinusoidal positional encoding.
+    Args:
+        d_model (int): embedding dim
+        dropout_rate (float): dropout rate
+        max_len (int): maximum input length
+        xscale (bool): whether to scale the input by sqrt(d_model)
+        dropout_rate_emb (float): dropout rate for the positional embeddings
+    """
+
+    def __init__(self, d_model, dropout_rate, max_len=5000, xscale=None, dropout_rate_emb=0.0):
+        """Construct an PositionalEncoding object."""
+        super(PositionalEncoding, self).__init__()
+        self.d_model = d_model
+        self.xscale = xscale
+        self.dropout = torch.nn.Dropout(p=dropout_rate)
+        self.extend_pe(length=max_len)
+        if dropout_rate_emb > 0:
+            self.dropout_emb = nn.Dropout(dropout_rate_emb)
+        else:
+            self.dropout_emb = None
+
+    def create_pe(self, positions):
+        pos_length = positions.size(0)
+        pe = torch.zeros(pos_length, self.d_model)
+        div_term = torch.exp(
+            torch.arange(0, self.d_model, 2, dtype=torch.float32) * -(math.log(10000.0) / self.d_model)
+        )
+        pe[:, 0::2] = torch.sin(positions * div_term)
+        pe[:, 1::2] = torch.cos(positions * div_term)
+        pe = pe.unsqueeze(0)
+        return pe
+
+    def extend_pe(self, length):
+        """Reset and extend the positional encodings if needed."""
+        if hasattr(self, 'pe') and self.pe.size(1) >= length:
+            return
+        positions = torch.arange(0, length, dtype=torch.float32).unsqueeze(1)
+        pe = self.create_pe(positions=positions)
+        if not hasattr(self, 'pe'):
+            self.register_buffer('pe', pe, persistent=False)
+        self.pe = pe
+
+    def forward(self, x: torch.Tensor):
+        """Adds positional encoding.
+        Args:
+            x (torch.Tensor): Input. Its shape is (batch, time, feature_size)
+        Returns:
+            x+pos_emb (torch.Tensor): Its shape is (batch, time, feature_size)
+            pos_emb (torch.Tensor): Its shape is (1, time, feature_size)
+        """
+        self.extend_pe(x)
+        if self.pe.dtype != x.dtype or self.pe.device != x.device:
+            self.pe = self.pe.to(device=x.device, dtype=x.dtype)
+
+        if self.xscale:
+            x = x * self.xscale
+        pos_emb = self.pe[:, : x.size(1)]
+        if self.dropout_emb:
+            pos_emb = self.dropout_emb(pos_emb)
+        x = x + pos_emb
+        return self.dropout(x), pos_emb
+
+
+class RelPositionalEncoding(PositionalEncoding):
+    """Relative positional encoding for TransformerXL's layers
+    See : Appendix B in https://arxiv.org/abs/1901.02860
+    Args:
+        d_model (int): embedding dim
+        dropout_rate (float): dropout rate
+        max_len (int): maximum input length
+        xscale (bool): whether to scale the input by sqrt(d_model)
+        dropout_rate_emb (float): dropout rate for the positional embeddings
+    """
+
+    def __init__(self, d_model, dropout_rate, max_len=5000, xscale=None, dropout_rate_emb=0.0):
+        super().__init__(d_model, dropout_rate, max_len, xscale=xscale)
+
+        if dropout_rate_emb > 0:
+            self.dropout_emb = nn.Dropout(dropout_rate_emb)
+        else:
+            self.dropout_emb = None
+
+        self.max_len = max_len
+
+    def extend_pe(self, length):
+        """Reset and extend the positional encodings if needed."""
+        needed_size = 2 * (length - 1) + 1
+        if hasattr(self, 'pe') and self.pe.size(1) >= needed_size:
+            return
+        # positions would be from negative numbers to positive
+        # positive positions would be used for left positions and negative for right positions
+        positions = torch.arange(length - 1, -length, -1, dtype=torch.float32).unsqueeze(1)
+        pe = self.create_pe(positions=positions)
+        if not hasattr(self, 'pe'):
+            self.register_buffer('pe', pe, persistent=False)
+        self.pe = pe
+
+    def forward(self, x):
+        """Compute positional encoding.
+        Args:
+            x (torch.Tensor): Input. Its shape is (batch, time, feature_size)
+        Returns:
+            x (torch.Tensor): Its shape is (batch, time, feature_size)
+            pos_emb (torch.Tensor): Its shape is (1, time, feature_size)
+        """
+        self.extend_pe(length=x.size(1))
+        if self.pe.dtype != x.dtype or self.pe.device != x.device:
+            self.pe = self.pe.to(device=x.device, dtype=x.dtype)
+
+        if self.xscale:
+            x = x * self.xscale
+
+        # center_pos would be the index of position 0
+        # negative positions would be used for right and positive for left tokens
+        # for input of length L, 2*L-1 positions are needed, positions from (L-1) to -(L-1)
+        center_pos = self.pe.size(1) // 2
+        start_pos = center_pos - x.size(1) + 1
+        end_pos = center_pos + x.size(1)
+        pos_emb = self.pe[:, start_pos:end_pos]
+        if self.dropout_emb:
+            pos_emb = self.dropout_emb(pos_emb)
+        return self.dropout(x), pos_emb

tts/nemo/collections/asr/parts/submodules/rnnt_beam_decoding.py

@@ -0,0 +1,791 @@
+# 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 2017 Johns Hopkins University (Shinji Watanabe)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+from dataclasses import dataclass
+from typing import List, Optional, Union
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from nemo.collections.asr.modules import rnnt_abstract
+from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis, NBestHypotheses
+from nemo.core.classes import Typing, typecheck
+from nemo.core.neural_types import AcousticEncodedRepresentation, HypothesisType, LengthsType, NeuralType
+
+
+class BeamRNNTInfer(Typing):
+    """
+    Beam Search implementation ported from ESPNet implementation -
+    https://github.com/espnet/espnet/blob/master/espnet/nets/beam_search_transducer.py
+
+    Sequence level beam decoding or batched-beam decoding, performed auto-repressively
+    depending on the search type chosen.
+
+    Args:
+        decoder_model: rnnt_utils.AbstractRNNTDecoder implementation.
+        joint_model: rnnt_utils.AbstractRNNTJoint implementation.
+
+        beam_size: number of beams for beam search. Must be a positive integer >= 1.
+            If beam size is 1, defaults to stateful greedy search.
+            This greedy search might result in slightly different results than
+            the greedy results obtained by GreedyRNNTInfer due to implementation differences.
+
+            For accurate greedy results, please use GreedyRNNTInfer or GreedyBatchedRNNTInfer.
+
+        search_type: str representing the type of beam search to perform.
+            Must be one of ['beam', 'tsd', 'alsd']. 'nsc' is currently not supported.
+
+            Algoritm used:
+            `beam` - basic beam search strategy. Larger beams generally result in better decoding,
+                however the time required for the search also grows steadily.
+
+            `tsd` - time synchronous decoding. Please refer to the paper:
+                [Alignment-Length Synchronous Decoding for RNN Transducer](https://ieeexplore.ieee.org/document/9053040)
+                for details on the algorithm implemented.
+
+                Time synchronous decoding (TSD) execution time grows by the factor T * max_symmetric_expansions.
+                For longer sequences, T is greater, and can therefore take a long time for beams to obtain
+                good results. This also requires greater memory to execute.
+
+            `alsd` - alignment-length synchronous decoding. Please refer to the paper:
+                [Alignment-Length Synchronous Decoding for RNN Transducer](https://ieeexplore.ieee.org/document/9053040)
+                for details on the algorithm implemented.
+
+                Alignment-length synchronous decoding (ALSD) execution time is faster than TSD, with growth
+                factor of T + U_max, where U_max is the maximum target length expected during execution.
+
+                Generally, T + U_max < T * max_symmetric_expansions. However, ALSD beams are non-unique,
+                therefore it is required to use larger beam sizes to achieve the same (or close to the same)
+                decoding accuracy as TSD.
+
+                For a given decoding accuracy, it is possible to attain faster decoding via ALSD than TSD.
+
+        score_norm: bool, whether to normalize the scores of the log probabilities.
+
+        return_best_hypothesis: bool, decides whether to return a single hypothesis (the best out of N),
+            or return all N hypothesis (sorted with best score first). The container class changes based
+            this flag -
+            When set to True (default), returns a single Hypothesis.
+            When set to False, returns a NBestHypotheses container, which contains a list of Hypothesis.
+
+        # The following arguments are specific to the chosen `search_type`
+
+        tsd_max_sym_exp_per_step: Used for `search_type=tsd`. The maximum symmetric expansions allowed
+            per timestep during beam search. Larger values should be used to attempt decoding of longer
+            sequences, but this in turn increases execution time and memory usage.
+
+        alsd_max_target_len: Used for `search_type=alsd`. The maximum expected target sequence length
+            during beam search. Larger values allow decoding of longer sequences at the expense of
+            execution time and memory.
+
+        # The following two flags are placeholders and unused until `nsc` implementation is stabilized.
+
+        nsc_max_timesteps_expansion: Unused int.
+
+        nsc_prefix_alpha: Unused int.
+
+        preserve_alignments: Bool flag which preserves the history of alignments generated during
+            beam decoding (sample). When set to true, the Hypothesis will contain
+            the non-null value for `alignments` in it. Here, `alignments` is a List of List of ints.
+
+            The length of the list corresponds to the Acoustic Length (T).
+            Each value in the list (Ti) is a torch.Tensor (U), representing 1 or more targets from a vocabulary.
+            U is the number of target tokens for the current timestep Ti.
+
+            NOTE: `preserve_alignments` is an invalid argument for any `search_type`
+            other than basic beam search.
+    """
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input ports.
+        """
+        return {
+            "encoder_output": NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()),
+            "encoded_lengths": NeuralType(tuple('B'), LengthsType()),
+        }
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output ports.
+        """
+        return {"predictions": [NeuralType(elements_type=HypothesisType())]}
+
+    def __init__(
+        self,
+        decoder_model: rnnt_abstract.AbstractRNNTDecoder,
+        joint_model: rnnt_abstract.AbstractRNNTJoint,
+        beam_size: int,
+        search_type: str = 'default',
+        score_norm: bool = True,
+        return_best_hypothesis: bool = True,
+        tsd_max_sym_exp_per_step: Optional[int] = 50,
+        alsd_max_target_len: Union[int, float] = 1.0,
+        nsc_max_timesteps_expansion: int = 1,
+        nsc_prefix_alpha: int = 1,
+        preserve_alignments: bool = False,
+    ):
+        self.decoder = decoder_model
+        self.joint = joint_model
+
+        self.blank = decoder_model.blank_idx
+        self.vocab_size = decoder_model.vocab_size
+        self.search_type = search_type
+        self.return_best_hypothesis = return_best_hypothesis
+
+        if beam_size < 1:
+            raise ValueError("Beam search size cannot be less than 1!")
+
+        self.beam_size = beam_size
+        self.score_norm = score_norm
+
+        if self.beam_size == 1:
+            self.search_algorithm = self.greedy_search
+        elif search_type == "default":
+            self.search_algorithm = self.default_beam_search
+        elif search_type == "tsd":
+            self.search_algorithm = self.time_sync_decoding
+        elif search_type == "alsd":
+            self.search_algorithm = self.align_length_sync_decoding
+        elif search_type == "nsc":
+            raise NotImplementedError("`nsc` (Constrained Beam Search) has not been implemented.")
+            # self.search_algorithm = self.nsc_beam_search
+        else:
+            raise NotImplementedError(
+                f"The search type ({search_type}) supplied is not supported!\n"
+                f"Please use one of : (default, tsd, alsd, nsc)"
+            )
+
+        if tsd_max_sym_exp_per_step is None:
+            tsd_max_sym_exp_per_step = -1
+
+        if search_type in ['tsd', 'alsd', 'nsc'] and not self.decoder.blank_as_pad:
+            raise ValueError(
+                f"Search type was chosen as '{search_type}', however the decoder module provided "
+                f"does not support the `blank` token as a pad value. {search_type} requires "
+                f"the blank token as pad value support in order to perform batched beam search."
+                f"Please chose one of the other beam search methods, or re-train your model "
+                f"with this support."
+            )
+
+        self.tsd_max_symmetric_expansion_per_step = tsd_max_sym_exp_per_step
+        self.alsd_max_target_length = alsd_max_target_len
+        self.nsc_max_timesteps_expansion = nsc_max_timesteps_expansion
+        self.nsc_prefix_alpha = nsc_prefix_alpha
+        self.preserve_alignments = preserve_alignments
+
+    @typecheck()
+    def __call__(
+        self, encoder_output: torch.Tensor, encoded_lengths: torch.Tensor
+    ) -> Union[Hypothesis, NBestHypotheses]:
+        """Perform general beam search.
+
+        Args:
+            encoder_output: Encoded speech features (B, T_max, D_enc)
+            encoded_lengths: Lengths of the encoder outputs
+
+        Returns:
+            Either a list containing a single Hypothesis (when `return_best_hypothesis=True`,
+            otherwise a list containing a single NBestHypotheses, which itself contains a list of
+            Hypothesis. This list is sorted such that the best hypothesis is the first element.
+        """
+        # Preserve decoder and joint training state
+        decoder_training_state = self.decoder.training
+        joint_training_state = self.joint.training
+
+        with torch.no_grad():
+            # Apply optional preprocessing
+            encoder_output = encoder_output.transpose(1, 2)  # (B, T, D)
+
+            self.decoder.eval()
+            self.joint.eval()
+
+            hypotheses = []
+            with tqdm(
+                range(encoder_output.size(0)),
+                desc='Beam search progress:',
+                total=encoder_output.size(0),
+                unit='sample',
+            ) as idx_gen:
+
+                # Freeze the decoder and joint to prevent recording of gradients
+                # during the beam loop.
+                with self.decoder.as_frozen(), self.joint.as_frozen():
+
+                    # Decode every sample in the batch independently.
+                    for batch_idx in idx_gen:
+                        inseq = encoder_output[batch_idx : batch_idx + 1, :, :]  # [1, T, D]
+                        logitlen = encoded_lengths[batch_idx]
+
+                        # Execute the specific search strategy
+                        nbest_hyps = self.search_algorithm(inseq, logitlen)  # sorted list of hypothesis
+
+                        # Pack the result
+                        if self.return_best_hypothesis:
+                            best_hypothesis = nbest_hyps[0]  # type: Hypothesis
+                        else:
+                            best_hypothesis = NBestHypotheses(nbest_hyps)  # type: NBestHypotheses
+                        hypotheses.append(best_hypothesis)
+
+        self.decoder.train(decoder_training_state)
+        self.joint.train(joint_training_state)
+
+        return (hypotheses,)
+
+    def sort_nbest(self, hyps: List[Hypothesis]) -> List[Hypothesis]:
+        """Sort hypotheses by score or score given sequence length.
+
+        Args:
+            hyps: list of hypotheses
+
+        Return:
+            hyps: sorted list of hypotheses
+        """
+        if self.score_norm:
+            return sorted(hyps, key=lambda x: x.score / len(x.y_sequence), reverse=True)
+        else:
+            return sorted(hyps, key=lambda x: x.score, reverse=True)
+
+    def greedy_search(self, h: torch.Tensor, encoded_lengths: torch.Tensor) -> List[Hypothesis]:
+        """Greedy search implementation for transducer.
+        Generic case when beam size = 1. Results might differ slightly due to implementation details
+        as compared to `GreedyRNNTInfer` and `GreedyBatchRNNTInfer`.
+
+        Args:
+            h: Encoded speech features (1, T_max, D_enc)
+
+        Returns:
+            hyp: 1-best decoding results
+        """
+        if self.preserve_alignments:
+            # Alignments is a 2-dimensional dangling list representing T x U
+            alignments = [[]]
+        else:
+            alignments = None
+
+        # Initialize zero state vectors
+        dec_state = self.decoder.initialize_state(h)
+
+        # Construct initial hypothesis
+        hyp = Hypothesis(
+            score=0.0, y_sequence=[self.blank], dec_state=dec_state, timestep=[-1], length=encoded_lengths
+        )
+        cache = {}
+
+        # Initialize state and first token
+        y, state, _ = self.decoder.score_hypothesis(hyp, cache)
+
+        for i in range(int(encoded_lengths)):
+            hi = h[:, i : i + 1, :]  # [1, 1, D]
+
+            not_blank = True
+            symbols_added = 0
+
+            while not_blank:
+                ytu = torch.log_softmax(self.joint.joint(hi, y), dim=-1)  # [1, 1, 1, V + 1]
+                ytu = ytu[0, 0, 0, :]  # [V + 1]
+
+                # max() requires float
+                if ytu.dtype != torch.float32:
+                    ytu = ytu.float()
+
+                logp, pred = torch.max(ytu, dim=-1)  # [1, 1]
+                pred = pred.item()
+
+                if self.preserve_alignments:
+                    # insert logits into last timestep
+                    alignments[-1].append(pred)
+
+                if pred == self.blank:
+                    not_blank = False
+
+                    if self.preserve_alignments:
+                        # convert Ti-th logits into a torch array
+                        alignments.append([])  # blank buffer for next timestep
+                else:
+                    # Update state and current sequence
+                    hyp.y_sequence.append(int(pred))
+                    hyp.score += float(logp)
+                    hyp.dec_state = state
+                    hyp.timestep.append(i)
+
+                    # Compute next state and token
+                    y, state, _ = self.decoder.score_hypothesis(hyp, cache)
+                symbols_added += 1
+
+        # Remove trailing empty list of alignments
+        if self.preserve_alignments:
+            if len(alignments[-1]) == 0:
+                del alignments[-1]
+
+        # attach alignments to hypothesis
+        hyp.alignments = alignments
+
+        return [hyp]
+
+    def default_beam_search(self, h: torch.Tensor, encoded_lengths: torch.Tensor) -> List[Hypothesis]:
+        """Beam search implementation.
+
+        Args:
+            x: Encoded speech features (1, T_max, D_enc)
+
+        Returns:
+            nbest_hyps: N-best decoding results
+        """
+        # Initialize states
+        beam = min(self.beam_size, self.vocab_size)
+        beam_k = min(beam, (self.vocab_size - 1))
+        blank_tensor = torch.tensor([self.blank], device=h.device, dtype=torch.long)
+
+        # Precompute some constants for blank position
+        ids = list(range(self.vocab_size + 1))
+        ids.remove(self.blank)
+
+        # Used when blank token is first vs last token
+        if self.blank == 0:
+            index_incr = 1
+        else:
+            index_incr = 0
+
+        # Initialize zero vector states
+        dec_state = self.decoder.initialize_state(h)
+
+        # Initialize first hypothesis for the beam (blank)
+        kept_hyps = [Hypothesis(score=0.0, y_sequence=[self.blank], dec_state=dec_state, timestep=[-1], length=0)]
+        cache = {}
+
+        if self.preserve_alignments:
+            kept_hyps[0].alignments = [[]]
+
+        for i in range(int(encoded_lengths)):
+            hi = h[:, i : i + 1, :]  # [1, 1, D]
+            hyps = kept_hyps
+            kept_hyps = []
+
+            while True:
+                max_hyp = max(hyps, key=lambda x: x.score)
+                hyps.remove(max_hyp)
+
+                # update decoder state and get next score
+                y, state, lm_tokens = self.decoder.score_hypothesis(max_hyp, cache)  # [1, 1, D]
+
+                # get next token
+                ytu = torch.log_softmax(self.joint.joint(hi, y), dim=-1)  # [1, 1, 1, V + 1]
+                ytu = ytu[0, 0, 0, :]  # [V + 1]
+
+                # remove blank token before top k
+                top_k = ytu[ids].topk(beam_k, dim=-1)
+
+                # Two possible steps - blank token or non-blank token predicted
+                ytu = (
+                    torch.cat((top_k[0], ytu[self.blank].unsqueeze(0))),
+                    torch.cat((top_k[1] + index_incr, blank_tensor)),
+                )
+
+                # for each possible step
+                for logp, k in zip(*ytu):
+                    # construct hypothesis for step
+                    new_hyp = Hypothesis(
+                        score=(max_hyp.score + float(logp)),
+                        y_sequence=max_hyp.y_sequence[:],
+                        dec_state=max_hyp.dec_state,
+                        lm_state=max_hyp.lm_state,
+                        timestep=max_hyp.timestep[:],
+                        length=encoded_lengths,
+                    )
+
+                    if self.preserve_alignments:
+                        new_hyp.alignments = copy.deepcopy(max_hyp.alignments)
+
+                    # if current token is blank, dont update sequence, just store the current hypothesis
+                    if k == self.blank:
+                        kept_hyps.append(new_hyp)
+                    else:
+                        # if non-blank token was predicted, update state and sequence and then search more hypothesis
+                        new_hyp.dec_state = state
+                        new_hyp.y_sequence.append(int(k))
+                        new_hyp.timestep.append(i)
+
+                        hyps.append(new_hyp)
+
+                    if self.preserve_alignments:
+                        if k == self.blank:
+                            new_hyp.alignments[-1].append(self.blank)
+                        else:
+                            new_hyp.alignments[-1].append(new_hyp.y_sequence[-1])
+
+                # keep those hypothesis that have scores greater than next search generation
+                hyps_max = float(max(hyps, key=lambda x: x.score).score)
+                kept_most_prob = sorted([hyp for hyp in kept_hyps if hyp.score > hyps_max], key=lambda x: x.score,)
+
+                # If enough hypothesis have scores greater than next search generation,
+                # stop beam search.
+                if len(kept_most_prob) >= beam:
+                    if self.preserve_alignments:
+                        # convert Ti-th logits into a torch array
+                        for kept_h in kept_most_prob:
+                            kept_h.alignments.append([])  # blank buffer for next timestep
+
+                    kept_hyps = kept_most_prob
+                    break
+
+        # Remove trailing empty list of alignments
+        if self.preserve_alignments:
+            for h in kept_hyps:
+                if len(h.alignments[-1]) == 0:
+                    del h.alignments[-1]
+
+        return self.sort_nbest(kept_hyps)
+
+    def time_sync_decoding(self, h: torch.Tensor, encoded_lengths: torch.Tensor) -> List[Hypothesis]:
+        """Time synchronous beam search implementation.
+        Based on https://ieeexplore.ieee.org/document/9053040
+
+        Args:
+            h: Encoded speech features (1, T_max, D_enc)
+
+        Returns:
+            nbest_hyps: N-best decoding results
+        """
+        if self.preserve_alignments:
+            raise NotImplementedError("`preseve_alignments` is not implemented for Time-Synchronous Decoding.")
+
+        # Precompute some constants for blank position
+        ids = list(range(self.vocab_size + 1))
+        ids.remove(self.blank)
+
+        # Used when blank token is first vs last token
+        if self.blank == 0:
+            index_incr = 1
+        else:
+            index_incr = 0
+
+        # prepare the batched beam states
+        beam = min(self.beam_size, self.vocab_size)
+        beam_state = self.decoder.initialize_state(
+            torch.zeros(beam, device=h.device, dtype=h.dtype)
+        )  # [L, B, H], [L, B, H] (for LSTMs)
+
+        # Initialize first hypothesis for the beam (blank)
+        B = [
+            Hypothesis(
+                y_sequence=[self.blank],
+                score=0.0,
+                dec_state=self.decoder.batch_select_state(beam_state, 0),
+                timestep=[-1],
+                length=0,
+            )
+        ]
+        cache = {}
+
+        for i in range(int(encoded_lengths)):
+            hi = h[:, i : i + 1, :]
+
+            # Update caches
+            A = []
+            C = B
+
+            h_enc = hi
+
+            # For a limited number of symmetric expansions per timestep "i"
+            for v in range(self.tsd_max_symmetric_expansion_per_step):
+                D = []
+
+                # Decode a batch of beam states and scores
+                beam_y, beam_state, beam_lm_tokens = self.decoder.batch_score_hypothesis(C, cache, beam_state)
+
+                # Extract the log probabilities and the predicted tokens
+                beam_logp = torch.log_softmax(self.joint.joint(h_enc, beam_y), dim=-1)  # [B, 1, 1, V + 1]
+                beam_logp = beam_logp[:, 0, 0, :]  # [B, V + 1]
+                beam_topk = beam_logp[:, ids].topk(beam, dim=-1)
+
+                seq_A = [h.y_sequence for h in A]
+
+                for j, hyp in enumerate(C):
+                    # create a new hypothesis in A
+                    if hyp.y_sequence not in seq_A:
+                        # If the sequence is not in seq_A, add it as the blank token
+                        # In this step, we dont add a token but simply update score
+                        A.append(
+                            Hypothesis(
+                                score=(hyp.score + float(beam_logp[j, self.blank])),
+                                y_sequence=hyp.y_sequence[:],
+                                dec_state=hyp.dec_state,
+                                lm_state=hyp.lm_state,
+                                timestep=hyp.timestep[:],
+                                length=encoded_lengths,
+                            )
+                        )
+                    else:
+                        # merge the existing blank hypothesis score with current score.
+                        dict_pos = seq_A.index(hyp.y_sequence)
+
+                        A[dict_pos].score = np.logaddexp(
+                            A[dict_pos].score, (hyp.score + float(beam_logp[j, self.blank]))
+                        )
+
+                if v < self.tsd_max_symmetric_expansion_per_step:
+                    for j, hyp in enumerate(C):
+                        # for each current hypothesis j
+                        # extract the top token score and top token id for the jth hypothesis
+                        for logp, k in zip(beam_topk[0][j], beam_topk[1][j] + index_incr):
+                            # create new hypothesis and store in D
+                            # Note: This loop does *not* include the blank token!
+                            new_hyp = Hypothesis(
+                                score=(hyp.score + float(logp)),
+                                y_sequence=(hyp.y_sequence + [int(k)]),
+                                dec_state=self.decoder.batch_select_state(beam_state, j),
+                                lm_state=hyp.lm_state,
+                                timestep=hyp.timestep[:] + [i],
+                                length=encoded_lengths,
+                            )
+
+                            D.append(new_hyp)
+
+                # Prune beam
+                C = sorted(D, key=lambda x: x.score, reverse=True)[:beam]
+
+            # Prune beam
+            B = sorted(A, key=lambda x: x.score, reverse=True)[:beam]
+
+        return self.sort_nbest(B)
+
+    def align_length_sync_decoding(self, h: torch.Tensor, encoded_lengths: torch.Tensor) -> List[Hypothesis]:
+        """Alignment-length synchronous beam search implementation.
+        Based on https://ieeexplore.ieee.org/document/9053040
+
+        Args:
+            h: Encoded speech features (1, T_max, D_enc)
+
+        Returns:
+            nbest_hyps: N-best decoding results
+        """
+        if self.preserve_alignments:
+            raise NotImplementedError(
+                "`preseve_alignments` is not implemented for Alignment-length Synchronous Decoding."
+            )
+
+        # Precompute some constants for blank position
+        ids = list(range(self.vocab_size + 1))
+        ids.remove(self.blank)
+
+        # Used when blank token is first vs last token
+        if self.blank == 0:
+            index_incr = 1
+        else:
+            index_incr = 0
+
+        # prepare the batched beam states
+        beam = min(self.beam_size, self.vocab_size)
+
+        h = h[0]  # [T, D]
+        h_length = int(encoded_lengths)
+        beam_state = self.decoder.initialize_state(
+            torch.zeros(beam, device=h.device, dtype=h.dtype)
+        )  # [L, B, H], [L, B, H] for LSTMS
+
+        # compute u_max as either a specific static limit,
+        # or a multiple of current `h_length` dynamically.
+        if type(self.alsd_max_target_length) == float:
+            u_max = int(self.alsd_max_target_length * h_length)
+        else:
+            u_max = int(self.alsd_max_target_length)
+
+        # Initialize first hypothesis for the beam (blank)
+        B = [
+            Hypothesis(
+                y_sequence=[self.blank],
+                score=0.0,
+                dec_state=self.decoder.batch_select_state(beam_state, 0),
+                timestep=[-1],
+                length=0,
+            )
+        ]
+
+        final = []
+        cache = {}
+
+        # ALSD runs for T + U_max steps
+        for i in range(h_length + u_max):
+            # Update caches
+            A = []
+            B_ = []
+            h_states = []
+
+            # preserve the list of batch indices which are added into the list
+            # and those which are removed from the list
+            # This is necessary to perform state updates in the correct batch indices later
+            batch_ids = list(range(len(B)))  # initialize as a list of all batch ids
+            batch_removal_ids = []  # update with sample ids which are removed
+
+            for bid, hyp in enumerate(B):
+                u = len(hyp.y_sequence) - 1
+                t = i - u + 1
+
+                if t > (h_length - 1):
+                    batch_removal_ids.append(bid)
+                    continue
+
+                B_.append(hyp)
+                h_states.append((t, h[t]))
+
+            if B_:
+                # Compute the subset of batch ids which were *not* removed from the list above
+                sub_batch_ids = None
+                if len(B_) != beam:
+                    sub_batch_ids = batch_ids
+                    for id in batch_removal_ids:
+                        # sub_batch_ids contains list of ids *that were not removed*
+                        sub_batch_ids.remove(id)
+
+                    # extract the states of the sub batch only.
+                    beam_state_ = [beam_state[state_id][:, sub_batch_ids, :] for state_id in range(len(beam_state))]
+                else:
+                    # If entire batch was used (none were removed), simply take all the states
+                    beam_state_ = beam_state
+
+                # Decode a batch/sub-batch of beam states and scores
+                beam_y, beam_state_, beam_lm_tokens = self.decoder.batch_score_hypothesis(B_, cache, beam_state_)
+
+                # If only a subset of batch ids were updated (some were removed)
+                if sub_batch_ids is not None:
+                    # For each state in the RNN (2 for LSTM)
+                    for state_id in range(len(beam_state)):
+                        # Update the current batch states with the sub-batch states (in the correct indices)
+                        # These indices are specified by sub_batch_ids, the ids of samples which were updated.
+                        beam_state[state_id][:, sub_batch_ids, :] = beam_state_[state_id][...]
+                else:
+                    # If entire batch was updated, simply update all the states
+                    beam_state = beam_state_
+
+                # h_states = list of [t, h[t]]
+                # so h[1] here is a h[t] of shape [D]
+                # Simply stack all of the h[t] within the sub_batch/batch (T <= beam)
+                h_enc = torch.stack([h[1] for h in h_states])  # [T=beam, D]
+                h_enc = h_enc.unsqueeze(1)  # [B=beam, T=1, D]; batch over the beams
+
+                # Extract the log probabilities and the predicted tokens
+                beam_logp = torch.log_softmax(self.joint.joint(h_enc, beam_y), dim=-1)  # [B=beam, 1, 1, V + 1]
+                beam_logp = beam_logp[:, 0, 0, :]  # [B=beam, V + 1]
+                beam_topk = beam_logp[:, ids].topk(beam, dim=-1)
+
+                for j, hyp in enumerate(B_):
+                    # For all updated samples in the batch, add it as the blank token
+                    # In this step, we dont add a token but simply update score
+                    new_hyp = Hypothesis(
+                        score=(hyp.score + float(beam_logp[j, self.blank])),
+                        y_sequence=hyp.y_sequence[:],
+                        dec_state=hyp.dec_state,
+                        lm_state=hyp.lm_state,
+                        timestep=hyp.timestep[:],
+                        length=i,
+                    )
+
+                    # Add blank prediction to A
+                    A.append(new_hyp)
+
+                    # If the prediction "timestep" t has reached the length of the input sequence
+                    # we can add it to the "finished" hypothesis list.
+                    if h_states[j][0] == (h_length - 1):
+                        final.append(new_hyp)
+
+                    # Here, we carefully select the indices of the states that we want to preserve
+                    # for the next token (non-blank) update.
+                    if sub_batch_ids is not None:
+                        h_states_idx = sub_batch_ids[j]
+                    else:
+                        h_states_idx = j
+
+                    # for each current hypothesis j
+                    # extract the top token score and top token id for the jth hypothesis
+                    for logp, k in zip(beam_topk[0][j], beam_topk[1][j] + index_incr):
+                        # create new hypothesis and store in A
+                        # Note: This loop does *not* include the blank token!
+                        new_hyp = Hypothesis(
+                            score=(hyp.score + float(logp)),
+                            y_sequence=(hyp.y_sequence[:] + [int(k)]),
+                            dec_state=self.decoder.batch_select_state(beam_state, h_states_idx),
+                            lm_state=hyp.lm_state,
+                            timestep=hyp.timestep[:] + [i],
+                            length=i,
+                        )
+
+                        A.append(new_hyp)
+
+                # Prune and recombine same hypothesis
+                # This may cause next beam to be smaller than max beam size
+                # Therefore larger beam sizes may be required for better decoding.
+                B = sorted(A, key=lambda x: x.score, reverse=True)[:beam]
+                B = self.recombine_hypotheses(B)
+
+            # If B_ is empty list, then we may be able to early exit
+            elif len(batch_ids) == len(batch_removal_ids):
+                break
+
+        if final:
+            return self.sort_nbest(final)
+        else:
+            return B
+
+    def recombine_hypotheses(self, hypotheses: List[Hypothesis]) -> List[Hypothesis]:
+        """Recombine hypotheses with equivalent output sequence.
+
+        Args:
+            hypotheses (list): list of hypotheses
+
+        Returns:
+           final (list): list of recombined hypotheses
+        """
+        final = []
+
+        for hyp in hypotheses:
+            seq_final = [f.y_sequence for f in final if f.y_sequence]
+
+            if hyp.y_sequence in seq_final:
+                seq_pos = seq_final.index(hyp.y_sequence)
+
+                final[seq_pos].score = np.logaddexp(final[seq_pos].score, hyp.score)
+            else:
+                final.append(hyp)
+
+        return hypotheses
+
+
+@dataclass
+class BeamRNNTInferConfig:
+    beam_size: int
+    search_type: str = 'default'
+    score_norm: bool = True
+    return_best_hypothesis: bool = True
+    tsd_max_sym_exp_per_step: Optional[int] = 50
+    alsd_max_target_len: float = 1.0
+    nsc_max_timesteps_expansion: int = 1
+    nsc_prefix_alpha: int = 1
+    preserve_alignments: bool = False

tts/nemo/collections/asr/parts/submodules/rnnt_greedy_decoding.py

@@ -0,0 +1,731 @@
+# 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 2017 Johns Hopkins University (Shinji Watanabe)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from typing import List, Optional, Union
+
+import torch
+
+from nemo.collections.asr.modules import rnnt_abstract
+from nemo.collections.asr.parts.utils import rnnt_utils
+from nemo.collections.common.parts.rnn import label_collate
+from nemo.core.classes import Typing, typecheck
+from nemo.core.neural_types import AcousticEncodedRepresentation, HypothesisType, LengthsType, NeuralType
+
+
+def pack_hypotheses(
+    hypotheses: List[List[int]],
+    timesteps: List[List[int]],
+    logitlen: torch.Tensor,
+    alignments: Optional[List[List[int]]] = None,
+) -> List[rnnt_utils.Hypothesis]:
+    logitlen_cpu = logitlen.to("cpu")
+    return [
+        rnnt_utils.Hypothesis(
+            y_sequence=torch.tensor(sent, dtype=torch.long),
+            score=-1.0,
+            timestep=timestep,
+            length=length,
+            alignments=alignments[idx] if alignments is not None else None,
+        )
+        for idx, (sent, timestep, length) in enumerate(zip(hypotheses, timesteps, logitlen_cpu))
+    ]
+
+
+class _GreedyRNNTInfer(Typing):
+    """A greedy transducer decoder.
+
+    Provides a common abstraction for sample level and batch level greedy decoding.
+
+    Args:
+        decoder_model: rnnt_utils.AbstractRNNTDecoder implementation.
+        joint_model: rnnt_utils.AbstractRNNTJoint implementation.
+        blank_index: int index of the blank token. Can be 0 or len(vocabulary).
+        max_symbols_per_step: Optional int. The maximum number of symbols that can be added
+            to a sequence in a single time step; if set to None then there is
+            no limit.
+        preserve_alignments: Bool flag which preserves the history of alignments generated during
+            greedy decoding (sample / batched). When set to true, the Hypothesis will contain
+            the non-null value for `alignments` in it. Here, `alignments` is a List of List of ints.
+
+            The length of the list corresponds to the Acoustic Length (T).
+            Each value in the list (Ti) is a torch.Tensor (U), representing 1 or more targets from a vocabulary.
+            U is the number of target tokens for the current timestep Ti.
+    """
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input ports.
+        """
+        return {
+            "encoder_output": NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()),
+            "encoded_lengths": NeuralType(tuple('B'), LengthsType()),
+        }
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output ports.
+        """
+        return {"predictions": [NeuralType(elements_type=HypothesisType())]}
+
+    def __init__(
+        self,
+        decoder_model: rnnt_abstract.AbstractRNNTDecoder,
+        joint_model: rnnt_abstract.AbstractRNNTJoint,
+        blank_index: int,
+        max_symbols_per_step: Optional[int] = None,
+        preserve_alignments: bool = False,
+    ):
+        super().__init__()
+        self.decoder = decoder_model
+        self.joint = joint_model
+
+        self._blank_index = blank_index
+        self._SOS = blank_index  # Start of single index
+        self.max_symbols = max_symbols_per_step
+        self.preserve_alignments = preserve_alignments
+
+    def __call__(self, *args, **kwargs):
+        return self.forward(*args, **kwargs)
+
+    @torch.no_grad()
+    def _pred_step(
+        self,
+        label: Union[torch.Tensor, int],
+        hidden: Optional[torch.Tensor],
+        add_sos: bool = False,
+        batch_size: Optional[int] = None,
+    ) -> (torch.Tensor, torch.Tensor):
+        """
+        Common prediction step based on the AbstractRNNTDecoder implementation.
+
+        Args:
+            label: (int/torch.Tensor): Label or "Start-of-Signal" token.
+            hidden: (Optional torch.Tensor): RNN State vector
+            add_sos (bool): Whether to add a zero vector at the begging as "start of sentence" token.
+            batch_size: Batch size of the output tensor.
+
+        Returns:
+            g: (B, U, H) if add_sos is false, else (B, U + 1, H)
+            hid: (h, c) where h is the final sequence hidden state and c is
+                the final cell state:
+                    h (tensor), shape (L, B, H)
+                    c (tensor), shape (L, B, H)
+        """
+        if isinstance(label, torch.Tensor):
+            # label: [batch, 1]
+            if label.dtype != torch.long:
+                label = label.long()
+
+        else:
+            # Label is an integer
+            if label == self._SOS:
+                return self.decoder.predict(None, hidden, add_sos=add_sos, batch_size=batch_size)
+
+            label = label_collate([[label]])
+
+        # output: [B, 1, K]
+        return self.decoder.predict(label, hidden, add_sos=add_sos, batch_size=batch_size)
+
+    def _joint_step(self, enc, pred, log_normalize: Optional[bool] = None):
+        """
+        Common joint step based on AbstractRNNTJoint implementation.
+
+        Args:
+            enc: Output of the Encoder model. A torch.Tensor of shape [B, 1, H1]
+            pred: Output of the Decoder model. A torch.Tensor of shape [B, 1, H2]
+            log_normalize: Whether to log normalize or not. None will log normalize only for CPU.
+
+        Returns:
+             logits of shape (B, T=1, U=1, V + 1)
+        """
+        with torch.no_grad():
+            logits = self.joint.joint(enc, pred)
+
+            if log_normalize is None:
+                if not logits.is_cuda:  # Use log softmax only if on CPU
+                    logits = logits.log_softmax(dim=len(logits.shape) - 1)
+            else:
+                if log_normalize:
+                    logits = logits.log_softmax(dim=len(logits.shape) - 1)
+
+        return logits
+
+
+class GreedyRNNTInfer(_GreedyRNNTInfer):
+    """A greedy transducer decoder.
+
+    Sequence level greedy decoding, performed auto-repressively.
+
+    Args:
+        decoder_model: rnnt_utils.AbstractRNNTDecoder implementation.
+        joint_model: rnnt_utils.AbstractRNNTJoint implementation.
+        blank_index: int index of the blank token. Can be 0 or len(vocabulary).
+        max_symbols_per_step: Optional int. The maximum number of symbols that can be added
+            to a sequence in a single time step; if set to None then there is
+            no limit.
+        preserve_alignments: Bool flag which preserves the history of alignments generated during
+            greedy decoding (sample / batched). When set to true, the Hypothesis will contain
+            the non-null value for `alignments` in it. Here, `alignments` is a List of List of ints.
+
+            The length of the list corresponds to the Acoustic Length (T).
+            Each value in the list (Ti) is a torch.Tensor (U), representing 1 or more targets from a vocabulary.
+            U is the number of target tokens for the current timestep Ti.
+    """
+
+    def __init__(
+        self,
+        decoder_model: rnnt_abstract.AbstractRNNTDecoder,
+        joint_model: rnnt_abstract.AbstractRNNTJoint,
+        blank_index: int,
+        max_symbols_per_step: Optional[int] = None,
+        preserve_alignments: bool = False,
+    ):
+        super().__init__(
+            decoder_model=decoder_model,
+            joint_model=joint_model,
+            blank_index=blank_index,
+            max_symbols_per_step=max_symbols_per_step,
+            preserve_alignments=preserve_alignments,
+        )
+
+    @typecheck()
+    def forward(self, encoder_output: torch.Tensor, encoded_lengths: torch.Tensor):
+        """Returns a list of hypotheses given an input batch of the encoder hidden embedding.
+        Output token is generated auto-repressively.
+
+        Args:
+            encoder_output: A tensor of size (batch, features, timesteps).
+            encoded_lengths: list of int representing the length of each sequence
+                output sequence.
+
+        Returns:
+            packed list containing batch number of sentences (Hypotheses).
+        """
+        # Preserve decoder and joint training state
+        decoder_training_state = self.decoder.training
+        joint_training_state = self.joint.training
+
+        with torch.no_grad():
+            # Apply optional preprocessing
+            encoder_output = encoder_output.transpose(1, 2)  # (B, T, D)
+
+            self.decoder.eval()
+            self.joint.eval()
+
+            hypotheses = []
+            timesteps = []
+            alignments = [] if self.preserve_alignments else None
+            # Process each sequence independently
+            with self.decoder.as_frozen(), self.joint.as_frozen():
+                for batch_idx in range(encoder_output.size(0)):
+                    inseq = encoder_output[batch_idx, :, :].unsqueeze(1)  # [T, 1, D]
+                    logitlen = encoded_lengths[batch_idx]
+                    sentence, timestep, alignment = self._greedy_decode(inseq, logitlen)
+                    hypotheses.append(sentence)
+                    timesteps.append(timestep)
+
+                    if self.preserve_alignments:
+                        alignments.append(alignment)
+
+            # Pack results into Hypotheses
+            packed_result = pack_hypotheses(hypotheses, timesteps, encoded_lengths, alignments=alignments)
+
+        self.decoder.train(decoder_training_state)
+        self.joint.train(joint_training_state)
+
+        return (packed_result,)
+
+    @torch.no_grad()
+    def _greedy_decode(self, x: torch.Tensor, out_len: torch.Tensor):
+        # x: [T, 1, D]
+        # out_len: [seq_len]
+
+        # Initialize blank state and empty label set
+        hidden = None
+        label = []
+        timesteps = []
+
+        if self.preserve_alignments:
+            # Alignments is a 2-dimensional dangling list representing T x U
+            alignments = [[]]
+        else:
+            alignments = None
+
+        # For timestep t in X_t
+        for time_idx in range(out_len):
+            # Extract encoder embedding at timestep t
+            # f = x[time_idx, :, :].unsqueeze(0)  # [1, 1, D]
+            f = x.narrow(dim=0, start=time_idx, length=1)
+
+            # Setup exit flags and counter
+            not_blank = True
+            symbols_added = 0
+
+            # While blank is not predicted, or we dont run out of max symbols per timestep
+            while not_blank and (self.max_symbols is None or symbols_added < self.max_symbols):
+                # In the first timestep, we initialize the network with RNNT Blank
+                # In later timesteps, we provide previous predicted label as input.
+                last_label = self._SOS if label == [] else label[-1]
+
+                # Perform prediction network and joint network steps.
+                g, hidden_prime = self._pred_step(last_label, hidden)
+                logp = self._joint_step(f, g, log_normalize=None)[0, 0, 0, :]
+
+                del g
+
+                # torch.max(0) op doesnt exist for FP 16.
+                if logp.dtype != torch.float32:
+                    logp = logp.float()
+
+                # get index k, of max prob
+                v, k = logp.max(0)
+                k = k.item()  # K is the label at timestep t_s in inner loop, s >= 0.
+
+                if self.preserve_alignments:
+                    # insert logits into last timestep
+                    alignments[-1].append(k)
+
+                del logp
+
+                # If blank token is predicted, exit inner loop, move onto next timestep t
+                if k == self._blank_index:
+                    not_blank = False
+
+                    if self.preserve_alignments:
+                        # convert Ti-th logits into a torch array
+                        alignments.append([])  # blank buffer for next timestep
+                else:
+                    # Append token to label set, update RNN state.
+                    label.append(k)
+                    timesteps.append(time_idx)
+                    hidden = hidden_prime
+
+                # Increment token counter.
+                symbols_added += 1
+
+        # Remove trailing empty list of Alignments
+        if self.preserve_alignments:
+            if len(alignments[-1]) == 0:
+                del alignments[-1]
+
+        return label, timesteps, alignments
+
+
+class GreedyBatchedRNNTInfer(_GreedyRNNTInfer):
+    """A batch level greedy transducer decoder.
+
+    Batch level greedy decoding, performed auto-repressively.
+
+    Args:
+        decoder_model: rnnt_utils.AbstractRNNTDecoder implementation.
+        joint_model: rnnt_utils.AbstractRNNTJoint implementation.
+        blank_index: int index of the blank token. Can be 0 or len(vocabulary).
+        max_symbols_per_step: Optional int. The maximum number of symbols that can be added
+            to a sequence in a single time step; if set to None then there is
+            no limit.
+        preserve_alignments: Bool flag which preserves the history of alignments generated during
+            greedy decoding (sample / batched). When set to true, the Hypothesis will contain
+            the non-null value for `alignments` in it. Here, `alignments` is a List of List of ints.
+
+            The length of the list corresponds to the Acoustic Length (T).
+            Each value in the list (Ti) is a torch.Tensor (U), representing 1 or more targets from a vocabulary.
+            U is the number of target tokens for the current timestep Ti.
+    """
+
+    def __init__(
+        self,
+        decoder_model: rnnt_abstract.AbstractRNNTDecoder,
+        joint_model: rnnt_abstract.AbstractRNNTJoint,
+        blank_index: int,
+        max_symbols_per_step: Optional[int] = None,
+        preserve_alignments: bool = False,
+    ):
+        super().__init__(
+            decoder_model=decoder_model,
+            joint_model=joint_model,
+            blank_index=blank_index,
+            max_symbols_per_step=max_symbols_per_step,
+            preserve_alignments=preserve_alignments,
+        )
+
+        # Depending on availability of `blank_as_pad` support
+        # switch between more efficient batch decoding technique
+        if self.decoder.blank_as_pad:
+            self._greedy_decode = self._greedy_decode_blank_as_pad
+        else:
+            self._greedy_decode = self._greedy_decode_masked
+
+    @typecheck()
+    def forward(self, encoder_output: torch.Tensor, encoded_lengths: torch.Tensor):
+        """Returns a list of hypotheses given an input batch of the encoder hidden embedding.
+        Output token is generated auto-repressively.
+
+        Args:
+            encoder_output: A tensor of size (batch, features, timesteps).
+            encoded_lengths: list of int representing the length of each sequence
+                output sequence.
+
+        Returns:
+            packed list containing batch number of sentences (Hypotheses).
+        """
+        # Preserve decoder and joint training state
+        decoder_training_state = self.decoder.training
+        joint_training_state = self.joint.training
+
+        with torch.no_grad():
+            # Apply optional preprocessing
+            encoder_output = encoder_output.transpose(1, 2)  # (B, T, D)
+            logitlen = encoded_lengths
+
+            self.decoder.eval()
+            self.joint.eval()
+
+            with self.decoder.as_frozen(), self.joint.as_frozen():
+                inseq = encoder_output  # [B, T, D]
+                hypotheses, timesteps, alignments = self._greedy_decode(inseq, logitlen, device=inseq.device)
+
+            # Pack the hypotheses results
+            packed_result = pack_hypotheses(hypotheses, timesteps, logitlen, alignments=alignments)
+
+            del hypotheses, timesteps
+
+        self.decoder.train(decoder_training_state)
+        self.joint.train(joint_training_state)
+
+        return (packed_result,)
+
+    def _greedy_decode_blank_as_pad(self, x: torch.Tensor, out_len: torch.Tensor, device: torch.device):
+        with torch.no_grad():
+            # x: [B, T, D]
+            # out_len: [B]
+            # device: torch.device
+
+            # Initialize state
+            hidden = None
+            batchsize = x.shape[0]
+
+            # Output string buffer
+            label = [[] for _ in range(batchsize)]
+            timesteps = [[] for _ in range(batchsize)]
+
+            # If alignments need to be preserved, register a danling list to hold the values
+            if self.preserve_alignments:
+                # alignments is a 3-dimensional dangling list representing B x T x U
+                alignments = []
+                for _ in range(batchsize):
+                    alignments.append([[]])
+            else:
+                alignments = None
+
+            # Last Label buffer + Last Label without blank buffer
+            # batch level equivalent of the last_label
+            last_label = torch.full([batchsize, 1], fill_value=self._blank_index, dtype=torch.long, device=device)
+
+            # Mask buffers
+            blank_mask = torch.full([batchsize], fill_value=0, dtype=torch.bool, device=device)
+
+            # Get max sequence length
+            max_out_len = out_len.max()
+            for time_idx in range(max_out_len):
+                f = x.narrow(dim=1, start=time_idx, length=1)  # [B, 1, D]
+
+                # Prepare t timestamp batch variables
+                not_blank = True
+                symbols_added = 0
+
+                # Reset blank mask
+                blank_mask.mul_(False)
+
+                # Update blank mask with time mask
+                # Batch: [B, T, D], but Bi may have seq len < max(seq_lens_in_batch)
+                # Forcibly mask with "blank" tokens, for all sample where current time step T > seq_len
+                blank_mask = time_idx >= out_len
+                # Start inner loop
+                while not_blank and (self.max_symbols is None or symbols_added < self.max_symbols):
+
+                    # Batch prediction and joint network steps
+                    # If very first prediction step, submit SOS tag (blank) to pred_step.
+                    # This feeds a zero tensor as input to AbstractRNNTDecoder to prime the state
+                    if time_idx == 0 and symbols_added == 0:
+                        g, hidden_prime = self._pred_step(self._SOS, hidden, batch_size=batchsize)
+                    else:
+                        # Perform batch step prediction of decoder, getting new states and scores ("g")
+                        g, hidden_prime = self._pred_step(last_label, hidden, batch_size=batchsize)
+
+                    # Batched joint step - Output = [B, V + 1]
+                    logp = self._joint_step(f, g, log_normalize=None)[:, 0, 0, :]
+
+                    if logp.dtype != torch.float32:
+                        logp = logp.float()
+
+                    # Get index k, of max prob for batch
+                    v, k = logp.max(1)
+                    del v, g
+
+                    # Update blank mask with current predicted blanks
+                    # This is accumulating blanks over all time steps T and all target steps min(max_symbols, U)
+                    k_is_blank = k == self._blank_index
+                    blank_mask.bitwise_or_(k_is_blank)
+
+                    del k_is_blank
+
+                    # If preserving alignments, check if sequence length of sample has been reached
+                    # before adding alignment
+                    if self.preserve_alignments:
+                        # Insert ids into last timestep per sample
+                        logp_vals = logp.to('cpu').max(1)[1]
+                        for batch_idx in range(batchsize):
+                            if time_idx < out_len[batch_idx]:
+                                alignments[batch_idx][-1].append(logp_vals[batch_idx])
+                        del logp_vals
+                    del logp
+
+                    # If all samples predict / have predicted prior blanks, exit loop early
+                    # This is equivalent to if single sample predicted k
+                    if blank_mask.all():
+                        not_blank = False
+
+                        # If preserving alignments, convert the current Uj alignments into a torch.Tensor
+                        # Then preserve U at current timestep Ti
+                        # Finally, forward the timestep history to Ti+1 for that sample
+                        # All of this should only be done iff the current time index <= sample-level AM length.
+                        # Otherwise ignore and move to next sample / next timestep.
+                        if self.preserve_alignments:
+
+                            # convert Ti-th logits into a torch array
+                            for batch_idx in range(batchsize):
+
+                                # this checks if current timestep <= sample-level AM length
+                                # If current timestep > sample-level AM length, no alignments will be added
+                                # Therefore the list of Uj alignments is empty here.
+                                if len(alignments[batch_idx][-1]) > 0:
+                                    alignments[batch_idx].append([])  # blank buffer for next timestep
+                    else:
+                        # Collect batch indices where blanks occurred now/past
+                        blank_indices = []
+                        if hidden is not None:
+                            blank_indices = (blank_mask == 1).nonzero(as_tuple=False)
+
+                        # Recover prior state for all samples which predicted blank now/past
+                        if hidden is not None:
+                            # LSTM has 2 states
+                            for state_id in range(len(hidden)):
+                                hidden_prime[state_id][:, blank_indices, :] = hidden[state_id][:, blank_indices, :]
+
+                        # Recover prior predicted label for all samples which predicted blank now/past
+                        k[blank_indices] = last_label[blank_indices, 0]
+
+                        # Update new label and hidden state for next iteration
+                        last_label = k.clone().view(-1, 1)
+                        hidden = hidden_prime
+
+                        # Update predicted labels, accounting for time mask
+                        # If blank was predicted even once, now or in the past,
+                        # Force the current predicted label to also be blank
+                        # This ensures that blanks propogate across all timesteps
+                        # once they have occured (normally stopping condition of sample level loop).
+                        for kidx, ki in enumerate(k):
+                            if blank_mask[kidx] == 0:
+                                label[kidx].append(ki)
+                                timesteps[kidx].append(time_idx)
+
+                        symbols_added += 1
+
+            # Remove trailing empty list of alignments at T_{am-len} x Uj
+            if self.preserve_alignments:
+                for batch_idx in range(batchsize):
+                    if len(alignments[batch_idx][-1]) == 0:
+                        del alignments[batch_idx][-1]
+
+        return label, timesteps, alignments
+
+    @torch.no_grad()
+    def _greedy_decode_masked(self, x: torch.Tensor, out_len: torch.Tensor, device: torch.device):
+        # x: [B, T, D]
+        # out_len: [B]
+        # device: torch.device
+
+        # Initialize state
+        hidden = None
+        batchsize = x.shape[0]
+
+        # Output string buffer
+        label = [[] for _ in range(batchsize)]
+        timesteps = [[] for _ in range(batchsize)]
+
+        # If alignments need to be preserved, register a danling list to hold the values
+        if self.preserve_alignments:
+            # alignments is a 3-dimensional dangling list representing B x T x U
+            alignments = []
+            for _ in range(batchsize):
+                alignments.append([[]])
+        else:
+            alignments = None
+
+        # Last Label buffer + Last Label without blank buffer
+        # batch level equivalent of the last_label
+        last_label = torch.full([batchsize, 1], fill_value=self._blank_index, dtype=torch.long, device=device)
+        last_label_without_blank = last_label.clone()
+
+        # Mask buffers
+        blank_mask = torch.full([batchsize], fill_value=0, dtype=torch.bool, device=device)
+
+        # Get max sequence length
+        max_out_len = out_len.max()
+        for time_idx in range(max_out_len):
+            f = x.narrow(dim=1, start=time_idx, length=1)  # [B, 1, D]
+
+            # Prepare t timestamp batch variables
+            not_blank = True
+            symbols_added = 0
+
+            # Reset blank mask
+            blank_mask.mul_(False)
+
+            # Update blank mask with time mask
+            # Batch: [B, T, D], but Bi may have seq len < max(seq_lens_in_batch)
+            # Forcibly mask with "blank" tokens, for all sample where current time step T > seq_len
+            blank_mask = time_idx >= out_len
+
+            # Start inner loop
+            while not_blank and (self.max_symbols is None or symbols_added < self.max_symbols):
+                # Batch prediction and joint network steps
+                # If very first prediction step, submit SOS tag (blank) to pred_step.
+                # This feeds a zero tensor as input to AbstractRNNTDecoder to prime the state
+                if time_idx == 0 and symbols_added == 0:
+                    g, hidden_prime = self._pred_step(self._SOS, hidden, batch_size=batchsize)
+                else:
+                    # Set a dummy label for the blank value
+                    # This value will be overwritten by "blank" again the last label update below
+                    # This is done as vocabulary of prediction network does not contain "blank" token of RNNT
+                    last_label_without_blank_mask = last_label == self._blank_index
+                    last_label_without_blank[last_label_without_blank_mask] = 0  # temp change of label
+                    last_label_without_blank[~last_label_without_blank_mask] = last_label[
+                        ~last_label_without_blank_mask
+                    ]
+
+                    # Perform batch step prediction of decoder, getting new states and scores ("g")
+                    g, hidden_prime = self._pred_step(last_label_without_blank, hidden, batch_size=batchsize)
+
+                # Batched joint step - Output = [B, V + 1]
+                logp = self._joint_step(f, g, log_normalize=None)[:, 0, 0, :]
+
+                if logp.dtype != torch.float32:
+                    logp = logp.float()
+
+                # Get index k, of max prob for batch
+                v, k = logp.max(1)
+                del v, g
+
+                # Update blank mask with current predicted blanks
+                # This is accumulating blanks over all time steps T and all target steps min(max_symbols, U)
+                k_is_blank = k == self._blank_index
+                blank_mask.bitwise_or_(k_is_blank)
+
+                # If preserving alignments, check if sequence length of sample has been reached
+                # before adding alignment
+                if self.preserve_alignments:
+                    # Insert ids into last timestep per sample
+                    logp_vals = logp.to('cpu').max(1)[1]
+                    for batch_idx in range(batchsize):
+                        if time_idx < out_len[batch_idx]:
+                            alignments[batch_idx][-1].append(logp_vals[batch_idx])
+                    del logp_vals
+                del logp
+
+                # If all samples predict / have predicted prior blanks, exit loop early
+                # This is equivalent to if single sample predicted k
+                if blank_mask.all():
+                    not_blank = False
+
+                    # If preserving alignments, convert the current Uj alignments into a torch.Tensor
+                    # Then preserve U at current timestep Ti
+                    # Finally, forward the timestep history to Ti+1 for that sample
+                    # All of this should only be done iff the current time index <= sample-level AM length.
+                    # Otherwise ignore and move to next sample / next timestep.
+                    if self.preserve_alignments:
+
+                        # convert Ti-th logits into a torch array
+                        for batch_idx in range(batchsize):
+
+                            # this checks if current timestep <= sample-level AM length
+                            # If current timestep > sample-level AM length, no alignments will be added
+                            # Therefore the list of Uj alignments is empty here.
+                            if len(alignments[batch_idx][-1]) > 0:
+                                alignments[batch_idx].append([])  # blank buffer for next timestep
+                else:
+                    # Collect batch indices where blanks occurred now/past
+                    blank_indices = []
+                    if hidden is not None:
+                        blank_indices = (blank_mask == 1).nonzero(as_tuple=False)
+
+                    # Recover prior state for all samples which predicted blank now/past
+                    if hidden is not None:
+                        # LSTM has 2 states
+                        for state_id in range(len(hidden)):
+                            hidden_prime[state_id][:, blank_indices, :] = hidden[state_id][:, blank_indices, :]
+
+                    # Recover prior predicted label for all samples which predicted blank now/past
+                    k[blank_indices] = last_label[blank_indices, 0]
+
+                    # Update new label and hidden state for next iteration
+                    last_label = k.view(-1, 1)
+                    hidden = hidden_prime
+
+                    # Update predicted labels, accounting for time mask
+                    # If blank was predicted even once, now or in the past,
+                    # Force the current predicted label to also be blank
+                    # This ensures that blanks propogate across all timesteps
+                    # once they have occured (normally stopping condition of sample level loop).
+                    for kidx, ki in enumerate(k):
+                        if blank_mask[kidx] == 0:
+                            label[kidx].append(ki)
+                            timesteps[kidx].append(time_idx)
+
+                symbols_added += 1
+
+        # Remove trailing empty list of alignments at T_{am-len} x Uj
+        if self.preserve_alignments:
+            for batch_idx in range(batchsize):
+                if len(alignments[batch_idx][-1]) == 0:
+                    del alignments[batch_idx][-1]
+
+        return label, timesteps, alignments
+
+
+@dataclass
+class GreedyRNNTInferConfig:
+    max_symbols_per_step: Optional[int] = None
+    preserve_alignments: bool = False
+
+
+@dataclass
+class GreedyBatchedRNNTInferConfig:
+    max_symbols_per_step: Optional[int] = None
+    preserve_alignments: bool = False

tts/nemo/collections/asr/parts/submodules/spectr_augment.py

@@ -0,0 +1,152 @@
+# 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 random
+
+import torch
+import torch.nn as nn
+
+from nemo.core.classes import Typing, typecheck
+from nemo.core.neural_types import NeuralType, SpectrogramType
+
+
+class SpecAugment(nn.Module, Typing):
+    """
+    Zeroes out(cuts) random continuous horisontal or
+    vertical segments of the spectrogram as described in
+    SpecAugment (https://arxiv.org/abs/1904.08779).
+
+    params:
+    freq_masks - how many frequency segments should be cut
+    time_masks - how many time segments should be cut
+    freq_width - maximum number of frequencies to be cut in one segment
+    time_width - maximum number of time steps to be cut in one segment.
+        Can be a positive integer or a float value in the range [0, 1].
+        If positive integer value, defines maximum number of time steps
+        to be cut in one segment.
+        If a float value, defines maximum percentage of timesteps that
+        are cut adaptively.
+    """
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input types
+        """
+        return {"input_spec": NeuralType(('B', 'D', 'T'), SpectrogramType())}
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output types
+        """
+        return {"augmented_spec": NeuralType(('B', 'D', 'T'), SpectrogramType())}
+
+    def __init__(
+        self, freq_masks=0, time_masks=0, freq_width=10, time_width=10, rng=None, mask_value=0.0,
+    ):
+        super(SpecAugment, self).__init__()
+
+        self._rng = random.Random() if rng is None else rng
+
+        self.freq_masks = freq_masks
+        self.time_masks = time_masks
+
+        self.freq_width = freq_width
+        self.time_width = time_width
+
+        self.mask_value = mask_value
+
+        if isinstance(time_width, int):
+            self.adaptive_temporal_width = False
+        else:
+            if time_width > 1.0 or time_width < 0.0:
+                raise ValueError("If `time_width` is a float value, must be in range [0, 1]")
+
+            self.adaptive_temporal_width = True
+
+    @typecheck()
+    @torch.no_grad()
+    def forward(self, input_spec):
+        sh = input_spec.shape
+
+        if self.adaptive_temporal_width:
+            time_width = max(1, int(sh[2] * self.time_width))
+        else:
+            time_width = self.time_width
+
+        for idx in range(sh[0]):
+            for i in range(self.freq_masks):
+                x_left = self._rng.randint(0, sh[1] - self.freq_width)
+
+                w = self._rng.randint(0, self.freq_width)
+
+                input_spec[idx, x_left : x_left + w, :] = self.mask_value
+
+            for i in range(self.time_masks):
+                y_left = self._rng.randint(0, sh[2] - time_width)
+
+                w = self._rng.randint(0, time_width)
+
+                input_spec[idx, :, y_left : y_left + w] = self.mask_value
+
+        return input_spec
+
+
+class SpecCutout(nn.Module, Typing):
+    """
+    Zeroes out(cuts) random rectangles in the spectrogram
+    as described in (https://arxiv.org/abs/1708.04552).
+
+    params:
+    rect_masks - how many rectangular masks should be cut
+    rect_freq - maximum size of cut rectangles along the frequency dimension
+    rect_time - maximum size of cut rectangles along the time dimension
+    """
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input types
+        """
+        return {"input_spec": NeuralType(('B', 'D', 'T'), SpectrogramType())}
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output types
+        """
+        return {"augmented_spec": NeuralType(('B', 'D', 'T'), SpectrogramType())}
+
+    def __init__(self, rect_masks=0, rect_time=5, rect_freq=20, rng=None):
+        super(SpecCutout, self).__init__()
+
+        self._rng = random.Random() if rng is None else rng
+
+        self.rect_masks = rect_masks
+        self.rect_time = rect_time
+        self.rect_freq = rect_freq
+
+    @typecheck()
+    @torch.no_grad()
+    def forward(self, input_spec):
+        sh = input_spec.shape
+
+        for idx in range(sh[0]):
+            for i in range(self.rect_masks):
+                rect_x = self._rng.randint(0, sh[1] - self.rect_freq)
+                rect_y = self._rng.randint(0, sh[2] - self.rect_time)
+
+                w_x = self._rng.randint(0, self.rect_freq)
+                w_y = self._rng.randint(0, self.rect_time)
+
+                input_spec[idx, rect_x : rect_x + w_x, rect_y : rect_y + w_y] = 0.0
+
+        return input_spec

tts/nemo/collections/asr/parts/submodules/subsampling.py

@@ -0,0 +1,138 @@
+# 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 math
+
+import torch
+import torch.nn as nn
+
+
+class ConvSubsampling(torch.nn.Module):
+    """Convolutional subsampling which supports VGGNet and striding approach introduced in:
+    VGGNet Subsampling: https://arxiv.org/pdf/1910.12977.pdf
+    Striding Subsampling:
+        "Speech-Transformer: A No-Recurrence Sequence-to-Sequence Model for Speech Recognition" by Linhao Dong et al.
+    Args:
+        subsampling (str): The subsampling technique from {"vggnet", "striding"}
+        subsampling_factor (int): The subsampling factor which should be a power of 2
+        feat_in (int): size of the input features
+        feat_out (int): size of the output features
+        conv_channels (int): Number of channels for the convolution layers.
+        activation (Module): activation function, default is nn.ReLU()
+    """
+
+    def __init__(self, subsampling, subsampling_factor, feat_in, feat_out, conv_channels, activation=nn.ReLU()):
+        super(ConvSubsampling, self).__init__()
+        self._subsampling = subsampling
+
+        if subsampling_factor % 2 != 0:
+            raise ValueError("Sampling factor should be a multiply of 2!")
+        self._sampling_num = int(math.log(subsampling_factor, 2))
+
+        in_channels = 1
+        layers = []
+        if subsampling == 'vggnet':
+            self._padding = 0
+            self._stride = 2
+            self._kernel_size = 2
+            self._ceil_mode = True
+
+            for i in range(self._sampling_num):
+                layers.append(
+                    torch.nn.Conv2d(
+                        in_channels=in_channels, out_channels=conv_channels, kernel_size=3, stride=1, padding=1
+                    )
+                )
+                layers.append(activation)
+                layers.append(
+                    torch.nn.Conv2d(
+                        in_channels=conv_channels, out_channels=conv_channels, kernel_size=3, stride=1, padding=1
+                    )
+                )
+                layers.append(activation)
+                layers.append(
+                    torch.nn.MaxPool2d(
+                        kernel_size=self._kernel_size,
+                        stride=self._stride,
+                        padding=self._padding,
+                        ceil_mode=self._ceil_mode,
+                    )
+                )
+                in_channels = conv_channels
+        elif subsampling == 'striding':
+            self._padding = 1
+            self._stride = 2
+            self._kernel_size = 3
+            self._ceil_mode = False
+
+            for i in range(self._sampling_num):
+                layers.append(
+                    torch.nn.Conv2d(
+                        in_channels=in_channels,
+                        out_channels=conv_channels,
+                        kernel_size=self._kernel_size,
+                        stride=self._stride,
+                        padding=self._padding,
+                    )
+                )
+                layers.append(activation)
+                in_channels = conv_channels
+        else:
+            raise ValueError(f"Not valid sub-sampling: {subsampling}!")
+
+        in_length = feat_in
+        for i in range(self._sampling_num):
+            out_length = calc_length(
+                length=int(in_length),
+                padding=self._padding,
+                kernel_size=self._kernel_size,
+                stride=self._stride,
+                ceil_mode=self._ceil_mode,
+            )
+            in_length = out_length
+
+        self.out = torch.nn.Linear(conv_channels * out_length, feat_out)
+        self.conv = torch.nn.Sequential(*layers)
+
+    def forward(self, x, lengths):
+        x = x.unsqueeze(1)
+        x = self.conv(x)
+        b, c, t, f = x.size()
+        x = self.out(x.transpose(1, 2).contiguous().view(b, t, c * f))
+
+        # TODO: improve the performance of length calculation
+        new_lengths = lengths
+        for i in range(self._sampling_num):
+            new_lengths = [
+                calc_length(
+                    length=int(length),
+                    padding=self._padding,
+                    kernel_size=self._kernel_size,
+                    stride=self._stride,
+                    ceil_mode=self._ceil_mode,
+                )
+                for length in new_lengths
+            ]
+
+        new_lengths = torch.IntTensor(new_lengths).to(lengths.device)
+        return x, new_lengths
+
+
+def calc_length(length, padding, kernel_size, stride, ceil_mode):
+    """ Calculates the output length of a Tensor passed through a convolution or max pooling layer"""
+    if ceil_mode:
+        length = math.ceil((length + (2 * padding) - (kernel_size - 1) - 1) / float(stride) + 1)
+    else:
+        length = math.floor((length + (2 * padding) - (kernel_size - 1) - 1) / float(stride) + 1)
+    return length

tts/nemo/collections/asr/parts/submodules/wav2vec.py

@@ -0,0 +1,222 @@
+# 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.
+
+import math
+from typing import List, Tuple
+
+import torch
+from torch import nn
+
+from nemo.collections.asr.models.wav2vec.wav2vec_config import Wav2VecConvExtractorMode, Wav2VecTransformerConfig
+
+
+class TransposeLast(torch.nn.Module):
+    """
+    Transposes last dimension. Useful for adding to a sequential block.
+    """
+
+    def forward(self, x):
+        return x.transpose(-2, -1)
+
+
+class SamePad(torch.nn.Module):
+    def __init__(self, kernel_size):
+        super().__init__()
+        self.remove = kernel_size % 2 == 0
+
+    def forward(self, x):
+        if self.remove:
+            x = x[:, :, :-1]
+        return x
+
+
+class ConvFeatureEncoder(nn.Module):
+    """
+        Converts input raw audio into features for downstream transformer model.
+        Uses 1D convolutional blocks with GeLU activation.
+    """
+
+    def __init__(
+        self,
+        conv_layers: List[Tuple[int, int, int]],
+        mode: Wav2VecConvExtractorMode = Wav2VecConvExtractorMode.default,
+        conv_bias: bool = False,
+    ):
+        super().__init__()
+
+        def block(
+            n_in, n_out, k, stride, is_layer_norm=False, is_group_norm=False, conv_bias=False,
+        ):
+            def make_conv():
+                conv = nn.Conv1d(n_in, n_out, k, stride=stride, bias=conv_bias)
+                nn.init.kaiming_normal_(conv.weight)
+                return conv
+
+            assert (is_layer_norm and is_group_norm) is False, "layer norm and group norm are exclusive"
+
+            if is_layer_norm:
+                return nn.Sequential(
+                    make_conv(),
+                    nn.Sequential(TransposeLast(), nn.LayerNorm(dim, elementwise_affine=True), TransposeLast()),
+                    nn.GELU(),
+                )
+            elif is_group_norm:
+                return nn.Sequential(make_conv(), nn.GroupNorm(dim, dim, affine=True), nn.GELU(),)
+            else:
+                return nn.Sequential(make_conv(), nn.GELU())
+
+        in_d = 1
+        self.conv_layers = nn.ModuleList()
+        for i, cl in enumerate(conv_layers):
+            assert len(cl) == 3, "invalid conv definition: " + str(cl)
+            (dim, k, stride) = cl
+
+            self.conv_layers.append(
+                block(
+                    in_d,
+                    dim,
+                    k,
+                    stride,
+                    is_layer_norm=mode is Wav2VecConvExtractorMode.layer_norm,
+                    is_group_norm=mode is Wav2VecConvExtractorMode.default and i == 0,
+                    conv_bias=conv_bias,
+                )
+            )
+            in_d = dim
+
+    def forward(self, x):
+        # BxT -> BxCxT
+        x = x.unsqueeze(1)
+        for conv in self.conv_layers:
+            x = conv(x)
+        return x
+
+
+class Wav2VecTransformerEncoder(nn.Module):
+    def __init__(self, cfg: Wav2VecTransformerConfig):
+        super().__init__()
+
+        conv_cfg = cfg.conv
+
+        self.dropout = cfg.dropout
+        self.embedding_dim = cfg.encoder.embedding_dim
+        self.layer_norm_first = cfg.encoder.layer_norm_first
+
+        # positional convolutional embeddings
+        self.pos_conv = nn.Conv1d(
+            self.embedding_dim,
+            self.embedding_dim,
+            kernel_size=conv_cfg.conv_pos,
+            padding=conv_cfg.conv_pos // 2,
+            groups=conv_cfg.conv_pos_groups,
+        )
+
+        self.feature_dropout = nn.Dropout(self.dropout)
+
+        dropout = 0
+        std = math.sqrt((4 * (1.0 - dropout)) / (conv_cfg.conv_pos * self.embedding_dim))
+        nn.init.normal_(self.pos_conv.weight, mean=0, std=std)
+        nn.init.constant_(self.pos_conv.bias, 0)
+
+        self.pos_conv = nn.utils.weight_norm(self.pos_conv, name="weight", dim=2)
+        self.pos_conv = nn.Sequential(self.pos_conv, SamePad(conv_cfg.conv_pos), nn.GELU())
+
+        encoder_cfg = cfg.encoder
+        self.transformer_encoder = nn.TransformerEncoder(
+            encoder_layer=nn.TransformerEncoderLayer(
+                d_model=self.embedding_dim,
+                nhead=encoder_cfg.num_attention_heads,
+                dim_feedforward=encoder_cfg.ffn_embedding_dim,
+                dropout=self.dropout,
+                activation=encoder_cfg.activation_fn.value,
+            ),
+            num_layers=encoder_cfg.encoder_layers,
+        )
+        self.layer_norm = nn.LayerNorm(self.embedding_dim)
+        self.apply(init_bert_params)
+
+    def forward(self, x, padding_mask=None):
+        x = self.extract_features(x, padding_mask)
+
+        if self.layer_norm_first:
+            x = self.layer_norm(x)
+
+        return x
+
+    def extract_features(self, x, padding_mask=None):
+
+        if padding_mask is not None:
+            x[padding_mask] = 0
+
+        x_conv = self.pos_conv(x.transpose(1, 2))
+        x_conv = x_conv.transpose(1, 2)
+        x += x_conv
+
+        if not self.layer_norm_first:
+            x = self.layer_norm(x)
+
+        x = self.feature_dropout(x)
+
+        # B x T x C -> T x B x C
+        x = x.transpose(0, 1)
+
+        x = self.transformer_encoder(x, src_key_padding_mask=padding_mask)
+
+        # T x B x C -> B x T x C
+        x = x.transpose(0, 1)
+
+        return x
+
+
+class GradMultiply(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, scale):
+        ctx.scale = scale
+        res = x.new(x)
+        return res
+
+    @staticmethod
+    def backward(ctx, grad):
+        return grad * ctx.scale, None
+
+
+def init_bert_params(module):
+    """
+    Initialize the weights specific to the BERT Model.
+    This overrides the default initializations depending on the specified arguments.
+        1. If normal_init_linear_weights is set then weights of linear
+           layer will be initialized using the normal distribution and
+           bias will be set to the specified value.
+        2. If normal_init_embed_weights is set then weights of embedding
+           layer will be initialized using the normal distribution.
+        3. If normal_init_proj_weights is set then weights of
+           in_project_weight for MultiHeadAttention initialized using
+           the normal distribution (to be validated).
+    """
+
+    if isinstance(module, nn.Linear):
+        module.weight.data.normal_(mean=0.0, std=0.02)
+        if module.bias is not None:
+            module.bias.data.zero_()
+    if isinstance(module, nn.Embedding):
+        module.weight.data.normal_(mean=0.0, std=0.02)
+        if module.padding_idx is not None:
+            module.weight.data[module.padding_idx].zero_()
+    if isinstance(module, nn.TransformerEncoderLayer):
+        module.self_attn.in_proj_weight.data.normal_(mean=0.0, std=0.02)

tts/nemo/collections/asr/parts/utils/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tts/nemo/collections/asr/parts/utils/activations.py

@@ -0,0 +1,24 @@
+# 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 torch.nn as nn
+
+__all__ = ['Swish']
+
+
+class Swish(nn.SiLU):
+    """
+    Swish activation function introduced in 'https://arxiv.org/abs/1710.05941'
+    Mathematically identical to SiLU. See note in nn.SiLU for references.
+    """

tts/nemo/collections/asr/parts/utils/asr_module_utils.py

@@ -0,0 +1,82 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional
+
+from omegaconf import DictConfig, open_dict
+
+from nemo.collections.asr.modules import conv_asr
+from nemo.collections.asr.parts.submodules import jasper
+from nemo.utils import logging
+
+
+def change_conv_asr_se_context_window(model: 'ASRModel', context_window: int, update_config: bool = True):
+    """
+    Update the context window of the SqueezeExcitation module if the provided model contains an
+    `encoder` which is an instance of `ConvASREncoder`.
+
+    Args:
+        model: A subclass of `ASRModel`, itself a subclass of `ModelPT`.
+        context_window:  An integer representing the number of input timeframes that will be used
+            to compute the context. Each timeframe corresponds to a single window stride of the
+            STFT features.
+
+            Say the window_stride = 0.01s, then a context window of 128 represents 128 * 0.01 s
+            of context to compute the Squeeze step.
+        update_config: Whether to update the config or not with the new context window.
+    """
+    if update_config and not hasattr(model.cfg, 'encoder'):
+        logging.info(
+            "Could not change the context window in SqueezeExcite module "
+            "since the model provided does not contain an `encoder` module in its config."
+        )
+        return
+
+    if not isinstance(model.encoder, conv_asr.ConvASREncoder):
+        logging.info(
+            f"Could not change the context window in SqueezeExcite module "
+            f"since the `encoder` module is not an instance of `ConvASREncoder`.\n"
+            f"Provided encoder class = {model.encoder.__class__.__name__}"
+        )
+        return
+
+    enc_cfg = model.cfg.encoder if update_config else None
+
+    if enc_cfg is not None:
+        with open_dict(enc_cfg):
+            _update_se_context_window(model, context_window, cfg=enc_cfg)
+    else:
+        _update_se_context_window(model, context_window)
+
+    # Update model config
+    if update_config:
+        model.cfg.encoder = enc_cfg
+
+
+def _update_se_context_window(model: 'ASRModel', context_window: int, cfg: Optional[DictConfig] = None):
+    jasper_block_counter = -1
+    for name, m in model.named_modules():
+        if type(m) == jasper.JasperBlock:
+            jasper_block_counter += 1
+
+        if type(m) == jasper.MaskedConv1d:
+            if m.conv.stride[0] > 1 and 'mconv' in name:
+                context_window = context_window // m.conv.stride[0]
+
+        if type(m) == jasper.SqueezeExcite:
+            m.change_context_window(context_window=context_window)
+
+            # update config
+            if cfg is not None:
+                cfg.jasper[jasper_block_counter].se_context_size = context_window

tts/nemo/collections/asr/parts/utils/nmse_clustering.py

@@ -0,0 +1,271 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Copyright (c) 2007-2020 The scikit-learn developers.
+
+# BSD 3-Clause License
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+# This file is part of https://github.com/scikit-learn/scikit-learn/blob/114616d9f6ce9eba7c1aacd3d4a254f868010e25/sklearn/manifold/_spectral_embedding.py and
+# https://github.com/tango4j/Auto-Tuning-Spectral-Clustering.
+
+
+import numpy as np
+import scipy
+from scipy import sparse
+from sklearn.cluster import SpectralClustering as sklearn_SpectralClustering
+from sklearn.metrics.pairwise import cosine_similarity
+from sklearn.preprocessing import MinMaxScaler
+
+scaler = MinMaxScaler(feature_range=(0, 1))
+
+
+def get_kneighbors_conn(X_dist, p_neighbors):
+    X_dist_out = np.zeros_like(X_dist)
+    for i, line in enumerate(X_dist):
+        sorted_idx = np.argsort(line)
+        sorted_idx = sorted_idx[::-1]
+        indices = sorted_idx[:p_neighbors]
+        X_dist_out[indices, i] = 1
+    return X_dist_out
+
+
+def get_X_conn_from_dist(X_dist_raw, p_neighbors):
+    X_r = get_kneighbors_conn(X_dist_raw, p_neighbors)
+    X_conn_from_dist = 0.5 * (X_r + X_r.T)
+    return X_conn_from_dist
+
+
+def isFullyConnected(X_conn_from_dist):
+    """ Return whether the graph is connected (True) or Not (False).
+    Parameters
+    ----------
+    graph : {array-like, sparse matrix} of shape (n_samples, n_samples)
+        Adjacency matrix of the graph, non-zero weight means an edge
+        between the nodes.
+    Returns
+    -------
+    is_connected : bool
+        True means the graph is fully connected and False means not.
+    """
+    gC = _graph_connected_component(X_conn_from_dist, 0).sum() == X_conn_from_dist.shape[0]
+    return gC
+
+
+def gc_thres_min_gc(mat, max_n, n_list):
+    p_neighbors, index = 1, 0
+    X_conn_from_dist = get_X_conn_from_dist(mat, p_neighbors)
+    fully_connected = isFullyConnected(X_conn_from_dist)
+    for i, p_neighbors in enumerate(n_list):
+        fully_connected = isFullyConnected(X_conn_from_dist)
+        X_conn_from_dist = get_X_conn_from_dist(mat, p_neighbors)
+        if fully_connected or p_neighbors > max_n:
+            break
+
+    return X_conn_from_dist, p_neighbors
+
+
+def _graph_connected_component(graph, node_id):
+    """Find the largest graph connected components that contains one
+    given node.
+    Parameters
+    ----------
+    graph : array-like of shape (n_samples, n_samples)
+        Adjacency matrix of the graph, non-zero weight means an edge
+        between the nodes.
+    node_id : int
+        The index of the query node of the graph.
+    Returns
+    -------
+    connected_components_matrix : array-like of shape (n_samples,)
+        An array of bool value indicating the indexes of the nodes
+        belonging to the largest connected components of the given query
+        node.
+    """
+    n_node = graph.shape[0]
+    if sparse.issparse(graph):
+        graph = graph.tocsr()
+    connected_nodes = np.zeros(n_node, dtype=np.bool)
+    nodes_to_explore = np.zeros(n_node, dtype=np.bool)
+    nodes_to_explore[node_id] = True
+    for _ in range(n_node):
+        last_num_component = connected_nodes.sum()
+        np.logical_or(connected_nodes, nodes_to_explore, out=connected_nodes)
+        if last_num_component >= connected_nodes.sum():
+            break
+        indices = np.where(nodes_to_explore)[0]
+        nodes_to_explore.fill(False)
+        for i in indices:
+            if sparse.issparse(graph):
+                neighbors = graph[i].toarray().ravel()
+            else:
+                neighbors = graph[i]
+            np.logical_or(nodes_to_explore, neighbors, out=nodes_to_explore)
+    return connected_nodes
+
+
+def getLaplacian(X):
+    """
+    returns laplacian matrix of X
+    """
+    X[np.diag_indices(X.shape[0])] = 0
+    A = X
+    D = np.sum(np.abs(A), axis=1)
+    D = np.diag(D)
+    L = D - A
+    return L
+
+
+def eig_decompose(L, k):
+    try:
+        lambdas, eig_vecs = scipy.linalg.eigh(L)
+    except:
+        try:
+            lambdas = scipy.linalg.eigvals(L)
+            eig_vecs = None
+        except:
+            lambdas, eig_vecs = scipy.sparse.linalg.eigsh(L)  ### Inaccurate results
+    return lambdas, eig_vecs
+
+
+def getLamdaGaplist(lambdas):
+    lambda_gap_list = []
+    for i in range(len(lambdas) - 1):
+        lambda_gap_list.append(float(lambdas[i + 1]) - float(lambdas[i]))
+    return lambda_gap_list
+
+
+def estimate_num_of_spkrs(X_conn, SPK_MAX):
+    """
+    Estimate number of speakers using eigen decompose on Laplacian Matrix
+    """
+    L = getLaplacian(X_conn)
+    lambdas, eig_vals = eig_decompose(L, k=X_conn.shape[0])
+    lambdas = np.sort(lambdas)
+    lambda_gap_list = getLamdaGaplist(lambdas)
+    num_of_spk = np.argmax(lambda_gap_list[: min(SPK_MAX, len(lambda_gap_list))]) + 1
+    return num_of_spk, lambdas, lambda_gap_list
+
+
+def NMEanalysis(mat, SPK_MAX, max_rp_threshold=0.250, sparse_search=True, search_p_volume=20, fixed_thres=None):
+    """
+    Normalized Eigengap analysis to get p_neighbours and estimated num of speakers
+    Reference: https://ieeexplore.ieee.org/stamp/stamp.jsp?arnumber=8937487&tag=1 
+    """
+    eps = 1e-10
+    eig_ratio_list = []
+    if fixed_thres:
+        p_neighbors_list = [int(mat.shape[0] * fixed_thres)]
+        max_N = p_neighbors_list[0]
+    else:
+        max_N = int(mat.shape[0] * max_rp_threshold)
+        if sparse_search:
+            N = min(max_N, search_p_volume)
+            p_neighbors_list = list(np.linspace(1, max_N, N, endpoint=True).astype(int))
+        else:
+            p_neighbors_list = list(range(1, max_N))
+
+    est_spk_n_dict = {}
+    for p_neighbors in p_neighbors_list:
+        X_conn_from_dist = get_X_conn_from_dist(mat, p_neighbors)
+        est_num_of_spk, lambdas, lambda_gap_list = estimate_num_of_spkrs(X_conn_from_dist, SPK_MAX)
+        est_spk_n_dict[p_neighbors] = (est_num_of_spk, lambdas)
+        arg_sorted_idx = np.argsort(lambda_gap_list[:SPK_MAX])[::-1]
+        max_key = arg_sorted_idx[0]
+        max_eig_gap = lambda_gap_list[max_key] / (max(lambdas) + eps)
+        eig_ratio_value = (p_neighbors / mat.shape[0]) / (max_eig_gap + eps)
+        eig_ratio_list.append(eig_ratio_value)
+
+    index_nn = np.argmin(eig_ratio_list)
+    rp_p_neighbors = p_neighbors_list[index_nn]
+    X_conn_from_dist = get_X_conn_from_dist(mat, rp_p_neighbors)
+    if not isFullyConnected(X_conn_from_dist):
+        X_conn_from_dist, rp_p_neighbors = gc_thres_min_gc(mat, max_N, p_neighbors_list)
+
+    return (
+        X_conn_from_dist,
+        float(rp_p_neighbors / mat.shape[0]),
+        est_spk_n_dict[rp_p_neighbors][0],
+        est_spk_n_dict[rp_p_neighbors][1],
+        rp_p_neighbors,
+    )
+
+
+def get_eigen_matrix(emb):
+    """
+    input
+    emb: embedding
+    
+    output
+    similarity matrix
+    """
+    sim_d = cosine_similarity(emb)
+    scaler.fit(sim_d)
+    sim_d = scaler.transform(sim_d)
+
+    return sim_d
+
+
+def COSclustering(key, emb, oracle_num_speakers=None, max_num_speaker=8, MIN_SAMPLES=6):
+    """
+    input:
+    key (str): speaker uniq name
+    emb (np array): speaker embedding
+    oracle_num_speaker (int or None): oracle number of speakers if known else None
+    max_num_speakers (int): maximum number of clusters to consider for each session
+    MIN_SAMPLES (int): minimum number of samples required for NME clustering, this avoids
+    zero p_neighbour_lists. Default of 6 is selected since  (1/rp_threshold) >= 4.
+    output:
+    Y (List[int]): speaker labels
+    """
+    est_num_spks_out_list = []
+    mat = get_eigen_matrix(emb)
+    if oracle_num_speakers:
+        max_num_speaker = oracle_num_speakers
+
+    X_conn_spkcount, rp_thres_spkcount, est_num_of_spk, lambdas, p_neigh = NMEanalysis(mat, max_num_speaker)
+
+    if emb.shape[0] > MIN_SAMPLES:
+        X_conn_from_dist = get_X_conn_from_dist(mat, p_neigh)
+    else:
+        X_conn_from_dist = mat
+
+    if oracle_num_speakers:
+        est_num_of_spk = oracle_num_speakers
+
+    est_num_spks_out_list.append([key, str(est_num_of_spk)])
+
+    # Perform spectral clustering
+    spectral_model = sklearn_SpectralClustering(
+        affinity='precomputed',
+        eigen_solver='amg',
+        random_state=0,
+        n_jobs=-1,
+        n_clusters=est_num_of_spk,
+        eigen_tol=1e-10,
+    )
+
+    Y = spectral_model.fit_predict(X_conn_from_dist)
+
+    return Y

tts/nemo/collections/asr/parts/utils/numba_utils.py

@@ -0,0 +1,88 @@
+# 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 numpy as np
+from numba import jit
+
+
+def phase_vocoder(D: np.ndarray, rate: float, phi_advance: np.ndarray, scale_buffer: np.ndarray):
+    """
+    Optimized implementation of phase vocoder from Librosa.
+    Reference implementation:
+        - https://librosa.github.io/librosa/generated/librosa.core.phase_vocoder.html
+    Args:
+        D: Complex spectograms of shape [d, t, complex=2].
+        rate: Speed rate, must be float greater than 0.
+        phi_advance: Precomputed phase advance buffer array of length [n_fft + 1]
+        scale_buffer: Precomputed numpy buffer array of length [n_fft + 1]
+    Returns:
+        Complex64 ndarray of shape [d, t / rate, complex=2]
+    """
+    time_steps = np.arange(0, D.shape[1], rate, dtype=np.float)
+
+    # Create an empty output array
+    d_stretch = np.zeros((D.shape[0], len(time_steps)), D.dtype, order='F')
+
+    # Phase accumulator; initialize to the first sample
+    phase_acc = np.angle(D[:, 0])
+
+    # Pad 0 columns to simplify boundary logic
+    D = np.pad(D, [(0, 0), (0, 2)], mode='constant')
+
+    d_stretch = _phase_vocoder_kernel(D, time_steps, phi_advance, d_stretch, phase_acc, scale_buffer)
+
+    return d_stretch
+
+
+@jit(nopython=True, nogil=True)
+def _phase_vocoder_kernel(D, time_steps, phi_advance, d_stretch, phase_acc, scale_buffer):
+    """
+    Numba optimized kernel to compute the phase vocoder step.
+    Args:
+        D: Complex spectograms of shape [d, t, complex=2].
+        rate: Speed rate, must be float greater than 0.
+        time_steps: Numpy ndarray of linearly spaced time steps, shape = [t]
+        phi_advance: Precomputed phase advance buffer array of length [n_fft + 1]
+        d_stretch: Output complex matrix of shape [d, t / rate, complex=2]
+        phase_acc: Phase accumulator initialized to first sample of shape [d, complex=2]
+        scale_buffer: Precomputed numpy buffer array of length [n_fft + 1]
+    Returns:
+        Complex64 ndarray of shape [d, t / rate, complex=2]
+    """
+    two_pi = 2.0 * np.pi
+
+    for (t, step) in enumerate(time_steps):
+        columns = D[:, int(step) : int(step + 2)]
+        columns_0 = columns[:, 0]
+        columns_1 = columns[:, 1]
+
+        # Weighting for linear magnitude interpolation
+        alpha = np.mod(step, 1.0)
+        mag = (1.0 - alpha) * np.abs(columns_0) + alpha * np.abs(columns_1)
+
+        # Store to output array
+        d_stretch[:, t] = mag * np.exp(1.0j * phase_acc)
+
+        # Compute phase advance
+        dphase = np.angle(columns_1) - np.angle(columns_0) - phi_advance
+
+        # Wrap to -pi:pi range
+        scale = dphase / two_pi
+        np.round(scale, 0, scale_buffer)
+
+        dphase = dphase - two_pi * scale_buffer
+
+        # Accumulate phase
+        phase_acc += phi_advance + dphase
+
+    return d_stretch

tts/nemo/collections/asr/parts/utils/rnnt_utils.py

@@ -0,0 +1,89 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Copyright 2017 Johns Hopkins University (Shinji Watanabe)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Union
+
+import torch
+
+
+@dataclass
+class Hypothesis:
+    """Hypothesis class for beam search algorithms.
+
+    score: A float score obtained from an AbstractRNNTDecoder module's score_hypothesis method.
+
+    y_sequence: Either a sequence of integer ids pointing to some vocabulary, or a packed torch.Tensor
+        behaving in the same manner. dtype must be torch.Long in the latter case.
+
+    dec_state: A list (or list of list) of LSTM-RNN decoder states. Can be None.
+
+    text: (Optional) A decoded string after processing via CTC / RNN-T decoding (removing the CTC/RNNT
+        `blank` tokens, and optionally merging word-pieces). Should be used as decoded string for
+        Word Error Rate calculation.
+
+    timestep: (Optional) A list of integer indices representing at which index in the decoding
+        process did the token appear. Should be of same length as the number of non-blank tokens.
+
+    alignments: (Optional) Represents the CTC / RNNT token alignments as integer tokens along an axis of
+        time T (for CTC) or Time x Target (TxU).
+        For CTC, represented as a single list of integer indices.
+        For RNNT, represented as a dangling list of list of integer indices.
+        Outer list represents Time dimension (T), inner list represents Target dimension (U).
+        The set of valid indices **includes** the CTC / RNNT blank token in order to represent alignments.
+
+    length: Represents the length of the sequence (the original length without padding), otherwise
+        defaults to 0.
+
+    y: (Unused) A list of torch.Tensors representing the list of hypotheses.
+
+    lm_state: (Unused) A dictionary state cache used by an external Language Model.
+
+    lm_scores: (Unused) Score of the external Language Model.
+
+    tokens: (Optional) A list of decoded tokens (can be characters or word-pieces.
+    """
+
+    score: float
+    y_sequence: Union[List[int], torch.Tensor]
+    dec_state: Optional[Union[List[List[torch.Tensor]], List[torch.Tensor]]] = None
+    text: Optional[str] = None
+    timestep: Union[List[int], torch.Tensor] = field(default_factory=list)
+    alignments: Optional[Union[List[int], List[List[int]]]] = None
+    length: int = 0
+    y: List[torch.tensor] = None
+    lm_state: Union[Dict[str, Any], List[Any]] = None
+    lm_scores: torch.Tensor = None
+    tokens: Optional[Union[List[int], torch.Tensor]] = None
+
+
+@dataclass
+class NBestHypotheses:
+    """List of N best hypotheses"""
+
+    n_best_hypotheses: Optional[List[Hypothesis]]

tts/nemo/collections/asr/parts/utils/speaker_utils.py

@@ -0,0 +1,400 @@
+# 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 math
+import os
+import pickle as pkl
+from copy import deepcopy
+
+import numpy as np
+from omegaconf import ListConfig
+from pyannote.core import Annotation, Segment
+from pyannote.metrics.diarization import DiarizationErrorRate
+from tqdm import tqdm
+
+from nemo.collections.asr.parts.utils.nmse_clustering import COSclustering
+from nemo.utils import logging
+
+
+"""
+This file contains all the utility functions required for speaker embeddings part in diarization scripts
+"""
+
+
+def audio_rttm_map(audio_file_list, rttm_file_list=None):
+    """
+    Returns a AUDIO_TO_RTTM dictionary thats maps all the unique file names with 
+    audio file paths and corresponding ground truth rttm files calculated from audio
+    file list and rttm file list
+
+    Args: 
+    audio_file_list(list,str): either list of audio file paths or file containing paths to audio files (required)
+    rttm_file_list(lisr,str): either list of rttm file paths or file containing paths to rttm files (optional) 
+    [Required if DER needs to be calculated]
+    Returns:
+    AUDIO_RTTM_MAP (dict): dictionary thats maps all the unique file names with 
+    audio file paths and corresponding ground truth rttm files 
+    """
+    rttm_notfound = False
+    if type(audio_file_list) in [list, ListConfig]:
+        audio_files = audio_file_list
+    else:
+        audio_pointer = open(audio_file_list, 'r')
+        audio_files = audio_pointer.read().splitlines()
+
+    if rttm_file_list:
+        if type(rttm_file_list) in [list, ListConfig]:
+            rttm_files = rttm_file_list
+        else:
+            rttm_pointer = open(rttm_file_list, 'r')
+            rttm_files = rttm_pointer.read().splitlines()
+    else:
+        rttm_notfound = True
+        rttm_files = ['-'] * len(audio_files)
+
+    assert len(audio_files) == len(rttm_files)
+
+    AUDIO_RTTM_MAP = {}
+    rttm_dict = {}
+    audio_dict = {}
+    for audio_file, rttm_file in zip(audio_files, rttm_files):
+        uniq_audio_name = audio_file.split('/')[-1].rsplit('.', 1)[0]
+        uniq_rttm_name = rttm_file.split('/')[-1].rsplit('.', 1)[0]
+
+        if rttm_notfound:
+            uniq_rttm_name = uniq_audio_name
+
+        audio_dict[uniq_audio_name] = audio_file
+        rttm_dict[uniq_rttm_name] = rttm_file
+
+    for key, value in audio_dict.items():
+
+        AUDIO_RTTM_MAP[key] = {'audio_path': audio_dict[key], 'rttm_path': rttm_dict[key]}
+
+    assert len(rttm_dict.items()) == len(audio_dict.items())
+
+    return AUDIO_RTTM_MAP
+
+
+def get_contiguous_stamps(stamps):
+    """
+    Return contiguous time stamps 
+    """
+    lines = deepcopy(stamps)
+    contiguous_stamps = []
+    for i in range(len(lines) - 1):
+        start, end, speaker = lines[i].split()
+        next_start, next_end, next_speaker = lines[i + 1].split()
+        if float(end) > float(next_start):
+            avg = str((float(next_start) + float(end)) / 2.0)
+            lines[i + 1] = ' '.join([avg, next_end, next_speaker])
+            contiguous_stamps.append(start + " " + avg + " " + speaker)
+        else:
+            contiguous_stamps.append(start + " " + end + " " + speaker)
+    start, end, speaker = lines[-1].split()
+    contiguous_stamps.append(start + " " + end + " " + speaker)
+    return contiguous_stamps
+
+
+def merge_stamps(lines):
+    """
+    merge time stamps of same speaker
+    """
+    stamps = deepcopy(lines)
+    overlap_stamps = []
+    for i in range(len(stamps) - 1):
+        start, end, speaker = stamps[i].split()
+        next_start, next_end, next_speaker = stamps[i + 1].split()
+        if float(end) == float(next_start) and speaker == next_speaker:
+            stamps[i + 1] = ' '.join([start, next_end, next_speaker])
+        else:
+            overlap_stamps.append(start + " " + end + " " + speaker)
+
+    start, end, speaker = stamps[-1].split()
+    overlap_stamps.append(start + " " + end + " " + speaker)
+
+    return overlap_stamps
+
+
+def labels_to_pyannote_object(labels):
+    """
+    converts labels to pyannote object to calculate DER and for visualization
+    """
+    annotation = Annotation()
+    for label in labels:
+        start, end, speaker = label.strip().split()
+        start, end = float(start), float(end)
+        annotation[Segment(start, end)] = speaker
+
+    return annotation
+
+
+def labels_to_rttmfile(labels, uniq_id, out_rttm_dir):
+    """
+    write rttm file with uniq_id name in out_rttm_dir with time_stamps in labels
+    """
+    filename = os.path.join(out_rttm_dir, uniq_id + '.rttm')
+    with open(filename, 'w') as f:
+        for line in labels:
+            line = line.strip()
+            start, end, speaker = line.split()
+            duration = float(end) - float(start)
+            start = float(start)
+            log = 'SPEAKER {} 1   {:.3f}   {:.3f} <NA> <NA> {} <NA> <NA>\n'.format(uniq_id, start, duration, speaker)
+            f.write(log)
+
+
+def rttm_to_labels(rttm_filename):
+    """
+    prepares time stamps label list from rttm file
+    """
+    labels = []
+    with open(rttm_filename, 'r') as f:
+        for line in f.readlines():
+            rttm = line.strip().split()
+            start, end, speaker = float(rttm[3]), float(rttm[4]) + float(rttm[3]), rttm[7]
+            labels.append('{} {} {}'.format(start, end, speaker))
+    return labels
+
+
+def get_time_stamps(embeddings_file, reco2num, manifest_path, sample_rate, window, shift):
+    """
+    Loads embedding file and generates time stamps based on window and shift for speaker embeddings 
+    clustering.
+
+    Args:
+    embeddings_file (str): Path to embeddings pickle file
+    reco2num (int,str): common integer number of speakers for every recording or path to 
+                        file that states unique_file id with number of speakers
+    manifest_path (str): path to manifest for time stamps matching
+    sample_rate (int): sample rate
+    window (float): window length in sec for speaker embeddings extraction
+    shift (float): shift length in sec for speaker embeddings extraction window shift
+
+    Returns:
+    embeddings (dict): Embeddings with key as unique_id
+    time_stamps (dict): time stamps list for each audio recording
+    speakers (dict): number of speaker for each audio recording
+    """
+    embeddings = pkl.load(open(embeddings_file, 'rb'))
+    all_uniq_files = list(embeddings.keys())
+    num_files = len(all_uniq_files)
+    logging.info("Number of files to diarize: {}".format(num_files))
+    sample = all_uniq_files[0]
+    logging.info("sample '{}' embeddings shape is {}\n".format(sample, embeddings[sample][0].shape))
+
+    speakers = {}
+    if isinstance(reco2num, int) or reco2num is None:
+        for key in embeddings.keys():
+            speakers[key] = reco2num
+    elif isinstance(reco2num, str) and os.path.exists(reco2num):
+        for key in open(reco2num).readlines():
+            key = key.strip()
+            wav_id, num = key.split()
+            speakers[wav_id] = int(num)
+    else:
+        raise TypeError("reco2num must be None, int or path to file containing number of speakers for each uniq id")
+
+    with open(manifest_path, 'r') as manifest:
+        time_stamps = {}
+        for line in manifest.readlines():
+            line = line.strip()
+            line = json.loads(line)
+            audio, offset, duration = line['audio_filepath'], line['offset'], line['duration']
+            audio = os.path.basename(audio).rsplit('.', 1)[0]
+            if audio not in time_stamps:
+                time_stamps[audio] = []
+            start = offset
+            slice_end = start + duration
+            base = math.ceil((duration - window) / shift)
+            slices = 1 if base < 0 else base + 1
+            for slice_id in range(slices):
+                end = start + window
+                if end > slice_end:
+                    end = slice_end
+                stamp = '{:.3f} {:.3f} '.format(start, end)
+                time_stamps[audio].append(stamp)
+                start = offset + (slice_id + 1) * shift
+
+    return embeddings, time_stamps, speakers
+
+
+def perform_clustering(embeddings, time_stamps, speakers, audio_rttm_map, out_rttm_dir, max_num_speakers=8):
+    """
+    performs spectral clustering on embeddings with time stamps generated from VAD output
+    Args:
+    
+    embeddings (dict): Embeddings with key as unique_id
+    time_stamps (dict): time stamps list for each audio recording
+    speakers (dict): number of speaker for each audio recording 
+    audio_rttm_map (dict): AUDIO_RTTM_MAP for mapping unique id with audio file path and rttm path
+    out_rttm_dir (str): Path to write predicted rttms
+    max_num_speakers (int): maximum number of speakers to consider for spectral clustering. Will be ignored if speakers['key'] is not None
+    
+    Returns:
+    all_reference (list[Annotation]): reference annotations for score calculation
+    all_hypothesis (list[Annotation]): hypothesis annotations for score calculation
+
+    """
+    all_hypothesis = []
+    all_reference = []
+    no_references = False
+
+    for uniq_key in tqdm(embeddings.keys()):
+        NUM_speakers = speakers[uniq_key]
+        emb = embeddings[uniq_key]
+        emb = np.asarray(emb)
+
+        cluster_labels = COSclustering(
+            uniq_key, emb, oracle_num_speakers=NUM_speakers, max_num_speaker=max_num_speakers
+        )
+
+        lines = time_stamps[uniq_key]
+        assert len(cluster_labels) == len(lines)
+        for idx, label in enumerate(cluster_labels):
+            tag = 'speaker_' + str(label)
+            lines[idx] += tag
+
+        a = get_contiguous_stamps(lines)
+        labels = merge_stamps(a)
+        if out_rttm_dir:
+            labels_to_rttmfile(labels, uniq_key, out_rttm_dir)
+        hypothesis = labels_to_pyannote_object(labels)
+        all_hypothesis.append(hypothesis)
+
+        rttm_file = audio_rttm_map[uniq_key]['rttm_path']
+        if os.path.exists(rttm_file) and not no_references:
+            ref_labels = rttm_to_labels(rttm_file)
+            reference = labels_to_pyannote_object(ref_labels)
+            all_reference.append(reference)
+        else:
+            no_references = True
+            all_reference = []
+
+    return all_reference, all_hypothesis
+
+
+def get_DER(all_reference, all_hypothesis):
+    """
+    calculates DER, CER, FA and MISS
+
+    Args:
+    all_reference (list[Annotation]): reference annotations for score calculation
+    all_hypothesis (list[Annotation]): hypothesis annotations for score calculation
+
+    Returns:
+    DER (float): Diarization Error Rate
+    CER (float): Confusion Error Rate
+    FA (float): False Alarm
+    Miss (float): Miss Detection 
+
+    """
+    metric = DiarizationErrorRate(collar=0.25, skip_overlap=True)
+    DER = 0
+
+    for reference, hypothesis in zip(all_reference, all_hypothesis):
+        metric(reference, hypothesis, detailed=True)
+
+    DER = abs(metric)
+    CER = metric['confusion'] / metric['total']
+    FA = metric['false alarm'] / metric['total']
+    MISS = metric['missed detection'] / metric['total']
+
+    metric.reset()
+
+    return DER, CER, FA, MISS
+
+
+def perform_diarization(
+    embeddings_file=None,
+    reco2num=2,
+    manifest_path=None,
+    sample_rate=16000,
+    window=1.5,
+    shift=0.75,
+    audio_rttm_map=None,
+    out_rttm_dir=None,
+    max_num_speakers=8,
+):
+    """
+    Performs diarization with embeddings generated based on VAD time stamps with recording 2 num of speakers (reco2num)
+    for spectral clustering 
+    """
+    embeddings, time_stamps, speakers = get_time_stamps(
+        embeddings_file, reco2num, manifest_path, sample_rate, window, shift
+    )
+    logging.info("Performing Clustering")
+    all_reference, all_hypothesis = perform_clustering(
+        embeddings, time_stamps, speakers, audio_rttm_map, out_rttm_dir, max_num_speakers
+    )
+
+    if len(all_reference) and len(all_hypothesis):
+        DER, CER, FA, MISS = get_DER(all_reference, all_hypothesis)
+        logging.info(
+            "Cumulative results of all the files:  \n FA: {:.3f}\t MISS {:.3f}\t \
+                Diarization ER: {:.3f}\t, Confusion ER:{:.3f}".format(
+                FA, MISS, DER, CER
+            )
+        )
+    else:
+        logging.warning("Please check if each ground truth RTTMs was present in provided path2groundtruth_rttm_files")
+        logging.warning("Skipping calculation of Diariazation Error Rate")
+
+
+def write_rttm2manifest(paths2audio_files, paths2rttm_files, manifest_file):
+    """
+    writes manifest file based on rttm files (or vad table out files). This manifest file would be used by 
+    speaker diarizer to compute embeddings and cluster them. This function also takes care of overlap time stamps
+
+    Args:
+    audio_file_list(list,str): either list of audio file paths or file containing paths to audio files (required)
+    rttm_file_list(lisr,str): either list of rttm file paths or file containing paths to rttm files (optional) 
+    manifest (str): path to write manifest file
+
+    Returns:
+    manifest (str): path to write manifest file 
+    """
+    AUDIO_RTTM_MAP = audio_rttm_map(paths2audio_files, paths2rttm_files)
+
+    with open(manifest_file, 'w') as outfile:
+        for key in AUDIO_RTTM_MAP:
+            f = open(AUDIO_RTTM_MAP[key]['rttm_path'], 'r')
+            audio_path = AUDIO_RTTM_MAP[key]['audio_path']
+            lines = f.readlines()
+            time_tup = (-1, -1)
+            for line in lines:
+                vad_out = line.strip().split()
+                if len(vad_out) > 3:
+                    start, dur, activity = float(vad_out[3]), float(vad_out[4]), vad_out[7]
+                else:
+                    start, dur, activity = float(vad_out[0]), float(vad_out[1]), vad_out[2]
+                start, dur = float("{:.3f}".format(start)), float("{:.3f}".format(dur))
+
+                if time_tup[0] >= 0 and start > time_tup[1]:
+                    dur2 = float("{:.3f}".format(time_tup[1] - time_tup[0]))
+                    meta = {"audio_filepath": audio_path, "offset": time_tup[0], "duration": dur2, "label": 'UNK'}
+                    json.dump(meta, outfile)
+                    outfile.write("\n")
+                    time_tup = (start, start + dur)
+                else:
+                    if time_tup[0] == -1:
+                        time_tup = (start, start + dur)
+                    else:
+                        time_tup = (min(time_tup[0], start), max(time_tup[1], start + dur))
+            dur2 = float("{:.3f}".format(time_tup[1] - time_tup[0]))
+            meta = {"audio_filepath": audio_path, "offset": time_tup[0], "duration": dur2, "label": 'UNK'}
+            json.dump(meta, outfile)
+            outfile.write("\n")
+            f.close()
+    return manifest_file

tts/nemo/collections/asr/parts/utils/vad_utils.py

@@ -0,0 +1,544 @@
+# 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 glob
+import json
+import os
+from itertools import repeat
+from multiprocessing import Pool
+
+import librosa
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+from pyannote.core import Annotation, Segment
+from pyannote.metrics import detection
+
+from nemo.utils import logging
+
+
+"""
+This file contains all the utility functions required for speaker embeddings part in diarization scripts
+"""
+
+
+def prepare_manifest(config):
+    """
+    Perform VAD on long audio snippet might cause CUDA out of memory issue. 
+    Automatically split manifest entry by split_duration to avoid the potential memory issue.
+    """
+    manifest_vad_input = config.get('manifest_vad_input', "manifest_vad_input.json")
+    input_audios = []
+    with open(config['manifest_filepath'], 'r') as manifest:
+        for line in manifest.readlines():
+            input_audios.append(json.loads(line.strip()))
+
+    p = Pool(processes=config['num_workers'])
+    args_func = {
+        'label': 'infer',
+        'split_duration': config['split_duration'],
+        'time_length': config['time_length'],
+    }
+    results = p.starmap(write_vad_infer_manifest, zip(input_audios, repeat(args_func)))
+    p.close()
+
+    if os.path.exists(manifest_vad_input):
+        logging.info("The prepared manifest file exists. Overwriting!")
+        os.remove(manifest_vad_input)
+
+    with open(manifest_vad_input, 'a') as fout:
+        for res in results:
+            for r in res:
+                json.dump(r, fout)
+                fout.write('\n')
+                fout.flush()
+
+    return manifest_vad_input
+
+
+def write_vad_infer_manifest(file, args_func):
+    """
+    Used by prepare_manifest.
+    Given a list of files, write them to manifest for dataloader with restrictions.
+    Args:
+        files (json) : file to be processed
+        label (str): label for audio snippet.
+        split_duration (float): max duration of each audio clip (each line in json)
+        time_length (float) : length of window for generating the frame. Used for taking care of joint. 
+    Returns:
+        res (list) : list of generated metadata line of json for file
+    """
+    res = []
+    label = args_func['label']
+    split_duration = args_func['split_duration']
+    time_length = args_func['time_length']
+
+    filepath = file['audio_filepath']
+    in_duration = file['duration']
+    in_offset = file['offset']
+
+    try:
+        sr = 16000
+        x, _sr = librosa.load(filepath, sr=sr, offset=in_offset, duration=in_duration)
+        duration = librosa.get_duration(x, sr=sr)
+        left = duration
+        current_offset = in_offset
+
+        status = 'single'
+        while left > 0:
+            if left <= split_duration:
+                status = 'end'
+                write_duration = left + time_length
+                current_offset -= time_length
+                offset_inc = left
+                left = 0
+            else:
+                if status == 'start' or status == 'next':
+                    status = 'next'
+                else:
+                    status = 'start'
+
+                if status == 'start':
+                    write_duration = split_duration
+                    offset_inc = split_duration
+                else:
+                    write_duration = split_duration + time_length
+                    current_offset -= time_length
+                    offset_inc = split_duration + time_length
+
+                left -= split_duration
+
+            metadata = {
+                'audio_filepath': filepath,
+                'duration': write_duration,
+                'label': label,
+                'text': '_',
+                'offset': current_offset,
+            }
+            res.append(metadata)
+
+            current_offset += offset_inc
+
+    except Exception as e:
+        err_file = "error.log"
+        with open(err_file, 'w') as fout:
+            fout.write(file + ":" + str(e))
+    return res
+
+
+def get_vad_stream_status(data):
+    """
+    Generate a list of status for each snippet in manifest. A snippet should be in single, start, next or end status. 
+    Used for concatenate to full audio file.
+    Args:
+        data (list): list of filepath of audio snippet
+    Returns:
+        status (list): list of status of each snippet.
+    """
+    if len(data) == 1:
+        return ['single']
+
+    status = [None] * len(data)
+    for i in range(len(data)):
+        if i == 0:
+            status[i] = 'start' if data[i] == data[i + 1] else 'single'
+        elif i == len(data) - 1:
+            status[i] = 'end' if data[i] == data[i - 1] else 'single'
+        else:
+            if data[i] != data[i - 1] and data[i] == data[i + 1]:
+                status[i] = 'start'
+            elif data[i] == data[i - 1] and data[i] == data[i + 1]:
+                status[i] = 'next'
+            elif data[i] == data[i - 1] and data[i] != data[i + 1]:
+                status[i] = 'end'
+            else:
+                status[i] = 'single'
+    return status
+
+
+def generate_overlap_vad_seq(frame_pred_dir, smoothing_method, overlap, seg_len, shift_len, num_workers):
+    """
+    Gnerate predictions with overlapping input windows/segments. Then a smoothing filter is applied to decide the label for a frame spanned by multiple windows. 
+    Two common smoothing filters are supported: majority vote (median) and average (mean).
+    Args:
+        frame_pred_dir (str): Directory of frame prediction file to be processed.
+        smoothing_method (str): median or mean smoothing filter.
+        overlap (float): amounts of overlap of adjacent windows.
+        seg_len (float): length of window for generating the frame.
+        shift_len (float): amount of shift of window for generating the frame.
+        out_dir (str): directory of generated predictions.
+        num_workers(float): number of process for multiprocessing
+    Returns:
+        overlap_out_dir(str): directory of generate predictions.
+    """
+
+    p = Pool(processes=num_workers)
+    frame_filepathlist = glob.glob(frame_pred_dir + "/*.frame")
+
+    overlap_out_dir = frame_pred_dir + "/overlap_smoothing_output" + "_" + smoothing_method + "_" + str(overlap)
+
+    if not os.path.exists(overlap_out_dir):
+        os.mkdir(overlap_out_dir)
+
+    per_args = {
+        "out_dir": overlap_out_dir,
+        "smoothing_method": smoothing_method,
+        "overlap": overlap,
+        "seg_len": seg_len,
+        "shift_len": shift_len,
+    }
+    p.starmap(generate_overlap_vad_seq_per_file, zip(frame_filepathlist, repeat(per_args)))
+    p.close()
+    p.join()
+
+    return overlap_out_dir
+
+
+def generate_overlap_vad_seq_per_file(frame_filepath, per_args):
+    """
+    Use generated frame prediction (generated by shifting window of shift_len (10ms)) to generate prediction with overlapping input window/segments
+    See discription in generate_overlap_vad_seq.
+    """
+    try:
+        smoothing_method = per_args['smoothing_method']
+        overlap = per_args['overlap']
+        seg_len = per_args['seg_len']
+        shift_len = per_args['shift_len']
+        out_dir = per_args['out_dir']
+
+        frame = np.loadtxt(frame_filepath)
+        name = os.path.basename(frame_filepath).split(".frame")[0] + "." + smoothing_method
+        overlap_filepath = os.path.join(out_dir, name)
+
+        shift = int(shift_len / 0.01)  # number of units of shift
+        seg = int((seg_len / 0.01 + 1))  # number of units of each window/segment
+
+        jump_on_target = int(seg * (1 - overlap))  # jump on target generated sequence
+        jump_on_frame = int(jump_on_target / shift)  # jump on input frame sequence
+
+        if jump_on_frame < 1:
+            raise ValueError(
+                f"Note we jump over frame sequence to generate overlapping input segments. \n \
+            Your input makes jump_on_fram={jump_on_frame} < 1 which is invalid because it cannot jump and will stuck.\n \
+            Please try different seg_len, shift_len and overlap choices. \n \
+            jump_on_target = int(seg * (1 - overlap)) \n \
+            jump_on_frame  = int(jump_on_frame/shift) "
+            )
+
+        target_len = int(len(frame) * shift)
+
+        if smoothing_method == 'mean':
+            preds = np.zeros(target_len)
+            pred_count = np.zeros(target_len)
+
+            for i, og_pred in enumerate(frame):
+                if i % jump_on_frame != 0:
+                    continue
+                start = i * shift
+                end = start + seg
+                preds[start:end] = preds[start:end] + og_pred
+                pred_count[start:end] = pred_count[start:end] + 1
+
+            preds = preds / pred_count
+            last_non_zero_pred = preds[pred_count != 0][-1]
+            preds[pred_count == 0] = last_non_zero_pred
+
+        elif smoothing_method == 'median':
+            preds = [[] for _ in range(target_len)]
+            for i, og_pred in enumerate(frame):
+                if i % jump_on_frame != 0:
+                    continue
+
+                start = i * shift
+                end = start + seg
+                for j in range(start, end):
+                    if j <= target_len - 1:
+                        preds[j].append(og_pred)
+
+            preds = np.array([np.median(l) for l in preds])
+            nan_idx = np.isnan(preds)
+            last_non_nan_pred = preds[~nan_idx][-1]
+            preds[nan_idx] = last_non_nan_pred
+
+        else:
+            raise ValueError("smoothing_method should be either mean or median")
+
+        round_final = np.round(preds, 4)
+        np.savetxt(overlap_filepath, round_final, delimiter='\n')
+
+    except Exception as e:
+        raise (e)
+
+
+def generate_vad_segment_table(vad_pred_dir, threshold, shift_len, num_workers):
+    """
+    Convert frame level prediction to speech segment in start and end times format.
+    And save to csv file  in rttm-like format
+            0, 10, speech
+            17,18, speech
+    Args:
+        vad_pred_dir (str): directory of prediction files to be processed.
+        threshold (float): threshold for prediction score (from 0 to 1).
+        shift_len (float): amount of shift of window for generating the frame.
+        out_dir (str): output dir of generated table/csv file.
+        num_workers(float): number of process for multiprocessing
+    Returns:
+        table_out_dir(str): directory of generate table.
+    """
+
+    p = Pool(processes=num_workers)
+    suffixes = ("frame", "mean", "median")
+    vad_pred_filepath_list = [os.path.join(vad_pred_dir, x) for x in os.listdir(vad_pred_dir) if x.endswith(suffixes)]
+
+    table_out_dir = os.path.join(vad_pred_dir, "table_output_" + str(threshold))
+    if not os.path.exists(table_out_dir):
+        os.mkdir(table_out_dir)
+
+    per_args = {
+        "threshold": threshold,
+        "shift_len": shift_len,
+        "out_dir": table_out_dir,
+    }
+
+    p.starmap(generate_vad_segment_table_per_file, zip(vad_pred_filepath_list, repeat(per_args)))
+    p.close()
+    p.join()
+
+    return table_out_dir
+
+
+def generate_vad_segment_table_per_file(pred_filepath, per_args):
+    """
+    See discription in generate_overlap_vad_seq.
+    """
+    threshold = per_args['threshold']
+    shift_len = per_args['shift_len']
+    out_dir = per_args['out_dir']
+
+    name = pred_filepath.split("/")[-1].rsplit(".", 1)[0]
+
+    sequence = np.loadtxt(pred_filepath)
+    start = 0
+    start_list = [0]
+    dur_list = []
+    state_list = []
+    current_state = "non-speech"
+    for i in range(len(sequence) - 1):
+        current_state = "non-speech" if sequence[i] <= threshold else "speech"
+        next_state = "non-speech" if sequence[i + 1] <= threshold else "speech"
+        if next_state != current_state:
+            dur = i * shift_len + shift_len - start  # shift_len for handling joint
+            state_list.append(current_state)
+            dur_list.append(dur)
+
+            start = (i + 1) * shift_len
+            start_list.append(start)
+
+    dur_list.append((i + 1) * shift_len + shift_len - start)
+    state_list.append(current_state)
+
+    seg_table = pd.DataFrame({'start': start_list, 'dur': dur_list, 'vad': state_list})
+    seg_speech_table = seg_table[seg_table['vad'] == 'speech']
+
+    save_name = name + ".txt"
+    save_path = os.path.join(out_dir, save_name)
+    seg_speech_table.to_csv(save_path, sep='\t', index=False, header=False)
+    return save_path
+
+
+def vad_construct_pyannote_object_per_file(vad_table_filepath, groundtruth_RTTM_file):
+    """
+    Construct pyannote object for evaluation.
+    Args:
+        vad_table_filepath(str) : path of vad rttm-like table.
+        groundtruth_RTTM_file(str): path of groundtruth rttm file.
+    Returns:
+        reference(pyannote.Annotation): groundtruth
+        hypothesise(pyannote.Annotation): prediction
+    """
+
+    pred = pd.read_csv(vad_table_filepath, sep="\t", header=None)
+    label = pd.read_csv(groundtruth_RTTM_file, sep=" ", delimiter=None, header=None)
+    label = label.rename(columns={3: "start", 4: "dur", 7: "speaker"})
+
+    # construct reference
+    reference = Annotation()
+    for index, row in label.iterrows():
+        reference[Segment(row['start'], row['start'] + row['dur'])] = row['speaker']
+
+    # construct hypothsis
+    hypothesis = Annotation()
+    for index, row in pred.iterrows():
+        hypothesis[Segment(row[0], row[0] + row[1])] = 'Speech'
+    return reference, hypothesis
+
+
+def vad_tune_threshold_on_dev(thresholds, vad_pred, groundtruth_RTTM, vad_pred_method="frame", focus_metric="DetER"):
+    """
+    Tune threshold on dev set. Return best threshold which gives the lowest detection error rate (DetER) in thresholds.
+    Args:
+        thresholds (list): list of thresholds.
+        vad_pred_method (str): suffix of prediction file. Use to locate file. Should be either in "frame", "mean" or "median".
+        vad_pred_dir (str): directory of vad predictions or a file contains the paths of them
+        groundtruth_RTTM_dir (str): directory of groundtruch rttm files or a file contains the paths of them.
+        focus_metric (str): metrics we care most when tuning threshold. Should be either in "DetER", "FA", "MISS"
+    Returns:
+        best_threhsold (float): threshold that gives lowest DetER.
+    """
+    threshold_perf = {}
+    best_threhsold = thresholds[0]
+    min_score = 100
+
+    try:
+        thresholds[0] >= 0 and thresholds[-1] <= 1
+    except:
+        raise ValueError("Invalid threshold! Should be in [0, 1]")
+
+    for threshold in thresholds:
+        metric = detection.DetectionErrorRate()
+        paired_filenames, groundtruth_RTTM_dict, vad_pred_dict = pred_rttm_map(
+            vad_pred, groundtruth_RTTM, vad_pred_method
+        )
+        print(paired_filenames)
+        for filename in paired_filenames:
+            vad_pred_filepath = vad_pred_dict[filename]
+            groundtruth_RTTM_file = groundtruth_RTTM_dict[filename]
+
+            if os.path.isdir(vad_pred):
+                table_out_dir = os.path.join(vad_pred, "table_output_" + str(threshold))
+            else:
+                table_out_dir = os.path.join("tmp_table_outputs", "table_output_" + str(threshold))
+
+            if not os.path.exists(table_out_dir):
+                os.makedirs(table_out_dir)
+
+            per_args = {"threshold": threshold, "shift_len": 0.01, "out_dir": table_out_dir}
+            vad_table_filepath = generate_vad_segment_table_per_file(vad_pred_filepath, per_args)
+
+            reference, hypothesis = vad_construct_pyannote_object_per_file(vad_table_filepath, groundtruth_RTTM_file)
+            metric(reference, hypothesis)  # accumulation
+
+        report = metric.report(display=False)
+        DetER = report.iloc[[-1]][('detection error rate', '%')].item()
+        FA = report.iloc[[-1]][('false alarm', '%')].item()
+        MISS = report.iloc[[-1]][('miss', '%')].item()
+
+        if focus_metric == "DetER":
+            score = DetER
+        elif focus_metric == "FA":
+            score = FA
+        elif focus_metric == "MISS":
+            score = MISS
+        else:
+            raise ValueError("Metric we care most should be only in 'DetER', 'FA'or 'MISS'!")
+
+        threshold_perf[threshold] = {'DetER (%)': DetER, 'FA (%)': FA, 'MISS (%)': MISS}
+        logging.info(f"threshold {threshold}, {threshold_perf[threshold]}")
+        del report
+        metric.reset()  # reset internal accumulator
+        if score < min_score:
+            min_score = score
+            best_threhsold = threshold
+    return best_threhsold
+
+
+def pred_rttm_map(vad_pred, groundtruth_RTTM, vad_pred_method="frame"):
+    """
+    Find paired files in vad_pred and groundtruth_RTTM
+    """
+    groundtruth_RTTM_dict = {}
+    if os.path.isfile(groundtruth_RTTM):
+        with open(groundtruth_RTTM, "r") as fp:
+            groundtruth_RTTM_files = fp.read().splitlines()
+    elif os.path.isdir(groundtruth_RTTM):
+        groundtruth_RTTM_files = glob.glob(os.path.join(groundtruth_RTTM, "*.rttm"))
+    else:
+        raise ValueError(
+            "groundtruth_RTTM should either be a directory contains rttm files or a file contains paths to them!"
+        )
+    for f in groundtruth_RTTM_files:
+        filename = os.path.basename(f).rsplit(".", 1)[0]
+        groundtruth_RTTM_dict[filename] = f
+
+    vad_pred_dict = {}
+    if os.path.isfile(vad_pred):
+        with open(vad_pred, "r") as fp:
+            vad_pred_files = fp.read().splitlines()
+    elif os.path.isdir(vad_pred):
+        vad_pred_files = glob.glob(os.path.join(vad_pred, "*." + vad_pred_method))
+    else:
+        raise ValueError(
+            "vad_pred should either be a directory contains vad pred files or a file contains paths to them!"
+        )
+    for f in vad_pred_files:
+        filename = os.path.basename(f).rsplit(".", 1)[0]
+        vad_pred_dict[filename] = f
+
+    paired_filenames = groundtruth_RTTM_dict.keys() & vad_pred_dict.keys()
+    return paired_filenames, groundtruth_RTTM_dict, vad_pred_dict
+
+
+def plot(path2audio_file, path2_vad_pred, path2ground_truth_label=None, threshold=0.85):
+    """
+    Plot VAD outputs for demonstration in tutorial
+    Args:
+        path2audio_file (str):  path to audio file.
+        path2_vad_pred (str): path to vad prediction file,
+        path2ground_truth_label(str): path to groundtruth label file.
+        threshold (float): threshold for prediction score (from 0 to 1).
+    """
+    plt.figure(figsize=[20, 2])
+    FRAME_LEN = 0.01
+    audio, sample_rate = librosa.load(path=path2audio_file, sr=16000, mono=True)
+    dur = librosa.get_duration(audio, sr=sample_rate)
+    time = np.arange(0, dur, FRAME_LEN)
+    frame = np.loadtxt(path2_vad_pred)
+    len_pred = len(frame)
+    ax1 = plt.subplot()
+    ax1.plot(np.arange(audio.size) / sample_rate, audio, 'gray')
+    ax1.set_xlim([0, int(dur) + 1])
+    ax1.tick_params(axis='y', labelcolor='b')
+    ax1.set_ylabel('Signal')
+    ax1.set_ylim([-1, 1])
+    ax2 = ax1.twinx()
+    prob = frame
+    pred = np.where(prob >= threshold, 1, 0)
+    if path2ground_truth_label:
+        label = extract_labels(path2ground_truth_label, time)
+        ax2.plot(np.arange(len_pred) * FRAME_LEN, label, 'r', label='label')
+    ax2.plot(np.arange(len_pred) * FRAME_LEN, pred, 'b', label='pred')
+    ax2.plot(np.arange(len_pred) * FRAME_LEN, prob, 'g--', label='speech prob')
+    ax2.tick_params(axis='y', labelcolor='r')
+    ax2.legend(loc='lower right', shadow=True)
+    ax2.set_ylabel('Preds and Probas')
+    ax2.set_ylim([-0.1, 1.1])
+    return None
+
+
+def extract_labels(path2ground_truth_label, time):
+    """
+    Extract groundtruth label for given time period.
+    path2ground_truth_label (str): path of groundtruth label file 
+    time (list) : a list of array represent time period.
+    """
+
+    data = pd.read_csv(path2ground_truth_label, sep=" ", delimiter=None, header=None)
+    data = data.rename(columns={3: "start", 4: "dur", 7: "speaker"})
+    labels = []
+    for pos in time:
+        line = data[(data["start"] <= pos) & (data["start"] + data["dur"] > pos)]
+        if len(line) >= 1:
+            labels.append(1)
+        else:
+            labels.append(0)
+    return labels

tts/nemo/collections/common/__init__.py

@@ -0,0 +1,26 @@
+# 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 nemo.collections.common.callbacks
+from nemo.collections.common import data, losses, parts, tokenizers
+from nemo.package_info import __version__
+
+# Set collection version equal to NeMo version.
+__version = __version__
+
+# Authorship.
+__author__ = "NVIDIA Corporation"
+
+# Set collection name.
+__description__ = "Common collection"

tts/nemo/collections/common/callbacks/__init__.py

@@ -0,0 +1,15 @@
+# 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.common.callbacks.callbacks import LogEpochTimeCallback

tts/nemo/collections/common/callbacks/callbacks.py

@@ -0,0 +1,97 @@
+# 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 time
+
+import numpy as np
+from pytorch_lightning.callbacks.base import Callback
+from pytorch_lightning.utilities import rank_zero_only
+
+# from sacrebleu import corpus_bleu
+
+
+class LogEpochTimeCallback(Callback):
+    """Simple callback that logs how long each epoch takes, in seconds, to a pytorch lightning log
+    """
+
+    @rank_zero_only
+    def on_train_epoch_start(self, trainer, pl_module):
+        self.epoch_start = time.time()
+
+    @rank_zero_only
+    def on_train_epoch_end(self, trainer, pl_module, outputs):
+        curr_time = time.time()
+        duration = curr_time - self.epoch_start
+        trainer.logger.log_metrics({"epoch_time": duration}, step=trainer.global_step)
+
+
+# class MachineTranslationLogEvalCallback(Callback):
+#     def _on_eval_end(self, trainer, pl_module, mode):
+#         counts = np.array(self._non_pad_tokens)
+#         eval_loss = np.sum(np.array(self._losses) * counts) / np.sum(counts)
+#         sacre_bleu = corpus_bleu(self._translations, [self._ground_truths], tokenize="13a")
+#         print(f"{mode} results for process with global rank {pl_module.global_rank}".upper())
+#         for i in range(pl_module.num_examples[mode]):
+#             print('\u0332'.join(f"EXAMPLE {i}:"))  # Underline output
+#             sent_id = np.random.randint(len(self._translations))
+#             print(f"Ground truth: {self._ground_truths[sent_id]}\n")
+#             print(f"Translation: {self._translations[sent_id]}\n")
+#             print()
+#         print("-" * 50)
+#         print(f"loss: {eval_loss:.3f}")
+#         print(f"SacreBLEU: {sacre_bleu}")
+#         print("-" * 50)
+
+#     @rank_zero_only
+#     def on_test_end(self, trainer, pl_module):
+#         self._on_eval_end(trainer, pl_module, "test")
+
+#     @rank_zero_only
+#     def on_validation_end(self, trainer, pl_module):
+#         self._on_eval_end(trainer, pl_module, "val")
+
+#     @rank_zero_only
+#     def on_sanity_check_end(self, trainer, pl_module):
+#         self._on_eval_end(trainer, pl_module, "val")
+
+#     def _on_eval_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx, mode):
+#         self._translations.extend(outputs['translations'])
+#         self._ground_truths.extend(outputs['ground_truths'])
+#         self._non_pad_tokens.append(outputs['num_non_pad_tokens'])
+#         self._losses.append(outputs[f'{mode}_loss'])
+
+#     @rank_zero_only
+#     def on_test_batch_end(self, trainer, pl_module, batch, outputs, batch_idx, dataloader_idx):
+#         self._on_eval_batch_end(trainer, pl_module, batch, outputs, batch_idx, dataloader_idx, 'test')
+
+#     @rank_zero_only
+#     def on_validation_batch_end(self, trainer, pl_module, batch, outputs, batch_idx, dataloader_idx):
+#         self._on_eval_batch_end(trainer, pl_module, batch, outputs, batch_idx, dataloader_idx, 'val')
+
+#     def _on_eval_start(self, trainer, pl_module):
+#         self._translations = []
+#         self._ground_truths = []
+#         self._losses = []
+#         self._non_pad_tokens = []
+
+#     @rank_zero_only
+#     def on_test_start(self, trainer, pl_module):
+#         self._on_eval_start(trainer, pl_module)
+
+#     @rank_zero_only
+#     def on_validation_start(self, trainer, pl_module):
+#         self._on_eval_start(trainer, pl_module)
+
+#     @rank_zero_only
+#     def on_sanity_check_start(self, trainer, pl_module):
+#         self._on_eval_start(trainer, pl_module)

tts/nemo/collections/common/data/__init__.py

@@ -0,0 +1,15 @@
+# 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.common.data.dataset import ConcatDataset

tts/nemo/collections/common/data/dataset.py

@@ -0,0 +1,178 @@
+# 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 typing import Any, List
+
+import numpy as np
+import torch.utils.data as pt_data
+from torch.utils.data import IterableDataset
+
+__all__ = ['ConcatDataset']
+
+
+class ConcatDataset(IterableDataset):
+    """
+    A dataset that accepts as argument multiple datasets and then samples from them based on the specified 
+    sampling technique.
+    Args:
+        datasets (list): A list of datasets to sample from.
+        shuffle (bool): Whether to shuffle individual datasets. Only works with non-iterable datasets. 
+            Defaults to True.
+        sampling_technique (str): Sampling technique to choose which dataset to draw a sample from.
+            Defaults to 'temperature'. Currently supports 'temperature', 'random' and 'round-robin'.
+        sampling_temperature (int): Temperature value for sampling. Only used when sampling_technique = 'temperature'.
+            Defaults to 5.
+        sampling_probabilities (list): Probability values for sampling. Only used when sampling_technique = 'random'.
+        global_rank (int): Worker rank, used for partitioning map style datasets. Defaults to 0.
+        world_size (int): Total number of processes, used for partitioning map style datasets. Defaults to 1.
+    """
+
+    def __init__(
+        self,
+        datasets: List[Any],
+        shuffle: bool = True,
+        sampling_technique: str = 'temperature',
+        sampling_temperature: int = 5,
+        sampling_probabilities: List[float] = None,
+        global_rank: int = 0,
+        world_size: int = 1,
+    ):
+        super().__init__()
+
+        supported_sampling_techniques = ['temperature', 'random', 'round-robin']
+        self.datasets = datasets
+        self.iterables = [None] * len(datasets)
+        self.shuffle = shuffle
+        self.global_rank = global_rank
+        self.world_size = world_size
+        self.sampling_kwargs = {}
+        if sampling_technique == 'temperature':
+            self.index_generator = ConcatDataset.temperature_generator
+            self.sampling_kwargs['temperature'] = sampling_temperature
+        elif sampling_technique == 'random':
+            self.index_generator = ConcatDataset.random_generator
+            self.sampling_kwargs['p'] = sampling_probabilities
+        elif sampling_technique == 'round-robin':
+            self.index_generator = ConcatDataset.round_robin_generator
+        else:
+            raise ValueError(f"Currently we only support sampling techniques in {supported_sampling_techniques}.")
+        self.length = 0
+
+        if isinstance(datasets[0], IterableDataset):
+            self.kind = 'iterable'
+        else:
+            self.kind = 'map'
+
+        for idx, dataset in enumerate(datasets):
+            isiterable = isinstance(dataset, IterableDataset)
+            if (isiterable and not self.kind == 'iterable') or (not isiterable and self.kind == 'iterable'):
+                raise ValueError("All datasets in ConcatDataset must be of the same kind (Iterable or Map).")
+
+            if self.kind == 'map':
+                self.length += len(dataset) // world_size
+            else:
+                self.length += len(dataset)
+
+    def get_iterable(self, dataset):
+        if isinstance(dataset, IterableDataset):
+            return dataset.__iter__()
+        else:
+            indices = np.arange(len(dataset))
+            if self.shuffle:
+                np.random.shuffle(indices)
+            return iter(indices)
+
+    def __iter__(self):
+        worker_info = pt_data.get_worker_info()
+        if worker_info is None:
+            max_elements = self.length
+            wid = 0
+            wnum = 1
+        else:
+            wid = worker_info.id
+            wnum = worker_info.num_workers
+            max_elements = len(range(wid, self.length, wnum))
+
+        if self.kind == 'map':
+            for idx in range(len(self.datasets)):
+                start_idx = (len(self.datasets[idx]) // self.world_size) * self.global_rank
+                end_idx = start_idx + (len(self.datasets[idx]) // self.world_size)
+                if self.global_rank == self.world_size - 1:
+                    end_idx = len(self.datasets[idx])
+                indices = range(start_idx + wid, end_idx, wnum)
+                self.datasets[idx] = pt_data.Subset(self.datasets[idx], indices)
+
+        for idx, dataset in enumerate(self.datasets):
+            iterable = self.get_iterable(dataset)
+            self.iterables[idx] = iterable
+
+        n = 0
+        ind_gen = self.index_generator(self.datasets, **self.sampling_kwargs)
+        while n < max_elements:
+            n += 1
+            try:
+                ind = next(ind_gen)
+            except StopIteration:
+                return
+            try:
+                val = next(self.iterables[ind])
+                if self.kind == 'map':
+                    val = self.datasets[ind][val]
+                yield val
+            except StopIteration:
+                self.iterables[ind] = self.get_iterable(self.datasets[ind])
+                n -= 1
+
+    def __len__(self):
+        return self.length
+
+    @staticmethod
+    def temperature_generator(datasets, **kwargs):
+        temp = kwargs.get('temperature')
+        if not temp:
+            raise ValueError("Temperature generator expects a 'temperature' keyowrd argument.")
+
+        lengths = []
+        num = len(datasets)
+        for dataset in datasets:
+            lengths.append(len(dataset))
+
+        p = np.array(lengths) / np.sum(lengths)
+        p = np.power(p, 1 / temp)
+        p = p / np.sum(p)
+
+        while True:
+            ind = np.random.choice(np.arange(num), p=p)
+            yield ind
+
+    @staticmethod
+    def round_robin_generator(datasets, **kwargs):
+        num = len(datasets)
+        while True:
+            for i in range(num):
+                yield i
+
+    @staticmethod
+    def random_generator(datasets, **kwargs):
+        p = kwargs.get('p')
+        if not p:
+            raise ValueError("Random generator expects a 'p' keyowrd argument for sampling probabilities.")
+
+        num = len(datasets)
+        if len(p) != num:
+            raise ValueError("Length of probabilities list must be equal to the number of datasets.")
+
+        while True:
+            ind = np.random.choice(np.arange(num), p=p)
+            yield ind

tts/nemo/collections/common/losses/__init__.py

@@ -0,0 +1,20 @@
+# 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.common.losses.aggregator import AggregatorLoss
+from nemo.collections.common.losses.cross_entropy import CrossEntropyLoss, NLLLoss
+from nemo.collections.common.losses.mse_loss import MSELoss
+from nemo.collections.common.losses.multi_similarity_loss import MultiSimilarityLoss
+from nemo.collections.common.losses.smoothed_cross_entropy import SmoothedCrossEntropyLoss
+from nemo.collections.common.losses.spanning_loss import SpanningLoss

tts/nemo/collections/common/losses/aggregator.py

@@ -0,0 +1,67 @@
+# 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 typing import List
+
+import torch
+
+from nemo.core.classes import Loss, typecheck
+from nemo.core.neural_types import LossType, NeuralType
+
+__all__ = ['AggregatorLoss']
+
+
+class AggregatorLoss(Loss):
+    """
+    Sums several losses into one.
+
+    Args:
+        num_inputs: number of input losses
+        weights: a list of coefficient for merging losses
+    """
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input ports.
+        """
+        input_types = {}
+        for i in range(self._num_losses):
+            input_types["loss_" + str(i + 1)] = NeuralType(elements_type=LossType())
+
+        return input_types
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output ports.
+        """
+        return {"loss": NeuralType(elements_type=LossType())}
+
+    def __init__(self, num_inputs: int = 2, weights: List[float] = None):
+        super().__init__()
+        self._num_losses = num_inputs
+        if weights is not None and len(weights) != num_inputs:
+            raise ValueError("Length of weights should be equal to the number of inputs (num_inputs)")
+
+        self._weights = weights
+
+    @typecheck()
+    def forward(self, **kwargs):
+        values = [kwargs[x] for x in sorted(kwargs.keys())]
+        loss = torch.zeros_like(values[0])
+        for loss_idx, loss_value in enumerate(values):
+            if self._weights is not None:
+                loss = loss.add(loss_value, alpha=self._weights[loss_idx])
+            else:
+                loss = loss.add(loss_value)
+        return loss

tts/nemo/collections/common/losses/cross_entropy.py

@@ -0,0 +1,138 @@
+# 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 torch
+from torch import nn
+
+from nemo.core.classes import Serialization, Typing, typecheck
+from nemo.core.neural_types import LabelsType, LogitsType, LogprobsType, LossType, MaskType, NeuralType
+
+__all__ = ['CrossEntropyLoss', 'NLLLoss']
+
+
+class CrossEntropyLoss(nn.CrossEntropyLoss, Serialization, Typing):
+    """
+    CrossEntropyLoss
+    """
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input ports.
+        """
+        return {
+            "logits": NeuralType(['B'] + ['ANY'] * (self._logits_dim - 1), LogitsType()),
+            "labels": NeuralType(['B'] + ['ANY'] * (self._logits_dim - 2), LabelsType()),
+            "loss_mask": NeuralType(['B'] + ['ANY'] * (self._logits_dim - 2), MaskType(), optional=True),
+        }
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output ports.
+        """
+        return {"loss": NeuralType(elements_type=LossType())}
+
+    def __init__(self, logits_ndim=2, weight=None, reduction='mean', ignore_index=-100):
+        """
+        Args:
+            logits_ndim (int): number of dimensions (or rank) of the logits tensor
+            weight (list): list of rescaling weight given to each class
+            reduction (str): type of the reduction over the batch
+        """
+        if weight is not None and not torch.is_tensor(weight):
+            weight = torch.FloatTensor(weight)
+        super().__init__(weight=weight, reduction=reduction, ignore_index=ignore_index)
+        self._logits_dim = logits_ndim
+
+    @typecheck()
+    def forward(self, logits, labels, loss_mask=None):
+        """
+        Args:
+            logits (float): output of the classifier
+            labels (long): ground truth labels
+            loss_mask (bool/float/int): tensor to specify the masking
+        """
+        logits_flatten = torch.flatten(logits, start_dim=0, end_dim=-2)
+        labels_flatten = torch.flatten(labels, start_dim=0, end_dim=-1)
+
+        if loss_mask is not None:
+            if loss_mask.dtype is not torch.bool:
+                loss_mask = loss_mask > 0.5
+            loss_mask_flatten = torch.flatten(loss_mask, start_dim=0, end_dim=-1)
+            logits_flatten = logits_flatten[loss_mask_flatten]
+            labels_flatten = labels_flatten[loss_mask_flatten]
+
+        if len(labels_flatten) == 0:
+            return super().forward(logits, torch.argmax(logits, dim=-1))
+
+        loss = super().forward(logits_flatten, labels_flatten)
+        return loss
+
+
+class NLLLoss(nn.NLLLoss, Serialization, Typing):
+    """
+    NLLLoss
+    """
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input ports.
+        """
+        return {
+            "log_probs": NeuralType(("B", "T", "D"), LogprobsType()),
+            "labels": NeuralType(("B", "T"), LabelsType()),
+            "output_mask": NeuralType(("B", "T"), MaskType(), optional=True),
+        }
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output ports.
+        """
+        return {"loss": NeuralType(elements_type=LossType())}
+
+    def __init__(self, log_probs_ndim=2, weight=None, reduction='mean', ignore_index=-100):
+        """
+        Args:
+            log_probs_ndim (int): number of dimensions (or rank) of the logprobs tensor
+            weight (list): list of rescaling weight given to each class
+            reduction (str): type of the reduction over the batch
+            ignore_index (int): mask out loss computation where labels = ignore_index
+        """
+        if weight is not None and not torch.is_tensor(weight):
+            weight = torch.FloatTensor(weight)
+        super().__init__(weight=weight, reduction=reduction, ignore_index=ignore_index)
+        self._log_probs_dim = log_probs_ndim
+
+    @typecheck()
+    def forward(self, log_probs, labels, loss_mask=None):
+        """
+        Args:
+            log_probs (float): output log probability tensor
+            labels (long): ground truth labels
+            loss_mask (bool/float/int): tensor to specify the masking
+        """
+        log_probs_flatten = torch.flatten(log_probs, start_dim=0, end_dim=-2)
+        labels_flatten = torch.flatten(labels, start_dim=0, end_dim=-1)
+
+        if loss_mask is not None:
+            if loss_mask.dtype is not torch.bool:
+                loss_mask = loss_mask > 0.5
+            loss_mask_flatten = torch.flatten(loss_mask, start_dim=0, end_dim=-1)
+            log_probs_flatten = log_probs_flatten[loss_mask_flatten]
+            labels_flatten = labels_flatten[loss_mask_flatten]
+
+        if len(labels_flatten) == 0:
+            return super().forward(log_probs, torch.argmax(log_probs, dim=-1))
+
+        loss = super().forward(log_probs_flatten, labels_flatten)
+        return loss

tts/nemo/collections/common/losses/mse_loss.py

@@ -0,0 +1,57 @@
+# 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 torch import Tensor, nn
+
+from nemo.core.classes import Serialization, Typing, typecheck
+from nemo.core.neural_types import LabelsType, LossType, NeuralType, RegressionValuesType
+
+__all__ = ['MSELoss']
+
+
+class MSELoss(nn.MSELoss, Serialization, Typing):
+    """
+    MSELoss
+    """
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input ports.
+        """
+        return {
+            "preds": NeuralType(tuple('B'), RegressionValuesType()),
+            "labels": NeuralType(tuple('B'), LabelsType()),
+        }
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output ports.
+        """
+        return {"loss": NeuralType(elements_type=LossType())}
+
+    def __init__(self, reduction: str = 'mean'):
+        """
+        Args:
+            reduction: type of the reduction over the batch
+        """
+        super().__init__(reduction=reduction)
+
+    @typecheck()
+    def forward(self, preds: Tensor, labels: Tensor) -> Tensor:
+        """
+        Args:
+            preds: output of the classifier
+            labels: ground truth labels
+        """
+        return super().forward(preds, labels)