Desafío: Graficar Mapas de Calor de Atención

Visualizar los pesos de atención con un mapa de calor ayuda a interpretar cómo un modelo transformer distribuye su enfoque a lo largo de una oración. Se utiliza matplotlib para graficar un mapa de calor donde tanto el eje x como el eje y representan los tokens de la oración. Cada celda en el mapa de calor muestra el peso de atención entre un par de tokens: la fila corresponde al token de consulta y la columna corresponde al token clave.

Comienza dividiendo la oración en tokens:

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

Luego, define la matriz de atención como un array de NumPy. Cada valor representa el peso de atención de un token a otro:

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 graficar el mapa de calor, utiliza el siguiente 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()

Colores más claros u oscuros indican valores de atención más altos o bajos, dependiendo del mapa de colores. Al observar el mapa de calor, se puede ver a qué palabras el modelo presta más atención al procesar cada token. Por ejemplo, si la celda en la fila focus y la columna important es clara, el modelo conecta fuertemente focus con important en su representación interna. Esta visualización ayuda a comprender qué partes de la oración de entrada se influyen entre sí y es útil para diagnosticar o interpretar el comportamiento del modelo en tareas de procesamiento de lenguaje natural.

Ahora, ejecuta el código para ver el mapa de calor resultante y luego escribe tu primer gráfico de visualización.