Added stratify option to train_test_split function.

[nandwalritik]

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.

[lhoestq]

Nice thank you ! This will be super useful :)

Could you also add some tests in test_arrow_dataset.py and add an example of usage in the Example: section of the train_test_split docstring ?

[nandwalritik]

Nice thank you ! This will be super useful :)

Could you also add some tests in test_arrow_dataset.py and add an example of usage in the Example: section of the train_test_split docstring ?

I will try to do it, is there any documentation for adding test cases? I have never done it before.

[lhoestq]

Thanks for the changes !

I will try to do it, is there any documentation for adding test cases? I have never done it before.

You can just add a function test_train_test_split_startify in test_arrow_dataset.py.

In this function you can define a dataset and make sure that train_test_split with the stratify argument works as expected.

You can do pytest tests/test_arrow_dataset.py::test_train_test_split_startify to run your test.

Feel free to get some inspiration from other tests like test_interleave_datasets for example

[nandwalritik]

I have added tests for stratified train_test_split in test_arrow_dataset.py file inside test_train_test_split_startify function. I have also added example usage with stratify arg in Example: section of the train_test_split docstring.
Results of tests:

(data) nandwalritik@hp:~/datasets$ pytest tests/test_arrow_dataset.py::test_train_test_split_startify -W ignore
============================================================================ test session starts ============================================================================
platform linux -- Python 3.9.5, pytest-7.1.2, pluggy-1.0.0
rootdir: /home/nandwalritik/datasets
plugins: datadir-1.3.1, forked-1.4.0, xdist-2.5.0
collected 1 item                                                                                                                                                            

tests/test_arrow_dataset.py .                                                                                                                                         [100%]

============================================================================= 1 passed in 0.12s =============================================================================

[mariosasko]

Very nice!

PS: Maybe we should put the _approximate_mode and _stratified_shuffle_split_generate_indices in a separate module @lhoestq (utils/stratify.py for instance)? And if we keep them here as methods, we should at least mark them with the @staticmethod decorator.

[lhoestq]

Thanks a lot !

utils/stratify.py sounds good yes :)

Also feel free to merge master into your branch to fix the CI ;)

[nandwalritik]

Added all the changes as were suggested and rebased with main.

[mariosasko]

A few more nits. Other than that looks good!

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[mariosasko]

Thanks! LGTM!

Pinging @lhoestq for the final review.

[lhoestq]

Thanks a lot !

[Damon03]

Hi, I encounter an error when I try to specify the stratify_by_column. However, I have a columns which specific the label of the row as a string. But an error showed when I try to do it. "ValueError: Stratifying by column is only supported for ClassLabel column, and column code is Value.".

[lhoestq]

Hi @Damon03 , you can change the type of your column to ClassLabel using

ds = ds.class_encode_column(column_name)

then you'll be free to use stratify :)

[Damon03]

Thank you so much. It worked.