"""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) # shared src/tgt/output 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) # [B,1,1,Ts] 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)) # [B,Tt,V] 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