mirror of
https://github.com/wassname/Open-Assistant.git
synced 2026-09-12 12:03:06 +08:00
[feature] Add OA private RM dataset
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user