Build an Employee Attrition Predictor

Train a Random Forest model to predict employee turnover using SHAP explainability.

Introduction

30 Second Summary

Every time a skilled employee leaves, companies lose months of productivity and thousands in hiring costs. What if you could predict who is about to leave before they hand in their resignation?

In this project, you will build a machine learning model that predicts employee attrition using real HR data. You will train the model, evaluate its accuracy, and use explainability tools to understand exactly why it flags certain employees as flight risks.

What You'll Build

A Python pipeline that loads IBM's HR Analytics dataset (1,470 employees, 35 features), preprocesses the data, trains a Random Forest classifier, and uses SHAP to explain predictions.

By the end of this project, you'll have:

  • A trained Random Forest model that predicts whether an employee will leave, with measurable accuracy scores.
  • SHAP visualizations that explain which factors drive employee turnover globally and for individual employees.
  • A saved model file (.pkl) ready to be integrated into a Flask API in a follow-on project.
  • Secret Mission: Build a risk scoring function that classifies individual employees as Low, Medium, or High risk and returns the top contributing factors.

Are there any prerequisites?

You need Python installed on a Windows machine and a code editor (such as VS Code). No prior machine learning experience is required.

Set Up Your Environment

Before you can train a machine learning model, you need the right tools installed and a workspace ready to go. This step gets your Python environment set up with the libraries you need for data manipulation, model training, and explainability.

You'll install packages like pandas for data handling, scikit-learn for model training, and SHAP for explaining predictions. Then you'll create a project folder with an empty script file where all your code will live.

In this step, get ready to:
  • Verify your Python and pip installation.
  • Install the required machine learning libraries.
  • Create your project folder and script file.
Verify your Python and pip installation

This project requires Python 3.11 or higher. Let's confirm your installation is ready before installing any packages.

  • Press the Windows key to open your search bar.
  • Type PowerShell and press Enter to open it.
  • Check your Python version by running this command:
python --version

You should see Python 3.11 or higher in the output.

  • Verify pip is available by running this command:
pip --version

You should see a version number and the path to your pip installation.

Not seeing Python 3.11 or higher?

If you see a version lower than 3.11, download the latest Python from python.org. During installation, make sure to check the box that says Add Python to PATH.

If pip --version returns an error, try python -m pip --version instead. If that works, use python -m pip in place of pip for the rest of this project.

Help me troubleshoot my Python or pip installation on Windows.

Install the ML libraries

Your project needs six packages working together. Each one handles a specific part of the machine learning pipeline.

  • Install all required packages by running this command:
pip install pandas scikit-learn shap matplotlib seaborn joblib

What are these packages?

  • pandas loads and manipulates tabular data (your employee dataset).
  • scikit-learn provides the machine learning algorithms and evaluation tools.
  • shap explains why the model makes specific predictions.
  • matplotlib and seaborn create charts and visualizations.
  • joblib saves your trained model to a file so you can reuse it later.

This may take a minute or two as pip downloads and installs each package along with its dependencies.

  • Verify the installation succeeded by running this command:
python -c "import pandas; import sklearn; import shap; print('All packages installed successfully')"

You should see All packages installed successfully printed in your terminal.

Seeing an import error?

If you see ModuleNotFoundError, the install may not have completed for one of the packages. Scroll up in your terminal and look for red error text during the install step.

Try installing the failing package individually (for example, pip install shap) to see a clearer error message. Make sure your internet connection is active.

Help me fix a package installation error with pip on Windows.

Create your project folder and script file

Now that your libraries are installed, you need a dedicated folder for your project files. You'll create this on your Desktop so it's easy to find.

  • Move to your Desktop by running this command:
cd ~/Desktop
  • Create your project folder and move into it by running these commands:
mkdir attrition-predictor
cd attrition-predictor
  • Create your script file by running this command:
New-Item employee_attrition.py
  • Confirm the file exists by running:
ls

You should see employee_attrition.py listed in the output.

  • Open the project folder in VS Code by running this command:
code .

What does this command do?

The code . command opens VS Code with the current folder as your workspace. The dot (.) means "this directory."

If you're using a different code editor, open the attrition-predictor folder from your Desktop using your editor's file open menu.

You should see VS Code open with employee_attrition.py listed in the Explorer sidebar on the left.

VS Code not opening from the terminal?

If you see an error saying code is not recognized, VS Code may not be added to your system PATH. Open VS Code manually from your Start menu, then click File and Open Folder to navigate to your attrition-predictor folder on the Desktop.

Help me open my project folder in VS Code on Windows.

Your environment is ready and your script file is waiting. Next up, you'll load a real HR dataset with 1,470 employee records and explore what the data looks like.

Load and Explore the Dataset

Your environment is ready and your empty script is open. Now it is time to bring in real data.

In this step, you will load IBM's HR Analytics dataset directly from GitHub into a pandas DataFrame. You will explore its structure, check for class imbalance in the target column, and create your first visualization.

