final model

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,4 @@
+__pycache__
+.venv
+gpt2-model
+

README.md

@@ -0,0 +1,13 @@
+---
+title: GPT2 Replica
+emoji: 👀
+colorFrom: indigo
+colorTo: indigo
+sdk: gradio
+sdk_version: 5.12.0
+app_file: app.py
+pinned: false
+short_description: 'A decoder trained starting with GPT2 weights. '
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

generate.py

@@ -0,0 +1,76 @@
+import torch
+import torch.nn.functional as F
+import tiktoken
+from train_get2_8_init import GPT, GPTConfig  # Import your model architecture
+
+# Device setup
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+print(f"Using device: {device}")
+
+# Initialize model and load trained weights
+model = GPT(GPTConfig())
+model.load_state_dict(torch.load('best_model.pt'))
+model.to(device)
+model.eval()  # Set to evaluation mode
+
+# Initialize tokenizer
+enc = tiktoken.get_encoding('gpt2')
+
+def generate_text(prompt, max_length=100, num_sequences=5, top_k=50, seed=42):
+    """
+    Generate text from a prompt using the trained model.
+    """
+    # Set random seed
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+    
+    # Encode the prompt
+    tokens = enc.encode(prompt)
+    x = torch.tensor(tokens).unsqueeze(0).repeat(num_sequences, 1).to(device)
+    
+    # Generate text
+    while x.size(1) < max_length:
+        with torch.no_grad():
+            logits = model(x)[0]
+            logits = logits[:, -1, :]
+            probs = F.softmax(logits, dim=-1)
+            
+            # Top-k sampling
+            topk_probs, topk_indices = torch.topk(probs, top_k, dim=-1)
+            ix = torch.multinomial(topk_probs, 1)
+            xcol = torch.gather(topk_indices, -1, ix)
+            
+            # Append to sequence
+            x = torch.cat((x, xcol), dim=1)
+    
+    # Decode and print results
+    print(f"\nPrompt: {prompt}")
+    print("\nGenerated sequences:")
+    print("-" * 50)
+    for i in range(num_sequences):
+        tokens = x[i, :].tolist()
+        decoded = enc.decode(tokens)
+        print(f"\n{i+1}. {decoded}")
+        print("-" * 50)
+
+# Interactive prompt loop
+if __name__ == "__main__":
+    while True:
+        prompt = input("\nEnter your prompt (or 'quit' to exit): ")
+        if prompt.lower() == 'quit':
+            break
+            
+        try:
+            max_len = int(input("Max length (default 100): ") or 100)
+            num_seq = int(input("Number of sequences (default 5): ") or 5)
+        except ValueError:
+            print("Using default values due to invalid input")
+            max_len = 100
+            num_seq = 5
+            
+        generate_text(
+            prompt=prompt,
+            max_length=max_len,
+            num_sequences=num_seq
+        ) 

logs.txt

