Skip to content

Commit f2a06ea

Browse files
author
s3ak6i-dev
committed
feat: add pre-trained churn model, metadata, and inference script
1 parent ee140ca commit f2a06ea

4 files changed

Lines changed: 204 additions & 9 deletions

File tree

README.md

Lines changed: 47 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -43,11 +43,20 @@
4343

4444
---
4545

46-
## Try It — Sample Dataset
46+
## Try It — Sample Dataset & Model
4747

48-
A ready-to-use demo dataset is included at [`sample_data/customer_churn_demo.csv`](sample_data/customer_churn_demo.csv).
48+
Three files in [`sample_data/`](sample_data/) let you test Datrix end-to-end and run inference immediately:
4949

50-
It's a realistic **SaaS customer churn dataset** (315 rows, 10 columns) with intentional quality issues embedded across all five dimensions the Quality Engine checks:
50+
| File | Description |
51+
|---|---|
52+
| [`customer_churn_demo.csv`](sample_data/customer_churn_demo.csv) | 315-row SaaS churn dataset with embedded quality issues |
53+
| [`churn_model.joblib`](sample_data/churn_model.joblib) | Pre-trained RandomForest churn predictor (scikit-learn pipeline) |
54+
| [`churn_model_metadata.json`](sample_data/churn_model_metadata.json) | Model card — features, metrics, class distribution |
55+
| [`predict.py`](sample_data/predict.py) | Inference script — batch predictions or interactive single-customer mode |
56+
57+
### Dataset — what Datrix detects
58+
59+
The CSV has intentional quality issues across all five Quality Engine dimensions:
5160

5261
| Issue | Column | Severity | What Datrix catches |
5362
|---|---|---|---|
@@ -60,12 +69,41 @@ It's a realistic **SaaS customer churn dataset** (315 rows, 10 columns) with int
6069
| Log-normal skew | `monthly_spend` | ℹ️ Info | Distribution — high skewness |
6170
| 5:1 class imbalance | `churned` | ⚠️ Warning | Label quality — class imbalance |
6271

