Introduce code from Rust monorepo. This is a mess

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

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

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

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

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

@@ -0,0 +1,129 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nemo.collections.asr.data import audio_to_label
+
+
+def get_classification_label_dataset(featurizer, config: dict) -> audio_to_label.AudioToClassificationLabelDataset:
+    """
+    Instantiates a Classification AudioLabelDataset.
+
+    Args:
+        config: Config of the AudioToClassificationLabelDataset.
+
+    Returns:
+        An instance of AudioToClassificationLabelDataset.
+    """
+    dataset = audio_to_label.AudioToClassificationLabelDataset(
+        manifest_filepath=config['manifest_filepath'],
+        labels=config['labels'],
+        featurizer=featurizer,
+        max_duration=config.get('max_duration', None),
+        min_duration=config.get('min_duration', None),
+        trim=config.get('trim_silence', False),
+        is_regression_task=config.get('is_regression_task', False),
+    )
+    return dataset
+
+
+def get_speech_label_dataset(featurizer, config: dict) -> audio_to_label.AudioToSpeechLabelDataset:
+    """
+    Instantiates a Speech Label (e.g. VAD, speaker recognition) AudioLabelDataset.
+
+    Args:
+        config: Config of the AudioToSpeechLabelDataSet.
+
+    Returns:
+        An instance of AudioToSpeechLabelDataset.
+    """
+    dataset = audio_to_label.AudioToSpeechLabelDataset(
+        manifest_filepath=config['manifest_filepath'],
+        labels=config['labels'],
+        featurizer=featurizer,
+        max_duration=config.get('max_duration', None),
+        min_duration=config.get('min_duration', None),
+        trim=config.get('trim_silence', False),
+        time_length=config.get('time_length', 0.31),
+        shift_length=config.get('shift_length', 0.01),
+        normalize_audio=config.get('normalize_audio', False),
+    )
+    return dataset
+
+
+def get_tarred_classification_label_dataset(
+    featurizer, config: dict, shuffle_n: int, global_rank: int, world_size: int
+) -> audio_to_label.TarredAudioToClassificationLabelDataset:
+    """
+    Instantiates a Classification TarredAudioLabelDataset.
+
+    Args:
+        config: Config of the TarredAudioToClassificationLabelDataset.
+        shuffle_n: How many samples to look ahead and load to be shuffled.
+            See WebDataset documentation for more details.
+        global_rank: Global rank of this device.
+        world_size: Global world size in the training method.
+
+    Returns:
+        An instance of TarredAudioToClassificationLabelDataset.
+    """
+    dataset = audio_to_label.TarredAudioToClassificationLabelDataset(
+        audio_tar_filepaths=config['tarred_audio_filepaths'],
+        manifest_filepath=config['manifest_filepath'],
+        labels=config['labels'],
+        featurizer=featurizer,
+        shuffle_n=shuffle_n,
+        max_duration=config.get('max_duration', None),
+        min_duration=config.get('min_duration', None),
+        trim=config.get('trim_silence', False),
+        shard_strategy=config.get('tarred_shard_strategy', 'scatter'),
+        global_rank=global_rank,
+        world_size=world_size,
+        is_regression_task=config.get('is_regression_task', False),
+    )
+    return dataset
+
+
+def get_tarred_speech_label_dataset(
+    featurizer, config: dict, shuffle_n: int, global_rank: int, world_size: int,
+) -> audio_to_label.TarredAudioToSpeechLabelDataset:
+    """
+    InInstantiates a Speech Label (e.g. VAD, speaker recognition) TarredAudioLabelDataset.
+
+    Args:
+        config: Config of the TarredAudioToSpeechLabelDataset.
+        shuffle_n: How many samples to look ahead and load to be shuffled.
+            See WebDataset documentation for more details.
+        global_rank: Global rank of this device.
+        world_size: Global world size in the training method.
+
+    Returns:
+        An instance of TarredAudioToSpeechLabelDataset.
+    """
+    dataset = audio_to_label.TarredAudioToSpeechLabelDataset(
+        audio_tar_filepaths=config['tarred_audio_filepaths'],
+        manifest_filepath=config['manifest_filepath'],
+        labels=config['labels'],
+        featurizer=featurizer,
+        shuffle_n=shuffle_n,
+        max_duration=config.get('max_duration', None),
+        min_duration=config.get('min_duration', None),
+        trim=config.get('trim_silence', False),
+        time_length=config.get('time_length', 0.31),
+        shift_length=config.get('shift_length', 0.01),
+        normalize_audio=config.get('normalize_audio', False),
+        shard_strategy=config.get('tarred_shard_strategy', 'scatter'),
+        global_rank=global_rank,
+        world_size=world_size,
+    )
+    return dataset

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

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

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