@@ -0,0 +1,111 @@
+loading weights from local path: ./S12/gpt2-model
+loaded 338025 tokens
+1 epoch = 2640 batches
+Epoch 1/5: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2640/2640 [03:07<00:00, 14.07batch/s, loss=3.5681, batch=2639/2640]
+Epoch 1 completed. Loss: 3.5681
+New best loss! Saving model...
+Epoch 2/5: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2640/2640 [03:07<00:00, 14.05batch/s, loss=2.8617, batch=2639/2640]
+Epoch 2 completed. Loss: 2.8617
+New best loss! Saving model...
+Epoch 3/5: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2640/2640 [03:07<00:00, 14.05batch/s, loss=1.9555, batch=2639/2640]
+Epoch 3 completed. Loss: 1.9555
+New best loss! Saving model...
+Epoch 4/5: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2640/2640 [03:07<00:00, 14.05batch/s, loss=1.2269, batch=2639/2640]
+Epoch 4 completed. Loss: 1.2269
+New best loss! Saving model...
+Epoch 5/5: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2640/2640 [03:08<00:00, 14.00batch/s, loss=0.7809, batch=2639/2640]
+Epoch 5 completed. Loss: 0.7809
+New best loss! Saving model...
+Training completed. Best loss: 0.7809
+tensor(1.1159, device='cuda:0', grad_fn=<NllLossBackward0>)
+
+/home/ubuntu/S12/train_get2-8-init.py:236: FutureWarning: You are using `torch.load` with `weights_only=False`...)
+
+loaded 338025 tokens
+1 epoch = 2640 batches
+Epoch 1/15: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2640/2640 [03:06<00:00, 14.17batch/s, loss=0.8081, batch=2639/2640]
+Epoch 1 completed. Loss: 0.8081
+New best loss! Saving model...
+Epoch 2/15: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2640/2640 [03:07<00:00, 14.06batch/s, loss=0.5428, batch=2639/2640]
+Epoch 2 completed. Loss: 0.5428
+New best loss! Saving model...
+Epoch 3/15: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2640/2640 [03:08<00:00, 14.02batch/s, loss=0.4627, batch=2639/2640]
+Epoch 3 completed. Loss: 0.4627
+New best loss! Saving model...
+Epoch 4/15: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2640/2640 [03:08<00:00, 13.99batch/s, loss=0.3558, batch=2639/2640]
+Epoch 4 completed. Loss: 0.3558
+New best loss! Saving model...
+Epoch 5/15: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2640/2640 [03:08<00:00, 13.99batch/s, loss=0.3293, batch=2639/2640]
+Epoch 5 completed. Loss: 0.3293
+New best loss! Saving model...
+Epoch 6/15: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2640/2640 [03:08<00:00, 13.99batch/s, loss=0.2828, batch=2639/2640]
+Epoch 6 completed. Loss: 0.2828
+New best loss! Saving model...
+Epoch 7/15: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2640/2640 [03:08<00:00, 14.02batch/s, loss=0.2875, batch=2639/2640]
+Epoch 7 completed. Loss: 0.2875
+Epoch 8/15: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2640/2640 [03:08<00:00, 14.01batch/s, loss=0.2383, batch=2639/2640]
+Epoch 8 completed. Loss: 0.2383
+New best loss! Saving model...
+Epoch 9/15: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2640/2640 [03:08<00:00, 14.02batch/s, loss=0.2269, batch=2639/2640]
+Epoch 9 completed. Loss: 0.2269
+New best loss! Saving model...
+Epoch 10/15: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2640/2640 [03:08<00:00, 14.02batch/s, loss=0.2363, batch=2639/2640]
+Epoch 10 completed. Loss: 0.2363
+Epoch 11/15: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2640/2640 [03:08<00:00, 14.03batch/s, loss=0.2088, batch=2639/2640]
+Epoch 11 completed. Loss: 0.2088
+New best loss! Saving model...
+Epoch 12/15: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2640/2640 [03:08<00:00, 14.02batch/s, loss=0.1982, batch=2639/2640]
+Epoch 12 completed. Loss: 0.1982
+New best loss! Saving model...
+Epoch 13/15: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2640/2640 [03:08<00:00, 14.03batch/s, loss=0.2130, batch=2639/2640]
+Epoch 13 completed. Loss: 0.2130
+Epoch 14/15: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2640/2640 [03:08<00:00, 14.02batch/s, loss=0.2038, batch=2639/2640]
+Epoch 14 completed. Loss: 0.2038
+Epoch 15/15: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2640/2640 [03:08<00:00, 14.01batch/s, loss=0.2160, batch=2639/2640]
+Epoch 15 completed. Loss: 0.2160
+Training completed. Best loss: 0.1982
+
+Increased max. tokens in a sequence to 64.
+
+using device: cuda
+/home/ubuntu/S12/train_get2-8-init.py:236: FutureWarning: You are using `torch.load` with `weights_only=False`...)
+
+loaded 338025 tokens
+1 epoch = 1320 batches
+Epoch 1/5: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1320/1320 [01:59<00:00, 11.01batch/s, loss=0.2214, batch=1319/1320]
+Epoch 1 completed. Loss: 0.2214
+New best loss! Saving model...
+Epoch 2/5: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1320/1320 [02:00<00:00, 10.95batch/s, loss=0.1412, batch=1319/1320]
+Epoch 2 completed. Loss: 0.1412
+New best loss! Saving model...
+Epoch 3/5: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1320/1320 [02:00<00:00, 10.95batch/s, loss=0.1462, batch=1319/1320]
+Epoch 3 completed. Loss: 0.1462
+Epoch 4/5: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1320/1320 [02:00<00:00, 10.94batch/s, loss=0.1165, batch=1319/1320]
+Epoch 4 completed. Loss: 0.1165
+New best loss! Saving model...
+Epoch 5/5: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1320/1320 [02:00<00:00, 10.94batch/s, loss=0.1074, batch=1319/1320]
+Epoch 5 completed. Loss: 0.1074
+New best loss! Saving model...
+Training completed. Best loss: 0.1074
+
+Increased max. tokens in a sequence to 128.
+
+using device: cuda
+/home/ubuntu/S12/train_get2-8-init.py:236: FutureWarning: You are using `torch.load` with `weights_only=False`...)
+
+loaded 338025 tokens
+1 epoch = 660 batches
+Epoch 1/5: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████| 660/660 [01:33<00:00,  7.03batch/s, loss=0.1102, batch=659/660]
+Epoch 1 completed. Loss: 0.1102
+New best loss! Saving model...
+Epoch 2/5: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████| 660/660 [01:33<00:00,  7.03batch/s, loss=0.0738, batch=659/660]
+Epoch 2 completed. Loss: 0.0738
+New best loss! Saving model...
+Epoch 3/5: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████| 660/660 [01:33<00:00,  7.02batch/s, loss=0.0526, batch=659/660]
+Epoch 3 completed. Loss: 0.0526
+New best loss! Saving model...
+Epoch 4/5: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████| 660/660 [01:33<00:00,  7.03batch/s, loss=0.0566, batch=659/660]
+Epoch 4 completed. Loss: 0.0566
+Epoch 5/5: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████| 660/660 [01:34<00:00,  7.01batch/s, loss=0.0587, batch=659/660]
+Epoch 5 completed. Loss: 0.0587
+Training completed. Best loss: 0.0526

