Skip to content

Ingesting SHAP Explanations for NLP

This notebook covers using TruEra's Python SDK to develop NLP models locally with SHAP explanations.

First, a sample Tensorflow 2 BERT model is trained on the Covid-19 Tweets dataset. Next, the SDK computes explanations for each record. Finally, the model, data, and explanations are ingested into a TruEra workspace, enabling analysis via TruEra's notebook widgets and Web App interface.

Each step in these instructions easily adapts to your own model or dataset.

Step 1. General Setup

Setup entails installing the required dependencies, installing the TruEra Python SDK (if not already installed), and setting custom size limits.

A. Install Dependencies

Install the following dependencies to your Python environment:

# Install visualization dependencies
pip3 install "plotly>5.5.0" "ipywidgets>7.7.0" "notebook>6.4.0"

You may need to restart your notebook after installing these dependencies.

B. Install the TruEra Python SDK and Download Sample Data

First, install the TruEra Python SDK with the command line: pip3 install "truera[nlp]"

Don't have access to pip?

Download and install the TruEra Python package from app.truera.net (see Download the TruEra SDK).

Next, download the NLP Quickstart Data in TAR or ZIP from the Account Settings > Resources page (see Downloading Samples).

Now extract the NLP Quickstart Data using either tar -xvf nlp_quickstart_data.tgz or unzip nlp_quickstart_data.zip

Then install the wheel in your Python environment using pip install 'truera-*.whl[nlp]'.

Careful

This will only work from the directory in which the whl is located.

C. Customizations

Although it is recommended that you use GPU whenever possible for NLP models, here we'll use small samples running on CPU for demo purposes only.

# Using small size limits for demo purposes only
demo_train_records = 100
demo_test_records = 100
truera_batch_size = 16
demo_influences = 4
train_epochs = 1

Step 2. Train your NLP Model

The example training code that follows is strictly for demonstration purposes. Adapt/replace it with your own model and code, when ready.

A. Import Dependencies

import os
import numpy as np
import pandas as pd
import nltk

nltk.download('stopwords')
from nltk.corpus import stopwords

B. Prep the Data

QUICKSTART_DATA_DIRECTORY = "<QUICKSTART_DATA_DIRECTORY>"

train_df_path = os.path.join(QUICKSTART_DATA_DIRECTORY, 'covid_tweets', 'train', 'dataset.csv')
test_df_path = os.path.join(QUICKSTART_DATA_DIRECTORY, 'covid_tweets', 'val_test', 'dataset.csv')

train_data = pd.read_csv(train_df_path, encoding='L1')
test_data = pd.read_csv(test_df_path,encoding='L1')

train_data.head()

def clean_data(df: pd.DataFrame, stopwords: set[str], n_records: int = None):
    """Clean the data by removing stopwords and converting sentiment labels to binary class IDs."""

    # Apply record limit
    if n_records is not None:
        df = df[:n_records].copy()

    # Remove rows with missing values
    df.dropna(subset=['OriginalTweet', 'Sentiment'], inplace=True)

    # Remove stopwords
    stopwords_check = lambda word: word.lower() not in stopwords
    text_preprocess_fn = lambda text: " ".join(filter(stopwords_check, str(text).split()))
    df['OriginalTweet'] = df['OriginalTweet'].apply(text_preprocess_fn)

    # Convert sentiment labels to binary class IDs
    sentiment_tol = {
        'Extremely Negative': 0,
        'Negative': 0,
        'Neutral': 0,
        'Positive': 1,
        'Extremely Positive': 1
    }
    process_labels_fn = lambda sentiments: [sentiment_tol[sentiment] for sentiment in sentiments]
    df['Sentiment'] = process_labels_fn(df['Sentiment'])
    return df

# Preprocess data
sw_nltk = set(stopwords.words('english'))
train_data = clean_data(train_data, sw_nltk, demo_train_records)
test_data = clean_data(test_data, sw_nltk, demo_test_records)

train_df = train_data.sample(frac=0.7, random_state=0)  # random state is a seed value
val_df = train_data.drop(train_df.index)
test_df = test_data

# Convert text to numpy arrays
train_texts = np.array(train_df['OriginalTweet'])
val_texts = np.array(val_df['OriginalTweet'])
test_texts = np.array(test_df['OriginalTweet'])

# Convert labels to numpy arrays
train_labels = np.array(train_df['Sentiment'])
val_labels = np.array(val_df['Sentiment'])
test_labels = np.array(test_df['Sentiment'])

