Update app.py

app.py

@@ -5,8 +5,6 @@ import joblib
 import os
 import warnings
 from huggingface_hub import hf_hub_download
-import xgboost  
-
 
 # Suppress warnings
 warnings.filterwarnings("ignore")
@@ -23,20 +21,27 @@ def load_model():
 
 model = load_model()
 
-# Prediction function
-def predict_employee_status(satisfaction_level, last_evaluation, number_project,
-                            average_monthly_hours, time_spend_company,
-                            work_accident, promotion_last_5years, salary, department, threshold=0.5):
-
-    departments = ['RandD', 'accounting', 'hr', 'management', 'marketing',
-                   'product_mng', 'sales', 'support', 'technical']
+# Define prediction function
+def predict_employee_status(
+    satisfaction_level, last_evaluation, number_project,
+    average_monthly_hours, time_spend_company,
+    work_accident, promotion_last_5years, salary, department, threshold=0.5
+):
+    departments = [
+        'sales', 'accounting', 'hr', 'technical', 'support',
+        'management', 'IT', 'product_mng', 'marketing', 'RandD'
+    ]
+    
+    # One-hot encode department (include department_IT explicitly)
     department_features = {f"department_{dept}": 0 for dept in departments}
     if department in departments:
         department_features[f"department_{department}"] = 1
 
+    # Interaction features
     satisfaction_evaluation = satisfaction_level * last_evaluation
     work_balance = average_monthly_hours / number_project
 
+    # Construct input DataFrame
     input_data = {
         "satisfaction_level": [satisfaction_level],
         "last_evaluation": [last_evaluation],
@@ -52,13 +57,17 @@ def predict_employee_status(satisfaction_level, last_evaluation, number_project,
     }
 
     input_df = pd.DataFrame(input_data)
-    prediction_prob = model.predict_proba(input_df)[0][1]
-    result = "✅ Employee is likely to quit." if prediction_prob >= threshold else "✅ Employee is likely to stay."
-    return f"{result} (Probability: {prediction_prob:.2%})"
 
-# Launch Gradio UI
+    try:
+        prob = model.predict_proba(input_df)[0][1]
+        result = "✅ Employee is likely to quit." if prob >= threshold else "✅ Employee is likely to stay."
+        return f"{result} (Probability: {prob:.2%})"
+    except Exception as e:
+        return f"❌ Prediction error: {str(e)}"
+
+# Gradio Interface
 def gradio_interface():
-    gr.Interface(
+    interface = gr.Interface(
         fn=predict_employee_status,
         inputs=[
             gr.Number(label="Satisfaction Level (0.0 - 1.0)"),
@@ -69,16 +78,18 @@ def gradio_interface():
             gr.Radio([0, 1], label="Work Accident (0 = No, 1 = Yes)"),
             gr.Radio([0, 1], label="Promotion in Last 5 Years (0 = No, 1 = Yes)"),
             gr.Radio([0, 1, 2], label="Salary (0 = Low, 1 = Medium, 2 = High)"),
-            gr.Dropdown(['RandD', 'accounting', 'hr', 'management', 'marketing',
-                         'product_mng', 'sales', 'support', 'technical'], label="Department"),
+            gr.Dropdown(
+                ['sales', 'accounting', 'hr', 'technical', 'support',
+                 'management', 'IT', 'product_mng', 'marketing', 'RandD'],
+                label="Department"
+            ),
             gr.Slider(0.1, 0.9, value=0.5, step=0.05, label="Prediction Threshold")
         ],
         outputs="text",
-        title="Employee Retention Prediction System (Ensemble from Hugging Face Hub)",
-        description="Predict whether an employee is likely to stay or quit based on their profile. Adjust the threshold for accurate predictions.",
+        title="Employee Retention Prediction System (Voting Ensemble)",
+        description="Predict whether an employee is likely to stay or quit based on their profile. Supports threshold adjustment.",
         theme="dark"
-    ).launch()
+    )
+    interface.launch()
 
 gradio_interface()
-
-