|
from datasets import DatasetDict, load_dataset |
|
|
|
from opencompass.registry import LOAD_DATASET |
|
|
|
from .base import BaseDataset |
|
|
|
|
|
@LOAD_DATASET.register_module() |
|
class CivilCommentsDataset(BaseDataset): |
|
|
|
@staticmethod |
|
def load(**kwargs): |
|
train_dataset = load_dataset(**kwargs, split='train') |
|
test_dataset = load_dataset(**kwargs, split='test') |
|
|
|
def pre_process(example): |
|
example['label'] = int(example['toxicity'] >= 0.5) |
|
example['choices'] = ['no', 'yes'] |
|
return example |
|
|
|
def remove_columns(dataset): |
|
return dataset.remove_columns([ |
|
'severe_toxicity', 'obscene', 'threat', 'insult', |
|
'identity_attack', 'sexual_explicit' |
|
]) |
|
|
|
train_dataset = remove_columns(train_dataset) |
|
test_dataset = remove_columns(test_dataset) |
|
test_dataset = test_dataset.shuffle(seed=42) |
|
test_dataset = test_dataset.select(list(range(10000))) |
|
test_dataset = test_dataset.map(pre_process) |
|
|
|
return DatasetDict({ |
|
'train': train_dataset, |
|
'test': test_dataset, |
|
}) |
|
|
