Introduce code from Rust monorepo. This is a mess

tts/nemo/collections/nlp/modules/common/__init__.py

@@ -0,0 +1,34 @@
+# Copyright 2018 The Google AI Language Team Authors and
+# The HuggingFace Inc. team.
+# 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.nlp.modules.common.bert_module import BertModule
+from nemo.collections.nlp.modules.common.huggingface import (
+    AlbertEncoder,
+    BertEncoder,
+    DistilBertEncoder,
+    RobertaEncoder,
+)
+from nemo.collections.nlp.modules.common.lm_utils import (
+    get_lm_model,
+    get_megatron_lm_models_list,
+    get_pretrained_lm_models_list,
+)
+from nemo.collections.nlp.modules.common.megatron import MegatronBertEncoder
+from nemo.collections.nlp.modules.common.sequence_classifier import SequenceClassifier
+from nemo.collections.nlp.modules.common.sequence_regression import SequenceRegression
+from nemo.collections.nlp.modules.common.sequence_token_classifier import SequenceTokenClassifier
+from nemo.collections.nlp.modules.common.token_classifier import BertPretrainingTokenClassifier, TokenClassifier
+from nemo.collections.nlp.modules.common.tokenizer_utils import get_tokenizer, get_tokenizer_list

tts/nemo/collections/nlp/modules/common/bert_module.py

@@ -0,0 +1,93 @@
+# 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, Optional
+
+import torch
+
+from nemo.core.classes import NeuralModule
+from nemo.core.classes.exportable import Exportable
+from nemo.core.neural_types import ChannelType, MaskType, NeuralType
+from nemo.utils import logging
+
+__all__ = ['BertModule']
+
+
+class BertModule(NeuralModule, Exportable):
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {
+            "input_ids": NeuralType(('B', 'T'), ChannelType()),
+            "attention_mask": NeuralType(('B', 'T'), MaskType(), optional=True),
+            "token_type_ids": NeuralType(('B', 'T'), ChannelType(), optional=True),
+        }
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {"last_hidden_states": NeuralType(('B', 'T', 'D'), ChannelType())}
+
+    @classmethod
+    def restore_from(cls, restore_path: str):
+        """Restores module/model with weights"""
+        pass
+
+    @classmethod
+    def save_to(self, save_path: str):
+        """Saves module/model with weights"""
+        pass
+
+    def restore_weights(self, restore_path: str):
+        """Restores module/model's weights"""
+        logging.info(f"Restoring weights from {restore_path}")
+
+        if not os.path.exists(restore_path):
+            logging.warning(f'Path {restore_path} not found')
+            return
+
+        pretrained_dict = torch.load(restore_path)
+
+        # backward compatibility with NeMo0.11
+        if "state_dict" in pretrained_dict.keys():
+            pretrained_dict = pretrained_dict["state_dict"]
+
+        # remove prefix from pretrained dict
+        m = re.match("^bert.*?\.", list(pretrained_dict.keys())[0])
+        if m:
+            prefix = m.group(0)
+            pretrained_dict = {k[len(prefix) :]: v for k, v in pretrained_dict.items()}
+        model_dict = self.state_dict()
+        pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
+
+        # starting with transformers 3.1.0, embeddings.position_ids is added to the model's state dict and could be
+        # missing in checkpoints trained with older transformers version
+        if 'embeddings.position_ids' in model_dict and 'embeddings.position_ids' not in pretrained_dict:
+            pretrained_dict['embeddings.position_ids'] = model_dict['embeddings.position_ids']
+
+        assert len(pretrained_dict) == len(model_dict)
+        model_dict.update(pretrained_dict)
+        self.load_state_dict(model_dict)
+        logging.info(f"Weights for {type(self).__name__} restored from {restore_path}")
+
+    def input_example(self):
+        """
+        Generates input examples for tracing etc.
+        Returns:
+            A tuple of input examples.
+        """
+        sample = next(self.parameters())
+        input_ids = torch.randint(low=0, high=2048, size=(2, 16), device=sample.device)
+        attention_mask = torch.randint(low=0, high=1, size=(2, 16), device=sample.device)
+        return tuple([input_ids, attention_mask, attention_mask])

tts/nemo/collections/nlp/modules/common/classifier.py

@@ -0,0 +1,87 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, Optional
+
+import torch
+from torch import nn as nn
+
+from nemo.collections.common.parts import transformer_weights_init
+from nemo.core.classes import Exportable, NeuralModule, typecheck
+from nemo.core.neural_types import ChannelType, NeuralType
+
+__all__ = ['Classifier']
+
+
+class Classifier(NeuralModule, Exportable):
+    """
+    A baseclass for modules to perform various classification tasks.
+    """
+
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        """
+        Returns definitions of module input ports.
+        We implement it here since all NLP classifiers have the same inputs
+        """
+        return {"hidden_states": NeuralType(('B', 'T', 'D'), ChannelType())}
+
+    def __init__(self, hidden_size: int, dropout: float = 0.0,) -> None:
+        """
+        Initializes the Classifier base module.
+        Args:
+            hidden_size: the size of the hidden dimension
+            dropout: dropout to apply to the input hidden states
+        """
+        super().__init__()
+        self._hidden_size = hidden_size
+        self.dropout = nn.Dropout(dropout)
+
+    def post_init(self, use_transformer_init: bool):
+        """
+        Common post-processing to be called at the end of concrete Classifiers init methods
+        Args:
+          use_transformer_init : whether or not to apply transformer_weights_init
+        """
+        if use_transformer_init:
+            self.apply(lambda module: transformer_weights_init(module, xavier=False))
+
+    def input_example(self):
+        """
+        Generates input examples for tracing etc.
+        Returns:
+            A tuple of input examples.
+        """
+        bs = 8
+        seq = 64
+        sample = next(self.parameters())
+        input_example = torch.randn(bs, seq, self._hidden_size).to(sample.device).to(sample.dtype)
+        return tuple([input_example])
+
+    def save_to(self, save_path: str):
+        """
+        Saves the module to the specified path.
+        Args:
+            save_path: Path to where to save the module.
+        """
+        pass
+
+    @classmethod
+    def restore_from(cls, restore_path: str):
+        """
+        Restores the module from the specified path.
+        Args:
+            restore_path: Path to restore the module from.
+        """
+        pass

tts/nemo/collections/nlp/modules/common/decoder_module.py

@@ -0,0 +1,58 @@
+# 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
+from typing import Any, Dict, Optional
+
+from nemo.core.classes import NeuralModule
+from nemo.core.neural_types import ChannelType, MaskType, NeuralType
+
+__all__ = ['DecoderModule']
+
+
+class DecoderModule(NeuralModule, ABC):
+    """ Base class for decoder neural module to be used in NLP models. """
+
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {
+            "input_ids": NeuralType(('B', 'T'), ChannelType()),
+            "decoder_mask": NeuralType(('B', 'T'), MaskType(), optional=True),
+            "encoder_mask": NeuralType(('B', 'T'), MaskType(), optional=True),
+            "encoder_embeddings": NeuralType(('B', 'T', 'D'), ChannelType(), optional=True),
+        }
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {"last_hidden_states": NeuralType(('B', 'T', 'D'), ChannelType())}
+
+    @property
+    def hidden_size(self) -> Optional[int]:
+        raise NotImplementedError
+
+    @property
+    def vocab_size(self) -> Optional[int]:
+        raise NotImplementedError
+
+    @property
+    def embedding(self) -> Optional[Any]:
+        raise NotImplementedError
+
+    @property
+    def decoder(self) -> Optional[Any]:
+        raise NotImplementedError
+
+    @property
+    def max_sequence_length(self) -> Optional[int]:
+        raise NotImplementedError

tts/nemo/collections/nlp/modules/common/encoder_module.py

@@ -0,0 +1,41 @@
+# 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
+from typing import Dict, Optional
+
+from nemo.core.classes import NeuralModule
+from nemo.core.neural_types import ChannelType, MaskType, NeuralType
+from nemo.utils import logging
+
+__all__ = ['EncoderModule']
+
+
+class EncoderModule(NeuralModule, ABC):
+    """ Base class for encoder neural module to be used in NLP models. """
+
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {
+            "input_ids": NeuralType(('B', 'T'), ChannelType()),
+            "encoder_mask": NeuralType(('B', 'T'), MaskType()),
+        }
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {"last_hidden_states": NeuralType(('B', 'T', 'D'), ChannelType())}
+
+    @property
+    def hidden_size(self) -> Optional[int]:
+        raise NotImplementedError

tts/nemo/collections/nlp/modules/common/huggingface/__init__.py

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

tts/nemo/collections/nlp/modules/common/huggingface/albert.py

@@ -0,0 +1,33 @@
+# Copyright 2018 The Google AI Language Team Authors and
+# The HuggingFace Inc. team.
+# 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 transformers import AlbertModel
+
+from nemo.collections.nlp.modules.common.bert_module import BertModule
+from nemo.core.classes import typecheck
+
+__all__ = ['AlbertEncoder']
+
+
+class AlbertEncoder(AlbertModel, BertModule):
+    """
+    Wraps around the Huggingface transformers implementation repository for easy use within NeMo.
+    """
+
+    @typecheck()
+    def forward(self, input_ids, attention_mask, token_type_ids):
+        res = super().forward(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)[0]
+        return res

tts/nemo/collections/nlp/modules/common/huggingface/bert.py

@@ -0,0 +1,33 @@
+# Copyright 2018 The Google AI Language Team Authors and
+# The HuggingFace Inc. team.
+# 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 transformers import BertModel
+
+from nemo.collections.nlp.modules.common.bert_module import BertModule
+from nemo.core.classes import typecheck
+
+__all__ = ['BertEncoder']
+
+
+class BertEncoder(BertModel, BertModule):
+    """
+    Wraps around the Huggingface transformers implementation repository for easy use within NeMo.
+    """
+
+    @typecheck()
+    def forward(self, input_ids, attention_mask, token_type_ids):
+        res = super().forward(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)[0]
+        return res

tts/nemo/collections/nlp/modules/common/huggingface/distilbert.py

@@ -0,0 +1,34 @@
+# Copyright 2020 The Google AI Language Team Authors and
+# The HuggingFace Inc. team.
+# 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 transformers import DistilBertModel
+
+from nemo.collections.nlp.modules.common.bert_module import BertModule
+from nemo.core.classes import typecheck
+
+__all__ = ['DistilBertEncoder']
+
+
+class DistilBertEncoder(DistilBertModel, BertModule):
+    """
+    Wraps around the Huggingface transformers implementation repository for easy use within NeMo.
+    """
+
+    @typecheck()
+    def forward(self, input_ids, attention_mask, token_type_ids=None):
+        # distilBert does not use token_type_ids as the most of the other Bert models
+        res = super().forward(input_ids=input_ids, attention_mask=attention_mask)[0]
+        return res

tts/nemo/collections/nlp/modules/common/huggingface/huggingface_decoder.py

