[bart-mnli] Fix class flipping bug (#5141)

This commit is contained in:
Sam Shleifer
2020-06-19 13:33:24 -04:00
committed by GitHub
parent e33929ef1e
commit f45e873910

View File

@@ -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]