attrition-prediction-engine

Attrition Prediction — Model B (XGBoost + Feature Engineering)

A binary classifier that predicts employee attrition risk using XGBoost trained on 14 engineered features derived from 8 raw HR signals.

---

Model at a Glance

---

Features

Raw Inputs (8)

Engineered Features (6)

---

Hyperparameters

XGBClassifier(
    n_estimators=400,
    max_depth=5,
    learning_rate=0.05,
    subsample=0.85,
    colsample_bytree=0.85,
    eval_metric="logloss",
    random_state=42
)

Pipeline: StandardScaler → XGBClassifier

---

Installation

pip install xgboost scikit-learn joblib pandas numpy

---

Usage

Load the bundle

import joblib
import pandas as pd
import numpy as np

bundle = joblib.load("attrition_model_bundle_XGB_FE.pkl")
pipeline   = bundle["pipeline"]
features   = bundle["features"]          # 14-element list
add_features = bundle["feature_engineering"]  # callable
thresholds = bundle["risk_thresholds"]   # {"high": 0.60, "medium": 0.35}

Feature engineering function

def add_features(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    df["stress_per_tenure"] = df["StressScore"] / (df["TenureMonths"] + 1)
    df["engagement_drop"]   = df["EngagementScore"] * (1 - df["EngagementTrend"])
    df["burnout_index"]     = (df["StressScore"] * df["ConflictScore"]) / (df["EngagementScore"] + 1)
    df["stress_log"]        = np.log1p(df["StressScore"])
    df["tenure_sqrt"]       = np.sqrt(df["TenureMonths"])
    df["fe_risk_score"]     = (
        0.4 * df["StressScore"]
        + 0.3 * df["ConflictScore"]
        - 0.3 * df["EngagementScore"]
    )
    return df

Predict a single employee

RAW_FEATURES = [
    "EngagementScore", "TenureMonths", "StressScore", "PromotionDelay",
    "EngagementTrend", "SentimentScore", "PerformanceScore", "ConflictScore"
]

def predict(raw_input: dict) -> dict:
    row = pd.DataFrame([raw_input])[RAW_FEATURES]
    row = add_features(row)[features]

    prob = pipeline.predict_proba(row)[0, 1]
    prob = np.clip(prob, 0.05, 0.95)

    tier = "High" if prob >= 0.60 else "Medium" if prob >= 0.35 else "Low"

    return {
        "resign_pct": round(prob * 100, 2),
        "stay_pct":   round((1 - prob) * 100, 2),
        "risk_tier":  tier,
        "probability": round(prob, 4),
    }

# Example
result = predict({
    "EngagementScore": 3.2,
    "TenureMonths": 28,
    "StressScore": 8.7,
    "PromotionDelay": 28,
    "EngagementTrend": -0.62,
    "SentimentScore": -0.30,
    "PerformanceScore": 6.8,
    "ConflictScore": 6.2,
})
print(result)
# {'resign_pct': 95.0, 'stay_pct': 5.0, 'risk_tier': 'High', 'probability': 0.95}

Batch prediction

df_raw = pd.read_csv("employees.csv")
df_fe  = add_features(df_raw)[features]

probs = pipeline.predict_proba(df_fe)[:, 1]
probs = np.clip(probs, 0.05, 0.95)

df_raw["resign_pct"] = (probs * 100).round(2)
df_raw["risk_tier"]  = pd.cut(
    probs,
    bins=[0, 0.35, 0.60, 1.0],
    labels=["Low", "Medium", "High"]
)

---

Risk Tiers

---

SHAP Explainability

import shap

scaler    = pipeline.named_steps["scaler"]
xgb_model = pipeline.named_steps["xgb"]

# Scale with feature names preserved — required for correct SHAP labels
X_scaled = pd.DataFrame(
    scaler.transform(X_sample),
    columns=X_sample.columns
)

# Background sample for interventional perturbation
background = pd.DataFrame(
    scaler.transform(X_train.sample(100, random_state=42)),
    columns=X_train.columns
)

explainer   = shap.TreeExplainer(xgb_model, background, feature_perturbation="interventional")
shap_values = explainer(X_scaled)

shap.plots.beeswarm(shap_values, max_display=14)

Important: always pass the scaled feature matrix to TreeExplainer, not raw values. Wrap the scaled array in a named DataFrame to preserve feature labels in plots.

---

Bundle Contents

{
    "pipeline":            <sklearn Pipeline>,  # StandardScaler + XGBClassifier
    "features":            [...],               # ordered list of 14 feature names
    "feature_engineering": <callable>,          # add_features() function
    "model_type":          "stack_xgb_fe",
    "risk_thresholds":     {"high": 0.60, "medium": 0.35},
    "shap_note":           "Scale before SHAP; wrap in named DataFrame"
}

---

Notes

• Probabilities are clipped to [0.05, 0.95] to prevent degenerate 0% / 100% outputs on edge-case inputs.
• The fe_risk_score feature is a useful standalone KPI for HR dashboards — it requires no model inference.
• For highest accuracy at the cost of inference speed, use Model C (stacking ensemble). Model B is recommended for real-time API deployments.
• Re-train quarterly on updated HR data to account for organisational changes.