Introduce code from Rust monorepo. This is a mess

tts/nemo/collections/common/__init__.py

@@ -0,0 +1,26 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import nemo.collections.common.callbacks
+from nemo.collections.common import data, losses, parts, tokenizers
+from nemo.package_info import __version__
+
+# Set collection version equal to NeMo version.
+__version = __version__
+
+# Authorship.
+__author__ = "NVIDIA Corporation"
+
+# Set collection name.
+__description__ = "Common collection"

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

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

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

@@ -0,0 +1,97 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import time
+
+import numpy as np
+from pytorch_lightning.callbacks.base import Callback
+from pytorch_lightning.utilities import rank_zero_only
+
+# from sacrebleu import corpus_bleu
+
+
+class LogEpochTimeCallback(Callback):
+    """Simple callback that logs how long each epoch takes, in seconds, to a pytorch lightning log
+    """
+
+    @rank_zero_only
+    def on_train_epoch_start(self, trainer, pl_module):
+        self.epoch_start = time.time()
+
+    @rank_zero_only
+    def on_train_epoch_end(self, trainer, pl_module, outputs):
+        curr_time = time.time()
+        duration = curr_time - self.epoch_start
+        trainer.logger.log_metrics({"epoch_time": duration}, step=trainer.global_step)
+
+
+# class MachineTranslationLogEvalCallback(Callback):
+#     def _on_eval_end(self, trainer, pl_module, mode):
+#         counts = np.array(self._non_pad_tokens)
+#         eval_loss = np.sum(np.array(self._losses) * counts) / np.sum(counts)
+#         sacre_bleu = corpus_bleu(self._translations, [self._ground_truths], tokenize="13a")
+#         print(f"{mode} results for process with global rank {pl_module.global_rank}".upper())
+#         for i in range(pl_module.num_examples[mode]):
+#             print('\u0332'.join(f"EXAMPLE {i}:"))  # Underline output
+#             sent_id = np.random.randint(len(self._translations))
+#             print(f"Ground truth: {self._ground_truths[sent_id]}\n")
+#             print(f"Translation: {self._translations[sent_id]}\n")
+#             print()
+#         print("-" * 50)
+#         print(f"loss: {eval_loss:.3f}")
+#         print(f"SacreBLEU: {sacre_bleu}")
+#         print("-" * 50)
+
+#     @rank_zero_only
+#     def on_test_end(self, trainer, pl_module):
+#         self._on_eval_end(trainer, pl_module, "test")
+
+#     @rank_zero_only
+#     def on_validation_end(self, trainer, pl_module):
+#         self._on_eval_end(trainer, pl_module, "val")
+
+#     @rank_zero_only
+#     def on_sanity_check_end(self, trainer, pl_module):
+#         self._on_eval_end(trainer, pl_module, "val")
+
+#     def _on_eval_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx, mode):
+#         self._translations.extend(outputs['translations'])
+#         self._ground_truths.extend(outputs['ground_truths'])
+#         self._non_pad_tokens.append(outputs['num_non_pad_tokens'])
+#         self._losses.append(outputs[f'{mode}_loss'])
+
+#     @rank_zero_only
+#     def on_test_batch_end(self, trainer, pl_module, batch, outputs, batch_idx, dataloader_idx):
+#         self._on_eval_batch_end(trainer, pl_module, batch, outputs, batch_idx, dataloader_idx, 'test')
+
+#     @rank_zero_only
+#     def on_validation_batch_end(self, trainer, pl_module, batch, outputs, batch_idx, dataloader_idx):
+#         self._on_eval_batch_end(trainer, pl_module, batch, outputs, batch_idx, dataloader_idx, 'val')
+
+#     def _on_eval_start(self, trainer, pl_module):
+#         self._translations = []
+#         self._ground_truths = []
+#         self._losses = []
+#         self._non_pad_tokens = []
+
+#     @rank_zero_only
+#     def on_test_start(self, trainer, pl_module):
+#         self._on_eval_start(trainer, pl_module)
+
+#     @rank_zero_only
+#     def on_validation_start(self, trainer, pl_module):
+#         self._on_eval_start(trainer, pl_module)
+
+#     @rank_zero_only
+#     def on_sanity_check_start(self, trainer, pl_module):
+#         self._on_eval_start(trainer, pl_module)

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

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

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

@@ -0,0 +1,178 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, List
+
+import numpy as np
+import torch.utils.data as pt_data
+from torch.utils.data import IterableDataset
+
+__all__ = ['ConcatDataset']
+
+
+class ConcatDataset(IterableDataset):
+    """
+    A dataset that accepts as argument multiple datasets and then samples from them based on the specified 
+    sampling technique.
+    Args:
+        datasets (list): A list of datasets to sample from.
+        shuffle (bool): Whether to shuffle individual datasets. Only works with non-iterable datasets. 
+            Defaults to True.
+        sampling_technique (str): Sampling technique to choose which dataset to draw a sample from.
+            Defaults to 'temperature'. Currently supports 'temperature', 'random' and 'round-robin'.
+        sampling_temperature (int): Temperature value for sampling. Only used when sampling_technique = 'temperature'.
+            Defaults to 5.
+        sampling_probabilities (list): Probability values for sampling. Only used when sampling_technique = 'random'.
+        global_rank (int): Worker rank, used for partitioning map style datasets. Defaults to 0.
+        world_size (int): Total number of processes, used for partitioning map style datasets. Defaults to 1.
+    """
+
+    def __init__(
+        self,
+        datasets: List[Any],
+        shuffle: bool = True,
+        sampling_technique: str = 'temperature',
+        sampling_temperature: int = 5,
+        sampling_probabilities: List[float] = None,
+        global_rank: int = 0,
+        world_size: int = 1,
+    ):
+        super().__init__()
+
+        supported_sampling_techniques = ['temperature', 'random', 'round-robin']
+        self.datasets = datasets
+        self.iterables = [None] * len(datasets)
+        self.shuffle = shuffle
+        self.global_rank = global_rank
+        self.world_size = world_size
+        self.sampling_kwargs = {}
+        if sampling_technique == 'temperature':
+            self.index_generator = ConcatDataset.temperature_generator
+            self.sampling_kwargs['temperature'] = sampling_temperature
+        elif sampling_technique == 'random':
+            self.index_generator = ConcatDataset.random_generator
+            self.sampling_kwargs['p'] = sampling_probabilities
+        elif sampling_technique == 'round-robin':
+            self.index_generator = ConcatDataset.round_robin_generator
+        else:
+            raise ValueError(f"Currently we only support sampling techniques in {supported_sampling_techniques}.")
+        self.length = 0
+
+        if isinstance(datasets[0], IterableDataset):
+            self.kind = 'iterable'
+        else:
+            self.kind = 'map'
+
+        for idx, dataset in enumerate(datasets):
+            isiterable = isinstance(dataset, IterableDataset)
+            if (isiterable and not self.kind == 'iterable') or (not isiterable and self.kind == 'iterable'):
+                raise ValueError("All datasets in ConcatDataset must be of the same kind (Iterable or Map).")
+
+            if self.kind == 'map':
+                self.length += len(dataset) // world_size
+            else:
+                self.length += len(dataset)
+
+    def get_iterable(self, dataset):
+        if isinstance(dataset, IterableDataset):
+            return dataset.__iter__()
+        else:
+            indices = np.arange(len(dataset))
+            if self.shuffle:
+                np.random.shuffle(indices)
+            return iter(indices)
+
+    def __iter__(self):
+        worker_info = pt_data.get_worker_info()
+        if worker_info is None:
+            max_elements = self.length
+            wid = 0
+            wnum = 1
+        else:
+            wid = worker_info.id
+            wnum = worker_info.num_workers
+            max_elements = len(range(wid, self.length, wnum))
+
+        if self.kind == 'map':
+            for idx in range(len(self.datasets)):
+                start_idx = (len(self.datasets[idx]) // self.world_size) * self.global_rank
+                end_idx = start_idx + (len(self.datasets[idx]) // self.world_size)
+                if self.global_rank == self.world_size - 1:
+                    end_idx = len(self.datasets[idx])
+                indices = range(start_idx + wid, end_idx, wnum)
+                self.datasets[idx] = pt_data.Subset(self.datasets[idx], indices)
+
+        for idx, dataset in enumerate(self.datasets):
+            iterable = self.get_iterable(dataset)
+            self.iterables[idx] = iterable
+
+        n = 0
+        ind_gen = self.index_generator(self.datasets, **self.sampling_kwargs)
+        while n < max_elements:
+            n += 1
+            try:
+                ind = next(ind_gen)
+            except StopIteration:
+                return
+            try:
+                val = next(self.iterables[ind])
+                if self.kind == 'map':
+                    val = self.datasets[ind][val]
+                yield val
+            except StopIteration:
+                self.iterables[ind] = self.get_iterable(self.datasets[ind])
+                n -= 1
+
+    def __len__(self):
+        return self.length
+
+    @staticmethod
+    def temperature_generator(datasets, **kwargs):
+        temp = kwargs.get('temperature')
+        if not temp:
+            raise ValueError("Temperature generator expects a 'temperature' keyowrd argument.")
+
+        lengths = []
+        num = len(datasets)
+        for dataset in datasets:
+            lengths.append(len(dataset))
+
+        p = np.array(lengths) / np.sum(lengths)
+        p = np.power(p, 1 / temp)
+        p = p / np.sum(p)
+
+        while True:
+            ind = np.random.choice(np.arange(num), p=p)
+            yield ind
+
+    @staticmethod
+    def round_robin_generator(datasets, **kwargs):
+        num = len(datasets)
+        while True:
+            for i in range(num):
+                yield i
+
+    @staticmethod
+    def random_generator(datasets, **kwargs):
+        p = kwargs.get('p')
+        if not p:
+            raise ValueError("Random generator expects a 'p' keyowrd argument for sampling probabilities.")
+
+        num = len(datasets)
+        if len(p) != num:
+            raise ValueError("Length of probabilities list must be equal to the number of datasets.")
+
+        while True:
+            ind = np.random.choice(np.arange(num), p=p)
+            yield ind

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

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

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

@@ -0,0 +1,67 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List
+
+import torch
+
+from nemo.core.classes import Loss, typecheck
+from nemo.core.neural_types import LossType, NeuralType
+
+__all__ = ['AggregatorLoss']
+
+
+class AggregatorLoss(Loss):
+    """
+    Sums several losses into one.
+
+    Args:
+        num_inputs: number of input losses
+        weights: a list of coefficient for merging losses
+    """
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input ports.
+        """
+        input_types = {}
+        for i in range(self._num_losses):
+            input_types["loss_" + str(i + 1)] = NeuralType(elements_type=LossType())
+
+        return input_types
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output ports.
+        """
+        return {"loss": NeuralType(elements_type=LossType())}
+
+    def __init__(self, num_inputs: int = 2, weights: List[float] = None):
+        super().__init__()
+        self._num_losses = num_inputs
+        if weights is not None and len(weights) != num_inputs:
+            raise ValueError("Length of weights should be equal to the number of inputs (num_inputs)")
+
+        self._weights = weights
+
+    @typecheck()
+    def forward(self, **kwargs):
+        values = [kwargs[x] for x in sorted(kwargs.keys())]
+        loss = torch.zeros_like(values[0])
+        for loss_idx, loss_value in enumerate(values):
+            if self._weights is not None:
+                loss = loss.add(loss_value, alpha=self._weights[loss_idx])
+            else:
+                loss = loss.add(loss_value)
+        return loss

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

@@ -0,0 +1,138 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from torch import nn
+
+from nemo.core.classes import Serialization, Typing, typecheck
+from nemo.core.neural_types import LabelsType, LogitsType, LogprobsType, LossType, MaskType, NeuralType
+
+__all__ = ['CrossEntropyLoss', 'NLLLoss']
+
+
+class CrossEntropyLoss(nn.CrossEntropyLoss, Serialization, Typing):
+    """
+    CrossEntropyLoss
+    """
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input ports.
+        """
+        return {
+            "logits": NeuralType(['B'] + ['ANY'] * (self._logits_dim - 1), LogitsType()),
+            "labels": NeuralType(['B'] + ['ANY'] * (self._logits_dim - 2), LabelsType()),
+            "loss_mask": NeuralType(['B'] + ['ANY'] * (self._logits_dim - 2), MaskType(), optional=True),
+        }
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output ports.
+        """
+        return {"loss": NeuralType(elements_type=LossType())}
+
+    def __init__(self, logits_ndim=2, weight=None, reduction='mean', ignore_index=-100):
+        """
+        Args:
+            logits_ndim (int): number of dimensions (or rank) of the logits tensor
+            weight (list): list of rescaling weight given to each class
+            reduction (str): type of the reduction over the batch
+        """
+        if weight is not None and not torch.is_tensor(weight):
+            weight = torch.FloatTensor(weight)
+        super().__init__(weight=weight, reduction=reduction, ignore_index=ignore_index)
+        self._logits_dim = logits_ndim
+
+    @typecheck()
+    def forward(self, logits, labels, loss_mask=None):
+        """
+        Args:
+            logits (float): output of the classifier
+            labels (long): ground truth labels
+            loss_mask (bool/float/int): tensor to specify the masking
+        """
+        logits_flatten = torch.flatten(logits, start_dim=0, end_dim=-2)
+        labels_flatten = torch.flatten(labels, start_dim=0, end_dim=-1)
+
+        if loss_mask is not None:
+            if loss_mask.dtype is not torch.bool:
+                loss_mask = loss_mask > 0.5
+            loss_mask_flatten = torch.flatten(loss_mask, start_dim=0, end_dim=-1)
+            logits_flatten = logits_flatten[loss_mask_flatten]
+            labels_flatten = labels_flatten[loss_mask_flatten]
+
+        if len(labels_flatten) == 0:
+            return super().forward(logits, torch.argmax(logits, dim=-1))
+
+        loss = super().forward(logits_flatten, labels_flatten)
+        return loss
+
+
+class NLLLoss(nn.NLLLoss, Serialization, Typing):
+    """
+    NLLLoss
+    """
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input ports.
+        """
+        return {
+            "log_probs": NeuralType(("B", "T", "D"), LogprobsType()),
+            "labels": NeuralType(("B", "T"), LabelsType()),
+            "output_mask": NeuralType(("B", "T"), MaskType(), optional=True),
+        }
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output ports.
+        """
+        return {"loss": NeuralType(elements_type=LossType())}
+
+    def __init__(self, log_probs_ndim=2, weight=None, reduction='mean', ignore_index=-100):
+        """
+        Args:
+            log_probs_ndim (int): number of dimensions (or rank) of the logprobs tensor
+            weight (list): list of rescaling weight given to each class
+            reduction (str): type of the reduction over the batch
+            ignore_index (int): mask out loss computation where labels = ignore_index
+        """
+        if weight is not None and not torch.is_tensor(weight):
+            weight = torch.FloatTensor(weight)
+        super().__init__(weight=weight, reduction=reduction, ignore_index=ignore_index)
+        self._log_probs_dim = log_probs_ndim
+
+    @typecheck()
+    def forward(self, log_probs, labels, loss_mask=None):
+        """
+        Args:
+            log_probs (float): output log probability tensor
+            labels (long): ground truth labels
+            loss_mask (bool/float/int): tensor to specify the masking
+        """
+        log_probs_flatten = torch.flatten(log_probs, start_dim=0, end_dim=-2)
+        labels_flatten = torch.flatten(labels, start_dim=0, end_dim=-1)
+
+        if loss_mask is not None:
+            if loss_mask.dtype is not torch.bool:
+                loss_mask = loss_mask > 0.5
+            loss_mask_flatten = torch.flatten(loss_mask, start_dim=0, end_dim=-1)
+            log_probs_flatten = log_probs_flatten[loss_mask_flatten]
+            labels_flatten = labels_flatten[loss_mask_flatten]
+
+        if len(labels_flatten) == 0:
+            return super().forward(log_probs, torch.argmax(log_probs, dim=-1))
+
+        loss = super().forward(log_probs_flatten, labels_flatten)
+        return loss

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

@@ -0,0 +1,57 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from torch import Tensor, nn
+
+from nemo.core.classes import Serialization, Typing, typecheck
+from nemo.core.neural_types import LabelsType, LossType, NeuralType, RegressionValuesType
+
+__all__ = ['MSELoss']
+
+
+class MSELoss(nn.MSELoss, Serialization, Typing):
+    """
+    MSELoss
+    """
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input ports.
+        """
+        return {
+            "preds": NeuralType(tuple('B'), RegressionValuesType()),
+            "labels": NeuralType(tuple('B'), LabelsType()),
+        }
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output ports.
+        """
+        return {"loss": NeuralType(elements_type=LossType())}
+
+    def __init__(self, reduction: str = 'mean'):
+        """
+        Args:
+            reduction: type of the reduction over the batch
+        """
+        super().__init__(reduction=reduction)
+
+    @typecheck()
+    def forward(self, preds: Tensor, labels: Tensor) -> Tensor:
+        """
+        Args:
+            preds: output of the classifier
+            labels: ground truth labels
+        """
+        return super().forward(preds, labels)

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

@@ -0,0 +1,95 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional
+
+import torch
+
+from nemo.core.classes import Loss
+from nemo.core.classes.common import typecheck
+from nemo.core.neural_types import LabelsType, LogitsType, LossType, NeuralType
+from nemo.utils import logging
+
+__all__ = ['MultiSimilarityLoss']
+
+
+class MultiSimilarityLoss(Loss):
+    @property
+    def input_types(self):
+        """Returns definitions of module input ports.
+        """
+        return {"logits": NeuralType(('B', 'D'), LogitsType()), "labels": NeuralType(('B'), LabelsType())}
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output ports.
+        """
+        return {"loss": NeuralType(elements_type=LossType())}
+
+    def __init__(
+        self,
+        scale_pos: Optional[float] = 2.0,  # Params found to work best in our experiments
+        scale_neg: Optional[float] = 40.0,
+        offset: Optional[float] = 0.5,
+        margin: Optional[float] = 0.1,
+    ):
+        super().__init__()
+        self._scale_pos = scale_pos
+        self._scale_neg = scale_neg
+        self._offset = offset
+        self._margin = margin
+        self._epsilon = 1e-5
+
+    @typecheck()
+    def forward(self, logits, labels):
+        cos_sim = torch.matmul(logits, torch.t(logits))
+        losses = []
+
+        for i in range(logits.size(0)):
+            # mine hard pairs relative to anchor i
+            positive_sims = cos_sim[i][labels.eq(labels[i])]
+            positive_sims = positive_sims[positive_sims.lt(1 - self._epsilon)]  # omit identical pairs
+            negative_sims = cos_sim[i][labels.ne(labels[i])]
+
+            if len(negative_sims) == 0 or len(positive_sims) == 0:
+                continue
+
+            # negatives that are more similar than the least-similar positive
+            hard_negatives = negative_sims[negative_sims.gt(min(positive_sims) - self._margin)]
+
+            # positives that are less similar than the most-similar negative
+            hard_positives = positive_sims[positive_sims.lt(max(negative_sims) + self._margin)]
+
+            if len(hard_negatives) == 0 or len(hard_positives) == 0:
+                continue
+
+            pos_term = (
+                1.0
+                / self._scale_pos
+                * torch.log(1 + torch.sum(torch.exp(-self._scale_pos * (hard_positives - self._offset))))
+            )
+            neg_term = (
+                1.0
+                / self._scale_neg
+                * torch.log(1 + torch.sum(torch.exp(self._scale_neg * (hard_negatives - self._offset))))
+            )
+            losses.append(pos_term + neg_term)
+
+        if len(losses) == 0:
+            loss = torch.zeros([], requires_grad=True).cuda()
+            logging.info(f'Encountered zero loss in multisimloss, loss = {loss}. No hard examples found in the batch')
+        else:
+            loss = torch.sum(torch.stack(losses)) / logits.size(0)
+
+        return loss

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

@@ -0,0 +1,103 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional
+
+import torch
+
+from nemo.core.classes import Loss, typecheck
+from nemo.core.neural_types import LabelsType, LogprobsType, LossType, MaskType, NeuralType
+
+__all__ = ["SmoothedCrossEntropyLoss"]
+
+
+class SmoothedCrossEntropyLoss(Loss):
+    """
+    Calculates Cross-entropy loss with label smoothing for a batch of sequences.
+
+    SmoothedCrossEntropyLoss:
+    1) excludes padding tokens from loss calculation
+    2) allows to use label smoothing regularization
+    3) allows to calculate loss for the desired number of last tokens
+
+    Args:
+        label_smoothing (float): label smoothing regularization coefficient
+        predict_last_k (int): parameter which sets the number of last tokens to calculate the loss for, for example
+            0: (default) calculate loss on the entire sequence (e.g., NMT)
+            1: calculate loss on the last token only (e.g., LM evaluation)
+            Intermediate values allow to control the trade-off between eval
+            time (proportional to the number of batches) and eval performance
+            (proportional to the number of context tokens)
+        pad_id (int): padding id
+        eps (float): the small eps number to avoid division buy zero
+    """
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input ports.
+        """
+        return {
+            "log_probs": NeuralType(("B", "T", "D"), LogprobsType()),
+            "labels": NeuralType(("B", "T"), LabelsType()),
+            "output_mask": NeuralType(("B", "T"), MaskType(), optional=True),
+        }
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output ports.
+        """
+        return {"loss": NeuralType(elements_type=LossType())}
+
+    def __init__(
+        self,
+        pad_id: Optional[int] = None,
+        label_smoothing: Optional[float] = 0.0,
+        predict_last_k: Optional[int] = 0,
+        eps: float = 1e-6,
+    ):
+        super().__init__()
+        self._pad_id = pad_id
+        self._eps = eps
+        self._predict_last_k = predict_last_k
+        self._label_smoothing = label_smoothing
+
+    @typecheck()
+    def forward(self, log_probs, labels, output_mask=None):
+        """
+        Args:
+            log_probs: float tensor of shape batch_size x seq_len x vocab_size, values should be log probabilities
+            labels: int tensor of shape batch_size x seq_len
+            output_mask: binary tensor of shape batch_size x seq_len
+            eps: epsilon param to avoid divide by zero in loss calculation
+        """
+        if output_mask is None and self._pad_id is None:
+            raise ValueError("Both output_mask and pad_id are None")
+        if output_mask is None and self._pad_id is not None:
+            output_mask = (labels != self._pad_id).to(log_probs.dtype)
+
+        if output_mask.dtype is not log_probs.dtype:
+            output_mask = output_mask.to(log_probs.dtype)
+
+        batch_size, seq_len, vocab_size = log_probs.size()
+        smoothing = vocab_size * self._label_smoothing / (vocab_size - 1)
+        target_log_probs = log_probs.gather(2, labels.unsqueeze(2)).squeeze(2)
+
+        smoothing_log_probs = log_probs.mean(dim=-1)
+        neg_log_likelihood = (1.0 - smoothing) * target_log_probs + smoothing * smoothing_log_probs
+        neg_log_likelihood = neg_log_likelihood[:, -self._predict_last_k :]
+        output_mask = output_mask[:, -self._predict_last_k :]
+        neg_log_likelihood = -torch.sum(neg_log_likelihood * output_mask)
+        neg_log_likelihood = neg_log_likelihood / (output_mask.sum() + self._eps)
+
+        return neg_log_likelihood

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

@@ -0,0 +1,79 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from torch import nn
+
+from nemo.core.classes import Loss, typecheck
+from nemo.core.neural_types import ChannelType, LogitsType, LossType, NeuralType
+
+__all__ = ['SpanningLoss']
+
+
+class SpanningLoss(Loss):
+    """
+    implements start and end loss of a span e.g. for Question Answering.
+    """
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input ports.
+        """
+        return {
+            "logits": NeuralType(('B', 'T', 'D'), LogitsType()),
+            "start_positions": NeuralType(tuple('B'), ChannelType()),
+            "end_positions": NeuralType(tuple('B'), ChannelType()),
+        }
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output ports.
+        """
+        return {
+            "loss": NeuralType(elements_type=LossType()),
+            "start_logits": NeuralType(('B', 'T'), LogitsType()),
+            "end_logits": NeuralType(('B', 'T'), LogitsType()),
+        }
+
+    def __init__(self,):
+        super().__init__()
+
+    @typecheck()
+    def forward(self, logits, start_positions, end_positions):
+        """
+        Args:
+            logits: Output of question answering head, which is a token classfier.
+            start_positions: Ground truth start positions of the answer w.r.t.
+                input sequence. If question is unanswerable, this will be
+                pointing to start token, e.g. [CLS], of the input sequence.
+            end_positions: Ground truth end positions of the answer w.r.t.
+                input sequence. If question is unanswerable, this will be
+                pointing to start token, e.g. [CLS], of the input sequence.
+        """
+        start_logits, end_logits = logits.split(1, dim=-1)
+        start_logits = start_logits.squeeze(-1)
+        end_logits = end_logits.squeeze(-1)
+        # If we are on multi-GPU, split add a dimension
+        if len(start_positions.size()) > 1:
+            start_positions = start_positions.squeeze(-1)
+        if len(end_positions.size()) > 1:
+            end_positions = end_positions.squeeze(-1)
+        ignored_index = start_logits.size(1)
+        start_positions.clamp_(0, ignored_index)
+        end_positions.clamp_(0, ignored_index)
+
+        loss_fct = nn.CrossEntropyLoss(ignore_index=ignored_index)
+        start_loss = loss_fct(start_logits, start_positions)
+        end_loss = loss_fct(end_logits, end_positions)
+        total_loss = (start_loss + end_loss) / 2
+        return total_loss, start_logits, end_logits

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

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

tts/nemo/collections/common/metrics/classification_accuracy.py

@@ -0,0 +1,152 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List
+
+import torch
+from torchmetrics import Metric
+
+__all__ = ['TopKClassificationAccuracy']
+
+
+class TopKClassificationAccuracy(Metric):
+    """
+    This metric computes numerator and denominator for Overall Accuracy between logits and labels.
+    When doing distributed training/evaluation the result of res=TopKClassificationAccuracy(logits, labels) calls
+    will be all-reduced between all workers using SUM operations.
+    Here contains two numbers res=[correctly_predicted, total_samples]. Accuracy=correctly_predicted/total_samples.
+
+    If used with PytorchLightning LightningModule, include correct_count and total_count inside validation_step results.
+    Then aggregate (sum) then at the end of validation epoch to correctly compute validation WER.
+
+    Example:
+        def validation_step(self, batch, batch_idx):
+            ...
+            correct_count, total_count = self._accuracy(logits, labels)
+            return {'val_loss': loss_value, 'val_correct_count': correct_count, 'val_total_count': total_count}
+
+        def validation_epoch_end(self, outputs):
+            ...
+            val_loss_mean = torch.stack([x['val_loss'] for x in outputs]).mean()
+            correct_counts = torch.stack([x['val_correct_counts'] for x in outputs])
+            total_counts = torch.stack([x['val_total_counts'] for x in outputs])
+
+            topk_scores = compute_topk_accuracy(correct_counts, total_counts)
+
+            tensorboard_log = {'val_loss': val_loss_mean}
+            for top_k, score in zip(self._accuracy.top_k, topk_scores):
+                tensorboard_log['val_epoch_top@{}'.format(top_k)] = score
+
+            return {'log': tensorboard_log}
+
+    Args:
+        top_k: Optional list of integers. Defaults to [1].
+
+    Returns:
+        res: a torch.Tensor object with two elements: [correct_count, total_count]. To correctly compute average
+        accuracy, compute acc=correct_count/total_count
+    """
+
+    def __init__(self, top_k=None, dist_sync_on_step=False):
+        super().__init__(dist_sync_on_step=dist_sync_on_step)
+
+        if top_k is None:
+            top_k = [1]
+
+        self.top_k = top_k
+        self.add_state(
+            "correct_counts_k", default=torch.zeros(len(self.top_k)), dist_reduce_fx='sum', persistent=False
+        )
+        self.add_state("total_counts_k", default=torch.zeros(len(self.top_k)), dist_reduce_fx='sum', persistent=False)
+
+    @torch.no_grad()
+    def top_k_predicted_labels(self, logits: torch.Tensor) -> torch.Tensor:
+        max_k = max(self.top_k)
+        _, predictions = logits.topk(max_k, dim=1, largest=True, sorted=True)
+        return predictions
+
+    def update(self, logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+        with torch.no_grad():
+            predictions = self.top_k_predicted_labels(logits)
+            predictions = predictions.t()
+            correct = predictions.eq(labels.view(1, -1)).expand_as(predictions)
+
+            correct_counts_k = []
+            total_counts_k = []
+
+            for k in self.top_k:
+                correct_k = correct[:k].reshape(-1).long().sum()
+                total_k = labels.shape[0]
+
+                correct_counts_k.append(correct_k)
+                total_counts_k.append(total_k)
+
+            self.correct_counts_k = torch.tensor(correct_counts_k, dtype=labels.dtype, device=labels.device)
+            self.total_counts_k = torch.tensor(total_counts_k, dtype=labels.dtype, device=labels.device)
+
+    def compute(self):
+        """
+        Computes the top-k accuracy.
+
+        Returns:
+            A list of length `K`, such that k-th index corresponds to top-k accuracy
+            over all distributed processes.
+        """
+        if not len(self.correct_counts_k) == len(self.top_k) == len(self.total_counts_k):
+            raise ValueError("length of counts must match to topk length")
+
+        if self.top_k == [1]:
+            return [self.correct_counts_k.float() / self.total_counts_k]
+
+        else:
+            top_k_scores = compute_topk_accuracy(self.correct_counts_k, self.total_counts_k)
+
+            return top_k_scores
+
+    @property
+    def top_k(self) -> List[int]:
+        return self._top_k
+
+    @top_k.setter
+    def top_k(self, value: List[int]):
+        if value is None:
+            value = [1]
+
+        if type(value) == int:
+            value = [value]
+
+        if type(value) != list:
+            value = list(value)
+
+        self._top_k = value
+
+
+def compute_topk_accuracy(correct_counts_k, total_counts_k):
+    """
+    Computes the top-k accuracy
+    Args:
+        correct_counts: Tensor of shape [K], K being the top-k parameter.
+        total_counts: Tensor of shape [K], and K being the top-k parameter.
+    Returns:
+        A list of length `K`, such that k-th index corresponds to top-k accuracy
+        over all distributed processes.
+    """
+    top_k_scores = []
+
+    for ki in range(len(correct_counts_k)):
+        correct_count = correct_counts_k[ki].item()
+        total_count = total_counts_k[ki].item()
+        top_k_scores.append(correct_count / float(total_count))
+
+    return top_k_scores

tts/nemo/collections/common/metrics/global_average_loss_metric.py

@@ -0,0 +1,73 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from torchmetrics import Metric
+
+__all__ = ['GlobalAverageLossMetric']
+
+
+class GlobalAverageLossMetric(Metric):
+    """
+    This class is for averaging loss across multiple processes if a distributed backend is used. True average is
+    computed not running average. It does not accumulate gradients so the averaged loss cannot be used for optimization.
+    If ``take_avg_loss`` is ``True``, the :meth:`update` method ``loss`` argument has to be a mean loss. If
+    ``take_avg_loss`` is ``False`` then the :meth:`update` method ``loss`` argument has to be a sum of losses.
+
+    See :doc:`PyTorch Lightning Metrics<pytorch-lightning:metrics>` for the metric usage instruction.
+
+    Args:
+        compute_on_step:
+            Forward only calls ``update()`` and returns ``None`` if this is set to ``False``. default: ``True``
+        dist_sync_on_step:
+            Synchronize metric state across processes at each ``forward()``
+            before returning the value at the step.
+        process_group:
+            Specify the process group on which synchronization is called. default: ``None`` (which selects the entire
+                world)
+        take_avg_loss:
+            If ``True`` values of :meth:`update` method ``loss`` argument has to be a mean loss. If ``False``
+            values of :meth:`update` method ``loss`` argument has to be a sum of losses. default: ``True``
+    """
+
+    def __init__(self, compute_on_step=True, dist_sync_on_step=False, process_group=None, take_avg_loss=True):
+        super().__init__(
+            compute_on_step=compute_on_step, dist_sync_on_step=dist_sync_on_step, process_group=process_group
+        )
+        self.add_state("loss_sum", torch.tensor(0.0, dtype=torch.float64), dist_reduce_fx='sum')
+        self.add_state("num_measurements", torch.tensor(0, dtype=torch.int64), dist_reduce_fx='sum')
+        self.take_avg_loss = take_avg_loss
+
+    def update(self, loss, num_measurements):
+        """
+        Updates :attr:`loss_sum` and :attr:`num_measurements`.
+
+        Args:
+            loss: A float zero dimensional ``torch.Tensor`` which is either sum or average of losses for processed
+                examples. See ``take_avg_loss`` parameter of :meth:`__init__`.
+            num_measurements: An integer zero dimensional ``torch.Tensor`` which contains a number of loss measurements.
+                The sum or mean of the results of these measurements are in the ``loss`` parameter.
+        """
+        if self.take_avg_loss:
+            loss *= num_measurements
+        self.loss_sum += loss
+        self.num_measurements += num_measurements
+
+    def compute(self):
+        """
+        Returns mean loss.
+        """
+        if self.num_measurements.eq(0):
+            return torch.tensor(float('nan'))
+        return self.loss_sum / self.num_measurements

tts/nemo/collections/common/metrics/perplexity.py

@@ -0,0 +1,71 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from torch.distributions.categorical import Categorical
+from torchmetrics import Metric
+
+__all__ = ['Perplexity']
+
+
+class Perplexity(Metric):
+    """
+    This class computes mean perplexity of distributions in the last dimension of inputs. It is a wrapper around
+    :doc:`torch.distributions.Categorical.perplexity<pytorch:distributions>` method. You have to provide either 
+    ``probs`` or ``logits`` to the :meth:`update` method. The class computes perplexities for distributions passed to 
+    :meth:`update` method in ``probs`` or ``logits`` arguments and averages the perplexities. Reducing results between
+    all workers is done via SUM operations.
+    See :doc:`PyTorch Lightning Metrics<pytorch-lightning:metrics>` for the metric usage instructions.
+    Args:
+        compute_on_step:
+            Forward only calls ``update()`` and returns ``None`` if this is set to ``False``. default: ``True``
+        dist_sync_on_step:
+            Synchronize metric state across processes at each ``forward()``
+            before returning the value at the step.
+        process_group:
+            Specify the process group on which synchronization is called. default: ``None`` (which selects the entire
+                world)
+        validate_args:
+            If ``True`` values of :meth:`update` method parameters are checked. ``logits`` has to not contain NaNs and
+            ``probs`` last dim has to be valid probability distribution.
+    """
+
+    def __init__(self, compute_on_step=True, dist_sync_on_step=False, process_group=None, validate_args=True):
+        super().__init__(
+            compute_on_step=compute_on_step, dist_sync_on_step=dist_sync_on_step, process_group=process_group
+        )
+        self.validate_args = validate_args
+        self.add_state('perplexities_sum', torch.tensor(0.0, dtype=torch.float64), dist_reduce_fx='sum')
+        # Total number of distributions seen since last reset
+        self.add_state('num_distributions', torch.tensor(0, dtype=torch.int64), dist_reduce_fx='sum')
+
+    def update(self, probs=None, logits=None):
+        """
+        Updates :attr:`perplexities_sum` and :attr:`num_distributions`.
+        Args:
+            probs: A ``torch.Tensor`` which innermost dimension is valid probability distribution.
+            logits: A ``torch.Tensor`` without NaNs.
+        """
+        d = Categorical(probs, logits, validate_args=self.validate_args)
+        ppl = d.perplexity()
+        self.num_distributions += ppl.numel()
+        self.perplexities_sum += ppl.sum()
+
+    def compute(self):
+        """
+        Returns perplexity across all workers and resets to 0 :attr:`perplexities_sum` and :attr:`num_distributions`.
+        """
+        if self.num_distributions.eq(0):
+            return None
+        return self.perplexities_sum / self.num_distributions

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

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

tts/nemo/collections/common/parts/multi_layer_perceptron.py

@@ -0,0 +1,61 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+
+
+class MultiLayerPerceptron(torch.nn.Module):
+    """
+    A simple MLP that can either be used independently or put on top
+    of pretrained models (such as BERT) and act as a classifier.
+    Args:
+        hidden_size (int): the size of each layer
+        num_classes (int): number of output classes
+        num_layers (int): number of layers
+        activation (str): type of activations for layers in between
+        log_softmax (bool): whether to add a log_softmax layer before output
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_classes: int,
+        num_layers: int = 2,
+        activation: str = 'relu',
+        log_softmax: bool = True,
+    ):
+        super().__init__()
+        self.layers = 0
+        for _ in range(num_layers - 1):
+            layer = torch.nn.Linear(hidden_size, hidden_size)
+            setattr(self, f'layer{self.layers}', layer)
+            setattr(self, f'layer{self.layers + 1}', getattr(torch, activation))
+            self.layers += 2
+        layer = torch.nn.Linear(hidden_size, num_classes)
+        setattr(self, f'layer{self.layers}', layer)
+        self.layers += 1
+        self.log_softmax = log_softmax
+
+    @property
+    def last_linear_layer(self):
+        return getattr(self, f'layer{self.layers - 1}')
+
+    def forward(self, hidden_states):
+        output_states = hidden_states[:]
+        for i in range(self.layers):
+            output_states = getattr(self, f'layer{i}')(output_states)
+
+        if self.log_softmax:
+            output_states = torch.log_softmax(output_states, dim=-1)
+        return output_states

