question_answering_documentai_vector_store_palm.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ur8xi4C7S06n"
   },
   "outputs": [],
   "source": [
    "# Copyright 2023 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": "JAPoU8Sm5E6e",
    "tags": []
   },
   "source": [
    "# Question answering with Documents using Document AI, Matching Engine, and PaLM\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/language/use-cases/document-qa/question_answering_documentai_vector_store_palm.ipynb\">\n",
    "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Run in Colab\n",
    "    </a>\n",
    "  </td>\n",
    "  <td style=\"text-align: center\">\n",
    "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/language/use-cases/document-qa/question_answering_documentai_vector_store_palm.ipynb\">\n",
    "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" 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/blob/main/language/use-cases/document-qa/question_answering_documentai_vector_store_palm.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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tvgnzT1CKxrO"
   },
   "source": [
    "## Overview\n",
    "\n",
    "\n",
    "Question answering for large documents is a challenging task because it requires the ability to understand and process a large amount of text.\n",
    "Question answering for large documents using Document AI and PaLM is a powerful approach that can be used to extract information from large amounts of text. Document AI can be used to extract structured data from documents, such as tables, lists, and images. PaLM can then be used to answer questions about the extracted data.\n",
    "\n",
    "\n",
    "[Document AI](https://cloud.google.com/document-ai) provides a scalable and managed way to extract data from documents using AI. In this notebook, you will use the [Document OCR processor](https://cloud.google.com/document-ai/docs/document-ocr), which is a pre-trained model that will extract text and layout information from document files.\n",
    "\n",
    "Here are some of the benefits of using Document AI and PaLM for question answering for large documents:\n",
    "\n",
    "It can be used to extract information from a variety of document formats.\n",
    "It can be used to answer questions about a variety of topics.\n",
    "It can be used to answer questions in natural language.\n",
    "It is a scalable solution that can be used to answer questions about large amounts of data.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "d975e698c9a4"
   },
   "source": [
    "### Objective\n",
    "\n",
    "In this notebook we will show how you can do the following:\n",
    "1. Extract text from pdf documents using the Document AI OCR processor.\n",
    "2. Use embedding model [Gecko](https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-text-embeddings) to generate embeddings for the extracted text\n",
    "1. Use PaLM `text-bison@001` model to answer questions on the embeddings datastore."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4j1gvi3jqG6U"
   },
   "source": [
    "### Costs\n",
    "\n",
    "This tutorial uses billable components of Google Cloud:\n",
    "\n",
    "* Vertex AI PaLM APIs offered by Generative AI studio\n",
    "* Document AI\n",
    "\n",
    "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing),\n",
    "Learn about [Document AI pricing](https://cloud.google.com/document-ai/pricing),\n",
    "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n",
    "to generate a cost estimate based on your projected usage."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QDU0XJ1xRDlL"
   },
   "source": [
    "## Getting Started"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2a5AEr0lkLKD"
   },
   "source": [
    "### Install Vertex AI SDK & Other dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GYACuZHAF3DQ"
   },
   "outputs": [],
   "source": [
    "%pip install --upgrade google-cloud-aiplatform==1.29.0 google-cloud-documentai==2.18.0 backoff==2.2.1  --user"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UG3WsMySF3DQ"
   },
   "source": [
    "**Colab only**: Run the following cell to restart the kernel. For Vertex AI Workbench you can restart the terminal using the button on top."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_Hsqwn4hkLKE"
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "if \"google.colab\" in sys.modules:\n",
    "    # Automatically 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)\n",
    "else:\n",
    "    # Otherwise, attempt to discover local credentials as described on https://cloud.google.com/docs/authentication/application-default-credentials\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5XtUQRDqlwTY"
   },
   "source": [
    "## Colab Only\n",
    "You will need to run the following cell to authenticates your Colab environment with your Google Cloud account."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "P_100VNRj-xn"
   },
   "outputs": [],
   "source": [
    "if \"google.colab\" in sys.modules:\n",
    "    from google.colab import auth\n",
    "\n",
    "    auth.authenticate_user()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "R-yyDsflbCc0"
   },
   "source": [
    "### Enable the APIs in your project"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ElgyIjJ7eSIy"
   },
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "import backoff\n",
    "from tenacity import retry, stop_after_attempt, wait_random_exponential\n",
    "from google.api_core.exceptions import ResourceExhausted\n",
    "from google.api_core.client_options import ClientOptions\n",
    "from google.api_core.exceptions import AlreadyExists\n",
    "from google.cloud import documentai\n",
    "import numpy as np\n",
    "import glob\n",
    "import os\n",
    "import textwrap\n",
    "from typing import Dict, List\n",
    "import pandas as pd\n",
    "from logging import error\n",
    "import re\n",
    "import textwrap\n",
    "from typing import Tuple, List\n",
    "import vertexai\n",
    "from vertexai.language_models import TextEmbeddingModel, TextGenerationModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ai62J2sLPDFX"
   },
   "outputs": [],
   "source": [
    "!gcloud config set project \"YOUR_PROJECT_ID\"\n",
    "\n",
    "!gcloud services enable documentai.googleapis.com storage.googleapis.com aiplatform.googleapis.com"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UP76a2la7O-a"
   },
   "source": [
    "## Document AI\n",
    "\n",
    "The following [limits](https://cloud.google.com/document-ai/quotas) apply for online processing with the Document OCR processor.\n",
    "\n",
    "| Limit             \t| Value \t|\n",
    "|-------------------\t|------:\t|\n",
    "| Maximum file size \t| 20 MB \t|\n",
    "| Maximum pages     \t| 15    \t|\n",
    "\n",
    "For documents that don't meet these limits, you can use [batch processing](https://cloud.google.com/document-ai/docs/send-request#batch-process) to extract the document text. (Not covered in this notebook.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HZkLDRTjTcfm"
   },
   "source": [
    "### Preparing data files\n",
    "\n",
    "To begin, you will need to download PDFs for the summarizing tasks below.\n",
    "For this notebook, you will be using Alphabet earnings report PDFs hosted in a public Google Cloud Storage bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "f-IWo-lb-gbn"
   },
   "outputs": [],
   "source": [
    "# Copying the files from the GCS bucket to local storage\n",
    "!gsutil -m cp -r gs://github-repo/documents/docai ."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FxmWYA05o4jj"
   },
   "source": [
    "### Create Document AI OCR Processor\n",
    "\n",
    "A [Document AI processor](https://cloud.google.com/document-ai/docs/overview#dai-processors) is an interface between a document file and a machine learning model that performs document processing actions. They can be used to classify, split, parse, or analyze a document. Each Google Cloud project needs to create its own processor instances.\n",
    "\n",
    "There are two types of Document AI processors:\n",
    "\n",
    "* Pre-trained processors: These processors are pre-trained on a large dataset of documents and can be used to perform common document processing tasks, such as Optical Character Recognition (OCR), form parsing, and entity extraction.\n",
    "* Custom processors: These processors can be trained on your own dataset of documents to perform specific tasks that are not covered by the pre-trained processors.\n",
    "\n",
    "Refer to [Full processor and detail list](https://cloud.google.com/document-ai/docs/processors-list) for all supported processors.\n",
    "\n",
    "Processors take a PDF or image file as input and output the data in the [`Document`](https://cloud.google.com/document-ai/docs/reference/rest/v1/Document) format."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "t9URyhLDo4jk"
   },
   "source": [
    "### Create a processor\n",
    "\n",
    "Set these variables below:\n",
    "\n",
    "1. Enter `YOUR_PROJECT_ID` in project_id\n",
    "2. Enter `YOUR_PROCESSOR_DISPLAY_NAME` for `processor_display_name`. For example `processor_display_name = \"my processor\"`, before running the code below.\n",
    "\n",
    "Note:  Run this code only once to create the processor.\n",
    "You cannot create multiple processors with the same display name. If you receieve an error, change the name of the processor and rerun."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7isig7e07O-a"
   },
   "outputs": [],
   "source": [
    "# TODO(developer): Edit these variables before running the code.\n",
    "project_id = \"YOUR_PROJECT_ID\"\n",
    "\n",
    "# See https://cloud.google.com/document-ai/docs/regions for all options.\n",
    "location = \"us\"\n",
    "\n",
    "# Must be unique per project, e.g.: \"My Processor\"\n",
    "processor_display_name = \"YOUR_PROCESSOR_DISPLAY_NAME\"\n",
    "\n",
    "# You must set the `api_endpoint` if you use a location other than \"us\".\n",
    "client_options = ClientOptions(api_endpoint=f\"{location}-documentai.googleapis.com\")\n",
    "\n",
    "\n",
    "def create_processor(\n",
    "    project_id: str, location: str, processor_display_name: str\n",
    ") -> documentai.Processor:\n",
    "    client = documentai.DocumentProcessorServiceClient(client_options=client_options)\n",
    "\n",
    "    # The full resource name of the location\n",
    "    # e.g.: projects/project_id/locations/location\n",
    "    parent = client.common_location_path(project_id, location)\n",
    "\n",
    "    # Create a processor\n",
    "    return client.create_processor(\n",
    "        parent=parent,\n",
    "        processor=documentai.Processor(\n",
    "            display_name=processor_display_name, type_=\"OCR_PROCESSOR\"\n",
    "        ),\n",
    "    )\n",
    "\n",
    "\n",
    "try:\n",
    "    processor = create_processor(project_id, location, processor_display_name)\n",
    "    print(f\"Created Processor {processor.name}\")\n",
    "except AlreadyExists as e:\n",
    "    print(\n",
    "        f\"Processor already exits, change the processor name and rerun this code. {e.message}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JkK74RRio4jk",
    "tags": []
   },
   "source": [
    "### Process the documents\n",
    "\n",
    "Process document takes the processor name and file path of the document and extracts the text from the document."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "L7bE7ZGtlcxv"
   },
   "outputs": [],
   "source": [
    "def process_document(\n",
    "    processor_name: str,\n",
    "    file_path: str,\n",
    ") -> documentai.Document:\n",
    "    client = documentai.DocumentProcessorServiceClient(client_options=client_options)\n",
    "\n",
    "    # Read the file into memory\n",
    "    with open(file_path, \"rb\") as image:\n",
    "        image_content = image.read()\n",
    "\n",
    "    # Load Binary Data into Document AI RawDocument Object\n",
    "    raw_document = documentai.RawDocument(\n",
    "        content=image_content, mime_type=\"application/pdf\"\n",
    "    )\n",
    "\n",
    "    # Configure the process request\n",
    "    request = documentai.ProcessRequest(name=processor_name, raw_document=raw_document)\n",
    "\n",
    "    result = client.process_document(request=request)\n",
    "\n",
    "    return result.document"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "usKPLAna5mLd"
   },
   "source": [
    "#### Create data chunks\n",
    "\n",
    "LLMs produce the best results when documents text is broken up into small \"chunks\" before being added to the prompt.\n",
    "\n",
    "Chunking is a technique used to break a document into smaller chunks that are easier to process. This can be done by dividing the document into sentences, paragraphs, or even sections.\n",
    "\n",
    "The current token size limit for PaLM is 8196 tokens. This means that a single request to the PaLM API can only process a document that is up to 8196 tokens long. If the document is longer than this, it will need to be broken up into smaller chunks.\n",
    "\n",
    "Chunking can be used to adjust for the token size limit in PaLM by breaking the document up into smaller chunks that are each less than 8196 tokens long. This allows the document to be processed in smaller, more manageable chunks.\n",
    "The best chunk size depends on the size of the documents. It is a good idea to experiment with different chunk sizes to see what works best for your specific dataset and application. For the provided documents, we are taking `5000` as a `chunk_value`. You should experiment with other values as well and see how it affects your summarization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nDf7Zw3SGuyW"
   },
   "outputs": [],
   "source": [
    "# If you already have a Document AI Processor in your project, assign the full processor resource name here.\n",
    "processor_name = processor.name\n",
    "chunk_size = 5000\n",
    "extracted_data: List[Dict] = []\n",
    "\n",
    "# Loop through each PDF file in the \"docai\" directory.\n",
    "for path in glob.glob(\"docai/*.pdf\"):\n",
    "    # Extract the file name and type from the path.\n",
    "    file_name, file_type = os.path.splitext(path)\n",
    "\n",
    "    print(f\"Processing {file_name}\")\n",
    "\n",
    "    # Process the document.\n",
    "    document = process_document(processor_name, file_path=path)\n",
    "\n",
    "    if document:\n",
    "        # Split the text into chunks of the specified size.\n",
    "        document_chunks = textwrap.wrap(text=document.text, width=chunk_size)\n",
    "\n",
    "        # Loop through each chunk and create a dictionary with metadata and content.\n",
    "        for chunk_number, chunk_content in enumerate(document_chunks, start=1):\n",
    "            # Append the chunk information to the extracted_data list.\n",
    "            extracted_data.append(\n",
    "                {\n",
    "                    \"file_name\": file_name,\n",
    "                    \"file_type\": file_type,\n",
    "                    \"chunk_number\": chunk_number,\n",
    "                    \"content\": chunk_content,\n",
    "                }\n",
    "            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zbPhNWEGw1hC"
   },
   "outputs": [],
   "source": [
    "# Convert extracted_data to a sorted Pandas DataFrame\n",
    "pdf_data = (\n",
    "    pd.DataFrame.from_dict(extracted_data)\n",
    "    .sort_values(by=[\"file_name\"])\n",
    "    .reset_index(drop=True)\n",
    ")\n",
    "\n",
    "pdf_data.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "y3r_0YLUw1hD"
   },
   "outputs": [],
   "source": [
    "# Define the maximum number of characters in each chunk.\n",
    "chunk_size = 5000\n",
    "\n",
    "pdf_data_sample = pdf_data.copy()\n",
    "\n",
    "# Remove all non-alphabets and numbers from the data to clean it up.\n",
    "# This is harsh cleaning. You can define your custom logic for cleansing here.\n",
    "pdf_data_sample[\"content\"] = pdf_data_sample[\"content\"].apply(\n",
    "    lambda x: re.sub(\"[^A-Za-z0-9]+\", \" \", x)\n",
    ")\n",
    "\n",
    "# Apply chunk splitting logic to each row of content in the DataFrame.\n",
    "pdf_data_sample[\"chunks\"] = pdf_data_sample[\"content\"].apply(\n",
    "    lambda row: textwrap.wrap(row, width=chunk_size)\n",
    ")\n",
    "\n",
    "# Now, each row in 'chunks' contains list of all chunks and hence we need to explode them into individual rows.\n",
    "pdf_data_sample = pdf_data_sample.explode(\"chunks\")\n",
    "\n",
    "# Sort and reset index\n",
    "pdf_data_sample = pdf_data_sample.sort_values(by=[\"file_name\"]).reset_index(drop=True)\n",
    "pdf_data_sample.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-Vp9LwVww1hE"
   },
   "outputs": [],
   "source": [
    "print(\"The original dataframe has :\", pdf_data.shape[0], \" rows without chunking\")\n",
    "print(\"The chunked dataframe has :\", pdf_data_sample.shape[0], \" rows with chunking\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Jj7KMupOPUX-",
    "tags": []
   },
   "source": [
    "## Question Answering using the [PaLM](https://ai.google/discover/palm2/) and Gecko Embedding model\n",
    "\n",
    "You have just used Document AI to extract text from PDF files.\n",
    "\n",
    "In the next section, you will perform question answering on the extracted text using the PaLM model and Gecko Embedding model with Vertex AI.\n",
    "\n",
    "Question answering using PaLM model and Gecko embedding model is a technique that can be used to improve the accuracy and efficiency of question answering.\n",
    "\n",
    "The Gecko embedding model is a pre-trained embedding model that has been trained on a large dataset of text. It can be used to represent text as vectors that can be used to measure the similarity between different pieces of text.\n",
    "\n",
    "The PaLM model is a large language model that has been trained on a massive dataset of text and code. It can be used to answer questions in a variety of ways, including factual question answering, open-ended question answering, and answering questions in context.\n",
    "\n",
    "To use the PaLM model and Gecko embedding model for question answering, you can use the following steps:\n",
    "\n",
    "\n",
    "1. Represent the documents that contain the answer as vectors using the Gecko embedding model.\n",
    "2. Compare the question to the document vector store.\n",
    "3. The document with the most similar vector to the question is the document that contains the answer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Xe7OuYuGkLKF"
   },
   "source": [
    "### Authenticating your notebook environment\n",
    "* If you are using **Colab** to run this notebook, run the cell below and continue.\n",
    "* If you are using **Vertex AI Workbench**, check out the setup instructions [here](https://github.com/GoogleCloudPlatform/generative-ai/tree/main/setup-env)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8YRKSFYOqSH4"
   },
   "outputs": [],
   "source": [
    "if \"google.colab\" in sys.modules:\n",
    "    PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
    "    vertexai.init(project=PROJECT_ID, location=\"us-central1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "H3edk4BiDkQ4"
   },
   "source": [
    "### Import  models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CKzpNLuzDeC4"
   },
   "outputs": [],
   "source": [
    "generation_model = TextGenerationModel.from_pretrained(\"text-bison@001\")\n",
    "embedding_model = TextEmbeddingModel.from_pretrained(\"textembedding-gecko@001\")\n",
    "\n",
    "\n",
    "# This decorator is used to handle exceptions and apply exponential backoff in case of ResourceExhausted errors.\n",
    "# It means the function will be retried with increasing time intervals in case of this specific exception.\n",
    "@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(5))\n",
    "def text_generation_model_with_backoff(**kwargs):\n",
    "    return generation_model.predict(**kwargs).text\n",
    "\n",
    "\n",
    "@retry(wait=wait_random_exponential(min=10, max=120), stop=stop_after_attempt(5))\n",
    "def embedding_model_with_backoff(text=[]):\n",
    "    embeddings = embedding_model.get_embeddings(text)\n",
    "    return [each.values for each in embeddings][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kU-s4MWnF3DV"
   },
   "source": [
    "You can start the implementation first by simply getting the embeddings for each chunk.\n",
    "\n",
    "This will add the embeddings (vector/number representation) of each chunk as a separate column."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PeJCuEx5w1hF"
   },
   "outputs": [],
   "source": [
    "pdf_data_sample[\"embedding\"] = pdf_data_sample[\"chunks\"].apply(\n",
    "    lambda x: embedding_model_with_backoff([x])\n",
    ")\n",
    "pdf_data_sample[\"embedding\"] = pdf_data_sample.embedding.apply(np.array)\n",
    "pdf_data_sample.head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8aGR1MrCw1hF"
   },
   "outputs": [],
   "source": [
    "def get_context_from_question(\n",
    "    question: str, vector_store: pd.DataFrame, sort_index_value: int = 2\n",
    ") -> Tuple[str, pd.DataFrame]:\n",
    "    query_vector = np.array(embedding_model_with_backoff([question]))\n",
    "    vector_store[\"dot_product\"] = vector_store[\"embedding\"].apply(\n",
    "        lambda row: np.dot(row, query_vector)\n",
    "    )\n",
    "    top_matched = vector_store.sort_values(by=\"dot_product\", ascending=False)[\n",
    "        :sort_index_value\n",
    "    ].index\n",
    "    top_matched_df = vector_store.loc[top_matched, [\"file_name\", \"chunks\"]]\n",
    "    context = \"\\n\".join(top_matched_df[\"chunks\"].values)\n",
    "    return context, top_matched_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "t_f5ZBqLF3DV"
   },
   "source": [
    "Now that you have a general function that always gets you custom relevant context for the question, you can call it with every new question."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nF7J1tanw1hG"
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "# your question for the documents\n",
    "question = \"When did google become carbon neutral?\"\n",
    "\n",
    "# get the custom relevant chunks from all the chunks in vector store.\n",
    "context, top_matched_df = get_context_from_question(\n",
    "    question,\n",
    "    vector_store=pdf_data_sample,\n",
    "    sort_index_value=2,  # Top N results to pick from embedding vector search\n",
    ")\n",
    "# top 5 data that has been picked by model based on user question. This becomes the context.\n",
    "top_matched_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KtYluejcw1hG"
   },
   "outputs": [],
   "source": [
    "# Prompt for Q&A which takes the custom context found in last step.\n",
    "prompt = f\"\"\" Answer the question as precise as possible using the provided context. \\n\\n\n",
    "            Context: \\n {context}?\\n\n",
    "            Question: \\n {question} \\n\n",
    "            Answer:\n",
    "          \"\"\"\n",
    "\n",
    "# Call the PaLM API on the prompt.\n",
    "print(\"PaLM Predicted:\", text_generation_model_with_backoff(prompt=prompt), \"\\n\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WLFuRGrtN-9l"
   },
   "source": [
    "# Conclusion\n",
    "\n",
    "In this notebook, you learned:\n",
    "\n",
    "1. Extract text from pdf documents using the Document AI OCR processor.\n",
    "2. Use Embedding model Gecko to generate embeddings for the extracted text\n",
    "3. Use PaLM chat-bison@latest model to answer questions on the embeddings datastore."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rhR1vTtpXj1q"
   },
   "source": [
    "## Clean Up\n",
    "\n",
    "If you no longer need the Document AI processor, you can delete it using the following code.\n",
    "\n",
    "Alternatively, you can use the Cloud Console to delete the processor as outlined in [Creating and managing processors > Delete a processor](https://cloud.google.com/document-ai/docs/create-processor#documentai_delete_processor-web)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EuTdzGhbe9C3"
   },
   "outputs": [],
   "source": [
    "def delete_processor(processor_name: str) -> None:\n",
    "    client = documentai.DocumentProcessorServiceClient(client_options=client_options)\n",
    "\n",
    "    # Delete a processor\n",
    "    operation = client.delete_processor(name=processor_name)\n",
    "    # Print operation details\n",
    "    print(operation.operation.name)\n",
    "    # Wait for operation to complete\n",
    "    operation.result()\n",
    "\n",
    "\n",
    "delete_processor(processor_name)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "environment": {
   "kernel": "python3",
   "name": "common-cpu.m108",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/base-cpu:m108"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}