notebooks/official/matching_engine/sdk_matching_engine_create_stack_overflow_embeddings_vertex.ipynb

{
  "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 Vector Search and Vertex AI embeddings for text for StackOverflow Questions \n",
        "![ ](https://www.google-analytics.com/collect?v=2&tid=G-L6X3ECH596&cid=1&en=page_view&sid=1&dt=sdk_matching_engine_create_stack_overflow_embeddings_vertex.ipynb&dl=notebooks%2Fofficial%2Fmatching_engine%2Fsdk_matching_engine_create_stack_overflow_embeddings_vertex.ipynb)\n",
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/matching_engine/sdk_matching_engine_create_stack_overflow_embeddings_vertex.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/matching_engine/sdk_matching_engine_create_stack_overflow_embeddings_vertex.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "      <td>\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/matching_engine/sdk_matching_engine_create_stack_overflow_embeddings_vertex\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b0a74aaf1481"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This example demonstrates how to encode text embeddings using the Vertex AI embeddings for text service and the StackOverflow dataset. These are uploaded to the Vertex AI Vector Search service, which is a high-scale, low-latency solution, for finding similar vectors from a large corpus. Vector Search is a fully managed offering, further reducing operational overhead. It's 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) and [Vertex AI embeddings for text](https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-text-embeddings)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "34a4b245e795"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this notebook, you learn how to encode 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 Vector Search\n",
        "- Vertex AI embeddings for text`\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "* Convert a BigQuery dataset to embeddings.\n",
        "* Create an index.\n",
        "* Upload embeddings to the index.\n",
        "* Create an index endpoint.\n",
        "* Deploy the index to the index endpoint.\n",
        "* Perform an online query."
      ]
    },
    {
      "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 the Vertex AI SDK for Python. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dfbccc635a17"
      },
      "outputs": [],
      "source": [
        "# Install the packages\n",
        "! pip3 install --upgrade google-cloud-aiplatform==1.36.0 \\\n",
        "                        google-cloud-storage \\\n",
        "                        'google-cloud-bigquery[pandas]' "
      ]
    },
    {
      "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": "42635d0c0019"
      },
      "outputs": [],
      "source": [
        "! gcloud config list"
      ]
    },
    {
      "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. To learn more about supported regions, see the [Vertex AI regions page](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 = 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 {REGION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lR6Wwv-hCCN-"
      },
      "source": [
        "## Prepare the data\n",
        "\n",
        "You use the [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.\n",
        "\n",
        "The BigQuery table is too large to fit into memory, so you need to write a generator called `query_bigquery_chunks` to yield chunks of the dataframe for processing. Additionally, an extra column `title_with_body` is added, which is a concatenation of the question title and body."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ed1b3f87c475"
      },
      "outputs": [],
      "source": [
        "import math\n",
        "from typing import Any, Generator\n",
        "\n",
        "import pandas as pd\n",
        "from google.cloud import bigquery\n",
        "\n",
        "client = bigquery.Client(project=PROJECT_ID)\n",
        "QUERY_TEMPLATE = \"\"\"\n",
        "        SELECT distinct q.id, q.title, q.body\n",
        "        FROM (SELECT * FROM `bigquery-public-data.stackoverflow.posts_questions` where Score>0 ORDER BY View_Count desc) AS q \n",
        "        LIMIT {limit} OFFSET {offset};\n",
        "        \"\"\"\n",
        "\n",
        "\n",
        "def query_bigquery_chunks(\n",
        "    max_rows: int, rows_per_chunk: int, start_chunk: int = 0\n",
        ") -> Generator[pd.DataFrame, Any, None]:\n",
        "    for offset in range(start_chunk, max_rows, rows_per_chunk):\n",
        "        query = QUERY_TEMPLATE.format(limit=rows_per_chunk, offset=offset)\n",
        "        query_job = client.query(query)\n",
        "        rows = query_job.result()\n",
        "        df = rows.to_dataframe()\n",
        "        df[\"title_with_body\"] = df.title + \"\\n\" + df.body\n",
        "        yield df"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b43937b6065d"
      },
      "outputs": [],
      "source": [
        "# Get a dataframe of 1000 rows for demonstration purposes\n",
        "df = next(query_bigquery_chunks(max_rows=1000, rows_per_chunk=1000))\n",
        "\n",
        "# Examine the data\n",
        "df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1124422cc200"
      },
      "source": [
        "#### Instantiate the text encoding model\n",
        "\n",
        "Use the [Vertex AI embeddings for text API](https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-text-embeddings) developed by Google for converting text to embeddings.\n",
        "\n",
        "> Text embeddings are a dense vector representation of a piece of content such that, if two pieces of content are semantically similar, their respective embeddings are located near each other in the embedding vector space. This representation can be used to solve common NLP tasks, such as:\n",
        "> - Semantic search: Search text ranked by semantic similarity.\n",
        "> - Recommendation: Return items with text attributes similar to the given text.\n",
        "> - Classification: Return the class of items whose text attributes are similar to the given text.\n",
        "> - Clustering: Cluster items whose text attributes are similar to the given text.\n",
        "> - Outlier Detection: Return items where text attributes are least related to the given text."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "43088937e820"
      },
      "source": [
        "#### Defining an encoding function\n",
        "\n",
        "Define a function to be used later that takes sentences and convert them to embeddings."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5c4520ae99f8"
      },
      "outputs": [],
      "source": [
        "from typing import List, Optional\n",
        "\n",
        "# Load the \"Vertex AI Embeddings for Text\" model\n",
        "from vertexai.preview.language_models import TextEmbeddingModel\n",
        "\n",
        "model = TextEmbeddingModel.from_pretrained(\"textembedding-gecko@001\")\n",
        "\n",
        "\n",
        "# Define an embedding method that uses the model\n",
        "def encode_texts_to_embeddings(sentences: List[str]) -> List[Optional[List[float]]]:\n",
        "    try:\n",
        "        embeddings = model.get_embeddings(sentences)\n",
        "        return [embedding.values for embedding in embeddings]\n",
        "    except Exception:\n",
        "        return [None for _ in range(len(sentences))]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eda80c5936ea"
      },
      "source": [
        "#### Define two more helper functions for converting text to embeddings\n",
        "\n",
        "- `generate_batches`: According to the documentation, each request can handle up to five text instances. Therefore, this method splits `sentences` into batches of five before sending to the embedding API.\n",
        "- `encode_text_to_embedding_batched`: This method calls `generate_batches` to handle batching and then calls the embedding API via `encode_texts_to_embeddings`. It also handles rate-limiting using `time.sleep`. For production use cases, a more sophisticated rate-limiting mechanism that takes retries into account is best."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a0370bd840d2"
      },
      "outputs": [],
      "source": [
        "import functools\n",
        "import time\n",
        "from concurrent.futures import ThreadPoolExecutor\n",
        "from typing import Generator, List, Tuple\n",
        "\n",
        "import numpy as np\n",
        "from tqdm.auto import tqdm\n",
        "\n",
        "\n",
        "# Generator function to yield batches of sentences\n",
        "def generate_batches(\n",
        "    sentences: List[str], batch_size: int\n",
        ") -> Generator[List[str], None, None]:\n",
        "    for i in range(0, len(sentences), batch_size):\n",
        "        yield sentences[i : i + batch_size]\n",
        "\n",
        "\n",
        "def encode_text_to_embedding_batched(\n",
        "    sentences: List[str], api_calls_per_second: int = 10, batch_size: int = 5\n",
        ") -> Tuple[List[bool], np.ndarray]:\n",
        "\n",
        "    embeddings_list: List[List[float]] = []\n",
        "\n",
        "    # Prepare the batches using a generator\n",
        "    batches = generate_batches(sentences, batch_size)\n",
        "\n",
        "    seconds_per_job = 1 / api_calls_per_second\n",
        "\n",
        "    with ThreadPoolExecutor() as executor:\n",
        "        futures = []\n",
        "        for batch in tqdm(\n",
        "            batches, total=math.ceil(len(sentences) / batch_size), position=0\n",
        "        ):\n",
        "            futures.append(\n",
        "                executor.submit(functools.partial(encode_texts_to_embeddings), batch)\n",
        "            )\n",
        "            time.sleep(seconds_per_job)\n",
        "\n",
        "        for future in futures:\n",
        "            embeddings_list.extend(future.result())\n",
        "\n",
        "    is_successful = [\n",
        "        embedding is not None for sentence, embedding in zip(sentences, embeddings_list)\n",
        "    ]\n",
        "    embeddings_list_successful = np.squeeze(\n",
        "        np.stack([embedding for embedding in embeddings_list if embedding is not None])\n",
        "    )\n",
        "    return is_successful, embeddings_list_successful"
      ]
    },
    {
      "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."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9b01baa906b5"
      },
      "outputs": [],
      "source": [
        "# Encode a subset of questions for validation\n",
        "questions = df.title.tolist()[:500]\n",
        "is_successful, question_embeddings = encode_text_to_embedding_batched(\n",
        "    sentences=df.title.tolist()[:500]\n",
        ")\n",
        "\n",
        "# Filter for successfully embedded sentences\n",
        "questions = np.array(questions)[is_successful]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d3761f56648b"
      },
      "source": [
        "Save the dimension size for later usage when creating the index."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d296e181205d"
      },
      "outputs": [],
      "source": [
        "DIMENSIONS = len(question_embeddings[0])\n",
        "\n",
        "print(DIMENSIONS)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d503db448252"
      },
      "source": [
        "##### Sort questions in order of similarity.\n",
        "\n",
        "According to the [embedding documentation](https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-text-embeddings#colab_example_of_semantic_search_using_embeddings), the similarity of embeddings is calculated using the dot-product.\n",
        "\n",
        "- Calculate vector similarity using `np.dot`.\n",
        "- Sort by similarity.\n",
        "- Print results for inspection."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "95e408daf219"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "\n",
        "question_index = random.randint(0, 99)\n",
        "\n",
        "print(f\"Query question = {questions[question_index]}\")\n",
        "\n",
        "# Get similarity scores for each embedding by using dot-product.\n",
        "scores = np.dot(question_embeddings[question_index], question_embeddings.T)\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 embeddings in JSONL format\n",
        "\n",
        "The data must be formatted in JSONL format, which means each embedding dictionary is written as an individual JSON object on its own line.\n",
        "\n",
        "For more information, see [Input data format and structure](https://cloud.google.com/vertex-ai/docs/matching-engine/match-eng-setup/format-structure#data-file-formats)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7c1193aca5d1"
      },
      "outputs": [],
      "source": [
        "import tempfile\n",
        "from pathlib import Path\n",
        "\n",
        "# Create temporary file to write embeddings to\n",
        "embeddings_file_path = Path(tempfile.mkdtemp())\n",
        "\n",
        "print(f\"Embeddings directory: {embeddings_file_path}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "279c0bbfc6bb"
      },
      "source": [
        "Write embeddings in batches to prevent out-of-memory errors"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "307f468a3ecd"
      },
      "outputs": [],
      "source": [
        "import gc\n",
        "import json\n",
        "\n",
        "BQ_NUM_ROWS = 50000\n",
        "BQ_CHUNK_SIZE = 1000\n",
        "BQ_NUM_CHUNKS = math.ceil(BQ_NUM_ROWS / BQ_CHUNK_SIZE)\n",
        "\n",
        "START_CHUNK = 0\n",
        "\n",
        "# Create a rate limit of 300 requests per minute. Adjust this depending on your quota.\n",
        "API_CALLS_PER_SECOND = 300 / 60\n",
        "# According to the docs, each request can process 5 instances per request\n",
        "ITEMS_PER_REQUEST = 5\n",
        "\n",
        "# Loop through each generated dataframe, convert\n",
        "for i, df in tqdm(\n",
        "    enumerate(\n",
        "        query_bigquery_chunks(\n",
        "            max_rows=BQ_NUM_ROWS, rows_per_chunk=BQ_CHUNK_SIZE, start_chunk=START_CHUNK\n",
        "        )\n",
        "    ),\n",
        "    total=BQ_NUM_CHUNKS - START_CHUNK,\n",
        "    position=-1,\n",
        "    desc=\"Chunk of rows from BigQuery\",\n",
        "):\n",
        "    # Create a unique output file for each chunk\n",
        "    chunk_path = embeddings_file_path.joinpath(\n",
        "        f\"{embeddings_file_path.stem}_{i+START_CHUNK}.json\"\n",
        "    )\n",
        "    with open(chunk_path, \"a\") as f:\n",
        "        id_chunk = df.id\n",
        "\n",
        "        # Convert batch to embeddings\n",
        "        is_successful, question_chunk_embeddings = encode_text_to_embedding_batched(\n",
        "            sentences=df.title_with_body,\n",
        "            api_calls_per_second=API_CALLS_PER_SECOND,\n",
        "            batch_size=ITEMS_PER_REQUEST,\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[is_successful], question_chunk_embeddings)\n",
        "        ]\n",
        "        f.writelines(embeddings_formatted)\n",
        "\n",
        "        # Delete the DataFrame and any other large data structures\n",
        "        del df\n",
        "        gc.collect()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QuVl8DrWG8NS"
      },
      "source": [
        "Upload the training data to a Google Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3PgsA_vbI8Vg"
      },
      "outputs": [],
      "source": [
        "remote_folder = f\"{BUCKET_URI}/{embeddings_file_path.stem}/\"\n",
        "! gsutil -m cp -r {embeddings_file_path}/* {remote_folder}"
      ]
    },
    {
      "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": [
        "DISPLAY_NAME = \"stack_overflow\"\n",
        "DESCRIPTION = \"question titles and bodies from stackoverflow\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "svLYiDf0OD2G"
      },
      "source": [
        "Create the index configuration\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": "dffb00b23f5a"
      },
      "outputs": [],
      "source": [
        "DIMENSIONS = 768\n",
        "\n",
        "tree_ah_index = aiplatform.MatchingEngineIndex.create_tree_ah_index(\n",
        "    display_name=DISPLAY_NAME,\n",
        "    contents_delta_uri=remote_folder,\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=80,\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": [
        "Use the resource name to 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"
      ]
    },
    {
      "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",
        "    public_endpoint_enabled=True,\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\"\n",
        "\n",
        "DEPLOYED_INDEX_ID"
      ]
    },
    {
      "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": "3cbfe4fd103a"
      },
      "source": [
        "#### Verify number of declared items matches the number of embeddings\n",
        "\n",
        "Each IndexEndpoint can have multiple indexes deployed to it. For each index, you can retrieved the number of deployed vectors using the `index_endpoint._gca_resource.index_stats.vectors_count`. The numbers may not match exactly due to potential failures using the embedding service."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "93f89a15f642"
      },
      "outputs": [],
      "source": [
        "number_of_vectors = sum(\n",
        "    aiplatform.MatchingEngineIndex(\n",
        "        deployed_index.index\n",
        "    )._gca_resource.index_stats.vectors_count\n",
        "    for deployed_index in my_index_endpoint.deployed_indexes\n",
        ")\n",
        "\n",
        "print(f\"Expected: {BQ_NUM_ROWS}, Actual: {number_of_vectors}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6LCGvBNvBd8D"
      },
      "source": [
        "## Create Online Queries\n",
        "\n",
        "After you build your indexes, you can query against the deployed index to find nearest neighbors.\n",
        "\n",
        "Note: For the `DOT_PRODUCT_DISTANCE` distance type, the \"distance\" property returned with each MatchNeighbor actually refers to the similarity."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ae9996f185fe"
      },
      "outputs": [],
      "source": [
        "test_embeddings = encode_texts_to_embeddings(sentences=[\"Install GPU for Tensorflow\"])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A3KYVw5HB-4v"
      },
      "outputs": [],
      "source": [
        "# Test query\n",
        "NUM_NEIGHBOURS = 10\n",
        "\n",
        "response = my_index_endpoint.find_neighbors(\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": "8a2879d3d9ca"
      },
      "source": [
        "Verify that the retrieved results are relevant by checking the StackOverflow links."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7c8682079e21"
      },
      "outputs": [],
      "source": [
        "for match_index, neighbor in enumerate(response[0]):\n",
        "    print(f\"https://stackoverflow.com/questions/{neighbor.id}\")"
      ]
    },
    {
      "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": [
        "import os\n",
        "\n",
        "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",
        "\n",
        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
        "    ! gsutil rm -rf {BUCKET_URI}"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "sdk_matching_engine_create_stack_overflow_embeddings_vertex.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}