tts/nemo/collections/common/parts/patch_utils.py

@@ -0,0 +1,132 @@
+# 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 packaging import version
+
+from nemo.utils import logging
+
+# Library version globals
+TORCH_VERSION = None
+TORCH_VERSION_MIN = version.Version('1.7')
+
+
+def stft_patch(
+    input: torch.Tensor,
+    n_fft: int,
+    hop_length: Optional[int] = None,
+    win_length: Optional[int] = None,
+    window: Optional[torch.Tensor] = None,
+    center: bool = True,
+    pad_mode: str = 'reflect',
+    normalized: bool = False,
+    onesided: Optional[bool] = None,
+    return_complex: Optional[bool] = None,
+):
+    """
+    Patch over torch.stft for PyTorch <= 1.6.
+    Arguments are same as torch.stft().
+
+    # TODO: Remove once PyTorch 1.7+ is a requirement.
+    """
+    global TORCH_VERSION
+    if TORCH_VERSION is None:
+        TORCH_VERSION = version.parse(torch.__version__)
+
+        logging.warning(
+            "torch.stft() signature has been updated for PyTorch 1.7+\n"
+            "Please update PyTorch to remain compatible with later versions of NeMo."
+        )
+
+    if TORCH_VERSION < TORCH_VERSION_MIN:
+        return torch.stft(
+            input=input,
+            n_fft=n_fft,
+            hop_length=hop_length,
+            win_length=win_length,
+            window=window,
+            center=center,
+            pad_mode=pad_mode,
+            normalized=normalized,
+            onesided=True,
+        )
+    else:
+        return torch.stft(
+            input=input,
+            n_fft=n_fft,
+            hop_length=hop_length,
+            win_length=win_length,
+            window=window,
+            center=center,
+            pad_mode=pad_mode,
+            normalized=normalized,
+            onesided=onesided,
+            return_complex=return_complex,
+        )
+
+
+def istft_patch(
+    input: torch.Tensor,
+    n_fft: int,
+    hop_length: Optional[int] = None,
+    win_length: Optional[int] = None,
+    window: Optional[torch.Tensor] = None,
+    center: bool = True,
+    normalized: bool = False,
+    onesided: Optional[bool] = None,
+    length: int = None,
+    return_complex: Optional[bool] = False,
+):
+    """
+    Patch over torch.stft for PyTorch <= 1.6.
+    Arguments are same as torch.stft().
+
+    # TODO: Remove once PyTorch 1.7+ is a requirement.
+    """
+    global TORCH_VERSION
+    if TORCH_VERSION is None:
+        TORCH_VERSION = version.parse(torch.__version__)
+
+        logging.warning(
+            "torch.stft() signature has been updated for PyTorch 1.7+\n"
+            "Please update PyTorch to remain compatible with later versions of NeMo."
+        )
+
+    if TORCH_VERSION < TORCH_VERSION_MIN:
+        return torch.istft(
+            input=input,
+            n_fft=n_fft,
+            hop_length=hop_length,
+            win_length=win_length,
+            window=window,
+            center=center,
+            length=length,
+            normalized=normalized,
+            onesided=True,
+        )
+    else:
+        return torch.istft(
+            input=input,
+            n_fft=n_fft,
+            hop_length=hop_length,
+            win_length=win_length,
+            window=window,
+            center=center,
+            length=length,
+            normalized=normalized,
+            onesided=onesided,
+            return_complex=return_complex,
+        )

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

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

tts/nemo/collections/common/parts/preprocessing/cleaners.py

@@ -0,0 +1,259 @@
+# 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 re
+
+import inflect
+from unidecode import unidecode
+
+from nemo.utils import logging
+
+NUM_CHECK = re.compile(r'([$]?)(^|\s)(\S*[0-9]\S*)(?=(\s|$)((\S*)(\s|$))?)')
+
+TIME_CHECK = re.compile(r'([0-9]{1,2}):([0-9]{2})(am|pm)?')
+CURRENCY_CHECK = re.compile(r'\$')
+ORD_CHECK = re.compile(r'([0-9]+)(st|nd|rd|th)')
+THREE_CHECK = re.compile(r'([0-9]{3})([.,][0-9]{1,2})?([!.?])?$')
+DECIMAL_CHECK = re.compile(r'([.,][0-9]{1,2})$')
+
+ABBREVIATIONS_COMMON = [
+    (re.compile('\\b%s\\.' % x[0]), x[1])
+    for x in [
+        ("ms", "miss"),
+        ("mrs", "misess"),
+        ("mr", "mister"),
+        ("messrs", "messeurs"),
+        ("dr", "doctor"),
+        ("drs", "doctors"),
+        ("st", "saint"),
+        ("co", "company"),
+        ("jr", "junior"),
+        ("sr", "senior"),
+        ("rev", "reverend"),
+        ("hon", "honorable"),
+        ("sgt", "sergeant"),
+        ("capt", "captain"),
+        ("maj", "major"),
+        ("col", "colonel"),
+        ("lt", "lieutenant"),
+        ("gen", "general"),
+        ("prof", "professor"),
+        ("lb", "pounds"),
+        ("rep", "representative"),
+        ("st", "street"),
+        ("ave", "avenue"),
+        ("etc", "et cetera"),
+        ("jan", "january"),
+        ("feb", "february"),
+        ("mar", "march"),
+        ("apr", "april"),
+        ("jun", "june"),
+        ("jul", "july"),
+        ("aug", "august"),
+        ("sep", "september"),
+        ("oct", "october"),
+        ("nov", "november"),
+        ("dec", "december"),
+    ]
+]
+
+ABBREVIATIONS_EXPANDED = [
+    (re.compile('\\b%s\\.' % x[0]), x[1])
+    for x in [
+        ("ltd", "limited"),
+        ("fig", "figure"),
+        ("figs", "figures"),
+        ("gent", "gentlemen"),
+        ("ft", "fort"),
+        ("esq", "esquire"),
+        ("prep", "preperation"),
+        ("bros", "brothers"),
+        ("ind", "independent"),
+        ("mme", "madame"),
+        ("pro", "professional"),
+        ("vs", "versus"),
+        ("inc", "include"),
+    ]
+]
+
+ABBREVIATIONS_TTS_FASTPITCH = [
+    (re.compile('\\b%s\\.' % x[0]), x[1])
+    for x in [
+        ("ms", "miss"),
+        ("mrs", "misess"),
+        ("mr", "mister"),
+        ("dr", "doctor"),
+        ("drs", "doctors"),
+        ("st", "saint"),
+        ("co", "company"),
+        ("jr", "junior"),
+        ("sr", "senior"),
+        ("rev", "reverend"),
+        ("hon", "honorable"),
+        ("sgt", "sergeant"),
+        ("capt", "captain"),
+        ("maj", "major"),
+        ("col", "colonel"),
+        ("lt", "lieutenant"),
+        ("gen", "general"),
+        ("prof", "professor"),
+        ("lb", "pounds"),
+        ("rep", "representative"),
+        ("st", "street"),
+        ("ave", "avenue"),
+        ("jan", "january"),
+        ("feb", "february"),
+        ("mar", "march"),
+        ("apr", "april"),
+        ("jun", "june"),
+        ("jul", "july"),
+        ("aug", "august"),
+        ("sep", "september"),
+        ("oct", "october"),
+        ("nov", "november"),
+        ("dec", "december"),
+        ("ltd", "limited"),
+        ("fig", "figure"),
+        ("figs", "figures"),
+        ("gent", "gentlemen"),
+        ("ft", "fort"),
+        ("esq", "esquire"),
+        ("prep", "preperation"),
+        ("bros", "brothers"),
+        ("ind", "independent"),
+        ("mme", "madame"),
+        ("pro", "professional"),
+        ("vs", "versus"),
+    ]
+]
+
+
+inflect = inflect.engine()
+
+
+def clean_text(string, table, punctuation_to_replace, abbreviation_version=None):
+    warn_common_chars(string)
+    string = unidecode(string)
+    string = string.lower()
+    string = re.sub(r'\s+', " ", string)
+    string = clean_numbers(string)
+    string = clean_abbreviations(string, version=abbreviation_version)
+    string = clean_punctuations(string, table, punctuation_to_replace)
+    string = re.sub(r'\s+', " ", string).strip()
+    return string
+
+
+def warn_common_chars(string):
+    if re.search(r'[£€]', string):
+        logging.warning("Your transcript contains one of '£' or '€' which we do not currently handle")
+
+
+def clean_numbers(string):
+    cleaner = NumberCleaner()
+    string = NUM_CHECK.sub(cleaner.clean, string)
+    return string
+
+
+def clean_abbreviations(string, version=None):
+    abbbreviations = ABBREVIATIONS_COMMON
+    if version == "fastpitch":
+        abbbreviations = ABBREVIATIONS_TTS_FASTPITCH
+    elif version == "expanded":
+        abbbreviations.extend = ABBREVIATIONS_EXPANDED
+    for regex, replacement in abbbreviations:
+        string = re.sub(regex, replacement, string)
+    return string
+
+
+def clean_punctuations(string, table, punctuation_to_replace):
+    for punc, replacement in punctuation_to_replace.items():
+        string = re.sub('\\{}'.format(punc), " {} ".format(replacement), string)
+    if table:
+        string = string.translate(table)
+    return string
+
+
+class NumberCleaner:
+    def __init__(self):
+        super().__init__()
+        self.reset()
+
+    def reset(self):
+        self.curr_num = []
+        self.currency = None
+
+    def format_final_number(self, whole_num, decimal):
+        if self.currency:
+            return_string = inflect.number_to_words(whole_num)
+            return_string += " dollar" if whole_num == 1 else " dollars"
+            if decimal:
+                return_string += " and " + inflect.number_to_words(decimal)
+                return_string += " cent" if whole_num == decimal else " cents"
+            self.reset()
+            return return_string
+
+        self.reset()
+        if decimal:
+            whole_num += "." + decimal
+            return inflect.number_to_words(whole_num)
+        else:
+            # Check if there are non-numbers
+            def convert_to_word(match):
+                return " " + inflect.number_to_words(match.group(0)) + " "
+
+            return re.sub(r'[0-9,]+', convert_to_word, whole_num)
+
+    def clean(self, match):
+        ws = match.group(2)
+        number = match.group(3)
+        _proceeding_symbol = match.group(7)
+
+        time_match = TIME_CHECK.match(number)
+        if time_match:
+            string = ws + inflect.number_to_words(time_match.group(1)) + "{}{}"
+            mins = int(time_match.group(2))
+            min_string = ""
+            if mins != 0:
+                min_string = " " + inflect.number_to_words(time_match.group(2))
+            ampm_string = ""
+            if time_match.group(3):
+                ampm_string = " " + time_match.group(3)
+            return string.format(min_string, ampm_string)
+
+        ord_match = ORD_CHECK.match(number)
+        if ORD_CHECK.match(number):
+            return ws + inflect.number_to_words(ord_match.group(0))
+
+        if self.currency is None:
+            # Check if it is a currency
+            self.currency = match.group(1) or CURRENCY_CHECK.match(number)
+
+        # Check to see if next symbol is a number
+        # If it is a number and it has 3 digits, then it is probably a
+        # continuation
+        three_match = THREE_CHECK.match(match.group(6))
+        if three_match:
+            self.curr_num.append(number)
+            return " "
+        # Else we can output
+        else:
+            # Check for decimals
+            whole_num = "".join(self.curr_num) + number
+            decimal = None
+            decimal_match = DECIMAL_CHECK.search(whole_num)
+            if decimal_match:
+                decimal = decimal_match.group(1)[1:]
+                whole_num = whole_num[: -len(decimal) - 1]
+            whole_num = re.sub(r'\.', '', whole_num)
+            return ws + self.format_final_number(whole_num, decimal)

tts/nemo/collections/common/parts/preprocessing/collections.py

