Building a Machine Learning Model with Scikit-Learn
This guide demonstrates the end-to-end process of building a machine learning classification model using scikit-learn. It covers data loading, preprocessing, training, evaluation, and saving the model for future use.
Modules Used:
- scikit-learn: For machine learning algorithms and utilities.
- joblib: For saving and loading the trained model.
- pandas: For data manipulation (optional, but recommended).
Installation
pip install scikit-learn pandas joblibThe Code
Save this as train_model.py. We will use the classic Wine dataset to classify different types of wines based on their chemical constituents.
import pandas as pd
from sklearn.datasets import load_wine
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, classification_report
import joblib
def train():
# 1. Load Dataset
print("Loading Wine dataset...")
data = load_wine()
X = data.data
y = data.target
feature_names = data.feature_names
target_names = data.target_names
# (Optional) View as DataFrame
df = pd.DataFrame(X, columns=feature_names)
df['target'] = y
print(f"Dataset shape: {df.shape}")
print(df.head())
# 2. Split Data (80% Train, 20% Test)
print("\nSplitting data...")
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=42
)
# 3. Preprocessing (Scaling)
# Many algorithms perform better when features are on the same scale
print("Scaling features...")
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)
# 4. Initialize and Train Model
print("Training Random Forest Classifier...")
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train_scaled, y_train)
# 5. Evaluate
print("Evaluating model...")
y_pred = model.predict(X_test_scaled)
accuracy = accuracy_score(y_test, y_pred)
print(f"\nAccuracy: {accuracy:.4f}")
print("\nClassification Report:")
print(classification_report(y_test, y_pred, target_names=target_names))
# 6. Save Model and Scaler
# We must save the scaler too, to preprocess new data exactly the same way
print("Saving artifacts...")
joblib.dump(model, 'wine_model.joblib')
joblib.dump(scaler, 'wine_scaler.joblib')
print("Model saved to 'wine_model.joblib'")
print("Scaler saved to 'wine_scaler.joblib'")
if __name__ == "__main__":
train()Making Predictions
Once the model is saved, you can load it in a separate script to make predictions on new data.
Save this as predict.py.
import joblib
import numpy as np
def predict_new_data():
# 1. Load Model and Scaler
try:
model = joblib.load('wine_model.joblib')
scaler = joblib.load('wine_scaler.joblib')
except FileNotFoundError:
print("Error: Model files not found. Run train_model.py first.")
return
# 2. New Data (Example: A single sample of wine features)
# In a real app, this might come from an API or user input
# This sample corresponds to class 0 (Gin)
new_sample = np.array([[
13.2, 1.78, 2.14, 11.2, 100.0, 2.65, 2.76, 0.26, 1.28, 4.38, 1.05, 3.40, 1050.0
]])
# 3. Preprocess (Scale)
new_sample_scaled = scaler.transform(new_sample)
# 4. Predict
prediction = model.predict(new_sample_scaled)
probabilities = model.predict_proba(new_sample_scaled)
# Map numeric prediction to class name
# (We know the names from the training step: ['class_0', 'class_1', 'class_2'])
class_names = ['Class 0', 'Class 1', 'Class 2']
predicted_class = class_names[prediction[0]]
print(f"Predicted Class: {predicted_class}")
print(f"Confidence: {probabilities[0][prediction[0]]:.2f}")
if __name__ == "__main__":
predict_new_data()Usage
-
Train the model:
python train_model.py -
Make a prediction:
python predict.py