A binary classifier that predicts employee attrition risk using XGBoost trained on 14 engineered features derived from 8 raw HR signals.
1XGBClassifier(
2 n_estimators=400,
3 max_depth=5,
4 learning_rate=0.05,
5 subsample=0.85,
6 colsample_bytree=0.85,
7 eval_metric="logloss",
8 random_state=42
9)
1import joblib
2import pandas as pd
3import numpy as np
4
5bundle = joblib.load("attrition_model_bundle_XGB_FE.pkl")
6pipeline = bundle["pipeline"]
7features = bundle["features"] # 14-element list
8add_features = bundle["feature_engineering"] # callable
9thresholds = bundle["risk_thresholds"] # {"high": 0.60, "medium": 0.35}
1def add_features(df: pd.DataFrame) -> pd.DataFrame:
2 df = df.copy()
3 df["stress_per_tenure"] = df["StressScore"] / (df["TenureMonths"] + 1)
4 df["engagement_drop"] = df["EngagementScore"] * (1 - df["EngagementTrend"])
5 df["burnout_index"] = (df["StressScore"] * df["ConflictScore"]) / (df["EngagementScore"] + 1)
6 df["stress_log"] = np.log1p(df["StressScore"])
7 df["tenure_sqrt"] = np.sqrt(df["TenureMonths"])
8 df["fe_risk_score"] = (
9 0.4 * df["StressScore"]
10 + 0.3 * df["ConflictScore"]
11 - 0.3 * df["EngagementScore"]
12 )
13 return df
1RAW_FEATURES = [
2 "EngagementScore", "TenureMonths", "StressScore", "PromotionDelay",
3 "EngagementTrend", "SentimentScore", "PerformanceScore", "ConflictScore"
4]
5
6def predict(raw_input: dict) -> dict:
7 row = pd.DataFrame([raw_input])[RAW_FEATURES]
8 row = add_features(row)[features]
9
10 prob = pipeline.predict_proba(row)[0, 1]
11 prob = np.clip(prob, 0.05, 0.95)
12
13 tier = "High" if prob >= 0.60 else "Medium" if prob >= 0.35 else "Low"
14
15 return {
16 "resign_pct": round(prob * 100, 2),
17 "stay_pct": round((1 - prob) * 100, 2),
18 "risk_tier": tier,
19 "probability": round(prob, 4),
20 }
21
22# Example
23result = predict({
24 "EngagementScore": 3.2,
25 "TenureMonths": 28,
26 "StressScore": 8.7,
27 "PromotionDelay": 28,
28 "EngagementTrend": -0.62,
29 "SentimentScore": -0.30,
30 "PerformanceScore": 6.8,
31 "ConflictScore": 6.2,
32})
33print(result)
34# {'resign_pct': 95.0, 'stay_pct': 5.0, 'risk_tier': 'High', 'probability': 0.95}
1df_raw = pd.read_csv("employees.csv")
2df_fe = add_features(df_raw)[features]
3
4probs = pipeline.predict_proba(df_fe)[:, 1]
5probs = np.clip(probs, 0.05, 0.95)
6
7df_raw["resign_pct"] = (probs * 100).round(2)
8df_raw["risk_tier"] = pd.cut(
9 probs,
10 bins=[0, 0.35, 0.60, 1.0],
11 labels=["Low", "Medium", "High"]
12)
1import shap
2
3scaler = pipeline.named_steps["scaler"]
4xgb_model = pipeline.named_steps["xgb"]
5
6# Scale with feature names preserved — required for correct SHAP labels
7X_scaled = pd.DataFrame(
8 scaler.transform(X_sample),
9 columns=X_sample.columns
10)
11
12# Background sample for interventional perturbation
13background = pd.DataFrame(
14 scaler.transform(X_train.sample(100, random_state=42)),
15 columns=X_train.columns
16)
17
18explainer = shap.TreeExplainer(xgb_model, background, feature_perturbation="interventional")
19shap_values = explainer(X_scaled)
20
21shap.plots.beeswarm(shap_values, max_display=14)
1{
2 "pipeline": <sklearn Pipeline>, # StandardScaler + XGBClassifier
3 "features": [...], # ordered list of 14 feature names
4 "feature_engineering": <callable>, # add_features() function
5 "model_type": "stack_xgb_fe",
6 "risk_thresholds": {"high": 0.60, "medium": 0.35},
7 "shap_note": "Scale before SHAP; wrap in named DataFrame"
8}