{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multimodel Retrieval using BGE-VL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we show a very simple case of using BGE-VL model on text to image and (image, text) to image retrieval tasks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make sure you have the following dependencies installed in your environment:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install numpy torch transformers faiss-cpu pillow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[CIRCO](https://github.com/miccunifi/CIRCO) is a well labeled multimodel dataset, which use images taken from the [COCO 2017 unlabeled set](https://cocodataset.org/#home).\n",
    "\n",
    "If you are interested in trying the full dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !wget http://images.cocodataset.org/zips/unlabeled2017.zip\n",
    "# !unzip unlabeled2017.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Due to the tremendous size of the original dataset, we prepare a subset of it as a small corpus in our repo. Clone the repo and you can find them in MegaPairs/assets/corpus/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "fid_to_oid = {}\n",
    "image_dir = \"./assets/corpus/\"\n",
    "image_paths = [os.path.join(image_dir, f) for f in os.listdir(image_dir) if f.endswith(\".jpg\")]\n",
    "\n",
    "# extract image id from the path, and store the {sequence id: image id}\n",
    "# sequence id will be used as FAISS's default id\n",
    "for i, img in enumerate(image_paths):\n",
    "    fid_to_oid[i] = int(img[-16:-4])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function that convert image id back to path for later use\n",
    "def id_to_name(id, image_dir):\n",
    "    return image_dir + str(id).zfill(12) + \".jpg\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Embedding and Indexing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, use BGE-VL to encode the whole corpus:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from transformers import AutoModel\n",
    "\n",
    "MODEL_NAME = \"BAAI/BGE-VL-large\" # or \"BAAI/BGE-VL-base\"\n",
    "\n",
    "model = AutoModel.from_pretrained(MODEL_NAME, trust_remote_code=True) # You must set trust_remote_code=True\n",
    "model.set_processor(MODEL_NAME)\n",
    "model.eval()\n",
    "\n",
    "with torch.no_grad():\n",
    "    embeddings = model.encode(image_paths).to('cpu')\n",
    "\n",
    "embeddings = np.array(embeddings).astype(np.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(11, 768)\n"
     ]
    }
   ],
   "source": [
    "print(embeddings.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then store all the embeddings in an FAISS index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import faiss\n",
    "\n",
    "dim = embeddings.shape[1]\n",
    "index = faiss.index_factory(dim, \"Flat\", faiss.METRIC_L2)\n",
    "\n",
    "index.add(embeddings)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "or you can directly load a saved index:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# index = faiss.read_index(\"./index.bin\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Text -> Image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first task is using a query text to retrieve relevant images. Let's have a try with the following queries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "queries = [\n",
    "    \"Find a picture of oranges in a blue basket\",\n",
    "    \"Find a picture of both oranges and bananas\",\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = 1\n",
    "\n",
    "with torch.no_grad():\n",
    "    queries_vec = model.encode(\n",
    "        text = queries,\n",
    "    ).to('cpu')\n",
    "\n",
    "D, I = index.search(queries_vec, k=k)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the results follow the instruction well in details:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "query text: Find a picture of oranges in a blue basket\n",
      "['./assets/corpus/000000275230.jpg']\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 359932 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "query text: Find a picture of both oranges and bananas\n",
      "['./assets/corpus/000000156031.jpg']\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 168216 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import Image, display\n",
    "\n",
    "for i in range(len(queries)):\n",
    "    print(f\"query text: {queries[i]}\")\n",
    "    \n",
    "    ids = [fid_to_oid[ind] for ind in I[i]]\n",
    "    names = list(map(id_to_name, ids, [image_dir]*k))\n",
    "    images = list(map(Image, names))\n",
    "    print(names)\n",
    "    display(*images)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## (Image, Text) -> Image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The second task is using the combination of image and text query to retrieve image:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 237300 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_img_pth = \"./assets/query/000000530944.jpg\"\n",
    "Image(filename=query_img_pth)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the following text query, we want an image that a motorbike carrying a bag of orange:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = 3\n",
    "\n",
    "with torch.no_grad():\n",
    "    query = model.encode(\n",
    "        images = query_img_pth,\n",
    "        text = \"Find a picture that also a motorbike but replace the load to a bag of orange\"\n",
    "    ).to('cpu')\n",
    "\n",
    "D, I = index.search(query, k=k)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the No.1 result perfectly match our requirements! And the other 2 in the top 3 results also show very relevant images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['./assets/corpus/000000545037.jpg', './assets/corpus/000000272130.jpg', './assets/corpus/000000098911.jpg']\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 198532 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 197860 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 215812 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ids = [fid_to_oid[ind] for ind in I[0]]\n",
    "names = list(map(id_to_name, ids, [image_dir]*k))\n",
    "images = list(map(Image, names))\n",
    "print(names)\n",
    "display(*images)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "test",
   "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.8.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
