nemo_rl.data.datasets.preference_datasets.preference_dataset#

Module Contents#

Classes#

PreferenceDataset

Dataset class for preference data which can be loaded from a JSON file.

API#

class nemo_rl.data.datasets.preference_datasets.preference_dataset.PreferenceDataset(
train_data_path: str,
val_data_path: Optional[str] = None,
train_split: Optional[str] = None,
val_split: Optional[str] = None,
)#

Dataset class for preference data which can be loaded from a JSON file.

This class handles loading of preference data for DPO and RM training. The input JSONL files should contain valid JSON objects formatted like this: { “context”: list of dicts, # The prompt message (including previous turns, if any) “completions”: list of dicts, # The list of completions { “rank”: int, # The rank of the completion (lower rank is preferred) “completion”: list of dicts, # The completion message(s) } }

Parameters:

  • train_data_path – Path to the JSON file containing training data

  • val_data_path – Path to the JSON file containing validation data

  • train_split – Split name for the training data, used for HuggingFace datasets, default is None

  • val_split – Split name for the validation data, used for HuggingFace datasets, default is None

Initialization