""" risk_model.py — CognitivePulse Trains and serves an XGBoost classifier on the Alzheimer's dataset to produce: - A calibrated risk score (0-100) for any input patient profile - SHAP values for per-feature contribution to that score - Global feature importance for the population-level dashboard SHAP note: shap.TreeExplainer receives model.get_booster() (the raw Booster object) rather than the sklearn XGBClassifier wrapper. This avoids a bug in XGBoost 2.x where base_score is serialised as scientific notation (e.g. '[5.003238E-1]') which SHAP's string-to-float parser cannot parse, raising: ValueError: could not convert string to float: '[5.003238E-1]' On GPU (Colab T4): pass device="cuda" to train_model() to use XGBoost's GPU-accelerated histogram method. """ from __future__ import annotations import pickle from pathlib import Path import numpy as np import pandas as pd import xgboost as xgb import shap from sklearn.model_selection import StratifiedKFold, cross_val_score from data_loader import FEATURE_COLS, TARGET_COL, FEATURE_META MODEL_PATH = Path(__file__).parent / "data" / "model.json" EXPLAINER_PATH = Path(__file__).parent / "data" / "shap_explainer.pkl" XGB_PARAMS = { "n_estimators": 300, "max_depth": 5, "learning_rate": 0.05, "subsample": 0.8, "colsample_bytree": 0.8, "min_child_weight": 3, "gamma": 0.1, "reg_alpha": 0.1, "reg_lambda": 1.0, "scale_pos_weight": 1.83, "base_score": 0.5, "eval_metric": "auc", "random_state": 42, } def train_model(df: pd.DataFrame, device: str = "cpu") -> tuple: """ Trains XGBoost + SHAP TreeExplainer. Returns (model, explainer, cv_results_dict). """ X = df[FEATURE_COLS].values.astype("float32") y = df[TARGET_COL].values params = XGB_PARAMS.copy() if device == "cuda": # XGBoost 1.7.x uses gpu_hist + gpu_id; 2.x uses device + hist import xgboost as _xgb xgb_major = int(_xgb.__version__.split(".")[0]) if xgb_major >= 2: params["device"] = "cuda" params["tree_method"] = "hist" else: params["tree_method"] = "gpu_hist" params["gpu_id"] = 0 model = xgb.XGBClassifier(**params) # 5-fold stratified cross-validation cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42) cv_auc = cross_val_score(model, X, y, cv=cv, scoring="roc_auc") cv_f1 = cross_val_score(model, X, y, cv=cv, scoring="f1") # Final fit on full dataset model.fit(X, y, verbose=False) # ------------------------------------------------------------------ # # FIX: pass the raw Booster to TreeExplainer, NOT the sklearn wrapper. # XGBoost 2.x serialises base_score as '[5.003238E-1]' inside the # sklearn wrapper's JSON, which SHAP cannot parse. The raw Booster # serialises the same value as a plain float and parses correctly. # ------------------------------------------------------------------ # explainer = shap.TreeExplainer(model.get_booster()) cv_results = { "auc_mean": round(float(cv_auc.mean()), 4), "auc_std": round(float(cv_auc.std()), 4), "f1_mean": round(float(cv_f1.mean()), 4), "f1_std": round(float(cv_f1.std()), 4), } # Cache to disk — use pickle for the sklearn wrapper (avoids XGBoost 1.7.x # _estimator_type bug in save_model()) and separately cache the booster JSON # so the explainer can reload without re-training. MODEL_PATH.parent.mkdir(exist_ok=True) with open(MODEL_PATH.with_suffix(".pkl"), "wb") as f: pickle.dump(model, f) with open(EXPLAINER_PATH, "wb") as f: pickle.dump(explainer, f) return model, explainer, cv_results def load_or_train(df: pd.DataFrame, force_retrain: bool = False) -> tuple: """Loads cached model/explainer if available; trains from scratch otherwise.""" model_pkl = MODEL_PATH.with_suffix(".pkl") if not force_retrain and model_pkl.exists() and EXPLAINER_PATH.exists(): with open(model_pkl, "rb") as f: model = pickle.load(f) with open(EXPLAINER_PATH, "rb") as f: explainer = pickle.load(f) return model, explainer, None return train_model(df) def predict_patient(model, explainer, patient: dict) -> dict: """ Returns risk_score (0-100), risk_band, and per-feature SHAP contributions for a single patient supplied as a feature-name -> value dict. """ row = pd.DataFrame([{f: patient.get(f, 0) for f in FEATURE_COLS}]) X = row[FEATURE_COLS].values.astype("float32") prob = float(model.predict_proba(X)[0, 1]) risk_score = round(prob * 100, 1) # float32 required when explainer wraps a raw Booster shap_vals = explainer.shap_values(X)[0] contributions = { FEATURE_COLS[i]: round(float(shap_vals[i]), 4) for i in range(len(FEATURE_COLS)) } return { "risk_score": risk_score, "risk_probability": round(prob, 4), "risk_band": _risk_band(prob), "shap_contributions": contributions, "top_drivers": _top_drivers(contributions), } def _risk_band(prob: float) -> str: if prob < 0.25: return "low" if prob < 0.50: return "moderate" if prob < 0.75: return "elevated" return "high" def _top_drivers(contributions: dict, n: int = 5) -> list: sorted_feats = sorted(contributions.items(), key=lambda x: abs(x[1]), reverse=True) return [ { "feature": feat, "label": FEATURE_META.get(feat, {}).get("label", feat), "shap_value": val, "direction": "increases risk" if val > 0 else "decreases risk", "modifiable": FEATURE_META.get(feat, {}).get("modifiable", False), } for feat, val in sorted_feats[:n] ] def global_feature_importance(model) -> pd.DataFrame: importance = model.get_booster().get_score(importance_type="gain") return pd.DataFrame( [(FEATURE_META.get(k, {}).get("label", k), round(v, 2)) for k, v in importance.items()], columns=["Feature", "Importance"], ).sort_values("Importance", ascending=False).reset_index(drop=True) if __name__ == "__main__": from data_loader import load_dataset df, source = load_dataset() print(f"Training on {source} data ({len(df)} rows)...") model, explainer, cv = train_model(df) print(f"CV AUC : {cv['auc_mean']:.4f} ± {cv['auc_std']:.4f}") print(f"CV F1 : {cv['f1_mean']:.4f} ± {cv['f1_std']:.4f}") sample = { "Age": 68, "Gender": 0, "Ethnicity": 0, "EducationLevel": 2, "BMI": 29.5, "Smoking": 0, "AlcoholConsumption": 5.0, "PhysicalActivity": 2.5, "DietQuality": 5.0, "SleepQuality": 6.0, "FamilyHistoryAlzheimers": 1, "CardiovascularDisease": 1, "Diabetes": 0, "Depression": 0, "HeadInjury": 0, "Hypertension": 1, "SystolicBP": 148, "DiastolicBP": 88, "CholesterolTotal": 240, "CholesterolLDL": 158, "CholesterolHDL": 45, "CholesterolTriglycerides": 185, "MMSE": 25, "FunctionalAssessment": 7.0, "MemoryComplaints": 1, "BehavioralProblems": 0, "ADL": 7.5, "Confusion": 0, "Disorientation": 0, "PersonalityChanges": 0, "DifficultyCompletingTasks": 0, "Forgetfulness": 1, } import json print(json.dumps(predict_patient(model, explainer, sample), indent=2))