Spaces:
Sleeping
Sleeping
| import pickle | |
| 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 Pickle | |
| with open("h22_xgb_Final.pkl", "wb") as f: | |
| loaded_model = pickle.load(f) | |
| # Ensure model is a Booster (handles both XGBClassifier & Booster cases) | |
| if isinstance(loaded_model, xgb.XGBClassifier): | |
| loaded_model = loaded_model.get_booster() | |
| # Setup SHAP Explainer for XGBoost | |
| explainer = shap.TreeExplainer(loaded_model) # Use TreeExplainer for XGBoost models | |
| # 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 new_row to DMatrix for XGBoost Booster | |
| dmatrix_new = xgb.DMatrix(new_row) | |
| # Predict probability for staying (XGBoost Booster returns only one class probability) | |
| prob = loaded_model.predict(dmatrix_new) | |
| # Compute SHAP values | |
| shap_values = explainer.shap_values(new_row) | |
| plot = shap.plots.bar(shap_values, 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() | |