Visualizar os pesos de atenção com um mapa de calor auxilia na interpretação de como um modelo transformer distribui seu foco ao longo de uma sentença. Utiliza-se o matplotlib para plotar um mapa de calor onde tanto o eixo x quanto o eixo y representam os tokens da sentença. Cada célula no mapa de calor exibe o peso de atenção entre um par de tokens: a linha corresponde ao token de consulta (query) e a coluna ao token de chave (key).

Comece dividindo sua sentença em tokens:

sentence = "Transformers help models focus on important words."
tokens = sentence.split()

Em seguida, defina sua matriz de atenção como um array NumPy. Cada valor representa o peso de atenção de um token para outro:

import numpy as np
attention = np.array([
    [0.20, 0.10, 0.05, 0.10, 0.25, 0.10, 0.20],
    [0.05, 0.30, 0.10, 0.10, 0.15, 0.20, 0.10],
    [0.10, 0.15, 0.35, 0.10, 0.10, 0.10, 0.10],
    [0.10, 0.10, 0.10, 0.30, 0.10, 0.15, 0.15],
    [0.15, 0.10, 0.10, 0.10, 0.30, 0.10, 0.15],
    [0.10, 0.10, 0.15, 0.10, 0.10, 0.35, 0.10],
    [0.20, 0.15, 0.10, 0.10, 0.10, 0.10, 0.25],
])

Para plotar o mapa de calor, utilize o seguinte código:

import matplotlib.pyplot as plt

fig, ax = plt.subplots(figsize=(8, 6))
im = ax.imshow(attention, cmap="viridis")

ax.set_xticks(np.arange(len(tokens)))
ax.set_yticks(np.arange(len(tokens)))
ax.set_xticklabels(tokens, rotation=45, ha="right")
ax.set_yticklabels(tokens)

cbar = plt.colorbar(im, ax=ax)
cbar.set_label("Attention Weight")

ax.set_title("Attention Heatmap")
ax.set_xlabel("Key Tokens")
ax.set_ylabel("Query Tokens")

plt.tight_layout()
plt.show()

Cores mais claras ou escuras indicam valores de atenção mais altos ou mais baixos, dependendo do mapa de cores. Ao observar o mapa de calor, é possível identificar para quais palavras o modelo dedica mais atenção ao processar cada token. Por exemplo, se a célula na linha focus e coluna important estiver clara, o modelo estabelece uma forte conexão entre focus e important em sua representação interna. Essa visualização permite compreender quais partes da sentença de entrada influenciam umas às outras e é útil para diagnosticar ou interpretar o comportamento do modelo em tarefas de processamento de linguagem natural.

Agora, execute o código para visualizar o mapa de calor resultante e, em seguida, escreva seu primeiro gráfico de visualização.