Building the transformer Model

Importing data

import torch
import math
from torch import nn

import sys
import os
sys.path.append(os.path.abspath(os.path.join("..", "seq2seq")))
from data_pipeline import (
    train_loader, test_pairs,pairs,
    english_word_to_index, english_index_to_word,
    arabic_word_to_index, arabic_index_to_word,
    tokenize, encode_source,
    PAD_IDX, SOS, EOS, UNK,
)
import sys
import os
sys.path.append(os.path.abspath(os.path.join("..", "seq2seq")))


SOS_IDX = arabic_word_to_index[SOS]
EOS_IDX = arabic_word_to_index[EOS]

src_batch, dec_in_batch, dec_tgt_batch, src_lengths = next(iter(train_loader))
lens = [len(tokenize(en)) for en, _ in test_pairs]
print(
    f"EN vocab {len(english_word_to_index)} | AR vocab {len(arabic_word_to_index)}")
print(f"PAD={PAD_IDX} SOS={SOS_IDX} EOS={EOS_IDX}")
print(
    f"src_batch {tuple(src_batch.shape)} | source length range {min(lens)}..{max(lens)} words")

EN vocab 793 | AR vocab 1068
PAD=0 SOS=1 EOS=2
src_batch (32, 14) | source length range 8..16 words

Below is the self attention mechanism that makes the model create relationships between token inputs. Creating Q, K, V for each token. The returned Output holds each token’s new representation — a blend of the other tokens most relevant to it. While the attention_weights stores how much each token attended to every other token.

class SelfAttention(nn.Module):
    def __init__(self, d_model):
        super().__init__()

        self.d_model = d_model

        # Learnable transformation
        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)

    def forward(self, x):
        # x: (batch, seq_len, d_model)

        # 1. Create Q, K, V
        Q = self.W_Q(x)
        K = self.W_K(x)
        V = self.W_V(x)

        # 2. QK^T
        scores = torch.matmul(Q, K.transpose(-2, -1))

        # 3. Scale
        scores = scores / math.sqrt(self.d_model)

        # 4. Softmax
        attention_weights = torch.softmax(scores, dim=-1)

        # 5. Weighted sum of V
        output = torch.matmul(attention_weights, V)

        return output, attention_weights

MultiHeadAttention class is doing the self attention between input tokens in encoder, and doing the cross attention between Q from the decoder, and K and V from the encoder’s output (the “memory”).

We will set the forward() to accept kv_input so we can later use it when building the decoder.

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, 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
        
        self.W_Q = nn.ModuleList([nn.Linear(self.d_k, self.d_k) for _ in range(num_heads)])
        self.W_K = nn.ModuleList([nn.Linear(self.d_k, self.d_k) for _ in range(num_heads)])
        self.W_V = nn.ModuleList([nn.Linear(self.d_k, self.d_k) for _ in range(num_heads)])
        
        # The weights to the ouptut, need to combine the result of the heads above before calling this.
        self.W_O = nn.Linear(d_model, d_model)

    def forward(self, query_input, kv_input=None, mask=None):
        
        if kv_input is None:
            kv_input = query_input   # plain self-attention: Q, K, V all from the same source
        
        # split the input into per-head chuncks
        query_chunks = query_input.split(self.d_k, dim=-1)
        kv_chunks = kv_input.split(self.d_k, dim=-1)
                
        
        head_outputs = []
        head_weights = []
        
        for i in range(self.num_heads):
            Q = self.W_Q[i](query_chunks[i])
            K = self.W_K[i](kv_chunks[i])
            V = self.W_V[i](kv_chunks[i])
            
            scores = Q @ K.transpose(-2, -1) / math.sqrt(self.d_k)
            # adding the mask padding (zero out scores for padding tokens)
            if mask is not None:
                if mask.dim() == 2:
                    mask = mask.unsqueeze(1)   # (batch, seq_k) -> (batch, 1, seq_k), broadcasts across all queries
                scores = scores.masked_fill(mask, float('-inf'))

            attention_weights = torch.softmax(scores, dim=-1)
            output = torch.matmul(attention_weights, V)
            
            #Store each head outputs and weights
            head_outputs.append(output)
            head_weights.append(attention_weights)
        
        # concatenate all heads outputs
        concat = torch.cat(head_outputs, dim=-1)   # (batch, seq_len, d_model)
        
        # Project the concatenated outputs to a d_model learnable
        output = self.W_O(concat)   # (batch, seq_len, d_model)

        # stack the attention weights (batch, num_heads, seq_len, seq_len)
        attention_weights = torch.stack(head_weights, dim=1)
        
        return output, attention_weights

