train_test_split()

Делит массивы признаков и целевой переменной на обучающую и тестовую подвыборки. Параметр test_size задаёт долю теста, random_state фиксирует случайность для воспроизводимости, а stratify=y сохраняет пропорции классов.

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
print(X_train.shape, X_test.shape)