fix: Support datasets saved with save_to_disk in ResponseDataset (#1610)

Signed-off-by: Sahger Lad <[email protected]>
Signed-off-by: sahgerlad <[email protected]>
Co-authored-by: Yuki Huang <[email protected]>

Author: sahgerlad

nemo_rl/data/datasets/response_datasets/response_dataset.py

@@ -56,11 +56,12 @@ def __init__(
         else:
             val_ds = None
 
-        # format the dataset
-        train_ds = train_ds.map(
-            self.add_messages_key, fn_kwargs={"task_name": self.task_name}
-        )
-        if val_ds:
+        # Only apply add_messages_key if 'messages' column doesn't exist
+        if "messages" not in train_ds.column_names:
+            train_ds = train_ds.map(
+                self.add_messages_key, fn_kwargs={"task_name": self.task_name}
+            )
+        if val_ds is not None and "messages" not in val_ds.column_names:
             val_ds = val_ds.map(
                 self.add_messages_key, fn_kwargs={"task_name": self.task_name}
             )

nemo_rl/data/datasets/utils.py

@@ -17,7 +17,7 @@
 from typing import Optional, Union
 
 import torch
-from datasets import DatasetDict, load_dataset
+from datasets import DatasetDict, load_dataset, load_from_disk
 from PIL import Image
 from transformers import AutoProcessor, PreTrainedTokenizerBase
 
@@ -62,7 +62,7 @@ def pil_to_base64(image: Image.Image, format: str = "PNG") -> str:
 
 
 def load_dataset_from_path(data_path: str, data_split: Optional[str] = "train"):
-    """Load a dataset from a json or huggingface dataset.
+    """Load a dataset from a json, huggingface dataset, or Arrow dataset (saved with save_to_disk).
 
     Args:
         data_path: The path to the dataset.
@@ -72,7 +72,13 @@ def load_dataset_from_path(data_path: str, data_split: Optional[str] = "train"):
     if suffix in [".json", ".jsonl"]:
         raw_dataset = load_dataset("json", data_files=data_path)
     else:
-        raw_dataset = load_dataset(data_path)
+        try:
+            raw_dataset = load_dataset(data_path)
+        except ValueError as e:
+            if "load_from_disk" in str(e):
+                raw_dataset = load_from_disk(data_path)
+            else:
+                raise e
 
     if data_split:
         raw_dataset = raw_dataset[data_split]

tests/unit/data/datasets/test_response_dataset.py

@@ -165,3 +165,75 @@ def test_squad_dataset():
             + " Answer: "
             + example["messages"][2]["content"]
         )
+
+
+def test_load_dataset_saved_with_save_to_disk():
+    """Test loading a dataset that was saved using HuggingFace's save_to_disk().
+
+    This tests the fix for datasets that already have a 'messages' column,
+    which should be preserved without applying add_messages_key again.
+    """
+    from datasets import Dataset
+
+    # Create a dataset with 'messages' column already present
+    train_data = [
+        {
+            "messages": [
+                {"role": "user", "content": "What is 2+2?"},
+                {"role": "assistant", "content": "4"},
+            ]
+        },
+        {
+            "messages": [
+                {"role": "user", "content": "What is the capital of France?"},
+                {"role": "assistant", "content": "Paris"},
+            ]
+        },
+    ]
+    val_data = [
+        {
+            "messages": [
+                {"role": "user", "content": "What is 3+3?"},
+                {"role": "assistant", "content": "6"},
+            ]
+        },
+    ]
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # Create HF datasets and save using save_to_disk
+        train_dataset = Dataset.from_list(train_data)
+        val_dataset = Dataset.from_list(val_data)
+
+        train_path = f"{tmpdir}/train"
+        val_path = f"{tmpdir}/val"
+
+        train_dataset.save_to_disk(train_path)
+        val_dataset.save_to_disk(val_path)
+
+        # Load using load_response_dataset
+        data_config = {
+            "dataset_name": "ResponseDataset",
+            "train_data_path": train_path,
+            "val_data_path": val_path,
+        }
+        dataset = load_response_dataset(data_config)
+
+        # Verify the dataset loaded correctly
+        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
+
+        # Verify messages are preserved correctly
+        first_train_example = dataset.formatted_ds["train"][0]
+        assert "messages" in first_train_example
+        assert len(first_train_example["messages"]) == 2
+        assert first_train_example["messages"][0]["role"] == "user"
+        assert first_train_example["messages"][0]["content"] == "What is 2+2?"
+        assert first_train_example["messages"][1]["role"] == "assistant"
+        assert first_train_example["messages"][1]["content"] == "4"
+
+        # Verify validation data
+        first_val_example = dataset.formatted_ds["validation"][0]
+        assert first_val_example["messages"][0]["content"] == "What is 3+3?"
+        assert first_val_example["messages"][1]["content"] == "6"