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(
data_path: str,
split: Optional[str] = None,
**kwargs,
)#

Bases: nemo_rl.data.datasets.raw_dataset.RawDataset

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[dict], # The prompt message (including previous turns, if any) “completions”: [ # The list of completions { “rank”: 0, # The rank of the completion (lower rank is preferred) “completion”: list[dict], # The completion message(s) }, { “rank”: 1, # The rank of the completion (lower rank is preferred) “completion”: list[dict], # The completion message(s) }, … # More completions can be added if needed ] } Please refer to https://github.com/NVIDIA-NeMo/RL/blob/main/docs/guides/dpo.md#datasets for more details.

Parameters:

  • data_path – Path to the dataset JSON file

  • split – Optional split name for the dataset, used for HuggingFace datasets

Initialization