Repository: allenai/tango (language: python)

Issue:
Add a `datasets::dataset_remix` step.
Should do the same thing as `tango.steps.dataset_remix.DatasetRemix`, but for `datasets.DatasetDict` rather than Tango's `DatasetDict.


Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9d9cc76f4..b189de960 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - If you try to import something from a tango integration that is not fully installed due to missing dependencies, an `IntegrationMissingError` will be raised
 instead of `ModuleNotFound`.
 
+### Added
+
+- Added a `datasets::dataset_remix` step that provides the split remixing functionality of `tango.steps.datasest_remix.DatasetRemixStep` now for Huggingface `DatasetDict`.
+
 ## [v0.8.0](https://github.com/allenai/tango/releases/tag/v0.8.0) - 2022-05-19
 
 ### Added
diff --git a/docs/source/api/integrations/datasets.rst b/docs/source/api/integrations/datasets.rst
index a2f1da5b9..10cd469e4 100644
--- a/docs/source/api/integrations/datasets.rst
+++ b/docs/source/api/integrations/datasets.rst
@@ -21,3 +21,6 @@ Reference
 
 .. autoclass:: tango.integrations.datasets.ConcatenateDatasets
    :members:
+
+.. autoclass:: tango.integrations.datasets.DatasetRemixStep
+   :members:
\ No newline at end of file
diff --git a/tango/integrations/datasets/__init__.py b/tango/integrations/datasets/__init__.py
index 29ca4a659..d5e81a8ba 100644
--- a/tango/integrations/datasets/__init__.py
+++ b/tango/integrations/datasets/__init__.py
@@ -23,8 +23,10 @@
 """
 
 
+import random
+import re
 from pathlib import Path
-from typing import Any, List, Optional, TypeVar, Union, overload
+from typing import Any, Dict, List, Optional, TypeVar, Union, overload
 
 from tango.common.aliases import PathOrStr
 from tango.common.dataset_dict import DatasetDict, IterableDatasetDict
@@ -44,6 +46,7 @@
     "convert_to_tango_dataset_dict",
     "InterleaveDatasets",
     "ConcatenateDatasets",
+    "DatasetRemixStep",
 ]
 
 
@@ -226,3 +229,127 @@ def run(  # type: ignore[override]
         Concatenate the list of datasets.
         """
         return ds.concatenate_datasets(datasets, info=info, split=split, axis=axis)
+
+
+@Step.register("datasets::dataset_remix")
+class DatasetRemixStep(Step):
+    """
+    This step can remix splits in a :class:`~datasets.DatasetDict` into new splits.
+
+    .. tip::
+
+        Registered as a :class:`~tango.step.Step` under the name "datasets::dataset_remix".
+
+    Examples
+    --------
+
+    .. testcode::
+        :hide:
+
+        from tango.common.logging import initialize_logging
+        initialize_logging(enable_cli_logs=True)
+        import datasets
+
+    .. testcode::
+
+        input = datasets.load_dataset("lhoestq/test")
+        new_splits = {
+            "all": "train + validation",
+            "crossval_train": "train[:1] + validation[1:]",
+            "crossval_test": "train[1:] + validation[:1]",
+        }
+        step = DatasetRemixStep()
+        remixed_dataset = step.run(input=input, new_splits=new_splits)
+
+    .. testoutput::
+        :hide:
+        :options: +ELLIPSIS
+
+        ...
+
+    """
+
+    DETERMINISTIC = True
+    CACHEABLE = True
+    VERSION = "001"
+
+    def run(  # type: ignore
+        self,
+        input: ds.DatasetDict,
+        new_splits: Dict[str, str],
+        keep_old_splits: bool = True,
+        shuffle_before: bool = False,
+        shuffle_after: bool = False,
+        random_seed: int = 1532637578,
+    ) -> ds.DatasetDict:
+        """
+        Remixes and shuffles a dataset. This is done eagerly with native 🤗 Datasets features.
+
+        :param input:
+            The input dataset that will be remixed.
+        :param new_splits:
+            Specifies the new splits that the output dataset should have. Keys are the name of the new
+            splits. Values refer to the original splits. You can refer to original splits in the following ways:
+
+            * Mention the original split name to copy it to a new name.
+            * Mention the original split name with Python's slicing syntax to select part of the original
+              split's instances. For example, ``"train[:1000]"`` selects the first 1000 instances from the
+              ``"train"`` split.
+            * ``"instances + instances"`` concatenates the instances into one split.
+
+            You can combine these possibilities.
+        :param keep_old_splits:
+            Whether to keep the splits from the input dataset in addition to the new ones given by
+            ``new_splits``.
+        :param shuffle_before:
+            Whether to shuffle the input splits before creating the new ones.
+
+            If you need shuffled instances and you're not sure the input is properly shuffled, use this.
+        :param shuffle_after:
+            Whether to shuffle the input splits after creating the new ones.
+
+            If you need shuffled instances and you're slicing or concatenating splits, use this.
+
+            If you want to be on the safe side, shuffle both before and after.
+        :param random_seed:
+            Random seed, affects shuffling
+
+        :returns:
+            Returns a new dataset that is appropriately remixed.
+        """
+
+        if shuffle_before:
+            input = input.shuffle(random_seed)
+
+        def get_slice(split_name: str) -> ds.Dataset:
+            slice_match = re.match(r"(.*)\[(-?[0-9]*:-?[0-9]*)\]", split_name)
+            if slice_match is None:
+                return input[split_name]
+            else:
+                split_name = slice_match[1]
+                slice_args = [int(a) if len(a) > 0 else None for a in slice_match[2].split(":")]
+                slice_indices = range(*slice(*slice_args).indices(len(input[split_name])))
+                return input[split_name].select(slice_indices)
+
+        def parse_split_spec(split_spec: str):
+            parts = [get_slice(name.strip()) for name in split_spec.split("+")]
+            if len(parts) == 1:
+                return parts[0]
+            else:
+                return ds.concatenate_datasets(parts)
+
+        if keep_old_splits:
+            result = ds.DatasetDict(input.items())
+        else:
+            result = ds.DatasetDict()
+        result.update(
+            {
+                new_split_name: parse_split_spec(new_split_spec)
+                for new_split_name, new_split_spec in new_splits.items()
+            }
+        )
+
+        if shuffle_after:
+            result = result.shuffle(random_seed)
+
+        return result


Test patch:
diff --git a/tests/integrations/datasets/dataset_test.py b/tests/integrations/datasets/dataset_test.py
index 578fa9623..2352c1027 100644
--- a/tests/integrations/datasets/dataset_test.py
+++ b/tests/integrations/datasets/dataset_test.py
@@ -3,6 +3,7 @@
 from tango.common.sequences import MappedSequence
 from tango.common.testing import TangoTestCase
 from tango.integrations.datasets import (
+    DatasetRemixStep,
     DatasetsFormat,
     LoadDataset,
     convert_to_tango_dataset_dict,
@@ -50,3 +51,19 @@ def test_mapped_sequence_of_dataset():
     assert len(ds) == len(mapped_ds)
     assert ds[0]["goal"] == mapped_ds[0]
     assert ds[0]["goal"] == mapped_ds[:10][0]
+
+
+def test_datasets_dataset_remix():
+    dataset_dict = datasets.load_dataset("lhoestq/test")
+    step = DatasetRemixStep()
+    result = step.run(
+        input=dataset_dict,
+        new_splits={
+            "all": "train + validation",
+            "crossval_train": "train[:1] + validation[1:]",
+            "crossval_test": "train[1:] + validation[:1]",
+        },
+    )
+    assert len(result["all"]) == len(dataset_dict["train"]) + len(dataset_dict["validation"])
+    assert len(result["crossval_train"]) == 3
+    assert len(result["crossval_test"]) == 2