Integrated-gradient on IMDB dataset (PyTorch)

This is an example of the integrated-gradient method on text classification with a PyTorch model. If using this explainer, please cite the original work: https://github.com/ankurtaly/Integrated-Gradients.

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import sklearn
from sklearn.datasets import fetch_20newsgroups

from omnixai.data.text import Text
from omnixai.preprocessing.text import Word2Id
from omnixai.explainers.tabular.agnostic.L2X.utils import Trainer, InputData, DataLoader
from omnixai.explainers.nlp.specific.ig import IntegratedGradientText

We apply a simple CNN model for this text classification task. Note that the method forward has two inputs inputs (token ids) and masks (the sentence masks). For IntegratedGradientText, the first input of the model must be the token ids.

class TextModel(nn.Module):

    def __init__(self, num_embeddings, num_classes, **kwargs):
        super().__init__()
        self.num_embeddings = num_embeddings
        self.embedding_size = kwargs.get("embedding_size", 50)
        self.embedding = nn.Embedding(self.num_embeddings, self.embedding_size)
        self.embedding.weight.data.normal_(mean=0.0, std=0.01)

        hidden_size = kwargs.get("hidden_size", 100)
        kernel_sizes = kwargs.get("kernel_sizes", [3, 4, 5])
        if type(kernel_sizes) == int:
            kernel_sizes = [kernel_sizes]

        self.activation = nn.ReLU()
        self.conv_layers = nn.ModuleList([
            nn.Conv1d(self.embedding_size, hidden_size, k, padding=k // 2) for k in kernel_sizes])
        self.dropout = nn.Dropout(0.2)
        self.output_layer = nn.Linear(len(kernel_sizes) * hidden_size, num_classes)

    def forward(self, inputs, masks):
        embeddings = self.embedding(inputs)
        x = embeddings * masks.unsqueeze(dim=-1)
        x = x.permute(0, 2, 1)
        x = [self.activation(layer(x).max(2)[0]) for layer in self.conv_layers]
        outputs = self.output_layer(self.dropout(torch.cat(x, dim=1)))
        if outputs.shape[1] == 1:
            outputs = outputs.squeeze(dim=1)
        return outputs

We use a Text object to represent a batch of texts/sentences. The package omnixai.preprocessing.text provides some transforms related to text data such as Tfidf and Word2Id.

# Load the training and test datasets
train_data = pd.read_csv('/home/ywz/data/imdb/labeledTrainData.tsv', sep='\t')
n = int(0.8 * len(train_data))
x_train = Text(train_data["review"].values[:n])
y_train = train_data["sentiment"].values[:n].astype(int)
x_test = Text(train_data["review"].values[n:])
y_test = train_data["sentiment"].values[n:].astype(int)
class_names = ["negative", "positive"]
# The transform for converting words/tokens to IDs
transform = Word2Id().fit(x_train)

The preprocessing function converts a batch of texts into token IDs and the masks. The outputs of the preprocessing function must fit the inputs of the model.

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

def preprocess(X: Text):
    samples = transform.transform(X)
    max_len = 0
    for i in range(len(samples)):
        max_len = max(max_len, len(samples[i]))
    max_len = min(max_len, max_length)
    inputs = np.zeros((len(samples), max_len), dtype=int)
    masks = np.zeros((len(samples), max_len), dtype=np.float32)
    for i in range(len(samples)):
        x = samples[i][:max_len]
        inputs[i, :len(x)] = x
        masks[i, :len(x)] = 1
    return inputs, masks

We now train the CNN model and evaluate its performance.

model = TextModel(
    num_embeddings=transform.vocab_size,
    num_classes=len(class_names)
).to(device)

Trainer(
    optimizer_class=torch.optim.AdamW,
    learning_rate=1e-3,
    batch_size=128,
    num_epochs=10,
).train(
    model=model,
    loss_func=nn.CrossEntropyLoss(),
    train_x=transform.transform(x_train),
    train_y=y_train,
    padding=True,
    max_length=max_length,
    verbose=True
)

 |████████████████████████████████████████| 100.0% Complete, Loss 0.0008

model.eval()
data = transform.transform(x_test)
data_loader = DataLoader(
    dataset=InputData(data, [0] * len(data), max_length),
    batch_size=32,
    collate_fn=InputData.collate_func,
    shuffle=False
)
outputs = []
for inputs in data_loader:
    value, mask, target = inputs
    y = model(value.to(device), mask.to(device))
    outputs.append(y.detach().cpu().numpy())
outputs = np.concatenate(outputs, axis=0)
predictions = np.argmax(outputs, axis=1)
print('Test accuracy: {}'.format(
    sklearn.metrics.f1_score(y_test, predictions, average='binary')))

Test accuracy: 0.8458027386386188

To initialize IntegratedGradientText, we need to set the following parameters:

• model: The model to explain, whose type is tf.keras.Model or torch.nn.Module.

• embedding_layer: The embedding layer in the model, which can be tf.keras.layers.Layer or torch.nn.Module.

• preprocess_function: The pre-processing function that converts the raw input data into the inputs of model. The first output of preprocess_function should be the token ids.

• mode: The task type, e.g., classification or regression.

• id2token: The mapping from token ids to tokens.

explainer = IntegratedGradientText(
    model=model,
    embedding_layer=model.embedding,
    preprocess_function=preprocess,
    id2token=transform.id_to_word
)
x = Text([
    "What a great movie! if you have no taste.",
    "it was a fantastic performance!",
    "best film ever",
    "such a great show!",
    "it was a horrible movie",
    "i've never watched something as bad"
])
explanations = explainer.explain(x)
explanations.ipython_plot(class_names=class_names)

Instance 0: Class positive

what a great movie if you have no taste

Instance 1: Class positive

it was a fantastic performance

Instance 2: Class positive

best film ever

Instance 3: Class positive

such a great show

Instance 4: Class negative

it was a horrible movie

Instance 5: Class negative

i never watched something as bad