llm-course/Decoding_Strategies_in_Large_Language Models.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/mlabonne/how-to-data-science/blob/main/Decoding_Strategies_in_Large_Language%C2%A0Models.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Decoding Strategies in Large Language Models\n",
        "\n",
        "> A Guide to Text Generation From Beam Search to Nucleus Sampling\n",
        "\n",
        "❤️ Created by [@maximelabonne](https://twitter.com/maximelabonne).\n",
        "\n",
        "Companion notebook to execute the code from the following article: https://mlabonne.github.io/blog/decoding/"
      ],
      "metadata": {
        "id": "qaLKx40NbTD6"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cBr2gYVBVDka"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "\n",
        "# Install transformers and graphviz\n",
        "!sudo apt-get install graphviz graphviz-dev\n",
        "!pip install transformers pygraphviz\n",
        "\n",
        "# Make sure we're using UTF-8 as encoding\n",
        "import locale\n",
        "locale.getpreferredencoding = lambda: \"UTF-8\"\n",
        "\n",
        "# Set seed\n",
        "import torch\n",
        "torch.manual_seed(42)\n",
        "torch.cuda.manual_seed(42)\n",
        "torch.cuda.manual_seed_all(42)\n",
        "torch.backends.cudnn.deterministic = True\n",
        "torch.backends.cudnn.benchmark = False"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 📚 Background"
      ],
      "metadata": {
        "id": "7R4IKg4lbNcN"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 229,
          "referenced_widgets": [
            "9713d401d53341e38141920ea03dbd5a",
            "7a4520eef6b046049ac015596f45e765",
            "f630193167224e92aef625d875d74f03",
            "8367416a0fae4c46887b947b4627c9ed",
            "3d8a0c7938274453989b93a819dfbaae",
            "24dfcd472872492ab9dd78704d01f391",
            "928ac9d1df6348da8b937dd4baed73bc",
            "df124877f6fd4d71be3dc3081577d963",
            "abdd1746c797421688fff22da1c24471",
            "69fe044f596f44d48ad092a6a605b8a7",
            "d30b4657252f44cc9f18dfce198fa490",
            "64efc321afdd4b3782765a02888a772b",
            "17831ee3579f4013a5381400fd70b38d",
            "9a35e7be906641fdbf6170cd73a15fa6",
            "5a78d9f4fa474d68a1dea3999bb961da",
            "ddb595c5abd54c81a6589063a0e9c472",
            "9ae75f0fed3846298fa003bc700cb30c",
            "974039b5318043569c2e5f04ab321efe",
            "d44421e59afd4cafa21cde90bd16bd1a",
            "b3dc8bf4a0814f2ba6ece0544d6ffba1",
            "6b27b86bcd4f440da024f952cde0fb86",
            "feda016e22184435b0beb5ee6a6c70a3",
            "5443950734244bb58cc0d16f3a8e6431",
            "7518763d64a643f88c6951d4a44da53c",
            "c705fa5b9cbf4234a1b6fb5f534b6c58",
            "08ef10859c1642c6a4062efdc8adb94d",
            "762ea8e309a94a3584228c7102d6468c",
            "33703304a4a74f608ca5387effe26df0",
            "698d554641de4a518ac8387ec70fbde0",
            "d1403cebed8f4963b1749398d17b1535",
            "9eead7b6eca74c5ab5a48a1917f21df8",
            "ef702a89b5b347ef93cffad017b14681",
            "a74ffb45ba3742c587b42a18344bf688",
            "a53f23ee62874333a78984b0043f04c0",
            "56ba1e0e45884b0dbb21c090f2969532",
            "78db0251ac9146809faf1401165d1b4e",
            "4e867868b7bc4b4c8fd7a91df446931a",
            "3a54459246c6490faf11ea96d1ee0d8e",
            "2fba37301c8a42e2b96c98028c78eee7",
            "464ad324599048089f36a5a78d27afd4",
            "1c19b681477d45e098ab58d3a85dcf31",
            "87ceaae75dda4979b97c63958929b3c6",
            "504924887692443e9791a861e79e3683",
            "a53a20a3248d460bb728b8931629328b",
            "85b0326d03b94a2a81a4af035d011cad",
            "285e6fce050c4bb591df0f93be22e4c3",
            "b00326a362f24a098ddee078671e2efe",
            "1e3601441cd0480aa6443b934d38bce1",
            "ce38696b9d844577ba0a930e50b65f26",
            "02a3d9c1a17141e589bc3cc938a60b05",
            "cae537a671ec4796b414e67987e40108",
            "6c7928af450e4cf89420dcafe21cef60",
            "9db081d87e174740b247f80a48065170",
            "0805e191fe564eaca6ff8545512c366b",
            "1833da9abbfe4090b09c8c9ddcce41a2"
          ]
        },
        "id": "1LS5sCUPwzaD",
        "outputId": "147addb9-bbf7-4eeb-bfc0-3b46b06b076d"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading (…)lve/main/config.json:   0%|          | 0.00/665 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "9713d401d53341e38141920ea03dbd5a"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading pytorch_model.bin:   0%|          | 0.00/548M [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "64efc321afdd4b3782765a02888a772b"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading (…)neration_config.json:   0%|          | 0.00/124 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "5443950734244bb58cc0d16f3a8e6431"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading (…)olve/main/vocab.json:   0%|          | 0.00/1.04M [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "a53f23ee62874333a78984b0043f04c0"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading (…)olve/main/merges.txt:   0%|          | 0.00/456k [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "85b0326d03b94a2a81a4af035d011cad"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
            "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Generated text: I have a dream of being a doctor.\n"
          ]
        }
      ],
      "source": [
        "from transformers import GPT2LMHeadModel, GPT2Tokenizer\n",
        "import torch\n",
        "\n",
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "model = GPT2LMHeadModel.from_pretrained('gpt2').to(device)\n",
        "tokenizer = GPT2Tokenizer.from_pretrained('gpt2')\n",
        "model.eval()\n",
        "\n",
        "text = \"I have a dream\"\n",
        "input_ids = tokenizer.encode(text, return_tensors='pt').to(device)\n",
        "\n",
        "outputs = model.generate(input_ids, max_length=len(input_ids.squeeze())+5)\n",
        "generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
        "print(f\"Generated text: {generated_text}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 🏃‍♂️ Greedy Search"
      ],
      "metadata": {
        "id": "xw7mElwjbPYW"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zm_AiUhK8yrP",
        "outputId": "74bef07c-cc93-4643-d67c-5ba67dfed8eb"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Generated text: I have a dream of being a doctor.\n"
          ]
        }
      ],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import networkx as nx\n",
        "import numpy as np\n",
        "import time\n",
        "\n",
        "def get_log_prob(logits, token_id):\n",
        "    # Compute the softmax of the logits\n",
        "    probabilities = torch.nn.functional.softmax(logits, dim=-1)\n",
        "    log_probabilities = torch.log(probabilities)\n",
        "    \n",
        "    # Get the log probability of the token\n",
        "    token_log_probability = log_probabilities[token_id].item()\n",
        "    return token_log_probability\n",
        "\n",
        "def greedy_search(input_ids, node, length=5):\n",
        "    if length == 0:\n",
        "        return input_ids\n",
        "\n",
        "    outputs = model(input_ids)\n",
        "    predictions = outputs.logits\n",
        "\n",
        "    # Get the predicted next sub-word (here we use top-k search)\n",
        "    logits = predictions[0, -1, :]\n",
        "    token_id = torch.argmax(logits).unsqueeze(0)\n",
        "\n",
        "    # Compute the score of the predicted token\n",
        "    token_score = get_log_prob(logits, token_id)\n",
        "\n",
        "    # Add the predicted token to the list of input ids\n",
        "    new_input_ids = torch.cat([input_ids, token_id.unsqueeze(0)], dim=-1)\n",
        "\n",
        "    # Add node and edge to graph\n",
        "    next_token = tokenizer.decode(token_id, skip_special_tokens=True)\n",
        "    current_node = list(graph.successors(node))[0]\n",
        "    graph.nodes[current_node]['tokenscore'] = np.exp(token_score) * 100\n",
        "    graph.nodes[current_node]['token'] = next_token + f\"_{length}\"\n",
        "\n",
        "    # Recursive call\n",
        "    input_ids = greedy_search(new_input_ids, current_node, length-1)\n",
        "    \n",
        "    return input_ids\n",
        "\n",
        "# Parameters\n",
        "length = 5\n",
        "beams = 1\n",
        "\n",
        "# Create a balanced tree with height 'length'\n",
        "graph = nx.balanced_tree(1, length, create_using=nx.DiGraph())\n",
        "\n",
        "# Add 'tokenscore', 'cumscore', and 'token' attributes to each node\n",
        "for node in graph.nodes:\n",
        "    graph.nodes[node]['tokenscore'] = 100\n",
        "    graph.nodes[node]['token'] = text\n",
        "\n",
        "# Start generating text\n",
        "output_ids = greedy_search(input_ids, 0, length=length)\n",
        "output = tokenizer.decode(output_ids.squeeze().tolist(), skip_special_tokens=True)\n",
        "print(f\"Generated text: {output}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "NeMNyl5maqTc",
        "outputId": "0d4a5d12-0c05-416f-e5d7-99721b7f7214"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 3633.75x2100 with 2 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAACo0AAAadCAYAAAAI95aAAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAC4jAAAuIwF4pT92AAEAAElEQVR4nOzdd5hdZb0+7mdPJo0EEkoCCS2hd0JVASGAWCmKIsoPjyA29BT9ehSVo4K9t6OioiIesaJSVXpRWujSa0IghZIGqZPJ7N8fexJm78mUPZlkzWTu+7rWxX7Xftd6P7N2hj/29cznLZXL5XIAAAAAAAAAAAAAWK81FF0AAAAAAAAAAAAAAGuf0CgAAAAAAAAAAADAACA0CgAAAAAAAAAAADAACI0CAAAAAAAAAAAADABCowAAAAAAAAAAAAADgNAoAAAAAAAAAAAAwAAgNAoAAAAAAAAAAAAwAAiNAgAAAAAAAAAAAAwAQqMAAAAAAAAAAAAAA4DQKAAAAAAAAAAAAMAAIDQKAAAAAAAAAAAAMAAIjQIAAAAAAAAAAAAMAEKjAAAAAAAAAAAAAAOA0CgAAAAAAAAAAADAACA0CgAAAAAAAAAAADAACI0CAAAAAAAAAAAADABCowAAAAAAAAAAAAADgNAoAAAAAAAAAAAAwAAgNAoAAAAAAAAAAAAwAAiNAgAAAAAAAAAAAAwAQqMAAAAAAAAAAAAAA4DQKAAAAAAAAAAAAMAAIDQKAAAAAAAAAAAAMAAIjQIAAAAAAAAAAAAMAEKjAAAAAAAAAAAAAAOA0CgAAAAAAAAAAADAACA0CgAAAAAAAAAAADAACI0CAAAAAAAAAAAADABCowAAAAAAAAAAAAADgNAoAAAAAAAAAAAAwADQWHQBAAAAAAAAAAAAwNqxdOnS3HzzzXn44Yczb968DBkyJFtttVVe8YpXZLvttuvVtZ544olMmTIlzzzzTJqamrLxxhtnl112yUEHHZRhw4b16lr0jNAoAAAAAAAAAAAArCMzZszIlClTctttt2XKlCm544478tJLL616f9ttt820adPWeJ3nn38+Z599dn75y19m0aJFq52z33775TOf+UyOO+64NVrroosuyhe+8IXcddddq31/5MiROeWUU/K5z30um2222RqtxZoplcvlctFFAAAAAAAAAAAAwPrqpptuyre+9a3cdtttmTlzZqdzeyM0ev311+eEE07ICy+80K35//Zv/5Zzzz03Q4YMqWudZcuW5bTTTssFF1zQrfljxozJhRdemEMPPbSudeg9DUUXAAAAAAAAAAAAAOuz22+/PX/5y1+6DIz2hn/+85954xvf2C4wOnr06Oyzzz6ZMGFCBg0aVPXer371q7zzne9MPT0oW1pacuKJJ7YLjA4aNCgTJ07MpEmTMmrUqKr3nn/++bzhDW/ILbfcUudPRW8RGgUAAAAAAAAAAICCjBw5stfuNW/evJx44olZsmTJqnPbbrttLrroosydOzd33XVXpk6dmmnTpuUDH/hA1bV//vOf853vfKfba33jG9/IxRdfXHXugx/8YKZPn54nn3wyd999d+bOnZs///nP2WabbVbNWbx4cd7+9rdnwYIFPfwpWRNCowAAAAAAAAAAALAObLjhhpk8eXI+/vGP549//GOmTZuWSy+9tNfu/41vfKOqm+nEiRNz880357jjjkupVFp1fquttsqPf/zjfOlLX6q6/vOf/3zmzZvX5Tpz5sxpd+1XvvKVnHPOORk/fvyqcw0NDXnLW96Sm2++ORMmTFh1/plnnsm3v/3ten88ekGpXE8/WQAAAAAAAAAAAKAuTzzxRJYtW5ZddtklDQ3VvR6vv/76HH744avG2267baZNm1b3Gs8//3y22267LFy4cNW5q6++OkceeWSH15TL5UyePDk33njjqnOf/vSn2wVCa51xxhn5+te/vmp86KGH5vrrr68Kpta65ppr8prXvGbVeMMNN8zUqVOz6aabdroWvUunUQAAAAAAAAAAAFiLtt9+++y2227tAqO96Xe/+11VYPTQQw/tNDCaJKVSKZ/73Oeqzv3iF79IZ70oW1pact5551WdO+usszoNjCbJkUcemVe/+tWrxi+99FL+8Ic/dHoNvU9oFAAAAAAAAAAAAPq5iy++uGp82mmndeu6ww8/PBMnTlw1nj17dm699dYO59988815/vnnV4232267TJ48uVtr1dZ00UUXdes6eo/QKAAAAAAAAAAAAPRjCxcurNpiPkle+9rXduvaUqlUtW18klx22WUdzr/88surxkcddVSXXUbbzm3r+uuvz6JFi7p1Lb1DaBQAAAAAAAAAAAD6sQceeCDLly9fNZ44cWK22GKLbl9/8MEHV43vueeeDufWvnfQQQd1e53x48dnwoQJq8ZNTU158MEHu309a05oFAAAAAAAAAAAAPqxhx56qGq822671XV97fza+xW1Fr1PaBQAAAAAAAAAAAD6sUceeaRqvPXWW9d1fe38p556KkuXLm03b8mSJZk+fXqvrlVbO2tXY9EFAAAAAAAAAAAAUG3p0qV54oknii5jvTBnzpzMmzcvW2+9dYYOHVrXtWPGjMnYsWPXUmW957nnnqsab7XVVnVdv/nmm6exsTHNzc1JkpaWlsyZMydbbrll1bwXXngh5XJ51Xjw4MF1P5/ae9bWztolNAoAAAAAAAAAANDHPPHEE9ljjz2KLmPA+9znPpezzjqr6DK6tHDhwqrxiBEj6rq+VCpl+PDheemllzq85+rObbDBBimVSnWtVVvb6tZh7bE9PQAAAAAAAAAAAPRjtcHLYcOG1X2P4cOHd3rPdbkOa4/QKAAAAAAAAAAAAPRjS5curRoPGTKk7nsMHTq0arxkyZLC1mHtsT09AAAAAAAAAABAX/eOJJsUXUQ/MzfJ714eXnTRRdlhhx3qusWYMWN6t6a1pLbjZ1NTU933WLZsWaf3XJfrsPYIjQIAAAAAAAAAAPR1myQZW3QR/dsOO+yQ3Xffvegy1oqRI0dWjWs7gnZHbcfP2nuuy3VYe2xPDwAAAAAAAAAAAP1YbfBy0aJFdV1fLpd7FBpdvHhxyuVyXWvV1iY0um4JjQIAAAAAAAAAAEA/NnZsdRvaZ555pq7rn3322TQ3N68aNzQ0ZLPNNms3b7PNNkupVFo1Xr58eZ577rm61poxY0bVuLZ21i6hUQAAAAAAAAAAgH6g5KjrGEh23nnnqvH06dPrur52/rbbbpthw4a1mzd8+PBss802vbrWLrvsUtf1rBmhUQAAAAAAAAAAAOjHaoOXDz74YF3XP/TQQ53er6i16H1CowAAAAAAAAAAANCP7b777hk8ePCq8bRp0zJr1qxuX3/TTTdVjSdNmtTh3Nr3br755m6vM2vWrEybNm3VePDgwdltt926fT1rTmgUAAAAAAAAAAAA+rENN9wwhx56aNW5q666qlvXlsvlXH311VXnjjnmmA7nH3300VXjq6++OuVyuVtrXXnllVXjww8/PCNHjuzWtfQOoVEAAAAAAAAAAADo54499tiq8c9//vNuXXfddddl6tSpq8abb755XvGKV3Q4/6CDDspmm222avzkk0/m+uuv79ZatTUdd9xx3bqO3iM0CgAAAAAAAAAAAP3cO97xjowYMWLV+MYbb8y1117b6TXlcjlnn3121blTTz01DQ0dRwsbGhpyyimnVJ07++yzu+w2es011+Qf//jHqvGGG26Yt7/97Z1eQ+8TGgUAAAAAAAAAAIB+buzYsfn3f//3qnPvfe97M3PmzA6v+cpXvpIbb7xx1XjUqFH5+Mc/3uVaZ5xxRtW28jfccEO+9rWvdTh/xowZee9731t17r/+67+qOpaybjQWXQAAAAAAAAAAAACdayglpVLRVfQv5VLSUnQRbdx0001ZsmRJu/P33ntv1Xjp0qW5+uqrV3uP8ePHZ7fddutwjU984hM5//zzM3v27CTJ1KlTc9BBB+X73/9+jjnmmJRa/xE988wz+eIXv5if/OQnVdefeeaZ2WSTTbr8WTbbbLN8+tOfzqc//elV5z71qU9l+
          },
          "metadata": {}
        }
      ],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import networkx as nx\n",
        "import matplotlib.colors as mcolors\n",
        "from matplotlib.colors import LinearSegmentedColormap\n",
        "\n",
        "def plot_graph(graph, length, beams, score):\n",
        "    fig, ax = plt.subplots(figsize=(3+1.2*beams**length, max(5, 2+length)), dpi=300, facecolor='white')\n",
        "\n",
        "    # Create positions for each node\n",
        "    pos = nx.nx_agraph.graphviz_layout(graph, prog=\"dot\")\n",
        "\n",
        "    # Normalize the colors along the range of token scores\n",
        "    if score == 'token':\n",
        "        scores = [data['tokenscore'] for _, data in graph.nodes(data=True) if data['token'] is not None]\n",
        "    elif score == 'sequence':\n",
        "        scores = [data['sequencescore'] for _, data in graph.nodes(data=True) if data['token'] is not None]\n",
        "    vmin = min(scores)\n",
        "    vmax = max(scores)\n",
        "    norm = mcolors.Normalize(vmin=vmin, vmax=vmax)\n",
        "    cmap = LinearSegmentedColormap.from_list('rg', [\"r\", \"y\", \"g\"], N=256) \n",
        "\n",
        "    # Draw the nodes\n",
        "    nx.draw_networkx_nodes(graph, pos, node_size=2000, node_shape='o', alpha=1, linewidths=4, \n",
        "                          node_color=scores, cmap=cmap)\n",
        "\n",
        "    # Draw the edges\n",
        "    nx.draw_networkx_edges(graph, pos)\n",
        "\n",
        "    # Draw the labels\n",
        "    if score == 'token':\n",
        "        labels = {node: data['token'].split('_')[0] + f\"\\n{data['tokenscore']:.2f}%\" for node, data in graph.nodes(data=True) if data['token'] is not None}\n",
        "    elif score == 'sequence':\n",
        "        labels = {node: data['token'].split('_')[0] + f\"\\n{data['sequencescore']:.2f}\" for node, data in graph.nodes(data=True) if data['token'] is not None}\n",
        "    nx.draw_networkx_labels(graph, pos, labels=labels, font_size=10)\n",
        "    plt.box(False)\n",
        "\n",
        "    # Add a colorbar\n",
        "    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)\n",
        "    sm.set_array([])\n",
        "    if score == 'token':\n",
        "        fig.colorbar(sm, ax=ax, orientation='vertical', pad=0, label='Token probability (%)')\n",
        "    elif score == 'sequence':\n",
        "        fig.colorbar(sm, ax=ax, orientation='vertical', pad=0, label='Sequence score')\n",
        "    plt.show()\n",
        "\n",
        "# Plot graph\n",
        "plot_graph(graph, length, 1.5, 'token')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J16z-Mw295Jx"
      },
      "source": [
        "## ⚖️ Beam Search"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49,
          "referenced_widgets": [
            "0d17978b46374b9ba937cbd1f23cc8a6",
            "b27f6d1a93ca40a2a2444a3e577de8ad",
            "167368b8456f46dba33398b8772dff77",
            "79436fe2590443cdac2d6bdb9d1fcc41",
            "a511af9eb3294b198914e93b9a25abd5",
            "ad49b293543841bfa47b063f0f309f72",
            "a53474be60f9448587413ff42a83ebdb",
            "e910d8d0679b404f87b71cc17a60b50d",
            "c5d544e4d2714f9d8d2aa9c840e41b4c",
            "c4989e8977b2418690d0072e6453644e",
            "51e4c1493e564978b312e66e558dff0b"
          ]
        },
        "id": "DTDkrVdutCk1",
        "outputId": "2a1e6009-b541-4dd6-ecb2-0c3f07bf0947"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/63 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "0d17978b46374b9ba937cbd1f23cc8a6"
            }
          },
          "metadata": {}
        }
      ],
      "source": [
        "from tqdm.notebook import tqdm\n",
        "\n",
        "def greedy_sampling(logits, beams):\n",
        "    return torch.topk(logits, beams).indices\n",
        "    \n",
        "def beam_search(input_ids, node, bar, length, beams, sampling, temperature=0.1):\n",
        "    if length == 0:\n",
        "        return None\n",
        "\n",
        "    outputs = model(input_ids)\n",
        "    predictions = outputs.logits\n",
        "\n",
        "    # Get the predicted next sub-word (here we use top-k search)\n",
        "    logits = predictions[0, -1, :]\n",
        "\n",
        "    if sampling == 'greedy':\n",
        "        top_token_ids = greedy_sampling(logits, beams)\n",
        "    elif sampling == 'top_k':\n",
        "        top_token_ids = top_k_sampling(logits, temperature, 20, beams)\n",
        "    elif sampling == 'nucleus':\n",
        "        top_token_ids = nucleus_sampling(logits, temperature, 0.5, beams)\n",
        "\n",
        "    for j, token_id in enumerate(top_token_ids):\n",
        "        bar.update(1)\n",
        "\n",
        "        # Compute the score of the predicted token\n",
        "        token_score = get_log_prob(logits, token_id)\n",
        "        cumulative_score = graph.nodes[node]['cumscore'] + token_score\n",
        "\n",
        "        # Add the predicted token to the list of input ids\n",
        "        new_input_ids = torch.cat([input_ids, token_id.unsqueeze(0).unsqueeze(0)], dim=-1)\n",
        "\n",
        "        # Add node and edge to graph\n",
        "        token = tokenizer.decode(token_id, skip_special_tokens=True)\n",
        "        current_node = list(graph.successors(node))[j]\n",
        "        graph.nodes[current_node]['tokenscore'] = np.exp(token_score) * 100\n",
        "        graph.nodes[current_node]['cumscore'] = cumulative_score\n",
        "        graph.nodes[current_node]['sequencescore'] = 1/(len(new_input_ids.squeeze())) * cumulative_score\n",
        "        graph.nodes[current_node]['token'] = token + f\"_{length}_{j}\"\n",
        "\n",
        "        # Recursive call\n",
        "        beam_search(new_input_ids, current_node, bar, length-1, beams, sampling, 1)\n",
        "\n",
        "# Parameters\n",
        "length = 5\n",
        "beams = 2\n",
        "\n",
        "# Create a balanced tree with height 'length' and branching factor 'k'\n",
        "graph = nx.balanced_tree(beams, length, create_using=nx.DiGraph())\n",
        "bar = tqdm(total=len(graph.nodes))\n",
        "\n",
        "# Add 'tokenscore', 'cumscore', and 'token' attributes to each node\n",
        "for node in graph.nodes:\n",
        "    graph.nodes[node]['tokenscore'] = 100\n",
        "    graph.nodes[node]['cumscore'] = 0\n",
        "    graph.nodes[node]['sequencescore'] = 0\n",
        "    graph.nodes[node]['token'] = text\n",
        "\n",
        "# Start generating text\n",
        "beam_search(input_ids, 0, bar, length, beams, 'greedy', 1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xusUHHBa4lt6",
        "outputId": "a6689abb-cf16-4cb9-b20b-e17c75db992e"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Generated text: I have a dream. I have a dream\n"
          ]
        }
      ],
      "source": [
        "def get_best_sequence(G):\n",
        "    # Create a list of leaf nodes\n",
        "    leaf_nodes = [node for node in G.nodes() if G.out_degree(node)==0]\n",
        "\n",
        "    # Get the leaf node with the highest cumscore\n",
        "    max_score_node = None\n",
        "    max_score = float('-inf')\n",
        "    for node in leaf_nodes:\n",
        "        if G.nodes[node]['sequencescore'] > max_score:\n",
        "            max_score = G.nodes[node]['sequencescore']\n",
        "            max_score_node = node\n",
        "\n",
        "    # Retrieve the sequence of nodes from this leaf node to the root node in a list\n",
        "    path = nx.shortest_path(G, source=0, target=max_score_node)\n",
        "\n",
        "    # Return the string of token attributes of this sequence\n",
        "    sequence = \"\".join([G.nodes[node]['token'].split('_')[0] for node in path])\n",
        "    \n",
        "    return sequence, max_score\n",
        "\n",
        "sequence, max_score = get_best_sequence(graph)\n",
        "print(f\"Generated text: {sequence}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 362
        },
        "id": "3gPbd4dw0J9k",
        "outputId": "1733ac6d-ce76-4344-d3c4-e84bffb1be5e"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 12420x2100 with 2 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAIVkAAAadCAYAAAD7sLEgAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAC4jAAAuIwF4pT92AAEAAElEQVR4nOzce5yWZb0v/s88DIiAeQJKPJ/CAyGllkIpZR5e5mHpSl1r7dZKs11Zu8POTGu1KqvdUnN5WP08pKm1q91Bd6lZbsHloZI8IaCAInISOQrKaWAYYJ7fH0zjzAAzzwwz3MPwfr9e94vne93f67o+93Xz8B9PVblcLgcAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAKCLlIoOAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA9GylogMAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA9W6noAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAQM9WKjoAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAADQs5WKDgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAPRspaIDAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAPVup6AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAEDPVio6AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA0LOVig4AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD0bKWiAwAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD1bqegAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABAz1YqOgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAANCzlYoOAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA9GylogMAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA9W6noAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAQM9WKjoAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAADQs5WKDgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAPRspaIDAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAPVup6AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAEDPVio6AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA0LOVig4AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD0bKWiAwAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD1bqegAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABAz1YqOgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAANCzlYoOAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA9GylogMAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA9W6noAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAQM9WKjoAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAADQs5WKDgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAPRspaIDAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAPVt10QEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAALY3M2bMyNNPP53XXnstdXV12X333XPYYYdl5MiR6du3b2G5yuVynnvuuUycODGLFy9Okrz97W/PUUcdlfe85z2pqqrqtL2WLl2aJ554IjNmzEhNTU369++fgw8+OKNGjcqee+7ZafvQM1QXHQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAYHtx77335rvf/W6ee+65zd4fMGBALrzwwnzrW9/KwIEDt1mudevW5cYbb8wNN9yQefPmbbZnn332yZe+9KV84QtfSO/evTu816RJk/LNb34zDzzwQOrr6ze536tXr3zkIx/Jd7/73QwfPrzD+9CzVJXL5XLRIQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAurO1a9fm4osvzi9+8YuK+gcNGpR77rknJ5xwQhcnS+bOnZuzzz47EyZMqKj/6KOPzn333Ze999673XvdeOON+cpXvpL169e32VtdXZ3rrrsun//859u9Dz1PVblcLhcdAgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACgu6qvr8+5556b++67r9l4r169st9++2XXXXfNrFmzsnz58mb3+/Xrl4cffjjHH398l2VbvHhxRo4cmRkzZjQb33nnnXPQQQelvr4+s2bNSm1tbbP7hx56aMaNG5eBAwdWvNd1112XSy+9dJPxvfbaK0OGDMn8+fOzYMGCTe7feOON+cIXvlDxPvRMpaIDAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAB0Zz/4wQ9y3333NRv7zGc+k1dffTUzZ87MhAkT8sYbb+S3v/1t9ttvv8ae1atX5/zzz8/y5cu7LNuFF16YGTNmNNZ9+/bNDTfckCVLlmTy5MmZOnVqlixZkuuuuy59+/Zt7Js+fXo+8YlPVLzPuHHj8tWvfrXZ2OjRozN+/PjMnz8/zz77bObPn59nnnkmJ554YrO+Sy+9NE8//XQHn5CeoqpcLpeLDgEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA0B0tXbo0Bx54YFauXNk49u///u+54oorNts/b968vP/978/s2bMbx775zW/myiuv7PRsY8aMyamnntpY9+7dOw8//HBOOOGEzfY//vjjOfnkk7Nu3brGsUceeSQf/OAH29xr1KhRGTduXGN95pln5p577kmfPn026a2rq8u5556bP/zhD41jJ5xwQh5//PGKnoueqVR0AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACA7uqaa67JypUrG+sTTjghl19++Rb799577/z4xz9uNnb99ddn6dKlnZ7t3/7t35rVV1xxRU444YQt9p944ombZP/GN77R5j4PPvhgxo0b11jvueeeueOOO9KnT5/N9vfp0yd33nln9txzz8axP/3pTxk7dmybe9FzlYoOAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAADQHdXX1+euu+5qNvbtb387VVVVrc476aST8oEPfKCxXrlyZX7zm990arYXXnghTz/9dGPdv3//XHbZZW3O++pXv5r+/fs31uPGjcuLL77Y6pwf//jHzerPfe5zGTRoUKtzBg8enM9+9rOtrsOOpVR0AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACA7mjcuHF5/fXXG+uDDjooo0ePrmjuxRdf3Ky+9957OzFZct999zWrzz///Oyyyy5tzttll11y3nnnNRtrLdvatWvz0EMPNRv7xCc+UVHGln0PPvhg6urqKppLz1MqOgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAQHf0hz/8oVl98sknp6qqqqK5J598crP6scceS01NTZdlO+WUUyqe2zLbAw88sMXelrmHDh2a/fffv6J9DjjggBx66KGN9cqVK/P4449XnJOepVR0AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACA7mjixInN6pEjR1Y8d8iQITnggAMa67q6ukydOrVTcpXL5Tz//PMdzjZq1Khm9aRJk1IulzfbuzVnsLm9Wq7HjqNUdAAAAAAAAAAAAAAAA
          },
          "metadata": {}
        }
      ],
      "source": [
        "# Plot graph\n",
        "plot_graph(graph, length, beams, 'sequence')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0vmObi_P6bR0"
      },
      "source": [
        "## 🎲 Top-k sampling"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xmsDJWQs-tav"
      },
      "outputs": [],
      "source": [
        "def plot_prob_distribution(probabilities, next_tokens, sampling, potential_nb, total_nb=50):\n",
        "    # Get top k tokens\n",
        "    top_k_prob, top_k_indices = torch.topk(probabilities, total_nb)\n",
        "    top_k_tokens = [tokenizer.decode([idx]) for idx in top_k_indices.tolist()]\n",
        "\n",
        "    # Get next tokens and their probabilities\n",
        "    next_tokens_list = [tokenizer.decode([idx]) for idx in next_tokens.tolist()]\n",
        "    next_token_prob = probabilities[next_tokens].tolist()\n",
        "\n",
        "    # Create figure\n",
        "    plt.figure(figsize=(0.4*total_nb, 5), dpi=300, facecolor='white')\n",
        "    plt.rc('axes', axisbelow=True)\n",
        "    plt.grid(axis='y', linestyle='-', alpha=0.5)\n",
        "    if potential_nb < total_nb:\n",
        "        plt.axvline(x=potential_nb-0.5, ls=':', color='grey', label='Sampled tokens')\n",
        "    plt.bar(top_k_tokens, top_k_prob.tolist(), color='blue')\n",
        "    plt.bar(next_tokens_list, next_token_prob, color='red', label='Selected tokens')\n",
        "    plt.xticks(rotation=45, ha='right', va='top')\n",
        "    plt.gca().spines['top'].set_visible(False)\n",
        "    plt.gca().spines['right'].set_visible(False)\n",
        "    if sampling == 'top_k':\n",
        "        plt.title('Probability distribution of predicted tokens with top-k sampling')\n",
        "    elif sampling == 'nucleus':\n",
        "        plt.title('Probability distribution of predicted tokens with nucleus sampling')\n",
        "    plt.legend()\n",
        "    plt.savefig(f'{sampling}_{time.time()}.png', dpi=300)\n",
        "    plt.close()\n",
        "\n",
        "def top_k_sampling(logits, temperature, top_k, beams, plot=True):\n",
        "    assert top_k >= 1\n",
        "    assert beams <= top_k\n",
        "\n",
        "    indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]\n",
        "    new_logits = torch.clone(logits)\n",
        "    new_logits[indices_to_remove] = float('-inf')\n",
        "\n",
        "    # Convert logits to probabilities\n",
        "    probabilities = torch.nn.functional.softmax(new_logits / temperature, dim=-1)\n",
        "\n",
        "    # Sample n tokens from the resulting distribution\n",
        "    next_tokens = torch.multinomial(probabilities, beams)\n",
        "\n",
        "    # Plot distribution\n",
        "    if plot:\n",
        "        total_prob = torch.nn.functional.softmax(logits / temperature, dim=-1)\n",
        "        plot_prob_distribution(total_prob, next_tokens, 'top_k', top_k)\n",
        "\n",
        "    return next_tokens\n",
        "\n",
        "# Start generating text\n",
        "beam_search(input_ids, 0, bar, length, beams, 'top_k', 1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sUrEPE899TsK",
        "outputId": "da72aba4-398d-4d73-b28c-298bec2fae4c"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Generated text: I have a dream: to be able to\n"
          ]
        }
      ],
      "source": [
        "sequence, max_score = get_best_sequence(graph)\n",
        "print(f\"Generated text: {sequence}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 363
        },
        "id": "qyop3NTp--4G",
        "outputId": "3347ffa8-7f1a-43b0-f611-f872b741ae0e"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 12420x2100 with 2 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAIT8AAAadCAYAAAAnScUXAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAC4jAAAuIwF4pT92AAEAAElEQVR4nOzceZxddZ3n//e9qSwkwbAkERIIi0FWQ5BFCIpRRH3wEBhoge4eu4WGUdFpdabbAW1bRR0b0FFof7IK6igtCq2AKGOgCUgTAQlJIAmEkH3f10oqlVTd3x8pi6oiSS2pyqlKns/H4zyoz7nf7zmvOofiP26pUqlUAgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA7HPKRQcAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAADFKBcdAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAFKNcdAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAFCMctEBAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABAMcpFBwAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAMUoFx0AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAUo1x0AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAUIxy0QEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAEAxykUHAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAxSgXHQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABSjXHQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABQjHLRAQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAQDHKRQcAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAADFKBcdAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAFKNcdAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAFCMctEBAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABAMcpFBwAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAMUoFx0AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAUo1x0AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAUIxy0QEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAEAxykUHAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAxSgXHQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABSjXHQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABQjHLRAQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAQDHKRQcAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAADFKBcdAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAFKNcdAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAFCMctEBAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABAMcpFBwAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAMUoFx0AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAUo6roAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACAnmbWrFl5/vnns3DhwtTW1ubAAw/McccdlzFjxqRfv36FdVUqlbz44ouZPHlyli9fniR561vfmpNPPjnvfOc7UyqVCmuje6oqOgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAoKd48MEH841vfCMvvvjiDj8fOHBgrrjiinz1q1/N4MGD91jX1q1bc8stt+Tmm2/OokWLdrjmsMMOy+c///l89rOfTe/evfdYG91bqVKpVIqOAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAADozrZs2ZKrrroq9957b5vWDxkyJA888EDOOeecLi5LFixYkIsuuiiTJk1q0/pTTz01Dz30UIYPH97FZfQEpUqlUik6AgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACgu6qvr88ll1yShx56qNn5Xr16ZcSIERk0aFDmzJmTdevWNfu8f//+efzxx3PWWWd1Wdvy5cszZsyYzJo1q9n5/fbbL0cffXTq6+szZ86c1NTUNPv8mGOOyYQJEzJ48OAua6NnKBcdAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAADQnX3729/OQw891Ozcpz71qcyfPz+zZ8/OpEmTsnr16vzqV7/KiBEjGtds2rQpl112WdatW9dlbVdccUVmzZrVOPfr1y8333xzVq5cmalTp2b69OlZuXJlvvvd76Zfv36N62bOnJm/+7u/67Iueo5SpVKpFB0BAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAANAdrVq1KkcddVQ2bNjQeO5f/uVfct111+1w/aJFi/Lud787c+fObTz3la98Jddff32nt40bNy4f+tCHGufevXvn8ccfzznnnLPD9U899VTOO++8bN26tfHcE088kfe9732d3kbPUS46AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACgu7rpppuyYcOGxvmcc87Jtddeu9P1w4cPzw9/+MNm5773ve9l1apVnd72z//8z83m6667Luecc85O17/3ve99U/uXv/zlTu+iZylVKpVK0REAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAN1NfX19DjnkkKxYsaLx3BNPPJH3ve99re4955xz8vTTTzfOt956a6655ppOa3v55ZczatSoxnnAgAFZsmRJ9t9//13u27BhQw499NBUV1c3nps+fXqOP/74TmujZykXHQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA0B1NmDAhK1asaJyPPvrojB07tk17r7rqqmbzgw8+2IllyUMPPdRsvuyyy7L//vu3um///ffPpZde2uxcZ7fRs5SLDgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA6I5++9vfNpvPO++8lEqlNu0977zzms1PPvlkqquru6ztgx/8YJv3tmx75JFHOqWJnqlcdAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAQHc0efLkZvOYMWPavHfYsGE58sgjG+fa2tpMnz69U7oqlUpeeumlDredffbZzeYpU6akUql0Shs9T7noAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACA7uiVV15pNp9wwgnt2t9yfcvrddS8efOyadOmxnnAgAEZMWJEm/cfccQR6d+/f+NcXV2dBQsWdEobPU+56AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAgO5m8+bNmT9/frNzhx9+eLuu0XL9jBkzdrtrR9dpb9eO9nRWGz1PVdEBAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAALA3qampyaxZs4rO2CusWrUqa9asyeGHH56+ffu2a++QIUMydOjQDt975cqVqVQqjXPv3r3bfb3hw4c3m5cvX97hnl1d57DDDmv3NYYPH54ZM2bs9JrsO6qKDgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACAvcmsWbNy0kknFZ2xz/vqV7+ar33tax3ev3HjxmZz//79UyqV2nWNAQMG7PKaHdXyOi3v0xZd1UbPUy46AAAAAAAAAAAAA
          },
          "metadata": {}
        }
      ],
      "source": [
        "# Plot graph\n",
        "plot_graph(graph, length, beams, 'sequence')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qDd_PePe_Itd"
      },
      "source": [
        "## 🔬 Nucleus sampling"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "exsgQ3NQ8oTS"
      },
      "outputs": [],
      "source": [
        "def nucleus_sampling(logits, temperature, p, beams, plot=True):\n",
        "    assert p > 0\n",
        "    assert p <= 1\n",
        "\n",
        "    # Sort the probabilities in descending order and compute cumulative probabilities\n",
        "    sorted_logits, sorted_indices = torch.sort(logits, descending=True)\n",
        "    probabilities = torch.nn.functional.softmax(sorted_logits / temperature, dim=-1)\n",
        "    cumulative_probabilities = torch.cumsum(probabilities, dim=-1)\n",
        "\n",
        "    # Create a mask for probabilities that are in the top-p\n",
        "    mask = cumulative_probabilities < p\n",
        "\n",
        "    # If there's not n index where cumulative_probabilities < p, we use the top n tokens instead\n",
        "    if mask.sum() > beams:\n",
        "        top_p_index_to_keep = torch.where(mask)[0][-1].detach().cpu().tolist()\n",
        "    else:\n",
        "        top_p_index_to_keep = beams\n",
        "\n",
        "    # Only keep top-p indices\n",
        "    indices_to_remove = sorted_indices[top_p_index_to_keep:]\n",
        "    sorted_logits[indices_to_remove] = float('-inf')\n",
        "\n",
        "    # Sample n tokens from the resulting distribution\n",
        "    probabilities = torch.nn.functional.softmax(sorted_logits / temperature, dim=-1)\n",
        "    next_tokens = torch.multinomial(probabilities, beams)\n",
        "\n",
        "    # Plot distribution\n",
        "    if plot:\n",
        "        total_prob = torch.nn.functional.softmax(logits / temperature, dim=-1)\n",
        "        plot_prob_distribution(total_prob, next_tokens, 'nucleus', top_p_index_to_keep)\n",
        "\n",
        "    return next_tokens\n",
        "\n",
        "# Start generating text\n",
        "beam_search(input_ids, 0, bar, length, beams, 'nucleus', 1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "p6uTqcMnk0qb",
        "outputId": "7bf6678b-e1a9-43aa-b935-6a92eccc2dbf"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Generated text: I have a dream. I have a dream\n"
          ]
        }
      ],
      "source": [
        "sequence, max_score = get_best_sequence(graph)\n",
        "print(f\"Generated text: {sequence}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 363
        },
        "id": "NPHx7ZRL2641",
        "outputId": "48b197b7-34ea-44c6-f7e6-46791083ccb5"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 12420x2100 with 2 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAIT8AAAadCAYAAAAnScUXAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAC4jAAAuIwF4pT92AAEAAElEQVR4nOzceZxddZ3n//e9qSwkwbAkERIIi0FWQ5BFCIpRRH3wEBhoge4eu4WGUdFpdabbAW1bRR0b0FFof7IK6igtCq2AKGOgCUgTAQlJIAmEkH3f10oqlVTd3x8pi6oiSS2pyqlKns/H4zyoz7nf7zmvOofiP26pUqlUAgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA7HPKRQcAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAADFKBcdAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAFKNcdAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAFCMctEBAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABAMcpFBwAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAMUoFx0AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAUo1x0AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAUIxy0QEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAEAxykUHAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAxSgXHQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABSjXHQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABQjHLRAQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAQDHKRQcAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAADFKBcdAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAFKNcdAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAFCMctEBAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABAMcpFBwAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAMUoFx0AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAUo1x0AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAUIxy0QEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAEAxykUHAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAxSgXHQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABSjXHQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABQjHLRAQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAQDHKRQcAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAADFKBcdAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAFKNcdAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAFCMctEBAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABAMcpFBwAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAMUoFx0AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAUo6roAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACAnmbWrFl5/vnns3DhwtTW1ubAAw/McccdlzFjxqRfv36FdVUqlbz44ouZPHlyli9fniR561vfmpNPPjnvfOc7UyqVCmuje6oqOgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAoKd48MEH841vfCMvvvjiDj8fOHBgrrjiinz1q1/N4MGD91jX1q1bc8stt+Tmm2/OokWLdrjmsMMOy+c///l89rOfTe/evfdYG91bqVKpVIqOAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAADozrZs2ZKrrroq9957b5vWDxkyJA888EDOOeecLi5LFixYkIsuuiiTJk1q0/pTTz01Dz30UIYPH97FZfQEpUqlUik6AgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACgu6qvr88ll1yShx56qNn5Xr16ZcSIERk0aFDmzJmTdevWNfu8f//+efzxx3PWWWd1Wdvy5cszZsyYzJo1q9n5/fbbL0cffXTq6+szZ86c1NTUNPv8mGOOyYQJEzJ48OAua6NnKBcdAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAADQnX3729/OQw891Ozcpz71qcyfPz+zZ8/OpEmTsnr16vzqV7/KiBEjGtds2rQpl112WdatW9dlbVdccUVmzZrVOPfr1y8333xzVq5cmalTp2b69OlZuXJlvvvd76Zfv36N62bOnJm/+7u/67Iueo5SpVKpFB0BAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAANAdrVq1KkcddVQ2bNjQeO5f/uVfct111+1w/aJFi/Lud787c+fObTz3la98Jddff32nt40bNy4f+tCHGufevXvn8ccfzznnnLPD9U899VTOO++8bN26tfHcE088kfe9732d3kbPUS46AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACgu7rpppuyYcOGxvmcc87Jtddeu9P1w4cPzw9/+MNm5773ve9l1apVnd72z//8z83m6667Luecc85O17/3ve99U/uXv/zlTu+iZylVKpVK0REAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAN1NfX19DjnkkKxYsaLx3BNPPJH3ve99re4955xz8vTTTzfOt956a6655ppOa3v55ZczatSoxnnAgAFZsmRJ9t9//13u27BhQw499NBUV1c3nps+fXqOP/74TmujZykXHQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA0B1NmDAhK1asaJyPPvrojB07tk17r7rqqmbzgw8+2IllyUMPPdRsvuyyy7L//vu3um///ffPpZde2uxcZ7fRs5SLDgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA6I5++9vfNpvPO++8lEqlNu0977zzms1PPvlkqquru6ztgx/8YJv3tmx75JFHOqWJnqlcdAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAQHc0efLkZvOYMWPavHfYsGE58sgjG+fa2tpMnz69U7oqlUpeeumlDredffbZzeYpU6akUql0Shs9T7noAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACA7uiVV15pNp9wwgnt2t9yfcvrddS8efOyadOmxnnAgAEZMWJEm/cfccQR6d+/f+NcXV2dBQsWdEobPU+56AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAgO5m8+bNmT9/frNzhx9+eLuu0XL9jBkzdrtrR9dpb9eO9nRWGz1PVdEBAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAALA3qampyaxZs4rO2CusWrUqa9asyeGHH56+ffu2a++QIUMydOjQDt975cqVqVQqjXPv3r3bfb3hw4c3m5cvX97hnl1d57DDDmv3NYYPH54ZM2bs9JrsO6qKDgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACAvcmsWbNy0kknFZ2xz/vqV7+ar33tax3ev3HjxmZz//79UyqV2nWNAQMG7PKaHdXyOi3v0xZd1UbPUy46AAAAAAAAAAAAA
          },
          "metadata": {}
        }
      ],
      "source": [
        "# Plot graph\n",
        "plot_graph(graph, length, beams, 'sequence')"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyNHyR+CsHtB4LQo1w5rfXe4",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "9713d401d53341e38141920ea03dbd5a": {
          "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_7a4520eef6b046049ac015596f45e765",
              "IPY_MODEL_f630193167224e92aef625d875d74f03",
              "IPY_MODEL_8367416a0fae4c46887b947b4627c9ed"
            ],
            "layout": "IPY_MODEL_3d8a0c7938274453989b93a819dfbaae"
          }
        },
        "7a4520eef6b046049ac015596f45e765": {
          "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_24dfcd472872492ab9dd78704d01f391",
            "placeholder": "​",
            "style": "IPY_MODEL_928ac9d1df6348da8b937dd4baed73bc",
            "value": "Downloading (…)lve/main/config.json: 100%"
          }
        },
        "f630193167224e92aef625d875d74f03": {
          "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_df124877f6fd4d71be3dc3081577d963",
            "max": 665,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_abdd1746c797421688fff22da1c24471",
            "value": 665
          }
        },
        "8367416a0fae4c46887b947b4627c9ed": {
          "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_69fe044f596f44d48ad092a6a605b8a7",
            "placeholder": "​",
            "style": "IPY_MODEL_d30b4657252f44cc9f18dfce198fa490",
            "value": " 665/665 [00:00&lt;00:00, 47.0kB/s]"
          }
        },
        "3d8a0c7938274453989b93a819dfbaae": {
          "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
          }
        },
        "24dfcd472872492ab9dd78704d01f391": {
          "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
          }
        },
        "928ac9d1df6348da8b937dd4baed73bc": {
          "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": ""
          }
        },
        "df124877f6fd4d71be3dc3081577d963": {
          "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
          }
        },
        "abdd1746c797421688fff22da1c24471": {
          "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": ""
          }
        },
        "69fe044f596f44d48ad092a6a605b8a7": {
          "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
          }
        },
        "d30b4657252f44cc9f18dfce198fa490": {
          "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": ""
          }
        },
        "64efc321afdd4b3782765a02888a772b": {
          "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_17831ee3579f4013a5381400fd70b38d",
              "IPY_MODEL_9a35e7be906641fdbf6170cd73a15fa6",
              "IPY_MODEL_5a78d9f4fa474d68a1dea3999bb961da"
            ],
            "layout": "IPY_MODEL_ddb595c5abd54c81a6589063a0e9c472"
          }
        },
        "17831ee3579f4013a5381400fd70b38d": {
          "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_9ae75f0fed3846298fa003bc700cb30c",
            "placeholder": "​",
            "style": "IPY_MODEL_974039b5318043569c2e5f04ab321efe",
            "value": "Downloading pytorch_model.bin: 100%"
          }
        },
        "9a35e7be906641fdbf6170cd73a15fa6": {
          "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_d44421e59afd4cafa21cde90bd16bd1a",
            "max": 548118077,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_b3dc8bf4a0814f2ba6ece0544d6ffba1",
            "value": 548118077
          }
        },
        "5a78d9f4fa474d68a1dea3999bb961da": {
          "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_6b27b86bcd4f440da024f952cde0fb86",
            "placeholder": "​",
            "style": "IPY_MODEL_feda016e22184435b0beb5ee6a6c70a3",
            "value": " 548M/548M [00:01&lt;00:00, 282MB/s]"
          }
        },
        "ddb595c5abd54c81a6589063a0e9c472": {
          "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
          }
        },
        "9ae75f0fed3846298fa003bc700cb30c": {
          "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
          }
        },
        "974039b5318043569c2e5f04ab321efe": {
          "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": ""
          }
        },
        "d44421e59afd4cafa21cde90bd16bd1a": {
          "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
          }
        },
        "b3dc8bf4a0814f2ba6ece0544d6ffba1": {
          "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": ""
          }
        },
        "6b27b86bcd4f440da024f952cde0fb86": {
          "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
          }
        },
        "feda016e22184435b0beb5ee6a6c70a3": {
          "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": ""
          }
        },
        "5443950734244bb58cc0d16f3a8e6431": {
          "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_7518763d64a643f88c6951d4a44da53c",
              "IPY_MODEL_c705fa5b9cbf4234a1b6fb5f534b6c58",
              "IPY_MODEL_08ef10859c1642c6a4062efdc8adb94d"
            ],
            "layout": "IPY_MODEL_762ea8e309a94a3584228c7102d6468c"
          }
        },
        "7518763d64a643f88c6951d4a44da53c": {
          "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_33703304a4a74f608ca5387effe26df0",
            "placeholder": "​",
            "style": "IPY_MODEL_698d554641de4a518ac8387ec70fbde0",
            "value": "Downloading (…)neration_config.json: 100%"
          }
        },
        "c705fa5b9cbf4234a1b6fb5f534b6c58": {
          "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_d1403cebed8f4963b1749398d17b1535",
            "max": 124,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_9eead7b6eca74c5ab5a48a1917f21df8",
            "value": 124
          }
        },
        "08ef10859c1642c6a4062efdc8adb94d": {
          "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_ef702a89b5b347ef93cffad017b14681",
            "placeholder": "​",
            "style": "IPY_MODEL_a74ffb45ba3742c587b42a18344bf688",
            "value": " 124/124 [00:00&lt;00:00, 10.5kB/s]"
          }
        },
        "762ea8e309a94a3584228c7102d6468c": {
          "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
          }
        },
        "33703304a4a74f608ca5387effe26df0": {
          "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
          }
        },
        "698d554641de4a518ac8387ec70fbde0": {
          "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": ""
          }
        },
        "d1403cebed8f4963b1749398d17b1535": {
          "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
          }
        },
        "9eead7b6eca74c5ab5a48a1917f21df8": {
          "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": ""
          }
        },
        "ef702a89b5b347ef93cffad017b14681": {
          "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
          }
        },
        "a74ffb45ba3742c587b42a18344bf688": {
          "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": ""
          }
        },
        "a53f23ee62874333a78984b0043f04c0": {
          "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_56ba1e0e45884b0dbb21c090f2969532",
              "IPY_MODEL_78db0251ac9146809faf1401165d1b4e",
              "IPY_MODEL_4e867868b7bc4b4c8fd7a91df446931a"
            ],
            "layout": "IPY_MODEL_3a54459246c6490faf11ea96d1ee0d8e"
          }
        },
        "56ba1e0e45884b0dbb21c090f2969532": {
          "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_2fba37301c8a42e2b96c98028c78eee7",
            "placeholder": "​",
            "style": "IPY_MODEL_464ad324599048089f36a5a78d27afd4",
            "value": "Downloading (…)olve/main/vocab.json: 100%"
          }
        },
        "78db0251ac9146809faf1401165d1b4e": {
          "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_1c19b681477d45e098ab58d3a85dcf31",
            "max": 1042301,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_87ceaae75dda4979b97c63958929b3c6",
            "value": 1042301
          }
        },
        "4e867868b7bc4b4c8fd7a91df446931a": {
          "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_504924887692443e9791a861e79e3683",
            "placeholder": "​",
            "style": "IPY_MODEL_a53a20a3248d460bb728b8931629328b",
            "value": " 1.04M/1.04M [00:00&lt;00:00, 23.0MB/s]"
          }
        },
        "3a54459246c6490faf11ea96d1ee0d8e": {
          "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
          }
        },
        "2fba37301c8a42e2b96c98028c78eee7": {
          "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
          }
        },
        "464ad324599048089f36a5a78d27afd4": {
          "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": ""
          }
        },
        "1c19b681477d45e098ab58d3a85dcf31": {
          "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
          }
        },
        "87ceaae75dda4979b97c63958929b3c6": {
          "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": ""
          }
        },
        "504924887692443e9791a861e79e3683": {
          "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
          }
        },
        "a53a20a3248d460bb728b8931629328b": {
          "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": ""
          }
        },
        "85b0326d03b94a2a81a4af035d011cad": {
          "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_285e6fce050c4bb591df0f93be22e4c3",
              "IPY_MODEL_b00326a362f24a098ddee078671e2efe",
              "IPY_MODEL_1e3601441cd0480aa6443b934d38bce1"
            ],
            "layout": "IPY_MODEL_ce38696b9d844577ba0a930e50b65f26"
          }
        },
        "285e6fce050c4bb591df0f93be22e4c3": {
          "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_02a3d9c1a17141e589bc3cc938a60b05",
            "placeholder": "​",
            "style": "IPY_MODEL_cae537a671ec4796b414e67987e40108",
            "value": "Downloading (…)olve/main/merges.txt: 100%"
          }
        },
        "b00326a362f24a098ddee078671e2efe": {
          "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_6c7928af450e4cf89420dcafe21cef60",
            "max": 456318,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_9db081d87e174740b247f80a48065170",
            "value": 456318
          }
        },
        "1e3601441cd0480aa6443b934d38bce1": {
          "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_0805e191fe564eaca6ff8545512c366b",
            "placeholder": "​",
            "style": "IPY_MODEL_1833da9abbfe4090b09c8c9ddcce41a2",
            "value": " 456k/456k [00:00&lt;00:00, 12.2MB/s]"
          }
        },
        "ce38696b9d844577ba0a930e50b65f26": {
          "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
          }
        },
        "02a3d9c1a17141e589bc3cc938a60b05": {
          "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
          }
        },
        "cae537a671ec4796b414e67987e40108": {
          "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": ""
          }
        },
        "6c7928af450e4cf89420dcafe21cef60": {
          "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
          }
        },
        "9db081d87e174740b247f80a48065170": {
          "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": ""
          }
        },
        "0805e191fe564eaca6ff8545512c366b": {
          "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
          }
        },
        "1833da9abbfe4090b09c8c9ddcce41a2": {
          "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": ""
          }
        },
        "0d17978b46374b9ba937cbd1f23cc8a6": {
          "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_b27f6d1a93ca40a2a2444a3e577de8ad",
              "IPY_MODEL_167368b8456f46dba33398b8772dff77",
              "IPY_MODEL_79436fe2590443cdac2d6bdb9d1fcc41"
            ],
            "layout": "IPY_MODEL_a511af9eb3294b198914e93b9a25abd5"
          }
        },
        "b27f6d1a93ca40a2a2444a3e577de8ad": {
          "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_ad49b293543841bfa47b063f0f309f72",
            "placeholder": "​",
            "style": "IPY_MODEL_a53474be60f9448587413ff42a83ebdb",
            "value": ""
          }
        },
        "167368b8456f46dba33398b8772dff77": {
          "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": "",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_e910d8d0679b404f87b71cc17a60b50d",
            "max": 63,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_c5d544e4d2714f9d8d2aa9c840e41b4c",
            "value": 63
          }
        },
        "79436fe2590443cdac2d6bdb9d1fcc41": {
          "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_c4989e8977b2418690d0072e6453644e",
            "placeholder": "​",
            "style": "IPY_MODEL_51e4c1493e564978b312e66e558dff0b",
            "value": " 185/? [01:16&lt;00:00,  2.19it/s]"
          }
        },
        "a511af9eb3294b198914e93b9a25abd5": {
          "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
          }
        },
        "ad49b293543841bfa47b063f0f309f72": {
          "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
          }
        },
        "a53474be60f9448587413ff42a83ebdb": {
          "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": ""
          }
        },
        "e910d8d0679b404f87b71cc17a60b50d": {
          "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
          }
        },
        "c5d544e4d2714f9d8d2aa9c840e41b4c": {
          "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": ""
          }
        },
        "c4989e8977b2418690d0072e6453644e": {
          "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
          }
        },
        "51e4c1493e564978b312e66e558dff0b": {
          "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
}