From f45e873910e60d89511ae0193711e71c5c710468 Mon Sep 17 00:00:00 2001 From: Sam Shleifer Date: Fri, 19 Jun 2020 13:33:24 -0400 Subject: [PATCH] [bart-mnli] Fix class flipping bug (#5141) --- src/transformers/data/datasets/glue.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/transformers/data/datasets/glue.py b/src/transformers/data/datasets/glue.py index 1775b93f9..f696f6824 100644 --- a/src/transformers/data/datasets/glue.py +++ b/src/transformers/data/datasets/glue.py @@ -9,6 +9,7 @@ import torch from filelock import FileLock from torch.utils.data.dataset import Dataset +from ...tokenization_bart import BartTokenizer, BartTokenizerFast from ...tokenization_roberta import RobertaTokenizer, RobertaTokenizerFast from ...tokenization_utils import PreTrainedTokenizer from ...tokenization_xlm_roberta import XLMRobertaTokenizer @@ -92,6 +93,8 @@ class GlueDataset(Dataset): RobertaTokenizer, RobertaTokenizerFast, XLMRobertaTokenizer, + BartTokenizer, + BartTokenizerFast, ): # HACK(label indices are swapped in RoBERTa pretrained model) label_list[1], label_list[2] = label_list[2], label_list[1]