Predicting Employee Exit with Keras: Building an ANN for Workforce Retention

import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, LabelEncoder

# Load the HR dataset
df = pd.read_csv('HR_Employee_Attrition.csv')

# Examine the target variable distribution
print(f"Total employees: {len(df)}")
print(f"Attrition rate: {df['Attrition'].value_counts(normalize=True)}")

# Check for class imbalance
attrition_counts = df['Attrition'].value_counts()
print(f"Stayed: {attrition_counts['No']} | Left: {attrition_counts['Yes']}")

# Encode categorical variables
categorical_columns = [
    'BusinessTravel', 'Department', 'EducationField',
    'Gender', 'JobRole', 'MaritalStatus', 'OverTime'
]

label_encoders = {}
for col in categorical_columns:
    le = LabelEncoder()
    df[col] = le.fit_transform(df[col])
    label_encoders[col] = le

# Convert target variable
df['Attrition'] = df['Attrition'].map({'Yes': 1, 'No': 0})

# Select features for the model
feature_columns = [
    'Age', 'BusinessTravel', 'DailyRate', 'Department',
    'DistanceFromHome', 'Education', 'EnvironmentSatisfaction',
    'Gender', 'HourlyRate', 'JobInvolvement', 'JobLevel',
    'JobRole', 'JobSatisfaction', 'MaritalStatus', 'MonthlyIncome',
    'MonthlyRate', 'NumCompaniesWorked', 'OverTime',
    'PercentSalaryHike', 'PerformanceRating', 'RelationshipSatisfaction',
    'StockOptionLevel', 'TotalWorkingYears', 'TrainingTimesLastYear',
    'WorkLifeBalance', 'YearsAtCompany', 'YearsInCurrentRole',
    'YearsSinceLastPromotion', 'YearsWithCurrManager'
]

X = df[feature_columns]
y = df['Attrition']

# Split the data
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

# Scale features for neural network
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

print(f"Training samples: {X_train_scaled.shape[0]}")
print(f"Test samples: {X_test_scaled.shape[0]}")
print(f"Features: {X_train_scaled.shape[1]}")

from keras.models import Sequential
from keras.layers import Dense, Dropout, BatchNormalization
from keras.optimizers import Adam
from keras.callbacks import EarlyStopping, ReduceLROnPlateau

def build_attrition_model(input_dim):
    """
    Build an ANN for employee attrition prediction.

    Architecture:
    - Input layer matching feature dimensions
    - Three hidden layers with decreasing neurons
    - Batch normalization for training stability
    - Dropout for regularization
    - Sigmoid output for binary classification
    """
    model = Sequential([
        # First hidden layer
        Dense(128, activation='relu', input_dim=input_dim),
        BatchNormalization(),
        Dropout(0.3),

        # Second hidden layer
        Dense(64, activation='relu'),
        BatchNormalization(),
        Dropout(0.3),

        # Third hidden layer
        Dense(32, activation='relu'),
        BatchNormalization(),
        Dropout(0.2),

        # Output layer
        Dense(1, activation='sigmoid')
    ])

    # Compile with binary crossentropy for classification
    model.compile(
        optimizer=Adam(learning_rate=0.001),
        loss='binary_crossentropy',
        metrics=['accuracy', 'AUC']
    )

    return model

# Create the model
model = build_attrition_model(X_train_scaled.shape[1])
model.summary()

from sklearn.utils.class_weight import compute_class_weight

# Calculate class weights to handle imbalance
class_weights = compute_class_weight(
    class_weight='balanced',
    classes=np.unique(y_train),
    y=y_train
)
class_weight_dict = dict(enumerate(class_weights))

print(f"Class weights: {class_weight_dict}")
# Typical output: {0: 0.58, 1: 2.89} - higher weight for minority class

# Define callbacks for optimal training
callbacks = [
    EarlyStopping(
        monitor='val_loss',
        patience=15,
        restore_best_weights=True,
        verbose=1
    ),
    ReduceLROnPlateau(
        monitor='val_loss',
        factor=0.5,
        patience=5,
        min_lr=0.0001,
        verbose=1
    )
]

# Train the model
history = model.fit(
    X_train_scaled,
    y_train,
    epochs=100,
    batch_size=32,
    validation_split=0.2,
    class_weight=class_weight_dict,
    callbacks=callbacks,
    verbose=1
)

import matplotlib.pyplot as plt

fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# Loss curves
axes[0].plot(history.history['loss'], label='Training Loss')
axes[0].plot(history.history['val_loss'], label='Validation Loss')
axes[0].set_xlabel('Epoch')
axes[0].set_ylabel('Loss')
axes[0].set_title('Model Loss Over Training')
axes[0].legend()

# AUC curves
axes[1].plot(history.history['auc'], label='Training AUC')
axes[1].plot(history.history['val_auc'], label='Validation AUC')
axes[1].set_xlabel('Epoch')
axes[1].set_ylabel('AUC')
axes[1].set_title('Model AUC Over Training')
axes[1].legend()

plt.tight_layout()
plt.savefig('training_history.png', dpi=150)

from sklearn.metrics import (
    classification_report, confusion_matrix,
    roc_auc_score, precision_recall_curve, roc_curve
)

# Make predictions
y_pred_proba = model.predict(X_test_scaled)
y_pred = (y_pred_proba > 0.5).astype(int)

# Classification report
print("Classification Report:")
print(classification_report(y_test, y_pred, target_names=['Stay', 'Leave']))

# ROC-AUC Score
roc_auc = roc_auc_score(y_test, y_pred_proba)
print(f"\nROC-AUC Score: {roc_auc:.4f}")

import seaborn as sns

cm = confusion_matrix(y_test, y_pred)

plt.figure(figsize=(8, 6))
sns.heatmap(
    cm, annot=True, fmt='d', cmap='Blues',
    xticklabels=['Predicted Stay', 'Predicted Leave'],
    yticklabels=['Actual Stay', 'Actual Leave']
)
plt.title('Confusion Matrix - Employee Attrition Prediction')
plt.ylabel('Actual')
plt.xlabel('Predicted')
plt.savefig('confusion_matrix.png', dpi=150)

# Calculate precision-recall curve
precision, recall, thresholds = precision_recall_curve(y_test, y_pred_proba)

# Find optimal threshold for recall-focused prediction
target_recall = 0.80  # Catch 80% of leavers
optimal_idx = np.argmin(np.abs(recall - target_recall))
optimal_threshold = thresholds[optimal_idx]

print(f"Optimal threshold for {target_recall*100}% recall: {optimal_threshold:.3f}")

# Apply optimized threshold
y_pred_optimized = (y_pred_proba > optimal_threshold).astype(int)
print("\nOptimized Classification Report:")
print(classification_report(y_test, y_pred_optimized, target_names=['Stay', 'Leave']))

from keras import backend as K

def get_feature_importance(model, X, feature_names):
    """
    Calculate feature importance using gradient-based method.
    """
    X_tensor = K.variable(X)

    with tf.GradientTape() as tape:
        tape.watch(X_tensor)
        predictions = model(X_tensor)

    gradients = tape.gradient(predictions, X_tensor)
    importance = np.mean(np.abs(gradients.numpy()), axis=0)

    return pd.DataFrame({
        'feature': feature_names,
        'importance': importance
    }).sort_values('importance', ascending=False)

# Get feature importance
importance_df = get_feature_importance(
    model, X_test_scaled[:100], feature_columns
)

# Display top factors
print("Top 10 Attrition Risk Factors:")
print(importance_df.head(10))

import joblib

# Save the trained model
model.save('employee_attrition_model.h5')

# Save the scaler for preprocessing
joblib.dump(scaler, 'feature_scaler.pkl')

# Save label encoders
joblib.dump(label_encoders, 'label_encoders.pkl')

print("Model and preprocessing artifacts saved successfully.")

from keras.models import load_model

class AttritionPredictor:
    """
    Production-ready attrition prediction class.
    """

    def __init__(self, model_path, scaler_path, encoders_path):
        self.model = load_model(model_path)
        self.scaler = joblib.load(scaler_path)
        self.encoders = joblib.load(encoders_path)
        self.threshold = 0.35  # Optimized threshold

    def preprocess(self, employee_data):
        """Preprocess employee data for prediction."""
        df = pd.DataFrame([employee_data])

        # Encode categorical variables
        for col, encoder in self.encoders.items():
            if col in df.columns:
                df[col] = encoder.transform(df[col])

        # Scale features
        return self.scaler.transform(df)

    def predict(self, employee_data):
        """
        Predict attrition probability for an employee.

        Returns:
            dict: risk_score, risk_level, recommendations
        """
        X = self.preprocess(employee_data)
        probability = self.model.predict(X)[0][0]

        # Determine risk level
        if probability < 0.3:
            risk_level = "Low"
        elif probability < 0.5:
            risk_level = "Medium"
        elif probability < 0.7:
            risk_level = "High"
        else:
            risk_level = "Critical"

        return {
            'risk_score': float(probability),
            'risk_level': risk_level,
            'at_risk': probability > self.threshold
        }

# Usage example
predictor = AttritionPredictor(
    'employee_attrition_model.h5',
    'feature_scaler.pkl',
    'label_encoders.pkl'
)

sample_employee = {
    'Age': 32,
    'BusinessTravel': 'Travel_Frequently',
    'Department': 'Sales',
    'OverTime': 'Yes',
    'MonthlyIncome': 5000,
    'YearsAtCompany': 3,
    'JobSatisfaction': 2,
    # ... other features
}

result = predictor.predict(sample_employee)
print(f"Risk Score: {result['risk_score']:.2%}")
print(f"Risk Level: {result['risk_level']}")

def calculate_retention_roi(
    total_employees,
    attrition_rate,
    avg_salary,
    replacement_cost_ratio,
    model_catch_rate,
    intervention_success_rate
):
    """
    Calculate ROI from ML-powered retention program.
    """
    # Expected leavers without intervention
    expected_leavers = total_employees * attrition_rate

    # Replacement cost per employee
    replacement_cost = avg_salary * replacement_cost_ratio

    # Leavers caught by model
    caught_by_model = expected_leavers * model_catch_rate

    # Successfully retained through intervention
    retained = caught_by_model * intervention_success_rate

    # Savings
    savings = retained * replacement_cost

    return {
        'expected_leavers': int(expected_leavers),
        'employees_retained': int(retained),
        'annual_savings': savings
    }

# Example calculation
roi = calculate_retention_roi(
    total_employees=5000,
    attrition_rate=0.15,
    avg_salary=75000,
    replacement_cost_ratio=0.75,
    model_catch_rate=0.80,
    intervention_success_rate=0.40
)

print(f"Expected leavers: {roi['expected_leavers']}")
print(f"Employees retained: {roi['employees_retained']}")
print(f"Annual savings: ${roi['annual_savings']:,.0f}")