C. Create a PyTorch Dataset and DataLoader

from torch.utils import data

class TextDataset(data.Dataset):
    def __init__(self, texts, labels):
        self.texts = texts
        self.labels = labels

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, index: int):
        return (index, self.texts[i], self.labels[i])

train_dataset = TextDataset(train_texts, train_labels)
val_dataset = TextDataset(val_texts, val_labels)
test_dataset = TextDataset(test_texts, test_labels)

BATCH_SIZE = 32

train_dl = data.DataLoader(train_dataset, batch_size=BATCH_SIZE)
val_dl = data.DataLoader(val_dataset, batch_size=BATCH_SIZE)
test_dl = data.DataLoader(test_dataset, batch_size=BATCH_SIZE)

D. Create a Tokenizer and a Model

import torch
from typing import Iterable, Callable
from transformers import BertTokenizerFast, BertModel

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

SEQ_LENGTH = 128

tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")

def tokenize_text(texts: Iterable[str]):
    return tokenizer(
        list(texts),
        padding='max_length',
        max_length=SEQ_LENGTH,
        truncation=True,
        return_tensors='pt'
    )['input_ids'].to(device)

import torch.nn as nn

MODEL_NAME = "covid_tweets_sentiment"
EMBEDDING_DIM = 512

class SentimentClassifier(nn.Module):

    def __init__(self, tokenizer: Callable):
        super(SentimentClassifier, self).__init__()

        self.name = MODEL_NAME
        self.bert = BertModel.from_pretrained(
            'google/bert_uncased_L-4_H-512_A-8')  # B x 128 x 512
        self.bert.resize_token_embeddings(len(tokenizer))

        self.flatten = torch.nn.Flatten()
        self.lin = torch.nn.Linear(in_features=128 * 512,
                                   out_features=1,
                                   bias=False)
        self.sigmoid = torch.nn.Sigmoid()

    def forward(self, *args, **kwargs):
        bert_output = self.bert(*args, **kwargs)
        flattened = self.flatten(bert_output[0])
        lin_output = self.lin(flattened)
        sigmoid = self.sigmoid(lin_output)

        return sigmoid

    def get_latent_embeddings(self, *args, **kwargs):
        return self.bert(*args, **kwargs)['pooler_output'] 

model = SentimentClassifier(tokenizer).to(device)

E. Define a Sentence Embedding Function (optional)

def get_embeddings(model: nn.Module, texts: Iterable[str]):
    model.eval()
    data = tokenize_text(texts).to(device)
    return model.get_latent_embeddings(data).float().cpu().detach().numpy().tolist()

F. Train

LEARNING_RATE = 1e-05

bce_loss = nn.BCELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)

def accuracy(output: torch.Tensor, target: torch.Tensor):
    acc = (output.argmax(-1) == target).float().cpu().numpy()
    return float(100 * acc.sum() / len(acc))

def train(model: nn.Module, train_loader: data.DataLoader):
    model.train()
    losses = []
    accs = []

    for batch_idx, batch in enumerate(train_loader):
        _, data, target = batch
        data = tokenize_text(data)
        data, target = data.to(device), target.to(device)

        output = model(data)
        optimizer.zero_grad()
        loss = bce_loss(torch.squeeze(output).float(), target.float())
        loss.backward()
        optimizer.step()

        losses.append(loss.item())
        accs.append(accuracy(output, target))

    print(f" - loss={np.mean(losses)}, acc={np.mean(accs)}")
    # return training loss, training accuracy
    return np.mean(losses), np.mean(accs)


def evaluate(model, test_loader):
    model.eval()
    losses = []
    accs = []

    with torch.no_grad():
        for batch in test_loader:
            _, data, target = batch
            data = tokenize_text(data)
            data, target = data.to(device), target.to(device)
            output = model(data)
            loss = bce_loss(torch.squeeze(output).float(), target.float())

            losses.append(loss.item())
            accs.append(accuracy(output, target))
    # return validation loss, validation accuracy
    return np.mean(losses), np.mean(accs)

metrics = {'train_loss': [], 'train_acc': [], 'val_loss': [], 'val_acc': []}
EPOCHS = 5

for i in range(EPOCHS):
    print(f"EPOCH: {i}")
    train_loss, train_acc = train(model, train_dl)
    val_loss, val_acc = evaluate(model, val_dl)
    for metric in metrics.keys():
        metrics[metric].append(eval(metric))
metrics

Step 3. Generate Explanations with the SHAP Explainer

