Packaging Decisions

A model returns a probability. But the application needs a decision. Call this customer, or not. Where should the threshold live? In the application, copied in five places? Or with the model, versioned, and tested together? Packaging the decision logic with the model is a small idea with big benefits. Let's do it with M L flow.

• Probabilities are not decisions
• Where does the threshold live?
• Package model and rule together

A model with its decision rule

Here is a custom M L flow model. A wrapper class takes our trained model, and a threshold. Its predict method returns a small table, with the churn probability, and a yes or no, call this customer. We train a model as usual, wrap it with the threshold we found earlier, and log it, together with a signature, that records the expected input columns. Then we load it back, by its address, and predict four customers. The output has both columns. The application receives decisions, and the threshold is versioned, and registered with the model.

import mlflow
import pandas as pd
from mlflow.models import infer_signature
from sklearn.linear_model import LogisticRegression
from churn import make_data

mlflow.set_tracking_uri("sqlite:///mlflow.db")

class ChurnWithDecision(mlflow.pyfunc.PythonModel):
    def __init__(self, model, threshold):
        self.model = model
        self.threshold = threshold

    def predict(self, context, model_input, params=None):
        probability = self.model.predict_proba(model_input)[:, 1]
        return pd.DataFrame({
            "churn_probability": probability.round(3),
            "call_customer": probability >= self.threshold,
        })

df = make_data(2000, seed=1)
X, y = df.drop(columns="churn"), df["churn"]
base = LogisticRegression(max_iter=500).fit(X, y)
wrapped = ChurnWithDecision(base, threshold=0.07)

with mlflow.start_run():
    info = mlflow.pyfunc.log_model(
        name="churn-with-decision",
        python_model=wrapped,
        signature=infer_signature(X, wrapped.predict(None, X.head(5))),
    )

loaded = mlflow.pyfunc.load_model(info.model_uri)
print(loaded.predict(X.head(4)).to_string(index=False))
print("expects columns:", [c.name for c in loaded.metadata.signature.inputs])

 churn_probability  call_customer
             0.122           True
             0.209           True
             0.047          False
             0.040          False
expects columns: ['tenure', 'charges', 'calls', 'contract']
  color_warning(
  color_warning(

import numpy as np, pandas as pd

def make_data(n, seed, shifted=False):
    rng = np.random.default_rng(seed)
    tenure = rng.integers(1, 72, n)
    charges = rng.normal(65 + (28 if shifted else 0), 18, n).clip(20, 140).round(2)
    calls = rng.poisson(1.5 + (1.5 if shifted else 0), n)
    contract = rng.choice([0, 1, 2], n, p=[.55, .25, .20] if not shifted else [.30, .30, .40])
    logit = -1.0 - 0.04 * tenure + 0.025 * (charges - 65) + 0.45 * calls - 0.8 * contract
    if shifted:
        logit = -0.2 - 0.01 * tenure - 0.01 * (charges - 65) + 0.2 * calls - 0.2 * contract
    churn = (rng.random(n) < 1 / (1 + np.exp(-logit))).astype(int)
    return pd.DataFrame({"tenure": tenure, "charges": charges, "calls": calls, "contract": contract, "churn": churn})

Why package the decision

Packaging the rule with the model has clear benefits. One place to change the threshold, instead of five. The threshold is versioned and tested, with the model, so rollbacks restore both. The signature documents the expected inputs, and catches mistakes early. The application stays simple, it just asks for a decision. And the registry, and its gate, can evaluate the whole decision, including the business cost, not only the model's score. Keep business rules out of the application, when they depend on the model.

Summary