Customer Retention Prediction

An end-to-end churn model: exploratory analysis of 7,043 customer records, compared and tuned ensemble classifiers, MLflow tracking, and a FastAPI and Gradio service with tests for training–serving consistency.

Focus: Machine Learning, MLOps
Core stack: XGBoost, MLflow, FastAPI
Links: Source on GitHub(opens in a new tab)

Decision threshold · tuned XGBoostRecorded metrics

Recall: 94.1%
Precision: 42.2%
ROC-AUC: 0.849

Measured at 0.25 on a split also used for tuning — optimistic, not a final estimate.

Framing

Which mistake costs more?

The question

Given a telecom customer's account, services and billing details, how likely are they to leave? A model that flags likely churners lets a retention team act before the customer goes.

The trade-off

Missing a customer who leaves is usually costlier than contacting one who would have stayed. That makes recall the priority — and makes the decision threshold, not just the model, a design choice.

Exploration

What the data says first

Exploratory analysis of the 7,043 records shows the classes are imbalanced and that churn concentrates in recognisable groups: month-to-month contracts, short tenure, fibre-optic internet and electronic-check payment are all associated with churn, while long tenure is associated with staying.

Heatmap of each feature's correlation with churn, from fibre-optic internet and electronic-check payment (about +0.3) down to tenure (about −0.35).

Bar chart of churned and retained customers by contract type; month-to-month contracts account for most churners.

Modelling

Compare, tune, track

After validation, preprocessing and a stratified train/test split, several ensemble classifiers were compared with class imbalance taken into account. XGBoost was tuned with Optuna, and every run — parameters, metrics and the resulting model artifact — was logged to MLflow, so the served model can be traced back to the experiment that produced it.

From raw records to a served model

Diagram

Telco records7,043 rows

Validate & cleanpreprocessing

Encodefixed category levels

Compare modelsXGBoost · LightGBM · RF

TuneOptuna

MLflowruns · saved model

ServeFastAPI + Gradio

Fig. 1Pipeline overview. The same encoding step is shared by training and serving.

Results

A recall-oriented operating point

Recall: 94.1%
Precision: 42.2%
ROC-AUC: 0.849
Threshold: 0.25

Tuned XGBoost, as recorded with the saved MLflow model. Recall and precision are measured at a 0.25 decision threshold.

Lowering the threshold to 0.25 catches about 94% of churners, at the cost of precision: fewer than half of the flagged customers actually churn. Whether that is a good trade depends on what a retention contact costs — which is exactly why the threshold is reported alongside the numbers.

An important caveat. The Optuna objective was scored on the same test split later used to report these results, so they are optimistic and should not be read as an independent estimate of performance on new customers.

Training–serving consistency

The bug that only appeared for one customer

A prediction request for a single customer could produce the wrong feature columns. With one-hot encoding and drop_first=True, a single row only contains one category per column — so the encoder can drop exactly the column the model needs. The serving path now fixes each categorical column's levels before encoding and selects the training feature list explicitly.

Before — one request, one row

row = {"Contract": "Two year", …}
pd.get_dummies(df, drop_first=True)

# only one category is present, so
# drop_first drops it — the column
# the model expects never appears

After — levels fixed up front

df[col] = pd.Categorical(
    df[col], categories=LEVELS[col])
X = pd.get_dummies(df, drop_first=True)
X = X[FEATURE_COLS]

# same columns for 1 row or 10,000

Fig. 4Simplified from the serving code. The fix is covered by a single-row regression test and row-by-row parity tests against training-time encoding.

Serving

An API and an interface

The MLflow model is served behind a FastAPI /predict endpoint, with a Gradio interface mounted on the same app for manual testing. A CI workflow builds and publishes the Docker image.

Gradio 'Customer Churn Predictor' form with demographic, phone and internet service inputs, a Predict Churn button, and a prediction result reading 'Not likely to churn', 'low-risk'.

Limitations

Read the numbers with these in mind

Technologies

Stack

Modelling: Python, Pandas, scikit-learn, XGBoost, LightGBM, Random Forest
Experimentation: Optuna, MLflow
Serving: FastAPI, Gradio
Engineering: pytest, Docker, GitHub Actions

Source on GitHub

All projects