WIP Parameterize gpt script, add CLI

.gitignore

@@ -4,6 +4,7 @@ datasets/
 
 # Training Tensorboard runs
 runs/
+flagged/
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

gpt.ipynb

@@ -10,7 +10,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 39,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -22,7 +22,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 40,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -32,7 +32,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 41,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -55,90 +55,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 42,
+   "execution_count": 5,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Requirement already satisfied: torch in ./venv/lib/python3.10/site-packages (2.0.0)\n",
-      "Requirement already satisfied: pandas in ./venv/lib/python3.10/site-packages (1.5.3)\n",
-      "Requirement already satisfied: numpy in ./venv/lib/python3.10/site-packages (1.24.1)\n",
-      "Requirement already satisfied: tensorboard in ./venv/lib/python3.10/site-packages (2.12.0)\n",
-      "Requirement already satisfied: nvidia-nccl-cu11==2.14.3 in ./venv/lib/python3.10/site-packages (from torch) (2.14.3)\n",
-      "Requirement already satisfied: nvidia-cudnn-cu11==8.5.0.96 in ./venv/lib/python3.10/site-packages (from torch) (8.5.0.96)\n",
-      "Requirement already satisfied: nvidia-cusolver-cu11==11.4.0.1 in ./venv/lib/python3.10/site-packages (from torch) (11.4.0.1)\n",
-      "Requirement already satisfied: nvidia-cuda-nvrtc-cu11==11.7.99 in ./venv/lib/python3.10/site-packages (from torch) (11.7.99)\n",
-      "Requirement already satisfied: networkx in ./venv/lib/python3.10/site-packages (from torch) (3.0)\n",
-      "Requirement already satisfied: nvidia-curand-cu11==10.2.10.91 in ./venv/lib/python3.10/site-packages (from torch) (10.2.10.91)\n",
-      "Requirement already satisfied: filelock in ./venv/lib/python3.10/site-packages (from torch) (3.10.4)\n",
-      "Requirement already satisfied: nvidia-nvtx-cu11==11.7.91 in ./venv/lib/python3.10/site-packages (from torch) (11.7.91)\n",
-      "Requirement already satisfied: typing-extensions in ./venv/lib/python3.10/site-packages (from torch) (4.5.0)\n",
-      "Requirement already satisfied: nvidia-cublas-cu11==11.10.3.66 in ./venv/lib/python3.10/site-packages (from torch) (11.10.3.66)\n",
-      "Requirement already satisfied: sympy in ./venv/lib/python3.10/site-packages (from torch) (1.11.1)\n",
-      "Requirement already satisfied: nvidia-cuda-runtime-cu11==11.7.99 in ./venv/lib/python3.10/site-packages (from torch) (11.7.99)\n",
-      "Requirement already satisfied: nvidia-cufft-cu11==10.9.0.58 in ./venv/lib/python3.10/site-packages (from torch) (10.9.0.58)\n",
-      "Requirement already satisfied: jinja2 in ./venv/lib/python3.10/site-packages (from torch) (3.1.2)\n",
-      "Requirement already satisfied: triton==2.0.0 in ./venv/lib/python3.10/site-packages (from torch) (2.0.0)\n",
-      "Requirement already satisfied: nvidia-cuda-cupti-cu11==11.7.101 in ./venv/lib/python3.10/site-packages (from torch) (11.7.101)\n",
-      "Requirement already satisfied: nvidia-cusparse-cu11==11.7.4.91 in ./venv/lib/python3.10/site-packages (from torch) (11.7.4.91)\n",
-      "Requirement already satisfied: wheel in ./venv/lib/python3.10/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch) (0.40.0)\n",
-      "Requirement already satisfied: setuptools in ./venv/lib/python3.10/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch) (65.5.0)\n",
-      "Requirement already satisfied: lit in ./venv/lib/python3.10/site-packages (from triton==2.0.0->torch) (16.0.0)\n",
-      "Requirement already satisfied: cmake in ./venv/lib/python3.10/site-packages (from triton==2.0.0->torch) (3.26.1)\n",
-      "Requirement already satisfied: python-dateutil>=2.8.1 in ./venv/lib/python3.10/site-packages (from pandas) (2.8.2)\n",
-      "Requirement already satisfied: pytz>=2020.1 in ./venv/lib/python3.10/site-packages (from pandas) (2023.2)\n",
-      "Requirement already satisfied: requests<3,>=2.21.0 in ./venv/lib/python3.10/site-packages (from tensorboard) (2.28.2)\n",
-      "Requirement already satisfied: werkzeug>=1.0.1 in ./venv/lib/python3.10/site-packages (from tensorboard) (2.2.3)\n",
-      "Requirement already satisfied: google-auth<3,>=1.6.3 in ./venv/lib/python3.10/site-packages (from tensorboard) (2.16.3)\n",
-      "Requirement already satisfied: protobuf>=3.19.6 in ./venv/lib/python3.10/site-packages (from tensorboard) (4.22.1)\n",
-      "Requirement already satisfied: markdown>=2.6.8 in ./venv/lib/python3.10/site-packages (from tensorboard) (3.4.3)\n",
-      "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in ./venv/lib/python3.10/site-packages (from tensorboard) (0.4.6)\n",
-      "Requirement already satisfied: grpcio>=1.48.2 in ./venv/lib/python3.10/site-packages (from tensorboard) (1.51.3)\n",
-      "Requirement already satisfied: tensorboard-data-server<0.8.0,>=0.7.0 in ./venv/lib/python3.10/site-packages (from tensorboard) (0.7.0)\n",
-      "Requirement already satisfied: absl-py>=0.4 in ./venv/lib/python3.10/site-packages (from tensorboard) (1.4.0)\n",
-      "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in ./venv/lib/python3.10/site-packages (from tensorboard) (1.8.1)\n",
-      "Requirement already satisfied: cachetools<6.0,>=2.0.0 in ./venv/lib/python3.10/site-packages (from google-auth<3,>=1.6.3->tensorboard) (5.3.0)\n",
-      "Requirement already satisfied: pyasn1-modules>=0.2.1 in ./venv/lib/python3.10/site-packages (from google-auth<3,>=1.6.3->tensorboard) (0.2.8)\n",
-      "Requirement already satisfied: six>=1.9.0 in ./venv/lib/python3.10/site-packages (from google-auth<3,>=1.6.3->tensorboard) (1.16.0)\n",
-      "Requirement already satisfied: rsa<5,>=3.1.4 in ./venv/lib/python3.10/site-packages (from google-auth<3,>=1.6.3->tensorboard) (4.9)\n",
-      "Requirement already satisfied: requests-oauthlib>=0.7.0 in ./venv/lib/python3.10/site-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard) (1.3.1)\n",
-      "Requirement already satisfied: idna<4,>=2.5 in ./venv/lib/python3.10/site-packages (from requests<3,>=2.21.0->tensorboard) (3.4)\n",
-      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in ./venv/lib/python3.10/site-packages (from requests<3,>=2.21.0->tensorboard) (1.26.15)\n",
-      "Requirement already satisfied: charset-normalizer<4,>=2 in ./venv/lib/python3.10/site-packages (from requests<3,>=2.21.0->tensorboard) (3.1.0)\n",
-      "Requirement already satisfied: certifi>=2017.4.17 in ./venv/lib/python3.10/site-packages (from requests<3,>=2.21.0->tensorboard) (2022.12.7)\n",
-      "Requirement already satisfied: MarkupSafe>=2.1.1 in ./venv/lib/python3.10/site-packages (from werkzeug>=1.0.1->tensorboard) (2.1.2)\n",
-      "Requirement already satisfied: mpmath>=0.19 in ./venv/lib/python3.10/site-packages (from sympy->torch) (1.3.0)\n",
-      "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in ./venv/lib/python3.10/site-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard) (0.4.8)\n",
-      "Requirement already satisfied: oauthlib>=3.0.0 in ./venv/lib/python3.10/site-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard) (3.2.2)\n",
-      "\n",
-      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip available: \u001b[0m\u001b[31;49m22.3.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.0.1\u001b[0m\n",
-      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
-      "Note: you may need to restart the kernel to use updated packages.\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
-    "%pip install torch pandas numpy tensorboard"
+    "#%pip install torch pandas numpy tensorboard gradio"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 43,
+   "execution_count": 6,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "<torch._C.Generator at 0x7fef50768610>"
-      ]
-     },
-     "execution_count": 43,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
     "import torch\n",
     "import torch.nn as nn\n",
@@ -150,12 +78,14 @@
     "import numpy as np\n",
     "import math\n",
     "\n",
-    "torch.manual_seed(1337)"
+    "torch.manual_seed(1337)\n",
+    "# Set device to CUDA if available, otherwise use CPU\n",
+    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 44,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -169,7 +99,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 45,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [
     {
@@ -202,7 +132,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 46,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -231,7 +161,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 66,
+   "execution_count": 10,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -298,7 +228,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 48,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -318,7 +248,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 60,
+   "execution_count": 12,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -344,7 +274,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 50,
+   "execution_count": 13,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -404,7 +334,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 51,
+   "execution_count": 14,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -419,7 +349,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 67,
+   "execution_count": 15,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -435,7 +365,7 @@
     "    context_size=BLOCK_SIZE,\n",
     "    )\n",
     "\n",
-    "model = model.to('cuda')\n",
+    "model = model.to(device)\n",
     "optimizer = optim.AdamW(model.parameters(), lr=LR)\n",
     "#scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1000, gamma=0.1)\n",
     "criterion = F.cross_entropy\n",
@@ -529,8 +459,8 @@
     "\n",
     "for epoch in range(EPOCHS):\n",
     "    for data, target in train_dataloader:\n",
-    "        data = data.to('cuda')\n",
-    "        target = target.to('cuda')\n",
+    "        data = data.to(device)\n",
+    "        target = target.to(device)\n",
     "\n",
     "        loss = compute_loss(model, criterion, data, target)\n",
     "\n",
@@ -551,8 +481,8 @@
     "            with torch.no_grad():\n",
     "                model.eval()\n",
     "                for x, y in test_dataloader:\n",
-    "                    x = x.to(\"cuda\")\n",
-    "                    y = y.to(\"cuda\")\n",
+    "                    x = x.to(device)\n",
+    "                    y = y.to(device)\n",
     "\n",
     "                    batch_loss = compute_loss(model, criterion, x, y)\n",
     "                    total_loss += batch_loss.item() * 512\n",
@@ -576,7 +506,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 69,
+   "execution_count": 16,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -600,7 +530,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 17,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -619,16 +549,16 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 26,
+   "execution_count": 21,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "841"
+       "3"
       ]
      },
-     "execution_count": 26,
+     "execution_count": 21,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -660,8 +590,8 @@
     "val_dataloader = DataLoader(val_dataset, batch_size=512, num_workers=4)\n",
     "with torch.no_grad():\n",
     "    for x, y in val_dataloader:\n",
-    "        x = x.to(\"cuda\")\n",
-    "        y = y.to(\"cuda\")\n",
+    "        x = x.to(device)\n",
+    "        y = y.to(device)\n",
     "\n",
     "        batch_loss = compute_loss(model, criterion, x, y)\n",
     "        total_loss += batch_loss.item() * x.size(0)\n",
@@ -704,94 +634,458 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 58,
+   "execution_count": 54,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Tutus, to Marcius, noble Marcius\n",
-      "Made to my voices! doing and hangs upon them!\n",
-      "Take it to down our foes and hates with stain,\n",
-      "Which thus follows slay with on I meland,\n",
-      "What I am after her to her fearful haunt it?\n",
       "\n",
-      "PAULINA:\n",
-      "But you are well to hold the king.\n",
+      "Shepherd:\n",
+      "What should be gone to your deeds to prince\n",
+      "That let us away the house of Lancasters.\n",
       "\n",
-      "ISABELLA:\n",
-      "And I will not go royalty to thy hand.\n",
+      "Second Murderer:\n",
+      "Alas, if I be not the fire rogued man\n",
+      "That hath dead by the fault of the good stars;\n",
+      "And the business of some many that be sounded.\n",
       "\n",
-      "LUCIO:\n",
-      "Since I do not well in such goodly talk of.\n",
-      "I think I have a stay of it!\n",
+      "AUTOLYCUS:\n",
+      "My soul are but not so dishonours.\n",
       "\n",
-      "HENRY BOLINGBROKE:\n",
-      "Who say I hate been a day's mind;\n",
-      "Till we here and so very little and way,\n",
-      "And wash the city has nest seen the feast.\n",
+      "BRUTUS:\n",
+      "I will do not be point, but it is a soul,\n",
+      "I would I see my heart fair state,\n",
+      "That he is in the grief and be contented\n",
+      "That you may made the father, the day doth nothing\n",
+      "And lie at like in his world moved me at him.\n",
       "\n",
-      "DUCHESS OF YORK:\n",
-      "No, by the matter.\n",
+      "First Senator:\n",
+      "The rotten far the points of our honour.\n",
       "\n",
-      "ISABELLA:\n",
-      "Flitter than desire never yet looks so.\n",
+      "AUTOLYCUS:\n",
+      "The king hath been you think the duke of the foes,\n",
+      "And had so long and confessal abroad,\n",
+      "And said in the light save of men so meeting\n",
+      "May not shall do be as we as here.\n",
       "\n",
-      "HENRY BOLINGBROKE:\n",
-      "I am not possible perceived\n",
-      "And both place, where I may not rafes,\n",
-      "And like me one air. What you'll your love day?\n",
+      "Second Servingman:\n",
+      "You are grace as the child and fight again,\n",
+      "If you do not shall not be much of the thing;\n",
+      "One but so the creature of the charge is were\n",
+      "Than the point of your several motion\n",
+      "May seem as her like a pride and as would all grave\n",
+      "The which of the blood"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "'\\nPlot histograms of the gradient values during training. If you notice a significant number of gradients are near zero (vanishing gradients) or very large values (exploding gradients), it could be a problem. TensorBoard is a useful tool for visualizing these histograms.\\n\\nShepherd:\\nWhat should be gone to your deeds to prince\\nThat let us away the house of Lancasters.\\n\\nSecond Murderer:\\nAlas, if I be not the fire rogued man\\nThat hath dead by the fault of the good stars;\\nAnd the business of some many that be sounded.\\n\\nAUTOLYCUS:\\nMy soul are but not so dishonours.\\n\\nBRUTUS:\\nI will do not be point, but it is a soul,\\nI would I see my heart fair state,\\nThat he is in the grief and be contented\\nThat you may made the father, the day doth nothing\\nAnd lie at like in his world moved me at him.\\n\\nFirst Senator:\\nThe rotten far the points of our honour.\\n\\nAUTOLYCUS:\\nThe king hath been you think the duke of the foes,\\nAnd had so long and confessal abroad,\\nAnd said in the light save of men so meeting\\nMay not shall do be as we as here.\\n\\nSecond Servingman:\\nYou are grace as the child and fight again,\\nIf you do not shall not be much of the thing;\\nOne but so the creature of the charge is were\\nThan the point of your several motion\\nMay seem as her like a pride and as would all grave\\nThe which of the blood'"
+      ]
+     },
+     "execution_count": 54,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "seed = \"\"\"\n",
+    "Plot histograms of the gradient values during training. If you notice a significant number of gradients are near zero (vanishing gradients) or very large values (exploding gradients), it could be a problem. TensorBoard is a useful tool for visualizing these histograms.\n",
+    "\"\"\"\n",
+    "GEN_LENGTH=1024\n",
+    "\n",
+    "def generate(prompt, gen_length, temp=1, top_k=10, top_p=None):\n",
+    "    g_cuda = torch.Generator(device=device)\n",
+    "    contexts = torch.tensor(encode_text(prompt), dtype=torch.int32).to(device)\n",
+    "\n",
+    "    model.eval()\n",
+    "    for i in range(GEN_LENGTH):\n",
+    "        transform = nn.LogSoftmax(1)\n",
+    "        x = contexts[-BLOCK_SIZE:]\n",
+    "        if x.size(0) < BLOCK_SIZE:\n",
+    "            x = F.pad(x, (BLOCK_SIZE - x.size(0), 0), \"constant\", 0).unsqueeze(0) # B*T\n",
+    "        else:\n",
+    "            x = x.unsqueeze(0)\n",
+    "\n",
+    "        preds = model.infer(x)\n",
+    "        preds = preds.squeeze(0)\n",
+    "        preds = preds / temp\n",
+    "        probs = F.softmax(preds, dim=-1)\n",
+    "\n",
+    "        if top_p is not None:\n",
+    "            # Apply top-p\n",
+    "            sorted_probs, sorted_indices = torch.sort(probs[-1, :], descending=True)\n",
+    "            cumulative_probs = torch.cumsum(sorted_probs, dim=-1)\n",
+    "            # find cutoff\n",
+    "            idx_top_p = (cumulative_probs < top_p).sum().item()\n",
+    "            top_probs = sorted_probs[:idx_top_p]\n",
+    "            top_indices = sorted_indices[:idx_top_p]\n",
+    "            # Null case\n",
+    "            if top_probs.size(0) == 0:\n",
+    "                top_probs = sorted_probs[:1]\n",
+    "                top_indices = sorted_indices[:1]\n",
+    "            \n",
+    "            next_char = torch.multinomial(top_probs, num_samples=1, generator=g_cuda)\n",
+    "            next_char = top_indices[next_char]\n",
+    "        elif top_k is not None:\n",
+    "            top_k_probs, top_k_indices = torch.topk(probs[-1, :], k=top_k)\n",
+    "            next_char = torch.multinomial(top_k_probs, num_samples=1, generator=g_cuda)\n",
+    "            next_char = top_k_indices[next_char]\n",
+    "        else:\n",
+    "            next_char = torch.multinomial(probs, num_samples=1, generator=g_cuda)\n",
+    "\n",
+    "\n",
+    "        contexts = torch.cat((contexts, next_char), dim=0)\n",
+    "        print(decode_text(next_char.cpu().numpy())[-1], end=\"\")\n",
+    "    \n",
+    "    return(\"\".join(decode_text(contexts.cpu().numpy())))\n",
+    "\n",
+    "\"\".join(generate(seed, GEN_LENGTH,temp=0.8,top_p=0.9))"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Gradio WebUI"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 55,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Running on local URL:  http://127.0.0.1:7866\n",
       "\n",