Self attention has no sense of order, if we shuffle the input words the attention scores wouldn’t change. PositionalEncoding fixes this by adding a fixed sin/cos pattern to each token’s embedding based on its position in the sequence, giving the model a way to tell where a token sits, not just what it is.

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=100):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2)
                             * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        # fixed, not learned — moves with .to(device) but isn't a Parameter
        self.register_buffer("pe", pe)

    def forward(self, x):
        # x: (batch, seq_len, d_model)
        return x + self.pe[:x.size(1)]

The ToyEncoder takes the self attention and wraps it with residual connection and a LayerNorm, followed by a position-wise feed-forward network. Making it trainable and stackable.

class ToyEncoder(nn.Module):
    def __init__(self, d_model=64, num_heads=4, ff_hidden=256):
        super().__init__()

        # first get the attention
        #self.self_Atten = SelfAttention(d_model)
        self.self_Atten =  MultiHeadAttention(d_model, num_heads)

        # Layer normalize
        self.norm1 = nn.LayerNorm(d_model)

        # feed forward net
        self.ff = nn.Sequential(
            nn.Linear(d_model, ff_hidden),
            nn.ReLU(),
            nn.Linear(ff_hidden, d_model)
        )

        # Layer normalization after feed-forward
        self.norm2 = nn.LayerNorm(d_model)
        
        #Define a drop out layer
        self.dropout = nn.Dropout(0.1)

    def forward(self, x, mask):

        # -------------------------
        # 1. Self-Attention
        # -------------------------
        attention_outputs, attention_weights = self.self_Atten(x, mask=mask)

        # residual connection +LayerNorm
        x = self.norm1(x + self.dropout(attention_outputs))

        # -------------------------
        # 2. Feed Forward
        # -------------------------
        ff_output = self.ff(x)

        # Residual connection + LayerNorm
        x = self.norm2(x + self.dropout(ff_output))

        return x, attention_weights

Encoder is the full encoder stack. It owns the token embedding and positional encoding, applied once at the very start, then passes the result through several stacked ToyEncoder layers. Each layer refines the token representations a little further, building on the context the previous layer already capture

class Encoder(nn.Module):
    def __init__(self, vocab_size, d_model=64, num_heads=4, ff_hidden=256, num_layers=2):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=PAD_IDX)
        self.pos_encoding = PositionalEncoding(d_model)
        self.layers = nn.ModuleList(
            [ToyEncoder(d_model, num_heads, ff_hidden) for _ in range(num_layers)])
        self.dropout = nn.Dropout(0.1)

    def forward(self, token_ids, mask):
        x = self.embedding(token_ids)
        x = self.pos_encoding(x)
        x = self.dropout(x)
        
        for layer in self.layers:
            x, attention_weights = layer(x, mask=mask)
        return x, attention_weights

padding_mask = (src_batch == PAD_IDX)
vocab_size = len(english_word_to_index)
encoder = Encoder(vocab_size, d_model=64, num_heads=4, ff_hidden=256, num_layers=2)
x, attention_weights = encoder(src_batch, padding_mask)

print(x.shape)
print("---------")
print(attention_weights.shape)

torch.Size([32, 15, 64])
---------
torch.Size([32, 4, 15, 15])

ex = 5
print(src_batch[ex])
# head 0, query position 0, over all 14 keys
print(attention_weights[ex, 0, 0, :])

tensor([ 18,  14,  55,  17, 518, 120,  60,  59, 581,  33,   2,   0,   0,   0,
          0])
tensor([0.0996, 0.0501, 0.0960, 0.0849, 0.1130, 0.1511, 0.0719, 0.1111, 0.0688,
        0.0770, 0.0767, 0.0000, 0.0000, 0.0000, 0.0000],
       grad_fn=<SelectBackward0>)

