refactor: Update pyproject.toml to include pytorch-gpu dependency and its source URL

bad_gpt.py

@@ -55,37 +55,43 @@ class BadGPTModel(nn.Module):
 
     # 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):
+    def generate(self, ctx: torch.Tensor, max_new_tokens: int):
+        for index in range(max_new_tokens):
             # Log progress so I don't go insane
-            if _ % 16 == 0:
-                logger.debug(f'Iteration {_} of {max_new_tokens}')
+            if index % 16 == 0:
+                logger.debug(f'Iteration {index} of {max_new_tokens}')
             # Crop out the last block_size tokens
-            cropped_idx = idx[:, -self.block_size:]
-            logits = self(cropped_idx)
+            cropped_ctx = ctx[:, -self.block_size:]
+            logits = self(cropped_ctx)
             # Logits has dimensions token, sentence, token_list
             # We want to make a new sentence, so only look at the last sentence
             logits = logits[:, -1, :]
             # Get possible next tokens and select one
             probabilities = F.softmax(logits, dim=-1)
-            idx_next = torch.multinomial(probabilities, num_samples=1)
+            ctx_next = torch.multinomial(probabilities, num_samples=1)
             # Add the new token to the end of the tensor
-            idx = torch.cat((idx, idx_next), dim=1)
-        return idx
+            ctx = torch.cat((ctx, ctx_next), dim=1)
+        return ctx
 
 
 @torch.no_grad()
-def estimate_loss(model: nn.Module, batcher: Batcher, eval_interval: int, device: Literal['cuda', 'cpu'] = 'cuda'):
+def estimate_loss(gpt: BadGPTModel, batcher: Batcher, eval_interval: int, device: Literal['cuda', 'cpu'] = 'cuda'):
     out = {}
-    model.eval()  # set to eval phase
+    gpt.eval()
     for split in ['train', 'val']:
         losses = torch.zeros(eval_interval)
-        for k in range(eval_interval):
-            x, y = batcher.get_batch(split=split)
-            logits, loss = model(x.to(device), y.to(device))
-            losses[k] = loss.item()
+        for epoch in range(eval_interval):
+            train, answer = batcher.get_batch(split='train')
+            logits = gpt.forward(train)
+            # Reformat pediction and answer so each entry can be compared
+            batch, block, vocab = logits.shape
+            logits = logits.view(batch * block, vocab)
+            answer = answer.view(batch * block)
+            # Compare entropy of predicted tokens to actual
+            loss = F.cross_entropy(logits, answer).item()
+            losses[epoch] = loss
         out[split] = losses.mean()
-    model.train()  # set back to training phase
+    gpt.train()
     return out
 
 
@@ -123,8 +129,14 @@ class BadGPTTrainer():
                 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(
+            logits = self.model(context_stack.to(
                 self.device), answer_stack.to(self.device))
+            batch, block, vocab = logits.shape
+            # Reformat logits and val so each entry can be compared
+            logits = logits.view(batch * block, vocab).to(self.device)
+            answer_stack = answer_stack.view(batch * block).to(self.device)
+            # Compare predicted tokens to actual
+            loss = F.cross_entropy(logits, answer_stack)
             self.optimizer.zero_grad(set_to_none=True)
             loss.backward()
             self.optimizer.step()
@@ -167,6 +179,8 @@ class BadGPT():
             learning_rate=lr
         )
         self._trainer.train()
+        # set to eval phase since we're only taking user input from here on
+        self._model.eval()
 
     def generate(self, prompt: str, response_size: int):
         start_ids = encode(prompt)
@@ -174,5 +188,5 @@ class BadGPT():
         # 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]
+            ctx=context, max_new_tokens=response_size)[0]
         return decode(encoded.tolist())

dataset.py

@@ -27,6 +27,7 @@ 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')
@@ -41,5 +42,5 @@ class Batcher():
         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])
+            [data[i+1:i+self.block_size+1] for i in random_indexes]).to(self.device)
         return context_stack, answer_stack

pyproject.toml

@@ -8,7 +8,7 @@ package-mode = false
 
 [tool.poetry.dependencies]
 python = "^3.10"
-torch = "^2.3.0"
+torch = { version = "^2.3.0", source = "pytorch-gpu" }
 numpy = "^1.26.4"
 datasets = "^2.19.0"
 tiktoken = "^0.6.0"
@@ -17,6 +17,12 @@ tiktoken = "^0.6.0"
 [tool.poetry.group.dev.dependencies]
 ipykernel = "^6.29.4"
 
+
+[[tool.poetry.source]]
+name = "pytorch-gpu"
+url = "https://download.pytorch.org/whl/cu118"
+priority = "supplemental"
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"