In this step, get ready to:
  • Load the IBM HR Analytics dataset and inspect its shape.
  • Analyze the Attrition target variable and calculate the attrition rate.
  • Create a count plot showing how overtime relates to attrition.
Load the dataset using pandas

Before you can explore the data, you need to import all the libraries your script will use throughout this project. This includes tools for data manipulation, visualization, machine learning, and model explainability.

  • In employee_attrition.py, paste the following import block:
import pandas as pd
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score
from sklearn.preprocessing import LabelEncoder
import shap
import joblib

What do these imports do?

  • pandas and numpy handle data loading and numeric operations.
  • matplotlib and seaborn create charts and visualizations.
  • matplotlib.use("Agg") tells matplotlib to save plots to files instead of opening display windows.
  • The sklearn imports bring in scikit-learn tools for machine learning (used in later steps).
  • shap provides model explainability through SHAP values.
  • joblib saves and loads trained models as .pkl files.
  • Save employee_attrition.py.
  • Run your script to confirm all packages import successfully by running:
python employee_attrition.py
  • You should see no output and no errors. The script exits silently because the imports loaded without problems.

Seeing a ModuleNotFoundError?

If you see an error like ModuleNotFoundError: No module named 'pandas', one or more packages from Step 1 did not install correctly.

  • Re-run the install command from Step 1: pip install pandas scikit-learn shap matplotlib seaborn joblib.
  • Make sure you are running python from the same environment where you installed the packages.

Help me fix a ModuleNotFoundError when running my employee_attrition.py script.

Now you can load the actual dataset. IBM published a synthetic HR dataset with 1,470 employee records and 35 features. You will pull it directly from GitHub so there is no manual download needed.

  • Add the following code below your imports in employee_attrition.py:
# --- Step 2: Load and Explore the Dataset ---
url = "https://raw.githubusercontent.com/IBM/employee-attrition-aif360/master/data/emp_attrition.csv"
df = pd.read_csv(url)

print("Dataset Shape:", df.shape)
print("\nFirst 5 rows:")
print(df.head())
print("\nDataset Info:")
df.info()

What does this code do?

  • url points to the raw CSV hosted on IBM's GitHub repository.
  • pd.read_csv(url) downloads the CSV and loads it into a DataFrame called df.
  • df.shape prints the number of rows and columns (1470, 35).
  • df.head() shows the first 5 rows so you can see what the data looks like.
  • df.info() lists every column with its data type and non-null count.
  • Save employee_attrition.py and run it again:
python employee_attrition.py
  • You should see Dataset Shape: (1470, 35) followed by a table of the first 5 rows and a full column listing.

Seeing a connection error?

If you see a URLError or timeout, your internet connection may be blocked or slow.

  • Check that you can reach GitHub from your network.
  • If you are behind a corporate firewall or VPN, try disconnecting and running again.

Help me troubleshoot a connection error when loading data from a GitHub URL in pandas.

Analyze the target variable

The column you want to predict is called Attrition. It contains "Yes" (the employee left) or "No" (the employee stayed). Checking its distribution tells you how balanced the dataset is.

  • Add this code below the df.info() line in employee_attrition.py:
print("\nAttrition Distribution:")
print(df["Attrition"].value_counts())
attrition_rate = df["Attrition"].value_counts(normalize=True)
print(f"\nAttrition Rate: {attrition_rate['Yes']:.2%}")

What does this code do?

  • value_counts() shows how many employees left vs stayed.
  • value_counts(normalize=True) converts those counts to percentages.
  • The dataset has approximately 16% attrition. This means it is imbalanced. Most employees stayed, and only a small fraction left.
  • Save employee_attrition.py and run it:
python employee_attrition.py
  • Scroll to the bottom of the output. You should see the attrition distribution (approximately 1233 "No" and 237 "Yes") and an attrition rate around 16.12%.

Seeing a KeyError on the Attrition column?

If you see KeyError: 'Attrition', double-check the column name is spelled exactly as "Attrition" with a capital A. Column names are case-sensitive.

Help me fix a KeyError when accessing the Attrition column.

Create a quick visualization

A count plot is a fast way to see relationships between categorical columns. You will plot attrition broken down by overtime status to see whether employees who work overtime are more likely to leave.

  • Add this code at the bottom of employee_attrition.py:
# Quick visualization
plt.figure(figsize=(8, 5))
sns.countplot(x="OverTime", hue="Attrition", data=df)
plt.title("Attrition by Overtime Status")
plt.tight_layout()
plt.savefig("attrition_overview.png")
plt.close()
print("\nSaved: attrition_overview.png")

What does this code do?

  • plt.figure(figsize=(8, 5)) creates a canvas sized 8 inches wide by 5 inches tall.
  • sns.countplot() uses seaborn to draw bars showing how many employees fall into each OverTime/Attrition combination.
  • plt.savefig("attrition_overview.png") writes the chart to a file in your project folder.
  • plt.close() frees the figure from memory so it does not interfere with future plots.
  • Save employee_attrition.py and run it:
