| """Encoder-decoder Transformer written from scratch (Vaswani et al., 2017), |
| with Pre-LN sublayers, sinusoidal positions and tied embeddings. |
| No torch.nn.Transformer / nn.MultiheadAttention is used. |
| |
| Masks are boolean tensors where True means "this key position may be attended to". |
| """ |
| import math |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| class PositionalEncoding(nn.Module): |
| def __init__(self, d_model, max_len, dropout): |
| super().__init__() |
| if d_model % 2 != 0: |
| raise ValueError("d_model must be even for sinusoidal positional encoding") |
| pe = torch.zeros(max_len, d_model) |
| pos = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1) |
| div = torch.exp(torch.arange(0, d_model, 2, dtype=torch.float32) * (-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), persistent=False) |
| self.max_len = max_len |
| self.drop = nn.Dropout(dropout) |
|
|
| def forward(self, x): |
| if x.size(1) > self.max_len: |
| raise ValueError(f"Sequence length {x.size(1)} exceeds max_seq_len {self.max_len}") |
| return self.drop(x + self.pe[:, : x.size(1)].to(x.dtype)) |
|
|
|
|
| class MultiHeadAttention(nn.Module): |
| def __init__(self, d_model, n_heads, dropout): |
| super().__init__() |
| if d_model % n_heads != 0: |
| raise ValueError("d_model must be divisible by n_heads") |
| self.h, self.dk = n_heads, d_model // n_heads |
| self.q_proj = nn.Linear(d_model, d_model) |
| self.k_proj = nn.Linear(d_model, d_model) |
| self.v_proj = nn.Linear(d_model, d_model) |
| self.o_proj = nn.Linear(d_model, d_model) |
| self.drop = nn.Dropout(dropout) |
|
|
| def forward(self, q, k, v, mask=None): |
| B, Tq, _ = q.shape |
| Tk = k.size(1) |
| Q = self.q_proj(q).view(B, Tq, self.h, self.dk).transpose(1, 2) |
| K = self.k_proj(k).view(B, Tk, self.h, self.dk).transpose(1, 2) |
| V = self.v_proj(v).view(B, Tk, self.h, self.dk).transpose(1, 2) |
| scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.dk) |
| if mask is not None: |
| scores = scores.masked_fill(~mask, torch.finfo(scores.dtype).min) |
| attn = self.drop(F.softmax(scores, dim=-1)) |
| out = torch.matmul(attn, V).transpose(1, 2).contiguous().view(B, Tq, self.h * self.dk) |
| return self.o_proj(out) |
|
|
|
|
| class FeedForward(nn.Module): |
| def __init__(self, d_model, d_ff, dropout): |
| super().__init__() |
| self.l1, self.l2 = nn.Linear(d_model, d_ff), nn.Linear(d_ff, d_model) |
| self.drop = nn.Dropout(dropout) |
|
|
| def forward(self, x): |
| return self.l2(self.drop(F.relu(self.l1(x)))) |
|
|
|
|
| class EncoderLayer(nn.Module): |
| def __init__(self, d_model, n_heads, d_ff, dropout): |
| super().__init__() |
| self.attn = MultiHeadAttention(d_model, n_heads, dropout) |
| self.ff = FeedForward(d_model, d_ff, dropout) |
| self.n1, self.n2 = nn.LayerNorm(d_model), nn.LayerNorm(d_model) |
| self.drop = nn.Dropout(dropout) |
|
|
| def forward(self, x, src_mask): |
| h = self.n1(x) |
| x = x + self.drop(self.attn(h, h, h, src_mask)) |
| return x + self.drop(self.ff(self.n2(x))) |
|
|
|
|
| class DecoderLayer(nn.Module): |
| def __init__(self, d_model, n_heads, d_ff, dropout): |
| 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.n1, self.n2, self.n3 = (nn.LayerNorm(d_model) for _ in range(3)) |
| self.drop = nn.Dropout(dropout) |
|
|
| def forward(self, x, memory, src_mask, tgt_mask): |
| h = self.n1(x) |
| x = x + self.drop(self.self_attn(h, h, h, tgt_mask)) |
| x = x + self.drop(self.cross_attn(self.n2(x), memory, memory, src_mask)) |
| return x + self.drop(self.ff(self.n3(x))) |
|
|
|
|
| class TransformerMT(nn.Module): |
| def __init__(self, vocab_size, d_model=256, n_heads=4, n_encoder_layers=3, n_decoder_layers=3, |
| d_ff=1024, dropout=0.3, max_seq_len=320, pad_id=0, tie_embeddings=True): |
| super().__init__() |
| self.d_model, self.pad_id, self.max_seq_len = d_model, pad_id, max_seq_len |
| self.embed = nn.Embedding(vocab_size, d_model, padding_idx=pad_id) |
| self.pos = PositionalEncoding(d_model, max_seq_len, dropout) |
| self.enc_layers = nn.ModuleList( |
| [EncoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(n_encoder_layers)]) |
| self.dec_layers = nn.ModuleList( |
| [DecoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(n_decoder_layers)]) |
| self.enc_norm, self.dec_norm = nn.LayerNorm(d_model), nn.LayerNorm(d_model) |
| self.out_proj = nn.Linear(d_model, vocab_size, bias=False) |
| if tie_embeddings: |
| self.out_proj.weight = self.embed.weight |
| self._init_weights() |
|
|
| def _init_weights(self): |
| for m in self.modules(): |
| if isinstance(m, nn.Linear): |
| nn.init.xavier_uniform_(m.weight) |
| if m.bias is not None: |
| nn.init.zeros_(m.bias) |
| nn.init.normal_(self.embed.weight, mean=0.0, std=self.d_model ** -0.5) |
| with torch.no_grad(): |
| self.embed.weight[self.pad_id].zero_() |
|
|
| def num_parameters(self): |
| return sum(p.numel() for p in self.parameters()) |
|
|
| def _embed(self, ids): |
| return self.pos(self.embed(ids) * math.sqrt(self.d_model)) |
|
|
| def make_src_mask(self, src): |
| return (src != self.pad_id).unsqueeze(1).unsqueeze(2) |
|
|
| def make_tgt_mask(self, tgt): |
| T = tgt.size(1) |
| causal = torch.tril(torch.ones(T, T, dtype=torch.bool, device=tgt.device)) |
| return (tgt != self.pad_id).unsqueeze(1).unsqueeze(2) & causal.unsqueeze(0).unsqueeze(0) |
|
|
| def encode(self, src, src_mask=None): |
| if src_mask is None: |
| src_mask = self.make_src_mask(src) |
| x = self._embed(src) |
| for layer in self.enc_layers: |
| x = layer(x, src_mask) |
| return self.enc_norm(x) |
|
|
| def decode(self, tgt_in, memory, src_mask, tgt_mask=None): |
| if tgt_mask is None: |
| tgt_mask = self.make_tgt_mask(tgt_in) |
| x = self._embed(tgt_in) |
| for layer in self.dec_layers: |
| x = layer(x, memory, src_mask, tgt_mask) |
| return self.out_proj(self.dec_norm(x)) |
|
|
| def forward(self, src, tgt_in): |
| src_mask = self.make_src_mask(src) |
| return self.decode(tgt_in, self.encode(src, src_mask), src_mask) |
|
|
|
|
| def load_model_from_checkpoint(path, device="cpu"): |
| """Loads a best.pt written by src/train.py. Returns (model_in_eval_mode, meta_dict).""" |
| ckpt = torch.load(path, map_location=device, weights_only=True) |
| model = TransformerMT(**ckpt["model_args"]).to(device) |
| model.load_state_dict(ckpt["model_state"]) |
| model.eval() |
| meta = {k: v for k, v in ckpt.items() if k != "model_state"} |
| return model, meta |
|
|
