RL/nemo_rl/data/datasets/response_datasets/response_dataset.py at main · JeffPengCoder/RL · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
# 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.raw_dataset import RawDataset
from nemo_rl.data.datasets.utils import load_dataset_from_path


class ResponseDataset(RawDataset):
    """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
    }
    Please refer to https://github.com/NVIDIA-NeMo/RL/blob/main/docs/guides/sft.md#datasets for more details.

    Args:
        data_path: Path to the dataset JSON file
        input_key: Key for the input text, default is "input"
        output_key: Key for the output text, default is "output"
        subset: Optional subset name for the dataset, used for HuggingFace datasets
        split: Optional split name for the dataset, used for HuggingFace datasets
        split_validation_size: Size of the validation data, default is 0
        seed: Seed for train/validation split when split_validation_size > 0, default is 42
    """

    def __init__(
        self,
        data_path: str,
        input_key: str = "input",
        output_key: str = "output",
        subset: Optional[str] = None,
        split: Optional[str] = None,
        split_validation_size: float = 0,
        seed: int = 42,
        **kwargs,
    ):
        self.input_key = input_key
        self.output_key = output_key

        self.task_name = "-".join(data_path.split("/")[-2:]).split(".")[0]
        if self.task_name[0] == "-":
            self.task_name = self.task_name[1:]

        # load from local or huggingface
        self.dataset = load_dataset_from_path(data_path, subset, split)

        # format the dataset
        if "messages" not in self.dataset.column_names:
            self.dataset = self.dataset.map(
                self.format_data,
                remove_columns=self.dataset.column_names,
            )
        else:
            self.dataset = self.dataset.add_column(
                "task_name", [self.task_name] * len(self.dataset)
            )

        # `self.val_dataset` is used (not None) only when current dataset is used for both training and validation
        self.val_dataset = None
        self.split_train_validation(split_validation_size, seed)

    def format_data(self, data: dict[str, Any]) -> dict[str, Any]:
        return {
            "messages": [
                {"role": "user", "content": data[self.input_key]},
                {"role": "assistant", "content": data[self.output_key]},
            ],
            "task_name": self.task_name,
        }