usage.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from tqdm import tqdm\n",
    "from pathlib import Path\n",
    "from model import MiniGPT\n",
    "from config import MiniGPTConfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_filepath = Path(\"./data.txt\")\n",
    "trained_model_filepaht = Path(\"./mini_gpt.pt\")\n",
    "batch_size = 32\n",
    "block_size = 8"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create Vocab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(dataset_filepath,\"r\") as file:\n",
    "    text = file.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create vocab for tokenizer\n",
    "chars = sorted(list(set(text)))\n",
    "vocab_size = len(chars)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create lookup table\n",
    "stoi = {ch: i for i, ch in enumerate(chars)}\n",
    "itos = {i: ch for i, ch in enumerate(chars)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create simple encoder and decoder in global namespace\n",
    "encode = lambda s: [stoi[c] for c in s]\n",
    "decode = lambda l: \"\".join([itos[i] for i in l])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create Simple Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Encode dataset and store it as a torch long tensor.\n",
    "data = torch.tensor(encode(text), dtype=torch.long)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_batch(split):\n",
    "    \"\"\"\n",
    "    generate a small batch of data of inputs x and targets y.\n",
    "    this helps you to get random chunk in the entire dataset\n",
    "    \"\"\"\n",
    "    # Create four random numbers.\n",
    "    ix = torch.randint(len(data) - block_size, (batch_size,)) # exp. tensor([ 76049, 234249, 934904, 560986])\n",
    "    # Create x and y\n",
    "    x = torch.stack([data[i:i+block_size] for i in ix])\n",
    "    y = torch.stack([data[i+1:i+block_size+1] for i in ix])\n",
    "    return x, y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check if encoder/decoder work properly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[46, 47, 47, 1, 58, 46, 43, 56, 43]\n",
      "hii there\n"
     ]
    }
   ],
   "source": [
    "print(encode(\"hii there\"))\n",
    "print(decode(encode(\"hii there\")))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_config = MiniGPTConfig()\n",
    "model = MiniGPT(model_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MiniGPT(\n",
       "  (token_embedder): Embedding(65, 768)\n",
       "  (positional_embedder): Embedding(1024, 768)\n",
       "  (dropout): Dropout(p=0.0, inplace=False)\n",
       "  (decoder_layers): ModuleList(\n",
       "    (0-11): 12 x Decoder(\n",
       "      (layer_norm1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (layer_norm2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (masked_multihead_attention): CausalSelfAttention(\n",
       "        (attention_projection): Linear(in_features=768, out_features=2304, bias=True)\n",
       "        (output_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        (residual_dropout): Dropout(p=0.0, inplace=False)\n",
       "      )\n",
       "      (feed_forward): FeedForward(\n",
       "        (linear1): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (gelu): GELU(approximate='none')\n",
       "        (linear2): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (layer_normalizar): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "  (language_model_head): Linear(in_features=768, out_features=65, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = torch.device(\"cuda\")\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Just some random text to generate to see it's total non-sense."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      ":$gq3fUeoCiuooMtT\n",
      "kyWhLEEocl\n",
      "'fybK\n",
      "QcEo.yOWR:'ynWCahe$T:yc?gwOK\n",
      "o-PHQunRu.Ux3.jdWRIlYAcEQ.l;j,\n",
      "?RhIv\n"
     ]
    }
   ],
   "source": [
    "# Text generation is going to start with new line character\n",
    "idx = torch.zeros((1,1), dtype= torch.long, device=device)\n",
    "# Use index 0 to unplug a sample\n",
    "print(decode(model.generate(idx, max_new_tokens=100)[0].tolist()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since this is just a simple demo we won't use tensorboard and only print out models loss before and after the training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Random sample model loss before training: 4.333423137664795\n"
     ]
    }
   ],
   "source": [
    "xb, yb = get_batch(\"train\")\n",
    "xb = xb.to(device)\n",
    "yb = yb.to(device)\n",
    "logits, loss = model(xb, yb)\n",
    "print(f\"Random sample model loss before training: {loss.item()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 50000/50000 [19:04<00:00, 43.69it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Random sample model loss after training: 1.572378158569336\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)\n",
    "model.train()\n",
    "for steps in tqdm(range(50000)):\n",
    "    xb, yb = get_batch(\"train\")\n",
    "    xb = xb.to(device)\n",
    "    yb = yb.to(device)\n",
    "    logits, loss = model(xb, yb)\n",
    "    optimizer.zero_grad(set_to_none=True)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "print(f\"Random sample model loss after training: {loss.item()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save trained model\n",
    "model.state_dict()\n",
    "torch.save(model.state_dict(), trained_model_filepaht)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate some text to see if model has learned a little or not (Don't expect anything crazy this is just to see if training is working or not.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "son, your thor twalllld whounourne ougeld y thinou\n"
     ]
    }
   ],
   "source": [
    "model.eval()\n",
    "with torch.inference_mode():\n",
    "    print(decode(model.generate(idx, max_new_tokens=50)[0].tolist())) # 0 is to unplug a sample"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ml",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}