python employee_attrition.py
  • At the very end of the output you should see Saved: attrition_overview.png.
  • Open the attrition-predictor folder in File Explorer and double-click attrition_overview.png to view the chart. You should see bars comparing attrition counts for employees with and without overtime.

Don't see the PNG file?

  • Make sure you saved employee_attrition.py before running it.
  • Confirm your terminal is in the attrition-predictor folder. The file saves to whatever directory the script runs from.

Help me find my saved plot file.

✔️ Awesome, I've got everything!

Great. Double check you have saved employee_attrition.py and that attrition_overview.png exists in your project folder.

ⓧ I'd like to double check the full code

Here is the complete employee_attrition.py file up to this point:

import pandas as pd
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score
from sklearn.preprocessing import LabelEncoder
import shap
import joblib

# --- Step 2: Load and Explore the Dataset ---
url = "https://raw.githubusercontent.com/IBM/employee-attrition-aif360/master/data/emp_attrition.csv"
df = pd.read_csv(url)

print("Dataset Shape:", df.shape)
print("\nFirst 5 rows:")
print(df.head())
print("\nDataset Info:")
df.info()

print("\nAttrition Distribution:")
print(df["Attrition"].value_counts())
attrition_rate = df["Attrition"].value_counts(normalize=True)
print(f"\nAttrition Rate: {attrition_rate['Yes']:.2%}")

# Quick visualization
plt.figure(figsize=(8, 5))
sns.countplot(x="OverTime", hue="Attrition", data=df)
plt.title("Attrition by Overtime Status")
plt.tight_layout()
plt.savefig("attrition_overview.png")
plt.close()
print("\nSaved: attrition_overview.png")

You now have real data loaded and your first visual insight: overtime appears to be a strong driver of attrition. Next up, you will preprocess this data so the machine learning model can work with it.

Preprocess the Data

In the previous step, you loaded the IBM HR dataset and confirmed it has 1,470 rows and 35 columns. You also spotted that overtime is a strong visual signal for attrition.

Now you need to transform this raw data into a format that a machine learning model can actually learn from. Models work with numbers, not text labels like "Yes" or "Sales". This step cleans out noise columns and converts everything to numeric values.

In this step, get ready to:
  • Drop uninformative columns and encode the target variable.
  • Convert all categorical text columns to numbers using LabelEncoder.
  • Split the data into training and test sets.
Drop uninformative columns and encode the target

Some columns in the dataset contain the same value for every row. These add zero information for the model and can actually confuse it.

  • In employee_attrition.py, scroll to the bottom of your existing code.
  • Add the following code below the print("\nSaved: attrition_overview.png") line:
# --- Step 3: Preprocess the Data ---
# Drop uninformative columns
df = df.drop(columns=["EmployeeCount", "EmployeeNumber", "Over18", "StandardHours"])

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

What does this code do?

  • The # --- Step 3 comment marks the new section in your script.
  • df.drop(columns=[...]) removes four columns that have zero predictive value. EmployeeCount is always 1, EmployeeNumber is just an ID, Over18 is always "Y", and StandardHours is always 80.
  • .map({"Yes": 1, "No": 0}) converts the text-based Attrition column into numeric format. The model needs 1 (left) and 0 (stayed) to compute predictions.

Why map instead of LabelEncoder for the target?

Using .map() gives you explicit control over which label gets which number. You know that 1 means "left" and 0 means "stayed". With LabelEncoder, the assignment is alphabetical, so "No" would become 0 and "Yes" would become 1 anyway. But explicit mapping makes your intent clear.

Encode categorical features

The dataset still has text columns like BusinessTravel, Department, and JobRole. Random Forest requires all inputs to be numeric, so you need to convert these text values into numbers.

  • Add this code directly below the target encoding line you just added:
# Encode categorical columns
categorical_cols = df.select_dtypes(include=["object"]).columns
print(f"\nCategorical columns to encode: {list(categorical_cols)}")

le = LabelEncoder()
for col in categorical_cols:
    df[col] = le.fit_transform(df[col])

print("\nAll columns are now numeric:")
print(df.dtypes.value_counts())

What does this code do?

  • select_dtypes(include=["object"]) finds every column that still holds text (strings). These are your remaining categorical features.
  • LabelEncoder() converts each unique text value to a number. For example, in the Department column, "Sales" might become 2, "R&D" might become 1, and "HR" becomes 0.
  • The for loop applies fit_transform() to each categorical column one at a time. This learns the unique values and replaces them with integers in a single call.
  • The final dtypes.value_counts() print confirms that every column is now int64 (numeric). No text columns remain.
  • Save employee_attrition.py.
  • Run the script in your terminal to confirm the encoding works:
python employee_attrition.py

Near the bottom of the output, you should see a list of categorical column names followed by a dtype summary showing only int64 types.

Seeing a KeyError or encoding error?

  • Make sure the drop and map lines come before the encoding loop. If you try to encode Over18 after it is already dropped, you will get an error.
  • Check that your .map() line uses "Yes" and "No" with capital letters. The dataset uses titlecase values.

Help me debug my LabelEncoder preprocessing code

