Update app.py

app.py

@@ -677,58 +677,56 @@ def train_and_save(
 
 
     # ---- Train survival model (CoxPH) ----
+    from sklearn.impute import SimpleImputer
+
     survival_trained = False
     surv_notes = None
     
     if time_days is not None and event01 is not None:
         try:
-            # Build survival training frame from ORIGINAL df
             df_surv = df[feature_cols].copy().replace({pd.NA: np.nan})
     
-            # Coerce numeric/categorical exactly like your main pipeline convention
+            # coerce numeric/cat like you already do
             for c in num_cols:
                 if c in df_surv.columns:
                     df_surv[c] = pd.to_numeric(df_surv[c], errors="coerce")
-    
             for c in cat_cols:
                 if c in df_surv.columns:
                     df_surv[c] = df_surv[c].astype("object")
                     df_surv.loc[df_surv[c].isna(), c] = np.nan
                     df_surv[c] = df_surv[c].map(lambda v: v if pd.isna(v) else str(v))
     
-            # Add survival targets
             df_surv["time_days"] = time_days
             df_surv["event"] = event01
     
-            # Keep only rows with valid survival targets
             df_surv = df_surv.dropna(subset=["time_days", "event"])
     
-            # One-hot categoricals (drop_first=True matches your intention)
-            df_surv_oh = pd.get_dummies(df_surv, columns=[c for c in cat_cols if c in df_surv.columns], drop_first=True)
+            # one-hot
+            df_surv_oh = pd.get_dummies(df_surv, columns=cat_cols, drop_first=True)
     
             duration_col = "time_days"
             event_col = "event"
             X_cols = [c for c in df_surv_oh.columns if c not in (duration_col, event_col)]
-    
-            # Ensure numeric matrix for Cox
-            df_surv_oh[X_cols] = df_surv_oh[X_cols].apply(pd.to_numeric, errors="coerce")
 
-            # Fit an imputer ONLY for Cox predictors and store it in the bundle
+            # numeric coercion
+            df_surv_oh[X_cols] = df_surv_oh[X_cols].apply(pd.to_numeric, errors="coerce")
+    
+            # impute predictors ONLY
             imp = SimpleImputer(strategy="median")
-            df_surv_oh[X_cols] = imp.fit_transform(df_surv_oh[X_cols])
+            X_imp = imp.fit_transform(df_surv_oh[X_cols])
     
-            # Final sanity: remove any remaining bad rows (rare)
-            df_surv_oh = df_surv_oh.dropna(subset=[duration_col, event_col])
+            # safe assignment back
+            df_surv_oh.loc[:, X_cols] = pd.DataFrame(X_imp, columns=X_cols, index=df_surv_oh.index)
     
-            # Fit penalized Cox
+            # fit Cox
             cph = CoxPHFitter(penalizer=0.1)
-            cph.fit(df_surv_oh, duration_col=duration_col, event_col=event_col)
-    
-            surv_columns = X_cols  # predictors used in Cox
+            cph.fit(df_surv_oh[[duration_col, event_col] + X_cols],
+                    duration_col=duration_col,
+                    event_col=event_col)
     
             bundle = {
                 "model": cph,
-                "columns": surv_columns,
+                "columns": X_cols,
                 "imputer": imp,
                 "cat_cols": cat_cols,
                 "num_cols": num_cols,
@@ -741,18 +739,17 @@ def train_and_save(
     
             survival_trained = True
             surv_notes = "Survival model trained successfully."
-    
+
         except Exception as e:
             survival_trained = False
             surv_notes = f"Survival model training failed: {e}"
     else:
         surv_notes = "Survival columns missing or could not be parsed; survival model not trained."
 
-        
-
+    
 
 
-    
+          
     
     joblib.dump(pipe, "model.joblib",compress=3)
 
@@ -1586,6 +1583,7 @@ with tab_train:
         else: 
     
             df = pd.read_excel(train_file, engine="openpyxl")
+            df.columns = [c.strip() for c in df.columns]
             feature_cols = get_feature_cols_from_df(df)
         
             st.dataframe(df.head(), use_container_width=True)