attention-visualiser 0.1.0__tar.gz

attention_visualiser-0.1.0/.github/workflows/tests.yml

@@ -0,0 +1,38 @@
+# This workflow will install Python dependencies, run tests and lint with a single version of Python
+# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python
+
+name: tests
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+permissions:
+  contents: read
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python 3.12
+      uses: actions/setup-python@v3
+      with:
+        python-version: "3.12"
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
+      with:
+        version: "0.6.16"
+
+    - name: Install dependencies
+      run: |
+        uv sync
+    - name: Test with pytest
+      run: |
+        source .venv/bin/activate
+        uv run pytest --cov --cov-report=xml
+    - name: Upload coverage reports to Codecov
+      uses: codecov/codecov-action@v5

attention_visualiser-0.1.0/.gitignore

@@ -0,0 +1,23 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+.mypy_cache/
+
+# Virtual environments
+.venv
+
+# IDE and Editor specific files
+.idea/
+.vscode/
+.code/
+
+# pytest
+.coverage
+coverage.xml
+
+# macos
+.DS_Store

attention_visualiser-0.1.0/.pre-commit-config.yaml

@@ -0,0 +1,10 @@
+repos:
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  # Ruff version.
+  rev: v0.11.6
+  hooks:
+    # Run the linter.
+    - id: ruff
+      args: [ --fix ]
+    # Run the formatter.
+    - id: ruff-format

attention_visualiser-0.1.0/.python-version

@@ -0,0 +1 @@
+3.12

attention_visualiser-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Shawon Ashraf
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

attention_visualiser-0.1.0/PKG-INFO

@@ -0,0 +1,60 @@
+Metadata-Version: 2.4
+Name: attention-visualiser
+Version: 0.1.0
+Summary: a module to visualise attention layer activations from transformer based models from huggingface
+License-File: LICENSE
+Requires-Python: >=3.11
+Requires-Dist: einops>=0.8.1
+Requires-Dist: loguru>=0.7.3
+Requires-Dist: seaborn>=0.13.2
+Requires-Dist: transformers>=4.51.3
+Description-Content-Type: text/markdown
+
+# attention-visualiser
+
+a module to visualise attention layer activations from transformer based models from huggingface
+
+## installation
+
+```bash
+pip install git+https://codeberg.org/rashomon/attention-visualiser
+```
+
+## usage
+
+```python
+from attention_visualiser import AttentionVisualiser
+from transformers import AutoModel, AutoTokenizer
+
+# visualising activations from gpt
+model_name = "openai-community/openai-gpt"
+
+model = AutoModel.from_pretrained(model_name)
+model.eval()
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+text = "Look on my Works, ye Mighty, and despair!"
+encoded_inputs = tokenizer.encode_plus(text, truncation=True, return_tensors="pt")
+
+visualiser = AttentionVisualiser(model, tokenizer)
+
+# visualise from the first attn layer
+visualiser.visualise_attn_layer(0, encoded_inputs)
+
+```
+
+
+## local dev
+
+```bash
+# env setup
+
+uv sync
+source .venv/bin/activate
+
+# tests
+uv run pytest
+
+# tests with coverage
+uv run pytest --cov --cov-report=xml
+```

attention_visualiser-0.1.0/README.md

@@ -0,0 +1,48 @@
+# attention-visualiser
+
+a module to visualise attention layer activations from transformer based models from huggingface
+
+## installation
+
+```bash
+pip install git+https://codeberg.org/rashomon/attention-visualiser
+```
+
+## usage
+
+```python
+from attention_visualiser import AttentionVisualiser
+from transformers import AutoModel, AutoTokenizer
+
+# visualising activations from gpt
+model_name = "openai-community/openai-gpt"
+
+model = AutoModel.from_pretrained(model_name)
+model.eval()
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+text = "Look on my Works, ye Mighty, and despair!"
+encoded_inputs = tokenizer.encode_plus(text, truncation=True, return_tensors="pt")
+
+visualiser = AttentionVisualiser(model, tokenizer)
+
+# visualise from the first attn layer
+visualiser.visualise_attn_layer(0, encoded_inputs)
+
+```
+
+
+## local dev
+
+```bash
+# env setup
+
+uv sync
+source .venv/bin/activate
+
+# tests
+uv run pytest
+
+# tests with coverage
+uv run pytest --cov --cov-report=xml
+```