@@ -0,0 +1,468 @@
+# 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 collections
+import json
+import os
+from typing import Any, Dict, List, Optional, Union
+
+import pandas as pd
+
+from nemo.collections.common.parts.preprocessing import manifest, parsers
+from nemo.utils import logging
+
+
+class _Collection(collections.UserList):
+    """List of parsed and preprocessed data."""
+
+    OUTPUT_TYPE = None  # Single element output type.
+
+
+class Text(_Collection):
+    """Simple list of preprocessed text entries, result in list of tokens."""
+
+    OUTPUT_TYPE = collections.namedtuple('TextEntity', 'tokens')
+
+    def __init__(self, texts: List[str], parser: parsers.CharParser):
+        """Instantiates text manifest and do the preprocessing step.
+
+        Args:
+            texts: List of raw texts strings.
+            parser: Instance of `CharParser` to convert string to tokens.
+        """
+
+        data, output_type = [], self.OUTPUT_TYPE
+        for text in texts:
+            tokens = parser(text)
+
+            if tokens is None:
+                logging.warning("Fail to parse '%s' text line.", text)
+                continue
+
+            data.append(output_type(tokens))
+
+        super().__init__(data)
+
+
+class FromFileText(Text):
+    """Another form of texts manifest with reading from file."""
+
+    def __init__(self, file: str, parser: parsers.CharParser):
+        """Instantiates text manifest and do the preprocessing step.
+
+        Args:
+            file: File path to read from.
+            parser: Instance of `CharParser` to convert string to tokens.
+        """
+
+        texts = self.__parse_texts(file)
+
+        super().__init__(texts, parser)
+
+    @staticmethod
+    def __parse_texts(file: str) -> List[str]:
+        if not os.path.exists(file):
+            raise ValueError('Provided texts file does not exists!')
+
+        _, ext = os.path.splitext(file)
+        if ext == '.csv':
+            texts = pd.read_csv(file)['transcript'].tolist()
+        elif ext == '.json':  # Not really a correct json.
+            texts = list(item['text'] for item in manifest.item_iter(file))
+        else:
+            with open(file, 'r') as f:
+                texts = f.readlines()
+
+        return texts
+
+
+class AudioText(_Collection):
+    """List of audio-transcript text correspondence with preprocessing."""
+
+    OUTPUT_TYPE = collections.namedtuple(
+        typename='AudioTextEntity', field_names='id audio_file duration text_tokens offset text_raw speaker orig_sr',
+    )
+
+    def __init__(
+        self,
+        ids: List[int],
+        audio_files: List[str],
+        durations: List[float],
+        texts: List[str],
+        offsets: List[str],
+        speakers: List[Optional[int]],
+        orig_sampling_rates: List[Optional[int]],
+        parser: parsers.CharParser,
+        min_duration: Optional[float] = None,
+        max_duration: Optional[float] = None,
+        max_number: Optional[int] = None,
+        do_sort_by_duration: bool = False,
+        index_by_file_id: bool = False,
+    ):
+        """Instantiates audio-text manifest with filters and preprocessing.
+
+        Args:
+            ids: List of examples positions.
+            audio_files: List of audio files.
+            durations: List of float durations.
+            texts: List of raw text transcripts.
+            offsets: List of duration offsets or None.
+            speakers: List of optional speakers ids.
+            orig_sampling_rates: List of original sampling rates of audio files.
+            parser: Instance of `CharParser` to convert string to tokens.
+            min_duration: Minimum duration to keep entry with (default: None).
+            max_duration: Maximum duration to keep entry with (default: None).
+            max_number: Maximum number of samples to collect.
+            do_sort_by_duration: True if sort samples list by duration. Not compatible with index_by_file_id.
+            index_by_file_id: If True, saves a mapping from filename base (ID) to index in data.
+        """
+
+        output_type = self.OUTPUT_TYPE
+        data, duration_filtered, num_filtered, total_duration = [], 0.0, 0, 0.0
+        if index_by_file_id:
+            self.mapping = {}
+
+        for id_, audio_file, duration, offset, text, speaker, orig_sr in zip(
+            ids, audio_files, durations, offsets, texts, speakers, orig_sampling_rates
+        ):
+            # Duration filters.
+            if min_duration is not None and duration < min_duration:
+                duration_filtered += duration
+                num_filtered += 1
+                continue
+
+            if max_duration is not None and duration > max_duration:
+                duration_filtered += duration
+                num_filtered += 1
+                continue
+
+            text_tokens = parser(text)
+            if text_tokens is None:
+                duration_filtered += duration
+                num_filtered += 1
+                continue
+
+            total_duration += duration
+
+            data.append(output_type(id_, audio_file, duration, text_tokens, offset, text, speaker, orig_sr))
+            if index_by_file_id:
+                file_id, _ = os.path.splitext(os.path.basename(audio_file))
+                self.mapping[file_id] = len(data) - 1
+
+            # Max number of entities filter.
+            if len(data) == max_number:
+                break
+
+        if do_sort_by_duration:
+            if index_by_file_id:
+                logging.warning("Tried to sort dataset by duration, but cannot since index_by_file_id is set.")
+            else:
+                data.sort(key=lambda entity: entity.duration)
+
+        logging.info("Dataset loaded with %d files totalling %.2f hours", len(data), total_duration / 3600)
+        logging.info("%d files were filtered totalling %.2f hours", num_filtered, duration_filtered / 3600)
+
+        super().__init__(data)
+
+
+class ASRAudioText(AudioText):
+    """`AudioText` collector from asr structured json files."""
+
+    def __init__(self, manifests_files: Union[str, List[str]], *args, **kwargs):
+        """Parse lists of audio files, durations and transcripts texts.
+
+        Args:
+            manifests_files: Either single string file or list of such -
+                manifests to yield items from.
+            *args: Args to pass to `AudioText` constructor.
+            **kwargs: Kwargs to pass to `AudioText` constructor.
+        """
+
+        ids, audio_files, durations, texts, offsets, speakers, orig_srs = [], [], [], [], [], [], []
+        for item in manifest.item_iter(manifests_files):
+            ids.append(item['id'])
+            audio_files.append(item['audio_file'])
+            durations.append(item['duration'])
+            texts.append(item['text'])
+            offsets.append(item['offset'])
+            speakers.append(item['speaker'])
+            orig_srs.append(item['orig_sr'])
+
+        super().__init__(ids, audio_files, durations, texts, offsets, speakers, orig_srs, *args, **kwargs)
+
+
+class SpeechLabel(_Collection):
+    """List of audio-label correspondence with preprocessing."""
+
+    OUTPUT_TYPE = collections.namedtuple(typename='SpeechLabelEntity', field_names='audio_file duration label offset',)
+
+    def __init__(
+        self,
+        audio_files: List[str],
+        durations: List[float],
+        labels: List[Union[int, str]],
+        offsets: List[Optional[float]],
+        min_duration: Optional[float] = None,
+        max_duration: Optional[float] = None,
+        max_number: Optional[int] = None,
+        do_sort_by_duration: bool = False,
+        index_by_file_id: bool = False,
+    ):
+        """Instantiates audio-label manifest with filters and preprocessing.
+
+        Args:
+            audio_files: List of audio files.
+            durations: List of float durations.
+            labels: List of labels.
+            offsets: List of offsets or None.
+            min_duration: Minimum duration to keep entry with (default: None).
+            max_duration: Maximum duration to keep entry with (default: None).
+            max_number: Maximum number of samples to collect.
+            do_sort_by_duration: True if sort samples list by duration.
+            index_by_file_id: If True, saves a mapping from filename base (ID) to index in data.
+        """
+
+        if index_by_file_id:
+            self.mapping = {}
+        output_type = self.OUTPUT_TYPE
+        data, duration_filtered = [], 0.0
+        for audio_file, duration, command, offset in zip(audio_files, durations, labels, offsets):
+            # Duration filters.
+            if min_duration is not None and duration < min_duration:
+                duration_filtered += duration
+                continue
+
+            if max_duration is not None and duration > max_duration:
+                duration_filtered += duration
+                continue
+
+            data.append(output_type(audio_file, duration, command, offset))
+
+            if index_by_file_id:
+                file_id, _ = os.path.splitext(os.path.basename(audio_file))
+                self.mapping[file_id] = len(data) - 1
+
+            # Max number of entities filter.
+            if len(data) == max_number:
+                break
+
+        if do_sort_by_duration:
+            if index_by_file_id:
+                logging.warning("Tried to sort dataset by duration, but cannot since index_by_file_id is set.")
+            else:
+                data.sort(key=lambda entity: entity.duration)
+
+        logging.info(
+            "Filtered duration for loading collection is %f.", duration_filtered,
+        )
+        self.uniq_labels = sorted(set(map(lambda x: x.label, data)))
+        logging.info("# {} files loaded accounting to # {} labels".format(len(data), len(self.uniq_labels)))
+
+        super().__init__(data)
+
+
+class ASRSpeechLabel(SpeechLabel):
+    """`SpeechLabel` collector from structured json files."""
+
+    def __init__(self, manifests_files: Union[str, List[str]], is_regression_task=False, *args, **kwargs):
+        """Parse lists of audio files, durations and transcripts texts.
+
+        Args:
+            manifests_files: Either single string file or list of such -
+                manifests to yield items from.
+            is_regression_task: It's a regression task
+            *args: Args to pass to `SpeechLabel` constructor.
+            **kwargs: Kwargs to pass to `SpeechLabel` constructor.
+        """
+        audio_files, durations, labels, offsets = [], [], [], []
+
+        for item in manifest.item_iter(manifests_files, parse_func=self.__parse_item):
+            audio_files.append(item['audio_file'])
+            durations.append(item['duration'])
+            if not is_regression_task:
+                labels.append(item['label'])
+            else:
+                labels.append(float(item['label']))
+
+            offsets.append(item['offset'])
+
+        super().__init__(audio_files, durations, labels, offsets, *args, **kwargs)
+
+    def __parse_item(self, line: str, manifest_file: str) -> Dict[str, Any]:
+        item = json.loads(line)
+
+        # Audio file
+        if 'audio_filename' in item:
+            item['audio_file'] = item.pop('audio_filename')
+        elif 'audio_filepath' in item:
+            item['audio_file'] = item.pop('audio_filepath')
+        else:
+            raise ValueError(
+                f"Manifest file has invalid json line " f"structure: {line} without proper audio file key."
+            )
+        item['audio_file'] = os.path.expanduser(item['audio_file'])
+
+        # Duration.
+        if 'duration' not in item:
+            raise ValueError(f"Manifest file has invalid json line " f"structure: {line} without proper duration key.")
+
+        # Label.
+        if 'command' in item:
+            item['label'] = item.pop('command')
+        elif 'target' in item:
+            item['label'] = item.pop('target')
+        elif 'label' in item:
+            pass
+        else:
+            raise ValueError(f"Manifest file has invalid json line " f"structure: {line} without proper label key.")
+
+        item = dict(
+            audio_file=item['audio_file'],
+            duration=item['duration'],
+            label=item['label'],
+            offset=item.get('offset', None),
+        )
+
+        return item
+
+
+class FeatureSequenceLabel(_Collection):
+    """List of feature sequence of label correspondence with preprocessing."""
+
+    OUTPUT_TYPE = collections.namedtuple(typename='FeatureSequenceLabelEntity', field_names='feature_file seq_label',)
+
+    def __init__(
+        self,
+        feature_files: List[str],
+        seq_labels: List[str],
+        max_number: Optional[int] = None,
+        index_by_file_id: bool = False,
+    ):
+        """Instantiates feature-SequenceLabel manifest with filters and preprocessing.
+
+        Args:
+            feature_files: List of feature files.
+            seq_labels: List of sequences of abels.
+            max_number: Maximum number of samples to collect.
+            index_by_file_id: If True, saves a mapping from filename base (ID) to index in data.
+        """
+
+        output_type = self.OUTPUT_TYPE
+        data, num_filtered = (
+            [],
+            0.0,
+        )
+        self.uniq_labels = set()
+
+        if index_by_file_id:
+            self.mapping = {}
+
+        for feature_file, seq_label in zip(feature_files, seq_labels):
+
+            label_tokens, uniq_labels_in_seq = self.relative_speaker_parser(seq_label)
+
+            data.append(output_type(feature_file, label_tokens))
+            self.uniq_labels |= uniq_labels_in_seq
+
+            if label_tokens is None:
+                num_filtered += 1
+                continue
+
+            if index_by_file_id:
+                file_id, _ = os.path.splitext(os.path.basename(feature_file))
+                self.mapping[feature_file] = len(data) - 1
+
+            # Max number of entities filter.
+            if len(data) == max_number:
+                break
+
+        logging.info("# {} files loaded including # {} unique labels".format(len(data), len(self.uniq_labels)))
+        super().__init__(data)
+
+    def relative_speaker_parser(self, seq_label):
+        """ Convert sequence of speaker labels to relative labels.
+        Convert sequence of absolute speaker to sequence of relative speaker [E A C A E E C] -> [0 1 2 1 0 0 2]
+        In this seq of label , if label do not appear before, assign new relative labels len(pos); else reuse previous assigned relative labels.
+        Args:
+            seq_label (str): A string of a sequence of labels.
+
+        Return:
+            relative_seq_label (List) : A list of relative sequence of labels
+            unique_labels_in_seq (Set): A set of unique labels in the sequence
+        """
+        seq = seq_label.split()
+        conversion_dict = dict()
+        relative_seq_label = []
+
+        for seg in seq:
+            if seg in conversion_dict:
+                converted = conversion_dict[seg]
+            else:
+                converted = len(conversion_dict)
+                conversion_dict[seg] = converted
+
+            relative_seq_label.append(converted)
+
+        unique_labels_in_seq = set(conversion_dict.keys())
+        return relative_seq_label, unique_labels_in_seq
+
+
+class ASRFeatureSequenceLabel(FeatureSequenceLabel):
+    """`FeatureSequenceLabel` collector from asr structured json files."""
+
+    def __init__(
+        self, manifests_files: Union[str, List[str]], max_number: Optional[int] = None, index_by_file_id: bool = False,
+    ):
+
+        """Parse lists of feature files and sequences of labels.
+
+        Args:
+            manifests_files: Either single string file or list of such -
+                manifests to yield items from.
+            max_number:  Maximum number of samples to collect; pass to `FeatureSequenceLabel` constructor.
+            index_by_file_id: If True, saves a mapping from filename base (ID) to index in data; pass to `FeatureSequenceLabel` constructor.
+        """
+
+        feature_files, seq_labels = [], []
+        for item in manifest.item_iter(manifests_files, parse_func=self._parse_item):
+            feature_files.append(item['feature_file'])
+            seq_labels.append(item['seq_label'])
+
+        super().__init__(feature_files, seq_labels, max_number, index_by_file_id)
+
+    def _parse_item(self, line: str, manifest_file: str) -> Dict[str, Any]:
+        item = json.loads(line)
+
+        # Feature file
+        if 'feature_filename' in item:
+            item['feature_file'] = item.pop('feature_filename')
+        elif 'feature_filepath' in item:
+            item['feature_file'] = item.pop('feature_filepath')
+        else:
+            raise ValueError(
+                f"Manifest file has invalid json line " f"structure: {line} without proper feature file key."
+            )
+        item['feature_file'] = os.path.expanduser(item['feature_file'])
+
+        # Seq of Label.
+        if 'seq_label' in item:
+            item['seq_label'] = item.pop('seq_label')
+        else:
+            raise ValueError(
+                f"Manifest file has invalid json line " f"structure: {line} without proper seq_label key."
+            )
+
+        item = dict(feature_file=item['feature_file'], seq_label=item['seq_label'],)
+
+        return item

tts/nemo/collections/common/parts/preprocessing/manifest.py

@@ -0,0 +1,120 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+from os.path import expanduser
+from typing import Any, Callable, Dict, Iterator, List, Optional, Union
+
+
+class ManifestBase:
+    def __init__(self, *args, **kwargs):
+        raise ValueError(
+            "This class is deprecated, look at https://github.com/NVIDIA/NeMo/pull/284 for correct behaviour."
+        )
+
+
+class ManifestEN:
+    def __init__(self, *args, **kwargs):
+        raise ValueError(
+            "This class is deprecated, look at https://github.com/NVIDIA/NeMo/pull/284 for correct behaviour."
+        )
+
+
+def item_iter(
+    manifests_files: Union[str, List[str]], parse_func: Callable[[str, Optional[str]], Dict[str, Any]] = None
+) -> Iterator[Dict[str, Any]]:
+    """Iterate through json lines of provided manifests.
+
+    NeMo ASR pipelines often assume certain manifest files structure. In
+    particular, each manifest file should consist of line-per-sample files with
+    each line being correct json dict. Each such json dict should have a field
+    for audio file string, a field for duration float and a field for text
+    string. Offset also could be additional field and is set to None by
+    default.
+
+    Args:
+        manifests_files: Either single string file or list of such -
+            manifests to yield items from.
+
+        parse_func: A callable function which accepts as input a single line
+            of a manifest and optionally the manifest file itself,
+            and parses it, returning a dictionary mapping from str -> Any.
+
+    Yields:
+        Parsed key to value item dicts.
+
+    Raises:
+        ValueError: If met invalid json line structure.
+    """
+
+    if isinstance(manifests_files, str):
+        manifests_files = [manifests_files]
+
+    if parse_func is None:
+        parse_func = __parse_item
+
+    k = -1
+    for manifest_file in manifests_files:
+        with open(expanduser(manifest_file), 'r') as f:
+            for line in f:
+                k += 1
+                item = parse_func(line, manifest_file)
+                item['id'] = k
+
+                yield item
+
+
+def __parse_item(line: str, manifest_file: str) -> Dict[str, Any]:
+    item = json.loads(line)
+
+    # Audio file
+    if 'audio_filename' in item:
+        item['audio_file'] = item.pop('audio_filename')
+    elif 'audio_filepath' in item:
+        item['audio_file'] = item.pop('audio_filepath')
+    else:
+        raise ValueError(
+            f"Manifest file {manifest_file} has invalid json line structure: {line} without proper audio file key."
+        )
+    item['audio_file'] = expanduser(item['audio_file'])
+
+    # Duration.
+    if 'duration' not in item:
+        raise ValueError(
+            f"Manifest file {manifest_file} has invalid json line structure: {line} without proper duration key."
+        )
+
+    # Text.
+    if 'text' in item:
+        pass
+    elif 'text_filepath' in item:
+        with open(item.pop('text_filepath'), 'r') as f:
+            item['text'] = f.read().replace('\n', '')
+    elif 'normalized_text' in item:
+        item['text'] = item['normalized_text']
+    else:
+        raise ValueError(
+            f"Manifest file {manifest_file} has invalid json line structure: {line} without proper text key."
+        )
+
+    item = dict(
+        audio_file=item['audio_file'],
+        duration=item['duration'],
+        text=item['text'],
+        offset=item.get('offset', None),
+        speaker=item.get('speaker', None),
+        orig_sr=item.get('orig_sample_rate', None),
+    )
+
+    return item

tts/nemo/collections/common/parts/preprocessing/parsers.py

@@ -0,0 +1,185 @@
+# 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 string
+from typing import List, Optional
+
+import frozendict
+
+from nemo.collections.common.parts.preprocessing import cleaners
+
+
+class CharParser:
+    """Functor for parsing raw strings into list of int tokens.
+
+    Examples:
+        >>> parser = CharParser(['a', 'b', 'c'])
+        >>> parser('abc')
+        [0, 1, 2]
+    """
+
+    def __init__(
+        self,
+        labels: List[str],
+        *,
+        unk_id: int = -1,
+        blank_id: int = -1,
+        do_normalize: bool = True,
+        do_lowercase: bool = True,
+    ):
+        """Creates simple mapping char parser.
+
+        Args:
+            labels: List of labels to allocate indexes for. Essentially,
+                this is a id to str mapping.
+            unk_id: Index to choose for OOV words (default: -1).
+            blank_id: Index to filter out from final list of tokens
+                (default: -1).
+            do_normalize: True if apply normalization step before tokenizing
+                (default: True).
+            do_lowercase: True if apply lowercasing at normalizing step
+                (default: True).
+        """
+
+        self._labels = labels
+        self._unk_id = unk_id
+        self._blank_id = blank_id
+        self._do_normalize = do_normalize
+        self._do_lowercase = do_lowercase
+
+        self._labels_map = {label: index for index, label in enumerate(labels)}
+        self._special_labels = set([label for label in labels if len(label) > 1])
+
+    def __call__(self, text: str) -> Optional[List[int]]:
+        if self._do_normalize:
+            text = self._normalize(text)
+            if text is None:
+                return None
+
+        text_tokens = self._tokenize(text)
+
+        return text_tokens
+
+    def _normalize(self, text: str) -> Optional[str]:
+        text = text.strip()
+
+        if self._do_lowercase:
+            text = text.lower()
+
+        return text
+
+    def _tokenize(self, text: str) -> List[int]:
+        tokens = []
+        # Split by word for find special labels.
+        for word_id, word in enumerate(text.split(' ')):
+            if word_id != 0:  # Not first word - so we insert space before.
+                tokens.append(self._labels_map.get(' ', self._unk_id))
+
+            if word in self._special_labels:
+                tokens.append(self._labels_map[word])
+                continue
+
+            for char in word:
+                tokens.append(self._labels_map.get(char, self._unk_id))
+
+        # If unk_id == blank_id, OOV tokens are removed.
+        tokens = [token for token in tokens if token != self._blank_id]
+
+        return tokens
+
+
+class ENCharParser(CharParser):
+    """Incorporates english-specific parsing logic."""
+
+    PUNCTUATION_TO_REPLACE = frozendict.frozendict({'+': 'plus', '&': 'and', '%': 'percent'})
+
+    def __init__(self, abbreviation_version=None, make_table=True, *args, **kwargs):
+        """Creates english-specific mapping char parser.
+
+        This class overrides normalizing implementation.
+
+        Args:
+            *args: Positional args to pass to `CharParser` constructor.
+            **kwargs: Key-value args to pass to `CharParser` constructor.
+        """
+
+        super().__init__(*args, **kwargs)
+
+        self._table = None
+        if make_table:
+            self._table = self.__make_trans_table()
+        self.abbreviation_version = abbreviation_version
+
+    def __make_trans_table(self):
+        punctuation = string.punctuation
+
+        for char in self.PUNCTUATION_TO_REPLACE:
+            punctuation = punctuation.replace(char, '')
+
+        for label in self._labels:
+            punctuation = punctuation.replace(label, '')
+
+        table = str.maketrans(punctuation, ' ' * len(punctuation))
+
+        return table
+
+    def _normalize(self, text: str) -> Optional[str]:
+        # noinspection PyBroadException
+        try:
+            text = cleaners.clean_text(
+                string=text,
+                table=self._table,
+                punctuation_to_replace=self.PUNCTUATION_TO_REPLACE,
+                abbreviation_version=self.abbreviation_version,
+            )
+        except Exception:
+            return None
+
+        return text
+
+
+NAME_TO_PARSER = frozendict.frozendict({'base': CharParser, 'en': ENCharParser})
+
+
+def make_parser(labels: Optional[List[str]] = None, name: str = 'base', **kwargs,) -> CharParser:
+    """Creates parser from labels, set of arguments and concise parser name.
+
+    Args:
+        labels: List of labels to allocate indexes for. If set to
+            None then labels would be ascii table list. Essentially, this is a
+            id to str mapping (default: None).
+        name: Concise name of parser to create (default: 'base').
+            (default: -1).
+        **kwargs: Other set of kwargs to pass to parser constructor.
+
+    Returns:
+        Instance of `CharParser`.
+
+    Raises:
+        ValueError: For invalid parser name.
+
+    Examples:
+        >>> type(make_parser(['a', 'b', 'c'], 'en'))
+        ENCharParser
+    """
+
+    if name not in NAME_TO_PARSER:
+        raise ValueError('Invalid parser name.')
+
+    if labels is None:
+        labels = list(string.printable)
+
+    parser_type = NAME_TO_PARSER[name]
+    parser = parser_type(labels=labels, **kwargs)
+
+    return parser

tts/nemo/collections/common/parts/rnn.py

