fix tgt col based on model selected

src/inference_utils.py

@@ -16,10 +16,11 @@ def compute_metrics(y_true, y_pred_proba, threshold=0.5):
 
 def add_predictions(df, probs):
     df['Predicted Probability'] = probs
-    df['GVHD Prediction'] = ['POSITIVE' if p > 0.5 else 'NEGATIVE' for p in probs]
+    pred_col = f"{st.session_state.target_col} Prediction"
+    df[pred_col] = ['POSITIVE' if p > 0.5 else 'NEGATIVE' for p in probs]
+
+    df_with_gt = df[['Predicted Probability', pred_col]].join(st.session_state.targets_df)
 
-    df_with_gt = df[['Predicted Probability', 'GVHD Prediction']].join(st.session_state.targets_df)
-    
     # Define cell-level styling
     def highlight_prediction(val):
         if val == "POSITIVE":
@@ -31,7 +32,7 @@ def add_predictions(df, probs):
     # Apply color and alignment
     df_styled = (
         df_with_gt.style
-        .applymap(highlight_prediction, subset=["GVHD Prediction"])
+        .applymap(highlight_prediction, subset=[pred_col])
         .set_properties(**{'text-align': 'center'})  # Apply center alignment to all cells
     )
 

src/pages/1_Individual_Predictions.py

@@ -190,14 +190,17 @@ if submitted:
         if "ensemble" in st.session_state.selected_model:
             # ensemble prediction
             models = load_model_ensemble(st.session_state.selected_model)
+            st.session_state.target_col = models.get("target_col", "UNKNOWN")
             models = models["model"]
             pred = ensemble_predict(models, X, cat_features)
         else:
             # single model prediction
             model = load_model(st.session_state.selected_model)
+            st.session_state.target_col = model.get("target_col", "UNKNOWN")
             model = model["model"]
             pred = model.predict_proba(X)[0][1]
 
+    st.warning(f"The model selected will only predict the target \"{st.session_state.target_col}\". Please choose a different model if you want to predict a different target.")
     result_df = pd.DataFrame()
     result_df = add_predictions(result_df, [pred])
 

src/pages/2_Bulk_Predictions.py

@@ -9,6 +9,29 @@ from sidebar import sidebar
 # Initialize sidebar
 sidebar()
 
+if "selected_model" in st.session_state:
+    if "ensemble" in st.session_state.selected_model:
+        # ensemble model
+        ensemble = True
+        try:
+            ensemble_data = load_model_ensemble(st.session_state.selected_model)
+            st.session_state.trained_models = ensemble_data["model"]
+            models = ensemble_data["model"]
+            st.session_state.best_iterations = ensemble_data.get("best_iterations", [])
+            st.session_state.fold_scores = ensemble_data.get("fold_scores", [])
+            target_col = ensemble_data.get("target_col", "UNKNOWN")
+
+        except Exception as e:
+            st.error(f"Error loading ensemble: {str(e)}")
+    else:
+        # single model
+        ensemble = False
+        model_dict = load_model(st.session_state.selected_model)
+        model = model_dict["model"]
+        target_col = model_dict.get("target_col", "UNKNOWN")
+
+    st.warning(f"The model selected will only predict the target \"{target_col}\". Please choose a different model if you want to predict a different target.")
+
 st.title("📊 Bulk Patient Predictions")
 
 uploaded_file = st.file_uploader("Upload CSV", type=["csv"])
@@ -27,41 +50,25 @@ if uploaded_file:
 
     # TODO: Define the target column (customize this based on your use case)
     # target_col = st.session_state.target_col # "GVHD"  # or "Acute GVHD(<100 days)", etc.
-    st.session_state.target_col = st.selectbox(
-        "Select target column to predict:",
-        options=[
-            "GVHD",
-            "Acute GVHD(<100 days)",
-            "Chronic GVHD>100 days",
-        ],
-        index=0
-    )
+
+    # # allow dynamic selection for target column
+    # st.session_state.target_col = st.selectbox(
+    #     "Select target column to predict:",
+    #     options=[
+    #         "GVHD",
+    #         "Acute GVHD(<100 days)",
+    #         "Chronic GVHD>100 days",
+    #     ],
+    #     index=0
+    # )
     
     if st.button("Predict"):
         if "bulk_input_df" not in st.session_state:
             st.warning("Please preprocess data first.")
         else:
-            if "ensemble" in st.session_state.selected_model:
-                # ensemble model
-                ensemble = True
-                try:
-                    ensemble_data = load_model_ensemble(st.session_state.selected_model)
-                    st.session_state.trained_models = ensemble_data["model"]
-                    models = ensemble_data["model"]
-                    st.session_state.best_iterations = ensemble_data.get("best_iterations", [])
-                    st.session_state.fold_scores = ensemble_data.get("fold_scores", [])
-
-                except Exception as e:
-                    st.error(f"Error loading ensemble: {str(e)}")
-            else:
-                # single model
-                ensemble = False
-                model_dict = load_model(st.session_state.selected_model)
-                model = model_dict["model"]
-
             df = st.session_state.bulk_input_df
 
-            target_col = st.session_state.target_col
+            # target_col = st.session_state.target_col
             
             # Optional filtering depending on target choice
             if target_col in ["Acute GVHD(<100 days)", "Chronic GVHD>100 days"]: