Repository: huggingface/datasets
Pull request: Fix dataset_dict.shuffle with single seed
Fix #1610 

I added support for single integer used in `DatasetDict.shuffle`. Previously only a dictionary of seed was allowed.
Moreover I added the missing `seed` parameter. Previously only `seeds` was allowed.

Problem statement:
shuffle does not accept seed 
Hi
I need to shuffle the dataset, but this needs to be based on epoch+seed to be consistent across the cores, when I pass seed to shuffle, this does not accept seed, could you assist me with this? thanks  @lhoestq

Hints from issue discussion:
Hi, did you check the doc on `shuffle`?
https://huggingface.co/docs/datasets/package_reference/main_classes.html?datasets.Dataset.shuffle#datasets.Dataset.shuffle
Hi Thomas
thanks for reponse, yes, I did checked it, but this does not work for me please see 

```
(internship) rkarimi@italix17:/idiap/user/rkarimi/dev$ python 
Python 3.7.9 (default, Aug 31 2020, 12:42:55) 
[GCC 7.3.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import datasets 
2020-12-20 01:48:50.766004: W tensorflow/stream_executor/platform/default/dso_loader.cc:60] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2020-12-20 01:48:50.766029: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.
>>> data = datasets.load_dataset("scitail", "snli_format")
cahce dir  /idiap/temp/rkarimi/cache_home_1/datasets
cahce dir  /idiap/temp/rkarimi/cache_home_1/datasets
Reusing dataset scitail (/idiap/temp/rkarimi/cache_home_1/datasets/scitail/snli_format/1.1.0/fd8ccdfc3134ce86eb4ef10ba7f21ee2a125c946e26bb1dd3625fe74f48d3b90)
>>> data.shuffle(seed=2)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: shuffle() got an unexpected keyword argument 'seed'

```

datasets version
`datasets                  1.1.2                     <pip>
`

Thanks for reporting ! 

Indeed it looks like an issue with `suffle` on `DatasetDict`. We're going to fix that.
In the meantime you can shuffle each split (train, validation, test) separately:
```python
shuffled_train_dataset = data["train"].shuffle(seed=42)
```

Commit message: missing test

Patch to review:
diff --git a/src/datasets/dataset_dict.py b/src/datasets/dataset_dict.py
index 1f7127b2a4d..12766afa95a 100644
--- a/src/datasets/dataset_dict.py
+++ b/src/datasets/dataset_dict.py
@@ -420,7 +420,8 @@ def sort(
 
     def shuffle(
         self,
-        seeds: Optional[Dict[str, int]] = None,
+        seeds: Optional[Union[int, Dict[str, int]]] = None,
+        seed: Optional[int] = None,
         generators: Optional[Dict[str, np.random.Generator]] = None,
         keep_in_memory: bool = False,
         load_from_cache_file: bool = True,
@@ -434,10 +435,11 @@ def shuffle(
         You can either supply a NumPy BitGenerator to use, or a seed to initiate NumPy's default random generator (PCG64).
 
         Args:
-            seeds (Optional `Dict[str, int]`): A seed to initialize the default BitGenerator if ``generator=None``.
+            seeds (Optional `Dict[str, int]` or `int`): A seed to initialize the default BitGenerator if ``generator=None``.
                 If None, then fresh, unpredictable entropy will be pulled from the OS.
                 If an int or array_like[ints] is passed, then it will be passed to SeedSequence to derive the initial BitGenerator state.
-                You have to provide one :obj:`seed` per dataset in the dataset dictionary.
+                You can provide one :obj:`seed` per dataset in the dataset dictionary.
+            seed (Optional `int`): A seed to initialize the default BitGenerator if ``generator=None``. Alias for seeds.
             generators (Optional `Dict[str, np.random.Generator]`): Numpy random Generator to use to compute the permutation of the dataset rows.
                 If ``generator=None`` (default), uses np.random.default_rng (the default BitGenerator (PCG64) of NumPy).
                 You have to provide one :obj:`generator` per dataset in the dataset dictionary.
@@ -451,8 +453,13 @@ def shuffle(
                 Higher value gives smaller cache files, lower value consume less temporary memory while running `.map()`.
         """
         self._check_values_type()
+        if seed is not None and seeds is not None:
+            raise ValueError("Please specify seed or seeds, but not both")
+        seeds = seed if seed is not None else seeds
         if seeds is None:
             seeds = {k: None for k in self}
+        elif not isinstance(seeds, dict):
+            seeds = {k: seeds for k in self}
         if generators is None:
             generators = {k: None for k in self}
         if indices_cache_file_names is None:
diff --git a/tests/test_dataset_dict.py b/tests/test_dataset_dict.py
index 3271ba7c5ae..3f1992128f6 100644
--- a/tests/test_dataset_dict.py
+++ b/tests/test_dataset_dict.py
@@ -283,7 +283,17 @@ def test_shuffle(self):
                 seeds=seeds, indices_cache_file_names=indices_cache_file_names_3, load_from_cache_file=False
             )
             self.assertNotEqual(dsets_shuffled_3["train"]["filename"], dsets_shuffled_3["test"]["filename"])
+
+            # other input types
+            dsets_shuffled_int = dsets.shuffle(42)
+            dsets_shuffled_alias = dsets.shuffle(seed=42)
+            dsets_shuffled_none = dsets.shuffle()
+            self.assertEqual(len(dsets_shuffled_int["train"]), 30)
+            self.assertEqual(len(dsets_shuffled_alias["train"]), 30)
+            self.assertEqual(len(dsets_shuffled_none["train"]), 30)
+
             del dsets, dsets_shuffled, dsets_shuffled_2, dsets_shuffled_3
+            del dsets_shuffled_int, dsets_shuffled_alias, dsets_shuffled_none
 
     def test_check_values_type(self):
         dsets = self._create_dummy_dataset_dict()


Code review comments:
diff_hunk: @@ -434,10 +435,11 @@ def shuffle(
         You can either supply a NumPy BitGenerator to use, or a seed to initiate NumPy's default random generator (PCG64).
 
         Args:
-            seeds (Optional `Dict[str, int]`): A seed to initialize the default BitGenerator if ``generator=None``.
+            seeds (Optional `Dict[str, int]` or `int`): A seed to initialize the default BitGenerator if ``generator=None``.
                 If None, then fresh, unpredictable entropy will be pulled from the OS.
                 If an int or array_like[ints] is passed, then it will be passed to SeedSequence to derive the initial BitGenerator state.
-                You have to provide one :obj:`seed` per dataset in the dataset dictionary.
+                You can provide one :obj:`seed` per dataset in the dataset dictionary.
+            seed (Optional `int`): A seed to initialize the default BitGenerator if ``generator=None``. Alias for seeds.
path: src/datasets/dataset_dict.py
text: @user1:
```suggestion
            seed (Optional `int`): A seed to initialize the default BitGenerator if ``generator=None``. Alias for seeds (the seed argument has priority over seeds if both arguments are provided).
```