Merge branch 'main' into imkc--matching-engine-filtering

notebooks/notebook_template.ipynb

@@ -402,7 +402,7 @@
       "source": [
         "### Initialize Vertex AI SDK for Python\n",
         "\n",
-        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
+        "Initialize the Vertex AI SDK for Python for your project."
       ]
     },
     {

notebooks/official/matching_engine/sdk_matching_engine_for_indexing.ipynb

@@ -98,171 +98,172 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "5e2eba58ad71"
+        "id": "f0f1bea346db"
       },
       "source": [
-        "### Costs \n",
+        "## Installation\n",
         "\n",
-        "This tutorial uses billable components of Google Cloud:\n",
-        "\n",
-        "* Vertex AI\n",
-        "* Cloud Storage\n",
-        "\n",
-        "Learn about [Vertex AI\n",
-        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
-        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
-        "Calculator](https://cloud.google.com/products/calculator/)\n",
-        "to generate a cost estimate based on your projected usage."
+        "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"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "d1e95a984673"
+        "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",
-        "### Set up your Google Cloud project\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",
-        "**The following steps are required, regardless of your notebook environment.**\n",
-        "\n",
-        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager).\n",
-        "\n",
-        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
-        "\n",
-        "1. [Enable the Vertex AI API and Compute Engine API, and Service Networking API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component,servicenetworking.googleapis.com).\n",
-        "\n",
-        "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
-        "Cloud SDK uses the right project for all the commands in this notebook.\n",
-        "\n",
-        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
+        "# Set the project id\n",
+        "! gcloud config set project {PROJECT_ID}"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "2b9daa35336a"
+        "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",
-        "**If you are using a Vertex AI Workbench notebook**, your environment is already\n",
-        "authenticated. Skip this step."
+        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "c6bed8c6a6b3"
+        "id": "b65b4ce80d9a"
       },
       "source": [
-        "**If you are using Colab**, run the cell below and follow the instructions\n",
-        "when prompted to authenticate your account via oAuth.\n",
-        "\n",
-        "**Otherwise**, follow these steps:\n",
-        "\n",
-        "1. In the Cloud Console, go to the [**Create service account key**\n",
-        "   page](https://console.cloud.google.com/apis/credentials/serviceaccountkey).\n",
-        "\n",
-        "2. Click **Create service account**.\n",
-        "\n",
-        "3. In the **Service account name** field, enter a name, and\n",
-        "   click **Create**.\n",
-        "\n",
-        "4. In the **Grant this service account access to project** section, click the **Role** drop-down list. Type \"Vertex AI\"\n",
-        "into the filter box, and select\n",
-        "   **Vertex AI Administrator**. Type \"Storage Object Admin\" into the filter box, and select **Storage Object Admin**.\n",
-        "\n",
-        "5. Click *Create*. A JSON file that contains your key downloads to your\n",
-        "local environment.\n",
-        "\n",
-        "6. Enter the path to your service account key as the\n",
-        "`GOOGLE_APPLICATION_CREDENTIALS` variable in the cell below and run the cell."
+        "**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": "3e2b43c2d2bf"
+        "id": "fbc9cd30cc4b"
       },
       "outputs": [],
       "source": [
-        "import os\n",
-        "import sys\n",
-        "\n",
-        "# If you are running this notebook in Colab, run this cell and follow the\n",
-        "# instructions to authenticate your GCP account. This provides access to your\n",
-        "# Cloud Storage bucket and lets you submit training jobs and prediction\n",
-        "# requests.\n",
-        "\n",
-        "# The Vertex AI Workbench notebook product has specific requirements\n",
-        "IS_VERTEX_AI_WORKBENCH_NOTEBOOK = os.path.exists(\n",
-        "    \"/opt/deeplearning/metadata/env_version\"\n",
-        ")\n",
-        "\n",
-        "# If on a Vertex AI Workbench notebook, then don't execute this code\n",
-        "if not IS_VERTEX_AI_WORKBENCH_NOTEBOOK:\n",
-        "    if \"google.colab\" in sys.modules:\n",
-        "        from google.colab import auth as google_auth\n",
-        "\n",
-        "        google_auth.authenticate_user()\n",
-        "\n",
-        "    # If you are running this notebook locally, log in using gcloud\n",
-        "    elif not os.getenv(\"IS_TESTING\"):\n",
-        "        ! gcloud auth login"
+        "# ! gcloud auth login"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "WReHDGG5g0XY"
+        "id": "79efab26ad02"
       },
       "source": [
-        "#### Set your project ID\n",
-        "\n",
-        "**If you don't know your project ID**, you may be able to get your project ID using `gcloud`."
+        "**3. Colab, uncomment and run:**"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "beb72f394541"
+        "id": "a336a05c6149"
       },
       "outputs": [],
       "source": [
-        "import os\n",
-        "\n",
-        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# Get your Google Cloud project ID from gcloud\n",
-        "if not os.getenv(\"IS_TESTING\"):\n",
-        "    shell_output = !gcloud config list --format 'value(core.project)' 2>/dev/null\n",
-        "    PROJECT_ID = shell_output[0]\n",
-        "    print(\"Project ID: \", PROJECT_ID)"
+        "# from google.colab import auth\n",
+        "# auth.authenticate_user()"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "f4c6d0a9e66c"
+        "id": "0c0a44fa330f"
       },
       "source": [
-        "Otherwise, set your project ID here."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "1dc3fa9ac4f7"
-      },
-      "outputs": [],
-      "source": [
-        "if PROJECT_ID == \"\" or PROJECT_ID is None:\n",
-        "    PROJECT_ID = \"<your_project_id>\"  # @param {type:\"string\"}"
+        "**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."
       ]
     },
     {
@@ -271,7 +272,8 @@
         "id": "4962667eec8e"
       },
       "source": [
-        "* **Prepare a VPC network**.  To reduce any network overhead that might lead to unnecessary increase in overhead latency, it is 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",
+        "### Prepare a VPC network\n",
+        "To reduce any network overhead that might lead to unnecessary increase in overhead latency, it is 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."
       ]