63-
**To test it:**
64-
1. Go to [datrix-test.vercel.app](https://datrix-test.vercel.app) and register a free account
65-
2. Navigate to **Datasets → Upload**
66-
3. Upload `sample_data/customer_churn_demo.csv`
67-
4. The Quality Engine scans automatically — you'll see an overall score and a ranked issue list within seconds
68-
5. Try the **Cleaning Wizard** to fix nulls and drop duplicates, then re-scan to watch the score improve
72+
### Model — quick start
73+
74+
```bash
75+
pip install scikit-learn pandas joblib
76+
77+
# Batch predictions on the demo CSV (saves *_predictions.csv)
78+
python sample_data/predict.py
79+
80+
# Run on your own CSV
81+
python sample_data/predict.py --file path/to/your_data.csv
82+
83+
# Interactive single-customer prediction
84+
python sample_data/predict.py --single
85+
```
86+
87+
Model card:
88+
89+
| | |
90+
|---|---|
91+
| **Algorithm** | RandomForestClassifier (200 trees, balanced class weights) |
92+
| **Target** | `churned` — binary yes/no |
93+
| **Features** | age, tenure, spend, tickets, last login, NPS, region, plan |
94+
| **Accuracy** | 80.6% |
95+
| **ROC-AUC** | 0.577 (on messy raw data — by design) |
96+
| **Top predictors** | tenure\_months, monthly\_spend, nps\_score, age |
97+
98+
> The modest AUC on the raw file is intentional — it demonstrates the data quality → model performance link. Upload to Datrix, fix the issues with the Cleaning Wizard, retrain, and watch the score improve.
99+
100+
### Testing it in Datrix
101+
102+
1. Go to [datrix-test.vercel.app](https://datrix-test.vercel.app) and register
103+
2. **Datasets → Upload** `customer_churn_demo.csv`
104+
3. The Quality Engine scans automatically — ranked issues appear within seconds
105+
4. Use the **Cleaning Wizard** to fix nulls and drop duplicates, then re-scan to watch the score climb
106+
5. Head to **Benchmark** to train your own model and compare against the pre-trained one
69107

70108
---
71109

sample_data/churn_model.joblib

703 KB
Binary file not shown.
Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
{
2+
"model_type": "RandomForestClassifier",
3+
"n_estimators": 200,
4+
"target_column": "churned",
5+
"features": [
6+
"age",
7+
"tenure_months",
8+
"monthly_spend",
9+
"support_tickets",
10+
"last_login_days",
11+
"nps_score",
12+
"region",
13+
"plan_type"
14+
],
15+
"performance": {
16+
"accuracy": 0.8065,
17+
"roc_auc": 0.5769,
18+
"cv_auc_mean": 0.5725,
19+
"cv_auc_std": 0.052
20+
},
21+
"training_rows": 247,
22+
"test_rows": 62,
23+
"trained_on": "customer_churn_demo.csv",
24+
"class_distribution": {
25+
"no_churn": 257,
26+
"churn": 52
27+
}
28+
}

sample_data/predict.py

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
"""
2+
Churn prediction inference script.
3+
Uses the pre-trained model in churn_model.joblib.
4+
5+
Usage:
6+
python predict.py # runs on customer_churn_demo.csv
7+
python predict.py --file path/to/data.csv # run on your own CSV
8+
python predict.py --single # predict one customer interactively
9+
10+
Requirements:
11+
pip install scikit-learn pandas joblib
12+
"""
13+
14+
import argparse
15+
import json
16+
import sys
17+
from pathlib import Path
18+
19+
import joblib
20+
import pandas as pd
21+
22+
MODEL_PATH = Path(__file__).parent / "churn_model.joblib"
23+
META_PATH = Path(__file__).parent / "churn_model_metadata.json"
24+
DATA_PATH = Path(__file__).parent / "customer_churn_demo.csv"
25+
26+
FEATURES = ["age", "tenure_months", "monthly_spend",
27+
"support_tickets", "last_login_days", "nps_score",
28+
"region", "plan_type"]
29+
30+
31+
def load_model():
32+
if not MODEL_PATH.exists():
33+
sys.exit(f"Model not found at {MODEL_PATH}")
34+
return joblib.load(MODEL_PATH)
35+
36+
37+
def predict_batch(model, csv_path: Path):
38+
df = pd.read_csv(csv_path)
39+
# Normalise mixed case (same as training)
40+
if "region" in df.columns:
41+
df["region"] = df["region"].str.title().str.strip()
42+
if "plan_type" in df.columns:
43+
df["plan_type"] = df["plan_type"].str.lower().str.strip()
44+
45+
X = df[FEATURES]
46+
probs = model.predict_proba(X)[:, 1]
47+
preds = model.predict(X)
48+
49+
df["churn_probability"] = probs.round(3)
50+
df["churn_prediction"] = ["yes" if p == 1 else "no" for p in preds]
51+
52+
# Summary
53+
n = len(df)
54+
flagged = (preds == 1).sum()
55+
print(f"\n{'─'*50}")
56+
print(f" Dataset : {csv_path.name}")
57+
print(f" Rows : {n}")
58+
print(f" At-risk : {flagged} customers ({flagged/n*100:.1f}%)")
59+
print(f"{'─'*50}")
60+
61+
# Top 10 highest risk
62+
top = df.nlargest(10, "churn_probability")[
63+
["customer_id", "churn_probability", "plan_type", "tenure_months", "nps_score"]
64+
] if "customer_id" in df.columns else df.nlargest(10, "churn_probability")[
65+
["churn_probability", "plan_type", "tenure_months", "nps_score"]
66+
]
67+
print("\nTop 10 highest-risk customers:")
68+
print(top.to_string(index=False))
69+
70+
out_path = csv_path.parent / (csv_path.stem + "_predictions.csv")
71+
df.to_csv(out_path, index=False)
72+
print(f"\nFull results saved → {out_path}\n")
73+
74+
75+
def predict_single(model):
76+
print("\nEnter customer details (press Enter to use median/default):\n")
77+
defaults = {
78+
"age": 35, "tenure_months": 18, "monthly_spend": 89.0,
79+
"support_tickets": 2, "last_login_days": 10, "nps_score": 6.5,
80+
"region": "North America", "plan_type": "pro",
81+
}
82+
data = {}
83+
for feat, default in defaults.items():
84+
val = input(f" {feat} [{default}]: ").strip()
85+
if val == "":
86+
data[feat] = default
87+
else:
88+
try:
89+
data[feat] = float(val) if "." in val or feat in ("monthly_spend", "nps_score") else int(val)
90+
except ValueError:
91+
data[feat] = val
92+
93+
X = pd.DataFrame([data])
94+
X["region"] = X["region"].str.title().str.strip()
95+
X["plan_type"] = X["plan_type"].str.lower().str.strip()
96+
97+
prob = model.predict_proba(X)[0, 1]
98+
pred = "WILL CHURN" if prob >= 0.5 else "will NOT churn"
99+
risk = "HIGH" if prob >= 0.6 else "MEDIUM" if prob >= 0.35 else "LOW"
100+
101+
print(f"\n{'─'*40}")
102+
print(f" Churn probability : {prob:.1%}")
103+
print(f" Prediction : {pred}")
104+
print(f" Risk level : {risk}")
105+
print(f"{'─'*40}\n")
106+
107+
108+
def print_metadata():
109+
if META_PATH.exists():
110+
meta = json.loads(META_PATH.read_text())
111+
perf = meta["performance"]
112+
print(f"\nModel: {meta['model_type']}")
113+
print(f"Trained on: {meta['trained_on']} ({meta['training_rows']} rows)")
114+
print(f"Accuracy: {perf['accuracy']:.1%} ROC-AUC: {perf['roc_auc']:.3f} CV-AUC: {perf['cv_auc_mean']:.3f} ± {perf['cv_auc_std']:.3f}")
115+
116+
117+
if __name__ == "__main__":
118+
parser = argparse.ArgumentParser(description="Datrix churn prediction demo")
119+
parser.add_argument("--file", type=Path, default=DATA_PATH)
120+
parser.add_argument("--single", action="store_true", help="Predict a single customer interactively")
121+
args = parser.parse_args()
122+
123+
print_metadata()
124+
model = load_model()
125+
126+
if args.single:
127+
predict_single(model)
128+
else:
129+
predict_batch(model, args.file)

0 commit comments

Comments
 (0)