#!/usr/bin/env python # # file: transformer_demo.py # # description: # This script provides an educational demonstration of the Transformer # architecture ("Attention Is All You Need", Vaswani et al., 2017) using # PyTorch. It builds a Transformer from scratch — including Scaled # Dot-Product Attention, Multi-Head Attention, Positional Encoding, # Encoder and Decoder stacks — and trains it on a simple integer-sequence # reversal task (e.g., [1,2,3,4] → [4,3,2,1]). # # The script prints per-epoch train/test losses and token-level accuracies, # then generates five educational plots: # 1. Training / test loss and accuracy curves # 2. Attention-weight heat-maps (one per head) from the last encoder layer # 3. Attention-weight heat-maps from the last decoder self-attention layer # 4. Attention-weight heat-maps from the decoder cross-attention layer # 5. A positional-encoding visualisation # # revision history: # 20260416 (AM): initial version # #------------------------------------------------------------------------------ # import system modules # import os import sys import math import numpy as np import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt import torch import torch.nn as nn import torch.optim as optim import torch.nn.functional as F from torch.utils.data import DataLoader, TensorDataset #------------------------------------------------------------------------------ # # global variables are listed here # #------------------------------------------------------------------------------ # set the filename using basename # __FILE__ = os.path.basename(__file__) # dataset / training hyper-parameters # DEF_SEQ_LEN = 8 # length of each integer sequence DEF_VOCAB_SIZE = 12 # integers 1..10 + PAD(0) + SOS(11) DEF_PAD_IDX = 0 # padding token index DEF_SOS_IDX = 11 # start-of-sequence token index DEF_NUM_TRAIN = 800 # number of training sequences DEF_NUM_TEST = 200 # number of test sequences DEF_BATCH_SIZE = 32 DEF_NUM_EPOCHS = 20 DEF_LEARNING_RATE = 1e-3 DEF_RANDOM_SEED = 42 # model hyper-parameters # DEF_D_MODEL = 64 # embedding / model dimension DEF_NUM_HEADS = 4 # number of attention heads DEF_NUM_LAYERS = 2 # number of encoder/decoder layers DEF_D_FF = 128 # feed-forward inner dimension DEF_DROPOUT = 0.1 # output file names # DEF_METRICS_FILE = "transformer_metrics.png" DEF_ENC_ATTN_FILE = "transformer_encoder_attention.png" DEF_DEC_SELF_FILE = "transformer_decoder_self_attention.png" DEF_DEC_CROSS_FILE = "transformer_decoder_cross_attention.png" DEF_PE_FILE = "transformer_positional_encoding.png" #------------------------------------------------------------------------------ # # helper / utility functions # #------------------------------------------------------------------------------ # lock all random seeds for reproducibility # def set_seed(seed=DEF_RANDOM_SEED): """ method: set_seed arguments: seed: integer random seed return: none description: Fixes NumPy and PyTorch random seeds so that synthetic data generation and weight initialisation are fully reproducible across runs. """ np.random.seed(seed) torch.manual_seed(seed) # exit gracefully # return #------------------------------------------------------------------------------ # # dataset generation # #------------------------------------------------------------------------------ # generate integer-sequence reversal pairs # def generate_reversal_data(n_train=DEF_NUM_TRAIN, n_test=DEF_NUM_TEST, seq_len=DEF_SEQ_LEN): """ method: generate_reversal_data arguments: n_train: number of training samples n_test: number of test samples seq_len: length of each source sequence return: train_loader: DataLoader for training data test_loader: DataLoader for test data description: Builds a sequence-reversal dataset. Each sample consists of: src — a random integer sequence of length seq_len drawn from [1, 10] tgt — SOS token followed by the reversed src (teacher-forcing input) gold — the reversed src followed by a PAD token (target output) This is a classical toy seq2seq task; a working Transformer should learn to reverse any unseen sequence after sufficient training. """ def _make_split(n): # source: random integers 1..10 # src = torch.randint(1, 11, (n, seq_len)) # target input (decoder input): [SOS, rev_1, rev_2, ..., rev_{n-1}] # target output (gold labels): [rev_1, rev_2, ..., rev_n, PAD] # rev = torch.flip(src, dims=[1]) sos = torch.full((n, 1), DEF_SOS_IDX, dtype=torch.long) pad = torch.full((n, 1), DEF_PAD_IDX, dtype=torch.long) tgt = torch.cat([sos, rev], dim=1) # (n, seq_len+1) gold = torch.cat([rev, pad], dim=1) # (n, seq_len+1) return TensorDataset(src, tgt, gold) train_ds = _make_split(n_train) test_ds = _make_split(n_test) train_loader = DataLoader(train_ds, batch_size=DEF_BATCH_SIZE, shuffle=True) test_loader = DataLoader(test_ds, batch_size=DEF_BATCH_SIZE, shuffle=False) # exit gracefully # return train_loader, test_loader #------------------------------------------------------------------------------ # # Transformer building blocks # #------------------------------------------------------------------------------ # sinusoidal positional encoding # class PositionalEncoding(nn.Module): """ class: PositionalEncoding description: Injects sequence-order information into token embeddings using fixed sinusoidal functions (Vaswani et al., 2017, Section 3.5): PE(pos, 2i) = sin( pos / 10000^(2i / d_model) ) PE(pos, 2i+1) = cos( pos / 10000^(2i / d_model) ) The encoding is added (not concatenated) to the embedding so that the model dimension d_model is preserved throughout. """ def __init__(self, d_model=DEF_D_MODEL, max_len=200, dropout=DEF_DROPOUT): super().__init__() self.dropout = nn.Dropout(dropout) # build (max_len, d_model) table of positional values # pe = torch.zeros(max_len, d_model) position = torch.arange(0, max_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) # register as a non-trainable buffer so it moves with .to(device) # self.register_buffer('pe', pe.unsqueeze(0)) # (1, max_len, d_model) def forward(self, x): # x: (batch, seq_len, d_model) # x = x + self.pe[:, :x.size(1)] return self.dropout(x) def get_encoding(self): """Return the raw PE matrix for visualisation (no dropout).""" return self.pe.squeeze(0).detach().cpu().numpy() # scaled dot-product attention # def scaled_dot_product_attention(Q, K, V, mask=None): """ function: scaled_dot_product_attention arguments: Q: query tensor (batch, heads, seq_q, d_k) K: key tensor (batch, heads, seq_k, d_k) V: value tensor (batch, heads, seq_k, d_v) mask: optional boolean mask — True positions are set to -inf return: output: context vectors (batch, heads, seq_q, d_v) weights: attention weights (batch, heads, seq_q, seq_k) description: Implements the core attention formula (Vaswani et al., 2017, Eq. 1): Attention(Q, K, V) = softmax( Q K^T / sqrt(d_k) ) V Dividing by sqrt(d_k) prevents the dot products from growing too large in magnitude, which would push softmax into regions with tiny gradients. """ d_k = Q.size(-1) scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k) if mask is not None: scores = scores.masked_fill(mask, float('-inf')) weights = F.softmax(scores, dim=-1) output = torch.matmul(weights, V) # exit gracefully # return output, weights # multi-head attention # class MultiHeadAttention(nn.Module): """ class: MultiHeadAttention description: Projects Q, K, V into h parallel "heads", applies scaled dot-product attention in each head independently, concatenates the results, and passes them through a final linear projection. MultiHead(Q,K,V) = Concat(head_1, ..., head_h) W^O head_i = Attention(Q W_i^Q, K W_i^K, V W_i^V) Using multiple heads allows the model to attend to information from different representation sub-spaces simultaneously. """ def __init__(self, d_model=DEF_D_MODEL, num_heads=DEF_NUM_HEADS): super().__init__() assert d_model % num_heads == 0, \ "d_model must be divisible by num_heads" self.num_heads = num_heads self.d_k = d_model // num_heads # single combined projection for Q, K, V (efficiency trick) # 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) # cache last attention weights for visualisation # self.last_attn_weights = None def _split_heads(self, x): """Reshape (batch, seq, d_model) → (batch, heads, seq, d_k).""" B, S, _ = x.size() return x.view(B, S, self.num_heads, self.d_k).transpose(1, 2) def forward(self, query, key, value, mask=None): B = query.size(0) Q = self._split_heads(self.W_q(query)) K = self._split_heads(self.W_k(key)) V = self._split_heads(self.W_v(value)) context, weights = scaled_dot_product_attention(Q, K, V, mask) self.last_attn_weights = weights.detach().cpu() # concatenate heads and project # context = context.transpose(1, 2).contiguous().view(B, -1, self.num_heads * self.d_k) # exit gracefully # return self.W_o(context) # position-wise feed-forward network # class FeedForward(nn.Module): """ class: FeedForward description: A two-layer MLP applied independently to each position: FFN(x) = ReLU( x W_1 + b_1 ) W_2 + b_2 The inner dimension d_ff is typically 4× d_model. This network provides the non-linear transformation capacity of the Transformer, complementing the linear attention mechanism. """ def __init__(self, d_model=DEF_D_MODEL, d_ff=DEF_D_FF, dropout=DEF_DROPOUT): super().__init__() self.net = nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Dropout(dropout), nn.Linear(d_ff, d_model), ) def forward(self, x): return self.net(x) # single encoder layer # class EncoderLayer(nn.Module): """ class: EncoderLayer description: One Transformer encoder layer (Vaswani et al., 2017, Figure 1, left): 1. Multi-Head Self-Attention (query = key = value = x) 2. Add & LayerNorm (residual connection) 3. Feed-Forward Network 4. Add & LayerNorm (residual connection) Residual connections ("Add") ease gradient flow through deep networks. Layer normalisation stabilises training. """ def __init__(self, d_model=DEF_D_MODEL, num_heads=DEF_NUM_HEADS, d_ff=DEF_D_FF, dropout=DEF_DROPOUT): super().__init__() self.self_attn = MultiHeadAttention(d_model, num_heads) self.ff = FeedForward(d_model, d_ff, dropout) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.drop = nn.Dropout(dropout) def forward(self, x, src_mask=None): # sub-layer 1: self-attention + residual # attn_out = self.self_attn(x, x, x, src_mask) x = self.norm1(x + self.drop(attn_out)) # sub-layer 2: feed-forward + residual # ff_out = self.ff(x) x = self.norm2(x + self.drop(ff_out)) # exit gracefully # return x # single decoder layer # class DecoderLayer(nn.Module): """ class: DecoderLayer description: One Transformer decoder layer (Vaswani et al., 2017, Figure 1, right): 1. Masked Multi-Head Self-Attention (prevents attending to future tokens) 2. Add & LayerNorm 3. Multi-Head Cross-Attention (queries from decoder, keys/values from encoder) 4. Add & LayerNorm 5. Feed-Forward Network 6. Add & LayerNorm The causal mask in step 1 ensures auto-regressive generation: position i can only attend to positions ≤ i. """ def __init__(self, d_model=DEF_D_MODEL, num_heads=DEF_NUM_HEADS, d_ff=DEF_D_FF, dropout=DEF_DROPOUT): super().__init__() self.self_attn = MultiHeadAttention(d_model, num_heads) # masked self self.cross_attn = MultiHeadAttention(d_model, num_heads) # cross 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.drop = nn.Dropout(dropout) def forward(self, tgt, memory, tgt_mask=None, memory_mask=None): # sub-layer 1: masked self-attention # sa_out = self.self_attn(tgt, tgt, tgt, tgt_mask) tgt = self.norm1(tgt + self.drop(sa_out)) # sub-layer 2: cross-attention (queries from decoder, K/V from encoder) # ca_out = self.cross_attn(tgt, memory, memory, memory_mask) tgt = self.norm2(tgt + self.drop(ca_out)) # sub-layer 3: feed-forward # ff_out = self.ff(tgt) tgt = self.norm3(tgt + self.drop(ff_out)) # exit gracefully # return tgt # full Transformer model # class Transformer(nn.Module): """ class: Transformer description: End-to-end Transformer for sequence-to-sequence tasks: - Source embedding + Positional Encoding → Encoder stack - Target embedding + Positional Encoding → Decoder stack - Linear projection + softmax → token probabilities The encoder reads the entire source sequence in parallel. The decoder generates output tokens one step at a time, conditioned on the encoder memory and previously generated tokens (during inference) or ground-truth tokens (teacher forcing, during training). """ def __init__(self, vocab_size=DEF_VOCAB_SIZE, d_model=DEF_D_MODEL, num_heads=DEF_NUM_HEADS, num_layers=DEF_NUM_LAYERS, d_ff=DEF_D_FF, dropout=DEF_DROPOUT, max_len=200): super().__init__() self.src_embed = nn.Embedding(vocab_size, d_model, padding_idx=DEF_PAD_IDX) self.tgt_embed = nn.Embedding(vocab_size, d_model, padding_idx=DEF_PAD_IDX) self.pos_enc = PositionalEncoding(d_model, max_len, dropout) self.encoder = nn.ModuleList([EncoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)]) self.decoder = nn.ModuleList([DecoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)]) self.output_proj = nn.Linear(d_model, vocab_size) # weight tying: share embedding weights with output projection # (reduces parameter count and improves generalisation) # self.output_proj.weight = self.tgt_embed.weight self._init_weights() def _init_weights(self): """Xavier uniform initialisation for all linear layers.""" for p in self.parameters(): if p.dim() > 1: nn.init.xavier_uniform_(p) @staticmethod def _make_causal_mask(size, device): """Upper-triangular mask to block future tokens in the decoder.""" return torch.triu(torch.ones(size, size, device=device), diagonal=1).bool() def encode(self, src): x = self.pos_enc(self.src_embed(src)) for layer in self.encoder: x = layer(x) return x def decode(self, tgt, memory): seq_len = tgt.size(1) causal = self._make_causal_mask(seq_len, tgt.device) x = self.pos_enc(self.tgt_embed(tgt)) for layer in self.decoder: x = layer(x, memory, tgt_mask=causal) return x def forward(self, src, tgt): memory = self.encode(src) out = self.decode(tgt, memory) return self.output_proj(out) # (batch, tgt_len, vocab_size) #------------------------------------------------------------------------------ # # training and evaluation # #------------------------------------------------------------------------------ # run one epoch of training # def train_epoch(model, loader, optimizer, criterion): """ method: train_epoch arguments: model: the Transformer model loader: DataLoader for training data optimizer: optimiser instance criterion: loss function (CrossEntropyLoss with ignore_index=PAD) return: avg_loss: mean loss over all batches accuracy: token-level accuracy (%, PAD tokens excluded) description: Performs one full pass through the training data using teacher forcing: the gold target sequence (shifted by one) is fed to the decoder at each step instead of the model's own previous prediction. """ model.train() total_loss = 0.0 total_tokens = 0 correct = 0 for src, tgt, gold in loader: optimizer.zero_grad() logits = model(src, tgt) # (B, T, V) B, T, V = logits.shape loss = criterion(logits.reshape(B * T, V), gold.reshape(B * T)) loss.backward() nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() # accumulate statistics (ignore PAD) # mask = gold.reshape(B * T) != DEF_PAD_IDX preds = logits.argmax(dim=-1).reshape(B * T) correct += (preds[mask] == gold.reshape(B * T)[mask]).sum().item() total_tokens += mask.sum().item() total_loss += loss.item() * mask.sum().item() # exit gracefully # return total_loss / total_tokens, 100.0 * correct / total_tokens # evaluate on a held-out split # def evaluate(model, loader, criterion): """ method: evaluate arguments: model: the Transformer model loader: DataLoader for test data criterion: loss function return: avg_loss: mean loss accuracy: token-level accuracy (%) description: Runs the model in inference mode (no gradient updates) and reports loss and accuracy on the provided dataset split. """ model.eval() total_loss = 0.0 total_tokens = 0 correct = 0 with torch.no_grad(): for src, tgt, gold in loader: logits = model(src, tgt) B, T, V = logits.shape loss = criterion(logits.reshape(B * T, V), gold.reshape(B * T)) mask = gold.reshape(B * T) != DEF_PAD_IDX preds = logits.argmax(dim=-1).reshape(B * T) correct += (preds[mask] == gold.reshape(B * T)[mask]).sum().item() total_tokens += mask.sum().item() total_loss += loss.item() * mask.sum().item() # exit gracefully # return total_loss / total_tokens, 100.0 * correct / total_tokens # main training loop # def train_model(model, train_loader, test_loader, epochs=DEF_NUM_EPOCHS, lr=DEF_LEARNING_RATE): """ method: train_model arguments: model: the Transformer to train train_loader: DataLoader for training data test_loader: DataLoader for test data epochs: number of training epochs lr: initial learning rate return: metrics: dict with lists of train/test loss and accuracy per epoch description: Runs the full training loop for the specified number of epochs. Uses Adam optimiser with a cosine annealing scheduler to decay the learning rate smoothly over training. Prints per-epoch stats. """ criterion = nn.CrossEntropyLoss(ignore_index=DEF_PAD_IDX) optimizer = optim.Adam(model.parameters(), lr=lr, betas=(0.9, 0.98), eps=1e-9) scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs) metrics = {'train_loss': [], 'train_acc': [], 'test_loss': [], 'test_acc': []} print("Beginning Transformer training for %d epochs..." % epochs) for epoch in range(1, epochs + 1): tr_loss, tr_acc = train_epoch(model, train_loader, optimizer, criterion) te_loss, te_acc = evaluate(model, test_loader, criterion) scheduler.step() metrics['train_loss'].append(tr_loss) metrics['train_acc'].append(tr_acc) metrics['test_loss'].append(te_loss) metrics['test_acc'].append(te_acc) print(" Epoch %2d/%d | Train Loss: %.4f Acc: %5.2f%% | " "Test Loss: %.4f Acc: %5.2f%%" % (epoch, epochs, tr_loss, tr_acc, te_loss, te_acc)) # exit gracefully # return metrics #------------------------------------------------------------------------------ # # plotting functions # #------------------------------------------------------------------------------ # plot training / test loss and accuracy curves # def plot_training_results(metrics, outfile=DEF_METRICS_FILE): """ method: plot_training_results arguments: metrics: dict with train/test loss and accuracy lists outfile: output file path return: status: True on success description: Two-panel figure showing cross-entropy loss (left) and token-level accuracy (right) for both the training and test splits over epochs. """ epochs_range = range(1, len(metrics['train_loss']) + 1) fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 5)) # panel 1 — loss # ax1.plot(epochs_range, metrics['train_loss'], marker='o', color='crimson', label='Train Loss') ax1.plot(epochs_range, metrics['test_loss'], marker='x', color='darkred', label='Test Loss', linestyle='--') ax1.set_title("Cross-Entropy Loss", fontsize=14) ax1.set_xlabel("Epoch") ax1.set_ylabel("Loss") ax1.legend() ax1.grid(True, linestyle='--', alpha=0.6) # panel 2 — accuracy # ax2.plot(epochs_range, metrics['train_acc'], marker='s', color='steelblue', label='Train Acc') ax2.plot(epochs_range, metrics['test_acc'], marker='^', color='navy', label='Test Acc', linestyle='--') ax2.set_title("Token-Level Accuracy", fontsize=14) ax2.set_xlabel("Epoch") ax2.set_ylabel("Accuracy (%)") ax2.set_ylim(0, 105) ax2.legend() ax2.grid(True, linestyle='--', alpha=0.6) fig.suptitle("Transformer Training Progression (Sequence Reversal Task)", fontsize=15, y=1.02) plt.tight_layout() try: plt.savefig(outfile, dpi=150, bbox_inches='tight') print("\nSaved metrics plot to: %s" % outfile) except Exception as e: print("**> Error saving metrics plot: %s" % str(e)) return False plt.close() # exit gracefully # return True # helper: draw a heat-map grid of attention weights # def _plot_attention_heatmaps(attn_weights, title, xlabel, ylabel, outfile): """ method: _plot_attention_heatmaps arguments: attn_weights: numpy array (num_heads, seq_q, seq_k) title: figure title xlabel: x-axis label (key positions) ylabel: y-axis label (query positions) outfile: output file path return: status: True on success description: Renders one heat-map per attention head in a grid layout. Brighter cells indicate higher attention weight, showing where each query position focuses when reading the key/value sequence. """ num_heads = attn_weights.shape[0] ncols = min(num_heads, 4) nrows = math.ceil(num_heads / ncols) fig, axes = plt.subplots(nrows, ncols, figsize=(4 * ncols, 4 * nrows)) axes = np.array(axes).reshape(-1) for h in range(num_heads): ax = axes[h] mat = attn_weights[h] im = ax.imshow(mat, cmap='viridis', aspect='auto', vmin=0, vmax=mat.max()) ax.set_title("Head %d" % (h + 1), fontsize=12) ax.set_xlabel(xlabel, fontsize=9) ax.set_ylabel(ylabel, fontsize=9) plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04) # hide unused subplots # for h in range(num_heads, len(axes)): axes[h].axis('off') fig.suptitle(title, fontsize=14, y=1.01) plt.tight_layout() try: plt.savefig(outfile, dpi=150, bbox_inches='tight') print("Saved attention plot to: %s" % outfile) except Exception as e: print("**> Error saving attention plot: %s" % str(e)) return False plt.close() # exit gracefully # return True # extract and plot attention weights from the encoder, decoder self-attention, # and decoder cross-attention for one sample sequence # def plot_attention_maps(model, train_loader, enc_file=DEF_ENC_ATTN_FILE, dec_self_file=DEF_DEC_SELF_FILE, dec_cross_file=DEF_DEC_CROSS_FILE): """ method: plot_attention_maps arguments: model: the trained Transformer train_loader: DataLoader (used to get a sample batch) enc_file: output file for encoder attention maps dec_self_file: output file for decoder self-attention maps dec_cross_file: output file for decoder cross-attention maps return: status: True on success description: Runs one sample through the model, then retrieves the cached attention weights from the last encoder layer, the last decoder self-attention sub-layer, and the last decoder cross-attention sub-layer. Three separate figures are produced, each showing a heat-map per head. These visualisations reveal which source tokens each decoder position attends to (cross-attention) and how tokens relate within the encoder (self-attention). """ model.eval() # grab a single batch # src, tgt, _ = next(iter(train_loader)) src, tgt = src[:1], tgt[:1] # just one sequence with torch.no_grad(): memory = model.encode(src) _ = model.decode(tgt, memory) seq_labels_src = [str(t.item()) for t in src[0]] seq_labels_tgt = ["SOS"] + [str(t.item()) for t in src[0].flip(0)] # 1. encoder self-attention (last encoder layer) # enc_w = model.encoder[-1].self_attn.last_attn_weights[0].numpy() # (H, S, S) status = _plot_attention_heatmaps( enc_w, title = "Encoder Self-Attention (Last Layer)\n" "— each row = one query token, each column = one key token", xlabel = "Key position (source token)", ylabel = "Query position (source token)", outfile = enc_file, ) # 2. decoder masked self-attention (last decoder layer) # dec_sa_w = model.decoder[-1].self_attn.last_attn_weights[0].numpy() # (H, T, T) status &= _plot_attention_heatmaps( dec_sa_w, title = "Decoder Masked Self-Attention (Last Layer)\n" "— causal mask zeroes the upper triangle (no future peeking)", xlabel = "Key position (target token)", ylabel = "Query position (target token)", outfile = dec_self_file, ) # 3. decoder cross-attention (last decoder layer) # dec_ca_w = model.decoder[-1].cross_attn.last_attn_weights[0].numpy() # (H, T, S) status &= _plot_attention_heatmaps( dec_ca_w, title = "Decoder Cross-Attention (Last Layer)\n" "— decoder queries attending to encoder (source) keys", xlabel = "Key position (source / encoder token)", ylabel = "Query position (target / decoder token)", outfile = dec_cross_file, ) # exit gracefully # return status # visualise the positional encoding matrix # def plot_positional_encoding(model, seq_len=DEF_SEQ_LEN + 1, outfile=DEF_PE_FILE): """ method: plot_positional_encoding arguments: model: the Transformer (provides access to the PE module) seq_len: number of positions to display outfile: output file path return: status: True on success description: Displays the positional encoding matrix as a heat-map (positions vs dimensions) and plots individual sine/cosine curves for the first few dimensions to show the frequency pattern. This helps students see why each position gets a unique and smoothly varying fingerprint. """ pe_full = model.pos_enc.get_encoding() # (max_len, d_model) pe = pe_full[:seq_len] # crop to display range fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5)) # panel 1 — heat-map # im = ax1.imshow(pe.T, cmap='RdBu', aspect='auto', vmin=-1, vmax=1, origin='lower') ax1.set_title("Positional Encoding Matrix\n(rows = embedding dims, cols = positions)", fontsize=12) ax1.set_xlabel("Position in sequence") ax1.set_ylabel("Embedding dimension") plt.colorbar(im, ax=ax1, fraction=0.046, pad=0.04) # panel 2 — individual PE curves for first 8 dims # positions = np.arange(seq_len) n_show = min(8, pe.shape[1]) colors = plt.cm.tab10(np.linspace(0, 1, n_show)) for d in range(n_show): label = "dim %d (%s)" % (d, "sin" if d % 2 == 0 else "cos") ax2.plot(positions, pe[:, d], color=colors[d], marker='o', markersize=4, label=label) ax2.set_title("PE Values for First %d Dimensions\n(alternating sin / cos frequencies)" % n_show, fontsize=12) ax2.set_xlabel("Position in sequence") ax2.set_ylabel("Encoding value") ax2.legend(fontsize=8, ncol=2) ax2.grid(True, linestyle='--', alpha=0.5) ax2.set_ylim(-1.1, 1.1) fig.suptitle("Sinusoidal Positional Encoding (d_model = %d)" % DEF_D_MODEL, fontsize=14, y=1.02) plt.tight_layout() try: plt.savefig(outfile, dpi=150, bbox_inches='tight') print("Saved positional encoding plot to: %s" % outfile) except Exception as e: print("**> Error saving PE plot: %s" % str(e)) return False plt.close() # exit gracefully # return True #------------------------------------------------------------------------------ # # quick sanity demo: greedily decode a handful of sequences # #------------------------------------------------------------------------------ # greedy decode one source sequence (no teacher forcing) # def greedy_decode(model, src_seq, max_len=DEF_SEQ_LEN + 1): """ method: greedy_decode arguments: model: the trained Transformer src_seq: 1-D LongTensor of source token indices max_len: maximum number of tokens to generate return: output_tokens: list of generated token indices description: Implements greedy auto-regressive decoding. At each step the model predicts a probability distribution over the vocabulary; the token with the highest probability is selected and appended to the decoder input for the next step. Generation stops after max_len tokens. (No EOS token is used in this toy task.) """ model.eval() src = src_seq.unsqueeze(0) # (1, S) memory = model.encode(src) tgt_tokens = [DEF_SOS_IDX] with torch.no_grad(): for _ in range(max_len): tgt = torch.tensor([tgt_tokens], dtype=torch.long) out = model.decode(tgt, memory) logits = model.output_proj(out[:, -1, :]) # last token logits next_t = logits.argmax(dim=-1).item() tgt_tokens.append(next_t) # exit gracefully (strip leading SOS) # return tgt_tokens[1:] # demonstrate greedy decoding on a few test sequences # def demo_greedy_decoding(model, test_loader, num_examples=6): """ method: demo_greedy_decoding arguments: model: the trained Transformer test_loader: DataLoader for test data num_examples: how many sequences to print return: none description: Samples a batch of test sequences, greedily decodes each one, and prints the source, expected (reversed) target, and predicted output side-by-side so students can verify the model has learned to reverse. """ print("\n--- Greedy Decoding Demo (no teacher forcing) ---") print("%-30s %-30s %-30s %s" % ("Source", "Expected (reversed)", "Predicted", "Correct?")) print("-" * 115) src_batch, _, gold_batch = next(iter(test_loader)) shown = 0 for i in range(src_batch.size(0)): if shown >= num_examples: break src_seq = src_batch[i] expected = gold_batch[i].tolist() # strip trailing PADs # expected = [t for t in expected if t != DEF_PAD_IDX] predicted = greedy_decode(model, src_seq)[:len(expected)] correct = "YES" if predicted == expected else "NO" print("%-30s %-30s %-30s %s" % ( str(src_seq.tolist()), str(expected), str(predicted), correct, )) shown += 1 # exit gracefully # return #------------------------------------------------------------------------------ # # entry point # #------------------------------------------------------------------------------ # function: main # def main(argv): print("--- Starting Transformer (Attention Is All You Need) Demonstration ---") # 1. reproducibility # set_seed() # 2. dataset # print("Generating synthetic sequence-reversal dataset...") train_loader, test_loader = generate_reversal_data() print(" Train samples: %d | Test samples: %d" % ( DEF_NUM_TRAIN, DEF_NUM_TEST)) # 3. model # print("\nBuilding Transformer model...") model = Transformer( vocab_size = DEF_VOCAB_SIZE, d_model = DEF_D_MODEL, num_heads = DEF_NUM_HEADS, num_layers = DEF_NUM_LAYERS, d_ff = DEF_D_FF, dropout = DEF_DROPOUT, ) n_params = sum(p.numel() for p in model.parameters() if p.requires_grad) print(" Trainable parameters: %d" % n_params) print(" Architecture: d_model=%d heads=%d layers=%d d_ff=%d" % (DEF_D_MODEL, DEF_NUM_HEADS, DEF_NUM_LAYERS, DEF_D_FF)) # 4. training # metrics = train_model(model, train_loader, test_loader) # 5. greedy decode a few examples # demo_greedy_decoding(model, test_loader) # 6. plots # print("\nGenerating educational plots...") ok = True ok &= plot_training_results(metrics) ok &= plot_positional_encoding(model) ok &= plot_attention_maps(model, train_loader) if not ok: print("**> One or more plots failed to save.") return False print("\n--- Demonstration Complete ---") print("Output files:") print(" %s — training curves" % DEF_METRICS_FILE) print(" %s — positional encoding" % DEF_PE_FILE) print(" %s — encoder self-attention" % DEF_ENC_ATTN_FILE) print(" %s — decoder masked self-attention" % DEF_DEC_SELF_FILE) print(" %s — decoder cross-attention" % DEF_DEC_CROSS_FILE) # exit gracefully # return True # begin gracefully # if __name__ == '__main__': main(sys.argv[0:]) # # end of file