{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Interacting with CLIP.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YPHN7PJgKOzb"
      },
      "source": [
        "# Interacting with CLIP\n",
        "\n",
        "This is a self-contained notebook that shows how to download and run CLIP models, calculate the similarity between arbitrary image and text inputs, and perform zero-shot image classifications."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "53N4k0pj_9qL"
      },
      "source": [
        "# Preparation for Colab\n",
        "\n",
        "Make sure you're running a GPU runtime; if not, select \"GPU\" as the hardware accelerator in Runtime > Change Runtime Type in the menu. The next cells will print the CUDA version of the runtime if it has a GPU, and install PyTorch 1.7.1."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0BpdJkdBssk9",
        "outputId": "f26d1899-1e21-427a-f730-9451be9a2572"
      },
      "source": [
        "import subprocess\n",
        "\n",
        "CUDA_version = [s for s in subprocess.check_output([\"nvcc\", \"--version\"]).decode(\"UTF-8\").split(\", \") if s.startswith(\"release\")][0].split(\" \")[-1]\n",
        "print(\"CUDA version:\", CUDA_version)\n",
        "\n",
        "if CUDA_version == \"10.0\":\n",
        "    torch_version_suffix = \"+cu100\"\n",
        "elif CUDA_version == \"10.1\":\n",
        "    torch_version_suffix = \"+cu101\"\n",
        "elif CUDA_version == \"10.2\":\n",
        "    torch_version_suffix = \"\"\n",
        "else:\n",
        "    torch_version_suffix = \"+cu110\""
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "CUDA version: 10.1\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RBVr18E5tse8",
        "outputId": "d4647553-46fa-4cc5-8a41-b152abc0b4d2"
      },
      "source": [
        "! pip install torch==1.7.1{torch_version_suffix} torchvision==0.8.2{torch_version_suffix} -f https://download.pytorch.org/whl/torch_stable.html ftfy regex"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Looking in links: https://download.pytorch.org/whl/torch_stable.html\n",
            "Requirement already satisfied: torch==1.7.1+cu101 in /usr/local/lib/python3.6/dist-packages (1.7.1+cu101)\n",
            "Requirement already satisfied: torchvision==0.8.2+cu101 in /usr/local/lib/python3.6/dist-packages (0.8.2+cu101)\n",
            "Requirement already satisfied: ftfy in /usr/local/lib/python3.6/dist-packages (5.8)\n",
            "Requirement already satisfied: regex in /usr/local/lib/python3.6/dist-packages (2019.12.20)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from torch==1.7.1+cu101) (1.19.4)\n",
            "Requirement already satisfied: dataclasses; python_version < \"3.7\" in /usr/local/lib/python3.6/dist-packages (from torch==1.7.1+cu101) (0.8)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.6/dist-packages (from torch==1.7.1+cu101) (3.7.4.3)\n",
            "Requirement already satisfied: pillow>=4.1.1 in /usr/local/lib/python3.6/dist-packages (from torchvision==0.8.2+cu101) (7.0.0)\n",
            "Requirement already satisfied: wcwidth in /usr/local/lib/python3.6/dist-packages (from ftfy) (0.2.5)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C1hkDT38hSaP",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "2fb746de-c8f1-4e5f-d611-e151934c0994"
      },
      "source": [
        "import numpy as np\n",
        "import torch\n",
        "\n",
        "print(\"Torch version:\", torch.__version__)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Torch version: 1.7.1+cu101\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eFxgLV5HAEEw"
      },
      "source": [
        "# Downloading the model\n",
        "\n",
        "CLIP models are distributed as TorchScript modules."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uLFS29hnhlY4"
      },
      "source": [
        "MODELS = {\n",
        "    \"RN50\": \"https://openaipublic.azureedge.net/clip/models/afeb0e10f9e5a86da6080e35cf09123aca3b358a0c3e3b6c78a7b63bc04b6762/RN50.pt\",\n",
        "    \"RN101\": \"https://openaipublic.azureedge.net/clip/models/8fa8567bab74a42d41c5915025a8e4538c3bdbe8804a470a72f30b0d94fab599/RN101.pt\",\n",
        "    \"RN50x4\": \"https://openaipublic.azureedge.net/clip/models/7e526bd135e493cef0776de27d5f42653e6b4c8bf9e0f653bb11773263205fdd/RN50x4.pt\",\n",
        "    \"ViT-B/32\": \"https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt\",    \n",
        "}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cboKZocQlSYX",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "c3b6fa02-3b17-4a54-d3c4-8970cba3de9d"
      },
      "source": [
        "! wget {MODELS[\"ViT-B/32\"]} -O model.pt"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "--2021-01-08 17:41:09--  https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt\n",
            "Resolving openaipublic.azureedge.net (openaipublic.azureedge.net)... 13.107.246.13, 2620:1ec:bdf::13\n",
            "Connecting to openaipublic.azureedge.net (openaipublic.azureedge.net)|13.107.246.13|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 353976522 (338M) [application/octet-stream]\n",
            "Saving to: ‘model.pt’\n",
            "\n",
            "model.pt            100%[===================>] 337.58M   125MB/s    in 2.7s    \n",
            "\n",
            "2021-01-08 17:41:12 (125 MB/s) - ‘model.pt’ saved [353976522/353976522]\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IBRVTY9lbGm8",
        "outputId": "768bce5b-e807-43fd-e559-c4f120f16dc7"
      },
      "source": [
        "model = torch.jit.load(\"model.pt\").cuda().eval()\n",
        "input_resolution = model.input_resolution.item()\n",
        "context_length = model.context_length.item()\n",
        "vocab_size = model.vocab_size.item()\n",
        "\n",
        "print(\"Model parameters:\", f\"{np.sum([int(np.prod(p.shape)) for p in model.parameters()]):,}\")\n",
        "print(\"Input resolution:\", input_resolution)\n",
        "print(\"Context length:\", context_length)\n",
        "print(\"Vocab size:\", vocab_size)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Model parameters: 151,277,313\n",
            "Input resolution: 224\n",
            "Context length: 77\n",
            "Vocab size: 49408\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "21slhZGCqANb"
      },
      "source": [
        "# Image Preprocessing\n",
        "\n",
        "We resize the input images and center-crop them to conform with the image resolution that the model expects. Before doing so, we will normalize the pixel intensity using the dataset mean and standard deviation.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "d6cpiIFHp9N6"
      },
      "source": [
        "from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize\n",
        "from PIL import Image\n",
        "\n",
        "preprocess = Compose([\n",
        "    Resize(input_resolution, interpolation=Image.BICUBIC),\n",
        "    CenterCrop(input_resolution),\n",
        "    ToTensor()\n",
        "])\n",
        "\n",
        "image_mean = torch.tensor([0.48145466, 0.4578275, 0.40821073]).cuda()\n",
        "image_std = torch.tensor([0.26862954, 0.26130258, 0.27577711]).cuda()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xwSB5jZki3Cj"
      },
      "source": [
        "# Text Preprocessing\n",
        "\n",
        "We use a case-insensitive tokenizer. The tokenizer code is hidden in the second cell below"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qGom156-i2kL",
        "outputId": "b61049f3-0ce5-4e95-c477-9c62b2fabca4"
      },
      "source": [
        "! pip install ftfy regex\n",
        "! wget https://openaipublic.azureedge.net/clip/bpe_simple_vocab_16e6.txt.gz -O bpe_simple_vocab_16e6.txt.gz"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: ftfy in /usr/local/lib/python3.6/dist-packages (5.8)\n",
            "Requirement already satisfied: regex in /usr/local/lib/python3.6/dist-packages (2019.12.20)\n",
            "Requirement already satisfied: wcwidth in /usr/local/lib/python3.6/dist-packages (from ftfy) (0.2.5)\n",
            "--2021-01-08 17:41:19--  https://openaipublic.azureedge.net/clip/bpe_simple_vocab_16e6.txt.gz\n",
            "Resolving openaipublic.azureedge.net (openaipublic.azureedge.net)... 13.107.246.13, 2620:1ec:bdf::13\n",
            "Connecting to openaipublic.azureedge.net (openaipublic.azureedge.net)|13.107.246.13|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 1356917 (1.3M) [application/octet-stream]\n",
            "Saving to: ‘bpe_simple_vocab_16e6.txt.gz’\n",
            "\n",
            "bpe_simple_vocab_16 100%[===================>]   1.29M  --.-KB/s    in 0.01s   \n",
            "\n",
            "2021-01-08 17:41:19 (93.6 MB/s) - ‘bpe_simple_vocab_16e6.txt.gz’ saved [1356917/1356917]\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "toGtcd-Ji_MD",
        "cellView": "form"
      },
      "source": [
        "#@title\n",
        "\n",
        "import gzip\n",
        "import html\n",
        "import os\n",
        "from functools import lru_cache\n",
        "\n",
        "import ftfy\n",
        "import regex as re\n",
        "\n",
        "\n",
        "@lru_cache()\n",
        "def bytes_to_unicode():\n",
        "    \"\"\"\n",
        "    Returns list of utf-8 byte and a corresponding list of unicode strings.\n",
        "    The reversible bpe codes work on unicode strings.\n",
        "    This means you need a large # of unicode characters in your vocab if you want to avoid UNKs.\n",
        "    When you're at something like a 10B token dataset you end up needing around 5K for decent coverage.\n",
        "    This is a signficant percentage of your normal, say, 32K bpe vocab.\n",
        "    To avoid that, we want lookup tables between utf-8 bytes and unicode strings.\n",
        "    And avoids mapping to whitespace/control characters the bpe code barfs on.\n",
        "    \"\"\"\n",
        "    bs = list(range(ord(\"!\"), ord(\"~\")+1))+list(range(ord(\"¡\"), ord(\"¬\")+1))+list(range(ord(\"®\"), ord(\"ÿ\")+1))\n",
        "    cs = bs[:]\n",
        "    n = 0\n",
        "    for b in range(2**8):\n",
        "        if b not in bs:\n",
        "            bs.append(b)\n",
        "            cs.append(2**8+n)\n",
        "            n += 1\n",
        "    cs = [chr(n) for n in cs]\n",
        "    return dict(zip(bs, cs))\n",
        "\n",
        "\n",
        "def get_pairs(word):\n",
        "    \"\"\"Return set of symbol pairs in a word.\n",
        "    Word is represented as tuple of symbols (symbols being variable-length strings).\n",
        "    \"\"\"\n",
        "    pairs = set()\n",
        "    prev_char = word[0]\n",
        "    for char in word[1:]:\n",
        "        pairs.add((prev_char, char))\n",
        "        prev_char = char\n",
        "    return pairs\n",
        "\n",
        "\n",
        "def basic_clean(text):\n",
        "    text = ftfy.fix_text(text)\n",
        "    text = html.unescape(html.unescape(text))\n",
        "    return text.strip()\n",
        "\n",
        "\n",
        "def whitespace_clean(text):\n",
        "    text = re.sub(r'\\s+', ' ', text)\n",
        "    text = text.strip()\n",
        "    return text\n",
        "\n",
        "\n",
        "class SimpleTokenizer(object):\n",
        "    def __init__(self, bpe_path: str = \"bpe_simple_vocab_16e6.txt.gz\"):\n",
        "        self.byte_encoder = bytes_to_unicode()\n",
        "        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}\n",
        "        merges = gzip.open(bpe_path).read().decode(\"utf-8\").split('\\n')\n",
        "        merges = merges[1:49152-256-2+1]\n",
        "        merges = [tuple(merge.split()) for merge in merges]\n",
        "        vocab = list(bytes_to_unicode().values())\n",
        "        vocab = vocab + [v+'</w>' for v in vocab]\n",
        "        for merge in merges:\n",
        "            vocab.append(''.join(merge))\n",
        "        vocab.extend(['<|startoftext|>', '<|endoftext|>'])\n",
        "        self.encoder = dict(zip(vocab, range(len(vocab))))\n",
        "        self.decoder = {v: k for k, v in self.encoder.items()}\n",
        "        self.bpe_ranks = dict(zip(merges, range(len(merges))))\n",
        "        self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'}\n",
        "        self.pat = re.compile(r\"\"\"<\\|startoftext\\|>|<\\|endoftext\\|>|'s|'t|'re|'ve|'m|'ll|'d|[\\p{L}]+|[\\p{N}]|[^\\s\\p{L}\\p{N}]+\"\"\", re.IGNORECASE)\n",
        "\n",
        "    def bpe(self, token):\n",
        "        if token in self.cache:\n",
        "            return self.cache[token]\n",
        "        word = tuple(token[:-1]) + ( token[-1] + '</w>',)\n",
        "        pairs = get_pairs(word)\n",
        "\n",
        "        if not pairs:\n",
        "            return token+'</w>'\n",
        "\n",
        "        while True:\n",
        "            bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf')))\n",
        "            if bigram not in self.bpe_ranks:\n",
        "                break\n",
        "            first, second = bigram\n",
        "            new_word = []\n",
        "            i = 0\n",
        "            while i < len(word):\n",
        "                try:\n",
        "                    j = word.index(first, i)\n",
        "                    new_word.extend(word[i:j])\n",
        "                    i = j\n",
        "                except:\n",
        "                    new_word.extend(word[i:])\n",
        "                    break\n",
        "\n",
        "                if word[i] == first and i < len(word)-1 and word[i+1] == second:\n",
        "                    new_word.append(first+second)\n",
        "                    i += 2\n",
        "                else:\n",
        "                    new_word.append(word[i])\n",
        "                    i += 1\n",
        "            new_word = tuple(new_word)\n",
        "            word = new_word\n",
        "            if len(word) == 1:\n",
        "                break\n",
        "            else:\n",
        "                pairs = get_pairs(word)\n",
        "        word = ' '.join(word)\n",
        "        self.cache[token] = word\n",
        "        return word\n",
        "\n",
        "    def encode(self, text):\n",
        "        bpe_tokens = []\n",
        "        text = whitespace_clean(basic_clean(text)).lower()\n",
        "        for token in re.findall(self.pat, text):\n",
        "            token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))\n",
        "            bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))\n",
        "        return bpe_tokens\n",
        "\n",
        "    def decode(self, tokens):\n",
        "        text = ''.join([self.decoder[token] for token in tokens])\n",
        "        text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors=\"replace\").replace('</w>', ' ')\n",
        "        return text\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4W8ARJVqBJXs"
      },
      "source": [
        "# Setting up input images and texts\n",
        "\n",
        "We are going to feed 8 example images and their textual descriptions to the model, and compare the similarity between the corresponding features.\n",
        "\n",
        "The tokenizer is case-insensitive, and we can freely give any suitable textual descriptions."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tMc1AXzBlhzm"
      },
      "source": [
        "import os\n",
        "import skimage\n",
        "import IPython.display\n",
        "import matplotlib.pyplot as plt\n",
        "from PIL import Image\n",
        "import numpy as np\n",
        "\n",
        "from collections import OrderedDict\n",
        "import torch\n",
        "\n",
        "%matplotlib inline\n",
        "%config InlineBackend.figure_format = 'retina'\n",
        "\n",
        "# images in skimage to use and their textual descriptions\n",
        "descriptions = {\n",
        "    \"page\": \"a page of text about segmentation\",\n",
        "    \"chelsea\": \"a facial photo of a tabby cat\",\n",
        "    \"astronaut\": \"a portrait of an astronaut with the American flag\",\n",
        "    \"rocket\": \"a rocket standing on a launchpad\",\n",
        "    \"motorcycle_right\": \"a red motorcycle standing in a garage\",\n",
        "    \"camera\": \"a person looking at a camera on a tripod\",\n",
        "    \"horse\": \"a black-and-white silhouette of a horse\", \n",
        "    \"coffee\": \"a cup of coffee on a saucer\"\n",
        "}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NSSrLY185jSf",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 368
        },
        "outputId": "f4094535-8b5e-469b-c2b7-e0d02f2d4e9c"
      },
      "source": [
        "images = []\n",
        "texts = []\n",
        "plt.figure(figsize=(16, 5))\n",
        "\n",
        "for filename in [filename for filename in os.listdir(skimage.data_dir) if filename.endswith(\".png\") or filename.endswith(\".jpg\")]:\n",
        "    name = os.path.splitext(filename)[0]\n",
        "    if name not in descriptions:\n",
        "        continue\n",
        "\n",
        "    image = preprocess(Image.open(os.path.join(skimage.data_dir, filename)).convert(\"RGB\"))\n",
        "    images.append(image)\n",
        "    texts.append(descriptions[name])\n",
        "\n",
        "    plt.subplot(2, 4, len(images))\n",
        "    plt.imshow(image.permute(1, 2, 0))\n",
        "    plt.title(f\"{filename}\\n{descriptions[name]}\")\n",
        "    plt.xticks([])\n",
        "    plt.yticks([])\n",
        "\n",
        "plt.tight_layout()\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 1092968 characters]\n",
            "text/plain": [
              "<Figure size 1152x360 with 8 Axes>"
            ]
          },
          "metadata": {
            "tags": [],
            "image/png": {
              "width": 1107,
              "height": 351
            }
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WEVKsji6WOIX"
      },
      "source": [
        "## Building features\n",
        "\n",
        "We normalize the images, tokenize each text input, and run the forward pass of the model to get the image and text features."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QwkkczUPBRMh"
      },
      "source": [
        "image_input = torch.tensor(np.stack(images)).cuda()\n",
        "image_input -= image_mean[:, None, None]\n",
        "image_input /= image_std[:, None, None]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HBgCanxi8JKw"
      },
      "source": [
        "tokenizer = SimpleTokenizer()\n",
        "text_tokens = [tokenizer.encode(\"This is \" + desc) for desc in texts]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "w1l_muuhZ_Nk"
      },
      "source": [
        "text_input = torch.zeros(len(text_tokens), model.context_length, dtype=torch.long)\n",
        "sot_token = tokenizer.encoder['<|startoftext|>']\n",
        "eot_token = tokenizer.encoder['<|endoftext|>']\n",
        "\n",
        "for i, tokens in enumerate(text_tokens):\n",
        "    tokens = [sot_token] + tokens + [eot_token]\n",
        "    text_input[i, :len(tokens)] = torch.tensor(tokens)\n",
        "\n",
        "text_input = text_input.cuda()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZN9I0nIBZ_vW"
      },
      "source": [
        "with torch.no_grad():\n",
        "    image_features = model.encode_image(image_input).float()\n",
        "    text_features = model.encode_text(text_input).float()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cuxm2Gt4Wvzt"
      },
      "source": [
        "## Calculating cosine similarity\n",
        "\n",
        "We normalize the features and calculate the dot product of each pair."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yKAxkQR7bf3A"
      },
      "source": [
        "image_features /= image_features.norm(dim=-1, keepdim=True)\n",
        "text_features /= text_features.norm(dim=-1, keepdim=True)\n",
        "similarity = text_features.cpu().numpy() @ image_features.cpu().numpy().T"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C5zvMxh8cU6m",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 830
        },
        "outputId": "7acf50ee-dead-4e8e-f77d-f3e8a596bf45"
      },
      "source": [
        "count = len(descriptions)\n",
        "\n",
        "plt.figure(figsize=(20, 14))\n",
        "plt.imshow(similarity, vmin=0.1, vmax=0.3)\n",
        "# plt.colorbar()\n",
        "plt.yticks(range(count), texts, fontsize=18)\n",
        "plt.xticks([])\n",
        "for i, image in enumerate(images):\n",
        "    plt.imshow(image.permute(1, 2, 0), extent=(i - 0.5, i + 0.5, -1.6, -0.6), origin=\"lower\")\n",
        "for x in range(similarity.shape[1]):\n",
        "    for y in range(similarity.shape[0]):\n",
        "        plt.text(x, y, f\"{similarity[y, x]:.2f}\", ha=\"center\", va=\"center\", size=12)\n",
        "\n",
        "for side in [\"left\", \"top\", \"right\", \"bottom\"]:\n",
        "  plt.gca().spines[side].set_visible(False)\n",
        "\n",
        "plt.xlim([-0.5, count - 0.5])\n",
        "plt.ylim([count + 0.5, -2])\n",
        "\n",
        "plt.title(\"Cosine similarity between text and image features\", size=20)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Text(0.5, 1.0, 'Cosine similarity between text and image features')"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 17
        },
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 609640 characters]\n",
            "text/plain": [
              "<Figure size 1440x1008 with 1 Axes>"
            ]
          },
          "metadata": {
            "tags": [],
            "image/png": {
              "width": 1038,
              "height": 796
            },
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "alePijoXy6AH"
      },
      "source": [
        "# Zero-Shot Image Classification\n",
        "\n",
        "You can classify images using the cosine similarity (times 100) as the logits to the softmax operation."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Nqu4GlfPfr-p",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "c5cd7575-79a3-49a2-cd51-36f7ee2da1b7"
      },
      "source": [
        "from torchvision.datasets import CIFAR100\n",
        "\n",
        "cifar100 = CIFAR100(os.path.expanduser(\"~/.cache\"), transform=preprocess, download=True)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Files already downloaded and verified\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "C4S__zCGy2MT",
        "outputId": "6e0c62da-e7ee-44f0-8e20-6d5d0020437d"
      },
      "source": [
        "text_descriptions = [f\"This is a photo of a {label}\" for label in cifar100.classes]\n",
        "text_tokens = [[sot_token] + tokenizer.encode(desc) + [eot_token] for desc in text_descriptions]\n",
        "text_input = torch.zeros(len(text_tokens), model.context_length, dtype=torch.long)\n",
        "\n",
        "for i, tokens in enumerate(text_tokens):\n",
        "    text_input[i, :len(tokens)] = torch.tensor(tokens)\n",
        "\n",
        "text_input = text_input.cuda()\n",
        "text_input.shape"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([100, 77])"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 19
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c4z1fm9vCpSR"
      },
      "source": [
        "with torch.no_grad():\n",
        "    text_features = model.encode_text(text_input).float()\n",
        "    text_features /= text_features.norm(dim=-1, keepdim=True)\n",
        "\n",
        "text_probs = (100.0 * image_features @ text_features.T).softmax(dim=-1)\n",
        "top_probs, top_labels = text_probs.cpu().topk(5, dim=-1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 931
        },
        "id": "T6Ju_6IBE2Iz",
        "outputId": "d604f8d8-ac64-4cbc-fb84-f4b6c2643686"
      },
      "source": [
        "plt.figure(figsize=(16, 16))\n",
        "\n",
        "for i, image in enumerate(images):\n",
        "    plt.subplot(4, 4, 2 * i + 1)\n",
        "    plt.imshow(image.permute(1, 2, 0))\n",
        "    plt.axis(\"off\")\n",
        "\n",
        "    plt.subplot(4, 4, 2 * i + 2)\n",
        "    y = np.arange(top_probs.shape[-1])\n",
        "    plt.grid()\n",
        "    plt.barh(y, top_probs[i])\n",
        "    plt.gca().invert_yaxis()\n",
        "    plt.gca().set_axisbelow(True)\n",
        "    plt.yticks(y, [cifar100.classes[index] for index in top_labels[i].numpy()])\n",
        "    plt.xlabel(\"probability\")\n",
        "\n",
        "plt.subplots_adjust(wspace=0.5)\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 1556656 characters]\n",
            "text/plain": [
              "<Figure size 1152x1152 with 16 Axes>"
            ]
          },
          "metadata": {
            "tags": [],
            "image/png": {
              "width": 912,
              "height": 914
            },
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0OENu-DQLzQY"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}