Features, Labels and train_test_split in scikit-learn

Every supervised learning dataset has two parts:

• Features (also called inputs, predictors or X): the information the model gets, such as size, location and age of a house.
• Label (also called target or y): the answer the model should predict, such as the house's price.

X and y from a DataFrame#

import pandas as pd

houses = pd.DataFrame({
    "area_sqft": [650, 900, 1200, 1500, 800, 1100],
    "bedrooms": [1, 2, 2, 3, 2, 3],
    "age_years": [15, 8, 5, 2, 20, 10],
    "price_lakh": [45, 68, 92, 120, 52, 85],
})
X = houses.drop(columns="price_lakh")     # everything except the answer
y = houses["price_lakh"]
print(X.shape, y.shape)
print(X.columns.tolist())

(6, 3) (6,)
['area_sqft', 'bedrooms', 'age_years']

By convention X is a capital letter (a 2D table) and y lowercase (a 1D column).

Why split the data?#

A model can memorise its training data. Testing it on the same data says nothing about how it handles new cases, just as a student who saw the exam answers in advance would score perfectly without understanding anything. So we hold back part of the data as a test set the model never sees during training:

from sklearn.datasets import load_diabetes
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeRegressor

X, y = load_diabetes(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=1)

tree = DecisionTreeRegressor(random_state=1).fit(X_train, y_train)
print("R² on training data:", round(tree.score(X_train, y_train), 3))
print("R² on test data:    ", round(tree.score(X_test, y_test), 3))

R² on training data: 1.0
R² on test data:     -0.258

The unlimited decision tree scores a perfect 1.0 on its training data, because it memorised every example, but poorly on new patients. (R² is 1 for perfect predictions and 0 for no better than predicting the average; it can even be negative.) The test score is the honest one. This gap between training and test performance is overfitting, covered in Overfitting and the Bias-Variance Trade-off.

train_test_split options#

import numpy as np
from sklearn.model_selection import train_test_split

X = np.arange(20).reshape(10, 2)
y = np.array([0, 0, 0, 0, 0, 0, 0, 1, 1, 1])

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=0, stratify=y)
print(len(X_train), len(X_test))
print("train labels:", y_train.tolist(), "test labels:", y_test.tolist())

7 3
train labels: [0, 0, 0, 1, 1, 0, 0] test labels: [0, 0, 1]

• test_size=0.3 holds back 30% (20–30% is typical; more data means you can hold back less).
• random_state fixes the shuffle so the split is reproducible.
• stratify=y keeps the class proportions the same in both parts, which matters when one class is rare.

A validation set#

When you compare many models or settings, using the test set to choose between them leaks information. Either split three ways (train / validation / test), or use cross-validation on the training data, which is usually better for small and medium datasets:

import numpy as np
from sklearn.model_selection import train_test_split

X = np.arange(100).reshape(50, 2)
y = np.arange(50)
X_temp, X_test, y_temp, y_test = train_test_split(X, y, test_size=0.2, random_state=0)
X_train, X_val, y_train, y_val = train_test_split(X_temp, y_temp, test_size=0.25, random_state=0)
print(len(X_train), len(X_val), len(X_test))

30 10 10

That gives a 60 / 20 / 20 split: train models on the training set, choose between them on the validation set, and report the test score once at the end.

Time-based splits#

For data over time (sales, sensor readings, prices), a random split lets the model "see the future". Train on the past and test on the most recent period instead:

import pandas as pd

sales = pd.DataFrame({"date": pd.date_range("2026-01-01", periods=10, freq="W"), "units": range(10)})
cutoff = pd.Timestamp("2026-02-15")
train = sales[sales["date"] < cutoff]
test = sales[sales["date"] >= cutoff]
print(len(train), len(test))
print(train["date"].max().date(), "<", test["date"].min().date())

6 4
2026-02-08 < 2026-02-15

scikit-learn's TimeSeriesSplit does this repeatedly for cross-validation.

Common mistakes#

Practice#

import numpy as np
from sklearn.datasets import load_wine
from sklearn.model_selection import train_test_split

X, y = load_wine(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, stratify=y, random_state=7)
print(len(X_train), len(X_test))
print(np.bincount(y_train), np.bincount(y_test))

133 45
[44 53 36] [15 18 12]

Key takeaway#

Next, the first algorithm: Linear Regression.