2_bigram_single_layer_nn.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install names-dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from names_dataset import NameDataset\n",
    "import re\n",
    "import matplotlib.pyplot as plt\n",
    "import torch.nn.functional as F\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SEED = 42"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nd = NameDataset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "country_codes = [\"US\", \"CA\", \"GB\", \"IN\"]            # USA, Canada, Great Britain, India"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = set()\n",
    "dataset_size = 50000\n",
    "for code in country_codes:\n",
    "    data = nd.get_top_names(n=dataset_size, country_alpha2=code)\n",
    "    dataset.update(set(data[code]['M'] + data[code]['F']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "43231"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "len(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# clean and make each name 1 word\n",
    "names = list()\n",
    "for name in dataset:\n",
    "    if re.match(r\"^[A-Za-z\\s]+$\", name):            # only consider names with english alphabets\n",
    "        for i in name.split():\n",
    "            if len(i)>1: names.append(i.lower())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Basic Neural Network model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'orla'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "names[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chars = sorted(list(set(''.join(names))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ctoi = {k:idx+1 for idx, k in enumerate(chars)}\n",
    "ctoi['.'] = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "itoc = {v:k for k,v in ctoi.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ".orla.\n",
      ". -----> o\n",
      "o -----> r\n",
      "r -----> l\n",
      "l -----> a\n",
      "a -----> .\n",
      "---------------\n",
      ".aboubacar.\n",
      ". -----> a\n",
      "a -----> b\n",
      "b -----> o\n",
      "o -----> u\n",
      "u -----> b\n",
      "b -----> a\n",
      "a -----> c\n",
      "c -----> a\n",
      "a -----> r\n",
      "r -----> .\n",
      "---------------\n"
     ]
    }
   ],
   "source": [
    "# xs are all the inputs to the nn, ys are the corresponding outputs\n",
    "xs,ys = [], []\n",
    "for name in names[:2]:\n",
    "    name = ['.'] + list(name) + ['.']\n",
    "    print(''.join(name))\n",
    "    for c1, c2 in zip(name, name[1:]):\n",
    "        xs.append(c1)\n",
    "        ys.append(c2)\n",
    "        print(f\"{c1} -----> {c2}\")\n",
    "    print(\"---------------\")\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xs,ys = [], []\n",
    "for name in names:\n",
    "    name = ['.'] + list(name) + ['.']\n",
    "    for c1, c2 in zip(name, name[1:]):\n",
    "        xs.append(ctoi[c1])\n",
    "        ys.append(ctoi[c2])\n",
    "\n",
    "xs = torch.tensor(xs)                   # torch.tensor - autodetects the dtype, torch.Tensor defaults to float\n",
    "ys = torch.tensor(ys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 0, 15, 18, 12,  1])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "xs[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([15, 18, 12,  1,  0])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ys[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## can't feed integer index to the nn, so create a vector -> common way is to create a onehot vector\n",
    "xenc = F.one_hot(xs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([306991, 27])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "xenc.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5, 27])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "xenc[:5].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x3294b0100>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAhYAAACHCAYAAABK4hAcAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAN7UlEQVR4nO3df2iV5cPH8c/Z2o4/Ojs6136cNufUUmpukrolkgkbTgvJ9A8r/1hDjOoozlHJAl1CsDAIqSQj6Os//krIJPliyHKTYP5gIibUHh3yeGRuS3k805lz7VzPH307z3NSt8527dw7x/cLbtjuc3HfHy4u2Gf3uc+5XcYYIwAAAAuSnA4AAAASB8UCAABYQ7EAAADWUCwAAIA1FAsAAGANxQIAAFhDsQAAANY8EsuThUIhtbe3y+PxyOVyxfLUAABgiIwxunnzpnw+n5KSBr4mEdNi0d7erry8vFieEgAAWBIIBJSbmzvgmJgWC4/HI0n67zNTlPbo8N6FefnJWTYiAQCAQfyhPv2kf4f/jg8kpsXir7c/0h5NUppneMXiEVeKjUgAAGAw/3n4xz+5jYGbNwEAgDUUCwAAYA3FAgAAWDOkYrFjxw5NmTJFY8aMUWlpqU6dOmU7FwAAiENRF4v9+/erpqZGdXV1OnPmjIqLi1VRUaGurq6RyAcAAOJI1MXik08+0dq1a1VVVaWnnnpKO3fu1Lhx4/T111+PRD4AABBHoioWd+/eVUtLi8rLy//vAElJKi8vV3Nz8z3je3t71d3dHbEBAIDEFVWxuHbtmvr7+5WVlRWxPysrSx0dHfeMr6+vl9frDW986yYAAIltRD8VUltbq2AwGN4CgcBIng4AADgsqm/ezMjIUHJysjo7OyP2d3Z2Kjs7+57xbrdbbrd7eAkBAEDciOqKRWpqqubMmaOGhobwvlAopIaGBs2fP996OAAAEF+iflZITU2NKisrNXfuXJWUlGj79u3q6elRVVXVSOQDAABxJOpisWrVKv3222/asmWLOjo6NHv2bB05cuSeGzoBAMDDx2WMMbE6WXd3t7xer/7nv6YO++mmFb7ZdkIBAIAB/WH61KhDCgaDSktLG3AszwoBAADWRP1WiA0vPzlLj7hSnDg1gATxQ/tZa8fiCihgD1csAACANRQLAABgDcUCAABYQ7EAAADWUCwAAIA1FAsAAGANxQIAAFhDsQAAANZQLAAAgDUUCwAAYA3FAgAAWEOxAAAA1lAsAACANRQLAABgDcUCAABYQ7EAAADWUCwAAIA1FAsAAGDNI04HAIChqPDNdjrCiPqh/ayV4yT6PGH04YoFAACwhmIBAACsoVgAAABrKBYAAMCaqIpFfX295s2bJ4/Ho8zMTC1fvlytra0jlQ0AAMSZqIpFU1OT/H6/Tpw4oaNHj6qvr0+LFy9WT0/PSOUDAABxJKqPmx45ciTi9127dikzM1MtLS1auHCh1WAAACD+DOt7LILBoCQpPT39vq/39vaqt7c3/Ht3d/dwTgcAAEa5Id+8GQqFVF1drQULFqiwsPC+Y+rr6+X1esNbXl7ekIMCAIDRb8jFwu/36/z589q3b98Dx9TW1ioYDIa3QCAw1NMBAIA4MKS3QtatW6fDhw/r+PHjys3NfeA4t9stt9s95HAAACC+RFUsjDFav369Dh48qMbGRhUUFIxULgAAEIeiKhZ+v1979uzRoUOH5PF41NHRIUnyer0aO3bsiAQEAADxI6p7LL744gsFg0EtWrRIOTk54W3//v0jlQ8AAMSRqN8KAQAAeBCeFQIAAKyhWAAAAGuG9c2bAEa3H9rPWjtWhW+2tWNhcMw34hVXLAAAgDUUCwAAYA3FAgAAWEOxAAAA1lAsAACANRQLAABgDcUCAABYQ7EAAADWUCwAAIA1FAsAAGANxQIAAFhDsQAAANZQLAAAgDUUCwAAYA3FAgAAWEOxAAAA1lAsAACANRQLAABgzSNOBxiqH9rPWjtWhW+2tWMBowlrG0CsccUCAABYQ7EAAADWUCwAAIA1FAsAAGDNsIrFRx99JJfLperqaktxAABAPBtysTh9+rS+/PJLFRUV2cwDAADi2JCKxa1bt7R69Wp99dVXmjhxou1MAAAgTg2pWPj9fr344osqLy8fcFxvb6+6u7sjNgAAkLii/oKsffv26cyZMzp9+vSgY+vr67V169YhBQMAAPEnqisWgUBAGzZs0O7duzVmzJhBx9fW1ioYDIa3QCAw5KAAAGD0i+qKRUtLi7q6uvTMM8+E9/X39+v48eP6/PPP1dvbq+Tk5PBrbrdbbrfbXloAADCqRVUsysrK9PPPP0fsq6qq0syZM7Vp06aIUgEAAB4+URULj8ejwsLCiH3jx4/XpEmT7tkPAAAePnzzJgAAsGbYj01vbGy0EAMAACQCrlgAAABrhn3FIhrGGEnSH+qTzPCO1X0zZCHRn/4wfdaOBQBAovlDf/6d/Ovv+EBc5p+MsuTKlSvKy8uL1ekAAIBFgUBAubm5A46JabEIhUJqb2+Xx+ORy+V64Lju7m7l5eUpEAgoLS0tVvEeWsx37DDXscV8xxbzHVuxnG9jjG7evCmfz6ekpIHvoojpWyFJSUmDNp3/Ly0tjcUZQ8x37DDXscV8xxbzHVuxmm+v1/uPxnHzJgAAsIZiAQAArBmVxcLtdquuro7njMQI8x07zHVsMd+xxXzH1mid75jevAkAABLbqLxiAQAA4hPFAgAAWEOxAAAA1lAsAACANRQLAABgzagrFjt27NCUKVM0ZswYlZaW6tSpU05HSkgffPCBXC5XxDZz5kynYyWM48ePa9myZfL5fHK5XPruu+8iXjfGaMuWLcrJydHYsWNVXl6uCxcuOBM2AQw236+//vo9633JkiXOhI1z9fX1mjdvnjwejzIzM7V8+XK1trZGjLlz5478fr8mTZqkRx99VCtXrlRnZ6dDiePbP5nvRYsW3bO+33zzTYcSj7JisX//ftXU1Kiurk5nzpxRcXGxKioq1NXV5XS0hPT000/r6tWr4e2nn35yOlLC6OnpUXFxsXbs2HHf17dt26ZPP/1UO3fu1MmTJzV+/HhVVFTozp07MU6aGAabb0lasmRJxHrfu3dvDBMmjqamJvn9fp04cUJHjx5VX1+fFi9erJ6envCYjRs36vvvv9eBAwfU1NSk9vZ2rVixwsHU8eufzLckrV27NmJ9b9u2zaHEkswoUlJSYvx+f/j3/v5+4/P5TH19vYOpElNdXZ0pLi52OsZDQZI5ePBg+PdQKGSys7PNxx9/HN5348YN43a7zd69ex1ImFj+Pt/GGFNZWWleeuklR/Ikuq6uLiPJNDU1GWP+XMspKSnmwIED4TG//PKLkWSam5udipkw/j7fxhjz/PPPmw0bNjgX6m9GzRWLu3fvqqWlReXl5eF9SUlJKi8vV3Nzs4PJEteFCxfk8/k0depUrV69WpcvX3Y60kPh0qVL6ujoiFjrXq9XpaWlrPUR1NjYqMzMTM2YMUNvvfWWrl+/7nSkhBAMBiVJ6enpkqSWlhb19fVFrO+ZM2dq8uTJrG8L/j7ff9m9e7cyMjJUWFio2tpa3b5924l4kmL8dNOBXLt2Tf39/crKyorYn5WVpV9//dWhVImrtLRUu3bt0owZM3T16lVt3bpVzz33nM6fPy+Px+N0vITW0dEhSfdd63+9BruWLFmiFStWqKCgQG1tbXr//fe1dOlSNTc3Kzk52el4cSsUCqm6uloLFixQYWGhpD/Xd2pqqiZMmBAxlvU9fPebb0l67bXXlJ+fL5/Pp3PnzmnTpk1qbW3Vt99+60jOUVMsEFtLly4N/1xUVKTS0lLl5+frm2++0Zo1axxMBtj3yiuvhH+eNWuWioqKNG3aNDU2NqqsrMzBZPHN7/fr/Pnz3J8VIw+a7zfeeCP886xZs5STk6OysjK1tbVp2rRpsY45em7ezMjIUHJy8j13Dnd2dio7O9uhVA+PCRMm6Mknn9TFixedjpLw/lrPrHXnTJ06VRkZGaz3YVi3bp0OHz6sY8eOKTc3N7w/Oztbd+/e1Y0bNyLGs76H50HzfT+lpaWS5Nj6HjXFIjU1VXPmzFFDQ0N4XygUUkNDg+bPn+9gsofDrVu31NbWppycHKejJLyCggJlZ2dHrPXu7m6dPHmStR4jV65c0fXr11nvQ2CM0bp163Tw4EH9+OOPKigoiHh9zpw5SklJiVjfra2tunz5Mut7CAab7/s5e/asJDm2vkfVWyE1NTWqrKzU3LlzVVJSou3bt6unp0dVVVVOR0s477zzjpYtW6b8/Hy1t7errq5OycnJevXVV52OlhBu3boV8d/CpUuXdPbsWaWnp2vy5Mmqrq7Whx9+qCeeeEIFBQXavHmzfD6fli9f7lzoODbQfKenp2vr1q1auXKlsrOz1dbWpvfee0/Tp09XRUWFg6njk9/v1549e3To0CF5PJ7wfRNer1djx46V1+vVmjVrVFNTo/T0dKWlpWn9+vWaP3++nn32WYfTx5/B5rutrU179uzRCy+8oEmTJuncuXPauHGjFi5cqKKiImdCO/2xlL/77LPPzOTJk01qaqopKSkxJ06ccDpSQlq1apXJyckxqamp5vHHHzerVq0yFy9edDpWwjh27JiRdM9WWVlpjPnzI6ebN282WVlZxu12m7KyMtPa2ups6Dg20Hzfvn3bLF682Dz22GMmJSXF5Ofnm7Vr15qOjg6nY8el+82zJPOvf/0rPOb33383b7/9tpk4caIZN26cefnll83Vq1edCx3HBpvvy5cvm4ULF5r09HTjdrvN9OnTzbvvvmuCwaBjmV3/CQ4AADBso+YeCwAAEP8oFgAAwBqKBQAAsIZiAQAArKFYAAAAaygWAADAGooFAACwhmIBAACsoVgAAABrKBYAAMAaigUAALDmfwEtsJFDJyE2FQAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.imshow(xenc[:5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.int64"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "xenc.dtype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.float32"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "## can't input integers to nn, so convert to float\n",
    "xenc = xenc.float()\n",
    "xenc.dtype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-1.0689],\n",
       "        [-0.2569],\n",
       "        [-0.3249],\n",
       "        [ 0.5182],\n",
       "        [-0.5608],\n",
       "        [ 1.2037],\n",
       "        [ 0.1602],\n",
       "        [-0.0324],\n",
       "        [ 0.1895],\n",
       "        [ 0.7312],\n",
       "        [ 0.2050],\n",
       "        [-0.5482],\n",
       "        [ 0.0922],\n",
       "        [ 0.8707],\n",
       "        [ 1.9514],\n",
       "        [-2.3171],\n",
       "        [ 0.4049],\n",
       "        [ 1.3411],\n",
       "        [-0.4866],\n",
       "        [ 0.7615],\n",
       "        [-0.3826],\n",
       "        [ 0.2325],\n",
       "        [-0.5492],\n",
       "        [-0.0530],\n",
       "        [-0.0723],\n",
       "        [ 0.4620],\n",
       "        [ 0.1041]])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "## let's create a simple single neuron with weights - W\n",
    "W = torch.randn((27,1))         # each input is (1,27) size, so W is (27,1) so that we can do 'Wx' and spit out the product corresponding to the input\n",
    "W"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-1.0689]])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# ex\n",
    "xenc[:1] @ W                # '@' matrix multiplication by torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-1.0689],\n",
       "        [-2.3171],\n",
       "        [-0.4866],\n",
       "        [ 0.0922],\n",
       "        [-0.2569]])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "## similarly for a batch of inputs\n",
    "xenc[:5] @ W          # 5,1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Let's create 27 neurons - so that we can output 27 dim vector (which can be given to softmax) to highlight the one out of 27 chars\n",
    "W = torch.randn((27,27), requires_grad=True)        # explicitly set requires_grad so that we can do backprop by calculating grads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-0.4744,  0.0674,  1.0105,  0.6498, -0.3584, -1.5673, -1.8801,  0.5003,\n",
       "        -0.7110, -0.9997,  0.3784,  1.1491,  0.9539, -0.0081, -0.0990, -1.7772,\n",
       "         0.5131, -0.3394, -1.6720,  0.2915,  0.2621, -0.0053, -0.6830, -0.2243,\n",
       "        -1.1239, -0.8245,  0.9831], grad_fn=<SqueezeBackward4>)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "xenc[0] @ W         # should give 27 dim vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5, 27])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# similarly batch\n",
    "(xenc[:5] @ W).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we want probability of each char given a prev char\n",
    "# currently, nn gives some -ve and +ve numbers as outputs\n",
    "# we had counts in the prev approach which are converted to probabilities\n",
    "# let's think that the nn outputs the log counts(logits), now this can be converted to probabilty\n",
    "# exp(inverse of log) the nn output, so that we'll get the counts, then convert the counts to probability by dividing with the sum\n",
    "# note that all these are differentiable operations, so that we can backprop to optimize the weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([306991, 27])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "logits = xenc @ W               # log-counts\n",
    "logits.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.4744,  0.0674,  1.0105,  0.6498, -0.3584, -1.5673, -1.8801,  0.5003,\n",
       "         -0.7110, -0.9997,  0.3784,  1.1491,  0.9539, -0.0081, -0.0990, -1.7772,\n",
       "          0.5131, -0.3394, -1.6720,  0.2915,  0.2621, -0.0053, -0.6830, -0.2243,\n",
       "         -1.1239, -0.8245,  0.9831],\n",
       "        [ 0.6808, -0.2723, -0.4479,  1.5656,  1.4405,  0.1476,  0.6211,  0.7006,\n",
       "          0.6008,  0.7934,  0.8319, -0.9536, -0.2047,  1.2169, -0.7423, -2.2152,\n",
       "         -0.6619,  0.2168, -0.0423,  2.4974,  0.0527,  1.2078,  2.2832, -0.6594,\n",
       "          0.6610,  0.2582, -0.5172]], grad_fn=<SliceBackward0>)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "logits[:2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.6222,  1.0697,  2.7470,  1.9152,  0.6988,  0.2086,  0.1526,  1.6493,\n",
       "          0.4911,  0.3680,  1.4600,  3.1553,  2.5957,  0.9920,  0.9057,  0.1691,\n",
       "          1.6704,  0.7122,  0.1879,  1.3385,  1.2996,  0.9947,  0.5051,  0.7991,\n",
       "          0.3250,  0.4385,  2.6727],\n",
       "        [ 1.9755,  0.7616,  0.6390,  4.7857,  4.2228,  1.1591,  1.8610,  2.0149,\n",
       "          1.8236,  2.2109,  2.2977,  0.3853,  0.8149,  3.3765,  0.4760,  0.1091,\n",
       "          0.5159,  1.2421,  0.9586, 12.1505,  1.0541,  3.3461,  9.8080,  0.5171,\n",
       "          1.9367,  1.2946,  0.5962]], grad_fn=<SliceBackward0>)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "counts = logits.exp()               # from 0 to inf\n",
    "counts[:2]                          # counts is equivalent to 'N' matrix in the prev approach"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.0206, 0.0355, 0.0911, 0.0635, 0.0232, 0.0069, 0.0051, 0.0547, 0.0163,\n",
       "         0.0122, 0.0484, 0.1047, 0.0861, 0.0329, 0.0300, 0.0056, 0.0554, 0.0236,\n",
       "         0.0062, 0.0444, 0.0431, 0.0330, 0.0168, 0.0265, 0.0108, 0.0145, 0.0887],\n",
       "        [0.0317, 0.0122, 0.0103, 0.0768, 0.0677, 0.0186, 0.0299, 0.0323, 0.0293,\n",
       "         0.0355, 0.0369, 0.0062, 0.0131, 0.0542, 0.0076, 0.0018, 0.0083, 0.0199,\n",
       "         0.0154, 0.1949, 0.0169, 0.0537, 0.1573, 0.0083, 0.0311, 0.0208, 0.0096]],\n",
       "       grad_fn=<SliceBackward0>)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "probs = counts/counts.sum(1, keepdim=True)\n",
    "probs[:2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above 2 cells correspond to the 'softmax' operation - converting logits into probabilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([27])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "probs[0].shape          # probability over all the 27 chars given the 0th training example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(1.0000, grad_fn=<SumBackward0>)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "probs[0].sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculating loss - for each prediction, take the probability that model assigns to the correct 'y'[likelihood]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 0, 15, 18, 12,  1])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "xs[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([15, 18, 12,  1,  0])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ys[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.0355, grad_fn=<SelectBackward0>)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "probs[0,1]  # probability of the correct index(y-1) that the model predicted for the 1st(0) training sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.0154, grad_fn=<SelectBackward0>)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "probs[1,18] # probability of the correct index(y-18) that the model predicted for the 2nd(1) training sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(0.0355, grad_fn=<SelectBackward0>),\n",
       " tensor(0.0154, grad_fn=<SelectBackward0>),\n",
       " tensor(0.0290, grad_fn=<SelectBackward0>),\n",
       " tensor(0.0032, grad_fn=<SelectBackward0>),\n",
       " tensor(0.0317, grad_fn=<SelectBackward0>))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "probs[0,1], probs[1,18], probs[2,23], probs[3,5], probs[4,14]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.0056, 0.0154, 0.0252, 0.0094, 0.0161], grad_fn=<IndexBackward0>)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# efficient way of retrieving the above\n",
    "range_ = torch.arange(5)    # ideally all the predictions\n",
    "probs[range_, ys[:5]]       # likelihood"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-5.1832, -4.1748, -3.6814, -4.6659, -4.1303], grad_fn=<LogBackward0>)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "probs[range_, ys[:5]].log() # log likelihood"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(-4.3671, grad_fn=<MeanBackward0>)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "probs[range_, ys[:5]].log().mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4.367127418518066"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "-probs[range_, ys[:5]].log().mean().item()     # avg -ve log likelihood ----> LOSS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss = -probs[range_, ys[:5]].log().mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(4.3671, grad_fn=<NegBackward0>)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Backprop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. set all gradients to zero"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "W.grad = None\n",
    "loss.backward()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.0041,  0.0071,  0.0182,  0.0127,  0.0046,  0.0014,  0.0010,  0.0109,\n",
       "          0.0033,  0.0024,  0.0097,  0.0209,  0.0172,  0.0066,  0.0060, -0.1989,\n",
       "          0.0111,  0.0047,  0.0012,  0.0089,  0.0086,  0.0066,  0.0034,  0.0053,\n",
       "          0.0022,  0.0029,  0.0177],\n",
       "        [-0.1968,  0.0587,  0.0030,  0.0129,  0.0025,  0.0061,  0.0018,  0.0004,\n",
       "          0.0035,  0.0054,  0.0092,  0.0122,  0.0060,  0.0017,  0.0063,  0.0002,\n",
       "          0.0011,  0.0058,  0.0062,  0.0027,  0.0059,  0.0072,  0.0024,  0.0011,\n",
       "          0.0233,  0.0052,  0.0058],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0040, -0.1981,  0.0112,  0.0291,  0.0019,  0.0006,  0.0028,  0.0016,\n",
       "          0.0021,  0.0265,  0.0026,  0.0023,  0.0096,  0.0048,  0.0071,  0.0037,\n",
       "          0.0033,  0.0070,  0.0181,  0.0026,  0.0182,  0.0012,  0.0079,  0.0105,\n",
       "          0.0016,  0.0050,  0.0128],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0063,  0.0024,  0.0021,  0.0154,  0.0135,  0.0037,  0.0060,  0.0065,\n",
       "          0.0059,  0.0071,  0.0074,  0.0012,  0.0026,  0.0108,  0.0015,  0.0004,\n",
       "          0.0017,  0.0040, -0.1969,  0.0390,  0.0034,  0.0107,  0.0315,  0.0017,\n",
       "          0.0062,  0.0042,  0.0019],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0034,  0.0052,  0.0061,  0.0033,  0.0059,  0.0147,  0.0073,  0.0480,\n",
       "          0.0210,  0.0014,  0.0007,  0.0019, -0.1950,  0.0042,  0.0016,  0.0148,\n",
       "          0.0011,  0.0120,  0.0084,  0.0038,  0.0068,  0.0040,  0.0032,  0.0058,\n",
       "          0.0033,  0.0020,  0.0051],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000]])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "W.grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([27, 27])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "W.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([27, 27])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "W.grad.shape        # each value corresponds to the +ve or -ve effect on loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.0041)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "W.grad[0,0]     # this has +ve effect of nudging loss to +ve (increase) if we add a +ve constant to this gradient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# update the weights(optimize) with the new grads\n",
    "W.data += -0.1 * W.grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logits = xenc @ W \n",
    "counts = logits.exp() \n",
    "probs = counts/counts.sum(1, keepdim=True)\n",
    "loss = -probs[range_, ys[:5]].log().mean()      # considering only first 5 samples to speed it up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(4.3460, grad_fn=<NegBackward0>)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "loss        # reduces after updation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "W.grad = None\n",
    "loss.backward()\n",
    "W.data += -0.1 * W.grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logits = xenc @ W \n",
    "counts = logits.exp() \n",
    "probs = counts/counts.sum(1, keepdim=True)\n",
    "loss = -probs[range_, ys[:5]].log().mean() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(4.3248, grad_fn=<NegBackward0>)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "loss    # reduces even further, now we're basically doing gradient descent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### lets run on entire training set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([306991]), torch.Size([306991]))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "xs.shape, ys.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total number of training examples - 306991\n"
     ]
    }
   ],
   "source": [
    "num_samples = xs.nelement()\n",
    "print(f\"Total number of training examples - {num_samples}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize the network\n",
    "g = torch.Generator().manual_seed(SEED)\n",
    "W = torch.randn((27,27), generator=g, requires_grad=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(num_epochs:int):\n",
    "    for _ in range(num_epochs):\n",
    "        # forward pass\n",
    "        logits = xenc @ W \n",
    "        counts = logits.exp() \n",
    "        probs = counts/counts.sum(1, keepdim=True)\n",
    "        loss = -probs[torch.arange(num_samples), ys].log().mean() \n",
    "        print(loss.item())\n",
    "        \n",
    "        # backward pass\n",
    "        W.grad = None\n",
    "        loss.backward()\n",
    "        \n",
    "        # update\n",
    "        W.data += -50 * W.grad      # took large laerning rate to have bigger steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.4069812297821045\n",
      "3.2142512798309326\n",
      "3.074751377105713\n",
      "2.9743690490722656\n",
      "2.901893377304077\n",
      "2.8478457927703857\n",
      "2.806061029434204\n",
      "2.7728171348571777\n",
      "2.7457494735717773\n",
      "2.723296642303467\n",
      "2.7043936252593994\n",
      "2.6882808208465576\n",
      "2.674398183822632\n",
      "2.6623215675354004\n",
      "2.651726245880127\n",
      "2.6423592567443848\n",
      "2.634021282196045\n",
      "2.6265530586242676\n",
      "2.619826078414917\n",
      "2.6137349605560303\n",
      "2.6081936359405518\n",
      "2.603130578994751\n",
      "2.5984861850738525\n",
      "2.5942108631134033\n",
      "2.5902628898620605\n",
      "2.5866057872772217\n",
      "2.583209991455078\n",
      "2.5800492763519287\n",
      "2.577101230621338\n",
      "2.574345588684082\n",
      "2.571765422821045\n",
      "2.569345474243164\n",
      "2.5670721530914307\n",
      "2.5649330615997314\n",
      "2.562917470932007\n",
      "2.5610156059265137\n",
      "2.559218645095825\n",
      "2.557518482208252\n",
      "2.555907726287842\n",
      "2.5543806552886963\n",
      "2.552929639816284\n",
      "2.55155086517334\n",
      "2.550238609313965\n",
      "2.5489883422851562\n",
      "2.5477964878082275\n",
      "2.5466582775115967\n",
      "2.5455715656280518\n",
      "2.544532060623169\n",
      "2.5435373783111572\n",
      "2.5425844192504883\n",
      "2.5416715145111084\n",
      "2.5407958030700684\n",
      "2.5399551391601562\n",
      "2.5391478538513184\n",
      "2.5383718013763428\n",
      "2.537625789642334\n",
      "2.536907911300659\n",
      "2.536216974258423\n",
      "2.5355513095855713\n",
      "2.534909725189209\n",
      "2.5342915058135986\n",
      "2.5336947441101074\n",
      "2.533118963241577\n",
      "2.5325632095336914\n",
      "2.5320262908935547\n",
      "2.5315072536468506\n",
      "2.531005859375\n",
      "2.530520439147949\n",
      "2.5300514698028564\n",
      "2.529597043991089\n",
      "2.5291574001312256\n",
      "2.528731107711792\n",
      "2.5283186435699463\n",
      "2.5279183387756348\n",
      "2.5275304317474365\n",
      "2.5271542072296143\n",
      "2.5267891883850098\n",
      "2.526435136795044\n",
      "2.5260908603668213\n",
      "2.525757312774658\n",
      "2.52543306350708\n",
      "2.5251173973083496\n",
      "2.524811029434204\n",
      "2.5245132446289062\n",
      "2.5242233276367188\n",
      "2.5239417552948\n",
      "2.523667335510254\n",
      "2.52340030670166\n",
      "2.5231406688690186\n",
      "2.5228874683380127\n",
      "2.522641181945801\n",
      "2.5224006175994873\n",
      "2.5221669673919678\n",
      "2.5219383239746094\n",
      "2.521716356277466\n",
      "2.5214993953704834\n",
      "2.521287679672241\n",
      "2.52108097076416\n",
      "2.5208792686462402\n",
      "2.5206828117370605\n"
     ]
    }
   ],
   "source": [
    "train(100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Matched the performance of the previous count model but with gradient based optimization\n",
    "\n",
    "advantage this time - we can take more advanced NN, and take more context length and repeat the same process to improve the performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# prediction\n",
    "g = torch.Generator().manual_seed(SEED)\n",
    "n = 5 # predict 5 words\n",
    "for _ in range(n):\n",
    "    out = []\n",
    "    ix = 0          # start with '.'\n",
    "    out.append(itoc[ix])\n",
    "    while True:\n",
    "        xenc = F.one_hot(torch.tensor([ix]), num_classes=27).float()\n",
    "        logits = xenc @ W \n",
    "        counts = logits.exp() \n",
    "        probs = counts/counts.sum(1, keepdim=True)\n",
    "        \n",
    "        ix = torch.multinomial(probs, generator=g, replacement=True, num_samples=1).item()\n",
    "        out.append(itoc[ix])\n",
    "        if ix==0:\n",
    "            break\n",
    "    print(''.join(out))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(W, \"single_layer.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_w = torch.load(\"single_layer.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# prediction\n",
    "g = torch.Generator().manual_seed(SEED)\n",
    "n = 5 # predict 5 words\n",
    "for _ in range(n):\n",
    "    out = []\n",
    "    ix = 0          # start with '.'\n",
    "    out.append(itoc[ix])\n",
    "    out.append(itoc[ix])\n",
    "    while True:\n",
    "        xenc = F.one_hot(torch.tensor([ix]), num_classes=27).float()\n",
    "        logits = xenc @ test_w \n",
    "        counts = logits.exp() \n",
    "        probs = counts/counts.sum(1, keepdim=True)\n",
    "        \n",
    "        ix = torch.multinomial(probs, generator=g, replacement=True, num_samples=1).item()\n",
    "        out.append(itoc[ix])\n",
    "        if ix==0:\n",
    "            break\n",
    "    print(''.join(out))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}