model.py

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from config_300M import ModelConfig


class SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, d_model: int, max_len: int = 2048):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        pos = torch.arange(0, max_len).unsqueeze(1).float()
        div = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(pos * div)
        pe[:, 1::2] = torch.cos(pos * div)
        self.register_buffer("pe", pe.unsqueeze(0))  # (1, max_len, d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.pe[:, :x.size(1)]


class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_k = d_model // n_heads
        self.n_heads = 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.dropout = nn.Dropout(dropout)

    def forward(self, q, k, v, mask=None):
        B, T, _ = q.shape
        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)

        out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask,
            dropout_p=self.dropout.p if self.training else 0.0)
        out = out.transpose(1, 2).contiguous().view(B, -1, self.n_heads * self.d_k)
        return self.W_o(out)

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

    def forward(self, x):
        return self.w2(self.dropout(F.gelu(self.w1(x))))


class EncoderLayer(nn.Module):
    def __init__(self, d_model: int, n_heads: int, d_ff: int, dropout: float):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.ff = FeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, src_mask=None):
        x = x + self.dropout(self.self_attn(self.norm1(x), self.norm1(x), self.norm1(x), src_mask))
        x = x + self.dropout(self.ff(self.norm2(x)))
        return x


class DecoderLayer(nn.Module):
    def __init__(self, d_model: int, n_heads: int, d_ff: int, dropout: float):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.cross_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.ff = FeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, enc_out, src_mask=None, tgt_mask=None):
        x = x + self.dropout(self.self_attn(self.norm1(x), self.norm1(x), self.norm1(x), tgt_mask))
        n2 = self.norm2(x)
        x = x + self.dropout(self.cross_attn(n2, enc_out, enc_out, src_mask))
        x = x + self.dropout(self.ff(self.norm3(x)))
        return x


class TranslationModel(nn.Module):
    def __init__(self, cfg: ModelConfig):
        super().__init__()
        self.cfg = cfg

        # shared embedding (encoder + decoder + output projection)
        self.shared_emb = nn.Embedding(cfg.vocab_size, cfg.d_model, padding_idx=cfg.pad_id)
        self.emb_scale = math.sqrt(cfg.d_model)
        self.pos_enc = SinusoidalPositionalEncoding(cfg.d_model, cfg.max_seq_len)
        self.emb_dropout = nn.Dropout(cfg.dropout)

        # language tag embeddings (one per supported language)
        self.lang_emb = nn.Embedding(cfg.n_langs, cfg.d_model)

        self.encoder = nn.ModuleList([
            EncoderLayer(cfg.d_model, cfg.n_heads, cfg.d_ff, cfg.dropout)
            for _ in range(cfg.encoder_layers)
        ])
        self.encoder_norm = nn.LayerNorm(cfg.d_model)

        self.decoder = nn.ModuleList([
            DecoderLayer(cfg.d_model, cfg.n_heads, cfg.d_ff, cfg.dropout)
            for _ in range(cfg.decoder_layers)
        ])
        self.decoder_norm = nn.LayerNorm(cfg.d_model)

        # output projection shares weights with embedding
        self.output_proj = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        self.output_proj.weight = self.shared_emb.weight

        self._init_weights()

    def _init_weights(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def encode(self, src_ids, src_lang_id, src_mask=None):
        x = self.shared_emb(src_ids) * self.emb_scale
        x = x + self.lang_emb(src_lang_id).unsqueeze(1)
        x = self.emb_dropout(self.pos_enc(x))
        for layer in self.encoder:
            x = layer(x, src_mask)
        return self.encoder_norm(x)

    def decode(self, tgt_ids, enc_out, tgt_lang_id, src_mask=None, tgt_mask=None):
        x = self.shared_emb(tgt_ids) * self.emb_scale
        x = x + self.lang_emb(tgt_lang_id).unsqueeze(1)
        x = self.emb_dropout(self.pos_enc(x))
        for layer in self.decoder:
            x = layer(x, enc_out, src_mask, tgt_mask)
        return self.decoder_norm(x)

    def forward(self, src_ids, tgt_ids, src_lang_id, tgt_lang_id, src_mask=None, tgt_mask=None):
        enc_out = self.encode(src_ids, src_lang_id, src_mask)
        dec_out = self.decode(tgt_ids, enc_out, tgt_lang_id, src_mask, tgt_mask)
        return self.output_proj(dec_out)

    @staticmethod
    def make_src_mask(src_ids, pad_id=0):
        # (B, 1, 1, S)
        return (src_ids != pad_id).unsqueeze(1).unsqueeze(2)

    @staticmethod
    def make_tgt_mask(tgt_ids, pad_id=0):
        B, T = tgt_ids.shape
        pad_mask = (tgt_ids != pad_id).unsqueeze(1).unsqueeze(2)  # (B, 1, 1, T)
        causal = torch.tril(torch.ones(T, T, device=tgt_ids.device)).bool()  # (T, T)
        return pad_mask & causal.unsqueeze(0).unsqueeze(0)

    def count_params(self):
        return sum(p.numel() for p in self.parameters())


if __name__ == "__main__":
    cfg = ModelConfig()
    model = TranslationModel(cfg)
    total = model.count_params()
    print(f"Total parameters: {total:,} ({total/1e6:.1f}M)")

    
    src = torch.randint(3, 100, (2, 32))
    tgt = torch.randint(3, 100, (2, 24))
    src_mask = model.make_src_mask(src)
    tgt_mask = model.make_tgt_mask(tgt)
    out = model(src, tgt, src_lang_id=0, tgt_lang_id=1, src_mask=src_mask, tgt_mask=tgt_mask)
    print(f"Output shape: {out.shape}")  # (2, 24, vocab_size)