mahagpt_dev.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/asrjy/mahaGPT/blob/main/mahagpt_dev.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "GbWmNMqtoDWo",
    "outputId": "b44df03a-b4dd-4b28-8ccd-91b454a4193a"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x7fc0f8308430>"
      ]
     },
     "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(420)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "jpRGj5vuXa-A"
   },
   "outputs": [],
   "source": [
    "with open('mahabharata.txt', 'r', encoding = 'utf-8') as f:\n",
    "  text = f.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "g-NznPM6YPIc",
    "outputId": "a323c56c-ed7b-4e42-dca0-c77d552c18cf"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "14929983"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "MNeNKE7jYQ1a",
    "outputId": "fdf5e30e-64fb-40a6-d6be-3619edf74b6d"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ADI PARVA\n",
      "\n",
      "SECTION I\n",
      "\n",
      "Om! Having bowed down to Narayana and Nara, the most exalted male being,\n",
      "and also to the goddess Saraswati, must the word Jaya be uttered.\n",
      "\n",
      "Ugrasrava, the son of Lomaharshana, surnamed Sauti, well-versed in the\n",
      "Puranas, bending with humility, one day approached the great sages of\n",
      "rigid vows, sitting at their ease, who had attended the twelve years'\n",
      "sacrifice of Saunaka, surnamed Kulapati, in the forest of Naimisha. Those\n",
      "ascetics, wishing to hear his wonderful narrations, presently began to\n",
      "address him who had thus arrived at that recluse abode of the inhabitants\n",
      "of the forest of Naimisha. Having been entertained with due respect by\n",
      "those holy men, he saluted those Munis (sages) with joined palms, even\n",
      "all of them, and inquired about the progress of their asceticism. Then\n",
      "all the ascetics being again seated, the son of Lomaharshana humbly\n",
      "occupied the seat that was assigned to him. Seeing that he was\n",
      "comfortably seated, and recovered from fatigue, one of the Rishi\n"
     ]
    }
   ],
   "source": [
    "print(text[:1000])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "AkIuEPNmYvNK",
    "outputId": "e5c765aa-5f11-4dc5-b0d6-dde3b53be3f7"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " !\"&'(),-.0123456789:;?ABCDEFGHIJKLMNOPQRSTUVWXYZ[]_abcdefghijklmnopqrstuvwxyz\n",
      "79\n"
     ]
    }
   ],
   "source": [
    "chars = sorted(list(set(text)))\n",
    "vocab_size = len(chars)\n",
    "print(''.join(chars))\n",
    "print(vocab_size)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "LGJuIkXhZECN"
   },
   "outputs": [],
   "source": [
    "stoi = {ch:i for i, ch in enumerate(chars)}\n",
    "itos = {i:ch for i, ch in enumerate(chars)}\n",
    "\n",
    "encode = lambda s: [stoi[c] for c in s]\n",
    "decode = lambda l: ''.join([itos[i] for i in l])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "AToGsY4SqHNm",
    "outputId": "aec590f9-f403-462b-bd02-b23cef2e3af9"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[77, 67, 67, 1, 75, 53, 71, 71, 73, 68]\n",
      "yoo wassup\n"
     ]
    }
   ],
   "source": [
    "print(encode(\"yoo wassup\"))\n",
    "print(decode(encode(\"yoo wassup\")))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "H-mxu4PSqM_b",
    "outputId": "e56a4160-58bb-4a30-91a7-1a76c7ad453a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([14929983]) torch.int64\n",
      "tensor([24, 27, 32,  1, 39, 24, 41, 45, 24,  0,  0, 42, 28, 26, 43, 32, 38, 37,\n",
      "         1, 32,  0,  0, 38, 65,  2,  1, 31, 53, 74, 61, 66, 59,  1, 54, 67, 75,\n",
      "        57, 56,  1, 56, 67, 75, 66,  1, 72, 67,  1, 37, 53, 70, 53, 77, 53, 66,\n",
      "        53,  1, 53, 66, 56,  1, 37, 53, 70, 53,  8,  1, 72, 60, 57,  1, 65, 67,\n",
      "        71, 72,  1, 57, 76, 53, 64, 72, 57, 56,  1, 65, 53, 64, 57,  1, 54, 57,\n",
      "        61, 66, 59,  8,  0, 53, 66, 56,  1, 53, 64, 71, 67,  1, 72, 67,  1, 72,\n",
      "        60, 57,  1, 59, 67, 56, 56, 57, 71, 71,  1, 42, 53, 70, 53, 71, 75, 53,\n",
      "        72, 61,  8,  1, 65, 73, 71, 72,  1, 72, 60, 57,  1, 75, 67, 70, 56,  1,\n",
      "        33, 53, 77, 53,  1, 54, 57,  1, 73, 72, 72, 57, 70, 57, 56, 10,  0,  0,\n",
      "        44, 59, 70, 53, 71, 70, 53, 74, 53,  8,  1, 72, 60, 57,  1, 71, 67, 66,\n",
      "         1, 67, 58,  1, 35, 67, 65, 53, 60, 53, 70, 71, 60, 53, 66, 53,  8,  1,\n",
      "        71, 73, 70, 66, 53, 65, 57, 56,  1, 42, 53, 73, 72, 61,  8,  1, 75, 57,\n",
      "        64, 64,  9, 74, 57, 70, 71, 57, 56,  1, 61, 66,  1, 72, 60, 57,  0, 39,\n",
      "        73, 70, 53, 66, 53, 71,  8,  1, 54, 57, 66, 56, 61, 66, 59,  1, 75, 61,\n",
      "        72, 60,  1, 60, 73, 65, 61, 64, 61, 72, 77,  8,  1, 67, 66, 57,  1, 56,\n",
      "        53, 77,  1, 53, 68, 68, 70, 67, 53, 55, 60, 57, 56,  1, 72, 60, 57,  1,\n",
      "        59, 70, 57, 53, 72,  1, 71, 53, 59, 57, 71,  1, 67, 58,  0, 70, 61, 59,\n",
      "        61, 56,  1, 74, 67, 75, 71,  8,  1, 71, 61, 72, 72, 61, 66, 59,  1, 53,\n",
      "        72,  1, 72, 60, 57, 61, 70,  1, 57, 53, 71, 57,  8,  1, 75, 60, 67,  1,\n",
      "        60, 53, 56,  1, 53, 72, 72, 57, 66, 56, 57, 56,  1, 72, 60, 57,  1, 72,\n",
      "        75, 57, 64, 74, 57,  1, 77, 57, 53, 70, 71,  5,  0, 71, 53, 55, 70, 61,\n",
      "        58, 61, 55, 57,  1, 67, 58,  1, 42, 53, 73, 66, 53, 63, 53,  8,  1, 71,\n",
      "        73, 70, 66, 53, 65, 57, 56,  1, 34, 73, 64, 53, 68, 53, 72, 61,  8,  1,\n",
      "        61, 66,  1, 72, 60, 57,  1, 58, 67, 70, 57, 71, 72,  1, 67, 58,  1, 37,\n",
      "        53, 61, 65, 61, 71, 60, 53, 10,  1, 43, 60, 67, 71, 57,  0, 53, 71, 55,\n",
      "        57, 72, 61, 55, 71,  8,  1, 75, 61, 71, 60, 61, 66, 59,  1, 72, 67,  1,\n",
      "        60, 57, 53, 70,  1, 60, 61, 71,  1, 75, 67, 66, 56, 57, 70, 58, 73, 64,\n",
      "         1, 66, 53, 70, 70, 53, 72, 61, 67, 66, 71,  8,  1, 68, 70, 57, 71, 57,\n",
      "        66, 72, 64, 77,  1, 54, 57, 59, 53, 66,  1, 72, 67,  0, 53, 56, 56, 70,\n",
      "        57, 71, 71,  1, 60, 61, 65,  1, 75, 60, 67,  1, 60, 53, 56,  1, 72, 60,\n",
      "        73, 71,  1, 53, 70, 70, 61, 74, 57, 56,  1, 53, 72,  1, 72, 60, 53, 72,\n",
      "         1, 70, 57, 55, 64, 73, 71, 57,  1, 53, 54, 67, 56, 57,  1, 67, 58,  1,\n",
      "        72, 60, 57,  1, 61, 66, 60, 53, 54, 61, 72, 53, 66, 72, 71,  0, 67, 58,\n",
      "         1, 72, 60, 57,  1, 58, 67, 70, 57, 71, 72,  1, 67, 58,  1, 37, 53, 61,\n",
      "        65, 61, 71, 60, 53, 10,  1, 31, 53, 74, 61, 66, 59,  1, 54, 57, 57, 66,\n",
      "         1, 57, 66, 72, 57, 70, 72, 53, 61, 66, 57, 56,  1, 75, 61, 72, 60,  1,\n",
      "        56, 73, 57,  1, 70, 57, 71, 68, 57, 55, 72,  1, 54, 77,  0, 72, 60, 67,\n",
      "        71, 57,  1, 60, 67, 64, 77,  1, 65, 57, 66,  8,  1, 60, 57,  1, 71, 53,\n",
      "        64, 73, 72, 57, 56,  1, 72, 60, 67, 71, 57,  1, 36, 73, 66, 61, 71,  1,\n",
      "         6, 71, 53, 59, 57, 71,  7,  1, 75, 61, 72, 60,  1, 62, 67, 61, 66, 57,\n",
      "        56,  1, 68, 53, 64, 65, 71,  8,  1, 57, 74, 57, 66,  0, 53, 64, 64,  1,\n",
      "        67, 58,  1, 72, 60, 57, 65,  8,  1, 53, 66, 56,  1, 61, 66, 69, 73, 61,\n",
      "        70, 57, 56,  1, 53, 54, 67, 73, 72,  1, 72, 60, 57,  1, 68, 70, 67, 59,\n",
      "        70, 57, 71, 71,  1, 67, 58,  1, 72, 60, 57, 61, 70,  1, 53, 71, 55, 57,\n",
      "        72, 61, 55, 61, 71, 65, 10,  1, 43, 60, 57, 66,  0, 53, 64, 64,  1, 72,\n",
      "        60, 57,  1, 53, 71, 55, 57, 72, 61, 55, 71,  1, 54, 57, 61, 66, 59,  1,\n",
      "        53, 59, 53, 61, 66,  1, 71, 57, 53, 72, 57, 56,  8,  1, 72, 60, 57,  1,\n",
      "        71, 67, 66,  1, 67, 58,  1, 35, 67, 65, 53, 60, 53, 70, 71, 60, 53, 66,\n",
      "        53,  1, 60, 73, 65, 54, 64, 77,  0, 67, 55, 55, 73, 68, 61, 57, 56,  1,\n",
      "        72, 60, 57,  1, 71, 57, 53, 72,  1, 72, 60, 53, 72,  1, 75, 53, 71,  1,\n",
      "        53, 71, 71, 61, 59, 66, 57, 56,  1, 72, 67,  1, 60, 61, 65, 10,  1, 42,\n",
      "        57, 57, 61, 66, 59,  1, 72, 60, 53, 72,  1, 60, 57,  1, 75, 53, 71,  0,\n",
      "        55, 67, 65, 58, 67, 70, 72, 53, 54, 64, 77,  1, 71, 57, 53, 72, 57, 56,\n",
      "         8,  1, 53, 66, 56,  1, 70, 57, 55, 67, 74, 57, 70, 57, 56,  1, 58, 70,\n",
      "        67, 65,  1, 58, 53, 72, 61, 59, 73, 57,  8,  1, 67, 66, 57,  1, 67, 58,\n",
      "         1, 72, 60, 57,  1, 41, 61, 71, 60, 61])\n"
     ]
    }
   ],
   "source": [
    "\n",
    "data = torch.tensor(encode(text), dtype = torch.long)\n",
    "print(data.shape, data.dtype)\n",
    "print(data[:1000])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "LxM5S8NwNhE6"
   },
   "outputs": [],
   "source": [
    "n = int(0.9 * len(data))\n",
    "train_data = data[:n]\n",
    "val_data = data[n:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "TwI3F3_Suldy",
    "outputId": "017f77a0-d82e-4b16-f3a0-9851106155b2"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([24, 27, 32,  1, 39, 24, 41, 45, 24])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "block_size = 8\n",
    "train_data[:block_size+1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Fpt-f8G2Yo0C",
    "outputId": "dee65c1e-03a4-45c1-a71b-d28e98989df9"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "when context is tensor([24]), target is 27\n",
      "when context is tensor([24, 27]), target is 32\n",
      "when context is tensor([24, 27, 32]), target is 1\n",
      "when context is tensor([24, 27, 32,  1]), target is 39\n",
      "when context is tensor([24, 27, 32,  1, 39]), target is 24\n",
      "when context is tensor([24, 27, 32,  1, 39, 24]), target is 41\n",
      "when context is tensor([24, 27, 32,  1, 39, 24, 41]), target is 45\n",
      "when context is tensor([24, 27, 32,  1, 39, 24, 41, 45]), target is 24\n"
     ]
    }
   ],
   "source": [
    "x = train_data[:block_size]\n",
    "y = train_data[1:block_size+1]\n",
    "for t in range(block_size):\n",
    "  context = x[:t+1]\n",
    "  target = y[t]\n",
    "  print(f\"when context is {context}, target is {target}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "IzXQ3Gk4JGrs",
    "outputId": "a1e2e405-b441-4f9c-da87-feede215fe26"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "13436984"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(train_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "O3fDJhSWY9rN",
    "outputId": "a0bfc999-9dd1-4699-e54b-0aa822255ff2"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "inputs:\n",
      "torch.Size([4, 8])\n",
      "tensor([[61, 67, 66, 71,  5,  1, 61, 71],\n",
      "        [10,  1, 32, 72,  1, 75, 53, 71],\n",
      "        [ 0, 65, 53, 77,  1, 54, 57,  1],\n",
      "        [57, 56,  1, 71, 53, 59, 57,  8]])\n",
      "targets:\n",
      "torch.Size([4, 8])\n",
      "tensor([[67, 66, 71,  5,  1, 61, 71,  1],\n",
      "        [ 1, 32, 72,  1, 75, 53, 71,  1],\n",
      "        [65, 53, 77,  1, 54, 57,  1, 68],\n",
      "        [56,  1, 71, 53, 59, 57,  8,  1]])\n",
      "--------\n",
      "When input is [61] the target is 67\n",
      "When input is [61, 67] the target is 66\n",
      "When input is [61, 67, 66] the target is 71\n",
      "When input is [61, 67, 66, 71] the target is 5\n",
      "When input is [61, 67, 66, 71, 5] the target is 1\n",
      "When input is [61, 67, 66, 71, 5, 1] the target is 61\n",
      "When input is [61, 67, 66, 71, 5, 1, 61] the target is 71\n",
      "When input is [61, 67, 66, 71, 5, 1, 61, 71] the target is 1\n",
      "When input is [10] the target is 1\n",
      "When input is [10, 1] the target is 32\n",
      "When input is [10, 1, 32] the target is 72\n",
      "When input is [10, 1, 32, 72] the target is 1\n",
      "When input is [10, 1, 32, 72, 1] the target is 75\n",
      "When input is [10, 1, 32, 72, 1, 75] the target is 53\n",
      "When input is [10, 1, 32, 72, 1, 75, 53] the target is 71\n",
      "When input is [10, 1, 32, 72, 1, 75, 53, 71] the target is 1\n",
      "When input is [0] the target is 65\n",
      "When input is [0, 65] the target is 53\n",
      "When input is [0, 65, 53] the target is 77\n",
      "When input is [0, 65, 53, 77] the target is 1\n",
      "When input is [0, 65, 53, 77, 1] the target is 54\n",
      "When input is [0, 65, 53, 77, 1, 54] the target is 57\n",
      "When input is [0, 65, 53, 77, 1, 54, 57] the target is 1\n",
      "When input is [0, 65, 53, 77, 1, 54, 57, 1] the target is 68\n",
      "When input is [57] the target is 56\n",
      "When input is [57, 56] the target is 1\n",
      "When input is [57, 56, 1] the target is 71\n",
      "When input is [57, 56, 1, 71] the target is 53\n",
      "When input is [57, 56, 1, 71, 53] the target is 59\n",
      "When input is [57, 56, 1, 71, 53, 59] the target is 57\n",
      "When input is [57, 56, 1, 71, 53, 59, 57] the target is 8\n",
      "When input is [57, 56, 1, 71, 53, 59, 57, 8] the target is 1\n"
     ]
    }
   ],
   "source": [
    "batch_size = 4\n",
    "block_size = 8\n",
    "\n",
    "def get_batch(split):\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",
    "  return x, y\n",
    "\n",
    "xb, yb = get_batch('train')\n",
    "\n",
    "print('inputs:')\n",
    "print(xb.shape)\n",
    "print(xb)\n",
    "print('targets:')\n",
    "print(yb.shape)\n",
    "print(yb)\n",
    "\n",
    "print('--------')\n",
    "\n",
    "for b in range(batch_size):\n",
    "  for t in range(block_size):\n",
    "    context = xb[b, :t+1]\n",
    "    target = yb[b, t]\n",
    "    print(f\"When input is {context.tolist()} the target is {target.tolist()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Qzk6gK89Jkf0",
    "outputId": "e14248cf-b135-4601-dae1-583020d64774"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[61, 67, 66, 71,  5,  1, 61, 71],\n",
       "        [10,  1, 32, 72,  1, 75, 53, 71],\n",
       "        [ 0, 65, 53, 77,  1, 54, 57,  1],\n",
       "        [57, 56,  1, 71, 53, 59, 57,  8]])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# input to the transformer\n",
    "xb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "bxEJ2O_1K35E",
    "outputId": "a5e0a7a9-1dac-489d-8e8e-10c8eb55db16"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([32, 79])\n",
      "tensor(5.0030, grad_fn=<NllLossBackward0>)\n",
      "\n",
      "yzI4QuIUJb(CJJSEXPctd](MadfMc7scTAjfDZFVdkO)Ep-_h0OGTE\"0E:lToxacFIXabHP;nqPeoL&5?Xb-rvYw)UJ2B5uCG,qJ\n"
     ]
    }
   ],
   "source": [
    "class BigramModel(nn.Module):\n",
    "\n",
    "  def __init__(self, vocab_size):\n",
    "    super().__init__()\n",
    "    self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)\n",
    "\n",
    "  def forward(self, idx, targets=None):\n",
    "    logits = self.token_embedding_table(idx)\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",
    "    # generate max_new_tokens new indices and concatenate to idx\n",
    "    # idx is (B, T) array of indices. row is number of batches, and column is context length\n",
    "    for _ in range(max_new_tokens):\n",
    "      logits, loss = self(idx)\n",
    "      # we only need the last value in the sequence to generate the next sequence, in this particular model\n",
    "      logits = logits[:, -1, :]\n",
    "      # getting the probabilites from the logits\n",
    "      probs = F.softmax(logits, dim = -1)\n",
    "      # sampling from the distrbution\n",
    "      idx_next = torch.multinomial(probs, num_samples = 1)\n",
    "      idx = torch.cat((idx, idx_next), dim = 1)\n",
    "    return idx\n",
    "\n",
    "\n",
    "m = BigramModel(vocab_size)\n",
    "logits, loss = m(xb, yb)\n",
    "print(logits.shape)\n",
    "print(loss)\n",
    "\n",
    "print(decode(m.generate(idx = torch.zeros((1, 1), dtype=torch.long), max_new_tokens = 100)[0].tolist()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "id": "_N9P7-v8RIf2"
   },
   "outputs": [],
   "source": [
    "optimizer = torch.optim.AdamW(m.parameters(), lr=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "fsFRYakOzq80",
    "outputId": "4eaa9af7-281f-4522-99a7-4fb6ec173c74"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.454437732696533\n"
     ]
    }
   ],
   "source": [
    "batch_size = 32\n",
    "for steps in range(10000):\n",
    "  xb, yb = get_batch('train')\n",
    "  logits, loss = m(xb, yb)\n",
    "  optimizer.zero_grad(set_to_none=True)\n",
    "  loss.backward()\n",
    "  optimizer.step()\n",
    "\n",
    "print(loss.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "1nhyZMMK1wrH",
    "outputId": "7c23e674-ee73-410b-83be-8e545abea6ce"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "sanse onof osquthiounglipedofomy PLed thesosuerje!'stache theain and wesheqlavQMavant!'Bhele, the woffe inesthoadene ofoupherowimiepare\n",
      "th \"Vaiffreren t Thas Thared ededathrim Vare t Pa ay wid. ove a aner An.\n",
      "the Gubear ly iontos icth watong tuved lallenalsprem s by s tolde. Thicay a aritedea ts thed acuiglothaucorete thoinests\n",
      "urtou, anevaigd osucavalea seaknome thancemYRegronthee dedsthee a )[18\n"
     ]
    }
   ],
   "source": [
    "print(decode(m.generate(idx = torch.zeros((1, 1), dtype=torch.long), max_new_tokens = 400)[0].tolist()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "id": "r6MwdHmW2IrW"
   },
   "outputs": [],
   "source": [
    "# ones tensor\n",
    "# tril on top to make it triangular\n",
    "# take weight vector and apply mask on it such that wherever ones tensor is zeros, weight vector is negative infinite\n",
    "# apply softmax on weight vector\n",
    "# result is attention weights. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 8, 16])"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "B, T, C = 4, 8, 32\n",
    "x = torch.randn(B, T, C)\n",
    "\n",
    "head_size = 16\n",
    "key = nn.Linear(C, head_size, bias = False)\n",
    "query = nn.Linear(C, head_size, bias = False)\n",
    "value = nn.Linear(C, head_size, bias = False)\n",
    "\n",
    "k = key(x) # shape: (batch_size, context_size, head_size)\n",
    "q = query(x) # shape: (batch_size, context_size, head_size)\n",
    "\n",
    "wei = q @ k.transpose(-2, -1) # (B, T, 16) @ (B, 16, T) --> (B, T, T)\n",
    "\n",
    "tril = torch.tril(torch.ones(T, T))\n",
    "wei = torch.zeros((T, T))\n",
    "wei = wei.masked_fill(tril == 0, float('-inf'))\n",
    "wei = F.softmax(wei, dim = -1)\n",
    "\n",
    "v = value(x)\n",
    "out = wei @ v\n",
    "\n",
    "out.shape\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "authorship_tag": "ABX9TyOzQ64nu6PGsXcOyeoXO87q",
   "include_colab_link": true,
   "provenance": []
  },
  "kernelspec": {
   "display_name": "saverx",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}