Can I explain a text model with Vertex AI? Yes, you can!
During this year, I presented Vertex Explainable AI several times. One of the most frequent questions I received was:
Can I explain a text model with Vertex AI?
In this blog post of the Vertex AI Q&A series , I will show you how to get feature-based explanations using Vertex Explainable AI for a recurrent neural network model with text data. The model is trained using IMDB large movie review dataset to predict the sentiment of the movie reviews — positive or negative. You can find more details of the TensorFlow model and training procedure in the following tutorial.
Recap: Explainable AI on Vertex AI
To support responsible AI applications, Vertex AI provides built-in explainability capabilities. Currently, Vertex Explainable AI supports both feature-based and example-based explanations.
With feature-based explanations, you can explain how much each feature in your model influenced the predictions for a given instance. Depending on the type of data (tabular, text, image), the type of model (differentiable or non-differentiable) and the task (classification, regression) you are trying to accomplish, Vertex Explainable AI offers three methods to use for feature attributions: Sampled Shapley, Integrated gradients, and XRAI. For a text classification model, you can use both Sampled Shapley and Integrated gradients.
Both Sampled Shapley and Integrated gradients are based on the Shapley value, a concept from game theory that assigns a value to each player in a game based on their contribution to the outcome of the game. With a text classification model, each word or token (feature) is treated as a “player” in the game. When you request explanations, Vertex Explainable AI indicates how much each word or token contributes to the predicted class.
For the rest of the blog, I will focus on Sampled Shapley and how to get these kinds of feature-based explanations on a custom text model using TensorFlow. To learn more about Sampled Shapley and Integrated Gradients and how they differ from each other, see the AI Explanations Whitepaper. Also, see the official Vertex AI documentation for a broad overview of Vertex Explainable AI.
Configure feature-based explanations for TensorFlow custom model
According to the Vertex AI documentation, getting feature-based explanations on a custom TensorFlow model using Vertex Explainable AI requires:
- finding the input and output tensor names,
- exporting your TensorFlow model as a TensorFlow SavedModel,
- finding the input and output tensor names from a TensorFlow SavedModel that has already been exported (if you don’t have the training code)
In our case, I assume that I downloaded a pre-trained TensorFlow model without any prerequisite of familiarity with the underlying code. To check whether the provided export directory contains a SavedModel, you can use the tf.saved_model.contains_saved_model method as shown below.
import tensorflow as tf
is_saved_model = tf.saved_model.contains_saved_model(model_uri)
print(is_saved_model)The method returns true if the export directory contains a SavedModel.
To find the input and output tensor names from a TensorFlow SavedModel, you can load the model and check signatures.
import tensorflow as tf
loaded_model = tf.saved_model.load(model_uri)
serving_input = list(
loaded_model.signatures["serving_default"].structured_input_signature[1].keys()
)[0]
serving_output = list(
loaded_model.signatures["serving_default"].structured_outputs.keys()
)[0]
print("Serving function input:", serving_input)
print("Serving function output:", serving_output)Finding the input and output tensor names it is important to configure explanations by providing an explanationSpec, which is the specification of the model explanation.
The explanationSpec requires the explanation_parameters and explanation_metadata. The explanation_parameters lets you define the method you want to use for getting explanations. Below find how to set the Sampled Shapley parameters using Vertex AI SDK
from google.cloud import aiplatform
XAI = "shapley"
PARAMETERS = {"sampled_shapley_attribution": {"path_count": 10}}
explanation_parameters = aiplatform.explain.ExplanationParameters(PARAMETERS)where path_count is the number of feature permutations for the Sampled Shapley attribution method.
The explanation_metadata contains the inputs and outputs metadata of your TensorFlow models. Below, find the explanation_metadata for our text classification model.
INPUT_METADATA = {
"my_input": aiplatform.explain.ExplanationMetadata.InputMetadata(
{
"input_tensor_name": serving_input,
}
),
}
OUTPUT_METADATA = {
"my_output": aiplatform.explain.ExplanationMetadata.OutputMetadata(
{"output_tensor_name": serving_output}
)
}
explanation_metadata = aiplatform.explain.ExplanationMetadata(
inputs=INPUT_METADATA, outputs=OUTPUT_METADATA
)Notice that explanation_metadata is optional for TensorFlow 2 models. If omitted, Vertex AI automatically infers the inputs and outputs from the model. Also Vertex AI SDK supports a method to build the explanation metadata of a TensorFlowmodel.
Next, you register your model as a Vertex AI Model resource to Vertex AI Model Registry with the explanation configuration you just defined. It’s important to mention that passing this configuration while you register the model is just one option. You can find other alternatives here.
from google.cloud import aiplatform
MODEL_DISPLAY_NAME = "tf_text_model_unique"
DEPLOY_IMAGE = "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-11:latest"
registered_model = aiplatform.Model.upload(
display_name=MODEL_DISPLAY_NAME,
artifact_uri=MODEL_DIR,
serving_container_image_uri=DEPLOY_IMAGE,
explanation_parameters=explanation_parameters,
explanation_metadata=explanation_metadata,
)To register your model, you must provide a serving image. A prebuilt container for serving predictions and explanations has been used . Alternatively, you can use a custom container. See the Vertex AI documentation for more on how to get explanations with a custom container.
Get your explained predictions
After you register your model, you are ready to deploy your model and get the explained predictions. Below you will find how to deploy your model
from google.cloud import aiplatform
DEPLOY_DISPLAY_NAME = "your_text_model"
DEPLOY_COMPUTE = "n1-standard-4"
endpoint = uploaded_model.deploy(
deployed_model_display_name=DEPLOY_DISPLAY_NAME,
machine_type=DEPLOY_COMPUTE,
accelerator_type=None,
accelerator_count=0,
)Once the deployment is complete, you get your explained predictions using the explain method.
example = ["The movie was cool. The animation and the graphics were out of this world. I would recommend this movie."]
encoded_example = encoder(example)
instances = encoded_example.numpy().tolist()
example_label = 'positive'
result = endpoint.explain(instances)
prediction = result.predictions[0][0]
attributions = result.explanations[0].attributions[0]where attributions is an instance of Attribution class containing feature attributions to explain model predictions.
Finally, you can visualize the attributions for the text instance by highlighting text with attribution scores using the highlight and colorize functions of the official notebook.
As explained in the notebook, words with high positive attribution are highlighted in shades of green and words with negative attribution in shades of pink. Stronger shading corresponds to higher attribution values. Positive attributions can be interpreted as increase in probability of the predicted class (positive in this case) while negative attributions correspond to decrease in probability of the predicted class.
As alternative visualization you can plot feature attributions directly using barh in matplotlib. I used the following code.
import matplotlib.pyplot as plt
# Create a dictionary with words and feature attribution values
importance_dict = {k:v for k, v in zip(words, feature_attributions)}
# Sort by feature attributions
sorted_feature_names = sorted(
importance_dict, key=importance_dict.get)
# List sorted values
sorted_attr_values = [importance_dict[key] for key in sorted_feature_names]
num_features = len(sorted_feature_names)
if num_features > 0:
x_pos = list(range(num_features))
plt.barh(x_pos, sorted_attr_values)
plt.yticks(x_pos, sorted_feature_names)
plt.title('Feature attributions')
plt.ylabel('Feature names')
plt.xlabel('Attribution value')
plt.show()And here you have the feature attributions plot.
Notice how words like recommend or cool increases the probability of the predicted class (positive) with their positive (Shapley sampled) attribution values rather than negative attribution values of words like out or would.
Summary
In this blog post of the Vertex AI Q&A series, I addressed the following question:
Can I explain a text model with Vertex AI?
Yes, you can. With a text classification model, you can use the Sampled Shapley method in Vertex Explainable AI to get feature-based explanations indicating how much each word or token of a given input contributes to the predicted class.
Notice that, in this blog, I presented the simplest scenario you can have. I didn’t cover how to improve feature-based explanations or how Explainable AI SDK can help to define explanation metadata, get explanations and visualize explanation results. See the official Vertex AI documentation for further guidance. Also, I tried to use the Integrated gradients method to generate explanations without success. Further explorations would be required.
In the meantime, I hope you found the article interesting. If so, clap or leave comments. And feel free to reach me on LinkedIn or Twitter for further discussion or if you have questions.
Thanks Eric Dong, Sheeraz Ahmad, Ann Farmer and Romin Irani for feedback and the great notebook content .
