Repository: huggingface/datasets
Pull request: Added stratify option to train_test_split function.
This PR adds `stratify` option to `train_test_split` method. I took reference from scikit-learn's `StratifiedShuffleSplit` class for implementing stratified split and integrated the changes as were suggested by @lhoestq.

It fixes #3452.

@lhoestq Please review and let me know, if any changes are required.


Problem statement:
why the stratify option is omitted from test_train_split function?
why the stratify option is omitted from test_train_split function?

is there any other way implement the stratify option while splitting the dataset? as it is important point to be considered while splitting the dataset.

Hints from issue discussion:
Hi ! It's simply not added yet :)

If someone wants to contribute to add the `stratify` parameter I'd be happy to give some pointers.

In the meantime, I guess you can use `sklearn` or other tools to do a stratified train/test split over the **indices** of your dataset and then do
```
train_dataset = dataset.select(train_indices)
test_dataset = dataset.select(test_indices)
```
Hi @lhoestq I would like to add `stratify` parameter, can you give me some pointers for adding the same ?
Hi ! Sure :)

The `train_test_split` method is defined here: 

https://github.com/huggingface/datasets/blob/dc62232fa1b3bcfe2fbddcb721f2d141f8908943/src/datasets/arrow_dataset.py#L3253-L3253

and inside `train_test_split ` we need to create the right `train_indices` and `test_indices` that are passed here to `.select()`:

https://github.com/huggingface/datasets/blob/dc62232fa1b3bcfe2fbddcb721f2d141f8908943/src/datasets/arrow_dataset.py#L3450-L3464

For example if your dataset is like
|    |   label |
|---:|--------:|
|  0 |       1 |
|  1 |       1 |
|  2 |       0 |
|  3 |       0 |

and the user passes `stratify=dataset["label"]`, then you should get indices that look like this
```
train_indices = [0, 2]
test_indices = [1, 3]
```

these indices will be passed to `.select` to return the stratified train and test splits :)

Feel free to îng me if you have any question !
@lhoestq 
I just added the implementation for `stratify` option here #4322

Commit message: Fix test cases according to ClassLabel class

Patch to review:
diff --git a/src/datasets/arrow_dataset.py b/src/datasets/arrow_dataset.py
index 01ee0ef9064..c65111e68e9 100644
--- a/src/datasets/arrow_dataset.py
+++ b/src/datasets/arrow_dataset.py
@@ -56,6 +56,8 @@
 from requests import HTTPError
 from tqdm.auto import tqdm
 
+import datasets
+
 from . import config
 from .arrow_reader import ArrowReader
 from .arrow_writer import ArrowWriter, OptimizedTypedSequence
@@ -91,6 +93,7 @@
 from .utils.file_utils import _retry, estimate_dataset_size
 from .utils.info_utils import is_small_dataset
 from .utils.py_utils import convert_file_size_to_int, temporary_assignment, unique_values
+from .utils.stratify import stratified_shuffle_split_generate_indices
 from .utils.streaming_download_manager import xgetsize
 from .utils.typing import PathLike
 
