Machine Learning (ML) · Beginner level

Beginner level · Module 2: Preparing Data for Machine Learning · Lesson 3 of 3

Splitting Data and Avoiding Leakage

Keep test data truly unseen and avoid features that secretly reveal the answer.

In this lesson you will

  • Split data into training and test sets with train_test_split
  • Identify and prevent data leakage

Key concepts

train_test_splitstratifydata leakagefit on training data only

The lesson

train_test_split(X, y, test_size=0.2, random_state=42, stratify=y) holds back 20% of the data for the final test. stratify keeps the same share of each class in both parts, and random_state makes the split repeatable.

Data leakage happens when information from the test set, or from the future, reaches the training process. A feature such as 'date the student left school' predicts dropout perfectly — because it is recorded after dropout.

Learn preparation steps such as the median for imputation from the training set only, then apply them to the test set. Calculating them on all the data leaks test information and makes results look better than they are.

Step by step

  1. Split the data with stratification.
  2. Find the leaking column in the provided data and remove it.
  3. Calculate the imputation median on the training set only.
  4. Apply it to the test set.
Troubleshooting tree

My model scores 99% — is it real?

A checklist for the most common reasons a suspiciously good result is wrong.

Was any preparation (scaling, imputation, feature selection) calculated on all the data before splitting?
  • Yes
    Leakage: refit preparation on training data only — put it in a pipeline — and evaluate again.
  • No
    Does any feature record something that happens after or because of the outcome (e.g. 'date left school' when predicting dropout)?
    • Yes
      Target leakage: remove the feature. It will not be available when predictions are needed.
    • No
      Is the data time-ordered, and was it split randomly?
      • Yes
        The model saw the future. Split by time: train on earlier periods, test on later ones.
      • No
        Is the positive class rare, with accuracy as the metric?
        • Yes
          Accuracy is misleading. Compare with the majority-class baseline and report precision and recall.
        • No
          Check for duplicate records across training and test sets, then ask a colleague to review the notebook before trusting the result.

Notes

In real projects, near-perfect scores are far more often a sign of a mistake than of a brilliant model.

Practice activity

Produce leakage-free training and test sets for the dropout project, with a note on each leak you removed.

Expected result: Correctly split, leakage-free data sets and a note naming each leaking feature.

Check your understanding

  1. What does stratify=y do?

    Show answer

    Keeps the class proportions the same in the training and test sets.

  2. Why calculate imputation values only from the training set?

    Show answer

    Using the test set would leak information and overstate performance.

Taught live by an instructor — online or in person. Free course lessons are for self-study and do not lead to a certificate — see registration options.