Update app.py

app.py

@@ -17,7 +17,8 @@ from sklearn.metrics import (
     precision_recall_curve, average_precision_score,
     brier_score_loss
 )
-
+import shap
+import matplotlib.pyplot as plt
 
 
 from sklearn.calibration import calibration_curve
@@ -39,7 +40,7 @@ from sklearn.model_selection import train_test_split
 
 # REPLACE make_fig with this (or add this and stop using plt.plot directly)
 def make_fig(figsize=(5.5, 3.6), dpi=120):
-    import matplotlib.pyplot as plt
+    
     fig, ax = plt.subplots(figsize=figsize, dpi=dpi)
     return fig, ax
 
@@ -63,7 +64,7 @@ def render_plot_with_download(
     export_dpi: int = 600,
     key: Optional[str] = None
 ):
-    import matplotlib.pyplot as plt  # lazy
+    
     png_bytes = fig_to_png_bytes(fig, dpi=export_dpi)
     st.pyplot(fig, clear_figure=False)
     st.download_button(
@@ -234,6 +235,9 @@ def train_survival_bundle(
 
     # one-hot
     df_surv_oh = pd.get_dummies(df_surv, columns=cat_cols, drop_first=True)
+    if duration_col not in df_surv_oh.columns or event_col not in df_surv_oh.columns:
+        raise ValueError("Survival DF missing duration/event columns after one-hot encoding.")
+
 
     # remove duplicate columns if any messy headers caused duplicates
     df_surv_oh = df_surv_oh.loc[:, ~df_surv_oh.columns.duplicated()].copy()
@@ -663,7 +667,7 @@ def train_and_save(
     l1_C: float,
     use_dimred: bool,
     svd_components: int,):
-    from lifelines import CoxPHFitter
+    
     X = df[feature_cols].copy()
     y_raw = df[LABEL_COL].copy()
 
@@ -863,7 +867,7 @@ def train_and_save(
 # SHAP
 # ============================================================
 def build_shap_explainer(pipe, X_bg, max_bg=200):
-    import shap  # lazy
+    
 
     if X_bg is None or len(X_bg) == 0:
         raise ValueError("SHAP background is empty.")
@@ -882,7 +886,7 @@ def build_shap_explainer(pipe, X_bg, max_bg=200):
     return explainer
 
 
-def safe_dense(Xt, max_rows: int = 2000):
+def safe_dense(Xt, max_rows: int = 200):
     """
     Convert sparse->dense carefully. Avoid converting huge matrices to dense.
     """
@@ -2612,6 +2616,9 @@ with tab_predict:
             try:
                 cph = bundle["model"]
                 surv_cols = bundle.get("columns", [])
+
+
+
         
                 # Build Cox input row (same preprocessing as Cox training)
                 df_one_surv = X_one[feature_cols].copy()
@@ -2627,12 +2634,21 @@ with tab_predict:
                         )
         
                 df_one_surv_oh = pd.get_dummies(df_one_surv, columns=cat_cols, drop_first=True)
+
+                
         
                 # Align to training predictor columns
                 for col in surv_cols:
                     if col not in df_one_surv_oh.columns:
                         df_one_surv_oh[col] = 0
                 df_one_surv_oh = df_one_surv_oh[surv_cols]
+
+                imp = bundle.get("imputer", None)
+                    if imp is not None:
+                        X_imp = imp.transform(df_one_surv_oh)
+                        df_one_surv_oh = pd.DataFrame(X_imp, columns=surv_cols, index=df_one_surv_oh.index)
+                    else:
+                        df_one_surv_oh = df_one_surv_oh.fillna(0)
         
                 # Predict survival function
                 surv_fn = cph.predict_survival_function(df_one_surv_oh)
@@ -2994,6 +3010,17 @@ with tab_predict:
                 imp = bundle.get("imputer", None)
         
                 df_surv_in = X_inf[feature_cols].copy()
+                for c in num_cols:
+                    if c in df_surv_in.columns:
+                        df_surv_in[c] = pd.to_numeric(df_surv_in[c], errors="coerce")
+                
+                for c in cat_cols:
+                    if c in df_surv_in.columns:
+                        df_surv_in[c] = df_surv_in[c].astype("object")
+                        df_surv_in.loc[df_surv_in[c].isna(), c] = np.nan
+                        df_surv_in[c] = df_surv_in[c].map(lambda v: v if pd.isna(v) else str(v))
+
+                
                 df_surv_in_oh = pd.get_dummies(df_surv_in, columns=cat_cols, drop_first=True)
         
                 # align columns