attention_visualiser-0.1.0/attention_visualiser/__init__.py

@@ -0,0 +1,3 @@
+from .pt import AttentionVisualiserPytorch as AttentionVisualiser
+
+__all__ = ["AttentionVisualiser"]

attention_visualiser-0.1.0/attention_visualiser/base.py

@@ -0,0 +1,160 @@
+import torch
+import matplotlib.pyplot as plt
+import seaborn as sns
+from transformers import AutoTokenizer, AutoModel, FlaxAutoModel
+from transformers import BatchEncoding
+from loguru import logger
+from einops import rearrange
+from abc import ABC, abstractmethod
+import numpy as np
+from typing import Optional
+
+
+class BaseAttentionVisualiser(ABC):
+    """Base abstract class for visualizing attention weights in transformer models.
+
+    This class provides the foundation for visualizing attention weights from
+    different transformer model implementations. Concrete subclasses must implement
+    methods for computing attention values and processing attention vectors.
+
+    Attributes:
+        model: A transformer model from the Hugging Face library
+        tokenizer: A tokenizer matching the model
+        config: Dictionary containing visualization configuration parameters
+    """
+
+    def __init__(
+        self,
+        model: AutoModel | FlaxAutoModel,
+        tokenizer: AutoTokenizer,
+        config: Optional[dict] = None,
+    ) -> None:
+        """Initialize the attention visualizer with a model and tokenizer.
+
+        Args:
+            model: A transformer model from Hugging Face (PyTorch or Flax)
+            tokenizer: A tokenizer matching the model
+            config: Optional dictionary with visualization parameters
+                   Default parameters include:
+                   - figsize: Tuple specifying figure dimensions
+                   - cmap: Colormap for the heatmap
+                   - annot: Whether to annotate heatmap cells with values
+                   - xlabel: Label for x-axis
+                   - ylabel: Label for y-axis
+        """
+        self.model = model
+        self.tokenizer = tokenizer
+
+        logger.info(f"Model config: {self.model.config}")  # type: ignore
+
+        if not config:
+            self.config = {
+                "figsize": (15, 15),
+                "cmap": "viridis",
+                "annot": True,
+                "xlabel": "",
+                "ylabel": "",
+            }
+            logger.info(f"Setting default visualiser config: {self.config}")
+        else:
+            logger.info(f"Visualiser config: {config}")
+            self.config = config
+
+        # a cache for storing already computed attention vectors
+        # these need to be updated by the `compute_attentions`
+        # method
+        self.current_input = None
+        self.cache = None
+
+    def id_to_tokens(self, encoded_input: BatchEncoding) -> list[str]:
+        """Convert token IDs to readable token strings.
+
+        Args:
+            encoded_input: The encoded input from the tokenizer
+
+        Returns:
+            List of token strings corresponding to the input IDs
+        """
+        tokens = self.tokenizer.convert_ids_to_tokens(encoded_input["input_ids"][0])  # type: ignore
+        return tokens
+
+    @abstractmethod
+    def compute_attentions(self, encoded_input: BatchEncoding) -> tuple:
+        """Compute attention weights for the given input.
+
+        This method must be implemented by concrete subclasses to compute
+        attention weights specific to the model implementation.
+
+        Args:
+            encoded_input: The encoded input from the tokenizer
+
+        Returns:
+            A tuple containing attention weights
+        """
+        pass
+
+    @abstractmethod
+    def get_attention_vector_mean(
+        self, attention: torch.Tensor, axis: int = 0
+    ) -> np.ndarray:
+        """Calculate mean of attention vectors along specified axis.
+
+        This method must be implemented by concrete subclasses to handle
+        either PyTorch or JAX tensors appropriately.
+
+        Args:
+            attention: Attention tensor from the model
+            axis: Axis along which to compute the mean (default: 0)
+
+        Returns:
+            NumPy array of mean attention values
+        """
+        pass
+
+    def visualise_attn_layer(self, idx: int, encoded_input: BatchEncoding) -> None:
+        """Visualize attention weights for a specific layer.
+
+        Creates a heatmap visualization of the attention weights for the specified
+        layer index.
+
+        Args:
+            idx: Index of the attention layer to visualize.
+                 Negative indices count from the end (-1 is the last layer).
+            encoded_input: The encoded input from the tokenizer
+
+        Raises:
+            AssertionError: If idx is outside the range of available attention layers
+        """
+        tokens = self.id_to_tokens(encoded_input)
+
+        attentions = self.compute_attentions(encoded_input)
+        n_attns = len(attentions)
+
+        # idx must no exceed attn_heads
+        assert idx < n_attns, (
+            f"index must be less than the number of attention outputs in the model, which is: {n_attns}"
+        )
+
+        # setting idx = -1 will get the last attention layer activations but
+        # the plot title will also show -1
+        if idx < 0:
+            idx = n_attns + idx
+
+        # get rid of the additional dimension since single input
+        attention = rearrange(attentions[idx], "1 a b c -> a b c")
+        # take mean over dim 0
+        attention = self.get_attention_vector_mean(attention)
+
+        plt.figure(figsize=self.config.get("figsize"))
+        sns.heatmap(
+            attention,
+            cmap=self.config.get("cmap"),
+            annot=self.config.get("annot"),
+            xticklabels=tokens,
+            yticklabels=tokens,
+        )
+
+        plt.title(f"Attention Weights for Layer idx: {idx}")
+        plt.xlabel(self.config.get("xlabel"))  # type: ignore
+        plt.ylabel(self.config.get("ylabel"))  # type: ignore
+        plt.show()

