longformer for seq2seq

ThilinaRajapakse · Jun 1, 2020 · 2bc8297 · 2bc8297
1 parent 7db02d1
commit 2bc8297
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/simpletransformers/seq2seq/seq2seq_model.py b/simpletransformers/seq2seq/seq2seq_model.py
@@ -36,6 +36,9 @@
     ElectraConfig,
     ElectraModel,
     ElectraTokenizer,
+    LongformerConfig,
+    LongformerModel,
+    LongformerTokenizer,
     PreTrainedModel,
     PreTrainedTokenizer,
     RobertaConfig,
@@ -60,13 +63,14 @@
 
 MODEL_CLASSES = {
     "auto": (AutoConfig, AutoModel, AutoTokenizer),
+    "bart": (BartConfig, BartForConditionalGeneration, BartTokenizer),
     "bert": (BertConfig, BertModel, BertTokenizer),
-    "roberta": (RobertaConfig, RobertaModel, RobertaTokenizer),
-    "distilbert": (DistilBertConfig, DistilBertModel, DistilBertTokenizer),
     "camembert": (CamembertConfig, CamembertModel, CamembertTokenizer),
+    "distilbert": (DistilBertConfig, DistilBertModel, DistilBertTokenizer),
     "electra": (ElectraConfig, ElectraModel, ElectraTokenizer),
-    "bart": (BartConfig, BartForConditionalGeneration, BartTokenizer),
+    "longformer": (LongformerConfig, LongformerModel, LongformerTokenizer),
     "marian": (MarianConfig, MarianMTModel, MarianTokenizer),
+    "roberta": (RobertaConfig, RobertaModel, RobertaTokenizer),
 }