@@ -0,0 +1,79 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional
+
+from hydra.utils import instantiate
+from transformers import AutoConfig, AutoModel
+
+from nemo.collections.nlp.modules.common.decoder_module import DecoderModule
+from nemo.collections.nlp.modules.common.huggingface.huggingface_utils import get_huggingface_pretrained_lm_models_list
+from nemo.utils import logging
+
+
+class HuggingFaceDecoderModule(DecoderModule):
+    """Gets HuggingFace based model to be used as an Decoder in NeMo NLP.
+    Use the model_name arg to get a named model architecture. 
+    Available model names can be found with get_huggingface_pretrained_lm_models_list() or
+    by going to https://huggingface.co/models.
+    Use the pretrained arg to get the named model architecture with or without pretrained weights.
+
+    If model_name is None, then we can pass in a custom configuration via the config_dict.
+    For example, to instantiate a HuggingFace BERT model with custom configuration we would do:
+        config_dict={
+            '_target_': 'transformers.BertConfig',
+            'hidden_size': 1536
+        } 
+
+
+    Args:
+        model_name (Optional[str]): Named model architecture from HuggingFace. Defaults to None.
+        pretrained (bool): Use True to get pretrained weights. 
+                                    False will use the same architecture but with randomly initialized weights.
+                                    Defaults to False.
+        config_dict (Optional[dict], optional): Use for custom configuration of the HuggingFace model. Defaults to None.
+        checkpoint_file (Optional[str], optional): Provide weights for the transformer from a local checkpoint. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        model_name: Optional[str] = None,
+        pretrained: bool = False,
+        config_dict: Optional[dict] = None,
+        checkpoint_file: Optional[str] = None,
+    ):
+        super().__init__()
+        model = None
+        if model_name is not None:
+            if model_name in get_huggingface_pretrained_lm_models_list():
+                if pretrained:
+                    model = AutoModel.from_pretrained(model_name)
+                else:
+                    cfg = AutoConfig.from_pretrained(model_name)
+                    model = AutoModel.from_config(cfg)
+            else:
+                logging.error(f'{model_name} not found in list of HuggingFace pretrained models')
+        else:
+            cfg = instantiate(config_dict)
+            model = AutoModel.from_config(cfg)
+        self._hidden_size = model.config.hidden_size
+        self._vocab_size = model.config.vocab_size
+
+    @property
+    def hidden_size(self) -> Optional[int]:
+        return self._hidden_size
+
+    @property
+    def vocab_size(self) -> Optional[int]:
+        return self._vocab_size

tts/nemo/collections/nlp/modules/common/huggingface/huggingface_encoder.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional
+
+from hydra.utils import instantiate
+from transformers import AutoConfig, AutoModel
+
+from nemo.collections.nlp.modules.common.encoder_module import EncoderModule
+from nemo.collections.nlp.modules.common.huggingface.huggingface_utils import get_huggingface_pretrained_lm_models_list
+from nemo.core.classes.common import typecheck
+from nemo.utils import logging
+
+
+class HuggingFaceEncoderModule(EncoderModule):
+    """ Class for using HuggingFace encoders in NeMo NLP."""
+
+    def __init__(
+        self,
+        model_name: Optional[str] = None,
+        pretrained: bool = False,
+        config_dict: Optional[dict] = None,
+        checkpoint_file: Optional[str] = None,
+    ):
+        """Gets HuggingFace based model to be used as an Encoder in NeMo NLP.
+        Use the model_name arg to get a named model architecture. 
+        Available model names can be found with get_huggingface_pretrained_lm_models_list() or
+        by going to https://huggingface.co/models.
+        Use the pretrained arg to get the named model architecture with or without pretrained weights.
+
+        If model_name is None, then we can pass in a custom configuration via the config_dict.
+        For example, to instantiate a HuggingFace BERT model with custom configuration we would do:
+            config_dict={
+                '_target_': 'transformers.BertConfig',
+                'hidden_size': 1536
+            } 
+
+
+        Args:
+            model_name (Optional[str]): Named model architecture from HuggingFace. Defaults to None.
+            pretrained (bool): Use True to get pretrained weights. 
+                                        False will use the same architecture but with randomly initialized weights.
+                                        Defaults to False.
+            config_dict (Optional[dict], optional): Use for custom configuration of the HuggingFace model. Defaults to None.
+            checkpoint_file (Optional[str], optional): Provide weights for the transformer from a local checkpoint. Defaults to None.
+        """
+        super().__init__()
+
+        if checkpoint_file:
+            raise NotImplementedError('Restoring from checkpoint file not implemented yet.')
+
+        model = None
+        if model_name is not None:
+            if model_name in get_huggingface_pretrained_lm_models_list(include_external=True):
+                if pretrained:
+                    config_dict.pop('vocab_size')
+                    if config_dict:
+                        raise ValueError(
+                            f'When using pretrained model, config_dict should be None or empty. Got: {config_dict}'
+                        )
+                    model = AutoModel.from_pretrained(model_name)
+                else:
+                    cfg = AutoConfig.from_pretrained(model_name)
+                    model = AutoModel.from_config(cfg)
+            else:
+                logging.error(f'{model_name} not found in list of HuggingFace pretrained models')
+        else:
+            if pretrained:
+                raise ValueError(f'If not using model_name, then pretrained should be False. Got: {pretrained}.')
+            cfg = instantiate(config_dict)
+            model = AutoModel.from_config(cfg)
+        self._hidden_size = model.config.hidden_size
+        self._vocab_size = model.config.vocab_size
+
+        self._encoder = model
+
+    @typecheck()
+    def forward(self, input_ids, encoder_mask):
+        encoder_hidden_states = self._encoder.forward(input_ids=input_ids, attention_mask=encoder_mask)[0]
+        return encoder_hidden_states
+
+    @property
+    def hidden_size(self) -> Optional[int]:
+        return self._hidden_size
+
+    @property
+    def vocab_size(self) -> Optional[int]:
+        return self._vocab_size

tts/nemo/collections/nlp/modules/common/huggingface/huggingface_utils.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.
+
+import os
+from typing import List, Optional
+
+from transformers import (
+    ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
+    ALL_PRETRAINED_CONFIG_ARCHIVE_MAP,
+    BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
+    DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
+    ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
+    AlbertConfig,
+    AutoModel,
+    BertConfig,
+    DistilBertConfig,
+    RobertaConfig,
+)
+
+from nemo.collections.nlp.modules.common.huggingface.albert import AlbertEncoder
+from nemo.collections.nlp.modules.common.huggingface.bert import BertEncoder
+from nemo.collections.nlp.modules.common.huggingface.distilbert import DistilBertEncoder
+from nemo.collections.nlp.modules.common.huggingface.roberta import RobertaEncoder
+from nemo.utils import logging
+
+__all__ = ["get_huggingface_lm_model", "get_huggingface_pretrained_lm_models_list"]
+
+
+HUGGINGFACE_MODELS = {
+    "BertModel": {
+        "default": "bert-base-uncased",
+        "class": BertEncoder,
+        "config": BertConfig,
+        "pretrained_model_list": BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
+    },
+    "DistilBertModel": {
+        "default": "distilbert-base-uncased",
+        "class": DistilBertEncoder,
+        "config": DistilBertConfig,
+        "pretrained_model_list": DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
+    },
+    "RobertaModel": {
+        "default": "roberta-base",
+        "class": RobertaEncoder,
+        "config": RobertaConfig,
+        "pretrained_model_list": ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
+    },
+    "AlbertModel": {
+        "default": "albert-base-v2",
+        "class": AlbertEncoder,
+        "config": AlbertConfig,
+        "pretrained_model_list": ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
+    },
+}
+
+
+def get_huggingface_lm_model(
+    pretrained_model_name: str, config_dict: Optional[dict] = None, config_file: Optional[str] = None,
+):
+    """
+    Returns lm model instantiated with Huggingface
+
+    Args:
+        pretrained_mode_name: specify this to instantiate pretrained model from Huggingface,
+            e.g. bert-base-cased. For entire list, see get_huggingface_pretrained_lm_models_list().
+        config_dict: model configuration dictionary used to instantiate Huggingface model from scratch
+        config_file: path to model configuration file used to instantiate Huggingface model from scratch
+
+    Returns:
+        BertModule
+    """
+
+    try:
+        automodel = AutoModel.from_pretrained(pretrained_model_name)
+    except Exception as e:
+        raise ValueError(f"{pretrained_model_name} is not supported by HuggingFace. {e}")
+
+    model_type = type(automodel).__name__
+    if model_type in HUGGINGFACE_MODELS:
+        model_class = HUGGINGFACE_MODELS[model_type]["class"]
+        if config_file:
+            if not os.path.exists(config_file):
+                logging.warning(
+                    f"Config file was not found at {config_file}. Will attempt to use config_dict or pretrained_model_name."
+                )
+            else:
+                config_class = HUGGINGFACE_MODELS[model_type]["config"]
+                return model_class(config_class.from_json_file(config_file))
+        if config_dict:
+            config_class = HUGGINGFACE_MODELS[model_type]["config"]
+            return model_class(config=config_class(**config_dict))
+        else:
+            return model_class.from_pretrained(pretrained_model_name)
+    else:
+        raise ValueError(f"Use HuffingFace API directly in NeMo for {pretrained_model_name}")
+
+
+def get_huggingface_pretrained_lm_models_list(include_external: bool = False,) -> List[str]:
+    """
+    Returns the list of pretrained HuggingFace language models
+    
+    Args:
+        include_external if true includes all HuggingFace model names, not only those supported language models in NeMo.
+    
+    Returns the list of HuggingFace models
+    """
+
+    huggingface_models = []
+    if include_external:
+        huggingface_models = list(ALL_PRETRAINED_CONFIG_ARCHIVE_MAP.keys())
+    else:
+        for model in HUGGINGFACE_MODELS:
+            model_names = HUGGINGFACE_MODELS[model]["pretrained_model_list"]
+            huggingface_models.extend(model_names)
+    return huggingface_models

tts/nemo/collections/nlp/modules/common/huggingface/roberta.py

@@ -0,0 +1,34 @@
+# Copyright 2018 The Google AI Language Team Authors and
+# The HuggingFace Inc. team.
+# 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 transformers import RobertaModel
+
+from nemo.collections.nlp.modules.common.bert_module import BertModule
+from nemo.core.classes import typecheck
+
+__all__ = ['RobertaEncoder']
+
+
+class RobertaEncoder(RobertaModel, BertModule):
+    """
+    Wraps around the Huggingface transformers implementation repository for easy use within NeMo.
+    """
+
+    @typecheck()
+    def forward(self, input_ids, token_type_ids, attention_mask):
+        res = super().forward(input_ids=input_ids, attention_mask=attention_mask)[0]
+        return res

tts/nemo/collections/nlp/modules/common/lm_utils.py

@@ -0,0 +1,179 @@
+# Copyright 2018 The Google AI Language Team Authors and
+# The HuggingFace Inc. team.
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+from dataclasses import dataclass
+from typing import List, Optional, Union
+
+from attr import asdict
+
+from nemo.collections.nlp.modules.common.bert_module import BertModule
+from nemo.collections.nlp.modules.common.decoder_module import DecoderModule
+from nemo.collections.nlp.modules.common.encoder_module import EncoderModule
+from nemo.collections.nlp.modules.common.huggingface.huggingface_utils import (
+    get_huggingface_lm_model,
+    get_huggingface_pretrained_lm_models_list,
+)
+from nemo.collections.nlp.modules.common.megatron.megatron_utils import (
+    get_megatron_lm_model,
+    get_megatron_lm_models_list,
+)
+from nemo.collections.nlp.modules.common.transformer.transformer import NeMoTransformerConfig
+from nemo.collections.nlp.modules.common.transformer.transformer_utils import (
+    get_huggingface_transformer,
+    get_nemo_transformer,
+)
+from nemo.utils import logging
+
+__all__ = ['get_pretrained_lm_models_list', 'get_lm_model']
+
+
+def get_pretrained_lm_models_list(include_external: bool = False) -> List[str]:
+    """
+    Returns the list of supported pretrained model names
+
+    Args:
+        include_external if true includes all HuggingFace model names, not only those supported language models in NeMo.
+
+    """
+    return get_megatron_lm_models_list() + get_huggingface_pretrained_lm_models_list(include_external=include_external)
+
+
+def get_lm_model(
+    pretrained_model_name: str,
+    config_dict: Optional[dict] = None,
+    config_file: Optional[str] = None,
+    checkpoint_file: Optional[str] = None,
+    vocab_file: Optional[str] = None,
+) -> BertModule:
+    """
+    Helper function to instantiate a language model encoder, either from scratch or a pretrained model.
+    If only pretrained_model_name are passed, a pretrained model is returned.
+    If a configuration is passed, whether as a file or dictionary, the model is initialized with random weights.
+
+    Args:
+        pretrained_model_name: pretrained model name, for example, bert-base-uncased or megatron-bert-cased.
+            See get_pretrained_lm_models_list() for full list.
+        config_dict: path to the model configuration dictionary
+        config_file: path to the model configuration file
+        checkpoint_file: path to the pretrained model checkpoint
+        vocab_file: path to vocab_file to be used with Megatron-LM
+
+    Returns:
+        Pretrained BertModule
+    """
+
+    # check valid model type
+    if not pretrained_model_name or pretrained_model_name not in get_pretrained_lm_models_list(include_external=False):
+        logging.warning(
+            f'{pretrained_model_name} is not in get_pretrained_lm_models_list(include_external=False), '
+            f'will be using AutoModel from HuggingFace.'
+        )
+
+    # warning when user passes both configuration dict and file
+    if config_dict and config_file:
+        logging.warning(
+            f"Both config_dict and config_file were found, defaulting to use config_file: {config_file} will be used."
+        )
+
+    if "megatron" in pretrained_model_name:
+        model, checkpoint_file = get_megatron_lm_model(
+            config_dict=config_dict,
+            config_file=config_file,
+            pretrained_model_name=pretrained_model_name,
+            checkpoint_file=checkpoint_file,
+            vocab_file=vocab_file,
+        )
+    else:
+        model = get_huggingface_lm_model(
+            config_dict=config_dict, config_file=config_file, pretrained_model_name=pretrained_model_name,
+        )
+
+    if checkpoint_file and os.path.exists(checkpoint_file):
+        model.restore_weights(restore_path=checkpoint_file)
+
+    return model
+
+
+# @dataclass
+# class TransformerConfig:
+#     library: str = 'nemo'
+#     model_name: Optional[str] = None
+#     pretrained: bool = False
+#     config_dict: Optional[dict] = None
+#     checkpoint_file: Optional[str] = None
+#     encoder: bool = True
+
+
+def get_transformer(
+    library: str = 'nemo',
+    model_name: Optional[str] = None,
+    pretrained: bool = False,
+    config_dict: Optional[dict] = None,
+    checkpoint_file: Optional[str] = None,
+    encoder: bool = True,
+    pre_ln_final_layer_norm=True,
+) -> Union[EncoderModule, DecoderModule]:
+    """Gets Transformer based model to be used as an Encoder or Decoder in NeMo NLP.
+       First choose the library to get the transformer from. This can be huggingface,
+       megatron, or nemo. Use the model_name arg to get a named model architecture
+       and use the pretrained arg to get the named model architecture with pretrained weights.
+
+       If model_name is None, then we can pass in a custom configuration via the config_dict.
+       For example, to instantiate a HuggingFace BERT model with custom configuration we would do:
+       encoder = get_transformer(library='huggingface',
+                                 config_dict={
+                                     '_target_': 'transformers.BertConfig',
+                                     'hidden_size': 1536
+                                 }) 
+
+
+    Args:
+        library (str, optional): Can be 'nemo', 'huggingface', or 'megatron'. Defaults to 'nemo'.
+        model_name (Optional[str], optional): Named model architecture from the chosen library. Defaults to None.
+        pretrained (bool, optional): Use True to get pretrained weights. 
+                                     False will use the same architecture but with randomly initialized weights.
+                                     Defaults to False.
+        config_dict (Optional[dict], optional): Use for custom configuration of transformer. Defaults to None.
+        checkpoint_file (Optional[str], optional): Provide weights for the transformer from a local checkpoint. Defaults to None.
+        encoder (bool, optional): True returns an EncoderModule, False returns a DecoderModule. Defaults to True.
+
+    Returns:
+        Union[EncoderModule, DecoderModule]: Ensures that Encoder/Decoder will work in EncDecNLPModel
+    """
+
+    model = None
+
+    if library == 'nemo':
+        if isinstance(config_dict, NeMoTransformerConfig):
+            config_dict = asdict(config_dict)
+        model = get_nemo_transformer(
+            model_name=model_name,
+            pretrained=pretrained,
+            config_dict=config_dict,
+            encoder=encoder,
+            pre_ln_final_layer_norm=pre_ln_final_layer_norm,
+        )
+
+        if checkpoint_file is not None:
+            if os.path.isfile(checkpoint_file):
+                raise ValueError(f'Loading transformer weights from checkpoint file has not been implemented yet.')
+
+    elif library == 'huggingface':
+        model = get_huggingface_transformer(
+            model_name=model_name, pretrained=pretrained, config_dict=config_dict, encoder=encoder
+        )
+
+    return model

tts/nemo/collections/nlp/modules/common/megatron/__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.nlp.modules.common.megatron.megatron_bert import MegatronBertEncoder
+from nemo.collections.nlp.modules.common.megatron.megatron_utils import (
+    get_megatron_checkpoint,
+    get_megatron_lm_models_list,
+)

tts/nemo/collections/nlp/modules/common/megatron/megatron_bert.py

@@ -0,0 +1,216 @@
+# Copyright 2018 The Google AI Language Team Authors and
+# The HuggingFace Inc. team.
+# 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 torch
+from megatron import get_args, initialize_megatron
+from megatron.checkpointing import set_checkpoint_version
+from megatron.model import get_language_model
+from megatron.model.bert_model import bert_attention_mask_func, bert_extended_attention_mask, bert_position_ids
+from megatron.mpu import (
+    get_model_parallel_group,
+    model_parallel_is_initialized,
+    set_pipeline_model_parallel_rank,
+    set_pipeline_model_parallel_world_size,
+)
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.modules.common.bert_module import BertModule
+from nemo.core.classes import typecheck
+from nemo.utils import logging
+from nemo.utils.app_state import AppState
+
+__all__ = ['MegatronBertEncoder']
+
+
+def complete_lazy_init(self):
+    # finish megatron-lm initialization
+    if hasattr(self, "_lazy_init_fn") and self._lazy_init_fn is not None:
+        self._lazy_init_fn()
+        self._lazy_init_fn = None
+
+
+class MegatronBertEncoder(BertModule):
+    """
+    MegatronBERT wraps around the Megatron Language model
+    from https://github.com/NVIDIA/Megatron-LM
+
+    Args:
+        config_file (str): path to model configuration file.
+        vocab_file (str): path to vocabulary file.
+        tokenizer_type (str): tokenizer type, currently only 'BertWordPieceLowerCase' supported.
+    """
+
+    def __init__(self, model_name, config, vocab_file, model_parallel_size=None, model_parallel_rank=None):
+
+        super().__init__()
+
+        self._model_parallel_size = model_parallel_size
+        self._model_parallel_rank = model_parallel_rank
+        self._restore_path = None
+        self._app_state = None
+        self._model_name = model_name
+
+        if not os.path.exists(vocab_file):
+            raise ValueError(f'Vocab file not found at {vocab_file}')
+
+        # convert config to dictionary
+        if isinstance(config, DictConfig):
+            config = OmegaConf.to_container(config)
+        config["vocab_file"] = vocab_file
+        config['tokenizer_type'] = 'BertWordPieceLowerCase'
+        config['lazy_mpu_init'] = True
+        config['onnx_safe'] = True
+
+        # if 'model_parallel_size' in config:
+        if self._model_parallel_size is not None:
+            app_state = AppState()
+            self._app_state = app_state
+
+            # must be set for model parallel megatron-lm
+            os.environ["WORLD_SIZE"] = str(app_state.world_size)
+            os.environ["RANK"] = str(self._model_parallel_rank)
+
+            extra_args_provider = self._update_megatron_args(tensor_model_parallel_size=self._model_parallel_size)
+
+        else:
+            extra_args_provider = self._update_megatron_args()
+
+        # configure globals for megatron
+        set_pipeline_model_parallel_rank(0)  # pipeline model parallelism not implemented in NeMo
+        set_pipeline_model_parallel_world_size(1)  # pipeline model parallelism not implemented in NeMo
+
+        # Initialize part of Megatron global state that is needed for its constructor.
+        # We set 'lazy_mpu_init' flag on to make Megatron do only the initialization that does not depend
+        # on ddp be initialized yet (and we don't want Megatron to initialize DDP itself either)
+        # and to return a hook for us to call after PTL has torch.distributed initialized.
+        # (or if no PTL in case of inference - then we'll initialize torch.distributed)
+        # We call and clear this hook on first call to forward()
+        self._lazy_init_fn = initialize_megatron(
+            extra_args_provider=extra_args_provider, args_defaults=config, ignore_unknown_args=True
+        )
+
+        # read Megatron arguments back
+        args = get_args()
+        logging.info(f'Megatron-lm argparse args: {args}')
+
+        self.language_model, self._language_model_key = get_language_model(
+            attention_mask_func=bert_attention_mask_func, num_tokentypes=2, add_pooler=False
+        )
+
+        self.config = OmegaConf.create(config)
+        # key used for checkpoints
+        self._hidden_size = self.language_model.hidden_size
+
+    def _update_megatron_args(
+        self,
+        micro_batch_size=1,
+        tensor_model_parallel_size=1,
+        scaled_masked_softmax_fusion=False,
+        bias_gelu_fusion=False,
+        bias_dropout_fusion=False,
+    ):
+        def extra_args_provider(parser):
+            parser.set_defaults(micro_batch_size=micro_batch_size)
+            parser.set_defaults(tensor_model_parallel_size=tensor_model_parallel_size)
+            parser.set_defaults(scaled_masked_softmax_fusion=scaled_masked_softmax_fusion)
+            parser.set_defaults(bias_gelu_fusion=bias_gelu_fusion)
+            parser.set_defaults(bias_dropout_fusion=bias_dropout_fusion)
+
+            return parser
+
+        return extra_args_provider
+
+    def complete_lazy_init(self):
+        # finish megatron-lm initialization
+        if hasattr(self, "_lazy_init_fn") and self._lazy_init_fn is not None:
+            self._lazy_init_fn()
+            self._lazy_init_fn = None
+
+    @property
+    def hidden_size(self):
+        """
+        Property returning hidden size.
+
+        Returns:
+            Hidden size.
+        """
+        return self._hidden_size
+
+    @typecheck()
+    def forward(self, input_ids, attention_mask, token_type_ids):
+        app_state = AppState()
+        if app_state.model_parallel_size is None:
+            self.complete_lazy_init()
+
+        extended_attention_mask = bert_extended_attention_mask(attention_mask)
+        position_ids = bert_position_ids(input_ids)
+
+        sequence_output = self.language_model(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            attention_mask=extended_attention_mask,
+            tokentype_ids=token_type_ids,
+        )
+        return sequence_output
+
+    def _load_checkpoint(self, filename):
+        """Helper function for loading megatron checkpoints.
+
+        Args:
+            filename (str): Path to megatron checkpoint.
+        """
+        state_dict = torch.load(filename, map_location='cpu')
+        if 'checkpoint_version' in state_dict:
+            if state_dict['checkpoint_version'] is not None:
+                set_checkpoint_version(state_dict['checkpoint_version'])
+                logging.info(
+                    f"Megatron-lm checkpoint version found. Setting checkpoint_version to {state_dict['checkpoint_version']}."
+                )
+        else:
+            logging.warning('Megatron-lm checkpoint version not found. Setting checkpoint_version to 0.')
+            set_checkpoint_version(0)
+        # to load from Megatron pretrained checkpoint
+        if 'model' in state_dict:
+            self.language_model.load_state_dict(state_dict['model'][self._language_model_key])
+        else:
+            self.load_state_dict(state_dict)
+
+        logging.info(f"Checkpoint loaded from from {filename}")
+
+    def restore_weights(self, restore_path: str):
+        """Restores module/model's weights.
+           For model parallel checkpoints the directory structure
+           should be restore_path/mp_rank_0X/model_optim_rng.pt
+
+        Args:
+            restore_path (str): restore_path should a file or a directory if using model parallel
+        """
+        self._restore_path = restore_path
+
+        if os.path.isfile(restore_path):
+            self._load_checkpoint(restore_path)
+        elif os.path.isdir(restore_path):
+            # need model parallel groups to restore model parallel checkpoints
+            if model_parallel_is_initialized():
+                model_parallel_rank = torch.distributed.get_rank(group=get_model_parallel_group())
+                mp_restore_path = f'{restore_path}/mp_rank_{model_parallel_rank:02d}/model_optim_rng.pt'
+                self._load_checkpoint(mp_restore_path)
+            else:
+                logging.info(f'torch.distributed not initialized yet. Will not restore model parallel checkpoint')
+        else:
+            logging.error(f'restore_path: {restore_path} must be a file or directory.')

tts/nemo/collections/nlp/modules/common/megatron/megatron_utils.py

@@ -0,0 +1,307 @@
+# Copyright 2020 The HuggingFace Inc. team.
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import glob
+import json
+import os
+from typing import Dict, List, Optional, Tuple
+
+import torch
+import wget
+from torch.hub import _get_torch_home
+
+from nemo.collections.nlp.modules.common.megatron.megatron_bert import MegatronBertEncoder
+from nemo.utils import AppState, logging
+
+__all__ = [
+    "get_megatron_lm_model",
+    "get_megatron_lm_models_list",
+    "get_megatron_checkpoint",
+    "is_lower_cased_megatron",
+    "get_megatron_tokenizer",
+]
+
+
+torch_home = _get_torch_home()
+
+if not isinstance(torch_home, str):
+    logging.info("Torch home not found, caching megatron in cwd")
+    torch_home = os.getcwd()
+
+MEGATRON_CACHE = os.path.join(torch_home, "megatron")
+
+
+CONFIGS = {"345m": {"hidden_size": 1024, "num_attention_heads": 16, "num_layers": 24, "max_position_embeddings": 512}}
+
+MEGATRON_CONFIG_MAP = {
+    "megatron-bert-345m-uncased": {
+        "config": CONFIGS["345m"],
+        "checkpoint": "https://api.ngc.nvidia.com/v2/models/nvidia/megatron_bert_345m/versions/v0.0/files/release/mp_rank_00/model_optim_rng.pt",
+        "vocab": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-vocab.txt",
+        "do_lower_case": True,
+        "tokenizer_name": "bert-large-uncased",
+    },
+    "megatron-bert-345m-cased": {
+        "config": CONFIGS["345m"],
+        "checkpoint": "https://api.ngc.nvidia.com/v2/models/nvidia/megatron_bert_345m/versions/v0.1_cased/files/release/mp_rank_00/model_optim_rng.pt",
+        "vocab": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-vocab.txt",
+        "do_lower_case": False,
+        "tokenizer_name": "bert-large-cased",
+    },
+    "megatron-bert-uncased": {
+        "config": None,
+        "checkpoint": None,
+        "vocab": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-vocab.txt",
+        "do_lower_case": True,
+        "tokenizer_name": "bert-large-uncased",
+    },
+    "megatron-bert-cased": {
+        "config": None,
+        "checkpoint": None,
+        "vocab": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-vocab.txt",
+        "do_lower_case": False,
+        "tokenizer_name": "bert-large-cased",
+    },
+    "biomegatron-bert-345m-uncased": {
+        "config": CONFIGS["345m"],
+        "checkpoint": "https://api.ngc.nvidia.com/v2/models/nvidia/biomegatron345muncased/versions/0/files/MegatronBERT.pt",
+        "vocab": "https://api.ngc.nvidia.com/v2/models/nvidia/biomegatron345muncased/versions/0/files/vocab.txt",
+        "do_lower_case": True,
+        "tokenizer_name": "bert-large-uncased",
+    },
+    "biomegatron-bert-345m-cased": {
+        "config": CONFIGS["345m"],
+        "checkpoint": "https://api.ngc.nvidia.com/v2/models/nvidia/biomegatron345mcased/versions/0/files/MegatronBERT.pt",
+        "vocab": "https://api.ngc.nvidia.com/v2/models/nvidia/biomegatron345mcased/versions/0/files/vocab.txt",
+        "do_lower_case": False,
+        "tokenizer_name": "bert-large-cased",
+    },
+}
+
+
+def get_megatron_lm_model(
+    pretrained_model_name: str,
+    config_dict: Optional[dict] = None,
+    config_file: Optional[str] = None,
+    checkpoint_file: Optional[str] = None,
+    vocab_file: Optional[str] = None,
+) -> Tuple[MegatronBertEncoder, str]:
+    """
+    Returns MegatronBertEncoder and a default or user specified path to the checkpoint file
+
+    Args:
+        pretrained_mode_name: model name from MEGATRON_CONFIG_MAP
+            for example: megatron-bert-cased
+        config_dict: model configuration parameters
+        config_file: path to model configuration file. Takes precedence over config_dict if both supplied.
+        checkpoint_file: path to checkpoint file or directory if using model parallel.
+        vocab_file: path to vocab file
+
+    Returns:
+        model: MegatronBertEncoder
+        checkpoint_file: path to checkpoint file or directory
+    """
+    config = None
+    # get default config and checkpoint
+    if config_file:
+        with open(config_file) as f:
+            config = json.load(f)
+            # replace dashes with underscores in config keys
+            fixed_config = {}
+            for key in config.keys():
+                fixed_key = key.replace("-", "_")
+                if fixed_key == 'max_seq_length':
+                    fixed_key = 'max_position_embeddings'
+                fixed_config[fixed_key] = config[key]
+            # 'vocab_size" no longer used.
+            if 'vocab_size' in fixed_config:
+                fixed_config.pop('vocab_size')
+            config = fixed_config
+    elif config_dict:
+        config = config_dict
+    elif pretrained_model_name in get_megatron_lm_models_list():
+        config = get_megatron_config(pretrained_model_name)
+    else:
+        raise ValueError(f"{pretrained_model_name} is not supported")
+
+    if config is None:
+        raise ValueError(f"config_file or config_dict is required for {pretrained_model_name}")
+
+    if not checkpoint_file:
+        checkpoint_file = get_megatron_checkpoint(pretrained_model_name)
+
+    if not vocab_file:
+        vocab_file = get_megatron_vocab_file(pretrained_model_name)
+
+    app_state = AppState()
+    if app_state.model_parallel_size is not None and app_state.model_parallel_rank is not None:
+        # model parallel already known from .nemo restore
+        model_parallel_size = app_state.model_parallel_size
+        model_parallel_rank = app_state.model_parallel_rank
+    elif os.path.isdir(checkpoint_file):
+        # starting training from megatron-lm checkpoint
+        mp_ranks = glob.glob(os.path.join(checkpoint_file, 'mp_rank*'))
+        model_parallel_size = len(mp_ranks)
+        app_state.model_parallel_size = model_parallel_size
+        logging.info(
+            (
+                f'restore_path: {checkpoint_file} is a directory. '
+                f'Assuming megatron model parallelism with '
+                f'model_parallel_size: {model_parallel_size}'
+            )
+        )
+        # try to get local rank from global
+        local_rank = None
+        try:
+            local_rank = int(os.environ['LOCAL_RANK'])
+        except:
+            logging.info('Global variable LOCAL_RANK not yet specified')
+        if local_rank is not None:
+            app_state.local_rank = local_rank
+        else:
+            # if local is None then we are on the main process
+            local_rank = 0
+        model_parallel_rank = compute_model_parallel_rank(local_rank, model_parallel_size)
+        app_state.model_parallel_rank = model_parallel_rank
+    else:
+        model_parallel_size = None
+        model_parallel_rank = None
+
+    model = MegatronBertEncoder(
+        model_name=pretrained_model_name,
+        config=config,
+        vocab_file=vocab_file,
+        model_parallel_size=model_parallel_size,
+        model_parallel_rank=model_parallel_rank,
+    )
+
+    return model, checkpoint_file
+
+
+def compute_model_parallel_rank(local_rank, model_parallel_size):
+    return local_rank % model_parallel_size
+
+
+def get_megatron_lm_models_list() -> List[str]:
+    """
+    Returns the list of supported Megatron-LM models
+    """
+    return list(MEGATRON_CONFIG_MAP.keys())
+
+
+def get_megatron_config(pretrained_model_name: str) -> Dict[str, int]:
+    """
+    Returns Megatron-LM model config file
+
+    Args:
+        pretrained_model_name (str): pretrained model name
+
+    Returns:
+        config (dict): contains model configuration: number of hidden layers, number of attention heads, etc
+    """
+    _check_megatron_name(pretrained_model_name)
+    return MEGATRON_CONFIG_MAP[pretrained_model_name]["config"]
+
+
+def _check_megatron_name(pretrained_model_name: str) -> None:
+    megatron_model_list = get_megatron_lm_models_list()
+    if pretrained_model_name not in megatron_model_list:
+        raise ValueError(f'For Megatron-LM models, choose from the following list: {megatron_model_list}')
+
+
+def get_megatron_vocab_file(pretrained_model_name: str) -> str:
+    """
+    Gets vocabulary file from cache or downloads it
+
+    Args:
+        pretrained_model_name: pretrained model name
+
+    Returns:
+        path: path to the vocab file
+    """
+    _check_megatron_name(pretrained_model_name)
+    url = MEGATRON_CONFIG_MAP[pretrained_model_name]["vocab"]
+
+    path = os.path.join(MEGATRON_CACHE, pretrained_model_name + "_vocab")
+    path = _download(path, url)
+    return path
+
+
+def get_megatron_checkpoint(pretrained_model_name: str) -> str:
+    """
+    Gets checkpoint file from cache or downloads it
+    Args:
+        pretrained_model_name: pretrained model name
+    Returns:
+        path: path to model checkpoint
+    """
+    _check_megatron_name(pretrained_model_name)
+    url = MEGATRON_CONFIG_MAP[pretrained_model_name]["checkpoint"]
+    path = os.path.join(MEGATRON_CACHE, pretrained_model_name)
+    return _download(path, url)
+
+
+def _download(path: str, url: str):
+    """
+    Gets a file from cache or downloads it
+
+    Args:
+        path: path to the file in cache
+        url: url to the file
+    Returns:
+        path: path to the file in cache
+    """
+    if url is None:
+        return None
+
+    if not os.path.exists(path):
+        master_device = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
+        if not os.path.exists(path):
+            if master_device:
+                os.makedirs(MEGATRON_CACHE, exist_ok=True)
+                logging.info(f"Downloading from {url}")
+                wget.download(url, path)
+            # wait until the master process downloads the file and writes it to the cache dir
+            if torch.distributed.is_initialized():
+                torch.distributed.barrier()
+
+    return path
+
+
+def is_lower_cased_megatron(pretrained_model_name):
+    """
+    Returns if the megatron is cased or uncased
+
+    Args:
+        pretrained_model_name (str): pretrained model name
+    Returns:
+        do_lower_cased (bool): whether the model uses lower cased data
+    """
+    _check_megatron_name(pretrained_model_name)
+    return MEGATRON_CONFIG_MAP[pretrained_model_name]["do_lower_case"]
+
+
+def get_megatron_tokenizer(pretrained_model_name: str):
+    """
+    Takes a pretrained_model_name for megatron such as "megatron-bert-cased" and returns the according 
+    tokenizer name for tokenizer instantiating.
+
+    Args:
+        pretrained_model_name: pretrained_model_name for megatron such as "megatron-bert-cased"
+    Returns: 
+        tokenizer name for tokenizer instantiating
+    """
+    _check_megatron_name(pretrained_model_name)
+    return MEGATRON_CONFIG_MAP[pretrained_model_name]["tokenizer_name"]

tts/nemo/collections/nlp/modules/common/sequence_classifier.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 Dict, Optional
+
+from nemo.collections.common.parts import MultiLayerPerceptron
+from nemo.collections.nlp.modules.common.classifier import Classifier
+from nemo.core.classes import typecheck
+from nemo.core.neural_types import LogitsType, LogprobsType, NeuralType
+
+__all__ = ['SequenceClassifier']
+
+
+class SequenceClassifier(Classifier):
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        if not self.log_softmax:
+            return {"logits": NeuralType(('B', 'D'), LogitsType())}
+        else:
+            return {"log_probs": NeuralType(('B', 'D'), LogprobsType())}
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_classes: int,
+        num_layers: int = 2,
+        activation: str = 'relu',
+        log_softmax: bool = True,
+        dropout: float = 0.0,
+        use_transformer_init: bool = True,
+        idx_conditioned_on: int = 0,
+    ):
+        """
+        Initializes the SequenceClassifier module.
+        Args:
+            hidden_size: the hidden size of the mlp head on the top of the encoder
+            num_classes: number of the classes to predict
+            num_layers: number of the linear layers of the mlp head on the top of the encoder
+            activation: type of activations between layers of the mlp head
+            log_softmax: applies the log softmax on the output
+            dropout: the dropout used for the mlp head
+            use_transformer_init: initializes the weights with the same approach used in Transformer
+            idx_conditioned_on: index of the token to use as the sequence representation for the classification task, default is the first token
+        """
+        super().__init__(hidden_size=hidden_size, dropout=dropout)
+        self.log_softmax = log_softmax
+        self._idx_conditioned_on = idx_conditioned_on
+        self.mlp = MultiLayerPerceptron(
+            hidden_size=hidden_size,
+            num_classes=num_classes,
+            num_layers=num_layers,
+            activation=activation,
+            log_softmax=log_softmax,
+        )
+        self.post_init(use_transformer_init=use_transformer_init)
+
+    @typecheck()
+    def forward(self, hidden_states):
+        hidden_states = self.dropout(hidden_states)
+        logits = self.mlp(hidden_states[:, self._idx_conditioned_on])
+        return logits

tts/nemo/collections/nlp/modules/common/sequence_regression.py

@@ -0,0 +1,69 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, Optional
+
+from torch import Tensor
+
+from nemo.collections.common.parts import MultiLayerPerceptron
+from nemo.collections.nlp.modules.common.classifier import Classifier
+from nemo.core.classes import typecheck
+from nemo.core.neural_types import NeuralType, RegressionValuesType
+
+__all__ = ['SequenceRegression']
+
+
+class SequenceRegression(Classifier):
+    """
+    Args:
+        hidden_size: the hidden size of the mlp head on the top of the encoder
+        num_layers: number of the linear layers of the mlp head on the top of the encoder
+        activation: type of activations between layers of the mlp head
+        dropout: the dropout used for the mlp head
+        use_transformer_init: initializes the weights with the same approach used in Transformer
+        idx_conditioned_on: index of the token to use as the sequence representation for the classification task,
+            default is the first token
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {"preds": NeuralType(tuple('B'), RegressionValuesType())}
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_layers: int = 2,
+        activation: str = 'relu',
+        dropout: float = 0.0,
+        use_transformer_init: bool = True,
+        idx_conditioned_on: int = 0,
+    ):
+        """ Initializes the SequenceRegression module. """
+        super().__init__(hidden_size=hidden_size, dropout=dropout)
+        self._idx_conditioned_on = idx_conditioned_on
+        self.mlp = MultiLayerPerceptron(
+            hidden_size, num_classes=1, num_layers=num_layers, activation=activation, log_softmax=False,
+        )
+        self.post_init(use_transformer_init=use_transformer_init)
+
+    @typecheck()
+    def forward(self, hidden_states: Tensor) -> Tensor:
+        """ Forward pass through the module.
+
+        Args:
+            hidden_states: hidden states for each token in a sequence, for example, BERT module output
+        """
+        hidden_states = self.dropout(hidden_states)
+        preds = self.mlp(hidden_states[:, self._idx_conditioned_on])
+        return preds.view(-1)

