transformer_II.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as data
import torch.optim as optim

from datasets import load_dataset
from transformers import AutoTokenizer

import matplotlib.pyplot as plt
import seaborn as sns

seq_len=10
embed_dim=16
num_heads=8
feed_forward_ratio=4
size=100
batch_size=2

device="cuda"

class MultiheadAttention(nn.Module):
    def __init__(self, embed_dim,num_heads):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

    def forward(self,q,k,v,mask=None): # q,k,v has shape [batch_size,seq_len,embed_dim]
        batch_size=q.shape[0]
        seq_len = q.shape[1]
        q = q.reshape(batch_size,seq_len,self.num_heads,self.head_dim).permute(0,2,1,3) # [batch_size,num_heads,seq_len,head_dim]
        kT = k.reshape(batch_size,seq_len,self.num_heads,self.head_dim).permute(0,2,3,1) # [batch_size,num_heads,head_dim,seq_len]
        v = v.reshape(batch_size,seq_len,self.num_heads,self.head_dim).permute(0,2,1,3) # [batch_size,num_heads,seq_len,head_dim]
        attention_logits=q@kT/torch.sqrt(torch.tensor(self.embed_dim)) # [batch_size,num_heads,seq_len,seq_len]
        if mask is not None:
            attention_logits=attention_logits.masked_fill(mask==0, -torch.inf)
        attn_weights=F.softmax(attention_logits, dim=-1)
        #print(f"attn_weights has shape {attn_weights.shape}")
        #print(f"v has shape {v.shape}")
        atten_values=attn_weights@v # [batch_size,num_heads,seq_len,head_dim]
        return atten_values.permute(0,2,1,3).reshape(batch_size,seq_len,self.embed_dim), attn_weights

class transformer(nn.Module):
    def __init__(self, vocab_size, positional_encoding, embed_dim, num_heads, feed_forward_ratio):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.feed_forward_ratio = feed_forward_ratio
        self.embed_layer=nn.Linear(embed_dim,3*embed_dim)
        self.MultiheadAttention=MultiheadAttention(embed_dim,num_heads)
        self.layer_norm=nn.LayerNorm(embed_dim)
        self.ffn_layers=nn.Sequential(
            nn.Linear(embed_dim,feed_forward_ratio*embed_dim),
            nn.ReLU(),
            nn.Linear(feed_forward_ratio*embed_dim,embed_dim)
        )
        self.input_embedding=nn.Embedding(vocab_size, embed_dim)
        self.output_embedding=nn.Linear(embed_dim, vocab_size)
        self.positional_encoding=positional_encoding(embed_dim)
       
    
    def forward(self, x): #x has shape [batch_size,seq_len] each element is an integer in range(0,vocab_size)
        x=self.input_embedding(x) # [batch_size,seq_len] -> [batch_size,seq_len,embed_dim]
        x=self.positional_encoding(x)
        #print(f"x has shape {x.shape}")
        qkv=self.embed_layer(x) # [batch_size,seq_len,embed_dim] -> [batch_size,seq_len,3*embed_dim] 
        q, k, v= qkv.chunk(3,dim=-1) 
        batch_size, seq_len, _=x.shape
        mask=torch.tril(torch.ones(seq_len,seq_len), diagonal=0).unsqueeze(0).unsqueeze(0).repeat(batch_size,self.num_heads,1,1).to(device) # [1,1,seq_len,seq_len]
        #print(f"mask has shape {mask.shape}")
        attn_value, _=self.MultiheadAttention(q,k,v,mask)
        #print(f"attn_value has shape {attn_value.shape}")
        #print(f"x has shape {x.shape}")
        x=x+attn_value
        x=self.layer_norm(x)
        x= self.ffn_layers(x)
        x=self.layer_norm(x)
        x=self.output_embedding(x) # [batch_size,seq_len,embed_dim] -> [batch_size,seq_len,vocab_size]
        return x # [batch_size,seq_len,vocab_size] 
    



class loss_function(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, pred_logits, label):
        batch_size, seq_len, embed_dim=pred_logits.shape
        return F.cross_entropy(pred_logits.reshape(batch_size*seq_len, embed_dim),label.reshape(batch_size*seq_len))

loss_function=loss_function()

