Update app.py

app.py

@@ -67,23 +67,12 @@ def main_func(Department, ChainScale, SupportiveGM, Merit, LearningDevelopment,
     shap_values = explainer(new_row)
 
     fig, ax = plt.subplots(figsize=(8, 4))
-    values = shap_values.values[0]  # Extract SHAP values
-    features = new_row.columns  # Feature names
-
-    colors = ['#FF0000' if v < 0 else '#1E4380' for v in values]
-    sorted_indices = np.argsort(np.abs(values))[-6:]  # Select top 6 features
-    sorted_values = values[sorted_indices]
-    sorted_features = [features[i] for i in sorted_indices]
-    sorted_colors = [colors[i] for i in sorted_indices]
-    ax.barh(sorted_features, sorted_values, color=sorted_colors)
-    ax.set_xlabel("SHAP Value Impact")
-    ax.set_title("Feature Importance (SHAP)")
-
+    shap.waterfall_plot(shap_values[0])  # Corrected to use waterfall plot
     plt.tight_layout()
     local_plot = plt.gcf()
     plt.close()
 
-    return {"Leave": float(prob[0][0]), "Stay": 1 - float(prob[0][0])}, local_plot
+    return {"Leave": float(prob[0][0]), "Stay": float(prob[0][1])}, local_plot
 
 # Create the UI
 title = "**Mod 3 Team 5: Employee Turnover Predictor & Interpreter**"
@@ -141,8 +130,8 @@ with gr.Blocks(title=title) as demo:
             submit_btn = gr.Button("Analyze")
         
         with gr.Column(visible=True, scale=1, min_width=600) as output_col:
-            label = gr.Label(label="Predicted Label")
-            local_plot = gr.Plot(label='SHAP Analysis')
+            label = gr.Label(label="Predicted Intent to Stay vs Leave")
+            local_plot = gr.Plot(label='SHAP Waterfall Analysis')
 
     submit_btn.click(
         main_func,
@@ -166,3 +155,4 @@ with gr.Blocks(title=title) as demo:
     )
 
 demo.launch()
+