Search⌘ K
AI Features
Home
Courses
Data Science Bootcamp

Saving, Packaging, and Serving a Model

Explore how to save, package, and serve machine learning models safely by preserving preprocessing steps within a pipeline. Understand the importance of consistent data transformations during training and inference, define a stable input-output contract, and validate model artifacts to ensure reliable predictions when deployed.

We'll cover the following...

A readmission model can score well on a held-out test set and still fail when a caller sends raw patient data. The failure usually comes from a mismatch between the data preparation used during training and the preparation used for prediction.

Training-time transformations are the changes learned while fitting the pipeline, such as calculating a numeric column’s scaling values or recording the categories seen in admission_type. Inference-time transformations are those same changes applied when new patient records arrive. The deployed artifact must apply identical transformations in the same order and with the same fitted parameters.

Saving only a classifier is unsafe when the classifier expects transformed columns rather than raw fields like age, prior_admissions, and admission_type. A complete Pipeline keeps the fitted preprocessing and the fitted classifier together, so one object accepts the same raw schema during evaluation and later prediction.

Fit one pipeline from raw hospital data

For a binary readmission decision, the pipeline below receives a small table with numeric fields and a categorical admission field. It separates the readmitted target from the input columns, splits the rows, fits preprocessing plus a classifier, and reports held-out accuracy with predicted probabilities.

Python
import io
import joblib
import numpy as np
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler
def fit_readmission_pipeline(random_state: int = 7) -> tuple[Pipeline, pd.DataFrame, pd.DataFrame, pd.Series, pd.Series]:
    df = pd.DataFrame(
        {
            "age": [
                23,
                45,
                67,
                52,
                36,
                74,
                58,
                41,
                29,
                63,
                55,
                48,
                31,
                70,
                60,
                39,
                27,
                66,
                50,
                44,
            ],
            "prior_admissions": [
                0,
                1,
                3,
                2,
                0,
                4,
                2,
                1,
                0,
                3,
                2,
                1,
                0,
                4,
                3,
                1,
                0,
                3,
                2,
                1,
            ],
            "admission_type": [
                "elective",
                "emergency",
                "emergency",
                "urgent",
                "elective",
                "emergency",
                "urgent",
                "urgent",
                "elective",
                "emergency",
                "urgent",
                "emergency",
                "elective",
                "emergency",
                "urgent",
                "urgent",
                "elective",
                "emergency",
                "urgent",
                "emergency",
            ],
            "readmitted": [
                0,
                0,
                1,
                1,
                0,
                1,
                1,
                0,
                0,
                1,
                1,
                0,
                0,
                1,
                1,
                0,
                0,
                1,
                1,
                0,
            ],
        }
    )
    X = df[["age", "prior_admissions", "admission_type"]]
    y = df["readmitted"]
    X_train, X_test, y_train, y_test = train_test_split(
        X,
        y,
        test_size=0.30,
        random_state=random_state,
        stratify=y,
    )
    numeric_features = ["age", "prior_admissions"]
    categorical_features = ["admission_type"]
    preprocessor = ColumnTransformer(
        transformers=[
            ("num", StandardScaler(), numeric_features),
            (
                "cat",
                OneHotEncoder(handle_unknown="ignore", sparse_output=False),
                categorical_features,
            ),
        ],
        remainder="drop",
    )
    pipeline = Pipeline(
        steps=[
            ("preprocessor", preprocessor),
            ("model", LogisticRegression(max_iter=1000, random_state=random_state)),
        ]
    )
    pipeline.fit(X_train, y_train)
    return pipeline, X_train, X_test, y_train, y_test
def serialize_roundtrip(pipeline: Pipeline) -> Pipeline:
    buffer = io.BytesIO()
    joblib.dump(pipeline, buffer)
    buffer.seek(0)
    reloaded = joblib.load(buffer)
    return reloaded
def predict_readmission(pipeline: Pipeline, patient: dict) -> dict:
    required_fields = {
        "age": (int, float, np.integer, np.floating),
        "prior_admissions": (int, float, np.integer, np.floating),
        "admission_type": (str,),
    }
    missing = [field for field in required_fields if field not in patient]
    if missing:
        raise KeyError(f"Missing required field(s): {missing}")
    for field, allowed_types in required_fields.items():
        if not isinstance(patient[field], allowed_types):
            raise TypeError(
                f"Field '{field}' must be of type(s) {allowed_types}, got {type(patient[field]).__name__}"
            )
    row = pd.DataFrame([patient], columns=["age", "prior_admissions", "admission_type"])
    proba_pos = float(pipeline.predict_proba(row)[0, 1])
    threshold = 0.5
    predicted_label = int(proba_pos >= threshold)
    return {
        "readmission_probability": proba_pos,
        "threshold": threshold,
        "predicted_label": predicted_label,
    }
def main() -> None:
    pipeline, _, X_test, _, y_test = fit_readmission_pipeline(random_state=7)
    test_proba_pos = pipeline.predict_proba(X_test)[:, 1]
    y_pred = pipeline.predict(X_test)
    test_accuracy = accuracy_score(y_test, y_pred)
    print(f"Test accuracy: {test_accuracy:.3f}")
    print("Predicted probabilities (positive class):")
    print(test_proba_pos)
    reloaded = serialize_roundtrip(pipeline)
    reloaded_proba_pos = reloaded.predict_proba(X_test)[:, 1]
    max_abs_diff = float(np.max(np.abs(test_proba_pos - reloaded_proba_pos)))
    print(f"Max absolute probability difference after reload: {max_abs_diff}")
    print("Reloaded predicted probabilities (positive class):")
    print(reloaded_proba_pos)
    valid_patient = {
        "age": 61,
        "prior_admissions": 2,
        "admission_type": "emergency",
    }
    print(predict_readmission(reloaded, valid_patient))
    invalid_patient_missing_key = {
        "age": 61,
        "admission_type": "emergency",
    }
    try:
        print(predict_readmission(reloaded, invalid_patient_missing_key))
    except KeyError as exc:
        print(f"Validation error: {exc}")
if __name__ == "__main__":
    main()

  • Lines 122–139: ColumnTransformer keeps numeric scaling and categorical one-hot encoding in a single fitted preprocessing step, and Pipeline packages that preprocessing ...