@@ -0,0 +1,247 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional
+
+import torch
+from omegaconf import DictConfig, open_dict
+
+from nemo.collections.asr.data import audio_to_text, audio_to_text_dali
+from nemo.utils import logging
+
+
+def inject_dataloader_value_from_model_config(model_cfg: dict, dataloader_cfg: dict, key: str):
+    """
+    Extracts the label set provided at the top level of the model, and propagates it to the dataloader
+    config.
+
+    Args:
+        model_cfg: A DictConfig representing the model's config.
+        dataloader_cfg: A DictConfig representing the individual data loader
+        key: A str value representing a key in the model_cfg whose value will be propagated to the
+            dataloader config.
+    """
+    if key not in model_cfg:
+        logging.info(
+            f"Model level config does not container `{key}`, please explicitly provide `{key}` to the dataloaders."
+        )
+        return
+
+    # If key exists in the data loader config (either set explicitly or as a placeholder (via None))
+    if key in dataloader_cfg:
+        # Dataloader `labels` is provided and is non-null
+        if dataloader_cfg[key] is not None and model_cfg[key] != dataloader_cfg[key]:
+            # Model level `labels` dont match Dataloader level `labels`
+            logging.warning(
+                f'`{key}` is explicitly provided to the data loader, and is different from '
+                f'the `{key}` provided at the model level config.\n'
+                f'If this is incorrect, please set the dataloader\'s `{key}` to None.'
+            )
+
+        else:
+            # Dataloader `key` is None or values match
+            # Propagate from model level `key` (even if they match)
+            with open_dict(dataloader_cfg):
+                dataloader_cfg[key] = model_cfg[key]
+
+    else:
+        # If key key doesnt even exist in dataloader_cfg, inject it explicitly
+        with open_dict(dataloader_cfg):
+            dataloader_cfg[key] = model_cfg[key]
+
+
+def get_char_dataset(config: dict, augmentor: Optional['AudioAugmentor'] = None) -> audio_to_text.AudioToCharDataset:
+    """
+    Instantiates a Character Encoding based AudioToCharDataset.
+
+    Args:
+        config: Config of the AudioToCharDataset.
+        augmentor: Optional AudioAugmentor object for augmentations on audio data.
+
+    Returns:
+        An instance of AudioToCharDataset.
+    """
+    dataset = audio_to_text.AudioToCharDataset(
+        manifest_filepath=config['manifest_filepath'],
+        labels=config['labels'],
+        sample_rate=config['sample_rate'],
+        int_values=config.get('int_values', False),
+        augmentor=augmentor,
+        max_duration=config.get('max_duration', None),
+        min_duration=config.get('min_duration', None),
+        max_utts=config.get('max_utts', 0),
+        blank_index=config.get('blank_index', -1),
+        unk_index=config.get('unk_index', -1),
+        normalize=config.get('normalize_transcripts', False),
+        trim=config.get('trim_silence', False),
+        parser=config.get('parser', 'en'),
+    )
+    return dataset
+
+
+def get_bpe_dataset(
+    config: dict, tokenizer: 'TokenizerSpec', augmentor: Optional['AudioAugmentor'] = None
+) -> audio_to_text.AudioToBPEDataset:
+    """
+    Instantiates a Byte Pair Encoding / Word Piece Encoding based AudioToBPEDataset.
+
+    Args:
+        config: Config of the AudioToBPEDataset.
+        tokenizer: An instance of a TokenizerSpec object.
+        augmentor: Optional AudioAugmentor object for augmentations on audio data.
+
+    Returns:
+        An instance of AudioToBPEDataset.
+    """
+    dataset = audio_to_text.AudioToBPEDataset(
+        manifest_filepath=config['manifest_filepath'],
+        tokenizer=tokenizer,
+        sample_rate=config['sample_rate'],
+        int_values=config.get('int_values', False),
+        augmentor=augmentor,
+        max_duration=config.get('max_duration', None),
+        min_duration=config.get('min_duration', None),
+        max_utts=config.get('max_utts', 0),
+        trim=config.get('trim_silence', False),
+        use_start_end_token=config.get('use_start_end_token', True),
+    )
+    return dataset
+
+
+def get_tarred_char_dataset(
+    config: dict, shuffle_n: int, global_rank: int, world_size: int, augmentor: Optional['AudioAugmentor'] = None
+) -> audio_to_text.TarredAudioToCharDataset:
+    """
+    Instantiates a Character Encoding based TarredAudioToCharDataset.
+
+    Args:
+        config: Config of the TarredAudioToCharDataset.
+        shuffle_n: How many samples to look ahead and load to be shuffled.
+            See WebDataset documentation for more details.
+        global_rank: Global rank of this device.
+        world_size: Global world size in the training method.
+        augmentor: Optional AudioAugmentor object for augmentations on audio data.
+
+    Returns:
+        An instance of TarredAudioToCharDataset.
+    """
+    dataset = audio_to_text.TarredAudioToCharDataset(
+        audio_tar_filepaths=config['tarred_audio_filepaths'],
+        manifest_filepath=config['manifest_filepath'],
+        labels=config['labels'],
+        sample_rate=config['sample_rate'],
+        int_values=config.get('int_values', False),
+        augmentor=augmentor,
+        shuffle_n=shuffle_n,
+        max_duration=config.get('max_duration', None),
+        min_duration=config.get('min_duration', None),
+        max_utts=config.get('max_utts', 0),
+        blank_index=config.get('blank_index', -1),
+        unk_index=config.get('unk_index', -1),
+        normalize=config.get('normalize_transcripts', False),
+        trim=config.get('trim_silence', False),
+        parser=config.get('parser', 'en'),
+        shard_strategy=config.get('tarred_shard_strategy', 'scatter'),
+        global_rank=global_rank,
+        world_size=world_size,
+    )
+    return dataset
+
+
+def get_tarred_bpe_dataset(
+    config: dict,
+    tokenizer: 'TokenizerSpec',
+    shuffle_n: int,
+    global_rank: int,
+    world_size: int,
+    augmentor: Optional['AudioAugmentor'] = None,
+) -> audio_to_text.TarredAudioToBPEDataset:
+    """
+    Instantiates a Byte Pair Encoding / Word Piece Encoding based TarredAudioToBPEDataset.
+
+    Args:
+        config: Config of the TarredAudioToBPEDataset.
+        tokenizer: An instance of a TokenizerSpec object.
+        shuffle_n: How many samples to look ahead and load to be shuffled.
+            See WebDataset documentation for more details.
+        global_rank: Global rank of this device.
+        world_size: Global world size in the training method.
+        augmentor: Optional AudioAugmentor object for augmentations on audio data.
+
+    Returns:
+        An instance of TarredAudioToBPEDataset.
+    """
+    dataset = audio_to_text.TarredAudioToBPEDataset(
+        audio_tar_filepaths=config['tarred_audio_filepaths'],
+        manifest_filepath=config['manifest_filepath'],
+        tokenizer=tokenizer,
+        sample_rate=config['sample_rate'],
+        int_values=config.get('int_values', False),
+        augmentor=augmentor,
+        shuffle_n=shuffle_n,
+        max_duration=config.get('max_duration', None),
+        min_duration=config.get('min_duration', None),
+        max_utts=config.get('max_utts', 0),
+        trim=config.get('trim_silence', False),
+        use_start_end_token=config.get('use_start_end_token', True),
+        shard_strategy=config.get('tarred_shard_strategy', 'scatter'),
+        global_rank=global_rank,
+        world_size=world_size,
+    )
+    return dataset
+
+
+def get_dali_char_dataset(
+    config: dict,
+    shuffle: bool,
+    device_id: int,
+    global_rank: int,
+    world_size: int,
+    preprocessor_cfg: Optional[DictConfig] = None,
+) -> audio_to_text_dali.AudioToCharDALIDataset:
+    """
+    Instantiates a Character Encoding based AudioToCharDALIDataset.
+
+    Args:
+        config: Config of the AudioToCharDALIDataset.
+        shuffle: Bool flag whether to shuffle the dataset.
+        device_id: Index of the GPU to be used (local_rank). Only applicable when device == 'gpu'. Defaults to 0.
+        global_rank: Global rank of this device.
+        world_size: Global world size in the training method.
+        augmentor: Optional AudioAugmentor object for augmentations on audio data.
+
+    Returns:
+        An instance of AudioToCharDALIDataset.
+    """
+    device = 'gpu' if torch.cuda.is_available() else 'cpu'
+    dataset = audio_to_text_dali.AudioToCharDALIDataset(
+        manifest_filepath=config['manifest_filepath'],
+        device=device,
+        batch_size=config['batch_size'],
+        labels=config['labels'],
+        sample_rate=config['sample_rate'],
+        max_duration=config.get('max_duration', None),
+        min_duration=config.get('min_duration', None),
+        blank_index=config.get('blank_index', -1),
+        unk_index=config.get('unk_index', -1),
+        normalize=config.get('normalize_transcripts', False),
+        trim=config.get('trim_silence', False),
+        parser=config.get('parser', 'en'),
+        shuffle=shuffle,
+        device_id=device_id,
+        global_rank=global_rank,
+        world_size=world_size,
+        preprocessor_cfg=preprocessor_cfg,
+    )
+    return dataset

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

