pyristix_gpt.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "uf80S4psYtSF",
        "P-zpd9qzY_EB"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Dataset Setup"
      ],
      "metadata": {
        "id": "uf80S4psYtSF"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Dependency Setup"
      ],
      "metadata": {
        "id": "P-zpd9qzY_EB"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "J-6dZuUHvXVa",
        "outputId": "bef0fc2d-9ea4-418e-d239-e767d7f52575"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting convokit\n",
            "  Downloading convokit-3.0.0.tar.gz (183 kB)\n",
            "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/183.2 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K     \u001b[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[91m╸\u001b[0m\u001b[90m━━\u001b[0m \u001b[32m174.1/183.2 kB\u001b[0m \u001b[31m5.9 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m183.2/183.2 kB\u001b[0m \u001b[31m4.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: matplotlib>=3.0.0 in /usr/local/lib/python3.10/dist-packages (from convokit) (3.7.1)\n",
            "Requirement already satisfied: pandas>=0.23.4 in /usr/local/lib/python3.10/dist-packages (from convokit) (1.5.3)\n",
            "Collecting msgpack-numpy>=0.4.3.2 (from convokit)\n",
            "  Downloading msgpack_numpy-0.4.8-py2.py3-none-any.whl (6.9 kB)\n",
            "Requirement already satisfied: spacy>=2.3.5 in /usr/local/lib/python3.10/dist-packages (from convokit) (3.7.3)\n",
            "Requirement already satisfied: scipy>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from convokit) (1.11.4)\n",
            "Requirement already satisfied: scikit-learn>=0.20.0 in /usr/local/lib/python3.10/dist-packages (from convokit) (1.2.2)\n",
            "Requirement already satisfied: nltk>=3.4 in /usr/local/lib/python3.10/dist-packages (from convokit) (3.8.1)\n",
            "Collecting dill>=0.2.9 (from convokit)\n",
            "  Downloading dill-0.3.8-py3-none-any.whl (116 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m116.3/116.3 kB\u001b[0m \u001b[31m11.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: joblib>=0.13.2 in /usr/local/lib/python3.10/dist-packages (from convokit) (1.3.2)\n",
            "Collecting clean-text>=0.6.0 (from convokit)\n",
            "  Downloading clean_text-0.6.0-py3-none-any.whl (11 kB)\n",
            "Collecting unidecode>=1.1.1 (from convokit)\n",
            "  Downloading Unidecode-1.3.8-py3-none-any.whl (235 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m235.5/235.5 kB\u001b[0m \u001b[31m13.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: tqdm>=4.64.0 in /usr/local/lib/python3.10/dist-packages (from convokit) (4.66.1)\n",
            "Collecting pymongo>=4.0 (from convokit)\n",
            "  Downloading pymongo-4.6.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (677 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m677.1/677.1 kB\u001b[0m \u001b[31m15.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: pyyaml>=5.4.1 in /usr/local/lib/python3.10/dist-packages (from convokit) (6.0.1)\n",
            "Collecting dnspython>=1.16.0 (from convokit)\n",
            "  Downloading dnspython-2.5.0-py3-none-any.whl (305 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m305.4/305.4 kB\u001b[0m \u001b[31m15.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting emoji<2.0.0,>=1.0.0 (from clean-text>=0.6.0->convokit)\n",
            "  Downloading emoji-1.7.0.tar.gz (175 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m175.4/175.4 kB\u001b[0m \u001b[31m12.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Collecting ftfy<7.0,>=6.0 (from clean-text>=0.6.0->convokit)\n",
            "  Downloading ftfy-6.1.3-py3-none-any.whl (53 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m53.4/53.4 kB\u001b[0m \u001b[31m8.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=3.0.0->convokit) (1.2.0)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=3.0.0->convokit) (0.12.1)\n",
            "Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=3.0.0->convokit) (4.48.1)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=3.0.0->convokit) (1.4.5)\n",
            "Requirement already satisfied: numpy>=1.20 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=3.0.0->convokit) (1.23.5)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=3.0.0->convokit) (23.2)\n",
            "Requirement already satisfied: pillow>=6.2.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=3.0.0->convokit) (9.4.0)\n",
            "Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=3.0.0->convokit) (3.1.1)\n",
            "Requirement already satisfied: python-dateutil>=2.7 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=3.0.0->convokit) (2.8.2)\n",
            "Requirement already satisfied: msgpack>=0.5.2 in /usr/local/lib/python3.10/dist-packages (from msgpack-numpy>=0.4.3.2->convokit) (1.0.7)\n",
            "Requirement already satisfied: click in /usr/local/lib/python3.10/dist-packages (from nltk>=3.4->convokit) (8.1.7)\n",
            "Requirement already satisfied: regex>=2021.8.3 in /usr/local/lib/python3.10/dist-packages (from nltk>=3.4->convokit) (2023.12.25)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas>=0.23.4->convokit) (2023.4)\n",
            "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from scikit-learn>=0.20.0->convokit) (3.2.0)\n",
            "Requirement already satisfied: spacy-legacy<3.1.0,>=3.0.11 in /usr/local/lib/python3.10/dist-packages (from spacy>=2.3.5->convokit) (3.0.12)\n",
            "Requirement already satisfied: spacy-loggers<2.0.0,>=1.0.0 in /usr/local/lib/python3.10/dist-packages (from spacy>=2.3.5->convokit) (1.0.5)\n",
            "Requirement already satisfied: murmurhash<1.1.0,>=0.28.0 in /usr/local/lib/python3.10/dist-packages (from spacy>=2.3.5->convokit) (1.0.10)\n",
            "Requirement already satisfied: cymem<2.1.0,>=2.0.2 in /usr/local/lib/python3.10/dist-packages (from spacy>=2.3.5->convokit) (2.0.8)\n",
            "Requirement already satisfied: preshed<3.1.0,>=3.0.2 in /usr/local/lib/python3.10/dist-packages (from spacy>=2.3.5->convokit) (3.0.9)\n",
            "Requirement already satisfied: thinc<8.3.0,>=8.2.2 in /usr/local/lib/python3.10/dist-packages (from spacy>=2.3.5->convokit) (8.2.3)\n",
            "Requirement already satisfied: wasabi<1.2.0,>=0.9.1 in /usr/local/lib/python3.10/dist-packages (from spacy>=2.3.5->convokit) (1.1.2)\n",
            "Requirement already satisfied: srsly<3.0.0,>=2.4.3 in /usr/local/lib/python3.10/dist-packages (from spacy>=2.3.5->convokit) (2.4.8)\n",
            "Requirement already satisfied: catalogue<2.1.0,>=2.0.6 in /usr/local/lib/python3.10/dist-packages (from spacy>=2.3.5->convokit) (2.0.10)\n",
            "Requirement already satisfied: weasel<0.4.0,>=0.1.0 in /usr/local/lib/python3.10/dist-packages (from spacy>=2.3.5->convokit) (0.3.4)\n",
            "Requirement already satisfied: typer<0.10.0,>=0.3.0 in /usr/local/lib/python3.10/dist-packages (from spacy>=2.3.5->convokit) (0.9.0)\n",
            "Requirement already satisfied: smart-open<7.0.0,>=5.2.1 in /usr/local/lib/python3.10/dist-packages (from spacy>=2.3.5->convokit) (6.4.0)\n",
            "Requirement already satisfied: requests<3.0.0,>=2.13.0 in /usr/local/lib/python3.10/dist-packages (from spacy>=2.3.5->convokit) (2.31.0)\n",
            "Requirement already satisfied: pydantic!=1.8,!=1.8.1,<3.0.0,>=1.7.4 in /usr/local/lib/python3.10/dist-packages (from spacy>=2.3.5->convokit) (2.6.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from spacy>=2.3.5->convokit) (3.1.3)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.10/dist-packages (from spacy>=2.3.5->convokit) (67.7.2)\n",
            "Requirement already satisfied: langcodes<4.0.0,>=3.2.0 in /usr/local/lib/python3.10/dist-packages (from spacy>=2.3.5->convokit) (3.3.0)\n",
            "Requirement already satisfied: wcwidth<0.3.0,>=0.2.12 in /usr/local/lib/python3.10/dist-packages (from ftfy<7.0,>=6.0->clean-text>=0.6.0->convokit) (0.2.13)\n",
            "Requirement already satisfied: annotated-types>=0.4.0 in /usr/local/lib/python3.10/dist-packages (from pydantic!=1.8,!=1.8.1,<3.0.0,>=1.7.4->spacy>=2.3.5->convokit) (0.6.0)\n",
            "Requirement already satisfied: pydantic-core==2.16.2 in /usr/local/lib/python3.10/dist-packages (from pydantic!=1.8,!=1.8.1,<3.0.0,>=1.7.4->spacy>=2.3.5->convokit) (2.16.2)\n",
            "Requirement already satisfied: typing-extensions>=4.6.1 in /usr/local/lib/python3.10/dist-packages (from pydantic!=1.8,!=1.8.1,<3.0.0,>=1.7.4->spacy>=2.3.5->convokit) (4.9.0)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.7->matplotlib>=3.0.0->convokit) (1.16.0)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0,>=2.13.0->spacy>=2.3.5->convokit) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0,>=2.13.0->spacy>=2.3.5->convokit) (3.6)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0,>=2.13.0->spacy>=2.3.5->convokit) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0,>=2.13.0->spacy>=2.3.5->convokit) (2024.2.2)\n",
            "Requirement already satisfied: blis<0.8.0,>=0.7.8 in /usr/local/lib/python3.10/dist-packages (from thinc<8.3.0,>=8.2.2->spacy>=2.3.5->convokit) (0.7.11)\n",
            "Requirement already satisfied: confection<1.0.0,>=0.0.1 in /usr/local/lib/python3.10/dist-packages (from thinc<8.3.0,>=8.2.2->spacy>=2.3.5->convokit) (0.1.4)\n",
            "Requirement already satisfied: cloudpathlib<0.17.0,>=0.7.0 in /usr/local/lib/python3.10/dist-packages (from weasel<0.4.0,>=0.1.0->spacy>=2.3.5->convokit) (0.16.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->spacy>=2.3.5->convokit) (2.1.5)\n",
            "Building wheels for collected packages: convokit, emoji\n",
            "  Building wheel for convokit (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for convokit: filename=convokit-3.0.0-py3-none-any.whl size=216707 sha256=51574f9c71de0f2053ad6c4b71c7df3a950367d85b558966eaa667fb80069516\n",
            "  Stored in directory: /root/.cache/pip/wheels/c4/89/8c/2677fdb888588b6f93cb6ac86bdfb020f1f1c33e0d5525b231\n",
            "  Building wheel for emoji (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for emoji: filename=emoji-1.7.0-py3-none-any.whl size=171033 sha256=2d09dc2dacc97eef013d5c5268e56867d3ce4ead0caa51042e71706fbda865d0\n",
            "  Stored in directory: /root/.cache/pip/wheels/31/8a/8c/315c9e5d7773f74b33d5ed33f075b49c6eaeb7cedbb86e2cf8\n",
            "Successfully built convokit emoji\n",
            "Installing collected packages: emoji, unidecode, msgpack-numpy, ftfy, dnspython, dill, pymongo, clean-text, convokit\n",
            "Successfully installed clean-text-0.6.0 convokit-3.0.0 dill-0.3.8 dnspython-2.5.0 emoji-1.7.0 ftfy-6.1.3 msgpack-numpy-0.4.8 pymongo-4.6.1 unidecode-1.3.8\n"
          ]
        }
      ],
      "source": [
        "!pip3 install convokit"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import nltk; nltk.download('punkt')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uSYbKizH2AGv",
        "outputId": "4ef0a60d-813b-42b1-8592-e938dfd00237"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "[nltk_data] Downloading package punkt to /root/nltk_data...\n",
            "[nltk_data]   Unzipping tokenizers/punkt.zip.\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {},
          "execution_count": 2
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!python3 -m spacy download en_core_web_sm"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Gq6OeigZ2IbP",
        "outputId": "d5ab4bdd-698b-41fa-a976-2e2cb755b504"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting en-core-web-sm==3.7.1\n",
            "  Downloading https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.7.1/en_core_web_sm-3.7.1-py3-none-any.whl (12.8 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m12.8/12.8 MB\u001b[0m \u001b[31m30.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: spacy<3.8.0,>=3.7.2 in /usr/local/lib/python3.10/dist-packages (from en-core-web-sm==3.7.1) (3.7.3)\n",
            "Requirement already satisfied: spacy-legacy<3.1.0,>=3.0.11 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (3.0.12)\n",
            "Requirement already satisfied: spacy-loggers<2.0.0,>=1.0.0 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (1.0.5)\n",
            "Requirement already satisfied: murmurhash<1.1.0,>=0.28.0 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (1.0.10)\n",
            "Requirement already satisfied: cymem<2.1.0,>=2.0.2 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (2.0.8)\n",
            "Requirement already satisfied: preshed<3.1.0,>=3.0.2 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (3.0.9)\n",
            "Requirement already satisfied: thinc<8.3.0,>=8.2.2 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (8.2.3)\n",
            "Requirement already satisfied: wasabi<1.2.0,>=0.9.1 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (1.1.2)\n",
            "Requirement already satisfied: srsly<3.0.0,>=2.4.3 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (2.4.8)\n",
            "Requirement already satisfied: catalogue<2.1.0,>=2.0.6 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (2.0.10)\n",
            "Requirement already satisfied: weasel<0.4.0,>=0.1.0 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (0.3.4)\n",
            "Requirement already satisfied: typer<0.10.0,>=0.3.0 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (0.9.0)\n",
            "Requirement already satisfied: smart-open<7.0.0,>=5.2.1 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (6.4.0)\n",
            "Requirement already satisfied: tqdm<5.0.0,>=4.38.0 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (4.66.1)\n",
            "Requirement already satisfied: requests<3.0.0,>=2.13.0 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (2.31.0)\n",
            "Requirement already satisfied: pydantic!=1.8,!=1.8.1,<3.0.0,>=1.7.4 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (2.6.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (3.1.3)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (67.7.2)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (23.2)\n",
            "Requirement already satisfied: langcodes<4.0.0,>=3.2.0 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (3.3.0)\n",
            "Requirement already satisfied: numpy>=1.19.0 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (1.23.5)\n",
            "Requirement already satisfied: annotated-types>=0.4.0 in /usr/local/lib/python3.10/dist-packages (from pydantic!=1.8,!=1.8.1,<3.0.0,>=1.7.4->spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (0.6.0)\n",
            "Requirement already satisfied: pydantic-core==2.16.2 in /usr/local/lib/python3.10/dist-packages (from pydantic!=1.8,!=1.8.1,<3.0.0,>=1.7.4->spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (2.16.2)\n",
            "Requirement already satisfied: typing-extensions>=4.6.1 in /usr/local/lib/python3.10/dist-packages (from pydantic!=1.8,!=1.8.1,<3.0.0,>=1.7.4->spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (4.9.0)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0,>=2.13.0->spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0,>=2.13.0->spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (3.6)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0,>=2.13.0->spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0,>=2.13.0->spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (2024.2.2)\n",
            "Requirement already satisfied: blis<0.8.0,>=0.7.8 in /usr/local/lib/python3.10/dist-packages (from thinc<8.3.0,>=8.2.2->spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (0.7.11)\n",
            "Requirement already satisfied: confection<1.0.0,>=0.0.1 in /usr/local/lib/python3.10/dist-packages (from thinc<8.3.0,>=8.2.2->spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (0.1.4)\n",
            "Requirement already satisfied: click<9.0.0,>=7.1.1 in /usr/local/lib/python3.10/dist-packages (from typer<0.10.0,>=0.3.0->spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (8.1.7)\n",
            "Requirement already satisfied: cloudpathlib<0.17.0,>=0.7.0 in /usr/local/lib/python3.10/dist-packages (from weasel<0.4.0,>=0.1.0->spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (0.16.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->spacy<3.8.0,>=3.7.2->en-core-web-sm==3.7.1) (2.1.5)\n",
            "\u001b[38;5;2m✔ Download and installation successful\u001b[0m\n",
            "You can now load the package via spacy.load('en_core_web_sm')\n",
            "\u001b[38;5;3m⚠ Restart to reload dependencies\u001b[0m\n",
            "If you are in a Jupyter or Colab notebook, you may need to restart Python in\n",
            "order to load all the package's dependencies. You can do this by selecting the\n",
            "'Restart kernel' or 'Restart runtime' option.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from convokit import Corpus, download\n",
        "corpus = Corpus(filename=download('movie-corpus'))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ic5gIK7t2bBa",
        "outputId": "eb0670d7-cd68-4cfb-ec4a-b6bc27cef59a"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading movie-corpus to /root/.convokit/downloads/movie-corpus\n",
            "Downloading movie-corpus from http://zissou.infosci.cornell.edu/convokit/datasets/movie-corpus/movie-corpus.zip (40.9MB)... Done\n",
            "No configuration file found at /root/.convokit/config.yml; writing with contents: \n",
            "# Default Backend Parameters\n",
            "db_host: localhost:27017\n",
            "data_directory: ~/.convokit/saved-corpora\n",
            "default_backend: mem\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Compile character dialogue"
      ],
      "metadata": {
        "id": "JPlsprzbZP59"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import json\n",
        "from pprint import pprint\n",
        "\n",
        "with open('../root/.convokit/downloads/movie-corpus/utterances.jsonl') as json_file:\n",
        "  utterances = list(json_file)\n",
        "\n",
        "  # Convert JSON strings to JSON, find starts of conversations, and map utterance IDs to next utterances\n",
        "  convo_starts = []\n",
        "  next_utterance_map = {}\n",
        "  for i in range(len(utterances)):\n",
        "    utterance = json.loads(utterances[i])\n",
        "    if not utterance['reply-to']:\n",
        "      convo_starts.append(utterance)\n",
        "    utterances[i] = utterance\n",
        "    next_utterance_map[utterance['reply-to']] = utterance\n",
        "\n",
        "  # Compile character dialogue into dialogue.txt file\n",
        "  with open('dialogue.txt', 'w') as output_file:\n",
        "    for convo_start in convo_starts:\n",
        "      current_utterance = convo_start\n",
        "      while True:\n",
        "        output_file.write(current_utterance['text'] + '\\n')\n",
        "        if current_utterance['id'] in next_utterance_map:\n",
        "          current_utterance = next_utterance_map[current_utterance['id']]\n",
        "        else:\n",
        "          break"
      ],
      "metadata": {
        "id": "8OXKauBF295c"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Code"
      ],
      "metadata": {
        "id": "sNV6toORZiOk"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Imports\n"
      ],
      "metadata": {
        "id": "2Cq_qOccZYkv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.nn import functional"
      ],
      "metadata": {
        "id": "kxdCz6NjkJwc"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Data"
      ],
      "metadata": {
        "id": "moXHsrRGTNmG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "with open('dialogue.txt') as dialogue_file:\n",
        "  dialogue = dialogue_file.read()\n",
        "\n",
        "char_set = sorted(list(set(dialogue)))\n",
        "vocab_size = len(char_set)\n",
        "print(char_set)\n",
        "print(vocab_size)"
      ],
      "metadata": {
        "id": "8EnncQJmZpyK",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "da18b94a-ecb3-4c42-bae3-d9abe799466c"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "['\\t', '\\n', ' ', '!', '\"', '#', '$', '%', '&', \"'\", ')', '*', '+', ',', '-', '.', '/', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', ':', ';', '<', '=', '>', '?', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', '[', ']', '^', '_', '`', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z', '{', '|', '}', '~']\n",
            "94\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "char_to_token = { char:index for index, char in enumerate(char_set) }\n",
        "token_to_char = { index:char for index, char in enumerate(char_set) }\n",
        "encode = lambda string: [char_to_token[char] for char in string]\n",
        "decode = lambda tokens: ''.join([token_to_char[token] for token in tokens])\n",
        "print(encode('testing'))\n",
        "print(decode(encode('testing')))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FWH8rIeRfmT8",
        "outputId": "5543a147-3636-4308-8241-a9f3a71af655"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[83, 68, 82, 83, 72, 77, 70]\n",
            "testing\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "dataset = torch.tensor(encode(dialogue), dtype=torch.long)\n",
        "print(dataset.shape, dataset.dtype)\n",
        "\n",
        "n = int(0.9*len(dataset))\n",
        "train_data = dataset[:n]\n",
        "val_data = dataset[n:]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tUg7vZNCliOb",
        "outputId": "4839f03b-87b2-4b35-c069-aee8677f5590"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([17143013]) torch.int64\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def get_batch(data_source):\n",
        "  data = train_data if data_source == 'train' else val_data\n",
        "  batch_starts = torch.randint(len(data) - max_context_length, (num_batches,))\n",
        "  input_tokens = torch.stack([data[batch_start:batch_start+max_context_length] for batch_start in batch_starts])\n",
        "  target_tokens = torch.stack([data[batch_start+1:batch_start+max_context_length+1] for batch_start in batch_starts])\n",
        "  input_tokens, target_tokens = input_tokens.to(device), target_tokens.to(device)\n",
        "  return input_tokens, target_tokens\n",
        "\n",
        "input_tokens, target_tokens = get_batch('train')\n",
        "print(input_tokens.shape)\n",
        "print(target_tokens.shape)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VQawa4hTnkUE",
        "outputId": "be4476a5-f12e-4686-d2c9-ce06c305b49a"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([8, 256])\n",
            "torch.Size([8, 256])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Transformer model"
      ],
      "metadata": {
        "id": "z_r6WHVNTRJ0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "@torch.no_grad()\n",
        "def estimate_loss():\n",
        "  out = {}\n",
        "  model.eval()    # Puts model into evaluation mode\n",
        "\n",
        "  for split in ['train', 'val']:\n",
        "    losses = torch.zeros(n_eval_iterations)\n",
        "    for k in range(n_eval_iterations):\n",
        "      input_tokens, target_tokens = get_batch(split)\n",
        "      logits, loss = model(input_tokens, target_tokens)\n",
        "      losses[k] = loss.item()\n",
        "    out[split] = losses.mean()\n",
        "\n",
        "  model.train()   # Puts model back into training mode\n",
        "  return out\n",
        "\n",
        "\n",
        "class SelfAttentionHead(nn.Module):\n",
        "  \"\"\" A head of self-attention \"\"\"\n",
        "\n",
        "  def __init__(self, head_size):\n",
        "    super().__init__()\n",
        "    self.key = nn.Linear(n_embedding_dimensions, head_size, bias=False)\n",
        "    self.query = nn.Linear(n_embedding_dimensions, head_size, bias=False)\n",
        "    self.value = nn.Linear(n_embedding_dimensions, head_size, bias=False)\n",
        "    self.register_buffer('tril', torch.tril(torch.ones(max_context_length, max_context_length)))\n",
        "    self.dropout = nn.Dropout(dropout_rate)\n",
        "\n",
        "  def forward(self, x):\n",
        "    # Input size: (batch, time-step, channels) a.k.a (batch, tokens, n_embedding_dimensions)\n",
        "    # Output size: (batch, time-step, head_size)\n",
        "    B,T,C = x.shape\n",
        "    k = self.key(x) # (B,T,head_size)     Interprets/linearly projects each token's embedding as a head_size-long 1D tensor. It's basically still an embedding though, just re-interpreted from an old one into a new one.\n",
        "    q = self.query(x) # (B,T,head_size)\n",
        "\n",
        "    # Compute attention scores\n",
        "    scores = q @ k.transpose(-2,-1) * k.shape[-1]**-0.5 # (B,T,head_size) @ (B,head_size,T) -> (B,T,T)     The result of this matrix multiplication shows how the key and query embeddings for each token interact with every other token's embeddings. The values at 1x2 and 2x1 of the result shows how the key and query embeddings of tokens 1 and 2 interact, etc. The k is tranposed so that q and k can be matrix multiplied despite being the same size. The result of each mat-mul is a square though, which is fixed later using the value tensor. The part with k.shape[-1]**-0.5 is used to normalize the output scores so that it doesn't end up as a situation where once we softmax the scores, the probabilities only focus on the highest score and the other tokens aren't used at all. Also remember that k.shape[-1] is just head_size\n",
        "    scores = scores.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # Mask the results of performing mat-mul with future tokens so that only the results of mut-mul with earlier tokens and the current token are taken into consideration when softmaxing each row/channel vector for probabilities.\n",
        "    scores = functional.softmax(scores, dim=-1) # Softmax each individual row/channel vector. This converts masked tokens from -infinity to 0 while also normalizing attention score values to values between 0 and 1.\n",
        "    scores = self.dropout(scores)\n",
        "\n",
        "    # Get output into the proper size again by matrix-multiplying with the value tensor. This step also serves to aggregate the scores into an overall head output with head_size # of values. This SA head's output will be concatenated with other SA head outputs afterwards.\n",
        "    v = self.value(x) # (B,T,head_size)\n",
        "    out = scores @ v # (B,T,T) @ (B,T,head_size) -> (B,T,head_size)\n",
        "    return out\n",
        "\n",
        "\n",
        "class MultiHeadSelfAttention(nn.Module):\n",
        "  \"\"\" Multiple self-attention heads in parallel \"\"\"\n",
        "\n",
        "  def __init__(self, num_heads, head_size):\n",
        "    super().__init__()\n",
        "    self.heads = nn.ModuleList([SelfAttentionHead(head_size) for _ in range(num_heads)])\n",
        "    self.linear_proj = nn.Linear(n_embedding_dimensions, n_embedding_dimensions)\n",
        "    self.dropout = nn.Dropout(dropout_rate)\n",
        "\n",
        "  def forward(self, x):\n",
        "    out = torch.cat([head(x) for head in self.heads], dim=-1) # num_heads * (B,T,head_size) -> (B,T,n_embedding_dimensions)  Concatenate the outputs of all the heads\n",
        "    out = self.linear_proj(out) # (B,T,n_embedding_dimensions)  A linear layer to interpret/linearly project the outputs of the heads into another tensor with the same shape\n",
        "    out = self.dropout(out)\n",
        "    return out\n",
        "\n",
        "\n",
        "class FeedForwardLayer(nn.Module):\n",
        "  \"\"\" Simple multi-layer perceptron \"\"\"\n",
        "\n",
        "  def __init__(self, n_embedding_dimensions):\n",
        "    super().__init__()\n",
        "    self.net = nn.Sequential(\n",
        "        nn.Linear(n_embedding_dimensions, 4 * n_embedding_dimensions), # The authors say to scale up to 4 * n_embedding_dimensions and then scale back down again in the original \"Atttention Is All You Need\" paper\n",
        "        nn.ReLU(),\n",
        "        nn.Linear(4 * n_embedding_dimensions, n_embedding_dimensions),\n",
        "        nn.Dropout(dropout_rate)\n",
        "    )\n",
        "\n",
        "  def forward(self, x):\n",
        "    return self.net(x)\n",
        "\n",
        "\n",
        "class TransformerBlock(nn.Module):\n",
        "  \"\"\" Transformer decoder block \"\"\"\n",
        "\n",
        "  def __init__(self, n_embedding_dimensions, n_heads):\n",
        "    super().__init__()\n",
        "    head_size = n_embedding_dimensions // n_heads\n",
        "    self.self_attention = MultiHeadSelfAttention(n_heads, head_size) # Since we don't use any encoders, we only use this self-attention block and don't have a cross-attention block where we receive input from an encoder block\n",
        "    self.feedforward = FeedForwardLayer(n_embedding_dimensions)\n",
        "    self.ln1 = nn.LayerNorm(n_embedding_dimensions)\n",
        "    self.ln2 = nn.LayerNorm(n_embedding_dimensions)\n",
        "\n",
        "  def forward(self, x):\n",
        "    # We normalize the inputs before each step.\n",
        "    # We also add the original vectors to the results of each step to act as residual/skip connections. This helps backpropagation train weights more efficiently in deep NNs by creating a path for the gradient to flow from end to beginning unimpeded by the SA blocks and FFNNs.\n",
        "    # In GPT models, layer norm is done on inputs instead of outputs\n",
        "    x = x + self.self_attention(self.ln1(x))\n",
        "    x = x + self.feedforward(self.ln2(x))\n",
        "    return x\n",
        "\n",
        "\n",
        "class GPTModel(nn.Module):\n",
        "\n",
        "  def __init__(self):\n",
        "    super().__init__()\n",
        "    self.token_embedding_table = nn.Embedding(vocab_size, n_embedding_dimensions) # Lookup table that stores embeddings for each token in the vocabulary\n",
        "    self.position_embedding_table = nn.Embedding(max_context_length, n_embedding_dimensions) # Lookup table that stores embeddings for each possible position in the context\n",
        "    self.transformer_blocks = nn.Sequential(*[TransformerBlock(n_embedding_dimensions, n_heads=n_heads) for _ in range(n_transformer_blocks)]) # Asterisk is for unpacking the list of transformer blocks into separate arguments\n",
        "    self.ln_final =  nn.LayerNorm(n_embedding_dimensions) # Final layer normalization\n",
        "    self.logit_generation = nn.Linear(n_embedding_dimensions, vocab_size) # We have a linear layer generate logits (raw probability scores) from the outputs of the transformer blocks instead of having the transformer blocks generate logits directly\n",
        "\n",
        "    self.apply(self._init_weights) # Initialize every parameter throughout entire model recursively\n",
        "\n",
        "  def _init_weights(self, module):\n",
        "    if isinstance(module, nn.Linear):\n",
        "      torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
        "      if module.bias is not None:\n",
        "        torch.nn.init.zeros_(module.bias)\n",
        "    elif isinstance(module, nn.Embedding):\n",
        "      torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
        "\n",
        "  def forward(self, input_tokens, targets=None):\n",
        "    B, T = input_tokens.shape # B is Batch, T is Time-step (tokens)\n",
        "\n",
        "    tok_emb = self.token_embedding_table(input_tokens) # (B,T,C) where C is Channel (n_embedding_dimensions)\n",
        "    pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T,n_embedding_dimensions)  The arange produces a tensor of integers from 0 to T-1\n",
        "    x = tok_emb + pos_emb # (B,T,n_embedding_dimensions) This step integrates both types of embeddings by adding pos_emb across all batches\n",
        "    x = self.transformer_blocks(x) # (B,T,n_embedding_dimensions) The output of this stack of transformer blocks is a vector for each individual time step, containing information about the sequence of previous tokens + the current token for that time step. A lot of embedding re-interpretation happens throughout this process.\n",
        "    x = self.ln_final(x) # (B,T,n_embedding_dimensions)\n",
        "    logits = self.logit_generation(x) # (B,T,vocab_size)\n",
        "\n",
        "    if targets is None:\n",
        "      loss = None\n",
        "    else:                               # In this forward function, logits and loss are calculated for all context lengths, including 1, 2, all the way to the max context length so that when the time does come that the # of input tokens is less than the max limit, like just 1 token, 2 tokens, etc., predictions for those situations will have already been trained to be accurate as well\n",
        "      B, T, C = logits.shape\n",
        "      logits = logits.view(B*T, C)      # Removes boundaries between different batches and converts to a 2D tensor showing time-steps on one axis and vocab_size logit probability scores on the other axis\n",
        "      targets = targets.view(B*T)       # Removed boundaries between different batches and converts to a 1D array showing the corresponding target token for a time-step\n",
        "      loss = functional.cross_entropy(logits, targets) # Calculate classification cross entropy loss with \"targets\" as the training classification labels\n",
        "\n",
        "    return logits, loss\n",
        "\n",
        "  def generate(self, input_tokens, max_new_tokens):\n",
        "    for _ in range(max_new_tokens):\n",
        "      cropped_input = input_tokens[:, -max_context_length:] # Only take in max_context_length amount of context for the next token prediction\n",
        "      logits, loss = self(cropped_input) # Get predictions\n",
        "\n",
        "      logits = logits[:, -1, :] # Only retrieve the prediction at the last T position, a.k.a. where all cropped_input tokens are used as context instead of just a subset of those tokens\n",
        "      probabilities = functional.softmax(logits, dim=-1) # (B,C) where C is vocab size\n",
        "      next_token = torch.multinomial(probabilities, num_samples=1) # (B,1)  Sample the next token from the softmaxed probabilities\n",
        "      print(decode([next_token.item()]), end='')\n",
        "      input_tokens = torch.cat((input_tokens, next_token), dim=1) # (B,T+1)  Add the new token to the context for the next token prediction\n",
        "    return input_tokens\n",
        "\n"
      ],
      "metadata": {
        "id": "wSg1luyNTUns"
      },
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Main code"
      ],
      "metadata": {
        "id": "D3Y03yeDFVZ3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# https://www.youtube.com/watch?v=kCc8FmEb1nY\n",
        "\n",
        "#\n",
        "# Configuration variables\n",
        "#\n",
        "\n",
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "max_context_length = 256\n",
        "num_batches = 8 #64\n",
        "max_iters = 50 #5000\n",
        "eval_interval = 2 #500\n",
        "n_eval_iterations = 5 #200\n",
        "learning_rate = 3e-4\n",
        "dropout_rate = 0.2\n",
        "\n",
        "n_embedding_dimensions = 384\n",
        "n_heads = 6\n",
        "n_transformer_blocks = 6"
      ],
      "metadata": {
        "id": "7JC5hzW6FU-S"
      },
      "execution_count": 52,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = GPTModel()\n",
        "m = model.to(device)\n",
        "\n",
        "# Print # of model parameters\n",
        "print(sum(param_tensor.numel() for param_tensor in m.parameters())/1e6, 'M parameters')\n",
        "\n",
        "# Create a PyTorch optimizer\n",
        "optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)"
      ],
      "metadata": {
        "id": "VCEVev39G3BK",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "b78ae7a2-9484-4147-a421-fc16e935e308"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "10.81123 M parameters\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Training loop\n",
        "for iter in range(max_iters):\n",
        "\n",
        "  # Evaluate and print loss for training and validation sets\n",
        "  if iter % eval_interval == 0 or iter == max_iters - 1:\n",
        "    losses = estimate_loss()\n",
        "    print(f'Step {iter}: training loss {losses[\"train\"]:.4f}, validation loss {losses[\"val\"]:.4f}')\n",
        "\n",
        "  # Sample a batch of data\n",
        "  input_tokens, target_tokens = get_batch('train')\n",
        "\n",
        "  # Evaluate loss and optimize\n",
        "  logits, loss = model(input_tokens, target_tokens)\n",
        "  optimizer.zero_grad(set_to_none=True)\n",
        "  loss.backward()\n",
        "  optimizer.step()"
      ],
      "metadata": {
        "id": "9RX6vlLlIvyX",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "a2f46859-8fec-4f28-d844-5657995d7535"
      },
      "execution_count": 53,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Step 0: training loss 1.5597, validation loss 1.5756\n",
            "Step 4: training loss 1.5305, validation loss 1.5584\n",
            "Step 6: training loss 1.5560, validation loss 1.5222\n",
            "Step 8: training loss 1.4930, validation loss 1.5061\n",
            "Step 10: training loss 1.5256, validation loss 1.5923\n",
            "Step 12: training loss 1.5177, validation loss 1.5522\n",
            "Step 14: training loss 1.5247, validation loss 1.6094\n",
            "Step 16: training loss 1.5981, validation loss 1.5525\n",
            "Step 18: training loss 1.5296, validation loss 1.5539\n",
            "Step 20: training loss 1.5719, validation loss 1.5680\n",
            "Step 22: training loss 1.6042, validation loss 1.5507\n",
            "Step 24: training loss 1.5188, validation loss 1.5145\n",
            "Step 26: training loss 1.5287, validation loss 1.5669\n",
            "Step 28: training loss 1.5280, validation loss 1.5436\n",
            "Step 30: training loss 1.5294, validation loss 1.5535\n",
            "Step 32: training loss 1.6144, validation loss 1.5928\n",
            "Step 34: training loss 1.5360, validation loss 1.5870\n",
            "Step 36: training loss 1.4809, validation loss 1.6149\n",
            "Step 38: training loss 1.5685, validation loss 1.6064\n",
            "Step 40: training loss 1.5690, validation loss 1.5534\n",
            "Step 42: training loss 1.5189, validation loss 1.5993\n",
            "Step 44: training loss 1.5215, validation loss 1.6176\n",
            "Step 46: training loss 1.5614, validation loss 1.6204\n",
            "Step 48: training loss 1.5584, validation loss 1.5821\n",
            "Step 49: training loss 1.5483, validation loss 1.5361\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Generate\n",
        "context = torch.ones((1, 1), dtype=torch.long, device=device)\n",
        "m.generate(context, max_new_tokens=100)                           #print(decode(m.generate(context, max_new_tokens=100)[0].tolist()))\n",
        "print()"
      ],
      "metadata": {
        "id": "CijmMufxYv35",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "6bdfc76b-2447-4e31-9ab9-de9b389d5cc0"
      },
      "execution_count": 63,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "This the sep and then make life at we do we wain't.\n",
            "I fruit. imptic.\n",
            "Oh, did you know like you shooo\n"
          ]
        }
      ]
    }
  ]
}