Package: bertviz
BertVizVisualize Attention in NLP ModelsQuick Tour•Getting Started•Colab Tutorial•Blog•Paper•CitationBertViz is an interactive tool for visualizing attention inTransformerlanguage models such as BERT, GPT2, or T5. It can be run inside a Jupyter or Colab
notebook through a simple Python API that supports mostHuggingface models. BertViz extends theTensor2Tensor visualization toolbyLlion Jones, providing multiple views that each offer a unique lens into the attention mechanism.For updates on BertViz and related projects, feel free to follow me onTwitter.🚀 Quick TourHead ViewThehead viewvisualizes attention for one or more attention heads in the same
layer. It is based on the excellentTensor2Tensor visualization toolbyLlion Jones.🕹 Try out the head view in theInteractive Colab Tutorial(all visualizations pre-loaded).Model ViewThemodel viewshows a bird's-eye view of attention across all layers and heads.🕹 Try out the model view in theInteractive Colab Tutorial(all visualizations pre-loaded).Neuron ViewTheneuron viewvisualizes individual neurons in the query and key vectors and shows how they are used to compute attention.🕹 Try out the neuron view in theInteractive Colab Tutorial(all visualizations pre-loaded).⚡️ Getting StartedRunning BertViz in a Jupyter NotebookFrom the command line:pipinstallbertvizYou must also have Jupyter Notebook and ipywidgets installed:pipinstalljupyterlab
pipinstallipywidgets(If you run into any issues installing Jupyter or ipywidgets, consult the documentationhereandhere.)To create a new Jupyter notebook, simply run:jupyternotebookThen clickNewand selectPython 3 (ipykernel)if prompted.Running BertViz in ColabTo run inColab, simply add the following cell at the beginning of your Colab notebook:!pip install bertvizSample codeRun the following code to load thextremedistil-l12-h384-uncasedmodel and display it in the model view:fromtransformersimportAutoTokenizer,AutoModel,utilsfrombertvizimportmodel_viewutils.logging.set_verbosity_error()# Suppress standard warningsmodel_name="microsoft/xtremedistil-l12-h384-uncased"# Find popular HuggingFace models here: https://huggingface.co/modelsinput_text="The cat sat on the mat"model=AutoModel.from_pretrained(model_name,output_attentions=True)# Configure model to return attention valuestokenizer=AutoTokenizer.from_pretrained(model_name)inputs=tokenizer.encode(input_text,return_tensors='pt')# Tokenize input textoutputs=model(inputs)# Run modelattention=outputs[-1]# Retrieve attention from model outputstokens=tokenizer.convert_ids_to_tokens(inputs[0])# Convert input ids to token stringsmodel_view(attention,tokens)# Display model viewThe visualization may take a few seconds to load. Feel free to experiment with different input texts andmodels.
SeeDocumentationfor additional use cases and examples, e.g., encoder-decoder models.Running sample notebooksYou may also run any of the samplenotebooksincluded with BertViz:gitclone--depth1git@github.com:jessevig/bertviz.gitcdbertviz/notebooks
jupyternotebook🕹 Interactive TutorialCheck out theInteractive Colab Tutorialto learn more about BertViz and try out the tool.Note: all visualizations are pre-loaded, so there is no need to execute any cells.📖 DocumentationTable of contentsSelf-attention models (BERT, GPT-2, etc.)Head and Model ViewsNeuron ViewEncoder-decoder models (BART, T5, etc.)Installing from sourceAdditional optionsDark / light modeFiltering layersSetting default layer/head(s)Visualizing sentence pairsObtain HTML representationsNon-Huggingface modelsSelf-attention models (BERT, GPT-2, etc.)Head and Model ViewsFirst load a Huggingface model, either a pre-trained model as shown below, or your own fine-tuned model.
Be sure to setoutput_attentions=True.fromtransformersimportAutoTokenizer,AutoModel,utilsutils.logging.set_verbosity_error()# Suppress standard warningstokenizer=AutoTokenizer.from_pretrained("bert-base-uncased")model=AutoModel.from_pretrained("bert-base-uncased",output_attentions=True)Then prepare inputs and compute attention:inputs=tokenizer.encode("The cat sat on the mat",return_tensors='pt')outputs=model(inputs)attention=outputs[-1]# Output includes attention weights when output_attentions=Truetokens=tokenizer.convert_ids_to_tokens(inputs[0])Finally, display the attention weights using thehead_viewormodel_viewfunctions:frombertvizimporthead_viewhead_view(attention,tokens)Examples: DistilBERT (Model View Notebook,Head View Notebook)For full API, please refer to the source code for thehead viewormodel view.Neuron ViewThe neuron view is invoked differently than the head view or model view, due to requiring access to the model's
query/key vectors, which are not returned through the Huggingface API. It is currently limited to custom versions of BERT, GPT-2, and
RoBERTa included with BertViz.# Import specialized versions of models (that return query/key vectors)frombertviz.transformers_neuron_viewimportBertModel,BertTokenizerfrombertviz.neuron_viewimportshowmodel_type='bert'model_version='bert-base-uncased'do_lower_case=Truesentence_a="The cat sat on the mat"sentence_b="The cat lay on the rug"model=BertModel.from_pretrained(model_version,output_attentions=True)tokenizer=BertTokenizer.from_pretrained(model_version,do_lower_case=do_lower_case)show(model,model_type,tokenizer,sentence_a,sentence_b,layer=2,head=0)Examples:
BERT (Notebook,Colab) •
GPT-2 (Notebook,Colab) •
RoBERTa
(Notebook)For full API, please refer to thesource.Encoder-decoder models (BART, T5, etc.)The head view and model view both support encoder-decoder models.First, load an encoder-decoder model:fromtransformersimportAutoTokenizer,AutoModeltokenizer=AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-de")model=AutoModel.from_pretrained("Helsinki-NLP/opus-mt-en-de",output_attentions=True)Then prepare the inputs and compute attention:encoder_input_ids=tokenizer("She sees the small elephant.",return_tensors="pt",add_special_tokens=True).input_idsdecoder_input_ids=tokenizer("Sie sieht den kleinen Elefanten.",return_tensors="pt",add_special_tokens=True).input_idsoutputs=model(input_ids=encoder_input_ids,decoder_input_ids=decoder_input_ids)encoder_text=tokenizer.convert_ids_to_tokens(encoder_input_ids[0])decoder_text=tokenizer.convert_ids_to_tokens(decoder_input_ids[0])Finally, display the visualization using eitherhead_viewormodel_view.frombertvizimportmodel_viewmodel_view(encoder_attention=outputs.encoder_attentions,decoder_attention=outputs.decoder_attentions,cross_attention=outputs.cross_attentions,encoder_tokens=encoder_text,decoder_tokens=decoder_text)You may selectEncoder,Decoder, orCrossattention from the drop-down in the upper left corner of the visualization.Examples: MarianMT (Notebook) • BART (Notebook)For full API, please refer to the source code for thehead viewormodel view.Installing from sourcegitclonehttps://github.com/jessevig/bertviz.gitcdbertviz
pythonsetup.pydevelopAdditional optionsDark / light modeThe model view and neuron view support dark (default) and light modes. You may set the mode using
thedisplay_modeparameter:model_view(attention,tokens,display_mode="light")Filtering layersTo improve the responsiveness of the tool when visualizing larger models or inputs, you may set theinclude_layersparameter to restrict the visualization to a subset of layers (zero-indexed). This option is available in the head view and model
view.Example:Render model view with only layers 5 and 6 displayedmodel_view(attention,tokens,include_layers=[5,6])For the model view, you may also restrict the visualization to a subset of attention heads (zero-indexed) by setting theinclude_headsparameter.Setting default layer/head(s)In the head view, you may choose a specificlayerand collection ofheadsas the default selection when the
visualization first renders. Note: this is different from theinclude_heads/include_layersparameter (above), which
removes layers and heads from the visualization completely.Example:Render head view with layer 2 and heads 3 and 5 pre-selectedhead_view(attention,tokens,layer=2,heads=[3,5])You may also pre-select a specificlayerand singleheadfor the neuron view.Visualizing sentence pairsSome models, e.g. BERT, accept a pair of sentences as input. BertViz optionally supports a drop-down menu that allows
user to filter attention based on which sentence the tokens are in, e.g. only show attention between tokens in first
sentence and tokens in second sentence.Head and model viewsTo enable this feature when invoking thehead_viewormodel_viewfunctions, set
thesentence_b_startparameter to the start index of the second sentence. Note that the method for computing this
index will depend on the model.Example (BERT):frombertvizimporthead_viewfromtransformersimportAutoTokenizer,AutoModel,utilsutils.logging.set_verbosity_error()# Suppress standard warnings# NOTE: This code is model-specificmodel_version='bert-base-uncased'model=AutoModel.from_pretrained(model_version,output_attentions=True)tokenizer=AutoTokenizer.from_pretrained(model_version)sentence_a="the rabbit quickly hopped"sentence_b="The turtle slowly crawled"inputs=tokenizer.encode_plus(sentence_a,sentence_b,return_tensors='pt')input_ids=inputs['input_ids']token_type_ids=inputs['token_type_ids']# token type id is 0 for Sentence A and 1 for Sentence Battention=model(input_ids,token_type_ids=token_type_ids)[-1]sentence_b_start=token_type_ids[0].tolist().index(1)# Sentence B starts at first index of token type id 1token_ids=input_ids[0].tolist()# Batch index 0tokens=tokenizer.convert_ids_to_tokens(token_ids)head_view(attention,tokens,sentence_b_start)Neuron viewTo enable this option in the neuron view, simply set thesentence_aandsentence_bparameters inneuron_view.show().Obtain HTML representationsSupport to retrieve the generated HTML representations has been added to head_view, model_view and neuron_view.Setting the 'html_action' parameter to 'return' will make the function call return a single HTML Python object that can be further processed. Remember you can access the HTML source using the data attribute of a Python HTML object.The default behavior for 'html_action' is 'view', which will display the visualization but won't return the HTML object.This functionality is useful if you need to:Save the representation as an independent HTML file that can be accessed via web browserUse custom display methods as the ones needed in Databricks to visualize HTML objectsExample (head and model views):fromtransformersimportAutoTokenizer,AutoModel,utilsfrombertvizimporthead_viewutils.logging.set_verbosity_error()# Suppress standard warningstokenizer=AutoTokenizer.from_pretrained("bert-base-uncased")model=AutoModel.from_pretrained("bert-base-uncased",output_attentions=True)inputs=tokenizer.encode("The cat sat on the mat",return_tensors='pt')outputs=model(inputs)attention=outputs[-1]# Output includes attention weights when output_attentions=Truetokens=tokenizer.convert_ids_to_tokens(inputs[0])html_head_view=head_view(attention,tokens,html_action='return')withopen("PATH_TO_YOUR_FILE/head_view.html",'w')asfile:file.write(html_head_view.data)Example (neuron view):# Import specialized versions of models (that return query/key vectors)frombertviz.transformers_neuron_viewimportBertModel,BertTokenizerfrombertviz.neuron_viewimportshowmodel_type='bert'model_version='bert-base-uncased'do_lower_case=Truesentence_a="The cat sat on the mat"sentence_b="The cat lay on the rug"model=BertModel.from_pretrained(model_version,output_attentions=True)tokenizer=BertTokenizer.from_pretrained(model_version,do_lower_case=do_lower_case)html_neuron_view=show(model,model_type,tokenizer,sentence_a,sentence_b,layer=2,head=0,html_action='return')withopen("PATH_TO_YOUR_FILE/neuron_view.html",'w')asfile:file.write(html_neuron_view.data)Non-Huggingface modelsThe head view and model view may be used to
visualize self-attention for any standard Transformer model,
as long as the attention weights are available and follow the format specified inhead_viewandmodel_view(which is the format
returned from Huggingface models). In some case, Tensorflow checkpoints may be loaded as Huggingface models as described
in theHuggingface docs.⚠️ LimitationsToolThis tool is designed for shorter inputs and may run slowly if the input text is very long and/or the model is very large.
To mitigate this, you may wish to filter the layers displayed by setting theinclude_layersparameter, as describedabove.When running on Colab, some of the visualizations will fail (runtime disconnection) when the input text is long.  To mitigate this, you may wish to filter the layers displayed by setting theinclude_layersparameter, as describedabove.Theneuron viewonly supports the custom BERT, GPT-2, and RoBERTa models included with the tool. This view needs access to the query and key vectors,
which required modifying the model code (seetransformers_neuron_viewdirectory), which has only been done for these three models.Attention as "explanation"?Visualizing attention weights illuminates one type of architecture within the model but does not
necessarily provide a directexplanationfor predictions [1,2,3].If you wish to understand how the input text influences output predictions more directly, considersaliency methodsprovided
by tools such as theLanguage Interpretability ToolkitorEcco.🔬 PaperA Multiscale Visualization of Attention in the Transformer Model(ACL System Demonstrations 2019).Citation@inproceedings{vig-2019-multiscale,title="A Multiscale Visualization of Attention in the Transformer Model",author="Vig, Jesse",booktitle="Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics: System Demonstrations",month=jul,year="2019",address="Florence, Italy",publisher="Association for Computational Linguistics",url="https://www.aclweb.org/anthology/P19-3007",doi="10.18653/v1/P19-3007",pages="37--42",}AuthorsJesse Vig(homepage)🙏 AcknowledgmentsWe are grateful to the authors of the following projects, which are incorporated into this repo:https://github.com/tensorflow/tensor2tensorhttps://github.com/huggingface/pytorch-pretrained-BERTLicenseThis project is licensed under the Apache 2.0 License - see theLICENSEfile for details