How do I split data into training and test sets in scikit-learn?

The short version

Use train_test_split to hold back some rows for testing so you can estimate how the model performs on data it has not seen. Set random_state so the split is reproducible, use stratify for classification to preserve class balance, and split before any preprocessing that learns from the data, to avoid leakage.

The basic call

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.25, random_state=0,
)

This holds back 25 per cent of the rows as a test set. The features and target are split together, so each row's X and y stay aligned. The order of the four returned values is train features, test features, train target, test target.

Why you need a test set at all

A model scored on the same data it trained on looks better than it is, because it can memorise those rows. Holding back a test set gives an honest estimate of performance on new data. The test set must stay untouched until the very end; if you make choices based on it, it stops being an honest measure.

Set random_state for reproducibility

Without random_state, the split is different every run, so your results are not repeatable and a colleague cannot recreate them. Fixing it to any number makes the split deterministic.

train_test_split(X, y, random_state=42)   # same split every time

Stratify for classification

For classification, especially with uneven class sizes, a plain random split can give the training and test sets different class proportions. stratify=y keeps them representative.

train_test_split(X, y, stratify=y, random_state=0)

Without this, a rare class might be almost absent from one side, which distorts both training and evaluation.

Split before preprocessing

The order matters for correctness. Split first, then fit any scaler, imputer or encoder on the training set alone. If you preprocess the whole dataset before splitting, the test set influences the transformation and leaks into training, inflating your scores. A pipeline makes this correct order automatic.

The three-way split

Test data is for the final check only. If you also need data to tune choices during development, hold out a validation set too, or use cross-validation on the training set. A common pattern is train plus validation for development, with a test set reserved strictly for the one final measurement.

X_train, X_temp, y_train, y_temp = train_test_split(X, y, test_size=0.3, random_state=0)
X_val, X_test, y_val, y_test = train_test_split(X_temp, y_temp, test_size=0.5, random_state=0)

This gives roughly 70 per cent train, 15 per cent validation and 15 per cent test, keeping the final test set genuinely unseen.