@@ -3254,6 +3257,7 @@ def train_test_split(
         test_size: Union[float, int, None] = None,
         train_size: Union[float, int, None] = None,
         shuffle: bool = True,
+        stratify: Optional[str] = None,
         seed: Optional[int] = None,
         generator: Optional[np.random.Generator] = None,
         keep_in_memory: bool = False,
@@ -3280,6 +3284,7 @@ def train_test_split(
                 If int, represents the absolute number of train samples.
                 If None, the value is automatically set to the complement of the test size.
             shuffle (:obj:`bool`, optional, default `True`): Whether or not to shuffle the data before splitting.
+            stratify (:obj:`str`, optional, default `None`):The column name of labels to use to stratify data.
             seed (:obj:`int`, optional): 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.
@@ -3319,6 +3324,24 @@ def train_test_split(
 
         # set a seed
         >>> ds = ds.train_test_split(test_size=0.2, seed=42)
+
+        # stratified split
+        >>> ds = load_dataset("imdb",split="train")
+        Dataset({
+            features: ['text', 'label'],
+            num_rows: 25000
+        })
+        >>> ds = ds.train_test_split(test_size=0.2, stratify="label")
+        DatasetDict({
+            train: Dataset({
+                features: ['text', 'label'],
+                num_rows: 20000
+            })
+            test: Dataset({
+                features: ['text', 'label'],
+                num_rows: 5000
+            })
+        })
         ```
         """
         from .dataset_dict import DatasetDict  # import here because of circular dependency
@@ -3436,15 +3459,26 @@ def train_test_split(
                         ),
                     }
                 )
-
         if not shuffle:
+            if stratify is not None:
+                raise ValueError("Stratified train/test split is not implemented for shuffle=False")
             train_indices = np.arange(n_train)
             test_indices = np.arange(n_train, n_train + n_test)
         else:
+            # stratified partition
+            if stratify is not None:
+                if not isinstance(self.features[stratify], datasets.ClassLabel):
+                    raise ValueError(f"Invalid value for stratify: {stratify} of type {type(stratify)}")
+                train_indices, test_indices = next(
+                    stratified_shuffle_split_generate_indices(
+                        self.with_format("numpy")[stratify], stratify, n_train, n_test, rng=generator
+                    )
+                )
             # random partition
-            permutation = generator.permutation(len(self))
-            test_indices = permutation[:n_test]
-            train_indices = permutation[n_test : (n_test + n_train)]
+            else:
+                permutation = generator.permutation(len(self))
+                test_indices = permutation[:n_test]
+                train_indices = permutation[n_test : (n_test + n_train)]
 
         train_split = self.select(
             indices=train_indices,
diff --git a/src/datasets/utils/stratify.py b/src/datasets/utils/stratify.py
new file mode 100644
index 00000000000..f3df7b6b088
--- /dev/null
+++ b/src/datasets/utils/stratify.py
@@ -0,0 +1,112 @@
+import numpy as np
+
+
+def approximate_mode(class_counts, n_draws, rng):
+    """Computes approximate mode of multivariate hypergeometric.
+    This is an approximation to the mode of the multivariate
+    hypergeometric given by class_counts and n_draws.
+    It shouldn't be off by more than one.
+    It is the mostly likely outcome of drawing n_draws many
+    samples from the population given by class_counts.
+    Args
+    ----------
+    class_counts : ndarray of int
+        Population per class.
+    n_draws : int
+        Number of draws (samples to draw) from the overall population.
+    rng : random state
+        Used to break ties.
+    Returns
+    -------
+    sampled_classes : ndarray of int
+        Number of samples drawn from each class.
+        np.sum(sampled_classes) == n_draws
+
+    """
+    # this computes a bad approximation to the mode of the
+    # multivariate hypergeometric given by class_counts and n_draws
+    continuous = n_draws * class_counts / class_counts.sum()
+    # floored means we don't overshoot n_samples, but probably undershoot
+    floored = np.floor(continuous)
+    # we add samples according to how much "left over" probability
+    # they had, until we arrive at n_samples
+    need_to_add = int(n_draws - floored.sum())
+    if need_to_add > 0:
+        remainder = continuous - floored
+        values = np.sort(np.unique(remainder))[::-1]
+        # add according to remainder, but break ties
+        # randomly to avoid biases
+        for value in values:
+            (inds,) = np.where(remainder == value)
+            # if we need_to_add less than what's in inds
+            # we draw randomly from them.
+            # if we need to add more, we add them all and
+            # go to the next value
+            add_now = min(len(inds), need_to_add)
+            inds = rng.choice(inds, size=add_now, replace=False)
+            floored[inds] += 1
+            need_to_add -= add_now
+            if need_to_add == 0:
+                break
+    return floored.astype(np.int)
+
+
+def stratified_shuffle_split_generate_indices(y, stratify, n_train, n_test, rng, n_splits=10):
+    """
+
+    Provides train/test indices to split data in train/test sets.
+    It's reference is taken from StratifiedShuffleSplit implementation
+    of scikit-learn library.
+
+    Args
+    ----------
+
+    n_train : int,
+        represents the absolute number of train samples.
+
+    n_test : int,
+        represents the absolute number of test samples.
+
+    random_state : int or RandomState instance, default=None
+        Controls the randomness of the training and testing indices produced.
+        Pass an int for reproducible output across multiple function calls.
+
+    n_splits : int, default=10
+        Number of re-shuffling & splitting iterations.
+    """
+    classes, y_indices = np.unique(y, return_inverse=True)
+    n_classes = classes.shape[0]
+    class_counts = np.bincount(y_indices)
+    if np.min(class_counts) < 2:
+        raise ValueError(
+            f"The least populated class in {stratify} column has only 1"
+            " member, which is too few. The minimum"
+            " number of groups for any class cannot"
+            " be less than 2."
+        )
+    if n_train < n_classes:
+        raise ValueError(
+            "The train_size = %d should be greater or " "equal to the number of classes = %d" % (n_train, n_classes)
+        )
+    if n_test < n_classes:
+        raise ValueError(
+            "The test_size = %d should be greater or " "equal to the number of classes = %d" % (n_test, n_classes)
+        )
+    class_indices = np.split(np.argsort(y_indices, kind="mergesort"), np.cumsum(class_counts)[:-1])
+    for _ in range(n_splits):
+        n_i = approximate_mode(class_counts, n_train, rng)
+        class_counts_remaining = class_counts - n_i
+        t_i = approximate_mode(class_counts_remaining, n_test, rng)
+
+        train = []
+        test = []
+
+        for i in range(n_classes):
+            permutation = rng.permutation(class_counts[i])
+            perm_indices_class_i = class_indices[i].take(permutation, mode="clip")
+            train.extend(perm_indices_class_i[: n_i[i]])
+            test.extend(perm_indices_class_i[n_i[i] : n_i[i] + t_i[i]])
+        train = rng.permutation(train)
+        test = rng.permutation(test)
+
+        yield train, test
diff --git a/tests/test_arrow_dataset.py b/tests/test_arrow_dataset.py
index 96fd19a83c1..7daf98f4f7f 100644
--- a/tests/test_arrow_dataset.py
+++ b/tests/test_arrow_dataset.py
@@ -11,6 +11,7 @@
 from unittest.mock import patch
 
 import numpy as np
+import numpy.testing as npt
 import pandas as pd
 import pyarrow as pa
 import pytest
@@ -3553,3 +3554,32 @@ def test_task_text_classification_when_columns_removed(self):
         with Dataset.from_dict(data, info=info) as dset:
             with dset.map(lambda x: {"new_column": 0}, remove_columns=dset.column_names) as dset:
                 self.assertDictEqual(dset.features, features_after_map)
+
+
+def test_train_test_split_startify():
+    ys = [
+        np.array([0, 0, 0, 0, 1, 1, 1, 2, 2, 2, 2, 2]),
+        np.array([0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3]),
+        np.array([0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2] * 2),
+        np.array([0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3]),
+        np.array([0] * 800 + [1] * 50),
+    ]
+    for y in ys:
+        features = Features({"text": Value("int64"), "label": ClassLabel(len(np.unique(y)))})
+        data = {"text": np.ones(len(y)), "label": y}
+        d1 = Dataset.from_dict(data, features=features)
+        d1 = d1.train_test_split(test_size=0.33, stratify="label")
+        y = np.asanyarray(y)  # To make it indexable for y[train]
+        test_size = np.ceil(0.33 * len(y))
+        train_size = len(y) - test_size
+        npt.assert_array_equal(np.unique(d1["train"]["label"]), np.unique(d1["test"]["label"]))
+
+        # checking classes proportion
+        p_train = np.bincount(np.unique(d1["train"]["label"], return_inverse=True)[1]) / float(
+            len(d1["train"]["label"])
+        )
+        p_test = np.bincount(np.unique(d1["test"]["label"], return_inverse=True)[1]) / float(len(d1["test"]["label"]))
+        npt.assert_array_almost_equal(p_train, p_test, 1)
+        assert len(d1["train"]["text"]) + len(d1["test"]["text"]) == y.size
+        assert len(d1["train"]["text"]) == train_size
+        assert len(d1["test"]["text"]) == test_size


Code review comments:
diff_hunk: @@ -3436,15 +3459,26 @@ def train_test_split(
                         ),
                     }
                 )
-
         if not shuffle:
+            if stratify is not None:
+                raise ValueError("Stratified train/test split is not implemented for shuffle=False")
path: src/datasets/arrow_dataset.py
text: @user1:
```suggestion
                raise ValueError("Stratified train/test split is not implemented for `shuffle=False`")
```

diff_hunk: @@ -0,0 +1,112 @@
+import numpy as np
+
+
+def approximate_mode(class_counts, n_draws, rng):
+    """Computes approximate mode of multivariate hypergeometric.
+    This is an approximation to the mode of the multivariate
+    hypergeometric given by class_counts and n_draws.
+    It shouldn't be off by more than one.
+    It is the mostly likely outcome of drawing n_draws many
+    samples from the population given by class_counts.
+    Args
+    ----------
+    class_counts : ndarray of int
+        Population per class.
+    n_draws : int
+        Number of draws (samples to draw) from the overall population.
+    rng : random state
+        Used to break ties.
+    Returns
+    -------
+    sampled_classes : ndarray of int
+        Number of samples drawn from each class.
+        np.sum(sampled_classes) == n_draws
+
+    """
+    # this computes a bad approximation to the mode of the
+    # multivariate hypergeometric given by class_counts and n_draws
+    continuous = n_draws * class_counts / class_counts.sum()
+    # floored means we don't overshoot n_samples, but probably undershoot
+    floored = np.floor(continuous)
+    # we add samples according to how much "left over" probability
+    # they had, until we arrive at n_samples
+    need_to_add = int(n_draws - floored.sum())
+    if need_to_add > 0:
+        remainder = continuous - floored
+        values = np.sort(np.unique(remainder))[::-1]
+        # add according to remainder, but break ties
+        # randomly to avoid biases
+        for value in values:
+            (inds,) = np.where(remainder == value)
+            # if we need_to_add less than what's in inds
+            # we draw randomly from them.
+            # if we need to add more, we add them all and
+            # go to the next value
+            add_now = min(len(inds), need_to_add)
+            inds = rng.choice(inds, size=add_now, replace=False)
+            floored[inds] += 1
+            need_to_add -= add_now
+            if need_to_add == 0:
+                break
+    return floored.astype(np.int)
+
+
+def stratified_shuffle_split_generate_indices(y, stratify, n_train, n_test, rng, n_splits=10):
+    """
+
+    Provides train/test indices to split data in train/test sets.
+    It's reference is taken from StratifiedShuffleSplit implementation
+    of scikit-learn library.
+
+    Args
+    ----------
+
+    n_train : int,
+        represents the absolute number of train samples.
+
+    n_test : int,
+        represents the absolute number of test samples.
+
+    random_state : int or RandomState instance, default=None
+        Controls the randomness of the training and testing indices produced.
+        Pass an int for reproducible output across multiple function calls.
+
+    n_splits : int, default=10
+        Number of re-shuffling & splitting iterations.
+    """
+    classes, y_indices = np.unique(y, return_inverse=True)
+    n_classes = classes.shape[0]
+    class_counts = np.bincount(y_indices)
+    if np.min(class_counts) < 2:
+        raise ValueError(
+            f"The least populated class in {stratify} column has only 1"
+            " member, which is too few. The minimum"
+            " number of groups for any class cannot"
+            " be less than 2."
+        )
path: src/datasets/utils/stratify.py
text: @user1:
Ideally, we want to avoid passing the `stratify` column name and only deal with the actual stratification procedure in this file. This can be achieved by throwing an error here, which you then catch in `train_test_split` with `try/except` around this function's call, updating the error message with the column name, and finally reraising the error. If it's unclear what I mean, I can refactor this part myself before merging.

@author:
It's clear to me, I will try adding as recommended.

@author:
I have made changes as you suggested, let me know if any more changes are required.

diff_hunk: @@ -3280,6 +3284,7 @@ def train_test_split(
                 If int, represents the absolute number of train samples.
                 If None, the value is automatically set to the complement of the test size.
             shuffle (:obj:`bool`, optional, default `True`): Whether or not to shuffle the data before splitting.
+            stratify (:obj:`str`, optional, default `None`):The column name of labels to use to stratify data.
path: src/datasets/arrow_dataset.py
text: @user1:
Can we improve the name of this param? Maybe `stratify_by_column`? I would expect `stratify` to be a boolean flag, but it works with column names.

diff_hunk: @@ -3436,15 +3459,26 @@ def train_test_split(
                         ),
                     }
                 )
-
         if not shuffle:
+            if stratify is not None:
+                raise ValueError("Stratified train/test split is not implemented for shuffle=False")
             train_indices = np.arange(n_train)
             test_indices = np.arange(n_train, n_train + n_test)
         else:
+            # stratified partition
+            if stratify is not None:
+                if not isinstance(self.features[stratify], datasets.ClassLabel):
+                    raise ValueError(f"Invalid value for stratify: {stratify} of type {type(stratify)}")
path: src/datasets/arrow_dataset.py
text: @user1:
`type(stratify)` is always `str` here. We want to print its feature type. Also, let's align the message with the existing error messages:
```suggestion
                    raise ValueError(f"Stratifying by column is  only supported for {ClassLabel.__name__} column, and column {stratify_by_column} is {type(self.features[stratify_by_column]).__name__}.")
```