Petals_and_Together_BLOOM_176B_GPU_Colab.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/R3gm/InsightSolver-Colab/blob/main/Petals_and_Together_BLOOM_176B_GPU_Colab.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "<div align=\"center\">\n",
        "<img src=\"https://camo.githubusercontent.com/473dd9f992924d27457650251786464f72e54121ac6e9210add0f483ca849277/68747470733a2f2f692e696d6775722e636f6d2f3765523750616e2e706e67\" width=\"40%\">  \n",
        "</div>\n"
      ],
      "metadata": {
        "id": "VsXHWJLuowcn"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "| Code Credits | Link |\n",
        "| ----------- | ---- |\n",
        "| 🎉 Repository | [![GitHub Repository](https://img.shields.io/github/stars/petals-infra/chat.petals.ml?style=social)](https://github.com/petals-infra/chat.petals.ml) |\n",
        "| Original Colab | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1Ervk6HPNS6AYVr3xVdQnY5a-TjjmLCdQ?usp=sharing) |\n",
        "| 🚀 Online inference for Guanaco-65B, LLaMA-65B, BLOOM and BLOOMZ | [Chat Petals](http://chat.petals.ml) |\n",
        "| 🚀 Alternative with api Together | [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/sambanovasystems/BLOOMChat) |\n",
        "| 🔥 Discover More Colab Notebooks | [![GitHub Repository](https://img.shields.io/badge/GitHub-Repository-black?style=flat-square&logo=github)](https://github.com/R3gm/Colab-resources/) |"
      ],
      "metadata": {
        "id": "zHRRq_83yCSI"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Petals is a project that leverages [Hivemind](https://github.com/learning-at-home/hivemind).\n",
        "\n",
        "Hivemind is a PyTorch library for decentralized deep learning, enabling distributed training without a master node, fault-tolerant backpropagation, and decentralized parameter averaging. It allows training large models on hundreds of computers across the Internet. You can install it with pip or from source and find examples in the documentation. Contributions are welcome.\n",
        "\n",
        "---\n",
        "\n"
      ],
      "metadata": {
        "id": "JWvw2ZOlJ5Eu"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "# Getting started with Petals\n",
        "\n",
        "This notebook will guide you through the basics of Petals &mdash; a system for inference and fine-tuning 100B+ language models without the need to have high-end GPUs. With Petals, you can join compute resources with other people over the Internet and run large language models such as Guanaco-65B, LLaMA-65B, 176B-parameter [BLOOM](https://huggingface.co/bigscience/bloom) or [BLOOMZ](https://huggingface.co/bigscience/bloomz), which are of the same size as GPT-3.\n",
        "\n",
        "💬 If you meet any issues while running this notebook, let us know in the **[#running-a-client](https://discord.gg/J29mCBNBvm)** channel of our Discord!\n",
        "\n",
        "So, let's get started! First, let's install [the Petals package](https://github.com/bigscience-workshop/petals):\n"
      ],
      "metadata": {
        "id": "M8PEszdKIu28"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pBC52TF3LVY1",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ddf4c14f-2b26-4057-faa6-d76c952e4932"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m87.1/87.1 kB\u001b[0m \u001b[31m4.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m98.2/98.2 MB\u001b[0m \u001b[31m10.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m219.1/219.1 kB\u001b[0m \u001b[31m25.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m236.8/236.8 kB\u001b[0m \u001b[31m27.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.1/7.1 MB\u001b[0m \u001b[31m110.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m9.2/9.2 MB\u001b[0m \u001b[31m111.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m86.8/86.8 kB\u001b[0m \u001b[31m12.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m4.1/4.1 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[32m2.4/2.4 MB\u001b[0m \u001b[31m108.9 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[31m43.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.4/2.4 MB\u001b[0m \u001b[31m48.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.4/2.4 MB\u001b[0m \u001b[31m105.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.4/2.4 MB\u001b[0m \u001b[31m105.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.4/2.4 MB\u001b[0m \u001b[31m97.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.4/2.4 MB\u001b[0m \u001b[31m109.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.4/2.4 MB\u001b[0m \u001b[31m113.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.4/2.4 MB\u001b[0m \u001b[31m99.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.4/2.4 MB\u001b[0m \u001b[31m114.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.9/1.9 MB\u001b[0m \u001b[31m96.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Building wheel for varint (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
          ]
        }
      ],
      "source": [
        "%pip install -q petals"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Step 1. The easiest way to generate text 🚀\n",
        "\n",
        "Let's start with the easiest task &mdash; creating a __`DistributedBloom`__ model and using it for generating text.\n",
        "\n",
        "This machine will download a small part of the model weights (~8 GB out of 352 GB) and rely on other computers in the network to run the rest of the model. Downloading the local part of the weights usually takes ~3 minutes.\n",
        "\n",
        "🧑‍🏫 __Note:__ We suggest to start with the regular BLOOM, but you can also use [BLOOMZ](https://huggingface.co/bigscience/bloomz) &mdash; a version of BLOOM fine-tuned to better follow human instructions in the zero-shot regime. You would need to set `MODEL_NAME = \"bigscience/bloomz-petals\"` to load this model."
      ],
      "metadata": {
        "id": "yEbot-oEXdpw"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2uuX1IMLLotQ",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 317,
          "referenced_widgets": [
            "d699cc8313944cc0b6fcb729a946319f",
            "b5d30f01ad0f49aabb1ea83735f70730",
            "29743cb700484e15a386b00c98aadcf9",
            "fe07a97af6c1469b8176a1a4c8888444",
            "b67922fe0f5543e399179cb4bcc2bc3a",
            "5c27fb179ce2476c86116f5620f4ef59",
            "efa3a5f4a9b8438c9ec211bdb1d24dbf",
            "e7cd84d56240480d897a808d2c8d8664",
            "8e238ad5782a40b58f5aeece220b1f45",
            "909b3e93de3445af866bc3a3eb3e67e7",
            "d640b7bddf6746f4bb7181d8746e518f",
            "a365d95a9e544b0aaf949dfc3c436fd8",
            "0b51448bcffb4c3881c42ef386ac8abe",
            "9ba0bd3f6a1146dbad9baa37ea15ad42",
            "a88549b6c53a45c09f1898aef972562a",
            "2c50dee6d3b740c58428ec5867302244",
            "311b129becc74fa787c96f8de2f4970f",
            "7f48720a85e743cd803351893d0a46c1",
            "3ade1584527c414193588cd88d05d68e",
            "70feae1bdc5e46efadc0ef3228b05847",
            "8e40327db5cc4912bfb1bba1c92aea30",
            "dd1abdbaf11046739d60dba36f38dfd7",
            "8b700d043a0041e2b3080b9a20ff8acd",
            "76ab08681ded4ac6ad38a73b6c7ead1b",
            "7ad076f3e73e4ac1a06c49b453f58429",
            "f49410314aad47b793fc1ad7eaa97107",
            "5b7c66755535426a8bde372a2863650f",
            "62271c6779614341a91c5279470ce2c5",
            "b8501505088c499e958bf251f3296bd8",
            "4204d9413bc945d786a16c488c4b4c6b",
            "c1dfd895f3874ab68a893321005e52a5",
            "46c94ebababc44e5b96c0639e7ae37d4",
            "0310bcdbd9d04daca17bf219479a83a2",
            "f80871b3d11d49d986961a8f76e8dcd9",
            "1fda61ef05424c3eafbfde69ea8a5c06",
            "e02b6473cf284fcf8a8e8765c794cd14",
            "154a1e132f004fa3b2a233d2e8ea2689",
            "6df211b873f9437ca34a80663b868207",
            "0d3d2bb3103f4b0c85084ff79a606997",
            "6146ecfb133e4656bcfb0ae7c618c50e",
            "6d87070ff58343bfaced6ee8977171b5",
            "f52b7b4d3fd84054ae8b1ba83acd9cd9",
            "b59b8fb5981f40e594083c2f4146477a",
            "7f186445fda344fc95934431cf8903aa",
            "0d3f25decf5a4067a9059b58c3220f88",
            "342ad5f476a541858eb8df963be00434",
            "57657e5d5cf04efcb0ef41d54d9cc48e",
            "92dfd8d47b80413bab1efc792f98f6e8",
            "6ea346f9fe674099b57b5500884bb325",
            "11cc3f308ab94f1583febaa4d77d46ec",
            "4ffd0800b6004fd1b8bfa495a6e4b610",
            "62396b1ef8b540faace8a512b5ec5c2b",
            "7e05f0ea127845eba8b89524c9a40465",
            "dd541d8439c244eb8d886fc817bf9551",
            "e5cb3e4076ef460595ce40de0fd16465"
          ]
        },
        "outputId": "7458e85d-5eab-48b3-aeb4-5a2f8f9d0959"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading (…)okenizer_config.json:   0%|          | 0.00/263 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "d699cc8313944cc0b6fcb729a946319f"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading tokenizer.json:   0%|          | 0.00/14.5M [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "a365d95a9e544b0aaf949dfc3c436fd8"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading (…)cial_tokens_map.json:   0%|          | 0.00/96.0 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "8b700d043a0041e2b3080b9a20ff8acd"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading (…)lve/main/config.json:   0%|          | 0.00/641 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "f80871b3d11d49d986961a8f76e8dcd9"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading pytorch_model.bin:   0%|          | 0.00/7.19G [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "0d3f25decf5a4067a9059b58c3220f88"
            }
          },
          "metadata": {}
        }
      ],
      "source": [
        "import torch\n",
        "from transformers import BloomTokenizerFast\n",
        "from petals import DistributedBloomForCausalLM\n",
        "\n",
        "MODEL_NAME = \"bigscience/bloom-petals\"\n",
        "tokenizer = BloomTokenizerFast.from_pretrained(MODEL_NAME)\n",
        "model = DistributedBloomForCausalLM.from_pretrained(MODEL_NAME)\n",
        "model = model.cuda()"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now, let's try to generate something by calling __`model.generate()`__ method.\n",
        "\n",
        "The first call to this method takes ~5 sec to connect to the Petals swarm. Once we do that, you should expect generation speed of 1&ndash;1.5 sec/token. If you don't have enough GPUs to host the entire model, this is much faster than what you get with other methods, such as offloading (which takes at least 10&ndash;20 sec/token)."
      ],
      "metadata": {
        "id": "zhyUxv13sfKO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "inputs = tokenizer('A cat in French is \"', return_tensors=\"pt\")[\"input_ids\"].cuda()\n",
        "outputs = model.generate(inputs, max_new_tokens=3)\n",
        "print(tokenizer.decode(outputs[0]))"
      ],
      "metadata": {
        "id": "1s1IrE1H8wwr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The `model.generate()` method runs **greedy** generation by default. However, you can choose other generation methods like **top-p/top-k sampling** or **beam search** by passing the corresponding parameters (you'll see an example in a bit). You can even implement custom generation methods (we'll cover that in **Step 5**).\n",
        "\n",
        "🔏 **Note:** Your data is processed by other people in the public swarm. Learn more about privacy [here](https://github.com/bigscience-workshop/petals/wiki/Security,-privacy,-and-AI-safety). For sensitive data, you can set up a [private swarm](https://github.com/bigscience-workshop/petals/wiki/Launch-your-own-swarm) among people you trust."
      ],
      "metadata": {
        "id": "02d0BDEAuUFQ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Step 2. Chat bots and interactive generation 💬\n",
        "\n",
        "If you'd like to talk to the model in an interactive way, you can use the __inference session__ interface. This interface provides a simple way to print generated tokens on the fly or make a chat bot that responds to human's phrases.\n",
        "\n",
        "The inference session looks for a sequence of servers that will run successive inference steps and store past attention caches. This way, you don't need to rerun previous tokens through the transformer to generate each phrase. If one of the remote servers fails, Petals will automatically find a replacement and regenerate a small part of the caches.\n",
        "\n",
        "Let's see how to use it to write a simple chat bot, showing tokens as soon as they are generated:"
      ],
      "metadata": {
        "id": "fZlzYVn0ApyF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "with model.inference_session(max_length=512) as sess:\n",
        "    while True:\n",
        "        prompt = input('Human: ')\n",
        "        if prompt == \"\":\n",
        "            break\n",
        "        prefix = f\"Human: {prompt}\\nFriendly AI:\"\n",
        "        prefix = tokenizer(prefix, return_tensors=\"pt\")[\"input_ids\"].cuda()\n",
        "        print(\"Friendly AI:\", end=\"\", flush=True)\n",
        "\n",
        "        while True:\n",
        "            outputs = model.generate(\n",
        "                prefix, max_new_tokens=1, do_sample=True, top_p=0.9, temperature=0.75, session=sess\n",
        "            )\n",
        "            outputs = tokenizer.decode(outputs[0, -1:])\n",
        "            print(outputs, end=\"\", flush=True)\n",
        "            if \"\\n\" in outputs:\n",
        "                break\n",
        "            prefix = None  # Prefix is passed only for the 1st token of the bot's response"
      ],
      "metadata": {
        "id": "m5bdMaYYAqYR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 📦 Making apps that use Petals\n",
        "\n",
        "If you develop a tool for other people, you can wrap up the code using Petals into a user-friendly web app, such as [chat.petals.ml](http://chat.petals.ml). Under the hood, this app may connect to a lightweight [HTTP endpoint](https://github.com/borzunov/petals-chat) for inference that forwards all requests to the Petals swarm.\n",
        "\n",
        "📋 **Note:** If you build an app running BLOOM with Petals, make sure it follows the BLOOM's [terms of use](https://huggingface.co/bigscience/bloom).\n",
        "\n",
        "<div align=\"center\">\n",
        "<br>\n",
        "<img src=\"https://i.imgur.com/p2nwiho.png\" width=\"40%\">  \n",
        "</div>"
      ],
      "metadata": {
        "id": "P0k_0PrTAr6B"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Step 3. How does it work? 🛠️\n",
        "\n",
        "The `model` you are running is the actual BLOOM-176B, but only a part of it is loaded into your machine's GPU. Let's have a look under the hood:"
      ],
      "metadata": {
        "id": "557VBCGS8Dpv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model.transformer"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FB-lCjKs8NEt",
        "outputId": "9b10a262-5ec5-43fa-c2a7-86bfa0accfde"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DistributedBloomModel(\n",
              "  (word_embeddings): Embedding(250880, 14336)\n",
              "  (word_embeddings_layernorm): LayerNorm((14336,), eps=1e-05, elementwise_affine=True)\n",
              "  (h): RemoteSequential(modules=bigscience/bloom-petals.0..bigscience/bloom-petals.69)\n",
              "  (ln_f): LayerNorm((14336,), eps=1e-05, elementwise_affine=True)\n",
              ")"
            ]
          },
          "metadata": {},
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "As you can see, word embeddings and some other layers are regular PyTorch modules hosted on your machine, but the rest of the model (e.g., transformers blocks) is encased in the __RemoteSequential__ class. This is an advanced PyTorch module that runs on a distributed swarm of other machines.\n",
        "\n",
        "Still, you can access individual layers and their outputs, as well as run forward/backward through them:"
      ],
      "metadata": {
        "id": "9EW3wBDO-aTJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "first_five_layers = model.transformer.h[0:5]\n",
        "first_five_layers"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OHYihUHH-Zo6",
        "outputId": "98237bb8-0c36-455d-a31a-6dd9c0ef2408"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "RemoteSequential(modules=bigscience/bloom-petals.0..bigscience/bloom-petals.4)"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "dummy_inputs = torch.randn(1, 3, 14336, dtype=torch.bfloat16, requires_grad=True)\n",
        "outputs = first_five_layers(dummy_inputs)\n",
        "outputs"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ilOwoxr47Sso",
        "outputId": "b9110e63-4e9e-44c0-d6f8-961380de08f4"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[[-2.8906, -0.6484, -1.9141,  ...,  3.5938,  1.2656, -2.0156],\n",
              "         [ 0.5156, -2.2031, -0.5820,  ...,  1.1250, -0.0127, -1.3594],\n",
              "         [-0.5508, -0.9766, -0.7695,  ...,  1.4766, -1.3594, -0.3086]]],\n",
              "       dtype=torch.bfloat16, grad_fn=<_RemoteSequentialAutogradFunctionBackward>)"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "loss = torch.mean((outputs - torch.ones_like(outputs)) ** 2)\n",
        "loss.backward()  # backpropagate through the internet\n",
        "print(\"Grad w.r.t. inputs:\", dummy_inputs.grad.flatten())"
      ],
      "metadata": {
        "id": "4XM6In8uArE5",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "7a98519e-787e-4872-f5cb-5019ee2ebb83"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Grad w.r.t. inputs: tensor([-0.0137,  0.0459,  0.0242,  ...,  0.0010, -0.0005,  0.0002],\n",
            "       dtype=torch.bfloat16)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "In general, you can mix and match distributed layers like in regular PyTorch and even insert and train your own layers (e.g., adapters) between the pre-trained ones."
      ],
      "metadata": {
        "id": "DxWOzWCwat6H"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "<div align=\"center\">\n",
        "<img src=\"https://camo.githubusercontent.com/58732a64488a9be928e25f3e60e3692b989ffe212ac86cb4902d8df20a042b03/68747470733a2f2f692e696d6775722e636f6d2f525459463379572e706e67\" width=\"80%\">\n",
        "</div>\n",
        "\n",
        "<p align=\"center\">📜 <b><a href=\"https://arxiv.org/pdf/2209.01188.pdf\">Read details in our paper</a></b></p>"
      ],
      "metadata": {
        "id": "0OZ5eXFrkfzt"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Step 4. Adding a trainable adapter 🏋️\n",
        "\n",
        "While the remotely hosted transformer blocks are **frozen** to keep the pretrained model the same for all users, using **parameter-efficient adapters** (small trainable layers added between the pretrained blocks of the model, such as [LoRA](https://arxiv.org/abs/2106.09685)) or **trainable prompts** (trainable inputs added before the inputs to the model, such as in [P-Tuning v2](https://arxiv.org/abs/2110.07602)) is usually enough to make BLOOM solve a variety of downstream tasks.\n",
        "\n",
        "Below, we show an example of how to add a basic **trainable** linear layer between 5th and 6th transformer blocks of the pretrained model. The layer's weights and the corresponding optimizer statistics will be stored locally:"
      ],
      "metadata": {
        "id": "lJt4OS2pIZe_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "\n",
        "class BloomBasedClassifier(nn.Module):\n",
        "  def __init__(self, model):\n",
        "    super().__init__()\n",
        "    self.distributed_layers = model.transformer.h\n",
        "    self.adapter = nn.Sequential(nn.Linear(14336, 32), nn.Linear(32, 14336))\n",
        "    self.head = nn.Sequential(nn.LayerNorm(14336), nn.Linear(14336, 2))\n",
        "\n",
        "  def forward(self, embeddings):\n",
        "    hidden_states = self.distributed_layers[0:6](embeddings)\n",
        "    hidden_states = self.adapter(hidden_states)\n",
        "    hidden_states = self.distributed_layers[6:10](hidden_states)\n",
        "    pooled_states = torch.mean(hidden_states, dim=1)\n",
        "    return self.head(pooled_states)"
      ],
      "metadata": {
        "id": "YHPFitSoIZO0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "classifier = BloomBasedClassifier(model).cuda()\n",
        "opt = torch.optim.Adam(classifier.parameters(), 3e-5)\n",
        "inputs = torch.randn(3, 2, 14336, device='cuda')\n",
        "labels = torch.tensor([1, 0, 1], device='cuda')\n",
        "\n",
        "for i in range(5):\n",
        "  loss = F.cross_entropy(classifier(inputs), labels)\n",
        "  print(f\"loss[{i}] = {loss.item():.3f}\")\n",
        "  opt.zero_grad()\n",
        "  loss.backward()\n",
        "  opt.step()\n",
        "\n",
        "print('predicted:', classifier(inputs).argmax(-1))  # l, o, l"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OJ7AFOumIYlY",
        "outputId": "eb5978a7-bfcb-43e2-8350-aa61282eb826"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "loss[0] = 0.706\n",
            "loss[1] = 0.554\n",
            "loss[2] = 0.446\n",
            "loss[3] = 0.362\n",
            "loss[4] = 0.294\n",
            "predicted: tensor([1, 0, 1], device='cuda:0')\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Step 5. Using custom sampling methods 🎰\n",
        "\n",
        "The __`model.inference_session()`__ interface in Petals also allows you to write custom inference code. You can use this to implement any sampling algorithms you want, or write a custom beam search algorithm that forbids the model from using swearwords.\n",
        "\n",
        "Below, let's see how we can reimplement the standard `model.generate()` interface by making forward passes through all the layers manually:"
      ],
      "metadata": {
        "id": "xz6EJ8VsYW2b"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "frLiu0yiL6Sn",
        "outputId": "ba745049-af2a-4115-f018-373af12d2b8e"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "What is AI? Answer: AI\n",
            "What is AI? Answer: AI is\n",
            "What is AI? Answer: AI is a\n",
            "What is AI? Answer: AI is a computer\n",
            "What is AI? Answer: AI is a computer program\n",
            "What is AI? Answer: AI is a computer program that\n",
            "What is AI? Answer: AI is a computer program that can\n",
            "What is AI? Answer: AI is a computer program that can think\n",
            "What is AI? Answer: AI is a computer program that can think and\n",
            "What is AI? Answer: AI is a computer program that can think and learn\n",
            "What is AI? Answer: AI is a computer program that can think and learn like\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human.\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human.\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human.\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human.\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human.\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human.\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer\n",
            "What is AI? Answer: AI is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program that can think and learn like a human. It is a computer program\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "\u001b[31m╭─\u001b[0m\u001b[31m──────────────────────────────\u001b[0m\u001b[31m \u001b[0m\u001b[1;31mTraceback \u001b[0m\u001b[1;2;31m(most recent call last)\u001b[0m\u001b[31m \u001b[0m\u001b[31m───────────────────────────────\u001b[0m\u001b[31m─╮\u001b[0m\n",
              "\u001b[31m│\u001b[0m in \u001b[92m<cell line: 4>\u001b[0m:\u001b[94m10\u001b[0m                                                                             \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m \u001b[2;33m/usr/local/lib/python3.10/dist-packages/petals/client/\u001b[0m\u001b[1;33minference_session.py\u001b[0m:\u001b[94m295\u001b[0m in \u001b[92mstep\u001b[0m           \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m292 \u001b[0m\u001b[2m│   │   │   │   │   \u001b[0m\u001b[94melse\u001b[0m:                                                                  \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m293 \u001b[0m\u001b[2m│   │   │   │   │   │   \u001b[0minputs = inputs[:, -n_input_tokens:]  \u001b[2m# No need to pass prefix f\u001b[0m   \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m294 \u001b[0m\u001b[2m│   │   │   │   │   \u001b[0m                                                                       \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m295 \u001b[2m│   │   │   │   │   \u001b[0moutputs = session.step(inputs, prompts[span.start : span.end], **kwa   \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m296 \u001b[0m\u001b[2m│   │   │   │   │   \u001b[0m\u001b[94massert\u001b[0m (                                                               \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m297 \u001b[0m\u001b[2m│   │   │   │   │   │   \u001b[0minputs.shape == outputs.shape                                      \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m298 \u001b[0m\u001b[2m│   │   │   │   │   \u001b[0m), \u001b[33mf\u001b[0m\u001b[33m\"\u001b[0m\u001b[33mShape mismatch: inputs.shape=\u001b[0m\u001b[33m{\u001b[0minputs.shape\u001b[33m}\u001b[0m\u001b[33m, outputs.shape=\u001b[0m\u001b[33m{\u001b[0mout   \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m \u001b[2;33m/usr/local/lib/python3.10/dist-packages/petals/client/\u001b[0m\u001b[1;33minference_session.py\u001b[0m:\u001b[94m107\u001b[0m in \u001b[92mstep\u001b[0m           \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m104 \u001b[0m\u001b[2m│   │   \u001b[0m                                                                                   \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m105 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[2m# serialize inputs and put them into the queue\u001b[0m                                     \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m106 \u001b[0m\u001b[2m│   │   \u001b[0minputs = (new_hidden_states, prompts, hypo_ids)                                    \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m107 \u001b[2m│   │   \u001b[0moutputs_serialized = RemoteExpertWorker.run_coroutine(                             \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m108 \u001b[0m\u001b[2m│   │   │   \u001b[0m\u001b[96mself\u001b[0m._step(                                                                    \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m109 \u001b[0m\u001b[2m│   │   │   │   \u001b[0mruntime_pb2.ExpertRequest(                                                 \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m110 \u001b[0m\u001b[2m│   │   │   │   │   \u001b[0muid=\u001b[96mself\u001b[0m.uid,                                                          \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m \u001b[2;33m/usr/local/lib/python3.10/dist-packages/hivemind/moe/client/\u001b[0m\u001b[1;33mremote_expert_worker.py\u001b[0m:\u001b[94m36\u001b[0m in        \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m \u001b[92mrun_coroutine\u001b[0m                                                                                    \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m33 \u001b[0m\u001b[2m│   │   \u001b[0m                                                                                    \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m34 \u001b[0m\u001b[2m│   │   \u001b[0mloop = \u001b[96mcls\u001b[0m._event_loop_fut.result()                                                 \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m35 \u001b[0m\u001b[2m│   │   \u001b[0mfuture = asyncio.run_coroutine_threadsafe(coro, loop)                               \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m36 \u001b[2m│   │   \u001b[0m\u001b[94mreturn\u001b[0m future \u001b[94mif\u001b[0m return_future \u001b[94melse\u001b[0m future.result()                                 \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m37 \u001b[0m                                                                                            \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m \u001b[2;33m/usr/lib/python3.10/concurrent/futures/\u001b[0m\u001b[1;33m_base.py\u001b[0m:\u001b[94m453\u001b[0m in \u001b[92mresult\u001b[0m                                    \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m450 \u001b[0m\u001b[2m│   │   │   │   \u001b[0m\u001b[94melif\u001b[0m \u001b[96mself\u001b[0m._state == FINISHED:                                              \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m451 \u001b[0m\u001b[2m│   │   │   │   │   \u001b[0m\u001b[94mreturn\u001b[0m \u001b[96mself\u001b[0m.__get_result()                                             \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m452 \u001b[0m\u001b[2m│   │   │   │   \u001b[0m                                                                           \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m453 \u001b[2m│   │   │   │   \u001b[0m\u001b[96mself\u001b[0m._condition.wait(timeout)                                              \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m454 \u001b[0m\u001b[2m│   │   │   │   \u001b[0m                                                                           \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m455 \u001b[0m\u001b[2m│   │   │   │   \u001b[0m\u001b[94mif\u001b[0m \u001b[96mself\u001b[0m._state \u001b[95min\u001b[0m [CANCELLED, CANCELLED_AND_NOTIFIED]:                     \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m456 \u001b[0m\u001b[2m│   │   │   │   │   \u001b[0m\u001b[94mraise\u001b[0m CancelledError()                                                 \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m \u001b[2;33m/usr/lib/python3.10/\u001b[0m\u001b[1;33mthreading.py\u001b[0m:\u001b[94m320\u001b[0m in \u001b[92mwait\u001b[0m                                                     \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m 317 \u001b[0m\u001b[2m│   │   \u001b[0mgotit = \u001b[94mFalse\u001b[0m                                                                     \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m 318 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mtry\u001b[0m:    \u001b[2m# restore state no matter what (e.g., KeyboardInterrupt)\u001b[0m                  \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m 319 \u001b[0m\u001b[2m│   │   │   \u001b[0m\u001b[94mif\u001b[0m timeout \u001b[95mis\u001b[0m \u001b[94mNone\u001b[0m:                                                           \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m 320 \u001b[2m│   │   │   │   \u001b[0mwaiter.acquire()                                                          \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m 321 \u001b[0m\u001b[2m│   │   │   │   \u001b[0mgotit = \u001b[94mTrue\u001b[0m                                                              \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m 322 \u001b[0m\u001b[2m│   │   │   \u001b[0m\u001b[94melse\u001b[0m:                                                                         \u001b[31m│\u001b[0m\n",
              "\u001b[31m│\u001b[0m   \u001b[2m 323 \u001b[0m\u001b[2m│   │   │   │   \u001b[0m\u001b[94mif\u001b[0m timeout > \u001b[94m0\u001b[0m:                                                           \u001b[31m│\u001b[0m\n",
              "\u001b[31m╰──────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n",
              "\u001b[1;91mKeyboardInterrupt\u001b[0m\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800000; text-decoration-color: #800000\">╭─────────────────────────────── </span><span style=\"color: #800000; text-decoration-color: #800000; font-weight: bold\">Traceback </span><span style=\"color: #bf7f7f; text-decoration-color: #bf7f7f; font-weight: bold\">(most recent call last)</span><span style=\"color: #800000; text-decoration-color: #800000\"> ────────────────────────────────╮</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">&lt;cell line: 4&gt;</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">10</span>                                                                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/usr/local/lib/python3.10/dist-packages/petals/client/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">inference_session.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">295</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">step</span>           <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">292 │   │   │   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">else</span>:                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">293 │   │   │   │   │   │   </span>inputs = inputs[:, -n_input_tokens:]  <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"># No need to pass prefix f</span>   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">294 │   │   │   │   │   </span>                                                                       <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>295 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   │   │   </span>outputs = session.step(inputs, prompts[span.start : span.end], **kwa   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">296 │   │   │   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">assert</span> (                                                               <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">297 │   │   │   │   │   │   </span>inputs.shape == outputs.shape                                      <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">298 │   │   │   │   │   </span>), <span style=\"color: #808000; text-decoration-color: #808000\">f\"Shape mismatch: inputs.shape={</span>inputs.shape<span style=\"color: #808000; text-decoration-color: #808000\">}, outputs.shape={</span>out   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/usr/local/lib/python3.10/dist-packages/petals/client/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">inference_session.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">107</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">step</span>           <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">104 │   │   </span>                                                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">105 │   │   # serialize inputs and put them into the queue</span>                                     <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">106 │   │   </span>inputs = (new_hidden_states, prompts, hypo_ids)                                    <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>107 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   </span>outputs_serialized = RemoteExpertWorker.run_coroutine(                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">108 │   │   │   </span><span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._step(                                                                    <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">109 │   │   │   │   </span>runtime_pb2.ExpertRequest(                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">110 │   │   │   │   │   </span>uid=<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.uid,                                                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/usr/local/lib/python3.10/dist-packages/hivemind/moe/client/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">remote_expert_worker.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">36</span> in        <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #00ff00; text-decoration-color: #00ff00\">run_coroutine</span>                                                                                    <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">33 │   │   </span>                                                                                    <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">34 │   │   </span>loop = <span style=\"color: #00ffff; text-decoration-color: #00ffff\">cls</span>._event_loop_fut.result()                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">35 │   │   </span>future = asyncio.run_coroutine_threadsafe(coro, loop)                               <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>36 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> future <span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> return_future <span style=\"color: #0000ff; text-decoration-color: #0000ff\">else</span> future.result()                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">37 </span>                                                                                            <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/usr/lib/python3.10/concurrent/futures/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">_base.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">453</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">result</span>                                    <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">450 │   │   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">elif</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._state == FINISHED:                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">451 │   │   │   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.__get_result()                                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">452 │   │   │   │   </span>                                                                           <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>453 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   │   </span><span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._condition.wait(timeout)                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">454 │   │   │   │   </span>                                                                           <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">455 │   │   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._state <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">in</span> [CANCELLED, CANCELLED_AND_NOTIFIED]:                     <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">456 │   │   │   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">raise</span> CancelledError()                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/usr/lib/python3.10/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">threading.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">320</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">wait</span>                                                     <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 317 │   │   </span>gotit = <span style=\"color: #0000ff; text-decoration-color: #0000ff\">False</span>                                                                     <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 318 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">try</span>:    <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"># restore state no matter what (e.g., KeyboardInterrupt)</span>                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 319 │   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> timeout <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">is</span> <span style=\"color: #0000ff; text-decoration-color: #0000ff\">None</span>:                                                           <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span> 320 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   │   </span>waiter.acquire()                                                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 321 │   │   │   │   </span>gotit = <span style=\"color: #0000ff; text-decoration-color: #0000ff\">True</span>                                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 322 │   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">else</span>:                                                                         <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 323 │   │   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> timeout &gt; <span style=\"color: #0000ff; text-decoration-color: #0000ff\">0</span>:                                                           <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
              "<span style=\"color: #800000; text-decoration-color: #800000\">╰──────────────────────────────────────────────────────────────────────────────────────────────────╯</span>\n",
              "<span style=\"color: #ff0000; text-decoration-color: #ff0000; font-weight: bold\">KeyboardInterrupt</span>\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ],
      "source": [
        "text = \"What is AI? Answer:\"\n",
        "token_ids = tokenizer(text, return_tensors=\"pt\")[\"input_ids\"].cuda()\n",
        "max_length = 1000\n",
        "with torch.inference_mode():\n",
        "    with model.inference_session(max_length=max_length) as sess:\n",
        "        while len(text) < max_length:\n",
        "            embs = model.transformer.word_embeddings(token_ids)\n",
        "            embs = model.transformer.word_embeddings_layernorm(embs)\n",
        "\n",
        "            h = sess.step(embs)\n",
        "            h_last = model.transformer.ln_f(h[:, -1])\n",
        "            logits = model.lm_head(h_last)\n",
        "\n",
        "            next_token = logits.argmax(dim=-1)\n",
        "            text += tokenizer.decode(next_token)\n",
        "            token_ids = next_token.reshape(1, 1)\n",
        "            print(text)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Step 6. Sharing is caring 🤗\n",
        "\n",
        "We developed Petals to be a community-run system, so we rely on people giving out their GPUs to increase the swarm’s capacity. If you have some GPUs that are not always busy, please **consider running a Petals server.** You can pause it any time if you want to use the GPUs for something else. As a bonus, people running a server get a certain speedup when using Petals, since a larger part of the model is hosted locally.\n",
        "\n",
        "<br>\n",
        "\n",
        "🐋 You can run our [Docker](https://www.docker.com) image (works on Linux, macOS, and Windows with [WSL2](https://learn.microsoft.com/en-us/windows/ai/directml/gpu-cuda-in-wsl)):\n",
        "\n",
        "```\n",
        "sudo docker run -p 31330:31330 --ipc host --gpus all --volume petals-cache:/cache --rm \\\n",
        "    learningathome/petals:main python -m petals.cli.run_server bigscience/bloom-petals --port 31330\n",
        "```\n",
        "\n",
        "🐍 Or run these commands in an [Anaconda](https://www.anaconda.com) env (requires Linux and Python 3.7+):\n",
        "\n",
        "```\n",
        "conda install pytorch pytorch-cuda=11.7 -c pytorch -c nvidia\n",
        "pip install -U petals\n",
        "python -m petals.cli.run_server bigscience/bloom-petals\n",
        "```\n",
        "\n",
        "<br>\n",
        "\n",
        "📚 See [FAQ](https://github.com/bigscience-workshop/petals/wiki/FAQ:-Frequently-asked-questions#running-a-server) to learn how to configure the server to use multiple GPUs, address common issues, etc.\n",
        "\n",
        "You can also host [BLOOMZ](https://huggingface.co/bigscience/bloomz), a version of BLOOM fine-tuned to follow human instructions in the zero-shot regime — just replace `bloom-petals` with `bloomz-petals`.\n",
        "\n",
        "🔒 Hosting a server does not allow others to run custom code on your computer. Learn more about security [here](https://github.com/bigscience-workshop/petals/wiki/Security,-privacy,-and-AI-safety).\n",
        "\n",
        "💬 If you have any issues or feedback, let us know on [our Discord server](https://discord.gg/D9MwApKgWa)!"
      ],
      "metadata": {
        "id": "iK_iT8J3zDC0"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Step 7. Using other fine-tuning and prompt-tuning methods\n",
        "\n",
        "While you can write your own custom adapters, Petals implements several [standard](https://arxiv.org/abs/2104.08691) [methods](https://arxiv.org/abs/2101.00190) for parameter-efficient fine-tuning. We provide a couple of advanced examples in our GitHub repository:\n",
        "\n",
        "- Training a personified chatbot: [notebook](https://github.com/bigscience-workshop/petals/blob/main/examples/prompt-tuning-personachat.ipynb)\n",
        "\n",
        "- Fine-tuning BLOOM for text semantic classification: [notebook](https://github.com/bigscience-workshop/petals/blob/main/examples/prompt-tuning-sst2.ipynb)"
      ],
      "metadata": {
        "id": "fwLMNcJ80ARs"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## What's next?\n",
        "\n",
        "Congratulations on finishing our tutorial! Now, you are familiar with how to use Petals for different tasks, how it works under the hood, and how to increase its capacity.\n",
        "\n",
        "You can find a few other helpful resources below:\n",
        "\n",
        "* __More about Petals.__ The [README](https://github.com/bigscience-workshop/petals#readme) file in our GitHub repository has links to more Petals-related materials, including instructions for starting your own swarm (possibly, with a model other than BLOOM).\n",
        "\n",
        "* __Discord server.__ If you have any feedback, questions, or technical issues, please [join our Discord server](https://discord.gg/D9MwApKgWa) and let us know. If you want to build something based on Petals, we'd be happy to hear what you are up to.\n",
        "\n",
        "* __Academic paper.__ We have released a [paper](https://arxiv.org/abs/2209.01188) that goes into details about our research and what happens in Petals under the hood."
      ],
      "metadata": {
        "id": "PGUpxhQxXVCF"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Other interesting projects\n",
        "\n",
        "- Together: A decentralized cloud for\n",
        "artificial intelligence.\n",
        "Open. Scalable.\n",
        "\n",
        "  https://www.together.xyz"
      ],
      "metadata": {
        "id": "X6pOudJGxrQS"
      }
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "d699cc8313944cc0b6fcb729a946319f": {
          "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_b5d30f01ad0f49aabb1ea83735f70730",
              "IPY_MODEL_29743cb700484e15a386b00c98aadcf9",
              "IPY_MODEL_fe07a97af6c1469b8176a1a4c8888444"
            ],
            "layout": "IPY_MODEL_b67922fe0f5543e399179cb4bcc2bc3a"
          }
        },
        "b5d30f01ad0f49aabb1ea83735f70730": {
          "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_5c27fb179ce2476c86116f5620f4ef59",
            "placeholder": "​",
            "style": "IPY_MODEL_efa3a5f4a9b8438c9ec211bdb1d24dbf",
            "value": "Downloading (…)okenizer_config.json: 100%"
          }
        },
        "29743cb700484e15a386b00c98aadcf9": {
          "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_e7cd84d56240480d897a808d2c8d8664",
            "max": 263,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_8e238ad5782a40b58f5aeece220b1f45",
            "value": 263
          }
        },
        "fe07a97af6c1469b8176a1a4c8888444": {
          "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_909b3e93de3445af866bc3a3eb3e67e7",
            "placeholder": "​",
            "style": "IPY_MODEL_d640b7bddf6746f4bb7181d8746e518f",
            "value": " 263/263 [00:00&lt;00:00, 11.7kB/s]"
          }
        },
        "b67922fe0f5543e399179cb4bcc2bc3a": {
          "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
          }
        },
        "5c27fb179ce2476c86116f5620f4ef59": {
          "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
          }
        },
        "efa3a5f4a9b8438c9ec211bdb1d24dbf": {
          "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": ""
          }
        },
        "e7cd84d56240480d897a808d2c8d8664": {
          "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
          }
        },
        "8e238ad5782a40b58f5aeece220b1f45": {
          "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": ""
          }
        },
        "909b3e93de3445af866bc3a3eb3e67e7": {
          "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
          }
        },
        "d640b7bddf6746f4bb7181d8746e518f": {
          "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": ""
          }
        },
        "a365d95a9e544b0aaf949dfc3c436fd8": {
          "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_0b51448bcffb4c3881c42ef386ac8abe",
              "IPY_MODEL_9ba0bd3f6a1146dbad9baa37ea15ad42",
              "IPY_MODEL_a88549b6c53a45c09f1898aef972562a"
            ],
            "layout": "IPY_MODEL_2c50dee6d3b740c58428ec5867302244"
          }
        },
        "0b51448bcffb4c3881c42ef386ac8abe": {
          "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_311b129becc74fa787c96f8de2f4970f",
            "placeholder": "​",
            "style": "IPY_MODEL_7f48720a85e743cd803351893d0a46c1",
            "value": "Downloading tokenizer.json: 100%"
          }
        },
        "9ba0bd3f6a1146dbad9baa37ea15ad42": {
          "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_3ade1584527c414193588cd88d05d68e",
            "max": 14500443,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_70feae1bdc5e46efadc0ef3228b05847",
            "value": 14500443
          }
        },
        "a88549b6c53a45c09f1898aef972562a": {
          "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_8e40327db5cc4912bfb1bba1c92aea30",
            "placeholder": "​",
            "style": "IPY_MODEL_dd1abdbaf11046739d60dba36f38dfd7",
            "value": " 14.5M/14.5M [00:00&lt;00:00, 40.3MB/s]"
          }
        },
        "2c50dee6d3b740c58428ec5867302244": {
          "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
          }
        },
        "311b129becc74fa787c96f8de2f4970f": {
          "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
          }
        },
        "7f48720a85e743cd803351893d0a46c1": {
          "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": ""
          }
        },
        "3ade1584527c414193588cd88d05d68e": {
          "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
          }
        },
        "70feae1bdc5e46efadc0ef3228b05847": {
          "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": ""
          }
        },
        "8e40327db5cc4912bfb1bba1c92aea30": {
          "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
          }
        },
        "dd1abdbaf11046739d60dba36f38dfd7": {
          "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": ""
          }
        },
        "8b700d043a0041e2b3080b9a20ff8acd": {
          "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_76ab08681ded4ac6ad38a73b6c7ead1b",
              "IPY_MODEL_7ad076f3e73e4ac1a06c49b453f58429",
              "IPY_MODEL_f49410314aad47b793fc1ad7eaa97107"
            ],
            "layout": "IPY_MODEL_5b7c66755535426a8bde372a2863650f"
          }
        },
        "76ab08681ded4ac6ad38a73b6c7ead1b": {
          "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_62271c6779614341a91c5279470ce2c5",
            "placeholder": "​",
            "style": "IPY_MODEL_b8501505088c499e958bf251f3296bd8",
            "value": "Downloading (…)cial_tokens_map.json: 100%"
          }
        },
        "7ad076f3e73e4ac1a06c49b453f58429": {
          "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_4204d9413bc945d786a16c488c4b4c6b",
            "max": 96,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_c1dfd895f3874ab68a893321005e52a5",
            "value": 96
          }
        },
        "f49410314aad47b793fc1ad7eaa97107": {
          "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_46c94ebababc44e5b96c0639e7ae37d4",
            "placeholder": "​",
            "style": "IPY_MODEL_0310bcdbd9d04daca17bf219479a83a2",
            "value": " 96.0/96.0 [00:00&lt;00:00, 7.09kB/s]"
          }
        },
        "5b7c66755535426a8bde372a2863650f": {
          "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
          }
        },
        "62271c6779614341a91c5279470ce2c5": {
          "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
          }
        },
        "b8501505088c499e958bf251f3296bd8": {
          "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": ""
          }
        },
        "4204d9413bc945d786a16c488c4b4c6b": {
          "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
          }
        },
        "c1dfd895f3874ab68a893321005e52a5": {
          "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": ""
          }
        },
        "46c94ebababc44e5b96c0639e7ae37d4": {
          "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
          }
        },
        "0310bcdbd9d04daca17bf219479a83a2": {
          "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": ""
          }
        },
        "f80871b3d11d49d986961a8f76e8dcd9": {
          "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_1fda61ef05424c3eafbfde69ea8a5c06",
              "IPY_MODEL_e02b6473cf284fcf8a8e8765c794cd14",
              "IPY_MODEL_154a1e132f004fa3b2a233d2e8ea2689"
            ],
            "layout": "IPY_MODEL_6df211b873f9437ca34a80663b868207"
          }
        },
        "1fda61ef05424c3eafbfde69ea8a5c06": {
          "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_0d3d2bb3103f4b0c85084ff79a606997",
            "placeholder": "​",
            "style": "IPY_MODEL_6146ecfb133e4656bcfb0ae7c618c50e",
            "value": "Downloading (…)lve/main/config.json: 100%"
          }
        },
        "e02b6473cf284fcf8a8e8765c794cd14": {
          "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_6d87070ff58343bfaced6ee8977171b5",
            "max": 641,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_f52b7b4d3fd84054ae8b1ba83acd9cd9",
            "value": 641
          }
        },
        "154a1e132f004fa3b2a233d2e8ea2689": {
          "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_b59b8fb5981f40e594083c2f4146477a",
            "placeholder": "​",
            "style": "IPY_MODEL_7f186445fda344fc95934431cf8903aa",
            "value": " 641/641 [00:00&lt;00:00, 29.3kB/s]"
          }
        },
        "6df211b873f9437ca34a80663b868207": {
          "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
          }
        },
        "0d3d2bb3103f4b0c85084ff79a606997": {
          "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
          }
        },
        "6146ecfb133e4656bcfb0ae7c618c50e": {
          "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": ""
          }
        },
        "6d87070ff58343bfaced6ee8977171b5": {
          "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
          }
        },
        "f52b7b4d3fd84054ae8b1ba83acd9cd9": {
          "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": ""
          }
        },
        "b59b8fb5981f40e594083c2f4146477a": {
          "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
          }
        },
        "7f186445fda344fc95934431cf8903aa": {
          "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": ""
          }
        },
        "0d3f25decf5a4067a9059b58c3220f88": {
          "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_342ad5f476a541858eb8df963be00434",
              "IPY_MODEL_57657e5d5cf04efcb0ef41d54d9cc48e",
              "IPY_MODEL_92dfd8d47b80413bab1efc792f98f6e8"
            ],
            "layout": "IPY_MODEL_6ea346f9fe674099b57b5500884bb325"
          }
        },
        "342ad5f476a541858eb8df963be00434": {
          "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_11cc3f308ab94f1583febaa4d77d46ec",
            "placeholder": "​",
            "style": "IPY_MODEL_4ffd0800b6004fd1b8bfa495a6e4b610",
            "value": "Downloading pytorch_model.bin: 100%"
          }
        },
        "57657e5d5cf04efcb0ef41d54d9cc48e": {
          "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_62396b1ef8b540faace8a512b5ec5c2b",
            "max": 7193348019,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_7e05f0ea127845eba8b89524c9a40465",
            "value": 7193348019
          }
        },
        "92dfd8d47b80413bab1efc792f98f6e8": {
          "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_dd541d8439c244eb8d886fc817bf9551",
            "placeholder": "​",
            "style": "IPY_MODEL_e5cb3e4076ef460595ce40de0fd16465",
            "value": " 7.19G/7.19G [00:48&lt;00:00, 143MB/s]"
          }
        },
        "6ea346f9fe674099b57b5500884bb325": {
          "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
          }
        },
        "11cc3f308ab94f1583febaa4d77d46ec": {
          "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
          }
        },
        "4ffd0800b6004fd1b8bfa495a6e4b610": {
          "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": ""
          }
        },
        "62396b1ef8b540faace8a512b5ec5c2b": {
          "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
          }
        },
        "7e05f0ea127845eba8b89524c9a40465": {
          "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": ""
          }
        },
        "dd541d8439c244eb8d886fc817bf9551": {
          "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
          }
        },
        "e5cb3e4076ef460595ce40de0fd16465": {
          "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
}