from dataclasses import dataclass import torch from torch.nn import functional as F import torch.nn as nn @dataclass class TranslationConfig: context_length = 32 ### Specifying the Context Length which the Transformer will be available to Process vocab_size = 8000 ### Vocab Size -- Tiktoken Vocab Size -- 50,257 also include <|endoftext|>, print(tiktoken.get_encoding('gpt2').n_vocab) n_layer = 4 ### No of the Transformer Decoder stacks n_head = 4 ### No of Attention Heads Used n_embd = 128 ### Dimensionality of the Word Embedding pad_idx = 0 bos_idx = 1 eos_idx = 2 class EncoderBlock(nn.Module): def __init__(self, config): super().__init__() self.config = config self.attn = nn.MultiheadAttention(embed_dim=config.n_embd, num_heads=config.n_head, batch_first=True, dropout=0.2) self.layer_norm1 = nn.LayerNorm(config.n_embd) self.multi_layer_perceptron = nn.Sequential( nn.Linear(config.n_embd, 4 * config.n_embd), nn.GELU(), nn.Linear(4 * config.n_embd, config.n_embd) ) self.layer_norm2 = nn.LayerNorm(config.n_embd) def forward(self, x): attn_out, _ = self.attn(x, x, x, attn_mask=None) x = x + attn_out x = self.layer_norm1(x) mlp_out = self.multi_layer_perceptron(x) x = x + mlp_out x = self.layer_norm2(x) return x class DecoderBlock(nn.Module): def __init__(self, config): super().__init__() self.config = config self.layer_norm1 = nn.LayerNorm(config.n_embd) self.attn = nn.MultiheadAttention(embed_dim=config.n_embd, num_heads=config.n_head, batch_first=True, dropout=0.2) self.layer_norm2 = nn.LayerNorm(config.n_embd) self.cross_attn = nn.MultiheadAttention(embed_dim=config.n_embd, num_heads=config.n_head, batch_first=True, dropout=0.2) self.layer_norm3 = nn.LayerNorm(config.n_embd) self.multi_layer_perceptron = nn.Sequential( nn.Linear(config.n_embd, 4 * config.n_embd), nn.GELU(), ### Activation Used as Gaussian Error Linear Unit nn.Linear(4 * config.n_embd, config.n_embd) ) def forward(self, x, encoder_output): B, T, C = x.shape ## Getting the Batch, time and Channel dimension from the input mask = torch.triu(torch.ones(T, T), diagonal=1).bool() ## Adding the Mask -- Disabling the Upper Triangle. attn_out, _ = self.attn(x, x, x, attn_mask=mask) ## It returns Attention Weights too x = x + attn_out ## Applied Residual Masked Multihead Attention with Causality x = self.layer_norm1(x) attn_out, _ = self.cross_attn(x, encoder_output, encoder_output, attn_mask=None) x = x + attn_out x = self.layer_norm2(x) mlp_out = self.multi_layer_perceptron(x) x = x + mlp_out x = self.layer_norm3(x) return x class Translation(nn.Module): def __init__(self, config): super().__init__() self.config = config self.transformer = nn.ModuleDict(dict( token_embedding_decoder = nn.Embedding(config.vocab_size, config.n_embd), pos_embedding_decoder = nn.Embedding(config.context_length, config.n_embd), decoder = nn.ModuleList([DecoderBlock(config) for _ in range(config.n_layer)]), layer_norm = nn.LayerNorm(config.n_embd), token_embedding_encoder = nn.Embedding(config.vocab_size, config.n_embd), pos_embedding_encoder = nn.Embedding(config.context_length, config.n_embd), encoder = nn.ModuleList([EncoderBlock(config) for _ in range(config.n_layer)]) )) self.lm_head = nn.Linear(config.n_embd, config.vocab_size) def num_parameters(self, only_trainable=False): parameters = self.parameters() if only_trainable: parameters = (p for p in parameters if p.requires_grad) return sum(p.numel() for p in parameters) def parameter_summary(self, only_trainable=False): summary = [] for name, parameter in self.named_parameters(): if only_trainable and not parameter.requires_grad: continue summary.append({ "name": name, "shape": list(parameter.shape), "num_parameters": parameter.numel(), "trainable": parameter.requires_grad, }) return summary def forward(self, src, tgt, targets=None): B, T_src = src.shape src_tok_emb = self.transformer.token_embedding_encoder(src) src_pos = torch.arange(0, T_src) src_pos_emb = self.transformer.pos_embedding_encoder(src_pos) enc_x = src_tok_emb + src_pos_emb for Block in self.transformer.encoder: enc_x = Block(enc_x) ### Passing through all 4 Encoder Block encoder_output = enc_x B, T_tgt = tgt.shape tgt_tok_emb = self.transformer.token_embedding_decoder(tgt) tgt_pos = torch.arange(0, T_tgt) tgt_pos_emb = self.transformer.pos_embedding_decoder(tgt_pos) dec_x = tgt_tok_emb + tgt_pos_emb for Block in self.transformer.decoder: dec_x = Block(dec_x, encoder_output) dec_x = self.transformer.layer_norm(dec_x) logits = self.lm_head(dec_x) loss = None if targets is not None: B, T, V = logits.shape logits = logits.view(B*T, V) targets = targets.view(B*T) loss = F.cross_entropy(logits, targets, ignore_index=self.config.pad_idx) return logits, loss @torch.no_grad() def generate(self, src, max_new_tokens=32): self.eval() # ---- Encoder ---- B, T_src = src.shape src_tok_emb = self.transformer.token_embedding_encoder(src) src_pos = torch.arange(0, T_src) src_pos_emb = self.transformer.pos_embedding_encoder(src_pos) enc_x = src_tok_emb + src_pos_emb for block in self.transformer.encoder: enc_x = block(enc_x) encoder_output = enc_x # ---- Decoder init ---- tgt = torch.full((B, 1), self.config.bos_idx, dtype=torch.long) for _ in range(max_new_tokens): T_tgt = tgt.shape[1] tgt_tok_emb = self.transformer.token_embedding_decoder(tgt) tgt_pos = torch.arange(0, T_tgt) tgt_pos_emb = self.transformer.pos_embedding_decoder(tgt_pos) dec_x = tgt_tok_emb + tgt_pos_emb for block in self.transformer.decoder: dec_x = block(dec_x, encoder_output) dec_x = self.transformer.layer_norm(dec_x) logits = self.lm_head(dec_x) # Take last token prediction next_token_logits = logits[:, -1, :] next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True) tgt = torch.cat([tgt, next_token], dim=1) # Stop if EOS generated if (next_token == self.config.eos_idx).all(): break return tgt