Update ner_model.py

ThilinaRajapakse · Jun 20, 2020 · d7992ad · d7992ad
1 parent 4549213
commit d7992ad
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/simpletransformers/ner/ner_model.py b/simpletransformers/ner/ner_model.py
@@ -47,6 +47,9 @@
     LongformerConfig,
     LongformerForTokenClassification,
     LongformerTokenizer,
+    MobileBertConfig,
+    MobileBertTokenizer,
+    MobileBertForTokenClassification,
     RobertaConfig,
     RobertaForTokenClassification,
     RobertaTokenizer,
@@ -130,6 +133,7 @@ def __init__(
             "distilbert": (DistilBertConfig, DistilBertForTokenClassification, DistilBertTokenizer),
             "electra": (ElectraConfig, ElectraForTokenClassification, ElectraTokenizer),
             "longformer": (LongformerConfig, LongformerForTokenClassification, LongformerTokenizer),
+            "mobilebert": (MobileBertConfig, MobileBertForTokenClassification, MobileBertTokenizer),
             "roberta": (RobertaConfig, RobertaForTokenClassification, RobertaTokenizer),
             "xlmroberta": (XLMRobertaConfig, XLMRobertaForTokenClassification, XLMRobertaTokenizer),
         }