Saving Models (joblib)

from sklearn.pipeline import Pipeline
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.linear_model import LogisticRegression
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
import joblib
import pickle
import json

# 1) Train a pipeline (not just the model)
X, y = load_breast_cancer(return_X_y=True, as_frame=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

pipe = Pipeline([
    ('scale', StandardScaler()),
    ('clf',   LogisticRegression(max_iter=1000)),
]).fit(X_train, y_train)
print(pipe.score(X_test, y_test))

# 2) joblib — sklearn's preferred format
joblib.dump(pipe, 'model.joblib', compress=3)
loaded = joblib.load('model.joblib')
print(loaded.predict(X_test[:3]))

# joblib compresses big NumPy arrays better than vanilla pickle.
# compress=3 → balance speed + size; compress=9 → smallest

# 3) pickle — generic Python
with open('model.pkl', 'wb') as f:
    pickle.dump(pipe, f, protocol=pickle.HIGHEST_PROTOCOL)
with open('model.pkl', 'rb') as f:
    model = pickle.load(f)

# CAVEAT: pickle is NOT a security boundary. Loading an untrusted .pkl executes arbitrary code.
# Only load files YOU produced or fully trust. Sign + verify in production pipelines.

# 4) Save the whole training context
import sklearn, joblib, sys, platform
metadata = {
    'sklearn_version': sklearn.__version__,
    'python_version':  sys.version,
    'platform':        platform.platform(),
    'feature_names':   list(X.columns),
    'target_classes':  ['benign', 'malignant'],
    'training_score':  pipe.score(X_train, y_train),
    'test_score':      pipe.score(X_test, y_test),
    'trained_at':      '2024-01-15T03:00:00Z',
}
with open('model.json', 'w') as f:
    json.dump(metadata, f, indent=2)

# Ship model.joblib + model.json together.

# 5) Load + validate version
meta = json.load(open('model.json'))
assert meta['sklearn_version'] == sklearn.__version__, 'sklearn version mismatch'

# 6) Pipeline with mixed columns (numeric + categorical)
numeric_cols = ['age', 'tenure']
cat_cols     = ['country', 'plan']

preprocess = ColumnTransformer([
    ('num', StandardScaler(),                                       numeric_cols),
    ('cat', OneHotEncoder(handle_unknown='ignore', sparse_output=False), cat_cols),
])

pipe = Pipeline([
    ('pre', preprocess),
    ('clf', LogisticRegression(max_iter=1000)),
]).fit(X_train, y_train)

joblib.dump(pipe, 'model.joblib')
# The fitted ColumnTransformer remembers the schema; serving code only needs to pass a DataFrame with the same cols.

# 7) Pickle compatibility caveats
# • Pickled sklearn models tied to scikit-learn version — load with the SAME version
# • Persist in train env, retrain in serve env if minor version differs
# • For cross-version portability, use ONNX or PMML

# 8) ONNX — cross-framework deployment
# pip install skl2onnx onnxruntime
import onnx
import onnxruntime as ort
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType

initial_type = [('float_input', FloatTensorType([None, X_train.shape[1]]))]
onx = convert_sklearn(pipe, initial_types=initial_type, target_opset=15)
with open('model.onnx', 'wb') as f:
    f.write(onx.SerializeToString())

# Load + predict
sess = ort.InferenceSession('model.onnx')
import numpy as np
proba = sess.run(['probabilities'], { 'float_input': X_test.values.astype(np.float32) })

# ONNX runs anywhere — Python, C++, .NET, mobile, browser via onnx-web.
# Convert from sklearn, PyTorch, TensorFlow, XGBoost, LightGBM.

# 9) PyTorch — torch.save with state_dict
import torch
import torch.nn as nn

model = MyNet()
torch.save(model.state_dict(), 'model.pt')

# Loading
model = MyNet()
model.load_state_dict(torch.load('model.pt', map_location='cpu'))
model.eval()

# Save with metadata
checkpoint = {
    'state_dict': model.state_dict(),
    'optimizer':  opt.state_dict(),
    'epoch':      10,
    'best_acc':   0.94,
}
torch.save(checkpoint, 'ckpt.pt')

# TorchScript — portable to C++ / mobile
scripted = torch.jit.script(model)
scripted.save('model.ts.pt')

# 10) TensorFlow / Keras
import tensorflow as tf
model.save('saved_model')                         # SavedModel (directory)
model.save('model.keras')                          # single file (newer)
model.save('model.h5')                             # legacy HDF5

loaded = tf.keras.models.load_model('model.keras')

# 11) Cloud + registry options
# • MLflow — tracking + model registry; one line to log + serve
# • Vertex AI Model Registry / SageMaker Model Registry
# • Weights & Biases artifacts
# • DVC for git-tracked model files
# Always store: model artefact + metrics + code commit + data snapshot

import mlflow
mlflow.sklearn.log_model(pipe, 'model')

# 12) Versioning strategies
# • SemVer per model: 1.0.0 → 1.0.1 (data drift), 1.1.0 (feature added), 2.0.0 (breaking schema)
# • Tag artefacts with git SHA + dataset hash
# • Keep last N versions deployable for instant rollback

# 13) Quantisation + pruning for serving
# • sklearn — usually doesn't matter; sklearn models are already small
# • PyTorch quantization — int8 cuts size 4x with minimal accuracy loss
# • TensorFlow Lite — mobile-ready conversions
# • ONNX Runtime + GraphOptimizationLevel

# 14) Inference services
# Wrap the model in:
#   • FastAPI + uvicorn — simple REST service
#   • BentoML — model serving framework
//   • TF Serving / TorchServe — production-grade
//   • Triton Inference Server — multi-framework + GPU
//   • SageMaker / Vertex AI endpoints — managed

# 15) Common bugs
# • Pickling the model only, not the preprocessing → inference output garbage; pickle the Pipeline
# • Loading pickle from an UNTRUSTED source → arbitrary code execution; sign + verify or use ONNX
# • sklearn version drift between train + serve → unpickling fails; pin versions in serve image
# • Forgetting to call model.eval() on PyTorch → dropout / batchnorm still active in inference
# • Serving without GPU → batch size matters; tune for CPU latency
# • No metadata → can't reproduce; embed feature names + classes
# • model.joblib size > 100 MB committed to git → use git LFS or model registry
# • Caching old model across deployments → versioned URLs + cache-bust
# • Mismatched feature ORDER between train + inference → silent wrong predictions; pass DataFrame with named columns

import joblib
joblib.dump(pipe, 'model.joblib')
pipe2 = joblib.load('model.joblib')
print(pipe2.predict(Xte[:5]))