The SHAP Explainer requires a model evaluation callable and a tokenizer callable. Click here to learn more.

# define a prediction function
def eval_model(x):
    model.eval()
    tv = tokenize_text(x)
    outputs = model(tv).float().cpu().detach().numpy().reshape(-1)
    return outputs

import shap
explainer = shap.Explainer(eval_model, tokenizer)
shap_infl = explainer(val_texts).values

# Influences need to be in shape batch_size x num_classes x num_tokens
shap_infl = [vals.T.reshape((1, -1)).tolist() for vals in shap_infl]

Step 4. Ingest SHAP Explanations into TruEra

A. Connect to the TruEra Endpoint

  • Provide your Truera deployment URL (http://app.truera.net).
  • Provide your generated token.
  • Truera Workspace creation will verify your connectivity to Truera services.
TRUERA_URL = "<TRUERA_URL>"
TOKEN = "<TRUERA_TOKEN>"

from truera.client.truera_workspace import TrueraWorkspace
from truera.client.truera_authentication import TokenAuthentication

# Replace with your URLs, credentials 
auth = TokenAuthentication(TOKEN)
tru = TrueraWorkspace(TRUERA_URL, auth)

B. Create a TruEra Project

Add a NLP project with a data collection and a model. 

from uuid import uuid4

PROJECT_NAME = f"NLP_Ingestion_{uuid4()}"
DATA_COLLECTION_NAME = "dc"

# Setup remote environment
tru.add_project(PROJECT_NAME, score_type="probits", input_type="text")
tru.add_data_collection(DATA_COLLECTION_NAME)
tru.add_model(MODEL_NAME)
tru

C. Prepare Split Data

from datetime import datetime

df = pd.DataFrame({
    "ids": val_df.index,
    "text": val_texts, 
    "labels": val_labels, 
    "tokens": [tokenizer.convert_ids_to_tokens(tokenizer.encode(v)) for v in val_texts],
    "embeddings": get_embeddings(model, val_texts), 
    "preds": eval_model(val_texts).tolist(), 
    "infl": shap_infl
})

df.head()

Specify NLPColumnSpec to map different columns used by TruEra to columns from your ingested dataframe to create a schema:

from truera.client.ingestion import NLPColumnSpec

column_spec = NLPColumnSpec(
    id_col_name='ids',
    text_col_name='text',
    tokens_col_name='tokens',
    sentence_embeddings_col_name='embeddings',
    prediction_col_name='preds',
    label_col_name='labels',
    token_influence_col_name='infl'
)

Store metadata about the model and influence type (if influences are being ingested) by creating a ModelOutputContext() object.

from truera.client.ingestion import ModelOutputContext

model_output_context = ModelOutputContext(
    model_name=MODEL_NAME,
    score_type="probits",
    influence_type="integrated-gradients"
)

Now ingest the data with add_data(), which requires a RemoteTrueraWorkspace, a unique split name, the prepared dataframe and the NLPColumnSpec() schema object, as well as the ModelOutputContext() object. 

SPLIT_NAME = f"covid_tweets_split"

tru.add_data(
    data=df,
    data_split_name=SPLIT_NAME,
    column_spec=column_spec,
    model_output_context=model_output_context,
)

print(f"Uploaded to {PROJECT_NAME} > {SPLIT_NAME}")

If no errors are encountered, the split is now be ingested into the TruEra workspace.

D. Query Back Split Data

We can now pick up where we left off by reconnecting to our workspace and configuring it to point to the project and split just ingested.

TRUERA_URL = "<TRUERA_URL>"
TOKEN = "<TRUERA_TOKEN>"

from truera.client.truera_workspace import TrueraWorkspace
from truera.client.truera_authentication import TokenAuthentication

# Replace with your URLs, credentials 
auth = TokenAuthentication(TOKEN)
tru = TrueraWorkspace(TRUERA_URL, auth)

# Setup remote environment
tru.set_project(PROJECT_NAME)
tru.set_data_collection(DATA_COLLECTION_NAME)
tru.set_model(MODEL_NAME)
tru.set_data_split(SPLIT_NAME)
tru

E. Query Back Raw Ingested Data

# Get ingested tokens
tru.get_tokens()

# Get ingested sentence embeddings
tru.get_embeddings()

# Get original text
tru.get_xs()

# Get labels
tru.get_ys()

# Get feature influences
tru.get_feature_influences()

F. Run Local and Global Explanation Widgets

tru.get_explainer().global_token_summary()

tru.get_explainer().record_explanations_attribution_tab()