Employee attrition, end to end

The IBM HR Analytics dataset: 1,470 employees, 30 columns, and a column saying who left. Below is the whole run — every line of code, and every line it printed.

Imports

One import from our package. The rest is the usual stack, plus XGBoost to compare against.

import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import balanced_accuracy_score
from xgboost import XGBClassifier
from white_box_builder import WhiteBoxBuilder

Load and split

The raw CSV goes in as it is, with its text columns and its numeric columns. It needs no encoding and no scaling. The target gets the two words that we want to read in the result. A quarter of the rows stay back for the test.

df = pd.read_csv("hr.csv")
X = df.drop(columns="Attrition")
y = df["Attrition"].map({"Yes": "LEAVES", "No": "STAYS"})
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.25, stratify=y, random_state=0)

Build the rule models

Two lines. The second one asks for five models of increasing size, so you can see what each extra condition buys you.

builder = WhiteBoxBuilder().fit(X_train, y_train)
models = builder.build("BalancedAccuracyDNF", k=range(1, 6))

The baseline

XGBoost on the same split. It needs the categorical columns one-hot encoded first.

encoded = pd.get_dummies(X)
xgb_model = XGBClassifier().fit(encoded.loc[X_train.index], y_train == "LEAVES")
predicted = xgb_model.predict(encoded.loc[X_test.index])
xgboost_score = balanced_accuracy_score(y_test == "LEAVES", predicted)

Score them

The rule model is not a summary of a black box or an approximation of one. It is the classifier, and this is what it says.

best = models.evaluate(X_test, y_test, metric="balanced_accuracy")[0]
print(f"glass box balanced accuracy: {best.scores['balanced_accuracy']:.3f}")
print(f"xgboost balanced accuracy:   {xgboost_score:.3f}")
print("\nThe glass box:\n")
print(f"{best}")

glass box balanced accuracy: 0.682
xgboost balanced accuracy:   0.636

The glass box:

IF
    NOT JobRole=Sales Representative AND StockOptionLevel=0 AND OverTime
    OR JobRole=Sales Representative AND NOT StockOptionLevel=0 AND NOT OverTime
    OR JobLevel=Entry Level AND OverTime
    OR JobLevel=Entry Level AND StockOptionLevel=0
THEN class = LEAVES
ELSE class = STAYS

All five, smallest first

You get the whole set, not one model that we chose for you. Read the models in order. You see what each new condition buys in accuracy. Then you decide where to stop.

for model in models.sort("k", ascending=True):
    score = model.scores["balanced_accuracy"]
    print(f"k={model.build_settings['k']}  balanced accuracy {score:.3f}")
    print(model, end="\n\n")

k=1  balanced accuracy 0.648
IF
    OverTime
THEN class = LEAVES
ELSE class = STAYS

k=3  balanced accuracy 0.664
IF
    NOT JobRole=Sales Representative AND OverTime
    OR JobRole=Sales Representative AND NOT OverTime
    OR JobLevel=Entry Level AND OverTime
THEN class = LEAVES
ELSE class = STAYS

k=4  balanced accuracy 0.682
IF
    NOT JobRole=Sales Representative AND StockOptionLevel=0 AND OverTime
    OR JobRole=Sales Representative AND NOT StockOptionLevel=0 AND NOT OverTime
    OR JobLevel=Entry Level AND OverTime
    OR JobLevel=Entry Level AND StockOptionLevel=0
THEN class = LEAVES
ELSE class = STAYS

k=5  balanced accuracy 0.682
IF
    NOT JobRole=Sales Representative AND StockOptionLevel=0 AND OverTime
    OR JobRole=Sales Representative AND NOT StockOptionLevel=0 AND NOT OverTime
    OR JobLevel=Entry Level AND NOT MaritalStatus=Single AND OverTime
    OR JobLevel=Entry Level AND StockOptionLevel=0
THEN class = LEAVES
ELSE class = STAYS

Become a pilot partner