refactoring

bigram.py → bad_gpt.py

@@ -1,70 +1,19 @@
+import os
 from typing import Literal
 import torch
 import torch.nn as nn
 from torch.nn import functional as F
 import numpy as np
-from datasets import load_dataset
+import logging
 
-from encoder import encode, decode, tokens
-from self_attention import Head, MultiHead
+from encoder import encode, decode
+from self_attention import Block
+from dataset import Batcher
 
-
-class Batcher():
-    def __init__(self, device: Literal['cuda', 'cpu'], batch_size: int, block_size: int):
-        self.device = device
-        self.batch_size = batch_size
-        self.block_size = block_size
-        from dataset import make_dataset
-        train_data = make_dataset('train')
-        val_data = make_dataset('validation')
-        self.train_data = torch.tensor(encode(train_data), dtype=torch.long)
-        self.val_data = torch.tensor(encode(val_data), dtype=torch.long)
-        self.vocab = tokens
-
-    def get_batch(self, split: str = 'val'):
-        data = self.train_data if split == 'train' else self.val_data
-        random_indexes = torch.randint(
-            len(data) - self.block_size, (self.batch_size,)).to(self.device)
-        context_stack = torch.stack(
-            [data[i:i+self.block_size] for i in random_indexes]).to(self.device)
-        answer_stack = torch.stack(
-            [data[i+1:i+self.block_size+1] for i in random_indexes])
-        return context_stack, answer_stack
-
-
-class FeedForward(nn.Module):
-    def __init__(self, n_embd: int, dropout: float):
-        super().__init__()
-        self.net = nn.Sequential(
-            # Scale out data before applying ReLU so we get more variance
-            nn.Linear(n_embd, n_embd * 4),
-            nn.ReLU(),
-            # Scale back down before returning, effectively averaging the variance from earlier
-            nn.Linear(n_embd * 4, n_embd),
-            nn.Dropout(dropout)
-        )
-
-    def forward(self, x: torch.Tensor):
-        return self.net(x)
-
-
-class Block(nn.Module):
-    def __init__(self, n_embd: int, block_size: int, n_head: int, dropout: float):
-        super().__init__()
-        head_size = n_embd // n_head
-        self.sa_head = MultiHead(
-            n_head, block_size, n_embd, head_size, dropout)
-        self.ffwd = FeedForward(n_embd, dropout)
-        self.norm1 = nn.LayerNorm(n_embd)
-        self.norm2 = nn.LayerNorm(n_embd)
-
-    def forward(self, x: torch.Tensor):
-        x = x + self.sa_head(self.norm1(x))
-        x = x + self.ffwd(self.norm2(x))
-        return x
+logger = logging.getLogger('bad_gpt').getChild(__name__)
 
 
-class BigramLanguageModel(nn.Module):
+class BadGPTModel(nn.Module):
     def __init__(
         self,
         device: Literal['cuda', 'cpu'],
@@ -102,24 +51,16 @@ class BigramLanguageModel(nn.Module):
         x: torch.Tensor = tok_emb + pos_emb
         x = self.blocks(x)
         logits: torch.Tensor = self.lm_head(x)
-        if targets is None:
-            loss = 0
-        else:
-            batch, block, vocab = logits.shape
-            # Reformat logits and targets so each entry can be compared
-            logits = logits.view(batch * block, vocab)
-            targets = targets.view(batch * block)
-            # Compare predicted tokens to actual
-            loss = F.cross_entropy(logits, targets)
-        return logits, loss
+        return logits
 
     # Given a 2d matrix of dimensions token and sentence
     # generate new tokens in the next sentence
     def generate(self, idx: torch.Tensor, max_new_tokens: int):
         for _ in range(max_new_tokens):
+            print(f'Iteration {_} of {max_new_tokens}')
             # Crop out the last block_size tokens
             cropped_idx = idx[:, -self.block_size:]
-            logits, _ = self(cropped_idx)
+            logits = self(cropped_idx)
             # Logits has dimensions token, sentence, token_list
             # We want to make a new sentence, so only look at the last sentence
             logits = logits[:, -1, :]
@@ -144,3 +85,92 @@ def estimate_loss(model: nn.Module, batcher: Batcher, eval_interval: int, device
         out[split] = losses.mean()
     model.train()  # set back to training phase
     return out
+
+
+class BadGPTTrainer():
+    def __init__(self, model: BadGPTModel, batcher: Batcher, eval_interval: int, iterations: int, learning_rate: float):
+        self.model = model
+        self.batcher = batcher
+        self.eval_interval = eval_interval
+        self.iterations = iterations
+        self.learning_rate = learning_rate
+        self.device = self.model.device
+        self.optimizer = torch.optim.AdamW(
+            self.model.parameters(), lr=self.learning_rate)
+
+    def train(self):
+        if os.path.exists('model.pth'):
+            logger.debug("Loading model from file...")
+            checkpoint = torch.load('model.pth', map_location=self.device)
+            self.model.load_state_dict(checkpoint['model_state_dict'])
+            logger.debug("Model loaded!")
+        else:
+            logger.debug("Training model...")
+            self._train()
+            torch.save({
+                'model_state_dict': self.model.state_dict(),
+                'optimizer_state_dict': self.optimizer.state_dict()
+            }, 'model.pth')
+            logger.debug("Training complete!")
+
+    def _train(self):
+        for i in range(self.iterations):
+            if i % self.eval_interval == 0:
+                losses = estimate_loss(
+                    self.model, self.batcher, self.eval_interval, self.device)
+                logger.debug(
+                    f"step {i}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")
+            context_stack, answer_stack = self.batcher.get_batch(split='train')
+            _, loss = self.model(context_stack.to(
+                self.device), answer_stack.to(self.device))
+            self.optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            self.optimizer.step()
+
+
+class BadGPT():
+    def __init__(
+        self,
+        device: Literal['cuda', 'cpu'],
+        block_size: int,
+        batch_size: int,
+        n_embd: int,
+        n_head: int,
+        n_layers: int,
+        dropout: float,
+        eval_interval: int,
+        iterations: int,
+        lr: float
+    ):
+        self.device = device
+        self._batcher = Batcher(
+            device=device,
+            batch_size=batch_size,
+            block_size=block_size
+        )
+        self._model = BadGPTModel(
+            device=device,
+            block_size=block_size,
+            vocab_size=len(self._batcher.vocab),
+            n_embd=n_embd,
+            n_head=n_head,
+            n_layers=n_layers,
+            dropout=dropout
+        ).to(device)
+        self._trainer = BadGPTTrainer(
+            model=self._model,
+            batcher=self._batcher,
+            eval_interval=eval_interval,
+            iterations=iterations,
+            learning_rate=lr
+        )
+        self._trainer.train()
+
+    def generate(self, prompt: str, response_size: int):
+        start_ids = encode(prompt)
+        context = torch.tensor(start_ids, dtype=torch.long, device=self.device)
+        # add batch dimension. it's just 1 batch, but we still need it cuz tensors
+        context = context[None, ...]
+        encoded = self._model.generate(
+            idx=context, max_new_tokens=response_size)[0]
+        return decode(encoded.tolist())

dataset.py

@@ -1,5 +1,7 @@
 from typing import Literal, Union
-from datasets import load_dataset, DatasetDict
+from datasets import load_dataset
+import torch
+from encoder import encode, tokens
 
 DatasetType = Union[None, str]
 
@@ -19,3 +21,25 @@ def make_dataset(split: Literal['train', 'validation', 'test'] = 'train'):
         out = str(list(ds)[0]['text'])
         _datasets[split] = out
     return str(_datasets[split])
+
+
+class Batcher():
+    def __init__(self, device: Literal['cuda', 'cpu'], batch_size: int, block_size: int):
+        self.device = device
+        self.batch_size = batch_size
+        from dataset import make_dataset
+        train_data = make_dataset('train')
+        val_data = make_dataset('validation')
+        self.train_data = torch.tensor(encode(train_data), dtype=torch.long)
+        self.val_data = torch.tensor(encode(val_data), dtype=torch.long)
+        self.vocab = tokens
+
+    def get_batch(self, split: str = 'val'):
+        data = self.train_data if split == 'train' else self.val_data
+        random_indexes = torch.randint(
+            len(data) - self.block_size, (self.batch_size,)).to(self.device)
+        context_stack = torch.stack(
+            [data[i:i+self.block_size] for i in random_indexes]).to(self.device)
+        answer_stack = torch.stack(
+            [data[i+1:i+self.block_size+1] for i in random_indexes])
+        return context_stack, answer_stack

main.py

@@ -1,12 +1,6 @@
-from typing import Literal
 import torch
-import torch.nn as nn
-from torch.nn import functional as F
-import numpy as np
-import os
 
-from encoder import encode, decode
-from bigram import BigramLanguageModel, Batcher, estimate_loss
+from bad_gpt import BadGPT
 
 # HYPERPARAMETERS #
 ### Impacts performance ###
@@ -24,60 +18,21 @@ DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 # --------------- #
 
 
-def train_model(model: nn.Module, batcher: Batcher, iterations=MAX_ITERS, lr=LEARNING_RATE):
-    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
-    for i in range(iterations):
-        if i % EVAL_INTERVAL == 0:
-            losses = estimate_loss(model, batcher, EVAL_INTERVAL, DEVICE)
-            print(
-                f"step {i}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")
-        context_stack, answer_stack = batcher.get_batch(split='train')
-        _, loss = model(context_stack.to(DEVICE), answer_stack.to(DEVICE))
-        optimizer.zero_grad(set_to_none=True)
-        loss.backward()
-        optimizer.step()
-    return optimizer
-
-
-b = Batcher(
-    device=DEVICE,
-    batch_size=BATCH_SIZE,
-    block_size=BLOCK_SIZE
-)
-m = BigramLanguageModel(
-    device=DEVICE,
-    block_size=BLOCK_SIZE,
-    vocab_size=len(b.vocab),
-    n_embd=NUM_EMBEDDING_DIMENSIONS,
-    n_head=NUM_HEADS,
-    n_layers=NUM_LAYERS,
-    dropout=DROPOUT_RATE
-).to(DEVICE)
-
-
-def run_model(model: nn.Module, response_size: int = BLOCK_SIZE, query: str = ''):
-    start_ids = encode(query)
-    context = torch.tensor(start_ids, dtype=torch.long, device=DEVICE)
-    # add batch dimension. it's just 1 batch, but we still need it cuz tensors
-    context = context[None, ...]
-    encoded = model.generate(
-        idx=context, max_new_tokens=response_size)[0]
-    return decode(encoded.tolist())
-
-
-if os.path.exists('model.pth'):
-    print("Loading model from file...")
-    checkpoint = torch.load('model.pth', map_location=DEVICE)
-    m.load_state_dict(checkpoint['model_state_dict'])
-    print("Model loaded!")
-else:
-    print("Training model...")
-    optimizer = train_model(m, b)
-    torch.save({
-        'model_state_dict': m.state_dict(),
-        'optimizer_state_dict': optimizer.state_dict()
-    }, 'model.pth')
-    print("Training complete!")
-print("Generating response...\n")
-resp = run_model(m, 256, 'JULIET:\nRomeo, Romeo, wherefore art thou Romeo?')
-print("Response:\n" + resp)
+if __name__ == '__main__':
+    bad_gpt = BadGPT(
+        device=DEVICE,
+        batch_size=BATCH_SIZE,
+        block_size=BLOCK_SIZE,
+        n_embd=NUM_EMBEDDING_DIMENSIONS,
+        n_head=NUM_HEADS,
+        n_layers=NUM_LAYERS,
+        dropout=DROPOUT_RATE,
+        eval_interval=EVAL_INTERVAL,
+        iterations=MAX_ITERS,
+        lr=LEARNING_RATE
+    )
+
+    print("Generating response...\n")
+    resp = bad_gpt.generate(
+        'JULIET:\nRomeo, Romeo, wherefore art thou Romeo?', 256)
+    print("Response:\n" + resp)

self_attention.py

@@ -2,6 +2,38 @@ import torch
 from torch import nn
 
 
+class FeedForward(nn.Module):
+    def __init__(self, n_embd: int, dropout: float):
+        super().__init__()
+        self.net = nn.Sequential(
+            # Scale out data before applying ReLU so we get more variance
+            nn.Linear(n_embd, n_embd * 4),
+            nn.ReLU(),
+            # Scale back down before returning, effectively averaging the variance from earlier
+            nn.Linear(n_embd * 4, n_embd),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x: torch.Tensor):
+        return self.net(x)
+
+
+class Block(nn.Module):
+    def __init__(self, n_embd: int, block_size: int, n_head: int, dropout: float):
+        super().__init__()
+        head_size = n_embd // n_head
+        self.sa_head = MultiHead(
+            n_head, block_size, n_embd, head_size, dropout)
+        self.ffwd = FeedForward(n_embd, dropout)
+        self.norm1 = nn.LayerNorm(n_embd)
+        self.norm2 = nn.LayerNorm(n_embd)
+
+    def forward(self, x: torch.Tensor):
+        x = x + self.sa_head(self.norm1(x))
+        x = x + self.ffwd(self.norm2(x))
+        return x
+
+
 class MultiHead(nn.Module):
     def __init__(self, num_heads: int, block_size: int, n_embd: int, head_size: int, dropout: float):
         super().__init__()