Stable Diffusion - Image to Prompts - Inference.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f66672f6",
   "metadata": {
    "papermill": {
     "duration": 0.020491,
     "end_time": "2023-05-15T14:21:29.964736",
     "exception": false,
     "start_time": "2023-05-15T14:21:29.944245",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 1. Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "886610a1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:21:30.002121Z",
     "iopub.status.busy": "2023-05-15T14:21:30.001212Z",
     "iopub.status.idle": "2023-05-15T14:21:30.963908Z",
     "shell.execute_reply": "2023-05-15T14:21:30.962527Z"
    },
    "papermill": {
     "duration": 0.984431,
     "end_time": "2023-05-15T14:21:30.966719",
     "exception": false,
     "start_time": "2023-05-15T14:21:29.982288",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "!mkdir ./packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "12555ae2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:21:31.004105Z",
     "iopub.status.busy": "2023-05-15T14:21:31.003086Z",
     "iopub.status.idle": "2023-05-15T14:22:02.120842Z",
     "shell.execute_reply": "2023-05-15T14:22:02.119463Z"
    },
    "papermill": {
     "duration": 31.139412,
     "end_time": "2023-05-15T14:22:02.123861",
     "exception": false,
     "start_time": "2023-05-15T14:21:30.984449",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Copy CLIP-Interrogator files to working directory\n",
    "!cp -r /kaggle/input/stable-diffusion-image-to-prompts-dataset/ci-preprocess ./\n",
    "\n",
    "# Copy packages to current working directory\n",
    "!cp -r /kaggle/input/stable-diffusion-image-to-prompts-dataset/packages/LAVIS-main/LAVIS-main ./packages\n",
    "!cp -r /kaggle/input/stable-diffusion-image-to-prompts-dataset/packages/clip-interrogator-0.6.0/clip-interrogator-0.6.0 ./packages\n",
    "!cp -r /kaggle/input/stable-diffusion-image-to-prompts-dataset/packages/safetensors-0.3.0/safetensors-0.3.0 ./packages\n",
    "!cp -r /kaggle/input/stable-diffusion-image-to-prompts-dataset/packages/sacremoses-0.0.53/sacremoses-0.0.53 ./packages\n",
    "!cp -r /kaggle/input/stable-diffusion-image-to-prompts-dataset/packages/antlr4-python3-runtime-4.9.3/antlr4-python3-runtime-4.9.3 ./packages\n",
    "!cp -r /kaggle/input/stable-diffusion-image-to-prompts-dataset/packages/contexttimer-0.3.3/contexttimer-0.3.3 ./packages\n",
    "!cp -r /kaggle/input/stable-diffusion-image-to-prompts-dataset/packages/fairscale-0.4.4/fairscale-0.4.4 ./packages\n",
    "!cp -r /kaggle/input/stable-diffusion-image-to-prompts-dataset/packages/iopath-0.1.10/iopath-0.1.10 ./packages\n",
    "!cp -r /kaggle/input/stable-diffusion-image-to-prompts-dataset/packages/pycocotools-2.0.6/pycocotools-2.0.6 ./packages\n",
    "!cp -r /kaggle/input/stable-diffusion-image-to-prompts-dataset/packages/validators-0.20.0/validators-0.20.0 ./packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8e1df4c6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:22:02.162506Z",
     "iopub.status.busy": "2023-05-15T14:22:02.161288Z",
     "iopub.status.idle": "2023-05-15T14:22:02.167368Z",
     "shell.execute_reply": "2023-05-15T14:22:02.166405Z"
    },
    "papermill": {
     "duration": 0.028126,
     "end_time": "2023-05-15T14:22:02.169623",
     "exception": false,
     "start_time": "2023-05-15T14:22:02.141497",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('/kaggle/input/stable-diffusion-image-to-prompts-dataset/packages/sentence-transformers-2.2.2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "13714cf5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:22:02.205910Z",
     "iopub.status.busy": "2023-05-15T14:22:02.205163Z",
     "iopub.status.idle": "2023-05-15T14:22:02.214771Z",
     "shell.execute_reply": "2023-05-15T14:22:02.213723Z"
    },
    "papermill": {
     "duration": 0.030154,
     "end_time": "2023-05-15T14:22:02.217011",
     "exception": false,
     "start_time": "2023-05-15T14:22:02.186857",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "INSTALL_OFA = False\n",
    "\n",
    "if INSTALL_OFA:\n",
    "    !pip install ./packages/sacremoses-0.0.53 --no-index --find-links /kaggle/input/stable-diffusion-image-to-prompts-dataset/packages\n",
    "    !pip install /kaggle/input/stable-diffusion-image-to-prompts-dataset/packages/transformers-4.18.0.dev0-py3-none-any.whl --no-index --find-links /kaggle/input/stable-diffusion-image-to-prompts-dataset/packages\n",
    "    \n",
    "    from transformers import OFAModel, OFATokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6c5a4e49",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:22:02.253296Z",
     "iopub.status.busy": "2023-05-15T14:22:02.252325Z",
     "iopub.status.idle": "2023-05-15T14:22:02.277898Z",
     "shell.execute_reply": "2023-05-15T14:22:02.276896Z"
    },
    "papermill": {
     "duration": 0.045833,
     "end_time": "2023-05-15T14:22:02.280109",
     "exception": false,
     "start_time": "2023-05-15T14:22:02.234276",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "INSTALL_LAVIS = False\n",
    "\n",
    "if INSTALL_LAVIS:\n",
    "    !pip install ./packages/antlr4-python3-runtime-4.9.3\n",
    "    !pip install ./packages/contexttimer-0.3.3\n",
    "    !pip install ./packages/fairscale-0.4.4 --find-links /kaggle/input/stable-diffusion-image-to-prompts-dataset/packages\n",
    "    !pip install ./packages/iopath-0.1.10\n",
    "    !pip install ./packages/pycocotools-2.0.6 --find-links /kaggle/input/stable-diffusion-image-to-prompts-dataset/packages\n",
    "    !pip install ./packages/validators-0.20.0\n",
    "    !pip install ./packages/LAVIS-main --no-index --find-links /kaggle/input/stable-diffusion-image-to-prompts-dataset/packages\n",
    "    \n",
    "    from lavis.models import load_model_and_preprocess, registry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2c5f56f9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:22:02.316199Z",
     "iopub.status.busy": "2023-05-15T14:22:02.315914Z",
     "iopub.status.idle": "2023-05-15T14:22:02.326976Z",
     "shell.execute_reply": "2023-05-15T14:22:02.325864Z"
    },
    "papermill": {
     "duration": 0.032369,
     "end_time": "2023-05-15T14:22:02.329652",
     "exception": false,
     "start_time": "2023-05-15T14:22:02.297283",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\n%%writefile /opt/conda/lib/python3.7/site-packages/lavis/models/blip_models/blip.py\\n\\n\"\"\"\\n Copyright (c) 2022, salesforce.com, inc.\\n All rights reserved.\\n SPDX-License-Identifier: BSD-3-Clause\\n For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause\\n\"\"\"\\n\\nimport logging\\nimport os\\n\\nimport torch\\nfrom lavis.common.dist_utils import download_cached_file\\nfrom lavis.common.utils import is_url\\nfrom lavis.models.base_model import BaseModel\\nfrom lavis.models.vit import interpolate_pos_embed\\nfrom transformers import BertTokenizer\\n\\n\\nclass BlipBase(BaseModel):\\n    @classmethod\\n    def init_tokenizer(cls):\\n        tokenizer = BertTokenizer.from_pretrained(\"/kaggle/input/huggingface-bert/bert-base-uncased\")\\n        tokenizer.add_special_tokens({\"bos_token\": \"[DEC]\"})\\n        tokenizer.add_special_tokens({\"additional_special_tokens\": [\"[ENC]\"]})\\n        tokenizer.enc_token_id = tokenizer.additional_special_tokens_ids[0]\\n        return tokenizer\\n\\n    def load_from_pretrained(self, url_or_filename):\\n        if is_url(url_or_filename):\\n            cached_file = download_cached_file(\\n                url_or_filename, check_hash=False, progress=True\\n            )\\n            checkpoint = torch.load(cached_file, map_location=\"cpu\")\\n        elif os.path.isfile(url_or_filename):\\n            checkpoint = torch.load(url_or_filename, map_location=\"cpu\")\\n        else:\\n            raise RuntimeError(\"checkpoint url or path is invalid\")\\n\\n        state_dict = checkpoint[\"model\"]\\n\\n        state_dict[\"visual_encoder.pos_embed\"] = interpolate_pos_embed(\\n            state_dict[\"visual_encoder.pos_embed\"], self.visual_encoder\\n        )\\n        if \"visual_encoder_m.pos_embed\" in self.state_dict().keys():\\n            state_dict[\"visual_encoder_m.pos_embed\"] = interpolate_pos_embed(\\n                state_dict[\"visual_encoder_m.pos_embed\"], self.visual_encoder_m\\n            )\\n\\n        for key in self.state_dict().keys():\\n            if key in state_dict.keys():\\n                if state_dict[key].shape != self.state_dict()[key].shape:\\n                    del state_dict[key]\\n\\n        msg = self.load_state_dict(state_dict, strict=False)\\n\\n        logging.info(\"Missing keys {}\".format(msg.missing_keys))\\n        logging.info(\"load checkpoint from %s\" % url_or_filename)\\n\\n        return msg\\n'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "'''\n",
    "%%writefile /opt/conda/lib/python3.7/site-packages/lavis/models/blip_models/blip.py\n",
    "\n",
    "\"\"\"\n",
    " Copyright (c) 2022, salesforce.com, inc.\n",
    " All rights reserved.\n",
    " SPDX-License-Identifier: BSD-3-Clause\n",
    " For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n",
    "\"\"\"\n",
    "\n",
    "import logging\n",
    "import os\n",
    "\n",
    "import torch\n",
    "from lavis.common.dist_utils import download_cached_file\n",
    "from lavis.common.utils import is_url\n",
    "from lavis.models.base_model import BaseModel\n",
    "from lavis.models.vit import interpolate_pos_embed\n",
    "from transformers import BertTokenizer\n",
    "\n",
    "\n",
    "class BlipBase(BaseModel):\n",
    "    @classmethod\n",
    "    def init_tokenizer(cls):\n",
    "        tokenizer = BertTokenizer.from_pretrained(\"/kaggle/input/huggingface-bert/bert-base-uncased\")\n",
    "        tokenizer.add_special_tokens({\"bos_token\": \"[DEC]\"})\n",
    "        tokenizer.add_special_tokens({\"additional_special_tokens\": [\"[ENC]\"]})\n",
    "        tokenizer.enc_token_id = tokenizer.additional_special_tokens_ids[0]\n",
    "        return tokenizer\n",
    "\n",
    "    def load_from_pretrained(self, url_or_filename):\n",
    "        if is_url(url_or_filename):\n",
    "            cached_file = download_cached_file(\n",
    "                url_or_filename, check_hash=False, progress=True\n",
    "            )\n",
    "            checkpoint = torch.load(cached_file, map_location=\"cpu\")\n",
    "        elif os.path.isfile(url_or_filename):\n",
    "            checkpoint = torch.load(url_or_filename, map_location=\"cpu\")\n",
    "        else:\n",
    "            raise RuntimeError(\"checkpoint url or path is invalid\")\n",
    "\n",
    "        state_dict = checkpoint[\"model\"]\n",
    "\n",
    "        state_dict[\"visual_encoder.pos_embed\"] = interpolate_pos_embed(\n",
    "            state_dict[\"visual_encoder.pos_embed\"], self.visual_encoder\n",
    "        )\n",
    "        if \"visual_encoder_m.pos_embed\" in self.state_dict().keys():\n",
    "            state_dict[\"visual_encoder_m.pos_embed\"] = interpolate_pos_embed(\n",
    "                state_dict[\"visual_encoder_m.pos_embed\"], self.visual_encoder_m\n",
    "            )\n",
    "\n",
    "        for key in self.state_dict().keys():\n",
    "            if key in state_dict.keys():\n",
    "                if state_dict[key].shape != self.state_dict()[key].shape:\n",
    "                    del state_dict[key]\n",
    "\n",
    "        msg = self.load_state_dict(state_dict, strict=False)\n",
    "\n",
    "        logging.info(\"Missing keys {}\".format(msg.missing_keys))\n",
    "        logging.info(\"load checkpoint from %s\" % url_or_filename)\n",
    "\n",
    "        return msg\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "bbf13313",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:22:02.366794Z",
     "iopub.status.busy": "2023-05-15T14:22:02.366204Z",
     "iopub.status.idle": "2023-05-15T14:22:47.893063Z",
     "shell.execute_reply": "2023-05-15T14:22:47.891779Z"
    },
    "papermill": {
     "duration": 45.548472,
     "end_time": "2023-05-15T14:22:47.896001",
     "exception": false,
     "start_time": "2023-05-15T14:22:02.347529",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in links: /kaggle/input/stable-diffusion-image-to-prompts-dataset/packages/\r\n",
      "Processing ./packages/clip-interrogator-0.6.0\r\n",
      "  Installing build dependencies ... \u001b[?25l-\b \b\\\b \b|\b \bdone\r\n",
      "\u001b[?25h  Getting requirements to build wheel ... \u001b[?25l-\b \bdone\r\n",
      "\u001b[?25h  Installing backend dependencies ... \u001b[?25l-\b \b\\\b \bdone\r\n",
      "\u001b[?25h  Preparing metadata (pyproject.toml) ... \u001b[?25l-\b \bdone\r\n",
      "\u001b[?25hRequirement already satisfied: torch in /opt/conda/lib/python3.7/site-packages (from clip-interrogator==0.6.0) (1.13.0)\r\n",
      "Requirement already satisfied: Pillow in /opt/conda/lib/python3.7/site-packages (from clip-interrogator==0.6.0) (9.4.0)\r\n",
      "Processing /kaggle/input/stable-diffusion-image-to-prompts-dataset/packages/safetensors-0.3.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl\r\n",
      "Requirement already satisfied: torchvision in /opt/conda/lib/python3.7/site-packages (from clip-interrogator==0.6.0) (0.14.0)\r\n",
      "Requirement already satisfied: requests in /opt/conda/lib/python3.7/site-packages (from clip-interrogator==0.6.0) (2.28.2)\r\n",
      "Requirement already satisfied: accelerate in /opt/conda/lib/python3.7/site-packages (from clip-interrogator==0.6.0) (0.12.0)\r\n",
      "Processing /kaggle/input/stable-diffusion-image-to-prompts-dataset/packages/open_clip_torch-2.16.0-py3-none-any.whl\r\n",
      "Processing /kaggle/input/stable-diffusion-image-to-prompts-dataset/packages/transformers-4.28.1-py3-none-any.whl\r\n",
      "Requirement already satisfied: tqdm in /opt/conda/lib/python3.7/site-packages (from clip-interrogator==0.6.0) (4.64.1)\r\n",
      "Requirement already satisfied: pyyaml>=5.1 in /opt/conda/lib/python3.7/site-packages (from transformers>=4.27.1->clip-interrogator==0.6.0) (6.0)\r\n",
      "Requirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in /opt/conda/lib/python3.7/site-packages (from transformers>=4.27.1->clip-interrogator==0.6.0) (0.13.2)\r\n",
      "Requirement already satisfied: regex!=2019.12.17 in /opt/conda/lib/python3.7/site-packages (from transformers>=4.27.1->clip-interrogator==0.6.0) (2021.11.10)\r\n",
      "Requirement already satisfied: huggingface-hub<1.0,>=0.11.0 in /opt/conda/lib/python3.7/site-packages (from transformers>=4.27.1->clip-interrogator==0.6.0) (0.12.1)\r\n",
      "Requirement already satisfied: packaging>=20.0 in /opt/conda/lib/python3.7/site-packages (from transformers>=4.27.1->clip-interrogator==0.6.0) (23.0)\r\n",
      "Requirement already satisfied: numpy>=1.17 in /opt/conda/lib/python3.7/site-packages (from transformers>=4.27.1->clip-interrogator==0.6.0) (1.21.6)\r\n",
      "Requirement already satisfied: importlib-metadata in /opt/conda/lib/python3.7/site-packages (from transformers>=4.27.1->clip-interrogator==0.6.0) (4.11.4)\r\n",
      "Requirement already satisfied: filelock in /opt/conda/lib/python3.7/site-packages (from transformers>=4.27.1->clip-interrogator==0.6.0) (3.9.0)\r\n",
      "Requirement already satisfied: psutil in /opt/conda/lib/python3.7/site-packages (from accelerate->clip-interrogator==0.6.0) (5.9.3)\r\n",
      "Requirement already satisfied: typing-extensions in /opt/conda/lib/python3.7/site-packages (from torch->clip-interrogator==0.6.0) (4.4.0)\r\n",
      "Requirement already satisfied: sentencepiece in /opt/conda/lib/python3.7/site-packages (from open-clip-torch->clip-interrogator==0.6.0) (0.1.97)\r\n",
      "Requirement already satisfied: protobuf<4 in /opt/conda/lib/python3.7/site-packages (from open-clip-torch->clip-interrogator==0.6.0) (3.20.3)\r\n",
      "Processing /kaggle/input/stable-diffusion-image-to-prompts-dataset/packages/ftfy-6.1.1-py3-none-any.whl\r\n",
      "Requirement already satisfied: timm in /opt/conda/lib/python3.7/site-packages (from open-clip-torch->clip-interrogator==0.6.0) (0.6.12)\r\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /opt/conda/lib/python3.7/site-packages (from requests->clip-interrogator==0.6.0) (1.26.14)\r\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/lib/python3.7/site-packages (from requests->clip-interrogator==0.6.0) (3.4)\r\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.7/site-packages (from requests->clip-interrogator==0.6.0) (2022.12.7)\r\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /opt/conda/lib/python3.7/site-packages (from requests->clip-interrogator==0.6.0) (2.1.1)\r\n",
      "Requirement already satisfied: wcwidth>=0.2.5 in /opt/conda/lib/python3.7/site-packages (from ftfy->open-clip-torch->clip-interrogator==0.6.0) (0.2.6)\r\n",
      "Requirement already satisfied: zipp>=0.5 in /opt/conda/lib/python3.7/site-packages (from importlib-metadata->transformers>=4.27.1->clip-interrogator==0.6.0) (3.11.0)\r\n",
      "Building wheels for collected packages: clip-interrogator\r\n",
      "  Building wheel for clip-interrogator (pyproject.toml) ... \u001b[?25l-\b \b\\\b \bdone\r\n",
      "\u001b[?25h  Created wheel for clip-interrogator: filename=clip_interrogator-0.6.0-py3-none-any.whl size=787803 sha256=bee212aeb36ca9f27ed10b029b9b1b03f155d6fdc72606a327630a67d23b58c3\r\n",
      "  Stored in directory: /root/.cache/pip/wheels/db/5a/0e/ee533281182af8b37a080411a279df3001b9cd1a6cb2025767\r\n",
      "Successfully built clip-interrogator\r\n",
      "Installing collected packages: safetensors, ftfy, transformers, open-clip-torch, clip-interrogator\r\n",
      "  Attempting uninstall: transformers\r\n",
      "    Found existing installation: transformers 4.26.1\r\n",
      "    Uninstalling transformers-4.26.1:\r\n",
      "      Successfully uninstalled transformers-4.26.1\r\n",
      "Successfully installed clip-interrogator-0.6.0 ftfy-6.1.1 open-clip-torch-2.16.0 safetensors-0.3.0 transformers-4.28.1\r\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\r\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!pip install ./packages/clip-interrogator-0.6.0 --no-index --find-links /kaggle/input/stable-diffusion-image-to-prompts-dataset/packages/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b99649fb",
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
    "execution": {
     "iopub.execute_input": "2023-05-15T14:22:47.938402Z",
     "iopub.status.busy": "2023-05-15T14:22:47.938060Z",
     "iopub.status.idle": "2023-05-15T14:22:57.137380Z",
     "shell.execute_reply": "2023-05-15T14:22:57.133450Z"
    },
    "papermill": {
     "duration": 9.224313,
     "end_time": "2023-05-15T14:22:57.141177",
     "exception": false,
     "start_time": "2023-05-15T14:22:47.916864",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "import yaml\n",
    "import gc\n",
    "from shutil import copy2\n",
    "from glob import glob\n",
    "from pathlib import Path\n",
    "from tqdm import tqdm\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from PIL import Image\n",
    "import cv2\n",
    "\n",
    "from sklearn.preprocessing import normalize\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset, DataLoader, SequentialSampler\n",
    "import timm\n",
    "from transformers import BlipForConditionalGeneration, BlipProcessor\n",
    "from sentence_transformers import SentenceTransformer\n",
    "import open_clip\n",
    "from clip_interrogator import *\n",
    "import albumentations as A\n",
    "from albumentations.pytorch.transforms import ToTensorV2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a810060d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:22:57.199477Z",
     "iopub.status.busy": "2023-05-15T14:22:57.198775Z",
     "iopub.status.idle": "2023-05-15T14:22:57.204520Z",
     "shell.execute_reply": "2023-05-15T14:22:57.203524Z"
    },
    "papermill": {
     "duration": 0.039644,
     "end_time": "2023-05-15T14:22:57.209535",
     "exception": false,
     "start_time": "2023-05-15T14:22:57.169891",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "competition_dataset = Path('/kaggle/input/stable-diffusion-image-to-prompts')\n",
    "external_dataset = Path('/kaggle/input/stable-diffusion-image-to-prompts-dataset')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "126a1703",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:22:57.266625Z",
     "iopub.status.busy": "2023-05-15T14:22:57.266198Z",
     "iopub.status.idle": "2023-05-15T14:22:57.285608Z",
     "shell.execute_reply": "2023-05-15T14:22:57.284562Z"
    },
    "papermill": {
     "duration": 0.051759,
     "end_time": "2023-05-15T14:22:57.289018",
     "exception": false,
     "start_time": "2023-05-15T14:22:57.237259",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Shape: (7, 1) - Test Embeddings Shape: (2688, 1)\n"
     ]
    }
   ],
   "source": [
    "image_paths = glob(str(competition_dataset / 'images' / '*.png'))\n",
    "image_ids = [image_path.split('/')[-1][:-4] for image_path in image_paths]\n",
    "df_test_embeddings = pd.DataFrame(dict(imgId_eId=[id + '_' + str(i) for id in image_ids for i in range(384)]))\n",
    "df_test = pd.DataFrame(image_ids, columns=['image_id'])\n",
    "\n",
    "print(f'Test Shape: {df_test.shape} - Test Embeddings Shape: {df_test_embeddings.shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa8b40b2",
   "metadata": {
    "papermill": {
     "duration": 0.027803,
     "end_time": "2023-05-15T14:22:57.347855",
     "exception": false,
     "start_time": "2023-05-15T14:22:57.320052",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 2. Utilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "bd64b383",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:22:57.405750Z",
     "iopub.status.busy": "2023-05-15T14:22:57.405281Z",
     "iopub.status.idle": "2023-05-15T14:22:57.421764Z",
     "shell.execute_reply": "2023-05-15T14:22:57.420744Z"
    },
    "papermill": {
     "duration": 0.048997,
     "end_time": "2023-05-15T14:22:57.425264",
     "exception": false,
     "start_time": "2023-05-15T14:22:57.376267",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def set_seed(seed, deterministic_cudnn=False):\n",
    "\n",
    "    \"\"\"\n",
    "    Set random seed for reproducible results\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    seed: int\n",
    "        Random seed\n",
    "\n",
    "    deterministic_cudnn: bool\n",
    "        Whether to set deterministic cuDNN or not\n",
    "    \"\"\"\n",
    "\n",
    "    if deterministic_cudnn:\n",
    "        torch.backends.cudnn.deterministic = True\n",
    "        torch.backends.cudnn.benchmark = False\n",
    "\n",
    "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed(seed)\n",
    "    torch.cuda.manual_seed_all(seed)\n",
    "\n",
    "\n",
    "def load_model(model_directory_path):\n",
    "    \n",
    "    \"\"\"\n",
    "    Load models and config file from the given directory path\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    model_directory_path: str\n",
    "        Path of the model directory\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    config: dict\n",
    "        Dictionary of model configurations\n",
    "    \n",
    "    model: torch.nn.Module\n",
    "        Model\n",
    "    \"\"\"\n",
    "    \n",
    "    config = yaml.load(open(f'{model_directory_path}/config.yaml', 'r'), Loader=yaml.FullLoader)\n",
    "    config['model']['model_args']['pretrained'] = False\n",
    "    model_path = glob(f'{model_directory_path}/*.pt')[0]\n",
    "        \n",
    "    model = eval(config['model']['model_class'])(**config['model']['model_args'])\n",
    "    model.load_state_dict(torch.load(model_path))\n",
    "    model = model.to(config['training']['device'])\n",
    "    model.eval()\n",
    "        \n",
    "    return config, model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c99d873e",
   "metadata": {
    "papermill": {
     "duration": 0.028191,
     "end_time": "2023-05-15T14:22:57.482348",
     "exception": false,
     "start_time": "2023-05-15T14:22:57.454157",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 3. Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "78715719",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:22:57.542914Z",
     "iopub.status.busy": "2023-05-15T14:22:57.542350Z",
     "iopub.status.idle": "2023-05-15T14:22:57.561984Z",
     "shell.execute_reply": "2023-05-15T14:22:57.560959Z"
    },
    "papermill": {
     "duration": 0.054223,
     "end_time": "2023-05-15T14:22:57.565042",
     "exception": false,
     "start_time": "2023-05-15T14:22:57.510819",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class ImageCaptionDataset(Dataset):\n",
    "\n",
    "    def __init__(self, image_paths, captions, transforms=None, image_reader='opencv'):\n",
    "\n",
    "        self.image_paths = image_paths\n",
    "        self.captions = captions\n",
    "        self.transforms = transforms\n",
    "        self.image_reader = image_reader\n",
    "\n",
    "    def __len__(self):\n",
    "\n",
    "        \"\"\"\n",
    "        Get the length the dataset\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        length: int\n",
    "            Length of the dataset\n",
    "        \"\"\"\n",
    "\n",
    "        return len(self.image_paths)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "\n",
    "        \"\"\"\n",
    "        Get the idxth element in the dataset\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        idx: int\n",
    "            Index of the sample (0 <= idx < length of the dataset)\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        images: torch.FloatTensor of shape (channel, height, width) or (height, width, channel)\n",
    "            Image tensor\n",
    "\n",
    "        captions: torch.Tensor of shape (1)\n",
    "            Caption tensor\n",
    "        \"\"\"\n",
    "\n",
    "        if self.image_reader == 'pil':\n",
    "            image = Image.open(self.image_paths[idx]).convert('RGB')\n",
    "        elif self.image_reader == 'opencv':\n",
    "            image = cv2.imread(str(self.image_paths[idx]))\n",
    "            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "        else:\n",
    "            raise ValueError(f'Invalid image_reader: {self.image_reader}')\n",
    "\n",
    "        if self.transforms is not None:\n",
    "            if isinstance(self.transforms, A.Compose):\n",
    "                image = self.transforms(image=np.array(image))['image'].float()\n",
    "            elif isinstance(self.transforms, BlipProcessor):\n",
    "                image = self.transforms(image, 'a photography of', return_tensors='pt').pixel_values\n",
    "                image = torch.squeeze(image, dim=0)\n",
    "            else:\n",
    "                image = self.transforms(image).float()\n",
    "\n",
    "        if self.captions is not None:\n",
    "            captions = torch.as_tensor(self.captions[idx])\n",
    "            return image, captions\n",
    "        else:\n",
    "            return image\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "fd1bbfc3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:22:57.628507Z",
     "iopub.status.busy": "2023-05-15T14:22:57.628078Z",
     "iopub.status.idle": "2023-05-15T14:22:57.643004Z",
     "shell.execute_reply": "2023-05-15T14:22:57.640878Z"
    },
    "papermill": {
     "duration": 0.052005,
     "end_time": "2023-05-15T14:22:57.646972",
     "exception": false,
     "start_time": "2023-05-15T14:22:57.594967",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class ImageEmbeddingDataset(Dataset):\n",
    "\n",
    "    def __init__(self, image_paths, embeddings, transforms=None):\n",
    "\n",
    "        self.image_paths = image_paths\n",
    "        self.embeddings = embeddings\n",
    "        self.transforms = transforms\n",
    "\n",
    "    def __len__(self):\n",
    "\n",
    "        \"\"\"\n",
    "        Get the length the dataset\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        length: int\n",
    "            Length of the dataset\n",
    "        \"\"\"\n",
    "\n",
    "        return len(self.image_paths)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "\n",
    "        \"\"\"\n",
    "        Get the idxth element in the dataset\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        idx: int\n",
    "            Index of the sample (0 <= idx < length of the dataset)\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        images: torch.FloatTensor of shape (channel, height, width) or (height, width, channel)\n",
    "            Image tensor or array\n",
    "\n",
    "        embeddings: torch.FloatTensor of shape (384)\n",
    "            Embedding tensor\n",
    "        \"\"\"\n",
    "\n",
    "        image = cv2.imread(str(self.image_paths[idx]))\n",
    "        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "\n",
    "        if self.transforms is not None:\n",
    "            image = self.transforms(image=np.array(image))['image'].float()\n",
    "        else:\n",
    "            image = torch.as_tensor(image, dtype=torch.float)\n",
    "\n",
    "        if self.embeddings is not None:\n",
    "            embeddings = torch.as_tensor(self.embeddings[idx], dtype=torch.float)\n",
    "            return image, embeddings\n",
    "        else:\n",
    "            return image\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87d8e341",
   "metadata": {
    "papermill": {
     "duration": 0.028296,
     "end_time": "2023-05-15T14:22:57.745561",
     "exception": false,
     "start_time": "2023-05-15T14:22:57.717265",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 4. OFA (transformers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "1debc391",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:22:57.806353Z",
     "iopub.status.busy": "2023-05-15T14:22:57.805530Z",
     "iopub.status.idle": "2023-05-15T14:22:57.810401Z",
     "shell.execute_reply": "2023-05-15T14:22:57.809446Z"
    },
    "papermill": {
     "duration": 0.041166,
     "end_time": "2023-05-15T14:22:57.815827",
     "exception": false,
     "start_time": "2023-05-15T14:22:57.774661",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "ofa_large_caption_inference = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "5bf783e3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:22:57.877277Z",
     "iopub.status.busy": "2023-05-15T14:22:57.876399Z",
     "iopub.status.idle": "2023-05-15T14:22:57.889332Z",
     "shell.execute_reply": "2023-05-15T14:22:57.888382Z"
    },
    "papermill": {
     "duration": 0.046991,
     "end_time": "2023-05-15T14:22:57.892268",
     "exception": false,
     "start_time": "2023-05-15T14:22:57.845277",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def load_ofa(model_directory_path, device):\n",
    "\n",
    "    \"\"\"\n",
    "    Load model and processor\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    model_directory_path: str\n",
    "        Path of the model directory\n",
    "\n",
    "    device: torch.device\n",
    "        Location of the model\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    model: transformers.OFAModel\n",
    "        Model\n",
    "\n",
    "    tokenizer: transformers.OFATokenizer\n",
    "        Tokenizer\n",
    "    \"\"\"\n",
    "\n",
    "    model = transformers.OFAModel.from_pretrained(model_directory_path)\n",
    "    model.to(device)\n",
    "    model.eval()\n",
    "\n",
    "    tokenizer = transformers.OFATokenizer.from_pretrained(model_directory_path)\n",
    "\n",
    "    return model, tokenizer\n",
    "\n",
    "\n",
    "def predict_ofa(inputs, model, tokenizer, prompt, num_beams, no_repeat_ngram_size, max_length, min_length, device):\n",
    "\n",
    "    \"\"\"\n",
    "    Predict given inputs with given model\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    inputs: torch.FloatTensor of shape (batch, channel, height, width) images\n",
    "        Inputs tensor\n",
    "\n",
    "    model: transformers.OFAModel\n",
    "        Model\n",
    "\n",
    "    tokenizer: transformers.OFATokenizer\n",
    "        Tokenizer\n",
    "\n",
    "    prompt: str\n",
    "        Prompt used for the image captioning\n",
    "\n",
    "    num_beams: int\n",
    "        Number of beams in beam search\n",
    "\n",
    "    no_repeat_ngram_size: int\n",
    "        N-gram size\n",
    "\n",
    "    max_length: int\n",
    "        Maximum length of the generated text\n",
    "        \n",
    "    min_length: int\n",
    "        Minimum length of the generated text\n",
    "\n",
    "    device: torch.device\n",
    "        Location of the model\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    outputs: list of shape (batch) texts\n",
    "        Outputs\n",
    "    \"\"\"\n",
    "\n",
    "    input_ids = tokenizer([prompt] * inputs.size(0), return_tensors='pt').input_ids.to(device)\n",
    "\n",
    "    with torch.no_grad(), torch.autocast(device_type=device.type, dtype=torch.float16):\n",
    "        outputs = model.generate(\n",
    "            input_ids,\n",
    "            patch_images=inputs,\n",
    "            num_beams=num_beams,\n",
    "            no_repeat_ngram_size=no_repeat_ngram_size,\n",
    "            max_length=max_length,\n",
    "            min_length=min_length\n",
    "        )\n",
    "\n",
    "    outputs = outputs.detach().cpu()\n",
    "    outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)\n",
    "\n",
    "    return outputs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "06aea735",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:22:57.952540Z",
     "iopub.status.busy": "2023-05-15T14:22:57.952142Z",
     "iopub.status.idle": "2023-05-15T14:22:57.970256Z",
     "shell.execute_reply": "2023-05-15T14:22:57.969325Z"
    },
    "papermill": {
     "duration": 0.052216,
     "end_time": "2023-05-15T14:22:57.973087",
     "exception": false,
     "start_time": "2023-05-15T14:22:57.920871",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "if ofa_large_caption_inference:\n",
    "\n",
    "    set_seed(42, deterministic_cudnn=False)\n",
    "    device = torch.device('cuda')\n",
    "\n",
    "    ofa_model, ofa_tokenizer = load_ofa(\n",
    "        model_directory_path=str(external_dataset / 'ofa-large-caption'),\n",
    "        device=device\n",
    "    )\n",
    "\n",
    "    ofa_transforms = A.Compose([\n",
    "        A.Resize(\n",
    "            height=480,\n",
    "            width=480,\n",
    "            interpolation=cv2.INTER_NEAREST,\n",
    "            always_apply=True\n",
    "        ),\n",
    "        A.Normalize(\n",
    "            mean=[0.5, 0.5, 0.5],\n",
    "            std=[0.5, 0.5, 0.5],\n",
    "            max_pixel_value=255,\n",
    "            always_apply=True\n",
    "        ),\n",
    "        ToTensorV2(always_apply=True)\n",
    "    ])\n",
    "\n",
    "    ofa_dataset = ImageCaptionDataset(\n",
    "        image_paths=image_paths,\n",
    "        captions=None,\n",
    "        transforms=ofa_transforms,\n",
    "        image_reader='opencv'\n",
    "    )\n",
    "\n",
    "    ofa_data_loader = DataLoader(\n",
    "        ofa_dataset,\n",
    "        batch_size=16,\n",
    "        sampler=SequentialSampler(ofa_dataset),\n",
    "        pin_memory=True,\n",
    "        drop_last=False,\n",
    "        num_workers=2,\n",
    "    )\n",
    "    \n",
    "    ofa_predictions = []\n",
    "\n",
    "    for inputs in tqdm(ofa_data_loader):\n",
    "        inputs = inputs.to(device)\n",
    "        predictions = predict_ofa(\n",
    "            inputs=inputs,\n",
    "            model=ofa_model,\n",
    "            tokenizer=ofa_tokenizer,\n",
    "            prompt='what does the image describe?',\n",
    "            num_beams=3,\n",
    "            no_repeat_ngram_size=4,\n",
    "            max_length=25,\n",
    "            min_length=5,\n",
    "            device=device\n",
    "        )\n",
    "        ofa_predictions.append(predictions)\n",
    "\n",
    "    ofa_predictions = np.concatenate(ofa_predictions)\n",
    "    df_test['ofa_predictions'] = ofa_predictions\n",
    "    \n",
    "    ofa_model.to('cpu')\n",
    "    del ofa_model, ofa_tokenizer, ofa_predictions\n",
    "    del ofa_transforms, ofa_dataset, ofa_data_loader\n",
    "    torch.cuda.empty_cache()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a83392a",
   "metadata": {
    "papermill": {
     "duration": 0.028886,
     "end_time": "2023-05-15T14:22:58.031179",
     "exception": false,
     "start_time": "2023-05-15T14:22:58.002293",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 5. CoCa (open_clip)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "d5c5b9a5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:22:58.090902Z",
     "iopub.status.busy": "2023-05-15T14:22:58.090472Z",
     "iopub.status.idle": "2023-05-15T14:22:58.096092Z",
     "shell.execute_reply": "2023-05-15T14:22:58.094985Z"
    },
    "papermill": {
     "duration": 0.041318,
     "end_time": "2023-05-15T14:22:58.101116",
     "exception": false,
     "start_time": "2023-05-15T14:22:58.059798",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "coca_vit_l_14_laion2B_s13B_b90k_inference = True\n",
    "mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_inference = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "1e1cbcc6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:22:58.159365Z",
     "iopub.status.busy": "2023-05-15T14:22:58.158995Z",
     "iopub.status.idle": "2023-05-15T14:22:58.173817Z",
     "shell.execute_reply": "2023-05-15T14:22:58.172793Z"
    },
    "papermill": {
     "duration": 0.047426,
     "end_time": "2023-05-15T14:22:58.176839",
     "exception": false,
     "start_time": "2023-05-15T14:22:58.129413",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def load_coca(model_name, model_checkpoint_path, device):\n",
    "\n",
    "    \"\"\"\n",
    "    Load model and processor\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    model_name: str\n",
    "        Name of the model\n",
    "\n",
    "    model_checkpoint_path: str\n",
    "        Path of the model checkpoint\n",
    "\n",
    "    device: torch.device\n",
    "        Location of the model\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    model: open_clip.CoCa\n",
    "        Model\n",
    "\n",
    "    processor: torchvision.transforms.Compose\n",
    "        Image processor\n",
    "    \"\"\"\n",
    "\n",
    "    model = open_clip.create_model(model_name=model_name, pretrained=model_checkpoint_path)\n",
    "    model.to(device)\n",
    "    model.eval()\n",
    "\n",
    "    processor = open_clip.image_transform(\n",
    "        model.visual.image_size,\n",
    "        is_train=False,\n",
    "        mean=getattr(model.visual, 'image_mean', None),\n",
    "        std=getattr(model.visual, 'image_std', None),\n",
    "    )\n",
    "\n",
    "    return model, processor\n",
    "\n",
    "\n",
    "def predict_coca(inputs, model, device):\n",
    "\n",
    "    \"\"\"\n",
    "    Predict given inputs with given model\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    inputs: torch.FloatTensor of shape (batch, height, width, channel)\n",
    "        Inputs tensor\n",
    "\n",
    "    model: transformers.VisionEncoderDecoderModel\n",
    "        Model\n",
    "\n",
    "    device: torch.device\n",
    "        Location of the model\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    outputs: torch.Tensor of shape (batch, max_length) or list of shape (batch)\n",
    "        Outputs\n",
    "    \"\"\"\n",
    "\n",
    "    with torch.no_grad(), torch.autocast(device_type=device.type, dtype=torch.float16):\n",
    "        outputs = model.generate(inputs)\n",
    "\n",
    "    outputs = outputs.detach().cpu()\n",
    "    outputs = [open_clip.decode(output).split('<end_of_text>')[0].replace('<start_of_text>', '').rstrip(' .,') for output in outputs]\n",
    "\n",
    "    return outputs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "b66f3d91",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:22:58.235333Z",
     "iopub.status.busy": "2023-05-15T14:22:58.234926Z",
     "iopub.status.idle": "2023-05-15T14:23:54.936671Z",
     "shell.execute_reply": "2023-05-15T14:23:54.935371Z"
    },
    "papermill": {
     "duration": 56.734453,
     "end_time": "2023-05-15T14:23:54.940064",
     "exception": false,
     "start_time": "2023-05-15T14:22:58.205611",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:07<00:00,  7.28s/it]\n"
     ]
    }
   ],
   "source": [
    "if coca_vit_l_14_laion2B_s13B_b90k_inference:\n",
    "\n",
    "    set_seed(42, deterministic_cudnn=False)\n",
    "    device = torch.device('cuda')\n",
    "\n",
    "    coca_model, coca_processor = load_coca(\n",
    "        model_name='coca_ViT-L-14',\n",
    "        model_checkpoint_path=str(external_dataset / 'CoCa-ViT-L-14-laion2B-s13B-b90k' / 'open_clip_pytorch_model.bin'),\n",
    "        device=device,\n",
    "    )\n",
    "\n",
    "    coca_dataset = ImageCaptionDataset(\n",
    "        image_paths=image_paths,\n",
    "        captions=None,\n",
    "        transforms=coca_processor,\n",
    "        image_reader='pil'\n",
    "    )\n",
    "\n",
    "    coca_data_loader = DataLoader(\n",
    "        coca_dataset,\n",
    "        batch_size=16,\n",
    "        sampler=SequentialSampler(coca_dataset),\n",
    "        pin_memory=True,\n",
    "        drop_last=False,\n",
    "        num_workers=2,\n",
    "    )\n",
    "    \n",
    "    coca_predictions = []\n",
    "\n",
    "    for inputs in tqdm(coca_data_loader):\n",
    "        inputs = inputs.to(device)\n",
    "        predictions = predict_coca(\n",
    "            inputs=inputs,\n",
    "            model=coca_model,\n",
    "            device=device\n",
    "        )\n",
    "        coca_predictions.append(predictions)\n",
    "\n",
    "    coca_predictions = np.concatenate(coca_predictions)\n",
    "    df_test['coca_vit_l_14_laion2B_s13B_b90k_prediction'] = coca_predictions\n",
    "    \n",
    "    coca_model.to('cpu')\n",
    "    del coca_model, coca_processor, coca_predictions\n",
    "    del coca_dataset, coca_data_loader\n",
    "    torch.cuda.empty_cache()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "fc97150a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:23:54.984908Z",
     "iopub.status.busy": "2023-05-15T14:23:54.984537Z",
     "iopub.status.idle": "2023-05-15T14:24:41.034384Z",
     "shell.execute_reply": "2023-05-15T14:24:41.032996Z"
    },
    "papermill": {
     "duration": 46.074627,
     "end_time": "2023-05-15T14:24:41.037333",
     "exception": false,
     "start_time": "2023-05-15T14:23:54.962706",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:03<00:00,  3.20s/it]\n"
     ]
    }
   ],
   "source": [
    "if mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_inference:\n",
    "\n",
    "    set_seed(42, deterministic_cudnn=False)\n",
    "    device = torch.device('cuda')\n",
    "\n",
    "    coca_model, coca_processor = load_coca(\n",
    "        model_name='coca_ViT-L-14',\n",
    "        model_checkpoint_path=str(external_dataset / 'mscoco_finetuned_CoCa-ViT-L-14-laion2B-s13B-b90k' / 'open_clip_pytorch_model.bin'),\n",
    "        device=device,\n",
    "    )\n",
    "\n",
    "    coca_dataset = ImageCaptionDataset(\n",
    "        image_paths=image_paths,\n",
    "        captions=None,\n",
    "        transforms=coca_processor,\n",
    "        image_reader='pil'\n",
    "    )\n",
    "\n",
    "    coca_data_loader = DataLoader(\n",
    "        coca_dataset,\n",
    "        batch_size=16,\n",
    "        sampler=SequentialSampler(coca_dataset),\n",
    "        pin_memory=True,\n",
    "        drop_last=False,\n",
    "        num_workers=2,\n",
    "    )\n",
    "    \n",
    "    coca_predictions = []\n",
    "\n",
    "    for inputs in tqdm(coca_data_loader):\n",
    "        inputs = inputs.to(device)\n",
    "        predictions = predict_coca(\n",
    "            inputs=inputs,\n",
    "            model=coca_model,\n",
    "            device=device\n",
    "        )\n",
    "        coca_predictions.append(predictions)\n",
    "\n",
    "    coca_predictions = np.concatenate(coca_predictions)\n",
    "    df_test['mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_prediction'] = coca_predictions\n",
    "    \n",
    "    coca_model.to('cpu')\n",
    "    del coca_model, coca_processor, coca_predictions\n",
    "    del coca_dataset, coca_data_loader\n",
    "    torch.cuda.empty_cache()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c9debce",
   "metadata": {
    "papermill": {
     "duration": 0.02049,
     "end_time": "2023-05-15T14:24:41.078827",
     "exception": false,
     "start_time": "2023-05-15T14:24:41.058337",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 6. BLIP (lavis)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "95f8bb9b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:24:41.122176Z",
     "iopub.status.busy": "2023-05-15T14:24:41.121839Z",
     "iopub.status.idle": "2023-05-15T14:24:41.127193Z",
     "shell.execute_reply": "2023-05-15T14:24:41.126110Z"
    },
    "papermill": {
     "duration": 0.029518,
     "end_time": "2023-05-15T14:24:41.129577",
     "exception": false,
     "start_time": "2023-05-15T14:24:41.100059",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "blip_caption_base_inference = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "b5f61b3c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:24:41.171812Z",
     "iopub.status.busy": "2023-05-15T14:24:41.171520Z",
     "iopub.status.idle": "2023-05-15T14:24:41.186614Z",
     "shell.execute_reply": "2023-05-15T14:24:41.185574Z"
    },
    "papermill": {
     "duration": 0.038844,
     "end_time": "2023-05-15T14:24:41.188823",
     "exception": false,
     "start_time": "2023-05-15T14:24:41.149979",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def load_lavis_blip(model_name, model_type, model_config_path, device):\n",
    "\n",
    "    \"\"\"\n",
    "    Load model and processor\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    model_name: str\n",
    "        Name of the model\n",
    "\n",
    "    model_type: str\n",
    "        Type of the model\n",
    "\n",
    "    model_config_path: str\n",
    "        Path of the model checkpoint\n",
    "        \n",
    "    device: torch.device\n",
    "        Location of the model\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    model: lavis.models.BlipCaption\n",
    "        Model\n",
    "\n",
    "    image_processor: dict\n",
    "        Image processor\n",
    "\n",
    "    text_processor: dict\n",
    "        Text processor\n",
    "    \"\"\"\n",
    "    \n",
    "    config = yaml.load(open(model_config_path, 'r'), Loader=yaml.FullLoader)\n",
    "    config['model']['pretrained'] = str(external_dataset / '/'.join(config['model']['pretrained'].split('/')[-2:]))\n",
    "    config['model']['finetuned'] = str(external_dataset / '/'.join(config['model']['finetuned'].split('/')[-2:]))\n",
    "    config['model']['med_config_path'] = str(external_dataset / '/'.join(config['model']['med_config_path'].split('/')[-2:]))\n",
    "    with open(str(model_config_path).split('/')[-1], mode='w') as f:\n",
    "        yaml.dump(config, f, default_flow_style=False)\n",
    "    \n",
    "    copy2(str(model_config_path).split('/')[-1], f'/opt/conda/lib/python3.7/site-packages/lavis/{str(model_config_path).split(\"/\")[-1]}')\n",
    "    \n",
    "    model_class = registry.get_model_class(model_name)\n",
    "    model_class.PRETRAINED_MODEL_CONFIG_DICT[model_type] = str(model_config_path).split('/')[-1]\n",
    "    model, image_processor, text_processor = load_model_and_preprocess(name=model_name, model_type=model_type)\n",
    "    \n",
    "    model.to(device)\n",
    "    model.eval()\n",
    "\n",
    "    return model, image_processor, text_processor\n",
    "\n",
    "\n",
    "def predict_lavis_blip(inputs, model, nucleus_sampling=False, num_beams=3, max_length=25, min_length=5):\n",
    "\n",
    "    \"\"\"\n",
    "    Predict given inputs with given model\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    inputs: torch.FloatTensor of shape (batch, height, width, channel) images\n",
    "        Inputs tensor\n",
    "\n",
    "    model: transformers.BlipForConditionalGeneration\n",
    "        Model\n",
    "\n",
    "    nucleus_sampling: bool\n",
    "        Whether to use nucleus sampling or not\n",
    "\n",
    "    num_beams: int\n",
    "        Number of beams in beam search\n",
    "\n",
    "    max_length: int\n",
    "        Max length of the generated sequence\n",
    "\n",
    "    min_length: int\n",
    "        Min length of the generated sequence\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    outputs: list of shape (batch) texts\n",
    "        Outputs\n",
    "    \"\"\"\n",
    "\n",
    "    with torch.no_grad(), torch.autocast(device_type=device.type, dtype=torch.float16):\n",
    "        outputs = model.generate(\n",
    "            samples={'image': inputs},\n",
    "            use_nucleus_sampling=nucleus_sampling,\n",
    "            num_beams=num_beams,\n",
    "            max_length=max_length,\n",
    "            min_length=min_length,\n",
    "            top_p=0.9,\n",
    "            repetition_penalty=1.0,\n",
    "            num_captions=1\n",
    "        )\n",
    "\n",
    "    return outputs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "fb7b17fc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:24:41.230751Z",
     "iopub.status.busy": "2023-05-15T14:24:41.230460Z",
     "iopub.status.idle": "2023-05-15T14:24:41.240252Z",
     "shell.execute_reply": "2023-05-15T14:24:41.239194Z"
    },
    "papermill": {
     "duration": 0.033408,
     "end_time": "2023-05-15T14:24:41.242575",
     "exception": false,
     "start_time": "2023-05-15T14:24:41.209167",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "if blip_caption_base_inference:\n",
    "\n",
    "    set_seed(42, deterministic_cudnn=False)\n",
    "    device = torch.device('cuda')\n",
    "\n",
    "    blip_model, blip_image_processor, blip_text_processor = load_lavis_blip(\n",
    "        model_name='blip_caption',\n",
    "        model_type='base_coco',\n",
    "        model_config_path=external_dataset / 'blip-caption-base' / 'blip_caption_base_coco.yaml',\n",
    "        device=device\n",
    "    )\n",
    "\n",
    "    blip_dataset = ImageCaptionDataset(\n",
    "        image_paths=image_paths,\n",
    "        captions=None,\n",
    "        transforms=blip_image_processor['eval'],\n",
    "        image_reader='pil'\n",
    "    )\n",
    "\n",
    "    blip_data_loader = DataLoader(\n",
    "        blip_dataset,\n",
    "        batch_size=16,\n",
    "        sampler=SequentialSampler(blip_dataset),\n",
    "        pin_memory=True,\n",
    "        drop_last=False,\n",
    "        num_workers=2,\n",
    "    )\n",
    "    \n",
    "    blip_predictions = []\n",
    "\n",
    "    for inputs in tqdm(blip_data_loader):\n",
    "        inputs = inputs.to(device)\n",
    "        predictions = predict_lavis_blip(\n",
    "            inputs=inputs,\n",
    "            model=blip_model\n",
    "        )\n",
    "        blip_predictions.append(predictions)\n",
    "\n",
    "    blip_predictions = np.concatenate(blip_predictions)\n",
    "    df_test['blip_caption_base_prediction'] = blip_predictions\n",
    "    \n",
    "    blip_model.to('cpu')\n",
    "    del blip_model, blip_image_processor, blip_text_processor, blip_predictions\n",
    "    del blip_dataset, blip_data_loader\n",
    "    torch.cuda.empty_cache()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d5e906f",
   "metadata": {
    "papermill": {
     "duration": 0.020991,
     "end_time": "2023-05-15T14:24:41.284900",
     "exception": false,
     "start_time": "2023-05-15T14:24:41.263909",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 7. BLIP (transformers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "2966d8c4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:24:41.327737Z",
     "iopub.status.busy": "2023-05-15T14:24:41.326832Z",
     "iopub.status.idle": "2023-05-15T14:24:41.331614Z",
     "shell.execute_reply": "2023-05-15T14:24:41.330635Z"
    },
    "papermill": {
     "duration": 0.02862,
     "end_time": "2023-05-15T14:24:41.333982",
     "exception": false,
     "start_time": "2023-05-15T14:24:41.305362",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "blip_image_captioning_large_inference = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "64d5a57d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:24:41.376853Z",
     "iopub.status.busy": "2023-05-15T14:24:41.376065Z",
     "iopub.status.idle": "2023-05-15T14:24:41.385558Z",
     "shell.execute_reply": "2023-05-15T14:24:41.384562Z"
    },
    "papermill": {
     "duration": 0.033427,
     "end_time": "2023-05-15T14:24:41.387808",
     "exception": false,
     "start_time": "2023-05-15T14:24:41.354381",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def load_transformers_blip(model_directory_path, device):\n",
    "\n",
    "    \"\"\"\n",
    "    Load model and processor\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    model_directory_path: str\n",
    "        Directory of the model\n",
    "\n",
    "    device: torch.device\n",
    "        Location of the model\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    model: transformers.BlipForConditionalGeneration\n",
    "        Model\n",
    "\n",
    "    processor: transformers.BlipProcessor\n",
    "        Processor\n",
    "    \"\"\"\n",
    "\n",
    "    model = BlipForConditionalGeneration.from_pretrained(model_directory_path)\n",
    "    model.to(device)\n",
    "    model.eval()\n",
    "\n",
    "    processor = BlipProcessor.from_pretrained(model_directory_path)\n",
    "\n",
    "    return model, processor\n",
    "\n",
    "\n",
    "def predict_transformers_blip(inputs, model, processor, device, num_beams, max_length, min_length):\n",
    "\n",
    "    \"\"\"\n",
    "    Predict given inputs with given model\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    inputs: torch.FloatTensor of shape (batch, height, width, channel) images\n",
    "        Inputs tensor\n",
    "\n",
    "    model: transformers.BlipForConditionalGeneration\n",
    "        Model\n",
    "        \n",
    "    processor: transformers.BlipProcessor\n",
    "        Processor\n",
    "        \n",
    "    device: torch.device\n",
    "        Location of the model\n",
    "\n",
    "    num_beams: int\n",
    "        Number of beams in beam search\n",
    "\n",
    "    max_length: int\n",
    "        Max length of the generated sequence\n",
    "\n",
    "    min_length: int\n",
    "        Min length of the generated sequence\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    outputs: list of shape (batch) texts\n",
    "        Outputs\n",
    "    \"\"\"\n",
    "\n",
    "    with torch.no_grad(), torch.autocast(device_type=device.type, dtype=torch.float16):\n",
    "        outputs = model.generate(pixel_values=inputs, num_beams=num_beams, max_length=max_length, min_length=min_length)\n",
    "\n",
    "    outputs = processor.batch_decode(outputs, skip_special_tokens=True)\n",
    "\n",
    "    return outputs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "270a7b37",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:24:41.431049Z",
     "iopub.status.busy": "2023-05-15T14:24:41.430283Z",
     "iopub.status.idle": "2023-05-15T14:25:25.028896Z",
     "shell.execute_reply": "2023-05-15T14:25:25.027498Z"
    },
    "papermill": {
     "duration": 43.623182,
     "end_time": "2023-05-15T14:25:25.031816",
     "exception": false,
     "start_time": "2023-05-15T14:24:41.408634",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:04<00:00,  4.53s/it]\n"
     ]
    }
   ],
   "source": [
    "if blip_image_captioning_large_inference:\n",
    "\n",
    "    set_seed(42, deterministic_cudnn=False)\n",
    "    device = torch.device('cuda')\n",
    "\n",
    "    blip_model, blip_processor = load_transformers_blip(\n",
    "        model_directory_path=external_dataset / 'blip-image-captioning-large',\n",
    "        device=device\n",
    "    )\n",
    "\n",
    "    blip_dataset = ImageCaptionDataset(\n",
    "        image_paths=image_paths,\n",
    "        captions=None,\n",
    "        transforms=blip_processor,\n",
    "        image_reader='pil'\n",
    "    )\n",
    "\n",
    "    blip_data_loader = DataLoader(\n",
    "        blip_dataset,\n",
    "        batch_size=16,\n",
    "        sampler=SequentialSampler(blip_dataset),\n",
    "        pin_memory=True,\n",
    "        drop_last=False,\n",
    "        num_workers=2,\n",
    "    )\n",
    "    \n",
    "    blip_predictions = []\n",
    "\n",
    "    for inputs in tqdm(blip_data_loader):\n",
    "        inputs = inputs.to(device)\n",
    "        predictions = predict_transformers_blip(\n",
    "            inputs=inputs,\n",
    "            model=blip_model,\n",
    "            processor=blip_processor,\n",
    "            device=device,\n",
    "            num_beams=1,\n",
    "            max_length=25,\n",
    "            min_length=5\n",
    "        )\n",
    "        blip_predictions.append(predictions)\n",
    "\n",
    "    blip_predictions = np.concatenate(blip_predictions)\n",
    "    df_test['blip_image_captioning_large_prediction'] = blip_predictions\n",
    "    \n",
    "    blip_model.to('cpu')\n",
    "    del blip_model, blip_processor, blip_predictions\n",
    "    del blip_dataset, blip_data_loader\n",
    "    torch.cuda.empty_cache()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e716dddf",
   "metadata": {
    "papermill": {
     "duration": 0.021469,
     "end_time": "2023-05-15T14:25:25.074873",
     "exception": false,
     "start_time": "2023-05-15T14:25:25.053404",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 8. CLIP-Interrogator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "b3978421",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:25:25.120097Z",
     "iopub.status.busy": "2023-05-15T14:25:25.119143Z",
     "iopub.status.idle": "2023-05-15T14:25:25.125006Z",
     "shell.execute_reply": "2023-05-15T14:25:25.124252Z"
    },
    "papermill": {
     "duration": 0.031015,
     "end_time": "2023-05-15T14:25:25.127162",
     "exception": false,
     "start_time": "2023-05-15T14:25:25.096147",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "coca_vit_l_14_laion2B_s13B_b90k_interrogation = False\n",
    "mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_interrogation = True\n",
    "blip_caption_base_interrogation = False\n",
    "blip_image_captioning_large_interrogation = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "1a068c48",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:25:25.171766Z",
     "iopub.status.busy": "2023-05-15T14:25:25.170094Z",
     "iopub.status.idle": "2023-05-15T14:25:25.178349Z",
     "shell.execute_reply": "2023-05-15T14:25:25.177198Z"
    },
    "papermill": {
     "duration": 0.03264,
     "end_time": "2023-05-15T14:25:25.180682",
     "exception": false,
     "start_time": "2023-05-15T14:25:25.148042",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def load_clip(model_name, model_checkpoint_path, device):\n",
    "\n",
    "    \"\"\"\n",
    "    Load model and processor\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    model_name: str\n",
    "        Name of the model\n",
    "\n",
    "    model_checkpoint_path: str\n",
    "        Path of the model checkpoint\n",
    "\n",
    "    device: torch.device\n",
    "        Location of the model\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    model: open_clip.CLIP\n",
    "        Model\n",
    "\n",
    "    processor: torchvision.transforms.Compose\n",
    "        Image processor\n",
    "    \"\"\"\n",
    "\n",
    "    model = open_clip.create_model(\n",
    "        model_name=model_name,\n",
    "        pretrained=model_checkpoint_path\n",
    "    )\n",
    "    model.to(device)\n",
    "    model.eval()\n",
    "\n",
    "    processor = open_clip.image_transform(\n",
    "        model.visual.image_size,\n",
    "        is_train=False,\n",
    "        mean=getattr(model.visual, 'image_mean', None),\n",
    "        std=getattr(model.visual, 'image_std', None),\n",
    "    )\n",
    "\n",
    "    return model, processor\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "2bc8e007",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:25:25.224250Z",
     "iopub.status.busy": "2023-05-15T14:25:25.223914Z",
     "iopub.status.idle": "2023-05-15T14:25:25.244912Z",
     "shell.execute_reply": "2023-05-15T14:25:25.243942Z"
    },
    "papermill": {
     "duration": 0.04582,
     "end_time": "2023-05-15T14:25:25.247324",
     "exception": false,
     "start_time": "2023-05-15T14:25:25.201504",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class InterrogatorWithoutForward(clip_interrogator.Interrogator):\n",
    "\n",
    "    def __init__(self, config):\n",
    "\n",
    "        super(InterrogatorWithoutForward, self).__init__(config)\n",
    "\n",
    "    def batch_image_to_features(self, inputs, device):\n",
    "\n",
    "        \"\"\"\n",
    "        Extract image features\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        inputs: torch.FloatTensor of shape (batch, height, width, channel)\n",
    "            Inputs tensor\n",
    "\n",
    "        device: torch.device\n",
    "            Location of the model\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        image_features: torch.FloatTensor of shape (batch, features)\n",
    "            Feature tensor\n",
    "        \"\"\"\n",
    "\n",
    "        with torch.no_grad(), torch.autocast(device_type=device.type, dtype=torch.float16):\n",
    "            image_features = self.clip_model.encode_image(inputs)\n",
    "\n",
    "        image_features = image_features / image_features.norm(dim=1).view(-1, 1)\n",
    "        image_features = image_features.detach().cpu()\n",
    "\n",
    "        return image_features\n",
    "\n",
    "    def interrogate_classic_without_forward(self, image_features, max_flavors, caption):\n",
    "\n",
    "        \"\"\"\n",
    "        Classical interrogation with single medium, artist, trending and movement and given number of flavors\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        image_features: torch.FloatTensor of shape (batch, features)\n",
    "            Feature tensor\n",
    "\n",
    "        max_flavors: int\n",
    "            Maximum number of flavors\n",
    "\n",
    "        caption: str\n",
    "            Caption of the image\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        caption: str\n",
    "            Caption after interrogation\n",
    "        \"\"\"\n",
    "\n",
    "        medium = self.mediums.rank(image_features, 1)[0]\n",
    "        movement = self.movements.rank(image_features, 1)[0]\n",
    "        flaves = \", \".join(self.flavors.rank(image_features, max_flavors))\n",
    "\n",
    "        if caption.startswith(medium):\n",
    "            prompt = f\"{caption}, {movement}, {flaves}\"\n",
    "        else:\n",
    "            prompt = f\"{caption}, {medium}, {movement}, {flaves}\"\n",
    "\n",
    "        return clip_interrogator._truncate_to_fit(prompt, self.tokenize)\n",
    "\n",
    "    def interrogate_fast_without_forward(self, image_features, max_flavors, caption):\n",
    "\n",
    "        \"\"\"\n",
    "        Fast interrogation with given number of flavors, mediums, artists, trendings and movements\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        image_features: torch.FloatTensor of shape (batch, features)\n",
    "            Feature tensor\n",
    "\n",
    "        max_flavors: int\n",
    "            Maximum number of flavors\n",
    "\n",
    "        caption: str\n",
    "            Caption of the image\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        caption: str\n",
    "            Caption after interrogation\n",
    "        \"\"\"\n",
    "\n",
    "        merged = clip_interrogator._merge_tables([self.artists, self.flavors, self.mediums, self.movements, self.trendings], self)\n",
    "        tops = merged.rank(image_features, max_flavors)\n",
    "\n",
    "        return clip_interrogator._truncate_to_fit(caption + \", \" + \", \".join(tops), self.tokenize)\n",
    "\n",
    "    def interrogate_negative_without_forward(self, image_features, max_flavors):\n",
    "\n",
    "        \"\"\"\n",
    "        Negative interrogation with given number of flavors\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        image_features: torch.FloatTensor of shape (batch, features)\n",
    "            Feature tensor\n",
    "\n",
    "        max_flavors: int\n",
    "            Maximum number of flavors\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        caption: str\n",
    "            Caption after interrogation\n",
    "        \"\"\"\n",
    "\n",
    "        flaves = self.flavors.rank(image_features, self.config.flavor_intermediate_count, reverse=True)\n",
    "        flaves = flaves + self.negative.labels\n",
    "\n",
    "        return self.chain(image_features, flaves, max_count=max_flavors, reverse=True)\n",
    "\n",
    "    def interrogate_without_forward(self, image_features, min_flavors, max_flavors, caption):\n",
    "\n",
    "        \"\"\"\n",
    "        Negative interrogation with given number of flavors\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        image_features: torch.FloatTensor of shape (batch, features)\n",
    "            Feature tensor\n",
    "\n",
    "        max_flavors: int\n",
    "            Maximum number of flavors\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        caption: str\n",
    "            Caption after interrogation\n",
    "        \"\"\"\n",
    "\n",
    "        merged = clip_interrogator._merge_tables([self.artists, self.flavors, self.mediums, self.movements, self.trendings], self)\n",
    "        flaves = merged.rank(image_features, self.config.flavor_intermediate_count)\n",
    "        best_prompt, best_sim = caption, self.similarity(image_features, caption)\n",
    "\n",
    "        best_prompt = self.chain(image_features, flaves, best_prompt, best_sim, min_count=min_flavors, max_count=max_flavors)\n",
    "\n",
    "        fast_prompt = self.interrogate_fast_without_forward(image_features, max_flavors, caption=caption)\n",
    "        classic_prompt = self.interrogate_classic_without_forward(image_features, max_flavors, caption=caption)\n",
    "        candidates = [caption, classic_prompt, fast_prompt, best_prompt]\n",
    "        \n",
    "        return candidates[np.argmax(self.similarities(image_features, candidates))]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "cfe9da8a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:25:25.291611Z",
     "iopub.status.busy": "2023-05-15T14:25:25.291194Z",
     "iopub.status.idle": "2023-05-15T14:26:32.701212Z",
     "shell.execute_reply": "2023-05-15T14:26:32.699930Z"
    },
    "papermill": {
     "duration": 67.436034,
     "end_time": "2023-05-15T14:26:32.704494",
     "exception": false,
     "start_time": "2023-05-15T14:25:25.268460",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:00<00:00,  1.12it/s]\n"
     ]
    }
   ],
   "source": [
    "if coca_vit_l_14_laion2B_s13B_b90k_interrogation or mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_interrogation or blip_caption_base_interrogation or blip_image_captioning_large_interrogation:\n",
    "\n",
    "    set_seed(42, deterministic_cudnn=False)\n",
    "    device = torch.device('cuda')\n",
    "\n",
    "    clip_model, clip_processor = load_clip(\n",
    "        model_name='ViT-H-14',\n",
    "        model_checkpoint_path=str(external_dataset / 'CLIP-ViT-H-14-laion2B-s32B-b79K' / 'open_clip_pytorch_model.bin'),\n",
    "        device=device\n",
    "    )\n",
    "    \n",
    "    clip_dataset = ImageCaptionDataset(\n",
    "        image_paths=image_paths,\n",
    "        captions=None,\n",
    "        transforms=clip_processor,\n",
    "        image_reader='pil'\n",
    "    )\n",
    "    \n",
    "    clip_data_loader = DataLoader(\n",
    "        clip_dataset,\n",
    "        batch_size=32,\n",
    "        sampler=SequentialSampler(clip_dataset),\n",
    "        pin_memory=False,\n",
    "        drop_last=False,\n",
    "        num_workers=2,\n",
    "    )\n",
    "    \n",
    "    clip_image_features = []\n",
    "\n",
    "    for inputs in tqdm(clip_data_loader):\n",
    "        \n",
    "        inputs = inputs.to(device)\n",
    "        \n",
    "        with torch.no_grad(), torch.autocast(device_type=device.type, dtype=torch.float16):\n",
    "            image_features = clip_model.encode_image(inputs)\n",
    "\n",
    "        image_features = image_features / image_features.norm(dim=1).view(-1, 1)\n",
    "        image_features = image_features.detach().cpu()\n",
    "        clip_image_features.append(image_features)\n",
    "\n",
    "    clip_image_features = torch.cat(clip_image_features, dim=0)\n",
    "    clip_model.to('cpu')\n",
    "    torch.cuda.empty_cache()\n",
    "\n",
    "    del clip_dataset, clip_data_loader\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "5d33f578",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:26:32.750226Z",
     "iopub.status.busy": "2023-05-15T14:26:32.749170Z",
     "iopub.status.idle": "2023-05-15T14:26:34.826571Z",
     "shell.execute_reply": "2023-05-15T14:26:34.825471Z"
    },
    "papermill": {
     "duration": 2.103047,
     "end_time": "2023-05-15T14:26:34.829414",
     "exception": false,
     "start_time": "2023-05-15T14:26:32.726367",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "if coca_vit_l_14_laion2B_s13B_b90k_interrogation or mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_interrogation or blip_caption_base_interrogation or blip_image_captioning_large_interrogation:\n",
    "    \n",
    "    ci_config = clip_interrogator.Config()\n",
    "    ci_config.caption_model = False\n",
    "    ci_config.caption_processor = False\n",
    "    ci_config.clip_model = clip_model\n",
    "    ci_config.clip_preprocess = clip_processor\n",
    "    ci_config.quiet = True\n",
    "    ci_config.download_cache = False\n",
    "    ci_config.clip_model_name = 'ViT-H-14/laion2b_s32b_b79k'\n",
    "    ci_config.cache_path = './ci-preprocess'\n",
    "    ci = InterrogatorWithoutForward(ci_config)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "fd6e0205",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:26:34.874546Z",
     "iopub.status.busy": "2023-05-15T14:26:34.874163Z",
     "iopub.status.idle": "2023-05-15T14:26:37.333854Z",
     "shell.execute_reply": "2023-05-15T14:26:37.332505Z"
    },
    "papermill": {
     "duration": 2.485859,
     "end_time": "2023-05-15T14:26:37.337139",
     "exception": false,
     "start_time": "2023-05-15T14:26:34.851280",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7/7 [00:02<00:00,  2.88it/s]\n"
     ]
    }
   ],
   "source": [
    "if blip_image_captioning_large_interrogation:\n",
    "    \n",
    "    for idx, (image_features, caption) in tqdm(\n",
    "        enumerate(zip(\n",
    "            clip_image_features,\n",
    "            df_test['blip_image_captioning_large_prediction'].values\n",
    "        )),\n",
    "        total=clip_image_features.shape[0]\n",
    "    ):\n",
    "        interrogated_caption = ci.interrogate_classic_without_forward(torch.unsqueeze(image_features, dim=0).to(device), 1, caption)\n",
    "        df_test.loc[idx, 'interrogated_blip_image_captioning_large_prediction'] = interrogated_caption\n",
    "\n",
    "if coca_vit_l_14_laion2B_s13B_b90k_interrogation:\n",
    "    \n",
    "    for idx, (image_features, caption) in tqdm(\n",
    "        enumerate(zip(\n",
    "            clip_image_features,\n",
    "            df_test['coca_vit_l_14_laion2B_s13B_b90k_prediction'].values\n",
    "        )),\n",
    "        total=clip_image_features.shape[0]\n",
    "    ):\n",
    "        interrogated_caption = ci.interrogate_classic_without_forward(torch.unsqueeze(image_features, dim=0).to(device), 1, caption)\n",
    "        df_test.loc[idx, 'interrogated_coca_vit_l_14_laion2B_s13B_b90k_prediction'] = interrogated_caption\n",
    "        \n",
    "        \n",
    "if mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_interrogation:\n",
    "    \n",
    "    for idx, (image_features, caption) in tqdm(\n",
    "        enumerate(zip(\n",
    "            clip_image_features,\n",
    "            df_test['mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_prediction'].values\n",
    "        )),\n",
    "        total=clip_image_features.shape[0]\n",
    "    ):\n",
    "        interrogated_caption = ci.interrogate_classic_without_forward(torch.unsqueeze(image_features, dim=0).to(device), 3, caption)\n",
    "        df_test.loc[idx, 'interrogated_mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_prediction'] = interrogated_caption\n",
    "\n",
    "if coca_vit_l_14_laion2B_s13B_b90k_interrogation or mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_interrogation or blip_caption_base_interrogation or blip_image_captioning_large_interrogation:\n",
    "    del clip_model, clip_processor\n",
    "    del ci, ci_config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "f8c6af7f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:26:37.384438Z",
     "iopub.status.busy": "2023-05-15T14:26:37.383413Z",
     "iopub.status.idle": "2023-05-15T14:26:40.497674Z",
     "shell.execute_reply": "2023-05-15T14:26:40.496665Z"
    },
    "papermill": {
     "duration": 3.140806,
     "end_time": "2023-05-15T14:26:40.500142",
     "exception": false,
     "start_time": "2023-05-15T14:26:37.359336",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bd443fca64b4421e8c13269a7481edf2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Batches:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ffab3e510fd5468c95cdfb92c108d7c8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Batches:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "13e78dab7b894644ad2a7195efedcf66",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Batches:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ac7214793e714f4fb75267db61a2f1d8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Batches:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "language_model = SentenceTransformer(external_dataset / 'sentence-transformers_all-MiniLM-L6-v2')\n",
    "\n",
    "coca_vit_l_14_laion2B_s13B_b90k_predictions_embeddings = language_model.encode(df_test['coca_vit_l_14_laion2B_s13B_b90k_prediction'])\n",
    "df_test_embeddings['coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding'] = coca_vit_l_14_laion2B_s13B_b90k_predictions_embeddings.flatten()\n",
    "\n",
    "mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_predictions_embeddings = language_model.encode(df_test['mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_prediction'])\n",
    "df_test_embeddings['mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding'] = mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_predictions_embeddings.flatten()\n",
    "\n",
    "interrogated_mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_predictions_embeddings = language_model.encode(df_test['interrogated_mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_prediction'])\n",
    "df_test_embeddings['interrogated_mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding'] = interrogated_mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_predictions_embeddings.flatten()\n",
    "\n",
    "blip_image_captioning_large_predictions_embeddings = language_model.encode(df_test['blip_image_captioning_large_prediction'])\n",
    "df_test_embeddings['blip_image_captioning_large_prediction_embedding'] = blip_image_captioning_large_predictions_embeddings.flatten()\n",
    "\n",
    "del language_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce82478c",
   "metadata": {
    "papermill": {
     "duration": 0.023276,
     "end_time": "2023-05-15T14:26:40.546849",
     "exception": false,
     "start_time": "2023-05-15T14:26:40.523573",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 9. KNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "b75706d4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:26:40.594878Z",
     "iopub.status.busy": "2023-05-15T14:26:40.594253Z",
     "iopub.status.idle": "2023-05-15T14:26:40.824945Z",
     "shell.execute_reply": "2023-05-15T14:26:40.823848Z"
    },
    "papermill": {
     "duration": 0.258076,
     "end_time": "2023-05-15T14:26:40.827828",
     "exception": false,
     "start_time": "2023-05-15T14:26:40.569752",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "objective_embedding_paths = []\n",
    "clip_embedding_paths = []\n",
    "\n",
    "objective_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-diffusiondb-14m-part1/all_minilm_l6_v2/*.npy'))\n",
    "objective_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-diffusiondb-14m-part2/all_minilm_l6_v2/*.npy'))\n",
    "objective_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-diffusiondb-14m-part3/all_minilm_l6_v2/*.npy'))\n",
    "objective_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-mscoco/all_minilm_l6_v2/*.npy'))\n",
    "objective_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-dataset80k/all_minilm_l6_v2/prompt_embeddings_allminilm_001.npy'))\n",
    "objective_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-dataset30k/all_minilm_l6_v2/prompt_embeddings_allminilm_001.npy'))\n",
    "objective_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-dataset900k/all_minilm_l6_v2/*.npy'))\n",
    "objective_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-conceptual-captions/all_minilm_l6_v2/*.npy'))\n",
    "objective_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-sd2gpt2/all_minilm_l6_v2/*.npy'))\n",
    "objective_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-sd2hardcode/all_minilm_l6_v2/*.npy'))\n",
    "objective_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-chatgpt/all_minilm_l6_v2/*.npy'))\n",
    "objective_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-laion2b-part0000-0004/all_minilm_l6_v2/*.npy'))\n",
    "objective_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-laion2b-part0005-0009/all_minilm_l6_v2/*.npy'))\n",
    "objective_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-laion2b-part0010-0014/all_minilm_l6_v2/*.npy'))\n",
    "objective_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-laion2b-part0015-0019/all_minilm_l6_v2/*.npy'))\n",
    "objective_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-laion2b-part0020-0024/all_minilm_l6_v2/*.npy'))\n",
    "objective_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-laion2b-part0025-0029/all_minilm_l6_v2/*.npy'))\n",
    "objective_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-laion2b-part0030-0034/all_minilm_l6_v2/*.npy'))\n",
    "objective_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-laion2b-part0035-0039/all_minilm_l6_v2/*.npy'))\n",
    "objective_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-laion2b-part0040-0044/all_minilm_l6_v2/*.npy'))\n",
    "objective_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-laion2b-part0045-0049/all_minilm_l6_v2/*.npy'))\n",
    "\n",
    "clip_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-diffusiondb-14m-part1/vith14/*.npy'))\n",
    "clip_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-diffusiondb-14m-part2/vith14/*.npy'))\n",
    "clip_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-diffusiondb-14m-part3/vith14/*.npy'))\n",
    "clip_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-mscoco/vith14/*.npy'))\n",
    "clip_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-dataset80k/vith14/prompt_embedding_vith14_001.npy'))\n",
    "clip_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-dataset30k/vith14/prompt_embedding_vith14_001.npy'))\n",
    "clip_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-dataset900k/vith14/*.npy'))\n",
    "clip_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-conceptual-captions/vith14/*.npy'))\n",
    "clip_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-sd2gpt2/vith14/*.npy'))\n",
    "clip_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-sd2hardcode/vith14/*.npy'))\n",
    "clip_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-chatgpt/vith14/*.npy'))\n",
    "clip_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-laion2b-part0000-0004/vith14/*.npy'))\n",
    "clip_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-laion2b-part0005-0009/vith14/*.npy'))\n",
    "clip_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-laion2b-part0010-0014/vith14/*.npy'))\n",
    "clip_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-laion2b-part0015-0019/vith14/*.npy'))\n",
    "clip_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-laion2b-part0020-0024/vith14/*.npy'))\n",
    "clip_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-laion2b-part0025-0029/vith14/*.npy'))\n",
    "clip_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-laion2b-part0030-0034/vith14/*.npy'))\n",
    "clip_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-laion2b-part0035-0039/vith14/*.npy'))\n",
    "clip_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-laion2b-part0040-0044/vith14/*.npy'))\n",
    "clip_embedding_paths += sorted(glob('/kaggle/input/pub-embeddings-laion2b-part0045-0049/vith14/*.npy'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "71272861",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:26:40.876259Z",
     "iopub.status.busy": "2023-05-15T14:26:40.875275Z",
     "iopub.status.idle": "2023-05-15T14:26:40.881221Z",
     "shell.execute_reply": "2023-05-15T14:26:40.880260Z"
    },
    "papermill": {
     "duration": 0.032888,
     "end_time": "2023-05-15T14:26:40.883771",
     "exception": false,
     "start_time": "2023-05-15T14:26:40.850883",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "if clip_image_features.shape[0] == 7:\n",
    "    objective_embedding_paths = objective_embedding_paths[:5]\n",
    "    clip_embedding_paths = clip_embedding_paths[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "83fe585f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:26:40.930882Z",
     "iopub.status.busy": "2023-05-15T14:26:40.930100Z",
     "iopub.status.idle": "2023-05-15T14:26:40.935449Z",
     "shell.execute_reply": "2023-05-15T14:26:40.934473Z"
    },
    "papermill": {
     "duration": 0.031429,
     "end_time": "2023-05-15T14:26:40.937788",
     "exception": false,
     "start_time": "2023-05-15T14:26:40.906359",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "n_neighbors = 125\n",
    "distance_dim = 15\n",
    "coef = 0.001"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "6fede96e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:26:40.985556Z",
     "iopub.status.busy": "2023-05-15T14:26:40.984666Z",
     "iopub.status.idle": "2023-05-15T14:26:40.998880Z",
     "shell.execute_reply": "2023-05-15T14:26:40.997914Z"
    },
    "papermill": {
     "duration": 0.040662,
     "end_time": "2023-05-15T14:26:41.001147",
     "exception": false,
     "start_time": "2023-05-15T14:26:40.960485",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def predict_local_knn(ref_x_embeddings, test_x_embeddings, n_neighbors=100, interval=1000, distance_dim=6, coef=0.001):\n",
    "    \n",
    "    ref_x_embeddings = torch.from_numpy(ref_x_embeddings).to('cuda')\n",
    "    ref_x_embeddings /= ref_x_embeddings.norm(dim=-1, keepdim=True)\n",
    "    \n",
    "    n_iter = test_x_embeddings.shape[0] // interval\n",
    "    if (test_x_embeddings.shape[0] % interval) != 0:\n",
    "        n_iter += 1\n",
    "        \n",
    "    dist_topk_store = []\n",
    "    idxs_topk_store = []\n",
    "    weights_store = []\n",
    "    preds = []\n",
    "    delta = 0.0001\n",
    "    \n",
    "    for i in range(n_iter):\n",
    "        batch_test_embeddings = torch.from_numpy(\n",
    "            test_x_embeddings[i * interval:(i + 1) * interval, :].copy()\n",
    "        ).to('cuda')\n",
    "        batch_test_embeddings /= batch_test_embeddings.norm(dim=-1, keepdim=True)\n",
    "        \n",
    "        dists = 1 - torch.mm(batch_test_embeddings, ref_x_embeddings.T)\n",
    "        del batch_test_embeddings\n",
    "        gc.collect()\n",
    "        \n",
    "        dist_topk, idxs_topk = torch.topk(dists, n_neighbors, largest=False, dim=-1)\n",
    "        dist_topk = dist_topk.to(torch.float64)        \n",
    "        \n",
    "        weights = 1 / (dist_topk ** distance_dim + delta) * coef\n",
    "        weights[dist_topk < 0] = delta\n",
    "        \n",
    "        dist_topk_store.append(dist_topk.to('cpu').detach().numpy().copy())\n",
    "        idxs_topk_store.append(idxs_topk.to('cpu').detach().numpy().copy())\n",
    "        weights_store.append(weights.to('cpu').detach().numpy().copy())\n",
    "                \n",
    "        del dists, weights, dist_topk, idxs_topk\n",
    "        torch.cuda.empty_cache()\n",
    "        gc.collect()\n",
    "        \n",
    "    del ref_x_embeddings\n",
    "    torch.cuda.empty_cache()\n",
    "    gc.collect()\n",
    "    return np.concatenate(dist_topk_store), np.concatenate(idxs_topk_store), np.concatenate(weights_store)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "f9306b9c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:26:41.049119Z",
     "iopub.status.busy": "2023-05-15T14:26:41.048168Z",
     "iopub.status.idle": "2023-05-15T14:27:06.958909Z",
     "shell.execute_reply": "2023-05-15T14:27:06.957613Z"
    },
    "papermill": {
     "duration": 25.937384,
     "end_time": "2023-05-15T14:27:06.961350",
     "exception": false,
     "start_time": "2023-05-15T14:26:41.023966",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5/5 [00:25<00:00,  5.18s/it]\n"
     ]
    }
   ],
   "source": [
    "for clip_embedding_idx, clip_embedding_path in enumerate(tqdm(clip_embedding_paths)):\n",
    "\n",
    "    clip_embeddings = np.load(clip_embedding_path).astype(np.float16)\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        local_dists, local_emb_indecies, local_weights = predict_local_knn(\n",
    "            clip_embeddings,\n",
    "            clip_image_features.numpy(),\n",
    "            n_neighbors=n_neighbors,\n",
    "            distance_dim=distance_dim,\n",
    "            coef=coef\n",
    "        )\n",
    "        \n",
    "    local_files = np.zeros(local_dists.shape, dtype=np.int32) + clip_embedding_idx\n",
    "    \n",
    "    if clip_embedding_idx == 0:\n",
    "        global_files = local_files\n",
    "        global_dists = local_dists\n",
    "        global_emb_indecies = local_emb_indecies\n",
    "        global_weights = local_weights\n",
    "    else:\n",
    "        global_files = np.concatenate([global_files, local_files], axis=-1)\n",
    "        global_dists = np.concatenate([global_dists, local_dists], axis=-1)\n",
    "        global_emb_indecies = np.concatenate([global_emb_indecies, local_emb_indecies], axis=-1)\n",
    "        global_weights = np.concatenate([global_weights, local_weights], axis=-1)\n",
    "\n",
    "        unsorted_min_indices = np.argpartition(global_dists, n_neighbors, axis=1)[:, :n_neighbors]\n",
    "\n",
    "        global_files = np.vstack([global_files[i, unsorted_min_indices[i,:]] for i in range(unsorted_min_indices.shape[0])])\n",
    "        global_dists = np.vstack([global_dists[i, unsorted_min_indices[i,:]] for i in range(unsorted_min_indices.shape[0])])\n",
    "        global_emb_indecies = np.vstack([global_emb_indecies[i, unsorted_min_indices[i,:]] for i in range(unsorted_min_indices.shape[0])])\n",
    "        global_weights = np.vstack([global_weights[i, unsorted_min_indices[i,:]] for i in range(unsorted_min_indices.shape[0])])\n",
    "    \n",
    "    gc.collect()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "5a844350",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:27:07.010384Z",
     "iopub.status.busy": "2023-05-15T14:27:07.009308Z",
     "iopub.status.idle": "2023-05-15T14:27:07.042039Z",
     "shell.execute_reply": "2023-05-15T14:27:07.040844Z"
    },
    "papermill": {
     "duration": 0.059847,
     "end_time": "2023-05-15T14:27:07.044856",
     "exception": false,
     "start_time": "2023-05-15T14:27:06.985009",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>file</th>\n",
       "      <th>dist</th>\n",
       "      <th>emb_index</th>\n",
       "      <th>test_index</th>\n",
       "      <th>weight</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>/kaggle/input/pub-embeddings-diffusiondb-14m-p...</td>\n",
       "      <td>0.746094</td>\n",
       "      <td>55979</td>\n",
       "      <td>0</td>\n",
       "      <td>0.080278</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>/kaggle/input/pub-embeddings-diffusiondb-14m-p...</td>\n",
       "      <td>0.739258</td>\n",
       "      <td>44239</td>\n",
       "      <td>0</td>\n",
       "      <td>0.092054</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>/kaggle/input/pub-embeddings-diffusiondb-14m-p...</td>\n",
       "      <td>0.683594</td>\n",
       "      <td>24001</td>\n",
       "      <td>0</td>\n",
       "      <td>0.291854</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>/kaggle/input/pub-embeddings-diffusiondb-14m-p...</td>\n",
       "      <td>0.683594</td>\n",
       "      <td>24004</td>\n",
       "      <td>0</td>\n",
       "      <td>0.291854</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>/kaggle/input/pub-embeddings-diffusiondb-14m-p...</td>\n",
       "      <td>0.683594</td>\n",
       "      <td>24005</td>\n",
       "      <td>0</td>\n",
       "      <td>0.291854</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>870</th>\n",
       "      <td>/kaggle/input/pub-embeddings-diffusiondb-14m-p...</td>\n",
       "      <td>0.678711</td>\n",
       "      <td>68350</td>\n",
       "      <td>6</td>\n",
       "      <td>0.323912</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>871</th>\n",
       "      <td>/kaggle/input/pub-embeddings-diffusiondb-14m-p...</td>\n",
       "      <td>0.678711</td>\n",
       "      <td>68349</td>\n",
       "      <td>6</td>\n",
       "      <td>0.323912</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>872</th>\n",
       "      <td>/kaggle/input/pub-embeddings-diffusiondb-14m-p...</td>\n",
       "      <td>0.678711</td>\n",
       "      <td>68354</td>\n",
       "      <td>6</td>\n",
       "      <td>0.323912</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>873</th>\n",
       "      <td>/kaggle/input/pub-embeddings-diffusiondb-14m-p...</td>\n",
       "      <td>0.678711</td>\n",
       "      <td>22053</td>\n",
       "      <td>6</td>\n",
       "      <td>0.323912</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>874</th>\n",
       "      <td>/kaggle/input/pub-embeddings-diffusiondb-14m-p...</td>\n",
       "      <td>0.678711</td>\n",
       "      <td>22056</td>\n",
       "      <td>6</td>\n",
       "      <td>0.323912</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>875 rows × 5 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                  file      dist  emb_index  \\\n",
       "0    /kaggle/input/pub-embeddings-diffusiondb-14m-p...  0.746094      55979   \n",
       "1    /kaggle/input/pub-embeddings-diffusiondb-14m-p...  0.739258      44239   \n",
       "2    /kaggle/input/pub-embeddings-diffusiondb-14m-p...  0.683594      24001   \n",
       "3    /kaggle/input/pub-embeddings-diffusiondb-14m-p...  0.683594      24004   \n",
       "4    /kaggle/input/pub-embeddings-diffusiondb-14m-p...  0.683594      24005   \n",
       "..                                                 ...       ...        ...   \n",
       "870  /kaggle/input/pub-embeddings-diffusiondb-14m-p...  0.678711      68350   \n",
       "871  /kaggle/input/pub-embeddings-diffusiondb-14m-p...  0.678711      68349   \n",
       "872  /kaggle/input/pub-embeddings-diffusiondb-14m-p...  0.678711      68354   \n",
       "873  /kaggle/input/pub-embeddings-diffusiondb-14m-p...  0.678711      22053   \n",
       "874  /kaggle/input/pub-embeddings-diffusiondb-14m-p...  0.678711      22056   \n",
       "\n",
       "     test_index    weight  \n",
       "0             0  0.080278  \n",
       "1             0  0.092054  \n",
       "2             0  0.291854  \n",
       "3             0  0.291854  \n",
       "4             0  0.291854  \n",
       "..          ...       ...  \n",
       "870           6  0.323912  \n",
       "871           6  0.323912  \n",
       "872           6  0.323912  \n",
       "873           6  0.323912  \n",
       "874           6  0.323912  \n",
       "\n",
       "[875 rows x 5 columns]"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_knn = pd.DataFrame()\n",
    "df_knn['file'] = global_files.flatten()\n",
    "df_knn['file'] = df_knn['file'].apply(lambda x: objective_embedding_paths[x])\n",
    "df_knn['dist'] = global_dists.flatten()\n",
    "df_knn['emb_index'] = global_emb_indecies.flatten()\n",
    "df_knn['test_index'] = np.array([[val] * n_neighbors for val in range(clip_image_features.shape[0])]).flatten()\n",
    "df_knn['weight'] = global_weights.flatten()\n",
    "df_knn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "6e095b36",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:27:07.095047Z",
     "iopub.status.busy": "2023-05-15T14:27:07.094019Z",
     "iopub.status.idle": "2023-05-15T14:27:10.248333Z",
     "shell.execute_reply": "2023-05-15T14:27:10.247319Z"
    },
    "papermill": {
     "duration": 3.18198,
     "end_time": "2023-05-15T14:27:10.251363",
     "exception": false,
     "start_time": "2023-05-15T14:27:07.069383",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5/5 [00:03<00:00,  1.60it/s]\n"
     ]
    }
   ],
   "source": [
    "test_prompt_embeddings = np.zeros((clip_image_features.shape[0], 384))\n",
    "\n",
    "for (objective_emb_file, gdf) in tqdm(df_knn.groupby('file')):\n",
    "    ref_objective_embeddings = np.load(objective_emb_file).astype(np.float16) \n",
    "    for _, r in gdf.iterrows():\n",
    "        test_prompt_embeddings[int(r.test_index), :] += r.weight * ref_objective_embeddings[int(r.emb_index), :]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "9fd4102d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:27:10.325520Z",
     "iopub.status.busy": "2023-05-15T14:27:10.324898Z",
     "iopub.status.idle": "2023-05-15T14:27:10.339466Z",
     "shell.execute_reply": "2023-05-15T14:27:10.338366Z"
    },
    "papermill": {
     "duration": 0.061069,
     "end_time": "2023-05-15T14:27:10.349512",
     "exception": false,
     "start_time": "2023-05-15T14:27:10.288443",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "BS = 1000\n",
    "num = test_prompt_embeddings.shape[0] // BS\n",
    "\n",
    "if test_prompt_embeddings.shape[0] % BS != 0:\n",
    "    num += 1\n",
    "    \n",
    "for i in range(num):\n",
    "    embeddings = test_prompt_embeddings[i * BS:(i + 1) * BS, :]\n",
    "    embeddings = embeddings / (np.abs(embeddings).max(axis=-1, keepdims=True) + 0.0000001)\n",
    "    embeddings = normalize(embeddings)\n",
    "    test_prompt_embeddings[i * BS:(i + 1) * BS, :] = embeddings\n",
    "\n",
    "df_test_embeddings['knn_prediction'] = test_prompt_embeddings.flatten()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3820387c",
   "metadata": {
    "papermill": {
     "duration": 0.035838,
     "end_time": "2023-05-15T14:27:10.420669",
     "exception": false,
     "start_time": "2023-05-15T14:27:10.384831",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 10. Representation Learning Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "afbceaf5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:27:10.494217Z",
     "iopub.status.busy": "2023-05-15T14:27:10.493794Z",
     "iopub.status.idle": "2023-05-15T14:27:10.513106Z",
     "shell.execute_reply": "2023-05-15T14:27:10.512094Z"
    },
    "papermill": {
     "duration": 0.058878,
     "end_time": "2023-05-15T14:27:10.516010",
     "exception": false,
     "start_time": "2023-05-15T14:27:10.457132",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def create_transformations(**transformation_parameters):\n",
    "\n",
    "    \"\"\"\n",
    "    Create training, validation and test sets transformations\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    transformation_parameters: dict\n",
    "        Dictionary of transformation parameters\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    transformations: dict\n",
    "        Transformations for training, validation and test sets\n",
    "    \"\"\"\n",
    "\n",
    "    train_transformations = A.Compose([\n",
    "        A.Resize(\n",
    "            height=transformation_parameters['resize_height'],\n",
    "            width=transformation_parameters['resize_width'],\n",
    "            interpolation=cv2.INTER_NEAREST,\n",
    "            always_apply=True\n",
    "        ),\n",
    "        A.HorizontalFlip(p=transformation_parameters['horizontal_flip_probability']),\n",
    "        A.Normalize(\n",
    "            mean=transformation_parameters['normalize_mean'],\n",
    "            std=transformation_parameters['normalize_std'],\n",
    "            max_pixel_value=transformation_parameters['normalize_max_pixel_value'],\n",
    "            always_apply=True\n",
    "        ),\n",
    "        ToTensorV2(always_apply=True)\n",
    "    ])\n",
    "\n",
    "    val_transformations = A.Compose([\n",
    "        A.Resize(\n",
    "            height=transformation_parameters['resize_height'],\n",
    "            width=transformation_parameters['resize_width'],\n",
    "            interpolation=cv2.INTER_NEAREST,\n",
    "            always_apply=True\n",
    "        ),\n",
    "        A.Normalize(\n",
    "            mean=transformation_parameters['normalize_mean'],\n",
    "            std=transformation_parameters['normalize_std'],\n",
    "            max_pixel_value=transformation_parameters['normalize_max_pixel_value'],\n",
    "            always_apply=True\n",
    "        ),\n",
    "        ToTensorV2(always_apply=True)\n",
    "    ])\n",
    "\n",
    "    test_transformations = A.Compose([\n",
    "        A.Resize(\n",
    "            height=transformation_parameters['resize_height'],\n",
    "            width=transformation_parameters['resize_width'],\n",
    "            interpolation=cv2.INTER_NEAREST,\n",
    "            always_apply=True\n",
    "        ),\n",
    "        A.Normalize(\n",
    "            mean=transformation_parameters['normalize_mean'],\n",
    "            std=transformation_parameters['normalize_std'],\n",
    "            max_pixel_value=transformation_parameters['normalize_max_pixel_value'],\n",
    "            always_apply=True\n",
    "        ),\n",
    "        ToTensorV2(always_apply=True)\n",
    "    ])\n",
    "\n",
    "    transformations = {'train': train_transformations, 'val': val_transformations, 'test': test_transformations}\n",
    "    return transformations\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "3466ab97",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:27:10.566403Z",
     "iopub.status.busy": "2023-05-15T14:27:10.566068Z",
     "iopub.status.idle": "2023-05-15T14:27:10.575786Z",
     "shell.execute_reply": "2023-05-15T14:27:10.574489Z"
    },
    "papermill": {
     "duration": 0.037948,
     "end_time": "2023-05-15T14:27:10.578510",
     "exception": false,
     "start_time": "2023-05-15T14:27:10.540562",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class TimmTransformerModel(nn.Module):\n",
    "\n",
    "    def __init__(self, model_name, pretrained, backbone_args, freeze_parameters, head_args):\n",
    "\n",
    "        super(TimmTransformerModel, self).__init__()\n",
    "\n",
    "        self.backbone = timm.create_model(\n",
    "            model_name=model_name,\n",
    "            pretrained=pretrained,\n",
    "            **backbone_args\n",
    "        )\n",
    "\n",
    "        if freeze_parameters:\n",
    "            for parameter in self.backbone.parameters():\n",
    "                parameter.requires_grad = False\n",
    "\n",
    "        input_features = self.backbone.get_classifier().in_features\n",
    "        self.backbone.head = nn.Identity()\n",
    "        self.head = nn.Linear(in_features=input_features, out_features=head_args['output_dimensions'], bias=True)\n",
    "\n",
    "    def forward(self, x):\n",
    "\n",
    "        x = self.backbone(x)\n",
    "        output = self.head(x)\n",
    "\n",
    "        return output\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "1ff47854",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:27:10.629554Z",
     "iopub.status.busy": "2023-05-15T14:27:10.629182Z",
     "iopub.status.idle": "2023-05-15T14:27:17.434925Z",
     "shell.execute_reply": "2023-05-15T14:27:17.432638Z"
    },
    "papermill": {
     "duration": 6.835001,
     "end_time": "2023-05-15T14:27:17.437964",
     "exception": false,
     "start_time": "2023-05-15T14:27:10.602963",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.7/site-packages/torch/functional.py:504: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at /usr/local/src/pytorch/aten/src/ATen/native/TensorShape.cpp:3190.)\n",
      "  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]\n",
      "100%|██████████| 1/1 [00:00<00:00,  1.05it/s]\n"
     ]
    }
   ],
   "source": [
    "swin_base_patch4_window12_384_config, swin_base_patch4_window12_384_model = load_model(\n",
    "    model_directory_path=str(external_dataset / 'swin_base_patch4_window12_384')\n",
    ")\n",
    "\n",
    "swin_base_patch4_window12_384_transformations = create_transformations(**swin_base_patch4_window12_384_config['transformations'])\n",
    "\n",
    "swin_base_patch4_window12_384_dataset = ImageEmbeddingDataset(\n",
    "    image_paths=image_paths,\n",
    "    embeddings=None,\n",
    "    transforms=swin_base_patch4_window12_384_transformations['test'],\n",
    ")\n",
    "\n",
    "swin_base_patch4_window12_384_data_loader = DataLoader(\n",
    "    swin_base_patch4_window12_384_dataset,\n",
    "    batch_size=16,\n",
    "    sampler=SequentialSampler(swin_base_patch4_window12_384_dataset),\n",
    "    pin_memory=False,\n",
    "    drop_last=False,\n",
    "    num_workers=2\n",
    ")\n",
    "\n",
    "swin_base_patch4_window12_384_predictions = []\n",
    "\n",
    "for inputs in tqdm(swin_base_patch4_window12_384_data_loader):\n",
    "    inputs = inputs.to('cuda')\n",
    "\n",
    "    with torch.no_grad():\n",
    "        outputs = swin_base_patch4_window12_384_model(inputs)\n",
    "\n",
    "    outputs = outputs.detach().cpu()\n",
    "\n",
    "    horizontal_flipped_inputs = inputs.flip(-1)\n",
    "    with torch.no_grad():\n",
    "        horizontal_flipped_outputs = swin_base_patch4_window12_384_model(horizontal_flipped_inputs)\n",
    "    horizontal_flipped_outputs = horizontal_flipped_outputs.detach().cpu()\n",
    "\n",
    "    outputs = torch.cat([\n",
    "        outputs.unsqueeze(dim=0),\n",
    "        horizontal_flipped_outputs.unsqueeze(dim=0)\n",
    "    ], dim=0).mean(dim=0)\n",
    "\n",
    "    swin_base_patch4_window12_384_predictions.append(outputs)\n",
    "\n",
    "swin_base_patch4_window12_384_predictions = torch.cat(swin_base_patch4_window12_384_predictions, dim=0).numpy()\n",
    "swin_base_patch4_window12_384_predictions = swin_base_patch4_window12_384_predictions / (np.abs(swin_base_patch4_window12_384_predictions).max(axis=-1, keepdims=True) + 0.0000001)\n",
    "swin_base_patch4_window12_384_predictions = normalize(swin_base_patch4_window12_384_predictions)\n",
    "df_test_embeddings['swin_base_patch4_window12_384_prediction'] = swin_base_patch4_window12_384_predictions.flatten()\n",
    "\n",
    "del swin_base_patch4_window12_384_config, swin_base_patch4_window12_384_model\n",
    "del swin_base_patch4_window12_384_transformations, swin_base_patch4_window12_384_dataset, swin_base_patch4_window12_384_data_loader\n",
    "del swin_base_patch4_window12_384_predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "c3ca5a39",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:27:17.490153Z",
     "iopub.status.busy": "2023-05-15T14:27:17.489793Z",
     "iopub.status.idle": "2023-05-15T14:27:33.592687Z",
     "shell.execute_reply": "2023-05-15T14:27:33.590330Z"
    },
    "papermill": {
     "duration": 16.131915,
     "end_time": "2023-05-15T14:27:33.595231",
     "exception": false,
     "start_time": "2023-05-15T14:27:17.463316",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:00<00:00,  1.56it/s]\n"
     ]
    }
   ],
   "source": [
    "swin_large_patch4_window7_224_config, swin_large_patch4_window7_224_model = load_model(\n",
    "    model_directory_path=str(external_dataset / 'swin_large_patch4_window7_224')\n",
    ")\n",
    "\n",
    "swin_large_patch4_window7_224_transformations = create_transformations(**swin_large_patch4_window7_224_config['transformations'])\n",
    "\n",
    "swin_large_patch4_window7_224_dataset = ImageEmbeddingDataset(\n",
    "    image_paths=image_paths,\n",
    "    embeddings=None,\n",
    "    transforms=swin_large_patch4_window7_224_transformations['test'],\n",
    ")\n",
    "\n",
    "swin_large_patch4_window7_224_data_loader = DataLoader(\n",
    "    swin_large_patch4_window7_224_dataset,\n",
    "    batch_size=16,\n",
    "    sampler=SequentialSampler(swin_large_patch4_window7_224_dataset),\n",
    "    pin_memory=False,\n",
    "    drop_last=False,\n",
    "    num_workers=2\n",
    ")\n",
    "\n",
    "swin_large_patch4_window7_224_predictions = []\n",
    "\n",
    "for inputs in tqdm(swin_large_patch4_window7_224_data_loader):\n",
    "    inputs = inputs.to('cuda')\n",
    "\n",
    "    with torch.no_grad():\n",
    "        outputs = swin_large_patch4_window7_224_model(inputs)\n",
    "\n",
    "    outputs = outputs.detach().cpu()\n",
    "\n",
    "    horizontal_flipped_inputs = inputs.flip(-1)\n",
    "    with torch.no_grad():\n",
    "        horizontal_flipped_outputs = swin_large_patch4_window7_224_model(horizontal_flipped_inputs)\n",
    "    horizontal_flipped_outputs = horizontal_flipped_outputs.detach().cpu()\n",
    "\n",
    "    outputs = torch.cat([\n",
    "        outputs.unsqueeze(dim=0),\n",
    "        horizontal_flipped_outputs.unsqueeze(dim=0)\n",
    "    ], dim=0).mean(dim=0)\n",
    "\n",
    "    swin_large_patch4_window7_224_predictions.append(outputs)\n",
    "\n",
    "swin_large_patch4_window7_224_predictions = torch.cat(swin_large_patch4_window7_224_predictions, dim=0).numpy()\n",
    "swin_large_patch4_window7_224_predictions = swin_large_patch4_window7_224_predictions / (np.abs(swin_large_patch4_window7_224_predictions).max(axis=-1, keepdims=True) + 0.0000001)\n",
    "swin_large_patch4_window7_224_predictions = normalize(swin_large_patch4_window7_224_predictions)\n",
    "df_test_embeddings['swin_large_patch4_window7_224_prediction'] = swin_large_patch4_window7_224_predictions.flatten()\n",
    "\n",
    "del swin_large_patch4_window7_224_config, swin_large_patch4_window7_224_model\n",
    "del swin_large_patch4_window7_224_transformations, swin_large_patch4_window7_224_dataset, swin_large_patch4_window7_224_data_loader\n",
    "del swin_large_patch4_window7_224_predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "fc724201",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:27:33.649144Z",
     "iopub.status.busy": "2023-05-15T14:27:33.647996Z",
     "iopub.status.idle": "2023-05-15T14:27:51.866475Z",
     "shell.execute_reply": "2023-05-15T14:27:51.864943Z"
    },
    "papermill": {
     "duration": 18.247982,
     "end_time": "2023-05-15T14:27:51.869303",
     "exception": false,
     "start_time": "2023-05-15T14:27:33.621321",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:01<00:00,  1.04s/it]\n"
     ]
    }
   ],
   "source": [
    "swin_large_patch4_window12_384_config, swin_large_patch4_window12_384_model = load_model(\n",
    "    model_directory_path=str(external_dataset / 'swin_large_patch4_window12_384')\n",
    ")\n",
    "\n",
    "swin_large_patch4_window12_384_transformations = create_transformations(**swin_large_patch4_window12_384_config['transformations'])\n",
    "\n",
    "swin_large_patch4_window12_384_dataset = ImageEmbeddingDataset(\n",
    "    image_paths=image_paths,\n",
    "    embeddings=None,\n",
    "    transforms=swin_large_patch4_window12_384_transformations['test'],\n",
    ")\n",
    "\n",
    "swin_large_patch4_window12_384_data_loader = DataLoader(\n",
    "    swin_large_patch4_window12_384_dataset,\n",
    "    batch_size=16,\n",
    "    sampler=SequentialSampler(swin_large_patch4_window12_384_dataset),\n",
    "    pin_memory=False,\n",
    "    drop_last=False,\n",
    "    num_workers=2\n",
    ")\n",
    "\n",
    "swin_large_patch4_window12_384_predictions = []\n",
    "\n",
    "for inputs in tqdm(swin_large_patch4_window12_384_data_loader):\n",
    "    inputs = inputs.to('cuda')\n",
    "\n",
    "    with torch.no_grad():\n",
    "        outputs = swin_large_patch4_window12_384_model(inputs)\n",
    "\n",
    "    outputs = outputs.detach().cpu()\n",
    "\n",
    "    horizontal_flipped_inputs = inputs.flip(-1)\n",
    "    with torch.no_grad():\n",
    "        horizontal_flipped_outputs = swin_large_patch4_window12_384_model(horizontal_flipped_inputs)\n",
    "    horizontal_flipped_outputs = horizontal_flipped_outputs.detach().cpu()\n",
    "\n",
    "    outputs = torch.cat([\n",
    "        outputs.unsqueeze(dim=0),\n",
    "        horizontal_flipped_outputs.unsqueeze(dim=0)\n",
    "    ], dim=0).mean(dim=0)\n",
    "\n",
    "    swin_large_patch4_window12_384_predictions.append(outputs)\n",
    "\n",
    "swin_large_patch4_window12_384_predictions = torch.cat(swin_large_patch4_window12_384_predictions, dim=0).numpy()\n",
    "swin_large_patch4_window12_384_predictions = swin_large_patch4_window12_384_predictions / (np.abs(swin_large_patch4_window12_384_predictions).max(axis=-1, keepdims=True) + 0.0000001)\n",
    "swin_large_patch4_window12_384_predictions = normalize(swin_large_patch4_window12_384_predictions)\n",
    "df_test_embeddings['swin_large_patch4_window12_384_prediction'] = swin_large_patch4_window12_384_predictions.flatten()\n",
    "\n",
    "del swin_large_patch4_window12_384_config, swin_large_patch4_window12_384_model\n",
    "del swin_large_patch4_window12_384_transformations, swin_large_patch4_window12_384_dataset, swin_large_patch4_window12_384_data_loader\n",
    "del swin_large_patch4_window12_384_predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "05d73010",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:27:51.922654Z",
     "iopub.status.busy": "2023-05-15T14:27:51.922287Z",
     "iopub.status.idle": "2023-05-15T14:28:16.548864Z",
     "shell.execute_reply": "2023-05-15T14:28:16.547405Z"
    },
    "papermill": {
     "duration": 24.656936,
     "end_time": "2023-05-15T14:28:16.552177",
     "exception": false,
     "start_time": "2023-05-15T14:27:51.895241",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:00<00:00,  1.28it/s]\n"
     ]
    }
   ],
   "source": [
    "beitv2_large_patch16_224_config, beitv2_large_patch16_224_model = load_model(\n",
    "    model_directory_path=str(external_dataset / 'beitv2_large_patch16_224')\n",
    ")\n",
    "\n",
    "beitv2_large_patch16_224_transformations = create_transformations(**beitv2_large_patch16_224_config['transformations'])\n",
    "\n",
    "beitv2_large_patch16_224_dataset = ImageEmbeddingDataset(\n",
    "    image_paths=image_paths,\n",
    "    embeddings=None,\n",
    "    transforms=beitv2_large_patch16_224_transformations['test'],\n",
    ")\n",
    "\n",
    "beitv2_large_patch16_224_data_loader = DataLoader(\n",
    "    beitv2_large_patch16_224_dataset,\n",
    "    batch_size=16,\n",
    "    sampler=SequentialSampler(beitv2_large_patch16_224_dataset),\n",
    "    pin_memory=False,\n",
    "    drop_last=False,\n",
    "    num_workers=2\n",
    ")\n",
    "\n",
    "beitv2_large_patch16_224_predictions = []\n",
    "\n",
    "for inputs in tqdm(beitv2_large_patch16_224_data_loader):\n",
    "    inputs = inputs.to('cuda')\n",
    "\n",
    "    with torch.no_grad():\n",
    "        outputs = beitv2_large_patch16_224_model(inputs)\n",
    "\n",
    "    outputs = outputs.detach().cpu()\n",
    "\n",
    "    horizontal_flipped_inputs = inputs.flip(-1)\n",
    "    with torch.no_grad():\n",
    "        horizontal_flipped_outputs = beitv2_large_patch16_224_model(horizontal_flipped_inputs)\n",
    "    horizontal_flipped_outputs = horizontal_flipped_outputs.detach().cpu()\n",
    "\n",
    "    outputs = torch.cat([\n",
    "        outputs.unsqueeze(dim=0),\n",
    "        horizontal_flipped_outputs.unsqueeze(dim=0)\n",
    "    ], dim=0).mean(dim=0)\n",
    "\n",
    "    beitv2_large_patch16_224_predictions.append(outputs)\n",
    "\n",
    "beitv2_large_patch16_224_predictions = torch.cat(beitv2_large_patch16_224_predictions, dim=0).numpy()\n",
    "beitv2_large_patch16_224_predictions = beitv2_large_patch16_224_predictions / (np.abs(beitv2_large_patch16_224_predictions).max(axis=-1, keepdims=True) + 0.0000001)\n",
    "beitv2_large_patch16_224_predictions = normalize(beitv2_large_patch16_224_predictions)\n",
    "df_test_embeddings['beitv2_large_patch16_224_prediction'] = beitv2_large_patch16_224_predictions.flatten()\n",
    "\n",
    "del beitv2_large_patch16_224_config, beitv2_large_patch16_224_model\n",
    "del beitv2_large_patch16_224_transformations, beitv2_large_patch16_224_dataset, beitv2_large_patch16_224_data_loader\n",
    "del beitv2_large_patch16_224_predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9112812",
   "metadata": {
    "papermill": {
     "duration": 0.059871,
     "end_time": "2023-05-15T14:28:16.639023",
     "exception": false,
     "start_time": "2023-05-15T14:28:16.579152",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 11. Ensemble"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "048e8cc3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:28:16.695038Z",
     "iopub.status.busy": "2023-05-15T14:28:16.694641Z",
     "iopub.status.idle": "2023-05-15T14:28:16.720763Z",
     "shell.execute_reply": "2023-05-15T14:28:16.719267Z"
    },
    "papermill": {
     "duration": 0.055765,
     "end_time": "2023-05-15T14:28:16.723337",
     "exception": false,
     "start_time": "2023-05-15T14:28:16.667572",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>imgId_eId</th>\n",
       "      <th>coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding</th>\n",
       "      <th>mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding</th>\n",
       "      <th>interrogated_mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding</th>\n",
       "      <th>blip_image_captioning_large_prediction_embedding</th>\n",
       "      <th>knn_prediction</th>\n",
       "      <th>swin_base_patch4_window12_384_prediction</th>\n",
       "      <th>swin_large_patch4_window7_224_prediction</th>\n",
       "      <th>swin_large_patch4_window12_384_prediction</th>\n",
       "      <th>beitv2_large_patch16_224_prediction</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>f27825b2c_0</td>\n",
       "      <td>-0.116310</td>\n",
       "      <td>0.029225</td>\n",
       "      <td>0.027709</td>\n",
       "      <td>-0.020460</td>\n",
       "      <td>-0.106356</td>\n",
       "      <td>-0.033799</td>\n",
       "      <td>-0.054163</td>\n",
       "      <td>-0.052756</td>\n",
       "      <td>-0.084640</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>f27825b2c_1</td>\n",
       "      <td>0.040138</td>\n",
       "      <td>0.053408</td>\n",
       "      <td>0.067174</td>\n",
       "      <td>0.083382</td>\n",
       "      <td>0.056334</td>\n",
       "      <td>0.043510</td>\n",
       "      <td>0.044165</td>\n",
       "      <td>0.033424</td>\n",
       "      <td>0.049386</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>f27825b2c_2</td>\n",
       "      <td>0.033303</td>\n",
       "      <td>-0.084613</td>\n",
       "      <td>-0.075604</td>\n",
       "      <td>-0.068206</td>\n",
       "      <td>0.009878</td>\n",
       "      <td>-0.007162</td>\n",
       "      <td>0.004634</td>\n",
       "      <td>0.001244</td>\n",
       "      <td>-0.019956</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>f27825b2c_3</td>\n",
       "      <td>-0.044670</td>\n",
       "      <td>0.017963</td>\n",
       "      <td>-0.041972</td>\n",
       "      <td>-0.043487</td>\n",
       "      <td>-0.060045</td>\n",
       "      <td>-0.008922</td>\n",
       "      <td>0.014103</td>\n",
       "      <td>-0.022813</td>\n",
       "      <td>-0.009255</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>f27825b2c_4</td>\n",
       "      <td>-0.060771</td>\n",
       "      <td>-0.075529</td>\n",
       "      <td>-0.067199</td>\n",
       "      <td>-0.005878</td>\n",
       "      <td>-0.067769</td>\n",
       "      <td>-0.024105</td>\n",
       "      <td>-0.047492</td>\n",
       "      <td>-0.039126</td>\n",
       "      <td>-0.068727</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2683</th>\n",
       "      <td>c98f79f71_379</td>\n",
       "      <td>-0.013048</td>\n",
       "      <td>-0.073867</td>\n",
       "      <td>-0.072286</td>\n",
       "      <td>-0.049304</td>\n",
       "      <td>-0.030339</td>\n",
       "      <td>0.025875</td>\n",
       "      <td>0.015397</td>\n",
       "      <td>0.016148</td>\n",
       "      <td>0.025705</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2684</th>\n",
       "      <td>c98f79f71_380</td>\n",
       "      <td>-0.016317</td>\n",
       "      <td>0.012851</td>\n",
       "      <td>0.044618</td>\n",
       "      <td>0.022202</td>\n",
       "      <td>0.125923</td>\n",
       "      <td>0.058460</td>\n",
       "      <td>0.036276</td>\n",
       "      <td>0.012421</td>\n",
       "      <td>0.069896</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2685</th>\n",
       "      <td>c98f79f71_381</td>\n",
       "      <td>0.019852</td>\n",
       "      <td>0.028804</td>\n",
       "      <td>0.056895</td>\n",
       "      <td>0.039194</td>\n",
       "      <td>-0.003771</td>\n",
       "      <td>0.036421</td>\n",
       "      <td>0.049093</td>\n",
       "      <td>0.058329</td>\n",
       "      <td>0.048571</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2686</th>\n",
       "      <td>c98f79f71_382</td>\n",
       "      <td>-0.074756</td>\n",
       "      <td>0.009999</td>\n",
       "      <td>-0.059861</td>\n",
       "      <td>-0.032646</td>\n",
       "      <td>-0.020211</td>\n",
       "      <td>-0.029689</td>\n",
       "      <td>-0.026619</td>\n",
       "      <td>-0.031885</td>\n",
       "      <td>-0.037808</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2687</th>\n",
       "      <td>c98f79f71_383</td>\n",
       "      <td>-0.028067</td>\n",
       "      <td>0.025200</td>\n",
       "      <td>0.075349</td>\n",
       "      <td>-0.010063</td>\n",
       "      <td>0.002076</td>\n",
       "      <td>0.017614</td>\n",
       "      <td>0.007647</td>\n",
       "      <td>-0.001963</td>\n",
       "      <td>0.018155</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>2688 rows × 10 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "          imgId_eId  coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding  \\\n",
       "0       f27825b2c_0                                          -0.116310      \n",
       "1       f27825b2c_1                                           0.040138      \n",
       "2       f27825b2c_2                                           0.033303      \n",
       "3       f27825b2c_3                                          -0.044670      \n",
       "4       f27825b2c_4                                          -0.060771      \n",
       "...             ...                                                ...      \n",
       "2683  c98f79f71_379                                          -0.013048      \n",
       "2684  c98f79f71_380                                          -0.016317      \n",
       "2685  c98f79f71_381                                           0.019852      \n",
       "2686  c98f79f71_382                                          -0.074756      \n",
       "2687  c98f79f71_383                                          -0.028067      \n",
       "\n",
       "      mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding  \\\n",
       "0                                              0.029225                       \n",
       "1                                              0.053408                       \n",
       "2                                             -0.084613                       \n",
       "3                                              0.017963                       \n",
       "4                                             -0.075529                       \n",
       "...                                                 ...                       \n",
       "2683                                          -0.073867                       \n",
       "2684                                           0.012851                       \n",
       "2685                                           0.028804                       \n",
       "2686                                           0.009999                       \n",
       "2687                                           0.025200                       \n",
       "\n",
       "      interrogated_mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding  \\\n",
       "0                                              0.027709                                    \n",
       "1                                              0.067174                                    \n",
       "2                                             -0.075604                                    \n",
       "3                                             -0.041972                                    \n",
       "4                                             -0.067199                                    \n",
       "...                                                 ...                                    \n",
       "2683                                          -0.072286                                    \n",
       "2684                                           0.044618                                    \n",
       "2685                                           0.056895                                    \n",
       "2686                                          -0.059861                                    \n",
       "2687                                           0.075349                                    \n",
       "\n",
       "      blip_image_captioning_large_prediction_embedding  knn_prediction  \\\n",
       "0                                            -0.020460       -0.106356   \n",
       "1                                             0.083382        0.056334   \n",
       "2                                            -0.068206        0.009878   \n",
       "3                                            -0.043487       -0.060045   \n",
       "4                                            -0.005878       -0.067769   \n",
       "...                                                ...             ...   \n",
       "2683                                         -0.049304       -0.030339   \n",
       "2684                                          0.022202        0.125923   \n",
       "2685                                          0.039194       -0.003771   \n",
       "2686                                         -0.032646       -0.020211   \n",
       "2687                                         -0.010063        0.002076   \n",
       "\n",
       "      swin_base_patch4_window12_384_prediction  \\\n",
       "0                                    -0.033799   \n",
       "1                                     0.043510   \n",
       "2                                    -0.007162   \n",
       "3                                    -0.008922   \n",
       "4                                    -0.024105   \n",
       "...                                        ...   \n",
       "2683                                  0.025875   \n",
       "2684                                  0.058460   \n",
       "2685                                  0.036421   \n",
       "2686                                 -0.029689   \n",
       "2687                                  0.017614   \n",
       "\n",
       "      swin_large_patch4_window7_224_prediction  \\\n",
       "0                                    -0.054163   \n",
       "1                                     0.044165   \n",
       "2                                     0.004634   \n",
       "3                                     0.014103   \n",
       "4                                    -0.047492   \n",
       "...                                        ...   \n",
       "2683                                  0.015397   \n",
       "2684                                  0.036276   \n",
       "2685                                  0.049093   \n",
       "2686                                 -0.026619   \n",
       "2687                                  0.007647   \n",
       "\n",
       "      swin_large_patch4_window12_384_prediction  \\\n",
       "0                                     -0.052756   \n",
       "1                                      0.033424   \n",
       "2                                      0.001244   \n",
       "3                                     -0.022813   \n",
       "4                                     -0.039126   \n",
       "...                                         ...   \n",
       "2683                                   0.016148   \n",
       "2684                                   0.012421   \n",
       "2685                                   0.058329   \n",
       "2686                                  -0.031885   \n",
       "2687                                  -0.001963   \n",
       "\n",
       "      beitv2_large_patch16_224_prediction  \n",
       "0                               -0.084640  \n",
       "1                                0.049386  \n",
       "2                               -0.019956  \n",
       "3                               -0.009255  \n",
       "4                               -0.068727  \n",
       "...                                   ...  \n",
       "2683                             0.025705  \n",
       "2684                             0.069896  \n",
       "2685                             0.048571  \n",
       "2686                            -0.037808  \n",
       "2687                             0.018155  \n",
       "\n",
       "[2688 rows x 10 columns]"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_test_embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "af154501",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:28:16.778040Z",
     "iopub.status.busy": "2023-05-15T14:28:16.777475Z",
     "iopub.status.idle": "2023-05-15T14:28:16.810690Z",
     "shell.execute_reply": "2023-05-15T14:28:16.809387Z"
    },
    "papermill": {
     "duration": 0.062897,
     "end_time": "2023-05-15T14:28:16.813232",
     "exception": false,
     "start_time": "2023-05-15T14:28:16.750335",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>imgId_eId</th>\n",
       "      <th>coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding</th>\n",
       "      <th>mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding</th>\n",
       "      <th>interrogated_mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding</th>\n",
       "      <th>blip_image_captioning_large_prediction_embedding</th>\n",
       "      <th>knn_prediction</th>\n",
       "      <th>swin_base_patch4_window12_384_prediction</th>\n",
       "      <th>swin_large_patch4_window7_224_prediction</th>\n",
       "      <th>swin_large_patch4_window12_384_prediction</th>\n",
       "      <th>beitv2_large_patch16_224_prediction</th>\n",
       "      <th>val</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>f27825b2c_0</td>\n",
       "      <td>-0.116310</td>\n",
       "      <td>0.029225</td>\n",
       "      <td>0.027709</td>\n",
       "      <td>-0.020460</td>\n",
       "      <td>-0.106356</td>\n",
       "      <td>-0.033799</td>\n",
       "      <td>-0.054163</td>\n",
       "      <td>-0.052756</td>\n",
       "      <td>-0.084640</td>\n",
       "      <td>-0.073501</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>f27825b2c_1</td>\n",
       "      <td>0.040138</td>\n",
       "      <td>0.053408</td>\n",
       "      <td>0.067174</td>\n",
       "      <td>0.083382</td>\n",
       "      <td>0.056334</td>\n",
       "      <td>0.043510</td>\n",
       "      <td>0.044165</td>\n",
       "      <td>0.033424</td>\n",
       "      <td>0.049386</td>\n",
       "      <td>0.051774</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>f27825b2c_2</td>\n",
       "      <td>0.033303</td>\n",
       "      <td>-0.084613</td>\n",
       "      <td>-0.075604</td>\n",
       "      <td>-0.068206</td>\n",
       "      <td>0.009878</td>\n",
       "      <td>-0.007162</td>\n",
       "      <td>0.004634</td>\n",
       "      <td>0.001244</td>\n",
       "      <td>-0.019956</td>\n",
       "      <td>-0.004998</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>f27825b2c_3</td>\n",
       "      <td>-0.044670</td>\n",
       "      <td>0.017963</td>\n",
       "      <td>-0.041972</td>\n",
       "      <td>-0.043487</td>\n",
       "      <td>-0.060045</td>\n",
       "      <td>-0.008922</td>\n",
       "      <td>0.014103</td>\n",
       "      <td>-0.022813</td>\n",
       "      <td>-0.009255</td>\n",
       "      <td>-0.035898</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>f27825b2c_4</td>\n",
       "      <td>-0.060771</td>\n",
       "      <td>-0.075529</td>\n",
       "      <td>-0.067199</td>\n",
       "      <td>-0.005878</td>\n",
       "      <td>-0.067769</td>\n",
       "      <td>-0.024105</td>\n",
       "      <td>-0.047492</td>\n",
       "      <td>-0.039126</td>\n",
       "      <td>-0.068727</td>\n",
       "      <td>-0.056508</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2683</th>\n",
       "      <td>c98f79f71_379</td>\n",
       "      <td>-0.013048</td>\n",
       "      <td>-0.073867</td>\n",
       "      <td>-0.072286</td>\n",
       "      <td>-0.049304</td>\n",
       "      <td>-0.030339</td>\n",
       "      <td>0.025875</td>\n",
       "      <td>0.015397</td>\n",
       "      <td>0.016148</td>\n",
       "      <td>0.025705</td>\n",
       "      <td>-0.015082</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2684</th>\n",
       "      <td>c98f79f71_380</td>\n",
       "      <td>-0.016317</td>\n",
       "      <td>0.012851</td>\n",
       "      <td>0.044618</td>\n",
       "      <td>0.022202</td>\n",
       "      <td>0.125923</td>\n",
       "      <td>0.058460</td>\n",
       "      <td>0.036276</td>\n",
       "      <td>0.012421</td>\n",
       "      <td>0.069896</td>\n",
       "      <td>0.077756</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2685</th>\n",
       "      <td>c98f79f71_381</td>\n",
       "      <td>0.019852</td>\n",
       "      <td>0.028804</td>\n",
       "      <td>0.056895</td>\n",
       "      <td>0.039194</td>\n",
       "      <td>-0.003771</td>\n",
       "      <td>0.036421</td>\n",
       "      <td>0.049093</td>\n",
       "      <td>0.058329</td>\n",
       "      <td>0.048571</td>\n",
       "      <td>0.022449</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2686</th>\n",
       "      <td>c98f79f71_382</td>\n",
       "      <td>-0.074756</td>\n",
       "      <td>0.009999</td>\n",
       "      <td>-0.059861</td>\n",
       "      <td>-0.032646</td>\n",
       "      <td>-0.020211</td>\n",
       "      <td>-0.029689</td>\n",
       "      <td>-0.026619</td>\n",
       "      <td>-0.031885</td>\n",
       "      <td>-0.037808</td>\n",
       "      <td>-0.027273</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2687</th>\n",
       "      <td>c98f79f71_383</td>\n",
       "      <td>-0.028067</td>\n",
       "      <td>0.025200</td>\n",
       "      <td>0.075349</td>\n",
       "      <td>-0.010063</td>\n",
       "      <td>0.002076</td>\n",
       "      <td>0.017614</td>\n",
       "      <td>0.007647</td>\n",
       "      <td>-0.001963</td>\n",
       "      <td>0.018155</td>\n",
       "      <td>0.007319</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>2688 rows × 11 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "          imgId_eId  coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding  \\\n",
       "0       f27825b2c_0                                          -0.116310      \n",
       "1       f27825b2c_1                                           0.040138      \n",
       "2       f27825b2c_2                                           0.033303      \n",
       "3       f27825b2c_3                                          -0.044670      \n",
       "4       f27825b2c_4                                          -0.060771      \n",
       "...             ...                                                ...      \n",
       "2683  c98f79f71_379                                          -0.013048      \n",
       "2684  c98f79f71_380                                          -0.016317      \n",
       "2685  c98f79f71_381                                           0.019852      \n",
       "2686  c98f79f71_382                                          -0.074756      \n",
       "2687  c98f79f71_383                                          -0.028067      \n",
       "\n",
       "      mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding  \\\n",
       "0                                              0.029225                       \n",
       "1                                              0.053408                       \n",
       "2                                             -0.084613                       \n",
       "3                                              0.017963                       \n",
       "4                                             -0.075529                       \n",
       "...                                                 ...                       \n",
       "2683                                          -0.073867                       \n",
       "2684                                           0.012851                       \n",
       "2685                                           0.028804                       \n",
       "2686                                           0.009999                       \n",
       "2687                                           0.025200                       \n",
       "\n",
       "      interrogated_mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding  \\\n",
       "0                                              0.027709                                    \n",
       "1                                              0.067174                                    \n",
       "2                                             -0.075604                                    \n",
       "3                                             -0.041972                                    \n",
       "4                                             -0.067199                                    \n",
       "...                                                 ...                                    \n",
       "2683                                          -0.072286                                    \n",
       "2684                                           0.044618                                    \n",
       "2685                                           0.056895                                    \n",
       "2686                                          -0.059861                                    \n",
       "2687                                           0.075349                                    \n",
       "\n",
       "      blip_image_captioning_large_prediction_embedding  knn_prediction  \\\n",
       "0                                            -0.020460       -0.106356   \n",
       "1                                             0.083382        0.056334   \n",
       "2                                            -0.068206        0.009878   \n",
       "3                                            -0.043487       -0.060045   \n",
       "4                                            -0.005878       -0.067769   \n",
       "...                                                ...             ...   \n",
       "2683                                         -0.049304       -0.030339   \n",
       "2684                                          0.022202        0.125923   \n",
       "2685                                          0.039194       -0.003771   \n",
       "2686                                         -0.032646       -0.020211   \n",
       "2687                                         -0.010063        0.002076   \n",
       "\n",
       "      swin_base_patch4_window12_384_prediction  \\\n",
       "0                                    -0.033799   \n",
       "1                                     0.043510   \n",
       "2                                    -0.007162   \n",
       "3                                    -0.008922   \n",
       "4                                    -0.024105   \n",
       "...                                        ...   \n",
       "2683                                  0.025875   \n",
       "2684                                  0.058460   \n",
       "2685                                  0.036421   \n",
       "2686                                 -0.029689   \n",
       "2687                                  0.017614   \n",
       "\n",
       "      swin_large_patch4_window7_224_prediction  \\\n",
       "0                                    -0.054163   \n",
       "1                                     0.044165   \n",
       "2                                     0.004634   \n",
       "3                                     0.014103   \n",
       "4                                    -0.047492   \n",
       "...                                        ...   \n",
       "2683                                  0.015397   \n",
       "2684                                  0.036276   \n",
       "2685                                  0.049093   \n",
       "2686                                 -0.026619   \n",
       "2687                                  0.007647   \n",
       "\n",
       "      swin_large_patch4_window12_384_prediction  \\\n",
       "0                                     -0.052756   \n",
       "1                                      0.033424   \n",
       "2                                      0.001244   \n",
       "3                                     -0.022813   \n",
       "4                                     -0.039126   \n",
       "...                                         ...   \n",
       "2683                                   0.016148   \n",
       "2684                                   0.012421   \n",
       "2685                                   0.058329   \n",
       "2686                                  -0.031885   \n",
       "2687                                  -0.001963   \n",
       "\n",
       "      beitv2_large_patch16_224_prediction       val  \n",
       "0                               -0.084640 -0.073501  \n",
       "1                                0.049386  0.051774  \n",
       "2                               -0.019956 -0.004998  \n",
       "3                               -0.009255 -0.035898  \n",
       "4                               -0.068727 -0.056508  \n",
       "...                                   ...       ...  \n",
       "2683                             0.025705 -0.015082  \n",
       "2684                             0.069896  0.077756  \n",
       "2685                             0.048571  0.022449  \n",
       "2686                            -0.037808 -0.027273  \n",
       "2687                             0.018155  0.007319  \n",
       "\n",
       "[2688 rows x 11 columns]"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_test_embeddings['val'] =\\\n",
    "(df_test_embeddings['blip_image_captioning_large_prediction_embedding'] * 0.04) +\\\n",
    "(df_test_embeddings['coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding'] * 0.03) +\\\n",
    "(df_test_embeddings['mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding'] * 0.035) +\\\n",
    "(df_test_embeddings['interrogated_mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding'] * 0.045) +\\\n",
    "(df_test_embeddings['knn_prediction'] * 0.475) +\\\n",
    "(df_test_embeddings['swin_base_patch4_window12_384_prediction'] * 0.08) +\\\n",
    "(df_test_embeddings['swin_large_patch4_window7_224_prediction'] * 0.09) +\\\n",
    "(df_test_embeddings['swin_large_patch4_window12_384_prediction'] * 0.125) +\\\n",
    "(df_test_embeddings['beitv2_large_patch16_224_prediction'] * 0.08)\n",
    "\n",
    "df_test_embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "720944f8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:28:16.868142Z",
     "iopub.status.busy": "2023-05-15T14:28:16.867828Z",
     "iopub.status.idle": "2023-05-15T14:28:16.892435Z",
     "shell.execute_reply": "2023-05-15T14:28:16.891274Z"
    },
    "papermill": {
     "duration": 0.055436,
     "end_time": "2023-05-15T14:28:16.895763",
     "exception": false,
     "start_time": "2023-05-15T14:28:16.840327",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding</th>\n",
       "      <th>mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding</th>\n",
       "      <th>interrogated_mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding</th>\n",
       "      <th>blip_image_captioning_large_prediction_embedding</th>\n",
       "      <th>knn_prediction</th>\n",
       "      <th>swin_base_patch4_window12_384_prediction</th>\n",
       "      <th>swin_large_patch4_window7_224_prediction</th>\n",
       "      <th>swin_large_patch4_window12_384_prediction</th>\n",
       "      <th>beitv2_large_patch16_224_prediction</th>\n",
       "      <th>val</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding</th>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.394830</td>\n",
       "      <td>0.415936</td>\n",
       "      <td>0.438780</td>\n",
       "      <td>0.496418</td>\n",
       "      <td>0.503803</td>\n",
       "      <td>0.487945</td>\n",
       "      <td>0.528896</td>\n",
       "      <td>0.491641</td>\n",
       "      <td>0.568876</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding</th>\n",
       "      <td>0.394830</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.777235</td>\n",
       "      <td>0.562500</td>\n",
       "      <td>0.498671</td>\n",
       "      <td>0.536388</td>\n",
       "      <td>0.541998</td>\n",
       "      <td>0.560076</td>\n",
       "      <td>0.548504</td>\n",
       "      <td>0.615261</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>interrogated_mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding</th>\n",
       "      <td>0.415936</td>\n",
       "      <td>0.777235</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.543048</td>\n",
       "      <td>0.626302</td>\n",
       "      <td>0.624522</td>\n",
       "      <td>0.627504</td>\n",
       "      <td>0.636011</td>\n",
       "      <td>0.645771</td>\n",
       "      <td>0.721854</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>blip_image_captioning_large_prediction_embedding</th>\n",
       "      <td>0.438780</td>\n",
       "      <td>0.562500</td>\n",
       "      <td>0.543048</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.507106</td>\n",
       "      <td>0.539150</td>\n",
       "      <td>0.534438</td>\n",
       "      <td>0.552370</td>\n",
       "      <td>0.541689</td>\n",
       "      <td>0.609641</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>knn_prediction</th>\n",
       "      <td>0.496418</td>\n",
       "      <td>0.498671</td>\n",
       "      <td>0.626302</td>\n",
       "      <td>0.507106</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.779868</td>\n",
       "      <td>0.778143</td>\n",
       "      <td>0.766443</td>\n",
       "      <td>0.774414</td>\n",
       "      <td>0.951969</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>swin_base_patch4_window12_384_prediction</th>\n",
       "      <td>0.503803</td>\n",
       "      <td>0.536388</td>\n",
       "      <td>0.624522</td>\n",
       "      <td>0.539150</td>\n",
       "      <td>0.779868</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.939095</td>\n",
       "      <td>0.944211</td>\n",
       "      <td>0.919918</td>\n",
       "      <td>0.911615</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>swin_large_patch4_window7_224_prediction</th>\n",
       "      <td>0.487945</td>\n",
       "      <td>0.541998</td>\n",
       "      <td>0.627504</td>\n",
       "      <td>0.534438</td>\n",
       "      <td>0.778143</td>\n",
       "      <td>0.939095</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.941544</td>\n",
       "      <td>0.933726</td>\n",
       "      <td>0.911871</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>swin_large_patch4_window12_384_prediction</th>\n",
       "      <td>0.528896</td>\n",
       "      <td>0.560076</td>\n",
       "      <td>0.636011</td>\n",
       "      <td>0.552370</td>\n",
       "      <td>0.766443</td>\n",
       "      <td>0.944211</td>\n",
       "      <td>0.941544</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.908267</td>\n",
       "      <td>0.909420</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>beitv2_large_patch16_224_prediction</th>\n",
       "      <td>0.491641</td>\n",
       "      <td>0.548504</td>\n",
       "      <td>0.645771</td>\n",
       "      <td>0.541689</td>\n",
       "      <td>0.774414</td>\n",
       "      <td>0.919918</td>\n",
       "      <td>0.933726</td>\n",
       "      <td>0.908267</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.904357</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>val</th>\n",
       "      <td>0.568876</td>\n",
       "      <td>0.615261</td>\n",
       "      <td>0.721854</td>\n",
       "      <td>0.609641</td>\n",
       "      <td>0.951969</td>\n",
       "      <td>0.911615</td>\n",
       "      <td>0.911871</td>\n",
       "      <td>0.909420</td>\n",
       "      <td>0.904357</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                    coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding  \\\n",
       "coca_vit_l_14_laion2B_s13B_b90k_prediction_embe...                                           1.000000      \n",
       "mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90...                                           0.394830      \n",
       "interrogated_mscoco_finetuned_coca_vit_l_14_lai...                                           0.415936      \n",
       "blip_image_captioning_large_prediction_embedding                                             0.438780      \n",
       "knn_prediction                                                                               0.496418      \n",
       "swin_base_patch4_window12_384_prediction                                                     0.503803      \n",
       "swin_large_patch4_window7_224_prediction                                                     0.487945      \n",
       "swin_large_patch4_window12_384_prediction                                                    0.528896      \n",
       "beitv2_large_patch16_224_prediction                                                          0.491641      \n",
       "val                                                                                          0.568876      \n",
       "\n",
       "                                                    mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding  \\\n",
       "coca_vit_l_14_laion2B_s13B_b90k_prediction_embe...                                           0.394830                       \n",
       "mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90...                                           1.000000                       \n",
       "interrogated_mscoco_finetuned_coca_vit_l_14_lai...                                           0.777235                       \n",
       "blip_image_captioning_large_prediction_embedding                                             0.562500                       \n",
       "knn_prediction                                                                               0.498671                       \n",
       "swin_base_patch4_window12_384_prediction                                                     0.536388                       \n",
       "swin_large_patch4_window7_224_prediction                                                     0.541998                       \n",
       "swin_large_patch4_window12_384_prediction                                                    0.560076                       \n",
       "beitv2_large_patch16_224_prediction                                                          0.548504                       \n",
       "val                                                                                          0.615261                       \n",
       "\n",
       "                                                    interrogated_mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90k_prediction_embedding  \\\n",
       "coca_vit_l_14_laion2B_s13B_b90k_prediction_embe...                                           0.415936                                    \n",
       "mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90...                                           0.777235                                    \n",
       "interrogated_mscoco_finetuned_coca_vit_l_14_lai...                                           1.000000                                    \n",
       "blip_image_captioning_large_prediction_embedding                                             0.543048                                    \n",
       "knn_prediction                                                                               0.626302                                    \n",
       "swin_base_patch4_window12_384_prediction                                                     0.624522                                    \n",
       "swin_large_patch4_window7_224_prediction                                                     0.627504                                    \n",
       "swin_large_patch4_window12_384_prediction                                                    0.636011                                    \n",
       "beitv2_large_patch16_224_prediction                                                          0.645771                                    \n",
       "val                                                                                          0.721854                                    \n",
       "\n",
       "                                                    blip_image_captioning_large_prediction_embedding  \\\n",
       "coca_vit_l_14_laion2B_s13B_b90k_prediction_embe...                                          0.438780   \n",
       "mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90...                                          0.562500   \n",
       "interrogated_mscoco_finetuned_coca_vit_l_14_lai...                                          0.543048   \n",
       "blip_image_captioning_large_prediction_embedding                                            1.000000   \n",
       "knn_prediction                                                                              0.507106   \n",
       "swin_base_patch4_window12_384_prediction                                                    0.539150   \n",
       "swin_large_patch4_window7_224_prediction                                                    0.534438   \n",
       "swin_large_patch4_window12_384_prediction                                                   0.552370   \n",
       "beitv2_large_patch16_224_prediction                                                         0.541689   \n",
       "val                                                                                         0.609641   \n",
       "\n",
       "                                                    knn_prediction  \\\n",
       "coca_vit_l_14_laion2B_s13B_b90k_prediction_embe...        0.496418   \n",
       "mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90...        0.498671   \n",
       "interrogated_mscoco_finetuned_coca_vit_l_14_lai...        0.626302   \n",
       "blip_image_captioning_large_prediction_embedding          0.507106   \n",
       "knn_prediction                                            1.000000   \n",
       "swin_base_patch4_window12_384_prediction                  0.779868   \n",
       "swin_large_patch4_window7_224_prediction                  0.778143   \n",
       "swin_large_patch4_window12_384_prediction                 0.766443   \n",
       "beitv2_large_patch16_224_prediction                       0.774414   \n",
       "val                                                       0.951969   \n",
       "\n",
       "                                                    swin_base_patch4_window12_384_prediction  \\\n",
       "coca_vit_l_14_laion2B_s13B_b90k_prediction_embe...                                  0.503803   \n",
       "mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90...                                  0.536388   \n",
       "interrogated_mscoco_finetuned_coca_vit_l_14_lai...                                  0.624522   \n",
       "blip_image_captioning_large_prediction_embedding                                    0.539150   \n",
       "knn_prediction                                                                      0.779868   \n",
       "swin_base_patch4_window12_384_prediction                                            1.000000   \n",
       "swin_large_patch4_window7_224_prediction                                            0.939095   \n",
       "swin_large_patch4_window12_384_prediction                                           0.944211   \n",
       "beitv2_large_patch16_224_prediction                                                 0.919918   \n",
       "val                                                                                 0.911615   \n",
       "\n",
       "                                                    swin_large_patch4_window7_224_prediction  \\\n",
       "coca_vit_l_14_laion2B_s13B_b90k_prediction_embe...                                  0.487945   \n",
       "mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90...                                  0.541998   \n",
       "interrogated_mscoco_finetuned_coca_vit_l_14_lai...                                  0.627504   \n",
       "blip_image_captioning_large_prediction_embedding                                    0.534438   \n",
       "knn_prediction                                                                      0.778143   \n",
       "swin_base_patch4_window12_384_prediction                                            0.939095   \n",
       "swin_large_patch4_window7_224_prediction                                            1.000000   \n",
       "swin_large_patch4_window12_384_prediction                                           0.941544   \n",
       "beitv2_large_patch16_224_prediction                                                 0.933726   \n",
       "val                                                                                 0.911871   \n",
       "\n",
       "                                                    swin_large_patch4_window12_384_prediction  \\\n",
       "coca_vit_l_14_laion2B_s13B_b90k_prediction_embe...                                   0.528896   \n",
       "mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90...                                   0.560076   \n",
       "interrogated_mscoco_finetuned_coca_vit_l_14_lai...                                   0.636011   \n",
       "blip_image_captioning_large_prediction_embedding                                     0.552370   \n",
       "knn_prediction                                                                       0.766443   \n",
       "swin_base_patch4_window12_384_prediction                                             0.944211   \n",
       "swin_large_patch4_window7_224_prediction                                             0.941544   \n",
       "swin_large_patch4_window12_384_prediction                                            1.000000   \n",
       "beitv2_large_patch16_224_prediction                                                  0.908267   \n",
       "val                                                                                  0.909420   \n",
       "\n",
       "                                                    beitv2_large_patch16_224_prediction  \\\n",
       "coca_vit_l_14_laion2B_s13B_b90k_prediction_embe...                             0.491641   \n",
       "mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90...                             0.548504   \n",
       "interrogated_mscoco_finetuned_coca_vit_l_14_lai...                             0.645771   \n",
       "blip_image_captioning_large_prediction_embedding                               0.541689   \n",
       "knn_prediction                                                                 0.774414   \n",
       "swin_base_patch4_window12_384_prediction                                       0.919918   \n",
       "swin_large_patch4_window7_224_prediction                                       0.933726   \n",
       "swin_large_patch4_window12_384_prediction                                      0.908267   \n",
       "beitv2_large_patch16_224_prediction                                            1.000000   \n",
       "val                                                                            0.904357   \n",
       "\n",
       "                                                         val  \n",
       "coca_vit_l_14_laion2B_s13B_b90k_prediction_embe...  0.568876  \n",
       "mscoco_finetuned_coca_vit_l_14_laion2B_s13B_b90...  0.615261  \n",
       "interrogated_mscoco_finetuned_coca_vit_l_14_lai...  0.721854  \n",
       "blip_image_captioning_large_prediction_embedding    0.609641  \n",
       "knn_prediction                                      0.951969  \n",
       "swin_base_patch4_window12_384_prediction            0.911615  \n",
       "swin_large_patch4_window7_224_prediction            0.911871  \n",
       "swin_large_patch4_window12_384_prediction           0.909420  \n",
       "beitv2_large_patch16_224_prediction                 0.904357  \n",
       "val                                                 1.000000  "
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_test_embeddings.corr()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aafb2744",
   "metadata": {
    "papermill": {
     "duration": 0.026602,
     "end_time": "2023-05-15T14:28:16.949655",
     "exception": false,
     "start_time": "2023-05-15T14:28:16.923053",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 12. Submission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "7ec7a769",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:28:17.005207Z",
     "iopub.status.busy": "2023-05-15T14:28:17.004859Z",
     "iopub.status.idle": "2023-05-15T14:28:17.026113Z",
     "shell.execute_reply": "2023-05-15T14:28:17.025064Z"
    },
    "papermill": {
     "duration": 0.051961,
     "end_time": "2023-05-15T14:28:17.028473",
     "exception": false,
     "start_time": "2023-05-15T14:28:16.976512",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_test_embeddings[['imgId_eId', 'val']].to_csv('submission.csv', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "6074c2e5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-15T14:28:17.085318Z",
     "iopub.status.busy": "2023-05-15T14:28:17.084398Z",
     "iopub.status.idle": "2023-05-15T14:28:19.371598Z",
     "shell.execute_reply": "2023-05-15T14:28:19.370017Z"
    },
    "papermill": {
     "duration": 2.318977,
     "end_time": "2023-05-15T14:28:19.374447",
     "exception": false,
     "start_time": "2023-05-15T14:28:17.055470",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "!rm -rf /kaggle/working/packages\n",
    "!rm -rf /kaggle/working/ci-preprocess"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 424.139785,
   "end_time": "2023-05-15T14:28:23.104034",
   "environment_variables": {},
   "exception": null,
   "input_path": "__notebook__.ipynb",
   "output_path": "__notebook__.ipynb",
   "parameters": {},
   "start_time": "2023-05-15T14:21:18.964249",
   "version": "2.4.0"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {
     "04b9740c09844aaf9fa103691f4ca1ad": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "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
      }
     },
     "13e78dab7b894644ad2a7195efedcf66": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HBoxModel",
      "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_19e8e2042fad4ef28427dc2fbbef8a06",
        "IPY_MODEL_33e6bb2204174bcf956af97b22f55121",
        "IPY_MODEL_ae47a0fa4ab548baba62a6be0b04c6a1"
       ],
       "layout": "IPY_MODEL_72894adc00b04a89ae46ecd0e6417986"
      }
     },
     "15e49787dd4d495d9df362731155d72b": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HTMLModel",
      "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_278072737d52436ab4bbefb15208258e",
       "placeholder": "​",
       "style": "IPY_MODEL_97f993eec6c04d2597cd136b89b649f9",
       "value": "Batches: 100%"
      }
     },
     "19e8e2042fad4ef28427dc2fbbef8a06": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HTMLModel",
      "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_04b9740c09844aaf9fa103691f4ca1ad",
       "placeholder": "​",
       "style": "IPY_MODEL_9162f48f259b483fbe0c3724ce344a73",
       "value": "Batches: 100%"
      }
     },
     "24eb031881974537b8278d46766ade02": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "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": ""
      }
     },
     "266450a1bac044b186bf47a55f5a2246": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "FloatProgressModel",
      "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_ac015b577df54738b9a8124cf81f3072",
       "max": 1.0,
       "min": 0.0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_79088089a91d42e4893355ad3efb442a",
       "value": 1.0
      }
     },
     "2702b92d078d492ab85fdf398935daf3": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "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
      }
     },
     "278072737d52436ab4bbefb15208258e": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "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
      }
     },
     "2d13640c4ee14a4c92bebc906ee03b5e": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "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
      }
     },
     "33e6bb2204174bcf956af97b22f55121": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "FloatProgressModel",
      "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_2702b92d078d492ab85fdf398935daf3",
       "max": 1.0,
       "min": 0.0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_9d6e78fd57dc48ffb0293fc94f3358bc",
       "value": 1.0
      }
     },
     "566422b75b25436d9c4e2ec42a02350a": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "ProgressStyleModel",
      "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": ""
      }
     },
     "646db25e5db8417698c18c9db7a44f97": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "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": ""
      }
     },
     "656d06da364346fda029068ec5af3cd1": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "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
      }
     },
     "68dc12f7ee7e46158fc7d93abe6b1682": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HTMLModel",
      "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_97dbf7d422b8416b993a3b13071a8d52",
       "placeholder": "​",
       "style": "IPY_MODEL_646db25e5db8417698c18c9db7a44f97",
       "value": "Batches: 100%"
      }
     },
     "6b3953993f7549bfb1527a18e628af75": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "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
      }
     },
     "6dfd451435a544f9b902105f40a2c817": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "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": ""
      }
     },
     "72894adc00b04a89ae46ecd0e6417986": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "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
      }
     },
     "768c932fa11e49f3ac4a4661b1324b10": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "FloatProgressModel",
      "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_656d06da364346fda029068ec5af3cd1",
       "max": 1.0,
       "min": 0.0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_566422b75b25436d9c4e2ec42a02350a",
       "value": 1.0
      }
     },
     "79088089a91d42e4893355ad3efb442a": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "ProgressStyleModel",
      "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": ""
      }
     },
     "7f2522e9b3f7451fb6fbb61bf3a50b7b": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "ProgressStyleModel",
      "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": ""
      }
     },
     "83825efd692441d69e03ea72c5c78343": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "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": ""
      }
     },
     "8e6a933639f1407784125d3c4a6c2b99": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "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
      }
     },
     "9162f48f259b483fbe0c3724ce344a73": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "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": ""
      }
     },
     "97dbf7d422b8416b993a3b13071a8d52": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "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
      }
     },
     "97f993eec6c04d2597cd136b89b649f9": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "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": ""
      }
     },
     "9d6e78fd57dc48ffb0293fc94f3358bc": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "ProgressStyleModel",
      "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": ""
      }
     },
     "9e0053f285a74e03a1ae099a18eeb63a": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "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": ""
      }
     },
     "ac015b577df54738b9a8124cf81f3072": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "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
      }
     },
     "ac7214793e714f4fb75267db61a2f1d8": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HBoxModel",
      "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_15e49787dd4d495d9df362731155d72b",
        "IPY_MODEL_266450a1bac044b186bf47a55f5a2246",
        "IPY_MODEL_e9e1ee8f550449db9651319ca4aedec6"
       ],
       "layout": "IPY_MODEL_f2c686af1cfe45ca990c384963c4b1b4"
      }
     },
     "ae47a0fa4ab548baba62a6be0b04c6a1": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HTMLModel",
      "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_2d13640c4ee14a4c92bebc906ee03b5e",
       "placeholder": "​",
       "style": "IPY_MODEL_83825efd692441d69e03ea72c5c78343",
       "value": " 1/1 [00:00&lt;00:00, 36.14it/s]"
      }
     },
     "b2adcd0d4ca3447a82c6f9dd4d2f1b81": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HTMLModel",
      "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_c2f01bff8dd64f2eac021b609b78985f",
       "placeholder": "​",
       "style": "IPY_MODEL_c137689f349b46b6bc18bada3532f6be",
       "value": " 1/1 [00:00&lt;00:00, 44.88it/s]"
      }
     },
     "b721d23531c84da39589edab197b0589": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HTMLModel",
      "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_6b3953993f7549bfb1527a18e628af75",
       "placeholder": "​",
       "style": "IPY_MODEL_9e0053f285a74e03a1ae099a18eeb63a",
       "value": " 1/1 [00:00&lt;00:00, 27.74it/s]"
      }
     },
     "ba4c07b7fd544fd3a5c977697dc410c3": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HTMLModel",
      "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_fe0c4532f74e42388fa609072991fa84",
       "placeholder": "​",
       "style": "IPY_MODEL_24eb031881974537b8278d46766ade02",
       "value": "Batches: 100%"
      }
     },
     "bc3bb8272e484878ba0feb38a798507a": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "FloatProgressModel",
      "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_ceb5a0bc78c14c4a9ff5d7dfe7926843",
       "max": 1.0,
       "min": 0.0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_7f2522e9b3f7451fb6fbb61bf3a50b7b",
       "value": 1.0
      }
     },
     "bc5de8dd66ce401587ef255c3bb14486": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "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
      }
     },
     "bd443fca64b4421e8c13269a7481edf2": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HBoxModel",
      "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_68dc12f7ee7e46158fc7d93abe6b1682",
        "IPY_MODEL_bc3bb8272e484878ba0feb38a798507a",
        "IPY_MODEL_b721d23531c84da39589edab197b0589"
       ],
       "layout": "IPY_MODEL_bc5de8dd66ce401587ef255c3bb14486"
      }
     },
     "bd8e1caf09f74f5981286b50d2bb41c2": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "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
      }
     },
     "c137689f349b46b6bc18bada3532f6be": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "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": ""
      }
     },
     "c2f01bff8dd64f2eac021b609b78985f": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "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
      }
     },
     "ceb5a0bc78c14c4a9ff5d7dfe7926843": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "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
      }
     },
     "e9e1ee8f550449db9651319ca4aedec6": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HTMLModel",
      "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_8e6a933639f1407784125d3c4a6c2b99",
       "placeholder": "​",
       "style": "IPY_MODEL_6dfd451435a544f9b902105f40a2c817",
       "value": " 1/1 [00:00&lt;00:00, 39.78it/s]"
      }
     },
     "f2c686af1cfe45ca990c384963c4b1b4": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "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
      }
     },
     "fe0c4532f74e42388fa609072991fa84": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "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
      }
     },
     "ffab3e510fd5468c95cdfb92c108d7c8": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HBoxModel",
      "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_ba4c07b7fd544fd3a5c977697dc410c3",
        "IPY_MODEL_768c932fa11e49f3ac4a4661b1324b10",
        "IPY_MODEL_b2adcd0d4ca3447a82c6f9dd4d2f1b81"
       ],
       "layout": "IPY_MODEL_bd8e1caf09f74f5981286b50d2bb41c2"
      }
     }
    },
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}