Split into features and target, then train/test split

With all columns numeric, you can now separate the data into what the model learns from (features) and what it tries to predict (target). You also need to hold back some data for testing so you can measure accuracy fairly.

  • Add this code below the encoding section:
# Split features and target
X = df.drop(columns=["Attrition"])
y = df["Attrition"]

# Train-test split
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

print(f"\nTraining set size: {X_train.shape[0]}")
print(f"Test set size: {X_test.shape[0]}")

What does this code do?

  • X = df.drop(columns=["Attrition"]) creates the features DataFrame. This contains every column except the one you are predicting.
  • y = df["Attrition"] is your target variable. The model will try to predict these 0/1 values.
  • train_test_split with test_size=0.2 holds out 20% of data for testing. The model never sees this during training, so it gives you an honest performance check.
  • random_state=42 makes the split reproducible. Every time you run the script, you get the same split.
  • stratify=y ensures the 16% attrition rate is preserved in both the training and test sets. Without this, random chance could give you a test set with almost no "Left" examples.
  • Save employee_attrition.py and run it again:
python employee_attrition.py

At the bottom of the output, you should see the training set size (approximately 1176) and test set size (approximately 294).

Seeing unexpected set sizes?

  • If both sizes are 0 or very small, confirm you did not accidentally overwrite df with an empty DataFrame. The drop line should use df = df.drop(columns=[...]), not df.drop(columns=[...], inplace=True) called twice.
  • Check that stratify=y uses lowercase y, not uppercase Y.

Help me fix my train_test_split output

✔️ Awesome, I've got everything!

Great. Double check you have saved employee_attrition.py before moving on.

ⓧ I'd like to double check the full code

Here is the complete employee_attrition.py file up to this point:

import pandas as pd
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score
from sklearn.preprocessing import LabelEncoder
import shap
import joblib

# --- Step 2: Load and Explore the Dataset ---
url = "https://raw.githubusercontent.com/IBM/employee-attrition-aif360/master/data/emp_attrition.csv"
df = pd.read_csv(url)

print("Dataset Shape:", df.shape)
print("\nFirst 5 rows:")
print(df.head())
print("\nDataset Info:")
df.info()

print("\nAttrition Distribution:")
print(df["Attrition"].value_counts())
attrition_rate = df["Attrition"].value_counts(normalize=True)
print(f"\nAttrition Rate: {attrition_rate['Yes']:.2%}")

# Quick visualization
plt.figure(figsize=(8, 5))
sns.countplot(x="OverTime", hue="Attrition", data=df)
plt.title("Attrition by Overtime Status")
plt.tight_layout()
plt.savefig("attrition_overview.png")
plt.close()
print("\nSaved: attrition_overview.png")

# --- Step 3: Preprocess the Data ---
# Drop uninformative columns
df = df.drop(columns=["EmployeeCount", "EmployeeNumber", "Over18", "StandardHours"])

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

# Encode categorical columns
categorical_cols = df.select_dtypes(include=["object"]).columns
print(f"\nCategorical columns to encode: {list(categorical_cols)}")

le = LabelEncoder()
for col in categorical_cols:
    df[col] = le.fit_transform(df[col])

print("\nAll columns are now numeric:")
print(df.dtypes.value_counts())

# Split features and target
X = df.drop(columns=["Attrition"])
y = df["Attrition"]

# Train-test split
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

print(f"\nTraining set size: {X_train.shape[0]}")
print(f"Test set size: {X_test.shape[0]}")

Your data is clean, numeric, and split into training and test sets. Next up, you will train a Random Forest classifier on this preprocessed data and measure how well it predicts employee attrition.

Train and Evaluate the Model

Your data is now clean, numeric, and split into training and test sets. The hard preprocessing work is done.

Now comes the exciting part: training an actual machine learning model on your data. You will train a Random Forest classifier to predict whether employees will leave, then measure how well it performs.

In this step, get ready to:
  • Train a Random Forest classifier on the training data.
  • Evaluate model accuracy with a classification report and confusion matrix.
  • Visualize the top 10 most important features as a bar chart.
Train the Random Forest model

A Random Forest works by creating many decision trees and combining their predictions. Each tree sees a slightly different version of the data, so the combined result is more reliable than any single tree.

  • In employee_attrition.py, scroll to the bottom of the file (after the train-test split).
  • Add the following code to train the model and make predictions:
# --- Step 4: Train and Evaluate the Model ---
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train, y_train)

y_pred = model.predict(X_test)

What does this code do?

RandomForestClassifier(n_estimators=100, random_state=42) creates a forest of 100 decision trees. The random_state=42 ensures you get the same results every time you run the script.

model.fit(X_train, y_train) trains the model by showing it the training data. model.predict(X_test) then uses the trained model to predict attrition on the test data it has never seen before.

  • Save employee_attrition.py.

Why Random Forest for this project?

Random Forest provides built-in feature importance scores, works well with mixed feature types, and handles the imbalanced dataset reasonably. It also integrates cleanly with SHAP, which you will use in the next step.

