diff --git a/notebooks/official/matching_engine/sdk_matching_engine_create_stack_overflow_embeddings.ipynb b/notebooks/official/matching_engine/sdk_matching_engine_create_stack_overflow_embeddings.ipynb index e6ea08aa6..80996339b 100644 --- a/notebooks/official/matching_engine/sdk_matching_engine_create_stack_overflow_embeddings.ipynb +++ b/notebooks/official/matching_engine/sdk_matching_engine_create_stack_overflow_embeddings.ipynb @@ -1,909 +1,1111 @@ { - "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": [ - "# Using Vertex AI Matching Engine for StackOverflow Questions\n", - "\n", - "\n", - " \n", - " \n", - " \n", - "
\n", - " \n", - " \"Colab Run in Colab\n", - " \n", - " \n", - " \n", - " \"GitHub\n", - " View on GitHub\n", - " \n", - " \n", - " \n", - " \"Vertex\n", - " Open in Vertex AI Workbench\n", - " \n", - "
" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "b0a74aaf1481" - }, - "source": [ - "## Overview\n", - "\n", - "This example demonstrates how to create text-to-image embeddings using the StackOverflow dataset and the sentence-T5 model. These are uploaded to Vertex AI Matching Engine service. It is a high scale, low latency solution, to find similar vectors (or more specifically \"embeddings\") for a large corpus. Moreover, it is a fully managed offering, further reducing operational overhead. It is built upon [Approximate Nearest Neighbor (ANN) technology](https://ai.googleblog.com/2020/07/announcing-scann-efficient-vector.html) developed by Google Research.\n", - "\n", - "**Pre-requisite**: This notebook requires you to already have a VPC network set up. See the \"Prepare a VPC network\" section in [Create Vertex AI Matching Engine index notebook](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/matching_engine/sdk_matching_engine_for_indexing.ipynb).\n", - "\n", - "Learn more about [Vertex AI Matching Engine](https://cloud.google.com/vertex-ai/docs/matching-engine/overview)." - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "34a4b245e795" - }, - "source": [ - "### Objective\n", - "\n", - "In this notebook, you learn how to encode custom text embeddings, create an Approximate Nearest Neighbor (ANN) Index, and query against indexes.\n", - "\n", - "This tutorial uses the following Google Cloud ML services:\n", - "\n", - "- `Vertex AI Matching Engine`\n", - "\n", - "The steps performed include:\n", - "\n", - "* Create ANN Index\n", - "* Create an IndexEndpoint with VPC Network\n", - "* Deploy ANN Index\n", - "* Perform online query\n" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "tvgnzT1CKxrO" - }, - "source": [ - "### Dataset\n", - "\n", - "The dataset used for this tutorial is the [StackOverflow dataset](https://console.cloud.google.com/marketplace/product/stack-exchange/stack-overflow).\n", - "\n", - "> Stack Overflow is the largest online community for programmers to learn, share their knowledge, and advance their careers. Updated on a quarterly basis, this BigQuery dataset includes an archive of Stack Overflow content, including posts, votes, tags, and badges. This dataset is updated to mirror the Stack Overflow content on the Internet Archive, and is also available through the Stack Exchange Data Explorer." - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "f0f1bea346db" - }, - "source": [ - "## Installation\n", - "\n", - "Install the latest version of Cloud Storage, BigQuery and Vertex AI SDKs for Python." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "dfbccc635a17" - }, - "outputs": [], - "source": [ - "# Install the packages\n", - "! pip3 install --upgrade google-cloud-aiplatform \\\n", - " google-cloud-storage --upgrade" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "5b08ba354c6e" - }, - "source": [ - "### Colab only: Uncomment the following cell to restart the kernel." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "bea801acf6b5" - }, - "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": "dd28c9e4f067" - }, - "source": [ - "## Before you begin\n", - "#### Set your project ID\n", - "\n", - "**If you don't know your project ID**, try the following:\n", - "* Run `gcloud config list`.\n", - "* Run `gcloud projects list`.\n", - "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "80c0215f05a0" - }, - "outputs": [], - "source": [ - "PROJECT_ID = \"[your-project-id]\" # @param {type:\"string\"}\n", - "\n", - "# Set the project id\n", - "! gcloud config set project {PROJECT_ID}" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "4f4512bf63b3" - }, - "source": [ - "#### Region\n", - "\n", - "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "474be5183c27" - }, - "outputs": [], - "source": [ - "REGION = \"us-central1\" # @param {type: \"string\"}" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "949271bfebe3" - }, - "source": [ - "### Authenticate your Google Cloud account\n", - "\n", - "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below." - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "b65b4ce80d9a" - }, - "source": [ - "**1. Vertex AI Workbench**\n", - "* Do nothing as you are already authenticated." - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "985cdbfe7372" - }, - "source": [ - "**2. Local JupyterLab instance, uncomment and run:**" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "fbc9cd30cc4b" - }, - "outputs": [], - "source": [ - "# ! gcloud auth login" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "79efab26ad02" - }, - "source": [ - "**3. Colab, uncomment and run:**" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "a336a05c6149" - }, - "outputs": [], - "source": [ - "# from google.colab import auth\n", - "# auth.authenticate_user()" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "0c0a44fa330f" - }, - "source": [ - "**4. Service account or other**\n", - "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples." - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "d3uj8x73nDX_" - }, - "source": [ - "* Authentication: Rerun the `gcloud auth login` command in the Vertex AI Workbench notebook terminal when you are logged out and need the credential again." - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "hhq5zEbGg0XX" - }, - "source": [ - "### Colab only: Uncomment the following cell to restart the kernel." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "EzrelQZ22IZj" - }, - "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": "zgPO1eR3CYjk" - }, - "source": [ - "### Create a Cloud Storage bucket\n", - "\n", - "Create a storage bucket to store intermediate artifacts such as datasets." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "MzGDU7TWdts_" - }, - "outputs": [], - "source": [ - "BUCKET_URI = \"gs://your-bucket-name-unique\" # @param {type:\"string\"}" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "-EcIXiGsCePi" - }, - "source": [ - "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "NIq7R4HZCfIc" - }, - "outputs": [], - "source": [ - "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "lR6Wwv-hCCN-" - }, - "source": [ - "## Prepare the data\n", - "\n", - "You will use [Stack Overflow dataset](https://console.cloud.google.com/marketplace/product/stack-exchange/stack-overflow) of question and answers hosted on BigQuery.\n", - "\n", - "> This public dataset is hosted in Google BigQuery and is included in BigQuery's 1TB/mo of free tier processing. This means that each user receives 1TB of free BigQuery processing every month, which can be used to run queries on this public dataset." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "9wzS85TeB9dG" - }, - "outputs": [], - "source": [ - "%%time\n", - "from google.cloud import bigquery\n", - "\n", - "client = bigquery.Client(project=PROJECT_ID)\n", - "\n", - "NUM_ROWS = 1000\n", - "\n", - "QUERY = f\"\"\"\n", - " SELECT distinct q.id, q.title, q.body, q.tags, a.body as answers, a.score \n", - " FROM (SELECT * FROM `bigquery-public-data.stackoverflow.posts_questions` where Score>0 ORDER BY View_Count desc) AS q \n", - " INNER JOIN (SELECT * FROM `bigquery-public-data.stackoverflow.posts_answers` where Score>0 ORDER BY Score desc) AS a ON q.id = a.parent_id \n", - " where q.tags like '%python%'\n", - " LIMIT {NUM_ROWS};\n", - " \"\"\"\n", - "\n", - "query_job = client.query(QUERY)\n", - "rows = query_job.result()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "b43937b6065d" - }, - "outputs": [], - "source": [ - "# Convert to a dataframe\n", - "df = rows.to_dataframe()\n", - "\n", - "# Examine the data\n", - "df.head()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "2cacd9869ee5" - }, - "outputs": [], - "source": [ - "# Extract the question ids and question text\n", - "ids = df.id.tolist()\n", - "questions = df.title.tolist()\n", - "\n", - "# Verify the length\n", - "len(ids)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "1124422cc200" - }, - "source": [ - "#### Instantiate the text encoding model\n", - "\n", - "Use the [sentence-t5 encoder](https://tfhub.dev/google/sentence-t5/st5-base/1) developed by Google for converting text to embeddings.\n", - "\n", - "> The sentence-T5 family of models encode text into high-dimensional vectors that can be used for text classification, semantic similarity, clustering and other natural language processing tasks.\n", - ">\n", - "> Our model is built on top of T5 (i.e. the Text-To-Text Transfer Transformer). It is trained on a variety of data sources and initialized from pre-trained T5 models with different model sizes as described in [1]. The input is variable-length English text and the output is a 768-dimensional vector. The sentence-T5 base model employs a 12-layer transformer architecture as the T5 base model does." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "ed41c7712930" - }, - "outputs": [], - "source": [ - "import tensorflow as tf\n", - "import tensorflow_hub as hub\n", - "# Registers the ops.\n", - "import tensorflow_text as text # noqa: F401\n", - "\n", - "hub_url = \"https://tfhub.dev/google/sentence-t5/st5-base/1\"\n", - "\n", - "encoder = hub.KerasLayer(hub_url)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "43088937e820" - }, - "source": [ - "#### Defining an encoding function\n", - "\n", - "Define a function to be used later that will take sentences and convert them to embeddings." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "a0370bd840d2" - }, - "outputs": [], - "source": [ - "from typing import List\n", - "\n", - "import numpy as np\n", - "\n", - "NUM_SENTENCES_IN_CHUNK = 100\n", - "\n", - "\n", - "def encode_text_to_embedding(\n", - " text_encoder: hub.KerasLayer, sentences: List[str]\n", - ") -> np.ndarray:\n", - " embeddings_list = []\n", - "\n", - " # The encoding models in TF hub have trouble processing too many strings at once, so we process them in chunks.\n", - " sentence_chunks = np.array_split(\n", - " sentences, (len(sentences) // NUM_SENTENCES_IN_CHUNK) + 1\n", - " )\n", - " for index, chunk in enumerate(sentence_chunks):\n", - " print(f\"\\tProcessing {(index+1)/len(sentence_chunks)*100:.1f}\")\n", - " embeddings_list.append(text_encoder(tf.constant(chunk)))\n", - "\n", - " return np.squeeze(np.column_stack(embeddings_list))" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "ba45d58bf96e" - }, - "source": [ - "#### Test the encoding function\n", - "\n", - "Encode a subset of data and see if the embeddings and distance metrics make sense.\n", - "\n", - "According to (sentence-T5 research paper)[https://arxiv.org/pdf/2108.08877.pdf], the similarity of embeddings is calculated using the dot-product. " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "9b01baa906b5" - }, - "outputs": [], - "source": [ - "# Encode 500 questions\n", - "questions = df.title.tolist()[:500]\n", - "question_embeddings = convert_text_to_embeddings(\n", - " text_encoder=encoder, sentences=questions\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "95e408daf219" - }, - "outputs": [], - "source": [ - "question_index = 10\n", - "\n", - "print(f\"Query question = {questions[question_index]}\")\n", - "scores = np.dot(question_embeddings[question_index], question_embeddings)\n", - "\n", - "# Print top 20 matches\n", - "for index, (question, score) in enumerate(\n", - " sorted(zip(questions, scores), key=lambda x: x[1], reverse=True)[:20]\n", - "):\n", - " print(f\"\\t{index}: {question}: {score}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "aQIQSyF9GtSv" - }, - "source": [ - "#### Save the train split in JSONL format.\n", - "\n", - "The data must be formatted in JSONL format, which means each embedding dictionary is written as a JSON string on its own line.\n", - "\n", - "Additionally, to demonstrate the filtering functionality, the `restricts` key is set such that each embedding has a different `class`, `even` or `odd`. These are used during the later matching step to filter for results.\n", - "See additional information of filtering here: https://cloud.google.com/vertex-ai/docs/matching-engine/filtering" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "307f468a3ecd" - }, - "outputs": [], - "source": [ - "import json\n", - "import tempfile\n", - "\n", - "NUM_ELEMENTS_IN_CHUNK = 1000\n", - "embeddings_file = tempfile.NamedTemporaryFile(suffix=\".json\", delete=False)\n", - "id_question_chunks = (\n", - " list(zip(ids, questions))[i : i + NUM_ELEMENTS_IN_CHUNK]\n", - " for i in range(0, len(questions), NUM_ELEMENTS_IN_CHUNK)\n", - ")\n", - "\n", - "with open(embeddings_file.name, \"a\") as f:\n", - " progress = 0\n", - " for chunk in id_question_chunks:\n", - " id_chunk, question_chunk = zip(*chunk)\n", - "\n", - " progress += 1\n", - " print(f\"Processing chunk {progress}/{len(questions)//NUM_ELEMENTS_IN_CHUNK}\")\n", - "\n", - " question_chunk_embeddings = encode_text_to_embedding(\n", - " text_encoder=encoder, sentences=question_chunk\n", - " )\n", - "\n", - " # Append to file\n", - " embeddings_formatted = [\n", - " json.dumps(\n", - " {\n", - " \"id\": str(id),\n", - " \"embedding\": [str(value) for value in embedding],\n", - " }\n", - " )\n", - " + \"\\n\"\n", - " for id, embedding in zip(id_chunk, question_chunk_embeddings)\n", - " ]\n", - " f.writelines(embeddings_formattedrmatted)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "QuVl8DrWG8NS" - }, - "source": [ - "Upload the training data to GCS." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "3PgsA_vbI8Vg" - }, - "outputs": [], - "source": [ - "EMBEDDINGS_INITIAL_URI = f\"{BUCKET_URI}/matching_engine/initial/embeddings_unique/\"\n", - "! gsutil cp {embeddings_file.name} {EMBEDDINGS_INITIAL_URI}" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "mglUPwHpJH98" - }, - "source": [ - "## Create Indexes\n" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "qhIBCQ7dDSbW" - }, - "source": [ - "### Create ANN Index (for Production Usage)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "qiIg9b5zJLi1" - }, - "outputs": [], - "source": [ - "DIMENSIONS = len(question_embeddings[0])\n", - "DISPLAY_NAME = \"stack_overflow\"\n", - "DESCRIPTION = \"questions from stackoverflow\"" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "svLYiDf0OD2G" - }, - "source": [ - "Create the ANN index configuration:\n", - "\n", - "To learn more about configuring the index, see [Input data format and structure](https://cloud.google.com/vertex-ai/docs/matching-engine/match-eng-setup#input-data-format).\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "Y4zooldkGoM4" - }, - "outputs": [], - "source": [ - "from google.cloud import aiplatform\n", - "\n", - "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "xzY7TpUSJcTV" - }, - "outputs": [], - "source": [ - "tree_ah_index = aiplatform.MatchingEngineIndex.create_tree_ah_index(\n", - " display_name=DISPLAY_NAME,\n", - " contents_delta_uri=EMBEDDINGS_INITIAL_URI,\n", - " dimensions=DIMENSIONS,\n", - " approximate_neighbors_count=150,\n", - " distance_measure_type=\"DOT_PRODUCT_DISTANCE\",\n", - " leaf_node_embedding_count=500,\n", - " leaf_nodes_to_search_percent=7,\n", - " description=DESCRIPTION,\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "17jrQi501QyX" - }, - "outputs": [], - "source": [ - "INDEX_RESOURCE_NAME = tree_ah_index.resource_name\n", - "INDEX_RESOURCE_NAME" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "0f1a9fbecabb" - }, - "source": [ - "Using the resource name, you can retrieve an existing MatchingEngineIndex." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "1ddb70647d98" - }, - "outputs": [], - "source": [ - "tree_ah_index = aiplatform.MatchingEngineIndex(index_name=INDEX_RESOURCE_NAME)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "qV2xjAnDDObD" - }, - "source": [ - "## Create an IndexEndpoint with VPC Network" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "BpZQoJyxDlbO" - }, - "outputs": [], - "source": [ - "# Retrieve the project number\n", - "PROJECT_NUMBER = !gcloud projects list --filter=\"PROJECT_ID:'{PROJECT_ID}'\" --format='value(PROJECT_NUMBER)'\n", - "PROJECT_NUMBER = PROJECT_NUMBER[0]\n", - "\n", - "VPC_NETWORK = \"[your-network-name]\"\n", - "VPC_NETWORK_FULL = \"projects/{}/global/networks/{}\".format(PROJECT_NUMBER, VPC_NETWORK)\n", - "VPC_NETWORK_FULL" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "QuARXzJVGyQX" - }, - "outputs": [], - "source": [ - "my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint.create(\n", - " display_name=DISPLAY_NAME,\n", - " description=DISPLAY_NAME,\n", - " network=VPC_NETWORK_FULL,\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "np2cgVuuIe9k" - }, - "source": [ - "## Deploy Indexes" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "8Ew1UgcIIiJG" - }, - "source": [ - "### Deploy ANN Index" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "nLOYTGygIlMK" - }, - "outputs": [], - "source": [ - "DEPLOYED_INDEX_ID = \"deployed_index_id_unique\"" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "_uK4WOgqN1NG" - }, - "outputs": [], - "source": [ - "my_index_endpoint = my_index_endpoint.deploy_index(\n", - " index=tree_ah_index, deployed_index_id=DEPLOYED_INDEX_ID\n", - ")\n", - "\n", - "my_index_endpoint.deployed_indexes" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "6LCGvBNvBd8D" - }, - "source": [ - "## Create Online Queries\n", - "\n", - "After you built your indexes, you may query against the deployed index to find nearest neighbours." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "ae9996f185fe" - }, - "outputs": [], - "source": [ - "test_embeddings = encode_text_to_embedding(\n", - " text_encoder=encoder, sentences=[\"How do I install tensorflow with GPU support?\"]\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "A3KYVw5HB-4v" - }, - "outputs": [], - "source": [ - "# Test query\n", - "response = my_index_endpoint.match(\n", - " deployed_index_id=DEPLOYED_INDEX_ID,\n", - " queries=test_embeddings,\n", - " num_neighbors=NUM_NEIGHBOURS,\n", - ")\n", - "\n", - "response" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "TpV-iwP9qw9c" - }, - "source": [ - "## Cleaning up\n", - "\n", - "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n", - "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n", - "You can also manually delete resources that you created by running the following code." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "sx_vKniMq9ZX" - }, - "outputs": [], - "source": [ - "# Force undeployment of indexes and delete endpoint\n", - "my_index_endpoint.delete(force=True)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "omj7N9iWv-Tq" - }, - "outputs": [], - "source": [ - "# Delete indexes\n", - "tree_ah_index.delete()" - ] - } - ], - "metadata": { - "colab": { - "collapsed_sections": [], - "name": "sdk_matching_engine_create_stack_overflow_embeddings.ipynb", - "toc_visible": true - }, - "kernelspec": { - "display_name": "Python 3", - "name": "python3" - } + "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." + ] }, - "nbformat": 4, - "nbformat_minor": 0 + { + "cell_type": "markdown", + "metadata": { + "id": "JAPoU8Sm5E6e" + }, + "source": [ + "# Using Vertex AI Matching Engine for StackOverflow Questions\n", + "\n", + "\n", + " \n", + " \n", + " \n", + "
\n", + " \n", + " \"Colab Run in Colab\n", + " \n", + " \n", + " \n", + " \"GitHub\n", + " View on GitHub\n", + " \n", + " \n", + " \n", + " \"Vertex\n", + " Open in Vertex AI Workbench\n", + " \n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "b0a74aaf1481" + }, + "source": [ + "## Overview\n", + "\n", + "This example demonstrates how to create text-to-image embeddings using the StackOverflow dataset and the sentence-T5 model. These are uploaded to Vertex AI Matching Engine service. It is a high scale, low latency solution, to find similar vectors (or more specifically \"embeddings\") for a large corpus. Moreover, it is a fully managed offering, further reducing operational overhead. It is built upon [Approximate Nearest Neighbor (ANN) technology](https://ai.googleblog.com/2020/07/announcing-scann-efficient-vector.html) developed by Google Research.\n", + "\n", + "**Pre-requisite**: This notebook requires you to already have a VPC network set up. See the \"Prepare a VPC network\" section in [Create Vertex AI Matching Engine index notebook](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/matching_engine/sdk_matching_engine_for_indexing.ipynb).\n", + "\n", + "Learn more about [Vertex AI Matching Engine](https://cloud.google.com/vertex-ai/docs/matching-engine/overview)." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "34a4b245e795" + }, + "source": [ + "### Objective\n", + "\n", + "In this notebook, you learn how to encode custom text embeddings, create an Approximate Nearest Neighbor (ANN) Index, and query against indexes.\n", + "\n", + "This tutorial uses the following Google Cloud ML services:\n", + "\n", + "- `Vertex AI Matching Engine`\n", + "\n", + "The steps performed include:\n", + "\n", + "* Create ANN Index\n", + "* Create an IndexEndpoint with VPC Network\n", + "* Deploy ANN Index\n", + "* Perform online query\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "tvgnzT1CKxrO" + }, + "source": [ + "### Dataset\n", + "\n", + "The dataset used for this tutorial is the [StackOverflow dataset](https://console.cloud.google.com/marketplace/product/stack-exchange/stack-overflow).\n", + "\n", + "> Stack Overflow is the largest online community for programmers to learn, share their knowledge, and advance their careers. Updated on a quarterly basis, this BigQuery dataset includes an archive of Stack Overflow content, including posts, votes, tags, and badges. This dataset is updated to mirror the Stack Overflow content on the Internet Archive, and is also available through the Stack Exchange Data Explorer." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "f0f1bea346db" + }, + "source": [ + "## Installation\n", + "\n", + "Install the latest version of Cloud Storage, BigQuery and Vertex AI SDKs for Python." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "dfbccc635a17" + }, + "outputs": [], + "source": [ + "# Install the packages\n", + "! pip3 install --upgrade 'google-cloud-aiplatform' \\\n", + " 'google-cloud-storage' \\\n", + " 'google-cloud-bigquery[pandas]'\n", + " " + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "5b08ba354c6e" + }, + "source": [ + "### Colab only: Uncomment the following cell to restart the kernel." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "bea801acf6b5" + }, + "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": "dd28c9e4f067" + }, + "source": [ + "## Before you begin\n", + "#### Set your project ID\n", + "\n", + "**If you don't know your project ID**, try the following:\n", + "* Run `gcloud config list`.\n", + "* Run `gcloud projects list`.\n", + "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "80c0215f05a0" + }, + "outputs": [], + "source": [ + "PROJECT_ID = \"[YOUR-PROJECT-ID]\"\n", + "\n", + "# Set the project id\n", + "! gcloud config set project {PROJECT_ID}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "4f4512bf63b3" + }, + "source": [ + "#### Region\n", + "\n", + "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": { + "id": "474be5183c27" + }, + "outputs": [], + "source": [ + "REGION = \"us-central1\" # @param {type: \"string\"}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "949271bfebe3" + }, + "source": [ + "### Authenticate your Google Cloud account\n", + "\n", + "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "b65b4ce80d9a" + }, + "source": [ + "**1. Vertex AI Workbench**\n", + "* Do nothing as you are already authenticated." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "985cdbfe7372" + }, + "source": [ + "**2. Local JupyterLab instance, uncomment and run:**" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": { + "id": "fbc9cd30cc4b" + }, + "outputs": [], + "source": [ + "# ! gcloud auth login" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "79efab26ad02" + }, + "source": [ + "**3. Colab, uncomment and run:**" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": { + "id": "a336a05c6149" + }, + "outputs": [], + "source": [ + "# from google.colab import auth\n", + "# auth.authenticate_user()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "0c0a44fa330f" + }, + "source": [ + "**4. Service account or other**\n", + "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "d3uj8x73nDX_" + }, + "source": [ + "* Authentication: Rerun the `gcloud auth login` command in the Vertex AI Workbench notebook terminal when you are logged out and need the credential again." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "hhq5zEbGg0XX" + }, + "source": [ + "### Colab only: Uncomment the following cell to restart the kernel." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": { + "id": "EzrelQZ22IZj" + }, + "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": "zgPO1eR3CYjk" + }, + "source": [ + "### Create a Cloud Storage bucket\n", + "\n", + "Create a storage bucket to store intermediate artifacts such as datasets." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": { + "id": "MzGDU7TWdts_" + }, + "outputs": [], + "source": [ + "BUCKET_URI = \"gs://your-bucket-name-unique\" # @param {type:\"string\"}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "-EcIXiGsCePi" + }, + "source": [ + "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "NIq7R4HZCfIc" + }, + "outputs": [], + "source": [ + "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "lR6Wwv-hCCN-" + }, + "source": [ + "## Prepare the data\n", + "\n", + "You will use [Stack Overflow dataset](https://console.cloud.google.com/marketplace/product/stack-exchange/stack-overflow) of question and answers hosted on BigQuery.\n", + "\n", + "> This public dataset is hosted in Google BigQuery and is included in BigQuery's 1TB/mo of free tier processing. This means that each user receives 1TB of free BigQuery processing every month, which can be used to run queries on this public dataset." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": { + "id": "9wzS85TeB9dG" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "CPU times: user 57.3 ms, sys: 6.84 ms, total: 64.1 ms\n", + "Wall time: 10.4 s\n" + ] + } + ], + "source": [ + "%%time\n", + "from google.cloud import bigquery\n", + "\n", + "client = bigquery.Client(project=PROJECT_ID)\n", + "\n", + "NUM_ROWS = 1000\n", + "\n", + "QUERY = f\"\"\"\n", + " SELECT distinct q.id, q.title, q.body, q.tags, a.body as answers, a.score \n", + " FROM (SELECT * FROM `bigquery-public-data.stackoverflow.posts_questions` where Score>0 ORDER BY View_Count desc) AS q \n", + " INNER JOIN (SELECT * FROM `bigquery-public-data.stackoverflow.posts_answers` where Score>0 ORDER BY Score desc) AS a ON q.id = a.parent_id \n", + " where q.tags like '%python%'\n", + " LIMIT {NUM_ROWS};\n", + " \"\"\"\n", + "\n", + "query_job = client.query(QUERY)\n", + "rows = query_job.result()" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": { + "id": "b43937b6065d" + }, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
idtitlebodytagsanswersscore
028819470Numbers of Day in Month<p>I have a data frame with a date time index,...python|pandas<p>There is now a <a href=\"https://pandas.pyda...11
135064304\"RuntimeError: Make sure the Graphviz executab...<p>I downloaded <code>Graphviz 2.38</code> MSI...python|installation|graphviz<p><strong>Step 1:</strong> Install Graphviz b...22
218783390python pip specify a library directory and an ...<p>I am using pip and trying to install a pyth...python|shared-libraries|pip|include-path|pyodbc<p>Another way to indicate the location of inc...19
31066933How to extract top-level domain name (TLD) fro...<p>how would you extract the domain name from ...python|url|parsing|dns|extract<p>No, there is no \"intrinsic\" way of knowing ...52
43817529syntax for creating a dictionary into another ...<blockquote>\\n <p><strong>Possible Duplicate:...python<p>You can declare a dictionary inside a dicti...106
\n", + "
" + ], + "text/plain": [ + " id title \\\n", + "0 28819470 Numbers of Day in Month \n", + "1 35064304 \"RuntimeError: Make sure the Graphviz executab... \n", + "2 18783390 python pip specify a library directory and an ... \n", + "3 1066933 How to extract top-level domain name (TLD) fro... \n", + "4 3817529 syntax for creating a dictionary into another ... \n", + "\n", + " body \\\n", + "0

I have a data frame with a date time index,... \n", + "1

I downloaded Graphviz 2.38 MSI... \n", + "2

I am using pip and trying to install a pyth... \n", + "3

how would you extract the domain name from ... \n", + "4

\\n

Possible Duplicate:... \n", + "\n", + " tags \\\n", + "0 python|pandas \n", + "1 python|installation|graphviz \n", + "2 python|shared-libraries|pip|include-path|pyodbc \n", + "3 python|url|parsing|dns|extract \n", + "4 python \n", + "\n", + " answers score \n", + "0

There is now a Step 1: Install Graphviz b... 22 \n", + "2

Another way to indicate the location of inc... 19 \n", + "3

No, there is no \"intrinsic\" way of knowing ... 52 \n", + "4

You can declare a dictionary inside a dicti... 106 " + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Convert to a dataframe\n", + "df = rows.to_dataframe()\n", + "\n", + "# Examine the data\n", + "df.head()" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": { + "id": "2cacd9869ee5" + }, + "outputs": [ + { + "data": { + "text/plain": [ + "1000" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Extract the question ids and question text\n", + "ids = df.id.tolist()\n", + "questions = df.title.tolist()\n", + "\n", + "# Verify the length\n", + "len(ids)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "1124422cc200" + }, + "source": [ + "#### Instantiate the text encoding model\n", + "\n", + "Use the [sentence-t5 encoder](https://tfhub.dev/google/sentence-t5/st5-base/1) developed by Google for converting text to embeddings.\n", + "\n", + "> The sentence-T5 family of models encode text into high-dimensional vectors that can be used for text classification, semantic similarity, clustering and other natural language processing tasks.\n", + ">\n", + "> Our model is built on top of T5 (i.e. the Text-To-Text Transfer Transformer). It is trained on a variety of data sources and initialized from pre-trained T5 models with different model sizes as described in [1]. The input is variable-length English text and the output is a 768-dimensional vector. The sentence-T5 base model employs a 12-layer transformer architecture as the T5 base model does." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "ed41c7712930" + }, + "outputs": [], + "source": [ + "import tensorflow as tf\n", + "import tensorflow_hub as hub\n", + "# Registers the ops.\n", + "import tensorflow_text as text # noqa: F401\n", + "\n", + "hub_url = \"https://tfhub.dev/google/sentence-t5/st5-base/1\"\n", + "\n", + "encoder = hub.KerasLayer(hub_url)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "43088937e820" + }, + "source": [ + "#### Defining an encoding function\n", + "\n", + "Define a function to be used later that will take sentences and convert them to embeddings." + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "metadata": { + "id": "a0370bd840d2" + }, + "outputs": [], + "source": [ + "from typing import List\n", + "\n", + "import numpy as np\n", + "from tqdm.auto import tqdm\n", + "\n", + "NUM_SENTENCES_IN_CHUNK = 100\n", + "\n", + "\n", + "def encode_text_to_embedding(\n", + " text_encoder: hub.KerasLayer, sentences: List[str]\n", + ") -> np.ndarray:\n", + " embeddings_list = []\n", + "\n", + " # The encoding models in TF hub have trouble processing too many strings at once, so we process them in chunks.\n", + " sentence_chunks = np.array_split(\n", + " sentences, (len(sentences) // NUM_SENTENCES_IN_CHUNK) + 1\n", + " )\n", + " for chunk in tqdm(sentence_chunks):\n", + " embeddings_list.append(text_encoder(tf.constant(chunk)))\n", + "\n", + " return np.squeeze(np.column_stack(embeddings_list))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ba45d58bf96e" + }, + "source": [ + "#### Test the encoding function\n", + "\n", + "Encode a subset of data and see if the embeddings and distance metrics make sense.\n", + "\n", + "According to (sentence-T5 research paper)[https://arxiv.org/pdf/2108.08877.pdf], the similarity of embeddings is calculated using the dot-product. " + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "metadata": { + "id": "9b01baa906b5" + }, + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "4cb8ca51b1864c189fd66209074debe5", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + " 0%| | 0/6 [00:00