Update app.py

app.py

@@ -63,16 +63,26 @@ def main_func(Department, ChainScale, SupportiveGM, Merit, LearningDevelopment,
     }).astype(float)
 
     prob = loaded_model.predict_proba(new_row) 
+    
+    # Ensure probabilities return correctly
+    if prob.shape[1] == 2:
+        leave_prob = float(prob[0][0])
+        stay_prob = float(prob[0][1])
+    else:
+        leave_prob = float(prob[0])
+        stay_prob = 1 - leave_prob
 
     shap_values = explainer(new_row)
 
     fig, ax = plt.subplots(figsize=(8, 4))
-    shap.waterfall_plot(shap_values[0])  # Corrected to use waterfall plot
+    shap.waterfall_plot(shap.Explanation(values=shap_values.values[0],
+                                         base_values=shap_values.base_values[0],
+                                         data=new_row.iloc[0]))  # Fix waterfall plot
     plt.tight_layout()
     local_plot = plt.gcf()
     plt.close()
 
-    return {"Leave": float(prob[0][0]), "Stay": float(prob[0][1])}, local_plot
+    return {"Leave": leave_prob, "Stay": stay_prob}, local_plot
 
 # Create the UI
 title = "**Mod 3 Team 5: Employee Turnover Predictor & Interpreter**"
@@ -154,5 +164,4 @@ with gr.Blocks(title=title) as demo:
         cache_examples=True
     )
 
-demo.launch()
-
+demo.launch()