tts/nemo/collections/nlp/modules/common/sequence_token_classifier.py

@@ -0,0 +1,80 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, Optional
+
+from nemo.collections.common.parts import MultiLayerPerceptron
+from nemo.collections.nlp.modules.common.classifier import Classifier
+from nemo.core.classes import typecheck
+from nemo.core.neural_types import LogitsType, NeuralType
+
+__all__ = ['SequenceTokenClassifier']
+
+
+class SequenceTokenClassifier(Classifier):
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {
+            "intent_logits": NeuralType(('B', 'D'), LogitsType()),
+            "slot_logits": NeuralType(('B', 'T', 'D'), LogitsType()),
+        }
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_intents: int,
+        num_slots: int,
+        num_layers: int = 2,
+        activation: str = 'relu',
+        log_softmax: bool = False,
+        dropout: float = 0.0,
+        use_transformer_init: bool = True,
+    ):
+        """
+        Initializes the SequenceTokenClassifier module, could be used for tasks that train sequence and
+        token classifiers jointly, for example, for intent detection and slot tagging task.
+        Args:
+            hidden_size: hidden size of the mlp head on the top of the encoder
+            num_intents: number of the intents to predict
+            num_slots: number of the slots to predict
+            num_layers: number of the linear layers of the mlp head on the top of the encoder
+            activation: type of activations between layers of the mlp head
+            log_softmax: applies the log softmax on the output
+            dropout: the dropout used for the mlp head
+            use_transformer_init: initializes the weights with the same approach used in Transformer
+        """
+        super().__init__(hidden_size=hidden_size, dropout=dropout)
+        self.intent_mlp = MultiLayerPerceptron(
+            hidden_size=hidden_size,
+            num_classes=num_intents,
+            num_layers=num_layers,
+            activation=activation,
+            log_softmax=log_softmax,
+        )
+        self.slot_mlp = MultiLayerPerceptron(
+            hidden_size=hidden_size,
+            num_classes=num_slots,
+            num_layers=num_layers,
+            activation=activation,
+            log_softmax=log_softmax,
+        )
+        self.post_init(use_transformer_init=use_transformer_init)
+
+    @typecheck()
+    def forward(self, hidden_states):
+        hidden_states = self.dropout(hidden_states)
+        # intent is classified by first hidden position
+        intent_logits = self.intent_mlp(hidden_states[:, 0])
+        slot_logits = self.slot_mlp(hidden_states)
+        return intent_logits, slot_logits

