# Extract attention weights for visualization
attention_matrix = weights[0, 0].detach().numpy()  # First head
words = ["the", "cat", "sat", "on", "the", "mat"]

print("Attention Matrix (first head):")
print("From -> To:")
for i, from_word in enumerate(words):
    print(f"{from_word:>4}: ", end="")
    for j, to_word in enumerate(words):
        print(f"{attention_matrix[i,j]:.2f} ", end="")
    print()