Fix it to work with BART (huggingface#6756)

fabiocapsouza · Nov 15, 2020 · 94a8802 · 94a8802
1 parent 00b3ab5
commit 94a8802
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/question-answering/run_squad.py b/examples/question-answering/run_squad.py
@@ -187,7 +187,7 @@ def train(args, train_dataset, model, tokenizer):
                 "end_positions": batch[4],
             }
 
-            if args.model_type in ["xlm", "roberta", "distilbert", "camembert"]:
+            if args.model_type in ["xlm", "roberta", "distilbert", "camembert", "bart"]:
                 del inputs["token_type_ids"]
 
             if args.model_type in ["xlnet", "xlm"]:
@@ -300,7 +300,7 @@ def evaluate(args, model, tokenizer, prefix=""):
                 "token_type_ids": batch[2],
             }
 
-            if args.model_type in ["xlm", "roberta", "distilbert", "camembert"]:
+            if args.model_type in ["xlm", "roberta", "distilbert", "camembert", "bart"]:
                 del inputs["token_type_ids"]
 
             feature_indices = batch[3]