@@ -284,7 +286,8 @@
       },
       "outputs": [],
       "source": [
-        "VPC_NETWORK = \"[your-vpc-network-name]\"  # @param {type:\"string\"}\n",
+        "# VPC_NETWORK = \"[your-vpc-network-name]\"  # @param {type:\"string\"}\n",
+        "VPC_NETWORK = \"matching-engine-test\"  # @param {type:\"string\"}\n",
         "\n",
         "PEERING_RANGE_NAME = \"ann-haystack-range\""
       ]
@@ -363,7 +366,7 @@
       },
       "outputs": [],
       "source": [
-        "! pip install -U google-cloud-aiplatform"
+        "! pip install --upgrade --quiet google-cloud-aiplatform grpcio-tools h5py"
       ]
     },
     {
@@ -375,27 +378,13 @@
         "Install the `h5py` to prepare sample dataset, and the `grpcio-tools` for querying against the index. "
       ]
     },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "-h5sqwOEZ5Yq"
-      },
-      "outputs": [],
-      "source": [
-        "! pip install -U grpcio-tools\n",
-        "! pip install -U h5py"
-      ]
-    },
     {
       "cell_type": "markdown",
       "metadata": {
         "id": "hhq5zEbGg0XX"
       },
       "source": [
-        "### Restart the kernel\n",
-        "\n",
-        "After you install the additional packages, you need to restart the notebook kernel so it can find the packages."
+        "### Colab only: Uncomment the following cell to restart the kernel."
       ]
     },
     {
@@ -406,37 +395,11 @@
       },
       "outputs": [],
       "source": [
-        "if not os.getenv(\"IS_TESTING\"):\n",
-        "    # Automatically restart kernel after installs\n",
-        "    import IPython\n",
+        "# 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": "q7tcBkCDI1_M"
-      },
-      "source": [
-        "### Random ID\n",
-        "\n",
-        "To avoid name collisions between users on resources created, create a random ID for each instance session, and append the id onto the name of resources you create in this tutorial."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "HpIK91y1IzDr"
-      },
-      "outputs": [],
-      "source": [
-        "import random\n",
-        "import string\n",
-        "\n",
-        "RANDOM_ID = \"\".join(random.choices(string.ascii_lowercase + string.digits, k=8))"
+        "# app = IPython.Application.instance()\n",
+        "# app.kernel.do_shutdown(True)"
       ]
     },
     {
@@ -447,15 +410,7 @@
       "source": [
         "### Create a Cloud Storage bucket\n",
         "\n",
-        "**The following steps are required, regardless of your notebook environment.**\n",
-        "\n",
-        "Set the name of your Cloud Storage bucket below. It must be unique across all\n",
-        "Cloud Storage buckets.\n",
-        "\n",
-        "You may also change the `REGION` variable, which is used for operations\n",
-        "throughout the rest of this notebook. Make sure to [choose a region where Vertex AI services are\n",
-        "available](https://cloud.google.com/vertex-ai/docs/general/locations#available_regions). You may\n",
-        "not use a Multi-Regional Storage bucket for training with Vertex AI."
+        "Create a storage bucket to store intermediate artifacts such as datasets."
       ]
     },
     {
@@ -466,23 +421,7 @@
       },
       "outputs": [],
       "source": [
-        "BUCKET_URI = \"gs://[your-bucket-name]\"  # @param {type:\"string\"}\n",
-        "REGION = \"[your-region]\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "cf221059d072"
-      },
-      "outputs": [],
-      "source": [
-        "if BUCKET_URI == \"\" or BUCKET_URI is None or BUCKET_URI == \"gs://[your-bucket-name]\":\n",
-        "    BUCKET_URI = \"gs://\" + PROJECT_ID + \"aip-\" + RANDOM_ID\n",
-        "\n",
-        "if REGION == \"[your-region]\":\n",
-        "    REGION = \"us-central1\""
+        "BUCKET_URI = \"gs://your-bucket-name-unique\"  # @param {type:\"string\"}"
       ]
     },
     {
@@ -505,84 +444,6 @@
         "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ucvCsknMCims"
-      },
-      "source": [
-        "Finally, validate access to your Cloud Storage bucket by examining its contents:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "vhOb7YnwClBb"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil ls -al $BUCKET_URI"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "XoEqT2Y4DJmf"
-      },
-      "source": [
-        "### Import libraries and define constants"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "Y9Uo3tifg1kx"
-      },
-      "source": [
-        "Import the Vertex AI (unified) client library into your Python environment. \n"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "f2d05ab4126a"
-      },
-      "outputs": [],
-      "source": [
-        "import h5py"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "76f7b9ffde0b"
-      },
-      "source": [
-        "Use gcloud to retrieve the project number."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "pRUOFELefqf1"
-      },
-      "outputs": [],
-      "source": [
-        "PROJECT_NUMBER = !gcloud projects list --filter=\"PROJECT_ID:'{PROJECT_ID}'\" --format='value(PROJECT_NUMBER)'\n",
-        "PROJECT_NUMBER = PROJECT_NUMBER[0]\n",
-        "\n",
-        "PARENT = \"projects/{}/locations/{}\".format(PROJECT_ID, REGION)\n",
-        "\n",
-        "print(\"PROJECT_ID: {}\".format(PROJECT_ID))\n",
-        "print(\"REGION: {}\".format(REGION))\n",
-        "\n",
-        "!gcloud config set project {PROJECT_ID} --quiet\n",
-        "!gcloud config set ai_platform/region {REGION} --quiet"
-      ]
-    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -627,6 +488,8 @@
       },
       "outputs": [],
       "source": [
+        "import h5py\n",
+        "\n",
         "# The number of nearest neighbors to be retrieved from database for each query.\n",
         "NUM_NEIGHBOURS = 10\n",
         "\n",
@@ -652,22 +515,42 @@
         "id": "aQIQSyF9GtSv"
       },
       "source": [
-        "Save the train split in JSONL format.\n"
+        "#### 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": "18wCiTwfG40P"
+        "id": "57fe2ce4b50f"
       },
       "outputs": [],
       "source": [
+        "import json\n",
+        "\n",
         "with open(\"glove100.json\", \"w\") as f:\n",
-        "    for i in range(len(train)):\n",
-        "        f.write('{\"id\":\"' + str(i) + '\",')\n",
-        "        f.write('\"embedding\":[' + \",\".join(str(x) for x in train[i]) + \"]}\")\n",
-        "        f.write(\"\\n\")"
+        "    embeddings_formatted = [\n",
+        "        json.dumps(\n",
+        "            {\n",
+        "                \"id\": str(index),\n",
+        "                \"embedding\": [str(value) for value 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",
+        "        for index, embedding in enumerate(train)\n",
+        "    ]\n",
+        "    f.writelines(embeddings_formatted)"
       ]
     },
     {
@@ -742,7 +625,6 @@
       "outputs": [],
       "source": [
         "import os\n",
-        "import sys\n",
         "\n",
         "from google.cloud import aiplatform\n",
         "\n",
@@ -878,8 +760,21 @@
       "outputs": [],
       "source": [
         "with open(\"glove100_incremental.json\", \"w\") as f:\n",
+        "    index = 0\n",
         "    f.write(\n",
-        "        '{\"id\":\"0\",\"embedding\":[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]}\\n'\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",
         "    )"
       ]
     },
@@ -965,6 +860,10 @@
       },
       "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"
