Unverified Commit f45e8739 authored by Sam Shleifer's avatar Sam Shleifer Committed by GitHub
Browse files

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

parent e33929ef
...@@ -9,6 +9,7 @@ import torch ...@@ -9,6 +9,7 @@ import torch
from filelock import FileLock from filelock import FileLock
from torch.utils.data.dataset import Dataset from torch.utils.data.dataset import Dataset
from ...tokenization_bart import BartTokenizer, BartTokenizerFast
from ...tokenization_roberta import RobertaTokenizer, RobertaTokenizerFast from ...tokenization_roberta import RobertaTokenizer, RobertaTokenizerFast
from ...tokenization_utils import PreTrainedTokenizer from ...tokenization_utils import PreTrainedTokenizer
from ...tokenization_xlm_roberta import XLMRobertaTokenizer from ...tokenization_xlm_roberta import XLMRobertaTokenizer
...@@ -92,6 +93,8 @@ class GlueDataset(Dataset): ...@@ -92,6 +93,8 @@ class GlueDataset(Dataset):
RobertaTokenizer, RobertaTokenizer,
RobertaTokenizerFast, RobertaTokenizerFast,
XLMRobertaTokenizer, XLMRobertaTokenizer,
BartTokenizer,
BartTokenizerFast,
): ):
# HACK(label indices are swapped in RoBERTa pretrained model) # HACK(label indices are swapped in RoBERTa pretrained model)
label_list[1], label_list[2] = label_list[2], label_list[1] label_list[1], label_list[2] = label_list[2], label_list[1]
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment