test(tests/test_prompt.py): eval code

README.md

@@ -92,6 +92,8 @@ python -m kan_gpt.prompt --prompt "Bangalore is often described as the " --model
   - [x] Train a dummy batch w/o any memory issues
 - [x] Mini training POC for MLP-GPT
 - [x] Train MLP-GPT on the webtext dataset as a baseline
+- [x] Train KAN-GPT on the webtext dataset as a baseline
+- [ ] Metrics comparing KAN-GPT and MLP-GPT
 - [x] Auto Save checkpoints
 - [x] Auto Save checkpoints to W&B
 - [ ] Auto Download model weights from git / huggingface

tests/test_prompt.py

@@ -0,0 +1,18 @@
+from kan_gpt.prompt import main
+
+VOCAB_SIZE = 8
+BLOCK_SIZE = 16
+MODEL_TYPE = "gpt-pico"
+
+
+def test_train():
+    class Args:
+        model_type = MODEL_TYPE
+        model_path = None
+        max_tokens = 3
+        prompt = "Bangalore is often described as the "
+        architecture = "KAN"
+        device = "cpu"
+
+    args = Args()
+    main(args)

tests/test_train.py

@@ -1,7 +1,7 @@
 import os
 import torch
 from kan_gpt.mingpt.model import GPT as MLP_GPT
-from kan_gpt.train import save_model
+from kan_gpt.train import save_model, main
 
 VOCAB_SIZE = 8
 BLOCK_SIZE = 16
@@ -33,3 +33,21 @@ def test_save_model():
                 f"Model not saved correctly at {save_path}, parameter "
                 f"{name} does not match original model"
             )
+
+
+# def test_train():
+
+#     # TODO: Download mini dataset for testing
+
+#     class Args:
+#         model_type = MODEL_TYPE
+#         dummy_dataset = True
+#         learning_rate = 5e-3
+#         max_iters = 1
+#         num_workers = 0
+#         batch_size = 1
+#         architecture = "KAN"
+#         device = "cpu"
+
+#     args = Args()
+#     main(args)