Exercises/makemore exercises/part1_bigrams._exercises.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "first 10 words['emma', 'olivia', 'ava', 'isabella', 'sophia', 'charlotte', 'mia', 'amelia', 'harper', 'evelyn']\n",
      "length of words: 32033\n",
      "min word length 2 and max word length 15\n"
     ]
    }
   ],
   "source": [
    "# reading the data\n",
    "words = open(\"../../data/names.txt\", \"r\").read().splitlines()\n",
    "\n",
    "# Exploring\n",
    "print(f\"first 10 words{words[:10]}\")\n",
    "print(f\"length of words: {len(words)}\")\n",
    "print(f\"min word length {min(len(w) for (w) in words)} and max word length {max(len(w) for (w) in words)}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# E01: Train a Trigram Language Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make a list of characters (a -> z)\n",
    "chars = sorted(list(set(\"\".join(words))))\n",
    "chars = [\".\"] + chars\n",
    "\n",
    "# # make a dictionary of character to index\n",
    "stoi = {ch: i for (i, ch) in enumerate(chars)}\n",
    "\n",
    "# # make a dictionary of index to character\n",
    "itos = {i: ch for (ch, i) in stoi.items()}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1- Using Counting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = torch.ones(27, 27, 27, dtype = torch.int32, device = device)\n",
    "N[0, 0, 0] = 0\n",
    "# getting the Bigrams\n",
    "for w in words:\n",
    "    # add start and end tokens\n",
    "    chs = [\".\"] + list(w) + [\".\"]\n",
    "    for ch1, ch2, ch3 in zip(chs, chs[1:], chs[2:]):\n",
    "        ix1 = stoi[ch1]\n",
    "        ix2 = stoi[ch2]\n",
    "        ix3 = stoi[ch3]\n",
    "\n",
    "        N[ix1, ix2, ix3] += 1\n",
    "\n",
    "P = N / N.sum(dim = 2, keepdim = True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loss Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def count_loss(input_list, verbose = False):\n",
    "    log_likelihood = 0.0\n",
    "    n = 0\n",
    "    for w in input_list:\n",
    "        chs = [\".\"] + list(w) + [\".\"]\n",
    "        for ch1, ch2, ch3 in zip(chs, chs[1:], chs[2:]):\n",
    "            ix1 = stoi[ch1]\n",
    "            ix2 = stoi[ch2]\n",
    "            ix3 = stoi[ch3]\n",
    "\n",
    "            prob = P[ix1, ix2, ix3]\n",
    "            logprob = torch.log(prob)\n",
    "            log_likelihood += logprob\n",
    "            n += 1\n",
    "            \n",
    "            if verbose:\n",
    "                print(f\"{ch1}{ch2} -> {prob:.4f} {logprob:.4f}\")\n",
    "\n",
    "    # higher the log likelihood (closer to 0) is better\n",
    "    print(f\"log Likelihood: {log_likelihood}\")\n",
    "\n",
    "    # but in loss function lower is better, so we negate it\n",
    "    nll = -log_likelihood\n",
    "    print(f\"Negative log likelihood: {nll}\")\n",
    "\n",
    "    # normalize it\n",
    "    print(f\"Normalized Negative log Likelihood: {(nll / n)}\") # we need to minimize this"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Loss\n",
      "log Likelihood: -410414.96875\n",
      "Negative log likelihood: 410414.96875\n",
      "Normalized Negative log Likelihood: 2.092747449874878\n"
     ]
    }
   ],
   "source": [
    "print(\"Training Loss\")\n",
    "count_loss(words)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['tori', 'oserdaloreta', 'blailyn', 'rajr', 'gruug', 'iva', 'hun', 'wyn', 'tyclaznpcnko', 'gren']\n",
      "Sampled words Loss\n",
      "log Likelihood: -139.654052734375\n",
      "Negative log likelihood: 139.654052734375\n",
      "Normalized Negative log Likelihood: 2.450071096420288\n"
     ]
    }
   ],
   "source": [
    "# Sampling\n",
    "names = []\n",
    "for i in range(10):\n",
    "    out = []\n",
    "    ix1, ix2 = 0, 0\n",
    "    while True:\n",
    "        p = P[ix1, ix2]\n",
    "        ix1 = ix2\n",
    "        ix2 = torch.multinomial(p, 1, replacement=True).item()\n",
    "        if ix2 == 0:\n",
    "            break\n",
    "        out.append(itos[ix2])\n",
    "\n",
    "    names.append(\"\".join(out))\n",
    "    \n",
    "print(names)\n",
    "print(\"Sampled words Loss\")\n",
    "count_loss(names)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2- Using MLP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare the training set\n",
    "xs , ys = [], []\n",
    "\n",
    "for w in words:\n",
    "    # add start and end tokens\n",
    "    chs = [\".\"] + list(w) + [\".\"]\n",
    "    for ch1, ch2, ch3 in zip(chs, chs[1:], chs[2:]):\n",
    "        ix1 = stoi[ch1]\n",
    "        ix2 = stoi[ch2]\n",
    "        ix3 = stoi[ch3]\n",
    "\n",
    "        xs.append([ix1, ix2])\n",
    "        ys.append(ix3)\n",
    "\n",
    "xs = torch.tensor(xs, dtype=torch.int64)\n",
    "ys = torch.tensor(ys, dtype=torch.int64)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0: 4.6715\n",
      "10: 2.6098\n",
      "20: 2.4700\n",
      "30: 2.4215\n",
      "40: 2.3986\n",
      "50: 2.3859\n",
      "60: 2.3784\n",
      "70: 2.3736\n",
      "80: 2.3705\n",
      "90: 2.3684\n",
      "100: 2.3670\n",
      "110: 2.3660\n",
      "120: 2.3654\n",
      "130: 2.3649\n",
      "140: 2.3646\n",
      "150: 2.3643\n",
      "160: 2.3642\n",
      "170: 2.3640\n",
      "180: 2.3640\n",
      "190: 2.3639\n"
     ]
    }
   ],
   "source": [
    "W = torch.randn((27*2,27), requires_grad = True, device = device)\n",
    "for k in range(200):\n",
    "    # forward pass\n",
    "    xenc = F.one_hot(xs, num_classes = 27).float().to(device)\n",
    "    xenc = xenc.view(-1, 27*2)\n",
    "    \n",
    "    # probs is softmax\n",
    "    logits = xenc @ W\n",
    "    counts = torch.exp(logits)\n",
    "    probs = counts / counts.sum(dim = 1, keepdim = True)\n",
    "    \n",
    "    # loss (normalized negative log likelihood)\n",
    "    loss = - probs[torch.arange(len(xs)), ys].log().mean()\n",
    "    # add regularization\n",
    "    loss += 0.2 * W.pow(2).mean()\n",
    "\n",
    "    if k % 10 == 0:\n",
    "        print(f\"{k}: {loss.item():.4f}\")\n",
    "\n",
    "    # backward pass\n",
    "    W.grad = None\n",
    "    loss.backward()\n",
    "\n",
    "    # update weights\n",
    "    with torch.no_grad():\n",
    "        W -= 50 * W.grad"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sample from the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "us.\n",
      "lelin.\n",
      "ra.\n",
      "viia.\n",
      "benon.\n",
      "lirisin.\n",
      "ejkylo.\n",
      "elina.\n",
      "ounnre.\n",
      "igpe.\n",
      "log Likelihood: -156.45962524414062\n",
      "Negative log likelihood: 156.45962524414062\n",
      "Normalized Negative log Likelihood: 2.793921947479248\n"
     ]
    }
   ],
   "source": [
    "names = []\n",
    "for i in range(10):\n",
    "    out = []\n",
    "    ix1, ix2 = 0, 0\n",
    "    while True:\n",
    "        # previosly we used P[ix]\n",
    "        # p = P[ix]\n",
    "\n",
    "        # now we use the softmax of the logits\n",
    "        xenc = F.one_hot(torch.tensor([ix1, ix2]).to(device), num_classes = 27).float().to(device)\n",
    "        xenc = xenc.view(-1, 27*2)\n",
    "        \n",
    "        logits = xenc @ W\n",
    "        counts = torch.exp(logits)\n",
    "        p = counts / counts.sum(dim = 1, keepdim = True)\n",
    "\n",
    "        ix1 = ix2\n",
    "        ix2 = torch.multinomial(p.to(device), num_samples = 1 , replacement = True).item()\n",
    "        out.append(itos[ix2])\n",
    "        if ix2 == 0:\n",
    "            break\n",
    "\n",
    "    names.append(\"\".join(out))\n",
    "    \n",
    "for name in names:\n",
    "    print(name)\n",
    "count_loss(names)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# E02: split up the dataset randomly into 80% train set, 10% dev set, 10% test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# prepare the dataset\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "words_train, words_test = train_test_split(words, test_size=0.2, random_state=1234)\n",
    "words_dev, words_test = train_test_split(words_test, test_size=0.5, random_state=1234)\n",
    "\n",
    "x_train, y_train, x_dev, y_dev, x_test, y_test = [], [], [], [], [], []\n",
    "for wgroup in [words_train, words_dev, words_test]:\n",
    "    xs , ys = [], []\n",
    "    for w in wgroup:\n",
    "        # add start and end tokens\n",
    "        chs = [\".\"] + list(w) + [\".\"]\n",
    "        for ch1, ch2, ch3 in zip(chs, chs[1:], chs[2:]):\n",
    "            ix1 = stoi[ch1]\n",
    "            ix2 = stoi[ch2]\n",
    "            ix3 = stoi[ch3]\n",
    "        \n",
    "            xs.append([ix1, ix2])\n",
    "            ys.append(ix3)\n",
    "\n",
    "    xs = torch.tensor(xs, dtype=torch.int64)\n",
    "    ys = torch.tensor(ys, dtype=torch.int64)\n",
    "\n",
    "    if wgroup == words_train:\n",
    "        x_train, y_train = xs, ys\n",
    "    elif wgroup == words_dev:\n",
    "        x_dev, y_dev = xs, ys\n",
    "    else:\n",
    "        x_test, y_test = xs, ys\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0: 4.1824\n",
      "10: 2.4879\n",
      "20: 2.3713\n",
      "30: 2.3283\n",
      "40: 2.3057\n",
      "50: 2.2916\n",
      "60: 2.2819\n",
      "70: 2.2749\n",
      "80: 2.2696\n",
      "90: 2.2655\n",
      "100: 2.2622\n",
      "110: 2.2595\n",
      "120: 2.2573\n",
      "130: 2.2555\n",
      "140: 2.2539\n",
      "150: 2.2525\n",
      "160: 2.2513\n",
      "170: 2.2503\n",
      "180: 2.2494\n",
      "190: 2.2486\n"
     ]
    }
   ],
   "source": [
    "W = torch.randn((27*2,27), requires_grad = True, device = device)\n",
    "for k in range(200):\n",
    "    # forward pass\n",
    "    xenc = F.one_hot(x_train, num_classes = 27).float().to(device)\n",
    "    xenc = xenc.view(-1, 27*2)\n",
    "    \n",
    "    # probs is softmax\n",
    "    logits = xenc @ W\n",
    "    counts = torch.exp(logits)\n",
    "    probs = counts / counts.sum(dim = 1, keepdim = True)\n",
    "    \n",
    "    # loss (normalized negative log likelihood)\n",
    "    loss = - probs[torch.arange(len(x_train)), y_train].log().mean()\n",
    "    # add regularization\n",
    "    # loss += 0.2 * W.pow(2).mean()\n",
    "\n",
    "    if k % 10 == 0:\n",
    "        print(f\"{k}: {loss.item():.4f}\")\n",
    "\n",
    "    # backward pass\n",
    "    W.grad = None\n",
    "    loss.backward()\n",
    "\n",
    "    # update weights\n",
    "    with torch.no_grad():\n",
    "        W -= 50 * W.grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def MLP_loss(x, y, W):\n",
    "    xenc = F.one_hot(x, num_classes = 27).float().to(device)\n",
    "    xenc = xenc.view(-1, 27*2)\n",
    "\n",
    "    # probs is softmax\n",
    "    logits = xenc @ W\n",
    "    counts = torch.exp(logits)\n",
    "    probs = counts / counts.sum(dim = 1, keepdim = True)\n",
    "\n",
    "    # loss (normalized negative log likelihood)\n",
    "    loss = - probs[torch.arange(len(x)), y].log().mean()\n",
    "    \n",
    "    return loss.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 2.2478\n",
      "Dev Loss: 2.2538\n",
      "Test Loss: 2.2511\n"
     ]
    }
   ],
   "source": [
    "print(f\"Train Loss: {MLP_loss(x_train, y_train, W):.4f}\")\n",
    "print(f\"Dev Loss: {MLP_loss(x_dev, y_dev, W):.4f}\")\n",
    "print(f\"Test Loss: {MLP_loss(x_test, y_test, W):.4f}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# E03: use the dev set \n",
    "to tune the strength of smoothing (or regularization) for the trigram model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0: Train Loss: 4.1959 | Dev Loss 4.1741\n",
      "10: Train Loss: 2.5150 | Dev Loss 2.5090\n",
      "20: Train Loss: 2.3855 | Dev Loss 2.3831\n",
      "30: Train Loss: 2.3360 | Dev Loss 2.3357\n",
      "40: Train Loss: 2.3102 | Dev Loss 2.3110\n",
      "50: Train Loss: 2.2947 | Dev Loss 2.2961\n",
      "60: Train Loss: 2.2843 | Dev Loss 2.2861\n",
      "70: Train Loss: 2.2769 | Dev Loss 2.2790\n",
      "80: Train Loss: 2.2713 | Dev Loss 2.2738\n",
      "90: Train Loss: 2.2670 | Dev Loss 2.2698\n",
      "100: Train Loss: 2.2636 | Dev Loss 2.2666\n",
      "110: Train Loss: 2.2608 | Dev Loss 2.2640\n",
      "120: Train Loss: 2.2585 | Dev Loss 2.2619\n",
      "130: Train Loss: 2.2566 | Dev Loss 2.2601\n",
      "140: Train Loss: 2.2549 | Dev Loss 2.2587\n",
      "150: Train Loss: 2.2535 | Dev Loss 2.2574\n",
      "160: Train Loss: 2.2522 | Dev Loss 2.2563\n",
      "170: Train Loss: 2.2511 | Dev Loss 2.2553\n",
      "180: Train Loss: 2.2501 | Dev Loss 2.2544\n",
      "190: Train Loss: 2.2492 | Dev Loss 2.2536\n"
     ]
    }
   ],
   "source": [
    "W = torch.randn((27*2,27), requires_grad = True, device = device)\n",
    "for k in range(200):\n",
    "    # forward pass\n",
    "    xenc = F.one_hot(x_train, num_classes = 27).float().to(device)\n",
    "    xenc = xenc.view(-1, 27*2)\n",
    "    \n",
    "    # probs is softmax\n",
    "    logits = xenc @ W\n",
    "    counts = torch.exp(logits)\n",
    "    probs = counts / counts.sum(dim = 1, keepdim = True)\n",
    "    \n",
    "    # loss (normalized negative log likelihood)\n",
    "    loss = - probs[torch.arange(len(x_train)), y_train].log().mean()\n",
    "    # add regularization\n",
    "    # loss += 0.05 * W.pow(2).mean()\n",
    "\n",
    "    if k % 10 == 0:\n",
    "        print(f\"{k}: Train Loss: {loss.item():.4f} | Dev Loss {MLP_loss(x_dev, y_dev, W):.4f}\")\n",
    "\n",
    "    # backward pass\n",
    "    W.grad = None\n",
    "    loss.backward()\n",
    "\n",
    "    # update weights\n",
    "    with torch.no_grad():\n",
    "        W -= 50 * W.grad"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "comment: no regularization is better"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# E04: Rewrite the MLP model without creating one hot vectors\n",
    "we saw that our 1-hot vectors merely select a row of W, so producing these vectors explicitly feels wasteful. Can you delete our use of F.one_hot in favor of simply indexing into rows of W?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0: 4.2724\n",
      "10: 2.5831\n",
      "20: 2.4545\n",
      "30: 2.4076\n",
      "40: 2.3848\n",
      "50: 2.3720\n",
      "60: 2.3644\n",
      "70: 2.3595\n",
      "80: 2.3564\n",
      "90: 2.3544\n",
      "100: 2.3530\n",
      "110: 2.3520\n",
      "120: 2.3514\n",
      "130: 2.3509\n",
      "140: 2.3506\n",
      "150: 2.3504\n",
      "160: 2.3503\n",
      "170: 2.3501\n",
      "180: 2.3501\n",
      "190: 2.3500\n"
     ]
    }
   ],
   "source": [
    "W = torch.randn((27*2,27), requires_grad = True, device = device)\n",
    "for k in range(200):\n",
    "    # forward pass\n",
    "    # ====================\n",
    "    # Previously: using onehot and multiplying by W \n",
    "    # xenc = F.one_hot(xs, num_classes = 27).float().to(device)\n",
    "    # xenc = xenc.view(-1, 27*2)\n",
    "    # logits = xenc @ W\n",
    "    # ====================\n",
    "\n",
    "    # ====================\n",
    "    # ✅ now: acess by xs indices directly\n",
    "    logits = W[xs[:,0]] + W[xs[:,1] + 27]\n",
    "    # ====================\n",
    "    \n",
    "    counts = torch.exp(logits)\n",
    "    probs = counts / counts.sum(dim = 1, keepdim = True)\n",
    "    \n",
    "    # loss (normalized negative log likelihood)\n",
    "    loss = - probs[torch.arange(len(xs)), ys].log().mean()\n",
    "    # add regularization\n",
    "    loss += 0.2 * W.pow(2).mean()\n",
    "\n",
    "    if k % 10 == 0:\n",
    "        print(f\"{k}: {loss.item():.4f}\")\n",
    "\n",
    "    # backward pass\n",
    "    W.grad = None\n",
    "    loss.backward()\n",
    "\n",
    "    # update weights\n",
    "    with torch.no_grad():\n",
    "        W -= 50 * W.grad"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# E05: look up and use F.cross_entropy instead\n",
    "nn.functonal.cross_entropy() takes the logits and the target class as input and returns the cross entropy loss directly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0: 4.1529\n",
      "10: 2.5952\n",
      "20: 2.4588\n",
      "30: 2.4093\n",
      "40: 2.3856\n",
      "50: 2.3725\n",
      "60: 2.3647\n",
      "70: 2.3598\n",
      "80: 2.3566\n",
      "90: 2.3545\n",
      "100: 2.3531\n",
      "110: 2.3521\n",
      "120: 2.3514\n",
      "130: 2.3510\n",
      "140: 2.3507\n",
      "150: 2.3504\n",
      "160: 2.3503\n",
      "170: 2.3501\n",
      "180: 2.3501\n",
      "190: 2.3500\n"
     ]
    }
   ],
   "source": [
    "W = torch.randn((27*2,27), requires_grad = True, device = device)\n",
    "for k in range(200):\n",
    "    # forward pass\n",
    "    xenc = F.one_hot(xs, num_classes = 27).float().to(device)\n",
    "    xenc = xenc.view(-1, 27*2)\n",
    "    logits = xenc @ W\n",
    "    \n",
    "    loss = torch.nn.functional.cross_entropy(logits, ys.to(device))\n",
    "    # add regularization\n",
    "    loss += 0.2 * W.pow(2).mean()\n",
    "\n",
    "    if k % 10 == 0:\n",
    "        print(f\"{k}: {loss.item():.4f}\")\n",
    "\n",
    "    # backward pass\n",
    "    W.grad = None\n",
    "    loss.backward()\n",
    "\n",
    "    # update weights\n",
    "    with torch.no_grad():\n",
    "        W -= 50 * W.grad"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# E06: meta-exercise! Think of a fun/interesting exercise and complete it\n",
    "we will reimplment the MLP model using pytorch nn.Module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 163,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLP(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.fc1 = torch.nn.Linear(27*2, 27)\n",
    "        # initialize weights with normal distribution with mean 0 and std 1\n",
    "        torch.nn.init.normal_(self.fc1.weight, mean = 0, std = 1)\n",
    "        \n",
    "    def forward(self, xs):\n",
    "        logits = W[xs[:,0]] + W[xs[:,1] + 27]\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 175,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0: 4.4746\n",
      "10: 2.6566\n",
      "20: 2.5239\n",
      "30: 2.4767\n",
      "40: 2.4558\n",
      "50: 2.4462\n",
      "60: 2.4419\n",
      "70: 2.4405\n",
      "80: 2.4408\n",
      "90: 2.4420\n",
      "100: 2.4438\n",
      "110: 2.4460\n",
      "120: 2.4485\n",
      "130: 2.4511\n",
      "140: 2.4539\n",
      "150: 2.4567\n",
      "160: 2.4596\n",
      "170: 2.4625\n",
      "180: 2.4654\n",
      "190: 2.4683\n"
     ]
    }
   ],
   "source": [
    "model = MLP().to(device)\n",
    "for k in range(200):\n",
    "    # forward pass\n",
    "    logits = model(xs)\n",
    "    \n",
    "    loss = torch.nn.functional.cross_entropy(logits, ys.to(device))\n",
    "    # add regularization\n",
    "    loss += 0.2 * model.fc1.weight.data.pow(2).mean()\n",
    "\n",
    "    if k % 10 == 0:\n",
    "        print(f\"{k}: {loss.item():.4f}\")\n",
    "\n",
    "    # backward pass\n",
    "    model.fc1.zero_grad() # reset the gradients of the layer\n",
    "    loss.backward()\n",
    "\n",
    "    # update weights\n",
    "    with torch.no_grad():\n",
    "        model.fc1.weight.data -= 50 * model.fc1.weight.grad # use linear.weight instead of W"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytor",
   "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.3"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}