Logistic Regression would be simpler but less accurate. XGBoost would add extra complexity for a first ML project.

Evaluate the model's performance

Training a model is only half the job. You need to measure how well it performs on data it has never seen. This tells you whether the model actually learned useful patterns or just memorized the training data.

  • Add the following evaluation code directly below the prediction line:
print(f"\nAccuracy: {accuracy_score(y_test, y_pred):.4f}")
print("\nClassification Report:")
print(classification_report(y_test, y_pred, target_names=["Stayed", "Left"]))
print("Confusion Matrix:")
print(confusion_matrix(y_test, y_pred))

What does this code do?

accuracy_score gives the overall percentage of correct predictions. This is a quick headline number, but it can be misleading when classes are imbalanced.

classification_report shows precision, recall, and F1-score for each class (Stayed and Left). Pay special attention to recall for the "Left" class. That tells you what percentage of employees who actually left the model caught.

confusion_matrix shows a 2x2 grid of correct vs incorrect predictions for each class. It reveals exactly where the model makes mistakes.

  • Save employee_attrition.py and run the script by running:
python employee_attrition.py
  • In your terminal output, you should see an accuracy score (likely around 0.85+), a classification report table with precision/recall/F1 for both "Stayed" and "Left", and a confusion matrix.

Why does recall for 'Left' matter most?

In HR analytics, missing an employee who is about to leave is far more costly than incorrectly flagging someone who stays. A low recall for the "Left" class means the model is missing at-risk employees. This is the metric HR teams care about most.

Seeing a ModuleNotFoundError?

Make sure you installed all packages in Step 1. Run pip install pandas scikit-learn shap matplotlib seaborn joblib again to confirm.

If you see a NameError for X_train or y_train, make sure your Step 3 code is above this new code in the file.

Help me debug my model training errors.

Visualize feature importance

Random Forest tracks how much each feature contributes to making correct predictions. Visualizing these importance scores reveals which employee attributes drive attrition the most.

  • Add the following code at the bottom of employee_attrition.py to extract and plot feature importance:
# Feature importance
feature_importance = pd.DataFrame({
    "Feature": X.columns,
    "Importance": model.feature_importances_
}).sort_values(by="Importance", ascending=False)

print("\nTop 10 Important Features:")
print(feature_importance.head(10).to_string(index=False))

plt.figure(figsize=(10, 6))
sns.barplot(x="Importance", y="Feature", data=feature_importance.head(10))
plt.title("Top 10 Features for Attrition Prediction")
plt.tight_layout()
plt.savefig("feature_importance.png")
plt.close()
print("\nSaved: feature_importance.png")

What does this code do?

model.feature_importances_ returns an array of scores showing how much each feature contributed to the model's decisions. Higher scores mean the feature was more useful for predicting attrition.

The code wraps these scores into a DataFrame, sorts them from highest to lowest, and plots the top 10 as a horizontal bar chart. This gives you a clear visual answer to the question: "What makes employees leave?"

  • Save employee_attrition.py and run the script again:
python employee_attrition.py
  • You should see a printed list of the top 10 features with their importance scores, and the terminal should confirm Saved: feature_importance.png.
  • Open feature_importance.png in your file explorer. You should see a horizontal bar chart with feature names on the y-axis and importance scores on the x-axis.

Don't see the chart file?

Make sure you are running the script from inside the attrition-predictor folder. The PNG file saves to whatever directory you run the script from.

If you see a matplotlib error, confirm that matplotlib.use("Agg") is near the top of your file, before any other matplotlib imports.

Help me fix my feature importance chart.

✔️ Awesome, I've got everything!

Great. Double check you have saved employee_attrition.py with all your Step 4 code.

ⓧ I'd like to double check the full code

Here is the complete employee_attrition.py file through the end of Step 4. Compare it against yours:

import pandas as pd
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score
from sklearn.preprocessing import LabelEncoder
import shap
import joblib

# --- Step 2: Load and Explore the Dataset ---
url = "https://raw.githubusercontent.com/IBM/employee-attrition-aif360/master/data/emp_attrition.csv"
df = pd.read_csv(url)

print("Dataset Shape:", df.shape)
print("\nFirst 5 rows:")
print(df.head())
print("\nDataset Info:")
df.info()

print("\nAttrition Distribution:")
print(df["Attrition"].value_counts())
attrition_rate = df["Attrition"].value_counts(normalize=True)
print(f"\nAttrition Rate: {attrition_rate['Yes']:.2%}")

# Quick visualization
plt.figure(figsize=(8, 5))
sns.countplot(x="OverTime", hue="Attrition", data=df)
plt.title("Attrition by Overtime Status")
plt.tight_layout()
plt.savefig("attrition_overview.png")
plt.close()
print("\nSaved: attrition_overview.png")

# --- Step 3: Preprocess the Data ---
# Drop uninformative columns
df = df.drop(columns=["EmployeeCount", "EmployeeNumber", "Over18", "StandardHours"])

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

