Update app.py

app.py

@@ -7,13 +7,14 @@ import os
 import warnings
 import shap
 import matplotlib.pyplot as plt
+from sklearn.metrics import roc_curve, precision_recall_curve
 
 # Suppress XGBoost warnings
 warnings.filterwarnings("ignore", category=UserWarning, message=".*WARNING.*")
 
 # Load your model (automatically detect XGBoost or joblib model)
 def load_model():
-    model_path = "best_model.json"  # Ensure this matches your file name
+    model_path = "xgboost_model.json"  # Ensure this matches your file name
     if os.path.exists(model_path):  
         model = xgb.Booster()
         model.load_model(model_path)
@@ -25,6 +26,16 @@ def load_model():
 
 model = load_model()
 
+# Automatically find the best threshold using ROC
+def optimize_threshold(X_test, y_test):
+    dmatrix = xgb.DMatrix(X_test)
+    y_prob = model.predict(dmatrix)
+    
+    fpr, tpr, thresholds = roc_curve(y_test, y_prob)
+    optimal_idx = np.argmax(tpr - fpr)
+    optimal_threshold = thresholds[optimal_idx]
+    return optimal_threshold
+
 # Prediction function with dynamic threshold
 def predict_employee_status(satisfaction_level, last_evaluation, number_project,
                             average_monthly_hours, time_spent_company, 
@@ -70,22 +81,25 @@ def predict_employee_status(satisfaction_level, last_evaluation, number_project,
     except Exception as e:
         return f"❌ Error: {str(e)}"
 
-# SHAP Explainability
+# SHAP Explainability (Directly Integrated)
 def explain_prediction(input_df):
-    explainer = shap.TreeExplainer(model)
-    shap_values = explainer.shap_values(input_df)
-    
-    # Generating SHAP explanation for this prediction
-    shap.initjs()
-    plt.figure()
-    shap.waterfall_plot(shap.Explanation(values=shap_values[0], 
-                                         base_values=explainer.expected_value,
-                                         data=input_df.iloc[0].values,
-                                         feature_names=input_df.columns))
-    plt.savefig("shap_explanation.png")
-    return "SHAP explanation generated for this prediction."
+    try:
+        explainer = shap.TreeExplainer(model)
+        shap_values = explainer.shap_values(input_df)
+        
+        # Generate and save SHAP explanation as an image
+        shap.initjs()
+        plt.figure()
+        shap.waterfall_plot(shap.Explanation(values=shap_values[0], 
+                                             base_values=explainer.expected_value,
+                                             data=input_df.iloc[0].values,
+                                             feature_names=input_df.columns))
+        plt.savefig("shap_explanation.png")
+        return "SHAP explanation generated for this prediction."
+    except Exception as e:
+        return f"❌ Error in SHAP: {str(e)}"
 
-# Gradio interface with dynamic threshold
+# Gradio interface with dynamic threshold and SHAP
 def gradio_interface():
     interface = gr.Interface(
         fn=predict_employee_status,
@@ -106,7 +120,7 @@ def gradio_interface():
             gr.Slider(0.1, 0.9, value=0.5, step=0.05, label="Prediction Threshold")
         ],
         outputs="text",
-        title="Employee Retention Prediction System (With SHAP Explainability)",
+        title="Employee Retention Prediction System (With SHAP & ROC Threshold)",
         description="Predict whether an employee is likely to stay or quit based on their profile. Adjust the threshold for accurate predictions.",
         theme="dark"
     )