notebooks/official/vector_search/sdk_vector_search_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_vector_search_create_stack_overflow_embeddings_vertex.ipynb&dl=notebooks%2Fofficial%2Fvector_search%2Fsdk_vector_search_create_stack_overflow_embeddings_vertex.ipynb)\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_create_stack_overflow_embeddings_vertex.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"><br> Run 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_create_stack_overflow_embeddings_vertex.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://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vector_search/sdk_vector_search_create_stack_overflow_embeddings_vertex.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br>\n",
        "      View on GitHub\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_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\"><br>\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 Vertex AI 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": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "81711614c199"
      },
      "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",
        "                        'google-cloud-bigquery[pandas]' "
      ]
    },
    {
      "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": "ffcde4d56c00"
      },
      "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": "dd28c9e4f067"
      },
      "source": [
        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\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). 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\"}\n",
        "\n",
        "import vertexai\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "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": "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.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",
        "# 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/vector-search/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.to_list(),\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=LOCATION, 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 Index resource."
      ]
    },
    {
      "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_vector_search_create_stack_overflow_embeddings_vertex.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}