<jupyter_start><jupyter_text>Load data<jupyter_code>import shap  # https://github.com/slundberg/shap
import shapreg  # https://github.com/iancovert/shapley-regression
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
# Load and split data
X_train, X_test, Y_train, Y_test = train_test_split(
    *shap.datasets.adult(), test_size=0.2, random_state=7)
X_train, X_val, Y_train, Y_val = train_test_split(
    X_train, Y_train, test_size=0.2, random_state=0)

# Data scaling
num_features = X_train.shape[1]
feature_names = X_train.columns.tolist()
ss = StandardScaler()
ss.fit(X_train)
X_train = ss.transform(X_train.values)
X_val = ss.transform(X_val.values)
X_test = ss.transform(X_test.values)<jupyter_output><empty_output><jupyter_text>Train model<jupyter_code>import pickle
import os.path
import lightgbm as lgb
if os.path.isfile('census model.pkl'):
    print('Loading saved model')
    with open('census model.pkl', 'rb') as f:
        model = pickle.load(f)

else:
    # Setup
    params = {
        "max_bin": 512,
        "learning_rate": 0.05,
        "boosting_type": "gbdt",
        "objective": "binary",
        "metric": "binary_logloss",
        "num_leaves": 10,
        "verbose": -1,
        "min_data": 100,
        "boost_from_average": True
    }

    # More setup
    d_train = lgb.Dataset(X_train, label=Y_train)
    d_val = lgb.Dataset(X_val, label=Y_val)

    # Train model
    model = lgb.train(params, d_train, 10000, valid_sets=[d_val],
                      early_stopping_rounds=50, verbose_eval=1000)
    
    # Save model
    with open('census model.pkl', 'wb') as f:
        pickle.dump(model, f)<jupyter_output>Loading saved model<jupyter_text>Train surrogate<jupyter_code>import torch
import torch.nn as nn
from fastshap.utils import MaskLayer1d
from fastshap import Surrogate, KLDivLoss
# Select device
device = torch.device('cuda')
# Check for model
if os.path.isfile('census surrogate.pt'):
    print('Loading saved surrogate model')
    surr = torch.load('census surrogate.pt').to(device)
    surrogate = Surrogate(surr, num_features)

else:
    # Create surrogate model
    surr = nn.Sequential(
        MaskLayer1d(value=0, append=True),
        nn.Linear(2 * num_features, 128),
        nn.ELU(inplace=True),
        nn.Linear(128, 128),
        nn.ELU(inplace=True),
        nn.Linear(128, 2)).to(device)

    # Set up surrogate object
    surrogate = Surrogate(surr, num_features)

    # Set up original model
    def original_model(x):
        pred = model.predict(x.cpu().numpy())
        pred = np.stack([1 - pred, pred]).T
        return torch.tensor(pred, dtype=torch.float32, device=x.device)

    # Train
    surrogate.train_original_model(
        X_train,
        X_val,
        original_model,
        batch_size=64,
        max_epochs=100,
        loss_fn=KLDivLoss(),
        validation_samples=10,
        validation_batch_size=10000,
        verbose=True)

    # Save surrogate
    surr.cpu()
    torch.save(surr, 'census surrogate.pt')
    surr.to(device)<jupyter_output>Loading saved surrogate model<jupyter_text>Train FastSHAP<jupyter_code>from fastshap import FastSHAP
# Check for model
if os.path.isfile('census explainer.pt'):
    print('Loading saved explainer model')
    explainer = torch.load('census explainer.pt').to(device)
    fastshap = FastSHAP(explainer, surrogate, normalization='additive',
                        link=nn.Softmax(dim=-1))

else:
    # Create explainer model
    explainer = nn.Sequential(
        nn.Linear(num_features, 128),
        nn.ReLU(inplace=True),
        nn.Linear(128, 128),
        nn.ReLU(inplace=True),
        nn.Linear(128, 2 * num_features)).to(device)

    # Set up FastSHAP object
    fastshap = FastSHAP(explainer, surrogate, normalization='additive',
                        link=nn.Softmax(dim=-1))

    # Train
    fastshap.train(
        X_train,
        X_val[:100],
        batch_size=32,
        num_samples=32,
        max_epochs=200,
        validation_samples=128,
        verbose=True)
    
    # Save explainer
    explainer.cpu()
    torch.save(explainer, 'census explainer.pt')
    explainer.to(device)<jupyter_output>Loading saved explainer model<jupyter_text>Compare with KernelSHAP<jupyter_code>import matplotlib.pyplot as plt
# Setup for KernelSHAP
def imputer(x, S):
    x = torch.tensor(x, dtype=torch.float32, device=device)
    S = torch.tensor(S, dtype=torch.float32, device=device)
    pred = surrogate(x, S).softmax(dim=-1)
    return pred.cpu().data.numpy()
# Select example
ind = np.random.choice(len(X_test))
x = X_test[ind:ind+1]
y = int(Y_test[ind])

# Run FastSHAP
fastshap_values = fastshap.shap_values(x)[0]

# Run KernelSHAP to convergence
game = shapreg.games.PredictionGame(imputer, x)
shap_values, all_results = shapreg.shapley.ShapleyRegression(
    game, batch_size=32, paired_sampling=False, detect_convergence=True,
    bar=True, return_all=True)
# Create figure
plt.figure(figsize=(9, 5.5))

# Bar chart
width = 0.75
kernelshap_iters = 128
plt.bar(np.arange(num_features) - width / 3, shap_values.values[:, y],
        width / 3, label='True SHAP values', color='tab:gray')
plt.bar(np.arange(num_features), fastshap_values[:, y],
        width / 3, label='FastSHAP', color='tab:green')
plt.bar(np.arange(num_features) + width / 3,
        all_results['values'][list(all_results['iters']).index(kernelshap_iters)][:, y],
        width / 3, label='KernelSHAP @ {}'.format(kernelshap_iters), color='tab:red')

# Annotations
plt.legend(fontsize=16)
plt.tick_params(labelsize=14)
plt.ylabel('SHAP Values', fontsize=16)
plt.title('Census Explanation Example', fontsize=18)
plt.xticks(np.arange(num_features), feature_names,
           rotation=35, rotation_mode='anchor', ha='right')

plt.tight_layout()
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12551159618008187, 'fraction_numerical': 0.020975443383356072, 'mean_word_length': 3.6290449881610103, 'pattern_counts': {'":': 9, '<': 18, '<?xml version=': 0, '>': 18, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '577701', 'license': 'MIT', 'n_tokens_mistral': 2195, 'n_tokens_neox': 2051, 'n_words': 395, 'path': 'notebooks/census.ipynb', 'repo_name': 'ghanasyam-rallapalli/fastshap', 'size': '5864'}