Before training any model, data must be split into separate portions for training and evaluation — testing a model on the same data it was trained on gives a wildly overoptimistic, misleading picture of its real performance.
from sklearn.model_selection import train_test_split
import pandas as pd
X = df.drop("target", axis=1) # features — everything except the column being predicted
y = df["target"] # the label/target being predicted
X_train, X_test, y_train, y_test = train_test_split(
X, y,
test_size=0.2, # 20% held out for testing, 80% for training
random_state=42, # fixes the random split so results are reproducible
stratify=y, # keeps class proportions consistent in both splits
)
print(X_train.shape, X_test.shape) # confirm the split sizes look right
By convention, X (capital, since it's typically a 2D table of multiple features) holds the input data and y (lowercase, a single column) holds the target being predicted — this naming convention appears throughout virtually all scikit-learn code and tutorials.
A model evaluated on data it was trained on will always look better than it actually is, since it's essentially being asked to "recall" answers it already memorized rather than genuinely generalize to new input — this is why a held-out test set, never touched during training, is non-negotiable for any honest performance estimate. random_state fixes the randomness in the split so re-running the same code produces the identical train/test division — essential for reproducible results and fair comparisons between different models or settings.
stratify=y ensures the proportion of each class in y is preserved in both the train and test splits — without it, a random split on an imbalanced dataset (say, 95% one class, 5% another) could by chance produce a test set with very few or even zero examples of the minority class, making evaluation on that class unreliable or impossible.