Spaces:
Sleeping
Sleeping
import gradio as gr | |
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 (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 | |
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', | |
'product_mng', 'sales', 'support', 'technical' | |
] | |
department_encoded = [1 if dept == department else 0 for dept in departments] | |
# Prepare the input with all 18 features | |
input_data = np.array([ | |
satisfaction_level, last_evaluation, number_project, | |
average_monthly_hours, time_spent_company, work_accident, | |
promotion_last_5years, salary | |
] + department_encoded).reshape(1, -1) | |
# Predict using the model | |
if model is None: | |
return "β No model found. Please upload the model file." | |
try: | |
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)}" | |
# Gradio interface | |
interface = gr.Interface( | |
fn=predict_employee_status, | |
inputs=[ | |
gr.Number(label="Satisfaction Level (0.0 - 1.0)"), | |
gr.Number(label="Last Evaluation (0.0 - 1.0)"), | |
gr.Number(label="Number of Projects (1 - 10)"), | |
gr.Number(label="Average Monthly Hours (80 - 320)"), | |
gr.Number(label="Time Spent at Company (Years)"), | |
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" | |
) | |
], | |
outputs="text", | |
title="Employee Retention Prediction System", | |
description="Predict whether an employee is likely to stay or quit based on their profile." | |
) | |
interface.launch() | |