المحولات والاهتمام الذاتي

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

def scaled_dot_product_attention(Q, K, V, mask=None):
    class="tok-str">"""
    Q, K, V: (batch, heads, seq_len, d_k)
    """
    d_k = Q.shape[-class="tok-num">1]
    class="tok-comment"># Attention scores
    scores = torch.matmul(Q, K.transpose(-class="tok-num">2, -class="tok-num">1)) / math.sqrt(d_k)
    class="tok-comment"># Causal mask (GPT-style)
    if mask is not None:
        scores = scores.masked_fill(mask == class="tok-num">0, float(class="tok-str">'-inf'))
    class="tok-comment"># Softmax over key dimension
    attn_weights = F.softmax(scores, dim=-class="tok-num">1)
    class="tok-comment"># Weighted sum of values
    return torch.matmul(attn_weights, V), attn_weights

class MultiHeadAttention(torch.nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.d_k = d_model // n_heads
        self.n_heads = n_heads
        self.W_q = torch.nn.Linear(d_model, d_model)
        self.W_k = torch.nn.Linear(d_model, d_model)
        self.W_v = torch.nn.Linear(d_model, d_model)
        self.W_o = torch.nn.Linear(d_model, d_model)

    def forward(self, Q, K, V, mask=None):
        B, T, D = Q.shape
        class="tok-comment"># Project + split into heads
        Q = self.W_q(Q).view(B, T, self.n_heads, self.d_k).transpose(class="tok-num">1, class="tok-num">2)
        K = self.W_k(K).view(B, T, self.n_heads, self.d_k).transpose(class="tok-num">1, class="tok-num">2)
        V = self.W_v(V).view(B, T, self.n_heads, self.d_k).transpose(class="tok-num">1, class="tok-num">2)
        x, weights = scaled_dot_product_attention(Q, K, V, mask)
        class="tok-comment"># Concat heads + project
        x = x.transpose(class="tok-num">1, class="tok-num">2).contiguous().view(B, T, D)
        return self.W_o(x), weights