attention_visualiser-0.1.0/attention_visualiser/pt.py

@@ -0,0 +1,76 @@
+import torch
+from transformers import AutoTokenizer, AutoModel
+from transformers import BatchEncoding
+from attention_visualiser.base import BaseAttentionVisualiser
+import numpy as np
+from typing import Optional
+
+
+class AttentionVisualiserPytorch(BaseAttentionVisualiser):
+    """Attention visualizer for PyTorch-based transformer models.
+
+    This class implements the abstract methods from BaseAttentionVisualiser
+    specifically for models implemented in PyTorch. It handles the extraction
+    and processing of attention weights from PyTorch transformer models.
+
+    Attributes:
+        model: A PyTorch-based transformer model from Hugging Face
+        tokenizer: A tokenizer matching the model
+        config: Dictionary containing visualization configuration parameters
+    """
+
+    def __init__(
+        self, model: AutoModel, tokenizer: AutoTokenizer, config: Optional[dict] = None
+    ) -> None:
+        """Initialize the PyTorch-specific attention visualizer.
+
+        Args:
+            model: A PyTorch-based transformer model from Hugging Face
+            tokenizer: A tokenizer matching the model
+            config: Optional dictionary with visualization parameters
+        """
+        super().__init__(model, tokenizer, config)
+
+    def compute_attentions(self, encoded_input: BatchEncoding) -> tuple:
+        """Compute attention weights for the given input using a PyTorch model.
+
+        Runs the PyTorch model in inference mode with output_attentions flag set to True
+        and extracts the attention weights from the model output.
+
+        Args:
+            encoded_input: The encoded input from the tokenizer
+
+        Returns:
+            A tuple containing attention weights from all layers of the model
+        """
+        if encoded_input == self.current_input:
+            # return from cache
+            return self.cache
+
+        # else recompute
+        with torch.no_grad():
+            output = self.model(**encoded_input, output_attentions=True)  # type: ignore
+
+        attentions = output.attentions
+
+        # update cache and current input
+        self.current_input = encoded_input
+        self.cache = attentions
+
+        return attentions
+
+    def get_attention_vector_mean(
+        self, attention: torch.Tensor, axis: int = 0
+    ) -> np.ndarray:
+        """Calculate mean of PyTorch attention vectors along specified axis.
+
+        Computes the mean of the attention tensor and converts it to a NumPy array.
+
+        Args:
+            attention: PyTorch tensor containing attention weights
+            axis: Axis along which to compute the mean (default: 0)
+
+        Returns:
+            NumPy array of mean attention values
+        """
+        return torch.mean(attention, dim=axis).detach().cpu().numpy()