# Encode categorical columns
categorical_cols = df.select_dtypes(include=["object"]).columns
print(f"\nCategorical columns to encode: {list(categorical_cols)}")

le = LabelEncoder()
for col in categorical_cols:
    df[col] = le.fit_transform(df[col])

print("\nAll columns are now numeric:")
print(df.dtypes.value_counts())

# Split features and target
X = df.drop(columns=["Attrition"])
y = df["Attrition"]

# Train-test split
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

print(f"\nTraining set size: {X_train.shape[0]}")
print(f"Test set size: {X_test.shape[0]}")

# --- Step 4: Train and Evaluate the Model ---
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train, y_train)

y_pred = model.predict(X_test)

print(f"\nAccuracy: {accuracy_score(y_test, y_pred):.4f}")
print("\nClassification Report:")
print(classification_report(y_test, y_pred, target_names=["Stayed", "Left"]))
print("Confusion Matrix:")
print(confusion_matrix(y_test, y_pred))

# Feature importance
feature_importance = pd.DataFrame({
    "Feature": X.columns,
    "Importance": model.feature_importances_
}).sort_values(by="Importance", ascending=False)

print("\nTop 10 Important Features:")
print(feature_importance.head(10).to_string(index=False))

plt.figure(figsize=(10, 6))
sns.barplot(x="Importance", y="Feature", data=feature_importance.head(10))
plt.title("Top 10 Features for Attrition Prediction")
plt.tight_layout()
plt.savefig("feature_importance.png")
plt.close()
print("\nSaved: feature_importance.png")

You now have a trained model and know which features drive attrition the most. Next, you will use SHAP to explain exactly why the model makes specific predictions for individual employees.

Explain Predictions with SHAP and Save the Model

Your model can now predict attrition, but prediction alone is not enough. An HR team needs to know why the model flags a specific employee as at-risk so they can take targeted action.

SHAP (SHapley Additive exPlanations) solves this by breaking down each prediction into per-feature contributions. In this step, you will generate SHAP explanations for your model, visualize what drives predictions globally and individually, then save the trained model to disk.

In this step, get ready to:
  • Generate SHAP values and a summary plot showing global feature impact.
  • Explain a single employee's prediction with a force plot.
  • Save the model and feature list as .pkl files for future use.
Generate SHAP values and a summary plot

SHAP's TreeExplainer is purpose-built for tree-based models like Random Forest. It computes how much each feature pushed a prediction toward "Leave" or "Stay" for every employee in the test set.

  • In employee_attrition.py, find the comment # --- Step 4: Train and Evaluate the Model --- section and scroll to the very end of the file (after the feature_importance.png save).
  • Add the following code at the bottom of the file to create a SHAP explainer and generate a summary plot:
# --- Step 5: Explain Predictions with SHAP and Save the Model ---
explainer = shap.TreeExplainer(model)
shap_values = explainer.shap_values(X_test)

# SHAP summary plot for class 1 (Left)
shap.summary_plot(shap_values[1], X_test, show=False)
plt.tight_layout()
plt.savefig("shap_summary.png", bbox_inches="tight")
plt.close()
print("\nSaved: shap_summary.png")

What does this code do?

This section computes how every feature influenced each prediction in your test set.

  • shap.TreeExplainer(model) builds an explainer optimized for tree-based algorithms.
  • explainer.shap_values(X_test) returns a list of two arrays. Index 0 is for "Stayed" and index 1 is for "Left".
  • The summary plot shows all features ranked by their overall impact on the "Left" class across the entire test set.
  • Save employee_attrition.py.
  • Run the script in your terminal by running:
python employee_attrition.py
  • You should see Saved: shap_summary.png printed near the end of the output.

SHAP taking a very long time?

Computing SHAP values for 294 test samples with 100 trees can take 30-60 seconds. This is normal. If it hangs for several minutes, make sure you are using shap.TreeExplainer (not shap.Explainer), which is faster for tree models.

Why is my SHAP summary plot not generating?

Explain a single employee prediction

A summary plot shows the big picture, but HR teams need explanations for individual employees. A force plot shows exactly which features pushed one specific prediction toward "Leave" or "Stay".

  • Add the following code below the print("\nSaved: shap_summary.png") line in employee_attrition.py:
# Explain a single prediction
sample_index = 0
sample = X_test.iloc[[sample_index]]
prediction = model.predict(sample)[0]
probability = model.predict_proba(sample)[0]

print(f"\nSample Employee Prediction: {'Will Leave' if prediction == 1 else 'Will Stay'}")
print(f"Probability of Leaving: {probability[1]:.2%}")
print(f"Probability of Staying: {probability[0]:.2%}")

# SHAP force plot for single prediction
shap.force_plot(
    explainer.expected_value[1],
    shap_values[1][sample_index],
    sample.iloc[0],
    matplotlib=True,
    show=False
)
plt.tight_layout()
plt.savefig("shap_individual.png", bbox_inches="tight")
plt.close()
print("Saved: shap_individual.png")

