attention-visualiser

A module to visualise attention layer activations from transformer based models from huggingface

installation

pip install git+https://github.com/shawonashraf/attention-visualiser

usage

from attention_visualiser.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(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

# env setup

uv sync
source .venv/bin/activate

# tests
uv run pytest

# tests with coverage
uv run pytest --cov --cov-report=xml