BertViz

repository·master·Indexed 27 days ago

https://github.com/jessevig/bertviz

An 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.

Tokens
6.2K
Snippets
17
Records
23
Agent score
92%

What's inside BertViz

  1. Visualize encoder-decoder models (BART, T5, etc.)

    master

    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
    )
  2. Install BertViz

    master

    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 ipywidgets
  3. Obtain HTML representations of visualizations

    master
    You can retrieve the generated visualization as an HTML object instead of displaying it immediately by setting html_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.
  4. Run BertViz sample notebooks

    master

    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 lab
  5. Filter layers and heads in visualizations

    master
    To improve performance with large models or long inputs, restrict the visualization to specific layers or heads using include_layers and include_heads (zero-indexed). These parameters remove the excluded layers/heads from the view entirely.
  6. Use the Neuron View to visualize attention computations

    master

    The Neuron View allows you to visualize attention, query values, key values, and intermediate computations for a specific attention head.

    Interaction Guide:

    1. Filter Attention: Hover over any token on the left side of the visualization to filter attention from that token.
    2. Expand View: Click the plus (+) icon revealed when hovering over a token. This displays query vectors, key vectors, and intermediate computations for attention weights (blue = positive, orange = negative).
    3. Inspect Computations: Once expanded, hover over any other token on the left to see its associated attention computations.
    4. Navigate Layers/Heads: Use the Layer or Head drop-down menus to switch between different model layers or attention heads (using zero-indexed values).
    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)
  7. Use the Neuron View to visualize attention and vector computations

    master

    The Neuron View allows you to visualize attention, query values, and key values within a specific attention head.

    Interactivity Guide:

    • Filter Attention: Hover over any token on the left side of the visualization to filter attention from that specific token.
    • Expand Computations: Click the plus (+) icon revealed when hovering over a token to show query vectors, key vectors, and intermediate computations for attention weights.
      • Blue indicates positive values.
      • Orange indicates negative values.
    • Inspect Related Computations: Once in the expanded view, hover over any other token on the left to see its associated attention computations.
    • Switch Layers/Heads: Use the Layer or Head drop-down menus to change the model layer or head (note: these are zero-indexed).
    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)
  8. Use the Neuron View to visualize attention head computations

    master

    The Neuron View allows you to visualize attention, query values, and key values within a specific attention head.

    Interactivity Guide:

    • Filter Attention: Hover over any token on the left side of the visualization to filter attention from that specific token.
    • Expand Computations: Click the plus (+) icon revealed when hovering over a token. This displays the query vectors, key vectors, and intermediate computations for the attention weights (where <span style='color:blue'>blue</span> represents positive values and <span style='color:#ff6318'>orange</span> represents negative values).
    • Inspect Associated Computations: Once in the expanded view, hover over any other token on the left to see its associated attention computations.
    • Change Layer/Head: Use the Layer or Head drop-down menus to switch between different model layers or attention heads (using zero-indexed values).
    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')
  9. BertViz Limitations

    master

    Tool Performance

    • Designed for shorter inputs; large models or long text may run slowly or cause Colab runtime disconnections. Use include_layers to mitigate this.
    • Neuron View is limited to custom BERT, GPT-2, and RoBERTa models included in the repository.

    Interpretability

    • Visualizing attention weights is an architectural illumination and does not necessarily provide a direct explanation for model predictions.