implement DatasetDict.shuffle

ArneBinder · Sep 9, 2024 · ba1168d · ba1168d
1 parent e1db8f3
commit ba1168d
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/src/pie_datasets/core/dataset_dict.py b/src/pie_datasets/core/dataset_dict.py
@@ -694,6 +694,15 @@ def cast_document_type(
         )
         return result
 
+    def shuffle(self, **kwargs):
+        result = DatasetDict.from_hf(super().shuffle(**kwargs), document_type=self.document_type)
+
+        # TODO: integrate into DatasetDict.from_hf
+        for split_name, split in result.items():
+            split.document_converters = self[split_name].document_converters
+
+        return result
+
 
 def load_dataset(*args, **kwargs) -> Union[DatasetDict, Dataset, IterableDataset]:
     dataset_or_dataset_dict = datasets.load_dataset(*args, **kwargs)