Confirm that the provided Dataset is in the DatasetDict format from HF

adap · Apr 4, 2024 · ba0cf03 · ba0cf03
1 parent 1fdc09f
commit ba0cf03
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/datasets/flwr_datasets/utils.py b/datasets/flwr_datasets/utils.py
@@ -79,13 +79,18 @@ def _instantiate_resplitter_if_needed(
     return cast(Optional[Resplitter], resplitter)
 
 
-def _check_if_dataset_tested(dataset: Union[str, Dataset]):
+def _check_if_dataset_tested(dataset: Union[str, DatasetDict]) -> None:
     """Check if the dataset is in the narrowed down list of the tested datasets."""
     if dataset not in tested_datasets and type(dataset) is str:
         warnings.warn(
             f"The currently tested dataset are {tested_datasets}. Given: {dataset}.",
             stacklevel=1,
         )
+    if type(dataset) is not DatasetDict:
+        warnings.warn(
+            f"The provided dataset is not in the DatasetDict format.",
+            stacklevel=1,
+        )
 
 
 def divide_dataset(