What does this code do?

  • X_test.iloc[[sample_index]] picks the first employee in the test set as our sample.
  • predict_proba(sample)[0] returns an array of two probabilities: index 0 for "Stay" and index 1 for "Leave".
  • shap.force_plot() visualizes how each feature value for this employee nudged the prediction up or down relative to the baseline.
  • explainer.expected_value[1] is the average model output (the baseline). Features push the prediction above or below this baseline.
  • Save employee_attrition.py and run the script again:
python employee_attrition.py
  • You should see output like Sample Employee Prediction: Will Stay (or Will Leave), followed by probabilities and Saved: shap_individual.png.

Seeing a matplotlib display error?

If you see an error about the display or backend, make sure matplotlib.use("Agg") is still at the top of your file (right after import matplotlib). This tells matplotlib to render without opening a window.

Help me fix my matplotlib display error when generating SHAP plots.

Save the model and feature list

A trained model is useless if it only lives in memory. Saving to a .pkl file means you can load it later without retraining. You also need to save the feature column names so a future API knows exactly which columns to expect.

  • Add the following code at the bottom of employee_attrition.py (after the shap_individual.png save):
# Save model and features
joblib.dump(model, "attrition_model.pkl")
joblib.dump(X.columns.tolist(), "model_features.pkl")
print("\nModel saved as 'attrition_model.pkl'")
print("Feature list saved as 'model_features.pkl'")

# Verify by loading
loaded_model = joblib.load("attrition_model.pkl")
test_pred = loaded_model.predict(sample)
print(f"Verification - Loaded model prediction: {'Will Leave' if test_pred[0] == 1 else 'Will Stay'}")

What does this code do?

  • joblib.dump(model, "attrition_model.pkl") serializes the trained model to a file on disk.
  • X.columns.tolist() captures the exact list of 30 feature column names the model was trained on. A future Flask API needs this list to arrange incoming data in the correct order.
  • The verification block loads the model back from the file and makes a prediction. If the output matches the earlier prediction, the save was successful.
  • Save employee_attrition.py and run the script one final time:
python employee_attrition.py
  • You should see Model saved as 'attrition_model.pkl' and Feature list saved as 'model_features.pkl' in the terminal output.
  • You should also see a Verification - Loaded model prediction: line confirming the saved model works.
  • Confirm the files exist by running:
ls
  • You should see attrition_model.pkl, model_features.pkl, shap_summary.png, and shap_individual.png alongside your existing files.

Missing .pkl files?

Make sure the joblib.dump() lines are not indented inside a function or if-block. They should sit at the top level of your script, at the same indentation as the rest of your Step 5 code.

Also confirm your terminal is in the attrition-predictor folder. The files save to whatever directory the terminal is currently in.

Help me figure out why my .pkl files are not saving.

Why save the feature list separately?

The model itself does not remember column names. When new employee data arrives through an API, you need to know the exact column order the model expects. Saving model_features.pkl ensures you can reconstruct the correct feature order without re-running the training pipeline.

✔️ Awesome, I've got everything!

Great. Double check you have saved employee_attrition.py and confirmed the .pkl files exist in your folder.

ⓧ I'd like to double check the full code

Here is the complete employee_attrition.py file including all steps up to this point:

import pandas as pd
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score
from sklearn.preprocessing import LabelEncoder
import shap
import joblib

# --- Step 2: Load and Explore the Dataset ---
url = "https://raw.githubusercontent.com/IBM/employee-attrition-aif360/master/data/emp_attrition.csv"
df = pd.read_csv(url)

print("Dataset Shape:", df.shape)
print("\nFirst 5 rows:")
print(df.head())
print("\nDataset Info:")
df.info()

print("\nAttrition Distribution:")
print(df["Attrition"].value_counts())
attrition_rate = df["Attrition"].value_counts(normalize=True)
print(f"\nAttrition Rate: {attrition_rate['Yes']:.2%}")

# Quick visualization
plt.figure(figsize=(8, 5))
sns.countplot(x="OverTime", hue="Attrition", data=df)
plt.title("Attrition by Overtime Status")
plt.tight_layout()
plt.savefig("attrition_overview.png")
plt.close()
print("\nSaved: attrition_overview.png")

# --- Step 3: Preprocess the Data ---
# Drop uninformative columns
df = df.drop(columns=["EmployeeCount", "EmployeeNumber", "Over18", "StandardHours"])

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

# Encode categorical columns
categorical_cols = df.select_dtypes(include=["object"]).columns
print(f"\nCategorical columns to encode: {list(categorical_cols)}")

le = LabelEncoder()
for col in categorical_cols:
    df[col] = le.fit_transform(df[col])

print("\nAll columns are now numeric:")
print(df.dtypes.value_counts())

# Split features and target
X = df.drop(columns=["Attrition"])
y = df["Attrition"]

# Train-test split
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

print(f"\nTraining set size: {X_train.shape[0]}")
print(f"Test set size: {X_test.shape[0]}")

# --- Step 4: Train and Evaluate the Model ---
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train, y_train)

y_pred = model.predict(X_test)