tts/nemo/collections/nlp/modules/common/token_classifier.py

@@ -0,0 +1,165 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from typing import Dict, Optional
+
+from omegaconf.omegaconf import MISSING
+from torch import nn as nn
+
+from nemo.collections.common.parts import MultiLayerPerceptron
+from nemo.collections.nlp.modules.common.classifier import Classifier
+from nemo.core.classes import typecheck
+from nemo.core.neural_types import LogitsType, LogprobsType, NeuralType
+
+__all__ = ['BertPretrainingTokenClassifier', 'TokenClassifier']
+
+ACT2FN = {"gelu": nn.functional.gelu, "relu": nn.functional.relu}
+
+
+@dataclass
+class TokenClassifierConfig:
+    num_layers: int = 1
+    activation: str = 'relu'
+    log_softmax: bool = True
+    dropout: float = 0.0
+    use_transformer_init: bool = True
+
+
+class TokenClassifier(Classifier):
+    """
+    A module to perform token level classification tasks such as Named entity recognition.
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """
+        Returns definitions of module output ports.
+        """
+        if not self.log_softmax:
+            return {"logits": NeuralType(('B', 'T', 'C'), LogitsType())}
+        else:
+            return {"log_probs": NeuralType(('B', 'T', 'C'), LogprobsType())}
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_classes: int,
+        num_layers: int = 1,
+        activation: str = 'relu',
+        log_softmax: bool = True,
+        dropout: float = 0.0,
+        use_transformer_init: bool = True,
+    ) -> None:
+
+        """
+        Initializes the Token Classifier module.
+
+        Args:
+            hidden_size: the size of the hidden dimension
+            num_classes: number of classes
+            num_layers: number of fully connected layers in the multilayer perceptron (MLP)
+            activation: activation to usee between fully connected layers in the MLP
+            log_softmax: whether to apply softmax to the output of the MLP
+            dropout: dropout to apply to the input hidden states
+            use_transformer_init: whether to initialize the weights of the classifier head with the same approach used in Transformer
+        """
+        super().__init__(hidden_size=hidden_size, dropout=dropout)
+        self.log_softmax = log_softmax
+        self.mlp = MultiLayerPerceptron(
+            hidden_size, num_classes, num_layers=num_layers, activation=activation, log_softmax=log_softmax
+        )
+        self.post_init(use_transformer_init=use_transformer_init)
+
+    @typecheck()
+    def forward(self, hidden_states):
+        """
+        Performs the forward step of the module.
+        Args:
+            hidden_states: batch of hidden states (for example, from the BERT encoder module)
+                [BATCH_SIZE x SEQ_LENGTH x HIDDEN_SIZE]
+        Returns: logits value for each class [BATCH_SIZE x SEQ_LENGTH x NUM_CLASSES]
+        """
+        hidden_states = self.dropout(hidden_states)
+        logits = self.mlp(hidden_states)
+        return logits
+
+
+class BertPretrainingTokenClassifier(Classifier):
+    """
+    A module to perform token level classification tasks for Bert pretraining.
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """
+        Returns definitions of module output ports.
+        """
+        if not self.log_softmax:
+            return {"logits": NeuralType(('B', 'T', 'C'), LogitsType())}
+        else:
+            return {"log_probs": NeuralType(('B', 'T', 'C'), LogprobsType())}
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_classes: int,
+        num_layers: int = 1,
+        activation: str = 'relu',
+        log_softmax: bool = True,
+        dropout: float = 0.0,
+        use_transformer_init: bool = True,
+    ) -> None:
+
+        """
+        Initializes the Token Classifier module.
+
+        Args:
+            hidden_size: the size of the hidden dimension
+            num_classes: number of classes
+            num_layers: number of fully connected layers in the multilayer perceptron (MLP)
+            activation: activation to usee between fully connected layers in the MLP
+            log_softmax: whether to apply softmax to the output of the MLP
+            dropout: dropout to apply to the input hidden states
+            use_transformer_init: whether to initialize the weights of the classifier head with the same approach used in Transformer
+        """
+        super().__init__(hidden_size=hidden_size, dropout=dropout)
+
+        self.log_softmax = log_softmax
+
+        if activation not in ACT2FN:
+            raise ValueError(f'activation "{activation}" not found')
+        self.dense = nn.Linear(hidden_size, hidden_size)
+        self.act = ACT2FN[activation]
+        self.norm = nn.LayerNorm(hidden_size, eps=1e-12)
+        self.mlp = MultiLayerPerceptron(
+            hidden_size, num_classes, num_layers=num_layers, activation=activation, log_softmax=log_softmax
+        )
+        self.post_init(use_transformer_init=use_transformer_init)
+
+    @typecheck()
+    def forward(self, hidden_states):
+        """
+        Performs the forward step of the module.
+        Args:
+            hidden_states: batch of hidden states (for example, from the BERT encoder module)
+                [BATCH_SIZE x SEQ_LENGTH x HIDDEN_SIZE]
+        Returns: logits value for each class [BATCH_SIZE x SEQ_LENGTH x NUM_CLASSES]
+        """
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.act(hidden_states)
+        transform = self.norm(hidden_states)
+        logits = self.mlp(transform)
+        return logits

tts/nemo/collections/nlp/modules/common/tokenizer_utils.py

@@ -0,0 +1,144 @@
+# 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.path
+from dataclasses import MISSING, dataclass
+from os import path
+from typing import Dict, List, Optional
+
+import nemo
+from nemo.collections.common.tokenizers.char_tokenizer import CharTokenizer
+from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer
+from nemo.collections.common.tokenizers.word_tokenizer import WordTokenizer
+from nemo.collections.common.tokenizers.youtokentome_tokenizer import YouTokenToMeTokenizer
+from nemo.collections.nlp.modules.common.huggingface.huggingface_utils import get_huggingface_pretrained_lm_models_list
+from nemo.collections.nlp.modules.common.lm_utils import get_pretrained_lm_models_list
+from nemo.collections.nlp.modules.common.megatron.megatron_utils import get_megatron_tokenizer
+from nemo.utils import logging
+
+__all__ = ['get_tokenizer', 'get_tokenizer_list']
+
+
+def get_tokenizer_list() -> List[str]:
+    """
+    Returns all all supported tokenizer names
+    """
+    s = set(get_pretrained_lm_models_list())
+    s.update(set(get_huggingface_pretrained_lm_models_list(include_external=True)))
+    return ["sentencepiece", "char", "word"] + list(s)
+
+
+@dataclass
+class TokenizerConfig:
+    library: str = MISSING
+    tokenizer_model: Optional[str] = None
+    vocab_size: Optional[int] = None
+    vocab_file: Optional[str] = None
+    special_tokens: Optional[Dict[str, str]] = None
+    bpe_dropout: Optional[float] = 0.0
+    coverage: Optional[float] = 0.999
+    training_sample_size: Optional[int] = None
+
+
+def get_tokenizer(
+    tokenizer_name: str,
+    tokenizer_model: Optional[str] = None,
+    vocab_file: Optional[str] = None,
+    special_tokens: Optional[Dict[str, str]] = None,
+    use_fast: Optional[bool] = False,
+    bpe_dropout: Optional[float] = 0.0,
+):
+    """
+    Args:
+        tokenizer_name: sentencepiece or pretrained model from the hugging face list,
+            for example: bert-base-cased
+            To see the list of all HuggingFace pretrained models, use: nemo_nlp.modules.common.get_huggingface_pretrained_lm_models_list()
+        tokenizer_model: tokenizer model file of sentencepiece or youtokentome
+        special_tokens: dict of special tokens
+        vocab_file: path to vocab file
+        use_fast: (only for HuggingFace AutoTokenizer) set to True to use fast HuggingFace tokenizer
+        bpe_dropout: (only supported by YTTM tokenizer) BPE dropout tries to corrupt the standard segmentation procedure of BPE to help
+            model better learn word compositionality and become robust to segmentation errors. 
+            It has emperically been shown to improve inference time BLEU scores.
+    """
+    if special_tokens is None:
+        special_tokens_dict = {}
+    else:
+        special_tokens_dict = special_tokens
+
+    if 'megatron' in tokenizer_name:
+        if vocab_file is None:
+            vocab_file = nemo.collections.nlp.modules.common.megatron.megatron_utils.get_megatron_vocab_file(
+                tokenizer_name
+            )
+        tokenizer_name = get_megatron_tokenizer(tokenizer_name)
+
+    if tokenizer_name == 'sentencepiece':
+        return nemo.collections.common.tokenizers.sentencepiece_tokenizer.SentencePieceTokenizer(
+            model_path=tokenizer_model, special_tokens=special_tokens, legacy=True
+        )
+    elif tokenizer_name == 'yttm':
+        return YouTokenToMeTokenizer(model_path=tokenizer_model, bpe_dropout=bpe_dropout)
+    elif tokenizer_name == 'word':
+        return WordTokenizer(vocab_file=vocab_file, **special_tokens_dict)
+    elif tokenizer_name == 'char':
+        return CharTokenizer(vocab_file=vocab_file, **special_tokens_dict)
+
+    return AutoTokenizer(
+        pretrained_model_name=tokenizer_name, vocab_file=vocab_file, **special_tokens_dict, use_fast=use_fast
+    )
+
+
+def get_nmt_tokenizer(
+    library: str = 'yttm',
+    model_name: Optional[str] = None,
+    tokenizer_model: Optional[str] = None,
+    vocab_file: Optional[str] = None,
+    special_tokens: Optional[Dict[str, str]] = None,
+    use_fast: Optional[bool] = False,
+    bpe_dropout: Optional[float] = 0.0,
+):
+    """
+    Args:
+        model_name: if using a pretrained model from NeMo or HuggingFace
+        tokenizer_model: tokenizer model file of sentencepiece or youtokentome
+        special_tokens: dict of special tokens
+        vocab_file: path to vocab file
+        use_fast: (only for HuggingFace AutoTokenizer) set to True to use fast HuggingFace tokenizer
+        bpe_dropout: (only supported by YTTM tokenizer) BPE dropout tries to corrupt the standard segmentation procedure of BPE to help
+            model better learn word compositionality and become robust to segmentation errors. 
+            It has emperically been shown to improve inference time BLEU scores.        
+    """
+    if special_tokens is None:
+        special_tokens_dict = {}
+    else:
+        special_tokens_dict = special_tokens
+
+    if library == 'yttm':
+        logging.info(f'Getting YouTokenToMeTokenizer with model: {tokenizer_model}.')
+        return YouTokenToMeTokenizer(model_path=tokenizer_model, bpe_dropout=bpe_dropout)
+    elif library == 'huggingface':
+        logging.info(f'Getting HuggingFace AutoTokenizer with pretrained_model_name: {model_name}')
+        return AutoTokenizer(
+            pretrained_model_name=model_name, vocab_file=vocab_file, **special_tokens_dict, use_fast=use_fast
+        )
+    elif library == 'sentencepiece':
+        logging.info(f'Getting SentencePiece with model: {model_name}')
+        return nemo.collections.common.tokenizers.sentencepiece_tokenizer.SentencePieceTokenizer(
+            model_path=tokenizer_model, special_tokens=special_tokens_dict
+        )
+    else:
+        raise NotImplementedError(
+            'Currently we only support "yttm", "huggingface", and "sentencepiece" tokenizer library.'
+        )

