Transformer_script.py

import torch
import torch.nn as nn
from torch.nn import functional as F

# Hyperparameters
batch_size = 64 # how many independent sequences will we process in parallel?
block_size = 256 # what is the maximum context length for predictions?
max_iters = 5000
eval_interval = 500
learning_rate = 3e-4
device = 'cuda' if torch.cuda.is_available() else 'cpu'
eval_iters = 200
n_embed=384
n_head=6    # therefore every head (head_size)= 384//6 = 64 dim
dropout=0.2
# ------------

torch.manual_seed(2506)

# wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt
with open(r'data\input.txt', 'r', encoding='utf-8') as f:
    text = f.read()

# here are all the unique characters that occur in this text
chars = sorted(list(set(text)))
vocab_size = len(chars)
# create a mapping from characters to integers
stoi = { ch:i for i,ch in enumerate(chars) }
itos = { i:ch for i,ch in enumerate(chars) }
encode = lambda s: [stoi[c] for c in s] # encoder: take a string, output a list of integers
decode = lambda l: ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string

# Train and test splits
data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9*len(data)) # first 90% will be train, rest val
train_data = data[:n]
val_data = data[n:]

# data loading
def get_batch(split):
    # generate a small batch of data of inputs x and targets y
    data = train_data if split == 'train' else val_data
    ix = torch.randint(len(data) - block_size, (batch_size,))
    x = torch.stack([data[i:i+block_size] for i in ix])
    y = torch.stack([data[i+1:i+block_size+1] for i in ix])
    x, y = x.to(device), y.to(device)
    return x, y

@torch.no_grad() # tells pyTorch not to do backpropagation
def estimate_loss(): # average outs the loss for eval_iterations
    out = {}
    model.eval()
    for split in ['train', 'val']:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            X, Y = get_batch(split)
            logits, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out


class Head(nn.Module):
    '''This class is responsible for creating the Attention or Communication Mechanism - that lets the tokens communicate with past timesteps'''
    def __init__(self, head_size):
        super().__init__()
        self.key=nn.Linear(n_embed, head_size, bias=False) # The attentione heads are like channels of dimension say (32,16) i.e 32 other tokens define the word and 16 other timesteps contributes to its semantic meaning
        self.query=nn.Linear(n_embed, head_size, bias =False)
        self.value=nn.Linear(n_embed, head_size, bias =False)
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size))) # buffer is a variable that is needed for use, but is not a model param. 
        self.drop=nn.Dropout(dropout)

    def forward(self, x):

        # Extracting tensor shape params
        B,T,C=x.shape
        # mapping key, query and value to the input
        k=self.key(x)   
        q=self.query(x)
        v=self.value(x)

        weights= q @ k.transpose(-2,-1) * C**(-0.5) # Scaled Dot product
        weights=weights.masked_fill(self.tril[:T,:T]==0, float('-inf')) # type: ignore
        weights=F.softmax(weights, dim=-1)
        weights = self.drop(weights)

        # Perform weighted aggregation of values
        out=weights @ v
        return out

# Improvement: I - Multi-head Attention
# Now that we are done with Single-Head Attention - lets implement Multi_head attention - replicas of how many heads we want.
# Each head is responsible for some definite and specific purpose - say head 1 captures relationship of how are vowels related, while say head 2 maps how the punctuation contributes to the next token 
class MultiHead(nn.Module):
    '''Multiple Heads of Self Attention in parallel'''
    def __init__(self, num_heads, head_size):
        super().__init__()
        self.heads=nn.ModuleList([Head(head_size) for _ in range(num_heads)]) # using the function from base class we can create a list of submodules within a list
        # we shall create another linear layer: projection that forks off from the main path, directly takes the multi_head concatenated vectors and perform a linear transformation on it and joins back to the main path
        self.proj = nn.Linear(n_embed, n_embed) # takes n_embed input and outputs linearly transformed n_embed....the dim remains same 
        self.drop=nn.Dropout(dropout)

    def forward(self,x):
        # here we are just concatenating all the separate heads over the last dimention(channel dim) into a single unit
        out = torch.cat([h(x) for h in self.heads], dim=-1) # take every single heads from the list of multiheads and for each of them compute k,q,v using x as argument and then concat all of them back
        out = self.proj(out)
        return self.drop(out) # here we again join the projecrtion back into the main residual pathway

