943 lines
35 KiB
Plaintext
943 lines
35 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "ur8xi4C7S06n"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Copyright 2024 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": [
|
|
"# Run RAG Pipelines in BigQuery with BQML and Vector Search\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/gemini/use-cases/retrieval-augmented-generation/rag_with_bigquery.ipynb\">\n",
|
|
" <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
|
|
" </a>\n",
|
|
" </td>\n",
|
|
" <td style=\"text-align: center\">\n",
|
|
" <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Fuse-cases%2Fretrieval-augmented-generation%2Frag_with_bigquery.ipynb\">\n",
|
|
" <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\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/main/gemini/use-cases/retrieval-augmented-generation/rag_with_bigquery.ipynb\">\n",
|
|
" <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
|
|
" </a>\n",
|
|
" </td>\n",
|
|
" <td style=\"text-align: center\">\n",
|
|
" <a href=\"https://console.cloud.google.com/bigquery/import?url=https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/retrieval-augmented-generation/rag_with_bigquery.ipynb\">\n",
|
|
" <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/bigquery/v1/32px.svg\" alt=\"BigQuery Studio logo\"><br> Open in BigQuery Studio\n",
|
|
" </a>\n",
|
|
" </td>\n",
|
|
" <td style=\"text-align: center\">\n",
|
|
" <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/retrieval-augmented-generation/rag_with_bigquery.ipynb\">\n",
|
|
" <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
|
|
" </a>\n",
|
|
" </td>\n",
|
|
"</table>\n",
|
|
"\n",
|
|
"<div style=\"clear: both;\"></div>\n",
|
|
"\n",
|
|
"<b>Share to:</b>\n",
|
|
"\n",
|
|
"<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/retrieval-augmented-generation/rag_with_bigquery.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
|
|
"</a>\n",
|
|
"\n",
|
|
"<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/retrieval-augmented-generation/rag_with_bigquery.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
|
|
"</a>\n",
|
|
"\n",
|
|
"<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/retrieval-augmented-generation/rag_with_bigquery.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
|
|
"</a>\n",
|
|
"\n",
|
|
"<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/retrieval-augmented-generation/rag_with_bigquery.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
|
|
"</a>\n",
|
|
"\n",
|
|
"<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/retrieval-augmented-generation/rag_with_bigquery.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
|
|
"</a> "
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "84f0f73a0f76"
|
|
},
|
|
"source": [
|
|
"| | |\n",
|
|
"|-|-|\n",
|
|
"| Author(s) | [Jeff Nelson](https://github.com/jeffonelson/), Eric Hao |"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "tvgnzT1CKxrO"
|
|
},
|
|
"source": [
|
|
"## Overview\n",
|
|
"\n",
|
|
"This notebook demonstrates a basic end-to-end retrieval-augmented generation (RAG) pipeline using [BigQuery](https://cloud.google.com/bigquery/) and [BigQuery ML](https://cloud.google.com/bigquery/docs/bqml-introduction) functions. To do so, we:\n",
|
|
"\n",
|
|
"* Complete setup steps to download sample data and access [Vertex AI](https://cloud.google.com/vertex-ai) from BigQuery\n",
|
|
"* Generate [object table](https://cloud.google.com/bigquery/docs/object-table-introduction) to access unstructured PDFs that reside in [Cloud Storage](https://cloud.google.com/storage)\n",
|
|
"* Create a remote model, so BigQuery can call [Document AI](https://cloud.google.com/document-ai) to parse the PDF inputs\n",
|
|
"* Parse response from Document AI into chunks and metadata, then generate vector embeddings for the chunks\n",
|
|
"* Run a [vector search](https://cloud.google.com/bigquery/docs/vector-search) against embeddings in BigQuery, return relevant chunks, and summarize them with Gemini"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "dc949afc1f08"
|
|
},
|
|
"source": [
|
|
"## How to open this notebook in BigQuery Studio\n",
|
|
"\n",
|
|
"This notebook was written to be compatible for use within BigQuery Studio. To open this notebook in BigQuery, click to [Run in Colab Enterprise](https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Fuse-cases%2Fretrieval-augmented-generation%2Frag_with_bigquery.ipynb). This will open a new window in the Cloud Console and prompt you to confirm import. Then, navigate to BigQuery, where you will find the notebook available in the Explorer pane under Notebooks."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "5ba5c12e483d"
|
|
},
|
|
"source": [
|
|
"## About the dataset\n",
|
|
"\n",
|
|
"This example uses [the Federal Reserve's 2023 Survey of Consumer Finances](https://www.federalreserve.gov/publications/files/scf23.pdf) (SCF) report. The document contains information around US family income, net worth, credit use, and other common household financial indicators."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "2ce33dbc8fde"
|
|
},
|
|
"source": [
|
|
"## Services and Costs\n",
|
|
"\n",
|
|
"This tutorial uses the following Google Cloud data analytics and ML services, they are billable components of Google Cloud:\n",
|
|
"\n",
|
|
"* BigQuery & BigQuery ML [(pricing)](https://cloud.google.com/bigquery/pricing)\n",
|
|
"* Vertex AI Generative AI models [(pricing)](https://cloud.google.com/vertex-ai/generative-ai/pricing)\n",
|
|
"* Document AI [(pricing)](https://cloud.google.com/document-ai/pricing)\n",
|
|
"* Cloud Storage [(pricing)](https://cloud.google.com/storage/pricing)\n",
|
|
"\n",
|
|
"Use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "61RBz8LLbxCR"
|
|
},
|
|
"source": [
|
|
"# Setup Steps to access Vertex AI models from BigQuery and enable APIs"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "ff210a6d4d21"
|
|
},
|
|
"source": [
|
|
"### Install Document AI SDK"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "2e9e2b9e1b1f"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"%pip install --quiet google-cloud-documentai==2.31.0"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "8ed31279f009"
|
|
},
|
|
"source": [
|
|
"### Restart runtime\n",
|
|
"\n",
|
|
"To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which restarts the current kernel.\n",
|
|
"\n",
|
|
"The restart might take a minute or longer. After it's restarted, continue to the next step."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "567212ff53a6"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"import IPython\n",
|
|
"\n",
|
|
"app = IPython.Application.instance()\n",
|
|
"app.kernel.do_shutdown(True)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "b96b39fd4d7b"
|
|
},
|
|
"source": [
|
|
"<div class=\"alert alert-block alert-warning\">\n",
|
|
"<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
|
|
"</div>"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "fa362c2ef5b5"
|
|
},
|
|
"source": [
|
|
"### Authenticate your notebook environment (Colab only)\n",
|
|
"\n",
|
|
"If you're running this notebook on Google Colab, run the cell below to authenticate your environment."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 3,
|
|
"metadata": {
|
|
"id": "9a07a9f9a4a9"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"import sys\n",
|
|
"\n",
|
|
"if \"google.colab\" in sys.modules:\n",
|
|
" from google.colab import auth\n",
|
|
"\n",
|
|
" auth.authenticate_user()\n",
|
|
" print(\"Authenticated\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "No17Cw5hgx12"
|
|
},
|
|
"source": [
|
|
"### Define your Google Cloud project"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 4,
|
|
"metadata": {
|
|
"id": "tFy3H3aPgx12"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"PROJECT_ID = \"your-project-id\" # @param {type: \"string\"}\n",
|
|
"PROJECT_NUMBER = \"your-project-number\" # @param {type: \"string\"}"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "04deeb11bbca"
|
|
},
|
|
"source": [
|
|
"### Enable Data Table Display\n",
|
|
"\n",
|
|
"This makes it easier to visualize tabular data within a Notebook environment later on."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "af9974f04f9f"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"%load_ext google.colab.data_table"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "b4256d07d596"
|
|
},
|
|
"source": [
|
|
"### Create a new dataset in BigQuery\n",
|
|
"\n",
|
|
"This will house any tables created throughout this notebook."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "8a4c1a356d10"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"!bq mk --location=us --dataset --project_id={PROJECT_ID} docai_demo"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "a100b689816b"
|
|
},
|
|
"source": [
|
|
"### Create a Cloud resource connection\n",
|
|
"\n",
|
|
"[Cloud resource connections](https://cloud.google.com/bigquery/docs/create-cloud-resource-connection) enable BigQuery to access other Cloud services, like Cloud Storage and Vertex AI."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "885da43402b1"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"!bq mk --connection --connection_type=CLOUD_RESOURCE --location=us --project_id={PROJECT_ID} \"demo_conn\"\n",
|
|
"!bq show --location=us --connection --project_id={PROJECT_ID} \"demo_conn\""
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "dd9f6cbe4393"
|
|
},
|
|
"source": [
|
|
"### Add permissions to Cloud resource connection service account\n",
|
|
"\n",
|
|
"The Cloud resource connection is associated with a service account. The following cell enables the service account to access services like Document AI, Cloud Storage, and Vertex AI.\n",
|
|
"\n",
|
|
"**Note:** Copy the service account ID from the prior cell and input it below. It will look like `your-copied-service-account@gcp-sa-bigquery-condel.iam.gserviceaccount.com`."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "16b193a840cd"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"connection_service_account = \"your-copied-service-account@gcp-sa-bigquery-condel.iam.gserviceaccount.com\" # @param {type: \"string\"}\n",
|
|
"connection_member = f\"serviceAccount:{connection_service_account}\"\n",
|
|
"\n",
|
|
"\n",
|
|
"!gcloud projects add-iam-policy-binding {PROJECT_ID} --member={connection_member} --role='roles/documentai.viewer' --condition=None --quiet\n",
|
|
"!gcloud projects add-iam-policy-binding {PROJECT_ID} --member={connection_member} --role='roles/storage.objectViewer' --condition=None --quiet\n",
|
|
"!gcloud projects add-iam-policy-binding {PROJECT_ID} --member={connection_member} --role='roles/aiplatform.user' --condition=None --quiet"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "ba09d9393559"
|
|
},
|
|
"source": [
|
|
"### Download the sample PDF used for this notebook and store it in a new Cloud Storage bucket"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "4605453a6675"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"import random\n",
|
|
"\n",
|
|
"# Create a unique Cloud Storage bucket name\n",
|
|
"bucket_name = f\"{PROJECT_ID}-{random.randint(10000, 99999)}\"\n",
|
|
"\n",
|
|
"# Create the bucket\n",
|
|
"!gsutil mb -l US -p {PROJECT_ID} gs://{bucket_name}\n",
|
|
"\n",
|
|
"# Download the PDF sample\n",
|
|
"!wget scf23.pdf \"https://www.federalreserve.gov/publications/files/scf23.pdf\"\n",
|
|
"\n",
|
|
"# Upload the PDF sample to the newly created Cloud Storage bucket\n",
|
|
"!gsutil cp scf23.pdf gs://{bucket_name}/\n",
|
|
"\n",
|
|
"# Print confirmation\n",
|
|
"print(f\"PDF uploaded to gs://{bucket_name}/scf23.pdf\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "6b0a0bd5c4fd"
|
|
},
|
|
"source": [
|
|
"## Create an object table\n",
|
|
"\n",
|
|
"An object table allows BigQuery to read unstructured data in Google Cloud Storage. This uses the BigQuery Python client library to continue using the `bucket_name` variable."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "94cc075094c4"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"from google.cloud import bigquery\n",
|
|
"\n",
|
|
"client = bigquery.Client(project=PROJECT_ID)\n",
|
|
"\n",
|
|
"query = f\"\"\"\n",
|
|
"CREATE OR REPLACE EXTERNAL TABLE `docai_demo.object_table`\n",
|
|
"WITH CONNECTION `us.demo_conn` -- Replace with your connection ID\n",
|
|
"OPTIONS (\n",
|
|
" uris = ['gs://{bucket_name}/scf23.pdf'],\n",
|
|
" object_metadata = 'DIRECTORY'\n",
|
|
");\n",
|
|
"\"\"\"\n",
|
|
"\n",
|
|
"query_job = client.query(query) # API request\n",
|
|
"query_job.result() # Waits for the query to complete\n",
|
|
"\n",
|
|
"print(\"External table docai_demo.object_table created or replaced successfully.\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "c17ec8736188"
|
|
},
|
|
"source": [
|
|
"### Show the object table\n",
|
|
"\n",
|
|
"Confirm that the results display the PDF document in your Cloud Storage bucket."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "9f471aa348b2"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"%%bigquery --project $PROJECT_ID\n",
|
|
"\n",
|
|
"SELECT * \n",
|
|
"FROM `docai_demo.object_table`;"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "ec9d2c49fd34"
|
|
},
|
|
"source": [
|
|
"## Use BQML and Document AI to parse documents"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "afc0a5902ef3"
|
|
},
|
|
"source": [
|
|
"### Create a Layout Parser Processor in Document AI\n",
|
|
"\n",
|
|
"[Create a new processor](https://cloud.google.com/document-ai/docs/create-processor#documentai_fetch_processor_types-python) in Document AI with the type `LAYOUT_PARSER_PROCESSOR`."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "519ea8a55496"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"from google.api_core.client_options import ClientOptions\n",
|
|
"from google.cloud import documentai\n",
|
|
"\n",
|
|
"location = \"us\"\n",
|
|
"processor_display_name = \"layout_parser_processor\"\n",
|
|
"processor_type = \"LAYOUT_PARSER_PROCESSOR\"\n",
|
|
"\n",
|
|
"\n",
|
|
"def create_processor_sample(\n",
|
|
" PROJECT_ID: str, location: str, processor_display_name: str, processor_type: str\n",
|
|
") -> None:\n",
|
|
" opts = ClientOptions(api_endpoint=f\"{location}-documentai.googleapis.com\")\n",
|
|
"\n",
|
|
" client = documentai.DocumentProcessorServiceClient(client_options=opts)\n",
|
|
"\n",
|
|
" # The full resource name of the location\n",
|
|
" parent = client.common_location_path(PROJECT_ID, location)\n",
|
|
"\n",
|
|
" # Create a processor\n",
|
|
" processor = client.create_processor(\n",
|
|
" parent=parent,\n",
|
|
" processor=documentai.Processor(\n",
|
|
" display_name=processor_display_name, type_=processor_type\n",
|
|
" ),\n",
|
|
" )\n",
|
|
"\n",
|
|
" # Return the processor ID needed for creating a BigQuery connection\n",
|
|
" return processor.name.split(\"/\")[-1]\n",
|
|
"\n",
|
|
"\n",
|
|
"# Call this function to create the processor and return its ID\n",
|
|
"processor_id = create_processor_sample(\n",
|
|
" PROJECT_ID, location, processor_display_name, processor_type\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "5cfcaa4c4584"
|
|
},
|
|
"source": [
|
|
"### Create a remote model in BigQuery that connects with your Document AI Layout Parser Processor\n",
|
|
"\n",
|
|
"This one-time setup step allows BigQuery to reference the Document AI Processor you just created.\n",
|
|
"\n",
|
|
"**Note:** If if you receive an 400 GET error \"permission denied for document processor\", you may need to wait a minute for permissions to propagate from earlier steps."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "57233716c232"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"query = f\"\"\"\n",
|
|
"CREATE OR REPLACE MODEL `docai_demo.layout_parser` \n",
|
|
"REMOTE WITH CONNECTION `us.demo_conn`\n",
|
|
"OPTIONS(remote_service_type=\"CLOUD_AI_DOCUMENT_V1\", document_processor=\"{processor_id}\")\n",
|
|
"\"\"\"\n",
|
|
"\n",
|
|
"query_job = client.query(query) # API request\n",
|
|
"query_job.result() # Waits for the query to complete\n",
|
|
"\n",
|
|
"print(\"Remote model docai_demo.layout_parser created or replaced successfully.\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "fc821b6f845d"
|
|
},
|
|
"source": [
|
|
"### Process the document using BigQuery ML\n",
|
|
"\n",
|
|
"Use the [`ML.PROCESS_DOCUMENT` function](https://cloud.google.com/bigquery/docs/process-document) from BigQuery to call your Document AI processor and pass through the PDF. This uses the Layout Parser configuration and chunks your document.\n",
|
|
"\n",
|
|
"**Note:** this may take a minute or so to complete."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "a489b3cb9e1d"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"%%bigquery --project $PROJECT_ID --location us\n",
|
|
"\n",
|
|
"CREATE or REPLACE TABLE docai_demo.demo_result AS (\n",
|
|
" SELECT * FROM ML.PROCESS_DOCUMENT(\n",
|
|
" MODEL docai_demo.layout_parser,\n",
|
|
" TABLE docai_demo.object_table,\n",
|
|
" PROCESS_OPTIONS => (JSON '{\"layout_config\": {\"chunking_config\": {\"chunk_size\": 250}}}')\n",
|
|
" )\n",
|
|
");"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "bfb0a1fa3266"
|
|
},
|
|
"source": [
|
|
"### Parse the JSON results returned to BigQuery\n",
|
|
"\n",
|
|
"The `ML.PROCESS_DOCUMENT` function parses the PDF from Cloud Storage and returns a JSON blob to BigQuery. In this step, we'll parse the JSON, extract document chunks and metadata, and return it to a new BigQuery table."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "2bc4dad2e399"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"%%bigquery --project $PROJECT_ID --location us\n",
|
|
"\n",
|
|
"CREATE OR REPLACE TABLE docai_demo.demo_result_parsed AS (\n",
|
|
"SELECT\n",
|
|
" uri,\n",
|
|
" JSON_EXTRACT_SCALAR(json , '$.chunkId') AS id,\n",
|
|
" JSON_EXTRACT_SCALAR(json , '$.content') AS content,\n",
|
|
" JSON_EXTRACT_SCALAR(json , '$.pageFooters[0].text') AS page_footers_text,\n",
|
|
" JSON_EXTRACT_SCALAR(json , '$.pageSpan.pageStart') AS page_span_start,\n",
|
|
" JSON_EXTRACT_SCALAR(json , '$.pageSpan.pageEnd') AS page_span_end\n",
|
|
"FROM docai_demo.demo_result, UNNEST(JSON_EXTRACT_ARRAY(ml_process_document_result.chunkedDocument.chunks, '$')) json\n",
|
|
");"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "adca53cc55d8"
|
|
},
|
|
"source": [
|
|
"### Display the parsed document chunks\n",
|
|
"\n",
|
|
"Show a preview of the parsed results and metadata."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "c60bcdc388c4"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"%%bigquery --project $PROJECT_ID --location us\n",
|
|
"\n",
|
|
"SELECT *\n",
|
|
"FROM docai_demo.demo_result_parsed\n",
|
|
"ORDER BY id\n",
|
|
"LIMIT 5;"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "a980e66443bc"
|
|
},
|
|
"source": [
|
|
"## Connect to Vertex AI embedding generation and Gemini access"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "eae56fa8c74c"
|
|
},
|
|
"source": [
|
|
"### Connect to a text embedding model\n",
|
|
"\n",
|
|
"[Create a remote model](https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-create-remote-model) allowing BigQuery access to a text embedding model hosted in Vertex AI."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "3c53a24e59a1"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"%%bigquery --project $PROJECT_ID\n",
|
|
"\n",
|
|
"CREATE OR REPLACE MODEL `docai_demo.embedding_model` \n",
|
|
"REMOTE WITH CONNECTION `us.demo_conn` OPTIONS(endpoint=\"text-embedding-005\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "45d1ccc016c8"
|
|
},
|
|
"source": [
|
|
"### Generate embeddings\n",
|
|
"\n",
|
|
"Use the [`ML.GENERATE_EMBEDDING` function](https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-generate-embedding) in BigQuery to generate embeddings for all text chunks in the document."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "63bf77f48b8c"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"%%bigquery --project $PROJECT_ID\n",
|
|
"\n",
|
|
"CREATE OR REPLACE TABLE `docai_demo.embeddings` AS\n",
|
|
"SELECT * FROM ML.GENERATE_EMBEDDING(\n",
|
|
" MODEL `docai_demo.embedding_model`,\n",
|
|
" TABLE `docai_demo.demo_result_parsed`\n",
|
|
");"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "e1ce3b78e01a"
|
|
},
|
|
"source": [
|
|
"### Connect to a Gemini LLM endpoint\n",
|
|
"\n",
|
|
"Create a remote model allowing BigQuery access to a Gemini foundation model hosted in Vertex AI."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "7b760c54502e"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"%%bigquery --project $PROJECT_ID\n",
|
|
"\n",
|
|
"CREATE OR REPLACE MODEL `docai_demo.gemini_flash` REMOTE\n",
|
|
"WITH CONNECTION `us.demo_conn` OPTIONS(endpoint=\"gemini-2.0-flash\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "879593c348c4"
|
|
},
|
|
"source": [
|
|
"## Run vector search, return results, and pass them to Gemini for text generation"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "2eb640b369a3"
|
|
},
|
|
"source": [
|
|
"### Sample BigQuery vector search\n",
|
|
"\n",
|
|
"Run a sample BigQuery vector search against your chunks. This query takes your text input, creates an embedding using the `ML.GENERATE_EMBEDDING` function, and then passes the embedding through to the [`VECTOR_SEARCH` function](https://cloud.google.com/bigquery/docs/reference/standard-sql/search_functions#vector_search). The results are the top ten chunks that are most semantically related to your input.\n",
|
|
"\n",
|
|
"In the search query below, the input text asks \"Did the typical family net worth increase? If so, by how much?\""
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "cf9fa689905d"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"%%bigquery --project $PROJECT_ID\n",
|
|
"\n",
|
|
"SELECT query.query, base.uri, base.id, base.content, distance\n",
|
|
" FROM\n",
|
|
" VECTOR_SEARCH( TABLE `docai_demo.embeddings`,\n",
|
|
" 'ml_generate_embedding_result',\n",
|
|
" (\n",
|
|
" SELECT\n",
|
|
" ml_generate_embedding_result,\n",
|
|
" content AS query\n",
|
|
" FROM\n",
|
|
" ML.GENERATE_EMBEDDING( MODEL `docai_demo.embedding_model`,\n",
|
|
" ( SELECT 'Did the typical family net worth increase? If so, by how much?' AS content)\n",
|
|
" ) \n",
|
|
" ),\n",
|
|
" top_k => 10,\n",
|
|
" OPTIONS => '{\"fraction_lists_to_search\": 0.01}') \n",
|
|
"ORDER BY distance DESC;"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "623765bd9154"
|
|
},
|
|
"source": [
|
|
"## Generate text augmented by vector search results\n",
|
|
"\n",
|
|
"This step builds upon the prior one - but instead of simply returning the top text chunks, it calls the `ML.GENERATE_TEXT` function to summarize them alongside the question we input.\n",
|
|
"\n",
|
|
"In this query you:\n",
|
|
"* **Retrieve** the closest chunks semantically using the `VECTOR_SEARCH` function (this is what was done in the prior query)\n",
|
|
"* **Augment** the Gemini LLM with this knowledge\n",
|
|
"* **Generate** a succinct answer using the `ML.GENERATE_TEXT` function"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "2f6f83f2eca7"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"%%bigquery --project $PROJECT_ID\n",
|
|
"\n",
|
|
"SELECT\n",
|
|
" ml_generate_text_llm_result AS generated,\n",
|
|
" -- prompt -- Commented out, but please feel free to uncomment if you would like to see the full context passed to the Gemini model\n",
|
|
"FROM\n",
|
|
" ML.GENERATE_TEXT( MODEL `docai_demo.gemini_flash`,\n",
|
|
" (\n",
|
|
" SELECT\n",
|
|
" CONCAT( 'Did the typical family net worth change? How does this compare the SCF survey a decade earlier? Be concise and use the following context:',\n",
|
|
" STRING_AGG(FORMAT(\"context: %s and reference: %s\", base.content, base.uri), ',\\n')) AS prompt,\n",
|
|
" FROM\n",
|
|
" VECTOR_SEARCH( TABLE \n",
|
|
" `docai_demo.embeddings`,\n",
|
|
" 'ml_generate_embedding_result',\n",
|
|
" (\n",
|
|
" SELECT\n",
|
|
" ml_generate_embedding_result,\n",
|
|
" content AS query\n",
|
|
" FROM\n",
|
|
" ML.GENERATE_EMBEDDING( MODEL `docai_demo.embedding_model`,\n",
|
|
" (\n",
|
|
" SELECT\n",
|
|
" 'Did the typical family net worth change? How does this compare the SCF survey a decade earlier?' AS content\n",
|
|
" )\n",
|
|
" ) \n",
|
|
" ),\n",
|
|
" top_k => 10,\n",
|
|
" OPTIONS => '{\"fraction_lists_to_search\": 0.01}') \n",
|
|
" ),\n",
|
|
" STRUCT(512 AS max_output_tokens, TRUE AS flatten_json_output)\n",
|
|
" );"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "945bbb4357b2"
|
|
},
|
|
"source": [
|
|
"### Sample questions to try out:\n",
|
|
"\n",
|
|
"Here are a list of a few other questions to spark your imagination. Feel free to try your own too!\n",
|
|
"* Did the amount of debt families own on their home increase between 2019 and 2022?\n",
|
|
"* Did younger or older families see their net worth increase more?\n",
|
|
"* How much did the median family income change between 2018 and 2021?"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "2a4e033321ad"
|
|
},
|
|
"source": [
|
|
"# Cleaning up\n",
|
|
"\n",
|
|
"To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
|
|
"\n",
|
|
"Otherwise, you can delete the individual resources you created in this tutorial by uncommenting the below:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "1ab59128be6c"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# # Deletes the BigQuery assets and Google Cloud Storage bucket\n",
|
|
"\n",
|
|
"# !bq rm -r -f $PROJECT_ID:docai_demo\n",
|
|
"# !bq rm --connection --project_id=$PROJECT_ID --location=us demo_conn\n",
|
|
"# !gsutil rm -r gs://{bucket_name}\n",
|
|
"\n",
|
|
"\n",
|
|
"# # Deletes the Document AI processor\n",
|
|
"# def delete_processor_sample(\n",
|
|
"# PROJECT_ID: str, location: str, processor_id: str\n",
|
|
"# ) -> None:\n",
|
|
"# \"\"\"Deletes a processor.\"\"\"\n",
|
|
"\n",
|
|
"# opts = ClientOptions(api_endpoint=f\"{location}-documentai.googleapis.com\")\n",
|
|
"\n",
|
|
"# client = documentai.DocumentProcessorServiceClient(client_options=opts)\n",
|
|
"\n",
|
|
"# # The full resource name of the processor\n",
|
|
"# name = f\"projects/{PROJECT_ID}/locations/{location}/processors/{processor_id}\"\n",
|
|
"\n",
|
|
"# try:\n",
|
|
"# client.delete_processor(name=name)\n",
|
|
"# print(f\"Processor {processor_id} deleted successfully.\")\n",
|
|
"# except Exception as e:\n",
|
|
"# print(f\"Error deleting processor: {e}\")\n",
|
|
"\n",
|
|
"\n",
|
|
"# # Call this function to delete the processor\n",
|
|
"# delete_processor_sample(PROJECT_ID, location, processor_id)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "aefa89207b70"
|
|
},
|
|
"source": [
|
|
"# Wrap up\n",
|
|
"\n",
|
|
"This notebook demonstrates an example of how to achieve a basic end-to-end retrieval-augmented generation pipeline using BigQuery. It integrates BigQuery ML functions like `ML.PROCESS_DOCUMENT` to call Document AI and parse PDFs, `ML.GENERATE_EMBEDDING` to generate embeddings on text chunks and input queries, and `ML.GENERATE_TEXT` to provide a concise answer. It also uses the `VECTOR_SEARCH` function to identify similar text (using embeddings) in BigQuery using familiar SQL syntax.\n",
|
|
"\n",
|
|
"To continue learn more, check out our documentation on [BigQuery ML](https://cloud.google.com/bigquery/docs/bqml-introduction) and [BigQuery Vector Search](https://cloud.google.com/bigquery/docs/vector-search)."
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"colab": {
|
|
"name": "rag_with_bigquery.ipynb",
|
|
"toc_visible": true
|
|
},
|
|
"kernelspec": {
|
|
"display_name": "Python 3",
|
|
"name": "python3"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 0
|
|
}
|