tts/nemo/collections/nlp/modules/common/transformer/__init__.py

@@ -0,0 +1,18 @@
+# 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.nlp.modules.common.transformer.transformer_decoders import *
+from nemo.collections.nlp.modules.common.transformer.transformer_encoders import *
+from nemo.collections.nlp.modules.common.transformer.transformer_generators import *
+from nemo.collections.nlp.modules.common.transformer.transformer_modules import *

tts/nemo/collections/nlp/modules/common/transformer/transformer.py

@@ -0,0 +1,248 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+from omegaconf.omegaconf import MISSING
+
+from nemo.collections.nlp.modules.common.decoder_module import DecoderModule
+from nemo.collections.nlp.modules.common.encoder_module import EncoderModule
+from nemo.collections.nlp.modules.common.transformer.transformer_decoders import TransformerDecoder
+from nemo.collections.nlp.modules.common.transformer.transformer_encoders import TransformerEncoder
+from nemo.collections.nlp.modules.common.transformer.transformer_modules import TransformerEmbedding
+from nemo.core.classes.common import typecheck
+from nemo.core.classes.exportable import Exportable
+
+# @dataclass
+# class TransformerConfig:
+#     # named model arguments
+#     library: str = 'nemo'
+#     model_name: Optional[str] = None
+#     pretrained: bool = False
+
+
+@dataclass
+class NeMoTransformerConfig:
+    # must be configured by the user
+    hidden_size: int = MISSING
+    num_layers: int = MISSING
+    inner_size: int = MISSING
+    num_attention_heads: int = MISSING
+
+    # embedding
+    max_sequence_length: int = 512
+    num_token_types: int = 2
+    embedding_dropout: float = 0.0
+    learn_positional_encodings: bool = False
+
+    # transformer
+    ffn_dropout: float = 0.0
+    attn_score_dropout: float = 0.0
+    attn_layer_dropout: float = 0.0
+    hidden_act: str = 'relu'
+    pre_ln: bool = False
+    pre_ln_final_layer_norm: bool = True
+
+    # named model arguments
+    library: str = 'nemo'
+    model_name: Optional[str] = None
+    pretrained: bool = False
+
+
+@dataclass
+class NeMoTransformerEncoderConfig(NeMoTransformerConfig):
+    mask_future: bool = False
+
+
+class TransformerEncoderNM(EncoderModule, Exportable):
+    def __init__(
+        self,
+        vocab_size: int,
+        hidden_size: int,
+        num_layers: int,
+        inner_size: int,
+        num_attention_heads: int,
+        max_sequence_length: int = 512,
+        num_token_types: int = 2,
+        embedding_dropout: float = 0.0,
+        learn_positional_encodings: bool = False,
+        ffn_dropout: float = 0.0,
+        attn_score_dropout: float = 0.0,
+        attn_layer_dropout: float = 0.0,
+        hidden_act: str = 'relu',
+        mask_future: bool = False,
+        pre_ln: bool = False,
+        pre_ln_final_layer_norm: bool = True,
+    ):
+        super().__init__()
+
+        self._vocab_size = vocab_size
+        self._hidden_size = hidden_size
+        self._max_sequence_length = max_sequence_length
+
+        self._embedding = TransformerEmbedding(
+            vocab_size=self._vocab_size,
+            hidden_size=self._hidden_size,
+            max_sequence_length=max_sequence_length,
+            num_token_types=num_token_types,
+            embedding_dropout=embedding_dropout,
+            learn_positional_encodings=learn_positional_encodings,
+        )
+
+        self._encoder = TransformerEncoder(
+            hidden_size=self._hidden_size,
+            num_layers=num_layers,
+            inner_size=inner_size,
+            num_attention_heads=num_attention_heads,
+            ffn_dropout=ffn_dropout,
+            attn_score_dropout=attn_score_dropout,
+            attn_layer_dropout=attn_layer_dropout,
+            hidden_act=hidden_act,
+            mask_future=mask_future,
+            pre_ln=pre_ln,
+            pre_ln_final_layer_norm=pre_ln_final_layer_norm,
+        )
+
+    @typecheck()
+    def forward(self, input_ids, encoder_mask):
+        embeddings = self._embedding(input_ids=input_ids)
+        encoder_hidden_states = self._encoder(encoder_states=embeddings, encoder_mask=encoder_mask)
+        return encoder_hidden_states
+
+    @property
+    def hidden_size(self):
+        return self._hidden_size
+
+    @property
+    def vocab_size(self):
+        return self._vocab_size
+
+    @property
+    def max_sequence_length(self):
+        return self._max_sequence_length
+
+    @property
+    def embedding(self):
+        return self._embedding
+
+    @property
+    def encoder(self):
+        return self._encoder
+
+    def input_example(self):
+        """
+        Generates input examples for tracing etc.
+        Returns:
+            A tuple of input examples.
+        """
+        sample = next(self.parameters())
+        input_ids = torch.randint(low=0, high=2048, size=(2, 16), device=sample.device)
+        encoder_mask = torch.randint(low=0, high=1, size=(2, 16), device=sample.device)
+        return tuple([input_ids, encoder_mask])
+
+
+class TransformerDecoderNM(DecoderModule, Exportable):
+    def __init__(
+        self,
+        vocab_size: int,
+        hidden_size: int,
+        num_layers: int,
+        inner_size: int,
+        num_attention_heads: int,
+        max_sequence_length: int = 512,
+        num_token_types: int = 2,
+        embedding_dropout: float = 0.0,
+        learn_positional_encodings: bool = False,
+        ffn_dropout: float = 0.0,
+        attn_score_dropout: float = 0.0,
+        attn_layer_dropout: float = 0.0,
+        hidden_act: str = 'relu',
+        pre_ln: bool = False,
+        pre_ln_final_layer_norm: bool = True,
+    ):
+        super().__init__()
+
+        self._vocab_size = vocab_size
+        self._hidden_size = hidden_size
+        self._max_sequence_length = max_sequence_length
+
+        self._embedding = TransformerEmbedding(
+            vocab_size=self.vocab_size,
+            hidden_size=self.hidden_size,
+            max_sequence_length=max_sequence_length,
+            num_token_types=num_token_types,
+            embedding_dropout=embedding_dropout,
+            learn_positional_encodings=learn_positional_encodings,
+        )
+
+        self._decoder = TransformerDecoder(
+            hidden_size=self.hidden_size,
+            num_layers=num_layers,
+            inner_size=inner_size,
+            num_attention_heads=num_attention_heads,
+            ffn_dropout=ffn_dropout,
+            attn_score_dropout=attn_score_dropout,
+            attn_layer_dropout=attn_layer_dropout,
+            hidden_act=hidden_act,
+            pre_ln=pre_ln,
+            pre_ln_final_layer_norm=pre_ln_final_layer_norm,
+        )
+
+    @typecheck()
+    def forward(self, input_ids, decoder_mask, encoder_embeddings, encoder_mask):
+        decoder_embeddings = self._embedding(input_ids=input_ids)
+        decoder_hidden_states = self._decoder(
+            decoder_states=decoder_embeddings,
+            decoder_mask=decoder_mask,
+            encoder_states=encoder_embeddings,
+            encoder_mask=encoder_mask,
+        )
+        return decoder_hidden_states
+
+    @property
+    def hidden_size(self):
+        return self._hidden_size
+
+    @property
+    def vocab_size(self):
+        return self._vocab_size
+
+    @property
+    def max_sequence_length(self):
+        return self._max_sequence_length
+
+    @property
+    def embedding(self):
+        return self._embedding
+
+    @property
+    def decoder(self):
+        return self._decoder
+
+    def input_example(self):
+        """
+        Generates input examples for tracing etc.
+        Returns:
+            A tuple of input examples.
+        """
+        sample = next(self.parameters())
+        input_ids = torch.randint(low=0, high=2048, size=(2, 16), device=sample.device)
+        encoder_mask = torch.randint(low=0, high=1, size=(2, 16), device=sample.device)
+        return tuple([input_ids, encoder_mask, self._embedding(input_ids), encoder_mask])
+
+    def _prepare_for_export(self, **kwargs):
+        self._decoder.diagonal = None
+        super()._prepare_for_export(**kwargs)

tts/nemo/collections/nlp/modules/common/transformer/transformer_decoders.py

