Update classification_model.py

ThilinaRajapakse · Jun 20, 2020 · 0293739 · 0293739
1 parent d7992ad
commit 0293739
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/simpletransformers/classification/classification_model.py b/simpletransformers/classification/classification_model.py
@@ -58,6 +58,9 @@
     LongformerConfig,
     LongformerForSequenceClassification,
     LongformerTokenizer,
+    MobileBertConfig,
+    MobileBertTokenizer,
+    MobileBertForSequenceClassification,
     RobertaConfig,
     RobertaTokenizer,
     XLMConfig,
@@ -106,6 +109,7 @@ def __init__(
             "electra": (ElectraConfig, ElectraForSequenceClassification, ElectraTokenizer),
             "flaubert": (FlaubertConfig, FlaubertForSequenceClassification, FlaubertTokenizer),
             "longformer": (LongformerConfig, LongformerForSequenceClassification, LongformerTokenizer),
+            "mobilebert": (MobileBertConfig, MobileBertForSequenceClassification, MobileBertTokenizer),
             "roberta": (RobertaConfig, RobertaForSequenceClassification, RobertaTokenizer),
             "xlnet": (XLNetConfig, XLNetForSequenceClassification, XLNetTokenizer),
             "xlm": (XLMConfig, XLMForSequenceClassification, XLMTokenizer),