@@ -0,0 +1,139 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Dict, List, Optional
+
+import torch
+
+from nemo.collections.common.parts.preprocessing import collections
+from nemo.core.classes import Dataset
+from nemo.core.neural_types import AcousticEncodedRepresentation, LabelsType, LengthsType, NeuralType
+from nemo.utils import logging
+
+
+def _feature_collate_fn(batch):
+    """collate batch of feat sig, feat len, tokens, tokens len
+    Args:
+        batch (FloatTensor, LongTensor, LongTensor, LongTensor):  A tuple of tuples of feature, feature lengths,
+               encoded tokens, and encoded tokens length. 
+    """
+    _, feat_lengths, _, tokens_lengths = zip(*batch)
+    feat_signal, tokens = [], []
+    for sig, sig_len, tokens_i, tokens_i_len in batch:
+        feat_signal.append(sig)
+        tokens.append(tokens_i)
+
+    feat_signal = torch.stack(feat_signal)
+    feat_lengths = torch.stack(feat_lengths)
+
+    tokens = torch.stack(tokens)
+    tokens_lengths = torch.stack(tokens_lengths)
+
+    return feat_signal, feat_lengths, tokens, tokens_lengths
+
+
+class _FeatureSeqSpeakerLabelDataset(Dataset):
+    """
+    Dataset that loads tensors via a json file containing paths to feature files, sequences of labels. 
+    Each new line is a different sample. Example below:
+    and their target labels. JSON files should be of the following format:
+        {"feature_filepath": "/path/to/feature_0.p", "seq_label": speakerA speakerB SpeakerA ....} \
+        ...
+        {"feature_filepath": "/path/to/feature_n.p", "seq_label": target_seq_label_n} 
+    target_seq_label_n is the string of sequence of speaker label, separated by space.
+
+    Args:
+        manifest_filepath (str): Dataset parameter. Path to JSON containing data.
+        labels (Optional[list]): Dataset parameter. List of unique labels collected from all samples.
+        feature_loader : Dataset parameter. Feature loader to load (external) feature.       
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Returns definitions of module output ports.
+        """
+        # TODO output type for external features
+        output_types = {
+            'external_feat': NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()),
+            'feat_length': NeuralType(tuple('B'), LengthsType()),
+        }
+
+        if self.is_speaker_emb:
+            output_types.update(
+                {
+                    'embs': NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()),
+                    'embs_length': NeuralType(tuple('B'), LengthsType()),
+                    'label': NeuralType(('B', 'T'), LabelsType()),
+                    'label_length': NeuralType(tuple('B'), LengthsType()),
+                }
+            )
+        else:
+            output_types.update(
+                {'label': NeuralType(('B', 'T'), LabelsType()), 'label_length': NeuralType(tuple('B'), LengthsType()),}
+            )
+
+        return output_types
+
+    def __init__(
+        self, *, manifest_filepath: str, labels: List[str], feature_loader, is_speaker_emb: bool = False,
+    ):
+        super().__init__()
+        self.collection = collections.ASRFeatureSequenceLabel(manifests_files=manifest_filepath.split(','),)
+
+        self.feature_loader = feature_loader
+        self.labels = labels if labels else self.collection.uniq_labels
+        self.is_speaker_emb = is_speaker_emb
+
+        self.label2id, self.id2label = {}, {}
+        for label_id, label in enumerate(self.labels):
+            self.label2id[label] = label_id
+            self.id2label[label_id] = label
+
+        for idx in range(len(self.labels[:5])):
+            logging.debug(" label id {} and its mapped label {}".format(idx, self.id2label[idx]))
+
+    def __len__(self):
+        return len(self.collection)
+
+    def __getitem__(self, index):
+        sample = self.collection[index]
+
+        features = self.feature_loader.process(sample.feature_file)
+        f, fl = features, torch.tensor(features.shape[0]).long()
+
+        t = torch.tensor(sample.seq_label).float()
+        tl = torch.tensor(len(sample.seq_label)).long()
+
+        return f, fl, t, tl
+
+
+class FeatureToSeqSpeakerLabelDataset(_FeatureSeqSpeakerLabelDataset):
+    """
+    Dataset that loads tensors via a json file containing paths to feature
+    files and sequence of speakers. Each new line is a
+    different sample. Example below:
+    {"feature_filepath": "/path/to/feature_0.p", "seq_label": speakerA speakerB SpeakerA ....} \
+    ...
+    {"feature_filepath": "/path/to/feature_n.p", "seq_label": target_seq_label_n} 
+    target_seq_label_n is the string of sequence of speaker label, separated by space.
+
+    Args:
+        manifest_filepath (str): Path to manifest json as described above. Canbe comma-separated paths.
+        labels (Optional[list]): String containing all the possible labels to map to
+            if None then automatically picks from ASRFeatureSequenceLabel collection.
+        feature_loader, Feature load to loader (external) feature.
+    
+    """
+
+    def _collate_fn(self, batch):
+        return _feature_collate_fn(batch)

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

@@ -0,0 +1,31 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from nemo.collections.asr.data import feature_to_label
+
+
+def get_feature_seq_speakerlabel_dataset(
+    feature_loader, config: dict
+) -> feature_to_label.FeatureToSeqSpeakerLabelDataset:
+    """
+    Instantiates a FeatureSeqSpeakerLabelDataset.
+    Args:
+        config: Config of the FeatureToSeqSpeakerLabelDataset.
+
+    Returns:
+        An instance of FeatureToSeqSpeakerLabelDataset.
+    """
+    dataset = feature_to_label.FeatureToSeqSpeakerLabelDataset(
+        manifest_filepath=config['manifest_filepath'], labels=config['labels'], feature_loader=feature_loader,
+    )
+    return dataset

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

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