@@ -0,0 +1,200 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from omegaconf.omegaconf import MISSING
+
+from nemo.collections.common.parts import form_attention_mask
+from nemo.collections.nlp.modules.common.transformer.transformer_modules import MultiHeadAttention, PositionWiseFF
+
+__all__ = ["TransformerDecoder"]
+
+
+class TransformerDecoderBlock(nn.Module):
+    """
+    Building block of Transformer decoder.
+
+    Args:
+        hidden_size: size of the embeddings in the model, also known as d_model
+        inner_size: number of neurons in the intermediate part of feed-forward
+            net, usually is (4-8 x hidden_size) in the papers
+        num_attention_heads: number of heads in multi-head attention
+        attn_score_dropout: probability of dropout applied to attention scores
+        attn_layer_dropout: probability of dropout applied to the output of the
+            attention layers, but before layer normalization
+        ffn_dropout: probability of dropout applied to FFN output
+        hidden_act: activation function used between two linear layers in FFN
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        inner_size: int,
+        num_attention_heads: int = 1,
+        attn_score_dropout: float = 0.0,
+        attn_layer_dropout: float = 0.0,
+        ffn_dropout: float = 0.0,
+        hidden_act: str = "relu",
+        pre_ln: bool = False,
+    ):
+        super().__init__()
+        self.pre_ln = pre_ln
+        self.layer_norm_1 = nn.LayerNorm(hidden_size, eps=1e-5)
+        self.first_sub_layer = MultiHeadAttention(
+            hidden_size, num_attention_heads, attn_score_dropout, attn_layer_dropout
+        )
+        self.layer_norm_2 = nn.LayerNorm(hidden_size, eps=1e-5)
+        self.second_sub_layer = MultiHeadAttention(
+            hidden_size, num_attention_heads, attn_score_dropout, attn_layer_dropout
+        )
+        self.layer_norm_3 = nn.LayerNorm(hidden_size, eps=1e-5)
+        self.third_sub_layer = PositionWiseFF(hidden_size, inner_size, ffn_dropout, hidden_act)
+
+    def forward_preln(self, decoder_query, decoder_mask, decoder_keys, encoder_states, encoder_mask):
+        """
+        Pre-LayerNorm block
+        Order of operations: LN -> Self-Attn -> Residual -> LN -> Cross-Attn -> Residual -> LN -> FFN
+        """
+        residual = decoder_query
+        decoder_query = self.layer_norm_1(decoder_query)
+        decoder_keys = self.layer_norm_1(decoder_keys)
+        self_attn_output = self.first_sub_layer(decoder_query, decoder_keys, decoder_keys, decoder_mask)
+        self_attn_output += residual
+
+        residual = self_attn_output
+        self_attn_output = self.layer_norm_2(self_attn_output)
+        enc_dec_attn_output = self.second_sub_layer(self_attn_output, encoder_states, encoder_states, encoder_mask)
+        enc_dec_attn_output += residual
+
+        residual = enc_dec_attn_output
+        enc_dec_attn_output = self.layer_norm_3(enc_dec_attn_output)
+        output_states = self.third_sub_layer(enc_dec_attn_output)
+        output_states += residual
+
+        return output_states
+
+    def forward_postln(self, decoder_query, decoder_mask, decoder_keys, encoder_states, encoder_mask):
+        """
+        Post-LayerNorm block
+        Order of operations: Self-Attn -> Residual -> LN -> Cross-Attn -> Residual -> LN -> FFN -> Residual -> LN
+        """
+        self_attn_output = self.first_sub_layer(decoder_query, decoder_keys, decoder_keys, decoder_mask)
+        self_attn_output += decoder_query
+        self_attn_output = self.layer_norm_1(self_attn_output)
+
+        enc_dec_attn_output = self.second_sub_layer(self_attn_output, encoder_states, encoder_states, encoder_mask)
+        enc_dec_attn_output += self_attn_output
+        enc_dec_attn_output = self.layer_norm_2(enc_dec_attn_output)
+
+        output_states = self.third_sub_layer(enc_dec_attn_output)
+        output_states += enc_dec_attn_output
+        return self.layer_norm_3(output_states)
+
+    def forward(self, decoder_query, decoder_mask, decoder_keys, encoder_states, encoder_mask):
+        if self.pre_ln:
+            return self.forward_preln(decoder_query, decoder_mask, decoder_keys, encoder_states, encoder_mask)
+        else:
+            return self.forward_postln(decoder_query, decoder_mask, decoder_keys, encoder_states, encoder_mask)
+
+
+class TransformerDecoder(nn.Module):
+    def __init__(
+        self,
+        num_layers: int,
+        hidden_size: int,
+        inner_size: int,
+        num_attention_heads: int = 1,
+        attn_score_dropout: float = 0.0,
+        attn_layer_dropout: float = 0.0,
+        ffn_dropout: float = 0.0,
+        hidden_act: str = "relu",
+        pre_ln: bool = False,
+        pre_ln_final_layer_norm: bool = True,
+    ):
+        super().__init__()
+
+        if pre_ln and pre_ln_final_layer_norm:
+            self.final_layer_norm = nn.LayerNorm(hidden_size, eps=1e-5)
+        else:
+            self.final_layer_norm = None
+
+        layer = TransformerDecoderBlock(
+            hidden_size,
+            inner_size,
+            num_attention_heads,
+            attn_score_dropout,
+            attn_layer_dropout,
+            ffn_dropout,
+            hidden_act,
+            pre_ln,
+        )
+        self.layers = nn.ModuleList([copy.deepcopy(layer) for _ in range(num_layers)])
+        self.diagonal = 0
+
+    def _get_memory_states(self, decoder_states, decoder_mems_list=None, i=0):
+        if decoder_mems_list is not None:
+            memory_states = torch.cat((decoder_mems_list[i], decoder_states), dim=1)
+        else:
+            memory_states = decoder_states
+        return memory_states
+
+    def forward(
+        self, decoder_states, decoder_mask, encoder_states, encoder_mask, decoder_mems_list=None, return_mems=False
+    ):
+        """
+        Args:
+            decoder_states: output of the embedding layer (B x L_dec x H)
+            decoder_mask: decoder inputs mask (B x L_dec)
+            encoder_states: output of the encoder (B x L_enc x H)
+            encoder_mask: encoder inputs mask (B x L_enc)
+            decoder_mems_list: list of the cached decoder hidden states
+                for fast autoregressive generation which will be used instead
+                of decoder_states as keys and values if not None
+            return_mems: bool, whether to return outputs of all decoder layers
+                or the last layer only
+        """
+        decoder_attn_mask = form_attention_mask(decoder_mask, diagonal=self.diagonal)
+        encoder_attn_mask = form_attention_mask(encoder_mask)
+        memory_states = self._get_memory_states(decoder_states, decoder_mems_list, 0)
+        cached_mems_list = [memory_states]
+
+        for i, layer in enumerate(self.layers):
+            decoder_states = layer(decoder_states, decoder_attn_mask, memory_states, encoder_states, encoder_attn_mask)
+            memory_states = self._get_memory_states(decoder_states, decoder_mems_list, i + 1)
+            cached_mems_list.append(memory_states)
+
+        if self.final_layer_norm is not None:
+            decoder_states = self.final_layer_norm(decoder_states)
+            memory_states = self._get_memory_states(decoder_states, decoder_mems_list, i + 1)
+            cached_mems_list.append(memory_states)
+
+        if return_mems:
+            return cached_mems_list
+        else:
+            return cached_mems_list[-1]
+
+    def input_example(self):
+        """
+        Generates input examples for tracing etc.
+        Returns:
+            A tuple of input examples.
+        """
+        sample = next(self.parameters())
+        input_ids = torch.randint(low=0, high=2048, size=(2, 16, 1024), device=sample.device)
+        encoder_mask = torch.randint(low=0, high=1, size=(2, 16), device=sample.device)
+        return tuple([input_ids, encoder_mask, input_ids, encoder_mask])

tts/nemo/collections/nlp/modules/common/transformer/transformer_encoders.py

@@ -0,0 +1,176 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from omegaconf.omegaconf import MISSING
+
+from nemo.collections.common.parts import form_attention_mask
+from nemo.collections.nlp.modules.common.transformer.transformer_modules import MultiHeadAttention, PositionWiseFF
+
+__all__ = ["TransformerEncoder"]
+
+
+class TransformerEncoderBlock(nn.Module):
+    """
+    Building block of Transformer encoder.
+
+    Args:
+        hidden_size: size of the embeddings in the model, also known as d_model
+        inner_size: number of neurons in the intermediate part of feed-forward
+            net, usually is (4-8 x hidden_size) in the papers
+        num_attention_heads: number of heads in multi-head attention
+        attn_score_dropout: probability of dropout applied to attention scores
+        attn_layer_dropout: probability of dropout applied to the output of the
+            attention layers, but before layer normalization
+        ffn_dropout: probability of dropout applied to FFN output
+        hidden_act: activation function used between two linear layers in FFN
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        inner_size: int,
+        num_attention_heads: int = 1,
+        attn_score_dropout: float = 0.0,
+        attn_layer_dropout: float = 0.0,
+        ffn_dropout: float = 0.0,
+        hidden_act: str = "relu",
+        pre_ln: bool = False,
+    ):
+        super().__init__()
+        self.pre_ln = pre_ln
+        self.layer_norm_1 = nn.LayerNorm(hidden_size, eps=1e-5)
+        self.first_sub_layer = MultiHeadAttention(
+            hidden_size, num_attention_heads, attn_score_dropout, attn_layer_dropout
+        )
+        self.layer_norm_2 = nn.LayerNorm(hidden_size, eps=1e-5)
+        self.second_sub_layer = PositionWiseFF(hidden_size, inner_size, ffn_dropout, hidden_act)
+
+    def forward_preln(self, encoder_query, encoder_mask, encoder_keys):
+        """
+        Pre-LayerNorm block
+        Order of operations: LN -> Self-Attn -> Residual -> LN -> Cross-Attn -> Residual -> LN -> FFN
+        """
+        residual = encoder_query
+        encoder_query = self.layer_norm_1(encoder_query)
+        encoder_keys = self.layer_norm_1(encoder_keys)
+        self_attn_output = self.first_sub_layer(encoder_query, encoder_keys, encoder_keys, encoder_mask)
+        self_attn_output += residual
+
+        residual = self_attn_output
+        self_attn_output = self.layer_norm_2(self_attn_output)
+        output_states = self.second_sub_layer(self_attn_output)
+        output_states += residual
+
+        return output_states
+
+    def forward_postln(self, encoder_query, encoder_mask, encoder_keys):
+        """
+        Post-LayerNorm block
+        Order of operations: Self-Attn -> Residual -> LN -> Cross-Attn -> Residual -> LN -> FFN -> Residual -> LN
+        """
+        self_attn_output = self.first_sub_layer(encoder_query, encoder_keys, encoder_keys, encoder_mask)
+        self_attn_output += encoder_query
+        self_attn_output = self.layer_norm_1(self_attn_output)
+
+        output_states = self.second_sub_layer(self_attn_output)
+        output_states += self_attn_output
+        output_states = self.layer_norm_2(output_states)
+
+        return output_states
+
+    def forward(self, encoder_query, encoder_mask, encoder_keys):
+        if self.pre_ln:
+            return self.forward_preln(encoder_query, encoder_mask, encoder_keys)
+        else:
+            return self.forward_postln(encoder_query, encoder_mask, encoder_keys)
+
+
+class TransformerEncoder(nn.Module):
+    def __init__(
+        self,
+        num_layers: int,
+        hidden_size: int,
+        inner_size: int,
+        mask_future: bool = False,
+        num_attention_heads: int = 1,
+        attn_score_dropout: float = 0.0,
+        attn_layer_dropout: float = 0.0,
+        ffn_dropout: float = 0.0,
+        hidden_act: str = "relu",
+        pre_ln: bool = False,
+        pre_ln_final_layer_norm: bool = True,
+    ):
+        super().__init__()
+
+        if pre_ln and pre_ln_final_layer_norm:
+            self.final_layer_norm = nn.LayerNorm(hidden_size, eps=1e-5)
+        else:
+            self.final_layer_norm = None
+
+        layer = TransformerEncoderBlock(
+            hidden_size,
+            inner_size,
+            num_attention_heads,
+            attn_score_dropout,
+            attn_layer_dropout,
+            ffn_dropout,
+            hidden_act,
+            pre_ln,
+        )
+        self.layers = nn.ModuleList([copy.deepcopy(layer) for _ in range(num_layers)])
+        self.diag = 0 if mask_future else None
+
+    def _get_memory_states(self, encoder_states, encoder_mems_list=None, i=0):
+        if encoder_mems_list is not None:
+            memory_states = torch.cat((encoder_mems_list[i], encoder_states), dim=1)
+        else:
+            memory_states = encoder_states
+        return memory_states
+
+    def forward(self, encoder_states, encoder_mask, encoder_mems_list=None, return_mems=False):
+        """
+        Args:
+            encoder_states: output of the embedding_layer (B x L_enc x H)
+            encoder_mask: encoder inputs mask (B x L_enc)
+            encoder_mems_list: list of the cached encoder hidden states
+                for fast autoregressive generation which will be used instead
+                of encoder_states as keys and values if not None
+            return_mems: bool, whether to return outputs of all encoder layers
+                or the last layer only
+        """
+
+        encoder_attn_mask = form_attention_mask(encoder_mask, self.diag)
+
+        memory_states = self._get_memory_states(encoder_states, encoder_mems_list, 0)
+        cached_mems_list = [memory_states]
+
+        for i, layer in enumerate(self.layers):
+            encoder_states = layer(encoder_states, encoder_attn_mask, memory_states)
+            memory_states = self._get_memory_states(encoder_states, encoder_mems_list, i + 1)
+            cached_mems_list.append(memory_states)
+
+        if self.final_layer_norm is not None:
+            encoder_states = self.final_layer_norm(encoder_states)
+            memory_states = self._get_memory_states(encoder_states, encoder_mems_list, i + 1)
+            cached_mems_list.append(memory_states)
+
+        if return_mems:
+            return cached_mems_list
+        else:
+            return cached_mems_list[-1]

tts/nemo/collections/nlp/modules/common/transformer/transformer_generators.py