Decoder

In the decoder, we will use a different structure. There will be two separate MultiHeadAttention instances: - Decoder attending to itself. - Decoder attending to the encoder’s memory.

class ToyDecoder(nn.Module):
    def __init__(self, d_model, num_heads, ff_hidden):
        super().__init__()
        # first get the attention
        self.masked_self_attn = MultiHeadAttention(d_model, num_heads) # Decoder attending to itself.
        # Layer normalize
        self.norm1 = nn.LayerNorm(d_model)
        

        # second attention layer
        # Decoder attending to the encoder's memory.
        self.cross_attn = MultiHeadAttention(d_model, num_heads)
        # Layer normalize
        self.norm2 = nn.LayerNorm(d_model)
        
        # Define dropout
        self.dropout = nn.Dropout(0.1)
        
        # feed forward net
        self.ff = nn.Sequential(
            nn.Linear(d_model, ff_hidden),
            nn.ReLU(),
            nn.Linear(ff_hidden, d_model)
        )
        # Layer normalization after feed-forward
        self.norm3 = nn.LayerNorm(d_model)
    
    def forward(self, x, memory, tgt_mask, src_mask):
        
        # Masked self-attention
        self_out, self_weights = self.masked_self_attn(x, mask=tgt_mask)
        x = self.norm1(x + self.dropout(self_out))
        
        # Cross attention
        cross_out, cross_weights = self.cross_attn(x, memory, mask=src_mask)
        x = self.norm2(x + self.dropout(cross_out))
        
        # Feed-forward
        ff_out = self.ff(x)
        x = self.norm3(x + self.dropout(ff_out))
        
        return x, self_weights, cross_weights

# (batch, src_len) — padding only
src_mask = (src_batch == PAD_IDX)
# (batch, tgt_len) — target padding
tgt_pad_mask = (dec_in_batch == PAD_IDX)

tgt_len = dec_in_batch.size(1)
causal_mask = torch.triu(torch.ones(
    tgt_len, tgt_len, dtype=torch.bool), diagonal=1)  # (tgt_len, tgt_len)

# block a (query, key) pair if the key is padding OR the key is a future position
tgt_mask = tgt_pad_mask.unsqueeze(1) | causal_mask.unsqueeze(
    0)   # (batch, tgt_len, tgt_len)

print(src_mask.shape)   # (32, 14)
print(tgt_mask.shape)   # (32, 11, 11)  <- already varies per query row

torch.Size([32, 15])
torch.Size([32, 15, 15])

# 1. Encoder memory (uses src_mask, no causal component — full source is visible)
memory, _ = encoder(src_batch, src_mask)          # (batch, src_len, d_model)

# 2. Embed the target side (normally Decoder wrapper owns this — done manually here)
d_model = 64
tgt_embedding = nn.Embedding(
    len(arabic_word_to_index), d_model, padding_idx=PAD_IDX)
tgt_pos_encoding = PositionalEncoding(d_model)

x = tgt_embedding(dec_in_batch)
x = tgt_pos_encoding(x)                            # (batch, tgt_len, d_model)

# 3. Run one decoder layer
decoder_layer = ToyDecoder(d_model=d_model, num_heads=4, ff_hidden=256)
out, self_weights, cross_weights = decoder_layer(x, memory, tgt_mask, src_mask)

print("Input batch: ", src_batch.shape)
print("Input target embedding: ", x.shape)
#print(out.shape)            # (32, 11, 64)   — one vector per target token
print("Self attending: ",self_weights.shape)
print("Cross attending: ",cross_weights.shape)

Input batch:  torch.Size([32, 15])
Input target embedding:  torch.Size([32, 15, 64])
Self attending:  torch.Size([32, 4, 15, 15])
Cross attending:  torch.Size([32, 4, 15, 15])

Self attending: (32, 4, 15, 15)

For each of the 32 sentences and each of the 4 heads, a 15×15 matrix: how much each of the 15 target (Arabic) positions attends to each of the other 15 target positions.

Cross attending: (32, 4, 15, 14)

For each sentence and head, a 15×14 matrix: how much each of the 15 target positions attends to each of the 14 source (English) positions.