-      "KING RICHARD II:\n",
-      "Then be thou--\n",
+      "To create a public link, set `share=True` in `launch()`.\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "<div><iframe src=\"http://127.0.0.1:7866/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": []
+     },
+     "execution_count": 55,
+     "metadata": {},
+     "output_type": "execute_result"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "If you do I am sure be your father's limbs,\n",
+      "Which is not drouble Paulina,\n",
+      "Where I come to me too much be to thus\n",
+      "As in a mind sorten but here hoped--followed as hung\n",
+      "bard at the time own that that he is death;\n",
+      "The stir weors shall be a new good groans.\n",
+      "\n",
+      "OXFORD:\n",
+      "And I say him to deed your city,\n",
+      "Thy heart advantage by mine at their world,\n",
+      "And be glood much show with upon the suddenly,\n",
+      "And in this free can company from here faults.\n",
+      "\n",
+      "YORK:\n",
+      "I bring thee, now I am again: it is a Montague.\n",
+      "\n",
+      "BRUTUS:\n",
+      "We while you not be the fault of the life.\n",
+      "\n",
+      "KING RICHARD III:\n",
+      "Come, sir, do I: in put I with so,\n",
+      "The rest redies with this ballad or suppher.\n",
+      "\n",
+      "Third Servingman:\n",
+      "Now, well do with kindness of power the that\n",
+      "ence of the ambassage; but to shall here,\n",
+      "But saving as the wounds both drawn dead.\n",
+      "\n",
+      "CLARENCE:\n",
+      "Peace, sir, he adieu.\n",
+      "\n",
+      "MISTRESS OVERDONE:\n",
+      "Cheep, comes thou lovest which I do well;\n",
+      "Now now as that protectors, in this parliament--\n",
+      "That nature i' the mark of her with old her,\n",
+      "And so here doth in circle and call'd the wind,\n",
+      "If you do I am sure be your father's limbs,\n",
+      "Which is not drouble Paulina,\n",
+      "Where I come to me too much be to thus\n",
+      "As in a mind sorten but here hoped--followed as hung\n",
+      "bard at the time own that that he is death;\n",
+      "The stir weors shall be a new good groans.\n",
+      "\n",
+      "OXFORD:\n",
+      "And I say him to deed your city,\n",
+      "Thy heart advantage by mine at their world,\n",
+      "And be glood much show with upon the suddenly,\n",
+      "And in this free can company from here faults.\n",
+      "\n",
+      "YORK:\n",
+      "I bring thee, now I am again: it is a Montague.\n",
+      "\n",
+      "BRUTUS:\n",
+      "We while you not be the fault of the life.\n",
+      "\n",
+      "KING RICHARD III:\n",
+      "Come, sir, do I: in put I with so,\n",
+      "The rest redies with this ballad or suppher.\n",
+      "\n",
+      "Third Servingman:\n",
+      "Now, well do with kindness of power the that\n",
+      "ence of the ambassage; but to shall here,\n",
+      "But saving as the wounds both drawn dead.\n",
+      "\n",
+      "CLARENCE:\n",
+      "Peace, sir, he adieu.\n",
+      "\n",
+      "MISTRESS OVERDONE:\n",
+      "Cheep, comes thou lovest which I do well;\n",
+      "Now now as that protectors, in this parliament--\n",
+      "That nature i' the mark of her with old her,\n",
+      "And so here doth in circle and call'd the wind,\n",
+      "If you do I am sure be your father's limbs,\n",
+      "Which is not drouble Paulina,\n",
+      "Where I come to me too much be to thus\n",
+      "As in a mind sorten but here hoped--followed as hung\n",
+      "bard at the time own that that he is death;\n",
+      "The stir weors shall be a new good groans.\n",
+      "\n",
+      "OXFORD:\n",
+      "And I say him to deed your city,\n",
+      "Thy heart advantage by mine at their world,\n",
+      "And be glood much show with upon the suddenly,\n",
+      "And in this free can company from here faults.\n",
+      "\n",
+      "YORK:\n",
+      "I bring thee, now I am again: it is a Montague.\n",
+      "\n",
+      "BRUTUS:\n",
+      "We while you not be the fault of the life.\n",
+      "\n",
+      "KING RICHARD III:\n",
+      "Come, sir, do I: in put I with so,\n",
+      "The rest redies with this ballad or suppher.\n",
+      "\n",
+      "Third Servingman:\n",
+      "Now, well do with kindness of power the that\n",
+      "ence of the ambassage; but to shall here,\n",
+      "But saving as the wounds both drawn dead.\n",
+      "\n",
+      "CLARENCE:\n",
+      "Peace, sir, he adieu.\n",
+      "\n",
+      "MISTRESS OVERDONE:\n",
+      "Cheep, comes thou lovest which I do well;\n",
+      "Now now as that protectors, in this parliament--\n",
+      "That nature i' the mark of her with old her,\n",
+      "And so here doth in circle and call'd the wind,\n",
+      "If you do I am sure be your father's limbs,\n",
+      "Which is not drouble Paulina,\n",
+      "Where I come to me too much be to thus\n",
+      "As in a mind sorten but here hoped--followed as hung\n",
+      "bard at the time own that that he is death;\n",
+      "The stir weors shall be a new good groans.\n",
+      "\n",
+      "OXFORD:\n",
+      "And I say him to deed your city,\n",
+      "Thy heart advantage by mine at their world,\n",
+      "And be glood much show with upon the suddenly,\n",
+      "And in this free can company from here faults.\n",
+      "\n",
+      "YORK:\n",
+      "I bring thee, now I am again: it is a Montague.\n",
+      "\n",
+      "BRUTUS:\n",
+      "We while you not be the fault of the life.\n",
+      "\n",
+      "KING RICHARD III:\n",
+      "Come, sir, do I: in put I with so,\n",
+      "The rest redies with this ballad or suppher.\n",
+      "\n",
+      "Third Servingman:\n",
+      "Now, well do with kindness of power the that\n",
+      "ence of the ambassage; but to shall here,\n",
+      "But saving as the wounds both drawn dead.\n",
+      "\n",
+      "CLARENCE:\n",
+      "Peace, sir, he adieu.\n",
+      "\n",
+      "MISTRESS OVERDONE:\n",
+      "Cheep, comes thou lovest which I do well;\n",
+      "Now now as that protectors, in this parliament--\n",
+      "That nature i' the mark of her with old her,\n",
+      "And so here doth in circle and call'd the wind,\n",
+      "If you do I am sure be your father's limbs,\n",
+      "Which is not drouble Paulina,\n",
+      "Where I come to me too much be to thus\n",
+      "As in a mind sorten but here hoped--followed as hung\n",
+      "bard at the time own that that he is death;\n",
+      "The stir weors shall be a new good groans.\n",
+      "\n",
+      "OXFORD:\n",
+      "And I say him to deed your city,\n",
+      "Thy heart advantage by mine at their world,\n",
+      "And be glood much show with upon the suddenly,\n",
+      "And in this free can company from here faults.\n",
+      "\n",
+      "YORK:\n",
+      "I bring thee, now I am again: it is a Montague.\n",
+      "\n",
+      "BRUTUS:\n",
+      "We while you not be the fault of the life.\n",
+      "\n",
+      "KING RICHARD III:\n",
+      "Come, sir, do I: in put I with so,\n",
+      "The rest redies with this ballad or suppher.\n",
+      "\n",
+      "Third Servingman:\n",
+      "Now, well do with kindness of power the that\n",
+      "ence of the ambassage; but to shall here,\n",
+      "But saving as the wounds both drawn dead.\n",
+      "\n",
+      "CLARENCE:\n",
+      "Peace, sir, he adieu.\n",
+      "\n",
+      "MISTRESS OVERDONE:\n",
+      "Cheep, comes thou lovest which I do well;\n",
+      "Now now as that protectors, in this parliament--\n",
+      "That nature i' the mark of her with old her,\n",
+      "And so here doth in circle and call'd the wind,\n",
+      "If you do I have some that to be blow.\n",
+      "\n",
+      "QUEEN ELIZABETH:\n",
+      "But shall we will not a man; for all one,\n",
+      "That you will not be but a fault and house,\n",
+      "I cannot strive at once from your father's son.\n",
       "\n",
       "GLOUCESTER:\n",