@@ -0,0 +1,541 @@
+# 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 contextlib import contextmanager
+
+import torch
+
+from nemo.collections.common.parts import NEG_INF, mask_padded_tokens
+
+__all__ = [
+    "GreedySequenceGenerator",
+    "TopKSequenceGenerator",
+    "BeamSearchSequenceGenerator",
+    "BeamSearchSequenceGeneratorWithLanguageModel",
+]
+
+
+class GreedySequenceGenerator:
+    """
+    Greedy sequence generator based on the decoder followed by log_softmax.
+
+    Args:
+        embedding: nn.Module, transforms input_ids into vector embeddings
+        decoder: nn.Module, takes embeddings and produces hidden_states
+        log_softmax: nn.Module, takes hidden_states and produces log_probs
+            which correspond to probability distribution of tokens (ids)
+        pad: index of padding token in the vocabulary
+        bos: index of beginning of sequence token in the vocabulary
+        eos: index of end of sequence token in the vocabulary
+        max_sequence_length: maximum allowed length for generated sequences
+        max_delta_length: in case of encoder-decoder generation (e.g. NMT),
+            forbids generated sequences to be longer than the length of
+            source sequences plus max_delta_length
+        batch_size: size of the batch of generated sequences if neither
+            source nor target starting sequences are provided
+    """
+
+    def __init__(
+        self,
+        embedding,
+        decoder,
+        log_softmax,
+        pad=0,
+        bos=1,
+        eos=2,
+        max_sequence_length=512,
+        max_delta_length=20,
+        batch_size=1,
+    ):
+        super().__init__()
+        self.embedding = embedding
+        self.decoder = decoder
+        self.log_softmax = log_softmax
+        self.pad, self.bos, self.eos = pad, bos, eos
+        self.max_seq_length = max_sequence_length
+        self.max_delta_len = max_delta_length
+        self.batch_size = batch_size
+
+    def _one_step_forward(
+        self,
+        decoder_input_ids=None,
+        encoder_hidden_states=None,
+        encoder_input_mask=None,
+        decoder_mems_list=None,
+        pos=0,
+    ):
+        """
+        One step of autoregressive output generation.
+
+        Args:
+            decoder_input_ids: starting sequence of tokens to generate from;
+                if None, generation will start from a batch of <bos> tokens
+            encoder_hidden_states: output of the encoder for conditional
+                sequence generation; if None, generator will use unconditional
+                mode (e.g., language modeling)
+            encoder_input_mask: input mask used in the encoder
+            decoder_mems_list: list of size num_layers with cached activations
+                of sequence (x[1], ..., x[k-1]) for fast generation of x[k]
+            pos: starting position in positional encoding
+        """
+
+        decoder_hidden_states = self.embedding.forward(decoder_input_ids, start_pos=pos)
+        decoder_input_mask = mask_padded_tokens(decoder_input_ids, self.pad).float()
+
+        if encoder_hidden_states is not None:
+            decoder_mems_list = self.decoder.forward(
+                decoder_hidden_states,
+                decoder_input_mask,
+                encoder_hidden_states,
+                encoder_input_mask,
+                decoder_mems_list,
+                return_mems=True,
+            )
+        else:
+            decoder_mems_list = self.decoder.forward(
+                decoder_hidden_states, decoder_input_mask, decoder_mems_list, return_mems=True
+            )
+        log_probs = self.log_softmax.forward(hidden_states=decoder_mems_list[-1][:, -1:])
+        return log_probs, decoder_mems_list
+
+    def _prepare_for_search(self, decoder_input_ids=None, encoder_hidden_states=None):
+        """
+        Helper function which defines starting sequence to begin generating
+        with and maximum allowed number of tokens to be generated.
+        """
+
+        decoder_parameter = next(self.decoder.parameters())
+        batch_size = self.batch_size
+
+        # for encoder-decoder generation, maximum length of generated sequence
+        # is min(max_sequence_length, src_len + max_delta_length)
+        if encoder_hidden_states is not None:
+            batch_size, src_len, _ = encoder_hidden_states.size()
+            max_seq_length = min(self.max_seq_length, src_len + self.max_delta_len)
+        else:
+            max_seq_length = self.max_seq_length
+
+        # if no input is provided, start with the batch of <bos> tokens
+        if decoder_input_ids is not None:
+            tgt = decoder_input_ids
+            batch_size, tgt_len = decoder_input_ids.size()
+        else:
+            tgt = torch.zeros(batch_size, 1).long().fill_(self.bos).to(decoder_parameter.device)
+            tgt_len = 1
+        max_generation_length = max_seq_length - tgt_len
+
+        return tgt, batch_size, max_generation_length
+
+    def _forward(self, decoder_input_ids=None, encoder_hidden_states=None, encoder_input_mask=None):
+
+        tgt, batch_size, max_generation_length = self._prepare_for_search(decoder_input_ids, encoder_hidden_states)
+
+        # pad profile tracks sequences ending with <eos> token to replace
+        # everything after <eos> with <pad> token
+        decoder_parameter = next(self.decoder.parameters())
+        pad_profile = torch.zeros(batch_size, 1).long().to(decoder_parameter.device)
+
+        decoder_mems_list = None
+        for i in range(max_generation_length):
+
+            log_probs, decoder_mems_list = self._one_step_forward(
+                tgt[:, -1:], encoder_hidden_states, encoder_input_mask, decoder_mems_list, i
+            )
+
+            next_tokens = torch.argmax(log_probs[:, -1], dim=-1, keepdim=True)
+            next_tokens = self.pad * pad_profile + next_tokens * (1 - pad_profile)
+            pad_profile = torch.max(pad_profile, (next_tokens == self.eos).long())
+            tgt = torch.cat((tgt, next_tokens), dim=-1)
+
+            # abort generation if all sequences end with <eos>
+            if pad_profile.sum() == batch_size:
+                break
+
+        return tgt
+
+    def __call__(self, decoder_input_ids=None, encoder_hidden_states=None, encoder_input_mask=None):
+        with self.as_frozen():
+            return self._forward(decoder_input_ids, encoder_hidden_states, encoder_input_mask)
+
+    def freeze(self) -> None:
+        """Freeze weights of embedding, decoder, and classification layers to prevent memory leak.
+        """
+        for param in self.embedding.parameters():
+            param.requires_grad = False
+        self.embedding.eval()
+        for param in self.decoder.parameters():
+            param.requires_grad = False
+        self.decoder.eval()
+        for param in self.log_softmax.parameters():
+            param.require_grad = False
+        self.log_softmax.eval()
+
+    def unfreeze(self) -> None:
+        """Unfreeze weights of embedding, decoder, and classification layers.
+        """
+        for param in self.embedding.parameters():
+            param.requires_grad = True
+        self.embedding.train()
+        for param in self.decoder.parameters():
+            param.requires_grad = True
+        self.decoder.train()
+        for param in self.log_softmax.parameters():
+            param.require_grad = True
+        self.log_softmax.train()
+
+    @contextmanager
+    def as_frozen(self):
+        """
+        Context manager which temporarily freezes embedding, decoder, and log_softmax modules,
+        yields control and finally unfreezes the modules.
+        """
+        self.freeze()
+
+        try:
+            yield
+        finally:
+            self.unfreeze()
+
+
+class TopKSequenceGenerator(GreedySequenceGenerator):
+    """
+    Top-k sequence generator based on the decoder followed by log_softmax.
+
+    Args:
+        *all args of GreedySequenceGenerator class
+        beam_size: size of the beam (parameter k in top-k)
+        temperature: temperature of top-k sampling, all logits are divided
+            by temperature before rescaling. High temperature leads to
+            uniform distribution, low leads to delta-like distribution.
+    Kwargs:
+        all remaining parameters of GreedySequenceGenerator class
+    """
+
+    def __init__(self, embedding, decoder, log_softmax, beam_size=1, temperature=1.0, **kwargs):
+        super().__init__(embedding, decoder, log_softmax, **kwargs)
+        self.beam_size = beam_size
+        self.temp = temperature
+
+    # @torch.no_grad()
+    def _one_step_forward(
+        self,
+        decoder_input_ids=None,
+        encoder_hidden_states=None,
+        encoder_input_mask=None,
+        decoder_mems_list=None,
+        pos=0,
+    ):
+        log_probs, decoder_mems_list = super()._one_step_forward(
+            decoder_input_ids, encoder_hidden_states, encoder_input_mask, decoder_mems_list, pos
+        )
+
+        batch_size, seq_len, vocab_size = log_probs.size()
+        scores, indices = torch.topk(log_probs, self.beam_size, dim=-1)
+
+        rescaled_logexp = torch.zeros_like(log_probs).scatter(-1, indices, scores.div(self.temp).exp())
+        probs = rescaled_logexp / rescaled_logexp.norm(1, -1, keepdim=True)
+
+        # We randomly sample next tokens from rescaled probability distribution
+        # over top-k candidates and return a binary tensor which indicates
+        # candidates that have been selected. We call this object
+        # `pseudo_log_probs` as genuine log_probs should have -infs instead of
+        # 0s and 0s instead of 1s.
+        ids = torch.multinomial(probs.view(-1, vocab_size), 1).view(-1, seq_len, 1)
+        pseudo_log_probs = torch.zeros_like(log_probs).scatter(-1, ids, 1.0)
+
+        return pseudo_log_probs, decoder_mems_list
+
+
+class BeamSearchSequenceGenerator(GreedySequenceGenerator):
+    def __init__(self, embedding, decoder, log_softmax, beam_size=1, len_pen=0, **kwargs):
+        """
+        Beam Search sequence generator based on the decoder followed by
+        log_softmax.
+
+        Args:
+            *all args of GreedySequenceGenerator class
+            beam_size: size of the beam
+            len_pen: length penalty parameter
+        Kwargs:
+            all remaining parameters of GreedySequenceGenerator class
+        """
+
+        super().__init__(embedding, decoder, log_softmax, **kwargs)
+        self.beam_size = beam_size
+        self.len_pen = len_pen
+
+    @staticmethod
+    def compute_len_penalty(lengths, alpha):
+        """Returns length penalty according to https://arxiv.org/pdf/1609.08144.pdf"""
+        return ((5 + lengths) / 6).pow(alpha)
+
+    def _forward(self, decoder_input_ids=None, encoder_hidden_states=None, encoder_input_mask=None):
+        tgt, batch_size, max_generation_length = self._prepare_for_search(decoder_input_ids, encoder_hidden_states)
+
+        # generate initial buffer of beam_size prefixes-hypotheses
+        log_probs, decoder_mems_list = self._one_step_forward(tgt, encoder_hidden_states, encoder_input_mask, None, 0)
+        scores, prefixes = torch.topk(log_probs.permute(0, 2, 1), self.beam_size, dim=1)
+        scores, prefixes = scores.view(-1, 1), prefixes.view(-1, 1)
+
+        # repeat init target prefixes and cached memory states beam_size times
+        prefixes = torch.cat((tgt.repeat(1, self.beam_size).view(-1, 1), prefixes), dim=1)
+        for j in range(len(decoder_mems_list)):
+            decoder_mems_list[j] = decoder_mems_list[j].repeat(self.beam_size, 1, 1)
+
+        # repeat source sequence beam_size times for beam search
+        if encoder_hidden_states is not None:
+            _, src_length, hidden_size = encoder_hidden_states.size()
+            encoder_input_mask = encoder_input_mask.repeat(1, self.beam_size).view(-1, src_length)
+            encoder_hidden_states = encoder_hidden_states.repeat(1, self.beam_size, 1).view(
+                -1, src_length, hidden_size
+            )
+        else:
+            hidden_size = decoder_mems_list[0].size(2)
+
+        # pad_profile tracks finished hypotheses to generate only <pad> tokens
+        # if <eos> or <pad> has been generated
+        pad_profile = torch.zeros_like(scores).long()
+
+        # prefixes_len tracks lengths of generated hypotheses to perform
+        # length penalty correction
+        prefixes_len = torch.zeros_like(scores).fill_(prefixes.size(1) + 1)
+
+        for i in range(max_generation_length):
+
+            # mask all finished hypotheses to exclude them from beam
+            pad_mask = pad_profile.repeat(1, self.beam_size)
+
+            # generate and score candidates for prefixes continuation
+            log_probs, decoder_mems_list = self._one_step_forward(
+                prefixes[:, -1:], encoder_hidden_states, encoder_input_mask, decoder_mems_list, i + 1
+            )
+            scores_i, prefixes_i = torch.topk(log_probs[:, -1, :], self.beam_size, dim=-1)
+
+            # for all prefixes ending with <eos> or <pad> replace generated
+            # continuations with <pad>
+            prefixes_i = self.pad * pad_mask + prefixes_i * (1 - pad_mask)
+
+            # force all hypotheses but one generated from already finished
+            # hypotheses to have extremely low score, so they will not be
+            # considered during beam re-ranking
+            pad_mask[:, 1:] = pad_mask[:, 1:] * NEG_INF
+            scores = scores + scores_i * (1 - pad_mask).to(scores.dtype)
+
+            # choose top-k hypotheses with length penalty applied
+            len_penalties = self.compute_len_penalty(prefixes_len, self.len_pen)
+            scores = scores / len_penalties
+            scores, indices_i = torch.topk(scores.view(-1, self.beam_size ** 2), self.beam_size, dim=1)
+            scores = scores.view(-1, 1) * len_penalties
+
+            # select prefixes which correspond to the chosen hypotheses
+            prefixes = prefixes.unsqueeze(1).repeat(1, self.beam_size, 1)
+            prefixes = torch.cat((prefixes, prefixes_i.unsqueeze(2)), dim=2)
+            prefixes = prefixes.view(batch_size, self.beam_size ** 2, -1)
+            p_len = prefixes.size(2)
+            prefixes_ids = indices_i.unsqueeze(2).repeat(1, 1, p_len)
+            prefixes = prefixes.gather(1, prefixes_ids).view(-1, p_len)
+
+            # reshuffle cached decoder memory states to restore the order
+            # of hypotheses broken after top-k selection
+            mems_ids = indices_i.unsqueeze(2).unsqueeze(3).repeat(1, 1, p_len - 1, hidden_size) // self.beam_size
+            for j in range(len(decoder_mems_list)):
+                decoder_mems_list[j] = (
+                    decoder_mems_list[j]
+                    .view(-1, self.beam_size, p_len - 1, hidden_size)
+                    .gather(1, mems_ids)
+                    .view(-1, p_len - 1, hidden_size)
+                )
+
+            # update prefixes_len and pad_profile
+            not_eos_pad = prefixes.ne(self.eos) & prefixes.ne(self.pad)
+            prefixes_len = 1 + not_eos_pad.sum(dim=1, keepdim=True).to(scores.dtype)
+            pad_profile = (~not_eos_pad[:, -1:]).long()
+
+            # if all hypotheses end with <eos> or <pad>, interrupt search
+            if pad_profile.sum() == batch_size * self.beam_size:
+                break
+
+        # select best performing hypotheses in each element of the batch
+        len_penalties = self.compute_len_penalty(prefixes_len, self.len_pen)
+        scores = scores / len_penalties
+        best_guesses = (
+            torch.argmax(scores.view(-1, self.beam_size), dim=1, keepdim=True).repeat(1, prefixes.size(1)).unsqueeze(1)
+        )
+        tgt = prefixes.view(batch_size, self.beam_size, -1).gather(1, best_guesses)
+
+        return tgt.squeeze(1)
+
+
+class BeamSearchSequenceGeneratorWithLanguageModel(GreedySequenceGenerator):
+    def __init__(
+        self, embedding, decoder, log_softmax, language_model, beam_size=1, len_pen=0, fusion_coef=0.0, **kwargs
+    ):
+        """
+        Beam Search sequence generator based on the decoder followed by log_softmax
+        with external language model fusion.
+
+        Args:
+            *all args of BeamSearchSequenceGenerator class
+            language_model: nemo TransformerLMModel
+            fusion_coef: coefficient before language model score, the resulting score is
+                score = log P_NMT(y|x) + fusion_coef * log P_LM(y|x)
+        Kwargs:
+            all remaining parameters of GreedySequenceGenerator class
+        """
+
+        super().__init__(embedding, decoder, log_softmax, **kwargs)
+        self.language_model = language_model
+        self.beam_size = beam_size
+        self.len_pen = len_pen
+        self.fusion_coef = fusion_coef
+
+    def _one_step_forward(
+        self,
+        decoder_input_ids=None,
+        encoder_hidden_states=None,
+        encoder_input_mask=None,
+        decoder_mems_list=None,
+        lm_mems_list=None,
+        pos=0,
+    ):
+
+        nmt_log_probs, decoder_mems_list = super()._one_step_forward(
+            decoder_input_ids, encoder_hidden_states, encoder_input_mask, decoder_mems_list, pos,
+        )
+        input_mask = mask_padded_tokens(decoder_input_ids, self.pad).float()
+        lm_hidden_states = self.language_model.encoder.embedding.forward(decoder_input_ids, start_pos=pos)
+
+        lm_mems_list = self.language_model.encoder.encoder.forward(
+            lm_hidden_states, input_mask, lm_mems_list, return_mems=True,
+        )
+        lm_log_probs = self.language_model.log_softmax.forward(hidden_states=lm_mems_list[-1][:, -1:])
+
+        log_probs = nmt_log_probs + self.fusion_coef * lm_log_probs
+
+        return log_probs, decoder_mems_list, lm_mems_list
+
+    @staticmethod
+    def compute_len_penalty(lengths, alpha):
+        """Returns length penalty according to https://arxiv.org/pdf/1609.08144.pdf"""
+        return ((5 + lengths) / 6).pow(alpha)
+
+    def _forward(self, decoder_input_ids=None, encoder_hidden_states=None, encoder_input_mask=None):
+
+        tgt, batch_size, max_generation_length = self._prepare_for_search(decoder_input_ids, encoder_hidden_states)
+
+        # generate initial buffer of beam_size prefixes-hypotheses
+        log_probs, decoder_mems_list, lm_mems_list = self._one_step_forward(
+            tgt, encoder_hidden_states, encoder_input_mask, None, None, 0
+        )
+        scores, prefixes = torch.topk(log_probs.permute(0, 2, 1), self.beam_size, dim=1)
+        scores, prefixes = scores.view(-1, 1), prefixes.view(-1, 1)
+
+        # repeat init target prefixes and cached memory states beam_size times
+        prefixes = torch.cat((tgt.repeat(1, self.beam_size).view(-1, 1), prefixes), dim=1)
+        for j in range(len(decoder_mems_list)):
+            decoder_mems_list[j] = decoder_mems_list[j].repeat(self.beam_size, 1, 1)
+        for j in range(len(lm_mems_list)):
+            lm_mems_list[j] = lm_mems_list[j].repeat(self.beam_size, 1, 1)
+
+        # repeat source sequence beam_size times for beam search
+        if encoder_hidden_states is not None:
+            _, src_length, hidden_size = encoder_hidden_states.size()
+            encoder_input_mask = encoder_input_mask.repeat(1, self.beam_size).view(-1, src_length)
+            encoder_hidden_states = encoder_hidden_states.repeat(1, self.beam_size, 1).view(
+                -1, src_length, hidden_size
+            )
+        else:
+            hidden_size = decoder_mems_list[0].size(2)
+        lm_hidden_size = lm_mems_list[0].size(2)
+
+        # pad_profile tracks finished hypotheses to generate only <pad> tokens
+        # if <eos> or <pad> has been generated
+        pad_profile = torch.zeros_like(scores).long()
+
+        # prefixes_len tracks lengths of generated hypotheses to perform
+        # length penalty correction
+        prefixes_len = torch.zeros_like(scores).fill_(prefixes.size(1) + 1)
+
+        for i in range(max_generation_length):
+
+            # mask all finished hypotheses to exclude them from beam
+            pad_mask = pad_profile.repeat(1, self.beam_size)
+
+            # generate and score candidates for prefixes continuation
+            log_probs, decoder_mems_list, lm_mems_list = self._one_step_forward(
+                prefixes[:, -1:], encoder_hidden_states, encoder_input_mask, decoder_mems_list, lm_mems_list, i + 1
+            )
+            scores_i, prefixes_i = torch.topk(log_probs[:, -1, :], self.beam_size, dim=-1)
+
+            # for all prefixes ending with <eos> or <pad> replace generated
+            # continuations with <pad>
+            prefixes_i = self.pad * pad_mask + prefixes_i * (1 - pad_mask)
+
+            # force all hypotheses but one generated from already finished
+            # hypotheses to have extremely low score, so they will not be
+            # considered during beam re-ranking
+            pad_mask[:, 1:] = pad_mask[:, 1:] * NEG_INF
+            scores = scores + scores_i * (1 - pad_mask).to(scores.dtype)
+
+            # choose top-k hypotheses with length penalty applied
+            len_penalties = self.compute_len_penalty(prefixes_len, self.len_pen)
+            scores = scores / len_penalties
+            scores, indices_i = torch.topk(scores.view(-1, self.beam_size ** 2), self.beam_size, dim=1)
+            scores = scores.view(-1, 1) * len_penalties
+
+            # select prefixes which correspond to the chosen hypotheses
+            prefixes = prefixes.unsqueeze(1).repeat(1, self.beam_size, 1)
+            prefixes = torch.cat((prefixes, prefixes_i.unsqueeze(2)), dim=2)
+            prefixes = prefixes.view(batch_size, self.beam_size ** 2, -1)
+            p_len = prefixes.size(2)
+            prefixes_ids = indices_i.unsqueeze(2).repeat(1, 1, p_len)
+            prefixes = prefixes.gather(1, prefixes_ids).view(-1, p_len)
+
+            # reshuffle cached decoder memory states to restore the order
+            # of hypotheses broken after top-k selection
+            mems_ids = indices_i.unsqueeze(2).unsqueeze(3).repeat(1, 1, p_len - 1, hidden_size) // self.beam_size
+            for j in range(len(decoder_mems_list)):
+                decoder_mems_list[j] = (
+                    decoder_mems_list[j]
+                    .view(-1, self.beam_size, p_len - 1, hidden_size)
+                    .gather(1, mems_ids)
+                    .view(-1, p_len - 1, hidden_size)
+                )
+            lm_mems_ids = indices_i.unsqueeze(2).unsqueeze(3).repeat(1, 1, p_len - 1, lm_hidden_size) // self.beam_size
+            for j in range(len(lm_mems_list)):
+                lm_mems_list[j] = (
+                    lm_mems_list[j]
+                    .view(-1, self.beam_size, p_len - 1, lm_hidden_size)
+                    .gather(1, lm_mems_ids)
+                    .view(-1, p_len - 1, lm_hidden_size)
+                )
+
+            # update prefixes_len and pad_profile
+            not_eos_pad = prefixes.ne(self.eos) & prefixes.ne(self.pad)
+            prefixes_len = 1 + not_eos_pad.sum(dim=1, keepdim=True).to(scores.dtype)
+            pad_profile = (~not_eos_pad[:, -1:]).long()
+
+            # if all hypotheses end with <eos> or <pad>, interrupt search
+            if pad_profile.sum() == batch_size * self.beam_size:
+                break
+
+        # select best performing hypotheses in each element of the batch
+        len_penalties = self.compute_len_penalty(prefixes_len, self.len_pen)
+        scores = scores / len_penalties
+        best_guesses = (
+            torch.argmax(scores.view(-1, self.beam_size), dim=1, keepdim=True).repeat(1, prefixes.size(1)).unsqueeze(1)
+        )
+        tgt = prefixes.view(batch_size, self.beam_size, -1).gather(1, best_guesses)
+
+        return tgt.squeeze(1)

