bulk_question_answering.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "359697d5"
   },
   "source": [
    "# Bulk Question Answering with Vertex AI Search\n",
    "\n",
    "Answer questions from a CSV using a Vertex AI Search data store.\n",
    "\n",
    "Created by: Ruchika Kharwar @rasalt with edits from @holtskinner\n",
    "\n",
    "<table align=\"left\">\n",
    "\n",
    "  <td>\n",
    "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/search/bulk-question-answering/bulk_question_answering.ipynb\">\n",
    "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
    "    </a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/search/bulk-question-answering/bulk_question_answering.ipynb\">\n",
    "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
    "      View on GitHub\n",
    "    </a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/blob/main/search/bulk-question-answering/bulk_question_answering.ipynb\">\n",
    "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
    "      Open in Vertex AI Workbench\n",
    "    </a>\n",
    "  </td>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pQbOO-Lc-2-7"
   },
   "source": [
    "## Install pre-requisites\n",
    "\n",
    "If running in Colab install the pre-requisites into the runtime. Otherwise it is assumed that the notebook is running in Vertex Workbench. In that case it is recommended to install the pre-requistes from a terminal using the `--user` option."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ohPUPez8imvE",
    "tags": []
   },
   "outputs": [],
   "source": [
    "%pip install google-cloud-discoveryengine pandas --upgrade --user\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "B-mPnZJdiwkg"
   },
   "source": [
    "---\n",
    "\n",
    "#### ⚠️ Do not forget to click the \"RESTART RUNTIME\" button above.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8vEczuYo_r-g"
   },
   "source": [
    "## Authenticate\n",
    "\n",
    "If running in Colab authenticate with `google.colab.google.auth` otherwise assume that running on Vertex Workbench."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "loTfn0KniwB2"
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "if \"google.colab\" in sys.modules:\n",
    "    from google.colab import auth as google_auth\n",
    "    google_auth.authenticate_user()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Definitions for the index values\n",
    "\n",
    "- \"Query\"\n",
    "- \"Golden Doc\"\n",
    "- \"Golden Doc Page Number\"\n",
    "- \"Golden Answer\"\n",
    "- \"Top 5 Docs\"\n",
    "- \"Top 5 extractive answers\"\n",
    "- \"Top 5 extractive segments\"\n",
    "- \"Answer / Summary\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import pandas as pd\n",
    "from google.api_core.client_options import ClientOptions\n",
    "from google.cloud import discoveryengine_v1beta as discoveryengine\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "66ff7614"
   },
   "source": [
    "### Set the following constants to reflect your environment\n",
    "* The queries used in the examples here relate to a GCS bucket containing Alphabet investor PDFs, but these should be customised to your own data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "uVxFSrppK8Oy"
   },
   "outputs": [],
   "source": [
    "PROJECT_ID = \"document-ai-test-337818\"\n",
    "LOCATION = \"global\"\n",
    "DATA_STORE_ID = \"alphabet-earnings-reports_1693942917632\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Function to search the Vertex AI Search data store"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def search_data_store(\n",
    "    project_id: str,\n",
    "    location: str,\n",
    "    data_store_id: str,\n",
    "    search_query: str,\n",
    ") -> discoveryengine.SearchResponse:\n",
    "    #  For more information, refer to:\n",
    "    # https://cloud.google.com/generative-ai-app-builder/docs/locations#specify_a_multi-region_for_your_data_store\n",
    "    client_options = (\n",
    "        ClientOptions(api_endpoint=f\"{location}-discoveryengine.googleapis.com\")\n",
    "        if location != \"global\"\n",
    "        else None\n",
    "    )\n",
    "\n",
    "    # Create a client\n",
    "    client = discoveryengine.SearchServiceClient(client_options=client_options)\n",
    "\n",
    "    # The full resource name of the search engine serving config\n",
    "    # e.g. projects/{project_id}/locations/{location}/dataStores/{data_store_id}/servingConfigs/{serving_config_id}\n",
    "    serving_config = client.serving_config_path(\n",
    "        project=project_id,\n",
    "        location=location,\n",
    "        data_store=data_store_id,\n",
    "        serving_config=\"default_config\",\n",
    "    )\n",
    "\n",
    "    # Optional: Configuration options for search\n",
    "    # Refer to the `ContentSearchSpec` reference for all supported fields:\n",
    "    # https://cloud.google.com/python/docs/reference/discoveryengine/latest/google.cloud.discoveryengine_v1.types.SearchRequest.ContentSearchSpec\n",
    "    content_search_spec = discoveryengine.SearchRequest.ContentSearchSpec(\n",
    "        # For information about snippets, refer to:\n",
    "        # https://cloud.google.com/generative-ai-app-builder/docs/snippets\n",
    "        snippet_spec=discoveryengine.SearchRequest.ContentSearchSpec.SnippetSpec(\n",
    "            return_snippet=True\n",
    "        ),\n",
    "        extractive_content_spec=discoveryengine.SearchRequest.ContentSearchSpec.ExtractiveContentSpec(\n",
    "            max_extractive_answer_count=5,\n",
    "            max_extractive_segment_count=1,\n",
    "        ),\n",
    "        # For information about search summaries, refer to:\n",
    "        # https://cloud.google.com/generative-ai-app-builder/docs/get-search-summaries\n",
    "        summary_spec=discoveryengine.SearchRequest.ContentSearchSpec.SummarySpec(\n",
    "            summary_result_count=5,\n",
    "            include_citations=True,\n",
    "            ignore_adversarial_query=False,\n",
    "            ignore_non_summary_seeking_query=False,\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    # Refer to the `SearchRequest` reference for all supported fields:\n",
    "    # https://cloud.google.com/python/docs/reference/discoveryengine/latest/google.cloud.discoveryengine_v1.types.SearchRequest\n",
    "    request = discoveryengine.SearchRequest(\n",
    "        serving_config=serving_config,\n",
    "        query=search_query,\n",
    "        page_size=5,\n",
    "        content_search_spec=content_search_spec,\n",
    "        query_expansion_spec=discoveryengine.SearchRequest.QueryExpansionSpec(\n",
    "            condition=discoveryengine.SearchRequest.QueryExpansionSpec.Condition.AUTO,\n",
    "        ),\n",
    "        spell_correction_spec=discoveryengine.SearchRequest.SpellCorrectionSpec(\n",
    "            mode=discoveryengine.SearchRequest.SpellCorrectionSpec.Mode.AUTO\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    response = client.search(request)\n",
    "    return response\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Function to load results into data frame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "def answer_questions(row, project_id: str, location: str, data_store_id: str, top_n: int = 5) -> None:\n",
    "    \"\"\"This function returns the top 5 docs, extractive segments, answers\"\"\"\n",
    "    # Perform search with Query\n",
    "    response = search_data_store(project_id, location, data_store_id, row[\"Query\"])\n",
    "    \n",
    "    # df.at[row.name, \"Answer / Summary\"] = response.summary.summary_text\n",
    "    row[\"Answer / Summary\"] = response.summary.summary_text\n",
    "\n",
    "    top5docs, top5answers, top5segments = [], [], []\n",
    "    ext_ans_cnt, ext_seg_cnt = 0, 0\n",
    "\n",
    "    for result in response.results:\n",
    "        doc_data = getattr(result.document, \"derived_struct_data\", None)\n",
    "        if not doc_data:\n",
    "            continue\n",
    "\n",
    "        # Process extractive answers\n",
    "        for chunk in doc_data.get(\"extractive_answers\", []):\n",
    "            content = chunk.get(\"content\", \"\").replace(\"\\n\", \"\")\n",
    "            top5answers.append(content)\n",
    "            top5docs.append(f\"Doc: {doc_data.get('link', '')}  Page: {chunk.get('pageNumber', '')}\")\n",
    "            ext_ans_cnt += 1\n",
    "\n",
    "        # Process extractive segments\n",
    "        for chunk in doc_data.get(\"extractive_segments\", []):\n",
    "            data = chunk.get(\"content\", \"\").replace(\"\\n\", \"\")\n",
    "            top5segments.append(data)\n",
    "            ext_seg_cnt += 1\n",
    "\n",
    "        if ext_ans_cnt >= top_n and ext_seg_cnt >= top_n:\n",
    "            break\n",
    "    \n",
    "    # df.at[row.name, \"Top 5 Docs\"] = \"\\n\\n\".join(top5docs)\n",
    "    # df.at[row.name, \"Top 5 extractive answers\"] = \"\\n\\n\".join(top5answers)\n",
    "    # df.at[row.name, \"Top 5 extractive segments\"] = \"\\n\\n\".join(top5segments)\n",
    "    row[\"Top 5 Docs\"] = \"\\n\\n\".join(top5docs)\n",
    "    row[\"Top 5 extractive answers\"] = \"\\n\\n\".join(top5answers)\n",
    "    row[\"Top 5 extractive segments\"] = \"\\n\\n\".join(top5segments)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gather all of the Vertex AI Search results\n",
    "\n",
    "- Read in CSV as Pandas DataFrame\n",
    "- Send questions to Vertex AI Search\n",
    "- Load Summary, top 5 docs, extractive answers, extractive segments to DataFrame\n",
    "- Output DataFrame to TSV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Query</th>\n",
       "      <th>Golden Doc</th>\n",
       "      <th>Golden Doc Page Number</th>\n",
       "      <th>Golden Answer</th>\n",
       "      <th>Top 5 Docs</th>\n",
       "      <th>Top 5 extractive answers</th>\n",
       "      <th>Top 5 extractive segments</th>\n",
       "      <th>Answer / Summary</th>\n",
       "      <th>Feedback from customer / account team about returned docs and answer</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>What was Google's revenue in 2021?</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Doc: gs://cloud-samples-data/gen-app-builder/s...</td>\n",
       "      <td>Google Cloud had an Operating Loss of $890 mil...</td>\n",
       "      <td>Within Other Revenues, we are pleased with the...</td>\n",
       "      <td>Google's revenue for the full year 2021 was $5...</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>What was Google's revenue in 2022?</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Doc: gs://cloud-samples-data/gen-app-builder/s...</td>\n",
       "      <td>Other Revenues were $8.2 billion, up 22%, driv...</td>\n",
       "      <td>Let me now turn to our segment financial resul...</td>\n",
       "      <td>Google's total revenue was $282.8 billion in 2...</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                Query Golden Doc Golden Doc Page Number  \\\n",
       "0  What was Google's revenue in 2021?        NaN                    NaN   \n",
       "1  What was Google's revenue in 2022?        NaN                    NaN   \n",
       "\n",
       "  Golden Answer                                         Top 5 Docs  \\\n",
       "0           NaN  Doc: gs://cloud-samples-data/gen-app-builder/s...   \n",
       "1           NaN  Doc: gs://cloud-samples-data/gen-app-builder/s...   \n",
       "\n",
       "                            Top 5 extractive answers  \\\n",
       "0  Google Cloud had an Operating Loss of $890 mil...   \n",
       "1  Other Revenues were $8.2 billion, up 22%, driv...   \n",
       "\n",
       "                           Top 5 extractive segments  \\\n",
       "0  Within Other Revenues, we are pleased with the...   \n",
       "1  Let me now turn to our segment financial resul...   \n",
       "\n",
       "                                    Answer / Summary  \\\n",
       "0  Google's revenue for the full year 2021 was $5...   \n",
       "1  Google's total revenue was $282.8 billion in 2...   \n",
       "\n",
       "  Feedback from customer / account team about returned docs and answer  \n",
       "0                                                NaN                    \n",
       "1                                                NaN                    "
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Open the CSV file and read column values\n",
    "df = pd.read_csv(\"bulk_question_answering_input.csv\", header=0, dtype=str)\n",
    "\n",
    "# Make Vertex AI Search request for each question\n",
    "df.apply(lambda row: answer_questions(row, PROJECT_ID, LOCATION, DATA_STORE_ID, top_n=5), axis=1)\n",
    "\n",
    "# Output results to new TSV file\n",
    "df.to_csv(\"bulk_question_answering_output.tsv\", index=False, sep=\"\\t\")\n",
    "\n",
    "df\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "name": "question_answering.ipynb",
   "toc_visible": true
  },
  "environment": {
   "kernel": "python3",
   "name": "common-cpu.m108",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/base-cpu: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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}