Update app.py

app.py

@@ -2930,45 +2930,47 @@ with tab_predict:
         # ---- Batch survival probabilities (if survival model loaded) ----
         # ---- Batch survival probabilities (if survival model loaded) ----
         bundle = st.session_state.get("surv_model", None)
+
         if isinstance(bundle, dict) and bundle.get("model") is not None:
-            cph = bundle["model"]
-            surv_cols = bundle["columns"]
-            imp = bundle.get("imputer", None)
+            try:
+                cph = bundle["model"]
+                surv_cols = bundle["columns"]
+                imp = bundle.get("imputer", None)
         
-            df_surv_in = X_inf[feature_cols].copy()
-            # coerce num/cat same as you already do
-            df_surv_in_oh = pd.get_dummies(df_surv_in, columns=cat_cols, drop_first=True)
+                df_surv_in = X_inf[feature_cols].copy()
+                df_surv_in_oh = pd.get_dummies(df_surv_in, columns=cat_cols, drop_first=True)
         
-            # align columns exactly
-            for col in surv_cols:
-                if col not in df_surv_in_oh.columns:
-                    df_surv_in_oh[col] = 0
-            df_surv_in_oh = df_surv_in_oh[surv_cols]
+                # align columns
+                for col in surv_cols:
+                    if col not in df_surv_in_oh.columns:
+                        df_surv_in_oh[col] = 0
+                df_surv_in_oh = df_surv_in_oh[surv_cols]
+        
+                # impute
+                if imp is not None:
+                    X_imp = imp.transform(df_surv_in_oh)
+                    df_surv_in_oh = pd.DataFrame(X_imp, columns=surv_cols, index=df_surv_in_oh.index)
+                else:
+                    df_surv_in_oh = df_surv_in_oh.fillna(0)
         
-            # impute with training imputer
-            if imp is not None:
-                X_imp = imp.transform(df_surv_in_oh)
-                df_surv_in_oh = pd.DataFrame(X_imp, columns=surv_cols, index=df_surv_in_oh.index)
-            else:
-                df_surv_in_oh = df_surv_in_oh.fillna(0)
+                surv_fn_all = cph.predict_survival_function(df_surv_in_oh)
         
-            surv_fn_all = cph.predict_survival_function(df_surv_in_oh)
-
+                def surv_vec_at(days: int):
+                    idx = surv_fn_all.index.values
+                    j = int(np.argmin(np.abs(idx - days)))
+                    return surv_fn_all.iloc[j, :].values.astype(float)
         
-            def surv_vec_at(days: int):
-                idx = surv_fn_all.index.values
-                j = int(np.argmin(np.abs(idx - days)))
-                return surv_fn_all.iloc[j, :].values.astype(float)
-    
-            df_out["survival_6m"] = surv_vec_at(180)
-            df_out["survival_1y"] = surv_vec_at(365)
-            df_out["survival_2y"] = surv_vec_at(730)
-            df_out["survival_3y"] = surv_vec_at(1095)
-    
-        except Exception as e:
-            st.warning(f"Batch survival probabilities could not be computed: {e}")
-    else:
-        st.info("Survival model not loaded/published yet (survival bundle missing).")
+                df_out["survival_6m"] = surv_vec_at(180)
+                df_out["survival_1y"] = surv_vec_at(365)
+                df_out["survival_2y"] = surv_vec_at(730)
+                df_out["survival_3y"] = surv_vec_at(1095)
+        
+            except Exception as e:
+                st.warning(f"Batch survival probabilities could not be computed: {e}")
+        
+        else:
+            st.info("Survival model not loaded/published yet (survival bundle missing).")
+