attention_visualiser.base

  1import torch
  2import matplotlib.pyplot as plt
  3import seaborn as sns
  4from transformers import AutoTokenizer, AutoModel
  5from transformers import BatchEncoding
  6from loguru import logger
  7from einops import rearrange
  8from abc import ABC, abstractmethod
  9import numpy as np
 10from typing import Optional
 11
 12
 13class BaseAttentionVisualiser(ABC):
 14    """Base abstract class for visualizing attention weights in transformer models.
 15
 16    This class provides the foundation for visualizing attention weights from
 17    different transformer model implementations. Concrete subclasses must implement
 18    methods for computing attention values and processing attention vectors.
 19
 20    Attributes:
 21        model: A transformer model from the Hugging Face library
 22        tokenizer: A tokenizer matching the model
 23        config: Dictionary containing visualization configuration parameters
 24    """
 25
 26    def __init__(
 27        self,
 28        model: AutoModel,
 29        tokenizer: AutoTokenizer,
 30        config: Optional[dict] = None,
 31    ) -> None:
 32        """Initialize the attention visualizer with a model and tokenizer.
 33
 34        Args:
 35            model: A transformer model from Hugging Face (PyTorch or Flax)
 36            tokenizer: A tokenizer matching the model
 37            config: Optional dictionary with visualization parameters
 38                   Default parameters include:
 39                   - figsize: Tuple specifying figure dimensions
 40                   - cmap: Colormap for the heatmap
 41                   - annot: Whether to annotate heatmap cells with values
 42                   - xlabel: Label for x-axis
 43                   - ylabel: Label for y-axis
 44        """
 45        self.model = model
 46        self.tokenizer = tokenizer
 47
 48        logger.info(f"Model config: {self.model.config}")  # type: ignore
 49
 50        if not config:
 51            self.config = {
 52                "figsize": (15, 15),
 53                "cmap": "viridis",
 54                "annot": True,
 55                "xlabel": "",
 56                "ylabel": "",
 57            }
 58            logger.info(f"Setting default visualiser config: {self.config}")
 59        else:
 60            logger.info(f"Visualiser config: {config}")
 61            self.config = config
 62
 63        # a cache for storing already computed attention vectors
 64        # these need to be updated by the `compute_attentions`
 65        # method
 66        self.current_input = None
 67        self.cache = None
 68
 69    def id_to_tokens(self, encoded_input: BatchEncoding) -> list[str]:
 70        """Convert token IDs to readable token strings.
 71
 72        Args:
 73            encoded_input: The encoded input from the tokenizer
 74
 75        Returns:
 76            List of token strings corresponding to the input IDs
 77        """
 78        tokens = self.tokenizer.convert_ids_to_tokens(encoded_input["input_ids"][0])  # type: ignore
 79        return tokens
 80
 81    @abstractmethod
 82    def compute_attentions(self, encoded_input: BatchEncoding) -> tuple:
 83        """Compute attention weights for the given input.
 84
 85        This method must be implemented by concrete subclasses to compute
 86        attention weights specific to the model implementation.
 87
 88        Args:
 89            encoded_input: The encoded input from the tokenizer
 90
 91        Returns:
 92            A tuple containing attention weights
 93        """
 94        pass
 95
 96    @abstractmethod
 97    def get_attention_vector_mean(
 98        self, attention: torch.Tensor, axis: int = 0
 99    ) -> np.ndarray:
100        """Calculate mean of attention vectors along specified axis.
101
102        This method must be implemented by concrete subclasses to handle
103        either PyTorch or JAX tensors appropriately.
104
105        Args:
106            attention: Attention tensor from the model
107            axis: Axis along which to compute the mean (default: 0)
108
109        Returns:
110            NumPy array of mean attention values
111        """
112        pass
113
114    def visualise_attn_layer(self, idx: int, encoded_input: BatchEncoding) -> None:
115        """Visualize attention weights for a specific layer.
116
117        Creates a heatmap visualization of the attention weights for the specified
118        layer index.
119
120        Args:
121            idx: Index of the attention layer to visualize.
122                 Negative indices count from the end (-1 is the last layer).
123            encoded_input: The encoded input from the tokenizer
124
125        Raises:
126            AssertionError: If idx is outside the range of available attention layers
127        """
128        tokens = self.id_to_tokens(encoded_input)
129
130        attentions = self.compute_attentions(encoded_input)
131        n_attns = len(attentions)
132
133        # idx must no exceed attn_heads
134        assert idx < n_attns, (
135            f"index must be less than the number of attention outputs in the model, which is: {n_attns}"
136        )
137
138        # setting idx = -1 will get the last attention layer activations but
139        # the plot title will also show -1
140        if idx < 0:
141            idx = n_attns + idx
142
143        # get rid of the additional dimension since single input
144        attention = rearrange(attentions[idx], "1 a b c -> a b c")
145        # take mean over dim 0
146        attention = self.get_attention_vector_mean(attention)
147
148        plt.figure(figsize=self.config.get("figsize"))
149        sns.heatmap(
150            attention,
151            cmap=self.config.get("cmap"),
152            annot=self.config.get("annot"),
153            xticklabels=tokens,
154            yticklabels=tokens,
155        )
156
157        plt.title(f"Attention Weights for Layer idx: {idx}")
158        plt.xlabel(self.config.get("xlabel"))  # type: ignore
159        plt.ylabel(self.config.get("ylabel"))  # type: ignore
160        plt.show()