"""Creation of attention heat maps

This script loads a pretrained BERT model, extracts its used attention for the embedding of a sentence specified
via command line argument and then creates a heat map for the amount of attention each token in the sentence receives.
"""
import torch
from transformers import BertConfig, BertTokenizer, BertModel
import seaborn as sns
import matplotlib.pyplot as plt
import sys

# load model and tokenizer
config = BertConfig.from_pretrained('../models/stsb-bert-large/0_BERT')
config.output_attentions = True
model = BertModel.from_pretrained('../models/stsb-bert-large/0_BERT', config=config).to('cpu')
tokenizer = BertTokenizer.from_pretrained('../models/stsb-bert-large/0_BERT')

# tokenize text given via command line argument
text = sys.argv[1]
tok = tokenizer.tokenize(text)
pos = 1  # change this to change which token you would like to see the attention for

# extract attention and change format for plotting
ids = torch.tensor(tokenizer.convert_tokens_to_ids(tok)).unsqueeze(0).to('cpu')
with torch.no_grad():
    output = model(ids)
attentions = torch.cat(output[2]).to('cpu')
attentions = attentions.permute(2, 1, 0, 3)
layers = len(attentions[0][0])
heads = len(attentions[0])
seqlen = len(attentions)

attentions_pos = attentions[pos]

# plot attention
avg_attention = attentions_pos.mean(dim=0)
sns.heatmap(avg_attention, vmin=0, vmax=1, xticklabels=tok)

plt.show()