prompts.txt

@@ -0,0 +1,116 @@
+Enter your prompt (or 'quit' to exit): O Brother, Where art thou?
+Max length (default 100): 50
+Number of sequences (default 5): 
+
+Prompt: O Brother, Where art thou?
+
+Generated sequences:
+--------------------------------------------------
+
+1. O Brother, Where art thou?
+
+MENENIUS:
+Being moved, what thou shouldst not,
+ being so far forthrong'd as the manner is, his wounds
+To the people, beg their stinking breaths.
+--------------------------------------------------
+
+2. O Brother, Where art thou?
+
+MENENIUS:
+Being a island, a whip, or a kept-anch!
+
+SICINIUS:
+Where is he, hear you?
+
+MENENIUS
+--------------------------------------------------
+
+3. O Brother, Where art thou?
+
+MENENIUS:
+Being moved, am sure, not against Aufidius.
+
+MARCIUS:
+That's my noble master!
+What shall I do? say what;
+--------------------------------------------------
+
+4. O Brother, Where art thou?
+
+MENENIUS:
+aw, what am I, that brought thee to this?
+
+TITAN:
+A single thing, as I am, that wonders to hear thee speak.
+
+
+--------------------------------------------------
+
+5. O Brother, Where art thou?
+
+CORIOLANUS:
+Thy memory is as theirs.
+
+SEBASTIAN:
+I am, as ne'er I heard thee go
+That thou mightst be malap
+--------------------------------------------------
+
+
+
+
+
+
+Enter your prompt (or 'quit' to exit): In Brutus, should we trust? 
+Max length (default 100): 50
+Number of sequences (default 5): 
+
+Prompt: In Brutus, should we trust? 
+
+Generated sequences:
+--------------------------------------------------
+
+1. In Brutus, should we trust?  goes not
+To him or our poorness?
+
+LUCENTIO:
+I pray, sir, let's away.
+
+MARCIUS:
+Nay, I will give thee
+--------------------------------------------------
+
+2. In Brutus, should we trust?  feed and pride.
+
+SICINIUS:
+As big as to love and good news the good news
+ fellowsies, which did befall'n us,
+That none of us
+--------------------------------------------------
+
+3. In Brutus, should we trust?  there's no
+ whatsoever you should come toved him o' the air or thine
+with him put it not a accusers, and thus answer'd:
+'True is it, my incorporate
+--------------------------------------------------
+
+4. In Brutus, should we trust?  Well, or I'll lean
+Of Trustwings or when the wars make us.
+
+First Officer:
+This isle with Calibans.
+
+COMINIUS:
+Spe
+--------------------------------------------------
+
+5. In Brutus, should we trust?  folly
+We three are married, but this is put forth,
+Like men born by the flood, and wind-footed rage,
+Then fair Milan's glory moves.
+
+MARiners:
+
+--------------------------------------------------
+

train_get2_8_init.py

