{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc383f7a",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "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",
   "id": "14dada32",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "# Summarize contracts (PDF files) using Spark and Gemini"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c177340-d3b4-4e5b-a927-31bcc8f031e7",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Overview"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb23a4a5-dbba-4370-b572-7b4382973e7b",
   "metadata": {
    "tags": []
   },
   "source": [
    "This notebook shows how to perform summarization using Gemini for a large number of contract PDF files in a GCS bucket\n",
    "\n",
    "#### **Steps**\n",
    "Using Spark, \n",
    "1) It reads the table of the [Contract Understanding Atticus Dataset (CUAD)](https://www.atticusprojectai.org/cuad) dataset located in the [gs://dataproc-metastore-public-binaries/cuad_v1/full_contract_pdf/](https://console.cloud.google.com/storage/browser/dataproc-metastore-public-binaries/cuad_v1)  \n",
    "   We will create a metadata table poiting to the paths of the image files in the bucket.  \n",
    "3) It calls [Vertex AI Gemini API](https://cloud.google.com/vertex-ai/docs/generative-ai/start/quickstarts/api-quickstart#try_text_prompts) to summarize the text.\n",
    "4) It saves the output to BigQuery\n",
    "\n",
    "#### Related content\n",
    "\n",
    "- [Design summarization prompts](https://cloud.google.com/vertex-ai/docs/generative-ai/text/summarization-prompts)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "239a33c4-5d5f-4a05-8ca2-3f30ddf23464",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69fb05c5-aac5-4f7e-9949-72e9890ec303",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Identity and Access Management (IAM)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a79ac509-8480-49ff-9891-f14e9a1346ac",
   "metadata": {},
   "source": [
    "Make sure the service account running this notebook has the required permissions:\n",
    "\n",
    "- **Run the notebook**\n",
    "  - AI Platform Notebooks Service Agent\n",
    "  - Notebooks Admin\n",
    "  - Vertex AI Administrator\n",
    "- **Read files from bucket**\n",
    "  - Storage Object Viewer\n",
    "- **Run Dataproc jobs**\n",
    "  - Dataproc Service Agent\n",
    "  - Dataproc Worker\n",
    "- **Call Google APIs (Gemini)**\n",
    "  - Service Usage Consumer\n",
    "  - VisionAI Admin\n",
    "- **BigQuery**\n",
    "  - BigQuery Data Editor"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a08dc3e-5364-443b-abcb-5a65dabe5d8e",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d323b55-1d2f-46d4-b2e1-d7c6b545ba15",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import udf\n",
    "\n",
    "import google.auth\n",
    "import google.auth.transport.requests\n",
    "import requests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37a9e9fc-8b54-48a6-9389-e3ed4ce3cb6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# When using Dataproc Serverless, installed packages are automatically available on all nodes\n",
    "!pip3 install --upgrade -q google-cloud-aiplatform google-genai \"protobuf~=4.25.3\" \"numpy~=1.26.4\" \n",
    "# When using a Dataproc cluster, you will need to install these packages during cluster creation: https://cloud.google.com/dataproc/docs/tutorials/python-configuration"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ec3e8cd-9bdc-4b20-ac45-f6fc9b8acc56",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Authentication"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ceea2972-fe50-4779-8c1b-12333855d3b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get credentials to authenticate with Google APIs\n",
    "credentials, project_id = google.auth.default()\n",
    "auth_req = google.auth.transport.requests.Request()\n",
    "credentials.refresh(auth_req)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c739bba-bad6-421a-91da-1f9916359edd",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Setup Spark Session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e03a39d4-d08f-4ceb-a11b-6cea8329ccb7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql import SparkSession"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a34ce061-651b-4fb9-88bb-bd59d4eaf2e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "spark = SparkSession.builder \\\n",
    "    .appName(\"PDF files summarization using Gemini\") \\\n",
    "    .enableHiveSupport() \\\n",
    "    .getOrCreate()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07afb76e-7298-4c11-b43b-105e0fe8673f",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac1459e1-f79b-440b-901a-1c9d1645a73e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change the maximum number of files you want to consider\n",
    "limit_files = 5\n",
    "# BigQuery\n",
    "output_dataset_bq = \"output_dataset\" # create the BigQuery dataset beforehand\n",
    "output_table_bq = \"summaries\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d49c961-0d83-467f-b107-ba24ebabb57c",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Read dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14fc5d32-a882-4b5d-ad98-5089647efec1",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Read CUAD V1 dataset from metastore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a29ee53c-cc25-4d47-9f7d-90355f6ec8f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "BINARIES_BUCKET_PATH = \"gs://dataproc-metastore-public-binaries/cuad_v1/full_contract_pdf/\"\n",
    "cuad_v1_df = spark.read.format(\"binaryFile\").option(\"recursiveFileLookup\", \"true\").load(BINARIES_BUCKET_PATH).limit(limit_files)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bb0deba-7fa0-42f3-8273-fb5079dd6825",
   "metadata": {},
   "source": [
    "|                path|    modificationTime| length|             content|\n",
    "|--------------------|--------------------|-------|--------------------|\n",
    "|gs://dataproc-met...|2023-05-15 20:53:...|3683550|[25 50 44 46 2D 3...|\n",
    "|gs://dataproc-met...|2023-05-15 20:53:...|2881262|[25 50 44 46 2D 3...|\n",
    "|gs://dataproc-met...|2023-05-15 20:54:...|1778356|[25 50 44 46 2D 3...|\n",
    "|gs://dataproc-met...|2023-05-15 20:53:...|1557129|[25 50 44 46 2D 3...|\n",
    "|gs://dataproc-met...|2023-05-15 20:53:...|1452180|[25 50 44 46 2D 3...|"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0fc3682-cb19-4717-a603-9bcea4935fca",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Summarize pages using Gemini API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1678db6b-3a2e-4a42-b411-34e7dbcb95e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gemini_predict(gcs_pdf_uri, model_name=\"gemini-2.0-flash\", max_retries=3, initial_delay=1):\n",
    "    \n",
    "    import time\n",
    "    from google import genai\n",
    "    from google.genai import types\n",
    "    \n",
    "    client = genai.Client(\n",
    "        vertexai=True,\n",
    "        project=project_id,\n",
    "        location=\"us-central1\"\n",
    "    )\n",
    "    \n",
    "    generate_content_config = types.GenerateContentConfig(\n",
    "        response_mime_type = \"text/plain\"\n",
    "    )\n",
    "\n",
    "    contents = [\n",
    "        types.Part.from_uri(\n",
    "            file_uri=gcs_pdf_uri,\n",
    "            mime_type='application/pdf',\n",
    "        ),\n",
    "        \"\"\" You an expert in reading contracts, articles, agreements, or text in general.\n",
    "            You are able to create concise summaries of the text provided to you.\n",
    "            Provide a summary about the attached pdf with about 3 sentences with the most important information from the text.\n",
    "            Summary:\n",
    "        \"\"\"\n",
    "  ]\n",
    "    \n",
    "    retries, delay = 0, initial_delay\n",
    "    while retries <= max_retries:\n",
    "        try:\n",
    "            response = client.models.generate_content(model=model_name,\n",
    "                                                      contents=contents,\n",
    "                                                      config=generate_content_config)\n",
    "            \n",
    "            return response.text\n",
    "        except Exception:\n",
    "            if retries == max_retries:\n",
    "                return\n",
    "            time.sleep(delay)\n",
    "            delay *= 2\n",
    "            retries += 1\n",
    "    return \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4024f4bb-874b-4907-9cc0-8a393be6f5f7",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "summarize_text = udf(gemini_predict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3184c463-6deb-42e8-9696-007d2c52f7f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "summaries_df = cuad_v1_df.withColumn(\"summary\", summarize_text(cuad_v1_df[\"path\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68b62926-1f4f-40ed-bb60-21e81cba905b",
   "metadata": {},
   "outputs": [],
   "source": [
    "summaries_df.show(5,50)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5dd1df3-02ac-4671-8668-c3bddd6e2de2",
   "metadata": {},
   "source": [
    "|                                              path|       modificationTime| length|                                           content|                                           summary|\n",
    "|--------------------------------------------------|-----------------------|-------|--------------------------------------------------|--------------------------------------------------|\n",
    "|gs://dataproc-metastore-public-binaries/cuad_v1...|2023-05-15 20:53:55.891|3683550|[25 50 44 46 2D 31 2E 34 0A 25 E2 E3 CF D3 0A 3...|Here is a summary of the provided document:\\n\\n...|\n",
    "|gs://dataproc-metastore-public-binaries/cuad_v1...|2023-05-15 20:53:57.195|2881262|[25 50 44 46 2D 31 2E 35 0A 25 E2 E3 CF D3 0A 0...|This document is a promotion and distribution a...|\n",
    "|gs://dataproc-metastore-public-binaries/cuad_v1...|2023-05-15 20:54:00.609|1778356|[25 50 44 46 2D 31 2E 35 0A 25 E2 E3 CF D3 0A 0...|This document is a strategic alliance agreement...|\n",
    "|gs://dataproc-metastore-public-binaries/cuad_v1...|2023-05-15 20:53:57.902|1557129|[25 50 44 46 2D 31 2E 35 0A 25 E2 E3 CF D3 0A 0...|This PDF is a collaboration agreement between t...|\n",
    "|gs://dataproc-metastore-public-binaries/cuad_v1...|2023-05-15 20:53:57.659|1452180|[25 50 44 46 2D 31 2E 34 0D 25 C8 C8 C8 C8 C8 C...|This is a Transportation Services Agreement bet...|"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac92387a-56e9-45aa-a741-9d2c0f66d866",
   "metadata": {},
   "source": [
    "## Save to BigQuery"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6c4d2c4-bb20-4575-96f9-48fbba382cd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "summaries_df.write \\\n",
    "            .format(\"bigquery\") \\\n",
    "            .option(\"table\", f\"{project_id}.{output_dataset_bq}.{output_table_bq}\") \\\n",
    "            .option(\"writeMethod\", \"direct\") \\\n",
    "            .mode(\"overwrite\") \\\n",
    "            .save()"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "9c39b79e5d2e7072beb4bd59-runtime",
   "name": "workbench-notebooks.m129",
   "type": "gcloud",
   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m129"
  },
  "kernelspec": {
   "display_name": "runtime on Serverless Spark (Remote)",
   "language": "python",
   "name": "9c39b79e5d2e7072beb4bd59-runtime"
  },
  "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}