Update logging

bad_gpt.py

@@ -57,7 +57,9 @@ class BadGPTModel(nn.Module):
     # 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}')
+            # Log progress so I don't go insane
+            if _ % 16 == 0:
+                logger.debug(f'Iteration {_} of {max_new_tokens}')
             # Crop out the last block_size tokens
             cropped_idx = idx[:, -self.block_size:]
             logits = self(cropped_idx)

main.py

@@ -1,6 +1,8 @@
-import torch
-
 from bad_gpt import BadGPT
+import torch
+import logging
+logging.basicConfig()
+logger = logging.getLogger('bad_gpt').getChild(__name__)
 
 # HYPERPARAMETERS #
 ### Impacts performance ###
@@ -19,6 +21,7 @@ DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 
 if __name__ == '__main__':
+    logging.getLogger('bad_gpt').setLevel(logging.DEBUG)
     bad_gpt = BadGPT(
         device=DEVICE,
         batch_size=BATCH_SIZE,
@@ -32,7 +35,7 @@ if __name__ == '__main__':
         lr=LEARNING_RATE
     )
 
-    print("Generating response...\n")
+    logger.info("Generating response...")
     resp = bad_gpt.generate(
         'JULIET:\nRomeo, Romeo, wherefore art thou Romeo?', 256)
-    print("Response:\n" + resp)
+    logger.info("Response:\n" + resp)