@@ -0,0 +1,294 @@
+# Solving for residual std scaling issue
+import os
+import math
+import time
+import inspect
+from dataclasses import dataclass
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from tqdm import tqdm
+
+class CausalSelfAttention(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0
+        # key, query, value projections for all heads, but in a batch
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
+        # output projection
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd)
+        self.c_proj.NANGPT_SCALE_INIT = 1
+        # regularization
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+        self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size)).view(1, 1, config.block_size, config.block_size))
+
+    def forward(self, x):
+        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)
+        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
+        # nh is "number of heads", hs is "head size", and C (number of channels) = nh * hs
+        # e.g. in GPT-2 (124M), n_head=12, hs=64, so nh*hs=C=768 channels in the Transformer
+        qkv = self.c_attn(x)
+        q, k, v = qkv.split(self.n_embd, dim=2)
+        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+
+        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+        att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf'))
+        att = F.softmax(att, dim=-1)
+        y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
+
+        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
+        # output projection
+        y = self.c_proj(y)
+        return y
+
+
+class MLP(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd)
+        self.gelu    = nn.GELU(approximate='tanh')
+        self.c_proj  = nn.Linear(4 * config.n_embd, config.n_embd)
+        self.c_proj.NANOGPT_SCALE_INIT = 1
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = self.gelu(x)
+        x = self.c_proj(x)
+        return x
+
+class Block(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.ln_1 = nn.LayerNorm(config.n_embd)
+        self.attn = CausalSelfAttention(config)
+        self.ln_2 = nn.LayerNorm(config.n_embd)
+        self.mlp = MLP(config)
+
+    def forward(self, x):
+        x = x + self.attn(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+@dataclass
+class GPTConfig:
+    block_size: int = 1024 # max sequence length
+    vocab_size: int = 50257 # number of tokens: 50,000 BPE merges + 256 bytes tokens + 1 <|endoftext|> token
+    n_layer: int = 12 # number of layers
+    n_head: int = 12 # number of heads
+    n_embd: int = 768 # embedding dimension
+
+
+class GPT(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+
+        self.transformer = nn.ModuleDict(dict(
+            wte = nn.Embedding(config.vocab_size, config.n_embd),
+            wpe = nn.Embedding(config.block_size, config.n_embd),
+            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
+            ln_f = nn.LayerNorm(config.n_embd),
+        ))
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+
+        # weight sharing
+        self.transformer.wte.weight = self.lm_head.weight
+
+        # weight initialization
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            std = 0.02
+            if hasattr(module, 'NANGPT_SCALE_INIT'):
+                std *= (2 * self.config.n_layer) ** -0.5
+            torch.nn.init.normal_(module.weight, mean = 0.0, std = std)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std = 0.02)
+
+
+
+    def forward(self, idx, targets=None):
+        # idx is of shape (B, T)
+        B, T = idx.size()
+        assert T <= self.config.block_size, f"Cannot forward sequence of length {T}, block size is only {self.config.block_size}"
+        # forward the token and posisition embeddings
+        pos = torch.arange(0, T, dtype=torch.long, device=idx.device) # shape (T)
+        pos_emb = self.transformer.wpe(pos) # position embeddings of shape (T, n_embd)
+        tok_emb = self.transformer.wte(idx) # token embeddings of shape (B, T, n_embd)
+        x = tok_emb + pos_emb
+        # forward the blocks of the transformer
+        for block in self.transformer.h:
+            x = block(x)
+        # forward the final layernorm and the classifier
+        x = self.transformer.ln_f(x)
+        logits = self.lm_head(x) # (B, T, vocab_size)
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
+        return logits, loss
+
+    @classmethod
+    def from_pretrained(cls, model_path):
+        """Loads pretrained GPT-2 model weights from local directory"""
+        from transformers import GPT2LMHeadModel, GPT2Config
+        print("loading weights from local path: %s" % model_path)
+
+        # n_layer, n_head and n_embd are determined from model_type
+        config_args = dict(n_layer=12, n_head=12, n_embd=768)  # 124M params
+        config_args['vocab_size'] = 50257 # always 50257 for GPT model checkpoints
+        config_args['block_size'] = 1024 # always 1024 for GPT model checkpoints
+        
+        # create a from-scratch initialized minGPT model
+        config = GPTConfig(**config_args)
+        model = GPT(config)
+        sd = model.state_dict()
+        sd_keys = sd.keys()
+        sd_keys = [k for k in sd_keys if not k.endswith('.attn.bias')]
+
+        # init a huggingface/transformers model
+        model_hf = GPT2LMHeadModel.from_pretrained(model_path, local_files_only=True)
+        sd_hf = model_hf.state_dict()
+
+        # copy while ensuring all of the parameters are aligned and match in names and shapes
+        sd_keys_hf = sd_hf.keys()
+        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.masked_bias')] # ignore these, just a buffer
+        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')] # same, just the mask (buffer)
+        transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']
+
+        assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}"
+        for k in sd_keys_hf:
+            if any(k.endswith(w) for w in transposed):
+                # special treatment for the Conv1D weights we need to transpose
+                assert sd_hf[k].shape[::-1] == sd[k].shape
+                with torch.no_grad():
+                    sd[k].copy_(sd_hf[k].t())
+            else:
+                # vanilla copy over the other parameters
+                assert sd_hf[k].shape == sd[k].shape
+                with torch.no_grad():
+                    sd[k].copy_(sd_hf[k])
+
+        return model
+
+device = 'cpu'
+if torch.cuda.is_available():
+    device = 'cuda'
+elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+    device = "mps"
+print(f"using device: {device}")
+
+# SEED
+torch.manual_seed(1337)
+if torch.cuda.is_available():
+    torch.cuda.manual_seed(1337)
+
+# STOP
+num_return_sequences = 5
+max_length = 30
+
+
+
+import tiktoken
+
+class DataLoaderLite:
+    def __init__(self, B, T):
+        self.B = B
+        self.T = T
+
+        # at init load tokens from disk and store them in memory
+        with open('input.txt', 'r') as f:
+            text = f.read()
+        enc = tiktoken.get_encoding('gpt2') 
+        tokens = enc.encode(text)
+        self.tokens = torch.tensor(tokens)
+        print(f'loaded {len(self.tokens)} tokens')
+        print(f'1 epoch = {len(self.tokens) // (B * T)} batches')
+
+        # state
+        self.current_position = 0
+    
+    def next_batch(self):
+        B, T = self.B, self.T
+        buf = self.tokens[self.current_position: self.current_position + B * T + 1]
+        x = (buf[:-1]).view(B, T) # inputs
+        y = (buf[1:]).view(B, T) # targets
+        # advance the position in the tensor
+        self.current_position += B*T
+        # if loading the next batch would be out of bounds, reset
+        if self.current_position + (B * T + 1) > len(self.tokens):
+            self.current_position = 0
+        return x, y
+
+
+# Move all training-specific code inside main
+if __name__ == "__main__":
+    # Device setup
+    device = 'cpu'
+    if torch.cuda.is_available():
+        device = 'cuda'
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        device = "mps"
+    print(f"using device: {device}")
+
+    # SEED
+    torch.manual_seed(1337)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(1337)
+
+    # Initialize model and training
+    model = GPT(GPTConfig())
+    model.load_state_dict(torch.load('best_model.pt'))
+    model.to(device)
+    
+    train_loader = DataLoaderLite(B = 4, T = 128)
+    
+    # Training loop
+    num_epochs = 5
+    optimizer = torch.optim.AdamW(model.parameters(), lr = 3e-4)
+    n_batches = len(train_loader.tokens) // (train_loader.B * train_loader.T)
+    best_loss = float('inf')
+    running_loss = 0.0
+    
+    for epoch in range(num_epochs):
+        # Train for one complete epoch with progress bar
+        progress_bar = tqdm(range(n_batches), 
+                           desc=f'Epoch {epoch+1}/{num_epochs}', 
+                           unit='batch')
+        
+        for i in progress_bar:
+            x, y = train_loader.next_batch()
+            x, y = x.to(device), y.to(device)
+            optimizer.zero_grad()
+            logits, loss = model(x, y)
+            loss.backward()
+            optimizer.step()
+            
+            # Update running loss and progress bar
+            running_loss = 0.9 * running_loss + 0.1 * loss.item()  # Smoothed loss
+            progress_bar.set_postfix({
+                'loss': f'{running_loss:.4f}',
+                'batch': f'{i}/{n_batches}'
+            })
+        
+        # Print epoch summary
+        print(f'Epoch {epoch+1} completed. Loss: {running_loss:.4f}')
+        
+        # Save best model with correct path
+        if running_loss < best_loss:
+            best_loss = running_loss
+            print(f'New best loss! Saving model...')
+            save_path = os.path.join('.', 'best_model.pt')
+            torch.save(model.state_dict(), save_path)
+
+    print(f'Training completed. Best loss: {best_loss:.4f}')