Spaces:
Running
Running
File size: 5,438 Bytes
dc3d345 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | """Train, evaluate, and persist the CKD screening model.
Trains three comparable classifiers (Logistic Regression, Random Forest,
Decision Tree), reports held-out and cross-validated metrics, and serializes the
Random Forest pipeline plus a metrics manifest for the API to load.
Run: python -m nephroscreen.train
"""
import json
import warnings
from datetime import datetime, timezone
import joblib
import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
accuracy_score,
f1_score,
precision_recall_curve,
precision_score,
recall_score,
roc_auc_score,
)
from sklearn.model_selection import StratifiedKFold, cross_val_score, train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.tree import DecisionTreeClassifier
from .config import (
DATA_PATH,
METRICS_PATH,
MODEL_COLUMNS,
MODEL_PATH,
MODELS_DIR,
RANDOM_STATE,
TEST_SIZE,
)
from .preprocessing import CKDPreprocessor, load_dataset
warnings.filterwarnings("ignore")
def build_pipeline(estimator, scale_numeric=False):
steps = [("preprocess", CKDPreprocessor())]
if scale_numeric:
steps.append(("scaler", StandardScaler()))
steps.append(("model", estimator))
return Pipeline(steps)
def score_model(model, X_test, y_test, X_train, y_train, cv):
pred = model.predict(X_test)
prob = model.predict_proba(X_test)[:, 1]
cv_scores = cross_val_score(model, X_train, y_train, cv=cv, scoring="accuracy")
return {
"accuracy": round(accuracy_score(y_test, pred) * 100, 2),
"precision": round(precision_score(y_test, pred) * 100, 2),
"recall": round(recall_score(y_test, pred) * 100, 2),
"f1": round(f1_score(y_test, pred) * 100, 2),
"roc_auc": round(roc_auc_score(y_test, prob), 4),
"cv_mean": round(cv_scores.mean() * 100, 2),
"cv_std": round(cv_scores.std() * 100, 2),
}
def recall_first_threshold(y_test, prob, min_recall=0.99):
"""Lowest decision threshold that still achieves `min_recall` on the test set.
For a screening tool a missed CKD case (false negative) is the costly error,
so we prefer a threshold that keeps recall very high rather than the default 0.5.
"""
precision, recall, thresholds = precision_recall_curve(y_test, prob)
# thresholds has len-1 vs precision/recall; align by dropping the last point.
best_t = 0.5
for p, r, t in zip(precision[:-1], recall[:-1], thresholds):
if r >= min_recall:
best_t = float(t)
return round(best_t, 3)
def main():
MODELS_DIR.mkdir(parents=True, exist_ok=True)
print(f"Loading dataset from {DATA_PATH} ...")
df = load_dataset(DATA_PATH)
print(f"Shape: {df.shape} | CKD: {int(df.target.sum())} | Healthy: {int((df.target == 0).sum())}")
X = df.drop(columns=["target"])
y = df["target"]
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=TEST_SIZE, random_state=RANDOM_STATE, stratify=y
)
cv = StratifiedKFold(n_splits=10, shuffle=True, random_state=RANDOM_STATE)
specs = {
"Logistic Regression": (
build_pipeline(LogisticRegression(C=1.0, max_iter=1000, random_state=RANDOM_STATE), True),
),
"Random Forest": (
build_pipeline(RandomForestClassifier(n_estimators=100, random_state=RANDOM_STATE, n_jobs=1)),
),
"Decision Tree": (
build_pipeline(DecisionTreeClassifier(max_depth=5, min_samples_split=5, random_state=RANDOM_STATE)),
),
}
results = {}
fitted = {}
for name, (pipe,) in specs.items():
pipe.fit(X_train, y_train)
fitted[name] = pipe
results[name] = score_model(pipe, X_test, y_test, X_train, y_train, cv)
print(f" {name:<22} acc={results[name]['accuracy']}% auc={results[name]['roc_auc']}")
# Random Forest is the serving model (best held-out performance).
serving = fitted["Random Forest"]
rf_prob = serving.predict_proba(X_test)[:, 1]
threshold = recall_first_threshold(y_test, rf_prob, min_recall=0.99)
importances = (
pd.Series(serving.named_steps["model"].feature_importances_, index=MODEL_COLUMNS)
.sort_values(ascending=False)
.round(4)
)
joblib.dump(serving, MODEL_PATH)
print(f"\nSaved serving model -> {MODEL_PATH}")
manifest = {
"model": "RandomForestClassifier(n_estimators=100)",
"trained_at": datetime.now(timezone.utc).isoformat(),
"random_state": RANDOM_STATE,
"test_size": TEST_SIZE,
"n_train": int(len(X_train)),
"n_test": int(len(X_test)),
"serving_threshold": threshold,
"metrics": results,
"top_features": importances.head(10).to_dict(),
"note": (
"Near-perfect scores reflect the high separability of this small UCI "
"dataset (400 patients, 80-patient test set) — not clinical deployment "
"readiness. See README for the honest interpretation."
),
}
with open(METRICS_PATH, "w", encoding="utf-8") as f:
json.dump(manifest, f, indent=2)
print(f"Saved metrics manifest -> {METRICS_PATH}")
print(f"Recall-first serving threshold (recall>=0.99): {threshold}")
if __name__ == "__main__":
main()
|