Add attention visualization tool

[ArthurZucker]

What does this PR do?

TODOS

• Add some tests
• make it work on all models

cc @yonigozlan and @zucchini-nlp @molbap

This will be available as:

from transformers.utils.attention_visualizer import AttentionMaskVisualizer
visualizer = AttentionMaskVisualizer("meta-llama/Llama-3.2-3B-Instruct")
visualizer("A normal attention mask")

visualizer = AttentionMaskVisualizer("mistralai/Mistral-Small-24B-Instruct-2501")
visualizer("A normal attention mask with a long text to see how it is displayed, and if it is displayed correctly")

visualizer = AttentionMaskVisualizer("google/paligemma2-3b-mix-224")
visualizer("<img> You are an assistant.", suffix = "What is on the image?")

visualizer = AttentionMaskVisualizer("google/gemma-2b")
visualizer("You are an assistant. Make sure you print me") # we should have slidiing on non sliding side by side

visualizer = AttentionMaskVisualizer("google/gemma-3-27b-it")
visualizer("<img>You are an assistant. Make sure you print me") # we should have slidiing on non sliding side by side

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[jmayank23]

Hi, I was trying to visualize the attention mask for Gemma 3 but I get

AttributeError: 'Gemma3ForConditionalGeneration' object has no attribute 'visualize_attention_mask'

For reference, sharing relevant excerpts of the code:

from transformers import AutoProcessor, Gemma3ForConditionalGeneration, AutoTokenizer
from PIL import Image
import json
import torch

model_id = "google/gemma-3-27b-it"

model = Gemma3ForConditionalGeneration.from_pretrained(
    model_id, device_map="auto"
).eval()

processor = AutoProcessor.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)
.
.
.
inputs = processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt"
    ).to(model.device, dtype=torch.bfloat16)

model.visualize_attention_mask(tokenizer, inputs)

[ArthurZucker]

Hey ! Will work a bit to make it more easy to use! THanks 🤗

[molbap]

Doing a first review, so happy about this

[molbap]

Amazing