[Megatron-LM] Add megatron-lm gpt2 (huggingface#9560)

* Add modeling_mlm_gpt2 init
soeque1 · Feb 22, 2021 · 245c99c · 245c99c
1 parent f51188c
commit 245c99c
Show file tree

Hide file tree

Showing 3 changed files with 1,293 additions and 0 deletions.
diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py
@@ -591,6 +591,15 @@
             "load_tf_weights_in_gpt2",
         ]
     )
+    _import_structure["models.gpt2"].extend(
+        [
+            "MLMGPT2DoubleHeadsModel",
+            "MLMGPT2ForSequenceClassification",
+            "MLMGPT2LMHeadModel",
+            "MLMGPT2Model",
+            "MLMGPT2PreTrainedModel",
+        ]
+    )
     _import_structure["models.layoutlm"].extend(
         [
             "LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST",
@@ -1668,6 +1677,10 @@
             GPT2PreTrainedModel,
             load_tf_weights_in_gpt2,
         )
+        from .models.gpt2 import (
+            MLMGPT2LMHeadModel,
+            MLMGPT2PreTrainedModel
+        )
         from .models.layoutlm import (
             LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST,
             LayoutLMForMaskedLM,

diff --git a/src/transformers/models/gpt2/__init__.py b/src/transformers/models/gpt2/__init__.py
@@ -39,6 +39,13 @@
         "GPT2PreTrainedModel",
         "load_tf_weights_in_gpt2",
     ]
+    _import_structure["modeling_mlm_gpt2"] = [
+        "MLMGPT2DoubleHeadsModel",
+        "MLMGPT2ForSequenceClassification",
+        "MLMGPT2LMHeadModel",
+        "MLMGPT2Model",
+        "MLMGPT2PreTrainedModel",
+    ]
 
 if is_tf_available():
     _import_structure["modeling_tf_gpt2"] = [
@@ -70,6 +77,13 @@
             load_tf_weights_in_gpt2,
         )
 
+        from .modeling_mlm_gpt2 import (
+            MLMGPT2LMHeadModel,
+            MLMGPT2Model,
+            MLMGPT2PreTrainedModel,
+        )
+
+
     if is_tf_available():
         from .modeling_tf_gpt2 import (
             TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST,