Using the attention to summarize a document about bertviz

[Hellisotherpeople]

I tried to implement this naively, but I may not fully understand the results of the attention mechanism

I simply summed the attention values in each layer and each head, and am trying to use the attention values that are highest in the summary (that is, highlight words which are attended to globally). This isn't as accurate as I'd hoped.

Can you see any places where I can improve my code? It's heavily based on your own

bert_version = 'bert-base-uncased'
model = BertModel.from_pretrained(bert_version)
tokenizer = BertTokenizer.from_pretrained(bert_version)
sentence_a = "Strawberries are bad for all humans. We will experience extinction from Strawberries soon if they are not stopped"
sentence_b = "Strawberries are yummy."
details_data = AttentionDetailsData(model, tokenizer)
tokens_a, tokens_b, queries, keys, atts = details_data.get_data(sentence_a, sentence_b)
print(tokens_a)
print(tokens_b)
total_tokens = tokens_a + tokens_b
print(len(tokens_a))
print(len(tokens_b))
print(len(tokens_a + tokens_b))
print(atts.shape)

head_summed = np.sum(atts, axis=2)
thing_summed = np.sum(head_summed, axis=1)
layer_summed = np.sum(head_summed, axis=0)
print("layer_summed")
compacted = layer_summed[0][1]
#np.sum(layer_summed, axis=1)
print(compacted.shape)
print(compacted)

for i in range(0, len(tokens_a)):
  if compacted[i] > 1:
    print(total_tokens[i] + ' | ' + str(compacted[i]))

[Hellisotherpeople]