Add TF implementation of GPT-J

docs/source/index.mdx

@@ -205,7 +205,7 @@ Flax), PyTorch, and/or TensorFlow.
 |     Funnel Transformer      |       ✅       |       ✅       |       ✅        |         ✅         |      ❌      |
 |            GLPN             |       ❌       |       ❌       |       ✅        |         ❌         |      ❌      |
 |           GPT Neo           |       ❌       |       ❌       |       ✅        |         ❌         |      ✅      |
-|            GPT-J            |       ❌       |       ❌       |       ✅        |         ❌         |      ✅      |
+|            GPT-J            |       ❌       |       ❌       |       ✅        |         ✅         |      ✅      |
 |           Hubert            |       ❌       |       ❌       |       ✅        |         ✅         |      ❌      |
 |           I-BERT            |       ❌       |       ❌       |       ✅        |         ❌         |      ❌      |
 |          ImageGPT           |       ❌       |       ❌       |       ✅        |         ❌         |      ❌      |

docs/source/model_doc/gptj.mdx

@@ -130,6 +130,26 @@ model.
 [[autodoc]] GPTJForQuestionAnswering
     - forward
 
+## TFGPTJModel
+
+[[autodoc]] TFGPTJModel
+    - call
+
+## TFGPTJForCausalLM
+
+[[autodoc]] TFGPTJForCausalLM
+    - call
+
+## TFGPTJForSequenceClassification
+
+[[autodoc]] TFGPTJForSequenceClassification
+    - call
+
+## TFGPTJForQuestionAnswering
+
+[[autodoc]] TFGPTJForQuestionAnswering
+    - call
+
 ## FlaxGPTJModel
 
 [[autodoc]] FlaxGPTJModel

src/transformers/__init__.py

@@ -1929,6 +1929,15 @@
             "TFGPT2PreTrainedModel",
         ]
     )
+    _import_structure["models.gptj"].extend(
+        [
+            "TFGPTJForCausalLM",
+            "TFGPTJForQuestionAnswering",
+            "TFGPTJForSequenceClassification",
+            "TFGPTJModel",
+            "TFGPTJPreTrainedModel",
+        ]
+    )
     _import_structure["models.hubert"].extend(
         [
             "TF_HUBERT_PRETRAINED_MODEL_ARCHIVE_LIST",
@@ -4003,6 +4012,13 @@
             TFGPT2Model,
             TFGPT2PreTrainedModel,
         )
+        from .models.gptj import (
+            TFGPTJForCausalLM,
+            TFGPTJForQuestionAnswering,
+            TFGPTJForSequenceClassification,
+            TFGPTJModel,
+            TFGPTJPreTrainedModel,
+        )
         from .models.hubert import (
             TF_HUBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
             TFHubertForCTC,

src/transformers/models/auto/modeling_tf_auto.py

@@ -52,6 +52,7 @@
         ("bert", "TFBertModel"),
         ("openai-gpt", "TFOpenAIGPTModel"),
         ("gpt2", "TFGPT2Model"),
+        ("gptj", "TFGPTJModel"),
         ("mobilebert", "TFMobileBertModel"),
         ("transfo-xl", "TFTransfoXLModel"),
         ("xlnet", "TFXLNetModel"),
@@ -123,6 +124,7 @@
         ("bert", "TFBertForMaskedLM"),
         ("openai-gpt", "TFOpenAIGPTLMHeadModel"),
         ("gpt2", "TFGPT2LMHeadModel"),
+        ("gptj", "TFGPTJForCausalLM"),
         ("mobilebert", "TFMobileBertForMaskedLM"),
         ("transfo-xl", "TFTransfoXLLMHeadModel"),
         ("xlnet", "TFXLNetLMHeadModel"),
@@ -146,6 +148,7 @@
         ("bert", "TFBertLMHeadModel"),
         ("openai-gpt", "TFOpenAIGPTLMHeadModel"),
         ("gpt2", "TFGPT2LMHeadModel"),
+        ("gptj", "TFGPTJForCausalLM"),
         ("transfo-xl", "TFTransfoXLLMHeadModel"),
         ("xlnet", "TFXLNetLMHeadModel"),
         ("xlm", "TFXLMWithLMHeadModel"),
@@ -239,6 +242,7 @@
         ("tapas", "TFTapasForSequenceClassification"),
         ("funnel", "TFFunnelForSequenceClassification"),
         ("gpt2", "TFGPT2ForSequenceClassification"),
+        ("gptj", "TFGPTJForSequenceClassification"),
         ("mpnet", "TFMPNetForSequenceClassification"),
         ("openai-gpt", "TFOpenAIGPTForSequenceClassification"),
         ("transfo-xl", "TFTransfoXLForSequenceClassification"),
@@ -267,6 +271,7 @@
         ("xlm", "TFXLMForQuestionAnsweringSimple"),
         ("electra", "TFElectraForQuestionAnswering"),
         ("funnel", "TFFunnelForQuestionAnswering"),
+        ("gptj", "TFGPTJForQuestionAnswering"),
         ("mpnet", "TFMPNetForQuestionAnswering"),
     ]
 )

src/transformers/models/gptj/__init__.py

@@ -17,7 +17,7 @@
 # limitations under the License.
 from typing import TYPE_CHECKING
 
-from ...utils import _LazyModule, is_flax_available, is_torch_available
+from ...utils import _LazyModule, is_flax_available, is_tf_available, is_torch_available
 
 
 _import_structure = {
@@ -34,6 +34,15 @@
         "GPTJPreTrainedModel",
     ]
 
+if is_tf_available():
+    _import_structure["modeling_tf_gptj"] = [
+        "TFGPTJForCausalLM",
+        "TFGPTJForQuestionAnswering",
+        "TFGPTJForSequenceClassification",
+        "TFGPTJModel",
+        "TFGPTJPreTrainedModel",
+    ]
+
 if is_flax_available():
     _import_structure["modeling_flax_gptj"] = [
         "FlaxGPTJForCausalLM",
@@ -55,6 +64,15 @@
             GPTJPreTrainedModel,
         )
 
+    if is_tf_available():
+        from .modeling_tf_gptj import (
+            TFGPTJForCausalLM,
+            TFGPTJForQuestionAnswering,
+            TFGPTJForSequenceClassification,
+            TFGPTJModel,
+            TFGPTJPreTrainedModel,
+        )
+
     if is_flax_available():
         from .modeling_flax_gptj import FlaxGPTJForCausalLM, FlaxGPTJModel, FlaxGPTJPreTrainedModel