# Improvement: II - FeedForward Network
#  Now by far whatever improvements we did, still we compute logits directly after the attention level. Though the loss decreased from 2.4 in ur first version to 2.2 in our latest- but still the logits were calculated directly, therefore the architecture don't get enough time to think upon what they found from other tokens
# Thus we need a linear single layer of network with relu non-linearity for propagation of the logits - with feedforward the loss further decreased from 2.28 to 2.23 
class FeedForward(nn.Module):
    '''a simnple linear layer followed by non-linearity'''

    def __init__(self, n_embed):
        super().__init__()
        self.ffnet=nn.Sequential(
            nn.Linear(n_embed, 4* n_embed),
            nn.ReLU(),
            nn.Linear(4*n_embed, n_embed), # similar skip connections we build here for optimization of the network
            nn.Dropout(dropout)
        )
# why the inner layer of feedforward net had 4 times the num of embeddings...i.e 4 X 32?? This is done purely based on the paper ("Attention is all You Need"). We do so to increase the computation and sensitivity upto 4 times in the forked branch and then again reduce it to normal n_embed while joining the forked branch to the main branch - After training we can see the loss further decreasd to 1.913(train) and 2.05(val)

    def forward(self, x):
        return self.ffnet(x) # executing each tokens on the feed forward neural net
        


# Improvement: III - Block Architecture
# Now by far whatever we did in the Decoder architecture (Pos encoding, sh, mh, scaled dot prod, feedforwd) can be trated as block of calculation & computation. Now to improve performance we can repeat these block k times 
class Block(nn.Module):
    '''Transformer Block: Communication (Attention) followed by Computation'''

    def __init__(self, n_embed, n_head):
        super().__init__()
        head_size=n_embed//n_head
        self.attention=MultiHead(n_head, head_size)
        self.ffwd=FeedForward(n_embed)
        self.ln1=nn.LayerNorm(n_embed) # creating 2 layernorm layers
        self.ln2=nn.LayerNorm(n_embed)

    def forward(self, x):
        x = x + self.attention(self.ln1(x)) # why x+ ?? This is an upgrade where we use skip connections for optimization, we fork off from the main bakprop branch, compute and then join the branch back again
        x = x + self.ffwd(self.ln2(x)) # This prevents 
        return x
# Block of Transformer creates Abstraction by judt defining the predefined methods and objects for Communication and Computation repeatedly
# But repeating blocks also maked the network deep, therefore we need to optimize the network - techniques: Residual(skip) Connection, batch norm, dropout  

# Improvement: IV - SKip Connection: 
# During Back Propagation - : Gradients (the error signals) travel backward through those same layers to tweak the weights via the Chain Rule. In a very deep network without skip connections, as the gradient travels backward through dozens of linear layers and activation functions, it gets multiplied over and over again: If the weights are slightly less than (1.0), the gradient continuously shrinks until it becomes virtually zero (0.0000...). The early layers of the model stop learning completely.If the weights are slightly greater than (1.0), the gradient multiplies exponentially until it hits infinity (NaN), crashing your training entirely.A skip connection modifies the architectural layout of a neural network block. Instead of forcing the data vector \(x\) to strictly pass through a layer \(F(x)\), the connection splits the path, allowing the original raw input to hop over the layer and be added directly to the output. Output = x + f(x) --------> Changes made under forard function in Block class and applied to MultiHead and Feedfwd class.


# Improvement: V - Layer Norm
# When we update a neural network's parameters via an optimizer, the weights of the earliest layers change. Consequently, the distributions of inputs to the deeper layers shift constantly throughout training. This shifting is called Internal Covariate Shift. Deeper layers are forced to chase a moving target, slowing down learning convergence and increasing the likelihood of vanishing or exploding gradients. 
# class BatchNorm1D:
#     '''BatchNorm1D -> LayerNorm. This is used as sequences in a batch may vary in length, therefore batchnorm would be messy - thus we use layer norm which is same in principle just applied for individual vectors independently along the feature axis, making it independent of sequence length and batch size.'''
#----------------------------------------------------------------------------------------------------------------------------------
#     def __init__(self, dim, eps=1e-5, momentum=0.1):
#         self.eps=eps
#         self.momentum=momentum
#         self.training=True

#         #parameters with backprop
#         self.gamma=torch.ones(dim)
#         self.beta=torch.zeros(dim)

#         #Buffers (trained with a running momentum update)
#         self.running_mean=torch.ones(dim)
#         self.running_var=torch.ones(dim)

#     def __call__(self, x):
#         if self.training:
#             xmean=x.mean(0, keepdim=True)
#             xvar=x.var(0, keepdim=True)
#         else:
#             xmean=self.running_mean
#             xvar=self.running_var
#         xhat=(x-xmean)/torch.sqrt(xvar+self.eps) # implying the formula to compute xhat; this is the Normalize to unite variance step that makes the inter layer architecture, gradients stable
#         self.out=self.gamma*xhat + self.beta
#         if self.training:
#             self.running_mean = (1-self.momentum)*self.running_mean + self.momentum* xmean
#             self.running_var = (1-self.momentum)*self.running_var + self.momentum* xvar
#         return self.out

