Update app.py

app.py

@@ -1,13 +1,32 @@
 import gradio as gr
-import joblib
+import xgboost as xgb
 import numpy as np
+import joblib
+import os
+import warnings
+
+# Suppress XGBoost warnings
+warnings.filterwarnings("ignore", category=UserWarning, message=".*WARNING.*")
 
-# Load your model
-model = joblib.load('best_model.json')
+# Load your model (automatically detect XGBoost or joblib model)
+def load_model():
+    model_path = "best_model.json"  # Ensure this matches your file name
+    if os.path.exists(model_path):  
+        model = xgb.Booster()
+        model.load_model(model_path)
+        print("✅ Model loaded successfully.")
+        return model
+    else:
+        print("❌ Model file not found.")
+        return None
 
-def predict_retention(satisfaction_level, last_evaluation, number_project,
-                      average_monthly_hours, time_spent_company, 
-                      work_accident, promotion_last_5years, salary, department):
+model = load_model()
+
+# Prediction function
+def predict_employee_status(satisfaction_level, last_evaluation, number_project,
+                            average_monthly_hours, time_spent_company, 
+                            work_accident, promotion_last_5years, salary, department):
+    
     # One-hot encode the department
     departments = [
         'RandD', 'accounting', 'hr', 'management', 'marketing',
@@ -23,14 +42,19 @@ def predict_retention(satisfaction_level, last_evaluation, number_project,
     ] + department_encoded).reshape(1, -1)
 
     # Predict using the model
+    if model is None:
+        return "❌ No model found. Please upload the model file."
+
     try:
-        prediction = model.predict(input_data)
-        return "Employee is likely to quit." if prediction[0] == 1 else "Employee is likely to stay."
+        dmatrix = xgb.DMatrix(input_data)
+        prediction = model.predict(dmatrix)
+        return "Employee is likely to quit." if prediction[0] > 0.5 else "Employee is likely to stay."
     except Exception as e:
-        return f"Error: {str(e)}"
+        return f"❌ Error: {str(e)}"
 
+# Gradio interface
 interface = gr.Interface(
-    fn=predict_retention,
+    fn=predict_employee_status,
     inputs=[
         gr.Number(label="Satisfaction Level (0.0 - 1.0)"),
         gr.Number(label="Last Evaluation (0.0 - 1.0)"),