Introduce code from Rust monorepo. This is a mess

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