In this lesson: Split data correctly for any problem type and detect leakage.
A model evaluated on the data it learned from will flatter you. It can memorise. The only honest measure is performance on rows it has never seen.
The basic split
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, # 20% held back
random_state=42, # reproducible
stratify=y, # keep the class balance in both halves
)
stratify=y matters for classification: without it a random split can leave your rare class barely represented in the test set, making the score meaningless. Always pass it for classification.
Three sets, not two
The moment you start comparing models or tuning settings, the test set starts leaking through you — each peek nudges your choices toward it. So:
- Train (60%) — the model learns here.
- Validation (20%) — you compare models and tune here, as often as you like.
- Test (20%) — touched once, at the very end, to report a number.
X_tmp, X_test, y_tmp, y_test = train_test_split(
X, y, test_size=0.2, random_state=42, stratify=y)
X_train, X_val, y_train, y_val = train_test_split(
X_tmp, y_tmp, test_size=0.25, random_state=42, stratify=y_tmp)
# 0.25 of the remaining 80% = 20% of the original
In practice cross-validation on the training set replaces the separate validation set — covered later — but the principle is unchanged: the test set is opened once.
Leakage: the mistake that produces 99% accuracy
Leakage is any information reaching the model that will not be available at prediction time. It is the single most common reason a model that looked excellent in a notebook fails in production. Four forms, in order of how often they occur:
1. Fitting a transformation before splitting.
# WRONG -- the scaler has seen the test rows
X_scaled = StandardScaler().fit_transform(X)
X_train, X_test = train_test_split(X_scaled, ...)
# RIGHT -- fit on train only, apply to test
X_train, X_test, y_train, y_test = train_test_split(X, y, ...)
scaler = StandardScaler().fit(X_train)
X_train_s = scaler.transform(X_train)
X_test_s = scaler.transform(X_test)
The same applies to imputing missing values, encoding categories and selecting features. Every one of them must learn from the training data only. This is precisely why pipelines exist, and the pipeline lesson makes it automatic.
2. A feature that encodes the answer. A final_grade_letter column when predicting the final mark. A discharge_date when predicting whether a patient was admitted. Suspiciously perfect scores almost always mean this.
3. Using the future to predict the past. For anything time-ordered, a random split trains on next month and tests on last month. Split by time:
train = df[df["date"] < "2026-01-01"]
test = df[df["date"] >= "2026-01-01"]
# or for cross-validation
from sklearn.model_selection import TimeSeriesSplit
cv = TimeSeriesSplit(n_splits=5)
4. Rows from the same entity in both halves. Three visits by the same patient, three terms for the same student, several photos of the same person. The model recognises the entity, not the pattern. Split by group:
from sklearn.model_selection import GroupShuffleSplit
splitter = GroupShuffleSplit(test_size=0.2, random_state=42)
train_idx, test_idx = next(splitter.split(X, y, groups=df["student_id"]))
The size of the split
- Under a few thousand rows: 80/20 and rely on cross-validation, because a single 20% test set is too noisy to compare models with.
- Tens of thousands: 80/20 is comfortable.
- Millions: 1% can be a plenty-large test set; you are limited by compute, not data.
And keep random_state fixed. Not for luck — so that a change in your score is caused by your change, not by a different shuffle.
Try it yourself
Split your data three ways with stratification and confirm the class proportions match in all three with y_train.value_counts(normalize=True). Then go through your feature list and mark each one: available before the prediction moment or not. Delete the second group.