{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sdBjRnWqLwWP"
      },
      "source": [
        "# Inference notenook for [CapDec](https://github.com/DavidHuji/CapDec)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GRfpGaz27IWs",
        "outputId": "5eb9c6ac-6803-49c3-9c90-26c9d5b35a9d"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: transformers in /usr/local/lib/python3.10/dist-packages (4.35.2)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from transformers) (3.13.1)\n",
            "Requirement already satisfied: huggingface-hub<1.0,>=0.16.4 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.19.4)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (1.23.5)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from transformers) (23.2)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (6.0.1)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (2023.6.3)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from transformers) (2.31.0)\n",
            "Requirement already satisfied: tokenizers<0.19,>=0.14 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.15.0)\n",
            "Requirement already satisfied: safetensors>=0.3.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.4.1)\n",
            "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.10/dist-packages (from transformers) (4.66.1)\n",
            "Requirement already satisfied: fsspec>=2023.5.0 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.16.4->transformers) (2023.6.0)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.16.4->transformers) (4.5.0)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (3.6)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2023.11.17)\n",
            "Collecting git+https://github.com/openai/CLIP.git\n",
            "  Cloning https://github.com/openai/CLIP.git to /tmp/pip-req-build-70bog_pn\n",
            "  Running command git clone --filter=blob:none --quiet https://github.com/openai/CLIP.git /tmp/pip-req-build-70bog_pn\n",
            "  Resolved https://github.com/openai/CLIP.git to commit a1d071733d7111c9c014f024669f959182114e33\n",
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Collecting ftfy (from clip==1.0)\n",
            "  Downloading ftfy-6.1.3-py3-none-any.whl (53 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m53.4/53.4 kB\u001b[0m \u001b[31m923.3 kB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: regex in /usr/local/lib/python3.10/dist-packages (from clip==1.0) (2023.6.3)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from clip==1.0) (4.66.1)\n",
            "Requirement already satisfied: torch in /usr/local/lib/python3.10/dist-packages (from clip==1.0) (2.1.0+cu121)\n",
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.10/dist-packages (from clip==1.0) (0.16.0+cu121)\n",
            "Requirement already satisfied: wcwidth<0.3.0,>=0.2.12 in /usr/local/lib/python3.10/dist-packages (from ftfy->clip==1.0) (0.2.12)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch->clip==1.0) (3.13.1)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from torch->clip==1.0) (4.5.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch->clip==1.0) (1.12)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch->clip==1.0) (3.2.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch->clip==1.0) (3.1.2)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch->clip==1.0) (2023.6.0)\n",
            "Requirement already satisfied: triton==2.1.0 in /usr/local/lib/python3.10/dist-packages (from torch->clip==1.0) (2.1.0)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from torchvision->clip==1.0) (1.23.5)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from torchvision->clip==1.0) (2.31.0)\n",
            "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /usr/local/lib/python3.10/dist-packages (from torchvision->clip==1.0) (9.4.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch->clip==1.0) (2.1.3)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->torchvision->clip==1.0) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->torchvision->clip==1.0) (3.6)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->torchvision->clip==1.0) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->torchvision->clip==1.0) (2023.11.17)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch->clip==1.0) (1.3.0)\n",
            "Building wheels for collected packages: clip\n",
            "  Building wheel for clip (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for clip: filename=clip-1.0-py3-none-any.whl size=1369497 sha256=74add828989105b4f396b328f4854460d45b9930495d4ab42eac55f65893eab0\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-suv503r9/wheels/da/2b/4c/d6691fa9597aac8bb85d2ac13b112deb897d5b50f5ad9a37e4\n",
            "Successfully built clip\n",
            "Installing collected packages: ftfy, clip\n",
            "Successfully installed clip-1.0 ftfy-6.1.3\n"
          ]
        }
      ],
      "source": [
        "#@title Install\n",
        "!pip install transformers\n",
        "! pip install git+https://github.com/openai/CLIP.git\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iqE3Fj5-uYSR",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "6c2330ce-cdcc-49ac-c49c-6e71c6139ff4"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:root:pydrive is deprecated and no longer maintained. We recommend that you migrate your projects to pydrive2, the maintained fork of pydrive\n"
          ]
        }
      ],
      "source": [
        "#@title Drive Downloader\n",
        "\n",
        "from pydrive.auth import GoogleAuth\n",
        "from pydrive.drive import GoogleDrive\n",
        "from google.colab import auth\n",
        "from oauth2client.client import GoogleCredentials\n",
        "\n",
        "download_with_pydrive = True #@param {type:\"boolean\"}\n",
        "\n",
        "class Downloader(object):\n",
        "    def __init__(self, use_pydrive):\n",
        "        self.use_pydrive = use_pydrive\n",
        "\n",
        "        if self.use_pydrive:\n",
        "            self.authenticate()\n",
        "\n",
        "    def authenticate(self):\n",
        "        auth.authenticate_user()\n",
        "        gauth = GoogleAuth()\n",
        "        gauth.credentials = GoogleCredentials.get_application_default()\n",
        "        self.drive = GoogleDrive(gauth)\n",
        "\n",
        "    def download_file(self, file_id, file_dst):\n",
        "        if self.use_pydrive:\n",
        "            downloaded = self.drive.CreateFile({'id':file_id})\n",
        "            downloaded.FetchMetadata(fetch_all=True)\n",
        "            downloaded.GetContentFile(file_dst)\n",
        "        else:\n",
        "            !gdown --id $file_id -O $file_dst\n",
        "\n",
        "downloader = Downloader(download_with_pydrive)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "OArDkm_24w4L"
      },
      "outputs": [],
      "source": [
        "#@title Imports\n",
        "\n",
        "import clip\n",
        "import os\n",
        "from torch import nn\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn.functional as nnf\n",
        "import sys\n",
        "from typing import Tuple, List, Union, Optional\n",
        "from transformers import GPT2Tokenizer, GPT2LMHeadModel, AdamW, get_linear_schedule_with_warmup\n",
        "from tqdm import tqdm, trange\n",
        "from google.colab import files\n",
        "import skimage.io as io\n",
        "import PIL.Image\n",
        "from IPython.display import Image\n",
        "\n",
        "\n",
        "N = type(None)\n",
        "V = np.array\n",
        "ARRAY = np.ndarray\n",
        "ARRAYS = Union[Tuple[ARRAY, ...], List[ARRAY]]\n",
        "VS = Union[Tuple[V, ...], List[V]]\n",
        "VN = Union[V, N]\n",
        "VNS = Union[VS, N]\n",
        "T = torch.Tensor\n",
        "TS = Union[Tuple[T, ...], List[T]]\n",
        "TN = Optional[T]\n",
        "TNS = Union[Tuple[TN, ...], List[TN]]\n",
        "TSN = Optional[TS]\n",
        "TA = Union[T, ARRAY]\n",
        "\n",
        "\n",
        "D = torch.device\n",
        "CPU = torch.device('cpu')\n",
        "\n",
        "\n",
        "def get_device(device_id: int) -> D:\n",
        "    if not torch.cuda.is_available():\n",
        "        return CPU\n",
        "    device_id = min(torch.cuda.device_count() - 1, device_id)\n",
        "    return torch.device(f'cuda:{device_id}')\n",
        "\n",
        "\n",
        "CUDA = get_device\n",
        "\n",
        "current_directory = os.getcwd()\n",
        "save_path = os.path.join(os.path.dirname(current_directory), \"pretrained_models\")\n",
        "os.makedirs(save_path, exist_ok=True)\n",
        "model_path = os.path.join(save_path, 'model_wieghts.pt')\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xE-uUStuv1Nl"
      },
      "outputs": [],
      "source": [
        "#@title Choose The Noise Variance of The Pre-trained Model\n",
        "\n",
        "\n",
        "pretrained_model_variance = \"0.0\" #@param [\"0.0\", \"0.0001\", \"0.001\", \"0.015\", \"0.1\", \"2.5\"]\n",
        "\n",
        "if pretrained_model_variance == \"0.0\":\n",
        "  downloader.download_file(\"1YwNp9GPSUej570E0qznB8HHQKELFYxUD\", model_path)\n",
        "elif pretrained_model_variance == \"0.0001\":\n",
        "  downloader.download_file(\"1YpacW4vm9Zz1C_aidnRX_28QmvQp8Iec\", model_path)\n",
        "elif pretrained_model_variance == \"0.001\":\n",
        "  downloader.download_file(\"1MR_hIJ_W5c7PNsqTA5cRaATHla5lzcqr\", model_path)\n",
        "elif pretrained_model_variance == \"0.015\":\n",
        "  downloader.download_file(\"1K_3ffRqRcF-ftI3ok89Nr5SBJf3ycmi-\", model_path)\n",
        "elif pretrained_model_variance == \"0.1\":\n",
        "  downloader.download_file(\"1TXr5wLWTQKbzz3ft8CI1e-sJl-okUm0a\", model_path)\n",
        "elif pretrained_model_variance == \"2.5\":\n",
        "  downloader.download_file(\"1Va6n7JshAlBJwjoOfMf23n6PPSImheM3\", model_path)\n",
        "else:\n",
        "  model_path=''"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rEobCOOCl45B",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "d5ccebbc-6595-4928-9e05-988d1cf162ff"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/pretrained_models\n"
          ]
        }
      ],
      "source": [
        "print(save_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "4ClW2ebek8DK"
      },
      "outputs": [],
      "source": [
        "#@title Model\n",
        "\n",
        "class ClipCaptionModel(nn.Module):\n",
        "\n",
        "    def get_dummy_token(self, batch_size: int, device: torch.device) -> torch.Tensor:\n",
        "        return torch.zeros(batch_size, self.prefix_length, dtype=torch.int64, device=device)\n",
        "\n",
        "    def forward(self, tokens: torch.Tensor, prefix: torch.Tensor, mask: Optional[torch.Tensor] = None,\n",
        "                labels: Optional[torch.Tensor] = None):\n",
        "        embedding_text = self.gpt.transformer.wte(tokens)\n",
        "        prefix_projections = self.clip_project(prefix).view(-1, self.prefix_length, self.gpt_embedding_size)\n",
        "        embedding_cat = torch.cat((prefix_projections, embedding_text), dim=1)\n",
        "        if labels is not None:\n",
        "            dummy_token = self.get_dummy_token(tokens.shape[0], tokens.device)\n",
        "            labels = torch.cat((dummy_token, tokens), dim=1)\n",
        "        out = self.gpt(inputs_embeds=embedding_cat, labels=labels, attention_mask=mask)\n",
        "        return out\n",
        "\n",
        "    def __init__(self):\n",
        "        super(ClipCaptionModel, self).__init__()\n",
        "        self.prefix_length = 40\n",
        "        self.gpt = GPT2LMHeadModel.from_pretrained('gpt2')\n",
        "        self.gpt_embedding_size = self.gpt.transformer.wte.weight.shape[1]\n",
        "        self.clip_project = TransformerMapper(640, self.gpt_embedding_size, 40,\n",
        "                                                                     40, 8)\n",
        "\n",
        "\n",
        "\n",
        "class MLP(nn.Module):\n",
        "\n",
        "    def forward(self, x: T) -> T:\n",
        "        return self.model(x)\n",
        "\n",
        "    def __init__(self, sizes: Tuple[int, ...], bias=True, act=nn.Tanh):\n",
        "        super(MLP, self).__init__()\n",
        "        layers = []\n",
        "        for i in range(len(sizes) -1):\n",
        "            layers.append(nn.Linear(sizes[i], sizes[i + 1], bias=bias))\n",
        "            if i < len(sizes) - 2:\n",
        "                layers.append(act())\n",
        "        self.model = nn.Sequential(*layers)\n",
        "\n",
        "\n",
        "\n",
        "class ClipCaptionPrefix(ClipCaptionModel):\n",
        "\n",
        "    def parameters(self, recurse: bool = True):\n",
        "        return self.clip_project.parameters()\n",
        "\n",
        "    def train(self, mode: bool = True):\n",
        "        super(ClipCaptionPrefix, self).train(mode)\n",
        "        self.gpt.eval()\n",
        "        return self"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "V7xocT3TUgey"
      },
      "outputs": [],
      "source": [
        "#@title Caption prediction\n",
        "\n",
        "def generate_beam(model, tokenizer, beam_size: int = 5, prompt=None, embed=None,\n",
        "                  entry_length=67, temperature=1., stop_token: str = '.'):\n",
        "\n",
        "    model.eval()\n",
        "    stop_token_index = tokenizer.encode(stop_token)[0]\n",
        "    tokens = None\n",
        "    scores = None\n",
        "    device = next(model.parameters()).device\n",
        "    seq_lengths = torch.ones(beam_size, device=device)\n",
        "    is_stopped = torch.zeros(beam_size, device=device, dtype=torch.bool)\n",
        "    with torch.no_grad():\n",
        "        if embed is not None:\n",
        "            generated = embed\n",
        "        else:\n",
        "            if tokens is None:\n",
        "                tokens = torch.tensor(tokenizer.encode(prompt))\n",
        "                tokens = tokens.unsqueeze(0).to(device)\n",
        "                generated = model.gpt.transformer.wte(tokens)\n",
        "        for i in range(entry_length):\n",
        "            outputs = model.gpt(inputs_embeds=generated)\n",
        "            logits = outputs.logits\n",
        "            logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0)\n",
        "            logits = logits.softmax(-1).log()\n",
        "            if scores is None:\n",
        "                scores, next_tokens = logits.topk(beam_size, -1)\n",
        "                generated = generated.expand(beam_size, *generated.shape[1:])\n",
        "                next_tokens, scores = next_tokens.permute(1, 0), scores.squeeze(0)\n",
        "                if tokens is None:\n",
        "                    tokens = next_tokens\n",
        "                else:\n",
        "                    tokens = tokens.expand(beam_size, *tokens.shape[1:])\n",
        "                    tokens = torch.cat((tokens, next_tokens), dim=1)\n",
        "            else:\n",
        "                logits[is_stopped] = -float(np.inf)\n",
        "                logits[is_stopped, 0] = 0\n",
        "                scores_sum = scores[:, None] + logits\n",
        "                seq_lengths[~is_stopped] += 1\n",
        "                scores_sum_average = scores_sum / seq_lengths[:, None]\n",
        "                scores_sum_average, next_tokens = scores_sum_average.view(-1).topk(beam_size, -1)\n",
        "                next_tokens_source = next_tokens // scores_sum.shape[1]\n",
        "                seq_lengths = seq_lengths[next_tokens_source]\n",
        "                next_tokens = next_tokens % scores_sum.shape[1]\n",
        "                next_tokens = next_tokens.unsqueeze(1)\n",
        "                tokens = tokens[next_tokens_source]\n",
        "                tokens = torch.cat((tokens, next_tokens), dim=1)\n",
        "                generated = generated[next_tokens_source]\n",
        "                scores = scores_sum_average * seq_lengths\n",
        "                is_stopped = is_stopped[next_tokens_source]\n",
        "            next_token_embed = model.gpt.transformer.wte(next_tokens.squeeze()).view(generated.shape[0], 1, -1)\n",
        "            generated = torch.cat((generated, next_token_embed), dim=1)\n",
        "            is_stopped = is_stopped + next_tokens.eq(stop_token_index).squeeze()\n",
        "            if is_stopped.all():\n",
        "                break\n",
        "    scores = scores / seq_lengths\n",
        "    output_list = tokens.cpu().numpy()\n",
        "    output_texts = [tokenizer.decode(output[:int(length)]) for output, length in zip(output_list, seq_lengths)]\n",
        "    order = scores.argsort(descending=True)\n",
        "    output_texts = [output_texts[i] for i in order]\n",
        "    return output_texts\n",
        "\n",
        "\n",
        "def generate2(\n",
        "        model,\n",
        "        tokenizer,\n",
        "        tokens=None,\n",
        "        prompt=None,\n",
        "        embed=None,\n",
        "        entry_count=1,\n",
        "        entry_length=67,  # maximum number of words\n",
        "        top_p=0.8,\n",
        "        temperature=1.,\n",
        "        stop_token: str = '.',\n",
        "):\n",
        "    model.eval()\n",
        "    generated_num = 0\n",
        "    generated_list = []\n",
        "    stop_token_index = tokenizer.encode(stop_token)[0]\n",
        "    filter_value = -float(\"Inf\")\n",
        "    device = next(model.parameters()).device\n",
        "\n",
        "    with torch.no_grad():\n",
        "\n",
        "        for entry_idx in trange(entry_count):\n",
        "            if embed is not None:\n",
        "                generated = embed\n",
        "            else:\n",
        "                if tokens is None:\n",
        "                    tokens = torch.tensor(tokenizer.encode(prompt))\n",
        "                    tokens = tokens.unsqueeze(0).to(device)\n",
        "\n",
        "                generated = model.gpt.transformer.wte(tokens)\n",
        "\n",
        "            for i in range(entry_length):\n",
        "\n",
        "                outputs = model.gpt(inputs_embeds=generated)\n",
        "                logits = outputs.logits\n",
        "                logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0)\n",
        "                sorted_logits, sorted_indices = torch.sort(logits, descending=True)\n",
        "                cumulative_probs = torch.cumsum(nnf.softmax(sorted_logits, dim=-1), dim=-1)\n",
        "                sorted_indices_to_remove = cumulative_probs > top_p\n",
        "                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[\n",
        "                                                    ..., :-1\n",
        "                                                    ].clone()\n",
        "                sorted_indices_to_remove[..., 0] = 0\n",
        "\n",
        "                indices_to_remove = sorted_indices[sorted_indices_to_remove]\n",
        "                logits[:, indices_to_remove] = filter_value\n",
        "                next_token = torch.argmax(logits, -1).unsqueeze(0)\n",
        "                next_token_embed = model.gpt.transformer.wte(next_token)\n",
        "                if tokens is None:\n",
        "                    tokens = next_token\n",
        "                else:\n",
        "                    tokens = torch.cat((tokens, next_token), dim=1)\n",
        "                generated = torch.cat((generated, next_token_embed), dim=1)\n",
        "                if stop_token_index == next_token.item():\n",
        "                    break\n",
        "\n",
        "            output_list = list(tokens.squeeze().cpu().numpy())\n",
        "            output_text = tokenizer.decode(output_list)\n",
        "            generated_list.append(output_text)\n",
        "\n",
        "    return generated_list[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "z3X2FkvhmIQu",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "9890b3bb-1d09-45cd-f52c-03d54354a4c9"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/pretrained_models/model_wieghts.pt\n"
          ]
        }
      ],
      "source": [
        "print(model_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7lCgFHSgr_ny"
      },
      "outputs": [],
      "source": [
        "#@title GPU/CPU\n",
        "\n",
        "\n",
        "is_gpu = True #@param {type:\"boolean\"}\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "6bi_2zQ3QD57",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 163,
          "referenced_widgets": [
            "823a2248d19e473a8457356a5c253c7b",
            "58e74106989c4ca68cd8b5db1beb28d9",
            "b0553d1b15a840c290932a2ef549fbf0",
            "502ca5a845fa4986ad63586a44e867b3",
            "d20be762e5ba434aa52e9cc029ccacf7",
            "d4386577eada433f92b57ce91c3153b8",
            "4ef5487b590e497ca719d501a397cb3d",
            "634c60cd8d934c55853e2811ae0a1a67",
            "69e36359b49d479b9f2898c2bc870ff0",
            "525544b3302843c6b579917a9b5df84c",
            "02a6dda06c344632a2b76a5fa14bea81",
            "dea69fba019c430181378d72844a77e5",
            "959964afad6c4a488cd6f1bb3e82753d",
            "2ce40954abb44071b0cbf0f50b400f37",
            "c64e5ac718e94a61a024d2b50d7a45c6",
            "682a3f51d720490a904b2289dff54cb1",
            "30e29170dd504d98980478676ecfbe99",
            "71ae935b227242bfa10a96cf33ae1050",
            "a06362dc06cb401d99b2615d16256b72",
            "00a3513cadfe4013af5ac601f166e43f",
            "ef0519825c654a2a8c05284e524e0dec",
            "3478562e7a3046cdb3da880aca125bc1",
            "caf41e797cd14ead854cbeb39cbc7bbc",
            "6cb22a4fcaf14125b13d3edba4cb92bb",
            "8d854b8c189e4e589e23d581d0d69378",
            "98877fbafecf4b0487cf7eaabb298f08",
            "da8af125ff7142bba7a577b7a4b3d1be",
            "8ca30619f84d47c4a6f614bce108df41",
            "f68e0625de544ef19a57c80927f060c6",
            "dee4883246dd40598540d4ed0ecb99b1",
            "bc2032f053564ff4bf0045caaa3eb8f0",
            "de4d5ce2b5d7425b89c8ed07fd924932",
            "81ac33e86b00495caff21bc98db2962a",
            "741d0e2852464492a51537a034feb29c",
            "c5e02cc9ecd34a77abd8ae9054475dba",
            "89af62a9be2d4ee4b10b0e7878a9b123",
            "4e6014de77d841f593223bac7ba5527b",
            "b5ba0c649ada4191b36ff542e3e5afd4",
            "143dc97614434be4bc17294a63591e4c",
            "392e47419e6d42a48cfff63de65e9fad",
            "2969fbff3bc8410cba1f85e4740afcdd",
            "a39a49927184438ab368fd5a9de0163a",
            "2d1a818d6a91432792b8d9d82b91fc62",
            "97f25d095fb24b8283fbaa193e1508a0"
          ]
        },
        "outputId": "070a7734-39bb-4424-c03d-b4b5cf02e02b"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|███████████████████████████████████████| 402M/402M [00:05<00:00, 82.4MiB/s]\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "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": "823a2248d19e473a8457356a5c253c7b"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "merges.txt:   0%|          | 0.00/456k [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "dea69fba019c430181378d72844a77e5"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "tokenizer.json:   0%|          | 0.00/1.36M [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "caf41e797cd14ead854cbeb39cbc7bbc"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "config.json:   0%|          | 0.00/665 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "741d0e2852464492a51537a034feb29c"
            }
          },
          "metadata": {}
        }
      ],
      "source": [
        "#@title CLIP model + GPT2 tokenizer\n",
        "\n",
        "device = CUDA(0) if is_gpu else \"cpu\"\n",
        "clip_model, preprocess = clip.load(\"RN50x4\", device=device, jit=False)\n",
        "tokenizer = GPT2Tokenizer.from_pretrained(\"gpt2\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "fIYC5J22_-jP"
      },
      "outputs": [],
      "source": [
        "\n",
        "#@title Mapper\n",
        "\n",
        "class MlpTransformer(nn.Module):\n",
        "    def __init__(self, in_dim, h_dim, out_d: Optional[int] = None, act=nnf.relu, dropout=0.):\n",
        "        super().__init__()\n",
        "        out_d = out_d if out_d is not None else in_dim\n",
        "        self.fc1 = nn.Linear(in_dim, h_dim)\n",
        "        self.act = act\n",
        "        self.fc2 = nn.Linear(h_dim, out_d)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.fc1(x)\n",
        "        x = self.act(x)\n",
        "        x = self.dropout(x)\n",
        "        x = self.fc2(x)\n",
        "        x = self.dropout(x)\n",
        "        return x\n",
        "\n",
        "class MultiHeadAttention(nn.Module):\n",
        "\n",
        "    def __init__(self, dim_self, dim_ref, num_heads, bias=True, dropout=0.):\n",
        "        super().__init__()\n",
        "        self.num_heads = num_heads\n",
        "        head_dim = dim_self // num_heads\n",
        "        self.scale = head_dim ** -0.5\n",
        "        self.to_queries = nn.Linear(dim_self, dim_self, bias=bias)\n",
        "        self.to_keys_values = nn.Linear(dim_ref, dim_self * 2, bias=bias)\n",
        "        self.project = nn.Linear(dim_self, dim_self)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x, y=None, mask=None):\n",
        "        y = y if y is not None else x\n",
        "        b, n, c = x.shape\n",
        "        _, m, d = y.shape\n",
        "        # b n h dh\n",
        "        queries = self.to_queries(x).reshape(b, n, self.num_heads, c // self.num_heads)\n",
        "        # b m 2 h dh\n",
        "        keys_values = self.to_keys_values(y).reshape(b, m, 2, self.num_heads, c // self.num_heads)\n",
        "        keys, values = keys_values[:, :, 0], keys_values[:, :, 1]\n",
        "        attention = torch.einsum('bnhd,bmhd->bnmh', queries, keys) * self.scale\n",
        "        if mask is not None:\n",
        "            if mask.dim() == 2:\n",
        "                mask = mask.unsqueeze(1)\n",
        "            attention = attention.masked_fill(mask.unsqueeze(3), float(\"-inf\"))\n",
        "        attention = attention.softmax(dim=2)\n",
        "        out = torch.einsum('bnmh,bmhd->bnhd', attention, values).reshape(b, n, c)\n",
        "        out = self.project(out)\n",
        "        return out, attention\n",
        "\n",
        "\n",
        "class TransformerLayer(nn.Module):\n",
        "\n",
        "    def forward_with_attention(self, x, y=None, mask=None):\n",
        "        x_, attention = self.attn(self.norm1(x), y, mask)\n",
        "        x = x + x_\n",
        "        x = x + self.mlp(self.norm2(x))\n",
        "        return x, attention\n",
        "\n",
        "    def forward(self, x, y=None, mask=None):\n",
        "        x = x + self.attn(self.norm1(x), y, mask)[0]\n",
        "        x = x + self.mlp(self.norm2(x))\n",
        "        return x\n",
        "\n",
        "    def __init__(self, dim_self, dim_ref, num_heads, mlp_ratio=4., bias=False, dropout=0., act=nnf.relu,\n",
        "                 norm_layer: nn.Module = nn.LayerNorm):\n",
        "        super().__init__()\n",
        "        self.norm1 = norm_layer(dim_self)\n",
        "        self.attn = MultiHeadAttention(dim_self, dim_ref, num_heads, bias=bias, dropout=dropout)\n",
        "        self.norm2 = norm_layer(dim_self)\n",
        "        self.mlp = MlpTransformer(dim_self, int(dim_self * mlp_ratio), act=act, dropout=dropout)\n",
        "\n",
        "\n",
        "class Transformer(nn.Module):\n",
        "\n",
        "    def forward_with_attention(self, x, y=None, mask=None):\n",
        "        attentions = []\n",
        "        for layer in self.layers:\n",
        "            x, att = layer.forward_with_attention(x, y, mask)\n",
        "            attentions.append(att)\n",
        "        return x, attentions\n",
        "\n",
        "    def forward(self, x, y=None, mask=None):\n",
        "        for i, layer in enumerate(self.layers):\n",
        "            if i % 2 == 0 and self.enc_dec: # cross\n",
        "                x = layer(x, y)\n",
        "            elif self.enc_dec:  # self\n",
        "                x = layer(x, x, mask)\n",
        "            else:  # self or cross\n",
        "                x = layer(x, y, mask)\n",
        "        return x\n",
        "\n",
        "    def __init__(self, dim_self: int, num_heads: int, num_layers: int, dim_ref: Optional[int] = None,\n",
        "                 mlp_ratio: float = 2., act=nnf.relu, norm_layer: nn.Module = nn.LayerNorm, enc_dec: bool = False):\n",
        "        super(Transformer, self).__init__()\n",
        "        dim_ref = dim_ref if dim_ref is not None else dim_self\n",
        "        self.enc_dec = enc_dec\n",
        "        if enc_dec:\n",
        "            num_layers = num_layers * 2\n",
        "        layers = []\n",
        "        for i in range(num_layers):\n",
        "            if i % 2 == 0 and enc_dec:  # cross\n",
        "                layers.append(TransformerLayer(dim_self, dim_ref, num_heads, mlp_ratio, act=act, norm_layer=norm_layer))\n",
        "            elif enc_dec:  # self\n",
        "                layers.append(TransformerLayer(dim_self, dim_self, num_heads, mlp_ratio, act=act, norm_layer=norm_layer))\n",
        "            else:  # self or cross\n",
        "                layers.append(TransformerLayer(dim_self, dim_ref, num_heads, mlp_ratio, act=act, norm_layer=norm_layer))\n",
        "        self.layers = nn.ModuleList(layers)\n",
        "\n",
        "\n",
        "class TransformerMapper(nn.Module):\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.linear(x).view(x.shape[0], self.clip_length, -1)\n",
        "        prefix = self.prefix_const.unsqueeze(0).expand(x.shape[0], *self.prefix_const.shape)\n",
        "        prefix = torch.cat((x, prefix), dim=1)\n",
        "        out = self.transformer(prefix)[:, self.clip_length:]\n",
        "        return out\n",
        "\n",
        "    def __init__(self, dim_clip: int, dim_embedding: int, prefix_length: int, clip_length: int, num_layers: int = 8):\n",
        "        super(TransformerMapper, self).__init__()\n",
        "        self.clip_length = clip_length\n",
        "        self.transformer = Transformer(dim_embedding, 8, num_layers)\n",
        "        self.linear = nn.Linear(dim_clip, clip_length * dim_embedding)\n",
        "        self.prefix_const = nn.Parameter(torch.randn(prefix_length, dim_embedding), requires_grad=True)\n",
        "\n",
        "class ClipCaptionModel(nn.Module):\n",
        "\n",
        "    def get_dummy_token(self, batch_size: int, device: torch.device) -> torch.Tensor:\n",
        "        return torch.zeros(batch_size, self.prefix_length, dtype=torch.int64, device=device)\n",
        "\n",
        "    def forward(self, tokens: torch.Tensor, prefix: torch.Tensor, mask: Optional[torch.Tensor] = None,\n",
        "                labels: Optional[torch.Tensor] = None):\n",
        "        embedding_text = self.gpt.transformer.wte(tokens)\n",
        "        prefix_projections = self.clip_project(prefix).view(-1, self.prefix_length, self.gpt_embedding_size)\n",
        "        embedding_cat = torch.cat((prefix_projections, embedding_text), dim=1)\n",
        "        if labels is not None:\n",
        "            dummy_token = self.get_dummy_token(tokens.shape[0], tokens.device)\n",
        "            labels = torch.cat((dummy_token, tokens), dim=1)\n",
        "        out = self.gpt(inputs_embeds=embedding_cat, labels=labels, attention_mask=mask)\n",
        "        return out\n",
        "\n",
        "    def __init__(self, prefix_size: int = 640,\n",
        "                 num_layers: int = 8):\n",
        "        super(ClipCaptionModel, self).__init__()\n",
        "        self.prefix_length = 40\n",
        "        self.gpt = GPT2LMHeadModel.from_pretrained('gpt2')\n",
        "        self.gpt_embedding_size = self.gpt.transformer.wte.weight.shape[1]\n",
        "        self.clip_project = TransformerMapper(prefix_size, self.gpt_embedding_size, 40,\n",
        "                                                                     40, num_layers)\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "glBzYsgIwhwF",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 564,
          "referenced_widgets": [
            "5bdb2ef4b36c47a6a62ea24514cde9c7",
            "e7d66287d3c6495bb25c1490474e978d",
            "024f7b6585b743c485a9d8ca2de05fe8",
            "10de39aa90fe41d48bd04203ec52219c",
            "ef956848925c462d8b5ed58c9e7d0c1c",
            "2437f113b53147b792fa1aae00892b6d",
            "1dbf72d51daf41ad914ced6447ab463a",
            "380211a9f3524c30ad43951c73f5a913",
            "3e59f79c3e6442d389b8109d4b8d7f46",
            "8050c6b705a34fc185f7b1050db65cc3",
            "7dea30e440d24a4a95d1ea5fdc970519",
            "278fdcfd4b8348dc87b30650b6d8d262",
            "3b65d2cf4ec14933b882e021ee7da9c9",
            "d166e7e33cb549ea9335398b194782b9",
            "7bb52e821c4d4c38af5dd893bc5c7344",
            "b21fe509fe2e40b6a076383c95f16fe8",
            "ded5eb90ef4d4023853433345195e965",
            "284d0667380644e2bdd84221f2d53cce",
            "d82416e512004483b5e03c4638a5f721",
            "81d8b2b204204c0d976f30b64f887e3f",
            "8c44c99f7e104bc6855b195237df6c01",
            "28c49bf2bcc74eac8d1b587fda22c290"
          ]
        },
        "outputId": "1a515f46-7b06-4c9c-91e1-5008ec13d4ed"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "model.safetensors:   0%|          | 0.00/548M [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "5bdb2ef4b36c47a6a62ea24514cde9c7"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "generation_config.json:   0%|          | 0.00/124 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "278fdcfd4b8348dc87b30650b6d8d262"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "error",
          "ename": "RuntimeError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-13-717d86197936>\u001b[0m in \u001b[0;36m<cell line: 5>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 5\u001b[0;31m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mload_state_dict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mload\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel_path\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmap_location\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mCPU\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      6\u001b[0m \u001b[0mmodel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0meval\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      7\u001b[0m \u001b[0mdevice\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mCUDA\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mis_gpu\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;34m\"cpu\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36mload_state_dict\u001b[0;34m(self, state_dict, strict, assign)\u001b[0m\n\u001b[1;32m   2150\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2151\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0merror_msgs\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2152\u001b[0;31m             raise RuntimeError('Error(s) in loading state_dict for {}:\\n\\t{}'.format(\n\u001b[0m\u001b[1;32m   2153\u001b[0m                                self.__class__.__name__, \"\\n\\t\".join(error_msgs)))\n\u001b[1;32m   2154\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0m_IncompatibleKeys\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmissing_keys\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0munexpected_keys\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mRuntimeError\u001b[0m: Error(s) in loading state_dict for ClipCaptionModel:\n\tUnexpected key(s) in state_dict: \"gpt.transformer.h.0.attn.bias\", \"gpt.transformer.h.0.attn.masked_bias\", \"gpt.transformer.h.1.attn.bias\", \"gpt.transformer.h.1.attn.masked_bias\", \"gpt.transformer.h.2.attn.bias\", \"gpt.transformer.h.2.attn.masked_bias\", \"gpt.transformer.h.3.attn.bias\", \"gpt.transformer.h.3.attn.masked_bias\", \"gpt.transformer.h.4.attn.bias\", \"gpt.transformer.h.4.attn.masked_bias\", \"gpt.transformer.h.5.attn.bias\", \"gpt.transformer.h.5.attn.masked_bias\", \"gpt.transformer.h.6.attn.bias\", \"gpt.transformer.h.6.attn.masked_bias\", \"gpt.transformer.h.7.attn.bias\", \"gpt.transformer.h.7.attn.masked_bias\", \"gpt.transformer.h.8.attn.bias\", \"gpt.transformer.h.8.attn.masked_bias\", \"gpt.transformer.h.9.attn.bias\", \"gpt.transformer.h.9.attn.masked_bias\", \"gpt.transformer.h.10.attn.bias\", \"gpt.transformer.h.10.attn.masked_bias\", \"gpt.transformer.h.11.attn.bias\", \"gpt.transformer.h.11.attn.masked_bias\". "
          ]
        }
      ],
      "source": [
        "#@title Load model weights\n",
        "model = ClipCaptionModel()\n",
        "\n",
        "\n",
        "model.load_state_dict(torch.load(model_path, map_location=CPU))\n",
        "model = model.eval()\n",
        "device = CUDA(0) if is_gpu else \"cpu\"\n",
        "model = model.to(device)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pohtQ8AfWNk_"
      },
      "outputs": [],
      "source": [
        "#@title Download random samples form COCO test set (Karpathy et al. split)\n",
        "\n",
        "IMAGE_NAME = '483108'  # @param ['562207', '579664', '060623', '165547', '334321', '483108', '386164', '354533']\n",
        "\n",
        "name_ = \"COCO_val2014_000000\" + IMAGE_NAME + \".jpg\"\n",
        "images_path = os.path.join(os.path.dirname(current_directory), \"images\")\n",
        "os.makedirs(images_path, exist_ok=True)\n",
        "UPLOADED_FILE = os.path.join(images_path, name_)\n",
        "\n",
        "if not os.path.isfile(UPLOADED_FILE):\n",
        "  download_path = os.path.join(images_path, \"images.zip\")\n",
        "  downloader.download_file(\"1BwJeBME-dpwcCT8IXYeWz7uaPkbexjNB\", download_path)\n",
        "\n",
        "  !unzip {download_path} -d {images_path}\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rRmcYnEfSMc_"
      },
      "outputs": [],
      "source": [
        "#@title Inference\n",
        "use_beam_search = True #@param {type:\"boolean\"}\n",
        "\n",
        "image = io.imread(UPLOADED_FILE)\n",
        "pil_image = PIL.Image.fromarray(image)\n",
        "#pil_img = Image(filename=UPLOADED_FILE)\n",
        "display(pil_image)\n",
        "\n",
        "image = preprocess(pil_image).unsqueeze(0).to(device)\n",
        "with torch.no_grad():\n",
        "    # if type(model) is ClipCaptionE2E:\n",
        "    #     prefix_embed = model.forward_image(image)\n",
        "    # else:\n",
        "    prefix = clip_model.encode_image(image).to(device, dtype=torch.float32)\n",
        "    prefix_embed = model.clip_project(prefix).reshape(1, 40, -1)\n",
        "if use_beam_search:\n",
        "    generated_text_prefix = generate_beam(model, tokenizer, embed=prefix_embed)[0]\n",
        "else:\n",
        "    generated_text_prefix = generate2(model, tokenizer, embed=prefix_embed)\n",
        "\n",
        "\n",
        "print('\\n')\n",
        "print(generated_text_prefix)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "m5jPDsEA5Kub"
      },
      "outputs": [],
      "source": [
        "#@title Or Upload Image\n",
        "\n",
        "uploaded = files.upload()\n",
        "\n",
        "if not uploaded:\n",
        "  UPLOADED_FILE = ''\n",
        "elif len(uploaded) == 1:\n",
        "  UPLOADED_FILE = list(uploaded.keys())[0]\n",
        "else:\n",
        "  raise AssertionError('Please upload one image at a time')\n",
        "\n",
        "print(UPLOADED_FILE)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YlusBDU0clra"
      },
      "outputs": [],
      "source": [
        "#@title Inference\n",
        "use_beam_search = True #@param {type:\"boolean\"}\n",
        "\n",
        "image = io.imread(UPLOADED_FILE)\n",
        "pil_image = PIL.Image.fromarray(image)\n",
        "#pil_img = Image(filename=UPLOADED_FILE)\n",
        "display(pil_image)\n",
        "\n",
        "image = preprocess(pil_image).unsqueeze(0).to(device)\n",
        "with torch.no_grad():\n",
        "    # if type(model) is ClipCaptionE2E:\n",
        "    #     prefix_embed = model.forward_image(image)\n",
        "    # else:\n",
        "    prefix = clip_model.encode_image(image).to(device, dtype=torch.float32)\n",
        "    prefix_embed = model.clip_project(prefix).reshape(1, 40, -1)\n",
        "if use_beam_search:\n",
        "    generated_text_prefix = generate_beam(model, tokenizer, embed=prefix_embed)[0]\n",
        "else:\n",
        "    generated_text_prefix = generate2(model, tokenizer, embed=prefix_embed)\n",
        "\n",
        "\n",
        "print('\\n')\n",
        "print(generated_text_prefix)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "823a2248d19e473a8457356a5c253c7b": {
          "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_58e74106989c4ca68cd8b5db1beb28d9",
              "IPY_MODEL_b0553d1b15a840c290932a2ef549fbf0",
              "IPY_MODEL_502ca5a845fa4986ad63586a44e867b3"
            ],
            "layout": "IPY_MODEL_d20be762e5ba434aa52e9cc029ccacf7"
          }
        },
        "58e74106989c4ca68cd8b5db1beb28d9": {
          "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_d4386577eada433f92b57ce91c3153b8",
            "placeholder": "​",
            "style": "IPY_MODEL_4ef5487b590e497ca719d501a397cb3d",
            "value": "vocab.json: 100%"
          }
        },
        "b0553d1b15a840c290932a2ef549fbf0": {
          "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_634c60cd8d934c55853e2811ae0a1a67",
            "max": 1042301,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_69e36359b49d479b9f2898c2bc870ff0",
            "value": 1042301
          }
        },
        "502ca5a845fa4986ad63586a44e867b3": {
          "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_525544b3302843c6b579917a9b5df84c",
            "placeholder": "​",
            "style": "IPY_MODEL_02a6dda06c344632a2b76a5fa14bea81",
            "value": " 1.04M/1.04M [00:00&lt;00:00, 11.2MB/s]"
          }
        },
        "d20be762e5ba434aa52e9cc029ccacf7": {
          "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
          }
        },
        "d4386577eada433f92b57ce91c3153b8": {
          "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
          }
        },
        "4ef5487b590e497ca719d501a397cb3d": {
          "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": ""
          }
        },
        "634c60cd8d934c55853e2811ae0a1a67": {
          "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
          }
        },
        "69e36359b49d479b9f2898c2bc870ff0": {
          "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": ""
          }
        },
        "525544b3302843c6b579917a9b5df84c": {
          "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
          }
        },
        "02a6dda06c344632a2b76a5fa14bea81": {
          "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": ""
          }
        },
        "dea69fba019c430181378d72844a77e5": {
          "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_959964afad6c4a488cd6f1bb3e82753d",
              "IPY_MODEL_2ce40954abb44071b0cbf0f50b400f37",
              "IPY_MODEL_c64e5ac718e94a61a024d2b50d7a45c6"
            ],
            "layout": "IPY_MODEL_682a3f51d720490a904b2289dff54cb1"
          }
        },
        "959964afad6c4a488cd6f1bb3e82753d": {
          "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_30e29170dd504d98980478676ecfbe99",
            "placeholder": "​",
            "style": "IPY_MODEL_71ae935b227242bfa10a96cf33ae1050",
            "value": "merges.txt: 100%"
          }
        },
        "2ce40954abb44071b0cbf0f50b400f37": {
          "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_a06362dc06cb401d99b2615d16256b72",
            "max": 456318,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_00a3513cadfe4013af5ac601f166e43f",
            "value": 456318
          }
        },
        "c64e5ac718e94a61a024d2b50d7a45c6": {
          "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_ef0519825c654a2a8c05284e524e0dec",
            "placeholder": "​",
            "style": "IPY_MODEL_3478562e7a3046cdb3da880aca125bc1",
            "value": " 456k/456k [00:00&lt;00:00, 11.1MB/s]"
          }
        },
        "682a3f51d720490a904b2289dff54cb1": {
          "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
          }
        },
        "30e29170dd504d98980478676ecfbe99": {
          "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
          }
        },
        "71ae935b227242bfa10a96cf33ae1050": {
          "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": ""
          }
        },
        "a06362dc06cb401d99b2615d16256b72": {
          "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
          }
        },
        "00a3513cadfe4013af5ac601f166e43f": {
          "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": ""
          }
        },
        "ef0519825c654a2a8c05284e524e0dec": {
          "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
          }
        },
        "3478562e7a3046cdb3da880aca125bc1": {
          "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": ""
          }
        },
        "caf41e797cd14ead854cbeb39cbc7bbc": {
          "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_6cb22a4fcaf14125b13d3edba4cb92bb",
              "IPY_MODEL_8d854b8c189e4e589e23d581d0d69378",
              "IPY_MODEL_98877fbafecf4b0487cf7eaabb298f08"
            ],
            "layout": "IPY_MODEL_da8af125ff7142bba7a577b7a4b3d1be"
          }
        },
        "6cb22a4fcaf14125b13d3edba4cb92bb": {
          "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_8ca30619f84d47c4a6f614bce108df41",
            "placeholder": "​",
            "style": "IPY_MODEL_f68e0625de544ef19a57c80927f060c6",
            "value": "tokenizer.json: 100%"
          }
        },
        "8d854b8c189e4e589e23d581d0d69378": {
          "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_dee4883246dd40598540d4ed0ecb99b1",
            "max": 1355256,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_bc2032f053564ff4bf0045caaa3eb8f0",
            "value": 1355256
          }
        },
        "98877fbafecf4b0487cf7eaabb298f08": {
          "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_de4d5ce2b5d7425b89c8ed07fd924932",
            "placeholder": "​",
            "style": "IPY_MODEL_81ac33e86b00495caff21bc98db2962a",
            "value": " 1.36M/1.36M [00:00&lt;00:00, 11.6MB/s]"
          }
        },
        "da8af125ff7142bba7a577b7a4b3d1be": {
          "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
          }
        },
        "8ca30619f84d47c4a6f614bce108df41": {
          "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
          }
        },
        "f68e0625de544ef19a57c80927f060c6": {
          "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": ""
          }
        },
        "dee4883246dd40598540d4ed0ecb99b1": {
          "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
          }
        },
        "bc2032f053564ff4bf0045caaa3eb8f0": {
          "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": ""
          }
        },
        "de4d5ce2b5d7425b89c8ed07fd924932": {
          "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
          }
        },
        "81ac33e86b00495caff21bc98db2962a": {
          "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": ""
          }
        },
        "741d0e2852464492a51537a034feb29c": {
          "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_c5e02cc9ecd34a77abd8ae9054475dba",
              "IPY_MODEL_89af62a9be2d4ee4b10b0e7878a9b123",
              "IPY_MODEL_4e6014de77d841f593223bac7ba5527b"
            ],
            "layout": "IPY_MODEL_b5ba0c649ada4191b36ff542e3e5afd4"
          }
        },
        "c5e02cc9ecd34a77abd8ae9054475dba": {
          "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_143dc97614434be4bc17294a63591e4c",
            "placeholder": "​",
            "style": "IPY_MODEL_392e47419e6d42a48cfff63de65e9fad",
            "value": "config.json: 100%"
          }
        },
        "89af62a9be2d4ee4b10b0e7878a9b123": {
          "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_2969fbff3bc8410cba1f85e4740afcdd",
            "max": 665,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_a39a49927184438ab368fd5a9de0163a",
            "value": 665
          }
        },
        "4e6014de77d841f593223bac7ba5527b": {
          "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_2d1a818d6a91432792b8d9d82b91fc62",
            "placeholder": "​",
            "style": "IPY_MODEL_97f25d095fb24b8283fbaa193e1508a0",
            "value": " 665/665 [00:00&lt;00:00, 36.8kB/s]"
          }
        },
        "b5ba0c649ada4191b36ff542e3e5afd4": {
          "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
          }
        },
        "143dc97614434be4bc17294a63591e4c": {
          "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
          }
        },
        "392e47419e6d42a48cfff63de65e9fad": {
          "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": ""
          }
        },
        "2969fbff3bc8410cba1f85e4740afcdd": {
          "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
          }
        },
        "a39a49927184438ab368fd5a9de0163a": {
          "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": ""
          }
        },
        "2d1a818d6a91432792b8d9d82b91fc62": {
          "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
          }
        },
        "97f25d095fb24b8283fbaa193e1508a0": {
          "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": ""
          }
        },
        "5bdb2ef4b36c47a6a62ea24514cde9c7": {
          "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_e7d66287d3c6495bb25c1490474e978d",
              "IPY_MODEL_024f7b6585b743c485a9d8ca2de05fe8",
              "IPY_MODEL_10de39aa90fe41d48bd04203ec52219c"
            ],
            "layout": "IPY_MODEL_ef956848925c462d8b5ed58c9e7d0c1c"
          }
        },
        "e7d66287d3c6495bb25c1490474e978d": {
          "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_2437f113b53147b792fa1aae00892b6d",
            "placeholder": "​",
            "style": "IPY_MODEL_1dbf72d51daf41ad914ced6447ab463a",
            "value": "model.safetensors: 100%"
          }
        },
        "024f7b6585b743c485a9d8ca2de05fe8": {
          "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_380211a9f3524c30ad43951c73f5a913",
            "max": 548105171,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_3e59f79c3e6442d389b8109d4b8d7f46",
            "value": 548105171
          }
        },
        "10de39aa90fe41d48bd04203ec52219c": {
          "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_8050c6b705a34fc185f7b1050db65cc3",
            "placeholder": "​",
            "style": "IPY_MODEL_7dea30e440d24a4a95d1ea5fdc970519",
            "value": " 548M/548M [00:04&lt;00:00, 132MB/s]"
          }
        },
        "ef956848925c462d8b5ed58c9e7d0c1c": {
          "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
          }
        },
        "2437f113b53147b792fa1aae00892b6d": {
          "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
          }
        },
        "1dbf72d51daf41ad914ced6447ab463a": {
          "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": ""
          }
        },
        "380211a9f3524c30ad43951c73f5a913": {
          "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
          }
        },
        "3e59f79c3e6442d389b8109d4b8d7f46": {
          "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": ""
          }
        },
        "8050c6b705a34fc185f7b1050db65cc3": {
          "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
          }
        },
        "7dea30e440d24a4a95d1ea5fdc970519": {
          "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": ""
          }
        },
        "278fdcfd4b8348dc87b30650b6d8d262": {
          "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_3b65d2cf4ec14933b882e021ee7da9c9",
              "IPY_MODEL_d166e7e33cb549ea9335398b194782b9",
              "IPY_MODEL_7bb52e821c4d4c38af5dd893bc5c7344"
            ],
            "layout": "IPY_MODEL_b21fe509fe2e40b6a076383c95f16fe8"
          }
        },
        "3b65d2cf4ec14933b882e021ee7da9c9": {
          "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_ded5eb90ef4d4023853433345195e965",
            "placeholder": "​",
            "style": "IPY_MODEL_284d0667380644e2bdd84221f2d53cce",
            "value": "generation_config.json: 100%"
          }
        },
        "d166e7e33cb549ea9335398b194782b9": {
          "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_d82416e512004483b5e03c4638a5f721",
            "max": 124,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_81d8b2b204204c0d976f30b64f887e3f",
            "value": 124
          }
        },
        "7bb52e821c4d4c38af5dd893bc5c7344": {
          "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_8c44c99f7e104bc6855b195237df6c01",
            "placeholder": "​",
            "style": "IPY_MODEL_28c49bf2bcc74eac8d1b587fda22c290",
            "value": " 124/124 [00:00&lt;00:00, 5.84kB/s]"
          }
        },
        "b21fe509fe2e40b6a076383c95f16fe8": {
          "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
          }
        },
        "ded5eb90ef4d4023853433345195e965": {
          "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
          }
        },
        "284d0667380644e2bdd84221f2d53cce": {
          "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": ""
          }
        },
        "d82416e512004483b5e03c4638a5f721": {
          "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
          }
        },
        "81d8b2b204204c0d976f30b64f887e3f": {
          "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": ""
          }
        },
        "8c44c99f7e104bc6855b195237df6c01": {
          "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
          }
        },
        "28c49bf2bcc74eac8d1b587fda22c290": {
          "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
}