class Decoder(nn.Module):
    def __init__(self, vocab_size, d_model, num_heads, ff_hidden, num_layers):
        super().__init__()
        
        self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=PAD_IDX)
        self.pos_encoding = PositionalEncoding(d_model)
        self.layers = nn.ModuleList([ToyDecoder(d_model, num_heads, ff_hidden) for _ in range(num_layers)])
        self.dropout = nn.Dropout(0.1)
    
    def forward(self, token_ids, memory, tgt_mask, src_mask):
        x = self.embedding(token_ids)
        x = self.pos_encoding(x)
        x = self.dropout(x)
        
        for layer in self.layers:
            x, self_weights, cross_weights = layer(x, memory, tgt_mask, src_mask)
        
        
        return x, self_weights, cross_weights

Transformer Wrapper

class Transformer(nn.Module):
    def __init__(self, src_vocab_size, tgt_vocab_size, d_model, num_heads, ff_hidden, num_layers):
        super().__init__()
        
        self.encoder = Encoder(src_vocab_size, d_model, num_heads, ff_hidden, num_layers)
        self.decoder = Decoder(tgt_vocab_size, d_model, num_heads, ff_hidden, num_layers)
        self.output_layer = nn.Linear(d_model, tgt_vocab_size)
        
    def forward(self, src, dec_in, src_mask, tgt_mask):
        memory, _ = self.encoder(src, src_mask)
        x, self_weights, cross_weights = self.decoder(dec_in, memory, tgt_mask, src_mask)
        logits = self.output_layer(x)
        return logits, self_weights, cross_weights

Training Loop

class NoamScheduler:
    def __init__(self, optimizer, d_model, warmup_steps=400):
        self.optimizer = optimizer
        self.d_model = d_model
        self.warmup_steps = warmup_steps
        self.step_num = 0

    def step(self):
        self.step_num += 1
        lr = (self.d_model ** -0.5) * min(self.step_num ** -
                                          0.5, self.step_num * self.warmup_steps ** -1.5)
        for group in self.optimizer.param_groups:
            group['lr'] = lr
        self.optimizer.step()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = Transformer(
    src_vocab_size=len(english_word_to_index),
    tgt_vocab_size=len(arabic_word_to_index),
    d_model=64,
    num_heads=4,
    ff_hidden=256,
    num_layers=2,
).to(device)

optimizer = torch.optim.Adam(
    model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9)
scheduler = NoamScheduler(optimizer, d_model=64, warmup_steps=400)

# padding positions don't count toward the loss
loss_fn = nn.CrossEntropyLoss(ignore_index=PAD_IDX)


def make_masks(src_batch, dec_in_batch):
    # (batch, src_len)
    src_mask = (src_batch == PAD_IDX)
    # (batch, tgt_len)
    tgt_pad_mask = (dec_in_batch == PAD_IDX)
    tgt_len = dec_in_batch.size(1)
    causal_mask = torch.triu(
        torch.ones(tgt_len, tgt_len, dtype=torch.bool, device=dec_in_batch.device), diagonal=1
    )
    tgt_mask = tgt_pad_mask.unsqueeze(1) | causal_mask.unsqueeze(0)
    return src_mask, tgt_mask


num_epochs = 50

for epoch in range(num_epochs):
    model.train()
    total_loss, num_batches = 0.0, 0

    for src_batch, dec_in_batch, dec_tgt_batch, src_lengths in train_loader:
        src_batch = src_batch.to(device)
        dec_in_batch = dec_in_batch.to(device)
        dec_tgt_batch = dec_tgt_batch.to(device)

        src_mask, tgt_mask = make_masks(src_batch, dec_in_batch)

        logits, _, _ = model(src_batch, dec_in_batch, src_mask, tgt_mask)
        # logits: (batch, tgt_len, vocab_size)

        loss = loss_fn(
            # (batch*tgt_len, vocab_size)
            logits.reshape(-1, logits.size(-1)),
            dec_tgt_batch.reshape(-1)               # (batch*tgt_len,)
        )

        optimizer.zero_grad()
        loss.backward()
        scheduler.step()

        total_loss += loss.item()
        num_batches += 1

    print(f"epoch {epoch+1:3d}  avg loss {total_loss/num_batches:.4f}")

