language/use-cases/document-qa/question_answering_documents_langchain.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 using LangChain 🦜🔗\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_documents_langchain.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_documents_langchain.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_documents_langchain.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": {},
   "source": [
    "| | |\n",
    "|-|-|\n",
    "|Author(s) | [Lavi Nigam](https://github.com/lavinigam-gcp) |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tvgnzT1CKxrO"
   },
   "source": [
    "## Overview\n",
    "\n",
    "This notebook demonstrates how to build a question-answering (Q&A) system using LangChain with Vertex AI PaLM API to extract information from large documents.\n",
    "\n",
    "The challenge with building a Q&A system over large documents is that Large Language Models, LLMs in short, have token limits that restrict how much context you can provide.\n",
    "\n",
    "There are several methods to provide the context. They can use similarity search or not. Also there are different methods to pass context to LLMs. This notebook covers the following methods or chains:\n",
    "\n",
    "- **Stuffing**: Push the whole document content as a context. This is the simplest method, but it can be inefficient for large documents.\n",
    "\n",
    "- **Map-Reduce**: Split documents into smaller chunks and process them in parallel. This is more efficient than stuffing, but it can be more complex to implement.\n",
    "\n",
    "- **Refine**: Run an initial prompt on a small chunk, generate an output and for each subsequent document, refine the output based on both output and new document. This is more accurate than Map-Reduce but less efficient.\n",
    "\n",
    "This notebook also shows **Map-Reduce with Similarity search** where you create embeddings of smaller chunks and use vector similarity search to find relevant context. This is the most efficient method, but it can be the most complex to implement.\n",
    "\n",
    "\n",
    "Learn more about [LangChain](https://python.langchain.com/en/latest/use_cases/question_answering.html) and [Vertex Generative AI](https://cloud.google.com/vertex-ai/docs/generative-ai/learn/overview)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "d975e698c9a4"
   },
   "source": [
    "### Objective\n",
    "\n",
    "In this tutorial, you learn how to:\n",
    "\n",
    "- Ingest documents which involves download the documents.\n",
    "- Extract text from the PDF by using LangChain `PyPDFLoader`.\n",
    "- Select context for identifying the relevant parts of the document that are needed to answer the question.\n",
    "- Design prompt for question-answering\n",
    "- Leverage chains for handling large contexts (with/without embeddings)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aed92deeb4a0"
   },
   "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": "1uSGoyR6RrTQ"
   },
   "source": [
    "## Getting Started"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "i7EUnXsZhAGF"
   },
   "source": [
    "### Install Vertex AI SDK, other packages and their dependencies\n",
    "\n",
    "Install the following packages required to execute this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9tvCecNMQpXk"
   },
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2b4ef9b72d43"
   },
   "outputs": [],
   "source": [
    "# Install the packages\n",
    "import os\n",
    "\n",
    "if not os.getenv(\"IS_TESTING\"):\n",
    "    USER = \"--user\"\n",
    "else:\n",
    "    USER = \"\"\n",
    "# Install Vertex AI LLM SDK, langchain and dependencies\n",
    "! pip install google-cloud-aiplatform langchain==0.0.323 chromadb==0.3.26 pydantic==1.10.8 typing-inspect==0.8.0 typing_extensions==4.5.0 pandas datasets google-api-python-client transformers==4.33.1 pypdf faiss-cpu config --user"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "58707a750154"
   },
   "source": [
    "### Colab only: Uncomment the following cell to restart the kernel."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yStiAHorWE3Q"
   },
   "source": [
    "***Colab only***: Run 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": "f200f10a1da3"
   },
   "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": "opUxT_k5TdgP"
   },
   "source": [
    "### Authenticating your notebook environment\n",
    "\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": "vbNgv4q1T2Mi"
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "if \"google.colab\" in sys.modules:\n",
    "    from google.colab import auth\n",
    "\n",
    "    auth.authenticate_user()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- If you are running this notebook in a local development environment:\n",
    "  - Install the [Google Cloud SDK](https://cloud.google.com/sdk).\n",
    "  - Obtain authentication credentials. Create local credentials by running the following command and following the oauth2 flow (read more about the command [here](https://cloud.google.com/sdk/gcloud/reference/beta/auth/application-default/login)):\n",
    "\n",
    "    ```bash\n",
    "    gcloud auth application-default login\n",
    "    ```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "960505627ddf"
   },
   "source": [
    "### Import libraries"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Colab only:** Run 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": "GjSsu6cmUdEx"
   },
   "outputs": [],
   "source": [
    "import vertexai\n",
    "\n",
    "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
    "REGION = \"us-central1\"\n",
    "\n",
    "vertexai.init(project=PROJECT_ID, location=REGION)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PyQmSRbKA8r-"
   },
   "outputs": [],
   "source": [
    "import urllib\n",
    "import warnings\n",
    "from pathlib import Path as p\n",
    "from pprint import pprint\n",
    "\n",
    "import pandas as pd\n",
    "from langchain import PromptTemplate\n",
    "from langchain.chains.question_answering import load_qa_chain\n",
    "from langchain.document_loaders import PyPDFLoader\n",
    "from langchain.embeddings import VertexAIEmbeddings\n",
    "from langchain.llms import VertexAI\n",
    "from langchain.text_splitter import CharacterTextSplitter\n",
    "from langchain.vectorstores import Chroma\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "# restart python kernel if issues with langchain import."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DAGaTjPVTmhP"
   },
   "source": [
    "### Import models\n",
    "\n",
    "You load the pre-trained text and embeddings generation model called `text-bison@001` and `textembedding-gecko@001` respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ITUmZiNZcMUW"
   },
   "outputs": [],
   "source": [
    "vertex_llm_text = VertexAI(model_name=\"text-bison@001\")\n",
    "vertex_embeddings = VertexAIEmbeddings(model_name=\"textembedding-gecko@001\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Dxgiio5qTSGG"
   },
   "source": [
    "## Question Answering with large documents\n",
    "\n",
    "Large language models (LLMs) are powerful tools that can be used to answer a wide range of questions about large document base. However, there are some challenges associated with using large language model (LLM) for question answering. One of these challenges is related with the limited knowledge of LLMs models, especially when documents are specific of some context.\n",
    "\n",
    "One way to address this limitation is to give more information about documents using retrieval augmented generation. Retrieval augmented generation is a technique for using a large language model (LLM) to answer questions about documents it was not trained on. The basic idea is to first retrieve any relevant documents from a corpus called context, then pass those documents along with the original question to the LLM. The LLM will then generate a response that is informed by the information in the retrieved documents.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HZkLDRTjTcfm"
   },
   "source": [
    "### Ingest documents\n",
    "\n",
    "To begin, you will need to download a few files that are required for the summarizing tasks below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7H0zINHpTaSu"
   },
   "outputs": [],
   "source": [
    "data_folder = p.cwd() / \"data\"\n",
    "p(data_folder).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "pdf_url = \"https://services.google.com/fh/files/misc/practitioners_guide_to_mlops_whitepaper.pdf\"\n",
    "pdf_file = str(p(data_folder, pdf_url.split(\"/\")[-1]))\n",
    "\n",
    "urllib.request.urlretrieve(pdf_url, pdf_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JELITHdBhnf0"
   },
   "source": [
    "### Extract text from the PDF\n",
    "\n",
    "You use an `PdfReader` to extract the text from our scanned documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "x3INtovxreI_"
   },
   "outputs": [],
   "source": [
    "pdf_loader = PyPDFLoader(pdf_file)\n",
    "pages = pdf_loader.load_and_split()\n",
    "print(pages[3].page_content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "soNkAis3F3DT"
   },
   "source": [
    "### Prompt Design\n",
    "\n",
    "In a Q&A system, you define a question and the associated prompt.\n",
    "\n",
    "The question is simply a string that represents the question that the application will be asked to answer. In this case, the question is ```\"What is Experimentation?\"```\n",
    "\n",
    "The prompt is a string that contains the context that the application will use to generate an answer to the question. In this case, the prompt is\n",
    "\n",
    "```\n",
    "Answer the question as precise as possible using the provided context.\n",
    "If the answer is not contained in the context, say \"answer not available in context\" \\n\\n\n",
    "\n",
    "Context: \\n {context}?\\n\n",
    "Question: \\n {question} \\n\n",
    "Answer:\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6EGBdXZWF3DT"
   },
   "outputs": [],
   "source": [
    "question = \"What is Experimentation?\"\n",
    "prompt_template = \"\"\"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",
    "                  \"\"\"\n",
    "\n",
    "prompt = PromptTemplate(\n",
    "    template=prompt_template, input_variables=[\"context\", \"question\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eLVkm8LPNxNb"
   },
   "source": [
    "### Q&A without similarity search\n",
    "\n",
    "About providing the context, you can provide it or you may use part of the text you are looking for answer.\n",
    "\n",
    "In this example, you select the first eight pages as context of your Q&A system."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "50nOEGm_F3DS"
   },
   "source": [
    "#### Context Selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yGKGGzcusLNP"
   },
   "outputs": [],
   "source": [
    "context = \"\\n\".join(str(p.page_content) for p in pages[:7])\n",
    "print(\"The total words in the context: \", len(context))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AsSKyi0vVnjj"
   },
   "source": [
    "#### Q&A Methods or Chains"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZDVwBFSjZ7ws"
   },
   "source": [
    "##### Method 1: Stuffing\n",
    "\n",
    "`Stuffing` is a simple method for applying large language models (LLMs) to question-answering. It involves providing the LLM with all of the relevant data as context in the prompt.\n",
    "\n",
    "In LangChain, you can use `StuffDocumentsChain` as part of the `load_qa_chain` method. What you need to do is setting `stuff` as `chain_type` of your chain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "T3baLTKgt7LU"
   },
   "outputs": [],
   "source": [
    "stuff_chain = load_qa_chain(vertex_llm_text, chain_type=\"stuff\", prompt=prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DSLlVLsMa-Kq"
   },
   "source": [
    "After you initialize a `load_qa_chain` chain, you can answer your question based on the input documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qHDja-L4bGFD"
   },
   "outputs": [],
   "source": [
    "stuff_answer = stuff_chain(\n",
    "    {\"input_documents\": pages[7:10], \"question\": question}, return_only_outputs=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FK1PO3Kzv_2v"
   },
   "outputs": [],
   "source": [
    "pprint(stuff_answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yJ4PJ96mX3kI"
   },
   "source": [
    "The `Stuffing` method has the advantage of only requiring a single call to the LLM, but it is limited by the LLM's context length and is not feasible for large amounts of data.\n",
    "\n",
    "Below you see an exception raising when the context reach the LLMs limit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-nYFv2Iyx9TD"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    print(\n",
    "        stuff_chain(\n",
    "            {\"input_documents\": pages[7:], \"question\": question},\n",
    "            return_only_outputs=True,\n",
    "        )\n",
    "    )\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": "_KDDYnK6yZOu"
   },
   "source": [
    "##### Method 2: MapReduce\n",
    "\n",
    "With `MapReduce`, you can overcome the context limit. It involves dividing the document into chunks, running an initial prompt on each chunk, and then combining the results of the initial prompts using a different prompt.\n",
    "\n",
    "In LangChain, you can use `MapReduceDocumentsChain` as part of the `load_qa_chain` method with `map_reduce` as `chain_type` of your chain.\n",
    "\n",
    "The `load_qa_chain` with `map_reduce` as `chain_type` requires two prompts, question and a combine prompts.\n",
    "\n",
    "The question prompt is used to ask the LLM to answer a question based on the provided context. In this case, the `question_prompt` is\n",
    "\n",
    "```\n",
    "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",
    "The combine prompt object is used to combine the extracted content and the question to create a final answer. In this case, the `combine_prompt` is\n",
    "\n",
    "```\n",
    "Given the extracted content and the question, create a final answer.\n",
    "If the answer is not contained in the context, say \"answer not available in context. \\n\\n\n",
    "Summaries: \\n {summaries}?\\n\n",
    "Question: \\n {question} \\n\n",
    "Answer:\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XJeB8gG-3OIB"
   },
   "outputs": [],
   "source": [
    "question_prompt_template = \"\"\"\n",
    "                    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",
    "question_prompt = PromptTemplate(\n",
    "    template=question_prompt_template, input_variables=[\"context\", \"question\"]\n",
    ")\n",
    "\n",
    "# summaries is required. a bit confusing.\n",
    "combine_prompt_template = \"\"\"Given the extracted content and the question, create a final answer.\n",
    "If the answer is not contained in the context, say \"answer not available in context. \\n\\n\n",
    "Summaries: \\n {summaries}?\\n\n",
    "Question: \\n {question} \\n\n",
    "Answer:\n",
    "\"\"\"\n",
    "combine_prompt = PromptTemplate(\n",
    "    template=combine_prompt_template, input_variables=[\"summaries\", \"question\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "K8IdKzmGade1"
   },
   "source": [
    "After you define expected prompt, you initialize a `load_qa_chain` chain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ZmkGiHXB2ID3"
   },
   "outputs": [],
   "source": [
    "map_reduce_chain = load_qa_chain(\n",
    "    vertex_llm_text,\n",
    "    chain_type=\"map_reduce\",\n",
    "    return_intermediate_steps=True,\n",
    "    question_prompt=question_prompt,\n",
    "    combine_prompt=combine_prompt,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_EjeKpt6bIiC"
   },
   "source": [
    "And you answer your question based on the input documents. Notice how you are passing entire document base."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "glkyJqiZ49Fq"
   },
   "outputs": [],
   "source": [
    "map_reduce_outputs = map_reduce_chain({\"input_documents\": pages, \"question\": question})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PiORwQpAbOAl"
   },
   "source": [
    "You can store answers in a Pandas dataframe for checking the `MapReduce` intermediate steps and the LLMs answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "r6FRSR7xRLew"
   },
   "outputs": [],
   "source": [
    "final_mp_data = []\n",
    "\n",
    "# for each document, extract metadata and intermediate steps of the MapReduce process\n",
    "for doc, out in zip(\n",
    "    map_reduce_outputs[\"input_documents\"], map_reduce_outputs[\"intermediate_steps\"]\n",
    "):\n",
    "    output = {}\n",
    "    output[\"file_name\"] = p(doc.metadata[\"source\"]).stem\n",
    "    output[\"file_type\"] = p(doc.metadata[\"source\"]).suffix\n",
    "    output[\"page_number\"] = doc.metadata[\"page\"]\n",
    "    output[\"chunks\"] = doc.page_content\n",
    "    output[\"answer\"] = out\n",
    "    final_mp_data.append(output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dA9cnh8YaNbF"
   },
   "outputs": [],
   "source": [
    "# create a dataframe from a dictionary\n",
    "pdf_mp_answers = pd.DataFrame.from_dict(final_mp_data)\n",
    "# sorting the dataframe by filename and page_number\n",
    "pdf_mp_answers = pdf_mp_answers.sort_values(by=[\"file_name\", \"page_number\"])\n",
    "pdf_mp_answers.reset_index(inplace=True, drop=True)\n",
    "pdf_mp_answers.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yA0eM1K3cvH2"
   },
   "outputs": [],
   "source": [
    "index = 3\n",
    "print(\"[Context]\")\n",
    "print(pdf_mp_answers[\"chunks\"].iloc[index])\n",
    "print(\"\\n\\n [Answer]\")\n",
    "print(pdf_mp_answers[\"answer\"].iloc[index])\n",
    "print(\"\\n\\n [Page number]\")\n",
    "print(pdf_mp_answers[\"page_number\"].iloc[index])\n",
    "print(\"\\n\\n [Source: file_name]\")\n",
    "print(pdf_mp_answers[\"file_name\"].iloc[index])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7YEQl-dnCIMh"
   },
   "outputs": [],
   "source": [
    "index = 5\n",
    "print(\"[Context]\")\n",
    "print(pdf_mp_answers[\"chunks\"].iloc[index])\n",
    "print(\"\\n\\n [Answer]\")\n",
    "print(pdf_mp_answers[\"answer\"].iloc[index])\n",
    "print(\"\\n\\n [Page number]\")\n",
    "print(pdf_mp_answers[\"page_number\"].iloc[index])\n",
    "print(\"\\n\\n [Source: file_name]\")\n",
    "print(pdf_mp_answers[\"file_name\"].iloc[index])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "heqYTfSocXcZ"
   },
   "source": [
    "**Consideration**: The `MapReduce` method has the advantage of being able to scale to larger amounts of data than the stuffing method, but it requires more calls to the LLM and may lose some information during the final combined call."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RzrghSSLC4Hr"
   },
   "source": [
    "##### Method 3: Refine\n",
    "\n",
    "With `Refine` method, you try to overcome the lost of `information` of `MapReduce` method. The method involves running an initial prompt on the first chunk of data, generating some output. For the remaining documents, that output is passed in, along with the next document, asking the LLM to refine the output based on the new document.\n",
    "\n",
    "In LangChain, you can use `MapReduceDocumentsChain` as part of the `load_qa_chain` method. What you need to do is setting `refine` as `chain_type` of your chain.\n",
    "\n",
    "The `load_qa_chain` with `refine` as chain_type requires two prompts, refine and a initial question prompts.\n",
    "\n",
    "The `refine prompt` is used to generate a prompt that asks the LLM to refine an existing answer based on the provided context. In this case, the `refine prompt` is:\n",
    "\n",
    "```\n",
    "The original question is: \\n {question} \\n\n",
    "The provided answer is: \\n {existing_answer}\\n\n",
    "Refine the existing answer if needed with the following context: \\n {context_str} \\n\n",
    "Given the extracted content and the question, create a final answer.\n",
    "If the answer is not contained in the context, say \"answer not available in context. \\n\\n\n",
    "```\n",
    "\n",
    "The `initial question` prompt is used to generate a prompt that asks the LLM to answer a question based on the provided context only. In this case, the `initial question prompt` is:\n",
    "\n",
    "```\n",
    "Answer the question as precise as possible using the provided context only. \\n\\n\n",
    "Context: \\n {context_str} \\n\n",
    "Question: \\n {question} \\n\n",
    "Answer:\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cG-9vfd3C99y"
   },
   "outputs": [],
   "source": [
    "refine_prompt_template = \"\"\"\n",
    "    The original question is: \\n {question} \\n\n",
    "    The provided answer is: \\n {existing_answer}\\n\n",
    "    Refine the existing answer if needed with the following context: \\n {context_str} \\n\n",
    "    Given the extracted content and the question, create a final answer.\n",
    "    If the answer is not contained in the context, say \"answer not available in context. \\n\\n\n",
    "\"\"\"\n",
    "refine_prompt = PromptTemplate(\n",
    "    input_variables=[\"question\", \"existing_answer\", \"context_str\"],\n",
    "    template=refine_prompt_template,\n",
    ")\n",
    "\n",
    "\n",
    "initial_question_prompt_template = \"\"\"\n",
    "    Answer the question as precise as possible using the provided context only. \\n\\n\n",
    "    Context: \\n {context_str} \\n\n",
    "    Question: \\n {question} \\n\n",
    "    Answer:\n",
    "\"\"\"\n",
    "\n",
    "initial_question_prompt = PromptTemplate(\n",
    "    input_variables=[\"context_str\", \"question\"],\n",
    "    template=initial_question_prompt_template,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "q4-NPHduDngj"
   },
   "source": [
    "After you define expected prompt, you initialize a `load_qa_chain` chain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "97O6F_xAI-9-"
   },
   "outputs": [],
   "source": [
    "refine_chain = load_qa_chain(\n",
    "    vertex_llm_text,\n",
    "    chain_type=\"refine\",\n",
    "    return_intermediate_steps=True,\n",
    "    question_prompt=initial_question_prompt,\n",
    "    refine_prompt=refine_prompt,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ddfy336FDs58"
   },
   "source": [
    "And you answer your question based on the input documents. Notice how you are passing entire document base."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xK4afrbsKTlT"
   },
   "outputs": [],
   "source": [
    "refine_outputs = refine_chain({\"input_documents\": pages, \"question\": question})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xTvanB3fELHt"
   },
   "source": [
    "You can store answers in a Pandas dataframe for checking the `Refine` intermediate steps and the LLMs answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "AhKyI5UAKzLY"
   },
   "outputs": [],
   "source": [
    "final_refine_data = []\n",
    "for doc, out in zip(\n",
    "    refine_outputs[\"input_documents\"], refine_outputs[\"intermediate_steps\"]\n",
    "):\n",
    "    output = {}\n",
    "    output[\"file_name\"] = p(doc.metadata[\"source\"]).stem\n",
    "    output[\"file_type\"] = p(doc.metadata[\"source\"]).suffix\n",
    "    output[\"page_number\"] = doc.metadata[\"page\"]\n",
    "    output[\"chunks\"] = doc.page_content\n",
    "    output[\"answer\"] = out\n",
    "    final_refine_data.append(output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fKBVwN7bKzLZ"
   },
   "outputs": [],
   "source": [
    "pdf_refine_answers = pd.DataFrame.from_dict(final_refine_data)\n",
    "pdf_refine_answers = pdf_refine_answers.sort_values(\n",
    "    by=[\"file_name\", \"page_number\"]\n",
    ")  # sorting the dataframe by filename and page_number\n",
    "pdf_refine_answers.reset_index(inplace=True, drop=True)\n",
    "pdf_refine_answers.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fzBcke2DKzLZ"
   },
   "outputs": [],
   "source": [
    "index = 3\n",
    "print(\"[Context]\")\n",
    "print(pdf_refine_answers[\"chunks\"].iloc[index])\n",
    "print(\"\\n\\n [Answer]\")\n",
    "print(pdf_refine_answers[\"answer\"].iloc[index])\n",
    "print(\"\\n\\n [Page number]\")\n",
    "print(pdf_refine_answers[\"page_number\"].iloc[index])\n",
    "print(\"\\n\\n [Source: file_name]\")\n",
    "print(pdf_refine_answers[\"file_name\"].iloc[index])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gbDPwfu5LJV6"
   },
   "outputs": [],
   "source": [
    "index = 5\n",
    "print(\"[Context]\")\n",
    "print(pdf_refine_answers[\"chunks\"].iloc[index])\n",
    "print(\"\\n\\n [Answer]\")\n",
    "print(pdf_refine_answers[\"answer\"].iloc[index])\n",
    "print(\"\\n\\n [Page number]\")\n",
    "print(pdf_refine_answers[\"page_number\"].iloc[index])\n",
    "print(\"\\n\\n [Source: file_name]\")\n",
    "print(pdf_refine_answers[\"file_name\"].iloc[index])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "O30DdYUZEP6p"
   },
   "source": [
    "**Consideration**: So far, you use both part of the document or the entire document as the context to answer your specific question. Both cases have several limitations, including incomplete context and slow to query, especially for large context.\n",
    "\n",
    "Similarity search over a vector database, is a newer approach that addresses these limitations.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1dw-aHBJN7DN"
   },
   "source": [
    "### Q&A with similarity search\n",
    "\n",
    "With similarity search over a vector database, each piece of context is represented as a vector. These vectors are then stored in a database. When a user asks a question, the system first calculates the similarity between the question and the vectors in the database. The most similar vectors are then used to fetch the context that is relevant to the question.\n",
    "\n",
    "This approach has several advantages including more accurate context with respect of the user's question.\n",
    "\n",
    "In this case, you use `Chroma` an in-memory open-source embedding database to create similarity search index."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ydX1I_P-ToEj"
   },
   "source": [
    "#### Context Selection"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oSNMYEZhKaZJ"
   },
   "source": [
    "Create the similarity search index using `Chroma`.\n",
    "\n",
    "`Chroma` works with Document Loaders like `PyPDFLoader`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WLAMaLkgNG5U"
   },
   "outputs": [],
   "source": [
    "vector_index = Chroma.from_documents(pages, vertex_embeddings).as_retriever()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "L1c8Av35K_sl"
   },
   "source": [
    "Next, retrieve relevant context using the original question."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qgxD6ezEOAgL"
   },
   "outputs": [],
   "source": [
    "docs = vector_index.get_relevant_documents(question)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aU4sV359LNTM"
   },
   "source": [
    "#### MapReduce method\n",
    "\n",
    "Finally you answer your question based on the context you retrieve with embeddings database and the input question.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PrO635SkPdkR"
   },
   "outputs": [],
   "source": [
    "map_reduce_embeddings_outputs = map_reduce_chain(\n",
    "    {\"input_documents\": docs, \"question\": question}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tQozAWI3lkXp"
   },
   "outputs": [],
   "source": [
    "print(map_reduce_embeddings_outputs[\"output_text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can store answers in a Pandas dataframe for checking the `MapReduce with similarity search` intermediate steps and the LLMs answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_mpe_data = []\n",
    "for doc, out in zip(\n",
    "    map_reduce_embeddings_outputs[\"input_documents\"],\n",
    "    map_reduce_embeddings_outputs[\"intermediate_steps\"],\n",
    "):\n",
    "    output = {}\n",
    "    output[\"file_name\"] = p(doc.metadata[\"source\"]).stem\n",
    "    output[\"file_type\"] = p(doc.metadata[\"source\"]).suffix\n",
    "    output[\"page_number\"] = doc.metadata[\"page\"]\n",
    "    output[\"chunks\"] = doc.page_content\n",
    "    output[\"answer\"] = out\n",
    "    final_mpe_data.append(output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pdf_mpe_answers = pd.DataFrame.from_dict(final_mpe_data)\n",
    "pdf_mpe_answers = pdf_mpe_answers.sort_values(\n",
    "    by=[\"file_name\", \"page_number\"]\n",
    ")  # sorting the dataframe by filename and page_number\n",
    "pdf_mpe_answers.reset_index(inplace=True, drop=True)\n",
    "pdf_mpe_answers.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can store answers in a Pandas dataframe for checking the `MapReduce with similarity search` intermediate steps and the LLM answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_mpe_data = []\n",
    "for doc, out in zip(\n",
    "    map_reduce_embeddings_outputs[\"input_documents\"],\n",
    "    map_reduce_embeddings_outputs[\"intermediate_steps\"],\n",
    "):\n",
    "    output = {}\n",
    "    output[\"file_name\"] = p(doc.metadata[\"source\"]).stem\n",
    "    output[\"file_type\"] = p(doc.metadata[\"source\"]).suffix\n",
    "    output[\"page_number\"] = doc.metadata[\"page\"]\n",
    "    output[\"chunks\"] = doc.page_content\n",
    "    output[\"answer\"] = out\n",
    "    final_mpe_data.append(output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pdf_mpe_answers = pd.DataFrame.from_dict(final_mpe_data)\n",
    "pdf_mpe_answers = pdf_mpe_answers.sort_values(\n",
    "    by=[\"file_name\", \"page_number\"]\n",
    ")  # sorting the dataframe by filename and page_number\n",
    "pdf_mpe_answers.reset_index(inplace=True, drop=True)\n",
    "pdf_mpe_answers.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TpV-iwP9qw9c"
   },
   "source": [
    "## Conclusion\n",
    "\n",
    "This notebook demonstrates how to build a question-answering (QA) system using LangChain with Vertex AI PaLM API to extract information from large documents.\n",
    "\n",
    "In this case, you use Chroma, an in-memory open-source embedding database to create similarity search index. But [LangChain](https://python.langchain.com/docs/integrations/vectorstores/matchingengine) supports Vertex AI Matching Engine, the Google Cloud high-scale low latency vector database. With Vertex AI Matching Engine, you have a fully managed service that can scale to meet the needs of even the most demanding applications. It provides high performance for both training and inference. And it has several features including support for multiple similarity metrics, batch inference, and online learning. These features can be important for applications that need to perform complex matching tasks or that need to be able to adapt to changing data."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "question_answering_large_documents_langchain.ipynb",
   "toc_visible": true
  },
  "environment": {
   "kernel": "python3",
   "name": "tf2-gpu.2-11.m111",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/tf2-gpu.2-11:m111"
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}