{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FI6OrHr9R6AA"
      },
      "source": [
        "# Document QnA using Gemini and Vertex AI Vector Search\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/gemini/use-cases/retrieval-augmented-generation/Document_QnA_using_gemini_and_vector_search.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%2Fgemini%2Fuse-cases%2Fretrieval-augmented-generation%2FDocument_QnA_using_gemini_and_vector_search.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/gemini/use-cases/retrieval-augmented-generation/Document_QnA_using_gemini_and_vector_search.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/gemini/use-cases/retrieval-augmented-generation/Document_QnA_using_gemini_and_vector_search.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" 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/gemini/use-cases/retrieval-augmented-generation/document_qna_using_gemini_and_vector_search.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/gemini/use-cases/retrieval-augmented-generation/document_qna_using_gemini_and_vector_search.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/gemini/use-cases/retrieval-augmented-generation/document_qna_using_gemini_and_vector_search.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/gemini/use-cases/retrieval-augmented-generation/document_qna_using_gemini_and_vector_search.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/gemini/use-cases/retrieval-augmented-generation/document_qna_using_gemini_and_vector_search.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": "-3K6twUOR-jD"
      },
      "source": [
        "| | |\n",
        "|-|-|\n",
        "|Author(s) | [Aakash Gouda](https://github.com/aksstar) , [Bhushan Garware](https://github.com/BhushanGarware)|"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "INIA40pMSYOn"
      },
      "source": [
        "## Overview\n",
        "\n",
        "Retrieval augmented generation (RAG) has become a popular paradigm for enabling LLMs to access external data and also as a mechanism for grounding to mitigate against hallucinations.\n",
        "\n",
        "In this notebook, you will learn how to perform RAG where you will perform Q&A over a document filled with both text and images.\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 2.0 model.\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 use multimodal RAG with Gemini API in Vertex AI, [text embeddings](https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/text-embeddings) to build a question answering system for a PDF document.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0M72MEscSasQ"
      },
      "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."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9VEOKEsLSdTx"
      },
      "source": [
        "### Objectives\n",
        "\n",
        "This notebook provides a guide to building a questions answering system using multimodal retrieval augmented generation (RAG).\n",
        "\n",
        "You will complete the following tasks:\n",
        "\n",
        "1. Extract data from documents containing both text and images using Gemini Vision Pro, and generate embeddings of the data, store it in vector store\n",
        "2. Search the vector store with text queries to find similar text data\n",
        "3. Using Text data as context, generate answer to the user query using Gemini Model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r11Gu7qNgx1p"
      },
      "source": [
        "## Getting Started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install Vertex AI SDK and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Hfn_UeXyTHaH"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --quiet pymupdf langchain gradio google-cloud-aiplatform langchain_google_vertexai"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j7UyNVSiyQ96"
      },
      "source": [
        "### Restart 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 restarts the current kernel.\n",
        "\n",
        "The restart might take a minute or longer. After its restarted, continue to the next step."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YmY9HVVGSBW5"
      },
      "outputs": [],
      "source": [
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EXQZrM5hQeKb"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ Wait for the kernel to finish restarting before you continue. ⚠️</b>\n",
        "</div>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you are running this notebook on Google Colab, run the cell below to authenticate your environment.\n",
        "\n",
        "This step is not required if you are using [Vertex AI Workbench](https://cloud.google.com/vertex-ai-workbench)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NyKGtVQjgx13"
      },
      "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": "DF4l8DTdWgPY"
      },
      "source": [
        "### Define Google Cloud project information and initialize Vertex AI\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "# Define project information\n",
        "# fmt: off\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
        "# fmt: on\n",
        "\n",
        "# Initialize Vertex AI\n",
        "import vertexai\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jXHfaVS66_01"
      },
      "source": [
        "### Import libraries\n",
        "Let's start by importing the libraries that we will need for this tutorial\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lslYAvw37JGQ"
      },
      "outputs": [],
      "source": [
        "# File system operations and displaying images\n",
        "import os\n",
        "\n",
        "# Import utility functions for timing and file handling\n",
        "import time\n",
        "\n",
        "# Libraries for downloading files, data manipulation, and creating a user interface\n",
        "import uuid\n",
        "from datetime import datetime\n",
        "\n",
        "import fitz\n",
        "import gradio as gr\n",
        "import pandas as pd\n",
        "from PIL import Image as PIL_Image\n",
        "\n",
        "# Initialize Vertex AI libraries for working with generative models\n",
        "from google.cloud import aiplatform\n",
        "from vertexai.generative_models import GenerativeModel, Image\n",
        "from vertexai.language_models import TextEmbeddingModel\n",
        "\n",
        "# Print Vertex AI SDK version\n",
        "print(f\"Vertex AI SDK version: {aiplatform.__version__}\")\n",
        "\n",
        "# Import LangChain components\n",
        "import langchain\n",
        "\n",
        "print(f\"LangChain version: {langchain.__version__}\")\n",
        "from langchain.text_splitter import CharacterTextSplitter\n",
        "from langchain_community.document_loaders import DataFrameLoader"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lP1EMJL4ddcA"
      },
      "source": [
        "### Initializing Gemini 2.0 and Text Embedding models"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sNOqAHc3dbDf"
      },
      "outputs": [],
      "source": [
        "# Loading Gemini 2.0 Model\n",
        "multimodal_model = GenerativeModel(\"gemini-2.0-flash\")\n",
        "\n",
        "# Initializing embedding model\n",
        "text_embedding_model = TextEmbeddingModel.from_pretrained(\"text-embedding-005\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bOZpB1hCtYtQ"
      },
      "source": [
        "### Download from internet a sample PDF file and default image to be shown when no results are found \n",
        "[Skip this step if you have uploaded your PDF file]\n",
        "\n",
        "---\n",
        "\n",
        "\n",
        "This document describes the importance of stable power grids in Japan, highlighting the recent failure of a generator step-up transformer at the Nakoso Power Station and the rapid restoration response undertaken to maintain power supply stability."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qCuvbLu9tDOa"
      },
      "outputs": [],
      "source": [
        "!wget https://www.hitachi.com/rev/archive/2023/r2023_04/pdf/04a02.pdf\n",
        "!wget https://img.freepik.com/free-vector/hand-drawn-no-data-illustration_23-2150696455.jpg\n",
        "\n",
        "# Create an \"Images\" directory if it doesn't exist\n",
        "Image_Path = \"./Images/\"\n",
        "if not os.path.exists(Image_Path):\n",
        "    os.makedirs(Image_Path)\n",
        "\n",
        "!mv hand-drawn-no-data-illustration_23-2150696455.jpg {Image_Path}/blank.jpg"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XPQ5eqaTcLP7"
      },
      "source": [
        "### Split PDF to images and extract data using Gemini Vision Pro\n",
        "This module processes a set of images, extracting text and tabular data using a multimodal model (Gemini 2.0).\n",
        "It handles potential errors, stores the extracted information in a DataFrame, and saves the results to a CSV file."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-41JZgAwsWXi"
      },
      "outputs": [],
      "source": [
        "# Run the following code for each file\n",
        "PDF_FILENAME = \"04a02.pdf\"  # Replace with your filename"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pSKL45w4O5ob"
      },
      "outputs": [],
      "source": [
        "# To get better resolution\n",
        "zoom_x = 2.0  # horizontal zoom\n",
        "zoom_y = 2.0  # vertical zoom\n",
        "mat = fitz.Matrix(zoom_x, zoom_y)  # zoom factor 2 in each dimension\n",
        "\n",
        "doc = fitz.open(PDF_FILENAME)  # open document\n",
        "for page in doc:  # iterate through the pages\n",
        "    pix = page.get_pixmap(matrix=mat)  # render page to an image\n",
        "    outpath = f\"./Images/{PDF_FILENAME}_{page.number}.jpg\"\n",
        "    pix.save(outpath)  # store image as a PNG\n",
        "\n",
        "# Define the path where images are located\n",
        "image_names = os.listdir(Image_Path)\n",
        "Max_images = len(image_names)\n",
        "\n",
        "# Create empty lists to store image information\n",
        "page_source = []\n",
        "page_content = []\n",
        "page_id = []\n",
        "\n",
        "p_id = 0  # Initialize image ID counter\n",
        "rest_count = 0  # Initialize counter for error handling\n",
        "\n",
        "while p_id < Max_images:\n",
        "    try:\n",
        "        # Construct the full path to the current image\n",
        "        image_path = Image_Path + image_names[p_id]\n",
        "\n",
        "        # Load the image\n",
        "        image = Image.load_from_file(image_path)\n",
        "\n",
        "        # Generate prompts for text and table extraction\n",
        "        prompt_text = \"Extract all text content in the image\"\n",
        "        prompt_table = (\n",
        "            \"Detect table in this image. Extract content maintaining the structure\"\n",
        "        )\n",
        "\n",
        "        # Extract text using your multimodal model\n",
        "        contents = [image, prompt_text]\n",
        "        response = multimodal_model.generate_content(contents)\n",
        "        text_content = response.text\n",
        "\n",
        "        # Extract table using your multimodal model\n",
        "        contents = [image, prompt_table]\n",
        "        response = multimodal_model.generate_content(contents)\n",
        "        table_content = response.text\n",
        "\n",
        "        # Log progress and store results\n",
        "        print(f\"processed image no: {p_id}\")\n",
        "        page_source.append(image_path)\n",
        "        page_content.append(text_content + \"\\n\" + table_content)\n",
        "        page_id.append(p_id)\n",
        "        p_id += 1\n",
        "\n",
        "    except Exception as err:\n",
        "        # Handle errors during processing\n",
        "        print(err)\n",
        "        print(\"Taking Some Rest\")\n",
        "        time.sleep(1)  # Pause execution for 1 second\n",
        "        rest_count += 1\n",
        "        if rest_count == 5:  # Limit consecutive error handling\n",
        "            rest_count = 0\n",
        "            print(f\"Cannot process image no: {image_path}\")\n",
        "            p_id += 1  # Move to the next image\n",
        "\n",
        "# Create a DataFrame to store extracted information\n",
        "df = pd.DataFrame(\n",
        "    {\"page_id\": page_id, \"page_source\": page_source, \"page_content\": page_content}\n",
        ")\n",
        "del page_id, page_source, page_content  # Conserve memory\n",
        "df.head()  # Preview the DataFrame"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ECB4ixV6VF0B"
      },
      "source": [
        "# Generate Text Embeddings\n",
        "Leverage a powerful language model text-embedding-005 to generate rich text embeddings that helps us find relevant information from a dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kHtZhmkYVCFP"
      },
      "outputs": [],
      "source": [
        "def generate_text_embedding(text) -> list:\n",
        "    \"\"\"Text embedding with a Large Language Model.\"\"\"\n",
        "    embeddings = text_embedding_model.get_embeddings([text])\n",
        "    vector = embeddings[0].values\n",
        "    return vector\n",
        "\n",
        "\n",
        "# Create a DataFrameLoader to prepare data for LangChain\n",
        "loader = DataFrameLoader(df, page_content_column=\"page_content\")\n",
        "\n",
        "# Load documents from the 'page_content' column of your DataFrame\n",
        "documents = loader.load()\n",
        "\n",
        "# Log the number of documents loaded\n",
        "print(f\"# of documents loaded (pre-chunking) = {len(documents)}\")\n",
        "\n",
        "# Create a text splitter to divide documents into smaller chunks\n",
        "text_splitter = CharacterTextSplitter(\n",
        "    chunk_size=10000,  # Target size of approximately 10000 characters per chunk\n",
        "    chunk_overlap=200,  # overlap between chunks\n",
        ")\n",
        "\n",
        "# Split the loaded documents\n",
        "doc_splits = text_splitter.split_documents(documents)\n",
        "\n",
        "# Add a 'chunk' ID to each document split's metadata for tracking\n",
        "for idx, split in enumerate(doc_splits):\n",
        "    split.metadata[\"chunk\"] = idx\n",
        "\n",
        "# Log the number of documents after splitting\n",
        "print(f\"# of documents = {len(doc_splits)}\")\n",
        "\n",
        "texts = [doc.page_content for doc in doc_splits]\n",
        "text_embeddings_list = []\n",
        "id_list = []\n",
        "page_source_list = []\n",
        "for doc in doc_splits:\n",
        "    id = uuid.uuid4()\n",
        "    text_embeddings_list.append(generate_text_embedding(doc.page_content))\n",
        "    id_list.append(str(id))\n",
        "    page_source_list.append(doc.metadata[\"page_source\"])\n",
        "    time.sleep(1)  # So that we don't run into Quota Issue\n",
        "\n",
        "# Creating a dataframe of ID, embeddings, page_source and text\n",
        "embedding_df = pd.DataFrame(\n",
        "    {\n",
        "        \"id\": id_list,\n",
        "        \"embedding\": text_embeddings_list,\n",
        "        \"page_source\": page_source_list,\n",
        "        \"text\": texts,\n",
        "    }\n",
        ")\n",
        "embedding_df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NQ2iLT3BG8vS"
      },
      "source": [
        "### Creating Vertex AI: Vector Search\n",
        "The code configures and deploys a vector search index on Google Cloud, making it ready to store and search through embeddings.\n",
        "\n",
        "Embedding size :  The number of values used to represent a piece of text in vector form. Larger dimensions mean a denser and potentially more expressive representation.\n",
        "\n",
        "\n",
        "Dimensions vs. Latency\n",
        "\n",
        "* Search: Higher-dimensional embeddings can make vector similarity searches slower, especially in large databases.\n",
        "* Computation: Calculations with larger vectors generally take more time during model training and inference.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WtCRJeKsYc6w"
      },
      "outputs": [],
      "source": [
        "VECTOR_SEARCH_REGION = \"us-central1\"\n",
        "VECTOR_SEARCH_INDEX_NAME = f\"{PROJECT_ID}-vector-search-index-ht\"\n",
        "VECTOR_SEARCH_EMBEDDING_DIR = f\"{PROJECT_ID}-vector-search-bucket-ht\"\n",
        "VECTOR_SEARCH_DIMENSIONS = 768"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6pu1a3zjfQ0D"
      },
      "source": [
        "### Save the embeddings in a JSON file\n",
        "To load the embeddings to Vector Search, we need to save them in JSON files with JSONL format. See more information in the docs at [Input data format and structure](https://cloud.google.com/vertex-ai/docs/matching-engine/match-eng-setup/format-structure#data-file-formats).\n",
        "\n",
        "First, export the `id` and `embedding` columns from the DataFrame in JSONL format, and save it.\n",
        "\n",
        "Then, create a new Cloud Storage bucket and copy the file to it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mSqVnRpqh7Yc"
      },
      "outputs": [],
      "source": [
        "# save id and embedding as a json file\n",
        "jsonl_string = embedding_df[[\"id\", \"embedding\"]].to_json(orient=\"records\", lines=True)\n",
        "with open(\"data.json\", \"w\") as f:\n",
        "    f.write(jsonl_string)\n",
        "\n",
        "# show the first few lines of the json file\n",
        "! head -n 3 data.json"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CzwDWJfzAk3n"
      },
      "outputs": [],
      "source": [
        "# Generates a unique ID for session\n",
        "UID = datetime.now().strftime(\"%m%d%H%M\")\n",
        "\n",
        "# Creates a GCS bucket\n",
        "BUCKET_URI = f\"gs://{VECTOR_SEARCH_EMBEDDING_DIR}-{UID}\"\n",
        "! gcloud storage buckets create {BUCKET_URI} --location=$LOCATION --project={PROJECT_ID}\n",
        "! gcloud storage cp data.json {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xxdbjKw1XDxl"
      },
      "source": [
        "### Create an Index\n",
        "\n",
        "Now it's ready to load the embeddings to Vector Search. Its APIs are available under the [aiplatform](https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform) package of the SDK."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xpMUXqWQ75uu"
      },
      "source": [
        "Create an [MatchingEngineIndex](https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform.MatchingEngineIndex) with its `create_tree_ah_index` function (Matching Engine is the previous name of Vector Search)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kKDw5VXMkXb3"
      },
      "outputs": [],
      "source": [
        "# create index\n",
        "my_index = aiplatform.MatchingEngineIndex.create_tree_ah_index(\n",
        "    display_name=f\"{VECTOR_SEARCH_INDEX_NAME}\",\n",
        "    contents_delta_uri=BUCKET_URI,\n",
        "    dimensions=768,\n",
        "    approximate_neighbors_count=20,\n",
        "    distance_measure_type=\"DOT_PRODUCT_DISTANCE\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2rFam_w9U0dI"
      },
      "source": [
        "By calling the `create_tree_ah_index` function, it starts building an Index. This will take under a few minutes if the dataset is small, otherwise about 50 minutes or more depending on the size of the dataset. You can check status of the index creation on [the Vector Search Console > INDEXES tab](https://console.cloud.google.com/vertex-ai/matching-engine/indexes).\n",
        "\n",
        "\n",
        "#### The parameters for creating index\n",
        "\n",
        "- `contents_delta_uri`: The URI of Cloud Storage directory where you stored the embedding JSON files\n",
        "- `dimensions`: Dimension size of each embedding. In this case, it is 768 as we are using the embeddings from the Text Embeddings API.\n",
        "- `approximate_neighbors_count`: how many similar items we want to retrieve in typical cases\n",
        "- `distance_measure_type`: what metrics to measure distance/similarity between embeddings. In this case it's `DOT_PRODUCT_DISTANCE`\n",
        "\n",
        "See [the document](https://cloud.google.com/vertex-ai/docs/vector-search/create-manage-index) for more details on creating Index and the parameters.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3fDs-FDS789e"
      },
      "source": [
        "### Create Index Endpoint and deploy the Index\n",
        "\n",
        "To use the Index, you need to create an [Index Endpoint](https://cloud.google.com/vertex-ai/docs/vector-search/deploy-index-public). It works as a server instance accepting query requests for your Index."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "peH6WpSj789m"
      },
      "outputs": [],
      "source": [
        "# create IndexEndpoint\n",
        "my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint.create(\n",
        "    display_name=f\"{VECTOR_SEARCH_INDEX_NAME}\",\n",
        "    public_endpoint_enabled=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5YtepoMX789m"
      },
      "source": [
        "This tutorial utilizes a [Public Endpoint](https://cloud.google.com/vertex-ai/docs/vector-search/setup/setup#choose-endpoint) and does not support [Virtual Private Cloud (VPC)](https://cloud.google.com/vpc/docs/private-services-access). Unless you have a specific requirement for VPC, we recommend using a Public Endpoint. Despite the term \"public\" in its name, it does not imply open access to the public internet. Rather, it functions like other endpoints in Vertex AI services, which are secured by default through IAM. Without explicit IAM permissions, as we have previously established, no one can access the endpoint."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ytZJmPMs8eav"
      },
      "source": [
        "With the Index Endpoint, deploy the Index by specifying an unique deployed index ID."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G6IKS9GD8ea7"
      },
      "outputs": [],
      "source": [
        "DEPLOYED_INDEX_NAME = VECTOR_SEARCH_INDEX_NAME.replace(\n",
        "    \"-\", \"_\"\n",
        ")  # Can't have - in deployment name, only alphanumeric and _ allowed\n",
        "DEPLOYED_INDEX_ID = f\"{DEPLOYED_INDEX_NAME}_{UID}\"\n",
        "# deploy the Index to the Index Endpoint\n",
        "my_index_endpoint.deploy_index(index=my_index, deployed_index_id=DEPLOYED_INDEX_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xu9ZmWcpXQ55"
      },
      "source": [
        "If it is the first time to deploy an Index to an Index Endpoint, it will take around 25 minutes to automatically build and initiate the backend for it. After the first deployment, it will finish in seconds. To see the status of the index deployment, open [the Vector Search Console > INDEX ENDPOINTS tab](https://console.cloud.google.com/vertex-ai/matching-engine/index-endpoints) and click the Index Endpoint."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fM3a9RQx4pQO"
      },
      "source": [
        "### Ask Questions to the PDF\n",
        "This code snippet establishes a question-answering (QA) system.  It leverages a vector search engine to find relevant information from a dataset and then uses the 'gemini-2.0-flash' LLM model to generate and refine the final answer to a user's query."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dRme7I2hhqC8"
      },
      "outputs": [],
      "source": [
        "def Test_LLM_Response(txt):\n",
        "    \"\"\"Determines whether a given text response generated by an LLM indicates a lack of information.\n",
        "\n",
        "    Args:\n",
        "        txt (str): The text response generated by the LLM.\n",
        "\n",
        "    Returns:\n",
        "        bool: True if the LLM's response suggests it was able to generate a meaningful answer,\n",
        "              False if the response indicates it could not find relevant information.\n",
        "\n",
        "    This function works by presenting a formatted classification prompt to the LLM.\n",
        "    The prompt includes the original text and specific categories indicating whether sufficient information was available.\n",
        "    The function analyzes the LLM's classification output to make the determination.\n",
        "    \"\"\"\n",
        "    classification_prompt = f\"\"\" Classify the text as one of the following categories:\n",
        "        -Information Present\n",
        "        -Information Not Present\n",
        "        Text=The provided context does not contain information.\n",
        "        Category:Information Not Present\n",
        "        Text=I cannot answer this question from the provided context.\n",
        "        Category:Information Not Present\n",
        "        Text:{txt}\n",
        "        Category:\"\"\"\n",
        "    classification_response = multimodal_model.generate_content(\n",
        "        classification_prompt\n",
        "    ).text\n",
        "\n",
        "    if \"Not Present\" in classification_response:\n",
        "        return False  # Indicates that the LLM couldn't provide an answer\n",
        "    return True  # Suggests the LLM generated a meaningful response\n",
        "\n",
        "\n",
        "def get_prompt_text(question, context):\n",
        "    \"\"\"Generates a formatted prompt string suitable for a language model, combining the provided question and context.\n",
        "\n",
        "    Args:\n",
        "        question (str): The user's original question.\n",
        "        context (str): The relevant text to be used as context for the answer.\n",
        "\n",
        "    Returns:\n",
        "        str: A formatted prompt string with placeholders for the question and context, designed to guide the language model's answer generation.\n",
        "    \"\"\"\n",
        "    prompt = f\"\"\"\n",
        "      Answer the question using the context below. Respond with only from the text provided\n",
        "      Question: {question}\n",
        "      Context : {context}\n",
        "      \"\"\"\n",
        "    return prompt\n",
        "\n",
        "\n",
        "def get_answer(query):\n",
        "    \"\"\"Retrieves an answer to a provided query using multimodal retrieval augmented generation (RAG).\n",
        "\n",
        "    This function leverages a vector search system to find relevant text documents from a\n",
        "    pre-indexed store of multimodal data. Then, it uses a large language model (LLM) to generate\n",
        "    an answer, using the retrieved documents as context.\n",
        "\n",
        "    Args:\n",
        "        query (str): The user's original query.\n",
        "\n",
        "    Returns:\n",
        "        dict: A dictionary containing the following keys:\n",
        "            * 'result' (str): The LLM-generated answer.\n",
        "            * 'neighbor_index' (int): The index of the most relevant document used for generation\n",
        "                                     (for fetching image path).\n",
        "\n",
        "    Raises:\n",
        "        RuntimeError: If no valid answer could be generated within the specified search attempts.\n",
        "    \"\"\"\n",
        "    neighbor_index = 0  # Initialize index for tracking the most relevant document\n",
        "    answer_found_flag = 0  # Flag to signal if an acceptable answer is found\n",
        "    result = \"\"  # Initialize the answer string\n",
        "    # Use a default image if the reference is not found\n",
        "    page_source = \"./Images/blank.jpg\"  # Initialize the blank image\n",
        "    query_embeddings = generate_text_embedding(\n",
        "        query\n",
        "    )  # Generate embeddings for the query\n",
        "\n",
        "    response = my_index_endpoint.find_neighbors(\n",
        "        deployed_index_id=DEPLOYED_INDEX_ID,\n",
        "        queries=[query_embeddings],\n",
        "        num_neighbors=5,\n",
        "    )  # Retrieve up to 5 relevant documents from the vector store\n",
        "\n",
        "    while answer_found_flag == 0 and neighbor_index < 4:\n",
        "        context = embedding_df[\n",
        "            embedding_df[\"id\"] == response[0][neighbor_index].id\n",
        "        ].text.values[0]  # Extract text context from the relevant document\n",
        "\n",
        "        prompt = get_prompt_text(\n",
        "            query, context\n",
        "        )  # Create a prompt using the question and context\n",
        "        result = multimodal_model.generate_content(\n",
        "            prompt\n",
        "        ).text  # Generate an answer with the LLM\n",
        "\n",
        "        if Test_LLM_Response(result):\n",
        "            answer_found_flag = 1  # Exit loop when getting a valid response\n",
        "        else:\n",
        "            neighbor_index += (\n",
        "                1  # Try the next retrieved document if the answer is unsatisfactory\n",
        "            )\n",
        "\n",
        "    if answer_found_flag == 1:\n",
        "        page_source = embedding_df[\n",
        "            embedding_df[\"id\"] == response[0][neighbor_index].id\n",
        "        ].page_source.values[0]  # Extract image_path from the relevant document\n",
        "    return result, page_source\n",
        "\n",
        "\n",
        "# fmt: off\n",
        "query = \"what is the 5th step of Transformer Manufacturing Flow ?\"  # @param {type:\"string\"}\n",
        "# fmt: on\n",
        "\n",
        "result, page_source = get_answer(query)\n",
        "print(result)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UeIDns8neIdO"
      },
      "source": [
        "# Ask Questions to the PDF using Gradio UI\n",
        " this code creates a web-based frontend for your question-answering system, allowing users to easily enter queries and see the results along with relevant images."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5dYlueAKbCJe"
      },
      "outputs": [],
      "source": [
        "def gradio_query(query):\n",
        "    # Retrieve the answer from your QA system\n",
        "    result, image_path = get_answer(query)\n",
        "\n",
        "    try:\n",
        "        # Attempt to fetch the source image reference\n",
        "        image = PIL_Image.open(image_path)  # Open the reference image\n",
        "    except:\n",
        "        # Use a default image if the reference is not found\n",
        "        image = PIL_Image.open(\"./Images/blank.jpg\")\n",
        "\n",
        "    return [result, image]  # Return both the text answer and the image\n",
        "\n",
        "\n",
        "gr.close_all()  # Ensure a clean Gradio interface\n",
        "with gr.Blocks() as demo:\n",
        "    with gr.Row():\n",
        "        with gr.Column():\n",
        "            # Input / Output Components\n",
        "            query = gr.Textbox(label=\"Query\", info=\"Enter your query\")\n",
        "            btn_enter = gr.Button(\"Process\")\n",
        "            answer = gr.Textbox(label=\"Response\", info=\"Enter your query\")\n",
        "            btn_clear = gr.Button(\"Clear\")\n",
        "        with gr.Column():\n",
        "            image = gr.Image(label=\"Reference\", visible=True)\n",
        "\n",
        "    # Button Click Event\n",
        "    btn_enter.click(fn=gradio_query, inputs=query, outputs=[answer, image])\n",
        "\n",
        "demo.launch(share=False, debug=True, inbrowser=True)  # Launch the Gradio app"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xbj4ooqjNUXv"
      },
      "source": [
        "### Close the demo\n",
        "\n",
        "Note: Stop the previous cell to close the Gradio server running locally then run this cell to free up the port utilised for running the server"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mgqWI-Z1xlVp"
      },
      "outputs": [],
      "source": [
        "demo.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eUac1wDXPDC2"
      },
      "source": [
        "### Cleaning up\n",
        "To clean up all Google Cloud resources used in this project, you can delete the Google Cloud project you used for the tutorial.\n",
        "\n",
        "Otherwise, you can delete the individual resources you created in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ItrB6IIUCMb1"
      },
      "outputs": [],
      "source": [
        "delete_bucket = False\n",
        "\n",
        "# Force undeployment of indexes and delete endpoint\n",
        "my_index_endpoint.delete(force=True)\n",
        "\n",
        "# Delete indexes\n",
        "my_index.delete()\n",
        "\n",
        "if delete_bucket:\n",
        "    ! gcloud storage rm --recursive --continue-on-error {BUCKET_URI}"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "Document_QnA_using_gemini_and_vector_search.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}