multimodal_rag_with_milvus.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WsUd_S1k2kiV"
      },
      "source": [
        "# Multimodal RAG with Milvus 🖼️\n",
        "\n",
        "<img src=\"https://raw.githubusercontent.com/milvus-io/bootcamp/master/bootcamp/tutorials/quickstart/apps/multimodal_rag_with_milvus/pics/step3.jpg\n",
        "\"/>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ShUQhrhP2kiW"
      },
      "source": [
        "This tutorial showcases the multimodal RAG powered by Milvus, [Visualized BGE model](https://github.com/FlagOpen/FlagEmbedding/tree/master/FlagEmbedding/visual), and [GPT-4o](https://openai.com/index/hello-gpt-4o/). With this system, users are able to upload an image and edit text instructions, which are processed by BGE's composed retrieval model to search for candidate images. GPT-4o then acts as a reranker, selecting the most suitable image and providing the rationale behind the choice. This powerful combination enables a seamless and intuitive image search experience, leveraging Milvus for efficient retrieval, BGE model for precise image processing and matching, and GPT-4o for advanced reranking."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cLGNLpS32kiX"
      },
      "source": [
        "## Preparation\n",
        "\n",
        "### Install Dependencies"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "scrolled": true,
        "id": "Tn2CVlNQwCoY",
        "outputId": "e440203b-de6c-4e75-8258-86b8ef44a5bd"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "ChatCompletion(id='chatcmpl-A8trRKQpYV5ZYzNYShwQRgfcUl9aw', choices=[Choice(finish_reason='stop', index=0, logprobs=None, message=ChatCompletionMessage(content='This is a test. How can I assist you further?', refusal=None, role='assistant', function_call=None, tool_calls=None))], created=1726684717, model='gpt-4o-mini-2024-07-18', object='chat.completion', service_tier=None, system_fingerprint='fp_71fab63729', usage=CompletionUsage(completion_tokens=12, prompt_tokens=12, total_tokens=24, completion_tokens_details=CompletionTokensDetails(reasoning_tokens=0)))\n"
          ]
        }
      ],
      "source": [
        "api_key = \"sk-\"\n",
        "\n",
        "client = OpenAI(\n",
        "  organization='',\n",
        "  project='',\n",
        "  api_key = api_key\n",
        ")\n",
        "\n",
        "response = client.chat.completions.create(\n",
        "    messages=[{\n",
        "        \"role\": \"user\",\n",
        "        \"content\": \"Say this is a test\",\n",
        "    }],\n",
        "    model=\"gpt-4o-mini\",\n",
        ")\n",
        "\n",
        "print(response)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WPd6Axtw2kiX"
      },
      "outputs": [],
      "source": [
        "#!pip install --upgrade pymilvus openai datasets opencv-python timm einops ftfy peft tqdm"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DSJRoacf2kiY"
      },
      "outputs": [],
      "source": [
        "#!git clone https://github.com/FlagOpen/FlagEmbedding.git\n",
        "#!pip install -e FlagEmbedding"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rbu3UzD12kiZ"
      },
      "source": [
        "> If you are using Google Colab, to enable dependencies just installed, you may need to **restart the runtime** (click on the \"Runtime\" menu at the top of the screen, and select \"Restart session\" from the dropdown menu)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "417rUTEawCob"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from tqdm import tqdm\n",
        "from glob import glob\n",
        "import torch\n",
        "from FlagEmbedding.visual.modeling import Visualized_BGE\n",
        "from pymilvus import MilvusClient\n",
        "import numpy as np\n",
        "import cv2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oLliMCew2kiZ"
      },
      "source": [
        "### Download Data\n",
        "\n",
        "The following command will download the example data and extract to a local folder \"./images_folder\" including:\n",
        "\n",
        "- **images**: A subset of [Amazon Reviews 2023](https://github.com/hyp1231/AmazonReviews2023) as containing approximately 900 images from the categories \"Appliance\", \"Cell_Phones_and_Accessories\", and \"Electronics\".\n",
        "- **leopard.jpg**: An example query image."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VdNf29Eu2kia"
      },
      "outputs": [],
      "source": [
        "#!wget https://github.com/milvus-io/bootcamp/releases/download/data/amazon_reviews_2023_subset.tar.gz\n",
        "#!tar -xzf amazon_reviews_2023_subset.tar.gz"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c7OYNESB2kia"
      },
      "source": [
        "### Load Embedding Model\n",
        "\n",
        "We will use the Visualized BGE model \"bge-visualized-base-en-v1.5\" to generate embeddings for both images and text.\n",
        "\n",
        "**1. Download weight**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OabFTATv2kia"
      },
      "outputs": [],
      "source": [
        "#!wget https://huggingface.co/BAAI/bge-visualized/resolve/main/Visualized_base_en_v1.5.pth"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4ZLd2bTS2kib"
      },
      "source": [
        "**2. Build encoder**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qxHZRzRD2kib",
        "outputId": "f5b4d20e-239c-49f7-fd1e-d5d9349c6ac7"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/home/suncell13/Documents/FlagEmbedding/FlagEmbedding/visual/modeling.py:105: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
            "  self.load_state_dict(torch.load(model_weight, map_location='cpu'))\n"
          ]
        }
      ],
      "source": [
        "class Encoder:\n",
        "    def __init__(self, model_name: str, model_path: str):\n",
        "        self.model = Visualized_BGE(model_name_bge=model_name, model_weight=model_path)\n",
        "        self.model.eval()\n",
        "\n",
        "    def encode_query(self, image_path: str, text: str) -> list[float]:\n",
        "        with torch.no_grad():\n",
        "            query_emb = self.model.encode(image=image_path, text=text)\n",
        "        return query_emb.tolist()[0]\n",
        "\n",
        "    def encode_image(self, image_path: str) -> list[float]:\n",
        "        with torch.no_grad():\n",
        "            query_emb = self.model.encode(image=image_path)\n",
        "        return query_emb.tolist()[0]\n",
        "\n",
        "\n",
        "model_name = \"BAAI/bge-base-en-v1.5\"\n",
        "model_path = \"./Visualized_base_en_v1.5.pth\"  # Change to your own value if using a different model path\n",
        "encoder = Encoder(model_name, model_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VObdRXaF2kic"
      },
      "source": [
        "## Load Data\n",
        "\n",
        "This section will load example images into the database with corresponding embeddings.\n",
        "\n",
        "### Generate embeddings\n",
        "\n",
        "Load all jpeg images from the data directory and apply the encoder to convert images to embeddings."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r2p9WzVqwCoc"
      },
      "outputs": [],
      "source": [
        "data_dir = (\n",
        "    \"./images_folder\"  # Change to your own value if using a different data directory\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RdAnf_cm2kic",
        "outputId": "8b763670-ee58-45c7-cbd5-c7e2c82f4e6c"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Generating image embeddings: 100%|████████████| 900/900 [47:55<00:00,  3.20s/it]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Number of encoded images: 900\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "# Generate embeddings for the image dataset\n",
        "\n",
        "image_list = glob(\n",
        "    os.path.join(data_dir, \"images\", \"*.jpg\")\n",
        ")  # We will only use images ending with \".jpg\"\n",
        "image_dict = {}\n",
        "for image_path in tqdm(image_list, desc=\"Generating image embeddings: \"):\n",
        "    try:\n",
        "        image_dict[image_path] = encoder.encode_image(image_path)\n",
        "    except Exception as e:\n",
        "        print(f\"Failed to generate embedding for {image_path}. Skipped.\")\n",
        "        continue\n",
        "print(\"Number of encoded images:\", len(image_dict))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O54NJ9Iv2kic"
      },
      "source": [
        "### Insert into Milvus\n",
        "\n",
        "Insert images with corresponding paths and embeddings into Milvus collection.\n",
        "\n",
        "> As for the argument of `MilvusClient`:\n",
        "> - Setting the `uri` as a local file, e.g. `./milvus_demo.db`, is the most convenient method, as it automatically utilizes [Milvus Lite](https://milvus.io/docs/milvus_lite.md) to store all data in this file.\n",
        "> - If you have large scale of data, you can set up a more performant Milvus server on [docker or kubernetes](https://milvus.io/docs/quickstart.md). In this setup, please use the server uri, e.g.`http://localhost:19530`, as your `uri`.\n",
        "> - If you want to use [Zilliz Cloud](https://zilliz.com/cloud), the fully managed cloud service for Milvus, adjust the `uri` and `token`, which correspond to the [Public Endpoint and Api key](https://docs.zilliz.com/docs/on-zilliz-cloud-console#free-cluster-details) in Zilliz Cloud."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lnTSeZp9wCod",
        "outputId": "23e652f3-77ed-4547-dbe6-cc0a842e7c81"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "DEBUG:pymilvus.milvus_client.milvus_client:Created new connection using: 5110845abd2c464f9dd98d41532ad676\n"
          ]
        }
      ],
      "source": [
        "# Connect to Milvus client given URI\n",
        "milvus_client = MilvusClient(uri=\"http://localhost:19530\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YjRz1HqowCod",
        "outputId": "a820845b-3147-4f2c-91b2-7e066d505b96"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "execution_count": 13,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Verify if the collection exists\n",
        "COLLECTION_NAME = \"multimodal_rag_demo\"\n",
        "milvus_client.has_collection(COLLECTION_NAME)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rG3TwtuGwCod",
        "outputId": "0898fd36-af06-4342-c2e2-35c5b0c2618d"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Collection has been deleted\n"
          ]
        }
      ],
      "source": [
        "# Run if you want to drop your old data\n",
        "try:\n",
        "    milvus_client.drop_collection(COLLECTION_NAME)\n",
        "    print(\"Collection has been deleted\")\n",
        "except:  # noqa: E722\n",
        "    passm"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "x91uZljl2kic",
        "outputId": "82fe1106-8105-4494-8871-b91c7bb82ff5"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "DEBUG:pymilvus.milvus_client.milvus_client:Successfully created collection: multimodal_rag_demo\n",
            "DEBUG:pymilvus.milvus_client.milvus_client:Successfully created an index on collection: multimodal_rag_demo\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "{'insert_count': 900, 'ids': [452639093245101000, 452639093245101001, 452639093245101002, 452639093245101003, 452639093245101004, 452639093245101005, 452639093245101006, 452639093245101007, 452639093245101008, 452639093245101009, 452639093245101010, 452639093245101011, 452639093245101012, 452639093245101013, 452639093245101014, 452639093245101015, 452639093245101016, 452639093245101017, 452639093245101018, 452639093245101019, 452639093245101020, 452639093245101021, 452639093245101022, 452639093245101023, 452639093245101024, 452639093245101025, 452639093245101026, 452639093245101027, 452639093245101028, 452639093245101029, 452639093245101030, 452639093245101031, 452639093245101032, 452639093245101033, 452639093245101034, 452639093245101035, 452639093245101036, 452639093245101037, 452639093245101038, 452639093245101039, 452639093245101040, 452639093245101041, 452639093245101042, 452639093245101043, 452639093245101044, 452639093245101045, 452639093245101046, 452639093245101047, 452639093245101048, 452639093245101049, 452639093245101050, 452639093245101051, 452639093245101052, 452639093245101053, 452639093245101054, 452639093245101055, 452639093245101056, 452639093245101057, 452639093245101058, 452639093245101059, 452639093245101060, 452639093245101061, 452639093245101062, 452639093245101063, 452639093245101064, 452639093245101065, 452639093245101066, 452639093245101067, 452639093245101068, 452639093245101069, 452639093245101070, 452639093245101071, 452639093245101072, 452639093245101073, 452639093245101074, 452639093245101075, 452639093245101076, 452639093245101077, 452639093245101078, 452639093245101079, 452639093245101080, 452639093245101081, 452639093245101082, 452639093245101083, 452639093245101084, 452639093245101085, 452639093245101086, 452639093245101087, 452639093245101088, 452639093245101089, 452639093245101090, 452639093245101091, 452639093245101092, 452639093245101093, 452639093245101094, 452639093245101095, 452639093245101096, 452639093245101097, 452639093245101098, 452639093245101099, 452639093245101100, 452639093245101101, 452639093245101102, 452639093245101103, 452639093245101104, 452639093245101105, 452639093245101106, 452639093245101107, 452639093245101108, 452639093245101109, 452639093245101110, 452639093245101111, 452639093245101112, 452639093245101113, 452639093245101114, 452639093245101115, 452639093245101116, 452639093245101117, 452639093245101118, 452639093245101119, 452639093245101120, 452639093245101121, 452639093245101122, 452639093245101123, 452639093245101124, 452639093245101125, 452639093245101126, 452639093245101127, 452639093245101128, 452639093245101129, 452639093245101130, 452639093245101131, 452639093245101132, 452639093245101133, 452639093245101134, 452639093245101135, 452639093245101136, 452639093245101137, 452639093245101138, 452639093245101139, 452639093245101140, 452639093245101141, 452639093245101142, 452639093245101143, 452639093245101144, 452639093245101145, 452639093245101146, 452639093245101147, 452639093245101148, 452639093245101149, 452639093245101150, 452639093245101151, 452639093245101152, 452639093245101153, 452639093245101154, 452639093245101155, 452639093245101156, 452639093245101157, 452639093245101158, 452639093245101159, 452639093245101160, 452639093245101161, 452639093245101162, 452639093245101163, 452639093245101164, 452639093245101165, 452639093245101166, 452639093245101167, 452639093245101168, 452639093245101169, 452639093245101170, 452639093245101171, 452639093245101172, 452639093245101173, 452639093245101174, 452639093245101175, 452639093245101176, 452639093245101177, 452639093245101178, 452639093245101179, 452639093245101180, 452639093245101181, 452639093245101182, 452639093245101183, 452639093245101184, 452639093245101185, 452639093245101186, 452639093245101187, 452639093245101188, 452639093245101189, 452639093245101190, 452639093245101191, 452639093245101192, 452639093245101193, 452639093245101194, 452639093245101195, 452639093245101196, 452639093245101197, 452639093245101198, 452639093245101199, 452639093245101200, 452639093245101201, 452639093245101202, 452639093245101203, 452639093245101204, 452639093245101205, 452639093245101206, 452639093245101207, 452639093245101208, 452639093245101209, 452639093245101210, 452639093245101211, 452639093245101212, 452639093245101213, 452639093245101214, 452639093245101215, 452639093245101216, 452639093245101217, 452639093245101218, 452639093245101219, 452639093245101220, 452639093245101221, 452639093245101222, 452639093245101223, 452639093245101224, 452639093245101225, 452639093245101226, 452639093245101227, 452639093245101228, 452639093245101229, 452639093245101230, 452639093245101231, 452639093245101232, 452639093245101233, 452639093245101234, 452639093245101235, 452639093245101236, 452639093245101237, 452639093245101238, 452639093245101239, 452639093245101240, 452639093245101241, 452639093245101242, 452639093245101243, 452639093245101244, 452639093245101245, 452639093245101246, 452639093245101247, 452639093245101248, 452639093245101249, 452639093245101250, 452639093245101251, 452639093245101252, 452639093245101253, 452639093245101254, 452639093245101255, 452639093245101256, 452639093245101257, 452639093245101258, 452639093245101259, 452639093245101260, 452639093245101261, 452639093245101262, 452639093245101263, 452639093245101264, 452639093245101265, 452639093245101266, 452639093245101267, 452639093245101268, 452639093245101269, 452639093245101270, 452639093245101271, 452639093245101272, 452639093245101273, 452639093245101274, 452639093245101275, 452639093245101276, 452639093245101277, 452639093245101278, 452639093245101279, 452639093245101280, 452639093245101281, 452639093245101282, 452639093245101283, 452639093245101284, 452639093245101285, 452639093245101286, 452639093245101287, 452639093245101288, 452639093245101289, 452639093245101290, 452639093245101291, 452639093245101292, 452639093245101293, 452639093245101294, 452639093245101295, 452639093245101296, 452639093245101297, 452639093245101298, 452639093245101299, 452639093245101300, 452639093245101301, 452639093245101302, 452639093245101303, 452639093245101304, 452639093245101305, 452639093245101306, 452639093245101307, 452639093245101308, 452639093245101309, 452639093245101310, 452639093245101311, 452639093245101312, 452639093245101313, 452639093245101314, 452639093245101315, 452639093245101316, 452639093245101317, 452639093245101318, 452639093245101319, 452639093245101320, 452639093245101321, 452639093245101322, 452639093245101323, 452639093245101324, 452639093245101325, 452639093245101326, 452639093245101327, 452639093245101328, 452639093245101329, 452639093245101330, 452639093245101331, 452639093245101332, 452639093245101333, 452639093245101334, 452639093245101335, 452639093245101336, 452639093245101337, 452639093245101338, 452639093245101339, 452639093245101340, 452639093245101341, 452639093245101342, 452639093245101343, 452639093245101344, 452639093245101345, 452639093245101346, 452639093245101347, 452639093245101348, 452639093245101349, 452639093245101350, 452639093245101351, 452639093245101352, 452639093245101353, 452639093245101354, 452639093245101355, 452639093245101356, 452639093245101357, 452639093245101358, 452639093245101359, 452639093245101360, 452639093245101361, 452639093245101362, 452639093245101363, 452639093245101364, 452639093245101365, 452639093245101366, 452639093245101367, 452639093245101368, 452639093245101369, 452639093245101370, 452639093245101371, 452639093245101372, 452639093245101373, 452639093245101374, 452639093245101375, 452639093245101376, 452639093245101377, 452639093245101378, 452639093245101379, 452639093245101380, 452639093245101381, 452639093245101382, 452639093245101383, 452639093245101384, 452639093245101385, 452639093245101386, 452639093245101387, 452639093245101388, 452639093245101389, 452639093245101390, 452639093245101391, 452639093245101392, 452639093245101393, 452639093245101394, 452639093245101395, 452639093245101396, 452639093245101397, 452639093245101398, 452639093245101399, 452639093245101400, 452639093245101401, 452639093245101402, 452639093245101403, 452639093245101404, 452639093245101405, 452639093245101406, 452639093245101407, 452639093245101408, 452639093245101409, 452639093245101410, 452639093245101411, 452639093245101412, 452639093245101413, 452639093245101414, 452639093245101415, 452639093245101416, 452639093245101417, 452639093245101418, 452639093245101419, 452639093245101420, 452639093245101421, 452639093245101422, 452639093245101423, 452639093245101424, 452639093245101425, 452639093245101426, 452639093245101427, 452639093245101428, 452639093245101429, 452639093245101430, 452639093245101431, 452639093245101432, 452639093245101433, 452639093245101434, 452639093245101435, 452639093245101436, 452639093245101437, 452639093245101438, 452639093245101439, 452639093245101440, 452639093245101441, 452639093245101442, 452639093245101443, 452639093245101444, 452639093245101445, 452639093245101446, 452639093245101447, 452639093245101448, 452639093245101449, 452639093245101450, 452639093245101451, 452639093245101452, 452639093245101453, 452639093245101454, 452639093245101455, 452639093245101456, 452639093245101457, 452639093245101458, 452639093245101459, 452639093245101460, 452639093245101461, 452639093245101462, 452639093245101463, 452639093245101464, 452639093245101465, 452639093245101466, 452639093245101467, 452639093245101468, 452639093245101469, 452639093245101470, 452639093245101471, 452639093245101472, 452639093245101473, 452639093245101474, 452639093245101475, 452639093245101476, 452639093245101477, 452639093245101478, 452639093245101479, 452639093245101480, 452639093245101481, 452639093245101482, 452639093245101483, 452639093245101484, 452639093245101485, 452639093245101486, 452639093245101487, 452639093245101488, 452639093245101489, 452639093245101490, 452639093245101491, 452639093245101492, 452639093245101493, 452639093245101494, 452639093245101495, 452639093245101496, 452639093245101497, 452639093245101498, 452639093245101499, 452639093245101500, 452639093245101501, 452639093245101502, 452639093245101503, 452639093245101504, 452639093245101505, 452639093245101506, 452639093245101507, 452639093245101508, 452639093245101509, 452639093245101510, 452639093245101511, 452639093245101512, 452639093245101513, 452639093245101514, 452639093245101515, 452639093245101516, 452639093245101517, 452639093245101518, 452639093245101519, 452639093245101520, 452639093245101521, 452639093245101522, 452639093245101523, 452639093245101524, 452639093245101525, 452639093245101526, 452639093245101527, 452639093245101528, 452639093245101529, 452639093245101530, 452639093245101531, 452639093245101532, 452639093245101533, 452639093245101534, 452639093245101535, 452639093245101536, 452639093245101537, 452639093245101538, 452639093245101539, 452639093245101540, 452639093245101541, 452639093245101542, 452639093245101543, 452639093245101544, 452639093245101545, 452639093245101546, 452639093245101547, 452639093245101548, 452639093245101549, 452639093245101550, 452639093245101551, 452639093245101552, 452639093245101553, 452639093245101554, 452639093245101555, 452639093245101556, 452639093245101557, 452639093245101558, 452639093245101559, 452639093245101560, 452639093245101561, 452639093245101562, 452639093245101563, 452639093245101564, 452639093245101565, 452639093245101566, 452639093245101567, 452639093245101568, 452639093245101569, 452639093245101570, 452639093245101571, 452639093245101572, 452639093245101573, 452639093245101574, 452639093245101575, 452639093245101576, 452639093245101577, 452639093245101578, 452639093245101579, 452639093245101580, 452639093245101581, 452639093245101582, 452639093245101583, 452639093245101584, 452639093245101585, 452639093245101586, 452639093245101587, 452639093245101588, 452639093245101589, 452639093245101590, 452639093245101591, 452639093245101592, 452639093245101593, 452639093245101594, 452639093245101595, 452639093245101596, 452639093245101597, 452639093245101598, 452639093245101599, 452639093245101600, 452639093245101601, 452639093245101602, 452639093245101603, 452639093245101604, 452639093245101605, 452639093245101606, 452639093245101607, 452639093245101608, 452639093245101609, 452639093245101610, 452639093245101611, 452639093245101612, 452639093245101613, 452639093245101614, 452639093245101615, 452639093245101616, 452639093245101617, 452639093245101618, 452639093245101619, 452639093245101620, 452639093245101621, 452639093245101622, 452639093245101623, 452639093245101624, 452639093245101625, 452639093245101626, 452639093245101627, 452639093245101628, 452639093245101629, 452639093245101630, 452639093245101631, 452639093245101632, 452639093245101633, 452639093245101634, 452639093245101635, 452639093245101636, 452639093245101637, 452639093245101638, 452639093245101639, 452639093245101640, 452639093245101641, 452639093245101642, 452639093245101643, 452639093245101644, 452639093245101645, 452639093245101646, 452639093245101647, 452639093245101648, 452639093245101649, 452639093245101650, 452639093245101651, 452639093245101652, 452639093245101653, 452639093245101654, 452639093245101655, 452639093245101656, 452639093245101657, 452639093245101658, 452639093245101659, 452639093245101660, 452639093245101661, 452639093245101662, 452639093245101663, 452639093245101664, 452639093245101665, 452639093245101666, 452639093245101667, 452639093245101668, 452639093245101669, 452639093245101670, 452639093245101671, 452639093245101672, 452639093245101673, 452639093245101674, 452639093245101675, 452639093245101676, 452639093245101677, 452639093245101678, 452639093245101679, 452639093245101680, 452639093245101681, 452639093245101682, 452639093245101683, 452639093245101684, 452639093245101685, 452639093245101686, 452639093245101687, 452639093245101688, 452639093245101689, 452639093245101690, 452639093245101691, 452639093245101692, 452639093245101693, 452639093245101694, 452639093245101695, 452639093245101696, 452639093245101697, 452639093245101698, 452639093245101699, 452639093245101700, 452639093245101701, 452639093245101702, 452639093245101703, 452639093245101704, 452639093245101705, 452639093245101706, 452639093245101707, 452639093245101708, 452639093245101709, 452639093245101710, 452639093245101711, 452639093245101712, 452639093245101713, 452639093245101714, 452639093245101715, 452639093245101716, 452639093245101717, 452639093245101718, 452639093245101719, 452639093245101720, 452639093245101721, 452639093245101722, 452639093245101723, 452639093245101724, 452639093245101725, 452639093245101726, 452639093245101727, 452639093245101728, 452639093245101729, 452639093245101730, 452639093245101731, 452639093245101732, 452639093245101733, 452639093245101734, 452639093245101735, 452639093245101736, 452639093245101737, 452639093245101738, 452639093245101739, 452639093245101740, 452639093245101741, 452639093245101742, 452639093245101743, 452639093245101744, 452639093245101745, 452639093245101746, 452639093245101747, 452639093245101748, 452639093245101749, 452639093245101750, 452639093245101751, 452639093245101752, 452639093245101753, 452639093245101754, 452639093245101755, 452639093245101756, 452639093245101757, 452639093245101758, 452639093245101759, 452639093245101760, 452639093245101761, 452639093245101762, 452639093245101763, 452639093245101764, 452639093245101765, 452639093245101766, 452639093245101767, 452639093245101768, 452639093245101769, 452639093245101770, 452639093245101771, 452639093245101772, 452639093245101773, 452639093245101774, 452639093245101775, 452639093245101776, 452639093245101777, 452639093245101778, 452639093245101779, 452639093245101780, 452639093245101781, 452639093245101782, 452639093245101783, 452639093245101784, 452639093245101785, 452639093245101786, 452639093245101787, 452639093245101788, 452639093245101789, 452639093245101790, 452639093245101791, 452639093245101792, 452639093245101793, 452639093245101794, 452639093245101795, 452639093245101796, 452639093245101797, 452639093245101798, 452639093245101799, 452639093245101800, 452639093245101801, 452639093245101802, 452639093245101803, 452639093245101804, 452639093245101805, 452639093245101806, 452639093245101807, 452639093245101808, 452639093245101809, 452639093245101810, 452639093245101811, 452639093245101812, 452639093245101813, 452639093245101814, 452639093245101815, 452639093245101816, 452639093245101817, 452639093245101818, 452639093245101819, 452639093245101820, 452639093245101821, 452639093245101822, 452639093245101823, 452639093245101824, 452639093245101825, 452639093245101826, 452639093245101827, 452639093245101828, 452639093245101829, 452639093245101830, 452639093245101831, 452639093245101832, 452639093245101833, 452639093245101834, 452639093245101835, 452639093245101836, 452639093245101837, 452639093245101838, 452639093245101839, 452639093245101840, 452639093245101841, 452639093245101842, 452639093245101843, 452639093245101844, 452639093245101845, 452639093245101846, 452639093245101847, 452639093245101848, 452639093245101849, 452639093245101850, 452639093245101851, 452639093245101852, 452639093245101853, 452639093245101854, 452639093245101855, 452639093245101856, 452639093245101857, 452639093245101858, 452639093245101859, 452639093245101860, 452639093245101861, 452639093245101862, 452639093245101863, 452639093245101864, 452639093245101865, 452639093245101866, 452639093245101867, 452639093245101868, 452639093245101869, 452639093245101870, 452639093245101871, 452639093245101872, 452639093245101873, 452639093245101874, 452639093245101875, 452639093245101876, 452639093245101877, 452639093245101878, 452639093245101879, 452639093245101880, 452639093245101881, 452639093245101882, 452639093245101883, 452639093245101884, 452639093245101885, 452639093245101886, 452639093245101887, 452639093245101888, 452639093245101889, 452639093245101890, 452639093245101891, 452639093245101892, 452639093245101893, 452639093245101894, 452639093245101895, 452639093245101896, 452639093245101897, 452639093245101898, 452639093245101899], 'cost': 0}"
            ]
          },
          "execution_count": 18,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "dim = len(list(image_dict.values())[0])\n",
        "\n",
        "# Create Milvus Collection\n",
        "# By default, vector field name is \"vector\"\n",
        "milvus_client.create_collection(\n",
        "    collection_name=COLLECTION_NAME,\n",
        "    auto_id=True,\n",
        "    dimension=dim,\n",
        "    enable_dynamic_field=True,\n",
        ")\n",
        "\n",
        "# Insert data into collection\n",
        "milvus_client.insert(\n",
        "    collection_name=COLLECTION_NAME,\n",
        "    data=[{\"image_path\": k, \"vector\": v} for k, v in image_dict.items()],\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sbtVftsp2kid"
      },
      "source": [
        "## Multimodal Search with Generative Reranker\n",
        "\n",
        "In this section, we will firstly search for relevant images by a multimodal query and then use LLM service to rerank the results and find the best one with explanation.\n",
        "\n",
        "### Run search\n",
        "\n",
        "Now we are ready to perform the advanced image search with query data composed of both image and text instruction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3OfqGION2kid",
        "outputId": "88c2903f-df71-413d-fa47-a320e65050cf"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "['./images_folder/images/518Gj1WQ-RL._AC_.jpg', './images_folder/images/41n00AOfWhL._AC_.jpg', './images_folder/images/51Wqge9HySL._AC_.jpg', './images_folder/images/51R2SZiywnL._AC_.jpg', './images_folder/images/516PebbMAcL._AC_.jpg', './images_folder/images/51RrgfYKUfL._AC_.jpg', './images_folder/images/51INtcXu9FL._AC_.jpg', './images_folder/images/41c0Jh8Mj7L._AC_.jpg', './images_folder/images/515DzQVKKwL._AC_.jpg']\n"
          ]
        }
      ],
      "source": [
        "import os\n",
        "query_image = os.path.join(\n",
        "    data_dir, \"leopard.jpg\"\n",
        ")  # Change to your own query image path\n",
        "query_text = \"phone case with this image theme\"\n",
        "\n",
        "# Generate query embedding given image and text instructions\n",
        "query_vec = encoder.encode_query(image_path=query_image, text=query_text)\n",
        "\n",
        "search_results = milvus_client.search(\n",
        "    collection_name=COLLECTION_NAME,\n",
        "    data=[query_vec],\n",
        "    output_fields=[\"image_path\"],\n",
        "    limit=9,  # Max number of search results to return\n",
        "    search_params={\"metric_type\": \"COSINE\", \"params\": {}},  # Search parameters\n",
        ")[0]\n",
        "\n",
        "retrieved_images = [hit.get(\"entity\").get(\"image_path\") for hit in search_results]\n",
        "print(retrieved_images)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QQR6jHR02kid"
      },
      "source": [
        "### Rerank with GPT-4o\n",
        "\n",
        "We will use an LLM to rank images and generate an explanation for the best result based on the user query and retrieved results.\n",
        "\n",
        "**1. Create a panoramic view**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oLTiiyBX2kid"
      },
      "outputs": [],
      "source": [
        "img_height = 300\n",
        "img_width = 300\n",
        "row_count = 3\n",
        "\n",
        "\n",
        "def create_panoramic_view(query_image_path: str, retrieved_images: list) -> np.ndarray:\n",
        "    \"\"\"\n",
        "    creates a 5x5 panoramic view image from a list of images\n",
        "\n",
        "    args:\n",
        "        images: list of images to be combined\n",
        "\n",
        "    returns:\n",
        "        np.ndarray: the panoramic view image\n",
        "    \"\"\"\n",
        "    panoramic_width = img_width * row_count\n",
        "    panoramic_height = img_height * row_count\n",
        "    panoramic_image = np.full(\n",
        "        (panoramic_height, panoramic_width, 3), 255, dtype=np.uint8\n",
        "    )\n",
        "\n",
        "    # create and resize the query image with a blue border\n",
        "    query_image_null = np.full((panoramic_height, img_width, 3), 255, dtype=np.uint8)\n",
        "    query_image = Image.open(query_image_path).convert(\"RGB\")\n",
        "    query_array = np.array(query_image)[:, :, ::-1]\n",
        "    resized_image = cv2.resize(query_array, (img_width, img_height))\n",
        "\n",
        "    border_size = 10\n",
        "    blue = (255, 0, 0)  # blue color in BGR\n",
        "    bordered_query_image = cv2.copyMakeBorder(\n",
        "        resized_image,\n",
        "        border_size,\n",
        "        border_size,\n",
        "        border_size,\n",
        "        border_size,\n",
        "        cv2.BORDER_CONSTANT,\n",
        "        value=blue,\n",
        "    )\n",
        "\n",
        "    query_image_null[img_height * 2 : img_height * 3, 0:img_width] = cv2.resize(\n",
        "        bordered_query_image, (img_width, img_height)\n",
        "    )\n",
        "\n",
        "    # add text \"query\" below the query image\n",
        "    text = \"query\"\n",
        "    font_scale = 1\n",
        "    font_thickness = 2\n",
        "    text_org = (10, img_height * 3 + 30)\n",
        "    cv2.putText(\n",
        "        query_image_null,\n",
        "        text,\n",
        "        text_org,\n",
        "        cv2.FONT_HERSHEY_SIMPLEX,\n",
        "        font_scale,\n",
        "        blue,\n",
        "        font_thickness,\n",
        "        cv2.LINE_AA,\n",
        "    )\n",
        "\n",
        "    # combine the rest of the images into the panoramic view\n",
        "    retrieved_imgs = [\n",
        "        np.array(Image.open(img).convert(\"RGB\"))[:, :, ::-1] for img in retrieved_images\n",
        "    ]\n",
        "    for i, image in enumerate(retrieved_imgs):\n",
        "        image = cv2.resize(image, (img_width - 4, img_height - 4))\n",
        "        row = i // row_count\n",
        "        col = i % row_count\n",
        "        start_row = row * img_height\n",
        "        start_col = col * img_width\n",
        "\n",
        "        border_size = 2\n",
        "        bordered_image = cv2.copyMakeBorder(\n",
        "            image,\n",
        "            border_size,\n",
        "            border_size,\n",
        "            border_size,\n",
        "            border_size,\n",
        "            cv2.BORDER_CONSTANT,\n",
        "            value=(0, 0, 0),\n",
        "        )\n",
        "        panoramic_image[\n",
        "            start_row : start_row + img_height, start_col : start_col + img_width\n",
        "        ] = bordered_image\n",
        "\n",
        "        # add red index numbers to each image\n",
        "        text = str(i)\n",
        "        org = (start_col + 50, start_row + 30)\n",
        "        (font_width, font_height), baseline = cv2.getTextSize(\n",
        "            text, cv2.FONT_HERSHEY_SIMPLEX, 1, 2\n",
        "        )\n",
        "\n",
        "        top_left = (org[0] - 48, start_row + 2)\n",
        "        bottom_right = (org[0] - 48 + font_width + 5, org[1] + baseline + 5)\n",
        "\n",
        "        cv2.rectangle(\n",
        "            panoramic_image, top_left, bottom_right, (255, 255, 255), cv2.FILLED\n",
        "        )\n",
        "        cv2.putText(\n",
        "            panoramic_image,\n",
        "            text,\n",
        "            (start_col + 10, start_row + 30),\n",
        "            cv2.FONT_HERSHEY_SIMPLEX,\n",
        "            1,\n",
        "            (0, 0, 255),\n",
        "            2,\n",
        "            cv2.LINE_AA,\n",
        "        )\n",
        "\n",
        "    # combine the query image with the panoramic view\n",
        "    panoramic_image = np.hstack([query_image_null, panoramic_image])\n",
        "    return panoramic_image"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s39-dqkT2kie"
      },
      "source": [
        "Combine the query image and retrieved images with indices in a panoramic view."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5xF5Hg0I2kie"
      },
      "outputs": [],
      "source": [
        "from PIL import Image\n",
        "\n",
        "combined_image_path = os.path.join(data_dir, \"combined_image.jpg\")\n",
        "panoramic_image = create_panoramic_view(query_image, retrieved_images)\n",
        "cv2.imwrite(combined_image_path, panoramic_image)\n",
        "\n",
        "combined_image = Image.open(combined_image_path)\n",
        "show_combined_image = combined_image.resize((300, 300))\n",
        "show_combined_image.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SUejRTtf2kif"
      },
      "source": [
        "**2. Rerank and explain**\n",
        "\n",
        "We will send the combined image to multimodal LLM service together with proper prompts to rank the retrieved results with explanation. To enable GPT-4o as the LLM, you need to prepare your [OpenAI API Key](https://platform.openai.com/docs/quickstart)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vzyEQdpu2kif"
      },
      "outputs": [],
      "source": [
        "import requests\n",
        "import base64\n",
        "\n",
        "\n",
        "openai_api_key = api_key  # Change to your OpenAI API Key\n",
        "\n",
        "\n",
        "def generate_ranking_explanation(\n",
        "    combined_image_path: str, caption: str, infos: dict = None\n",
        ") -> tuple[list[int], str]:\n",
        "    with open(combined_image_path, \"rb\") as image_file:\n",
        "        base64_image = base64.b64encode(image_file.read()).decode(\"utf-8\")\n",
        "\n",
        "    information = (\n",
        "        \"You are responsible for ranking results for a Composed Image Retrieval. \"\n",
        "        \"The user retrieves an image with an 'instruction' indicating their retrieval intent. \"\n",
        "        \"For example, if the user queries a red car with the instruction 'change this car to blue,' a similar type of car in blue would be ranked higher in the results. \"\n",
        "        \"Now you would receive instruction and query image with blue border. Every item has its red index number in its top left. Do not misunderstand it. \"\n",
        "        f\"User instruction: {caption} \\n\\n\"\n",
        "    )\n",
        "\n",
        "    # add additional information for each image\n",
        "    if infos:\n",
        "        for i, info in enumerate(infos[\"product\"]):\n",
        "            information += f\"{i}. {info}\\n\"\n",
        "\n",
        "    information += (\n",
        "        \"Provide a new ranked list of indices from most suitable to least suitable, followed by an explanation for the top 1 most suitable item only. \"\n",
        "        \"The format of the response has to be 'Ranked list: []' with the indices in brackets as integers, followed by 'Reasons:' plus the explanation why this most fit user's query intent.\"\n",
        "    )\n",
        "\n",
        "    headers = {\n",
        "        \"Content-Type\": \"application/json\",\n",
        "        \"Authorization\": f\"Bearer {openai_api_key}\",\n",
        "    }\n",
        "\n",
        "    payload = {\n",
        "        \"model\": \"gpt-4o\",\n",
        "        \"messages\": [\n",
        "            {\n",
        "                \"role\": \"user\",\n",
        "                \"content\": [\n",
        "                    {\"type\": \"text\", \"text\": information},\n",
        "                    {\n",
        "                        \"type\": \"image_url\",\n",
        "                        \"image_url\": {\"url\": f\"data:image/jpeg;base64,{base64_image}\"},\n",
        "                    },\n",
        "                ],\n",
        "            }\n",
        "        ],\n",
        "        \"max_tokens\": 300,\n",
        "    }\n",
        "\n",
        "    response = requests.post(\"https://api.openai.com/v1/chat/completions\", headers=headers, json=payload)\n",
        "    result = response.json()[\"choices\"][0][\"message\"][\"content\"]\n",
        "\n",
        "\n",
        "    # parse the ranked indices from the response\n",
        "    start_idx = result.find(\"[\")\n",
        "    end_idx = result.find(\"]\")\n",
        "    ranked_indices_str = result[start_idx + 1 : end_idx].split(\",\")\n",
        "    ranked_indices = [int(index.strip()) for index in ranked_indices_str]\n",
        "\n",
        "    # extract explanation\n",
        "    explanation = result[end_idx + 1 :].strip()\n",
        "\n",
        "    return ranked_indices, explanation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lcb6SttP2kif"
      },
      "source": [
        "Get the image indices after ranking and the reason for the best result:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QoZzRAKO2kif"
      },
      "outputs": [],
      "source": [
        "ranked_indices, explanation = generate_ranking_explanation(\n",
        "    combined_image_path, query_text\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hvqRekfD2kig"
      },
      "source": [
        "**3. Display the best result with explanation**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MHHjGqTJ2kig",
        "outputId": "c5860f58-ad1c-4dc0-9cc0-61593366e579"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Reasons: The top 1 most suitable item is index 8. This phone case features a leopard print design, which directly corresponds to the leopard theme in the user's query image. The pattern and color closely resemble the animal, making it the most suitable match for the user's desired phone case theme.\n"
          ]
        }
      ],
      "source": [
        "print(explanation)\n",
        "\n",
        "best_index = ranked_indices[0]\n",
        "best_img = Image.open(retrieved_images[best_index])\n",
        "best_img = best_img.resize((150, 150))\n",
        "best_img.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-M3hKRis2kig"
      },
      "source": [
        "### Quick Deploy\n",
        "\n",
        "To learn about how to start an online demo with this tutorial, please refer to [the example application](https://github.com/milvus-io/bootcamp/tree/master/bootcamp/tutorials/quickstart/apps/multimodal_rag_with_milvus)."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "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.9.19"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}