Added stratify option to train_test_split function.
{'login': 'nandwalritik', 'id': 48522685, 'node_id': 'MDQ6VXNlcjQ4NTIyNjg1', 'avatar_url': 'https://avatars.githubusercontent.com/u/48522685?v=4', 'gravatar_id': '', 'url': 'https://api.github.com/users/nandwalritik', 'html_url': 'https://github.com/nandwalritik', 'followers_url': 'https://api.github.com/users/nandwalritik/followers', 'following_url': 'https://api.github.com/users/nandwalritik/following{/other_user}', 'gists_url': 'https://api.github.com/users/nandwalritik/gists{/gist_id}', 'starred_url': 'https://api.github.com/users/nandwalritik/starred{/owner}{/repo}', 'subscriptions_url': 'https://api.github.com/users/nandwalritik/subscriptions', 'organizations_url': 'https://api.github.com/users/nandwalritik/orgs', 'repos_url': 'https://api.github.com/users/nandwalritik/repos', 'events_url': 'https://api.github.com/users/nandwalritik/events{/privacy}', 'received_events_url': 'https://api.github.com/users/nandwalritik/received_events', 'type': 'User', 'site_admin': False}
[]
[]
['> Nice thank you ! This will be super useful :)\r\n> \r\n> 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 ?\r\n\r\nI will try to do it, is there any documentation for adding test cases? I have never done it before.'
 'Thanks for the changes !\r\n\r\n> I will try to do it, is there any documentation for adding test cases? I have never done it before.\r\n\r\nYou can just add a function `test_train_test_split_startify` in `test_arrow_dataset.py`.\r\n\r\nIn this function you can define a dataset and make sure that `train_test_split` with the `stratify` argument works as expected.\r\n\r\nYou can do `pytest tests/test_arrow_dataset.py::test_train_test_split_startify` to run your test.\r\n\r\nFeel free to get some inspiration from other tests like `test_interleave_datasets` for example'
 '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.\r\nResults of tests:\r\n```\r\n(data) nandwalritik@hp:~/datasets$ pytest tests/test_arrow_dataset.py::test_train_test_split_startify -W ignore\r\n============================================================================ test session starts ============================================================================\r\nplatform linux -- Python 3.9.5, pytest-7.1.2, pluggy-1.0.0\r\nrootdir: /home/nandwalritik/datasets\r\nplugins: datadir-1.3.1, forked-1.4.0, xdist-2.5.0\r\ncollected 1 item                                                                                                                                                            \r\n\r\ntests/test_arrow_dataset.py .                                                                                                                                         [100%]\r\n\r\n============================================================================= 1 passed in 0.12s =============================================================================\r\n\r\n```'
 'Thanks a lot !\r\n\r\n`utils/stratify.py` sounds good yes :)\r\n\r\nAlso feel free to merge `master` into your branch to fix the CI ;)'
 'Added all the changes as were suggested and rebased with `main`.'
 '_The documentation is not available anymore as the PR was closed or merged._']
{'url': 'https://api.github.com/repos/huggingface/datasets/pulls/4322', 'html_url': 'https://github.com/huggingface/datasets/pull/4322', 'diff_url': 'https://github.com/huggingface/datasets/pull/4322.diff', 'patch_url': 'https://github.com/huggingface/datasets/pull/4322.patch', 'merged_at': datetime.datetime(2022, 5, 25, 20, 43, 51)}
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.

{'url': 'https://api.github.com/repos/huggingface/datasets/issues/4322/reactions', 'total_count': 3, '+1': 2, '-1': 0, 'laugh': 0, 'hooray': 0, 'confused': 0, 'heart': 1, 'rocket': 0, 'eyes': 0}