-      "No, Lord Hastings:\n",
-      "Else queen, though my trowbers grands me to-morrow\n",
-      "Here to Bolingbroke's match;\n",
-      "When the your life and spur at homely speak.\n",
+      "\n",
+      "CLARENCE:\n",
+      "That is the field to be put in his name.\n",
+      "\n",
+      "KING RICHARD III:\n",
+      "Ay, sir.\n",
+      "\n",
+      "QUEEN ELIZABETH:\n",
+      "How farewell.\n",
+      "\n",
+      "POLIXENES:\n",
+      "She hath stand and better me for you.\n",
+      "\n",
+      "QUEEN ELIZABETH:\n",
+      "I will be patient me.\n",
+      "\n",
+      "CLARENCE:\n",
+      "The been are the lawful strength, and he made\n",
+      "That me was upon the which he comes with a holy face.\n",
       "\n",
       "BUCKINGHAM:\n",
-      "My father was I follow: if you be your your kingdom,\n",
-      "My approbations an"
+      "The noble and of my tent.\n",
+      "\n",
+      "GLOUCESTER:\n",
+      "My gracious lord, he did soul the news,\n",
+      "Which he had of pretty with the fair traitor\n",
+      "Than a service hath some said for his prince.\n",
+      "\n",
+      "QUEEN MARGARET:\n",
+      "Stay the sad is mine execution,\n",
+      "Which can slain with you have some so death;\n",
+      "For I have sent done the many of the city\n",
+      "That you more so much hold me speak.\n",
+      "\n",
+      "GLOUCESTER:\n",
+      "\n",
+      "GLOUCESTER:\n",
+      "So doth fair lord, that be straight in his wit\n",
+      "The matter, the city distance.\n",
+      "\n",
+      "LEONTES:\n",
+      "The brother about the fire--thou sha\n",
+      "If you do I am sure be so princely false;\n",
+      "And the prince may do me for the house.\n",
+      "\n",
+      "GLOUCESTER:\n",
+      "Why, what wilt me, when we strike a mine eyes;\n",
+      "Of he doth deficers, in my fair deeds,\n",
+      "That thou wert from me hence his pardon lives.\n",
+      "\n",
+      "LORD FITZWATER:\n",
+      "Say, I will be ground to some of my body;\n",
+      "And what thou wilt part so supposed with him.\n",
+      "\n",
+      "CLARENCE:\n",
+      "You might be breathe in me?\n",
+      "\n",
+      "CAMILLO:\n",
+      "Stay now the first is no more resolved\n",
+      "With all the lawful and shed the gates word.\n",
+      "\n",
+      "CLARENCE:\n",
+      "Thou that, when my fair that I would not for heaven,\n",
+      "That fortune but not like her for the mind,\n",
+      "When I do be my tongues and her love.\n",
+      "\n",
+      "ANGELO:\n",
+      "She shall we not have done by his house.\n",
+      "\n",
+      "HENRY BOLINGBROKE:\n",
+      "Have you no more by my fortune should as heard.\n",
+      "\n",
+      "QUEEN ELIZABETH:\n",
+      "O, bring forth thy brother Gloucester,\n",
+      "But both an earthly of love from spoke me.\n",
+      "\n",
+      "QUEEN ELIZABETH:\n",
+      "Thou art thou as you shalt be dish'd in their head.\n",
+      "\n",
+      "Second Murderer:\n",
+      "The prince and my execution and heard it.\n",
+      "\n",
+      "KING RICHARD III:\n",
+      "Say, well I am a house in my heart;\n",
+      "But I am"
      ]
     }
    ],
    "source": [
-    "g_cuda = torch.Generator(device='cuda')\n",
+    "import gradio as gr\n",
     "\n",
-    "seed = \"\"\"\n",
-    "Plot histograms of the gradient values during training. If you notice a significant number of gradients are near zero (vanishing gradients) or very large values (exploding gradients), it could be a problem. TensorBoard is a useful tool for visualizing these histograms.\n",
-    "\"\"\"\n",
-    "\n",
-    "contexts = torch.tensor(encode_text(seed), dtype=torch.int32).to('cuda')\n",
-    "GEN_LENGTH=1024\n",
+    "demo = gr.Interface(\n",
+    "    fn=generate, \n",
+    "    inputs=[\n",
+    "        gr.Textbox(lines=2, placeholder=\"Prompt here...\"),\n",
+    "        gr.Number(value=256),\n",
+    "        gr.Number(value=0.8),\n",
+    "        gr.Slider(maximum=128,value=10),\n",
+    "        gr.Slider(maximum=1,value=1)\n",
+    "    ], \n",
+    "    outputs=\"text\",\n",
+    "    title=\"Shakespeare-GPT\",\n",
+    "    description=\"Putting theater kids out of their nonexistent jobs since 2023\"\n",
+    ")\n",
     "\n",
-    "model.eval()\n",
-    "for i in range(GEN_LENGTH):\n",
-    "    transform = nn.LogSoftmax(1)\n",
-    "    # What happens if GEN_LENGTH > CONTEXT? don't worry about it\n",
-    "    #x = F.pad(contexts[:, -BLOCK_SIZE:], (0, BLOCK_SIZE - contexts.size(0)), \"constant\", 0)\n",
-    "    x = contexts[-BLOCK_SIZE:]\n",
-    "    if x.size(0) < BLOCK_SIZE:\n",
-    "        x = F.pad(x, (0, BLOCK_SIZE - x.size(0)), \"constant\", 0).unsqueeze(0) # B*T\n",
-    "    else:\n",
-    "        x = x.unsqueeze(0)\n",
-    "\n",
-    "    preds = model.infer(x)\n",
-    "    preds = preds.squeeze(0)\n",
-    "    probs = torch.softmax(preds, dim=-1)\n",
-    "\n",
-    "    # TODO: Broken because of bug with the trailing 0s. FIX THIS\n",
-    "    # next_char = torch.multinomial(torch.exp(preds[(-1 if i >= BLOCK_SIZE else i), :]), num_samples=1, generator=g_cuda)\n",
-    "    next_char = torch.multinomial(torch.exp(preds[-1, :]), num_samples=1, generator=g_cuda)\n",
-    "\n",
-    "    #context = torch.cat(context, next_char)\n",
-    "    contexts = torch.cat((contexts, next_char), dim=0)\n",
-    "    print(decode_text(next_char.cpu().numpy())[-1], end=\"\")\n",
-    "\n",
-    "#print(\"\".join(decode_text(contexts.cpu().numpy())))"
+    "demo.launch()\n"
    ]
   }
  ],