#     def parameters(self):
#         return [self.gamma, self.beta]

# torch.manual_seed(2506)
# module=BatchNorm1D(dim=100)
# x=torch.randn(32, 100)
# x=module(x)
# x.shape
#----------------------------------------------------------------------------------------------------------------------------------
# We now know why not batchnorm and why layernorm applies to our Transformer architecture - thus layernorm implemented in Block class
# Present status after LayerNorm: step 4800: train loss 1.9299, val loss 2.0375


# Improvement VI: We add DropOut layers before the forked branch connects back ot residuals - This drops a % of neurons from each layer thereby generalizing the model - DropOut added at head, MultiHead, D
# status: step 4800: train loss 2.0543, val loss 2.1065

# Gradually improved architeture of a naive Bigram model implemnted with M-H attention, FeedForward class
class BigramLanguageModel(nn.Module):

    # Configurations 
    def __init__(self):
        super().__init__()
        # each token directly reads off the logits for the next token from a lookup table
        self.token_embedding_table = nn.Embedding(vocab_size, n_embed)
        # We besides creating the embedding table to store the embedding vectors for every token also create a positional_embeddding table for positional encoding
        self.positional_embedding_table=nn.Embedding(block_size,n_embed)
        # self.sa = Head(n_embed) # creating the self attention head - commented out since we have improved to Multi_head attention
        # Now that we have evolved to Multiple heads, lets comment out Self-head and implement Multi-head attention
        # self.mha = MultiHead(num_heads=4, head_size=n_embed//4) # we are dividing the entire num of embedd. into 4 sets for 4 heads of uniform size. (4 * 8Dim = 32Dim embedd vector)
        # self.feedfwd = FeedForward(n_embed)

        # We now comment out all the isolated unit level operations, since we have created a Block containing the same operations and we can just Sequentially repeat them
        self.blocks = nn.Sequential(
            Block(n_embed=n_embed, n_head=n_head),
            Block(n_embed=n_embed, n_head=n_head),
            Block(n_embed=n_embed, n_head=n_head),
            Block(n_embed=n_embed, n_head=n_head),
            Block(n_embed=n_embed, n_head=n_head),
            Block(n_embed=n_embed, n_head=n_head),
            nn.LayerNorm(n_embed), # We also normalize each vectors independently before providing it to the final linear layer
        )
        self.lm_head = nn.Linear(n_embed, vocab_size)
                

    def forward(self, idx, targets=None):
        B,T = idx.shape
        # idx and targets are both (B,T) tensor of integers
        tok_embd = self.token_embedding_table(idx) # (B,T,C)
        pos_embd = self.positional_embedding_table(torch.arange(T, device=device)) # creating pos embd vectors from 0 to range T-1 of shape (T,C)
        x=tok_embd + pos_embd

        # x = self.mha(x)
        # s=self.feedfwd(x) # (B,T,C)

        x=self.blocks(x)
        logits=self.lm_head(x)

        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets)

        return logits, loss

    def generate(self, idx, max_new_tokens):
        # idx is (B, T) array of indices in the current context
        for _ in range(max_new_tokens):

            # idx can go out of range that's not present in the scope of Embeddings table, thus we need to filter it
            idx_cond=idx[:, -block_size:]
            # get the predictions
            logits, loss = self(idx_cond) # forward(idx_cond)
            # focus only on the last time step
            logits = logits[:, -1, :] # becomes (B, C)
            # apply softmax to get probabilities
            probs = F.softmax(logits, dim=-1) # (B, C)
            # sample from the distribution
            idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)
            # append sampled index to the running sequence
            idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)
        return idx


model = BigramLanguageModel()
m = model.to(device)
print(sum(p.numel() for p in m.parameters())/1e6, 'M parameters') # Roughly around 10M params on a 1M char dataset

# create a PyTorch optimizer
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

for iter in range(max_iters):

    # every once in a while evaluate the loss on train and val sets
    if iter % eval_interval == 0:
        losses = estimate_loss()
        print(f"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")

    # sample a batch of data
    xb, yb = get_batch('train')

    # evaluate the loss
    logits, loss = model(xb, yb)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

# generate from the model
context = torch.zeros((1, 1), dtype=torch.long, device=device)
print(decode(m.generate(context, max_new_tokens=500)[0].tolist()))