Attention Mechanism

import torch
import torch.nn as nn
import torch.nn.functional as F
import math


class ScaledDotProductAttention(nn.Module):
    """Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V."""

    def forward(self, q, k, v, mask=None):
        # q: [B, H, L_q, d_k]
        # k: [B, H, L_k, d_k]
        # v: [B, H, L_k, d_v]
        d_k = q.size(-1)

        # 1) Raw scores = Q · K^T / sqrt(d_k)
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)
        # scores: [B, H, L_q, L_k]

        # 2) Optional causal mask (decoder)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float("-inf"))

        # 3) Softmax → attention weights
        attn = F.softmax(scores, dim=-1)  # tổng theo L_k = 1

        # 4) Weighted sum of V
        out = torch.matmul(attn, v)
        # out: [B, H, L_q, d_v]

        return out, attn


class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

        self.attn = ScaledDotProductAttention()

    def forward(self, q, k, v, mask=None):
        B = q.size(0)

        # 1) Linear projection → split into heads
        #    [B, L, d_model] → [B, H, L, d_k]
        q = self.W_q(q).view(B, -1, self.n_heads, self.d_k).transpose(1, 2)
        k = self.W_k(k).view(B, -1, self.n_heads, self.d_k).transpose(1, 2)
        v = self.W_v(v).view(B, -1, self.n_heads, self.d_k).transpose(1, 2)

        # 2) Scaled dot-product attention trên tất cả các đầu song song
        out, attn = self.attn(q, k, v, mask=mask)

        # 3) Concat các đầu
        out = out.transpose(1, 2).contiguous().view(B, -1, self.d_model)

        # 4) Projection cuối
        return self.W_o(out), attn


# ─── Ví dụ nhanh trên câu "Tôi yêu Việt Nam quá"
if __name__ == "__main__":
    seq_len, d_model, n_heads = 5, 64, 4
    mha = MultiHeadAttention(d_model, n_heads)

    # Giả lập embedding của 5 token
    x = torch.randn(1, seq_len, d_model)  # [B=1, L=5, d=64]

    # Self-attention: Q, K, V đều = x
    context, attn = mha(x, x, x)
    print("Context shape:", context.shape)        # [1, 5, 64]
    print("Attention shape:", attn.shape)         # [1, 4, 5, 5]
    print("Sum per row (đầu 0):",
          attn[0, 0].sum(dim=-1))                 # tensor([1., 1., 1., 1., 1.])

import matplotlib.pyplot as plt
import seaborn as sns
import torch

# Cho ma trận attention đã có từ model (shape [L, L])
def plot_attention(attn_matrix, tokens):
    """Vẽ heatmap attention như visualization ở phía trên."""
    fig, ax = plt.subplots(figsize=(6, 5))
    sns.heatmap(
        attn_matrix.cpu().detach(),
        xticklabels=tokens,
        yticklabels=tokens,
        cmap="YlOrBr",      # amber scale — giống UI của bài này
        annot=True,
        fmt=".2f",
        cbar_kws={"label": "attention weight"},
        ax=ax,
    )
    ax.set_xlabel("Key (từ được nhìn)")
    ax.set_ylabel("Query (từ đang hỏi)")
    ax.set_title("Attention matrix")
    plt.tight_layout()
    return fig


# Pipeline đầy đủ cho 1 câu
def attention_for_sentence(model, tokenizer, sentence):
    tokens = tokenizer.tokenize(sentence)
    ids = torch.tensor([tokenizer.convert_tokens_to_ids(tokens)])

    with torch.no_grad():
        outputs = model(ids, output_attentions=True)

    # outputs.attentions: tuple[L_layers] of [B, H, L, L]
    # Lấy layer cuối, trung bình qua các đầu
    attn = outputs.attentions[-1][0].mean(dim=0)  # [L, L]

    return plot_attention(attn, tokens)


# sentence = "Tôi yêu Việt Nam quá"
# attention_for_sentence(model, tokenizer, sentence)


# ─── Bonus: so sánh attention ở các layer khác nhau ───
def attention_across_layers(model, tokenizer, sentence):
    """Tạo subplot cho mỗi layer — xem attention evolving theo chiều sâu."""
    tokens = tokenizer.tokenize(sentence)
    ids = torch.tensor([tokenizer.convert_tokens_to_ids(tokens)])

    with torch.no_grad():
        outputs = model(ids, output_attentions=True)

    n_layers = len(outputs.attentions)
    fig, axes = plt.subplots(
        1, n_layers, figsize=(3 * n_layers, 3), sharey=True
    )

    for layer_idx, attn_layer in enumerate(outputs.attentions):
        # Trung bình qua các đầu của layer này
        attn = attn_layer[0].mean(dim=0).cpu()
        ax = axes[layer_idx] if n_layers > 1 else axes
        sns.heatmap(
            attn,
            xticklabels=tokens,
            yticklabels=tokens if layer_idx == 0 else False,
            cmap="YlOrBr",
            cbar=False,
            ax=ax,
            square=True,
        )
        ax.set_title(f"Layer {layer_idx}")

    fig.suptitle("Attention từ nông đến sâu")
    plt.tight_layout()
    return fig


# ─── Bonus 2: temperature sweep (giống slider ở phần viz) ───
def temperature_sweep(raw_scores, tokens, temperatures=(0.2, 1.0, 5.0)):
    """Vẽ 3 heatmap với T khác nhau để so sánh độ 'quyết đoán'."""
    fig, axes = plt.subplots(1, len(temperatures), figsize=(4 * len(temperatures), 4))

    for ax, T in zip(axes, temperatures):
        scaled = raw_scores / T
        attn = torch.softmax(scaled, dim=-1)

        sns.heatmap(
            attn,
            xticklabels=tokens,
            yticklabels=tokens,
            cmap="YlOrBr",
            annot=True,
            fmt=".2f",
            cbar=False,
            ax=ax,
        )
        ax.set_title(f"T = {T}")

    return fig