Add mobilebert next sentence prediction

huggingface · Nov 10, 2020 · 3fd78dd · 3fd78dd
1 parent a30c3f8
commit 3fd78dd
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/transformers/modeling_tf_auto.py b/src/transformers/modeling_tf_auto.py
@@ -121,6 +121,7 @@
 from .modeling_tf_mobilebert import (
     TFMobileBertForMaskedLM,
     TFMobileBertForMultipleChoice,
+    TFMobileBertForNextSentencePrediction,
     TFMobileBertForPreTraining,
     TFMobileBertForQuestionAnswering,
     TFMobileBertForSequenceClassification,
@@ -359,6 +360,7 @@
 TF_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING = OrderedDict(
     [
         (BertConfig, TFBertForNextSentencePrediction),
+        (MobileBertConfig, TFMobileBertForNextSentencePrediction),
     ]
 )
 

diff --git a/utils/check_repo.py b/utils/check_repo.py
@@ -87,10 +87,8 @@
     "RagSequenceForGeneration",
     "RagTokenForGeneration",
     "T5Stack",
-    "TFBertForNextSentencePrediction",
     "TFFunnelBaseModel",
     "TFGPT2DoubleHeadsModel",
-    "TFMobileBertForNextSentencePrediction",
     "TFOpenAIGPTDoubleHeadsModel",
     "XLMForQuestionAnswering",
     "XLMProphetNetDecoder",