Update app.py

app.py

@@ -1,3 +1,4 @@
+import pickle
 import xgboost as xgb
 import pandas as pd
 import shap
@@ -6,36 +7,43 @@ import gradio as gr
 import numpy as np
 import matplotlib.pyplot as plt
 
-# Load the XGBoost model from JSON
-loaded_model = xgb.Booster()
-loaded_model.load_model("h22_xgb_Final.json")
+# Load the XGBoost model from Pickle
+with open("h22_xgb_Final.pkl", "rb") as f:
+    loaded_model = pickle.load(f)
 
-# Setup SHAP Explainer
-explainer = shap.Explainer(loaded_model)  # PLEASE DO NOT CHANGE THIS.
+# Ensure model is a Booster (handles both XGBClassifier & Booster cases)
+if isinstance(loaded_model, xgb.XGBClassifier):
+    loaded_model = loaded_model.get_booster()
+
+# Setup SHAP Explainer for XGBoost
+explainer = shap.TreeExplainer(loaded_model)  # Use TreeExplainer for XGBoost models
 
 # Define the prediction function
 def main_func(SupportiveGM, Merit, LearningDevelopment, WorkEnvironment, Engagement, WellBeing):
-    new_row = pd.DataFrame.from_dict({'SupportiveGM': SupportiveGM, 'Merit': Merit,
-              'LearningDevelopment': LearningDevelopment, 'WorkEnvironment': WorkEnvironment, 
-              'Engagement': Engagement, 'WellBeing': WellBeing}, orient='index').transpose()
+    new_row = pd.DataFrame.from_dict({
+        'SupportiveGM': SupportiveGM, 
+        'Merit': Merit,
+        'LearningDevelopment': LearningDevelopment, 
+        'WorkEnvironment': WorkEnvironment, 
+        'Engagement': Engagement, 
+        'WellBeing': WellBeing
+    }, orient='index').transpose()
     
-    # Convert input to DMatrix format (needed for XGBoost Booster)
+    # Convert new_row to DMatrix 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)
-    
-    # Generate SHAP Plot
-    plot = shap.plots.bar(shap_values[0], max_display=6, order=shap.Explanation.abs, show_data='auto', show=False)
-    
+    # Predict probability for staying (XGBoost Booster returns only one class probability)
+    prob = loaded_model.predict(dmatrix_new)
+
+    # Compute SHAP values
+    shap_values = explainer.shap_values(new_row)
+    plot = shap.plots.bar(shap_values, 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": 1 - float(prob[0]), "Stay": float(prob[0])}, local_plot
 
 # Create the UI
@@ -84,3 +92,4 @@ with gr.Blocks(title=title) as demo:
                 [label, local_plot], main_func, cache_examples=True)
 
 demo.launch()
+