Fix unpairing IterableDatasetDict preference datasets

tests/test_data_utils.py

@@ -18,7 +18,7 @@
 
 import pytest
 import transformers
-from datasets import Dataset, DatasetDict
+from datasets import Dataset, DatasetDict, IterableDatasetDict
 from packaging.version import Version
 from transformers import AutoProcessor, AutoTokenizer, is_vision_available
 
@@ -1090,6 +1090,15 @@ def test_unpair_preference_dataset_dict(self):
             "The paired dataset should be converted to unpaired."
         )
 
+    def test_unpair_preference_iterable_dataset_dict(self):
+        # Test that a paired iterable dataset dict is correctly converted to 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_maybe_unpair_preference_dataset(self):
         # Test that a paired dataset is correctly converted to unpaired with maybe_unpair_preference_dataset
         unpaired_dataset = maybe_unpair_preference_dataset(self.paired_dataset)

trl/data_utils.py

@@ -494,7 +494,7 @@ def unpair_preference_dataset(
     {'prompt': 'The sky is', 'completion': ' blue.', 'label': True}
     ```
     """
-    if isinstance(dataset, DatasetDict):
+    if isinstance(dataset, (DatasetDict, IterableDatasetDict)):
         column_names = next(iter(dataset.values())).column_names
     elif isinstance(dataset, Dataset):
         column_names = dataset.column_names