diff --git a/docs/guides/dpo.md b/docs/guides/dpo.md index ce685469930..f00dde9f122 100644 --- a/docs/guides/dpo.md +++ b/docs/guides/dpo.md @@ -91,31 +91,38 @@ DPO training supports only two completions (where the lowest rank is preferred a } ``` -NeMo RL provides a DPO-compatible implementation of the [HelpSteer3](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/data/hf_datasets/helpsteer3.py) dataset as an example. This dataset is downloaded from Hugging Face and preprocessed on-the-fly, so there's no need to provide a path to any datasets on disk. +By default, NeMo RL has support for [HelpSteer3](../../nemo_rl/data/datasets/preference_datasets/helpsteer3.py) and [Tulu3Preference](../../nemo_rl/data/datasets/preference_datasets/tulu3.py) datasets. Both of these datasets are downloaded from HuggingFace and preprocessed on-the-fly, so there's no need to provide a path to any datasets on disk. -We also provide a [PreferenceDataset](../../nemo_rl/data/hf_datasets/preference_dataset.py) class that is compatible with JSONL-formatted preference datasets. You can modify your config as follows to use such a custom preference dataset: +We provide a [PreferenceDataset](../../nemo_rl/data/datasets/preference_datasets/preference_dataset.py) class that is compatible with jsonl-formatted preference datasets for loading datasets from local path or HuggingFace. You can modify your config as follows to use such a custom preference dataset: ```yaml data: dataset_name: PreferenceDataset - train_data_path: + train_data_path: # e.g., /path/to/local/dataset.jsonl or hf_org/hf_dataset_name (HuggingFace) + # multiple validation sets is supported val_data_paths: - : + : + : + train_split: , default is None # used for HuggingFace datasets + val_split: , default is None # used for HuggingFace datasets ``` -with support for multiple validation sets achieved with: + +We also provide a [BinaryPreferenceDataset](../../nemo_rl/data/datasets/preference_datasets/binary_preference_dataset.py) class, which is a simplified version of PreferenceDataset for pairwise ranked preference with single turn completions. You can use `prompt_key`, `chosen_key` and `rejected_key` to specify which fields in your data correspond to the question, chosen answer and rejected answer respectively. Here's an example configuration: ```yaml data: - dataset_name: PreferenceDataset - train_data_path: - val_data_paths: - : - : + dataset_name: BinaryPreferenceDataset + train_data_path: # e.g., /path/to/local/dataset.jsonl or hf_org/hf_dataset_name (HuggingFace) + val_data_path: + prompt_key: , default is "prompt" + chosen_key: , default is "chosen" + rejected_key: , default is "rejected" + train_split: , default is None # used for HuggingFace datasets + val_split: , default is None # used for HuggingFace datasets ``` + Please note: - If you are using a logger, the prefix used for each validation set will be `validation-`. The total validation time, summed across all validation sets, is reported under `timing/validation/total_validation_time`. - If you are doing checkpointing, the `metric_name` value in your `checkpointing` config should reflect the metric and validation set to be tracked. For example, `validation-_loss`. -The older [DPODataset](../../nemo_rl/data/hf_datasets/dpo.py) class is deprecated. This class is also compatible with JSONL-formatted preference datsets. It assumes train and validation datasets have been split and processed into the expected format offline. The JSONL files should consist of examples with `prompt`, `chosen_response`, and `rejected_response` keys. - ## DPO-Specific Parameters The DPO implementation in NeMo RL supports several key parameters that can be adjusted: diff --git a/docs/guides/eval.md b/docs/guides/eval.md index f5737913b72..7fb456e64b7 100644 --- a/docs/guides/eval.md +++ b/docs/guides/eval.md @@ -4,10 +4,10 @@ This document explains how to use an evaluation script for assessing model capab ## Prepare for Evaluation -To prepare for evaluation, first ensure your model is in the correct format, which may involve an optional conversion of PyTorch DCP checkpoints to the Hugging Face format. Following this, you need to prepare the evaluation configuration, which includes defining prompt templates and any custom settings required to run the evaluation. +To prepare for evaluation, first ensure your model is in the correct format, which may involve an optional conversion of PyTorch DCP checkpoints to the HuggingFace format. Following this, you need to prepare the evaluation configuration, which includes defining prompt templates and any custom settings required to run the evaluation. ### Convert DCP to HF (Optional) -If you have trained a model and saved the checkpoint in the Pytorch DCP format, you first need to convert it to the Hugging Face format before running evaluation. +If you have trained a model and saved the checkpoint in the Pytorch DCP format, you first need to convert it to the HuggingFace format before running evaluation. Use the `examples/converters/convert_dcp_to_hf.py` script. You'll need the path to the training configuration file (`config.yaml`), the DCP checkpoint directory, and specify an output path for the HF format model. @@ -37,9 +37,9 @@ For open-source models, we recommend setting `tokenizer.chat_template=default`, ## Run the Evaluation Script -We will use the `run_eval.py` script to run an evaluation using a model directly from the Hugging Face Hub or from a local path that is already in Hugging Face format. +We will use the `run_eval.py` script to run an evaluation using a model directly from the HuggingFace Hub or from a local path that is already in HuggingFace format. -Note that the evaluation script only supports the Hugging Face format model. If you haven't converted your DCP format model, you should back to [Convert DCP to HF](#convert-dcp-to-hf-optional) and follow the guide to convert your model. +Note that the evaluation script only supports the HuggingFace format model. If you haven't converted your DCP format model, you should back to [Convert DCP to HF](#convert-dcp-to-hf-optional) and follow the guide to convert your model. ```sh # Run evaluation script with default config (examples/configs/evals/eval.yaml) @@ -89,12 +89,10 @@ score=0.1000 (3.0/30) ## List of currently supported benchmarks -- [AIME-2024](../../nemo_rl/data/eval_datasets/aime2024.py): the corresponding `data.dataset_name` is `"aime2024"`. -- [AIME-2025](../../nemo_rl/data/eval_datasets/aime2025.py): the corresponding `data.dataset_name` is `"aime2025"`. -- [GPQA and GPQA-diamond](../../nemo_rl/data/eval_datasets/gpqa.py): the corresponding `data.dataset_name` are `"gpqa"` and `"gpqa-diamond"`. -- [MATH and MATH-500](../../nemo_rl/data/eval_datasets/math.py): the corresponding `data.dataset_name` are `"math"` and `"math500"`. -- [MMLU](../../nemo_rl/data/eval_datasets/mmlu.py): this also includes MMMLU (Multilingual MMLU), a total of 14 languages. When `data.dataset_name` is set to `mmlu`, the English version is used. If one wants to run evaluation on another language, `data.dataset_name` should be set to `mmlu_{language}` where `language` is one of following 14 values, `["AR-XY", "BN-BD", "DE-DE", "ES-LA", "FR-FR", "HI-IN", "ID-ID", "IT-IT", "JA-JP", "KO-KR", "PT-BR", "ZH-CN", "SW-KE", "YO-NG"]`. -- [MMLU-Pro](../../nemo_rl/data/eval_datasets/mmlu_pro.py): the corresponding `data.dataset_name` is `"mmlu_pro"`. - -More details can be found in [load_eval_dataset](../../nemo_rl/data/eval_datasets/__init__.py). +- [AIME-2024 and AIME-2025](../../nemo_rl/data/datasets/eval_datasets/aime.py): the corresponding `data.dataset_name` are `"aime2024"` and `"aime2025"`. +- [GPQA and GPQA-diamond](../../nemo_rl/data/datasets/eval_datasets/gpqa.py): the corresponding `data.dataset_name` are `"gpqa"` and `"gpqa_diamond"`. +- [MATH and MATH-500](../../nemo_rl/data/datasets/eval_datasets/math.py): the corresponding `data.dataset_name` are `"math"` and `"math500"`. +- [MMLU](../../nemo_rl/data/datasets/eval_datasets/mmlu.py): this also includes MMMLU (Multilingual MMLU), a total of 14 languages. When `data.dataset_name` is set to `mmlu`, the English version is used. If one wants to run evaluation on another language, `data.dataset_name` should be set to `mmlu_{language}` where `language` is one of following 14 values, `["AR-XY", "BN-BD", "DE-DE", "ES-LA", "FR-FR", "HI-IN", "ID-ID", "IT-IT", "JA-JP", "KO-KR", "PT-BR", "ZH-CN", "SW-KE", "YO-NG"]`. +- [MMLU-Pro](../../nemo_rl/data/datasets/eval_datasets/mmlu_pro.py): the corresponding `data.dataset_name` is `"mmlu_pro"`. +More details can be found in [load_eval_dataset](../../nemo_rl/data/datasets/eval_datasets/__init__.py). diff --git a/docs/guides/grpo.md b/docs/guides/grpo.md index e5e0bfe431f..3c3bd7aff69 100644 --- a/docs/guides/grpo.md +++ b/docs/guides/grpo.md @@ -36,6 +36,22 @@ To support this, we need to know: * Which data should go to which environments * How to prepare the data from your dataset into a form we can use +#### Dataset + +By default, NeMo RL has support for [OpenMathInstruct-2](../../nemo_rl/data/datasets/response_datasets/openmathinstruct2.py) and [DeepScaler](../../nemo_rl/data/datasets/response_datasets/deepscaler.py) datasets. Both of these datasets are downloaded from HuggingFace and preprocessed on-the-fly, so there's no need to provide a path to any datasets on disk. + +We provide a [ResponseDataset](../../nemo_rl/data/datasets/response_datasets/response_dataset.py) class that is compatible with jsonl-formatted response datasets for loading datasets from local path or HuggingFace. You can use `input_key`, `output_key` to specify which fields in your data correspond to the question and answer respectively. Here's an example configuration: +```yaml +data: + dataset_name: ResponseDataset + train_data_path: # e.g., /path/to/local/dataset.jsonl or hf_org/hf_dataset_name (HuggingFace) + val_data_path: + input_key: , default is "input" + output_key: , default is "output" + train_split: , default is None # used for HuggingFace datasets + val_split: , default is None # used for HuggingFace datasets +``` + #### Common Data Format We define a [DatumSpec](../../nemo_rl/data/interfaces.py) that holds all relevant information for each training example: diff --git a/docs/guides/rm.md b/docs/guides/rm.md index f1843cd92c3..7ee370bbdf7 100644 --- a/docs/guides/rm.md +++ b/docs/guides/rm.md @@ -80,25 +80,34 @@ Currently, RM training supports only two completions (where the lowest rank is p } ``` -NeMo RL provides a RM-compatible implementation of the [HelpSteer3](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/data/hf_datasets/helpsteer3.py) dataset as an example. This dataset is downloaded from Hugging Face and preprocessed on-the-fly, so there's no need to provide a path to any datasets on disk. +By default, NeMo RL has support for [HelpSteer3](../../nemo_rl/data/datasets/preference_datasets/helpsteer3.py) and [Tulu3Preference](../../nemo_rl/data/datasets/preference_datasets/tulu3.py) datasets. Both of these datasets are downloaded from HuggingFace and preprocessed on-the-fly, so there's no need to provide a path to any datasets on disk. -We also provide a [PreferenceDataset](../../nemo_rl/data/hf_datasets/preference_dataset.py) class that is compatible with JSONL-formatted preference datasets. You can modify your config as follows to use such a custom preference dataset: +We provide a [PreferenceDataset](../../nemo_rl/data/datasets/preference_datasets/preference_dataset.py) class that is compatible with jsonl-formatted preference datasets for loading datasets from local path or HuggingFace.. You can modify your config as follows to use such a custom preference dataset: ```yaml data: dataset_name: PreferenceDataset - train_data_path: + train_data_path: # e.g., /path/to/local/dataset.jsonl or hf_org/hf_dataset_name (HuggingFace) + # multiple validation sets is supported val_data_paths: - : + : + : + train_split: , default is None # used for HuggingFace datasets + val_split: , default is None # used for HuggingFace datasets ``` -with support for multiple validation sets achieved with: + +We also provide a [BinaryPreferenceDataset](../../nemo_rl/data/datasets/preference_datasets/binary_preference_dataset.py) class, which is a simplified version of PreferenceDataset for pairwise ranked preference with single turn completions. You can use `prompt_key`, `chosen_key` and `rejected_key` to specify which fields in your data correspond to the question, chosen answer and rejected answer respectively. Here's an example configuration: ```yaml data: - dataset_name: PreferenceDataset - train_data_path: - val_data_paths: - : - : + dataset_name: BinaryPreferenceDataset + train_data_path: # e.g., /path/to/local/dataset.jsonl or hf_org/hf_dataset_name (HuggingFace) + val_data_path: + prompt_key: , default is "prompt" + chosen_key: , default is "chosen" + rejected_key: , default is "rejected" + train_split: , default is None # used for HuggingFace datasets + val_split: , default is None # used for HuggingFace datasets ``` + Please note: - If you are using a logger, the prefix used for each validation set will be `validation-`. The total validation time, summed across all validation sets, is reported under `timing/validation/total_validation_time`. -- If you are doing checkpointing, the `metric_name` value in your `checkpointing` config should reflect the metric and validation set to be tracked. For example, `validation-_loss`. \ No newline at end of file +- If you are doing checkpointing, the `metric_name` value in your `checkpointing` config should reflect the metric and validation set to be tracked. For example, `validation-_loss`. diff --git a/docs/guides/sft.md b/docs/guides/sft.md index e272c6d9d69..da64fa1ab32 100644 --- a/docs/guides/sft.md +++ b/docs/guides/sft.md @@ -34,7 +34,7 @@ SFT datasets in NeMo RL are encapsulated using classes. Each SFT data class is e 1. `formatted_ds`: The dictionary of formatted datasets. This dictionary should contain `train` and `validation` splits, and each split should conform to the format described below. 2. `task_spec`: The `TaskDataSpec` for this dataset. This should specify the name you choose for this dataset. -SFT datasets are expected to follow the Hugging Face chat format. Refer to the [chat dataset document](../design-docs/chat-datasets.md) for details. If your data is not in the correct format, simply write a preprocessing script to convert the data into this format. [data/hf_datasets/squad.py](../../nemo_rl/data/hf_datasets/squad.py) has an example: +SFT datasets are expected to follow the HuggingFace chat format. Refer to the [chat dataset document](../design-docs/chat-datasets.md) for details. If your data is not in the correct format, simply write a preprocessing script to convert the data into this format. [response_datasets/squad.py](../../nemo_rl/data/datasets/response_datasets/squad.py) has an example: ```python def format_squad(data): @@ -56,7 +56,7 @@ def format_squad(data): } ``` -NeMo RL SFT uses Hugging Face chat templates to format the individual examples. Three types of chat templates are supported, which can be configured via `tokenizer.chat_template` in your yaml config (see [sft.yaml](../../examples/configs/sft.yaml) for an example): +NeMo RL SFT uses HuggingFace chat templates to format the individual examples. Three types of chat templates are supported, which can be configured via `tokenizer.chat_template` in your yaml config (see [sft.yaml](../../examples/configs/sft.yaml) for an example): 1. Apply the tokenizer's default chat template. To use the tokenizer's default, either omit `tokenizer.chat_template` from the config altogether, or set `tokenizer.chat_template="default"`. 2. Use a "passthrough" template which simply concatenates all messages. This is desirable if the chat template has been applied to your dataset as an offline preprocessing step. In this case, you should set `tokenizer.chat_template` to None as follows: @@ -71,8 +71,19 @@ NeMo RL SFT uses Hugging Face chat templates to format the individual examples. custom_template: "{% for message in messages %}{%- if message['role'] == 'system' %}{{'Context: ' + message['content'].strip()}}{%- elif message['role'] == 'user' %}{{' Question: ' + message['content'].strip() + ' Answer: '}}{%- elif message['role'] == 'assistant' %}{{message['content'].strip()}}{%- endif %}{% endfor %}" ``` - -By default, NeMo RL has support for `Squad` and `OpenAssistant` datasets. Both of these datasets are downloaded from Hugging Face and preprocessed on-the-fly, so there's no need to provide a path to any datasets on disk. +By default, NeMo RL has support for [OpenAssistant](../../nemo_rl/data/datasets/response_datasets/oasst.py), [Squad](../../nemo_rl/data/datasets/response_datasets/squad.py) and [OpenMathInstruct-2](../../nemo_rl/data/datasets/response_datasets/openmathinstruct2.py) datasets. All of these datasets are downloaded from HuggingFace and preprocessed on-the-fly, so there's no need to provide a path to any datasets on disk. + +We provide a [ResponseDataset](../../nemo_rl/data/datasets/response_datasets/response_dataset.py) class that is compatible with jsonl-formatted response datasets for loading datasets from local path or HuggingFace. You can use `input_key`, `output_key` to specify which fields in your data correspond to the question and answer respectively. Here's an example configuration: +```yaml +data: + dataset_name: ResponseDataset + train_data_path: # e.g., /path/to/local/dataset.jsonl or hf_org/hf_dataset_name (HuggingFace) + val_data_path: + input_key: , default is "input" + output_key: , default is "output" + train_split: , default is None # used for HuggingFace datasets + val_split: , default is None # used for HuggingFace datasets +``` Adding a new dataset is a straightforward process. As long as your custom dataset has the `formatted_ds` and `task_spec` attributes described above, it can serve as a drop-in replacement for Squad and OpenAssistant. diff --git a/examples/configs/dpo.yaml b/examples/configs/dpo.yaml index e4a296be06f..7d9d0ba7b2d 100755 --- a/examples/configs/dpo.yaml +++ b/examples/configs/dpo.yaml @@ -157,12 +157,27 @@ data: dataset_name: HelpSteer3 # You can use custom preference datasets for training and validation. For example: + # 1. PreferenceDataset # data: # dataset_name: PreferenceDataset - # train_data_path: + # train_data_path: # e.g., /path/to/local/dataset.jsonl or hf_org/hf_dataset_name (HuggingFace) # val_data_paths: - # : + # : # ... + # train_split: , default is None # used for HuggingFace datasets + # val_split: , default is None # used for HuggingFace datasets + # 2. BinaryPreferenceDataset + # data: + # dataset_name: BinaryPreferenceDataset + # train_data_path: # e.g., /path/to/local/dataset.jsonl or hf_org/hf_dataset_name (HuggingFace) + # val_data_path: + # prompt_key: , default is "prompt" + # chosen_key: , default is "chosen" + # rejected_key: , default is "rejected" + # train_split: , default is None # used for HuggingFace datasets + # val_split: , default is None # used for HuggingFace datasets + # See https://github.com/NVIDIA-NeMo/RL/blob/main/docs/guides/dpo.md#datasets for more details. + # If you are doing checkpointing, `metric_name` should reflect the metric and validation set to be tracked. For example: # checkpointing: # metric_name: "validation-_loss" diff --git a/examples/configs/evals/local_eval.yaml b/examples/configs/evals/local_eval.yaml index ad9def2112d..b97180dd5af 100644 --- a/examples/configs/evals/local_eval.yaml +++ b/examples/configs/evals/local_eval.yaml @@ -6,9 +6,8 @@ generation: data: prompt_file: "examples/prompts/cot.txt" - dataset_name: "local" + dataset_name: "https://openaipublic.blob.core.windows.net/simple-evals/math_500_test.csv" problem_key: "Question" solution_key: "Answer" split: "train" - data_paths: "https:\/\/openaipublic.blob.core.windows.net\/simple-evals\/math_500_test.csv" file_format: "csv" diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 6ba7e4d54b5..19dc8a2cf65 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -202,9 +202,20 @@ data: max_input_seq_length: ${policy.max_total_sequence_length} # upper bound, real truncation occurs at vllm.max_model_len prompt_file: "examples/prompts/cot.txt" system_prompt_file: null - dataset_name: "OpenMathInstruct-2" shuffle: true + dataset_name: "OpenMathInstruct-2" + # You can use custom response datasets for training and validation. For example: + # data: + # dataset_name: ResponseDataset + # train_data_path: # e.g., /path/to/local/dataset.jsonl or hf_org/hf_dataset_name (HuggingFace) + # val_data_path: + # input_key: , default is "input" + # output_key: , default is "output" + # train_split: , default is None # used for HuggingFace datasets + # val_split: , default is None # used for HuggingFace datasets + # See https://github.com/NVIDIA-NeMo/RL/blob/main/docs/guides/grpo.md#datasets for more details. + env: math: num_workers: 8 diff --git a/examples/configs/rm.yaml b/examples/configs/rm.yaml index 0395037a23e..0302cd8236a 100644 --- a/examples/configs/rm.yaml +++ b/examples/configs/rm.yaml @@ -128,12 +128,27 @@ data: dataset_name: HelpSteer3 # You can use custom preference datasets for training and validation. For example: + # 1. PreferenceDataset # data: # dataset_name: PreferenceDataset - # train_data_path: + # train_data_path: # e.g., /path/to/local/dataset.jsonl or hf_org/hf_dataset_name (HuggingFace) # val_data_paths: - # : + # : # ... + # train_split: , default is None # used for HuggingFace datasets + # val_split: , default is None # used for HuggingFace datasets + # 2. BinaryPreferenceDataset + # data: + # dataset_name: BinaryPreferenceDataset + # train_data_path: # e.g., /path/to/local/dataset.jsonl or hf_org/hf_dataset_name (HuggingFace) + # val_data_path: + # prompt_key: , default is "prompt" + # chosen_key: , default is "chosen" + # rejected_key: , default is "rejected" + # train_split: , default is None # used for HuggingFace datasets + # val_split: , default is None # used for HuggingFace datasets + # See https://github.com/NVIDIA-NeMo/RL/blob/main/docs/guides/rm.md#datasets for more details. + # If you are doing checkpointing, `metric_name` should reflect the metric and validation set to be tracked. For example: # checkpointing: # metric_name: "validation-_loss" diff --git a/examples/configs/sft.yaml b/examples/configs/sft.yaml index 0db82b6a2fe..5d83c09080d 100644 --- a/examples/configs/sft.yaml +++ b/examples/configs/sft.yaml @@ -132,12 +132,23 @@ policy: data: max_input_seq_length: ${policy.max_total_sequence_length} - dataset_name: "squad" add_bos: true add_eos: true add_generation_prompt: false shuffle: true + dataset_name: "squad" + # You can use custom response datasets for training and validation. For example: + # data: + # dataset_name: ResponseDataset + # train_data_path: # e.g., /path/to/local/dataset.jsonl or hf_org/hf_dataset_name (HuggingFace) + # val_data_path: + # input_key: , default is "input" + # output_key: , default is "output" + # train_split: , default is None # used for HuggingFace datasets + # val_split: , default is None # used for HuggingFace datasets + # See https://github.com/NVIDIA-NeMo/RL/blob/main/docs/guides/sft.md#datasets for more details. + ## unused with squad dataset prompt_file: null split: null diff --git a/examples/run_dpo.py b/examples/run_dpo.py index b9b31cfcf68..ea84d1c0daa 100644 --- a/examples/run_dpo.py +++ b/examples/run_dpo.py @@ -22,8 +22,9 @@ from nemo_rl.algorithms.dpo import MasterConfig, dpo_train, setup from nemo_rl.algorithms.utils import get_tokenizer -from nemo_rl.data import DataConfig, hf_datasets -from nemo_rl.data.datasets import AllTaskProcessedDataset +from nemo_rl.data import DataConfig +from nemo_rl.data.datasets import AllTaskProcessedDataset, load_preference_dataset +from nemo_rl.data.datasets.preference_datasets import PreferenceDataset from nemo_rl.data.interfaces import DatumSpec, TaskDataSpec from nemo_rl.data.llm_message_utils import get_formatted_message_log from nemo_rl.distributed.virtual_cluster import init_ray @@ -118,7 +119,7 @@ def dpo_preprocessor( rejected_completion = datum_dict["completions"][0] else: raise NotImplementedError( - "Ties are not supported yet. You can use the following command to filter out ties: `cat | jq 'select(.completions[0].rank != .completions[1].rank)'`." + "Ties are not supported yet. You can use the following command to filter out ties: `cat | jq 'select(.completions[0].rank != .completions[1].rank)'`." ) messages_chosen = datum_dict["context"] + chosen_completion["completion"] @@ -164,41 +165,15 @@ def dpo_preprocessor( def setup_data(data_config: DataConfig, policy_config: PolicyConfig): print("\n▶ Setting up data...") - data_cls = data_config["dataset_name"] - - if data_cls == "PreferenceDataset": - data_path = data_config["train_data_path"] - data = hf_datasets.PreferenceDataset(data_path, split="train") - train_dataset = data.formatted_ds["train"] - val_dataset = None - elif data_cls == "HelpSteer3": - data = hf_datasets.HelpSteer3Dataset() - train_dataset = data.formatted_ds["train"] - val_dataset = data.formatted_ds["validation"] - elif data_cls == "Tulu3Preference": - data = hf_datasets.Tulu3PreferenceDataset() - train_dataset = data.formatted_ds["train"] - val_dataset = None - elif data_cls == "DPODataset": - data = hf_datasets.DPODataset( - train_data_path=data_config["train_data_path"], - val_data_path=data_config["val_data_path"], - ) - train_dataset = data.formatted_ds["train"] - val_dataset = data.formatted_ds["validation"] - else: - raise ValueError( - f"Unknown dataset class: {data_cls}. Supported datasets are: PreferenceDataset, HelpSteer3, Tulu3Preference, and DPODataset (deprecated)." - ) - if train_dataset: - print( - f" ✓ Training dataset loaded with {len(data.formatted_ds['train'])} samples." - ) + # load dataset + data = load_preference_dataset(data_config) + train_dataset = data.formatted_ds["train"] + val_dataset = data.formatted_ds["validation"] + + print(f" ✓ Training dataset loaded with {len(train_dataset)} samples.") if val_dataset: - print( - f" ✓ Validation dataset loaded with {len(data.formatted_ds['validation'])} samples." - ) + print(f" ✓ Validation dataset loaded with {len(val_dataset)} samples.") dpo_task_spec = data.task_spec @@ -211,15 +186,10 @@ def setup_data(data_config: DataConfig, policy_config: PolicyConfig): max_seq_length=data_config["max_input_seq_length"], ) - if data_cls == "PreferenceDataset": + # TODO @yukih: unify the code when support multiple datasets for other algorithms + if "val_data_paths" in data_config and data_config["val_data_paths"]: val_dataset = {} - assert "val_data_path" not in data_config, ( - "`val_data_path` cannot be provided for PreferenceDataset. You should use `val_data_paths` instead." - ) - assert "val_data_paths" in data_config, ( - "`val_data_paths` must be provided for PreferenceDataset" - ) assert isinstance(data_config["val_data_paths"], dict), ( f"Invalid type for val_data_paths: {type(data_config['val_data_paths'])}. val_data_paths must be a dictionary." ) @@ -227,14 +197,12 @@ def setup_data(data_config: DataConfig, policy_config: PolicyConfig): for val_dataset_name, val_dataset_path in val_data_paths.items(): assert val_dataset_name not in val_dataset - val_data = hf_datasets.PreferenceDataset( - val_dataset_path, split="validation" - ) + val_data = PreferenceDataset(val_dataset_path) print( - f" ✓ Validation dataset '{val_dataset_name}' loaded with {len(val_data.formatted_ds['validation'])} samples." + f" ✓ Validation dataset '{val_dataset_name}' loaded with {len(val_data.formatted_ds['train'])} samples." ) val_dataset[val_dataset_name] = AllTaskProcessedDataset( - val_data.formatted_ds["validation"], + val_data.formatted_ds["train"], tokenizer, val_data.task_spec, dpo_preprocessor, diff --git a/examples/run_eval.py b/examples/run_eval.py index 89e2ede3952..89669386325 100644 --- a/examples/run_eval.py +++ b/examples/run_eval.py @@ -23,11 +23,8 @@ from transformers import AutoTokenizer, PreTrainedTokenizerBase from nemo_rl.algorithms.utils import get_tokenizer -from nemo_rl.data.datasets import AllTaskProcessedDataset -from nemo_rl.data.eval_datasets import load_eval_dataset -from nemo_rl.distributed.ray_actor_environment_registry import ( - get_actor_python_env, -) +from nemo_rl.data.datasets import AllTaskProcessedDataset, load_eval_dataset +from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env from nemo_rl.distributed.virtual_cluster import init_ray from nemo_rl.environments.math_environment import MathEnvironment from nemo_rl.evals.eval import MasterConfig, run_env_eval, setup diff --git a/examples/run_grpo_math.py b/examples/run_grpo_math.py index f31c2c212cc..17c6513e08f 100644 --- a/examples/run_grpo_math.py +++ b/examples/run_grpo_math.py @@ -24,9 +24,7 @@ from nemo_rl.algorithms.grpo import MasterConfig, grpo_train, setup from nemo_rl.algorithms.utils import get_tokenizer from nemo_rl.data import DataConfig -from nemo_rl.data.datasets import AllTaskProcessedDataset -from nemo_rl.data.hf_datasets.deepscaler import DeepScalerDataset -from nemo_rl.data.hf_datasets.openmathinstruct2 import OpenMathInstruct2Dataset +from nemo_rl.data.datasets import AllTaskProcessedDataset, load_response_dataset from nemo_rl.data.interfaces import ( DatumSpec, LLMMessageLogType, @@ -73,7 +71,7 @@ def hf_data_processor( max_seq_length: int, idx: int, ) -> DatumSpec: - """Process a datum dictionary (directly loaded from data/hf_datasets/openmathinstruct2.py) into a DatumSpec for the Math Environment.""" + """Process a datum dictionary (directly loaded from response_datasets/.py) into a DatumSpec for the Math Environment.""" user_message = datum_dict["messages"] problem = user_message[0]["content"] extra_env_info = {"ground_truth": user_message[1]["content"]} @@ -138,23 +136,16 @@ def setup_data( system_prompt_file=data_config["system_prompt_file"], ) - # Load OpenMathInstruct2Dataset using nemo rl datasets - if data_config["dataset_name"] == "OpenMathInstruct-2": - print("Loading nvidia/OpenMathInstruct2Dataset for training and validation") - data: Any = OpenMathInstruct2Dataset(seed=seed) - elif data_config["dataset_name"] == "DeepScaler": - print( - "Loading agentica-org/DeepScaleR-Preview-Dataset for training and validation" - ) - data: Any = DeepScalerDataset(seed=seed) - else: - raise ValueError(f"No processor for dataset {data_config['dataset_name']}.") + # load dataset + data: Any = load_response_dataset(data_config, seed) + # data processor task_data_processors: dict[str, tuple[TaskDataSpec, TaskDataProcessFnCallable]] = ( defaultdict(lambda: (math_task_spec, hf_data_processor)) ) task_data_processors["math"] = (math_task_spec, hf_data_processor) + # setup math environment math_env = MathEnvironment.options( # type: ignore # it's wrapped with ray.remote runtime_env={ "py_executable": get_actor_python_env( @@ -163,6 +154,7 @@ def setup_data( "env_vars": dict(os.environ), # Pass thru all user environment variables } ).remote(env_configs["math"]) + dataset = AllTaskProcessedDataset( data.formatted_ds["train"], tokenizer, diff --git a/examples/run_rm.py b/examples/run_rm.py index 0adf84490d1..bf9607841be 100644 --- a/examples/run_rm.py +++ b/examples/run_rm.py @@ -23,8 +23,9 @@ from nemo_rl.algorithms.rm import MasterConfig, rm_train, setup from nemo_rl.algorithms.utils import get_tokenizer -from nemo_rl.data import DataConfig, hf_datasets -from nemo_rl.data.datasets import AllTaskProcessedDataset +from nemo_rl.data import DataConfig +from nemo_rl.data.datasets import AllTaskProcessedDataset, load_preference_dataset +from nemo_rl.data.datasets.preference_datasets import PreferenceDataset from nemo_rl.data.interfaces import DatumSpec, TaskDataSpec from nemo_rl.data.llm_message_utils import get_formatted_message_log from nemo_rl.distributed.virtual_cluster import init_ray @@ -68,7 +69,7 @@ def rm_preprocessor( rejected_completion = datum_dict["completions"][0] else: raise NotImplementedError( - "Ties are not supported yet. You can use the following command to filter out ties: `cat | jq 'select(.completions[0].rank != .completions[1].rank)'`." + "Ties are not supported yet. You can use the following command to filter out ties: `cat | jq 'select(.completions[0].rank != .completions[1].rank)'`." ) messages_chosen = datum_dict["context"] + chosen_completion["completion"] @@ -120,34 +121,15 @@ def rm_preprocessor( def setup_data(tokenizer: AutoTokenizer, data_config: DataConfig): print("\n▶ Setting up data...") - data_cls = data_config["dataset_name"] - - if data_cls == "PreferenceDataset": - data_path = data_config["train_data_path"] - data = hf_datasets.PreferenceDataset(data_path, split="train") - train_dataset = data.formatted_ds["train"] - val_dataset = None - elif data_cls == "HelpSteer3": - data = hf_datasets.HelpSteer3Dataset() - train_dataset = data.formatted_ds["train"] - val_dataset = data.formatted_ds["validation"] - elif data_cls == "Tulu3Preference": - data = hf_datasets.Tulu3PreferenceDataset() - train_dataset = data.formatted_ds["train"] - val_dataset = None - else: - raise ValueError( - f"Unknown dataset class: {data_cls}. Supported datasets are: PreferenceDataset, HelpSteer3, and Tulu3Preference." - ) - if train_dataset: - print( - f" ✓ Training dataset loaded with {len(data.formatted_ds['train'])} samples." - ) + # load dataset + data = load_preference_dataset(data_config) + train_dataset = data.formatted_ds["train"] + val_dataset = data.formatted_ds["validation"] + + print(f" ✓ Training dataset loaded with {len(train_dataset)} samples.") if val_dataset: - print( - f" ✓ Validation dataset loaded with {len(data.formatted_ds['validation'])} samples." - ) + print(f" ✓ Validation dataset loaded with {len(val_dataset)} samples.") rm_task_spec = data.task_spec @@ -159,15 +141,10 @@ def setup_data(tokenizer: AutoTokenizer, data_config: DataConfig): max_seq_length=data_config["max_input_seq_length"], ) - if data_cls == "PreferenceDataset": + # TODO @yukih: unify the code when support multiple datasets for other algorithms + if "val_data_paths" in data_config and data_config["val_data_paths"]: val_dataset = {} - assert "val_data_path" not in data_config, ( - "`val_data_path` cannot be provided for PreferenceDataset. You should use `val_data_paths` instead." - ) - assert "val_data_paths" in data_config, ( - "`val_data_paths` must be provided for PreferenceDataset" - ) assert isinstance(data_config["val_data_paths"], dict), ( f"Invalid type for val_data_paths: {type(data_config['val_data_paths'])}. val_data_paths must be a dictionary." ) @@ -175,14 +152,12 @@ def setup_data(tokenizer: AutoTokenizer, data_config: DataConfig): for val_dataset_name, val_dataset_path in val_data_paths.items(): assert val_dataset_name not in val_dataset - val_data = hf_datasets.PreferenceDataset( - val_dataset_path, split="validation" - ) + val_data = PreferenceDataset(val_dataset_path) print( - f" ✓ Validation dataset '{val_dataset_name}' loaded with {len(val_data.formatted_ds['validation'])} samples." + f" ✓ Validation dataset '{val_dataset_name}' loaded with {len(val_data.formatted_ds['train'])} samples." ) val_dataset[val_dataset_name] = AllTaskProcessedDataset( - val_data.formatted_ds["validation"], + val_data.formatted_ds["train"], tokenizer, val_data.task_spec, rm_preprocessor, diff --git a/examples/run_sft.py b/examples/run_sft.py index 392dfad645e..1281d775b97 100644 --- a/examples/run_sft.py +++ b/examples/run_sft.py @@ -23,8 +23,8 @@ from nemo_rl.algorithms.sft import MasterConfig, setup, sft_train from nemo_rl.algorithms.utils import get_tokenizer -from nemo_rl.data import DataConfig, hf_datasets -from nemo_rl.data.datasets import AllTaskProcessedDataset +from nemo_rl.data import DataConfig +from nemo_rl.data.datasets import AllTaskProcessedDataset, load_response_dataset from nemo_rl.data.interfaces import DatumSpec, TaskDataSpec from nemo_rl.data.llm_message_utils import get_formatted_message_log from nemo_rl.distributed.virtual_cluster import init_ray @@ -98,55 +98,24 @@ def sft_preprocessor( def setup_data(tokenizer: AutoTokenizer, data_config: DataConfig, seed: int): print("\n▶ Setting up data...") - data_cls = data_config["dataset_name"] + # load dataset + data = load_response_dataset(data_config, seed) + train_dataset = data.formatted_ds["train"] + val_dataset = data.formatted_ds["validation"] + sft_task_spec = data.task_spec + print( + f" ✓ Training and validation datasets loaded with {len(train_dataset)} and {len(val_dataset)} samples, respectively." + ) + + # add preprocessor if needed datum_preprocessor = None - if data_cls == "open_assistant": - data = hf_datasets.OasstDataset( - output_dir="/tmp/open_assistant", - seed=seed, - ) - elif data_cls == "squad": - data = hf_datasets.SquadDataset() - elif data_cls == "prompt_response_dataset": - data = hf_datasets.PromptResponseDataset( - data_config["train_data_path"], - data_config["val_data_path"], - data_config["input_key"], - data_config["output_key"], - ) - elif data_cls == "openmathinstruct2": - data = hf_datasets.OpenMathInstruct2Dataset( - split=data_config["split"], - output_key=data_config["output_key"], - prompt_file=data_config["prompt_file"], - seed=seed, + if "dataset_name" in data_config and data_config["dataset_name"] == "clevr_cogent": + from nemo_rl.data.datasets.response_datasets.clevr import ( + format_clevr_cogent_dataset, ) - elif data_cls == "openai_format": - data = hf_datasets.OpenAIFormatDataset( - data_config["train_data_path"], - data_config["val_data_path"], - data_config["chat_key"], - data_config["system_key"], - data_config["system_prompt"], - ) - elif data_cls == "clevr_cogent": - from nemo_rl.data.hf_datasets.clevr import format_clevr_cogent_dataset - data = hf_datasets.CLEVRCoGenTDataset( - split=data_config["split"], - prompt_file=data_config["prompt_file"], - ) datum_preprocessor = partial(format_clevr_cogent_dataset, return_pil=True) - else: - raise ValueError(f"Unknown dataset class: {data_cls}") - print( - f" ✓ Training and validation datasets loaded with {len(data.formatted_ds['train'])} and {len(data.formatted_ds['validation'])} samples, respectively." - ) - - train_dataset = data.formatted_ds["train"] - val_dataset = data.formatted_ds["validation"] - sft_task_spec = data.task_spec train_dataset = AllTaskProcessedDataset( train_dataset, diff --git a/examples/run_vlm_grpo.py b/examples/run_vlm_grpo.py index fdb5fa26e99..ef69d425282 100644 --- a/examples/run_vlm_grpo.py +++ b/examples/run_vlm_grpo.py @@ -28,16 +28,10 @@ from nemo_rl.algorithms.grpo import MasterConfig, grpo_train, setup from nemo_rl.algorithms.utils import get_tokenizer from nemo_rl.data import DataConfig -from nemo_rl.data.datasets import AllTaskProcessedDataset -from nemo_rl.data.hf_datasets.clevr import ( - CLEVRCoGenTDataset, - format_clevr_cogent_dataset, -) -from nemo_rl.data.hf_datasets.geometry3k import ( - Geometry3KDataset, - format_geometry3k_dataset, -) -from nemo_rl.data.hf_datasets.refcoco import RefCOCODataset, format_refcoco_dataset +from nemo_rl.data.datasets import AllTaskProcessedDataset, load_response_dataset +from nemo_rl.data.datasets.response_datasets.clevr import format_clevr_cogent_dataset +from nemo_rl.data.datasets.response_datasets.geometry3k import format_geometry3k_dataset +from nemo_rl.data.datasets.response_datasets.refcoco import format_refcoco_dataset from nemo_rl.data.interfaces import ( DatumSpec, LLMMessageLogType, @@ -112,7 +106,7 @@ def hf_data_processor( max_seq_length: int, idx: int, ) -> DatumSpec: - """Process a datum dictionary (directly loaded from data/hf_datasets/.py) into a DatumSpec for the VLM Environment.""" + """Process a datum dictionary (directly loaded from response_datasets/.py) into a DatumSpec for the VLM Environment.""" # depending on the task, format the data differently if task_data_spec.task_name == "clevr-cogent": datum_dict = format_clevr_cogent_dataset(datum_dict) @@ -230,6 +224,7 @@ def setup_data( processor: AutoProcessor, data_config: DataConfig, env_configs: dict[str, Any], + seed: int, ) -> tuple[ AllTaskProcessedDataset, Optional[AllTaskProcessedDataset], @@ -241,23 +236,10 @@ def setup_data( task_spec contains the task name as well as prompt and system prompt modifiers that can be used by data processor """ print("\n▶ Setting up data...") - # Load CLEVR-CoGenT dataset using nemo rl datasets - # other VLM datasets can be added here - if data_config["dataset_name"] == "clevr-cogent": - data: Any = CLEVRCoGenTDataset( - split=data_config["split"], - ) - elif data_config["dataset_name"] == "refcoco": - data: Any = RefCOCODataset( - split=data_config["split"], - download_dir=data_config["download_dir"], - ) - elif data_config["dataset_name"] == "geometry3k": - data: Any = Geometry3KDataset( - split=data_config["split"], - ) - else: - raise ValueError(f"No processor for dataset {data_config['dataset_name']}.") + + # load dataset + # TODO @yukih: currently seed is not used for vlm datasets + data: Any = load_response_dataset(data_config, seed) task_name = data.task_name vlm_task_spec = TaskDataSpec( @@ -357,7 +339,7 @@ def main() -> None: val_dataset, task_to_env, val_task_to_env, - ) = setup_data(processor, config["data"], config["env"]) + ) = setup_data(processor, config["data"], config["env"], config["grpo"]["seed"]) ( policy, diff --git a/nemo_rl/algorithms/dpo.py b/nemo_rl/algorithms/dpo.py index b22f059f745..dff2c7a6f64 100644 --- a/nemo_rl/algorithms/dpo.py +++ b/nemo_rl/algorithms/dpo.py @@ -28,7 +28,8 @@ ) from nemo_rl.algorithms.utils import maybe_pad_last_batch, set_seed from nemo_rl.data import DataConfig -from nemo_rl.data.datasets import AllTaskProcessedDataset, preference_collate_fn +from nemo_rl.data.collate_fn import preference_collate_fn +from nemo_rl.data.datasets import AllTaskProcessedDataset from nemo_rl.distributed.virtual_cluster import ClusterConfig, RayVirtualCluster from nemo_rl.models.policy import PolicyConfig from nemo_rl.models.policy.interfaces import PolicyInterface diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index e9bc749d60f..ff751f2c666 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -32,29 +32,21 @@ ) from nemo_rl.algorithms.utils import calculate_baseline_and_std_per_prompt, set_seed from nemo_rl.data import DataConfig -from nemo_rl.data.datasets import AllTaskProcessedDataset, rl_collate_fn -from nemo_rl.data.interfaces import ( - DatumSpec, -) +from nemo_rl.data.collate_fn import rl_collate_fn +from nemo_rl.data.datasets import AllTaskProcessedDataset +from nemo_rl.data.interfaces import DatumSpec from nemo_rl.data.llm_message_utils import ( batched_message_log_to_flat_message, get_keys_from_message_log, ) from nemo_rl.distributed.batched_data_dict import BatchedDataDict -from nemo_rl.distributed.virtual_cluster import ( - ClusterConfig, - RayVirtualCluster, -) -from nemo_rl.environments.interfaces import ( - EnvironmentInterface, -) +from nemo_rl.distributed.virtual_cluster import ClusterConfig, RayVirtualCluster +from nemo_rl.environments.interfaces import EnvironmentInterface from nemo_rl.experience.rollouts import ( run_async_multi_turn_rollout, run_multi_turn_rollout, ) -from nemo_rl.models.generation.interfaces import ( - GenerationInterface, -) +from nemo_rl.models.generation.interfaces import GenerationInterface from nemo_rl.models.generation.vllm import VllmConfig, VllmGeneration from nemo_rl.models.policy import PolicyConfig from nemo_rl.models.policy.interfaces import ColocatablePolicyInterface diff --git a/nemo_rl/algorithms/rm.py b/nemo_rl/algorithms/rm.py index ad646d0021c..c480cff358c 100644 --- a/nemo_rl/algorithms/rm.py +++ b/nemo_rl/algorithms/rm.py @@ -28,10 +28,8 @@ ) from nemo_rl.algorithms.utils import maybe_pad_last_batch, set_seed from nemo_rl.data import DataConfig -from nemo_rl.data.datasets import ( - AllTaskProcessedDataset, - preference_collate_fn, -) +from nemo_rl.data.collate_fn import preference_collate_fn +from nemo_rl.data.datasets import AllTaskProcessedDataset from nemo_rl.data.interfaces import TaskDataSpec from nemo_rl.distributed.virtual_cluster import ClusterConfig, RayVirtualCluster from nemo_rl.models.policy import PolicyConfig diff --git a/nemo_rl/algorithms/sft.py b/nemo_rl/algorithms/sft.py index 71853294d94..879f7e9088e 100644 --- a/nemo_rl/algorithms/sft.py +++ b/nemo_rl/algorithms/sft.py @@ -26,7 +26,8 @@ ) from nemo_rl.algorithms.utils import maybe_pad_last_batch, set_seed from nemo_rl.data import DataConfig -from nemo_rl.data.datasets import AllTaskProcessedDataset, rl_collate_fn +from nemo_rl.data.collate_fn import rl_collate_fn +from nemo_rl.data.datasets import AllTaskProcessedDataset from nemo_rl.data.interfaces import TaskDataSpec from nemo_rl.data.llm_message_utils import ( add_loss_mask_to_message_log, diff --git a/nemo_rl/algorithms/utils.py b/nemo_rl/algorithms/utils.py index 40bcf1475ba..b41c9db5490 100644 --- a/nemo_rl/algorithms/utils.py +++ b/nemo_rl/algorithms/utils.py @@ -26,7 +26,7 @@ PreTrainedTokenizerBase, ) -from nemo_rl.data import hf_datasets +from nemo_rl.data.chat_templates import COMMON_CHAT_TEMPLATES from nemo_rl.models.policy import TokenizerConfig @@ -242,9 +242,7 @@ def get_tokenizer( if "chat_template" in tokenizer_config: if tokenizer_config["chat_template"] is None: print("Using passthrough chat template") - tokenizer.chat_template = ( - hf_datasets.COMMON_CHAT_TEMPLATES.passthrough_prompt_response - ) + tokenizer.chat_template = COMMON_CHAT_TEMPLATES.passthrough_prompt_response elif tokenizer_config["chat_template"].lower() == "default": print("Using tokenizer's default chat template") else: diff --git a/nemo_rl/data/hf_datasets/chat_templates.py b/nemo_rl/data/chat_templates.py similarity index 100% rename from nemo_rl/data/hf_datasets/chat_templates.py rename to nemo_rl/data/chat_templates.py diff --git a/nemo_rl/data/datasets.py b/nemo_rl/data/collate_fn.py similarity index 58% rename from nemo_rl/data/datasets.py rename to nemo_rl/data/collate_fn.py index 60dd0e091cd..e1ca4894550 100644 --- a/nemo_rl/data/datasets.py +++ b/nemo_rl/data/collate_fn.py @@ -11,18 +11,12 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Optional, Union +from typing import Any, Union import torch -from datasets import Dataset from transformers import AutoProcessor, PreTrainedTokenizerBase -from nemo_rl.data.interfaces import ( - DatumSpec, - DPODatumSpec, - TaskDataProcessFnCallable, - TaskDataSpec, -) +from nemo_rl.data.interfaces import DatumSpec, DPODatumSpec from nemo_rl.data.llm_message_utils import ( add_loss_mask_to_message_log, batched_message_log_to_flat_message, @@ -32,105 +26,6 @@ TokenizerType = Union[PreTrainedTokenizerBase, AutoProcessor] -# TODO @sahilj handle too-long prompts and masking them out throughout the whole process and renormalizing on loss -class AllTaskProcessedDataset: - """Dataset for processing single or multi-task data with task-specific tokenization and processing. - - Args: - dataset: Input dataset containing raw data - tokenizer: Tokenizer for text processing - default_task_data_spec: Default task processing specifications. - In the case of single-task, this is the spec used for processing all entries. - In the case of multi-task, any values not specified in the task-specific specs will be taken from the default spec. - task_data_processors: Either a single TaskDataProcessFnCallable for single-task, - or a dict mapping task names to (TaskDataSpec, TaskDataProcessFnCallable) for multi-task - max_seq_length: Maximum sequence length for tokenized outputs - """ - - def __init__( - self, - dataset: Dataset | Any, - tokenizer: TokenizerType, - default_task_data_spec: TaskDataSpec, - task_data_processors: ( - dict[str, tuple[TaskDataSpec, TaskDataProcessFnCallable]] - | TaskDataProcessFnCallable - ), - max_seq_length: Optional[int] = None, - ): - self.dataset = dataset - self.tokenizer = tokenizer - self.default_task_data_spec = default_task_data_spec - self.task_data_processors = task_data_processors - self.max_seq_length = max_seq_length - self._bos_checked = False - - if isinstance(task_data_processors, dict): - # apply defaults to all task data specs - for task_name, ( - task_data_spec, - task_data_processor, - ) in task_data_processors.items(): - task_data_spec.copy_defaults(self.default_task_data_spec) - - def __len__(self) -> int: - return len(self.dataset) - - def encode_single( - self, text: Union[str, list[str]] - ) -> tuple[list[int] | torch.Tensor, int]: - """Takes either a single string or a list of strings that represent multiple turns for the same conversation. - - Returns a single (concatenated) list of tokenized ids and the length of the tokenized ids. - """ - if isinstance(text, str): - text_ids = self.tokenizer.text_to_ids(text) - return text_ids, len(text_ids) - elif isinstance(text, list): - text_ids = [self.tokenizer.text_to_ids(t) for t in text] - return torch.cat(text_ids), sum(len(t) for t in text_ids) - else: - raise ValueError( - f"text must be a string or a list of strings, got {type(text)}" - ) - - def __getitem__(self, idx: int) -> DatumSpec: - """Return a single prompt.""" - entry = self.dataset[idx] - - if isinstance(self.task_data_processors, dict): - task_name = entry["task_name"] - - assert task_name in self.task_data_processors, ( - f"task processor not provided for {task_name}. Provided processors: {self.task_data_processors.keys()}" - ) - task_data_spec, task_data_processor = self.task_data_processors[task_name] - else: - task_data_spec = self.default_task_data_spec - task_data_processor = self.task_data_processors - - datum_spec = task_data_processor( - entry, task_data_spec, self.tokenizer, self.max_seq_length, idx - ) - - # Check the first processed item for BOS token assertion - if ( - not self._bos_checked - and "message_log" in datum_spec - and datum_spec["message_log"] - ): - first_message = datum_spec["message_log"][0] - if "token_ids" in first_message: - token_ids = first_message["token_ids"] - assert isinstance(token_ids, torch.Tensor), ( - f"token_ids must be a torch.Tensor, got {type(token_ids)}" - ) - assert_no_double_bos(token_ids, self.tokenizer) - self._bos_checked = True - - return datum_spec - - def rl_collate_fn(data_batch: list[DatumSpec]) -> BatchedDataDict[Any]: """Collate function for RL training.""" message_log = [datum_spec["message_log"] for datum_spec in data_batch] @@ -193,7 +88,7 @@ def eval_collate_fn(data_batch: list[DatumSpec]) -> BatchedDataDict[Any]: Examples: ```{doctest} >>> import torch - >>> from nemo_rl.data.datasets import eval_collate_fn + >>> from nemo_rl.data.collate_fn import eval_collate_fn >>> from nemo_rl.data.interfaces import DatumSpec >>> data_batch = [ ... DatumSpec( @@ -300,24 +195,3 @@ def preference_collate_fn( data["token_mask"] = cat_and_padded["token_loss_mask"] return data - - -def assert_no_double_bos(token_ids: torch.Tensor, tokenizer: TokenizerType) -> None: - """Assert that there are no double starting BOS tokens in the message. - - Args: - token_ids: List of token IDs - tokenizer: Tokenizer - """ - if tokenizer.bos_token_id is not None: - token_ids_list = token_ids.tolist() - if len(token_ids_list) > 1: - assert not ( - token_ids_list[0] == tokenizer.bos_token_id - and token_ids_list[1] == tokenizer.bos_token_id - ), "Found double BOS token in the first two positions of the message." - else: - # `name_or_path` is not available for AutoProcessor, temp fix in get_tokenizer - print( - f"skip assert_start_single_bos since Tokenizer {tokenizer.name_or_path} has no BOS token" - ) diff --git a/nemo_rl/data/datasets/__init__.py b/nemo_rl/data/datasets/__init__.py new file mode 100644 index 00000000000..f859705dbaf --- /dev/null +++ b/nemo_rl/data/datasets/__init__.py @@ -0,0 +1,26 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from nemo_rl.data.datasets.eval_datasets import load_eval_dataset +from nemo_rl.data.datasets.preference_datasets import load_preference_dataset +from nemo_rl.data.datasets.processed_dataset import AllTaskProcessedDataset +from nemo_rl.data.datasets.response_datasets import load_response_dataset +from nemo_rl.data.datasets.utils import assert_no_double_bos + +__all__ = [ + "AllTaskProcessedDataset", + "load_eval_dataset", + "load_preference_dataset", + "load_response_dataset", + "assert_no_double_bos", +] diff --git a/nemo_rl/data/eval_datasets/__init__.py b/nemo_rl/data/datasets/eval_datasets/__init__.py similarity index 81% rename from nemo_rl/data/eval_datasets/__init__.py rename to nemo_rl/data/datasets/eval_datasets/__init__.py index c735bf2e188..8386286c832 100644 --- a/nemo_rl/data/eval_datasets/__init__.py +++ b/nemo_rl/data/datasets/eval_datasets/__init__.py @@ -12,18 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. -from nemo_rl.data.eval_datasets.aime2024 import AIME2024Dataset -from nemo_rl.data.eval_datasets.aime2025 import AIME2025Dataset -from nemo_rl.data.eval_datasets.gpqa import GPQADataset -from nemo_rl.data.eval_datasets.local_math_dataset import LocalMathDataset -from nemo_rl.data.eval_datasets.math import MathDataset -from nemo_rl.data.eval_datasets.mmlu import MMLUDataset -from nemo_rl.data.eval_datasets.mmlu_pro import MMLUProDataset +from nemo_rl.data.datasets.eval_datasets.aime import AIMEDataset +from nemo_rl.data.datasets.eval_datasets.gpqa import GPQADataset +from nemo_rl.data.datasets.eval_datasets.local_math_dataset import LocalMathDataset +from nemo_rl.data.datasets.eval_datasets.math import MathDataset +from nemo_rl.data.datasets.eval_datasets.mmlu import MMLUDataset +from nemo_rl.data.datasets.eval_datasets.mmlu_pro import MMLUProDataset def load_eval_dataset(data_config): """Loads evaluation dataset.""" dataset_name = data_config["dataset_name"] + + # mmlu if dataset_name.startswith("mmlu") and dataset_name != "mmlu_pro": if dataset_name == "mmlu": base_dataset = MMLUDataset( @@ -37,16 +38,25 @@ def load_eval_dataset(data_config): prompt_file=data_config["prompt_file"], system_prompt_file=data_config["system_prompt_file"], ) + elif dataset_name == "mmlu_pro": + base_dataset = MMLUProDataset( + prompt_file=data_config["prompt_file"], + system_prompt_file=data_config["system_prompt_file"], + ) + # aime elif dataset_name == "aime2024": - base_dataset = AIME2024Dataset( + base_dataset = AIMEDataset( + variant="2024", prompt_file=data_config["prompt_file"], system_prompt_file=data_config["system_prompt_file"], ) elif dataset_name == "aime2025": - base_dataset = AIME2025Dataset( + base_dataset = AIMEDataset( + variant="2025", prompt_file=data_config["prompt_file"], system_prompt_file=data_config["system_prompt_file"], ) + # gpqa elif dataset_name == "gpqa": base_dataset = GPQADataset( variant="main", @@ -59,11 +69,7 @@ def load_eval_dataset(data_config): prompt_file=data_config["prompt_file"], system_prompt_file=data_config["system_prompt_file"], ) - elif dataset_name == "mmlu_pro": - base_dataset = MMLUProDataset( - prompt_file=data_config["prompt_file"], - system_prompt_file=data_config["system_prompt_file"], - ) + # math elif dataset_name == "math": base_dataset = MathDataset( variant="math_test", @@ -76,10 +82,11 @@ def load_eval_dataset(data_config): prompt_file=data_config["prompt_file"], system_prompt_file=data_config["system_prompt_file"], ) - elif dataset_name == "local": + # fall back to local dataset + else: + print(f"Loading dataset from {dataset_name}...") base_dataset = LocalMathDataset( - name=dataset_name, - data_paths=data_config["data_paths"], + data_path=dataset_name, problem_key=data_config["problem_key"], solution_key=data_config["solution_key"], file_format=data_config["file_format"], @@ -87,14 +94,12 @@ def load_eval_dataset(data_config): prompt_file=data_config["prompt_file"], system_prompt_file=data_config["system_prompt_file"], ) - else: - raise ValueError(f"Unknown dataset {dataset_name}.") + return base_dataset __all__ = [ - "AIME2024Dataset", - "AIME2025Dataset", + "AIMEDataset", "GPQADataset", "LocalMathDataset", "MathDataset", diff --git a/nemo_rl/data/eval_datasets/aime2025.py b/nemo_rl/data/datasets/eval_datasets/aime.py similarity index 63% rename from nemo_rl/data/eval_datasets/aime2025.py rename to nemo_rl/data/datasets/eval_datasets/aime.py index 43db69cf1bf..671b410a69e 100644 --- a/nemo_rl/data/eval_datasets/aime2025.py +++ b/nemo_rl/data/datasets/eval_datasets/aime.py @@ -12,9 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""AIME 2025 dataset.""" +"""AIME dataset.""" -from typing import Any, Optional +from typing import Any, Literal, Optional from datasets import concatenate_datasets, load_dataset @@ -22,18 +22,27 @@ from nemo_rl.data.interfaces import TaskDataSpec -class AIME2025Dataset: +class AIMEDataset: def __init__( self, + variant: Literal["2024", "2025"] = "2025", prompt_file: Optional[str] = None, system_prompt_file: Optional[str] = None, ): - ds0 = load_dataset("opencompass/AIME2025", "AIME2025-I", split="test") - ds1 = load_dataset("opencompass/AIME2025", "AIME2025-II", split="test") - ds = concatenate_datasets([ds0, ds1]) + if variant == "2024": + ds = load_dataset("HuggingFaceH4/aime_2024", split="train") + self.input_key = "problem" + elif variant == "2025": + ds0 = load_dataset("opencompass/AIME2025", "AIME2025-I", split="test") + ds1 = load_dataset("opencompass/AIME2025", "AIME2025-II", split="test") + ds = concatenate_datasets([ds0, ds1]) + self.input_key = "question" + else: + raise ValueError(f"Invalid variant for aime dataset: aime{variant}") + self.rekeyed_ds = ds.map(self._rekey, remove_columns=ds.column_names) self.task_spec = TaskDataSpec( - task_name="aime2025", + task_name=f"aime{variant}", prompt_file=prompt_file, system_prompt_file=system_prompt_file, ) @@ -41,6 +50,6 @@ def __init__( def _rekey(self, data: dict[str, Any]): return { - "problem": data["question"], + "problem": data[self.input_key], "expected_answer": data["answer"], } diff --git a/nemo_rl/data/eval_datasets/gpqa.py b/nemo_rl/data/datasets/eval_datasets/gpqa.py similarity index 100% rename from nemo_rl/data/eval_datasets/gpqa.py rename to nemo_rl/data/datasets/eval_datasets/gpqa.py diff --git a/nemo_rl/data/eval_datasets/local_math_dataset.py b/nemo_rl/data/datasets/eval_datasets/local_math_dataset.py similarity index 91% rename from nemo_rl/data/eval_datasets/local_math_dataset.py rename to nemo_rl/data/datasets/eval_datasets/local_math_dataset.py index 2810899b4aa..7ea90ca498a 100644 --- a/nemo_rl/data/eval_datasets/local_math_dataset.py +++ b/nemo_rl/data/datasets/eval_datasets/local_math_dataset.py @@ -14,6 +14,7 @@ """Local math dataset.""" +import os from typing import Any, Literal, Optional from datasets import load_dataset @@ -25,23 +26,22 @@ class LocalMathDataset: def __init__( self, - data_paths: str | list[str], + data_path: str, problem_key: str, solution_key: str, - name: str, split: Optional[str] = None, file_format: Literal["csv", "json"] = "csv", prompt_file: Optional[str] = None, system_prompt_file: Optional[str] = None, ): - ds = load_dataset(file_format, data_files=data_paths) + ds = load_dataset(file_format, data_files=data_path) if split is not None: ds = ds[split] self._problem_key = problem_key self._solution_key = solution_key self.rekeyed_ds = ds.map(self._rekey, remove_columns=ds.column_names) self.task_spec = TaskDataSpec( - task_name=name, + task_name=os.path.basename(data_path).split(".")[0], prompt_file=prompt_file, system_prompt_file=system_prompt_file, ) diff --git a/nemo_rl/data/eval_datasets/math.py b/nemo_rl/data/datasets/eval_datasets/math.py similarity index 100% rename from nemo_rl/data/eval_datasets/math.py rename to nemo_rl/data/datasets/eval_datasets/math.py diff --git a/nemo_rl/data/eval_datasets/mmlu.py b/nemo_rl/data/datasets/eval_datasets/mmlu.py similarity index 100% rename from nemo_rl/data/eval_datasets/mmlu.py rename to nemo_rl/data/datasets/eval_datasets/mmlu.py diff --git a/nemo_rl/data/eval_datasets/mmlu_pro.py b/nemo_rl/data/datasets/eval_datasets/mmlu_pro.py similarity index 100% rename from nemo_rl/data/eval_datasets/mmlu_pro.py rename to nemo_rl/data/datasets/eval_datasets/mmlu_pro.py diff --git a/nemo_rl/data/datasets/preference_datasets/__init__.py b/nemo_rl/data/datasets/preference_datasets/__init__.py new file mode 100644 index 00000000000..6593d0bd7a5 --- /dev/null +++ b/nemo_rl/data/datasets/preference_datasets/__init__.py @@ -0,0 +1,86 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from nemo_rl.data.datasets.preference_datasets.binary_preference_dataset import ( + BinaryPreferenceDataset, +) +from nemo_rl.data.datasets.preference_datasets.helpsteer3 import HelpSteer3Dataset +from nemo_rl.data.datasets.preference_datasets.preference_dataset import ( + PreferenceDataset, +) +from nemo_rl.data.datasets.preference_datasets.tulu3 import Tulu3PreferenceDataset +from nemo_rl.data.datasets.utils import get_extra_kwargs + + +def load_preference_dataset(data_config): + """Loads preference dataset.""" + dataset_name = data_config["dataset_name"] + + if dataset_name == "HelpSteer3": + base_dataset = HelpSteer3Dataset() + elif dataset_name == "Tulu3Preference": + base_dataset = Tulu3PreferenceDataset() + # fall back to load from JSON file + elif dataset_name == "BinaryPreferenceDataset": + if "train_data_path" not in data_config: + raise ValueError( + "train_data_path is required for dataset_name=BinaryPreferenceDataset." + ) + extra_kwargs = get_extra_kwargs( + data_config, + [ + "val_data_path", + "prompt_key", + "chosen_key", + "rejected_key", + "train_split", + "val_split", + ], + ) + base_dataset = BinaryPreferenceDataset( + train_data_path=data_config["train_data_path"], + **extra_kwargs, + ) + elif dataset_name == "PreferenceDataset": + if "train_data_path" not in data_config: + raise ValueError( + "train_data_path is required for dataset_name=PreferenceDataset." + ) + extra_kwargs = get_extra_kwargs( + data_config, + [ + "val_data_path", + "train_split", + "val_split", + ], + ) + base_dataset = PreferenceDataset( + train_data_path=data_config["train_data_path"], + **extra_kwargs, + ) + else: + raise ValueError( + f"Unsupported {dataset_name=}. " + "Please either set dataset_name in {'HelpSteer3', 'Tulu3Preference'} to use a built-in dataset " + "or set dataset_name in {'PreferenceDataset', 'BinaryPreferenceDataset'} to load from local JSONL file or HuggingFace." + ) + + return base_dataset + + +__all__ = [ + "BinaryPreferenceDataset", + "HelpSteer3Dataset", + "PreferenceDataset", + "Tulu3PreferenceDataset", +] diff --git a/nemo_rl/data/datasets/preference_datasets/binary_preference_dataset.py b/nemo_rl/data/datasets/preference_datasets/binary_preference_dataset.py new file mode 100644 index 00000000000..1a166cdd7e2 --- /dev/null +++ b/nemo_rl/data/datasets/preference_datasets/binary_preference_dataset.py @@ -0,0 +1,110 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import Any, Optional + +from nemo_rl.data.datasets.utils import load_dataset_from_path +from nemo_rl.data.interfaces import TaskDataSpec + + +def to_preference_data_format( + data: dict[str, Any], prompt_key: str, chosen_key: str, rejected_key: str +) -> dict[str, list[dict[str, Any]]]: + return { + "context": data[prompt_key] + if isinstance(data[prompt_key], list) + else [{"role": "user", "content": data[prompt_key]}], + "completions": [ + { + "rank": 0, + "completion": [{"role": "assistant", "content": data[chosen_key]}], + }, + { + "rank": 1, + "completion": [{"role": "assistant", "content": data[rejected_key]}], + }, + ], + } + + +class BinaryPreferenceDataset: + """Dataset class for binary preference data which can be loaded from a JSON file. + + This class handles loading of preference data for DPO and RM training. + It will be converted to the format of PreferenceDataset through the `to_preference_data_format` function. + + The input JSONL files should contain valid JSON objects formatted like this: + { + prompt_key: str, # The input prompt/context + chosen_key: str, # The preferred/winning response + rejected_key: str, # The non-preferred/losing response + } + + Args: + train_data_path: Path to the JSON file containing training data + val_data_path: Path to the JSON file containing validation data + prompt_key: Key for the input prompt/context, default is "prompt" + chosen_key: Key for the preferred/winning response, default is "chosen" + rejected_key: Key for the non-preferred/losing response, default is "rejected" + 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 + """ + + def __init__( + self, + train_data_path: str, + val_data_path: Optional[str] = None, + prompt_key: str = "prompt", + chosen_key: str = "chosen", + rejected_key: str = "rejected", + train_split: Optional[str] = None, + val_split: Optional[str] = None, + ): + self.prompt_key = prompt_key + self.chosen_key = chosen_key + self.rejected_key = rejected_key + + # load from json file or huggingface + train_ds = load_dataset_from_path(train_data_path, train_split) + if val_data_path: + val_ds = load_dataset_from_path(val_data_path, val_split) + else: + val_ds = None + + # format the dataset + # convert to PreferenceDataset format + train_ds = train_ds.map( + to_preference_data_format, + fn_kwargs={ + "prompt_key": prompt_key, + "chosen_key": chosen_key, + "rejected_key": rejected_key, + }, + ) + if val_ds: + val_ds = val_ds.map( + to_preference_data_format, + fn_kwargs={ + "prompt_key": prompt_key, + "chosen_key": chosen_key, + "rejected_key": rejected_key, + }, + ) + + # store the formatted dataset + self.formatted_ds = { + "train": train_ds, + "validation": val_ds, + } + + self.task_spec = TaskDataSpec(task_name="BinaryPreferenceDataset") diff --git a/nemo_rl/data/hf_datasets/helpsteer3.py b/nemo_rl/data/datasets/preference_datasets/helpsteer3.py similarity index 100% rename from nemo_rl/data/hf_datasets/helpsteer3.py rename to nemo_rl/data/datasets/preference_datasets/helpsteer3.py diff --git a/nemo_rl/data/datasets/preference_datasets/preference_dataset.py b/nemo_rl/data/datasets/preference_datasets/preference_dataset.py new file mode 100644 index 00000000000..a168fc23601 --- /dev/null +++ b/nemo_rl/data/datasets/preference_datasets/preference_dataset.py @@ -0,0 +1,61 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import Optional + +from nemo_rl.data.datasets.utils import load_dataset_from_path +from nemo_rl.data.interfaces import TaskDataSpec + + +class PreferenceDataset: + """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) + } + } + + Args: + 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 + """ + + def __init__( + self, + train_data_path: str, + val_data_path: Optional[str] = None, + train_split: Optional[str] = None, + val_split: Optional[str] = None, + ): + # load from json file or huggingface + train_ds = load_dataset_from_path(train_data_path, train_split) + if val_data_path: + val_ds = load_dataset_from_path(val_data_path, val_split) + else: + val_ds = None + + # store the formatted dataset + self.formatted_ds = { + "train": train_ds, + "validation": val_ds, + } + + self.task_spec = TaskDataSpec(task_name="PreferenceDataset") diff --git a/nemo_rl/data/hf_datasets/tulu3.py b/nemo_rl/data/datasets/preference_datasets/tulu3.py similarity index 96% rename from nemo_rl/data/hf_datasets/tulu3.py rename to nemo_rl/data/datasets/preference_datasets/tulu3.py index 266daea186c..20e381da7e1 100644 --- a/nemo_rl/data/hf_datasets/tulu3.py +++ b/nemo_rl/data/datasets/preference_datasets/tulu3.py @@ -73,6 +73,8 @@ def __init__(self) -> None: trust_remote_code=True, ) self.formatted_ds = ds.map(to_preference_data_format) + # Tulu3 preference dataset has no validation set + self.formatted_ds["validation"] = None self.task_spec = TaskDataSpec( task_name="Tulu3Preference", diff --git a/nemo_rl/data/datasets/processed_dataset.py b/nemo_rl/data/datasets/processed_dataset.py new file mode 100644 index 00000000000..906ab591fc7 --- /dev/null +++ b/nemo_rl/data/datasets/processed_dataset.py @@ -0,0 +1,126 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import Any, Optional, Union + +import torch +from datasets import Dataset +from transformers import AutoProcessor, PreTrainedTokenizerBase + +from nemo_rl.data.datasets.utils import assert_no_double_bos +from nemo_rl.data.interfaces import ( + DatumSpec, + TaskDataProcessFnCallable, + TaskDataSpec, +) + +TokenizerType = Union[PreTrainedTokenizerBase, AutoProcessor] + + +# TODO @sahilj handle too-long prompts and masking them out throughout the whole process and renormalizing on loss +class AllTaskProcessedDataset: + """Dataset for processing single or multi-task data with task-specific tokenization and processing. + + Args: + dataset: Input dataset containing raw data + tokenizer: Tokenizer for text processing + default_task_data_spec: Default task processing specifications. + In the case of single-task, this is the spec used for processing all entries. + In the case of multi-task, any values not specified in the task-specific specs will be taken from the default spec. + task_data_processors: Either a single TaskDataProcessFnCallable for single-task, + or a dict mapping task names to (TaskDataSpec, TaskDataProcessFnCallable) for multi-task + max_seq_length: Maximum sequence length for tokenized outputs + """ + + def __init__( + self, + dataset: Dataset | Any, + tokenizer: TokenizerType, + default_task_data_spec: TaskDataSpec, + task_data_processors: ( + dict[str, tuple[TaskDataSpec, TaskDataProcessFnCallable]] + | TaskDataProcessFnCallable + ), + max_seq_length: Optional[int] = None, + ): + self.dataset = dataset + self.tokenizer = tokenizer + self.default_task_data_spec = default_task_data_spec + self.task_data_processors = task_data_processors + self.max_seq_length = max_seq_length + self._bos_checked = False + + if isinstance(task_data_processors, dict): + # apply defaults to all task data specs + for task_name, ( + task_data_spec, + task_data_processor, + ) in task_data_processors.items(): + task_data_spec.copy_defaults(self.default_task_data_spec) + + def __len__(self) -> int: + return len(self.dataset) + + def encode_single( + self, text: Union[str, list[str]] + ) -> tuple[list[int] | torch.Tensor, int]: + """Takes either a single string or a list of strings that represent multiple turns for the same conversation. + + Returns a single (concatenated) list of tokenized ids and the length of the tokenized ids. + """ + if isinstance(text, str): + text_ids = self.tokenizer.text_to_ids(text) + return text_ids, len(text_ids) + elif isinstance(text, list): + text_ids = [self.tokenizer.text_to_ids(t) for t in text] + return torch.cat(text_ids), sum(len(t) for t in text_ids) + else: + raise ValueError( + f"text must be a string or a list of strings, got {type(text)}" + ) + + def __getitem__(self, idx: int) -> DatumSpec: + """Return a single prompt.""" + entry = self.dataset[idx] + + if isinstance(self.task_data_processors, dict): + task_name = entry["task_name"] + + assert task_name in self.task_data_processors, ( + f"task processor not provided for {task_name}. Provided processors: {self.task_data_processors.keys()}" + ) + task_data_spec, task_data_processor = self.task_data_processors[task_name] + else: + task_data_spec = self.default_task_data_spec + task_data_processor = self.task_data_processors + + datum_spec = task_data_processor( + entry, task_data_spec, self.tokenizer, self.max_seq_length, idx + ) + + # Check the first processed item for BOS token assertion + if ( + not self._bos_checked + and "message_log" in datum_spec + and datum_spec["message_log"] + ): + first_message = datum_spec["message_log"][0] + if "token_ids" in first_message: + token_ids = first_message["token_ids"] + assert isinstance(token_ids, torch.Tensor), ( + f"token_ids must be a torch.Tensor, got {type(token_ids)}" + ) + assert_no_double_bos(token_ids, self.tokenizer) + self._bos_checked = True + + return datum_spec diff --git a/nemo_rl/data/datasets/response_datasets/__init__.py b/nemo_rl/data/datasets/response_datasets/__init__.py new file mode 100644 index 00000000000..d3e63b9ddac --- /dev/null +++ b/nemo_rl/data/datasets/response_datasets/__init__.py @@ -0,0 +1,128 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import Any + +from nemo_rl.data.datasets.response_datasets.clevr import CLEVRCoGenTDataset +from nemo_rl.data.datasets.response_datasets.deepscaler import DeepScalerDataset +from nemo_rl.data.datasets.response_datasets.geometry3k import Geometry3KDataset +from nemo_rl.data.datasets.response_datasets.oai_format_dataset import ( + OpenAIFormatDataset, +) +from nemo_rl.data.datasets.response_datasets.oasst import OasstDataset +from nemo_rl.data.datasets.response_datasets.openmathinstruct2 import ( + OpenMathInstruct2Dataset, +) +from nemo_rl.data.datasets.response_datasets.refcoco import RefCOCODataset +from nemo_rl.data.datasets.response_datasets.response_dataset import ResponseDataset +from nemo_rl.data.datasets.response_datasets.squad import SquadDataset +from nemo_rl.data.datasets.utils import get_extra_kwargs + + +def load_response_dataset(data_config, seed: int = 42): + """Loads response dataset.""" + dataset_name = data_config["dataset_name"] + + # TODO @yukih: remove duplicated dataset_name (openmathinstruct2, clevr_cogent) + # for sft training + if dataset_name == "open_assistant": + base_dataset = OasstDataset( + output_dir="/tmp/open_assistant", + seed=seed, + ) + elif dataset_name == "squad": + base_dataset = SquadDataset() + elif dataset_name == "openmathinstruct2": + base_dataset = OpenMathInstruct2Dataset( + split=data_config["split"], + output_key=data_config["output_key"], + prompt_file=data_config["prompt_file"], + seed=seed, + ) + elif dataset_name == "clevr_cogent": + base_dataset = CLEVRCoGenTDataset( + split=data_config["split"], + prompt_file=data_config["prompt_file"], + ) + elif dataset_name == "openai_format": + base_dataset = OpenAIFormatDataset( + data_config["train_data_path"], + data_config["val_data_path"], + data_config["chat_key"], + data_config["system_key"], + data_config["system_prompt"], + ) + # for rl training + elif dataset_name == "OpenMathInstruct-2": + print("Loading nvidia/OpenMathInstruct2Dataset for training and validation") + base_dataset: Any = OpenMathInstruct2Dataset(seed=seed) + elif dataset_name == "DeepScaler": + print( + "Loading agentica-org/DeepScaleR-Preview-Dataset for training and validation" + ) + base_dataset: Any = DeepScalerDataset(seed=seed) + # for vlm rl training + elif dataset_name == "clevr-cogent": + base_dataset: Any = CLEVRCoGenTDataset( + split=data_config["split"], + ) + elif dataset_name == "refcoco": + base_dataset: Any = RefCOCODataset( + split=data_config["split"], + download_dir=data_config["download_dir"], + ) + elif dataset_name == "geometry3k": + base_dataset: Any = Geometry3KDataset( + split=data_config["split"], + ) + # fall back to load from JSON file + elif dataset_name == "ResponseDataset": + if "train_data_path" not in data_config: + raise ValueError( + "train_data_path is required when dataset_name is not one of the built-ins." + ) + extra_kwargs = get_extra_kwargs( + data_config, + [ + "val_data_path", + "input_key", + "output_key", + "train_split", + "val_split", + ], + ) + base_dataset = ResponseDataset( + train_data_path=data_config["train_data_path"], + **extra_kwargs, + ) + else: + raise ValueError( + f"Unsupported {dataset_name=}. " + "Please either use a built-in dataset " + "or set dataset_name=ResponseDataset to load from local JSONL file or HuggingFace." + ) + + return base_dataset + + +__all__ = [ + "CLEVRCoGenTDataset", + "DeepScalerDataset", + "Geometry3KDataset", + "OpenAIFormatDataset", + "OasstDataset", + "OpenMathInstruct2Dataset", + "RefCOCODataset", + "ResponseDataset", + "SquadDataset", +] diff --git a/nemo_rl/data/hf_datasets/clevr.py b/nemo_rl/data/datasets/response_datasets/clevr.py similarity index 87% rename from nemo_rl/data/hf_datasets/clevr.py rename to nemo_rl/data/datasets/response_datasets/clevr.py index b3281f586f6..843bf34414d 100644 --- a/nemo_rl/data/hf_datasets/clevr.py +++ b/nemo_rl/data/datasets/response_datasets/clevr.py @@ -12,32 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -import base64 -import io from typing import Any, Optional from datasets import load_dataset -from PIL import Image +from nemo_rl.data.datasets.utils import pil_to_base64 from nemo_rl.data.interfaces import TaskDataSpec -def pil_to_base64(image: Image.Image, format: str = "PNG") -> str: - """Converts a PIL Image object to a base64 encoded string. - - Args: - image: The PIL Image object to convert. - format: The image format (e.g., "PNG", "JPEG"). Defaults to "PNG". - - Returns: - A base64 encoded string representation of the image. - """ - buffered = io.BytesIO() - image.save(buffered, format=format) - img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") - return f"data:image/png;base64,{img_str}" - - def format_answer_fromtags(answer: str) -> str: """Extract content between tags and strip whitespace.""" import re diff --git a/nemo_rl/data/hf_datasets/deepscaler.py b/nemo_rl/data/datasets/response_datasets/deepscaler.py similarity index 100% rename from nemo_rl/data/hf_datasets/deepscaler.py rename to nemo_rl/data/datasets/response_datasets/deepscaler.py diff --git a/nemo_rl/data/hf_datasets/geometry3k.py b/nemo_rl/data/datasets/response_datasets/geometry3k.py similarity index 98% rename from nemo_rl/data/hf_datasets/geometry3k.py rename to nemo_rl/data/datasets/response_datasets/geometry3k.py index b1501477984..deaac1e8e76 100644 --- a/nemo_rl/data/hf_datasets/geometry3k.py +++ b/nemo_rl/data/datasets/response_datasets/geometry3k.py @@ -15,7 +15,7 @@ from datasets import load_dataset -from nemo_rl.data.hf_datasets.clevr import pil_to_base64 +from nemo_rl.data.datasets.utils import pil_to_base64 from nemo_rl.data.interfaces import TaskDataSpec diff --git a/nemo_rl/data/hf_datasets/oai_format_dataset.py b/nemo_rl/data/datasets/response_datasets/oai_format_dataset.py similarity index 100% rename from nemo_rl/data/hf_datasets/oai_format_dataset.py rename to nemo_rl/data/datasets/response_datasets/oai_format_dataset.py diff --git a/nemo_rl/data/hf_datasets/oasst.py b/nemo_rl/data/datasets/response_datasets/oasst.py similarity index 100% rename from nemo_rl/data/hf_datasets/oasst.py rename to nemo_rl/data/datasets/response_datasets/oasst.py diff --git a/nemo_rl/data/hf_datasets/openmathinstruct2.py b/nemo_rl/data/datasets/response_datasets/openmathinstruct2.py similarity index 100% rename from nemo_rl/data/hf_datasets/openmathinstruct2.py rename to nemo_rl/data/datasets/response_datasets/openmathinstruct2.py diff --git a/nemo_rl/data/hf_datasets/refcoco.py b/nemo_rl/data/datasets/response_datasets/refcoco.py similarity index 99% rename from nemo_rl/data/hf_datasets/refcoco.py rename to nemo_rl/data/datasets/response_datasets/refcoco.py index b5de9f68912..87d5d7372f4 100644 --- a/nemo_rl/data/hf_datasets/refcoco.py +++ b/nemo_rl/data/datasets/response_datasets/refcoco.py @@ -23,7 +23,7 @@ from PIL import Image from tqdm import tqdm # Using tqdm for progress bar, install with: pip install tqdm -from nemo_rl.data.hf_datasets.clevr import pil_to_base64 +from nemo_rl.data.datasets.utils import pil_to_base64 from nemo_rl.data.interfaces import TaskDataSpec diff --git a/nemo_rl/data/datasets/response_datasets/response_dataset.py b/nemo_rl/data/datasets/response_datasets/response_dataset.py new file mode 100644 index 00000000000..6dbd7752b84 --- /dev/null +++ b/nemo_rl/data/datasets/response_datasets/response_dataset.py @@ -0,0 +1,80 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, Optional + +from nemo_rl.data.datasets.utils import load_dataset_from_path +from nemo_rl.data.interfaces import TaskDataSpec + + +class ResponseDataset: + """Dataset class for response data which can be loaded from a JSON file. + + This class handles loading of response data for SFT and RL training. + The input JSONL files should contain valid JSON objects formatted like this: + { + input_key: str, # The input prompt/context + output_key: str, # The output response/answer + } + + Args: + train_data_path: Path to the JSON file containing training data + val_data_path: Path to the JSON file containing validation data + input_key: Key for the input text + output_key: Key for the output text + 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 + """ + + def __init__( + self, + train_data_path: str, + val_data_path: Optional[str] = None, + input_key: str = "input", + output_key: str = "output", + train_split: Optional[str] = None, + val_split: Optional[str] = None, + ): + self.input_key = input_key + self.output_key = output_key + + # load from json file or huggingface + train_ds = load_dataset_from_path(train_data_path, train_split) + if val_data_path: + val_ds = load_dataset_from_path(val_data_path, val_split) + else: + val_ds = None + + # format the dataset + train_ds = train_ds.map(self.add_messages_key) + if val_ds: + val_ds = val_ds.map(self.add_messages_key) + + # store the formatted dataset + self.formatted_ds = { + "train": train_ds, + "validation": val_ds, + } + + self.task_spec = TaskDataSpec(task_name="ResponseDataset") + + def add_messages_key( + self, example: dict[str, Any] + ) -> dict[str, list[dict[str, Any]]]: + return { + "messages": [ + {"role": "user", "content": example[self.input_key]}, + {"role": "assistant", "content": example[self.output_key]}, + ] + } diff --git a/nemo_rl/data/hf_datasets/squad.py b/nemo_rl/data/datasets/response_datasets/squad.py similarity index 100% rename from nemo_rl/data/hf_datasets/squad.py rename to nemo_rl/data/datasets/response_datasets/squad.py diff --git a/nemo_rl/data/datasets/utils.py b/nemo_rl/data/datasets/utils.py new file mode 100644 index 00000000000..12c65ef14f0 --- /dev/null +++ b/nemo_rl/data/datasets/utils.py @@ -0,0 +1,102 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import base64 +import io +import os +from typing import Optional, Union + +import torch +from datasets import DatasetDict, load_dataset +from PIL import Image +from transformers import AutoProcessor, PreTrainedTokenizerBase + +TokenizerType = Union[PreTrainedTokenizerBase, AutoProcessor] + + +def assert_no_double_bos(token_ids: torch.Tensor, tokenizer: TokenizerType) -> None: + """Assert that there are no double starting BOS tokens in the message. + + Args: + token_ids: List of token IDs + tokenizer: Tokenizer + """ + if tokenizer.bos_token_id is not None: + token_ids_list = token_ids.tolist() + if len(token_ids_list) > 1: + assert not ( + token_ids_list[0] == tokenizer.bos_token_id + and token_ids_list[1] == tokenizer.bos_token_id + ), "Found double BOS token in the first two positions of the message." + else: + # `name_or_path` is not available for AutoProcessor, temp fix in get_tokenizer + print( + f"skip assert_start_single_bos since Tokenizer {tokenizer.name_or_path} has no BOS token" + ) + + +def pil_to_base64(image: Image.Image, format: str = "PNG") -> str: + """Converts a PIL Image object to a base64 encoded string. + + Args: + image: The PIL Image object to convert. + format: The image format (e.g., "PNG", "JPEG"). Defaults to "PNG". + + Returns: + A base64 encoded string representation of the image. + """ + buffered = io.BytesIO() + image.save(buffered, format=format) + img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") + return f"data:image/png;base64,{img_str}" + + +def load_dataset_from_path(data_path: str, data_split: Optional[str] = "train"): + """Load a dataset from a json or huggingface dataset. + + Args: + data_path: The path to the dataset. + data_split: The split to load from the dataset. + """ + suffix = os.path.splitext(data_path)[-1] + if suffix in [".json", ".jsonl"]: + raw_dataset = load_dataset("json", data_files=data_path) + else: + raw_dataset = load_dataset(data_path) + + if data_split: + raw_dataset = raw_dataset[data_split] + # if the dataset doesn't contain split, load_dataset will use "train" as default + elif isinstance(raw_dataset, DatasetDict) and "train" in raw_dataset: + raw_dataset = raw_dataset["train"] + + return raw_dataset + + +def get_extra_kwargs(data_config: dict, keys: list[str]) -> dict: + """Get extra kwargs from the data config. + + If the key is not in the data config, it will be ignored. + + Args: + data_config: The data config. + keys: The keys to get from the data config. + + Returns: + The extra kwargs. + """ + extra_kwargs = {} + for key in keys: + if key in data_config: + extra_kwargs[key] = data_config[key] + return extra_kwargs diff --git a/nemo_rl/data/eval_datasets/aime2024.py b/nemo_rl/data/eval_datasets/aime2024.py deleted file mode 100644 index 9e585bb5119..00000000000 --- a/nemo_rl/data/eval_datasets/aime2024.py +++ /dev/null @@ -1,44 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""AIME 2024 dataset.""" - -from typing import Any, Optional - -from datasets import load_dataset - -from nemo_rl.data import processors -from nemo_rl.data.interfaces import TaskDataSpec - - -class AIME2024Dataset: - def __init__( - self, - prompt_file: Optional[str] = None, - system_prompt_file: Optional[str] = None, - ): - ds = load_dataset("HuggingFaceH4/aime_2024", split="train") - self.rekeyed_ds = ds.map(self._rekey, remove_columns=ds.column_names) - self.task_spec = TaskDataSpec( - task_name="aime2024", - prompt_file=prompt_file, - system_prompt_file=system_prompt_file, - ) - self.processor = processors.math_data_processor - - def _rekey(self, data: dict[str, Any]): - return { - "problem": data["problem"], - "expected_answer": data["answer"], - } diff --git a/nemo_rl/data/hf_datasets/__init__.py b/nemo_rl/data/hf_datasets/__init__.py deleted file mode 100644 index 3c3d7c91fa7..00000000000 --- a/nemo_rl/data/hf_datasets/__init__.py +++ /dev/null @@ -1,41 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from nemo_rl.data.hf_datasets.chat_templates import COMMON_CHAT_TEMPLATES -from nemo_rl.data.hf_datasets.clevr import CLEVRCoGenTDataset -from nemo_rl.data.hf_datasets.dpo import DPODataset -from nemo_rl.data.hf_datasets.helpsteer3 import HelpSteer3Dataset -from nemo_rl.data.hf_datasets.oai_format_dataset import OpenAIFormatDataset -from nemo_rl.data.hf_datasets.oasst import OasstDataset -from nemo_rl.data.hf_datasets.openmathinstruct2 import OpenMathInstruct2Dataset -from nemo_rl.data.hf_datasets.preference_dataset import PreferenceDataset -from nemo_rl.data.hf_datasets.prompt_response_dataset import ( - PromptResponseDataset, -) -from nemo_rl.data.hf_datasets.squad import SquadDataset -from nemo_rl.data.hf_datasets.tulu3 import Tulu3PreferenceDataset - -__all__ = [ - "DPODataset", - "HelpSteer3Dataset", - "OasstDataset", - "OpenAIFormatDataset", - "OpenMathInstruct2Dataset", - "PreferenceDataset", - "PromptResponseDataset", - "SquadDataset", - "Tulu3PreferenceDataset", - "COMMON_CHAT_TEMPLATES", - "CLEVRCoGenTDataset", -] diff --git a/nemo_rl/data/hf_datasets/dpo.py b/nemo_rl/data/hf_datasets/dpo.py deleted file mode 100644 index d0c96a7e217..00000000000 --- a/nemo_rl/data/hf_datasets/dpo.py +++ /dev/null @@ -1,81 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import warnings -from typing import Any - -from datasets import load_dataset - -from nemo_rl.data.interfaces import TaskDataSpec - - -def to_preference_data_format(data: dict[str, Any]) -> dict[str, list[dict[str, Any]]]: - return { - "context": data["prompt"] - if isinstance(data["prompt"], list) - else [{"role": "user", "content": data["prompt"]}], - "completions": [ - { - "rank": 0, - "completion": [ - {"role": "assistant", "content": data["chosen_response"]} - ], - }, - { - "rank": 1, - "completion": [ - {"role": "assistant", "content": data["rejected_response"]} - ], - }, - ], - } - - -class DPODataset: - """Dataset class for Direct Preference Optimization (DPO) training. - - This class is deprecated and will be removed in a future version. Use PreferenceDataset instead. - - This class handles loading of preference data for DPO training. - The input JSON files should contain examples with the following structure: - { - "prompt": str, # The input prompt/context - "chosen_response": str, # The preferred/winning response - "rejected_response": str # The non-preferred/losing response - } - - Args: - train_data_path (str): Path to the JSON file containing training data - val_data_path (str): Path to the JSON file containing validation data - - """ - - def __init__(self, train_data_path: str, val_data_path: str): - warnings.warn( - "DPODataset is deprecated and will be removed in a future version. Use PreferenceDataset instead (see function `to_preference_data_format()` on how to convert your data to this new format).", - category=DeprecationWarning, - stacklevel=2, - ) - - self.formatted_ds = { - "train": load_dataset( - "json", data_files=train_data_path, split="train" - ).map(to_preference_data_format), - "validation": load_dataset( - "json", data_files=val_data_path, split="train" - ).map(to_preference_data_format), - } - - self.task_spec = TaskDataSpec( - task_name="DPO", - ) diff --git a/nemo_rl/data/hf_datasets/preference_dataset.py b/nemo_rl/data/hf_datasets/preference_dataset.py deleted file mode 100644 index dea4d7213d6..00000000000 --- a/nemo_rl/data/hf_datasets/preference_dataset.py +++ /dev/null @@ -1,43 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from datasets import DatasetDict, load_dataset - -from nemo_rl.data.interfaces import TaskDataSpec - - -class PreferenceDataset: - """Preference dataset. - - This class handles loading of custom preference data. - The input JSONL file 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) - } - } - """ - - def __init__(self, dataset_path: str, split: str) -> None: - # Specifying split="train" returns Dataset instead of DatasetDict({"train": Dataset}) - self.formatted_ds = DatasetDict( - {split: load_dataset("json", data_files=dataset_path, split="train")} - ) - - self.task_spec = TaskDataSpec( - task_name="PreferenceDataset", - ) diff --git a/nemo_rl/data/hf_datasets/prompt_response_dataset.py b/nemo_rl/data/hf_datasets/prompt_response_dataset.py deleted file mode 100644 index 398459d2b29..00000000000 --- a/nemo_rl/data/hf_datasets/prompt_response_dataset.py +++ /dev/null @@ -1,56 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from typing import Any - -from datasets import load_dataset - -from nemo_rl.data.interfaces import TaskDataSpec - - -class PromptResponseDataset: - def __init__( - self, - train_ds_path: str, - val_ds_path: str, - input_key: str = "input", - output_key: str = "output", - ): - train_original_dataset = load_dataset("json", data_files=train_ds_path)["train"] - val_original_dataset = load_dataset("json", data_files=val_ds_path)["train"] - - self.input_key = input_key - self.output_key = output_key - - formatted_train_dataset = train_original_dataset.map(self.add_messages_key) - formatted_val_dataset = val_original_dataset.map(self.add_messages_key) - - self.formatted_ds = { - "train": formatted_train_dataset, - "validation": formatted_val_dataset, - } - - self.task_spec = TaskDataSpec( - "json_dataset", - ) - - def add_messages_key( - self, example: dict[str, Any] - ) -> dict[str, list[dict[str, Any]]]: - return { - "messages": [ - {"role": "user", "content": example[self.input_key]}, - {"role": "assistant", "content": example[self.output_key]}, - ] - } diff --git a/nemo_rl/evals/eval.py b/nemo_rl/evals/eval.py index 32881d77066..0ea3cf64aa0 100644 --- a/nemo_rl/evals/eval.py +++ b/nemo_rl/evals/eval.py @@ -26,7 +26,8 @@ from nemo_rl.algorithms.utils import set_seed from nemo_rl.data import MathDataConfig -from nemo_rl.data.datasets import AllTaskProcessedDataset, eval_collate_fn +from nemo_rl.data.collate_fn import eval_collate_fn +from nemo_rl.data.datasets import AllTaskProcessedDataset from nemo_rl.data.llm_message_utils import get_keys_from_message_log from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.distributed.virtual_cluster import ClusterConfig, RayVirtualCluster diff --git a/pyrefly.toml b/pyrefly.toml index 51b55740020..e11b1fd6c2c 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -41,32 +41,35 @@ project-includes = [ "nemo_rl/converters/megatron/__init__.py", "nemo_rl/converters/megatron/vllm_export.py", "nemo_rl/data/__init__.py", - "nemo_rl/data/datasets.py", - "nemo_rl/data/eval_datasets/__init__.py", - "nemo_rl/data/eval_datasets/aime2024.py", - "nemo_rl/data/eval_datasets/aime2025.py", - "nemo_rl/data/eval_datasets/gpqa.py", - "nemo_rl/data/eval_datasets/local_math_dataset.py", - "nemo_rl/data/eval_datasets/math.py", - "nemo_rl/data/eval_datasets/mmlu.py", - "nemo_rl/data/eval_datasets/mmlu_pro.py", - "nemo_rl/data/hf_datasets/__init__.py", - "nemo_rl/data/hf_datasets/clevr.py", - "nemo_rl/data/hf_datasets/geometry3k.py", - "nemo_rl/data/hf_datasets/refcoco.py", - "nemo_rl/data/hf_datasets/chat_templates.py", - "nemo_rl/data/hf_datasets/deepscaler.py", - "nemo_rl/data/hf_datasets/dpo.py", - "nemo_rl/data/hf_datasets/helpsteer3.py", - "nemo_rl/data/hf_datasets/oai_format_dataset.py", - "nemo_rl/data/hf_datasets/oasst.py", - "nemo_rl/data/hf_datasets/openmathinstruct2.py", - "nemo_rl/data/hf_datasets/preference_dataset.py", - "nemo_rl/data/hf_datasets/prompt_response_dataset.py", - "nemo_rl/data/hf_datasets/squad.py", - "nemo_rl/data/hf_datasets/tulu3.py", - "nemo_rl/data/multimodal_utils.py", + "nemo_rl/data/chat_templates.py", + "nemo_rl/data/collate_fn.py", + "nemo_rl/data/datasets/__init__.py", + "nemo_rl/data/datasets/eval_datasets/__init__.py", + "nemo_rl/data/datasets/eval_datasets/aime.py", + "nemo_rl/data/datasets/eval_datasets/gpqa.py", + "nemo_rl/data/datasets/eval_datasets/local_math_dataset.py", + "nemo_rl/data/datasets/eval_datasets/math.py", + "nemo_rl/data/datasets/eval_datasets/mmlu.py", + "nemo_rl/data/datasets/eval_datasets/mmlu_pro.py", + "nemo_rl/data/datasets/preference_datasets/__init__.py", + "nemo_rl/data/datasets/preference_datasets/binary_preference_dataset.py", + "nemo_rl/data/datasets/preference_datasets/helpsteer3.py", + "nemo_rl/data/datasets/preference_datasets/preference_dataset.py", + "nemo_rl/data/datasets/preference_datasets/tulu3.py", + "nemo_rl/data/datasets/processed_dataset.py", + "nemo_rl/data/datasets/response_datasets/__init__.py", + "nemo_rl/data/datasets/response_datasets/clevr.py", + "nemo_rl/data/datasets/response_datasets/deepscaler.py", + "nemo_rl/data/datasets/response_datasets/geometry3k.py", + "nemo_rl/data/datasets/response_datasets/oai_format_dataset.py", + "nemo_rl/data/datasets/response_datasets/oasst.py", + "nemo_rl/data/datasets/response_datasets/openmathinstruct2.py", + "nemo_rl/data/datasets/response_datasets/refcoco.py", + "nemo_rl/data/datasets/response_datasets/response_dataset.py", + "nemo_rl/data/datasets/response_datasets/squad.py", + "nemo_rl/data/datasets/utils.py", "nemo_rl/data/interfaces.py", + "nemo_rl/data/multimodal_utils.py", "nemo_rl/data/packing/__init__.py", "nemo_rl/data/processors.py", "nemo_rl/distributed/__init__.py", @@ -76,13 +79,13 @@ project-includes = [ "nemo_rl/distributed/virtual_cluster.py", "nemo_rl/distributed/worker_group_utils.py", "nemo_rl/environments/__init__.py", - "nemo_rl/environments/rewards.py", - "nemo_rl/environments/vlm_environment.py", "nemo_rl/environments/games/sliding_puzzle.py", "nemo_rl/environments/interfaces.py", "nemo_rl/environments/math_environment.py", "nemo_rl/environments/metrics.py", + "nemo_rl/environments/rewards.py", "nemo_rl/environments/utils.py", + "nemo_rl/environments/vlm_environment.py", "nemo_rl/evals/__init__.py", "nemo_rl/evals/answer_parsing.py", "nemo_rl/experience/__init__.py", diff --git a/tests/unit/algorithms/test_utils.py b/tests/unit/algorithms/test_utils.py index a2c8a5036e6..7562e541edd 100755 --- a/tests/unit/algorithms/test_utils.py +++ b/tests/unit/algorithms/test_utils.py @@ -18,7 +18,7 @@ import torch from nemo_rl.algorithms.utils import get_tokenizer, maybe_pad_last_batch -from nemo_rl.data.hf_datasets.chat_templates import COMMON_CHAT_TEMPLATES +from nemo_rl.data.chat_templates import COMMON_CHAT_TEMPLATES from nemo_rl.distributed.batched_data_dict import BatchedDataDict diff --git a/tests/unit/data/datasets/test_eval_dataset.py b/tests/unit/data/datasets/test_eval_dataset.py new file mode 100644 index 00000000000..dc567d09d75 --- /dev/null +++ b/tests/unit/data/datasets/test_eval_dataset.py @@ -0,0 +1,116 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import pytest +from transformers import AutoTokenizer + +from nemo_rl.data.datasets import load_eval_dataset + + +@pytest.mark.skip(reason="dataset download is flaky") +def test_gpqa_dataset(): + # load the dataset + data_config = { + "dataset_name": "gpqa", + "prompt_file": None, + "system_prompt_file": None, + } + gpqa_dataset = load_eval_dataset(data_config) + + # load the tokenizer + tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-1.5B-Instruct") + + # check that the dataset is formatted correctly + for example in gpqa_dataset.rekeyed_ds.take(5): + assert "question" in example + assert "options" in example + assert "answer" in example + + ## check that applying chat template works as expected + default_templated = tokenizer.apply_chat_template( + [{"role": "user", "content": example["question"]}], + tokenize=False, + add_generation_prompt=False, + add_special_tokens=False, + ) + + assert ( + default_templated + == f"<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n<|im_start|>user\n{example['question']}<|im_end|>\n" + ) + + +@pytest.mark.skip(reason="dataset download is flaky") +def test_math_dataset(): + # load the dataset + data_config = { + "dataset_name": "math", + "prompt_file": None, + "system_prompt_file": None, + } + math_dataset = load_eval_dataset(data_config) + + # load the tokenizer + tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-1.5B-Instruct") + + # check that the dataset is formatted correctly + for example in math_dataset.rekeyed_ds.take(5): + assert "problem" in example + assert "expected_answer" in example + + ## check that applying chat template works as expected + default_templated = tokenizer.apply_chat_template( + [{"role": "user", "content": example["problem"]}], + tokenize=False, + add_generation_prompt=False, + add_special_tokens=False, + ) + + assert ( + default_templated + == f"<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n<|im_start|>user\n{example['problem']}<|im_end|>\n" + ) + + +@pytest.mark.skip(reason="dataset download is flaky") +def test_mmlu_dataset(): + # load the dataset + data_config = { + "dataset_name": "mmlu", + "prompt_file": None, + "system_prompt_file": None, + } + mmlu_dataset = load_eval_dataset(data_config) + + # load the tokenizer + tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-1.5B-Instruct") + + # check that the dataset is formatted correctly + for example in mmlu_dataset.rekeyed_ds.take(5): + assert "question" in example + assert "options" in example + assert "answer" in example + assert "subject" in example + + ## check that applying chat template works as expected + default_templated = tokenizer.apply_chat_template( + [{"role": "user", "content": example["question"]}], + tokenize=False, + add_generation_prompt=False, + add_special_tokens=False, + ) + + assert ( + default_templated + == f"<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n<|im_start|>user\n{example['question']}<|im_end|>\n" + ) diff --git a/tests/unit/data/hf_datasets/test_helpsteer.py b/tests/unit/data/datasets/test_helpsteer.py similarity index 97% rename from tests/unit/data/hf_datasets/test_helpsteer.py rename to tests/unit/data/datasets/test_helpsteer.py index 4015a83b72f..337a2b630d0 100644 --- a/tests/unit/data/hf_datasets/test_helpsteer.py +++ b/tests/unit/data/datasets/test_helpsteer.py @@ -15,8 +15,8 @@ import pytest -from nemo_rl.data.hf_datasets.helpsteer3 import ( - HelpSteer3Dataset, +from nemo_rl.data.datasets.preference_datasets import HelpSteer3Dataset +from nemo_rl.data.datasets.preference_datasets.helpsteer3 import ( to_preference_data_format, ) diff --git a/tests/unit/data/hf_datasets/test_oai_format_dataset.py b/tests/unit/data/datasets/test_oai_format_dataset.py similarity index 96% rename from tests/unit/data/hf_datasets/test_oai_format_dataset.py rename to tests/unit/data/datasets/test_oai_format_dataset.py index ae6b8787797..aad989ed15a 100644 --- a/tests/unit/data/hf_datasets/test_oai_format_dataset.py +++ b/tests/unit/data/datasets/test_oai_format_dataset.py @@ -18,10 +18,8 @@ import pytest from transformers import AutoTokenizer -from nemo_rl.data.hf_datasets.chat_templates import COMMON_CHAT_TEMPLATES -from nemo_rl.data.hf_datasets.oai_format_dataset import ( - OpenAIFormatDataset, -) +from nemo_rl.data.chat_templates import COMMON_CHAT_TEMPLATES +from nemo_rl.data.datasets.response_datasets import OpenAIFormatDataset @pytest.fixture diff --git a/tests/unit/data/datasets/test_preference_dataset.py b/tests/unit/data/datasets/test_preference_dataset.py new file mode 100644 index 00000000000..c4633e08b86 --- /dev/null +++ b/tests/unit/data/datasets/test_preference_dataset.py @@ -0,0 +1,240 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import os +import tempfile + +import pytest + +from nemo_rl.data.datasets import load_preference_dataset + + +@pytest.fixture +def mock_preference_data(): + """Create temporary preference dataset files with sample data.""" + preference_data = [ + { + "context": [{"role": "user", "content": "What is 2+2?"}], + "completions": [ + { + "rank": 1, + "completion": [ + {"role": "assistant", "content": "The answer is 4."} + ], + }, + { + "rank": 2, + "completion": [{"role": "assistant", "content": "I don't know."}], + }, + ], + }, + { + "context": [{"role": "user", "content": "What is the capital of France?"}], + "completions": [ + { + "rank": 1, + "completion": [ + { + "role": "assistant", + "content": "The capital of France is Paris.", + } + ], + }, + { + "rank": 2, + "completion": [ + { + "role": "assistant", + "content": "The capital of France is London.", + } + ], + }, + ], + }, + ] + + with tempfile.NamedTemporaryFile( + mode="w", suffix=".json", delete=False + ) as preference_file: + json.dump(preference_data, preference_file) + preference_path = preference_file.name + + try: + yield preference_path + finally: + # Cleanup + os.unlink(preference_path) + + +def test_preference_dataset_initialization(mock_preference_data): + """Test that PreferenceDataset initializes correctly with valid data files.""" + # Load the dataset + data_config = { + "dataset_name": "PreferenceDataset", + "train_data_path": mock_preference_data, + } + dataset = load_preference_dataset(data_config) + + # Verify dataset initialization + assert dataset.task_spec.task_name == "PreferenceDataset" + + # Verify formatted_ds structure + assert "train" in dataset.formatted_ds + assert len(dataset.formatted_ds["train"]) == 2 + + +def test_preference_dataset_data_format(mock_preference_data): + """Test that PreferenceDataset correctly loads and formats the data.""" + # Load the dataset + data_config = { + "dataset_name": "PreferenceDataset", + "train_data_path": mock_preference_data, + } + dataset = load_preference_dataset(data_config) + + # Verify data format + sample = dataset.formatted_ds["train"][0] + assert "context" in sample + assert "completions" in sample + + # Verify context structure + assert isinstance(sample["context"], list) + assert len(sample["context"]) == 1 + assert "role" in sample["context"][0] + assert "content" in sample["context"][0] + + # Verify completions structure + assert isinstance(sample["completions"], list) + assert len(sample["completions"]) == 2 + + for completion in sample["completions"]: + assert "rank" in completion + assert "completion" in completion + assert isinstance(completion["rank"], int) + assert isinstance(completion["completion"], list) + + +@pytest.fixture +def mock_binary_preference_data(): + """Create temporary chosen_rejected dataset files with sample data.""" + train_data = [ + { + "prompt": "What is 2+2?", + "chosen_response": "The answer is 4.", + "rejected_response": "I don't know.", + }, + { + "prompt": "What is the capital of France?", + "chosen_response": "The capital of France is Paris.", + "rejected_response": "The capital of France is London.", + }, + ] + + val_data = [ + { + "prompt": "What is 3*3?", + "chosen_response": "The answer is 9.", + "rejected_response": "The answer is 6.", + } + ] + + with tempfile.NamedTemporaryFile( + mode="w", suffix=".json", delete=False + ) as train_file: + json.dump(train_data, train_file) + train_path = train_file.name + + with tempfile.NamedTemporaryFile( + mode="w", suffix=".json", delete=False + ) as val_file: + json.dump(val_data, val_file) + val_path = val_file.name + + try: + yield train_path, val_path + finally: + # Cleanup + os.unlink(train_path) + os.unlink(val_path) + + +def test_binary_preference_dataset_initialization(mock_binary_preference_data): + """Test that PreferenceDataset initializes correctly with valid data files.""" + # Load the dataset + train_path, val_path = mock_binary_preference_data + data_config = { + "dataset_name": "BinaryPreferenceDataset", + "train_data_path": train_path, + "val_data_path": val_path, + "prompt_key": "prompt", + "chosen_key": "chosen_response", + "rejected_key": "rejected_response", + } + dataset = load_preference_dataset(data_config) + + # Verify dataset initialization + assert dataset.task_spec.task_name == "BinaryPreferenceDataset" + + # Verify formatted_ds structure + assert "train" in dataset.formatted_ds + assert "validation" in dataset.formatted_ds + + assert len(dataset.formatted_ds["train"]) == 2 + assert len(dataset.formatted_ds["validation"]) == 1 + + +def test_binary_preference_dataset_invalid_files(): + """Test that PreferenceDataset raises appropriate errors with invalid files.""" + with pytest.raises(FileNotFoundError): + data_config = { + "dataset_name": "BinaryPreferenceDataset", + "train_data_path": "nonexistent.json", + "val_data_path": "nonexistent.json", + "prompt_key": "prompt", + "chosen_key": "chosen_response", + "rejected_key": "rejected_response", + } + load_preference_dataset(data_config) + + +def test_binary_preference_dataset_data_format(mock_binary_preference_data): + """Test that PreferenceDataset correctly formats the data.""" + # Load the dataset + train_path, val_path = mock_binary_preference_data + data_config = { + "dataset_name": "BinaryPreferenceDataset", + "train_data_path": train_path, + "val_data_path": val_path, + "prompt_key": "prompt", + "chosen_key": "chosen_response", + "rejected_key": "rejected_response", + } + dataset = load_preference_dataset(data_config) + + # Verify data format + train_sample = dataset.formatted_ds["train"][0] + assert "context" in train_sample + assert "completions" in train_sample + + # Verify data content + print(train_sample["completions"]) + assert train_sample["context"] == [{"content": "What is 2+2?", "role": "user"}] + assert train_sample["completions"] == [ + { + "completion": [{"content": "The answer is 4.", "role": "assistant"}], + "rank": 0, + }, + {"completion": [{"content": "I don't know.", "role": "assistant"}], "rank": 1}, + ] diff --git a/tests/unit/data/hf_datasets/test_prompt_response.py b/tests/unit/data/datasets/test_response_dataset.py similarity index 54% rename from tests/unit/data/hf_datasets/test_prompt_response.py rename to tests/unit/data/datasets/test_response_dataset.py index cbf18977a4a..e8c935da4f1 100644 --- a/tests/unit/data/hf_datasets/test_prompt_response.py +++ b/tests/unit/data/datasets/test_response_dataset.py @@ -18,10 +18,8 @@ import pytest from transformers import AutoTokenizer -from nemo_rl.data.hf_datasets.chat_templates import COMMON_CHAT_TEMPLATES -from nemo_rl.data.hf_datasets.prompt_response_dataset import ( - PromptResponseDataset, -) +from nemo_rl.data.chat_templates import COMMON_CHAT_TEMPLATES +from nemo_rl.data.datasets import load_response_dataset @pytest.fixture @@ -56,8 +54,14 @@ def sample_data(request): @pytest.mark.parametrize("sample_data", [("input", "output")], indirect=True) def test_dataset_initialization(sample_data): + # load the dataset train_path, val_path = sample_data - dataset = PromptResponseDataset(train_path, val_path) + data_config = { + "dataset_name": "ResponseDataset", + "train_data_path": train_path, + "val_data_path": val_path, + } + dataset = load_response_dataset(data_config) assert dataset.input_key == "input" assert dataset.output_key == "output" @@ -67,10 +71,16 @@ def test_dataset_initialization(sample_data): @pytest.mark.parametrize("sample_data", [("question", "answer")], indirect=True) def test_custom_keys(sample_data): + # load the dataset train_path, val_path = sample_data - dataset = PromptResponseDataset( - train_path, val_path, input_key="question", output_key="answer" - ) + data_config = { + "dataset_name": "ResponseDataset", + "train_data_path": train_path, + "val_data_path": val_path, + "input_key": "question", + "output_key": "answer", + } + dataset = load_response_dataset(data_config) assert dataset.input_key == "question" assert dataset.output_key == "answer" @@ -79,10 +89,16 @@ def test_custom_keys(sample_data): @pytest.mark.hf_gated @pytest.mark.parametrize("sample_data", [("question", "answer")], indirect=True) def test_message_formatting(sample_data): + # load the dataset train_path, val_path = sample_data - dataset = PromptResponseDataset( - train_path, val_path, input_key="question", output_key="answer" - ) + data_config = { + "dataset_name": "ResponseDataset", + "train_data_path": train_path, + "val_data_path": val_path, + "input_key": "question", + "output_key": "answer", + } + dataset = load_response_dataset(data_config) first_example = dataset.formatted_ds["train"][0] @@ -105,3 +121,47 @@ def test_message_formatting(sample_data): assert combined_message == "".join( message["content"] for message in first_example["messages"] ) + + +@pytest.mark.hf_gated +@pytest.mark.skip(reason="dataset download is flaky") +def test_squad_dataset(): + # load the dataset + data_config = { + "dataset_name": "squad", + "prompt_file": None, + "system_prompt_file": None, + } + squad_dataset = load_response_dataset(data_config) + + # load the tokenizer + tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct") + + # check that the dataset is formatted correctly + for example in squad_dataset.formatted_ds["train"].take(5): + assert "messages" in example + assert len(example["messages"]) == 3 + + assert example["messages"][0]["role"] == "system" + assert example["messages"][1]["role"] == "user" + assert example["messages"][2]["role"] == "assistant" + + template = "{% for message in messages %}{%- if message['role'] == 'system' %}{{'Context: ' + message['content'].strip()}}{%- elif message['role'] == 'user' %}{{' Question: ' + message['content'].strip() + ' Answer:'}}{%- elif message['role'] == 'assistant' %}{{' ' + message['content'].strip()}}{%- endif %}{% endfor %}" + + ## check that applying chat template works as expected + default_templated = tokenizer.apply_chat_template( + example["messages"], + chat_template=template, + tokenize=False, + add_generation_prompt=False, + add_special_tokens=False, + ) + + assert default_templated == ( + "Context: " + + example["messages"][0]["content"] + + " Question: " + + example["messages"][1]["content"] + + " Answer: " + + example["messages"][2]["content"] + ) diff --git a/tests/unit/data/hf_datasets/test_tulu3.py b/tests/unit/data/datasets/test_tulu3.py similarity index 93% rename from tests/unit/data/hf_datasets/test_tulu3.py rename to tests/unit/data/datasets/test_tulu3.py index d5ccf2d2546..2379bee6465 100644 --- a/tests/unit/data/hf_datasets/test_tulu3.py +++ b/tests/unit/data/datasets/test_tulu3.py @@ -15,10 +15,8 @@ import pytest -from nemo_rl.data.hf_datasets.tulu3 import ( - Tulu3PreferenceDataset, - to_preference_data_format, -) +from nemo_rl.data.datasets.preference_datasets import Tulu3PreferenceDataset +from nemo_rl.data.datasets.preference_datasets.tulu3 import to_preference_data_format @pytest.fixture(scope="module") diff --git a/tests/unit/data/eval_datasets/test_gpqa.py b/tests/unit/data/eval_datasets/test_gpqa.py deleted file mode 100644 index 3441f119742..00000000000 --- a/tests/unit/data/eval_datasets/test_gpqa.py +++ /dev/null @@ -1,42 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import pytest -from transformers import AutoTokenizer - -from nemo_rl.data.eval_datasets.gpqa import GPQADataset - - -@pytest.mark.skip(reason="dataset download is flaky") -def test_gpqa_dataset(): - tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-1.5B-Instruct") - gpqa_dataset = GPQADataset() - - # check that the dataset is formatted correctly - for example in gpqa_dataset.rekeyed_ds.take(5): - assert "question" in example - assert "options" in example - assert "answer" in example - - ## check that applying chat template works as expected - default_templated = tokenizer.apply_chat_template( - [{"role": "user", "content": example["question"]}], - tokenize=False, - add_generation_prompt=False, - add_special_tokens=False, - ) - - assert ( - default_templated - == f"<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n<|im_start|>user\n{example['question']}<|im_end|>\n" - ) diff --git a/tests/unit/data/eval_datasets/test_math.py b/tests/unit/data/eval_datasets/test_math.py deleted file mode 100644 index 3bab184f1a7..00000000000 --- a/tests/unit/data/eval_datasets/test_math.py +++ /dev/null @@ -1,41 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import pytest -from transformers import AutoTokenizer - -from nemo_rl.data.eval_datasets.math import MathDataset - - -@pytest.mark.skip(reason="dataset download is flaky") -def test_math_dataset(): - tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-1.5B-Instruct") - math_dataset = MathDataset() - - # check that the dataset is formatted correctly - for example in math_dataset.rekeyed_ds.take(5): - assert "problem" in example - assert "expected_answer" in example - - ## check that applying chat template works as expected - default_templated = tokenizer.apply_chat_template( - [{"role": "user", "content": example["problem"]}], - tokenize=False, - add_generation_prompt=False, - add_special_tokens=False, - ) - - assert ( - default_templated - == f"<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n<|im_start|>user\n{example['problem']}<|im_end|>\n" - ) diff --git a/tests/unit/data/eval_datasets/test_mmlu.py b/tests/unit/data/eval_datasets/test_mmlu.py deleted file mode 100644 index 02c19360030..00000000000 --- a/tests/unit/data/eval_datasets/test_mmlu.py +++ /dev/null @@ -1,43 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import pytest -from transformers import AutoTokenizer - -from nemo_rl.data.eval_datasets.mmlu import MMLUDataset - - -@pytest.mark.skip(reason="dataset download is flaky") -def test_mmlu_dataset(): - tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-1.5B-Instruct") - mmlu_dataset = MMLUDataset() - - # check that the dataset is formatted correctly - for example in mmlu_dataset.rekeyed_ds.take(5): - assert "question" in example - assert "options" in example - assert "answer" in example - assert "subject" in example - - ## check that applying chat template works as expected - default_templated = tokenizer.apply_chat_template( - [{"role": "user", "content": example["question"]}], - tokenize=False, - add_generation_prompt=False, - add_special_tokens=False, - ) - - assert ( - default_templated - == f"<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n<|im_start|>user\n{example['question']}<|im_end|>\n" - ) diff --git a/tests/unit/data/hf_datasets/test_dpo_dataset.py b/tests/unit/data/hf_datasets/test_dpo_dataset.py deleted file mode 100644 index 85261ff9583..00000000000 --- a/tests/unit/data/hf_datasets/test_dpo_dataset.py +++ /dev/null @@ -1,109 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import json -import os -import tempfile - -import pytest - -from nemo_rl.data.hf_datasets.dpo import DPODataset - - -@pytest.fixture -def mock_dpo_data(): - """Create temporary DPO dataset files with sample data.""" - train_data = [ - { - "prompt": "What is 2+2?", - "chosen_response": "The answer is 4.", - "rejected_response": "I don't know.", - }, - { - "prompt": "What is the capital of France?", - "chosen_response": "The capital of France is Paris.", - "rejected_response": "The capital of France is London.", - }, - ] - - val_data = [ - { - "prompt": "What is 3*3?", - "chosen_response": "The answer is 9.", - "rejected_response": "The answer is 6.", - } - ] - - train_ctx = tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) - val_ctx = tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) - - with tempfile.NamedTemporaryFile( - mode="w", suffix=".json", delete=False - ) as train_file: - json.dump(train_data, train_file) - train_path = train_file.name - with tempfile.NamedTemporaryFile( - mode="w", suffix=".json", delete=False - ) as val_file: - json.dump(val_data, val_file) - val_path = val_file.name - yield train_path, val_path - # Cleanup - os.unlink(train_path) - os.unlink(val_path) - - -def test_dpo_dataset_initialization(mock_dpo_data): - """Test that DPODataset initializes correctly with valid data files.""" - train_path, val_path = mock_dpo_data - - dataset = DPODataset(train_data_path=train_path, val_data_path=val_path) - - # Verify dataset initialization - assert dataset.task_spec.task_name == "DPO" - - # Verify formatted_ds structure - assert "train" in dataset.formatted_ds - assert "validation" in dataset.formatted_ds - - assert len(dataset.formatted_ds["train"]) == 2 - assert len(dataset.formatted_ds["validation"]) == 1 - - -def test_dpo_dataset_invalid_files(): - """Test that DPODataset raises appropriate errors with invalid files.""" - with pytest.raises(FileNotFoundError): - DPODataset(train_data_path="nonexistent.json", val_data_path="nonexistent.json") - - -def test_dpo_dataset_data_format(mock_dpo_data): - """Test that DPODataset correctly formats the data.""" - train_path, val_path = mock_dpo_data - dataset = DPODataset(train_data_path=train_path, val_data_path=val_path) - - # Verify data format - train_sample = dataset.formatted_ds["train"][0] - assert "context" in train_sample - assert "completions" in train_sample - - # Verify data content - print(train_sample["completions"]) - assert train_sample["context"] == [{"content": "What is 2+2?", "role": "user"}] - assert train_sample["completions"] == [ - { - "completion": [{"content": "The answer is 4.", "role": "assistant"}], - "rank": 0, - }, - {"completion": [{"content": "I don't know.", "role": "assistant"}], "rank": 1}, - ] diff --git a/tests/unit/data/hf_datasets/test_preference_dataset.py b/tests/unit/data/hf_datasets/test_preference_dataset.py deleted file mode 100644 index 955a91809f7..00000000000 --- a/tests/unit/data/hf_datasets/test_preference_dataset.py +++ /dev/null @@ -1,119 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import json -import os -import tempfile - -import pytest - -from nemo_rl.data.hf_datasets.preference_dataset import PreferenceDataset - - -@pytest.fixture -def mock_preference_data(): - """Create temporary preference dataset files with sample data.""" - preference_data = [ - { - "context": [{"role": "user", "content": "What is 2+2?"}], - "completions": [ - { - "rank": 1, - "completion": [ - {"role": "assistant", "content": "The answer is 4."} - ], - }, - { - "rank": 2, - "completion": [{"role": "assistant", "content": "I don't know."}], - }, - ], - }, - { - "context": [{"role": "user", "content": "What is the capital of France?"}], - "completions": [ - { - "rank": 1, - "completion": [ - { - "role": "assistant", - "content": "The capital of France is Paris.", - } - ], - }, - { - "rank": 2, - "completion": [ - { - "role": "assistant", - "content": "The capital of France is London.", - } - ], - }, - ], - }, - ] - - with tempfile.NamedTemporaryFile( - mode="w", suffix=".json", delete=False - ) as preference_file: - json.dump(preference_data, preference_file) - preference_path = preference_file.name - - try: - yield preference_path - finally: - # Cleanup - os.unlink(preference_path) - - -def test_preference_dataset_initialization(mock_preference_data): - """Test that PreferenceDataset initializes correctly with valid data files.""" - preference_path = mock_preference_data - - dataset = PreferenceDataset(dataset_path=preference_path, split="train") - - # Verify dataset initialization - assert dataset.task_spec.task_name == "PreferenceDataset" - - # Verify formatted_ds structure - assert "train" in dataset.formatted_ds - assert len(dataset.formatted_ds["train"]) == 2 - - -def test_preference_dataset_data_format(mock_preference_data): - """Test that PreferenceDataset correctly loads and formats the data.""" - preference_path = mock_preference_data - dataset = PreferenceDataset(dataset_path=preference_path, split="train") - - # Verify data format - sample = dataset.formatted_ds["train"][0] - assert "context" in sample - assert "completions" in sample - - # Verify context structure - assert isinstance(sample["context"], list) - assert len(sample["context"]) == 1 - assert "role" in sample["context"][0] - assert "content" in sample["context"][0] - - # Verify completions structure - assert isinstance(sample["completions"], list) - assert len(sample["completions"]) == 2 - - for completion in sample["completions"]: - assert "rank" in completion - assert "completion" in completion - assert isinstance(completion["rank"], int) - assert isinstance(completion["completion"], list) diff --git a/tests/unit/data/hf_datasets/test_squad.py b/tests/unit/data/hf_datasets/test_squad.py deleted file mode 100644 index f5e01b250a6..00000000000 --- a/tests/unit/data/hf_datasets/test_squad.py +++ /dev/null @@ -1,53 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import pytest -from transformers import AutoTokenizer - -from nemo_rl.data.hf_datasets.squad import SquadDataset - - -@pytest.mark.hf_gated -@pytest.mark.skip(reason="dataset download is flaky") -def test_squad_dataset(): - tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct") - squad_dataset = SquadDataset() - - # check that the dataset is formatted correctly - for example in squad_dataset.formatted_ds["train"].take(5): - assert "messages" in example - assert len(example["messages"]) == 3 - - assert example["messages"][0]["role"] == "system" - assert example["messages"][1]["role"] == "user" - assert example["messages"][2]["role"] == "assistant" - - template = "{% for message in messages %}{%- if message['role'] == 'system' %}{{'Context: ' + message['content'].strip()}}{%- elif message['role'] == 'user' %}{{' Question: ' + message['content'].strip() + ' Answer:'}}{%- elif message['role'] == 'assistant' %}{{' ' + message['content'].strip()}}{%- endif %}{% endfor %}" - - ## check that applying chat template works as expected - default_templated = tokenizer.apply_chat_template( - example["messages"], - chat_template=template, - tokenize=False, - add_generation_prompt=False, - add_special_tokens=False, - ) - - assert default_templated == ( - "Context: " - + example["messages"][0]["content"] - + " Question: " - + example["messages"][1]["content"] - + " Answer: " - + example["messages"][2]["content"] - ) diff --git a/tests/unit/data/test_datasets.py b/tests/unit/data/test_collate_fn.py similarity index 98% rename from tests/unit/data/test_datasets.py rename to tests/unit/data/test_collate_fn.py index 6cdd8203b57..e531f4cb5e7 100755 --- a/tests/unit/data/test_datasets.py +++ b/tests/unit/data/test_collate_fn.py @@ -16,7 +16,7 @@ import torch -from nemo_rl.data.datasets import preference_collate_fn +from nemo_rl.data.collate_fn import preference_collate_fn from nemo_rl.data.interfaces import DatumSpec from nemo_rl.distributed.batched_data_dict import BatchedDataDict diff --git a/tests/unit/data/test_data_processor.py b/tests/unit/data/test_data_processor.py index f161d1f7a57..f68f6fbf68e 100644 --- a/tests/unit/data/test_data_processor.py +++ b/tests/unit/data/test_data_processor.py @@ -26,15 +26,16 @@ from examples.run_grpo_math import hf_data_processor from nemo_rl.algorithms.utils import get_tokenizer from nemo_rl.data.datasets import AllTaskProcessedDataset -from nemo_rl.data.eval_datasets import ( - AIME2024Dataset, - AIME2025Dataset, +from nemo_rl.data.datasets.eval_datasets import ( + AIMEDataset, GPQADataset, MathDataset, MMLUDataset, ) -from nemo_rl.data.hf_datasets.deepscaler import DeepScalerDataset -from nemo_rl.data.hf_datasets.openmathinstruct2 import OpenMathInstruct2Dataset +from nemo_rl.data.datasets.response_datasets import ( + DeepScalerDataset, + OpenMathInstruct2Dataset, +) from nemo_rl.data.interfaces import TaskDataProcessFnCallable, TaskDataSpec from nemo_rl.data.processors import math_data_processor from nemo_rl.models.policy import TokenizerConfig @@ -154,11 +155,10 @@ def system_prompt_file(request): @pytest.mark.parametrize( "dataset_cls", [ - MMLUDataset, + AIMEDataset, GPQADataset, MathDataset, - AIME2024Dataset, - AIME2025Dataset, + MMLUDataset, ], ) @pytest.mark.parametrize( diff --git a/tests/unit/data/test_data_shuffle_reproducity.py b/tests/unit/data/test_data_shuffle_reproducity.py index 3821423d165..49e48d578c4 100644 --- a/tests/unit/data/test_data_shuffle_reproducity.py +++ b/tests/unit/data/test_data_shuffle_reproducity.py @@ -22,8 +22,9 @@ from examples.run_grpo_math import hf_data_processor from nemo_rl.algorithms.utils import get_tokenizer, set_seed -from nemo_rl.data.datasets import AllTaskProcessedDataset, rl_collate_fn -from nemo_rl.data.hf_datasets.openmathinstruct2 import OpenMathInstruct2Dataset +from nemo_rl.data.collate_fn import rl_collate_fn +from nemo_rl.data.datasets import AllTaskProcessedDataset +from nemo_rl.data.datasets.response_datasets import OpenMathInstruct2Dataset from nemo_rl.data.interfaces import TaskDataProcessFnCallable, TaskDataSpec from nemo_rl.models.policy import TokenizerConfig diff --git a/tests/unit/data/test_llm_message_utils.py b/tests/unit/data/test_llm_message_utils.py index afb1f71398b..07e206af00d 100644 --- a/tests/unit/data/test_llm_message_utils.py +++ b/tests/unit/data/test_llm_message_utils.py @@ -20,7 +20,7 @@ from PIL import Image from transformers import AutoProcessor, AutoTokenizer -from nemo_rl.data.hf_datasets import COMMON_CHAT_TEMPLATES +from nemo_rl.data.chat_templates import COMMON_CHAT_TEMPLATES from nemo_rl.data.interfaces import LLMMessageLogType, TaskDataSpec from nemo_rl.data.llm_message_utils import ( _validate_tensor_consistency,