NLP101_DaRMoD_demo.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/patrickabadi/python-pytorch-notebooks/blob/main/NLP101_DaRMoD_demo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TqENYWPfkj2Q"
      },
      "source": [
        "## Importing a LLM : GPT2-XL\n",
        "\n",
        "**Change the runtime type to a GPU before downloading gpt2**\n",
        "\n",
        "**Note: We are running a quantized version (8-bit integer) of the model so that we can fit it into a free google colab isntance (T4) . See https://huggingface.co/docs/transformers/main_classes/quantization and https://colab.research.google.com/drive/1qOjXfQIAULfKvZqwCen8-MoWKGdSatZ4#scrollTo=W8tQtyjp75O for more details**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zz2K5lLGkYek"
      },
      "outputs": [],
      "source": [
        "!mkdir -p /content/llm/gpt2-xl/\n",
        "llm_dir = \"/content/llm/gpt2-xl/\"\n",
        "\n",
        "!export LC_ALL=C.UTF-8 # So that gdown works for data download\n",
        "!export LANG=C.UTF-8"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fBDVcKm3mxbA",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "e41dbc49-7240-4750-8b3d-5a7db923cd06"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--2023-07-11 14:55:51--  https://huggingface.co/gpt2-xl/resolve/main/pytorch_model.bin\n",
            "Resolving huggingface.co (huggingface.co)... 18.172.134.88, 18.172.134.124, 18.172.134.4, ...\n",
            "Connecting to huggingface.co (huggingface.co)|18.172.134.88|:443... connected.\n",
            "HTTP request sent, awaiting response... 302 Found\n",
            "Location: https://cdn-lfs.huggingface.co/gpt2-xl/cd2a29e31040ef64d9362cb96801969c9f67b9e0bdbd6e00b9dda57cdbe17435?response-content-disposition=attachment%3B+filename*%3DUTF-8%27%27pytorch_model.bin%3B+filename%3D%22pytorch_model.bin%22%3B&response-content-type=application%2Foctet-stream&Expires=1689346551&Policy=eyJTdGF0ZW1lbnQiOlt7IkNvbmRpdGlvbiI6eyJEYXRlTGVzc1RoYW4iOnsiQVdTOkVwb2NoVGltZSI6MTY4OTM0NjU1MX19LCJSZXNvdXJjZSI6Imh0dHBzOi8vY2RuLWxmcy5odWdnaW5nZmFjZS5jby9ncHQyLXhsL2NkMmEyOWUzMTA0MGVmNjRkOTM2MmNiOTY4MDE5NjljOWY2N2I5ZTBiZGJkNmUwMGI5ZGRhNTdjZGJlMTc0MzU%7EcmVzcG9uc2UtY29udGVudC1kaXNwb3NpdGlvbj0qJnJlc3BvbnNlLWNvbnRlbnQtdHlwZT0qIn1dfQ__&Signature=wYQ0PL5OZVnTQBemUijX6bbD%7ESHvdytiHHRvxEcPAvZoRoSV-Jjh5rR3kfUH8kq%7EntoiQf7gJ%7E1d42zcBODuJluoZ0YCjQ3LkvcP13ocVG9EgcCELJuuRSZJZ2euelEqDiMMxthl5jzr-2aarOd6Su-Nmu1tLl%7EK-oMbbvkKOjDZnqzAOaxfOHDVoP9t3xswWXLxsNRIbzJ4pzMA3vax596dZ62gNG4KASkz00D3zk9ftZ9iZ2LvugYUEbxbyPIAEhBjIekxkGBZ-AmKD-oRyjpGOoP6f3YEkUY6IJCXEkojmjSxP9nY156NPG5K0E7i8Ss09tMZ6bswn24KQiS86Q__&Key-Pair-Id=KVTP0A1DKRTAX [following]\n",
            "--2023-07-11 14:55:51--  https://cdn-lfs.huggingface.co/gpt2-xl/cd2a29e31040ef64d9362cb96801969c9f67b9e0bdbd6e00b9dda57cdbe17435?response-content-disposition=attachment%3B+filename*%3DUTF-8%27%27pytorch_model.bin%3B+filename%3D%22pytorch_model.bin%22%3B&response-content-type=application%2Foctet-stream&Expires=1689346551&Policy=eyJTdGF0ZW1lbnQiOlt7IkNvbmRpdGlvbiI6eyJEYXRlTGVzc1RoYW4iOnsiQVdTOkVwb2NoVGltZSI6MTY4OTM0NjU1MX19LCJSZXNvdXJjZSI6Imh0dHBzOi8vY2RuLWxmcy5odWdnaW5nZmFjZS5jby9ncHQyLXhsL2NkMmEyOWUzMTA0MGVmNjRkOTM2MmNiOTY4MDE5NjljOWY2N2I5ZTBiZGJkNmUwMGI5ZGRhNTdjZGJlMTc0MzU%7EcmVzcG9uc2UtY29udGVudC1kaXNwb3NpdGlvbj0qJnJlc3BvbnNlLWNvbnRlbnQtdHlwZT0qIn1dfQ__&Signature=wYQ0PL5OZVnTQBemUijX6bbD%7ESHvdytiHHRvxEcPAvZoRoSV-Jjh5rR3kfUH8kq%7EntoiQf7gJ%7E1d42zcBODuJluoZ0YCjQ3LkvcP13ocVG9EgcCELJuuRSZJZ2euelEqDiMMxthl5jzr-2aarOd6Su-Nmu1tLl%7EK-oMbbvkKOjDZnqzAOaxfOHDVoP9t3xswWXLxsNRIbzJ4pzMA3vax596dZ62gNG4KASkz00D3zk9ftZ9iZ2LvugYUEbxbyPIAEhBjIekxkGBZ-AmKD-oRyjpGOoP6f3YEkUY6IJCXEkojmjSxP9nY156NPG5K0E7i8Ss09tMZ6bswn24KQiS86Q__&Key-Pair-Id=KVTP0A1DKRTAX\n",
            "Resolving cdn-lfs.huggingface.co (cdn-lfs.huggingface.co)... 18.154.185.27, 18.154.185.94, 18.154.185.26, ...\n",
            "Connecting to cdn-lfs.huggingface.co (cdn-lfs.huggingface.co)|18.154.185.27|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 6431878936 (6.0G) [application/octet-stream]\n",
            "Saving to: ‘/content/llm/gpt2-xl/pytorch_model.bin’\n",
            "\n",
            "pytorch_model.bin   100%[===================>]   5.99G   184MB/s    in 32s     \n",
            "\n",
            "2023-07-11 14:56:23 (191 MB/s) - ‘/content/llm/gpt2-xl/pytorch_model.bin’ saved [6431878936/6431878936]\n",
            "\n",
            "--2023-07-11 14:56:23--  https://huggingface.co/gpt2-xl/resolve/main/config.json\n",
            "Resolving huggingface.co (huggingface.co)... 99.84.160.57, 99.84.160.43, 99.84.160.9, ...\n",
            "Connecting to huggingface.co (huggingface.co)|99.84.160.57|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 689 [text/plain]\n",
            "Saving to: ‘/content/llm/gpt2-xl/config.json’\n",
            "\n",
            "config.json         100%[===================>]     689  --.-KB/s    in 0s      \n",
            "\n",
            "2023-07-11 14:56:23 (148 MB/s) - ‘/content/llm/gpt2-xl/config.json’ saved [689/689]\n",
            "\n",
            "--2023-07-11 14:56:24--  https://huggingface.co/gpt2-xl/resolve/main/merges.txt\n",
            "Resolving huggingface.co (huggingface.co)... 99.84.160.57, 99.84.160.43, 99.84.160.9, ...\n",
            "Connecting to huggingface.co (huggingface.co)|99.84.160.57|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 456318 (446K) [text/plain]\n",
            "Saving to: ‘/content/llm/gpt2-xl/merges.txt’\n",
            "\n",
            "merges.txt          100%[===================>] 445.62K  --.-KB/s    in 0.05s   \n",
            "\n",
            "2023-07-11 14:56:24 (8.96 MB/s) - ‘/content/llm/gpt2-xl/merges.txt’ saved [456318/456318]\n",
            "\n",
            "--2023-07-11 14:56:24--  https://huggingface.co/gpt2-xl/resolve/main/tokenizer.json\n",
            "Resolving huggingface.co (huggingface.co)... 99.84.160.57, 99.84.160.43, 99.84.160.9, ...\n",
            "Connecting to huggingface.co (huggingface.co)|99.84.160.57|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 1355256 (1.3M) [text/plain]\n",
            "Saving to: ‘/content/llm/gpt2-xl/tokenizer.json’\n",
            "\n",
            "tokenizer.json      100%[===================>]   1.29M  --.-KB/s    in 0.1s    \n",
            "\n",
            "2023-07-11 14:56:24 (12.1 MB/s) - ‘/content/llm/gpt2-xl/tokenizer.json’ saved [1355256/1355256]\n",
            "\n",
            "--2023-07-11 14:56:24--  https://huggingface.co/gpt2-xl/resolve/main/vocab.json\n",
            "Resolving huggingface.co (huggingface.co)... 99.84.160.57, 99.84.160.43, 99.84.160.9, ...\n",
            "Connecting to huggingface.co (huggingface.co)|99.84.160.57|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 1042301 (1018K) [text/plain]\n",
            "Saving to: ‘/content/llm/gpt2-xl/vocab.json’\n",
            "\n",
            "vocab.json          100%[===================>]   1018K  --.-KB/s    in 0.05s   \n",
            "\n",
            "2023-07-11 14:56:24 (18.3 MB/s) - ‘/content/llm/gpt2-xl/vocab.json’ saved [1042301/1042301]\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# Download the required files from HugginFace (https://huggingface.co/hivemind/gpt-j-6B-8bit/tree/main):\n",
        "# Possible ETA ~ 10 min for the 6G model; usually < 1 min\n",
        "!wget -P /content/llm/gpt2-xl/ https://huggingface.co/gpt2-xl/resolve/main/pytorch_model.bin\n",
        "!wget -P /content/llm/gpt2-xl/ https://huggingface.co/gpt2-xl/resolve/main/config.json\n",
        "!wget -P /content/llm/gpt2-xl/ https://huggingface.co/gpt2-xl/resolve/main/merges.txt\n",
        "!wget -P /content/llm/gpt2-xl/ https://huggingface.co/gpt2-xl/resolve/main/tokenizer.json\n",
        "!wget -P /content/llm/gpt2-xl/ https://huggingface.co/gpt2-xl/resolve/main/vocab.json"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "06hXae-coaEl",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "711016af-df5a-4b9b-c06f-a97bb831a8b0"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.2/7.2 MB\u001b[0m \u001b[31m54.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m268.8/268.8 kB\u001b[0m \u001b[31m27.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.8/7.8 MB\u001b[0m \u001b[31m105.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m76.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m227.6/227.6 kB\u001b[0m \u001b[31m6.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m101.8/101.8 MB\u001b[0m \u001b[31m8.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ],
      "source": [
        "# Install HF transformers library:\n",
        "!pip -q install transformers\n",
        "# Accelerate and bitsandbytes for 8bit models version\n",
        "!pip -q install accelerate\n",
        "!pip -q install bitsandbytes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NNivSsBQoC7L",
        "outputId": "241f7148-7ece-404f-e59f-67dc5bf8c596"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "===================================BUG REPORT===================================\n",
            "Welcome to bitsandbytes. For bug reports, please run\n",
            "\n",
            "python -m bitsandbytes\n",
            "\n",
            " and submit this information together with your error trace to: https://github.com/TimDettmers/bitsandbytes/issues\n",
            "================================================================================\n",
            "bin /usr/local/lib/python3.10/dist-packages/bitsandbytes/libbitsandbytes_cuda118.so\n",
            "CUDA_SETUP: WARNING! libcudart.so not found in any environmental path. Searching in backup paths...\n",
            "CUDA SETUP: CUDA runtime path found: /usr/local/cuda/lib64/libcudart.so\n",
            "CUDA SETUP: Highest compute capability among GPUs detected: 7.5\n",
            "CUDA SETUP: Detected CUDA version 118\n",
            "CUDA SETUP: Loading binary /usr/local/lib/python3.10/dist-packages/bitsandbytes/libbitsandbytes_cuda118.so...\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/bitsandbytes/cuda_setup/main.py:149: UserWarning: /usr/lib64-nvidia did not contain ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] as expected! Searching further paths...\n",
            "  warn(msg)\n",
            "/usr/local/lib/python3.10/dist-packages/bitsandbytes/cuda_setup/main.py:149: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('/sys/fs/cgroup/memory.events /var/colab/cgroup/jupyter-children/memory.events')}\n",
            "  warn(msg)\n",
            "/usr/local/lib/python3.10/dist-packages/bitsandbytes/cuda_setup/main.py:149: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('http'), PosixPath('//172.28.0.1'), PosixPath('8013')}\n",
            "  warn(msg)\n",
            "/usr/local/lib/python3.10/dist-packages/bitsandbytes/cuda_setup/main.py:149: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('--logtostderr --listen_host=172.28.0.12 --target_host=172.28.0.12 --tunnel_background_save_url=https'), PosixPath('//colab.research.google.com/tun/m/cc48301118ce562b961b3c22d803539adc1e0c19/gpu-t4-s-1fxadzwgtftv9 --tunnel_background_save_delay=10s --tunnel_periodic_background_save_frequency=30m0s --enable_output_coalescing=true --output_coalescing_required=true')}\n",
            "  warn(msg)\n",
            "/usr/local/lib/python3.10/dist-packages/bitsandbytes/cuda_setup/main.py:149: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('/env/python')}\n",
            "  warn(msg)\n",
            "/usr/local/lib/python3.10/dist-packages/bitsandbytes/cuda_setup/main.py:149: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('module'), PosixPath('//ipykernel.pylab.backend_inline')}\n",
            "  warn(msg)\n",
            "/usr/local/lib/python3.10/dist-packages/bitsandbytes/cuda_setup/main.py:149: UserWarning: Found duplicate ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] files: {PosixPath('/usr/local/cuda/lib64/libcudart.so'), PosixPath('/usr/local/cuda/lib64/libcudart.so.11.0')}.. We'll flip a coin and try one of these, in order to fail forward.\n",
            "Either way, this might cause trouble in the future:\n",
            "If you get `CUDA error: invalid device function` errors, the above might be the cause and the solution is to make sure only one ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] in the paths that we search based on your env.\n",
            "  warn(msg)\n",
            "You are loading your model in 8bit or 4bit but no linear modules were found in your model. this can happen for some architectures such as gpt2 that uses Conv1D instead of Linear layers. Please double check your model architecture, or submit an issue on github if you think this is a bug.\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "GPT2LMHeadModel(\n",
              "  (transformer): GPT2Model(\n",
              "    (wte): Embedding(50257, 1600)\n",
              "    (wpe): Embedding(1024, 1600)\n",
              "    (drop): Dropout(p=0.1, inplace=False)\n",
              "    (h): ModuleList(\n",
              "      (0-47): 48 x GPT2Block(\n",
              "        (ln_1): LayerNorm((1600,), eps=1e-05, elementwise_affine=True)\n",
              "        (attn): GPT2Attention(\n",
              "          (c_attn): Conv1D()\n",
              "          (c_proj): Conv1D()\n",
              "          (attn_dropout): Dropout(p=0.1, inplace=False)\n",
              "          (resid_dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "        (ln_2): LayerNorm((1600,), eps=1e-05, elementwise_affine=True)\n",
              "        (mlp): GPT2MLP(\n",
              "          (c_fc): Conv1D()\n",
              "          (c_proj): Conv1D()\n",
              "          (act): NewGELUActivation()\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "    )\n",
              "    (ln_f): LayerNorm((1600,), eps=1e-05, elementwise_affine=True)\n",
              "  )\n",
              "  (lm_head): Linear(in_features=1600, out_features=50257, bias=False)\n",
              ")"
            ]
          },
          "metadata": {},
          "execution_count": 4
        }
      ],
      "source": [
        "# Mount the model\n",
        "import json\n",
        "import random\n",
        "from datetime import datetime\n",
        "from time import sleep\n",
        "import logging\n",
        "import argparse\n",
        "from tqdm.notebook import tqdm\n",
        "import csv\n",
        "import os\n",
        "\n",
        "import torch\n",
        "from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM\n",
        "\n",
        "llm_dir = \"/content/llm/gpt2-xl/\"\n",
        "\n",
        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(llm_dir, use_fast=False)\n",
        "# set pad token ids for batched inference cus gpt2 does not have one\n",
        "tokenizer.padding_side = \"left\"\n",
        "tokenizer.pad_token = tokenizer.eos_token\n",
        "tokenizer.pad_token_id = tokenizer.eos_token_id\n",
        "model_config = AutoConfig.from_pretrained(llm_dir)\n",
        "model = AutoModelForCausalLM.from_pretrained(llm_dir, load_in_8bit=True) # Load 8bit int model to fit in Gcolab memory\n",
        "#model.to(device)\n",
        "model.eval()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FfQf-Bv8pUXu",
        "outputId": "44dd8342-3a98-4ef5-9de1-737a016fdde8"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
            "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Once upon a time, a venture founder was taking a class about ML where he learned about the Data Mining library. He didn't have a big enough data set for a data mining task and he was curious about how to get started. He knew that the only way to do data mining was to do it in python. He also knew that data mining would be a relatively slow process. So that's what he did!\n",
            "\n",
            "The data was not very large but it was enough for him to learn the basics of data mining. But he had no idea how to scale it up.\n",
            "\n",
            "After a couple of months of coding, he had some ideas of how to scale his python data mining job and he wrote a blog post about it. The post got a lot of traction and that's what brought him to the idea of Data Science jobs. Here's what he writes:\n",
            "\n",
            "\"After a couple of months of coding, I had some ideas of how to scale my python data mining job\n"
          ]
        }
      ],
      "source": [
        "# Test and play a little bit with gpt2-XL\n",
        "\n",
        "# Define the input text\n",
        "input_text = \"Once upon a time, a venture founder was taking a class about ML where\" # Feel free to modify!\n",
        "\n",
        "# Tokenize the input text\n",
        "input_ids = tokenizer.encode(input_text, return_tensors=\"pt\")\n",
        "\n",
        "# Move the input IDs to the same device as the model\n",
        "input_ids = input_ids.to(device)\n",
        "\n",
        "# Generate text from the model\n",
        "output = model.generate(input_ids, max_length=200, do_sample=True, temperature=0.7)  # Can play with max_length of tokens to generate smaller/longer outputs\n",
        "\n",
        "# Decode the output IDs to text\n",
        "output_text = tokenizer.decode(output[0], skip_special_tokens=True)\n",
        "\n",
        "print(output_text)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5iWJmUlfk0I2"
      },
      "source": [
        "## Basic Prompt Engineering (In-Context Learning (ICL)) and KNN Prompting"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0DOxsQbrk551"
      },
      "outputs": [],
      "source": [
        "# Download some test data:\n",
        "!gdown --id 1Yh2blPkJvMtdm5xWKoHr2fLp2i2Bn5Ir\n",
        "!unzip data.zip"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6t83AZjcDkxT"
      },
      "outputs": [],
      "source": [
        "# Above fails from time to time ... Can try to restart runtime and run:\n",
        "!export LC_ALL=C.UTF-8 # So that gdown works for data download\n",
        "!export LANG=C.UTF-8"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LodGmLbQxcls"
      },
      "outputs": [],
      "source": [
        "# This colab is extensively reusing stuff from https://github.com/BenfengXu/KNNPrompting/tree/main (https://openreview.net/pdf?id=fe2S7736sNS)\n",
        "# Let's install the repo locally\n",
        "!git clone https://github.com/BenfengXu/KNNPrompting.git\n",
        "import sys\n",
        "sys.path.append('/content/KNNPrompting')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c1yRayaTHusW"
      },
      "source": [
        "Note: The cell below, that buils the embeddings to use with knn classifier, takes about ~13 min to execute on a T4 instance."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_952dhvqvGpv"
      },
      "outputs": [],
      "source": [
        "from utils.dataset import *\n",
        "from utils.anchor import AnchorStore\n",
        "from utils.template import *\n",
        "\n",
        "# We will play with sst2 dataset: https://huggingface.co/datasets/sst2\n",
        "datadir = \"/content/data/sst2/\"\n",
        "AutoDataset = SST2Dataset\n",
        "\n",
        "train_data = AutoDataset(datadir, mode='train')\n",
        "dev_data = AutoDataset(datadir, mode='dev')\n",
        "\n",
        "anchor_data = AutoDataset(datadir, mode='train')\n",
        "\n",
        "knn=3 # number of k nearest neighbord to look for classification\n",
        "max_context_len = 1024\n",
        "n_demo_shot = 32\n",
        "n_anchor_shot = max_context_len - n_demo_shot\n",
        "seed = 43\n",
        "\n",
        "def llm_gen(model, prompt, tokenizer, max_context_len):\n",
        "    inputs = tokenizer.encode_plus(prompt, return_tensors=\"pt\", padding=True).to(device=model.device)\n",
        "    if inputs['input_ids'].shape[1] > max_context_len:\n",
        "        inputs['input_ids'] = inputs['input_ids'][:, -max_context_len:]\n",
        "        inputs['attention_mask'] = inputs['attention_mask'][:, -max_context_len:]\n",
        "    with torch.no_grad():\n",
        "        logits = model.forward(input_ids=inputs['input_ids'],\n",
        "                               attention_mask=inputs['attention_mask'],\n",
        "                               return_dict=True).logits.detach().cpu()\n",
        "    # the output prob is shifted by -1, so we should use the output at the last input token position\n",
        "    # gen_logits.shape = [1, 50257]\n",
        "    gen_logits = logits[:, -1, :]\n",
        "\n",
        "    return gen_logits\n",
        "\n",
        "# Stage1: Meta Test -> This populates the anchor_store datastore with the target embeddings\n",
        "train_data.subsamplebyshot(n_demo_shot, seed) # (demo_shots, seed)\n",
        "prompt_prefix = make_prompt(train_data, \"sst2\", mode='train')\n",
        "anchor_data.subsamplebyshot(n_anchor_shot, seed, exclude=train_data.data)\n",
        "label2id = dev_data.label2id\n",
        "id2verb = train_data.id2verb\n",
        "anchor_store = AnchorStore(K=anchor_data.__len__(),\n",
        "                            dim=model_config.vocab_size,\n",
        "                            knn=knn,\n",
        "                            n_class=len(label2id))\n",
        "for ins in tqdm(anchor_data.data, total=anchor_data.__len__()):  # This could be parallelize on GPU; left as an exercise ;-)\n",
        "    labels = label2id[ins['label']]\n",
        "    prompt = prompt_prefix + make_prompt(ins, 'sst2', mode='inference')\n",
        "    gen_logits = llm_gen(model, prompt, tokenizer, max_context_len)\n",
        "    anchor_store.enqueue(torch.softmax(gen_logits.float(), dim=-1), torch.tensor(labels))"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# That took a long time to generate the datastore! Let's pickle it for later usage\n",
        "import pickle\n",
        "\n",
        "with open('anchor_store.pkl', 'wb') as f:\n",
        "    pickle.dump(anchor_store, f)"
      ],
      "metadata": {
        "id": "WzeqVDAjkq5J"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# To load it back:\n",
        "with open('anchor_store.pkl', 'rb') as f:\n",
        "    anchor_store = pickle.load(f)"
      ],
      "metadata": {
        "id": "UShnfVSGl9zp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Can explore the dataset\n",
        "\n",
        "# Take data instance you want by modifying index\n",
        "index=1\n",
        "example_ins = dev_data.data[index]\n",
        "\n",
        "# Make an example label\n",
        "example_label = label2id[example_ins['label']]\n",
        "\n",
        "# Make query part of the example prompt\n",
        "example_prompt_query = make_prompt(example_ins, 'sst2', mode='inference')\n",
        "\n",
        "# Append context for full example prompt (for ICL)\n",
        "example_prompt = prompt_prefix + example_prompt_query\n",
        "\n",
        "print(f\"Example to classify: \\n{example_prompt_query}\")\n",
        "print()\n",
        "print(f\"Example true label: {example_label} (1 for positive / 0 for negative)\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HdT1hV9M2D8e",
        "outputId": "ff31b460-cb4c-42c0-e2d5-bb9dd700aa21"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Example to classify: \n",
            "Review: it 's worth seeing just on the basis of the wisdom , and at times , the startling optimism , of the children . \n",
            "Sentiment:\n",
            "\n",
            "Example true label: 1 (1 for positive / 0 for negative)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Let's compare vanilla prompting to ICL"
      ],
      "metadata": {
        "id": "H9-_B-mJt9bi"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Let's see what basic prompting gives as an answer:\n",
        "\n",
        "# Make the basic prompt\n",
        "base_prompt = \"Is the following review expressing a positive or negative sentiment? \\n\" + make_prompt(example_ins, 'sst2', mode='inference')\n",
        "\n",
        "inputs = tokenizer.encode_plus(base_prompt, return_tensors=\"pt\", padding=True).to(device=model.device)\n",
        "if inputs['input_ids'].shape[1] > max_context_len:\n",
        "        inputs['input_ids'] = inputs['input_ids'][:, -max_context_len:]\n",
        "        inputs['attention_mask'] = inputs['attention_mask'][:, -max_context_len:]\n",
        "output = model.generate(input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'], max_length=1, do_sample=False, num_beams=1)  # Can play with max_length of tokens to generate smaller/longer outputs\n",
        "output_text = tokenizer.decode(output[0], skip_special_tokens=True)\n",
        "\n",
        "print()\n",
        "print(\"model output: \\n\" + output_text)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Zml5CA8AtNp2",
        "outputId": "089892d0-63b9-438f-a1ae-7093b1500539"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n",
            "Input length of input_ids is 45, but `max_length` is set to 1. This can lead to unexpected behavior. You should consider increasing `max_new_tokens`.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "model output: \n",
            "Is the following review expressing a positive or negative sentiment? \n",
            "Review: it's worth seeing just on the basis of the wisdom, and at times, the startling optimism, of the children. \n",
            "Sentiment: \n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NpB6auw_wq4h",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "44c9eb18-35f7-4c28-c3db-2c00aea42d7e"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n",
            "Input length of input_ids is 1024, but `max_length` is set to 1. This can lead to unexpected behavior. You should consider increasing `max_new_tokens`.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "#################\n",
            "Model answer\n",
            "#################\n",
            "\n",
            "model output:  positive\n",
            "\n",
            "Review: '' is a sweet, honest, and enjoyable comedy-drama about a young woman who wants many things in life, but fears she 'll become her mother before she gets to fulfill her dreams. \n",
            "Sentiment: positive\n",
            "\n",
            "Review: inane and unimaginative \n",
            "Sentiment: negative\n",
            "\n",
            "Review: and not in a good way \n",
            "Sentiment: negative\n",
            "\n",
            "Review: personal low \n",
            "Sentiment: negative\n",
            "\n",
            "Review: '' has the right stuff for silly summer entertainment and has enough laughs to sustain interest to the end. \n",
            "Sentiment: positive\n",
            "\n",
            "Review: close to losing my lunch \n",
            "Sentiment: negative\n",
            "\n",
            "Review: memorable zingers \n",
            "Sentiment: positive\n",
            "\n",
            "Review: a graceful, moving tribute to the courage of new york's finest and a nicely understated expression of the grief \n",
            "Sentiment: positive\n",
            "\n",
            "Review: cheesy b-movie playing \n",
            "Sentiment: negative\n",
            "\n",
            "Review: that would make it the darling of many a kids-and-family-oriented cable channel \n",
            "Sentiment: positive\n",
            "\n",
            "Review: wasted. \n",
            "Sentiment: negative\n",
            "\n",
            "Review: classic moral-condundrum drama : what would you have done to survive? \n",
            "Sentiment: negative\n",
            "\n",
            "Review: cliched dialogue and perverse escapism \n",
            "Sentiment: negative\n",
            "\n",
            "Review: that the material is so second-rate \n",
            "Sentiment: negative\n",
            "\n",
            "Review:, enigma offers all the pleasure of a handsome and well-made entertainment. \n",
            "Sentiment: positive\n",
            "\n",
            "Review: there are some laughs in this movie \n",
            "Sentiment: positive\n",
            "\n",
            "Review: pass without reminding audiences that it's only a movie \n",
            "Sentiment: negative\n",
            "\n",
            "Review: to make the formula feel fresh \n",
            "Sentiment: positive\n",
            "\n",
            "Review: if you don't understand what on earth is going on \n",
            "Sentiment: negative\n",
            "\n",
            "Review: started with a great premise \n",
            "Sentiment: positive\n",
            "\n",
            "Review: noise, mayhem and stupidity \n",
            "Sentiment: negative\n",
            "\n",
            "Review: assume he had a bad run in the market or a costly divorce, because there is no earthly reason other than money why this distinguished actor would stoop so low \n",
            "Sentiment: negative\n",
            "\n",
            "Review: a convincing brogue \n",
            "Sentiment: positive\n",
            "\n",
            "Review: make this delicate coming-of-age tale a treat \n",
            "Sentiment: positive\n",
            "\n",
            "Review: cut their losses \n",
            "Sentiment: negative\n",
            "\n",
            "Review:, it's just grating \n",
            "Sentiment: negative\n",
            "\n",
            "Review: this is a film living far too much in its own head. \n",
            "Sentiment: negative\n",
            "\n",
            "Review: the only thing avary seems to care about \n",
            "Sentiment: negative\n",
            "\n",
            "Review: a funny and touching film that is gorgeously acted by a british cast to rival gosford park's. \n",
            "Sentiment: positive\n",
            "\n",
            "Review: out-of-kilter \n",
            "Sentiment: negative\n",
            "\n",
            "Review: a term paper \n",
            "Sentiment: negative\n",
            "\n",
            "Review:, the wild thornberrys movie doesn't offer much more than the series \n",
            "Sentiment: negative\n",
            "\n",
            "Review: have been kind enough to share it \n",
            "Sentiment: positive\n",
            "\n",
            "Review: a striking style \n",
            "Sentiment: positive\n",
            "\n",
            "Review: is just too bad the film's story does not live up to its style \n",
            "Sentiment: negative\n",
            "\n",
            "Review: literate \n",
            "Sentiment: positive\n",
            "\n",
            "Review: juiced \n",
            "Sentiment: positive\n",
            "\n",
            "Review: show blind date, only less \n",
            "Sentiment: negative\n",
            "\n",
            "Review: wry humor \n",
            "Sentiment: positive\n",
            "\n",
            "Review: viewers are asked so often to suspend belief that were it not for holm's performance, the film would be a total washout. \n",
            "Sentiment: negative\n",
            "\n",
            "Review: hitting below the belt \n",
            "Sentiment: negative\n",
            "\n",
            "Review: piece to watch with kids and use to introduce video as art \n",
            "Sentiment: positive\n",
            "\n",
            "Review: moved to the edge of their seats by the dynamic first act \n",
            "Sentiment: positive\n",
            "\n",
            "Review: irritating films \n",
            "Sentiment: negative\n",
            "\n",
            "Review: uninvolving storytelling \n",
            "Sentiment: negative\n",
            "\n",
            "Review:... is at once playful and haunting, an in-depth portrait of an iconoclastic artist who was fundamentally unknowable even to his closest friends. \n",
            "Sentiment: positive\n",
            "\n",
            "Review: the dark theater \n",
            "Sentiment: negative\n",
            "\n",
            "Review: chabrol's most intense psychological mysteries \n",
            "Sentiment: positive\n",
            "\n",
            "Review: less than adorable \n",
            "Sentiment: negative\n",
            "\n",
            "Review: it's worth seeing just on the basis of the wisdom, and at times, the startling optimism, of the children. \n",
            "Sentiment: positive\n"
          ]
        }
      ],
      "source": [
        "# Can we do better with ICL ?\n",
        "print(\"#################\")\n",
        "print('Model answer')\n",
        "print(\"#################\")\n",
        "inputs = tokenizer.encode_plus(example_prompt, return_tensors=\"pt\", padding=True).to(device=model.device)\n",
        "if inputs['input_ids'].shape[1] > max_context_len:\n",
        "        inputs['input_ids'] = inputs['input_ids'][:, -max_context_len:]\n",
        "        inputs['attention_mask'] = inputs['attention_mask'][:, -max_context_len:]\n",
        "output = model.generate(input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'], max_length=1, do_sample=False, num_beams=1)  # Can play with max_length of tokens to generate smaller/longer outputs\n",
        "output_text = tokenizer.decode(output[0], skip_special_tokens=True)\n",
        "\n",
        "print()\n",
        "print(\"model output:\", output_text)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### What can we do with KNN prompting?"
      ],
      "metadata": {
        "id": "sa5GBfp7uMBo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Train a KNN classifier with the anchors embeddings, using scikit-learn.\n",
        "import numpy as np\n",
        "from sklearn.neighbors import KNeighborsClassifier\n",
        "\n",
        "def kl_divergence(p, q):\n",
        "    epsilon = np.finfo(float).eps # Add a small constant to avoid division by zero or taking log(0)\n",
        "\n",
        "    p += epsilon\n",
        "    q += epsilon\n",
        "    return np.mean(p * (np.log(p) - np.log(q)))\n",
        "\n",
        "# Instantiate the knn classifier with kl as a metric\n",
        "knn_classifier = KNeighborsClassifier(n_neighbors=knn, metric=kl_divergence)\n",
        "\n",
        "# Fit the model to the data in the datastore\n",
        "knn_classifier.fit(anchor_store.queue_anchor.cpu().numpy(), anchor_store.queue_label.cpu().numpy())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 92
        },
        "id": "C8VyzIqWbqum",
        "outputId": "7af8ef77-bb7a-49d5-e18d-d04d0fa2eb99"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "KNeighborsClassifier(metric=<function kl_divergence at 0x7f0b2b4b71c0>,\n",
              "                     n_neighbors=3)"
            ],
            "text/html": [
              "<style>#sk-container-id-1 {color: black;background-color: white;}#sk-container-id-1 pre{padding: 0;}#sk-container-id-1 div.sk-toggleable {background-color: white;}#sk-container-id-1 label.sk-toggleable__label {cursor: pointer;display: block;width: 100%;margin-bottom: 0;padding: 0.3em;box-sizing: border-box;text-align: center;}#sk-container-id-1 label.sk-toggleable__label-arrow:before {content: \"▸\";float: left;margin-right: 0.25em;color: #696969;}#sk-container-id-1 label.sk-toggleable__label-arrow:hover:before {color: black;}#sk-container-id-1 div.sk-estimator:hover label.sk-toggleable__label-arrow:before {color: black;}#sk-container-id-1 div.sk-toggleable__content {max-height: 0;max-width: 0;overflow: hidden;text-align: left;background-color: #f0f8ff;}#sk-container-id-1 div.sk-toggleable__content pre {margin: 0.2em;color: black;border-radius: 0.25em;background-color: #f0f8ff;}#sk-container-id-1 input.sk-toggleable__control:checked~div.sk-toggleable__content {max-height: 200px;max-width: 100%;overflow: auto;}#sk-container-id-1 input.sk-toggleable__control:checked~label.sk-toggleable__label-arrow:before {content: \"▾\";}#sk-container-id-1 div.sk-estimator input.sk-toggleable__control:checked~label.sk-toggleable__label {background-color: #d4ebff;}#sk-container-id-1 div.sk-label input.sk-toggleable__control:checked~label.sk-toggleable__label {background-color: #d4ebff;}#sk-container-id-1 input.sk-hidden--visually {border: 0;clip: rect(1px 1px 1px 1px);clip: rect(1px, 1px, 1px, 1px);height: 1px;margin: -1px;overflow: hidden;padding: 0;position: absolute;width: 1px;}#sk-container-id-1 div.sk-estimator {font-family: monospace;background-color: #f0f8ff;border: 1px dotted black;border-radius: 0.25em;box-sizing: border-box;margin-bottom: 0.5em;}#sk-container-id-1 div.sk-estimator:hover {background-color: #d4ebff;}#sk-container-id-1 div.sk-parallel-item::after {content: \"\";width: 100%;border-bottom: 1px solid gray;flex-grow: 1;}#sk-container-id-1 div.sk-label:hover label.sk-toggleable__label {background-color: #d4ebff;}#sk-container-id-1 div.sk-serial::before {content: \"\";position: absolute;border-left: 1px solid gray;box-sizing: border-box;top: 0;bottom: 0;left: 50%;z-index: 0;}#sk-container-id-1 div.sk-serial {display: flex;flex-direction: column;align-items: center;background-color: white;padding-right: 0.2em;padding-left: 0.2em;position: relative;}#sk-container-id-1 div.sk-item {position: relative;z-index: 1;}#sk-container-id-1 div.sk-parallel {display: flex;align-items: stretch;justify-content: center;background-color: white;position: relative;}#sk-container-id-1 div.sk-item::before, #sk-container-id-1 div.sk-parallel-item::before {content: \"\";position: absolute;border-left: 1px solid gray;box-sizing: border-box;top: 0;bottom: 0;left: 50%;z-index: -1;}#sk-container-id-1 div.sk-parallel-item {display: flex;flex-direction: column;z-index: 1;position: relative;background-color: white;}#sk-container-id-1 div.sk-parallel-item:first-child::after {align-self: flex-end;width: 50%;}#sk-container-id-1 div.sk-parallel-item:last-child::after {align-self: flex-start;width: 50%;}#sk-container-id-1 div.sk-parallel-item:only-child::after {width: 0;}#sk-container-id-1 div.sk-dashed-wrapped {border: 1px dashed gray;margin: 0 0.4em 0.5em 0.4em;box-sizing: border-box;padding-bottom: 0.4em;background-color: white;}#sk-container-id-1 div.sk-label label {font-family: monospace;font-weight: bold;display: inline-block;line-height: 1.2em;}#sk-container-id-1 div.sk-label-container {text-align: center;}#sk-container-id-1 div.sk-container {/* jupyter's `normalize.less` sets `[hidden] { display: none; }` but bootstrap.min.css set `[hidden] { display: none !important; }` so we also need the `!important` here to be able to override the default hidden behavior on the sphinx rendered scikit-learn.org. See: https://github.com/scikit-learn/scikit-learn/issues/21755 */display: inline-block !important;position: relative;}#sk-container-id-1 div.sk-text-repr-fallback {display: none;}</style><div id=\"sk-container-id-1\" class=\"sk-top-container\"><div class=\"sk-text-repr-fallback\"><pre>KNeighborsClassifier(metric=&lt;function kl_divergence at 0x7f0b2b4b71c0&gt;,\n",
              "                     n_neighbors=3)</pre><b>In a Jupyter environment, please rerun this cell to show the HTML representation or trust the notebook. <br />On GitHub, the HTML representation is unable to render, please try loading this page with nbviewer.org.</b></div><div class=\"sk-container\" hidden><div class=\"sk-item\"><div class=\"sk-estimator sk-toggleable\"><input class=\"sk-toggleable__control sk-hidden--visually\" id=\"sk-estimator-id-1\" type=\"checkbox\" checked><label for=\"sk-estimator-id-1\" class=\"sk-toggleable__label sk-toggleable__label-arrow\">KNeighborsClassifier</label><div class=\"sk-toggleable__content\"><pre>KNeighborsClassifier(metric=&lt;function kl_divergence at 0x7f0b2b4b71c0&gt;,\n",
              "                     n_neighbors=3)</pre></div></div></div></div></div>"
            ]
          },
          "metadata": {},
          "execution_count": 19
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# We can compare ICL predictions to kNN Prompting prediction\n",
        "\n",
        "def knn_predict(prompt):\n",
        "    query = llm_gen(model, prompt, tokenizer, max_context_len)\n",
        "    query_np = torch.softmax(query.float(), dim=-1).cpu().numpy()\n",
        "    predicted_labels = knn_classifier.predict(query_np)\n",
        "    return predicted_labels\n",
        "\n",
        "example_knn_prediction = knn_predict(example_prompt)\n",
        "print(f\"Predicted label: {example_knn_prediction[0]} / True label {example_label}\")\n",
        "print(f\"Id to label mapping: 0 = {id2verb[0]} / 1 = {id2verb[1]} \")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kGT1eW4EiBBy",
        "outputId": "c86b9a9b-8cf0-4b7c-ebb3-39f8d9ab49e0"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Predicted label: 1 / True label 1\n",
            "Id to label mapping: 0 = negative / 1 = positive \n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from tqdm.notebook import tqdm as tqdm_not\n",
        "\n",
        "# Evaluate across the whole test dataset\n",
        "dev_labels = []\n",
        "dev_pred = []\n",
        "for ins in tqdm_not(dev_data.data, total=dev_data.__len__()):\n",
        "    dev_labels.append(label2id[ins['label']])\n",
        "    prompt = prompt_prefix + make_prompt(ins, 'sst2', mode='inference')\n",
        "    dev_pred.extend(knn_predict(prompt))\n",
        "\n",
        "dev_correct = [1 if dev_labels[i] == dev_pred[i] else 0 for i in range(len(dev_labels))]\n",
        "acc = sum(dev_correct) / len(dev_labels)\n",
        "print(f\"Prediction accuracy of KNN Prompting across the validation set: {acc*100:.1f}%\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 66,
          "referenced_widgets": [
            "37bebe115def4928b1530a65af6a64fc",
            "fd43fbc747f648c6ad94801bf14b7652",
            "70fda9ad89d74d4ba9ac45fc2209abde",
            "8766075c4a5b4394a860da66b653f774",
            "2faa4ce6b9f64fa59f27531c73911990",
            "d2affff983bb4a1895fd6b4d365598d3",
            "a6211fd182ea4521bc1adcc85bb7274f",
            "8455ce072bef4dd4974625c720e09fbe",
            "73c2ace9974e4da5b8c9c9864809fd72",
            "185563dc7e4a403f821da2f366ba4f87",
            "fb1a8a78a868431fa06775f2fdb25f79"
          ]
        },
        "id": "NirOlx2nm0aL",
        "outputId": "beb36d42-20e9-4edf-9321-873bf0e1d370"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/256 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "37bebe115def4928b1530a65af6a64fc"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Prediction accuracy of KNN Prompting across the validation set: 89.8%\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Let's compare with ICL: (Reusing code from https://github.com/BenfengXu/KNNPrompting/blob/main/icl.py)"
      ],
      "metadata": {
        "id": "4RpsTQvJvVo9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from utils.dataset import *\n",
        "from utils.anchor import AnchorStore\n",
        "from utils.template import *\n",
        "\n",
        "# Still using sst2 dataset: https://huggingface.co/datasets/sst2\n",
        "AutoDataset = SST2Dataset\n",
        "\n",
        "def llm_gen(model, prompt, tokenizer, max_context_len):\n",
        "    inputs = tokenizer.encode_plus(prompt, return_tensors=\"pt\", padding=True).to(device=model.device)\n",
        "    if inputs['input_ids'].shape[1] > max_context_len:\n",
        "        inputs['input_ids'] = inputs['input_ids'][:, -max_context_len:]\n",
        "        inputs['attention_mask'] = inputs['attention_mask'][:, -max_context_len:]\n",
        "    with torch.no_grad():\n",
        "        logits = model.forward(input_ids=inputs['input_ids'],\n",
        "                               attention_mask=inputs['attention_mask'],\n",
        "                               return_dict=True).logits.detach().cpu()\n",
        "    # the output prob is shifted by -1, so we should use the output at the last input token position\n",
        "    # gen_logits.shape = [1, 50257]\n",
        "    gen_logits = logits[:, -1, :]\n",
        "\n",
        "    return gen_logits\n",
        "\n",
        "\n",
        "def parse_response(gen_logits, tokenizer, id2verb):\n",
        "    gen_prob = torch.softmax(gen_logits.float(), dim=-1)\n",
        "    prob_per_cls = []\n",
        "    for label_verb in id2verb:\n",
        "        label_verb_token_id = tokenizer.encode(' ' + label_verb)[-1] # note the space before label word\n",
        "        prob_per_cls.append(gen_prob[:, label_verb_token_id])\n",
        "    pred = torch.argmax(torch.cat(prob_per_cls, dim=0)).tolist()\n",
        "    return pred\n",
        "\n",
        "\n",
        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "dataset_dir = \"/content/data/sst2/\"\n",
        "train_data = AutoDataset(dataset_dir, mode='train')\n",
        "dev_data = AutoDataset(dataset_dir, mode='dev')\n",
        "\n",
        "max_context_len = 1024\n",
        "n_demo_shot = 32\n",
        "seed = 43\n",
        "\n",
        "# inference\n",
        "train_data.subsamplebyshot(n_demo_shot, seed)\n",
        "prompt_prefix = make_prompt(train_data, \"sst2\", mode='train')\n",
        "dev_labels = []\n",
        "dev_pred = []\n",
        "label2id = dev_data.label2id\n",
        "id2verb = train_data.id2verb\n",
        "for ins in tqdm(dev_data.data, total=dev_data.__len__()):\n",
        "    dev_labels.append(label2id[ins['label']])\n",
        "    prompt = prompt_prefix + make_prompt(ins, \"sst2\", mode='inference')\n",
        "    gen_logits = llm_gen(model, prompt, tokenizer, max_context_len)\n",
        "    dev_pred.append(parse_response(gen_logits, tokenizer, id2verb))\n",
        "\n",
        "dev_correct = [1 if dev_labels[i] == dev_pred[i] else 0 for i in range(len(dev_labels))]\n",
        "acc = sum(dev_correct) / len(dev_labels)\n",
        "print(f\"Prediction accuracy for ICL across the validation set: {acc*100:.1f}%\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 66,
          "referenced_widgets": [
            "c50af065971b4b73a3a493ce9697225f",
            "4830f83d08a04b9f9d2d692848c661a4",
            "82db694856ac4587ac2a01f4b0b70dcd",
            "410991c1df0948318bbce477c31af473",
            "8686ae8c73844024b3673a0ecf8d6db2",
            "137282885f074b6782fe279508d6c4e7",
            "231a041f24824d56a0d4a492ff441314",
            "c7163e646190420c918de6bfbe66653c",
            "ac8b043917524154aa8522539b398831",
            "19fd8eb9ac354c998ba5eb4a0ee6f755",
            "edbc64b2accc4ef18e68ff3aaedae879"
          ]
        },
        "id": "khMs0ECipf_N",
        "outputId": "acdf9688-def5-4d97-82c1-0c21c5ebfb61"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/256 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "c50af065971b4b73a3a493ce9697225f"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Prediction accuracy for ICL across the validation set: 77.7%\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "acxpBPqTxmDH"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "c50af065971b4b73a3a493ce9697225f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_4830f83d08a04b9f9d2d692848c661a4",
              "IPY_MODEL_82db694856ac4587ac2a01f4b0b70dcd",
              "IPY_MODEL_410991c1df0948318bbce477c31af473"
            ],
            "layout": "IPY_MODEL_8686ae8c73844024b3673a0ecf8d6db2"
          }
        },
        "4830f83d08a04b9f9d2d692848c661a4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_137282885f074b6782fe279508d6c4e7",
            "placeholder": "​",
            "style": "IPY_MODEL_231a041f24824d56a0d4a492ff441314",
            "value": "100%"
          }
        },
        "82db694856ac4587ac2a01f4b0b70dcd": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_c7163e646190420c918de6bfbe66653c",
            "max": 256,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_ac8b043917524154aa8522539b398831",
            "value": 256
          }
        },
        "410991c1df0948318bbce477c31af473": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_19fd8eb9ac354c998ba5eb4a0ee6f755",
            "placeholder": "​",
            "style": "IPY_MODEL_edbc64b2accc4ef18e68ff3aaedae879",
            "value": " 256/256 [01:47&lt;00:00,  2.44it/s]"
          }
        },
        "8686ae8c73844024b3673a0ecf8d6db2": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "137282885f074b6782fe279508d6c4e7": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "231a041f24824d56a0d4a492ff441314": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "c7163e646190420c918de6bfbe66653c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "ac8b043917524154aa8522539b398831": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "19fd8eb9ac354c998ba5eb4a0ee6f755": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "edbc64b2accc4ef18e68ff3aaedae879": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "37bebe115def4928b1530a65af6a64fc": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_fd43fbc747f648c6ad94801bf14b7652",
              "IPY_MODEL_70fda9ad89d74d4ba9ac45fc2209abde",
              "IPY_MODEL_8766075c4a5b4394a860da66b653f774"
            ],
            "layout": "IPY_MODEL_2faa4ce6b9f64fa59f27531c73911990"
          }
        },
        "fd43fbc747f648c6ad94801bf14b7652": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_d2affff983bb4a1895fd6b4d365598d3",
            "placeholder": "​",
            "style": "IPY_MODEL_a6211fd182ea4521bc1adcc85bb7274f",
            "value": "100%"
          }
        },
        "70fda9ad89d74d4ba9ac45fc2209abde": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_8455ce072bef4dd4974625c720e09fbe",
            "max": 256,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_73c2ace9974e4da5b8c9c9864809fd72",
            "value": 256
          }
        },
        "8766075c4a5b4394a860da66b653f774": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_185563dc7e4a403f821da2f366ba4f87",
            "placeholder": "​",
            "style": "IPY_MODEL_fb1a8a78a868431fa06775f2fdb25f79",
            "value": " 256/256 [04:40&lt;00:00,  1.08s/it]"
          }
        },
        "2faa4ce6b9f64fa59f27531c73911990": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "d2affff983bb4a1895fd6b4d365598d3": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "a6211fd182ea4521bc1adcc85bb7274f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "8455ce072bef4dd4974625c720e09fbe": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "73c2ace9974e4da5b8c9c9864809fd72": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "185563dc7e4a403f821da2f366ba4f87": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "fb1a8a78a868431fa06775f2fdb25f79": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}