39. End-to-End Machine Learning Project (India Dataset)
39.2 Load, Clean and Split
Why. Dirty X ruins every model. Load, peek, form X/y, split.
What. Build the fictional table in pandas (in real life you would read_csv). Confirm dtypes, check class balance, stratify split.
How.
Prepare the table
- Create the DataFrame.
- Define X and y.
- value_counts on y.
- train_test_split with stratify and random_state.
import pandas as pd
from sklearn.model_selection import train_test_split
df = pd.DataFrame({
"shop": [f"S{i}" for i in range(1, 21)],
"litres_last_week": [120, 80, 200, 60, 180, 90, 220, 70, 150, 110,
95, 210, 130, 75, 160, 100, 240, 85, 140, 190],
"skus": [8, 5, 12, 4, 11, 6, 14, 5, 9, 7, 6, 13, 8, 5, 10, 7, 15, 5, 9, 12],
"distance_km": [12, 40, 8, 55, 10, 35, 6, 48, 15, 22, 30, 9, 18, 42, 14, 28, 5, 50, 16, 11],
"reorder_this_week": [1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 0, 1, 1, 0, 1, 0, 1, 0, 1, 1],
})
X = df[["litres_last_week", "skus", "distance_km"]]
y = df["reorder_this_week"]
print(y.value_counts())
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.3, random_state=42, stratify=y
)
print("train", len(X_train), "test", len(X_test))
What you should see. Both classes present; train larger than test. Keep this split for the rest of the chapter.