Mod3_Team5 / app.py
Mpodszus's picture
Update app.py
a3ffeca verified
raw
history blame
3.75 kB
import xgboost as xgb
import pandas as pd
import shap
from shap.plots._force_matplotlib import draw_additive_plot
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")
# Setup SHAP Explainer
explainer = shap.Explainer(loaded_model) # PLEASE DO NOT CHANGE THIS.
# 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()
# Convert input to DMatrix format (needed 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)
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
title = "**Mod 3 Team 5: Employee Turnover Predictor**"
description1 = """
This app takes six inputs about employees' satisfaction with different aspects of their work (such as work-life balance, ...)
and predicts whether the employee intends to stay with the employer or leave. The outputs include:
1. The predicted probability of staying or leaving.
2. A SHAP plot that visualizes how different factors impact the prediction.
"""
description2 = """
To use the app, adjust the values of the six employee satisfaction factors and click **Analyze**. ✨
"""
with gr.Blocks(title=title) as demo:
gr.Markdown(f"## {title}")
gr.Markdown(description1)
gr.Markdown("""---""")
gr.Markdown(description2)
gr.Markdown("""---""")
with gr.Row():
with gr.Column():
SupportiveGM = gr.Slider(label="Supportive GM Score", minimum=1, maximum=5, value=4, step=0.1)
Merit = gr.Slider(label="Merit Score", minimum=1, maximum=5, value=4, step=0.1)
LearningDevelopment = gr.Slider(label="Learning & Development Score", minimum=1, maximum=5, value=4, step=0.1)
WorkEnvironment = gr.Slider(label="Work Environment Score", minimum=1, maximum=5, value=4, step=0.1)
Engagement = gr.Slider(label="Engagement Score", minimum=1, maximum=5, value=4, step=0.1)
WellBeing = gr.Slider(label="Well-Being Score", minimum=1, maximum=5, value=4, step=0.1)
submit_btn = gr.Button("Analyze")
with gr.Column(visible=True, scale=1, min_width=600) as output_col:
label = gr.Label(label="Predicted Turnover Probability")
local_plot = gr.Plot(label="SHAP Plot:")
submit_btn.click(
main_func,
[SupportiveGM, Merit, LearningDevelopment, WorkEnvironment, Engagement, WellBeing],
[label, local_plot], api_name="Employee_Turnover"
)
gr.Markdown("### Click on an example below to see how it works:")
gr.Examples([[4, 4, 4, 4, 5, 5], [5, 4, 5, 4, 4, 4]],
[SupportiveGM, Merit, LearningDevelopment, WorkEnvironment, Engagement, WellBeing],
[label, local_plot], main_func, cache_examples=True)
demo.launch()