{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ijGzTHJJUCPY"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NDsTUvKjwHBW"
      },
      "source": [
        "# Building a DIY Multimodal Question Answering System with Vertex AI (A Beginner's Guide - Multimodal RAG)\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/workshops/qa-ops/building_diy_multimodal_qa_system_with_mrag.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fworkshops%2Fqa-ops%2Fbuilding_diy_multimodal_qa_system_with_mrag.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Run in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/workshops/qa-ops/building_diy_multimodal_qa_system_with_mrag.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/workshops/qa-ops/building_diy_multimodal_qa_system_with_mrag.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>    \n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/workshops/qa-ops/building_diy_multimodal_qa_system_with_mrag.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/workshops/qa-ops/building_diy_multimodal_qa_system_with_mrag.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/workshops/qa-ops/building_diy_multimodal_qa_system_with_mrag.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/workshops/qa-ops/building_diy_multimodal_qa_system_with_mrag.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/workshops/qa-ops/building_diy_multimodal_qa_system_with_mrag.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>            "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YR65ni7TRNYG"
      },
      "source": [
        "| | |\n",
        "|-|-|\n",
        "|Author(s) | [Lavi Nigam](https://github.com/lavinigam-gcp) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QduB55tHOa8N"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ This is a new version of the old mRAG notebook with modifications and new data. You can refer to the old notebook here:  ⚠️</b>\n",
        "</div>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "70SQT5lKO9dp"
      },
      "source": [
        "[**intro_multimodal_rag.ipynb**](https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/retrieval-augmented-generation/intro_multimodal_rag.ipynb)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VK1Q5ZYdVL4Y"
      },
      "source": [
        "## Overview\n",
        "\n",
        "\n",
        "This guide is your hands-on introduction to creating a question answering system that understands both text and images. We'll build this system from the ground up using Google's Vertex AI, giving you a clear understanding of how it works without relying on complex third-party tools.\n",
        "\n",
        "\n",
        "## Why Build It Yourself?\n",
        "\n",
        "Large Language Models (LLMs) are powerful, but they can seem like a \"black box\". By building our own system, we'll break open that box and explore the core concepts. This will give you the knowledge to customize and optimize every aspect of your question answering system, whether you ultimately choose to code everything yourself or use external libraries.\n",
        "\n",
        "\n",
        "## What We'll Do:\n",
        "\n",
        "* **Focus on Fundamentals**: We'll start with the essential design pattern of \"Retrieval Augmented Generation\" (RAG) – a way to find and use relevant information to answer questions.\n",
        "\n",
        "* **Work with Text and Images**: We'll expand RAG to handle both text and images found in PDF documents. Future guides in this series will explore even more types of data, like videos and audio.\n",
        "\n",
        "* **Use Vertex AI**: We'll only use Google's Vertex AI Embeddings API and Gemini API, ensuring you have complete control and understanding of the building blocks.\n",
        "\n",
        "\n",
        "By the end of this guide, you'll have a solid foundation in building multimodal question answering systems, empowering you to create smarter applications that can understand and respond to a wider range of information.\n",
        "\n",
        "\n",
        "### Gemini\n",
        "\n",
        "Gemini is a family of generative AI models developed by Google DeepMind that is designed for multimodal use cases. The Gemini API gives you access to the Gemini models.\n",
        "\n",
        "### Comparing text-based and multimodal RAG\n",
        "\n",
        "Multimodal RAG offers several advantages over text-based RAG:\n",
        "\n",
        "1. **Enhanced knowledge access:** Multimodal RAG can access and process both textual and visual information, providing a richer and more comprehensive knowledge base for the LLM.\n",
        "2. **Improved reasoning capabilities:** By incorporating visual cues, multimodal RAG can make better informed inferences across different types of data modalities.\n",
        "\n",
        "This notebook shows you how to implement DIY RAG with Gemini API in Vertex AI\n",
        " and Vertex AI Embeddings API; [text embeddings](https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/text-embeddings), and [multimodal embeddings](https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/multimodal-embeddings), to build a document search engine.\n",
        "\n",
        "Through hands-on examples, you will discover how to construct a multimedia-rich metadata repository of your document sources, enabling search, comparison, and reasoning across diverse information streams."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RQT500QqVPIb"
      },
      "source": [
        "### Objectives\n",
        "\n",
        "This notebook provides a guide to building a document search engine using multimodal retrieval augmented generation (RAG), step by step:\n",
        "\n",
        "1. Extract and store metadata of documents containing both text and images, and generate embeddings the documents\n",
        "2. Search the metadata with text queries to find similar text or images\n",
        "3. Search the metadata with image queries to find similar images\n",
        "4. Using a text query as input, search for contextual answers using both text and images"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KnpYxfesh2rI"
      },
      "source": [
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "- Vertex AI\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DXJpXzKrh2rJ"
      },
      "source": [
        "## Getting Started\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N5afkyDMSBW5"
      },
      "source": [
        "### Install Vertex AI SDK for Python and other dependencies\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kc4WxYmLSBW5"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --user google-cloud-aiplatform pymupdf rich"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R5Xep4W9lq-Z"
      },
      "source": [
        "### Restart current runtime\n",
        "\n",
        "To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which will restart the current kernel."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
        "# Restart kernel after installs so that your environment can access the new packages\n",
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SbmM4z7FOBpM"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Please wait until it is finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FtsU9Bw9h2rL"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you are running this notebook on Google Colab, run the following cell to authenticate your environment. This step is not required if you are using [Vertex AI Workbench](https://cloud.google.com/vertex-ai-workbench).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GpYEyLsOh2rL"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "# Additional authentication is required for Google Colab\n",
        "if \"google.colab\" in sys.modules:\n",
        "    # Authenticate user to Google Cloud\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O1vKZZoEh2rL"
      },
      "source": [
        "### Define Google Cloud project information\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gJqZ76rJh2rM"
      },
      "outputs": [],
      "source": [
        "# Define project information\n",
        "\n",
        "import sys\n",
        "\n",
        "PROJECT_ID = \"YOUR_PROJECT_ID\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
        "\n",
        "# if not running on Colab, try to get the PROJECT_ID automatically\n",
        "if \"google.colab\" not in sys.modules:\n",
        "    import subprocess\n",
        "\n",
        "    PROJECT_ID = subprocess.check_output(\n",
        "        [\"gcloud\", \"config\", \"get-value\", \"project\"], text=True\n",
        "    ).strip()\n",
        "\n",
        "print(f\"Your project ID is: {PROJECT_ID}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D48gUW5-h2rM"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "# Initialize Vertex AI\n",
        "import vertexai\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BuQwwRiniVFG"
      },
      "source": [
        "### Import libraries\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rtMowvm-yQ97"
      },
      "outputs": [],
      "source": [
        "from IPython.display import display\n",
        "from rich import print as rich_print\n",
        "from rich.markdown import Markdown as rich_Markdown\n",
        "from vertexai.generative_models import (\n",
        "    GenerationConfig,\n",
        "    GenerativeModel,\n",
        "    HarmBlockThreshold,\n",
        "    HarmCategory,\n",
        "    Image,\n",
        ")\n",
        "from vertexai.language_models import TextEmbeddingModel\n",
        "from vertexai.vision_models import MultiModalEmbeddingModel"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r-TX_R_xh2rM"
      },
      "source": [
        "### Load the Gemini model\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HySgZekYzCpW"
      },
      "source": [
        "Learn more about each models and their differences: [here](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/send-multimodal-prompts)\n",
        "\n",
        "Learn about the quotas: [here](https://cloud.google.com/vertex-ai/generative-ai/docs/quotas)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SvMwSRJJh2rM"
      },
      "outputs": [],
      "source": [
        "# Multimodal models: Choose based on your performance/cost needs\n",
        "multimodal_model_2 = GenerativeModel(\n",
        "    \"gemini-2.0-flash\"\n",
        ")  # works with text, code, images, video(with or without audio) and audio(mp3) with 1M input context - complex reasoning\n",
        "\n",
        "# Load text embedding model from pre-trained source\n",
        "text_embedding_model = TextEmbeddingModel.from_pretrained(\"text-embedding-005\")\n",
        "\n",
        "# Load multimodal embedding model from pre-trained source\n",
        "multimodal_embedding_model = MultiModalEmbeddingModel.from_pretrained(\n",
        "    \"multimodalembedding\"\n",
        ")  # works with image, image with caption(~32 words), video, video with caption(~32 words)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g7bKCQMFT7JT"
      },
      "source": [
        "#### Get documents and images from GCS"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KwbL89zcY39N"
      },
      "outputs": [],
      "source": [
        "# download documents and images used in this notebook - will take ~30 sec\n",
        "!gsutil -m -q rsync -r gs://github-repo/rag/intro_multimodal_rag/intro_multimodal_rag_v2 .\n",
        "print(\"Download completed\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ps1G-cCfpibN"
      },
      "source": [
        "## Building metadata of documents containing text and images"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V7uv_PVR1T6B"
      },
      "source": [
        "### The data\n",
        "\n",
        "The source data that you will use in this notebook are:\n",
        "\n",
        "\n",
        "* [Google Cloud TPU Scaling blog](https://storage.googleapis.com/github-repo/rag/intro_multimodal_rag/data/Google%20Cloud%20TPU%20blog.pdf)\n",
        "* [Gemini 2.0 Technical Report](https://storage.googleapis.com/github-repo/rag/intro_multimodal_rag/data/gemini_v1_5_report_technical.pdf)\n",
        "* [Google Gemma Technical Paper](https://storage.googleapis.com/github-repo/rag/intro_multimodal_rag/data/gemma_technical_paper.pdf)\n",
        "* [Med-Gemini Technical Paper](https://storage.googleapis.com/github-repo/rag/intro_multimodal_rag/data/med_gemini.pdf)\n",
        "\n",
        "\n",
        "\n",
        "You can also use your data, by first deleting the current files and then placing your files in the `data/` folder."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zvt0sus5KSNX"
      },
      "source": [
        "### Import helper functions to build metadata\n",
        "\n",
        "Before building the Multimodal Question Answering System with Vertex AI, it's important to have metadata of all the text and images in the document. For references and citations purposes, the metadata should contain essential elements, including page number, file name, image counter, and so on. Hence, as a next step, you will generate embeddings from the metadata, which is required to perform similarity search when querying the data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tStqXX32RNYK"
      },
      "outputs": [],
      "source": [
        "from multimodal_qa_with_rag_utils import get_document_metadata, set_global_variable\n",
        "\n",
        "set_global_variable(\"text_embedding_model\", text_embedding_model)\n",
        "set_global_variable(\"multimodal_embedding_model\", multimodal_embedding_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WtaDuBXhFmkL"
      },
      "source": [
        " You can also view the code (`multimodal_qa_with_rag_utils`) [directly](https://storage.googleapis.com/github-repo/rag/intro_multimodal_rag/utils/multimodal_qa_with_rag_utils.py)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5BOAkYN0KlSL"
      },
      "source": [
        "### Extract and store metadata of text and images from a document"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q9hBPPWs5CMd"
      },
      "source": [
        "You just imported a function called `get_document_metadata()`. This function extracts text and image metadata from a document, and returns two dataframes, namely *text_metadata* and *image_metadata*, as outputs. If you want to find out more about how `get_document_metadata()` function is implemented using Gemini and the embedding models, you can take look at the [source code](https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/use-cases/retrieval-augmented-generation/utils/intro_multimodal_rag_utils.py) directly.\n",
        "\n",
        "The reason for extraction and storing both text metadata and image metadata is that just by using either of the two alone is not sufficient to come out with a relevent answer. For example, the relevant answers could be in visual form within a document, but text-based RAG won't be able to take into consideration of the visual images. You will also be exploring this example later in this notebook.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PnKru0sBh2rN"
      },
      "source": [
        "At the next step, you will use the function to extract and store metadata of text and images from a document. Please note that the following cell may take a few minutes to complete:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x8h0XSG_7e5M"
      },
      "source": [
        "**NOTE: Given that we are loading 4 files with roughly 200 pages and approximately 84 images, the cell below will take approximately 7 minutes to run. We recommend loading pre-computed metadata instead.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "X8hE0tWD-lf8"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "# Specify the PDF folder with multiple PDF ~7m\n",
        "\n",
        "print(\n",
        "    \"Removing pre-existing images folder, since you are running the logic from scratch\"\n",
        ")\n",
        "! rm -rf images/\n",
        "\n",
        "pdf_folder_path = \"data/\"  # if running in Vertex AI Workbench.\n",
        "\n",
        "# Specify the image description prompt. Change it\n",
        "# image_description_prompt = \"\"\"Explain what is going on in the image.\n",
        "# If it's a table, extract all elements of the table.\n",
        "# If it's a graph, explain the findings in the graph.\n",
        "# Do not include any numbers that are not mentioned in the image.\n",
        "# \"\"\"\n",
        "\n",
        "image_description_prompt = \"\"\"You are a technical image analysis expert. You will be provided with various types of images extracted from documents like research papers, technical blogs, and more.\n",
        "Your task is to generate concise, accurate descriptions of the images without adding any information you are not confident about.\n",
        "Focus on capturing the key details, trends, or relationships depicted in the image.\n",
        "\n",
        "Important Guidelines:\n",
        "* Prioritize accuracy:  If you are uncertain about any detail, state \"Unknown\" or \"Not visible\" instead of guessing.\n",
        "* Avoid hallucinations: Do not add information that is not directly supported by the image.\n",
        "* Be specific: Use precise language to describe shapes, colors, textures, and any interactions depicted.\n",
        "* Consider context: If the image is a screenshot or contains text, incorporate that information into your description.\n",
        "\"\"\"\n",
        "\n",
        "\n",
        "# Extract text and image metadata from the PDF document\n",
        "text_metadata_df, image_metadata_df = get_document_metadata(\n",
        "    multimodal_model_2,  # we are passing Gemini 2.0 Flash\n",
        "    pdf_folder_path,\n",
        "    image_save_dir=\"images\",\n",
        "    image_description_prompt=image_description_prompt,\n",
        "    embedding_size=1408,\n",
        "    # add_sleep_after_page = True, # Uncomment this if you are running into API quota issues\n",
        "    # sleep_time_after_page = 5,\n",
        "    add_sleep_after_document=True,  # Uncomment this if you are running into API quota issues\n",
        "    sleep_time_after_document=5,  # Increase the value in seconds, if you are still getting quota issues. It will slow down the processing.\n",
        "    # generation_config = # see next cell\n",
        "    # safety_settings =  # see next cell\n",
        ")\n",
        "\n",
        "print(\"\\n\\n --- Completed processing. ---\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sWtOx1wb86Az"
      },
      "source": [
        "If you would like to pass additional parameters to Gemini while building metadata, here are some options:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QlNK0o2DRNYK"
      },
      "outputs": [],
      "source": [
        "# # Parameters for Gemini API call.\n",
        "# # reference for parameters: https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/gemini\n",
        "\n",
        "# generation_config=  GenerationConfig(temperature=0.2, max_output_tokens=2048)\n",
        "\n",
        "# # Set the safety settings if Gemini is blocking your content or you are facing \"ValueError(\"Content has no parts\")\" error or \"Exception occurred\" in your data.\n",
        "# # ref for settings and thresholds: https://cloud.google.com/vertex-ai/docs/generative-ai/multimodal/configure-safety-attributes\n",
        "\n",
        "# safety_settings = {\n",
        "#                   HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,\n",
        "#                   HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,\n",
        "#                   HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,\n",
        "#                   HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,\n",
        "#                   }\n",
        "\n",
        "# # You can also pass parameters and safety_setting to \"get_gemini_response\" function"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tW3Ci1IL8wSW"
      },
      "source": [
        "### Load pre-computed metadata of text and images from source document"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c10dCy6Cig3H"
      },
      "source": [
        "**If you are facing constant issues with Quota or want to focus on the outputs, you should load pre-computed metadata.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "w1AGYOYb0In7"
      },
      "outputs": [],
      "source": [
        "# import pickle\n",
        "\n",
        "# # Load the pickle file\n",
        "# with open(\"mrag_metadata.pkl\", \"rb\") as f:\n",
        "#     data = pickle.load(f)\n",
        "\n",
        "# # Extract the DataFrames\n",
        "# text_metadata_df = data[\"text_metadata\"]\n",
        "# image_metadata_df = data[\"image_metadata\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "miBBoEXwh2rN"
      },
      "source": [
        "#### Inspect the processed text metadata\n",
        "\n",
        "\n",
        "The following cell will produce a metadata table which describes the different parts of text metadata, including:\n",
        "\n",
        "- **text**: the original text from the page\n",
        "- **text_embedding_page**: the embedding of the original text from the page\n",
        "- **chunk_text**: the original text divided into smaller chunks\n",
        "- **chunk_number**: the index of each text chunk\n",
        "- **text_embedding_chunk**: the embedding of each text chunk"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6t3AIGFar8Mo"
      },
      "outputs": [],
      "source": [
        "text_metadata_df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NjIQYI3mh2rO"
      },
      "source": [
        "#### Inspect the processed image metadata\n",
        "\n",
        "The following cell will produce a metadata table which describes the different parts of image metadata, including:\n",
        "* **img_desc**: Gemini-generated textual description of the image.\n",
        "* **mm_embedding_from_text_desc_and_img**: Combined embedding of image and its description, capturing both visual and textual information.\n",
        "* **mm_embedding_from_img_only**: Image embedding without description, for comparison with description-based analysis.\n",
        "* **text_embedding_from_image_description**: Separate text embedding of the generated description, enabling textual analysis and comparison."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tkHtAYIK-y-q"
      },
      "outputs": [],
      "source": [
        "image_metadata_df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iBhoOkutUtPr"
      },
      "source": [
        "### Import the helper functions to implement RAG\n",
        "\n",
        "You will be importing the following functions which will be used in the remainder of this notebook to implement RAG:\n",
        "\n",
        "* **get_similar_text_from_query():** Given a text query, finds text from the document which are relevant, using cosine similarity algorithm. It uses text embeddings from the metadata to compute and the results can be filtered by top score, page/chunk number, or embedding size.\n",
        "* **print_text_to_text_citation():** Prints the source (citation) and details of the retrieved text from the `get_similar_text_from_query()` function.\n",
        "* **get_similar_image_from_query():** Given an image path or an image, finds images from the document which are relevant. It uses image embeddings from the metadata.\n",
        "* **print_text_to_image_citation():** Prints the source (citation) and the details of retrieved images from the `get_similar_image_from_query()` function.\n",
        "* **get_gemini_response():** Interacts with a Gemini model to answer questions based on a combination of text and image inputs.\n",
        "* **display_images():**  Displays a series of images provided as paths or PIL Image objects."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Tngn_vrIKdE1"
      },
      "outputs": [],
      "source": [
        "from multimodal_qa_with_rag_utils import (\n",
        "    display_images,\n",
        "    get_answer_from_qa_system,\n",
        "    get_gemini_response,\n",
        "    get_similar_image_from_query,\n",
        "    get_similar_text_from_query,\n",
        "    print_text_to_image_citation,\n",
        "    print_text_to_text_citation,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c9jGEj6DY1Rj"
      },
      "source": [
        "Before implementing a Multimodal Question Answering System with Vertex AI, let's explore what you can achieve with just text or image embeddings. This will set the foundation for implementing a multimodal Retrieval Augmented Generation (RAG) system, which you will do later in this notebook.\n",
        "\n",
        "You can also use these essential elements together to build applications for multimodal use cases, extracting meaningful information from documents."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KHuLlEvSKFWt"
      },
      "source": [
        "## Text Search\n",
        "\n",
        "Let's start the search with a simple question and see if the simple text search using text embeddings can answer it. The expected answer is to show the value of basic and diluted net income per share of Google for different share types.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5mrFVhtCut7t"
      },
      "outputs": [],
      "source": [
        "query = \"What are various med-gemini medical benchmarks that shows its performance relative to other models?\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XWw7-AIar-S8"
      },
      "source": [
        "### Search similar text with text query"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eEzP6Yyv7N-G"
      },
      "outputs": [],
      "source": [
        "# Matching user text query with \"chunk_embedding\" to find relevant chunks.\n",
        "matching_results_text = get_similar_text_from_query(\n",
        "    query,\n",
        "    text_metadata_df,\n",
        "    column_name=\"text_embedding_chunk\",\n",
        "    top_n=3,\n",
        "    chunk_text=True,\n",
        ")\n",
        "\n",
        "# Print the matched text citations\n",
        "print_text_to_text_citation(\n",
        "    matching_results_text, print_top=True, chunk_text=True\n",
        ")  # print_top=False to see all text matches"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KY1J2sHr-N8f"
      },
      "source": [
        "### Get answer with text-RAG"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ORCistIdDWoE"
      },
      "outputs": [],
      "source": [
        "# All relevant text chunk found across documents based on user query\n",
        "context = \"\\n\".join(\n",
        "    [value[\"chunk_text\"] for key, value in matching_results_text.items()]\n",
        ")\n",
        "\n",
        "prompt = f\"\"\"Answer the question with the given context. If the specific answer is not in the context, please answer \"I don't know\".\n",
        "Question: {query}\n",
        "Context: {context}\n",
        "Answer:\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GKD-Ew8IM287"
      },
      "outputs": [],
      "source": [
        "safety_settings = {\n",
        "    HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,\n",
        "    HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,\n",
        "    HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,\n",
        "    HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,\n",
        "}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m8jyc5_SAwlF"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "# Generate response with Gemini\n",
        "print(\"\\n **** Result: ***** \\n\")\n",
        "\n",
        "rich_Markdown(\n",
        "    get_gemini_response(\n",
        "        multimodal_model_2,\n",
        "        model_input=prompt,\n",
        "        stream=True,\n",
        "        safety_settings=safety_settings,\n",
        "        generation_config=GenerationConfig(temperature=1, max_output_tokens=8192),\n",
        "    )\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0rKLLKt1An97"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "# Generate response with Gemini\n",
        "rich_Markdown(\n",
        "    get_gemini_response(\n",
        "        multimodal_model_2,\n",
        "        model_input=prompt,\n",
        "        stream=True,\n",
        "        safety_settings=safety_settings,\n",
        "        generation_config=GenerationConfig(temperature=0.1),\n",
        "    )\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uXm271jdD-Rl"
      },
      "source": [
        "### Search similar images with text query"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oPxwfyVrr9-G"
      },
      "source": [
        "Since plain text search and RAG didn't provide the detailed answer and the information may be visually represented in a table or another image format, you can use multimodal capability of Gemini model for the similar task.\n",
        "\n",
        "The goal here also is to find an image similar to the text query. You may also print the citations to verify."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nwXpqMSq4ppr"
      },
      "outputs": [],
      "source": [
        "query = \"What are various med-gemini medical benchmarks that shows its performance relative to other models?\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "knj4qQ4xni24"
      },
      "outputs": [],
      "source": [
        "matching_results_image = get_similar_image_from_query(\n",
        "    text_metadata_df,\n",
        "    image_metadata_df,\n",
        "    query=query,\n",
        "    column_name=\"text_embedding_from_image_description\",  # Use image description text embedding\n",
        "    image_emb=False,  # Use text embedding instead of image embedding\n",
        "    top_n=5,\n",
        "    embedding_size=1408,\n",
        ")\n",
        "\n",
        "# Markdown(print_text_to_image_citation(matching_results_image, print_top=True))\n",
        "print(\"\\n **** Result: ***** \\n\")\n",
        "\n",
        "# Display the top matching image\n",
        "display_images(\n",
        "    [\n",
        "        matching_results_image[0][\"img_path\"],\n",
        "        matching_results_image[1][\"img_path\"],\n",
        "        matching_results_image[2][\"img_path\"],\n",
        "        matching_results_image[3][\"img_path\"],\n",
        "    ],\n",
        "    resize_ratio=0.3,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HuRD1lZ8RNYP"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "\n",
        "print(\"\\n **** Result: ***** \\n\")\n",
        "\n",
        "instruction = f\"\"\"Answer the question and explain results with the given Image:\n",
        "Question: {query}\n",
        "Image:\n",
        "\"\"\"\n",
        "\n",
        "# Prepare the model input\n",
        "model_input = [\n",
        "    instruction,\n",
        "    # passing all matched images to Gemini\n",
        "    \"Image:\",\n",
        "    matching_results_image[0][\"image_object\"],\n",
        "    \"Description:\",\n",
        "    matching_results_image[0][\"image_description\"],\n",
        "    \"Image:\",\n",
        "    matching_results_image[1][\"image_object\"],\n",
        "    \"Description:\",\n",
        "    matching_results_image[1][\"image_description\"],\n",
        "    \"Image:\",\n",
        "    matching_results_image[2][\"image_object\"],\n",
        "    \"Description:\",\n",
        "    matching_results_image[2][\"image_description\"],\n",
        "    \"Image:\",\n",
        "    matching_results_image[3][\"image_object\"],\n",
        "    \"Description:\",\n",
        "    matching_results_image[3][\"image_description\"],\n",
        "]\n",
        "\n",
        "# Generate Gemini response with streaming output\n",
        "rich_Markdown(\n",
        "    get_gemini_response(\n",
        "        multimodal_model_2,\n",
        "        model_input=model_input,\n",
        "        stream=True,\n",
        "        safety_settings=safety_settings,\n",
        "        generation_config=GenerationConfig(temperature=1, max_output_tokens=8192),\n",
        "    )\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uykUaQvIRNYP"
      },
      "outputs": [],
      "source": [
        "## you can check the citations to probe further.\n",
        "## check the \"image description:\" which is a description extracted through Gemini which helped search our query.\n",
        "rich_print(print_text_to_image_citation(matching_results_image, print_top=True))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oDd9rE4NrRod"
      },
      "source": [
        "## Image Search"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pJL6ElyEy4mc"
      },
      "source": [
        "### Search similar image with image input [using multimodal image embeddings]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ReKjHleFxUu9"
      },
      "source": [
        "Imagine searching for images, but instead of typing words, you use an actual image as the clue.\n",
        "\n",
        "Think of it like searching with a mini-map instead of a written address.\n",
        "It's a different way to ask, \"Show me more stuff like this\".\n",
        "\n",
        "So, instead of typing \"various example of Gemini 2.0 long context\", you show a picture of that image and say, \"Find me more like this\"\n",
        "\n",
        "For demonstration purposes, we will only be finding similar images that show the various features of Gemini in a single document below. However, you can scale this design pattern to match (find relevant images) across multiple documents."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DJhhS5eZw7QI"
      },
      "outputs": [],
      "source": [
        "# You can find a similar image as per the images you have in the metadata.\n",
        "\n",
        "image_query_path = \"images/gemini_v1_5_report_technical.pdf_image_5_0_148.jpeg\"\n",
        "\n",
        "# Print a message indicating the input image\n",
        "print(\"***Input image from user:***\")\n",
        "\n",
        "# Display the input image\n",
        "Image.load_from_file(image_query_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3zBTtGChTmrd"
      },
      "source": [
        "You expect to find images that are similar in terms of \"long context prompts for Gemini 2.0\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nZcU7vZC-8vr"
      },
      "outputs": [],
      "source": [
        "# Search for Similar Images Based on Input Image and Image Embedding\n",
        "\n",
        "matching_results_image = get_similar_image_from_query(\n",
        "    text_metadata_df,\n",
        "    image_metadata_df,\n",
        "    query=query,  # Use query text for additional filtering (optional)\n",
        "    column_name=\"mm_embedding_from_img_only\",  # Use image embedding for similarity calculation\n",
        "    image_emb=True,\n",
        "    image_query_path=image_query_path,  # Use input image for similarity calculation\n",
        "    top_n=3,  # Retrieve top 3 matching images\n",
        "    embedding_size=1408,  # Use embedding size of 1408\n",
        ")\n",
        "\n",
        "print(\"\\n **** Result: ***** \\n\")\n",
        "\n",
        "# Display the Top Matching Image\n",
        "display(\n",
        "    matching_results_image[0][\"image_object\"]\n",
        ")  # Display the top matching image object (Pillow Image)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uhT17rke15XY"
      },
      "source": [
        "You can also print the citation to see what it has matched."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mksXQoezweg0"
      },
      "outputs": [],
      "source": [
        "# Display citation details for the top matching image\n",
        "print_text_to_image_citation(\n",
        "    matching_results_image, print_top=True\n",
        ")  # Print citation details for the top matching image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VJWnhDJwI-uO"
      },
      "outputs": [],
      "source": [
        "# Check Other Matched Images (Optional)\n",
        "# You can access the other two matched images using:\n",
        "\n",
        "print(\"---------------Matched Images------------------\\n\")\n",
        "display_images(\n",
        "    [\n",
        "        matching_results_image[0][\"img_path\"],\n",
        "        matching_results_image[1][\"img_path\"],\n",
        "        matching_results_image[2][\"img_path\"],\n",
        "    ],\n",
        "    resize_ratio=0.2,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JvwZIgD84CNc"
      },
      "source": [
        "The ability to identify similar text and images based on user input, using Gemini and embeddings, forms a crucial foundation for development of Multimodal Question Answering System with multimodal RAG design pattern, which you will explore in the coming sections."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lUnsv5Co6pJF"
      },
      "source": [
        "### Comparative reasoning"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1AFbqHiz5vvo"
      },
      "source": [
        "Next, let's apply what you have done so far in doing comparative reasoning.\n",
        "\n",
        "For this example:\n",
        "\n",
        "* **Step 1:** You will search all the images for a specific query\n",
        "\n",
        "* **Step 2:** Send those images to Gemini to ask multiple questions, where it has to compare among those images and provide you with answers."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "E6AHCSwojyX0"
      },
      "outputs": [],
      "source": [
        "matching_results_image_query_1 = get_similar_image_from_query(\n",
        "    text_metadata_df,\n",
        "    image_metadata_df,\n",
        "    query=\"Show me all the images that can describe LLMs and TPU v5e scaling\",\n",
        "    column_name=\"text_embedding_from_image_description\",  # Use image description text embedding # mm_embedding_from_img_only text_embedding_from_image_description\n",
        "    image_emb=False,  # Use text embedding instead of image embedding\n",
        "    top_n=5,\n",
        "    embedding_size=1408,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jja-9iGFRNYQ"
      },
      "outputs": [],
      "source": [
        "# Check Matched Images\n",
        "# You can access the other two matched images using:\n",
        "\n",
        "print(\"---------------Matched Images------------------\\n\")\n",
        "display_images(\n",
        "    [\n",
        "        matching_results_image_query_1[0][\"img_path\"],\n",
        "        matching_results_image_query_1[1][\"img_path\"],\n",
        "        matching_results_image_query_1[2][\"img_path\"],\n",
        "        matching_results_image_query_1[3][\"img_path\"],\n",
        "        matching_results_image_query_1[4][\"img_path\"],\n",
        "    ],\n",
        "    resize_ratio=0.2,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fSR_JWkSC_7p"
      },
      "outputs": [],
      "source": [
        "prompt = f\"\"\"Task: Answer the following questions in detail, providing clear reasoning and evidence from the images in bullet points.\n",
        "Instructions:\n",
        "1. Analyze the provided images focusing on the relationship between TPU v5e scaling efficiency, LLM model size growth, performance metrics, and quantization effects.\n",
        "2. Answer the following questions in detail, providing clear reasoning and evidence from the images in bullet points\n",
        "3. Cite the image sources to support your explanations. Mention the file name.\n",
        "\n",
        "Additional Considerations:\n",
        "* Clearly define any technical terms (e.g., EMFU, TFLOP/chip/s) within your answers for better understanding.\n",
        "* Use specific examples and data points from the images to support your explanations.\n",
        "* Feel free to request additional information or clarification if the images are unclear or ambiguous.\n",
        "\n",
        "Question:\n",
        " - How does the scaling efficiency of TPU v5e compare to the overall growth in LLM model size over time?\n",
        " - How does the model size impact the observed Per-chip performance and EMFU for a fixed number of TPU v5e chips (e.g., 256)?\n",
        " - For the INT8 Quant training with 32B parameters, how does its high EMFU relate to the observed TFLOP/chip/s?\n",
        " - how does the \"per device batch (seq)\" for a 16B model compare to a 128B model, and how does this affect the \"Total observed Perf\"?\n",
        " - how might the MFU be impacted by increasing LLM model size?\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wYkzpB4PTSfm"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "# Generate response with Gemini\n",
        "print(\"\\n **** Result: ***** \\n\")\n",
        "rich_Markdown(\n",
        "    get_gemini_response(\n",
        "        multimodal_model_2,\n",
        "        model_input=[\n",
        "            prompt,\n",
        "            \"Images:\",\n",
        "            matching_results_image_query_1[0][\"image_object\"],\n",
        "            matching_results_image_query_1[1][\"image_object\"],\n",
        "            matching_results_image_query_1[2][\"image_object\"],\n",
        "            matching_results_image_query_1[3][\"image_object\"],\n",
        "            matching_results_image_query_1[4][\"image_object\"],\n",
        "        ],\n",
        "        stream=True,\n",
        "        safety_settings=safety_settings,\n",
        "        generation_config=GenerationConfig(temperature=1, max_output_tokens=8192),\n",
        "    )\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "efJPPrzRhvIT"
      },
      "source": [
        "## Building Multimodal QA System with retrieval augmented generation (mRAG)\n",
        "\n",
        "Let's bring everything together to implement multimodal RAG. You will use all the elements that you've explored in previous sections to implement the multimodal RAG. These are the steps:\n",
        "\n",
        "* **Step 1:** The user gives a query in text format where the expected information is available in the document and is embedded in images and text.\n",
        "* **Step 2:** Find all text chunks from the pages in the documents using a method similar to the one you explored in `Text Search`.\n",
        "* **Step 3:** Find all similar images from the pages based on the user query matched with `image_description` using a method identical to the one you explored in `Image Search`.\n",
        "* **Step 4:** Combine all similar text and images found in steps 2 and 3 as `context_text` and `context_images`.\n",
        "* **Step 5:** With the help of Gemini, we can pass the user query with text and image context found in steps 2 & 3. You can also add a specific instruction the model should remember while answering the user query.\n",
        "* **Step 6:** Gemini produces the answer, and you can print the citations to check all relevant text and images used to address the query."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EI62Hzuw_0_b"
      },
      "source": [
        "### Step 1: User query"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XvTKFwOPHLQ_"
      },
      "outputs": [],
      "source": [
        "# this time we are not passing any images, but just a simple text query.\n",
        "\n",
        "query = \"\"\"- How does the scaling efficiency of TPU v5e compare to the overall growth in LLM model size over time?\n",
        " - How does the model size impact the observed Per-chip performance and EMFU for a fixed number of TPU v5e chips (e.g., 256)?\n",
        " - For the INT8 Quant training with 32B parameters, how does its high EMFU relate to the observed TFLOP/chip/s?\n",
        " - how does the \"per device batch (seq)\" for a 16B model compare to a 128B model, and how does this affect the \"Total observed Perf\"?\n",
        " - how might the MFU be impacted by increasing LLM model size?\n",
        " \"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UUqlkKUaYvZA"
      },
      "source": [
        "### Step 2: Get all relevant text chunks"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r65yBb5gR_NG"
      },
      "outputs": [],
      "source": [
        "# Retrieve relevant chunks of text based on the query\n",
        "matching_results_chunks_data = get_similar_text_from_query(\n",
        "    query,\n",
        "    text_metadata_df,\n",
        "    column_name=\"text_embedding_chunk\",\n",
        "    top_n=20,\n",
        "    chunk_text=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mIgXgVIpYzxj"
      },
      "source": [
        "### Step 3: Get all relevant images"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wzu5Gf4yR_J4"
      },
      "outputs": [],
      "source": [
        "# Get all relevant images based on user query\n",
        "matching_results_image_fromdescription_data = get_similar_image_from_query(\n",
        "    text_metadata_df,\n",
        "    image_metadata_df,\n",
        "    query=query,\n",
        "    column_name=\"text_embedding_from_image_description\",\n",
        "    image_emb=False,\n",
        "    top_n=10,\n",
        "    embedding_size=1408,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RhUpWlGAY2uG"
      },
      "source": [
        "### Step 4: Create context_text and context_images"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "B_EEuuLCe6Y5"
      },
      "outputs": [],
      "source": [
        "instruction = \"\"\"Task: Answer the following questions in detail, providing clear reasoning and evidence from the images and text in bullet points.\n",
        "Instructions:\n",
        "\n",
        "1. **Analyze:** Carefully examine the provided images and text context.\n",
        "2. **Synthesize:** Integrate information from both the visual and textual elements.\n",
        "3. **Reason:**  Deduce logical connections and inferences to address the question.\n",
        "4. **Respond:** Provide a concise, accurate answer in the following format:\n",
        "\n",
        "   * **Question:** [Question]\n",
        "   * **Answer:** [Direct response to the question]\n",
        "   * **Explanation:** [Bullet-point reasoning steps if applicable]\n",
        "   * **Source** [name of the file, page, image from where the information is citied]\n",
        "\n",
        "5. **Ambiguity:** If the context is insufficient to answer, respond \"Not enough context to answer.\"\n",
        "\n",
        "\"\"\"\n",
        "\n",
        "# combine all the selected relevant text chunks\n",
        "context_text = [\"Text Context: \"]\n",
        "for key, value in matching_results_chunks_data.items():\n",
        "    context_text.extend(\n",
        "        [\n",
        "            \"Text Source: \",\n",
        "            f\"\"\"file_name: \"{value[\"file_name\"]}\" Page: \"{value[\"page_num\"]}\"\"\",\n",
        "            \"Text\",\n",
        "            value[\"chunk_text\"],\n",
        "        ]\n",
        "    )\n",
        "\n",
        "# combine all the selected relevant images\n",
        "gemini_content = [\n",
        "    instruction,\n",
        "    \"Questions: \",\n",
        "    query,\n",
        "    \"Image Context: \",\n",
        "]\n",
        "for key, value in matching_results_image_fromdescription_data.items():\n",
        "    gemini_content.extend(\n",
        "        [\n",
        "            \"Image Path: \",\n",
        "            value[\"img_path\"],\n",
        "            \"Image Description: \",\n",
        "            value[\"image_description\"],\n",
        "            \"Image:\",\n",
        "            value[\"image_object\"],\n",
        "        ]\n",
        "    )\n",
        "gemini_content.extend(context_text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rHrtodcBAEu9"
      },
      "source": [
        "### Step 5: Pass context to Gemini"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aZuhtJu7fW4n"
      },
      "outputs": [],
      "source": [
        "# Generate Gemini response with streaming output\n",
        "rich_Markdown(\n",
        "    get_gemini_response(\n",
        "        multimodal_model_2,\n",
        "        model_input=gemini_content,\n",
        "        stream=True,\n",
        "        safety_settings=safety_settings,\n",
        "        generation_config=GenerationConfig(temperature=1, max_output_tokens=8192),\n",
        "    )\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L0FtXYl1fzKh"
      },
      "source": [
        "### Step 6: Print citations and references [Optional]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "voThgteH-Tm8"
      },
      "source": [
        "**Optional:** Uncomment to see the detailed citations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IYRLQ47or1I8"
      },
      "outputs": [],
      "source": [
        "# print(\"---------------Matched Images------------------\\n\")\n",
        "# display_images(\n",
        "#     [\n",
        "#         matching_results_image_fromdescription_data[0][\"img_path\"],\n",
        "#         matching_results_image_fromdescription_data[1][\"img_path\"],\n",
        "#     ],\n",
        "#     resize_ratio=0.2,\n",
        "# )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "buwd_gp6HJ5K"
      },
      "outputs": [],
      "source": [
        "# # Image citations. You can check how Gemini generated metadata helped in grounding the answer.\n",
        "\n",
        "# print_text_to_image_citation(\n",
        "#     matching_results_image_fromdescription_data, print_top=True\n",
        "# )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "06vYM4MOHJ1-"
      },
      "outputs": [],
      "source": [
        "# # Text citations\n",
        "\n",
        "# print_text_to_text_citation(\n",
        "#     matching_results_chunks_data,\n",
        "#     print_top=True,\n",
        "#     chunk_text=True,\n",
        "# )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gIVF-QHuGVDD"
      },
      "source": [
        "### Multimodal RAG"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U82wS4nIB8IS"
      },
      "source": [
        "### More questions with Multimodal QA System"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZAyQVQ54B8X0"
      },
      "outputs": [],
      "source": [
        "# Some questions to try\n",
        "# this time we are not passing any images, but just a simple text query.\n",
        "query = \"\"\"Question 1: Imagine a patient presents with new onset prurigo nodularis.\n",
        "Could Med-Gemini-M 1.5 be used to analyze dermatological images of the patient's lesions in conjunction with a comprehensive history taken\n",
        "from an EHR dialogue to help a clinician reach a diagnosis and develop a treatment plan?\n",
        "What are the limitations and potential ethical considerations of using the model in this way?\n",
        "\n",
        "Question 2: The paper focuses on uncertainty-guided search for text-based reasoning tasks.\n",
        "How could this approach be extended to multimodal tasks?\n",
        "For instance, if Med-Gemini-M 1.5 encounters uncertainty when analyzing a dermatology image, could it generate queries to\n",
        "search for relevant visual examples or supplemental clinical information to refine its interpretation?\n",
        "\n",
        "Question 3:  Considering the potential benefits and risks highlighted in the paper, what specific steps should be taken during the development,\n",
        "validation, and deployment of Med-Gemini models to ensure they are used safely, fairly, and effectively in real-world clinical settings?\n",
        "How can these steps be informed by ongoing collaboration between researchers, clinicians, regulators, and patient communities?\n",
        " \"\"\"\n",
        "\n",
        "(\n",
        "    response,\n",
        "    matching_results_chunks_data,\n",
        "    matching_results_image_fromdescription_data,\n",
        ") = get_answer_from_qa_system(\n",
        "    query,\n",
        "    text_metadata_df,\n",
        "    image_metadata_df,\n",
        "    top_n_text=10,\n",
        "    top_n_image=5,\n",
        "    model=multimodal_model_2,\n",
        "    safety_settings=safety_settings,\n",
        "    generation_config=GenerationConfig(temperature=1, max_output_tokens=8192),\n",
        ")\n",
        "\n",
        "rich_Markdown(response)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dTHdui6jCHzc"
      },
      "outputs": [],
      "source": [
        "# Some questions to try\n",
        "\n",
        "query = \"\"\"Question 1: How does the mixture-of-experts architecture in Gemini 2.0 contribute to its ability to handle long\n",
        "context while maintaining performance on core capabilities? Discuss the potential trade-offs involved.\n",
        "\n",
        "Question 2: Gemini 2.0 incorporates various safety mitigations, including supervised fine-tuning and reinforcement learning.\n",
        "Discuss the effectiveness of these mitigations in addressing content safety and representational harms in both text-to-text and\n",
        "image-to-text modalities. How can these evaluations be improved?\n",
        "\n",
        "Question 3: Gemini 2.0 demonstrates surprising in-context language learning capabilities for Kalamang,\n",
        "a low-resource language. What are the implications of this finding for language preservation and revitalization?\n",
        "What challenges need to be addressed for broader applicability of this approach?\n",
        "\"\"\"\n",
        "(\n",
        "    response,\n",
        "    matching_results_chunks_data,\n",
        "    matching_results_image_fromdescription_data,\n",
        ") = get_answer_from_qa_system(\n",
        "    query,\n",
        "    text_metadata_df,\n",
        "    image_metadata_df,\n",
        "    top_n_text=10,\n",
        "    top_n_image=5,\n",
        "    model=multimodal_model_2,\n",
        "    safety_settings=safety_settings,\n",
        "    generation_config=GenerationConfig(temperature=1, max_output_tokens=8192),\n",
        ")\n",
        "\n",
        "rich_Markdown(response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KwNrHCqbi3xi"
      },
      "source": [
        "## Conclusions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "05jynhZnkgxn"
      },
      "source": [
        "Congratulations on making it through this multimodal RAG notebook!\n",
        "\n",
        "While multimodal RAG can be quite powerful, note that it can face some limitations:\n",
        "\n",
        "* **Data dependency:** Needs high-quality paired text and visuals.\n",
        "* **Computationally demanding:** Processing multimodal data is resource-intensive.\n",
        "* **Domain specific:** Models trained on general data may not shine in specialized fields like medicine.\n",
        "* **Black box:** Understanding how these models work can be tricky, hindering trust and adoption.\n",
        "\n",
        "\n",
        "Despite these challenges, multimodal RAG represents a significant step towards search and retrieval systems that can handle diverse, multimodal data."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "building_DIY_multimodal_qa_system_with_mRAG.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
