Introduce code from Rust monorepo. This is a mess

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