@@ -1023,7 +922,7 @@
       },
       "outputs": [],
       "source": [
-        "DEPLOYED_INDEX_ID = f\"tree_ah_glove_deployed_{RANDOM_ID}\""
+        "DEPLOYED_INDEX_ID = \"tree_ah_glove_deployed_unique\""
       ]
     },
     {
@@ -1058,7 +957,7 @@
       },
       "outputs": [],
       "source": [
-        "DEPLOYED_BRUTE_FORCE_INDEX_ID = f\"glove_brute_force_deployed_{RANDOM_ID}\""
+        "DEPLOYED_BRUTE_FORCE_INDEX_ID = \"glove_brute_force_deployed_unique\""
       ]
     },
     {
@@ -1084,16 +983,9 @@
       "source": [
         "## Create Online Queries\n",
         "\n",
-        "After you 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).  "
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "IcXa9lSuB9AT"
-      },
-      "source": [
-        "Test your query:"
+        "After you 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."
       ]
     },
     {
@@ -1104,9 +996,16 @@
       },
       "outputs": [],
       "source": [
+        "# Test query\n",
+        "from google.cloud.aiplatform.matching_engine.matching_engine_index_endpoint import \\\n",
+        "    Namespace\n",
+        "\n",
         "# Test query\n",
         "response = my_index_endpoint.match(\n",
-        "    deployed_index_id=DEPLOYED_INDEX_ID, queries=test[:1], num_neighbors=NUM_NEIGHBOURS\n",
+        "    deployed_index_id=DEPLOYED_INDEX_ID,\n",
+        "    queries=test[:1],\n",
+        "    num_neighbors=NUM_NEIGHBOURS,\n",
+        "    filter=[Namespace(\"class\", [\"even\"])],\n",
         ")\n",
         "\n",
         "response"