gpt.py

@@ -0,0 +1,502 @@
+import argparse
+import gradio as gr
+import gc
+from torch.utils.tensorboard import SummaryWriter
+import os
+import json
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader, TensorDataset, random_split
+import pandas as pd
+import numpy as np
+import math
+from typing import Type
+
+
+# Define the command-line arguments
+parser = argparse.ArgumentParser(description='GPT CLI')
+parser.add_argument('--gui', action='store_true', help='Enable Gradio UI mode')
+parser.add_argument('--config', default='./gpt_config.json',
+                    help='Path to the config file')
+subparsers = parser.add_subparsers(dest='command', help='Choose a command')
+
+# Define the training command
+train_parser = subparsers.add_parser('train', help='Train the model')
+train_parser.add_argument('--load-from-restore', action='store_true',
+                          help='Load from restore path instead of training from scratch')
+
+# Define the evaluation command
+eval_parser = subparsers.add_parser('eval', help='Evaluate the model')
+eval_parser.add_argument('--data', default='./data/evaluation_data.txt',
+                         help='Path to the evaluation data file')
+
+# Define the inference command
+infer_parser = subparsers.add_parser(
+    'infer', help='Generate text from the model')
+infer_parser.add_argument('--text', type=str, required=True,
+                          help='Input text for generating continuation')
+infer_parser.add_argument('--length', type=int,
+                          default=100, help='Number of characters to generate')
+
+torch.manual_seed(1337)
+# Set device to CUDA if available, otherwise use CPU
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class GPTConfig:
+    def __init__(self, config_file_path):
+        with open(config_file_path, 'r') as f:
+            config = json.load(f)
+
+        # Architecture configuration
+        architecture_config = config['architecture']
+        self.embedding_dim = architecture_config['embedding_dim']
+        self.vocab_size = architecture_config['vocab_size']
+        self.context_size = architecture_config['context_size']
+        self.num_heads = architecture_config['num_heads']
+        self.num_layers = architecture_config['num_layers']
+
+        # Training configuration
+        training_config = config['training']
+        self.batch_size = training_config['batch_size']
+        self.training_data_path = training_config['training_data_path']
+        self.learning_rate = training_config['learning_rate']
+        self.num_steps = training_config['num_steps']
+        self.val_interval = training_config['val_interval']
+
+        # Checkpoint restore configuration
+        self.restore_path = config['restore_path']
+
+
+def encode_text(text):
+    # Simple dumb ASCII character-level "encoding" since all training data is ASCII.
+    # Consider better tokenization if moving off character-level
+    return ([ord(t) for t in text])
+
+
+def decode_text(indices):
+    return ([chr(x) for x in indices])
+
+
+class TextDataset(Dataset):
+    def __init__(self, data_tensor, context_size):
+        self.data_tensor = data_tensor
+        self.context_size = context_size
+
+    def __len__(self):
+        return len(self.data_tensor) - self.context_size
+
+    def __getitem__(self, index):
+        x = self.data_tensor[index:index + self.context_size]
+        y = self.data_tensor[index + 1:index + self.context_size + 1]
+
+        return x, y
+
+
+def load_dataset(data_path, val, context_size):
+    with open(data_path, 'r', encoding='utf-8') as f:
+        text = f.read()
+
+    # Tensorify data, put it in dataset
+    data = torch.tensor(encode_text(text), dtype=torch.int32)
+
+    test_split_idx = int(0.8 * len(data))
+    val_split_idx = int(0.9 * len(data))
+    train_data = data[:test_split_idx]
+    test_data = data[test_split_idx:val_split_idx]
+    val_data = data[val_split_idx:]
+    # print(f"{len(data)} chars of data")
+
+    train_dataset = TextDataset(train_data, context_size)
+    test_dataset = TextDataset(test_data, context_size)
+    val_dataset = TextDataset(test_data, context_size)
+    return ((train_dataset, test_dataset, val_dataset))
+
+
+class MultiheadAttention(nn.Module):
+    def __init__(self, embed_dim, num_heads, dropout=0.0, bias=True, device=None, dtype=None):
+        super(MultiheadAttention, self).__init__()
+
+        # Save variables
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.d_k = embed_dim // num_heads
+
+        self.Q = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.K = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.V = nn.Linear(embed_dim, embed_dim, bias=False)
+
+        self.dropout = nn.Dropout(dropout)
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+
+    def forward(self, query, key, value, attn_mask=None):
+        batch_size = query.size(0)
+
+        # Apply linear layers
+        q = self.Q(query)  # [B, C, E]
+        k = self.K(key)  # [B, C, E]
+        v = self.V(value)  # [B, C, E]
+
+        # Mutate dimensions so the attention matmul can get rid of the inner d_k
+        # [batch_size, num_heads, C, d_k]
+        q = q.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
+        # [batch_size, num_heads, C, d_k]
+        k = k.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
+        # [batch_size, num_heads, C, d_k]
+        v = v.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
+
+        # Get raw attention scores
+        scores = torch.matmul(q, k.transpose(-2, -1)) / \
+            math.sqrt(self.d_k)  # [B, num_heads, C, C]
+
+        # Apply mask, if necessary
+        if attn_mask is not None:
+            scores = scores.masked_fill(attn_mask, float('-inf'))
+
+        # Scale by sqrt(k)
+        attn = F.softmax(scores, dim=-1)
+        attn = self.dropout(attn)
+        out = attn @ v  # [B, num_heads, C, d_k]
+
+        # Concat and project
+        # Swap C and num_heads, force memory to coalesce, then fuse back num_heads and d_k together
+        out = out.transpose(1, 2).contiguous().view(
+            batch_size, -1, self.embed_dim)
+        # Project: give attention "time to think". Maybe this should be part of a different module but whatever
+        out = self.out_proj(out)
+        return ((out, None))
+
+
+class FeedForward(nn.Module):
+    def __init__(self, embed_dim, dropout):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(embed_dim, 4 * embed_dim),
+            nn.GELU(),
+            nn.Linear(4 * embed_dim, embed_dim),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x):
+        return (self.net(x))
+
+
+class Block(nn.Module):
+    """Self-attention"""
+
+    def __init__(self, embed_dim, num_heads, mask, dropout=0.2):
+        super(Block, self).__init__()
+        self.register_buffer("mask", mask)
+        self.head = MultiheadAttention(
+            embed_dim=embed_dim, num_heads=num_heads, dropout=dropout)
+        # self.head = nn.MultiheadAttention(embed_dim=embed_dim, num_heads=num_heads, dropout=dropout, batch_first=True)
+        self.ffwd = FeedForward(embed_dim=embed_dim, dropout=dropout)
+        self.ln1 = nn.LayerNorm(embed_dim)
+        self.ln2 = nn.LayerNorm(embed_dim)
+
+    def forward(self, x):
+        # Residual connections
+        x = self.ln1(x)
+        attn_output, _ = self.head(x, x, x, attn_mask=self.mask)
+        x = x + attn_output
+        out = x + self.ffwd(self.ln2(x))
+        return out
+
+
+class GPT(nn.Module):
+    def __init__(self, embedding_dim, vocab_size, context_size):
+        super(GPT, self).__init__()
+
+        self.embedding_dim = embedding_dim
+        self.output_dim = vocab_size
+        self.context_size = context_size
+
+        NUM_HEADS = 4
+        NUM_LAYERS = 4
+
+        # Initialize layers
+        self.tok_embed = nn.Embedding(vocab_size, embedding_dim)
+        self.pos_embed = nn.Embedding(context_size, embedding_dim)
+
+        mask = torch.tril(torch.ones(
+            self.context_size, self.context_size)).bool()
+        mask = ~mask
+        self.register_buffer("mask", mask)
+
+        self.blocks = nn.Sequential(
+            *[Block(embed_dim=embedding_dim, num_heads=NUM_HEADS, mask=mask, dropout=0.2) for _ in range(NUM_LAYERS)]
+        )
+
+        self.ln_f = nn.LayerNorm(self.embedding_dim)
+        # Final feed-forward layer from embeddings
+        self.ffwd = nn.Linear(
+            embedding_dim, out_features=vocab_size, bias=False)
+
+    def forward(self, x):
+        tok_embed = self.tok_embed(x)
+        pos_embed = self.pos_embed(
+            torch.arange(0, self.context_size, device="cuda")
+        )
+        x = tok_embed + pos_embed
+
+        x = self.blocks(x)
+        x = self.ln_f(x)
+
+        logits = self.ffwd(x)
+        return (logits)
+
+    def infer(self, x):
+        with torch.no_grad():
+            self.eval()
+            res = self.forward(x)
+            return (res)
+
+    def num_params(self):
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)
+
+
+def load_checkpoint(model, optimizer, path):
+    """
+    Loads a saved checkpoint file into the model and optimizer.
+
+    Args:
+        model (nn.Module): The PyTorch model to load the checkpoint into.
+        optimizer (torch.optim.Optimizer): The PyTorch optimizer to load the checkpoint into.
+        path (str): The path to the saved checkpoint file.
+
+    Returns:
+        Tuple[nn.Module, torch.optim.Optimizer]: The model and optimizer, loaded with the checkpoint state.
+    """
+    checkpoint = torch.load(path)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    return (model, optimizer)
+
+
+def save_checkpoint(model, optimizer, path, steps):
+    """
+    Saves a checkpoint of the model and optimizer to disk.
+
+    Args:
+        model (nn.Module): The PyTorch model to save the checkpoint of.
+        optimizer (torch.optim.Optimizer): The PyTorch optimizer to save the checkpoint of.
+        path (str): The path to save the checkpoint file.
+        steps (int): The number of training steps that have been completed.
+
+    Returns:
+        None
+    """
+    torch.save({
+        'steps': steps,
+        'model_state_dict': model.state_dict(),
+        'optimizer_state_dict': optimizer.state_dict(),
+    }, path)
+
+
+def compute_loss(model, criterion, x, y):
+    logits = model(x)
+    B, C, V = logits.shape
+    logits = logits.view(B*C, V)
+    y = y.view(B*C)
+    loss = F.cross_entropy(logits, y.long())
+    return loss
+
+
+def train(model, optimizer, config: Type[GPTConfig]):
+    model = model.to(device)
+    criterion = F.cross_entropy
+
+    global_step = 0
+
+    train_dataset, val_dataset = load_dataset(
+        config.training_data_path, None, model.context_size)
+
+    train_dataloader = DataLoader(
+        train_dataset,
+        batch_size=config.batch_size,
+        shuffle=True,
+        num_workers=4
+    )
+
+    test_dataloader = DataLoader(
+        test_dataset, batch_size=512, num_workers=4, shuffle=True)
+
+    model.train()
+
+    EPOCHS = 1
+    STEPS = config.num_steps
+    VAL_INTERVAL = 100
+
+    writer = SummaryWriter()
+
+    step = 0
+
+    for epoch in range(EPOCHS):
+        for data, target in train_dataloader:
+            data = data.to(device)
+            target = target.to(device)
+
+            loss = compute_loss(model, criterion, data, target)
+
+            # Backward pass
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            writer.add_scalar(
+                "Loss/train", loss.cpu().detach().numpy(), global_step)
+            global_step += 1
+
+            if step % VAL_INTERVAL == 0:
+                total_loss = 0
+                total_samples = 0
+
+                with torch.no_grad():
+                    model.eval()
+                    for x, y in test_dataloader:
+                        x = x.to(device)
+                        y = y.to(device)
+
+                        batch_loss = compute_loss(model, criterion, x, y)
+                        total_loss += batch_loss.item() * 512
+                        total_samples += 512
+                        if total_samples > 10:
+                            break
+
+                model.train()
+                average_loss = total_loss / total_samples
+
+                print(f"Step {step}; loss: {average_loss}")
+                writer.add_scalar("Loss/val", average_loss, global_step)
+
+            step += 1
+            if step >= STEPS:
+                break
+
+    writer.close()
+
+
+def evaluate_model(model, val_dataset, block_size=512, max_samples=100000):
+    model.eval()
+    total_loss = 0.0
+    total_samples = 0
+    criterion = F.cross_entropy
+
+    val_dataloader = DataLoader(
+        val_dataset, batch_size=block_size, num_workers=4)
+    with torch.no_grad():
+        for inputs, targets in val_dataloader:
+            inputs = inputs.to(device)
+            targets = targets.to(device)
+
+            batch_loss = compute_loss(model, criterion, inputs, targets)
+            total_loss += batch_loss.item() * inputs.size(0)
+            total_samples += inputs.size(0)
+            if total_samples > max_samples:
+                break
+
+    average_loss = total_loss / total_samples
+    return average_loss
+
+
+def generate(model, config, prompt, gen_length, temp=1, top_k=10, top_p=None):
+    g_cuda = torch.Generator(device=device)
+    contexts = torch.tensor(encode_text(prompt), dtype=torch.int32).to(device)
+
+    model.eval()
+    for i in range(gen_length):
+        transform = nn.LogSoftmax(1)
+        x = contexts[-config.context_size:]
+        if x.size(0) < config.context_size:
+            x = F.pad(x, (config.context_size - x.size(0), 0),
+                      "constant", 0).unsqueeze(0)  # B*T
+        else:
+            x = x.unsqueeze(0)
+
+        preds = model.infer(x)
+        preds = preds.squeeze(0)
+        preds = preds / temp
+        probs = F.softmax(preds, dim=-1)
+
+        if top_p is not None:
+            # Apply top-p
+            sorted_probs, sorted_indices = torch.sort(
+                probs[-1, :], descending=True)
+            cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+            # find cutoff
+            idx_top_p = (cumulative_probs < top_p).sum().item()
+            top_probs = sorted_probs[:idx_top_p]
+            top_indices = sorted_indices[:idx_top_p]
+            # Null case
+            if top_probs.size(0) == 0:
+                top_probs = sorted_probs[:1]
+                top_indices = sorted_indices[:1]
+
+            next_char = torch.multinomial(
+                top_probs, num_samples=1, generator=g_cuda)
+            next_char = top_indices[next_char]
+        elif top_k is not None:
+            top_k_probs, top_k_indices = torch.topk(probs[-1, :], k=top_k)
+            next_char = torch.multinomial(
+                top_k_probs, num_samples=1, generator=g_cuda)
+            next_char = top_k_indices[next_char]
+        else:
+            next_char = torch.multinomial(
+                probs, num_samples=1, generator=g_cuda)
+
+        contexts = torch.cat((contexts, next_char), dim=0)
+        print(decode_text(next_char.cpu().numpy())[-1], end="")
+
+    return ("".join(decode_text(contexts.cpu().numpy())))
+
+
+def main():
+    # Parse the command-line arguments
+    args = parser.parse_args()
+
+    config = GPTConfig(args.config)
+    # Create the GPT model
+    model = GPT(
+        vocab_size=config.vocab_size,
+        context_size=config.context_size,
+        embedding_dim=config.embedding_dim
+    )
+    model.to(device)
+
+    optimizer = optim.AdamW(model.parameters(), lr=config.learning_rate)
+    if args.gui:
+        load_checkpoint(model, optimizer, config.restore_path)
+        demo = gr.Interface(
+            fn=lambda *args: generate(model, config, *args),
+            inputs=[
+                gr.Textbox(lines=2, placeholder="Prompt here..."),
+                gr.Number(precision=0, value=256),
+                gr.Number(value=0.8),
+                gr.Slider(maximum=128, value=10),
+                gr.Slider(maximum=1, value=1)
+            ],
+            outputs="text",
+            title="Shakespeare-GPT",
+            description="Putting theater kids out of their nonexistent jobs since 2023"
+        )
+
+        demo.launch()
+    elif args.command == "train":
+        if args.load_from_restore:
+            load_checkpoint(model, optimizer, path)
+
+        train(model, config)
+    elif args.command == "eval":
+        _, _, test_dataset = load_dataset(
+            config.training_data_path, None, model.context_size)
+        evaluate_model(model, test_dataset)
+    elif args.command == "infer":
+        prompt = args.text
+        generated_text = generate(model, config, prompt, args.length)
+        print(generated_text)
+
+
+if __name__ == "__main__":
+    main()

gpt_config.json

@@ -0,0 +1,17 @@
+{
+    "architecture": {
+        "embedding_dim": 256,
+        "vocab_size": 128,
+        "context_size": 256,
+        "num_heads": 4,
+        "num_layers": 4
+    },
+    "training": {
+        "batch_size": 64,
+        "training_data_path": "datasets/training_data.txt",
+        "learning_rate": 3e-4,
+        "num_steps": 5000, 
+        "val_interval": 100
+    },
+    "restore_path": "checkpoints/model.pt"
+}