Enterprise Customer Churn API

A churn-prediction pipeline and serving API: training with experiment tracking, calibrated probabilities, a cost-optimized decision threshold, and a FastAPI service that returns a churn probability plus the top factors driving each at-risk customer.

It is built on the public IBM Telco Customer Churn dataset (7,032 customers after cleaning, 26.6% churn). The goal is to demonstrate the engineering around a model (validation, calibration, explainability, testing, containerization), not to claim results on proprietary enterprise data.

What it does

• Cost-based threshold. Instead of the default 0.5, the decision threshold minimizes expected business cost: every flagged customer receives a $20 retention offer, and every missed churner costs $100. The threshold is tuned on out-of-fold predictions from the training set only, so the test set is touched once.
• Calibrated probabilities. The tuned model is wrapped in CalibratedClassifierCV (sigmoid), and the API serves the calibrated probability, so 0.30 means roughly a 30% churn risk.
• Per-customer explanations. For customers flagged as likely to churn, the API returns up to 3 SHAP-based risk drivers.
• Strict, documented input schema. Pydantic Literal types reject invalid categories with a 422. Secondary fields are optional and imputed by the sklearn pipeline.
• Non-blocking serving. CPU-bound inference and SHAP run in a threadpool; drift logging runs as a background task.
• Tested and containerized. An 11-test pytest suite covers model robustness and the API; a Dockerfile and GitHub Actions workflow handle CI and the container.

Results

Evaluated once on the held-out 20% test split (1,407 customers, stratified), using the threshold chosen on training out-of-fold predictions.

Metric	Value
ROC-AUC	0.833
PR-AUC	0.642
Log loss (calibrated)	0.433
F1 at chosen threshold	0.606
Chosen threshold	0.194
Selected model	RandomForestClassifier

Business cost on the test set (lower is better; assumes every retention offer succeeds):

Strategy	Total cost
Flag nobody	37400
Flag everybody	28140
This model	19240

Project structure

├── .github/workflows/        # CI workflow (runs tests, builds the image)
├── data_load_prep.py         # Downloads Telco data, creates train/test splits
├── model_train_test_track.py # GridSearchCV, calibration, threshold tuning, MLflow logging
├── main.py                   # FastAPI app: /predict, /health, SHAP explanations
├── test_pipeline.py          # Pytest suite (model robustness + API behavior)
├── requirements.txt          # Dependencies used by the Docker image
├── pyproject.toml            # uv project configuration
├── Dockerfile
└── *.joblib                  # Trained artifacts (model, calibrated model, threshold)

Getting started

Uses uv for environments (plain pip works too).

1. Install

git clone https://github.com/Tahamah01/Enterprise-Churn-Prediction-API.git
cd Enterprise-Churn-Prediction-API
uv venv
uv pip install -r requirements.txt

2. Prepare data and train

# Default: stratified 80/20 split (use this for headline metrics)
python data_load_prep.py

# Grid search over RandomForest and XGBoost, calibrate, tune threshold, save artifacts
python model_train_test_track.py

View runs, metrics, and the confusion matrix with mlflow ui (run it from the repository root).

3. Test

pytest test_pipeline.py -v

4. Serve

uvicorn main:app --reload

Interactive docs are available at http://localhost:8000/docs.

Optional: tenure-shift stress test

python data_load_prep.py --split time
python model_train_test_track.py

The Telco dataset has no real dates, so signup_date is synthesized from tenure. Splitting on it trains on long-tenure customers and tests on short-tenure ones. That is a useful covariate-shift stress test, but it is not true out-of-time validation, so headline metrics use the stratified split.

API

POST /predict

Required fields: tenure (int, 0–240), MonthlyCharges (≥ 0), TotalCharges (≥ 0), Contract, InternetService, PaymentMethod, DeviceProtection, TechSupport.

Optional fields (imputed if omitted): gender, SeniorCitizen, Partner, Dependents, PhoneService, MultipleLines, OnlineSecurity, OnlineBackup, StreamingTV, StreamingMovies, PaperlessBilling.

