Quick Tips · Tip 04 · 0:21
Train/test split with a rare class? Always stratify
With 8% positives in 200 rows, a plain split left 1 positive in 50 test rows, just 2%. Adding stratify=y to train_test_split keeps the real 8%.
Key moments
What you’ll needPython 3.10+, scikit-learn
What you’ll need
- python.org/downloads →
Python 3.10 or newerRuns the code. Check yours with python3 --version. - code.visualstudio.com →
VS Code or CursorEither works; so does any editor you like. - git-scm.com/downloads →
gitOptional for a tip: to keep your own copy of the code. 
scikit-learn 1.4+Model selection and evaluation toolspip install "scikit-learn>=1.4"
Set it up
- 1Make a virtual environmentIt keeps this project's packages separate from the rest of your computer.
python3 -m venv .venv source .venv/bin/activate
- 2Install the package
pip install "scikit-learn>=1.4"
- 3Run itSave the snippet below as split.py.
python3 split.py
split.pyOpen on GitHub →
from sklearn.model_selection import (
train_test_split as split)
y = [1] * 16 + [0] * 184
_, plain = split(y, random_state=59)
_, strat = split(y, random_state=59, stratify=y)
for s in y, plain, strat: print(sum(s) / len(s))Transcript
Your test set might miss the rare class. 16 positives, 200 rows. Plain split: one positive in 50 test rows. Add stratify equals y.
Ratio kept. Print each rate: 8, 2, 8 percent. Rare class? Always stratify.
Read the lesson on GitHub →- #scikitlearn
- #machinelearning
- #python
- #datascience