epoch   1  avg loss 6.5310
epoch   2  avg loss 5.5122
epoch   3  avg loss 5.0480
epoch   4  avg loss 4.8401
epoch   5  avg loss 4.5144
epoch   6  avg loss 4.1822
epoch   7  avg loss 3.7923
epoch   8  avg loss 3.5297
epoch   9  avg loss 3.2323
epoch  10  avg loss 2.9492
epoch  11  avg loss 2.6797
epoch  12  avg loss 2.4027
epoch  13  avg loss 2.1495
epoch  14  avg loss 1.9248
epoch  15  avg loss 1.7699
epoch  16  avg loss 1.5914
epoch  17  avg loss 1.4522
epoch  18  avg loss 1.3394
epoch  19  avg loss 1.2425
epoch  20  avg loss 1.1539
epoch  21  avg loss 1.0921
epoch  22  avg loss 0.9647
epoch  23  avg loss 0.9167
epoch  24  avg loss 0.8884
epoch  25  avg loss 0.8359
epoch  26  avg loss 0.7675
epoch  27  avg loss 0.7356
epoch  28  avg loss 0.6936
epoch  29  avg loss 0.6781
epoch  30  avg loss 0.6705
epoch  31  avg loss 0.6416
epoch  32  avg loss 0.5704
epoch  33  avg loss 0.5726
epoch  34  avg loss 0.5906
epoch  35  avg loss 0.5719
epoch  36  avg loss 0.5489
epoch  37  avg loss 0.5103
epoch  38  avg loss 0.4615
epoch  39  avg loss 0.4349
epoch  40  avg loss 0.4674
epoch  41  avg loss 0.4585
epoch  42  avg loss 0.4273
epoch  43  avg loss 0.3965
epoch  44  avg loss 0.3855
epoch  45  avg loss 0.3731
epoch  46  avg loss 0.4120
epoch  47  avg loss 0.3976
epoch  48  avg loss 0.3640
epoch  49  avg loss 0.3773
epoch  50  avg loss 0.3716

def translate(sentence, model, max_len=20, device=device):
    model.eval()
    with torch.no_grad():
        # 1. Encode the source sentence once
        src_ids = torch.tensor(encode_source(tokenize(
            # (1, src_len)
            sentence)), dtype=torch.long, device=device).unsqueeze(0)
        # (1, src_len)
        src_mask = (src_ids == PAD_IDX)

        # (1, src_len, d_model) — computed once, reused every step
        memory, _ = model.encoder(src_ids, src_mask)

        # 2. Autoregressive decoding, starting from SOS
        dec_in = torch.tensor(
            [[SOS_IDX]], device=device)                 # (1, 1)

        for _ in range(max_len):
            tgt_len = dec_in.size(1)
            causal_mask = torch.triu(
                torch.ones(tgt_len, tgt_len, dtype=torch.bool, device=device), diagonal=1
            )
            # (1, tgt_len, tgt_len) — no padding yet, every position generated so far is real
            tgt_mask = causal_mask.unsqueeze(0)

            dec_out, _, _ = model.decoder(dec_in, memory, tgt_mask, src_mask)
            # (1, tgt_len, vocab_size)
            logits = model.output_layer(dec_out)

            # prediction for the position just generated
            next_token = logits[0, -1, :].argmax(-1).item()
            dec_in = torch.cat([dec_in, torch.tensor(
                [[next_token]], device=device)], dim=1)

            if next_token == EOS_IDX:
                break

        # 3. Convert ids back to words (skip the leading SOS, stop at EOS)
        words = []
        for idx in dec_in[0, 1:].tolist():
            if idx == EOS_IDX:
                break
            words.append(arabic_index_to_word[idx])

        return " ".join(words)

#en_sentence, ar_reference = test_pairs[5]
en_sentence, ar_reference = pairs[28]
print("EN     :", en_sentence)
print("Target :", ar_reference)
print("Model  :", translate(en_sentence, model))

EN     : I am going to play tennis in the afternoon
Target : سألعب التنس بعد الظهر
Model  : سألعب التنس بعد الظهر