add bart for QA

ThilinaRajapakse · Jun 14, 2020 · 28c9399 · 28c9399
1 parent 16db7f0
commit 28c9399
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/simpletransformers/question_answering/question_answering_model.py b/simpletransformers/question_answering/question_answering_model.py
@@ -47,6 +47,9 @@
     AlbertConfig,
     AlbertForQuestionAnswering,
     AlbertTokenizer,
+    BartConfig,
+    BartForQuestionAnswering,
+    BartTokenizer,
     BertConfig,
     BertForQuestionAnswering,
     BertTokenizer,
@@ -101,6 +104,7 @@ def __init__(self, model_type, model_name, args=None, use_cuda=True, cuda_device
         MODEL_CLASSES = {
             "albert": (AlbertConfig, AlbertForQuestionAnswering, AlbertTokenizer),
             "auto": (AutoConfig, AutoTokenizer, AutoModelForQuestionAnswering),
+            "bart": (BartConfig, BartForQuestionAnswering, BartTokenizer),
             "bert": (BertConfig, BertForQuestionAnswering, BertTokenizer),
             "distilbert": (DistilBertConfig, DistilBertForQuestionAnswering, DistilBertTokenizer),
             "electra": (ElectraConfig, ElectraForQuestionAnswering, ElectraTokenizer),