{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "363b8c2c-2350-48dc-a419-56f536411ba1",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/dbamman/anlp25/blob/main/12.multimodal/Multimodal.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef74cf81-b111-4099-b00b-dac83511b698",
   "metadata": {},
   "source": [
    "# Multimodal embeddings with CLIP\n",
    "\n",
    "In this notebook, we will explore multimodal embeddings using the CLIP model, which includes an image embedder and text encoder that project both into the same embedding space.\n",
    "\n",
    "We will use CLIP to explore a small subset of images from the National Gallery of Art."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58ae8a50-589d-4ece-a655-93a1c812637b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import islice\n",
    "from pathlib import Path\n",
    "\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from torchvision.io import read_image\n",
    "from transformers import SiglipModel, SiglipProcessor, CLIPModel, AutoProcessor\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c02fa869-d761-4a4e-9b5e-f66d92c13998",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "model = CLIPModel.from_pretrained(\"openai/clip-vit-base-patch32\").to(device)\n",
    "processor = AutoProcessor.from_pretrained(\"openai/clip-vit-base-patch32\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "611e9329-a21b-43f5-812f-55bce7ed830b",
   "metadata": {},
   "source": [
    "## Load the data\n",
    "\n",
    "We've compiled some public-domain images from [the National Gallery of Art](https://www.nga.gov/artworks/free-images-and-open-access#a-section-header-p102746). Let's download the data and take a look at a few of the images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6727bb79-6f93-4aee-95e3-eff012229cc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://github.com/dbamman/anlp25/raw/refs/heads/main/data/nga.tar.gz -O nga.tar.gz\n",
    "!tar -xzf nga.tar.gz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8936e17d-e129-4c04-a0d3-dd6b89b271e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_paths = sorted(list(Path(\"images/\").glob(\"*.jpg\")))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b9e6a6f-ae05-490a-9193-674dc347f820",
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_image(path):\n",
    "    plt.imshow(read_image(path).permute(1, 2, 0))\n",
    "    plt.axis('off')\n",
    "\n",
    "def show_images(paths, num_per_row=5):\n",
    "    num_images = len(paths)\n",
    "    num_rows = (num_images + num_per_row - 1) // num_per_row  # Ceiling division\n",
    "    \n",
    "    fig, axes = plt.subplots(num_rows, num_per_row, figsize=(num_per_row * 3, num_rows * 3))\n",
    "    \n",
    "    # Flatten axes array for easier indexing (handles both 1D and 2D cases)\n",
    "    if num_rows == 1 and num_per_row == 1:\n",
    "        axes = np.array([axes])\n",
    "    elif num_rows == 1 or num_per_row == 1:\n",
    "        axes = axes.flatten()\n",
    "    else:\n",
    "        axes = axes.flatten()\n",
    "    \n",
    "    for idx, path in enumerate(paths):\n",
    "        axes[idx].imshow(read_image(path).permute(1, 2, 0))\n",
    "        axes[idx].axis('off')  # Hide axes for cleaner display\n",
    "    \n",
    "    # Hide any unused subplots\n",
    "    for idx in range(num_images, len(axes)):\n",
    "        axes[idx].axis('off')\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ef264f0-ec8f-4f8c-86ad-a28426b34bfb",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_images(image_paths[:15])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "099895d0-acb4-41ba-9b02-9b224b77f5e5",
   "metadata": {},
   "source": [
    "## Embedding images\n",
    "\n",
    "We will use the CLIP model to get the image embeddings for all of the images in our dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5565d38-6cce-4a0c-ba35-4ad38495fd29",
   "metadata": {},
   "outputs": [],
   "source": [
    "def batched(iterable, n, *, strict=False):\n",
    "    # batched('ABCDEFG', 2) → AB CD EF G\n",
    "    if n < 1:\n",
    "        raise ValueError('n must be at least one')\n",
    "    iterator = iter(iterable)\n",
    "    while batch := tuple(islice(iterator, n)):\n",
    "        if strict and len(batch) != n:\n",
    "            raise ValueError('batched(): incomplete batch')\n",
    "        yield batch\n",
    "\n",
    "def get_image_embeddings():\n",
    "    with torch.no_grad():\n",
    "        all_outputs = []\n",
    "        for batch in batched(tqdm(image_paths), 32):\n",
    "            batch_images = [read_image(path) for path in batch]\n",
    "            inputs = processor.image_processor(batch_images, return_tensors=\"pt\").to(model.device)\n",
    "            outputs = model.get_image_features(**inputs)\n",
    "            all_outputs.append(outputs.cpu())\n",
    "    return torch.vstack(all_outputs)\n",
    "\n",
    "    \n",
    "embeds = get_image_embeddings()\n",
    "embeds.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b16d94b-85ad-4fdb-8fdb-d94d664d542b",
   "metadata": {},
   "source": [
    "## Querying against the embeddings\n",
    "\n",
    "Recall our previous experiments with word and sentence embeddings, where we queried for nearest neighbors based on cosine similarity. We can do the same here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07974622-7328-4339-a605-420376551bef",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_nn(query_vec, n=10):\n",
    "    if len(query_vec.shape) < 2:\n",
    "        # if query_vec is a single vector, make it a batch of size 1\n",
    "        query_vec = query_vec.unsqueeze(0)\n",
    "    sims = torch.cosine_similarity(query_vec, embeds)\n",
    "    return sims.argsort()[-n:].tolist()[::-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "981a9560-665c-42f3-8ff6-cea15dfea95f",
   "metadata": {},
   "source": [
    "For example, let's find the 10 most similar images to this image of a bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f4942e7-2b6a-4be9-81da-7d19934a01c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_image(image_paths[754])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3f497f6-d62b-4cc5-9899-2cf942da1e3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "neighbors = get_nn(embeds[754])\n",
    "show_images([image_paths[i] for i in neighbors])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a72b314-71e9-4dac-8daf-d939e681410b",
   "metadata": {},
   "source": [
    "The advantage of a multimodal model like CLIP is that we can query using not only images, but also text. Here, let's implement a search by using the text encoder of the CLIP model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbed5a5e-ef6b-4b93-97ca-f31df8724ed4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def search(query, n=10):\n",
    "    processed = processor.tokenizer([query], return_tensors=\"pt\").to(model.device)\n",
    "    with torch.no_grad():\n",
    "        out = model.get_text_features(**processed)\n",
    "        embedding = out.cpu()\n",
    "    nns = get_nn(embedding, n)\n",
    "    return nns"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54744763-ab7a-407a-a03d-101da5443e07",
   "metadata": {},
   "source": [
    "Now we can use natural language to search our images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c10ac63c-493a-45be-8591-74c83a1af433",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_images([image_paths[i] for i in search(\"A serene lake\", n=6)], num_per_row=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1798f735-6815-496f-bd81-2841ce22974f",
   "metadata": {},
   "source": [
    "## Explore\n",
    "\n",
    "What else can you find in this dataset?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}