[feature] Add OA private RM dataset

This commit is contained in:
theblackcat102
2023-02-03 15:07:05 +00:00
parent 8b2080559c
commit 0be4d88605
2 changed files with 49 additions and 1 deletions
+7 -1
View File
@@ -115,7 +115,7 @@ def argument_parsing(parser):
def get_datasets(dataset_list: List[AnyStr], tokenizer):
from rank_datasets import AnthropicRLHF, GPTJSynthetic, HFSummary, WebGPT
from rank_datasets import AnthropicRLHF, GPTJSynthetic, HFSummary, OAPrivate, WebGPT
from torch.utils.data import ConcatDataset
train_datasets, evals = [], {}
@@ -141,5 +141,11 @@ def get_datasets(dataset_list: List[AnyStr], tokenizer):
eval = AnthropicRLHF("test", tokenizer.sep_token)
train_datasets.append(train)
evals["anthropic_rlhf"] = eval
elif "oa_private" == dataset_name:
train = OAPrivate(split="train", sep_token=tokenizer.sep_token)
eval = OAPrivate(split="val", sep_token=tokenizer.sep_token)
train_datasets.append(train)
evals["oa_private"] = eval
train = ConcatDataset(train_datasets)
return train, evals