@@ -0,0 +1,553 @@
+# Copyright (c) 2019, Myrtle Software Limited. All rights reserved.
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from nemo.utils import logging
+
+
+def rnn(
+    input_size: int,
+    hidden_size: int,
+    num_layers: int,
+    norm: Optional[str] = None,
+    forget_gate_bias: Optional[float] = 1.0,
+    dropout: Optional[float] = 0.0,
+    norm_first_rnn: Optional[bool] = None,
+    t_max: Optional[int] = None,
+    weights_init_scale: float = 1.0,
+    hidden_hidden_bias_scale: float = 0.0,
+) -> torch.nn.Module:
+    """
+    Utility function to provide unified interface to common LSTM RNN modules.
+
+    Args:
+        input_size: Input dimension.
+
+        hidden_size: Hidden dimension of the RNN.
+
+        num_layers: Number of RNN layers.
+
+        norm: Optional string representing type of normalization to apply to the RNN.
+            Supported values are None, batch and layer.
+
+        forget_gate_bias: float, set by default to 1.0, which constructs a forget gate
+                initialized to 1.0.
+                Reference:
+                [An Empirical Exploration of Recurrent Network Architectures](http://proceedings.mlr.press/v37/jozefowicz15.pdf)
+
+        dropout: Optional dropout to apply to end of multi-layered RNN.
+
+        norm_first_rnn: Whether to normalize the first RNN layer.
+
+        t_max: int value, set to None by default. If an int is specified, performs Chrono Initialization
+            of the LSTM network, based on the maximum number of timesteps `t_max` expected during the course
+            of training.
+            Reference:
+            [Can recurrent neural networks warp time?](https://openreview.net/forum?id=SJcKhk-Ab)
+
+        weights_init_scale: Float scale of the weights after initialization. Setting to lower than one
+            sometimes helps reduce variance between runs.
+
+        hidden_hidden_bias_scale: Float scale for the hidden-to-hidden bias scale. Set to 0.0 for
+            the default behaviour.
+
+    Returns:
+        A RNN module
+    """
+    if norm not in [None, "batch", "layer"]:
+        raise ValueError(f"unknown norm={norm}")
+
+    if norm is None:
+        return LSTMDropout(
+            input_size=input_size,
+            hidden_size=hidden_size,
+            num_layers=num_layers,
+            dropout=dropout,
+            forget_gate_bias=forget_gate_bias,
+            t_max=t_max,
+            weights_init_scale=weights_init_scale,
+            hidden_hidden_bias_scale=hidden_hidden_bias_scale,
+        )
+
+    if norm == "batch":
+        return BNRNNSum(
+            input_size=input_size,
+            hidden_size=hidden_size,
+            rnn_layers=num_layers,
+            batch_norm=True,
+            dropout=dropout,
+            forget_gate_bias=forget_gate_bias,
+            t_max=t_max,
+            norm_first_rnn=norm_first_rnn,
+            weights_init_scale=weights_init_scale,
+            hidden_hidden_bias_scale=hidden_hidden_bias_scale,
+        )
+
+    if norm == "layer":
+        return torch.jit.script(
+            ln_lstm(  # torch.jit.script(
+                input_size=input_size,
+                hidden_size=hidden_size,
+                num_layers=num_layers,
+                dropout=dropout,
+                forget_gate_bias=forget_gate_bias,
+                t_max=t_max,
+                weights_init_scale=weights_init_scale,
+                hidden_hidden_bias_scale=hidden_hidden_bias_scale,
+            )
+        )
+
+
+class OverLastDim(torch.nn.Module):
+    """Collapses a tensor to 2D, applies a module, and (re-)expands the tensor.
+    An n-dimensional tensor of shape (s_1, s_2, ..., s_n) is first collapsed to
+    a tensor with shape (s_1*s_2*...*s_n-1, s_n). The module is called with
+    this as input producing (s_1*s_2*...*s_n-1, s_n') --- note that the final
+    dimension can change. This is expanded to (s_1, s_2, ..., s_n-1, s_n') and
+    returned.
+    Args:
+        module (torch.nn.Module): Module to apply. Must accept a 2D tensor as
+            input and produce a 2D tensor as output, optionally changing the
+            size of the last dimension.
+    """
+
+    def __init__(self, module: torch.nn.Module):
+        super().__init__()
+        self.module = module
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        *dims, _ = x.size()
+
+        reduced_dims = 1
+        for dim in dims:
+            reduced_dims *= dim
+
+        x = x.view(reduced_dims, -1)
+        x = self.module(x)
+        x = x.view(*dims, -1)
+        return x
+
+
+class LSTMDropout(torch.nn.Module):
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        num_layers: int,
+        dropout: Optional[float],
+        forget_gate_bias: Optional[float],
+        t_max: Optional[int] = None,
+        weights_init_scale: float = 1.0,
+        hidden_hidden_bias_scale: float = 0.0,
+    ):
+        """Returns an LSTM with forget gate bias init to `forget_gate_bias`.
+        Args:
+            input_size: See `torch.nn.LSTM`.
+            hidden_size: See `torch.nn.LSTM`.
+            num_layers: See `torch.nn.LSTM`.
+            dropout: See `torch.nn.LSTM`.
+
+            forget_gate_bias: float, set by default to 1.0, which constructs a forget gate
+                initialized to 1.0.
+                Reference:
+                [An Empirical Exploration of Recurrent Network Architectures](http://proceedings.mlr.press/v37/jozefowicz15.pdf)
+
+            t_max: int value, set to None by default. If an int is specified, performs Chrono Initialization
+                of the LSTM network, based on the maximum number of timesteps `t_max` expected during the course
+                of training.
+                Reference:
+                [Can recurrent neural networks warp time?](https://openreview.net/forum?id=SJcKhk-Ab)
+
+            weights_init_scale: Float scale of the weights after initialization. Setting to lower than one
+                sometimes helps reduce variance between runs.
+
+            hidden_hidden_bias_scale: Float scale for the hidden-to-hidden bias scale. Set to 0.0 for
+                the default behaviour.
+
+        Returns:
+            A `torch.nn.LSTM`.
+        """
+        super(LSTMDropout, self).__init__()
+
+        self.lstm = torch.nn.LSTM(
+            input_size=input_size, hidden_size=hidden_size, num_layers=num_layers, dropout=dropout,
+        )
+
+        if t_max is not None:
+            # apply chrono init
+            for name, v in self.lstm.named_parameters():
+                if 'bias' in name:
+                    p = getattr(self.lstm, name)
+                    n = p.nelement()
+                    hidden_size = n // 4
+                    p.data.fill_(0)
+                    p.data[hidden_size : 2 * hidden_size] = torch.log(
+                        torch.nn.init.uniform_(p.data[0:hidden_size], 1, t_max - 1)
+                    )
+                    # forget gate biases = log(uniform(1, Tmax-1))
+                    p.data[0:hidden_size] = -p.data[hidden_size : 2 * hidden_size]
+                    # input gate biases = -(forget gate biases)
+
+        elif forget_gate_bias is not None:
+            for name, v in self.lstm.named_parameters():
+                if "bias_ih" in name:
+                    bias = getattr(self.lstm, name)
+                    bias.data[hidden_size : 2 * hidden_size].fill_(forget_gate_bias)
+                if "bias_hh" in name:
+                    bias = getattr(self.lstm, name)
+                    bias.data[hidden_size : 2 * hidden_size] *= float(hidden_hidden_bias_scale)
+
+        self.dropout = torch.nn.Dropout(dropout) if dropout else None
+
+        for name, v in self.named_parameters():
+            if 'weight' in name or 'bias' in name:
+                v.data *= float(weights_init_scale)
+
+    def forward(
+        self, x: torch.Tensor, h: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
+    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        x, h = self.lstm(x, h)
+
+        if self.dropout:
+            x = self.dropout(x)
+
+        return x, h
+
+
+class RNNLayer(torch.nn.Module):
+    """A single RNNLayer with optional batch norm."""
+
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        rnn_type: torch.nn.Module = torch.nn.LSTM,
+        batch_norm: bool = True,
+        forget_gate_bias: Optional[float] = 1.0,
+        t_max: Optional[int] = None,
+        weights_init_scale: float = 1.0,
+        hidden_hidden_bias_scale: float = 0.0,
+    ):
+        super().__init__()
+
+        if batch_norm:
+            self.bn = OverLastDim(torch.nn.BatchNorm1d(input_size))
+
+        if isinstance(rnn_type, torch.nn.LSTM) and not batch_norm:
+            # batch_norm will apply bias, no need to add a second to LSTM
+            self.rnn = LSTMDropout(
+                input_size=input_size,
+                hidden_size=hidden_size,
+                num_layers=1,
+                dropout=0.0,
+                forget_gate_bias=forget_gate_bias,
+                t_max=t_max,
+                weights_init_scale=weights_init_scale,
+                hidden_hidden_bias_scale=hidden_hidden_bias_scale,
+            )
+        else:
+            self.rnn = rnn_type(input_size=input_size, hidden_size=hidden_size, bias=not batch_norm)
+
+    def forward(
+        self, x: torch.Tensor, hx: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None
+    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        if hasattr(self, 'bn'):
+            x = x.contiguous()
+            x = self.bn(x)
+        x, h = self.rnn(x, hx=hx)
+        return x, h
+
+    def _flatten_parameters(self):
+        self.rnn.flatten_parameters()
+
+
+class BNRNNSum(torch.nn.Module):
+    """RNN wrapper with optional batch norm.
+    Instantiates an RNN. If it is an LSTM it initialises the forget gate
+    bias =`lstm_gate_bias`. Optionally applies a batch normalisation layer to
+    the input with the statistics computed over all time steps.  If dropout > 0
+    then it is applied to all layer outputs except the last.
+    """
+
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        rnn_type: torch.nn.Module = torch.nn.LSTM,
+        rnn_layers: int = 1,
+        batch_norm: bool = True,
+        dropout: Optional[float] = 0.0,
+        forget_gate_bias: Optional[float] = 1.0,
+        norm_first_rnn: bool = False,
+        t_max: Optional[int] = None,
+        weights_init_scale: float = 1.0,
+        hidden_hidden_bias_scale: float = 0.0,
+    ):
+        super().__init__()
+        self.rnn_layers = rnn_layers
+
+        self.layers = torch.nn.ModuleList()
+        for i in range(rnn_layers):
+            final_layer = (rnn_layers - 1) == i
+
+            self.layers.append(
+                RNNLayer(
+                    input_size,
+                    hidden_size,
+                    rnn_type=rnn_type,
+                    batch_norm=batch_norm and (norm_first_rnn or i > 0),
+                    forget_gate_bias=forget_gate_bias,
+                    t_max=t_max,
+                    weights_init_scale=weights_init_scale,
+                    hidden_hidden_bias_scale=hidden_hidden_bias_scale,
+                )
+            )
+
+            if dropout is not None and dropout > 0.0 and not final_layer:
+                self.layers.append(torch.nn.Dropout(dropout))
+
+            input_size = hidden_size
+
+    def forward(
+        self, x: torch.Tensor, hx: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None
+    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        hx = self._parse_hidden_state(hx)
+
+        hs = []
+        cs = []
+        rnn_idx = 0
+        for layer in self.layers:
+            if isinstance(layer, torch.nn.Dropout):
+                x = layer(x)
+            else:
+                x, h_out = layer(x, hx=hx[rnn_idx])
+                hs.append(h_out[0])
+                cs.append(h_out[1])
+                rnn_idx += 1
+                del h_out
+
+        h_0 = torch.stack(hs, dim=0)
+        c_0 = torch.stack(cs, dim=0)
+        return x, (h_0, c_0)
+
+    def _parse_hidden_state(
+        self, hx: Optional[Tuple[torch.Tensor, torch.Tensor]]
+    ) -> Union[List[None], List[Tuple[torch.Tensor, torch.Tensor]]]:
+        """
+        Dealing w. hidden state:
+        Typically in pytorch: (h_0, c_0)
+            h_0 = ``[num_layers * num_directions, batch, hidden_size]``
+            c_0 = ``[num_layers * num_directions, batch, hidden_size]``
+        """
+        if hx is None:
+            return [None] * self.rnn_layers
+        else:
+            h_0, c_0 = hx
+
+            if h_0.shape[0] != self.rnn_layers:
+                raise ValueError(
+                    'Provided initial state value `h_0` must be of shape : '
+                    '[num_layers * num_directions, batch, hidden_size]'
+                )
+
+            return [(h_0[i], c_0[i]) for i in range(h_0.shape[0])]
+
+    def _flatten_parameters(self):
+        for layer in self.layers:
+            if isinstance(layer, (torch.nn.LSTM, torch.nn.GRU, torch.nn.RNN)):
+                layer._flatten_parameters()
+
+
+class StackTime(torch.nn.Module):
+    """
+    Stacks time within the feature dim, so as to behave as a downsampling operation.
+    """
+
+    def __init__(self, factor: int):
+        super().__init__()
+        self.factor = int(factor)
+
+    def forward(self, x: List[Tuple[torch.Tensor]]) -> (torch.Tensor, torch.Tensor):
+        # T, B, U
+        x, x_lens = x
+        seq = [x]
+        for i in range(1, self.factor):
+            tmp = torch.zeros_like(x)
+            tmp[:-i, :, :] = x[i:, :, :]
+            seq.append(tmp)
+        x_lens = torch.ceil(x_lens.float() / self.factor).int()
+        return torch.cat(seq, dim=2)[:: self.factor, :, :], x_lens
+
+
+def ln_lstm(
+    input_size: int,
+    hidden_size: int,
+    num_layers: int,
+    dropout: Optional[float],
+    forget_gate_bias: Optional[float],
+    t_max: Optional[int],
+    weights_init_scale: Optional[float] = None,  # ignored
+    hidden_hidden_bias_scale: Optional[float] = None,  # ignored
+) -> torch.nn.Module:
+    """Returns a ScriptModule that mimics a PyTorch native LSTM."""
+    # The following are not implemented.
+    if dropout is not None and dropout != 0.0:
+        raise ValueError('`dropout` not supported with LayerNormLSTM')
+
+    if t_max is not None:
+        logging.warning("LayerNormLSTM does not support chrono init via `t_max`")
+
+    if weights_init_scale is not None:
+        logging.warning("`weights_init_scale` is ignored for LayerNormLSTM")
+
+    if hidden_hidden_bias_scale is not None:
+        logging.warning("`hidden_hidden_bias_scale` is ignored for LayerNormLSTM")
+
+    return StackedLSTM(
+        num_layers,
+        LSTMLayer,
+        first_layer_args=[LayerNormLSTMCell, input_size, hidden_size, forget_gate_bias],
+        other_layer_args=[LayerNormLSTMCell, hidden_size, hidden_size, forget_gate_bias],
+    )
+
+
+class LSTMLayer(torch.nn.Module):
+    def __init__(self, cell, *cell_args):
+        super(LSTMLayer, self).__init__()
+        self.cell = cell(*cell_args)
+
+    def forward(
+        self, input: torch.Tensor, state: Tuple[torch.Tensor, torch.Tensor]
+    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        inputs = input.unbind(0)
+        outputs = []
+        for i in range(len(inputs)):
+            out, state = self.cell(inputs[i], state)
+            outputs += [out]
+        return torch.stack(outputs), state
+
+
+class LayerNormLSTMCell(torch.nn.Module):
+    def __init__(self, input_size, hidden_size, forget_gate_bias):
+        super().__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.weight_ih = torch.nn.Parameter(torch.randn(4 * hidden_size, input_size))
+        self.weight_hh = torch.nn.Parameter(torch.randn(4 * hidden_size, hidden_size))
+
+        # LayerNorm provide learnable biases
+        self.layernorm_i = torch.nn.LayerNorm(4 * hidden_size)
+        self.layernorm_h = torch.nn.LayerNorm(4 * hidden_size)
+        self.layernorm_c = torch.nn.LayerNorm(hidden_size)
+
+        self.reset_parameters()
+
+        self.layernorm_i.bias.data[hidden_size : 2 * hidden_size].fill_(0.0)
+        self.layernorm_h.bias.data[hidden_size : 2 * hidden_size].fill_(forget_gate_bias)
+
+    def reset_parameters(self):
+        stdv = 1.0 / math.sqrt(self.hidden_size)
+        for weight in self.parameters():
+            torch.nn.init.uniform_(weight, -stdv, stdv)
+
+    def forward(
+        self, input: torch.Tensor, state: Tuple[torch.Tensor, torch.Tensor]
+    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        hx, cx = state
+        igates = self.layernorm_i(torch.mm(input, self.weight_ih.t()))
+        hgates = self.layernorm_h(torch.mm(hx, self.weight_hh.t()))
+        gates = igates + hgates
+        ingate, forgetgate, cellgate, outgate = gates.chunk(4, 1)
+
+        ingate = torch.sigmoid(ingate)
+        forgetgate = torch.sigmoid(forgetgate)
+        cellgate = torch.tanh(cellgate)
+        outgate = torch.sigmoid(outgate)
+
+        cy = self.layernorm_c((forgetgate * cx) + (ingate * cellgate))
+        hy = outgate * torch.tanh(cy)
+
+        return hy, (hy, cy)
+
+
+def init_stacked_lstm(
+    num_layers: int, layer: torch.nn.Module, first_layer_args: List, other_layer_args: List
+) -> torch.nn.ModuleList:
+    layers = [layer(*first_layer_args)] + [layer(*other_layer_args) for _ in range(num_layers - 1)]
+    return torch.nn.ModuleList(layers)
+
+
+class StackedLSTM(torch.nn.Module):
+    def __init__(self, num_layers: int, layer: torch.nn.Module, first_layer_args: List, other_layer_args: List):
+        super(StackedLSTM, self).__init__()
+        self.layers: torch.nn.ModuleList = init_stacked_lstm(num_layers, layer, first_layer_args, other_layer_args)
+
+    def forward(
+        self, input: torch.Tensor, states: Optional[List[Tuple[torch.Tensor, torch.Tensor]]]
+    ) -> Tuple[torch.Tensor, List[Tuple[torch.Tensor, torch.Tensor]]]:
+        if states is None:
+            temp_states: List[Tuple[torch.Tensor, torch.Tensor]] = []
+            batch = input.size(1)
+            for layer in self.layers:
+                temp_states.append(
+                    (
+                        torch.zeros(batch, layer.cell.hidden_size, dtype=input.dtype, device=input.device),
+                        torch.zeros(batch, layer.cell.hidden_size, dtype=input.dtype, device=input.device),
+                    )
+                )
+
+            states = temp_states
+
+        output_states: List[Tuple[torch.Tensor, torch.Tensor]] = []
+        output = input
+        for i, rnn_layer in enumerate(self.layers):
+            state = states[i]
+            output, out_state = rnn_layer(output, state)
+            output_states.append(out_state)
+            i += 1
+        return output, output_states
+
+
+def label_collate(labels, device=None):
+    """Collates the label inputs for the rnn-t prediction network.
+    If `labels` is already in torch.Tensor form this is a no-op.
+
+    Args:
+        labels: A torch.Tensor List of label indexes or a torch.Tensor.
+        device: Optional torch device to place the label on.
+
+    Returns:
+        A padded torch.Tensor of shape (batch, max_seq_len).
+    """
+
+    if isinstance(labels, torch.Tensor):
+        return labels.type(torch.int64)
+    if not isinstance(labels, (list, tuple)):
+        raise ValueError(f"`labels` should be a list or tensor not {type(labels)}")
+
+    batch_size = len(labels)
+    max_len = max(len(label) for label in labels)
+
+    cat_labels = np.full((batch_size, max_len), fill_value=0.0, dtype=np.int32)
+    for e, l in enumerate(labels):
+        cat_labels[e, : len(l)] = l
+    labels = torch.tensor(cat_labels, dtype=torch.int64, device=device)
+
+    return labels

tts/nemo/collections/common/parts/transformer_utils.py

@@ -0,0 +1,79 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import torch.nn as nn
+
+__all__ = ['NEG_INF', 'form_attention_mask', 'transformer_weights_init', 'mask_padded_tokens']
+
+NEG_INF = -10000.0
+
+
+def form_attention_mask(input_mask, diagonal=None):
+    """
+    Build attention mask with optional masking of future tokens we forbid
+    to attend to (e.g. as it is in Transformer decoder).
+
+    Args:
+        input_mask: binary mask of size B x L with 1s corresponding to valid
+            tokens and 0s corresponding to padding tokens
+        diagonal: diagonal where triangular future mask starts
+            None -- do not mask anything
+            0 -- regular translation or language modeling future masking
+            1 -- query stream masking as in XLNet architecture
+    Returns:
+        attention_mask: mask of size B x 1 x L x L with 0s corresponding to
+            tokens we plan to attend to and -10000 otherwise
+    """
+
+    if input_mask is None:
+        return None
+    attn_shape = (1, input_mask.shape[1], input_mask.shape[1])
+    attn_mask = input_mask.to(dtype=bool).unsqueeze(1)
+    if diagonal is not None:
+        future_mask = torch.tril(torch.ones(attn_shape, dtype=torch.bool, device=input_mask.device), diagonal)
+        attn_mask = attn_mask & future_mask
+    attention_mask = (1 - attn_mask.to(torch.float)) * NEG_INF
+    return attention_mask.unsqueeze(1)
+
+
+def transformer_weights_init(module, std_init_range=0.02, xavier=True):
+    """
+    Initialize different weights in Transformer model.
+
+    Args:
+        module: torch.nn.Module to be initialized
+        std_init_range: standard deviation of normal initializer
+        xavier: if True, xavier initializer will be used in Linear layers
+            as was proposed in AIAYN paper, otherwise normal initializer
+            will be used (like in BERT paper)
+    """
+
+    if isinstance(module, nn.Linear):
+        if xavier:
+            nn.init.xavier_uniform_(module.weight)
+        else:
+            nn.init.normal_(module.weight, mean=0.0, std=std_init_range)
+        if module.bias is not None:
+            nn.init.constant_(module.bias, 0.0)
+    elif isinstance(module, nn.Embedding):
+        nn.init.normal_(module.weight, mean=0.0, std=std_init_range)
+    elif isinstance(module, nn.LayerNorm):
+        nn.init.constant_(module.weight, 1.0)
+        nn.init.constant_(module.bias, 0.0)
+
+
+def mask_padded_tokens(tokens, pad_id):
+    mask = tokens != pad_id
+    return mask

tts/nemo/collections/common/parts/utils.py

@@ -0,0 +1,57 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+import os
+from typing import List
+
+__all__ = ['if_exist', '_compute_softmax']
+
+
+def if_exist(outfold: str, files: List[str]):
+    """
+    Returns true if all given files exist in the given folder
+    Args:
+        outfold: folder path
+        files: list of file names relative to outfold
+    """
+    if not os.path.exists(outfold):
+        return False
+    for file in files:
+        if not os.path.exists(f'{outfold}/{file}'):
+            return False
+    return True
+
+
+def _compute_softmax(scores):
+    """Compute softmax probability over raw logits."""
+    if not scores:
+        return []
+
+    max_score = None
+    for score in scores:
+        if max_score is None or score > max_score:
+            max_score = score
+
+    exp_scores = []
+    total_sum = 0.0
+    for score in scores:
+        x = math.exp(score - max_score)
+        exp_scores.append(x)
+        total_sum += x
+
+    probs = []
+    for score in exp_scores:
+        probs.append(score / total_sum)
+    return probs

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

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

tts/nemo/collections/common/tokenizers/char_tokenizer.py

@@ -0,0 +1,142 @@
+# 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
+
+from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
+
+__all__ = ['CharTokenizer']
+
+
+class CharTokenizer(TokenizerSpec):
+    "Tokenizes each character"
+
+    def __init__(
+        self,
+        vocab_file: str,
+        mask_token: Optional[str] = None,
+        bos_token: Optional[str] = None,
+        eos_token: Optional[str] = None,
+        pad_token: Optional[str] = None,
+        sep_token: Optional[str] = None,
+        cls_token: Optional[str] = None,
+        unk_token: Optional[str] = None,
+    ):
+        """
+        Args:
+            vocab_file: path to file with vocabulary which consists
+                of characters separated by \n
+            mask_token: mask token 
+            bos_token: the beginning of sequence token
+            eos_token: the end of sequence token. Usually equal to sep_token
+            pad_token: token to use for padding
+            sep_token: token used for separating sequences
+            cls_token: class token. Usually equal to bos_token
+            unk_token: token to use for unknown tokens
+        """
+
+        vocab_list = open(vocab_file, "r").readlines()
+        self.vocab = {vocab_list[i].strip(): i for i in range(len(vocab_list))}
+
+        special_tokens_dict = {}
+        if unk_token:
+            special_tokens_dict["unk_token"] = unk_token
+        if sep_token:
+            special_tokens_dict["sep_token"] = sep_token
+        if mask_token:
+            special_tokens_dict["mask_token"] = mask_token
+        if bos_token:
+            special_tokens_dict["bos_token"] = bos_token
+        if eos_token:
+            special_tokens_dict["eos_token"] = eos_token
+        if pad_token:
+            special_tokens_dict["pad_token"] = pad_token
+        if cls_token:
+            special_tokens_dict["cls_token"] = cls_token
+
+        self.add_special_tokens(special_tokens_dict)
+        self.inv_vocab = {v: k for k, v in self.vocab.items()}
+        self.vocab_size = len(self.vocab)
+        self.special_tokens = self.tokens_to_ids(special_tokens_dict.values())
+
+    def add_special_tokens(self, special_tokens_dict: dict) -> int:
+        """
+        Adds a dictionary of special tokens (eos, pad, cls...).
+        If special tokens are NOT in the vocabulary, they are added
+        to it (indexed starting from the last index of the current vocabulary).
+        Args:
+            special_tokens_dict: dict of special tokens
+        """
+        for token in special_tokens_dict:
+            token_str = special_tokens_dict[token]
+            if token_str not in self.vocab:
+                self.vocab[token_str] = len(self.vocab)
+            setattr(self, token, token_str)
+
+    def text_to_tokens(self, text):
+        token_candidates = [char for char in text]
+        tokens = []
+        for token in token_candidates:
+            if token in self.vocab:
+                tokens.append(token)
+            else:
+                tokens.append(self.unk_token)
+        return tokens
+
+    def tokens_to_text(self, tokens):
+        return self.ids_to_text(self.tokens_to_ids(tokens))
+
+    def text_to_ids(self, text):
+        return [self.vocab[token] for token in self.text_to_tokens(text)]
+
+    def ids_to_text(self, ids):
+        ids_ = [id_ for id_ in ids if id_ not in self.special_tokens]
+        return "".join(self.ids_to_tokens(ids_))
+
+    def tokens_to_ids(self, tokens):
+        return [self.vocab[token] for token in tokens]
+
+    def token_to_id(self, token):
+        return self.vocab[token]
+
+    def ids_to_tokens(self, ids):
+        return [self.inv_vocab[id] for id in ids]
+
+    @property
+    def pad_id(self):
+        return self.vocab[self.pad_token]
+
+    @property
+    def bos_id(self):
+        return self.vocab[self.bos_token]
+
+    @property
+    def eos_id(self):
+        return self.vocab[self.eos_token]
+
+    @property
+    def unk_id(self):
+        return self.vocab[self.unk_token]
+
+    @property
+    def mask_id(self):
+        return self.vocab[self.mask_token]
+
+    @property
+    def sep_id(self):
+        return self.vocab[self.sep_token]
+
+    @property
+    def cls_id(self):
+        return self.vocab[self.cls_token]

tts/nemo/collections/common/tokenizers/chinese_tokenizers.py

@@ -0,0 +1,63 @@
+# 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.
+
+# The MIT License (MIT)
+# Copyright (c) 2016 The-Orizon
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+
+# The detokenize function is based on : https://github.com/The-Orizon/nlputils/blob/master/detokenizer.py
+
+import re
+from typing import List
+
+import jieba
+import opencc
+from pangu import spacing
+
+
+class ChineseProcessor:
+    """
+    Tokenizer, Detokenizer and Normalizer utilities for Chinese.
+    """
+
+    def __init__(self):
+        self.normalizer = opencc.OpenCC('t2s.json')
+
+    def normalize(self, text: str) -> str:
+        return self.normalizer.convert(text)
+
+    def detokenize(self, text: List[str]) -> str:
+        RE_WS_IN_FW = re.compile(
+            r'([\u2018\u2019\u201c\u201d\u2e80-\u312f\u3200-\u32ff\u3400-\u4dbf\u4e00-\u9fff\uf900-\ufaff\uff00-\uffef])\s+(?=[\u2018\u2019\u201c\u201d\u2e80-\u312f\u3200-\u32ff\u3400-\u4dbf\u4e00-\u9fff\uf900-\ufaff\uff00-\uffef])'
+        )
+
+        detokenize = lambda s: spacing(RE_WS_IN_FW.sub(r'\1', s)).strip()
+        return detokenize(' '.join(text))
+
+    def tokenize(self, text: str) -> str:
+        text = jieba.cut(text)
+        return ' '.join(text)

tts/nemo/collections/common/tokenizers/en_ja_tokenizers.py

@@ -0,0 +1,57 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import List
+
+from sacremoses import MosesDetokenizer, MosesPunctNormalizer, MosesTokenizer
+
+
+class EnJaProcessor:
+    """
+    Tokenizer, Detokenizer and Normalizer utilities for Japanese & English
+    Args:
+        lang_id: One of ['en', 'ja'].
+    """
+
+    def __init__(self, lang_id: str):
+        self.lang_id = lang_id
+        self.moses_tokenizer = MosesTokenizer(lang=lang_id)
+        self.moses_detokenizer = MosesDetokenizer(lang=lang_id)
+        self.normalizer = MosesPunctNormalizer(
+            lang=lang_id, pre_replace_unicode_punct=True, post_remove_control_chars=True
+        )
+
+    def detokenize(self, tokens: List[str]) -> str:
+        """
+        Detokenizes a list of tokens
+        Args:
+            tokens: list of strings as tokens
+        Returns:
+            detokenized Japanese or English string
+        """
+        return self.moses_detokenizer.detokenize(tokens)
+
+    def tokenize(self, text) -> str:
+        """
+        Tokenizes text using Moses. Returns a string of tokens.
+        """
+        tokens = self.moses_tokenizer.tokenize(text)
+        return ' '.join(tokens)
+
+    def normalize(self, text) -> str:
+        # Normalization doesn't handle Japanese periods correctly;
+        # '。'becomes '.'.
+        if self.lang_id == 'en':
+            return self.normalizer.normalize(text)
+        else:
+            return text

tts/nemo/collections/common/tokenizers/fairseq_tokenizer.py

@@ -0,0 +1,126 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Copyright (c) 2017-present, Facebook, Inc.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the LICENSE file in
+# the root directory of this source tree. An additional grant of patent rights
+# can be found in the PATENTS file in the same directory.
+
+""" Code from
+https://github.com/NVIDIA/DeepLearningExamples/blob/
+master/PyTorch/Translation/Transformer/fairseq/tokenizer.py
+"""
+
+import re
+import sys
+import unicodedata
+from collections import defaultdict
+
+__all__ = ['get_unicode_categories', 'tokenize_en']
+
+
+def get_unicode_categories():
+    cats = defaultdict(list)
+    for c in map(chr, range(sys.maxunicode + 1)):
+        cats[unicodedata.category(c)].append(c)
+    return cats
+
+
+NUMERICS = ''.join(get_unicode_categories()['No'])
+
+
+def tokenize_en(line):
+    line = line.strip()
+    line = ' ' + line + ' '
+    # remove ASCII junk
+    line = re.sub(r'\s+', ' ', line)
+    line = re.sub(r'[\x00-\x1F]', '', line)
+    # fix whitespaces
+    line = re.sub(r'\ +', ' ', line)
+    line = re.sub('^ ', '', line)
+    line = re.sub(' $', '', line)
+    # separate other special characters
+    line = re.sub(r'([^\s\.\'\`\,\-\w]|[_' + NUMERICS + '])', r' \g<1> ', line)
+    line = re.sub(r'(\w)\-(?=\w)', r'\g<1> @-@ ', line)
+
+    # multidots stay together
+    line = re.sub(r'\.([\.]+)', r' DOTMULTI\g<1>', line)
+    while re.search(r'DOTMULTI\.', line):
+        line = re.sub(r'DOTMULTI\.([^\.])', r'DOTDOTMULTI \g<1>', line)
+        line = re.sub(r'DOTMULTI\.', r'DOTDOTMULTI', line)
+
+    # separate out "," except if within numbers (5,300)
+    line = re.sub(r'([\D])[,]', r'\g<1> , ', line)
+    line = re.sub(r'[,]([\D])', r' , \g<1>', line)
+
+    # separate "," after a number if it's the end of sentence
+    line = re.sub(r'(\d)[,]$', r'\g<1> ,', line)
+
+    # split contractions right
+    line = re.sub(r'([\W\d])[\']([\W\d])', r'\g<1> \' \g<2>', line)
+    line = re.sub(r'(\W)[\']([\w\D])', r'\g<1> \' \g<2>', line)
+    line = re.sub(r'([\w\D])[\']([\W\d])', r'\g<1> \' \g<2>', line)
+    line = re.sub(r'([\w\D])[\']([\w\D])', r'\g<1> \'\g<2>', line)
+    # special case for "1990's"
+    line = re.sub(r'([\W\d])[\']([s])', r'\g<1> \'\g<2>', line)
+
+    # apply nonbreaking prefixes
+    words = line.split()
+    line = ''
+    for i in range(len(words)):
+        word = words[i]
+        match = re.search(r'^(\S+)\.$', word)
+        if match:
+            pre = match.group(1)
+            if i == len(words) - 1:
+                """split last words independently as they are unlikely
+                to be non-breaking prefixes"""
+                word = pre + ' .'
+            else:
+                word = pre + ' .'
+
+        word += ' '
+        line += word
+
+    # clean up extraneous spaces
+    line = re.sub(' +', ' ', line)
+    line = re.sub('^ ', '', line)
+    line = re.sub(' $', '', line)
+
+    # .' at end of sentence is missed
+    line = re.sub(r'\.\' ?$', ' . \' ', line)
+
+    # restore multi-dots
+    while re.search('DOTDOTMULTI', line):
+        line = re.sub('DOTDOTMULTI', 'DOTMULTI.', line)
+
+    line = re.sub('DOTMULTI', '.', line)
+
+    # escape special characters
+    line = re.sub(r'\&', r'&amp;', line)
+    line = re.sub(r'\|', r'&#124;', line)
+    line = re.sub(r'\<', r'&lt;', line)
+    line = re.sub(r'\>', r'&gt;', line)
+    line = re.sub(r'\'', r'&apos;', line)
+    line = re.sub(r'\"', r'&quot;', line)
+    line = re.sub(r'\[', r'&#91;', line)
+    line = re.sub(r'\]', r'&#93;', line)
+
+    # ensure final line breaks
+    # if line[-1] is not '\n':
+    #     line += '\n'
+
+    return line

tts/nemo/collections/common/tokenizers/huggingface/__init__.py

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

tts/nemo/collections/common/tokenizers/huggingface/auto_tokenizer.py

@@ -0,0 +1,260 @@
+# 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 re
+from typing import Optional
+
+from transformers import AutoTokenizer as AUTOTOKENIZER
+
+from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
+from nemo.utils import logging
+
+__all__ = [
+    'AutoTokenizer',
+]
+
+
+def handle_quotes(text):
+    text_ = ""
+    quote = 0
+    i = 0
+    while i < len(text):
+        if text[i] == "\"":
+            if quote % 2:
+                text_ = text_[:-1] + "\""
+            else:
+                text_ += "\""
+                i += 1
+            quote += 1
+        else:
+            text_ += text[i]
+        i += 1
+    return text_
+
+
+def remove_spaces(text):
+    text = text.replace("( ", "(")
+    text = text.replace(" )", ")")
+    text = text.replace("[ ", "[")
+    text = text.replace(" ]", "]")
+    text = text.replace(" / ", "/")
+    text = text.replace("„ ", "„")
+    text = text.replace(" - ", "-")
+    text = text.replace(" ' ", "'")
+    text = re.sub(r'([0-9])( )([\.,])', '\\1\\3', text)
+    text = re.sub(r'([\.,])( )([0-9])', '\\1\\3', text)
+    text = re.sub(r'([0-9])(:)( )([0-9])', '\\1\\2\\4', text)
+    text = text.replace(" %", "%")
+    text = text.replace("$ ", "$")
+    text = re.sub(r'([^0-9])(,)([0-9])', '\\1\\2 \\3', text)
+    return text
+
+
+class AutoTokenizer(TokenizerSpec):
+    '''
+        Wrapper of HuggingFace AutoTokenizer https://huggingface.co/transformers/model_doc/auto.html#autotokenizer.
+    '''
+
+    def __init__(
+        self,
+        pretrained_model_name: str,
+        vocab_file: Optional[str] = None,
+        mask_token: Optional[str] = None,
+        bos_token: Optional[str] = None,
+        eos_token: Optional[str] = None,
+        pad_token: Optional[str] = None,
+        sep_token: Optional[str] = None,
+        cls_token: Optional[str] = None,
+        unk_token: Optional[str] = None,
+        use_fast: Optional[bool] = False,
+    ):
+
+        """
+        Args:
+            pretrained_model_name: corresponds to HuggingFace-AutoTokenizer's 'pretrained_model_name_or_path' input argument. 
+                For more details please refer to https://huggingface.co/transformers/_modules/transformers/tokenization_auto.html#AutoTokenizer.from_pretrained. 
+                The list of all supported models can be found here: ALL_PRETRAINED_CONFIG_ARCHIVE_MAP
+            vocab_file: path to file with vocabulary which consists
+                of characters separated by '\n'.
+            mask_token: mask token 
+            bos_token: the beginning of sequence token
+            eos_token: the end of sequence token. Usually equal to sep_token
+            pad_token: token to use for padding
+            sep_token: token used for separating sequences
+            cls_token: class token. Usually equal to bos_token
+            unk_token: token to use for unknown tokens
+            use_fast: whether to use fast HuggingFace tokenizer
+        """
+        try:
+            if vocab_file is not None:
+                message = 'Using "slow" HuggingFace tokenizer'
+                if use_fast:
+                    message += f'{vocab_file} is ignored in "fast" tokenizers, using a "slow" version'
+                self.tokenizer = AUTOTOKENIZER.from_pretrained(
+                    pretrained_model_name_or_path=pretrained_model_name, vocab_file=vocab_file, use_fast=False
+                )
+            else:
+                self.tokenizer = AUTOTOKENIZER.from_pretrained(
+                    pretrained_model_name_or_path=pretrained_model_name, use_fast=use_fast
+                )
+        except Exception as e:
+            raise ValueError(f'{pretrained_model_name} is not supported by HuggingFace. {e}')
+
+        special_tokens_dict = {}
+
+        # # setting special tokens, by default the default model's special tokens will be preserved
+        # # unless passes new values to the special tokens
+        if unk_token is not None:
+            special_tokens_dict["unk_token"] = unk_token
+        if mask_token is not None:
+            special_tokens_dict["mask_token"] = mask_token
+        if pad_token is not None:
+            special_tokens_dict["pad_token"] = pad_token
+
+        if sep_token is not None:
+            special_tokens_dict["sep_token"] = sep_token
+        elif self.tokenizer.sep_token is None and self.tokenizer.eos_token:
+            special_tokens_dict["sep_token"] = self.tokenizer.eos_token
+        if bos_token is not None:
+            special_tokens_dict["bos_token"] = bos_token
+        elif self.tokenizer.bos_token is None and self.tokenizer.cls_token:
+            special_tokens_dict["bos_token"] = self.tokenizer.cls_token
+        if eos_token is not None:
+            special_tokens_dict["eos_token"] = eos_token
+        elif self.tokenizer.eos_token is None and self.tokenizer.sep_token:
+            special_tokens_dict["eos_token"] = self.tokenizer.sep_token
+        if cls_token is not None:
+            special_tokens_dict["cls_token"] = cls_token
+        elif self.tokenizer.cls_token is None and self.tokenizer.bos_token:
+            special_tokens_dict["cls_token"] = self.tokenizer.bos_token
+
+        new_tokens_in_vocab = []
+        for token in [mask_token, bos_token, eos_token, pad_token, sep_token, cls_token, unk_token]:
+            if token is not None and token not in self.tokenizer.get_vocab():
+                new_tokens_in_vocab.append(token)
+
+        if len(new_tokens_in_vocab) > 0:
+            """
+            Special tokens that were not previously included in the tokenizer's vocabulary file will be added to 
+            the vocabulary and, as a result, the model should be resized, for example:
+            
+            # define your model
+            pretrained_model_name = 'roberta-base'
+            model = nemo_nlp.modules.get_lm_model(pretrained_model_name=pretrained_model_name)
+            
+            # define pretrained tokenizer
+            tokenizer_default = nemo_nlp.modules.get_tokenizer(tokenizer_name=pretrained_model_name)
+            
+            special_tokens = {'bos_token': '<BOS>',
+                              'cls_token': '<CSL>',
+                              'additional_special_tokens': ['<MY_NER_TOKEN>', '<ANOTHER_TOKEN>']}
+            tokenizer_default.add_special_tokens(special_tokens_dict=special_tokens)
+            
+            # resize your model so that the embeddings for newly added tokens are updated during training/finetuning
+            model.resize_token_embeddings(tokenizer_default.vocab_size)
+            
+            See NLP_Tokenizers.ipynb for more details.
+            """
+            logging.warning(
+                f'{new_tokens_in_vocab} \n will be added to the vocabulary.\n'
+                f'Please resize your model accordingly, '
+                f'see NLP_Tokenizers.ipynb for more details.'
+            )
+        self.add_special_tokens(special_tokens_dict)
+
+    @property
+    def vocab_size(self):
+        return len(self.tokenizer)
+
+    def add_special_tokens(self, special_tokens_dict: dict) -> int:
+        """
+        Adds a dictionary of special tokens (eos, pad, cls...). If special tokens are NOT in the vocabulary, they are added
+        to it (indexed starting from the last index of the current vocabulary).
+        Args:
+            special_tokens_dict: dict of string. Keys should be in the list of predefined special attributes:
+                [``bos_token``, ``eos_token``, ``unk_token``, ``sep_token``, ``pad_token``, ``cls_token``, ``mask_token``,
+                ``additional_special_tokens``].
+            Tokens are only added if they are not already in the vocabulary.
+        Returns:
+            Number of tokens added to the vocabulary.
+        """
+        num_tokens_added = self.tokenizer.add_special_tokens(special_tokens_dict)
+
+        if num_tokens_added > 0:
+            logging.info(f'{num_tokens_added} special tokens added, resize your model accordingly.')
+        for k in self.tokenizer.SPECIAL_TOKENS_ATTRIBUTES:
+            setattr(self, k, getattr(self.tokenizer, k, None))
+        return num_tokens_added
+
+    def text_to_tokens(self, text):
+        tokens = self.tokenizer.tokenize(text)
+        return tokens
+
+    def tokens_to_text(self, tokens):
+        text = self.tokenizer.convert_tokens_to_string(tokens)
+        return remove_spaces(handle_quotes(text.strip()))
+
+    def token_to_id(self, token):
+        return self.tokens_to_ids([token])[0]
+
+    def tokens_to_ids(self, tokens):
+        ids = self.tokenizer.convert_tokens_to_ids(tokens)
+        return ids
+
+    def ids_to_tokens(self, ids):
+        tokens = self.tokenizer.convert_ids_to_tokens(ids)
+        return tokens
+
+    def text_to_ids(self, text):
+        tokens = self.text_to_tokens(text)
+        ids = self.tokens_to_ids(tokens)
+        return ids
+
+    def ids_to_text(self, ids):
+        tokens = self.ids_to_tokens(ids)
+        tokens_clean = [t for t in tokens if t not in self.tokenizer.all_special_tokens]
+        text = self.tokens_to_text(tokens_clean)
+        return text
+
+    @property
+    def pad_id(self):
+        return self.tokens_to_ids([getattr(self, 'pad_token')])[0]
+
+    @property
+    def bos_id(self):
+        return self.tokens_to_ids([getattr(self, 'bos_token')])[0]
+
+    @property
+    def eos_id(self):
+        return self.tokens_to_ids([getattr(self, 'eos_token')])[0]
+
+    @property
+    def sep_id(self):
+        return self.tokens_to_ids([getattr(self, 'sep_token')])[0]
+
+    @property
+    def cls_id(self):
+        return self.tokens_to_ids([getattr(self, 'cls_token')])[0]
+
+    @property
+    def unk_id(self):
+        return self.tokens_to_ids([getattr(self, 'unk_token')])[0]
+
+    @property
+    def mask_id(self):
+        return self.tokens_to_ids([getattr(self, 'mask_token')])[0]
+
+    @property
+    def name(self):
+        return type(self.tokenizer).__name__

tts/nemo/collections/common/tokenizers/moses_tokenizers.py

@@ -0,0 +1,47 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List
+
+from sacremoses import MosesDetokenizer, MosesPunctNormalizer, MosesTokenizer
+
+
+class MosesProcessor:
+    """
+    Tokenizer, Detokenizer and Normalizer utilities in Moses
+    """
+
+    def __init__(self, lang_id: str):
+        self.moses_tokenizer = MosesTokenizer(lang=lang_id)
+        self.moses_detokenizer = MosesDetokenizer(lang=lang_id)
+        self.normalizer = MosesPunctNormalizer(lang=lang_id)
+
+    def detokenize(self, tokens: List[str]) -> str:
+        """
+        Detokenizes a list of tokens
+        Args:
+            tokens: list of strings as tokens
+        Returns:
+            detokenized string
+        """
+        return self.moses_detokenizer.detokenize(tokens)
+
+    def tokenize(self, text: str):
+        """
+        Tokenizes text using Moses -> Sentencepiece.
+        """
+        return self.moses_tokenizer.tokenize(text, escape=False, return_str=True)
+
+    def normalize(self, text: str):
+        return self.normalizer.normalize(text)

tts/nemo/collections/common/tokenizers/sentencepiece_tokenizer.py

@@ -0,0 +1,353 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import re
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+import sentencepiece
+
+from nemo.collections.common.parts.utils import if_exist
+from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
+from nemo.utils import logging
+
+__all__ = ['SentencePieceTokenizer', 'create_spt_model']
+
+
+class SentencePieceTokenizer(TokenizerSpec):
+    '''
+    Sentencepiecetokenizer https://github.com/google/sentencepiece.
+        Args:
+        model_path: path to sentence piece tokenizer model. To create the model use create_spt_model()
+        special_tokens: either list of special tokens or dictionary of token name to token value
+        legacy: when set to True, the previous behavior of the SentecePiece wrapper will be restored, 
+            including the possibility to add special tokens inside wrapper.
+    '''
+
+    def __init__(
+        self, model_path: str, special_tokens: Optional[Union[Dict[str, str], List[str]]] = None, legacy: bool = False
+    ):
+        if not model_path or not os.path.exists(model_path):
+            raise ValueError(f"model_path: {model_path} is invalid")
+        self.tokenizer = sentencepiece.SentencePieceProcessor()
+        self.tokenizer.Load(model_path)
+        self.original_vocab_size = self.tokenizer.get_piece_size()
+        self.vocab_size = self.tokenizer.get_piece_size()
+        self.legacy = legacy
+        self.special_token_to_id = {}
+        self.id_to_special_token = {}
+        if special_tokens:
+            if not self.legacy:
+                raise ValueError(
+                    "Special tokens must be None when legacy is set to False. Provide special tokens at train time."
+                )
+            self.add_special_tokens(special_tokens)
+
+    def text_to_tokens(self, text):
+        if self.legacy:
+            tokens = []
+            idx = 0
+            last_idx = 0
+
+            while 1:
+                indices = {}
+
+                for token in self.special_token_to_id:
+                    try:
+                        indices[token] = text[idx:].index(token)
+                    except ValueError:
+                        continue
+
+                if len(indices) == 0:
+                    break
+
+                next_token = min(indices, key=indices.get)
+                next_idx = idx + indices[next_token]
+
+                tokens.extend(self.tokenizer.encode_as_pieces(text[idx:next_idx]))
+                tokens.append(next_token)
+                idx = next_idx + len(next_token)
+
+            tokens.extend(self.tokenizer.encode_as_pieces(text[idx:]))
+            return tokens
+
+        return self.tokenizer.encode_as_pieces(text)
+
+    def text_to_ids(self, text):
+        if self.legacy:
+            ids = []
+            idx = 0
+            last_idx = 0
+
+            while 1:
+                indices = {}
+
+                for token in self.special_token_to_id:
+                    try:
+                        indices[token] = text[idx:].index(token)
+                    except ValueError:
+                        continue
+
+                if len(indices) == 0:
+                    break
+
+                next_token = min(indices, key=indices.get)
+                next_idx = idx + indices[next_token]
+
+                ids.extend(self.tokenizer.encode_as_ids(text[idx:next_idx]))
+                ids.append(self.special_token_to_id[next_token])
+                idx = next_idx + len(next_token)
+
+            ids.extend(self.tokenizer.encode_as_ids(text[idx:]))
+            return ids
+
+        return self.tokenizer.encode_as_ids(text)
+
+    def tokens_to_text(self, tokens):
+        if isinstance(tokens, np.ndarray):
+            tokens = tokens.tolist()
+
+        return self.tokenizer.decode_pieces(tokens)
+
+    def ids_to_text(self, ids):
+        if isinstance(ids, np.ndarray):
+            ids = ids.tolist()
+
+        if self.legacy:
+            text = ""
+            last_i = 0
+
+            for i, id in enumerate(ids):
+                if id in self.id_to_special_token:
+                    text += self.tokenizer.decode_ids(ids[last_i:i]) + " "
+                    text += self.id_to_special_token[id] + " "
+                    last_i = i + 1
+
+            text += self.tokenizer.decode_ids(ids[last_i:])
+            return text.strip()
+
+        return self.tokenizer.decode_ids(ids)
+
+    def token_to_id(self, token):
+        if self.legacy and token in self.special_token_to_id:
+            return self.special_token_to_id[token]
+
+        return self.tokenizer.piece_to_id(token)
+
+    def ids_to_tokens(self, ids):
+        tokens = []
+        for id in ids:
+            if id >= self.original_vocab_size:
+                tokens.append(self.id_to_special_token[id])
+            else:
+                tokens.append(self.tokenizer.id_to_piece(id))
+        return tokens
+
+    def tokens_to_ids(self, tokens: Union[str, List[str]]) -> Union[int, List[int]]:
+        if isinstance(tokens, str):
+            tokens = [tokens]
+        ids = []
+        for token in tokens:
+            ids.append(self.token_to_id(token))
+        return ids
+
+    def add_special_tokens(self, special_tokens):
+        if not self.legacy:
+            raise AttributeError("Special Token addition does not work when legacy is set to False.")
+
+        if isinstance(special_tokens, list):
+            for token in special_tokens:
+                if (
+                    self.tokenizer.piece_to_id(token) == self.tokenizer.unk_id()
+                    and token not in self.special_token_to_id
+                ):
+                    self.special_token_to_id[token] = self.vocab_size
+                    self.id_to_special_token[self.vocab_size] = token
+                    self.vocab_size += 1
+        elif isinstance(special_tokens, dict):
+            for token_name, token in special_tokens.items():
+                setattr(self, token_name, token)
+                if (
+                    self.tokenizer.piece_to_id(token) == self.tokenizer.unk_id()
+                    and token not in self.special_token_to_id
+                ):
+                    self.special_token_to_id[token] = self.vocab_size
+                    self.id_to_special_token[self.vocab_size] = token
+                    self.vocab_size += 1
+
+    @property
+    def pad_id(self):
+        if self.legacy:
+            pad_id = self.tokens_to_ids([self.pad_token])[0]
+        else:
+            pad_id = self.tokenizer.pad_id()
+        return pad_id
+
+    @property
+    def bos_id(self):
+        if self.legacy:
+            bos_id = self.tokens_to_ids([self.bos_token])[0]
+        else:
+            bos_id = self.tokenizer.bos_id()
+        return bos_id
+
+    @property
+    def eos_id(self):
+        if self.legacy:
+            eos_id = self.tokens_to_ids([self.eos_token])[0]
+        else:
+            eos_id = self.tokenizer.eos_id()
+        return eos_id
+
+    @property
+    def sep_id(self):
+        if self.legacy:
+            return self.tokens_to_ids([self.sep_token])[0]
+        else:
+            raise NameError("Use function token_to_id to retrieve special tokens other than unk, pad, bos, and eos.")
+
+    @property
+    def cls_id(self):
+        if self.legacy:
+            return self.tokens_to_ids([self.cls_token])[0]
+        else:
+            raise NameError("Use function token_to_id to retrieve special tokens other than unk, pad, bos, and eos.")
+
+    @property
+    def unk_id(self):
+        return self.tokenizer.unk_id()
+
+
+def create_spt_model(
+    data_file: str,
+    vocab_size: int,
+    sample_size: int,
+    do_lower_case: bool,
+    tokenizer_type: str = 'unigram',
+    output_dir: Optional[str] = None,
+    character_coverage: float = 1.0,
+    train_extremely_large_corpus: bool = False,
+    max_sentencepiece_length: int = -1,
+    bos: bool = False,
+    eos: bool = False,
+    pad: bool = False,
+    control_symbols: List[str] = None,
+    user_defined_symbols: List[str] = None,
+):
+    """
+    Creates sentence piece tokenizer model from data file.
+    Args:
+        data_file: data file
+        vocab_size: vocabulary size
+        sample_size: maximum size of sentences the trainer loads
+        do_lower_case: if text should be lower cased before tokenizer model is created
+        character_coverage: float value between 0 and 1 (as a percentage). For languages with a vast charset,
+            can be < 1.0, but for all other languages, it should be set as 1.0
+        output_dir: folder to save created tokenizer model. If not specified will store model at data_file/../spt folder
+        train_extremely_large_corpus: If training on huge datasets, pass this flag to allow SentencePiece
+            to build the tokenizer.
+        max_sentencepiece_length: Limits the maximum length of the SentencePiece subword that can be constructed.
+            By default, no limit is placed.
+        bos: when True, bos token "<s>" is added to the vocabulary.
+        eos: when True, eos token "</s>" is added to the vocabulary.
+        pad: when True, pad token "<pad>" is added to the vocabulary.
+        control_symbols: control symbols to add to tokenizer, as defined by sentencepiece.
+            These tokens get removed at decode time and are not encoded from the text - can only be added to the input programatically.
+        user_defined_symbols: user symbols to add to tokenizer, as defined by sentencepiece.
+            These tokens remain in the decoded text and are encoded automatically when present in the input text.
+    """
+
+    if not data_file or not os.path.exists(data_file):
+        raise ValueError(f"data_file must be valid file path, but got {data_file}")
+    data_dir = os.path.dirname(data_file)
+    vocab = []
+    special_tokens = ["<s>", "</s>", "<pad>", "<unk>"]
+    if not output_dir:
+        output_dir = f'{data_dir}/spt'
+    if if_exist(output_dir, ['tokenizer.model']):
+        logging.info(f"tokenizer model {output_dir}/tokenizer.model already exists")
+        return f'{output_dir}/tokenizer.model', f'{output_dir}/vocab.txt'
+    logging.info(f'Processing {data_file} and store at {output_dir}')
+    os.makedirs(output_dir, exist_ok=True)
+
+    cmd = (
+        f"--input={data_file} --model_prefix={output_dir}/tokenizer "
+        f"--vocab_size={vocab_size} "
+        f"--shuffle_input_sentence=true --hard_vocab_limit=false "
+        f"--model_type={tokenizer_type} "
+        f"--character_coverage={character_coverage}"
+    )
+
+    pad_id = 3
+    if not bos:
+        pad_id -= 1
+        cmd += " --bos_id=-1"
+
+    if not eos:
+        pad_id -= 1
+        cmd += " --eos_id=-1"
+
+    if pad:
+        cmd += f" --pad_id={pad_id}"
+
+    if control_symbols:
+        control_string = (",").join(control_symbols)
+        cmd += f" --control_symbols={control_string}"
+        special_tokens += control_symbols
+
+    if user_defined_symbols:
+        user_string = (",").join(user_defined_symbols)
+        cmd += f" --user_defined_symbols={user_string}"
+        special_tokens += user_defined_symbols
+
+    if do_lower_case:
+        cmd += " --normalization_rule_name=nmt_nfkc_cf"
+
+    if sample_size > 0:
+        cmd += f" --input_sentence_size={sample_size}"
+
+    if train_extremely_large_corpus:
+        cmd += " --train_extremely_large_corpus=true"
+
+    if max_sentencepiece_length >= 0:
+        cmd += f" --max_sentencepiece_length={max_sentencepiece_length}"
+
+    sentencepiece.SentencePieceTrainer.Train(cmd)
+
+    # Add BERT control symbols
+    tokens = []
+
+    with open(f"{output_dir}/tokenizer.vocab", "r") as f:
+        # Read tokens from each line and parse for vocab
+        for line in f:
+            piece = line.split("\t")[0]
+            if piece in special_tokens:
+                # skip special tokens
+                continue
+            token = piece[1:] if piece.startswith("▁") else f"##{piece}"
+
+            if len(token) > 0:
+                tokens.append(token)
+            else:
+                tokens.append(piece[0])
+
+    vocab.extend(tokens)
+
+    # Save vocabulary to output file
+    vocab_file = f'{output_dir}/vocab.txt'
+    with open(vocab_file, "w") as f:
+        for token in vocab:
+            f.write(f"{token}\n")
+    return f'{output_dir}/tokenizer.model', vocab_file

tts/nemo/collections/common/tokenizers/tokenizer_spec.py

@@ -0,0 +1,55 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from abc import ABC, abstractmethod
+from typing import List
+
+__all__ = ['TokenizerSpec']
+
+
+class TokenizerSpec(ABC):
+    """
+    Inherit this class to implement a new tokenizer.
+    """
+
+    @abstractmethod
+    def text_to_tokens(self, text):
+        pass
+
+    @abstractmethod
+    def tokens_to_text(self, tokens):
+        pass
+
+    @abstractmethod
+    def tokens_to_ids(self, tokens):
+        pass
+
+    @abstractmethod
+    def ids_to_tokens(self, ids):
+        pass
+
+    @abstractmethod
+    def text_to_ids(self, text):
+        pass
+
+    @abstractmethod
+    def ids_to_text(self, ids):
+        pass
+
+    def add_special_tokens(self, special_tokens: List[str]):
+        raise NotImplementedError("To be implemented")
+
+    @property
+    def name(self):
+        return type(self).__name__

tts/nemo/collections/common/tokenizers/word_tokenizer.py

@@ -0,0 +1,72 @@
+# 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
+
+from nemo.collections.common.tokenizers.char_tokenizer import CharTokenizer
+
+__all__ = ['WordTokenizer']
+
+
+class WordTokenizer(CharTokenizer):
+    "Tokenizes at word boundary"
+
+    def __init__(
+        self,
+        vocab_file: str,
+        mask_token: Optional[str] = None,
+        bos_token: Optional[str] = None,
+        eos_token: Optional[str] = None,
+        pad_token: Optional[str] = None,
+        sep_token: Optional[str] = None,
+        cls_token: Optional[str] = None,
+        unk_token: Optional[str] = None,
+    ):
+        """
+        Args:
+            vocab_file: path to file with vocabulary which consists
+                of characters separated by \n
+            mask_token: mask token 
+            bos_token: the beginning of sequence token
+            eos_token: the end of sequence token. Usually equal to sep_token
+            pad_token: token to use for padding
+            sep_token: token used for separating sequences
+            cls_token: class token. Usually equal to bos_token
+            unk_token: token to use for unknown tokens
+        """
+
+        super().__init__(
+            vocab_file=vocab_file,
+            mask_token=mask_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            pad_token=pad_token,
+            unk_token=unk_token,
+            sep_token=sep_token,
+            cls_token=cls_token,
+        )
+
+    def text_to_tokens(self, text):
+        token_candidates = text.strip().split()
+        tokens = []
+        for token in token_candidates:
+            if token in self.vocab:
+                tokens.append(token)
+            else:
+                tokens.append(self.unk_token)
+        return tokens
+
+    def ids_to_text(self, ids):
+        ids_ = [id_ for id_ in ids if id_ not in self.special_tokens]
+        return " ".join(self.ids_to_tokens(ids_))

tts/nemo/collections/common/tokenizers/youtokentome_tokenizer.py

@@ -0,0 +1,70 @@
+# 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 pathlib import Path
+
+import youtokentome as yttm
+
+from nemo.collections.common.tokenizers import TokenizerSpec
+
+__all__ = ['YouTokenToMeTokenizer']
+
+
+class YouTokenToMeTokenizer(TokenizerSpec):
+    def __init__(self, model_path, bpe_dropout=0.0, legacy=False):
+        model_path = Path(model_path).expanduser()
+        self.tokenizer = yttm.BPE(model=str(model_path))
+        self.vocab_size = len(self.tokenizer.vocab())
+        self.special_tokens = self.tokens_to_ids(["<PAD>", "<UNK>", "<BOS>", "<EOS>"])
+        self.bpe_dropout = bpe_dropout
+        self.legacy = legacy
+
+    def text_to_tokens(self, text):
+        return self.tokenizer.encode(text, output_type=yttm.OutputType.SUBWORD, dropout_prob=self.bpe_dropout)
+
+    def tokens_to_text(self, tokens):
+        return self.ids_to_text(self.tokens_to_ids(tokens))
+
+    def text_to_ids(self, text):
+        return self.tokenizer.encode(text, output_type=yttm.OutputType.ID, dropout_prob=self.bpe_dropout)
+
+    def ids_to_text(self, ids):
+        ids_ = [id_ for id_ in ids if id_ not in self.special_tokens]
+        return self.tokenizer.decode([ids_])[0]
+
+    def tokens_to_ids(self, tokens):
+        return [self.tokenizer.subword_to_id(token) for token in tokens]
+
+    def ids_to_tokens(self, ids):
+        if self.legacy:
+            ids_ = [id_ for id_ in ids if id_ not in self.special_tokens]
+        else:
+            ids_ = ids
+        return [self.tokenizer.id_to_subword(id_) for id_ in ids_]
+
+    @property
+    def pad_id(self):
+        return self.tokenizer.subword_to_id("<PAD>")
+
+    @property
+    def bos_id(self):
+        return self.tokenizer.subword_to_id("<BOS>")
+
+    @property
+    def eos_id(self):
+        return self.tokenizer.subword_to_id("<EOS>")
+
+    @property
+    def unk_id(self):
+        return self.tokenizer.subword_to_id("<UNK>")