language/examples/document-qa/question_answering_large_documents.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"
   },
   "source": [
    "# Question Answering with Large Documents\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/examples/document-qa/question_answering_large_documents.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/examples/document-qa/question_answering_large_documents.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/examples/document-qa/question_answering_large_documents.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",
    "This notebook shows how you can build a question-answering (Q&A) system (or \"bot\") over multiple large documents so that Vertex AI PaLM API can answer any questions about the contents of those documents.\n",
    "\n",
    "Many companies have lots of information stored in documents, but retrieving that information easily and quickly can be challenging. To solve this, you will build a question-answering system powered by PaLM API to enable users to extract or query important details from those documents, which could be in any standard doc format such as .pdf, .doc, .docx, .txt, .pptx, or .html.\n",
    "\n",
    "The challenge with building a Q&A system over large documents is that you must do more than just pass the entire documents, into the prompts themselves, as the prompt context. This is because LLMs, including [Vertex PaLM API](https://cloud.google.com/vertex-ai/docs/generative-ai/learn/models), have token limits that [restrict how much context you can provide](https://ai.google/static/documents/palm2techreport.pdf).\n",
    "\n",
    "So how can you build a Q&A system with restrictions on token lengths? To solve this, in addition to your question (your prompt), you will need to provide just the relevant context; context that comes from your closed-domain sources (i.e. the large documents).\n",
    "\n",
    "In this notebook, you will see three methods that can address the large context challenge, known as:\n",
    "\n",
    "* **Stuffing** - pushing whole document content as a context.\n",
    "* **Map-Reduce** - splitting documents in smaller chunks.\n",
    "* **Map-Reduce - embedding** - creating embeddings of smaller chunks and using vector similarity search to find relevant context.\n",
    "\n",
    "The notebook introduces you to the fundamental approach towards handling huge documents for building a question-answering bot using Vertex PaLM API and finding relevant context for a user query, keeping the context limitation in check.\n",
    "\n",
    "In addition, there can be open source or Google Cloud drop-in replacement of steps, which will be discussed later in the notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "d975e698c9a4"
   },
   "source": [
    "### Objective\n",
    "\n",
    "By the end of the notebook, you will learn how to build a question-answering system that can handle large documents using the PaLM API.\n",
    "\n",
    "You will also learn the conceptual implementation of two methods to help you embed large contexts from many documents.\n",
    "\n",
    "At a high level, here are the topics that will be covered in this notebook\"\n",
    "\n",
    "* Install Vertex AI SDK & Other dependencies\n",
    "* Authenticating your notebook environment\n",
    "* Import libraries and Load models\n",
    "* Introduction to chains and index chains\n",
    "* Method 1: Stuffing\n",
    "* Method 2: Map Reduce\n",
    "* Method 3: Map Reduce with embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4j1gvi3jqG6U"
   },
   "source": [
    "### Costs\n",
    "\n",
    "This tutorial uses billable components of Google Cloud:\n",
    "\n",
    "* Vertex AI Generative AI Studio\n",
    "\n",
    "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-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": [
    "# Base system dependencies\n",
    "!sudo apt -y -qq install tesseract-ocr libtesseract-dev\n",
    "\n",
    "# required by PyPDF2 for page count and other pdf utilities\n",
    "!sudo apt-get -y -qq install poppler-utils python-dev libxml2-dev libxslt1-dev antiword unrtf poppler-utils pstotext tesseract-ocr flac ffmpeg lame libmad0 libsox-fmt-mp3 sox libjpeg-dev swig\n",
    "\n",
    "# Python dependencies\n",
    "!pip install google-cloud-aiplatform pytesseract PyPDF2 textract --upgrade --quiet --user"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UG3WsMySF3DQ"
   },
   "source": [
    "***Colab only***: Uncomment the following cell to restart the kernel or use the button 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": [
    "# # 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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Xe7OuYuGkLKF"
   },
   "source": [
    "### Authenticating your notebook environment\n",
    "* If you are using **Colab** to run this notebook, uncomment 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": "U9Gx2SAZkLKF"
   },
   "outputs": [],
   "source": [
    "# from google.colab import auth\n",
    "# auth.authenticate_user()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "960505627ddf"
   },
   "source": [
    "### Import libraries"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dQeD2K6eqP-W"
   },
   "source": [
    "**Colab only:** Uncomment the following cell to initialize the Vertex AI SDK. For Vertex AI Workbench, you don't need to run this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8YRKSFYOqSH4"
   },
   "outputs": [],
   "source": [
    "# import vertexai\n",
    "\n",
    "# PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
    "# vertexai.init(project=PROJECT_ID, location=\"us-central1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PyQmSRbKA8r-"
   },
   "outputs": [],
   "source": [
    "import glob\n",
    "import os\n",
    "import re\n",
    "import warnings\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import textract\n",
    "from PyPDF2 import PdfReader\n",
    "from tenacity import retry, stop_after_attempt, wait_random_exponential\n",
    "from vertexai.language_models import (TextEmbeddingModel,\n",
    "                                              TextGenerationModel)\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UP76a2la7O-a"
   },
   "source": [
    "### Import models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7isig7e07O-a"
   },
   "outputs": [],
   "source": [
    "generation_model = TextGenerationModel.from_pretrained(\"text-bison@001\")\n",
    "embedding_model = TextEmbeddingModel.from_pretrained(\"textembedding-gecko@001\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ubc1G1PlF3DR"
   },
   "source": [
    "To make PaLM API calls more resilient and comply with [API quotas](https://cloud.google.com/vertex-ai/docs/quotas), you can use an  [exponential backoff](https://en.wikipedia.org/wiki/Exponential_backoff) mechanism that keeps trying the API to ensure the call is successful without over-calling the API and adhering to the quotas.\n",
    "\n",
    "If you need your API quotas to be increased, refer [here](https://cloud.google.com/docs/quota_detail/view_manage#requesting_higher_quota).\n",
    "\n",
    "You can find API guide for the current method [here](https://tenacity.readthedocs.io/en/latest/api.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HWWog_UbF3DS"
   },
   "outputs": [],
   "source": [
    "@retry(wait=wait_random_exponential(min=1, max=20), stop=stop_after_attempt(3))\n",
    "def text_generation_model_with_backoff(**kwargs):\n",
    "    return generation_model.predict(**kwargs).text\n",
    "\n",
    "\n",
    "@retry(wait=wait_random_exponential(min=1, max=20), stop=stop_after_attempt(3))\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": "fIPcn5dZ7O-b"
   },
   "source": [
    "## Question Answering with large documents\n",
    "\n",
    "One of the most commonly used methods across the industry is `chains` to solve question-answering with large and multiple documents using LLMs.\n",
    "\n",
    "A chain is a sequence of steps an LLM takes to complete a task. For example, a chain might start with the LLM reading a document, then asking a question about the document, and finally generating a response to the question.\n",
    "\n",
    "\n",
    "An index chain is a special type of chain that uses an index to store and retrieve information. An index is a data structure that allows the LLM to quickly find relevant information for a given task. For example, an index chain might use an index to store the names of all the people mentioned in a document so that the LLM can quickly find the information it needs to answer a question about those people. Or it can store document path, name, page number, and other metadata.\n",
    "\n",
    "The idea is to create a simple index of all source documents so that LLMs can search through vast information easily. Index chains are helpful in question-answering, summarization, and chatbot\n",
    "\n",
    "Foundationally, there are four [index-related chains](https://docs.langchain.com/docs/components/chains/index_related_chains):\n",
    "* Stuffing\n",
    "* Map Reduce\n",
    "* Refine\n",
    "* Map-Rerank\n",
    "\n",
    "In this notebook, you will see three methods; Stuffing, Map Reduce and Map Reduce with embedding."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "s27oLm7LF3DS"
   },
   "source": [
    "### Method 1: Stuffing\n",
    "\n",
    "Stuffing is the simplest way to pass data to a large language model (LLM). You simply combine all of the data into a single prompt, and then pass that prompt to the LLM. This method has two advantages:\n",
    "\n",
    "* It only makes a single call to the LLM, which can improve performance.\n",
    "* The LLM has access to all of the data at once, which can improve the quality of the generated text.\n",
    "\n",
    "However, stuffing has one major disadvantage: it only works with small amounts of data. If you have a large dataset, stuffing will not be feasible.\n",
    "\n",
    "Before you dive deeper into possible methods for large document question-answering, you can explore the primary process of stuffing and how it fails with larger files and context.\n",
    "\n",
    "Here is the flow of stuffing:\n",
    "\n",
    "* **Document Loader**: Loading the required document from the source to your bucket or local storage.\n",
    "* **Document Processing**: Processing the documents by extracting content and other metadata.\n",
    "* **Context**: Building the context to pass the entire content extracted in the previous step.\n",
    "* **Prompt Engineering**: Building a question-answering prompt that takes the context built in the previous step and adds instructions to perform specific tasks.\n",
    "* **Vertex PaLM API**: Finally, with the prompt and the context, call the PaLM API to get the expected answer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-yCTqnJIF3DS"
   },
   "source": [
    "#### Document Loader\n",
    "You start copying the documents from a Cloud Bucket and store them in your project bucket or locally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TZgvF9ceF3DS"
   },
   "outputs": [],
   "source": [
    "# Copying the files from the GCS bucket to local\n",
    "!mkdir documents\n",
    "!gsutil -m cp -r gs://github-repo/documents ."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z8tr3zklo7hJ"
   },
   "source": [
    "You can view one of the documents here:\n",
    "https://storage.googleapis.com/github-repo/documents/20230426_alphabet_10Q.pdf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9tWvliZ2F3DS"
   },
   "source": [
    "#### Document Processing\n",
    "\n",
    "When you have documents, you need to process them for downstream consumption. In the processing phase, you aim to read the documents and convert them into a format that the downstream logic can easily use. While reading, you should keep as much metadata as possible from the original document.\n",
    "\n",
    "In this case, you are loading different file types, such as .pdf, .txt, .docx, and .json. Each file type has its reader, and you can use a simple open-source library called [textract](https://textract.readthedocs.io/en/stable/) and [PyPDF2](https://pypdf2.readthedocs.io/en/3.0.0/) to load them. You can save the file name, file type, page number (shown only for pdf), and content for each file.\n",
    "\n",
    "This metadata will be essential for quoting the source of information when sending it as a context and answering queries later on.\n",
    "\n",
    "The metadata and content extracted and processed are necessary because:\n",
    "* Quote the source of information when sending it as a context.\n",
    "* Answer queries about the documents.\n",
    "* Track changes to the documents.\n",
    "* Identify duplicate documents.\n",
    "* Organize the documents.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "L7bE7ZGtlcxv"
   },
   "outputs": [],
   "source": [
    "def create_data_packet(file_name, file_type, page_number, file_content):\n",
    "    \"\"\"Creating a simple dictionary to store all information (content and metadata)\n",
    "    extracted from the document\"\"\"\n",
    "    data_packet = {}\n",
    "    data_packet[\"file_name\"] = file_name\n",
    "    data_packet[\"file_type\"] = file_type\n",
    "    data_packet[\"page_number\"] = page_number\n",
    "    data_packet[\"content\"] = file_content\n",
    "    return data_packet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "MHkdhkCXF3DS"
   },
   "outputs": [],
   "source": [
    "final_data = []\n",
    "\n",
    "def files(path):\n",
    "    '''\n",
    "    Function that returns only filenames (and not folder names)\n",
    "    '''\n",
    "    for file in os.listdir(path):\n",
    "        if os.path.isfile(os.path.join(path, file)):\n",
    "            yield file\n",
    "            \n",
    "for file_name in files(\"documents/\"):\n",
    "    path = f\"documents/{file_name}\"\n",
    "    _, file_type = os.path.splitext(path)\n",
    "    if file_type == \".pdf\":\n",
    "        # loading pdf files, with page numbers as metadata.\n",
    "        reader = PdfReader(path)\n",
    "        for i, page in enumerate(reader.pages):\n",
    "            text = page.extract_text()\n",
    "            if text:\n",
    "                packet = create_data_packet(\n",
    "                    file_name, file_type, page_number=int(i + 1), file_content=text\n",
    "                )\n",
    "\n",
    "                final_data.append(packet)\n",
    "    else:\n",
    "        # loading other file types\n",
    "        text = textract.process(path).decode(\"utf-8\")\n",
    "        packet = create_data_packet(\n",
    "            file_name, file_type, page_number=None, file_content=text\n",
    "        )\n",
    "        final_data.append(packet)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z1QRxdrgF3DS"
   },
   "source": [
    "While extracting the content and metadata from the documents, you can store them in the pandas dataframe for easy downstream integration for citing the source of answer extraction. In addition, applying a text chunking process (splitting input text into smaller strings to fit into the token limit)  in the pandas dataframe will also be helpful."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9MISf5b7F3DS"
   },
   "outputs": [],
   "source": [
    "# converting the data that has been read from GCS to Pandas DataFrame for easy readibility and downstream logic\n",
    "pdf_data = pd.DataFrame.from_dict(final_data)\n",
    "pdf_data = pdf_data.sort_values(\n",
    "    by=[\"file_name\", \"page_number\"]\n",
    ")  # sorting the datafram by filename and page_number\n",
    "pdf_data.reset_index(inplace=True, drop=True)\n",
    "pdf_data.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EncFtHp7F3DS"
   },
   "outputs": [],
   "source": [
    "# you can check how many different file type you have in our datafrmae.\n",
    "print(\"Data has these different file types : \\n\", pdf_data[\"file_type\"].value_counts())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "50nOEGm_F3DS"
   },
   "source": [
    "#### Context Selection"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LiLlglWMF3DS"
   },
   "source": [
    "Now, the next step in the conventional method is to pass the context to PaLM API while asking the question.\n",
    "\n",
    "You don't know which document will be helpful, so you can go ahead and use all the document's text present in `content` column as context."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eXU6FzMZF3DS"
   },
   "outputs": [],
   "source": [
    "# combining all the content of the PDF as single string such that it can be passed as context.\n",
    "context = \"\\n\".join(str(v) for v in pdf_data[\"content\"].values)\n",
    "print(\"The total words in the context: \", len(context))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "soNkAis3F3DT"
   },
   "source": [
    "#### Prompt Engineering"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HEfMsjYpF3DT"
   },
   "source": [
    "Next, you can write a simple prompt along with the question. Then, you can preempt the prompt by making it follow some basic instructions. In the prompt, you only ask to answer if it finds the answer in the given `context`.\n",
    "\n",
    "You are dynamically passing the context and the question so that you can change it as per requirements and experimentations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6EGBdXZWF3DT"
   },
   "outputs": [],
   "source": [
    "question = \"What is the effect of change in accounting estimate for google in 2020?\"\n",
    "prompt = f\"\"\"Answer the question as precise as possible using the provided context. If the answer is\n",
    "              not contained in the context, say \"answer not available in context\" \\n\\n\n",
    "            Context: \\n {context}?\\n\n",
    "            Question: \\n {question} \\n\n",
    "            Answer:\n",
    "          \"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Dl7riuA3F3DT"
   },
   "source": [
    "#### Vertex PaLM API - Answer Extraction & Evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pDMEdg_fF3DT"
   },
   "source": [
    "In your prompt, you are passing so many words as context (roughly all documents).\n",
    "\n",
    "You already know that you have a input\n",
    "(prompt) token limit of [8192 tokens](https://cloud.google.com/vertex-ai/docs/generative-ai/learn/models) for the `text-bison@001` model, so your PaLM API call should fail. Because, as per ~8k token limit, the PaLM model is expecting ~6k words (input token). However, you are sending  ~ `1531642` words just as a prompt.\n",
    "\n",
    "As a reminder, a single token may be smaller than a word. A token is approximately four characters. Therefore, 100 tokens correspond to roughly 60-80 words.\n",
    "\n",
    "Hence, you know why conventional methods would not work when you want to do question-answering on large documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yglMnpZIF3DT"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    print(\"PaLM Predicted:\", generation_model.predict(prompt).text)\n",
    "except Exception as e:\n",
    "    print(\n",
    "        \"The code failed since it won't be able to run inference on such a huge context and throws this exception: \",\n",
    "        e,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-QcmWM4ZF3DT"
   },
   "source": [
    "However, you can still run the code, if you restrict the context to first 5000 words or something which is lesser than the token limit for PaLM API. But there is a good chance you will miss getting the expected answer, since your context might be missing in the first 5000 words."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SdqcCqlXF3DT"
   },
   "outputs": [],
   "source": [
    "prompt = f\"\"\"Answer the question as precise as possible using the provided context. If the answer is\n",
    "              not contained in the context, say \"answer not available in context\" \\n\\n\n",
    "            Context: \\n {context[:5000]}?\\n\n",
    "            Question: \\n {question} \\n\n",
    "            Answer:\n",
    "          \"\"\"\n",
    "print(\"the words in the prompt: \", len(prompt))\n",
    "print(\"PaLM Predicted:\", generation_model.predict(prompt).text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XTlsvjeMF3DT"
   },
   "source": [
    "So, now you have seen how stuffing the whole document content of so many files is not a very promising method to build question-answering systems. There are many different methods to address this limitation, but as discussed in the overview section, you will see two foundational and important methods:\n",
    "\n",
    "* Map-Reduce\n",
    "* Map-Reduce with embedding: Q&A"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iSotaftwF3DT"
   },
   "source": [
    "### Method 2: Map Reduce"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EHG-jxFPF3DT"
   },
   "source": [
    "[Map Reduce](https://docs.langchain.com/docs/components/chains/index_related_chains) Chains is a method for processing large amounts of data with a large language model (LLM). It works by breaking the data into smaller chunks, running an initial prompt on each chunk, and then combining the results of the initial prompts with a different prompt.\n",
    "\n",
    "For example, for question-answering, you can run initial prompt on each chunk to extract the answer and then finally combine the answers of the individual chunk with a different prompt.\n",
    "\n",
    "\n",
    "The typical flow for this method goes like this:\n",
    "\n",
    "* You take N documents from your source.\n",
    "* Split documents into N chunks (let's say 1000 words for each chunk)\n",
    "* Each chunk should be passed as context to the question-answer prompt\n",
    "* Summarize the answers from all chunk by using a seperate prompt."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CbhSQwhbQS4m"
   },
   "source": [
    "![Embedding Learning](https://storage.googleapis.com/github-repo/img/reference-architecture%20/map_reduce_flow_new.jpeg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mNSN2hXDF3DT"
   },
   "source": [
    "You can start by writing a simple function `get_chunks_iter` that takes a long string `text` and the size of the chunk as `maxlength`.\n",
    "\n",
    "This function aims to divide input string `text` into the size of `maxlength` - which are total words in that chunk and, save all the individual chunks into a list, and return `final_chunk` list."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qxLRYZSnF3DT"
   },
   "outputs": [],
   "source": [
    "# The function get_chunks_iter() can be used to split a piece of text into smaller chunks,\n",
    "# each of which is at most maxlength characters long.\n",
    "# This can be useful for tasks such as summarization, question answering, and translation.\n",
    "def get_chunks_iter(text, maxlength):\n",
    "    \"\"\"\n",
    "    Get chunks of text, each of which is at most maxlength characters long.\n",
    "\n",
    "    Args:\n",
    "        text: The text to be chunked.\n",
    "        maxlength: The maximum length of each chunk.\n",
    "\n",
    "    Returns:\n",
    "        An iterator over the chunks of text.\n",
    "    \"\"\"\n",
    "    start = 0\n",
    "    end = 0\n",
    "    final_chunk = []\n",
    "    while start + maxlength < len(text) and end != -1:\n",
    "        end = text.rfind(\" \", start, start + maxlength + 1)\n",
    "        final_chunk.append(text[start:end])\n",
    "        start = end + 1\n",
    "    final_chunk.append(text[start:])\n",
    "    return final_chunk\n",
    "\n",
    "\n",
    "# function to apply \"get_chunks_iter\" function on each row of dataframe.\n",
    "# currently each row here for file_type=pdf is content of each page and for other file_type its the whole document.\n",
    "def split_text(row):\n",
    "    chunk_iter = get_chunks_iter(row, chunk_size)\n",
    "    return chunk_iter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mt2FZVf0AGfc"
   },
   "source": [
    "The `global` keyword is used to declare a variable as global. This means that the variable can be accessed from any scope within the program. The `chunk_size` variable is declared as global because it will be used by other functions in the program.\n",
    "\n",
    "The `pdf_data_sample` variable is a copy of the `pdf_data` variable. This is done because the `pdf_data` variable will be modified by other functions in the program. By creating a copy of the variable, you can ensure that the original data is not modified."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8UagZAxsF3DT"
   },
   "outputs": [],
   "source": [
    "global chunk_size\n",
    "# you can define how many words should be there in a given chunk.\n",
    "chunk_size = 5000\n",
    "\n",
    "pdf_data_sample = pdf_data.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "n9ETtaMdwWKT"
   },
   "outputs": [],
   "source": [
    "# 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",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-r9sB_33BRru"
   },
   "source": [
    "The `split_text` function is a function that splits a string into a list of chunks, where each chunk is a continuous sequence of characters. In the second line of code below,\n",
    "```\n",
    "pdf_data_sample = pdf_data_sample.explode(\"chunks\")\n",
    "\n",
    "```\n",
    " explodes the chunks column into individual rows. This means that each row in the pdf_data_sample dataframe will now represent a single chunk of text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6BbTA1ezwZRS"
   },
   "outputs": [],
   "source": [
    "# Apply the chunk splitting logic here on each row of content in dataframe.\n",
    "pdf_data_sample[\"chunks\"] = pdf_data_sample[\"content\"].apply(split_text)\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\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BV23SywkwbRW"
   },
   "outputs": [],
   "source": [
    "# Sort and reset index\n",
    "pdf_data_sample = pdf_data_sample.sort_values(by=[\"file_name\", \"page_number\"])\n",
    "pdf_data_sample.reset_index(inplace=True, drop=True)\n",
    "pdf_data_sample.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1nKNvvTNF3DU"
   },
   "source": [
    "You can observe how a single page in the `20210203_alphabet_10K.pdf` file is divided into three chunks.\n",
    "\n",
    "You have three pages with the same \"1\" indicating that a page has been divided into three subsets (chunks). This is important because now you have a manageable chunk to send as context, rather than whole document as seen before.\n",
    "\n",
    "This will increase the total number of rows in the dataframe as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HF4f2fLuF3DU"
   },
   "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": "baKOz4SzF3DU"
   },
   "source": [
    "Now you can define the prompt and pass each chunk as the context."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GScEca0rF3DU"
   },
   "outputs": [],
   "source": [
    "# function to pass in the apply function on dataframe to extract answer for specific question on each row.\n",
    "def get_answer(df):\n",
    "    prompt = f\"\"\"Answer the question as precise as possible using the provided context. If the answer is\n",
    "                 not contained in the context, say \"answer not available in context\" \\n\\n\n",
    "                  Context: \\n {df['chunks']}?\\n\n",
    "                  Question: \\n {question} \\n\n",
    "                  Answer:\n",
    "            \"\"\"\n",
    "\n",
    "    pred = text_generation_model_with_backoff(prompt=prompt)\n",
    "    return pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iiBY30qDF3DU"
   },
   "outputs": [],
   "source": [
    "# we can take a small sample of the whole dataframe to avoid making too many calls to the API.\n",
    "pdf_data_sample_head = pdf_data_sample.head(10)\n",
    "\n",
    "question = \"What is the effect of change in accounting estimate for google in 2020?\"\n",
    "pdf_data_sample_head[\"predicted_answer\"] = pdf_data_sample_head.apply(\n",
    "    get_answer, axis=1\n",
    ")\n",
    "pdf_data_sample_head.head(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5hxh11yqGE5H"
   },
   "source": [
    "After you have asked the question-answering prompt to each chunk, combine all the answers into a new context. Then, send this new context to the final prompt. In the prompt you used for each chunk, you have told the model to return \"answer not available in context\" if it doesn't find any answers.\n",
    "\n",
    "This will help you remove the chunks where the model responded with \"answer not available in context\". The remaining chunks will be the new context."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QdAVZaaghwUg"
   },
   "outputs": [],
   "source": [
    "context_map_reduce = [\n",
    "    eachanswer\n",
    "    for eachanswer in pdf_data_sample_head[\"predicted_answer\"].values\n",
    "    if eachanswer != \"answer not available in context\"\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "E-l9HdEZiTHm"
   },
   "outputs": [],
   "source": [
    "prompt = f\"\"\"Answer the question as precise as possible using the provided context. If the answer is\n",
    "              not contained in the context, say \"answer not available in context\" \\n\\n\n",
    "            Context: \\n {context_map_reduce}?\\n\n",
    "            Question: \\n {question} \\n\n",
    "            Answer:\n",
    "          \"\"\"\n",
    "print(\"the words in the prompt: \", len(prompt))\n",
    "print(\"PaLM Predicted:\", generation_model.predict(prompt).text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NearlQHAF3DU"
   },
   "source": [
    "Now, let's look into this method's various pros and cons to summarize what you have done.\n",
    "\n",
    "**Pros:**\n",
    "\n",
    "* Increased precision due to chunking\n",
    "* Most helpful for extracting entities across different document levels.\n",
    "* It can scale to larger documents and more documents than other methods because chunks can be parallelized.\n",
    "\n",
    "\n",
    "**Cons:**\n",
    "\n",
    "* Multiple API calls, which can be costly and time-consuming\n",
    "* Slow, as it searches through all chunks even if the answer is found early\n",
    "* Conflicting answers, which can be difficult to resolve\n",
    "\n",
    "Moving forward, let's explore the following method, which addresses some of the shortcomings of Method 1."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "c7EV1fCdF3DU"
   },
   "source": [
    "### Method 3: Map Reduce with embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Y5B7BqbLF3DV"
   },
   "source": [
    "The previous method for question answering was inefficient because it required calling the PaLM API on all chunks of text. A more efficient approach is to create embeddings of the chunks and then use vector mathematics to find similar chunks. This allows you to find the relevant context from all the chunks in the dataframe where your answer may exist.\n",
    "\n",
    "The typical flow for this method is as follows:\n",
    "* Split documents into chunks.\n",
    "* Create embeddings for each chunk.\n",
    "* Convert the question to embeddings.\n",
    "* Perform a cosine similarity between the question and chunk embeddings to find the closest chunks.\n",
    "* Use the closest chunks as context for the PaLM API.\n",
    "\n",
    "This method is more efficient because it only calls the PaLM API on the relevant chunks.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6H5wJYviF3DV"
   },
   "source": [
    "![Embedding Learning](https://storage.googleapis.com/github-repo/img/reference-architecture%20/map_reduce_embedding.jpeg)\n"
   ]
  },
  {
   "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": "-G-9z2owF3DV"
   },
   "outputs": [],
   "source": [
    "pdf_data_sample_head[\"embedding\"] = pdf_data_sample_head[\"chunks\"].apply(\n",
    "    lambda x: embedding_model_with_backoff([x])\n",
    ")\n",
    "pdf_data_sample_head[\"embedding\"] = pdf_data_sample_head.embedding.apply(np.array)\n",
    "pdf_data_sample_head.head(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ocJjpRmoF3DV"
   },
   "source": [
    "Now comes the heart of this method. First, you can define a function `get_context_from_question`, which takes the:\n",
    "* `question` user wants to ask,\n",
    "* `vector_store`: vector db store, which you created in the last step and,\n",
    "* `sort_index_value`: The value defines how many chunks will be picked after running the sort on the cosine similarity score.\n",
    "\n",
    "The function will take the `valid_question`, create the embeddings, and do the dot product (cosine similarity) with all the chunks you passed in the vector store. Once you have the score, you can sort the results in decreasing order and pick chunks per the `sort_index_value` value as a combined string.\n",
    "\n",
    "This will become your context for the question asked."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rAJiEjbkF3DV"
   },
   "outputs": [],
   "source": [
    "def get_dot_product(row):\n",
    "    return np.dot(row, query_vector)\n",
    "\n",
    "\n",
    "def get_context_from_question(question, vector_store, sort_index_value=2):\n",
    "    global query_vector\n",
    "    query_vector = np.array(embedding_model_with_backoff([question]))\n",
    "    top_matched = (\n",
    "        vector_store[\"embedding\"]\n",
    "        .apply(get_dot_product)\n",
    "        .sort_values(ascending=False)[:sort_index_value]\n",
    "        .index\n",
    "    )\n",
    "    top_matched_df = vector_store[vector_store.index.isin(top_matched)][\n",
    "        [\"file_name\", \"page_number\", \"chunks\"]\n",
    "    ]\n",
    "    context = \" \".join(\n",
    "        vector_store[vector_store.index.isin(top_matched)][\"chunks\"].values\n",
    "    )\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": "q6jny8BMF3DV"
   },
   "outputs": [],
   "source": [
    "# your question for the documents\n",
    "question = \"What efforts have been taken by Google to safeguard their intellectual property in 2020?\"\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_head,\n",
    "    sort_index_value=5,  # 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": "UxdMFd5ZKU-z"
   },
   "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))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DacSlg_lF3DV"
   },
   "source": [
    "As you can see, the best part of this method is that you don't have to call the API multiple times. Instead, just one time, and it figured out the answers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "e2zcH_WCF3DV"
   },
   "source": [
    "Now, let's look into this method's various pros and cons to summarize what you have done.\n",
    "\n",
    "\n",
    "*  Fast: this is fast since it doesn't require the API to be executed on all the chunks.\n",
    "\n",
    "\n",
    "*  The dataframe can run into a vast length, and cosine similarity and basic mathematics can become slow.\n",
    "\n",
    "**Pros:**\n",
    "\n",
    "\n",
    "* Can be used to efficiently compute complex operations on large datasets of embeddings.\n",
    "* Can be used to learn representations of words and phrases that are more informative than traditional bag-of-words representations.\n",
    "* Can be used to improve the performance of a variety of natural language processing tasks, such as text classification, machine translation, and question answering.\n",
    "\n",
    "\n",
    "**Cons:**\n",
    "\n",
    "* Can be computationally expensive to compute vector similairty.\n",
    "* Can be sensitive to the choice of embeddings."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "question_answering_large_documents.ipynb",
   "toc_visible": true
  },
  "environment": {
   "kernel": "python3",
   "name": "tf2-gpu.2-11.m108",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/tf2-gpu.2-11:m108"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}