"""A small, readable encoder-decoder Transformer inspired by Vaswani et al. (2017). Run with Python and PyTorch. CPU is sufficient. This is an educational forward / backward example, not a trained translator or the paper's full training recipe. All custom attention masks use True = allowed, False = blocked. """ import math import torch from torch import nn import torch.nn.functional as F def attention(q, k, v, allow=None): """Last axes: q[..., Tq, Dk], k[..., Tk, Dk], v[..., Tk, Dv].""" scores = (q @ k.transpose(-2, -1)) / math.sqrt(q.size(-1)) if allow is not None: if not allow.any(dim=-1).all(): raise ValueError('Every query needs at least one allowed key') scores = scores.masked_fill(~allow, float('-inf')) weights = scores.softmax(dim=-1) return weights @ v, weights class MultiHead(nn.Module): def __init__(self, d_model, heads): super().__init__() if d_model % heads: raise ValueError('d_model must be divisible by heads') self.heads, self.d_head = heads, d_model // 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.out_proj = nn.Linear(d_model, d_model) def forward(self, query, key, value, allow=None): batch = query.size(0) def split(x, projection): return projection(x).reshape(batch, -1, self.heads, self.d_head).transpose(1, 2) q = split(query, self.q_proj) k = split(key, self.k_proj) v = split(value, self.v_proj) out, weights = attention(q, k, v, allow) out = out.transpose(1, 2).contiguous().reshape(batch, query.size(1), -1) return self.out_proj(out), weights class Positions(nn.Module): def __init__(self, d_model, max_length=128): super().__init__() if d_model % 2: raise ValueError('This compact implementation needs even d_model') position = torch.arange(max_length).unsqueeze(1) frequencies = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.) / d_model)) pe = torch.zeros(max_length, d_model) pe[:, 0::2] = torch.sin(position * frequencies) pe[:, 1::2] = torch.cos(position * frequencies) self.register_buffer('pe', pe.unsqueeze(0)) def forward(self, x): if x.size(1) > self.pe.size(1): raise ValueError('Sequence exceeds configured positional encoding length') return x + self.pe[:, :x.size(1)] def feed_forward(d_model, d_ff): return nn.Sequential(nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model)) class EncoderLayer(nn.Module): def __init__(self, d_model, heads, d_ff, dropout): super().__init__() self.self_attention = MultiHead(d_model, heads) self.ffn = feed_forward(d_model, d_ff) self.norm1, self.norm2 = nn.LayerNorm(d_model), nn.LayerNorm(d_model) self.drop = nn.Dropout(dropout) def forward(self, x, source_allow): context, _ = self.self_attention(x, x, x, source_allow) x = self.norm1(x + self.drop(context)) # Original paper: post-norm. return self.norm2(x + self.drop(self.ffn(x))) class DecoderLayer(nn.Module): def __init__(self, d_model, heads, d_ff, dropout): super().__init__() self.self_attention = MultiHead(d_model, heads) self.cross_attention = MultiHead(d_model, heads) self.ffn = feed_forward(d_model, d_ff) self.norm1, self.norm2, self.norm3 = [nn.LayerNorm(d_model) for _ in range(3)] self.drop = nn.Dropout(dropout) def forward(self, x, memory, target_allow, source_allow): context, _ = self.self_attention(x, x, x, target_allow) x = self.norm1(x + self.drop(context)) context, _ = self.cross_attention(x, memory, memory, source_allow) x = self.norm2(x + self.drop(context)) return self.norm3(x + self.drop(self.ffn(x))) class TinyTransformer(nn.Module): def __init__(self, vocab_size=16, d_model=32, heads=4, d_ff=64, layers=2, pad_id=0, dropout=.1): super().__init__() self.pad_id, self.scale = pad_id, math.sqrt(d_model) # A shared toy vocabulary makes sharing source/target embeddings possible. self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=pad_id) self.positions = Positions(d_model) self.drop = nn.Dropout(dropout) self.encoder = nn.ModuleList([EncoderLayer(d_model, heads, d_ff, dropout) for _ in range(layers)]) self.decoder = nn.ModuleList([DecoderLayer(d_model, heads, d_ff, dropout) for _ in range(layers)]) self.output = nn.Linear(d_model, vocab_size, bias=False) self.output.weight = self.embedding.weight def forward(self, source, target_input): # Padding keys must never be read. Padded query rows can exist; their # outputs are ignored by the loss. Reject fully padded sources explicitly. source_allow = (source != self.pad_id)[:, None, None, :] n = target_input.size(1) causal = torch.ones(n, n, dtype=torch.bool, device=target_input.device).tril() target_allow = causal[None, None] & (target_input != self.pad_id)[:, None, None, :] memory = self.drop(self.positions(self.embedding(source) * self.scale)) x = self.drop(self.positions(self.embedding(target_input) * self.scale)) for layer in self.encoder: memory = layer(memory, source_allow) for layer in self.decoder: x = layer(x, memory, target_allow, source_allow) return self.output(x) # Raw logits: CrossEntropyLoss supplies log-softmax. def demo(): torch.manual_seed(7) q = torch.tensor([[1., 0.]]) k = torch.tensor([[1., 0.], [0., 1.], [1., 1.]]) v = torch.tensor([[2., 0.], [0., 2.], [2., 2.]]) out, weights = attention(q, k, v) print('Attention weights:', weights.round(decimals=4)) print('Weighted output:', out.round(decimals=4)) # PAD=0, BOS=1, EOS=2; all remaining IDs are toy vocabulary entries. source = torch.tensor([[3, 4, 2, 0], [5, 6, 7, 2]]) target = torch.tensor([[1, 8, 9, 2, 0], [1, 10, 11, 12, 2]]) model = TinyTransformer() logits = model(source, target[:, :-1]) labels = target[:, 1:] loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), labels.reshape(-1), ignore_index=0) optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) optimizer.zero_grad() loss.backward() optimizer.step() print('Logit shape:', tuple(logits.shape)) # (2, 4, 16) print('One toy training step completed; this model is not a trained translator.') if __name__ == '__main__': demo()