print(f"\nAccuracy: {accuracy_score(y_test, y_pred):.4f}")
print("\nClassification Report:")
print(classification_report(y_test, y_pred, target_names=["Stayed", "Left"]))
print("Confusion Matrix:")
print(confusion_matrix(y_test, y_pred))

# Feature importance
feature_importance = pd.DataFrame({
    "Feature": X.columns,
    "Importance": model.feature_importances_
}).sort_values(by="Importance", ascending=False)

print("\nTop 10 Important Features:")
print(feature_importance.head(10).to_string(index=False))

plt.figure(figsize=(10, 6))
sns.barplot(x="Importance", y="Feature", data=feature_importance.head(10))
plt.title("Top 10 Features for Attrition Prediction")
plt.tight_layout()
plt.savefig("feature_importance.png")
plt.close()
print("\nSaved: feature_importance.png")

# --- Step 5: Explain Predictions with SHAP and Save the Model ---
explainer = shap.TreeExplainer(model)
shap_values = explainer.shap_values(X_test)

# SHAP summary plot for class 1 (Left)
shap.summary_plot(shap_values[1], X_test, show=False)
plt.tight_layout()
plt.savefig("shap_summary.png", bbox_inches="tight")
plt.close()
print("\nSaved: shap_summary.png")

# Explain a single prediction
sample_index = 0
sample = X_test.iloc[[sample_index]]
prediction = model.predict(sample)[0]
probability = model.predict_proba(sample)[0]

print(f"\nSample Employee Prediction: {'Will Leave' if prediction == 1 else 'Will Stay'}")
print(f"Probability of Leaving: {probability[1]:.2%}")
print(f"Probability of Staying: {probability[0]:.2%}")

# SHAP force plot for single prediction
shap.force_plot(
    explainer.expected_value[1],
    shap_values[1][sample_index],
    sample.iloc[0],
    matplotlib=True,
    show=False
)
plt.tight_layout()
plt.savefig("shap_individual.png", bbox_inches="tight")
plt.close()
print("Saved: shap_individual.png")

# Save model and features
joblib.dump(model, "attrition_model.pkl")
joblib.dump(X.columns.tolist(), "model_features.pkl")
print("\nModel saved as 'attrition_model.pkl'")
print("Feature list saved as 'model_features.pkl'")

# Verify by loading
loaded_model = joblib.load("attrition_model.pkl")
test_pred = loaded_model.predict(sample)
print(f"Verification - Loaded model prediction: {'Will Leave' if test_pred[0] == 1 else 'Will Stay'}")

Secret mission

Build a Risk Scoring Function

Your model can predict attrition across an entire test set. But what if HR needs to assess a single employee right now? Build a function that takes one employee's data, returns a risk classification (Low, Medium, or High), and explains why with the top contributing factors from SHAP.

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely locally, so there are no ongoing costs. Decide whether to keep your resources for the follow-on Flask API project, or delete them entirely.

Resources you used:

  • The attrition-predictor folder on your Desktop (contains your script, model files, and plot images).
  • Python packages installed via pip (pandas, scikit-learn, shap, matplotlib, seaborn, joblib).

Keep everything running

No action needed. Choose this if you plan to continue with the follow-on project where you wrap this model in a Flask API.

  • Your trained model (attrition_model.pkl) and feature list (model_features.pkl) remain in the attrition-predictor folder, ready for Project 2.
  • The pip packages stay available for any future ML work.

Pause - I'll come back to this later

Since everything is local, there is nothing to pause or shut down. Your files sit idle on disk until you return.

  • Leave the attrition-predictor folder in place. Nothing consumes resources while the script is not running.
  • When you come back, open the folder in your editor and pick up where you left off.

Delete - I don't want to use this again

Remove all project files from your machine. Only do this if you do not plan to continue with the Flask API project.

  • Delete the project folder by running the command for your operating system.

macOS

cd ~/Desktop
rm -rf attrition-predictor

Windows

cd ~\Desktop
rmdir /s /q attrition-predictor
  • Optionally, uninstall the Python packages you installed for this project:
pip uninstall pandas scikit-learn shap matplotlib seaborn joblib -y

Should I uninstall these packages?

These are general-purpose data science libraries. If you plan to do more Python ML work, keep them installed. Only uninstall if you want a completely clean slate.

Nice Work!

Nice Work!

You just built a complete machine learning pipeline that predicts employee attrition and explains its own reasoning. That is a real HR Analytics skill you can walk into an interview and demonstrate.

Here's what you learned:

  • Trained a Random Forest classifier on IBM's HR Analytics dataset (1,470 employees, 30 features after preprocessing) to predict employee attrition.
  • Evaluated model performance using a classification report and confusion matrix, focusing on recall for the minority class that HR cares about most.
  • Used SHAP to explain which features drive attrition predictions globally and for individual employees, then saved the model as a .pkl file ready for API integration.
  • Secret Mission: Built a risk scoring function that classifies employees as Low/Medium/High risk with the top contributing factors.

Ready to quiz yourself?