app/model_def.py

import torch
import torch.nn as nn
import math


class InputEmbedding(nn.Module):
    def __init__(self, d_model: int, vocab_size: int):
        super().__init__()
        self.d_model = d_model
        self.vocab_size = vocab_size
        self.embed = nn.Embedding(vocab_size, d_model)

    def forward(self, x):
        return self.embed(x) * math.sqrt(self.d_model)


class PositionalEncoding(nn.Module):
    def __init__(self, d_model: int, seq_len: int, dropout: float):
        super().__init__()
        self.d_model = d_model
        self.seq_len = seq_len
        self.dropout = nn.Dropout(dropout)
        pe = torch.zeros(seq_len, d_model)
        position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer("pe", pe)

    def forward(self, x):
        x = x + (self.pe[:, : x.shape[1], :]).requires_grad_(False)
        return self.dropout(x)


class LayerNorm(nn.Module):
    def __init__(self, d_model: int, epsilon: float = 1e-6):
        super().__init__()
        self.epsilon = epsilon
        self.gamma = nn.Parameter(torch.ones(d_model))
        self.beta = nn.Parameter(torch.zeros(d_model))

    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        std = x.std(dim=-1, keepdim=True)
        return self.gamma * (x - mean) / (std + self.epsilon) + self.beta


class FeedForward(nn.Module):
    def __init__(self, d_model: int, d_ff: int, dropout: float):
        super().__init__()
        self.layer1 = nn.Linear(d_model, d_ff)
        self.layer2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.layer2(self.dropout(torch.relu(self.layer1(x))))


class MHA(nn.Module):
    def __init__(self, d_model: int, h: int, dropout: float):
        super().__init__()
        self.d_model = d_model
        self.h = h
        assert d_model % h == 0, "d_model must be divisible by h"
        self.d_k = d_model // h
        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.dropout = nn.Dropout(dropout)

    @staticmethod
    def attention(query, key, value, mask, dropout: nn.Dropout):
        d_k = query.shape[-1]
        attention_scores = (query @ key.transpose(-2, -1)) / math.sqrt(d_k)
        if mask is not None:
            attention_scores = attention_scores.masked_fill(mask == 0, -1e9)
        attention_scores = attention_scores.softmax(dim=-1)
        if dropout is not None:
            attention_scores = dropout(attention_scores)
        return (attention_scores @ value), attention_scores

    def forward(self, q, k, v, mask):
        query = self.w_q(q)
        key = self.w_k(k)
        value = self.w_v(v)
        query = query.view(query.shape[0], query.shape[1], self.h, self.d_k).transpose(
            1, 2
        )
        key = key.view(key.shape[0], key.shape[1], self.h, self.d_k).transpose(1, 2)
        value = value.view(value.shape[0], value.shape[1], self.h, self.d_k).transpose(
            1, 2
        )
        x, self.attention_scores = MHA.attention(query, key, value, mask, self.dropout)
        x = x.transpose(1, 2).contiguous().view(x.shape[0], -1, self.h * self.d_k)
        return self.w_o(x)


class SkipConnection(nn.Module):
    def __init__(self, d_model: int, dropout: float):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        self.norm = LayerNorm(d_model)

    def forward(self, x, sublayer):
        return x + self.dropout(sublayer(self.norm(x)))


class EncoderBlock(nn.Module):
    def __init__(
        self, self_attention: MHA, ffn: FeedForward, d_model: int, dropout: float
    ):
        super().__init__()
        # Name required by the saved model file
        self.attention = self_attention
        self.ffn = ffn
        # Name required by the saved model file
        self.residual = nn.ModuleList(
            [SkipConnection(d_model, dropout) for _ in range(2)]
        )

    def forward(self, x, src_mask):
        x = self.residual[0](x, lambda x: self.attention(x, x, x, src_mask))
        x = self.residual[1](x, self.ffn)
        return x


class Encoder(nn.Module):
    def __init__(self, d_model: int, layers: nn.ModuleList):
        super().__init__()
        self.layers = layers
        self.norm = LayerNorm(d_model)

    def forward(self, x, mask):
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)


class DecoderBlock(nn.Module):
    def __init__(
        self,
        self_attention: MHA,
        cross_attention: MHA,
        ffn: FeedForward,
        d_model: int,
        dropout: float,
    ):
        super().__init__()
        # Name required by the saved model file
        self.self_attention = self_attention
        self.cross_attention = cross_attention
        self.ffn = ffn
        # Name required by the saved model file
        self.residual = nn.ModuleList(
            [SkipConnection(d_model, dropout) for _ in range(3)]
        )

    def forward(self, x, encoder_output, src_mask, trg_mask):
        x = self.residual[0](x, lambda x: self.self_attention(x, x, x, trg_mask))
        x = self.residual[1](
            x,
            lambda x: self.cross_attention(x, encoder_output, encoder_output, src_mask),
        )
        x = self.residual[2](x, self.ffn)
        return x


