encoder-decoder attention?

Dear Author,

Thank you very much for the visualization tool. It is super helpful.
Now, we need to visualize the encoder-decoder attention instead of the self-attention. For instance, while using the following code, we only can visualize the self-attention one. It would be great if you can give some clue on how to visualize the encoder-decoder attention of BART.

from bertviz import model_view
from bertviz import head_view
from transformers import BartForConditionalGeneration, BartTokenizer
model = BartForConditionalGeneration.from_pretrained("facebook/bart-large", force_bos_token_to_be_generated=True)
tokenizer = BartTokenizer.from_pretrained('facebook/bart-large')
sentence_a = "The cat sat on the mat"
sentence_b = "The cat lay on the rug"
inputs = tokenizer.encode_plus(sentence_a, sentence_b, return_tensors='pt', add_special_tokens=True)
input_ids = inputs['input_ids']
attention = model(input_ids,output_attentions=True)[2] ####
#print (attention.size())
input_id_list = input_ids[0].tolist() # Batch index 0
tokens = tokenizer.convert_ids_to_tokens(input_id_list)
model_view(attention, tokens)
head_view(attention, tokens)

Thank you! Best, Shirley

[jessevig]

Hi @xwuShirley, I’ve started a branch to visualize encoder-decoder attention: https://github.com/jessevig/bertviz/tree/encoder-decoder

I’ve added the encoder_decoder.ipynb notebook.

git clone  -b encoder-decoder [e-mail]:jessevig/bertviz.git
cd bertviz
jupyter notebook encoder_decoder.ipynb

Please let me know if that works for you. I’ll work on adding support for the model view and pushing a new version to pypi as well soon.

[chen-yifu]

Super helpful!! Thanks!