Visualize sentence pairs in Head and Model views
mastersentence_b_start parameter to the index where the second sentence begins.repository·master·Indexed 27 days ago
https://github.com/jessevig/bertvizAn interactive tool for visualizing attention mechanisms in Transformer language models. It supports most HuggingFace models and provides specialized views including head_view for self-attention, model_view for aggregated attention across layers and heads, and neuron_view for inspecting individual neurons. It supports encoder-decoder architectures (such as BART and T5) and provides options for dark/light display modes, layer/head filtering, and HTML export.
sentence_b_start parameter to the index where the second sentence begins.For encoder-decoder architectures, head_view and model_view support visualizing encoder, decoder, and cross-attention. You must provide the specific attention tensors for each component.
from bertviz import model_view
model_view(
encoder_attention=outputs.encoder_attentions,
decoder_attention=outputs.decoder_attentions,
cross_attention=outputs.cross_attentions,
encoder_tokens=encoder_text,
decoder_tokens=decoder_text
)To use BertViz in a local environment, install the package via pip. You must also have jupyterlab and ipywidgets installed to view the interactive visualizations in a notebook.
pip install bertviz
pip install jupyterlab
pip install ipywidgetshtml_action='return'. This is useful for saving visualizations to files or using them in environments like Databricks. Access the raw HTML string via the .data attribute of the returned object.You can explore the included sample notebooks by cloning the repository and launching Jupyter Lab from the notebooks directory:
git clone --depth 1 git@github.com:jessevig/bertviz.git
cd bertviz/notebooks
jupyter labTo use BertViz in a Google Colab notebook, run the following command in a cell at the beginning of your notebook:
!pip install bertvizSet the visual theme for model_view and neuron_view using the display_mode parameter. The default is dark.
model_view(attention, tokens, display_mode="light")include_layers and include_heads (zero-indexed). These parameters remove the excluded layers/heads from the view entirely.The Neuron View allows you to visualize attention, query values, key values, and intermediate computations for a specific attention head.
Interaction Guide:
from bertviz.transformers_neuron_view import RobertaModel, RobertaTokenizer
from bertviz.neuron_view import show
model_type = 'roberta'
model_version = 'roberta-base'
model = RobertaModel.from_pretrained(model_version)
tokenizer = RobertaTokenizer.from_pretrained(model_version)
sentence_a = "The cat sat on the mat"
sentence_b = "The cat lay on the rug"
show(model, model_type, tokenizer, sentence_a, sentence_b)The Neuron View allows you to visualize attention, query values, and key values within a specific attention head.
Interactivity Guide:
from bertviz.transformers_neuron_view import BertModel, BertTokenizer
from bertviz.neuron_view import show
model_type = 'bert'
model_version = 'bert-base-uncased'
do_lower_case = True
model = BertModel.from_pretrained(model_version)
tokenizer = BertTokenizer.from_pretrained(model_version, do_lower_case=do_lower_case)
sentence_a = "The cat sat on the mat"
sentence_b = "The cat lay on the rug"
show(model, model_type, tokenizer, sentence_a, sentence_b, display_mode='dark', layer=2, head=0)The Neuron View allows you to visualize attention, query values, and key values within a specific attention head.
Interactivity Guide:
from bertviz.transformers_neuron_view import GPT2Model, GPT2Tokenizer
from bertviz.neuron_view import show
model_type = 'gpt2'
model_version = 'gpt2'
model = GPT2Model.from_pretrained(model_version)
tokenizer = GPT2Tokenizer.from_pretrained(model_version)
text = "At the store, she bought apples, oranges, bananas,"
show(model, model_type, tokenizer, text, display_mode='dark')include_layers to mitigate this.