A Visual Analysis Tool to Explore Learned Representations in Transformers Models
by Ben Hoover, Hendrik Strobelt, Sebastian Gehrmann
from IBM Research and Harvard NLP
Link to pre-paper and demo: exbert.net
exBERT is a tool that enables users to explore the learned attention weights and contextual representations of different huggingface Transformer models. Input a sentence, and exBERT will pass the tokenized input through the specified model, visualizing attentions as curved lines and allowing the embeddings to be searchable across an annotated corpus, if available.
- Support for many of the common Transformer models (e.g., BERT, GPT2, DistilBERT, DistilGPT2, Roberta, ALBERT, )
- Toggle visibility of the attentions to the
[SEP]tokens that often assume the functionality of a
no-operationof a particular head.
- For masked language models, interactively mask particular tokens and observe how this affects the attention patterns
- Observe what a model would predict at (bidirectional models) or following (autoregressive models) a particular token
- View the attention patterns of single heads or any linear combination of heads
- Search for the contextual representation of any token at the output of any layer across a corpus annotated by that model
- Discover linguistic features (e.g., Part of Speech, Dependencies, Entities) learned by particular heads.
- exBERT does not visualize the attention between pairs of input sequences. All inputs are assumed to come from a single sequence.
- Has only been developed to support English
Install and Getting Started
Setting up the Environment
make env from the root directory
- From the root of this project, create a new conda directory with
conda env create -f environment.yml. This will create an environment named
- (optional) Intstall development dependencies with
conda env update -f environment-dev.yml
- Activate this environment with
conda activate exbert
pip install -e server/transformers
pip install -e server/spacyface
pip install -e server
- Install English support for spaCy
python -m spacy download en_core_web_sm
Attention Only Vis
Support for visualizing the attentions only is available out of the box. This minimal visualization is available with select deployed models on Huggingface's model page here.
Attention + Corpus
Significant preprocessing needs to be performed to allow corpus searching. Please see the instructions here
Starting the backend:
conda activate exbert python server/main.py
Because exBERT lazy-loads the large annotated corpus and models, the first call to search across the corpus will be slow.
If you want to make custom changes to the code, these are some hints to get you started.
Use as package
Some find it useful to expose the code inside
server for development in an environment like Jupyter Notebooks. From the root folder with the
exbert environment active:
conda env update -f environment-dev.yml pip install -e ./server
exbert environment should allow the server code to be accessible in any folder so long as there are no additional module name clashes in the environment.
Compiling the frontend
cd client/src npm install #installs all necessary node packages npm run build #This will create the static files living in `client/dist`.
Running a development environment
You can run a client server that automatically recompiles the frontend with
npm run watch. After making a change, you should be able to refresh the browser window to see your most recent changes.
Because the backend has to load in a lot of data for inference, we do not allow auto-backend refresh on every saved change in flask even though the framework supports it.
Uploading your own model locally
Uploading your own model consists of the following steps:
- Save your pretrained huggingface model according to the naming conventions specified in the
modeling_auto.pyof the original Transformers repo (as of v2.8):
The model class to instantiate is selected as the first pattern matching in the `pretrained_model_name_or_path` string (in the following order): - contains `t5`: T5Model (T5 model) - contains `distilbert`: DistilBertModel (DistilBERT model) - contains `albert`: AlbertModel (ALBERT model) - contains `camembert`: CamembertModel (CamemBERT model) - contains `xlm-roberta`: XLMRobertaModel (XLM-RoBERTa model) - contains `roberta`: RobertaModel (RoBERTa model) - contains `bert`: BertModel (Bert model) - contains `openai-gpt`: OpenAIGPTModel (OpenAI GPT model) - contains `gpt2`: GPT2Model (OpenAI GPT-2 model) - contains `transfo-xl`: TransfoXLModel (Transformer-XL model) - contains `xlnet`: XLNetModel (XLNet model) - contains `xlm`: XLMModel (XLM model) - contains `ctrl`: CTRLModel (Salesforce CTRL model)
The deployed server supports only BERT, RoBERTa, GPT2, and DistilBERT for context searching.
- Create the reference corpus. Warning: Depending on the number of layers and size of the hidden dimension in the model, this step could take many gigabytes on your computer to store the hidden representations and attentions at every layer.
Notes on SubRepo Usage
This project makes use of two public pip repositories (
spacyface), both of which needed modification as this project was being developed. The
git-subrepo tool was used to achieve this workflow with a forked repository of both transformers and spacyface. However, this introduces the following steps when setting up the environment:
- From the
pip install -e .
- Repeat for the
This project was inspired in part by the original BertViz by Jesse Vig.
- If you get a
No module named '_swigfaiss'error, check that
libompis installed on your system. If you are on a mac, this is as simple as
brew install libomp.