fix: unpair IterableDatasetDict inputs

Author: YZJF

tests/test_data_utils.py

@@ -18,7 +18,7 @@
 
 import pytest
 import transformers
-from datasets import Dataset, DatasetDict
+from datasets import Dataset, DatasetDict, IterableDataset, IterableDatasetDict
 from packaging.version import Version
 from transformers import AutoProcessor, AutoTokenizer, is_vision_available
 
@@ -1082,6 +1082,36 @@ def test_unpair_preference_dataset_iterable_extra_columns(self):
             dict(zip(expected.keys(), vals, strict=False)) for vals in zip(*expected.values(), strict=False)
         ]
 
+    def test_unpair_preference_dataset_iterable_dict(self):
+        # Test that an IterableDatasetDict is correctly unpaired
+        paired_dataset_dict = IterableDatasetDict({"abc": self.paired_dataset.to_iterable_dataset()})
+        unpaired_dataset_dict = unpair_preference_dataset(paired_dataset_dict)
+        assert list(unpaired_dataset_dict["abc"]) == [
+            dict(zip(self.unpaired_dataset.column_names, vals, strict=False))
+            for vals in zip(*self.unpaired_dataset.to_dict().values(), strict=False)
+        ]
+
+    def test_unpair_preference_dataset_iterable_dict_without_column_names(self):
+        # Test that an IterableDatasetDict without schema metadata is correctly unpaired
+        def generate_examples():
+            for prompt, chosen, rejected in zip(
+                self.paired_dataset["prompt"],
+                self.paired_dataset["chosen"],
+                self.paired_dataset["rejected"],
+                strict=True,
+            ):
+                yield {"prompt": prompt, "chosen": chosen, "rejected": rejected}
+
+        paired_dataset_dict = IterableDatasetDict({"abc": IterableDataset.from_generator(generate_examples)})
+        assert paired_dataset_dict["abc"].column_names is None
+
+        unpaired_dataset_dict = unpair_preference_dataset(paired_dataset_dict)
+
+        assert list(unpaired_dataset_dict["abc"]) == [
+            dict(zip(self.unpaired_dataset.column_names, vals, strict=False))
+            for vals in zip(*self.unpaired_dataset.to_dict().values(), strict=False)
+        ]
+
     def test_unpair_preference_dataset_dict(self):
         # Test that a paired dataset dict is correctly converted to unpaired
         paired_dataset_dict = DatasetDict({"abc": self.paired_dataset})

trl/data_utils.py

@@ -494,8 +494,9 @@ def unpair_preference_dataset(
     {'prompt': 'The sky is', 'completion': ' blue.', 'label': True}
     ```
     """
-    if isinstance(dataset, DatasetDict):
-        column_names = next(iter(dataset.values())).column_names
+    if isinstance(dataset, (DatasetDict, IterableDatasetDict)):
+        first_split = next(iter(dataset.values()))
+        column_names = first_split.column_names or list(next(iter(first_split)).keys())
     elif isinstance(dataset, Dataset):
         column_names = dataset.column_names
     else:  # IterableDataset