Step 4 — Build & Train the Model

Phase: Machine Learning

“This is where we actually teach the computer to predict diamond prices. We start simple, then go bigger.”

4a. Split the Data

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(

    X, y, test_size=0.2, random_state=42

)

 

# Training set: 43,152 diamonds

# Test set:     10,788 diamonds

 

Think of this as a practice exam vs. the real exam. The model studies on 43,152 diamonds and gets tested on 10,788 it has never seen before.

4b. Baseline — Linear Regression

from sklearn.linear_model import LinearRegression

lr = LinearRegression()

lr.fit(X_train, y_train)

 

4c. Better Models

Linear regression assumes the relationship between features and price is perfectly linear. In reality, the relationship is curved and features interact with each other. Tree-based models handle this naturally:

from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor

rf  = RandomForestRegressor(n_estimators=100, random_state=42)

gbr = GradientBoostingRegressor(n_estimators=100, random_state=42)

 

rf.fit(X_train, y_train)

gbr.fit(X_train, y_train)

 

💡 Why Random Forest? It builds hundreds of decision trees and averages their predictions. Think of it as polling hundreds of experienced gemologists and averaging their valuations — much more reliable than asking just one.