Update app.py

app.py

@@ -1,4 +1,4 @@
-import pickle 
+import xgboost as xgb
 import pandas as pd
 import shap
 from shap.plots._force_matplotlib import draw_additive_plot
@@ -6,9 +6,9 @@ import gradio as gr
 import numpy as np
 import matplotlib.pyplot as plt
 
-# Load the XGBoost model
-with open("h22_xgb_Final.json", "rb") as f:
-    loaded_model = pickle.load(f)
+# Load the XGBoost model from JSON
+loaded_model = xgb.Booster()
+loaded_model.load_model("h22_xgb_Final.json")
 
 # Setup SHAP Explainer
 explainer = shap.Explainer(loaded_model)  # PLEASE DO NOT CHANGE THIS.
@@ -19,19 +19,24 @@ def main_func(SupportiveGM, Merit, LearningDevelopment, WorkEnvironment, Engagem
               'LearningDevelopment': LearningDevelopment, 'WorkEnvironment': WorkEnvironment, 
               'Engagement': Engagement, 'WellBeing': WellBeing}, orient='index').transpose()
     
-    # Predict probabilities
-    prob = loaded_model.predict_proba(new_row)
+    # Convert input to DMatrix format (needed for XGBoost Booster)
+    dmatrix_new = xgb.DMatrix(new_row)
+
+    # Predict probabilities (assuming binary classification)
+    prob = loaded_model.predict(dmatrix_new)  # Returns only probabilities for class 1
+    
+    # Ensure proper formatting for SHAP
+    shap_values = explainer(dmatrix_new)
     
-    # Compute SHAP values
-    shap_values = explainer(new_row)
+    # Generate SHAP Plot
     plot = shap.plots.bar(shap_values[0], max_display=6, order=shap.Explanation.abs, show_data='auto', show=False)
-
+    
     plt.tight_layout()
     local_plot = plt.gcf()
     plt.rcParams['figure.figsize'] = (6, 4)
     plt.close()
     
-    return {"Leave": float(prob[0][0]), "Stay": 1 - float(prob[0][0])}, local_plot
+    return {"Leave": 1 - float(prob[0]), "Stay": float(prob[0])}, local_plot
 
 # Create the UI
 title = "**Mod 3 Team 5: Employee Turnover Predictor**"
@@ -79,5 +84,3 @@ with gr.Blocks(title=title) as demo:
                 [label, local_plot], main_func, cache_examples=True)
 
 demo.launch()
-
-