class Decoder(nn.Module):
    def __init__(self, d_model: int, layers: nn.ModuleList):
        super().__init__()
        self.layers = layers
        self.norm = LayerNorm(d_model)

    def forward(self, x, encoder_output, src_mask, trg_mask):
        for layer in self.layers:
            x = layer(x, encoder_output, src_mask, trg_mask)
        return self.norm(x)


class Output(nn.Module):
    def __init__(self, d_model: int, vocab_size: int):
        super().__init__()
        self.proj = nn.Linear(d_model, vocab_size)

    def forward(self, x):
        return self.proj(x)


class Transformer(nn.Module):
    def __init__(
        self,
        encoder: Encoder,
        decoder: Decoder,
        src_embed: InputEmbedding,
        trg_embed: InputEmbedding,
        src_pos: PositionalEncoding,
        trg_pos: PositionalEncoding,
        output: Output,
    ):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = src_embed
        self.trg_embed = trg_embed
        self.src_pos = src_pos
        self.trg_pos = trg_pos
        self.output_layer = output

    def encode(self, src, src_mask):
        src = self.src_embed(src)
        src = self.src_pos(src)
        return self.encoder(src, src_mask)

    def decode(self, encoder_output, src_mask, trg, trg_mask):
        trg = self.trg_embed(trg)
        trg = self.trg_pos(trg)
        return self.decoder(trg, encoder_output, src_mask, trg_mask)

    def project(self, x):
        return self.output_layer(x)

    def get_cross_attention_weights(self):
        """Collect cross-attention weights from all decoder layers.
        Returns list of 6 tensors, each [heads, src_len, tgt_len].
        Must be called after a forward pass through decode()."""
        weights = []
        for block in self.decoder.layers:
            attn_scores = block.cross_attention.attention_scores
            if attn_scores is not None:
                weights.append(attn_scores.squeeze(0))
        return weights

    def get_all_attention_weights(self):
        """Collect ALL attention weights from encoder and decoder.
        Returns dict with lists of layer tensors."""
        enc_weights = []
        dec_self_weights = []
        cross_weights = []
        for block in self.encoder.layers:
            scores = block.attention.attention_scores
            if scores is not None:
                enc_weights.append(scores.squeeze(0))
        for block in self.decoder.layers:
            scores = block.self_attention.attention_scores
            if scores is not None:
                dec_self_weights.append(scores.squeeze(0))
            scores = block.cross_attention.attention_scores
            if scores is not None:
                cross_weights.append(scores.squeeze(0).transpose(-2, -1))
        return {
            "encoder_self_attn": enc_weights,
            "decoder_self_attn": dec_self_weights,
            "cross_attn": cross_weights,
        }


def BuildTransformer(
    src_vocab_size: int,
    trg_vocab_size: int,
    src_seq_len: int,
    trg_seq_len: int,
    d_model: int = 512,
    N: int = 6,
    h: int = 8,
    dropout: float = 0.1,
    d_ff: int = 2048,
) -> Transformer:
    src_embed = InputEmbedding(d_model, src_vocab_size)
    trg_embed = InputEmbedding(d_model, trg_vocab_size)
    src_pos = PositionalEncoding(d_model, src_seq_len, dropout)
    trg_pos = PositionalEncoding(d_model, trg_seq_len, dropout)
    encoder_blocks = []
    for _ in range(N):
        encoder_self_attention = MHA(d_model, h, dropout)
        ffn = FeedForward(d_model, d_ff, dropout)
        encoder_block = EncoderBlock(encoder_self_attention, ffn, d_model, dropout)
        encoder_blocks.append(encoder_block)
    decoder_blocks = []
    for _ in range(N):
        decoder_self_attention = MHA(d_model, h, dropout)
        cross_attention = MHA(d_model, h, dropout)
        ffn = FeedForward(d_model, d_ff, dropout)
        decoder_block = DecoderBlock(
            decoder_self_attention, cross_attention, ffn, d_model, dropout
        )
        decoder_blocks.append(decoder_block)
    encoder = Encoder(d_model, nn.ModuleList(encoder_blocks))
    decoder = Decoder(d_model, nn.ModuleList(decoder_blocks))
    projection = Output(d_model, trg_vocab_size)
    transformer = Transformer(
        encoder, decoder, src_embed, trg_embed, src_pos, trg_pos, projection
    )
    for p in transformer.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)
    return transformer