fix bart for QA

ThilinaRajapakse · Jun 14, 2020 · 4c624a0 · 4c624a0
1 parent 28c9399
commit 4c624a0
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/simpletransformers/question_answering/question_answering_model.py b/simpletransformers/question_answering/question_answering_model.py
@@ -704,7 +704,7 @@ def evaluate(self, eval_data, output_dir, verbose_logging=False):
                     "token_type_ids": batch[2],
                 }
 
-                if self.args["model_type"] in ["xlm", "roberta", "distilbert", "camembert", "electra", "xlmroberta"]:
+                if self.args["model_type"] in ["xlm", "roberta", "distilbert", "camembert", "electra", "xlmroberta","bart]:
                     del inputs["token_type_ids"]
 
                 example_indices = batch[3]
@@ -834,7 +834,7 @@ def predict(self, to_predict, n_best_size=None):
                     "token_type_ids": batch[2],
                 }
 
-                if self.args["model_type"] in ["xlm", "roberta", "distilbert", "camembert", "electra", "xlmroberta"]:
+                if self.args["model_type"] in ["xlm", "roberta", "distilbert", "camembert", "electra", "xlmroberta","bart"]:
                     del inputs["token_type_ids"]
 
                 example_indices = batch[3]
@@ -958,7 +958,7 @@ def _get_inputs_dict(self, batch):
             "end_positions": batch[4],
         }
 
-        if self.args["model_type"] in ["xlm", "roberta", "distilbert", "camembert", "electra", "xlmroberta"]:
+        if self.args["model_type"] in ["xlm", "roberta", "distilbert", "camembert", "electra", "xlmroberta","bart"]:
             del inputs["token_type_ids"]
 
         if self.args["model_type"] in ["xlnet", "xlm"]: