pruthvimv/DecoderonlyTransformer

1 2import torch3import torch.nn as nn4import torch.nn.functional as F5from dataclasses import dataclass6from transformers import GPT2Tokenizer7import tiktoken8import gradio as gr9 10@dataclass11class Config:12    vocab_size: int = 5025713    max_seq_len: int = 204814    dim: int = 76815    num_layers: int = 1216    num_heads: int = 1217    dropout: float = 0.118 19class MultiHeadAttention(nn.Module):20    def __init__(self, config):21        super().__init__()22        self.config = config23        self.n_head = config.num_heads24        self.n_embd = config.dim25        26        # Linear projections for Q, K, V27        self.c_attn = nn.Linear(config.dim, 3 * config.dim) # [n_embd, 3 * n_embd]28        self.c_proj = nn.Linear(config.dim, config.dim) # [n_embd, n_embd]29        30        self.attn_dropout = nn.Dropout(config.dropout)31        self.resid_dropout = nn.Dropout(config.dropout)32 33    def forward(self, x):34        B, T, C = x.size() # [B, T, n_embd]35        36        # Linear projection and split into Q, K, V37        q, k, v = self.c_attn(x).split(self.n_embd, dim=2) # [B, T, n_embd] each38        39        # Reshape for multi-head attention40        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # [B, n_head, T, n_embd/n_head]41        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # [B, n_head, T, n_embd/n_head]42        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # [B, n_head, T, n_embd/n_head]43        44        # Attention scores45        att = (q @ k.transpose(-2, -1)) * (1.0 / (k.size(-1) ** 0.5)) # [B, n_head, T, T]46        att = F.softmax(att, dim=-1) # [B, n_head, T, T]47        att = self.attn_dropout(att) # [B, n_head, T, T]48        49        # Weighted sum of values50        y = att @ v # [B, n_head, T, n_embd/n_head]51        52        # Reshape and project53        y = y.transpose(1, 2).contiguous().view(B, T, C) # [B, T, n_embd]54        y = self.c_proj(y) # [B, T, n_embd]55        y = self.resid_dropout(y) # [B, T, n_embd]56        57        return y58 59class FeedForward(nn.Module):60    def __init__(self, config):61        super().__init__()62        self.c_fc = nn.Linear(config.dim, 4 * config.dim) # [n_embd, 4 * n_embd]63        self.c_proj = nn.Linear(4 * config.dim, config.dim) # [4 * n_embd, n_embd]64        self.dropout = nn.Dropout(config.dropout)65 66    def forward(self, x):67        x = self.c_fc(x) # [B, T, 4 * n_embd]68        x = F.gelu(x) # [B, T, 4 * n_embd]69        x = self.c_proj(x) # [B, T, n_embd]70        x = self.dropout(x) # [B, T, n_embd]71        return x72 73class TransformerBlock(nn.Module):74    def __init__(self, config):75        super().__init__()76        self.ln_1 = nn.LayerNorm(config.dim) # [n_embd]77        self.attn = MultiHeadAttention(config)78        self.ln_2 = nn.LayerNorm(config.dim) # [n_embd]79        self.mlp = FeedForward(config)80 81    def forward(self, x):82        x = x + self.attn(self.ln_1(x)) # [B, T, n_embd]83 84class DecoderOnlyTransformer(nn.Module):85    def __init__(self, config):86        super().__init__()87        self.config = config88        self.wte = nn.Embedding(config.vocab_size, config.dim) # [vocab_size, n_embd]89        self.wpe = nn.Embedding(config.max_seq_len, config.dim) # [max_seq_len, n_embd]90        self.drop = nn.Dropout(config.dropout)91        self.blocks = nn.ModuleList([TransformerBlock(config) for _ in range(config.num_layers)])92        self.ln_f = nn.LayerNorm(config.dim) # [n_embd]93        self.lm_head = nn.Linear(config.dim, config.vocab_size, bias=False) # [n_embd, vocab_size]94        95        self.apply(self._init_weights)96 97    def _init_weights(self, module):98        if isinstance(module, (nn.Linear, nn.Embedding)):99            module.weight.data.normal_(mean=0.0, std=0.02)100            if isinstance(module, nn.Linear) and module.bias is not None:101                module.bias.data.zero_()102        elif isinstance(module, nn.LayerNorm):103            module.bias.data.zero_()104            module.weight.data.fill_(1.0)105 106    def forward(self, idx):107        B, T = idx.size() # [B, T]108        109        # Positional embeddings110        pos = torch.arange(0, T, dtype=torch.long, device=idx.device).unsqueeze(0) # [1, T]111        112        # Token and position embeddings113        tok_emb = self.wte(idx) # [B, T, n_embd]114        pos_emb = self.wpe(pos) # [1, T, n_embd]115        116        # Combine embeddings and apply dropout117        x = self.drop(tok_emb + pos_emb) # [B, T, n_embd]118        119        # Transformer blocks120        for block in self.blocks:121            x = block(x) # [B, T, n_embd]122        123        # Final layer norm and linear projection124        x = self.ln_f(x) # [B, T, n_embd]125        logits = self.lm_head(x) # [B, T, vocab_size]126        127        return logits128 129class DataLoaderLite:130   def __init__(self, B, T):131     self.B = B132     self.T = T133     # at init load tokens from disk and store them in memory134     with open('input.txt', 'r') as f:135       text = f.read()136     enc = tiktoken.get_encoding('gpt2')137     tokens = enc.encode(text)138     self.tokens = torch.tensor(tokens)139     print(f'loaded {len(self.tokens)} tokens')140     print(f'1 epoch = {len(self.tokens) // (B * T)} batches')141     # state142     self.current_position = 0143   def next_batch(self):144      B, T = self.B, self.T145      buf = self.tokens[self.current_position: self.current_position + B * T + 1]146      x = (buf[:-1]).view(B, T) # inputs147      y = (buf[1:]).view(B, T) # targets148      # advance the position in the tensor149      self.current_position += B*T150      # if loading the next batch would be out of bounds, reset151      if self.current_position + (B * T + 1) > len(self.tokens):152        self.current_position = 0153      return x, y154 155def train_model():156    use_cuda = torch.cuda.is_available()157    device = torch.device("cuda" if use_cuda else "cpu")158    config = Config()159    model = DecoderOnlyTransformer(config)160    model.to(device)161 162    # Train the model163    train_loader = DataLoaderLite(B = 4, T = 128)164    # NEW CODE165    optimizer = torch.optim.AdamW(model.parameters(), lr = 3e-4)166    loss_fn = nn.CrossEntropyLoss() # Define loss_fn here167    for i in range(5000):168       x, y = train_loader.next_batch()169       x, y = x.to(device), y.to(device)170       optimizer.zero_grad()171       logits = model(x)172       loss = loss_fn(logits.reshape(-1, logits.size(-1)), y.reshape(-1)) # Calculate loss using logits and target173       loss.backward()174       optimizer.step()175    return model176 177def preprocess_predictoutput(model, input):178    use_cuda = torch.cuda.is_available()179    device = torch.device("cuda" if use_cuda else "cpu")180    prompt = input181    tokenizer = GPT2Tokenizer.from_pretrained('gpt2')182    input_ids = tokenizer(prompt, return_tensors="pt").input_ids183    input_ids = input_ids.to(device)184    output = model(input_ids)185    predicted_ids = torch.argmax(output, dim=-1)186    predicted_text = tokenizer.decode(predicted_ids[0])187    return predicted_text188 189def TrainPredict(input):190  model = train_model()191  input = input192  predicted_text = preprocess_predictoutput(model, input)193  return predicted_text194 195 196iface = gr.Interface(197    fn=TrainPredict,198    inputs="textbox",199    outputs="textbox",200    title="Next word predictor",201    description="Next word being predicted ."202)203 204if __name__ == "__main__":205  iface.launch()