attention_visualiser-0.1.0/codecov.yml

@@ -0,0 +1,5 @@
+coverage:
+  status:
+    project:
+      default:
+        target: 80%

attention_visualiser-0.1.0/main.py

@@ -0,0 +1,18 @@
+from attention_visualiser import AttentionVisualiser
+from transformers import AutoModel, AutoTokenizer
+
+if __name__ == "__main__":
+    # visualising activations from gpt
+    model_name = "openai-community/openai-gpt"
+
+    model = AutoModel.from_pretrained(model_name)
+    model.eval()
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    text = "Look on my Works, ye Mighty, and despair!"
+    encoded_inputs = tokenizer.encode_plus(text, truncation=True, return_tensors="pt")
+
+    visualiser = AttentionVisualiser(model, tokenizer)
+
+    # visualise from the first attn layer
+    visualiser.visualise_attn_layer(0, encoded_inputs)

attention_visualiser-0.1.0/pyproject.toml

@@ -0,0 +1,28 @@
+[project]
+name = "attention-visualiser"
+version = "0.1.0"
+description = "a module to visualise attention layer activations from transformer based models from huggingface"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "einops>=0.8.1",
+    "loguru>=0.7.3",
+    "seaborn>=0.13.2",
+    "transformers>=4.51.3",
+]
+
+[dependency-groups]
+dev = [
+    "pre-commit>=4.2.0",
+    "pytest>=8.3.5",
+    "pytest-cov>=6.1.1",
+    "ruff>=0.11.6",
+    "torch>=2.7.0",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["attention-visualiser"]

attention_visualiser-0.1.0/tests/test_pt.py

@@ -0,0 +1,167 @@
+import pytest
+import numpy as np
+import torch
+from attention_visualiser import AttentionVisualiser
+
+
+class TestAttentionVisualiser:
+    @pytest.fixture
+    def mock_model(self):
+        class MockConfig:
+            num_attention_heads = 12
+            num_hidden_layers = 12
+
+        class MockModel:
+            def __init__(self):
+                self.config = MockConfig()
+                self.call_count = 0
+
+            def __call__(self, **kwargs):
+                self.call_count += 1
+                self.last_kwargs = kwargs
+                return self.output  # type: ignore
+
+        model = MockModel()
+        return model
+
+    @pytest.fixture
+    def mock_tokenizer(self):
+        class MockTokenizer:
+            def __init__(self):
+                self.call_count = 0
+
+            def convert_ids_to_tokens(self, ids):
+                self.call_count += 1
+                return ["[CLS]", "Hello", "world", "[SEP]"]
+
+        return MockTokenizer()
+
+    @pytest.fixture
+    def visualiser(self, mock_model, mock_tokenizer):
+        return AttentionVisualiser(
+            model=mock_model,
+            tokenizer=mock_tokenizer,
+        )
+
+    @pytest.fixture
+    def mock_encoded_input(self):
+        class MockEncodedInput:
+            def __init__(self):
+                self.data = {"input_ids": torch.tensor([[101, 7592, 2088, 102]])}
+
+            def __getitem__(self, key):
+                return self.data.get(key)
+
+            def __eq__(self, other):
+                return id(self) == id(other)
+
+            # Add these methods to support ** unpacking
+            def keys(self):
+                return self.data.keys()
+
+            def __iter__(self):
+                return iter(self.data)
+
+            def get(self, key, default=None):
+                return self.data.get(key, default)
+
+        return MockEncodedInput()
+
+    @pytest.fixture
+    def mock_attention_data(self):
+        # Create mock attention data: (batch_size, num_heads, seq_len, seq_len)
+        # Shape: (1, 12, 4, 4) - 1 batch, 12 attention heads, sequence length 4
+        return torch.ones((1, 12, 4, 4)) * 0.25
+
+    def test_init(self, mock_model, mock_tokenizer):
+        """Test initialization of AttentionVisualiser."""
+        visualiser = AttentionVisualiser(mock_model, mock_tokenizer)
+
+        assert visualiser.model == mock_model
+        assert visualiser.tokenizer == mock_tokenizer
+        assert visualiser.config is not None
+        assert visualiser.current_input is None
+        assert visualiser.cache is None
+
+    def test_id_to_tokens(self, visualiser, mock_encoded_input):
+        """Test the id_to_tokens method."""
+        tokens = visualiser.id_to_tokens(mock_encoded_input)
+
+        assert visualiser.tokenizer.call_count == 1
+        assert tokens == ["[CLS]", "Hello", "world", "[SEP]"]
+
+    def test_compute_attentions_new_input(
+        self, visualiser, mock_encoded_input, mock_attention_data
+    ):
+        """Test compute_attentions with new input."""
+
+        # Setup mock return value for model call
+        class MockOutput:
+            def __init__(self, attn_data):
+                self.attentions = attn_data
+
+        # Add the output attribute to the model
+        output = MockOutput(mock_attention_data)
+        visualiser.model.output = output
+
+        # Initialize call_count if not present
+        if not hasattr(visualiser.model, "call_count"):
+            visualiser.model.call_count = 0
+        if not hasattr(visualiser.model, "last_kwargs"):
+            visualiser.model.last_kwargs = {}
+
+        # Call the method
+        attentions = visualiser.compute_attentions(mock_encoded_input)
+
+        # Verify the model was called
+        assert visualiser.model.call_count > 0
+        assert visualiser.model.last_kwargs.get("output_attentions")
+
+        # Verify the return value and cache update
+        assert attentions is mock_attention_data
+        assert visualiser.current_input == mock_encoded_input
+        assert visualiser.cache is mock_attention_data
+
+    def test_compute_attentions_cached_input(
+        self, visualiser, mock_encoded_input, mock_attention_data
+    ):
+        """Test compute_attentions with cached input."""
+        # Setup cache
+        visualiser.current_input = mock_encoded_input
+        visualiser.cache = mock_attention_data
+        visualiser.model.call_count = 0
+
+        # Call the method
+        attentions = visualiser.compute_attentions(mock_encoded_input)
+
+        # Verify model wasn't called and cached data was returned
+        assert visualiser.model.call_count == 0
+        assert attentions is mock_attention_data
+
+    def test_get_attention_vector_mean(self, visualiser):
+        """Test get_attention_vector_mean method."""
+        # Create test data
+        test_data = torch.tensor([[[0.1, 0.2], [0.3, 0.4]], [[0.5, 0.6], [0.7, 0.8]]])
+
+        # Expected result: mean along axis 0
+        expected = np.array([[0.3, 0.4], [0.5, 0.6]])
+
+        # Call the method
+        result = visualiser.get_attention_vector_mean(test_data)
+
+        # Verify
+        np.testing.assert_allclose(result, expected)
+
+    def test_get_attention_vector_mean_different_axis(self, visualiser):
+        """Test get_attention_vector_mean with a different axis."""
+        # Create test data
+        test_data = torch.tensor([[[0.1, 0.2], [0.3, 0.4]], [[0.5, 0.6], [0.7, 0.8]]])
+
+        # Expected result: mean along axis 1
+        expected = np.array([[0.2, 0.3], [0.6, 0.7]])
+
+        # Call the method
+        result = visualiser.get_attention_vector_mean(test_data, axis=1)
+
+        # Verify
+        np.testing.assert_allclose(result, expected)