notebooks/official/vector_search/sdk_vector_search_for_indexing.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2022 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": [
        "# Create Vertex AI Vector Search index\n",
        "\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vector_search/sdk_vector_search_for_indexing.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" 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%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fvector_search%2Fsdk_vector_search_for_indexing.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" 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/vertex-ai-samples/main/notebooks/official/vector_search/sdk_vector_search_for_indexing.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vector_search/sdk_vector_search_for_indexing.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b0a74aaf1481"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This example demonstrates how to use the Vertex AI ANN Service. It's a high scale, low latency solution, to find similar vectors (or more specifically \"embeddings\") for a large corpus. Moreover, it's a fully managed offering, further reducing operational overhead. The Vertex AI ANN Service 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",
        "Learn more about [Vertex AI Vector Search](https://cloud.google.com/vertex-ai/docs/vector-search/overview)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "34a4b245e795"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this notebook, you learn how to create Approximate Nearest Neighbor (ANN) Index, query against indexes, and validate the performance of the index. \n",
        "\n",
        "This tutorial uses the following Vertex AI services:\n",
        "\n",
        "- Vertex AI Vector Search\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "* Create ANN Index and Brute Force Index.\n",
        "* Create an IndexEndpoint with VPC Network.\n",
        "* Deploy ANN Index and Brute Force Index.\n",
        "* Perform online query.\n",
        "* Compute recall.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "The dataset used for this tutorial is the [GloVe dataset](https://nlp.stanford.edu/projects/glove/).\n",
        "\n",
        "GloVe is an unsupervised learning algorithm for obtaining vector representations for words. Training is performed on aggregated global word-word co-occurrence statistics from a corpus. The resulting representations showcase interesting linear substructures of the word vector space.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8925ff9e165e"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f0f1bea346db"
      },
      "source": [
        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
      "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",
        "                         grpcio-tools \\\n",
        "                         h5py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5b08ba354c6e"
      },
      "source": [
        "### Restart runtime (Colab only)\n",
        "\n",
        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bea801acf6b5"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7b49231643e4"
      },
      "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>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7176ea64999b"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7de6ef0fac42"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ee2dcc922f3d"
      },
      "source": [
        "### Set Google Cloud project information\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "80c0215f05a0"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4962667eec8e"
      },
      "source": [
        "### Prepare a VPC network\n",
        "To reduce any network overhead that might lead to unnecessary increase in overhead latency, it's best to call the ANN endpoints from your VPC via a direct [VPC Peering](https://cloud.google.com/vertex-ai/docs/general/vpc-peering) connection. \n",
        "  * The following section describes how to setup a VPC Peering connection if you don't have one. \n",
        "  * This is a one-time initial setup task. You can also reuse existing VPC network and skip this section."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KDH8CgQiSxhv"
      },
      "outputs": [],
      "source": [
        "VPC_NETWORK = \"[your-vpc-network-name]\"  # @param {type:\"string\"}\n",
        "\n",
        "PEERING_RANGE_NAME = \"ann-haystack-range\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lW2LneA5mmmP"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# Remove the if condition to run the encapsulated code\n",
        "if not os.getenv(\"IS_TESTING\"):\n",
        "    # Create a VPC network\n",
        "    ! gcloud compute networks create {VPC_NETWORK} --bgp-routing-mode=regional --subnet-mode=auto --project={PROJECT_ID}\n",
        "\n",
        "    # Add necessary firewall rules\n",
        "    ! gcloud compute firewall-rules create {VPC_NETWORK}-allow-icmp --network {VPC_NETWORK} --priority 65534 --project {PROJECT_ID} --allow icmp\n",
        "\n",
        "    ! gcloud compute firewall-rules create {VPC_NETWORK}-allow-internal --network {VPC_NETWORK} --priority 65534 --project {PROJECT_ID} --allow all --source-ranges 10.128.0.0/9\n",
        "\n",
        "    ! gcloud compute firewall-rules create {VPC_NETWORK}-allow-rdp --network {VPC_NETWORK} --priority 65534 --project {PROJECT_ID} --allow tcp:3389\n",
        "\n",
        "    ! gcloud compute firewall-rules create {VPC_NETWORK}-allow-ssh --network {VPC_NETWORK} --priority 65534 --project {PROJECT_ID} --allow tcp:22\n",
        "\n",
        "    # Reserve IP range\n",
        "    ! gcloud compute addresses create {PEERING_RANGE_NAME} --global --prefix-length=16 --network={VPC_NETWORK} --purpose=VPC_PEERING --project={PROJECT_ID} --description=\"peering range\"\n",
        "\n",
        "    # Set up peering with service networking\n",
        "    # Your account must have the \"Compute Network Admin\" role to run the following.\n",
        "    ! gcloud services vpc-peerings connect --service=servicenetworking.googleapis.com --network={VPC_NETWORK} --ranges={PEERING_RANGE_NAME} --project={PROJECT_ID}"
      ]
    },
    {
      "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": "d5de53b31bf1"
      },
      "source": [
        "**WARNING:** The `MatchingIndexEndpoint.match` method (to create online queries against your deployed index) has to be executed in a Vertex AI Workbench notebook instance that is created with the following requirements:\n",
        "  * **In the same region where your ANN service is deployed** (for example, if you set `LOCATION = \"us-central1\"` as same as the tutorial, the notebook instance has to be in `us-central1`).\n",
        "  \n",
        "  * **Make sure you select the VPC network you created for ANN service** (instead of using the \"default\" one). That is, you have to create a new notebook instance that uses the VPC network you created earlier. Run the rest of the tutorial from that notebook instance.\n",
        "  * If you run it in a Colab or a Vertex AI Workbench notebook in a different VPC network or region, \"Create Online Queries\" section is expected to fail."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "17ba55acc776"
      },
      "source": [
        "## Import the required libraries\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "56182a4d44c4"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "\n",
        "import h5py\n",
        "from google.cloud import aiplatform\n",
        "from google.cloud.aiplatform.matching_engine.matching_engine_index_endpoint import \\\n",
        "    Namespace"
      ]
    },
    {
      "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 = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-EcIXiGsCePi"
      },
      "source": [
        "**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 {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d5868f2942b8"
      },
      "source": [
        "## Initialize Vertex AI SDK\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0a2809ebac11"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lR6Wwv-hCCN-"
      },
      "source": [
        "## Prepare the data\n",
        "\n",
        "The GloVe dataset consists of a set of pre-trained embeddings. The embeddings are split into a \"train\" split, and a \"test\" split. You create a vector search index from the \"train\" split, and use the embedding vectors in the \"test\" split as query vectors to test the vector search index.\n",
        "\n",
        "**Note:** While the data split uses the term \"train\", these are pre-trained embeddings and therefore are ready to be indexed for search. The terms \"train\" and \"test\" split are used just to be consistent with machine learning terminology.\n",
        "\n",
        "Download the GloVe dataset.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9wzS85TeB9dG"
      },
      "outputs": [],
      "source": [
        "! gsutil cp gs://cloud-samples-data/vertex-ai/matching_engine/glove-100-angular.hdf5 ."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4fAO9CMoCNtq"
      },
      "source": [
        "Read the data into memory.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lZ3JQTS6CN-3"
      },
      "outputs": [],
      "source": [
        "# The number of nearest neighbors to be retrieved from database for each query.\n",
        "NUM_NEIGHBOURS = 10\n",
        "\n",
        "h5 = h5py.File(\"glove-100-angular.hdf5\", \"r\")\n",
        "train = h5[\"train\"]\n",
        "test = h5[\"test\"]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pE6bBBo7GjJK"
      },
      "outputs": [],
      "source": [
        "# check the first record\n",
        "train[0]"
      ]
    },
    {
      "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",
        "For additional information of filtering, see [Filter vector matches](https://cloud.google.com/vertex-ai/docs/vector-search/filtering)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "57fe2ce4b50f"
      },
      "outputs": [],
      "source": [
        "with open(\"glove100.json\", \"w\") as f:\n",
        "    embeddings_formatted = [\n",
        "        json.dumps(\n",
        "            {\n",
        "                \"id\": str(index),\n",
        "                \"embedding\": [str(value) for value in embedding],\n",
        "                \"restricts\": [\n",
        "                    {\n",
        "                        \"namespace\": \"class\",\n",
        "                        \"allow_list\": [\"even\" if index % 2 == 0 else \"odd\"],\n",
        "                    }\n",
        "                ],\n",
        "            }\n",
        "        )\n",
        "        + \"\\n\"\n",
        "        for index, embedding in enumerate(train)\n",
        "    ]\n",
        "    f.writelines(embeddings_formatted)"
      ]
    },
    {
      "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}/vector_search/initial/\"\n",
        "! gsutil cp glove100.json {EMBEDDINGS_INITIAL_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mglUPwHpJH98"
      },
      "source": [
        "## Create Indexes\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qiIg9b5zJLi1"
      },
      "outputs": [],
      "source": [
        "# set no.of dimensions for your embeddings\n",
        "DIMENSIONS = 100\n",
        "# set the dispaly name for ann index\n",
        "DISPLAY_NAME = \"glove_100_1\"\n",
        "# set the display name for brute force index\n",
        "DISPLAY_NAME_BRUTE_FORCE = DISPLAY_NAME + \"_brute_force\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qhIBCQ7dDSbW"
      },
      "source": [
        "### Create ANN Index (for Production Usage)"
      ]
    },
    {
      "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/vector-search/setup/format-structure).\n"
      ]
    },
    {
      "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=\"Glove 100 ANN index\",\n",
        "    labels={\"label_name\": \"label_value\"},\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 Vector Search Index."
      ]
    },
    {
      "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": "kSsqZuyoA1SG"
      },
      "source": [
        "### Create Brute Force Index (for Ground Truth)\n",
        "\n",
        "The brute force index uses a naive brute force method to find the nearest neighbors. This method isn't fast or efficient. Hence, brute force indices are not recommended for production usage. They are to be used to find the \"ground truth\" set of neighbors so that the \"ground truth\" set can be used to measure recall of the indices being tuned for production usage. To ensure an \"apples to apples\" comparison, the `distanceMeasureType` and `dimensions` of the brute force index should match those of the production indices being tuned.\n",
        "\n",
        "Create the brute force index configuration:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DXnBLqjXBsv8"
      },
      "outputs": [],
      "source": [
        "brute_force_index = aiplatform.MatchingEngineIndex.create_brute_force_index(\n",
        "    display_name=DISPLAY_NAME_BRUTE_FORCE,\n",
        "    contents_delta_uri=EMBEDDINGS_INITIAL_URI,\n",
        "    dimensions=DIMENSIONS,\n",
        "    distance_measure_type=\"DOT_PRODUCT_DISTANCE\",\n",
        "    description=\"Glove 100 index (brute force)\",\n",
        "    labels={\"label_name\": \"label_value\"},\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_oD5SieYJbbW"
      },
      "outputs": [],
      "source": [
        "INDEX_BRUTE_FORCE_RESOURCE_NAME = brute_force_index.resource_name\n",
        "INDEX_BRUTE_FORCE_RESOURCE_NAME"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "865fcad494d7"
      },
      "outputs": [],
      "source": [
        "brute_force_index = aiplatform.MatchingEngineIndex(\n",
        "    index_name=INDEX_BRUTE_FORCE_RESOURCE_NAME\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "omlgEZ-sGoM5"
      },
      "source": [
        "## Update Indexes\n",
        "\n",
        "Create an incremental data file.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DDAvm_mj_BVs"
      },
      "outputs": [],
      "source": [
        "with open(\"glove100_incremental.json\", \"w\") as f:\n",
        "    index = 0\n",
        "    f.write(\n",
        "        json.dumps(\n",
        "            {\n",
        "                \"id\": str(index),\n",
        "                \"embedding\": [str(0) for _ in train[index]],\n",
        "                \"restricts\": [\n",
        "                    {\n",
        "                        \"namespace\": \"class\",\n",
        "                        \"allow_list\": [\"even\" if index % 2 == 0 else \"odd\"],\n",
        "                    }\n",
        "                ],\n",
        "            }\n",
        "        )\n",
        "        + \"\\n\"\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZU7TU7C7GoM6"
      },
      "source": [
        "Copy the incremental data file to a new subdirectory.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RLWcDvNLGoM6"
      },
      "outputs": [],
      "source": [
        "EMBEDDINGS_UPDATE_URI = f\"{BUCKET_URI}/vector-search/incremental/\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FgpEDX0oGoM6"
      },
      "outputs": [],
      "source": [
        "! gsutil cp glove100_incremental.json {EMBEDDINGS_UPDATE_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aiXtF_x0GoM6"
      },
      "source": [
        "Create update index request.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tvedBONtGoM6"
      },
      "outputs": [],
      "source": [
        "tree_ah_index = tree_ah_index.update_embeddings(\n",
        "    contents_delta_uri=EMBEDDINGS_UPDATE_URI,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HKPDojFpGoM6"
      },
      "outputs": [],
      "source": [
        "INDEX_RESOURCE_NAME = tree_ah_index.resource_name\n",
        "INDEX_RESOURCE_NAME"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qV2xjAnDDObD"
      },
      "source": [
        "## Create an IndexEndpoint in your 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",
        "# Get the full network resource 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": [
        "# Create your IndexEndpoint\n",
        "my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint.create(\n",
        "    display_name=\"index_endpoint_for_demo\",\n",
        "    description=\"index endpoint description\",\n",
        "    network=VPC_NETWORK_FULL,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PJ3bcZqi-cfM"
      },
      "outputs": [],
      "source": [
        "INDEX_ENDPOINT_NAME = my_index_endpoint.resource_name\n",
        "INDEX_ENDPOINT_NAME"
      ]
    },
    {
      "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": [
        "# Set an id for your ann index deployment\n",
        "DEPLOYED_INDEX_ID = \"tree_ah_glove_deployed_unique\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_uK4WOgqN1NG"
      },
      "outputs": [],
      "source": [
        "# Deploy your ann index\n",
        "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": "RNZnXmO5AhDO"
      },
      "source": [
        "### Deploy Brute Force Index"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3p9e4828AkSv"
      },
      "outputs": [],
      "source": [
        "# Set an id for your brute force index deployment\n",
        "DEPLOYED_BRUTE_FORCE_INDEX_ID = \"glove_brute_force_deployed_unique\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-2kgd01SA4rk"
      },
      "outputs": [],
      "source": [
        "# Deploy your brute force index\n",
        "my_index_endpoint = my_index_endpoint.deploy_index(\n",
        "    index=brute_force_index, deployed_index_id=DEPLOYED_BRUTE_FORCE_INDEX_ID\n",
        ")\n",
        "\n",
        "my_index_endpoint.deployed_indexes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6LCGvBNvBd8D"
      },
      "source": [
        "## Create online queries\n",
        "\n",
        "After you've built your indexes, you may query against the deployed index through the online querying gRPC API (Match service) within the virtual machine instances from the same region (for example, 'us-central1' in this tutorial).\n",
        "\n",
        "The `filter` parameter is an optional way to filter for a subset of embeddings. In this case, only embeddings that have the `class` set as `even` are returned."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A3KYVw5HB-4v"
      },
      "outputs": [],
      "source": [
        "# Use match service with a test query\n",
        "response = my_index_endpoint.match(\n",
        "    deployed_index_id=DEPLOYED_INDEX_ID,\n",
        "    queries=test[:1].tolist(),\n",
        "    num_neighbors=NUM_NEIGHBOURS,\n",
        "    filter=[Namespace(\"class\", [\"even\"])],\n",
        ")\n",
        "\n",
        "response"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xeUZO3bAGoM-"
      },
      "source": [
        "### Compute recall\n",
        "\n",
        "Use the deployed brute force Index as the ground truth to calculate the recall of ANN Index. You can run multiple queries in a single match call."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U9dNIbkEGoM-"
      },
      "outputs": [],
      "source": [
        "# Retrieve nearest neighbors for both the tree-AH index and the brute force index\n",
        "tree_ah_response_test = my_index_endpoint.match(\n",
        "    deployed_index_id=DEPLOYED_INDEX_ID,\n",
        "    queries=test[:].tolist(),\n",
        "    num_neighbors=NUM_NEIGHBOURS,\n",
        ")\n",
        "brute_force_response_test = my_index_endpoint.match(\n",
        "    deployed_index_id=DEPLOYED_BRUTE_FORCE_INDEX_ID,\n",
        "    queries=test[:].tolist(),\n",
        "    num_neighbors=NUM_NEIGHBOURS,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V-eMF05UGoM-"
      },
      "outputs": [],
      "source": [
        "# Calculate recall by determining how many neighbors were correctly retrieved as compared to the brute force option.\n",
        "recalled_neighbors = 0\n",
        "for tree_ah_neighbors, brute_force_neighbors in zip(\n",
        "    tree_ah_response_test, brute_force_response_test\n",
        "):\n",
        "    tree_ah_neighbor_ids = [neighbor.id for neighbor in tree_ah_neighbors]\n",
        "    brute_force_neighbor_ids = [neighbor.id for neighbor in brute_force_neighbors]\n",
        "\n",
        "    recalled_neighbors += len(\n",
        "        set(tree_ah_neighbor_ids).intersection(brute_force_neighbor_ids)\n",
        "    )\n",
        "\n",
        "recall = recalled_neighbors / len(\n",
        "    [neighbor for neighbors in brute_force_response_test for neighbor in neighbors]\n",
        ")\n",
        "\n",
        "print(\"Recall: {}\".format(recall))"
      ]
    },
    {
      "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": [
        "delete_bucket = False\n",
        "\n",
        "# Force undeployment of indexes and delete endpoint\n",
        "my_index_endpoint.delete(force=True)\n",
        "\n",
        "# Delete indexes\n",
        "tree_ah_index.delete()\n",
        "brute_force_index.delete()\n",
        "\n",
        "if delete_bucket:\n",
        "    ! gsutil rm -rf {BUCKET_URI}"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "sdk_vector_search_for_indexing.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}