class positional_encoding(nn.Module):
    def __init__(self,embed_dim):
        super().__init__()
        self.embed_dim=embed_dim
    
    def forward(self,x): # x has shape [batch_size, seq_len, seq_len]
        
        batch_size, seq_len, embed_dim=x.shape
        pe=torch.arange(0,seq_len).unsqueeze(1) # [seq_len,1]
        embed=torch.arange(embed_dim)
        embed1=torch.where(embed%2==0,0,1)*torch.sin(pe*(100**(-embed/embed_dim)).unsqueeze(0)) # [seq_len,embed_dim]
        embed2=torch.where(embed%2==0,1,0)*torch.sin(pe*(100**(-embed/embed_dim)).unsqueeze(0)) # [seq_len,embed_dim]
        pe_embed=(embed1+embed2).unsqueeze(0).repeat(batch_size,1,1).to(device) # [batch_size,seq_len,embed_dim]
        
        return x+pe_embed


# load dataset and tokenizer from Huggingface
dataset_name="tiny_shakespeare"
tokenizer_name="gpt2"
dataset=load_dataset(dataset_name, split="train", trust_remote_code=True)
tokenizer=AutoTokenizer.from_pretrained(tokenizer_name)
if tokenizer.pad_token is None:
    tokenizer.pad_token=tokenizer.eos_token
vocab_size=tokenizer.vocab_size

# print(dir(tokenizer))  # all the methods and attributes of the tokenizer
# print(f"max_len for the tokenizer: {tokenizer.model_max_length}")

#sample_data=next(iter(dataset)) # it is a continuous text
#for item in sample_data.keys():
    #print(item) # to see the keys - here only key is "text"

class tokenized_text_dataset(data.Dataset):
    def __init__(self, dataset,tokenizer, seq_len):
        super().__init__()
        self.dataset=dataset
        self.tokenizer=tokenizer
        self.seq_len=seq_len
        self.tokenized_dataset=[]
        for item in dataset: # all the text is made into a continuous text
            self.tokenized_dataset.extend(tokenizer(item["text"], truncation=True, max_length=tokenizer.model_max_length)["input_ids"])
        self.size=max(0,len(self.tokenized_dataset)-seq_len-1)

    def __len__(self):
        return self.size
    
    def __getitem__(self,idx):
        input_seq=torch.tensor(self.tokenized_dataset[idx:idx+self.seq_len], dtype=torch.long) # important to make it a tensor explicitly
        output_seq=torch.tensor(self.tokenized_dataset[idx+1:idx+self.seq_len+1], dtype=torch.long) # important to make it a tensor explicitly
        return input_seq, output_seq
    
text_dataset=tokenized_text_dataset(dataset,tokenizer,seq_len)
dataloader=data.DataLoader(text_dataset, batch_size=batch_size, shuffle=False)
input_seq, output_seq=next(iter(dataloader)) # (batch_size, seq_len), (batch_size, seq_len)


model=transformer(vocab_size, positional_encoding, embed_dim, num_heads, feed_forward_ratio)
print(model)
optimizer=optim.Adam(model.parameters(),lr=0.1)
scheduler=optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda step:min((step+1)/100,1/(step+1)**(1/2)))


class trainer(nn.Module):
    def __init__(self, model, optimizer,scheduler, loss_function, dataloader,positional_encoding, embed_dim, vocab_size, num_epochs):
        super().__init__()
        self.model=model
        self.optimizer=optimizer
        self.loss_function=loss_function
        self.dataloader=dataloader
        self.num_epochs=num_epochs
        self.positional_encoding=positional_encoding(embed_dim)
        self.scheduler=scheduler
        self.vocab_size=vocab_size
        self.embed_dim=embed_dim
        

    def forward(self):
        model=self.model.to(device)
        model.train()
        for epoch in range(self.num_epochs):
            epoch_loss=0
            for data, label in dataloader:
                data=data.to(device)
                label=label.to(device)
                self.optimizer.zero_grad()
                pred_logits=model(data)
                loss=self.loss_function(pred_logits,label)
                loss.backward()
                nn.utils.clip_grad_norm_(model.parameters(), max_norm=10)
                self.optimizer.step()
                self.scheduler.step()
                epoch_loss+=loss.item()
            print(f"Epoch: {epoch+1}, Loss: {epoch_loss/len(dataloader)}")


trainer=trainer(model, optimizer,scheduler, loss_function, dataloader,positional_encoding,embed_dim, vocab_size, num_epochs=10)
trainer()