train_test_split and stratify breaks with Numpy 2.0

[davebulaval]

Describe the bug

As stated in the title, since Numpy changed in version >2.0 with copy, the stratify parameters break.

e.g. all_dataset.train_test_split(test_size=0.2,stratify_by_column="label") returns a Numpy error.

It works if you downgrade Numpy to a version lower than 2.0.

Steps to reproduce the bug

1. Numpy > 2.0
2. all_dataset.train_test_split(test_size=0.2,stratify_by_column="label")

Expected behavior

It returns a stratified split as per the results of Numpy < 2.0

Environment info

• datasets version: 2.14.4
• Platform: Linux-6.8.0-85-generic-x86_64-with-glibc2.35
• Python version: 3.13.7
• Huggingface_hub version: 0.34.4
• PyArrow version: 19.0.0
• Pandas version: 2.3.2

[CloseChoice]

I can't reproduce this. Could you pls provide an example with a public dataset/artificial dataset and show how you loaded that?

This works for me:

import numpy as np
from datasets import Dataset, Features, ClassLabel, Value

data = {"text": [f"sample_{i}" for i in range(100)], "label": [i % 3 for i in range(100)]}
features = Features({"text": Value("string"),
                     "label": ClassLabel(names=["class_0", "class_1", "class_2"])})
dataset = Dataset.from_dict(data, features=features)
splits = dataset.train_test_split(test_size=0.2, stratify_by_column="label")
print(f"Success with numpy {np.__version__}")

but it also works for numpy<2