Categorical fields accept only the values seen in the Telco data (for example Contract ∈ Month-to-month, One year, Two year). See the CustomerData model in main.py or /docs for the full list.

curl -X POST http://localhost:8000/predict \
  -H "Content-Type: application/json" \
  -d '{
    "tenure": 1,
    "MonthlyCharges": 95.0,
    "TotalCharges": 95.0,
    "Contract": "Month-to-month",
    "InternetService": "Fiber optic",
    "PaymentMethod": "Electronic check",
    "DeviceProtection": "No",
    "TechSupport": "No"
  }'

Example response (illustrative; exact values depend on the trained model):

{
  "churn_probability": 0.71,
  "will_churn": true,
  "threshold_used": 0.20,
  "retention_action_drivers": [
    {"feature": "tenure", "impact_weight": 0.92},
    {"feature": "Contract_Month-to-month", "impact_weight": 0.61},
    {"feature": "MonthlyCharges", "impact_weight": 0.24}
  ],
  "model_version": "2.1.0"
}

• churn_probability is the calibrated probability.
• will_churn is churn_probability >= threshold_used.
• retention_action_drivers lists up to 3 features pushing the prediction toward churn, and is empty when will_churn is false. Weights are SHAP values of the underlying tree model (log-odds units for XGBoost). Use them to rank drivers, not as percentage points of probability.
• Invalid or unknown input returns 422. Unexpected failures return a generic 500; details are logged server-side only.

GET /health

Returns {"status": "ok", "model_version": "..."}.

Docker

docker build -t enterprise-churn-api .
docker run -d -p 8000:8000 --name churn_service enterprise-churn-api

Prediction logs are written to a JSONL file (DRIFT_LOG_PATH, default production_drift_logs.jsonl) and are lost when the container is removed unless you mount a volume:

docker run -d -p 8000:8000 -e DRIFT_LOG_PATH=/data/drift.jsonl -v churn_logs:/data \
  --name churn_service enterprise-churn-api

Design decisions

Decision	Why
Calibrate, then tune the threshold on calibrated out-of-fold probabilities	Class weighting distorts raw probabilities; the cost trade-off only makes sense on calibrated ones.
Threshold chosen on train OOF predictions, not on test	Avoids optimistic test metrics from choosing the threshold on the data you report on.
Serve the calibrated model, explain with the base model	Calibration is a monotonic rescaling, so driver rankings stay meaningful, but SHAP values do not sum to the calibrated probability.
SHAP only for flagged customers	Keeps latency low for the majority of requests.
Literal validation at the API, handle_unknown="ignore" in the pipeline	The API rejects unknown categories loudly; the pipeline stays safe if used directly.
Sync model calls in a threadpool	Inference is CPU-bound and would otherwise block the event loop.

Monitoring and drift

Every successful prediction is appended to the JSONL log (timestamp, input features, probability, decision). The request schema contains no direct identifiers. The log is meant as input for offline drift analysis (for example with Evidently); no automated drift report is included in this repository.

Limitations

• Public benchmark data. 7k rows of a telecom dataset, not an enterprise-scale system.
• Simplified cost model. Fixed $20 / $100 costs and a 100% offer success rate. Real retention offers succeed only some of the time, which would shift the optimal threshold; change cost_offer, cost_loss, or add a success rate in business_cost().
• No true out-of-time validation. The dataset has no real timestamps (see the stress-test note above).
• Pickled artifacts. .joblib files must be loaded with the same scikit-learn, XGBoost, and SHAP versions used for training. Versions are logged to MLflow as ver_* params; pin exact versions (==) in requirements.txt and retrain when you upgrade.
• Not hardened for the public internet. No authentication, rate limiting, or batch endpoint, and retraining is manual.

Tech stack

scikit-learn · XGBoost · SHAP · pandas · NumPy · MLflow · FastAPI · Pydantic · Uvicorn · pytest · Docker · GitHub Actions · uv