GPT_Model.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyP3eVet4OA3vSB1ZXISo7XY",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/alephpt/AI_Projects/blob/main/GPT_Model.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kwjBAG2zQ3-5"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.nn import functional as fn\n",
        "import requests"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "batch_size = 16 # Number of Sequences to Process in Parallel\n",
        "block_size = 32 # Max Content Length for Predictions\n",
        "max_iterations = 5000\n",
        "eval_interval = 100\n",
        "learning_rate = 1e-3\n",
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "eval_iterations = 200 # Number of iterations before saving state\n",
        "n_embed = 64 # Maps the input sequence into high-dimensional space\n",
        "n_head = 4 # Attention head for multi-head attention layers\n",
        "n_layers = 4\n",
        "dropout = 0.0"
      ],
      "metadata": {
        "id": "QwE7Io8fS-e-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.manual_seed(1337)\n",
        "bible_url = 'https://www.o-bible.com/download/kjv.txt'\n",
        "bible = requests.get(bible_url)\n",
        "bible_text = bible.text.split(\"\\n\")\n",
        "bible_text.remove(bible_text[0])"
      ],
      "metadata": {
        "id": "cNYvFvEvWQjO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(bible_text[0])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_dghIlZ7iVoL",
        "outputId": "1af54159-f6e5-404e-bf54-898db8995170"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Ge1:1 In the beginning God created the heaven and the earth.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "vocab = sorted(list(set(bible.text)))\n",
        "vocab_size = len(vocab)\n",
        "print(\"Vocab Size:\", vocab_size)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_3Q4MPZ6iqfF",
        "outputId": "2cd62f95-822b-4542-f17f-c0663b17ccab"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Vocab Size: 73\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "stoi = { ch: i for i, ch in enumerate(vocab)}\n",
        "itos = { i: ch for i, ch in enumerate(vocab)}\n",
        "encode = lambda s: [stoi[c] for c in s]\n",
        "decode = lambda l: \"\".join([itos[i] for i in l])\n",
        "\n",
        "nums = encode(\"Today is the day.\")\n",
        "words = decode(nums)\n",
        "\n",
        "print(\"nums\", nums)\n",
        "print(\"words\", words)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "eU1VGta-mTL0",
        "outputId": "61d7f8e8-4d88-44a4-8e0e-51ff568947f0"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "nums [41, 61, 50, 47, 71, 1, 55, 65, 1, 66, 54, 51, 1, 50, 47, 71, 8]\n",
            "words Today is the day.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "data = torch.tensor(encode(bible.text), dtype=torch.long)\n",
        "print(data.shape, data.dtype)\n",
        "print(data[:100])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uZGhGt1cpw4q",
        "outputId": "50f7585c-1872-489c-f894-fa9327396632"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([4404443]) torch.int64\n",
            "tensor([29, 61, 58, 71,  1, 23, 55, 48, 58, 51,  6,  1, 22, 67, 66, 54, 61, 64,\n",
            "        55, 72, 51, 50,  1,  4, 32, 55, 60, 53,  1, 31, 47, 59, 51, 65,  5,  1,\n",
            "        43, 51, 64, 65, 55, 61, 60,  6,  1, 41, 51, 70, 66, 52, 55, 58, 51,  1,\n",
            "        18, 12,  9, 10,  9, 14,  8,  0, 28, 51, 10, 19, 10,  1, 30, 60,  1, 66,\n",
            "        54, 51,  1, 48, 51, 53, 55, 60, 60, 55, 60, 53,  1, 28, 61, 50,  1, 49,\n",
            "        64, 51, 47, 66, 51, 50,  1, 66, 54, 51])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "n = int(0.9 * len(data)) # train on the first 90%\n",
        "training_data = data[:n]\n",
        "validation_data = data[n:]\n",
        "block_size = 8\n",
        "training_data[:block_size + 1]\n",
        "x = training_data[:block_size]\n",
        "y = training_data[1:block_size + 1]\n",
        "\n",
        "for t in range(block_size):\n",
        "  context = x[:t+1]\n",
        "  target = y[t]\n",
        "  print(f\"when input is {context} the target: {target}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UmlC4J0NqOWk",
        "outputId": "91510f92-15c2-48a4-a694-4addbf55d958"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "when input is tensor([29]) the target: 61\n",
            "when input is tensor([29, 61]) the target: 58\n",
            "when input is tensor([29, 61, 58]) the target: 71\n",
            "when input is tensor([29, 61, 58, 71]) the target: 1\n",
            "when input is tensor([29, 61, 58, 71,  1]) the target: 23\n",
            "when input is tensor([29, 61, 58, 71,  1, 23]) the target: 55\n",
            "when input is tensor([29, 61, 58, 71,  1, 23, 55]) the target: 48\n",
            "when input is tensor([29, 61, 58, 71,  1, 23, 55, 48]) the target: 58\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def get_batch(split):\n",
        "  data = training_data if split == \"train\" else validation_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"
      ],
      "metadata": {
        "id": "Hqv1tKTxqgdO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "xb, yb = get_batch('train')\n",
        "print(\"inputs:\")\n",
        "print(xb.shape)\n",
        "print(xb)\n",
        "print(\"targets:\")\n",
        "print(yb.shape)\n",
        "print(yb)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4clRbxrTsZ0N",
        "outputId": "2eff93df-8f50-44e5-dd9a-45231c0db768"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "inputs:\n",
            "torch.Size([16, 8])\n",
            "tensor([[65, 47, 64,  6,  1, 61, 64,  1],\n",
            "        [ 1, 54, 51, 47, 64, 20,  1, 65],\n",
            "        [71, 19,  1, 52, 61, 64,  1, 66],\n",
            "        [48, 61, 69, 51, 50,  1, 54, 55],\n",
            "        [51, 51,  6,  1, 47, 60, 50,  1],\n",
            "        [69, 51, 64, 51,  1, 60, 47, 59],\n",
            "        [67, 60, 50, 47, 66, 55, 61, 60],\n",
            "        [58, 51, 66,  1, 66, 54, 55, 60],\n",
            "        [ 1, 55, 52,  1, 66, 54, 61, 67],\n",
            "        [60, 50,  1, 66, 54, 51, 71,  1],\n",
            "        [54, 51, 64, 48, 19,  1, 47, 60],\n",
            "        [51,  1, 65, 61, 60, 65,  1, 61],\n",
            "        [64,  1, 28, 61, 50,  8,  0, 29],\n",
            "        [19, 18,  1, 22, 60, 50,  1, 28],\n",
            "        [33, 36, 39, 25,  1, 61, 67, 64],\n",
            "        [71, 62, 66,  6,  1, 52, 64, 61]])\n",
            "targets:\n",
            "torch.Size([16, 8])\n",
            "tensor([[47, 64,  6,  1, 61, 64,  1, 60],\n",
            "        [54, 51, 47, 64, 20,  1, 65, 47],\n",
            "        [19,  1, 52, 61, 64,  1, 66, 54],\n",
            "        [61, 69, 51, 50,  1, 54, 55, 59],\n",
            "        [51,  6,  1, 47, 60, 50,  1, 49],\n",
            "        [51, 64, 51,  1, 60, 47, 59, 51],\n",
            "        [60, 50, 47, 66, 55, 61, 60, 65],\n",
            "        [51, 66,  1, 66, 54, 55, 60, 51],\n",
            "        [55, 52,  1, 66, 54, 61, 67,  1],\n",
            "        [50,  1, 66, 54, 51, 71,  1, 65],\n",
            "        [51, 64, 48, 19,  1, 47, 60, 50],\n",
            "        [ 1, 65, 61, 60, 65,  1, 61, 52],\n",
            "        [ 1, 28, 61, 50,  8,  0, 29, 61],\n",
            "        [18,  1, 22, 60, 50,  1, 28, 61],\n",
            "        [36, 39, 25,  1, 61, 67, 64,  1],\n",
            "        [62, 66,  6,  1, 52, 64, 61, 59]])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "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: {target}\")"
      ],
      "metadata": {
        "id": "e_s-uB1hsz10"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "@torch.no_grad()\n",
        "def estimate_loss():\n",
        "    out = {}\n",
        "    model.eval()\n",
        "    for split in ['train', 'val']:\n",
        "        losses = torch.zeros(eval_iterations)\n",
        "        for k in range(eval_iterations):\n",
        "            X, Y = get_batch(split)\n",
        "            logits, loss = model(X, Y)\n",
        "            losses[k] = loss.item()\n",
        "        out[split] = losses.mean()\n",
        "    model.train()\n",
        "    return out"
      ],
      "metadata": {
        "id": "mDWDYISdtLWC"
      },
      "execution_count": 35,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Head(nn.Module):\n",
        "  def __init__(self, head_size):\n",
        "    super().__init__()\n",
        "    self.key = nn.Linear(n_embed, head_size, bias=False)\n",
        "    self.query = nn.Linear(n_embed, head_size, bias=False)\n",
        "    self.value = nn.Linear(n_embed, head_size, bias=False)\n",
        "    self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))\n",
        "    self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "  def forward(self, x):\n",
        "    B, T, C = x.shape\n",
        "    k = self.key(x)\n",
        "    q = self.query(x)\n",
        "    \n",
        "    wei = q @ k.transpose(-2, -1) * C**-0.5\n",
        "    wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))\n",
        "    wei = fn.softmax(wei, -1)\n",
        "    wei = self.dropout(wei)\n",
        "\n",
        "    v = self.value(x)\n",
        "\n",
        "    return wei @ v"
      ],
      "metadata": {
        "id": "4A_XYuJyHitQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class MultiHeadAtt(nn.Module):\n",
        "  def __init__(self, n_heads, head_size):\n",
        "    super().__init__()\n",
        "    self.heads = nn.ModuleList([Head(head_size) for _ in range(n_heads)])\n",
        "    self.proj = nn.Linear(n_embed, n_embed)\n",
        "    self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "  def forward(self, x):\n",
        "    out = torch.cat([h(x) for h in self.heads], -1)\n",
        "    out = self.dropout(self.proj(out))\n",
        "    return out"
      ],
      "metadata": {
        "id": "PD7v5jFdOTRP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class FeedForward(nn.Module):\n",
        "  def __init__(self, n_embed):\n",
        "    super().__init__()\n",
        "    self.net = nn.Sequential(\n",
        "        nn.Linear(n_embed, n_embed),\n",
        "        nn.ReLU(),\n",
        "    )\n",
        "  \n",
        "  def forward(self, x):\n",
        "    return self.net(x)"
      ],
      "metadata": {
        "id": "vpO9OqbYuxdt"
      },
      "execution_count": 21,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Block(nn.Module):\n",
        "  def __init__(self, n_embed, n_head):\n",
        "    super().__init__()\n",
        "    head_size = n_embed // n_head\n",
        "    self.sa = MultiHeadAtt(n_head, head_size)\n",
        "    self.ffwd = FeedForward(n_embed)\n",
        "\n",
        "  def forward(self):\n",
        "    x = x + self.sa(x)\n",
        "    x = x + self.ffwd(x)\n",
        "    return x"
      ],
      "metadata": {
        "id": "nn_6rLOBbCtp"
      },
      "execution_count": 23,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Bigram(nn.Module):\n",
        "  def __init__(self, n_embed, n_head):\n",
        "    super().__init__()\n",
        "    self.token_embedding_table = nn.Embedding(vocab_size, n_embed)\n",
        "    self.position_embedding_table = nn.Embedding(block_size, n_embed)\n",
        "    self.blocks = nn.Sequential(*[Block(n_embed, n_head) for _ in range(n_layers)])\n",
        "    self.ln_f = nn.LayerNorm(n_embed)\n",
        "    self.lm_head = nn.Linear(n_embed, vocab_size)\n",
        "\n",
        "  def forward(self, idx, targets=None):\n",
        "    B, T = idx.shape\n",
        "\n",
        "    tok_emb = self.token_embedding_table(idx)\n",
        "    pos_emb = self.position_embedding_table(torch.arange(T, device=device))\n",
        "\n",
        "    x = tok_emb + pos_emb\n",
        "    x = self.blocks(x)\n",
        "    x = self.ln_f(x)\n",
        "\n",
        "    logits = self.lm_head(x)\n",
        "\n",
        "    if targets == None:\n",
        "      loss = targets\n",
        "    else:\n",
        "      B, T, C = logits.shape\n",
        "      logits = logits.view(B*T, c)\n",
        "      targets = targets.view(B*T)\n",
        "      loss = fn.cross_entropy(logits, targets)\n",
        "\n",
        "    return logits, loss\n",
        "  \n",
        "  def generate(self, idx, max_new_tokens):\n",
        "    for _ in range(max_new_tokens):\n",
        "      idx_cond = idx[:, -block_size:]\n",
        "      logits, loss = self(idx_cond)\n",
        "      logits = logits[:, -1, :]\n",
        "      probs = fn.softmax(logits, -1)\n",
        "      idx_next = torch.multinomial(probs, 1)\n",
        "      idx = torch.cat((idx, idx_next), 1)\n",
        "    return idx"
      ],
      "metadata": {
        "id": "rGMxq9mUSg71"
      },
      "execution_count": 39,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = Bigram(n_embed, n_head)\n",
        "m = model.to(device)\n",
        "print(sum(p.numel() for p in m.parameters())/1e3, 'K parameters')\n",
        "\n",
        "optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)\n",
        "\n",
        "for i in range(max_iterations):\n",
        "  if i % eval_iterations == 0 or i == max_iterations - 1:\n",
        "    losses = estimate_loss()\n",
        "    print(f\"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}\")\n",
        "\n",
        "  xb, yb = get_batch('train')\n",
        "\n",
        "  logits, loss = model(xb, yb)\n",
        "  optimizer.zero_grad(set_to_none=True)\n",
        "  loss.backwards()\n",
        "  optimizer.step()\n",
        "\n",
        "context = torch.zero((1, 1), dtype=torch.long, device=device)\n",
        "print(decode(m.generate(context, max_new_tokens=2000)[0].tolist()))\n",
        "\n",
        "  "
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 407
        },
        "id": "nwkiwWMqp0oB",
        "outputId": "42b42552-2bd0-4096-9b0c-cc550f940c66"
      },
      "execution_count": 42,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "92.489 K parameters\n"
          ]
        },
        {
          "output_type": "error",
          "ename": "TypeError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-42-cabc6e42a1ab>\u001b[0m in \u001b[0;36m<cell line: 7>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      7\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmax_iterations\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      8\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0mi\u001b[0m \u001b[0;34m%\u001b[0m \u001b[0meval_iterations\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m0\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0mmax_iterations\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 9\u001b[0;31m     \u001b[0mlosses\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mestimate_loss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     10\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/torch/utils/_contextlib.py\u001b[0m in \u001b[0;36mdecorate_context\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    113\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mdecorate_context\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    114\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0mctx_factory\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 115\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    116\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    117\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mdecorate_context\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-35-f08c3538c4bb>\u001b[0m in \u001b[0;36mestimate_loss\u001b[0;34m()\u001b[0m\n\u001b[1;32m      7\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mk\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0meval_iterations\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      8\u001b[0m             \u001b[0mX\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mY\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_batch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msplit\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 9\u001b[0;31m             \u001b[0mlogits\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mY\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     10\u001b[0m             \u001b[0mlosses\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m         \u001b[0mout\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0msplit\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlosses\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmean\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1499\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1500\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1502\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1503\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-39-9e7abd1dbd41>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, idx, targets)\u001b[0m\n\u001b[1;32m     15\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     16\u001b[0m     \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtok_emb\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mpos_emb\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 17\u001b[0;31m     \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mblocks\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     18\u001b[0m     \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mln_f\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     19\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1499\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1500\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1502\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1503\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/torch/nn/modules/container.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    215\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    216\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 217\u001b[0;31m             \u001b[0minput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    218\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    219\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1499\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1500\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1502\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1503\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mTypeError\u001b[0m: forward() takes 1 positional argument but 2 were given"
          ]
        }
      ]
    }
  ]
}