tts/nemo/collections/nlp/modules/common/transformer/transformer_modules.py

@@ -0,0 +1,210 @@
+# Copyright 2018 The Google AI Language Team Authors and
+# The HuggingFace Inc. team.
+# 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 dataclasses import dataclass
+
+import torch
+from omegaconf.omegaconf import MISSING
+from torch import nn
+from torch.nn.functional import gelu
+
+__all__ = ["TransformerEmbedding"]
+
+
+class FixedPositionalEncoding(nn.Module):
+    """
+    Fixed positional encoding (embedding layer) from sine and cosine functions
+    of different frequencies according to https://arxiv.org/abs/1706.03762
+
+    Args:
+        hidden_size: size of the embeddings in the model, also known as d_model
+        max_sequence_length: maximum allowed length of the input sequence
+    """
+
+    def __init__(self, hidden_size, max_sequence_length=512):
+        super().__init__()
+
+        pos_enc = torch.zeros(max_sequence_length, hidden_size)
+        position = torch.arange(0.0, max_sequence_length).unsqueeze(1)
+        coef = -math.log(10000.0) / hidden_size
+        div_term = torch.exp(coef * torch.arange(0.0, hidden_size, 2))
+        pos_enc[:, 0::2] = torch.sin(position * div_term)
+        pos_enc[:, 1::2] = torch.cos(position * div_term)
+        pos_enc.div_(math.sqrt(hidden_size))
+        self.register_buffer('pos_enc', pos_enc)
+
+    def forward(self, position_ids):
+        return torch.embedding(self.pos_enc, position_ids)
+
+
+class TransformerEmbedding(nn.Module):
+    """
+    Embedding from token and position embeddings.
+    Optionally add token_type embedding (e.g. type of the sentence in BERT).
+
+    Args:
+        vocab_size: size of the vocabulary
+        hidden_size: size of the embeddings in the model, also known as d_model
+        max_sequence_length: maximum allowed length of the input sequence
+        num_token_types: number of different token types
+            (e.g. tokens of sentence A and tokens of sentence B in BERT)
+        embedding_dropout: probability of dropout applied to embeddings
+        learn_positional_encodings: whether to learn positional encodings or
+            use fixed (sine-cosine) ones
+    """
+
+    def __init__(
+        self,
+        vocab_size,
+        hidden_size,
+        max_sequence_length=512,
+        num_token_types=2,
+        embedding_dropout=0.0,
+        learn_positional_encodings=False,
+    ):
+        super().__init__()
+
+        self.max_sequence_length = max_sequence_length
+        self.token_embedding = nn.Embedding(vocab_size, hidden_size, padding_idx=0)
+        if learn_positional_encodings:
+            self.position_embedding = nn.Embedding(max_sequence_length, hidden_size)
+        else:
+            self.position_embedding = FixedPositionalEncoding(hidden_size, max_sequence_length)
+        if num_token_types > 0:
+            self.token_type_embedding = nn.Embedding(num_token_types, hidden_size)
+        self.layer_norm = nn.LayerNorm(hidden_size, eps=1e-5)
+        self.dropout = nn.Dropout(embedding_dropout)
+
+    def forward(self, input_ids, token_type_ids=None, start_pos=0):
+        seq_length = input_ids.size(1)
+        if seq_length > self.max_sequence_length:
+            raise ValueError(
+                f"Input sequence is longer than maximum allowed sequence length for positional encoding. "
+                f"Got {seq_length} and {self.max_sequence_length}"
+            )
+        position_ids = torch.arange(
+            start=start_pos, end=start_pos + seq_length, dtype=torch.long, device=input_ids.device
+        )
+        position_ids = position_ids.unsqueeze(0).expand_as(input_ids)
+
+        token_embeddings = self.token_embedding(input_ids)
+        position_embeddings = self.position_embedding(position_ids)
+        embeddings = token_embeddings + position_embeddings
+
+        if token_type_ids is not None:
+            token_type_embeddings = self.token_type_embedding(token_type_ids)
+            embeddings = embeddings + token_type_embeddings
+
+        embeddings = self.layer_norm(embeddings)
+        embeddings = self.dropout(embeddings)
+
+        return embeddings
+
+
+class MultiHeadAttention(nn.Module):
+    """
+    Multi-head scaled dot-product attention layer.
+
+    Args:
+        hidden_size: size of the embeddings in the model, also known as d_model
+        num_attention_heads: number of heads in multi-head attention
+        attn_score_dropout: probability of dropout applied to attention scores
+        attn_layer_dropout: probability of dropout applied to the output of the
+            whole layer, but before layer normalization
+    """
+
+    def __init__(self, hidden_size, num_attention_heads, attn_score_dropout=0.0, attn_layer_dropout=0.0):
+        super().__init__()
+        if hidden_size % num_attention_heads != 0:
+            raise ValueError(
+                "The hidden size (%d) is not a multiple of the number "
+                "of attention heads (%d)" % (hidden_size, num_attention_heads)
+            )
+        self.hidden_size = hidden_size
+        self.num_attention_heads = num_attention_heads
+        self.attn_head_size = int(hidden_size / num_attention_heads)
+        self.attn_scale = math.sqrt(math.sqrt(self.attn_head_size))
+
+        self.query_net = nn.Linear(hidden_size, hidden_size)
+        self.key_net = nn.Linear(hidden_size, hidden_size)
+        self.value_net = nn.Linear(hidden_size, hidden_size)
+        self.out_projection = nn.Linear(hidden_size, hidden_size)
+
+        self.attn_dropout = nn.Dropout(attn_score_dropout)
+        self.layer_dropout = nn.Dropout(attn_layer_dropout)
+
+    def transpose_for_scores(self, x):
+        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attn_head_size)
+        x = x.view(*new_x_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def forward(self, queries, keys, values, attention_mask):
+
+        # attention_mask is needed to hide the tokens which correspond to [PAD]
+        # in the case of BERT, or to hide the future tokens in the case of
+        # vanilla language modeling and translation
+        query = self.query_net(queries)
+        key = self.key_net(keys)
+        value = self.value_net(values)
+        query = self.transpose_for_scores(query) / self.attn_scale
+        key = self.transpose_for_scores(key) / self.attn_scale
+        value = self.transpose_for_scores(value)
+
+        # for numerical stability we pre-divide query and key by sqrt(sqrt(d))
+        attention_scores = torch.matmul(query, key.transpose(-1, -2))
+        if attention_mask is not None:
+            attention_scores = attention_scores + attention_mask.to(attention_scores.dtype)
+        attention_probs = torch.softmax(attention_scores, dim=-1)
+        attention_probs = self.attn_dropout(attention_probs)
+
+        context = torch.matmul(attention_probs, value)
+        context = context.permute(0, 2, 1, 3).contiguous()
+        new_context_shape = context.size()[:-2] + (self.hidden_size,)
+        context = context.view(*new_context_shape)
+
+        # output projection
+        output_states = self.out_projection(context)
+        output_states = self.layer_dropout(output_states)
+        return output_states
+
+
+class PositionWiseFF(nn.Module):
+    """
+    Position-wise feed-forward network of Transformer block.
+
+    Args:
+        hidden_size: size of the embeddings in the model, also known as d_model
+        inner_size: number of neurons in the intermediate part of feed-forward
+            net, usually is (4-8 x hidden_size) in the papers
+        ffn_dropout: probability of dropout applied to net output
+        hidden_act: activation function used between two linear layers
+    """
+
+    def __init__(self, hidden_size, inner_size, ffn_dropout=0.0, hidden_act="relu"):
+        super().__init__()
+        self.dense_in = nn.Linear(hidden_size, inner_size)
+        self.dense_out = nn.Linear(inner_size, hidden_size)
+        self.layer_dropout = nn.Dropout(ffn_dropout)
+        ACT2FN = {"gelu": gelu, "relu": torch.relu}
+        self.act_fn = ACT2FN[hidden_act]
+
+    def forward(self, hidden_states):
+        output_states = self.dense_in(hidden_states)
+        output_states = self.act_fn(output_states)
+        output_states = self.dense_out(output_states)
+        output_states = self.layer_dropout(output_states)
+        return output_states

tts/nemo/collections/nlp/modules/common/transformer/transformer_utils.py

@@ -0,0 +1,118 @@
+# 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, Union
+
+from omegaconf.dictconfig import DictConfig
+
+from nemo.collections.nlp.modules.common.huggingface.huggingface_decoder import HuggingFaceDecoderModule
+from nemo.collections.nlp.modules.common.huggingface.huggingface_encoder import HuggingFaceEncoderModule
+from nemo.collections.nlp.modules.common.transformer.transformer import TransformerDecoderNM, TransformerEncoderNM
+
+
+def get_nemo_transformer(
+    model_name: Optional[str] = None,
+    pretrained: bool = False,
+    config_dict: Optional[Union[dict, DictConfig]] = None,
+    encoder: bool = True,
+    pre_ln_final_layer_norm: bool = True,
+) -> Union[TransformerEncoderNM, TransformerDecoderNM]:
+    """Returns NeMo transformer.
+    The following configurations are mandatory:
+        vocab_size: int
+        hidden_size: int
+        num_layers: int
+        inner_size: int
+    and must be specified if using config_dict.
+
+    Args:
+        model_name (Optional[str]): model name to download from NGC
+        pretrained: (bool): False will instantiate the named model architecture with random weights.
+        config_dict (Optional[dict], optional): model configuration parameters. Defaults to None.
+        config_file (Optional[str], optional): path to json file containing model configuration. Defaults to None.
+        checkpoint_file (Optional[str], optional): load weights from path to local checkpoint. Defaults to None.
+        encoder (bool, optional): True will use EncoderTransformerNM, False will use DecoderTransformerNM. Defaults to True.
+    """
+    if model_name is not None:
+        raise ValueError(f'NeMo transformers cannot be loaded from NGC yet. model_name should be None')
+
+    if pretrained:
+        raise ValueError(f'NeMo transformers cannot be loaded from NGC yet. pretrained should be False')
+
+    cfg = None
+
+    if not pretrained:
+        assert (
+            config_dict.get('vocab_size') is not None
+            and config_dict.get('hidden_size') is not None
+            and config_dict.get('num_layers') is not None
+            and config_dict.get('inner_size') is not None
+        ), f'Using config_dict: {config_dict}. vocab_size, hidden_size, num_layers, and inner_size must are mandatory arguments'
+
+        cfg = config_dict
+
+    if encoder:
+        model = TransformerEncoderNM(
+            vocab_size=cfg.get('vocab_size'),
+            hidden_size=cfg.get('hidden_size'),
+            num_layers=cfg.get('num_layers'),
+            inner_size=cfg.get('inner_size'),
+            max_sequence_length=cfg.get('max_sequence_length', 512),
+            embedding_dropout=cfg.get('embedding_dropout', 0.0),
+            learn_positional_encodings=cfg.get('learn_positional_encodings', False),
+            num_attention_heads=cfg.get('num_attention_heads'),
+            ffn_dropout=cfg.get('ffn_dropout', 0.0),
+            attn_score_dropout=cfg.get('attn_score_dropout', 0.0),
+            attn_layer_dropout=cfg.get('attn_layer_dropout', 0.0),
+            hidden_act=cfg.get('hidden_act', 'relu'),
+            mask_future=cfg.get('mask_future', False),
+            pre_ln=cfg.get('pre_ln', False),
+            pre_ln_final_layer_norm=pre_ln_final_layer_norm,
+            num_token_types=cfg.get('num_token_types', 2),
+        )
+    else:
+        model = TransformerDecoderNM(
+            vocab_size=cfg.get('vocab_size'),
+            hidden_size=cfg.get('hidden_size'),
+            num_layers=cfg.get('num_layers'),
+            inner_size=cfg.get('inner_size'),
+            max_sequence_length=cfg.get('max_sequence_length', 512),
+            embedding_dropout=cfg.get('embedding_dropout', 0.0),
+            learn_positional_encodings=cfg.get('learn_positional_encodings', False),
+            num_attention_heads=cfg.get('num_attention_heads'),
+            ffn_dropout=cfg.get('ffn_dropout', 0.0),
+            attn_score_dropout=cfg.get('attn_score_dropout', 0.0),
+            attn_layer_dropout=cfg.get('attn_layer_dropout', 0.0),
+            hidden_act=cfg.get('hidden_act', 'relu'),
+            pre_ln=cfg.get('pre_ln', False),
+            pre_ln_final_layer_norm=pre_ln_final_layer_norm,
+            num_token_types=cfg.get('num_token_types', 2),
+        )
+
+    return model
+
+
+def get_huggingface_transformer(
+    model_name: Optional[str] = None,
+    pretrained: bool = False,
+    config_dict: Optional[Union[dict, DictConfig]] = None,
+    encoder: bool = True,
+) -> Union[HuggingFaceEncoderModule, HuggingFaceDecoderModule]:
+    if encoder:
+        model = HuggingFaceEncoderModule(model_name, pretrained, config_dict)
+    else:
+        model = HuggingFaceDecoderModule(model_name, pretrained, config_dict)
+
+    return model