Complete_v2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "005c77f7-afbf-4c22-80ef-d67997fe5dc1",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2823cbd6-819a-4a65-94a5-0cc4161bfec9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x7f849c020a50>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "torch.manual_seed(1337)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c9dac65-9892-4386-a82f-22ec4b6a6f34",
   "metadata": {},
   "source": [
    "## Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b99dc9a5-8ff8-4326-8328-f771978c6ad0",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 64\n",
    "block_size = 256\n",
    "max_iters = 5000\n",
    "eval_interval = 500\n",
    "learning_rate = 3e-4\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "eval_iters = 200\n",
    "n_embd = 384\n",
    "n_head = 6\n",
    "n_layer = 6\n",
    "dropout = 0.2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d52e6d8-e37a-4c7e-ac8d-21cce05a1585",
   "metadata": {},
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba1e8ea7-757c-40b0-b8ce-fb1c9d067ab4",
   "metadata": {},
   "source": [
    "### Get and read data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a346777e-794f-434a-a362-5b58bdc9beb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We always start with a dataset to train on. Let's download the tiny shakespeare dataset\n",
    "# !wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt\n",
    "with open(\"input.txt\", \"r\", encoding=\"utf-8\") as f:\n",
    "    text = f.read()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b793e7ff-373c-4a01-ae8e-6d0165125353",
   "metadata": {},
   "source": [
    "### Minor exploration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f4eb7731-a1e2-4ef3-85de-450ea82e6032",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Characters occuring in text = ['\\n', ' ', '!', '$', '&', \"'\", ',', '-', '.', '3', ':', ';', '?', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']\n",
      "Vocabulary size = 65\n"
     ]
    }
   ],
   "source": [
    "# All unique characters that occur in this text\n",
    "chars = sorted(list(set(text)))\n",
    "vocab_size = len(chars)\n",
    "print (f\"Characters occuring in text = {chars}\\nVocabulary size = {vocab_size}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "533656d3-19d3-4e72-ad2a-933a9f026dcd",
   "metadata": {},
   "source": [
    "### Create a mapping (for encoding and decoding)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "47de8609-a8b9-451e-ae7e-c438f46d5683",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a mapping from characters to integers and vice versa\n",
    "stoi = {ch: i for i, ch in enumerate(chars)}\n",
    "itos = {i: ch for i, ch in enumerate(chars)}\n",
    "def encode(s):\n",
    "    return [stoi[c] for c in s]\n",
    "def decode(l):\n",
    "    return \"\".join([itos[i] for i in l])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbbfff9b-c956-4421-8143-ce03ff788908",
   "metadata": {},
   "source": [
    "### Splits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31ebe967-3de0-40e5-a3e1-545e7b97a106",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train and test splits\n",
    "data = torch.tensor(encode(text), dtype=torch.long)\n",
    "n = int(0.9 * len(data))\n",
    "train_data = data[:n]\n",
    "val_data = data[n:]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50bbe8e2-af49-4246-822a-98a39d5fc0e4",
   "metadata": {},
   "source": [
    "### Data loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e8bf080-a41e-4a02-8e00-5cc1914ff99b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_batch(split):\n",
    "    # Generate a small batch of data of inputs x and targets y\n",
    "    data = train_data if split == \"train\" else val_data\n",
    "    ix = torch.randint(len(data) - block_size, (batch_size,))\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",
    "    x, y = x.to(device), y.to(device)\n",
    "    return x, y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e98a0344-b21b-4df7-9acf-9c2ebb2933be",
   "metadata": {},
   "source": [
    "### Loss function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "accd0e0e-5bf5-4cb2-bcd7-6e7018795f60",
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def estimate_loss():\n",
    "    out = {}\n",
    "    model.eval()\n",
    "    for split in [\"train\", \"val\"]:\n",
    "        losses = torch.zeros(eval_iters)\n",
    "        for k in range(eval_iters):\n",
    "            X, Y = get_batch(split)\n",
    "            _, loss = model(X, Y)\n",
    "            losses[k] = loss.item()\n",
    "        out[split] = losses.mean()\n",
    "    model.train()\n",
    "    return out\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00fde2e3-5b48-46b7-bb7f-f39d399e7af6",
   "metadata": {},
   "source": [
    "### Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "adc9b087-3d9f-40fa-b4d4-3bda821a80ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Head(nn.Module):\n",
    "    \"\"\"One head of self-attention\"\"\"\n",
    "\n",
    "    def __init__(self, head_size) -> None:\n",
    "        super().__init__()\n",
    "        self.key = nn.Linear(n_embd, head_size, bias=False)\n",
    "        self.query = nn.Linear(n_embd, head_size, bias=False)\n",
    "        self.value = nn.Linear(n_embd, head_size, bias=False)\n",
    "        self.register_buffer(\"tril\", torch.tril(torch.ones(block_size, block_size)))\n",
    "\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        B, T, C = x.shape\n",
    "        k = self.key(x)  # (B, T, C)\n",
    "        q = self.query(x)  # (B, T, C)\n",
    "\n",
    "        # Compute attention scores (\"affininities\")\n",
    "        wei = q @ k.transpose(-2, -1) * C**-0.5  # (B, T, C) @ (B, C, T) --> (B, T, T)\n",
    "        wei = wei.masked_fill(self.tril[:T, :T] == 0, float(\"-inf\"))\n",
    "        wei = F.softmax(wei, dim=-1)  # (B, T, T)\n",
    "        wei = self.dropout(wei)\n",
    "\n",
    "        # Perform weighted aggregation of the values\n",
    "        v = self.value(x)  # (B, T, C)\n",
    "        out = wei @ v  # (B, T, T) @ (B, T, C) --> (B, T, C)\n",
    "        return out\n",
    "\n",
    "\n",
    "class MultiHeadAttention(nn.Module):\n",
    "    \"\"\"Multiple heads of self-attention in parallel\"\"\"\n",
    "\n",
    "    def __init__(self, num_heads, head_size) -> None:\n",
    "        super().__init__()\n",
    "        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])\n",
    "        self.proj = nn.Linear(n_embd, n_embd)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = torch.cat([h(x) for h in self.heads], dim=-1)\n",
    "        out = self.dropout(self.proj(out))\n",
    "        return out\n",
    "\n",
    "\n",
    "class FeedForward(nn.Module):\n",
    "    \"\"\"A simple linear layer followed by a non-linearity\"\"\"\n",
    "\n",
    "    def __init__(self, n_embd) -> None:\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(n_embd, 4 * n_embd),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(4 * n_embd, n_embd),\n",
    "            nn.Dropout(dropout),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)\n",
    "\n",
    "\n",
    "class Block(nn.Module):\n",
    "    \"\"\"Transformer block: Communication followed by computation\"\"\"\n",
    "\n",
    "    def __init__(self, n_embd, n_head) -> None:\n",
    "        # n_embd: Embedding dimension, n_head: The number of heads we'd like\n",
    "        super().__init__()\n",
    "        head_size = n_embd // n_head\n",
    "        self.sa = MultiHeadAttention(n_head, head_size)\n",
    "        self.ffwd = FeedForward(n_embd)\n",
    "        self.ln1 = nn.LayerNorm(n_embd)\n",
    "        self.ln2 = nn.LayerNorm(n_embd)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x + self.sa(self.ln1(x))\n",
    "        x = x + self.ffwd(self.ln2(x))\n",
    "        return x\n",
    "\n",
    "\n",
    "class BigramLanguageModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        # Each token directly reads off the logits for the next token from a lookup table\n",
    "        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)\n",
    "        self.position_embedding_table = nn.Embedding(block_size, n_embd)\n",
    "        self.blocks = nn.Sequential(\n",
    "            *[Block(n_embd, n_head=n_head) for _ in range(n_layer)]\n",
    "        )\n",
    "        self.ln_f = nn.LayerNorm(n_embd)  # Final layer norm\n",
    "        self.lm_head = nn.Linear(n_embd, vocab_size)\n",
    "\n",
    "    def forward(self, idx, targets=None):\n",
    "        B, T = idx.shape\n",
    "\n",
    "        # idx and targets are both (B,T) tensor of integers.\n",
    "        tok_emb = self.token_embedding_table(idx)  # (B,T,C)\n",
    "        pos_emb = self.position_embedding_table(\n",
    "            torch.arange(T, device=device)\n",
    "        )  # (T, C)\n",
    "        x = tok_emb + pos_emb  # (B, T, C)\n",
    "        x = self.blocks(x)  # (B, T, C)\n",
    "        logits = self.lm_head(x)  # (B, T, vocab_size)\n",
    "\n",
    "        if targets is None:\n",
    "            loss = None\n",
    "        else:\n",
    "            B, T, C = logits.shape\n",
    "            logits = logits.view(B * T, C)\n",
    "            targets = targets.view(B * T)\n",
    "            loss = F.cross_entropy(logits, targets)\n",
    "\n",
    "        return logits, loss\n",
    "\n",
    "    def generate(self, idx, max_new_tokens):\n",
    "        # idx is (B,T) array of indices in the current context\n",
    "        for _ in range(max_new_tokens):\n",
    "            # Crop idx to the last block_size tokens\n",
    "            idx_cond = idx[:, -block_size:]\n",
    "            # Get the predictions\n",
    "            logits, loss = self.forward(idx_cond)\n",
    "            # focus only on the last time step\n",
    "            logits = logits[:, -1, :]  # Becomes (B,C)\n",
    "            # Apply softmax to get probabilities\n",
    "            probs = F.softmax(logits, dim=-1)  # (B,C)\n",
    "            # Sample from the distribution\n",
    "            idx_next = torch.multinomial(probs, num_samples=1)  # (B,1)\n",
    "            # Append sampled index to the running sequence\n",
    "            idx = torch.cat((idx, idx_next), dim=1)  # (B, T+1)\n",
    "        return idx\n",
    "\n",
    "model = BigramLanguageModel()\n",
    "m = model.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff4ed84d-dff5-4852-868d-1aebf7e4b188",
   "metadata": {},
   "source": [
    "### Optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9fe64e4f-2a4c-45e6-aff4-eefa2b6fa7b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a PyTorch optimizer\n",
    "# optimizer = torch.optim.AdamW(m.parameters(), lr=learning_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e25a8938-cb72-4c97-9ed6-5896ad0c34dc",
   "metadata": {},
   "source": [
    "### Training & Saving Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ad297e40-42f8-4875-9f95-f80536fd2545",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# for iter in range(max_iters):\n",
    "#     if iter % eval_interval == 0:\n",
    "#         losses = estimate_loss()\n",
    "#         print(\n",
    "#             f\"Step {iter}: Train loss = {losses['train']:.4f}, Val loss = {losses['val']:.4f}\"\n",
    "#         )\n",
    "\n",
    "#     xb, yb = get_batch(\"train\")\n",
    "\n",
    "#     # Evaluate the loss\n",
    "#     logits, loss = m(xb, yb)\n",
    "#     optimizer.zero_grad(set_to_none=True)\n",
    "#     loss.backward()\n",
    "#     optimizer.step()\n",
    "\n",
    "# # Saving Model\n",
    "# torch.save(model.state_dict(), model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fc94ef1-3a5b-4219-a612-e7ec8c9f24b1",
   "metadata": {},
   "source": [
    "### Load model (from state_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "04f89b71-ebdf-4f20-8bb0-3c718c2ba203",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model_path = 'normal_model.pth'\n",
    "model_path = 'overfitted_model.pth'\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    model.load_state_dict(torch.load(model_path))\n",
    "else:\n",
    "    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))   "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac13412c-daf8-4c82-9380-91f5385a2c8e",
   "metadata": {},
   "source": [
    "### Run Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "69a229ab-827d-4d37-a5f5-c83080dbd6e8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Cold and gates i' the fuel Hout way the lunity\n",
      "That lack no more is deaful father, like an\n",
      "Ershalm, and freed in dream, and kneel devise.\n",
      "Stanley, my ghous in my friend,\n",
      "Speak like lane hath budgeneral in this victory.\n",
      "Give me a brief, that cracking impatiety,\n",
      "Corred o' the certail state, to returneys.\n",
      "\n",
      "ROMEO:\n",
      "I will but the forsworn: if ever either stake,\n",
      "Give nothing to his grief, and vigility\n",
      "A man neight that did pust the gate.\n",
      "But O this time to present York's mine and mind\n",
      "Shall benefit th\n"
     ]
    }
   ],
   "source": [
    "context = torch.zeros((1, 1), dtype=torch.long).to(device)\n",
    "print(decode(m.generate(context, max_new_tokens=500)[0].tolist()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21394978-48cb-4317-bd36-02efc01e326e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}