Fit on train. Score on test. train_test_split shuffles, then cuts. random_state makes the cut repeatable. stratify=y keeps the same label mix in both sides.

Goal

Split a dataset, print shapes and label counts, and see why stratify matters.

A 75 / 25 cut

from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split

X, y = make_classification(n_samples=200, n_features=4, random_state=0)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.25, random_state=0
)
print("train", X_train.shape, "test", X_test.shape)
print("train labels", np.bincount(y_train))
print("test labels", np.bincount(y_test))

test_size=0.25 means a quarter of the rows are test. You can also pass train_size.

Stratify

from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split

X, y = make_classification(
    n_samples=200,
    n_features=4,
    weights=[0.8, 0.2],
    random_state=0,
)
X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.25, random_state=1)
X_s, X_st, y_s, y_st = train_test_split(
    X, y, test_size=0.25, random_state=1, stratify=y
)
print("no stratify  train", np.bincount(y_tr), "test", np.bincount(y_te))
print("stratify     train", np.bincount(y_s), "test", np.bincount(y_st))

With a rare class, an unlucky shuffle can leave almost none of it in test. stratify=y avoids that.

Split a DataFrame together

from sklearn.model_selection import train_test_split

df = pd.DataFrame(
    {
        "units": np.arange(20),
        "price": np.linspace(10, 30, 20),
        "high": [0, 1] * 10,
    }
)
train, test = train_test_split(df, test_size=0.25, random_state=0, stratify=df["high"])
print("train rows", len(train), "test rows", len(test))
print(train.head())

Passing the whole frame keeps columns aligned. Then X_train = train[["units", "price"]] and y_train = train["high"].

Pitfall

Do not scale, encode, or pick features using the test set. Split first, then fit transformers on train only — the Scaling and Pipeline chapters.