ROC-AUC is the default metric for a binary classifier, and it’s a ranking metric. A model with ROC-AUC of 0.82 can order flights from most-likely-delayed to least-likely-delayed reasonably well. What it can’t tell you is whether the model’s predicted probability of 40% corresponds to flights that are actually delayed 40% of the time.
For a classification benchmark, that doesn’t matter much. For a risk gauge on a dashboard where someone is looking at a number and deciding whether to rebook, it matters a lot. A model that always predicts 40% delay probability regardless of the true underlying rate will have fine ROC-AUC and completely uninformative probabilities. Isotonic calibration is the step that closes that gap — it remaps raw model scores to honest probability estimates, measured by Brier score rather than AUC.
That’s the central concern of this tutorial: building a model pipeline where the probability output means what it claims to mean, each prediction comes with a SHAP-based explanation, and the model artifact can be served without the training framework installed.
What You Need Coming In
- Comfortable with Python, pandas, and scikit-learn’s API conventions
- Familiar with gradient boosting conceptually (trees, learning rate, iterations)
- Basic probability theory — what “calibration” means in a statistical sense
LightGBM on Imbalanced Data
The prediction target is dep_del15 — whether a departure delay exceeds 15 minutes. Around 20% of flights in the BTS dataset qualify, which means a naive model that always predicts “on time” scores 80% accuracy. That’s not useful.
LightGBM handles class imbalance through scale_pos_weight, which upweights the minority class (delayed flights) during training:
def train(X_train, y_train, X_valid, y_valid, cat_features, config):
pos = int(y_train.sum())
neg = len(y_train) - pos
params = {
"objective": "binary",
"metric": ["binary_logloss", "auc"],
"scale_pos_weight": neg / pos, # compensate for ~4:1 imbalance
"learning_rate": config.get("learning_rate", 0.05),
"num_leaves": config.get("num_leaves", 63),
"min_data_in_leaf": config.get("min_data_in_leaf", 200),
"verbose": -1,
}
train_ds = lgb.Dataset(X_train, label=y_train,
categorical_feature=cat_features, free_raw_data=False)
valid_ds = lgb.Dataset(X_valid, label=y_valid,
categorical_feature=cat_features, free_raw_data=False)
booster = lgb.train(
params,
train_ds,
num_boost_round=2000,
valid_sets=[valid_ds],
callbacks=[lgb.early_stopping(50, verbose=False)],
)
return TrainedModel(booster=booster, params=params, n_features=X_train.shape[1])scale_pos_weight = neg / pos tells LightGBM to weight each delayed flight’s gradient contribution by the class ratio — roughly 4x for this dataset. Early stopping uses the held-out validation set (the final months of training data, never test data) and stops when binary logloss stops improving for 50 rounds.
categorical_feature passes the feature names for categorical columns directly to LightGBM rather than one-hot encoding. LightGBM’s native categorical handling uses histogram-based splitting on category codes, which is faster and often better than dummy variables for high-cardinality features like airport codes.
Calibration as a Separate Step
The booster trained above produces raw scores, not calibrated probabilities. The sigmoid of the booster’s raw output is a probability in the mathematical sense (between 0 and 1), but it’s not necessarily a calibrated probability in the statistical sense (a model that outputs 0.4 for a set of flights should see 40% of those flights actually delayed).
Isotonic calibration fits a piecewise-constant monotone function from raw scores to observed frequencies on a held-out calibration slice — the last 10% of training data by date:
from sklearn.isotonic import IsotonicRegression
from sklearn.metrics import brier_score_loss
def fit_calibrator(booster, X_calib, y_calib):
raw_scores = booster.predict(X_calib) # uncalibrated probabilities
# Brier score measures calibration quality (lower = better)
brier_before = brier_score_loss(y_calib, raw_scores)
print(f"Brier before calibration: {brier_before:.4f}")
# Fit isotonic regression: maps raw scores → honest frequencies
iso = IsotonicRegression(out_of_bounds="clip")
iso.fit(raw_scores, y_calib)
calibrated = iso.transform(raw_scores)
brier_after = brier_score_loss(y_calib, calibrated)
print(f"Brier after calibration: {brier_after:.4f}")
return isoThe calibration slice has to be temporally after the training data — months the model hasn’t seen — so that the calibrator learns to correct the model’s overconfidence or underconfidence on genuinely out-of-sample data. Using in-sample data to calibrate would produce a calibrator that looks correct on the calibration set and overcorrects on unseen data.
Isotonic regression is the right choice over Platt scaling (sigmoid) here because the miscalibration pattern in LightGBM outputs tends to be nonlinear. Platt scaling fits a two-parameter logistic curve, which works well when the miscalibration is roughly sigmoid-shaped. Isotonic regression makes no assumption about the shape — it only requires monotonicity, which is a weaker constraint.
At serving time, every prediction runs through both steps:
raw = booster.predict(X)
calibrated = calibrator.transform(raw)The Brier score after calibration is the number to track across retraining runs. A model where Brier score is drifting upward between retrains is losing calibration — probably because the underlying delay patterns have shifted.
SHAP Explanations Per Prediction
The risk gauge shows a number. The SHAP factor bars show why. Both come from the same model; one is the scalar prediction, the other is the per-feature decomposition of that prediction.
The explainer is built from the booster directly:
import shap
def build_explainer(booster):
# TreeExplainer handles LightGBM natively — no sampling, exact computation
explainer = shap.TreeExplainer(booster)
return explainerTreeExplainer uses the model’s tree structure to compute exact SHAP values — it doesn’t sample or approximate. For tree ensembles on moderately sized feature sets (15-20 features), this is fast enough to run per-request.
At serving time, SHAP values decompose a single prediction into contributions from each feature:
def top_factors(explainer, X_row, feature_names, k=5):
"""Return top-k factors driving the prediction, signed."""
shap_vals = explainer.shap_values(X_row)
# SHAP output shape varies by version: (1, n_features) or (n_features,)
vals = shap_vals[0] if shap_vals.ndim > 1 else shap_vals
factors = [
{
"feature": feature_names[i],
"value": float(X_row.iloc[0, i]),
"contribution": float(vals[i]),
"direction": "up" if vals[i] > 0 else "down",
}
for i in range(len(vals))
]
# Sort by absolute contribution, return top-k
return sorted(factors, key=lambda f: abs(f["contribution"]), reverse=True)[:k]A positive contribution means this feature pushed the delay probability up relative to the baseline (the average prediction across training data). Negative means it pulled it down. For a flight departing at 6am from a low-congestion airport with a carrier that has a clean recent record, you’d expect dep_hour, origin_congestion, and carrier_prior_delay_rate all to have negative contributions — each one is evidence against delay.
The factors go directly into the API response:
{
"delay_probability": 0.31,
"factors": [
{ "feature": "route_prior_delay_rate", "value": 0.28, "contribution": 0.12, "direction": "up" },
{ "feature": "origin_congestion", "value": 1.4, "contribution": 0.08, "direction": "up" },
{ "feature": "dep_hour", "value": 7, "contribution": -0.06, "direction": "down" }
]
}This is the difference between a prediction that asks for trust and one that earns it. 31% with no explanation is just a number. 31% because this route historically delays 28% of flights and the airport is currently congested, offset by an early-morning departure, is information.
ONNX Export and the Parity Test
The serving container doesn’t need LightGBM installed. The model exports to ONNX and runs under onnxruntime, which is a smaller dependency with better cross-platform portability:
from lightgbm import Booster
from onnxmltools import convert_lightgbm
from onnxmltools.convert.common.data_types import FloatTensorType
def export_onnx(booster: Booster, n_features: int, output_path: str):
initial_type = [("X", FloatTensorType([None, n_features]))]
onnx_model = convert_lightgbm(booster, initial_types=initial_type)
with open(output_path, "wb") as f:
f.write(onnx_model.SerializeToString())The export itself is fast. The important step is the parity test that runs immediately after and validates the ONNX runtime produces the same outputs as the native booster:
import numpy as np
import onnxruntime as rt
def portability_test(booster, onnx_path, X_sample, atol=1e-4):
"""Assert ONNX runtime matches native booster within floating-point tolerance."""
# Native LightGBM predictions
native_probs = booster.predict(X_sample)
# ONNX runtime predictions
sess = rt.InferenceSession(onnx_path)
X_float32 = X_sample.values.astype(np.float32)
onnx_out = sess.run(None, {"X": X_float32})
onnx_probs = onnx_out[1][:, 1] # column 1 is P(delayed)
max_diff = np.abs(native_probs - onnx_probs).max()
if max_diff > atol:
raise AssertionError(
f"ONNX parity test failed — max difference {max_diff:.6f} exceeds {atol}"
)
print(f"ONNX parity: max diff = {max_diff:.2e} ✓")The tolerance of 1e-4 accounts for float32 vs float64 precision differences between the native booster and the ONNX runtime. If the conversion produced a semantically different model — wrong tree structure, wrong feature ordering, wrong output mapping — the differences would be orders of magnitude larger.
This test runs as part of the build pipeline, not just once manually. A model that fails the parity test doesn’t get registered in MLflow and doesn’t get deployed. The serving layer loads the ONNX bundle and runs the calibrator on top:
class ModelBundle:
def __init__(self, onnx_path, calibrator, explainer, feature_names):
self.session = rt.InferenceSession(onnx_path)
self.calibrator = calibrator
self.explainer = explainer
self.features = feature_names
def predict(self, X):
raw = self.session.run(None, {"X": X.values.astype(np.float32)})[1][:, 1]
return self.calibrator.transform(raw)
def explain(self, X_row):
return top_factors(self.explainer, X_row, self.features)ModelBundle is what FastAPI loads at startup. Inference is a session run (no LightGBM import) plus an isotonic transform (a numpy lookup). SHAP computation still uses the original booster — the TreeExplainer requires the tree structure, which isn’t preserved in ONNX. In practice this means the serving container needs both onnxruntime and the serialized booster object, but not the lightgbm package itself.
Evaluation as Documentation
The evaluation script writes a metrics.json that gets logged to MLflow alongside the model artifact:
metrics = {
"roc_auc": roc_auc_score(y_test, probs_cal),
"pr_auc": average_precision_score(y_test, probs_cal),
"brier_raw": brier_score_loss(y_test, probs_raw),
"brier_calibrated": brier_score_loss(y_test, probs_cal),
"baseline_brier": brier_score_loss(y_test, np.full_like(y_test, y_test.mean(), dtype=float)),
"n_test": len(y_test),
"test_years": TEST_YEARS,
}The baseline Brier score is computed from the trivial “always predict the base rate” model. If brier_calibrated is close to baseline_brier, calibration was still poor — the model isn’t adding much over just knowing the average delay rate. If it’s substantially lower, calibration worked and the model has useful per-flight discrimination.
test_years goes into the artifact metadata because it’s not obvious from the metrics alone. A future retraining run that changes TEST_YEARS produces different numbers that are not directly comparable to the previous run. Recording which years went into the test set makes the comparison interpretable.
