embeddings/hybrid-search.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wDrRkb1v_LMi"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 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": "pNK49xFFj8M1"
      },
      "source": [
        "# Combining Semantic & Keyword Search: A Hybrid Search Tutorial with Vertex AI Vector Search"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fArkSYhZ8Vr1"
      },
      "source": [
        "Vector Search supports hybrid search, a popular architecture pattern in information retrieval (IR) that combines both semantic search and keyword search (also called token-based search). With hybrid search, developers can take advantage of the best of the two approaches, effectively providing higher search quality.\n",
        "\n",
        "This tutorial explains the concepts of hybrid search, semantic search, and token-based search, and includes examples of how to set up token-based search and hybrid search."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uUVFcXFpCMxc"
      },
      "source": [
        "This tutorial takes about 1 - 1.5 hours, mostly on waiting for Vector Search deployments (30 minutes x 2 times)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JAPoU8Sm5E6e"
      },
      "source": [
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/embeddings/hybrid-search.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fembeddings%2Fhybrid-search.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/embeddings/hybrid-search.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/embeddings/hybrid-search.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://upload.wikimedia.org/wikipedia/commons/9/91/Octicons-mark-github.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rsdWWUgC8nfO"
      },
      "source": [
        "# Why does hybrid search matter?\n",
        "\n",
        "As described in [Overview of Vector Search](https://cloud.google.com/vertex-ai/docs/vector-search/overview), semantic search with Vector Search can find items with semantic similarity by using queries.\n",
        "\n",
        "Embedding models such as Vertex AI embeddings build a vector space as a map of content meanings. Each text or multimodal embedding is a location in the map that represents the meaning of some content. As a simplified example, when an embedding model takes a text that discusses movies for 10%, music for 2%, and actors for 30%, it could represent this text with an embedding `[0.1, 0.02, 0.3]`. With Vector Search, you can quickly find other embeddings in its neighborhood. This searching by content meaning is called semantic search.\n",
        "\n",
        "![Semantic search image](https://cloud.google.com/vertex-ai/docs/vector-search/images/semantic_search.gif)\n",
        "\n",
        "Semantic search with embeddings and vector search can help make IT systems as smart as experienced librarians or shop staff. Embeddings can be used for tying different business data with their meanings; for example, queries and search results; texts and images; user activities and recommended products; English texts and Japanese texts; or sensor data and alerting conditions. With this capability, there's a wide variety of use cases for embeddings.\n",
        "\n",
        "## Why combine semantic search with keyword-based search?\n",
        "\n",
        "Semantic search doesn't cover all the possible requirements for information retrieval applications, such as [Retrieval-Augmented Generation (RAG)](https://cloud.google.com/use-cases/retrieval-augmented-generation). Semantic search can only find data that the embedding model can make sense of. For example, queries or datasets with arbitrary product numbers or SKUs, brand new product names that were added recently, and corporate proprietary codenames don't work with semantic search because they aren't included in the training dataset of the embedding model. This is called \"out of domain\" data.\n",
        "\n",
        "In such cases, you would need to combine semantic search with keyword-based (also called token-based) search to form a hybrid search. With hybrid search, you can take advantage of both semantic and token-based search to achieve higher search quality.\n",
        "\n",
        "One of the most popular hybrid search systems is Google Search. The service incorporated [semantic search in 2015 with RankBrain model](https://blog.google/products/search/how-ai-powers-great-search-results/), in addition to its token-based keyword search algorithm. With the introduction of hybrid search, Google Search was able to improve the search quality significantly by addressing the two requirements: search by meaning and search by keyword.\n",
        "\n",
        "In the past, building a hybrid search engine was a complex task. Just like with Google Search, you have to build and operate two different kinds of search engines (semantic search and token-based search) and merge and rank the results from them. With hybrid search support in Vector Search, you can build your own hybrid search system with a single Vector Search index, customized to your business requirements.\n",
        "\n",
        "## How token-based search works\n",
        "\n",
        "How does token-based search in Vector Search work? After [splitting the text into tokens](https://huggingface.co/docs/transformers/en/tokenizer_summary) (such as words or sub-words), you can use popular sparse embedding algorithms such as [TF-IDF](https://en.wikipedia.org/wiki/Tf%E2%80%93idf), [BM25](https://en.wikipedia.org/wiki/Okapi_BM25), or [SPLADE](https://en.wikipedia.org/wiki/Learned_sparse_retrieval) to generate sparse embedding for the text.\n",
        "\n",
        "A simplified explanation of sparse embeddings is that they are vectors that represent how many times each word or sub-word appears in the text. Typical sparse embeddings don't take semantics of the text into account.\n",
        "\n",
        "![Sparse embeddings image](https://cloud.google.com/vertex-ai/docs/vector-search/images/sparse_embeddings.png)\n",
        "\n",
        "There could be thousands of different words used in texts. Thus, this embedding usually has tens of thousands of dimensions, with only a few dimensions in them having non-zero values. This is why they're called \"sparse\" embeddings. The majority of their values are zeroes. This sparse embedding space works as a map of keywords, similar to an index of books.\n",
        "\n",
        "In this sparse embedding space, you can find similar embeddings by looking at the neighborhood of a query embedding. These embeddings are similar in terms of the distribution of the keywords used in their texts.\n",
        "\n",
        "![Token search image](https://cloud.google.com/vertex-ai/docs/vector-search/images/token_search.gif)\n",
        "\n",
        "This is the basic mechanism of the token-based search with sparse embeddings. With hybrid search in Vector Search, you can mix both dense and sparse embeddings into a single vector index and run queries with dense embeddings, sparse embeddings, or both. The result is a combination of semantic search and token-based search results.\n",
        "\n",
        "Hybrid search also provides shorter query latency compared to a token-based search engine with an [inverted index](https://en.wikipedia.org/wiki/Inverted_index) design. Just like vector search for semantic search, each query with dense or sparse embeddings finishes within milliseconds, even with millions or billions of items.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KliK1dHx98_f"
      },
      "source": [
        "# Example: How to use token-based search\n",
        "\n",
        "To explain how to use token-based search, the following sections include code examples that generate sparse embeddings and build an index with them on Vector Search.\n",
        "\n",
        "## Setup\n",
        "Let's start setting up the SDK and environment variables."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "axMumYbIaUpk"
      },
      "source": [
        "### Install Python SDK\n",
        "\n",
        "This tutorial uses Vertex AI SDK and Cloud Storage SDK.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sZGqrhg5aZnj"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --quiet --user google-cloud-aiplatform google-cloud-storage"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xn6t9lVFafUz"
      },
      "source": [
        "To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which will restart the current kernel."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7FvjSJBTacMr"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'status': 'ok', 'restart': True}"
            ]
          },
          "execution_count": 2,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# 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": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you're running this notebook on Google Colab, run the cell below to authenticate your environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NyKGtVQjgx13"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DF4l8DTdWgPY"
      },
      "source": [
        "### Set Google Cloud project information and initialize Vertex AI SDK\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HPfsNnW1e5mm"
      },
      "outputs": [],
      "source": [
        "# get project ID\n",
        "PROJECT_ID = ! gcloud config get project\n",
        "PROJECT_ID = PROJECT_ID[0]\n",
        "LOCATION = \"us-central1\"\n",
        "if PROJECT_ID == \"(unset)\":\n",
        "    print(f\"Please set the project ID manually below\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xR81WstRZyaH"
      },
      "outputs": [],
      "source": [
        "# define project information\n",
        "if PROJECT_ID == \"(unset)\":\n",
        "    PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GXr6PQWRZv-q"
      },
      "outputs": [],
      "source": [
        "# generate an unique id for this session\n",
        "from datetime import datetime\n",
        "\n",
        "UID = datetime.now().strftime(\"%m%d%H%M\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QulDaFf7Yodp"
      },
      "source": [
        "## Prepare a sample dataset\n",
        "\n",
        "The first step is to prepare a data file to build an index for sparse embeddings, based on the data format described in [Input data format and structure](https://cloud.google.com/vertex-ai/docs/vector-search/setup/format-structure).\n",
        "\n",
        "In JSON, the data file looks like this:\n",
        "\n",
        "```\n",
        "{\"id\": \"3\", \"sparse_embedding\": {\"values\": [0.1, 0.2], \"dimensions\": [1, 4]}}\n",
        "{\"id\": \"4\", \"sparse_embedding\": {\"values\": [-0.4, 0.2, -1.3], \"dimensions\": [10, 20, 20]}}\n",
        "```\n",
        "\n",
        "Each item should have a `sparse_embedding` property that has `values` and `dimensions` properties. Sparse embeddings have thousands of dimensions with a few non-zero values. This data format works efficiently because it contains the non-zero values only with their positions in the space.\n",
        "\n",
        "As a sample dataset, we'll use the [Google Merch Shop](https://shop.merch.google/) dataset, which has about 200 rows of Google-branded goods."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KAa1Ff5zaNKf"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "0                          Google Sticker\n",
              "1                    Google Cloud Sticker\n",
              "2                       Android Black Pen\n",
              "3                   Google Ombre Lime Pen\n",
              "4                    For Everyone Eco Pen\n",
              "                      ...                \n",
              "197        Google Recycled Black Backpack\n",
              "198    Google Cascades Unisex Zip Sweater\n",
              "199    Google Cascades Womens Zip Sweater\n",
              "200         Google Cloud Skyline Backpack\n",
              "201       Google City Black Tote Backpack\n",
              "Name: title, Length: 202, dtype: object"
            ]
          },
          "execution_count": 4,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "import pandas as pd\n",
        "\n",
        "CSV_URL = \"https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/sample-apps/photo-discovery/ag-web/google_merch_shop_items.csv\"\n",
        "\n",
        "# Load the CSV file into a DataFrame\n",
        "df = pd.read_csv(CSV_URL)\n",
        "df[\"title\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AkcwG37lAfoS"
      },
      "source": [
        "## How to use Token-based Search\n",
        "\n",
        "With the dataset, we will create sparse embeddings for implementing a token-based search with Vector Search."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UBzaDGpYYodr"
      },
      "source": [
        "### Prepare a TF-IDF vectorizer\n",
        "With this dataset, we'll train a vectorizer, a model that generates sparse embeddings from a text. This example uses [TfidfVectorizer in scikit-learn](https://scikit-learn.org/stable/modules/generated/sklearn.feature_extraction.text.TfidfVectorizer.html), which is a basic vectorizer that uses the [TF-IDF algorithm](https://en.wikipedia.org/wiki/Tf%E2%80%93idf)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DfQhFEa7Yodr"
      },
      "outputs": [],
      "source": [
        "from sklearn.feature_extraction.text import TfidfVectorizer\n",
        "\n",
        "# Sample Text Data\n",
        "corpus = df.title.tolist()\n",
        "\n",
        "# Initialize TfidfVectorizer\n",
        "vectorizer = TfidfVectorizer()\n",
        "\n",
        "# Fit and Transform\n",
        "vectorizer.fit_transform(corpus)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VCFLoP78B7G2"
      },
      "source": [
        "The variable `corpus` holds a list of the 200 item names, such as \"Google Sticker\" or \"Chrome Dino Pin\". Then, the code passes them to the vectorizer by calling the `fit_transform()` function. With that, the vectorizer gets ready to generate sparse embeddings.\n",
        "\n",
        "TF-IDF vectorizer tries to give higher weight to signature words in the dataset (such as \"Shirts\" or \"Dino\") compared to trivial words (such as \"The\", \"a\", or \"of\"), and counts how many times those signature words are used in the specified document. Each value of a sparse embedding represents a frequency of each word based on the counts. For more information about TF-IDF, see the later section \"How do TF-IDF and TfidfVectorizer work?\".\n",
        "\n",
        "In this example, we use the basic word-level tokenization and TF-IDF vectorization for simplicity. In production development, you can choose any other options for tokenizations and vectorizations for generating sparse embeddings based on your requirements. For tokenizers, in many cases [subword tokenizers](https://www.tensorflow.org/text/guide/subwords_tokenizer) perform well compared to the word-level tokenization and are popular choices. For vectorizers, [BM25](https://en.wikipedia.org/wiki/Okapi_BM25) is popular as an improved version of TF-IDF. [SPLADE](https://en.wikipedia.org/wiki/Learned_sparse_retrieval) is another popular vectorization algorithm that takes some semantics for the sparse embedding.\n",
        "\n",
        "### Get a sparse embedding\n",
        "To make the vectorizer easier to use with Vector Search, we'll define a wrapper function, get_sparse_embedding():"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "as6pjlwBYodr"
      },
      "outputs": [],
      "source": [
        "def get_sparse_embedding(text):\n",
        "    # Transform Text into TF-IDF Sparse Vector\n",
        "    tfidf_vector = vectorizer.transform([text])\n",
        "\n",
        "    # Create Sparse Embedding for the New Text\n",
        "    values = []\n",
        "    dims = []\n",
        "    for i, tfidf_value in enumerate(tfidf_vector.data):\n",
        "        values.append(float(tfidf_value))\n",
        "        dims.append(int(tfidf_vector.indices[i]))\n",
        "    return {\"values\": values, \"dimensions\": dims}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SxPHuqkwCdzW"
      },
      "source": [
        "This function passes the parameter \"text\" to the vectorizer to generate a sparse embedding. Then convert it to the `{\"values\": ...., \"dimensions\": ...}` format mentioned earlier for building a Vector Search sparse index.\n",
        "\n",
        "You can test this function:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nC2JBCZAYodr"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'values': [0.6756557405747007, 0.5212913389979028, 0.5212913389979028],\n",
              " 'dimensions': [157, 48, 33]}"
            ]
          },
          "execution_count": 7,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "text_text = \"Chrome Dino Pin\"\n",
        "get_sparse_embedding(text_text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O6SfR7Z6hFIC"
      },
      "source": [
        "### Create an input data file\n",
        "For this example, we'll generate sparse embeddings for all 200 items."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BizNf7S2f298"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Creating gs://gcp-samples-ic0-vs-hybridsearch-07300053/...\n"
          ]
        }
      ],
      "source": [
        "# create bucket\n",
        "BUCKET_URI = f\"gs://{PROJECT_ID}-vs-hybridsearch-{UID}\"\n",
        "! gsutil mb -l $LOCATION -p $PROJECT_ID $BUCKET_URI"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wvTkcSl0hgVf"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "[{'id': 0,\n",
              "  'title': 'Google Sticker',\n",
              "  'sparse_embedding': {'values': [0.933008728540452, 0.359853737603667],\n",
              "   'dimensions': [191, 78]}},\n",
              " {'id': 1,\n",
              "  'title': 'Google Cloud Sticker',\n",
              "  'sparse_embedding': {'values': [0.6884926145791601,\n",
              "    0.26554589800720163,\n",
              "    0.6748802084233179],\n",
              "   'dimensions': [191, 78, 37]}},\n",
              " {'id': 2,\n",
              "  'title': 'Android Black Pen',\n",
              "  'sparse_embedding': {'values': [0.578954308910934,\n",
              "    0.5468824106727572,\n",
              "    0.6047574200373373],\n",
              "   'dimensions': [153, 16, 2]}},\n",
              " {'id': 3,\n",
              "  'title': 'Google Ombre Lime Pen',\n",
              "  'sparse_embedding': {'values': [0.4417396140901719,\n",
              "    0.580646486216207,\n",
              "    0.6623352815578535,\n",
              "    0.17037530980725205],\n",
              "   'dimensions': [153, 141, 119, 78]}},\n",
              " {'id': 4,\n",
              "  'title': 'For Everyone Eco Pen',\n",
              "  'sparse_embedding': {'values': [0.4749050808211265,\n",
              "    0.5533375082080382,\n",
              "    0.5533375082080382,\n",
              "    0.4026169000806076],\n",
              "   'dimensions': [153, 67, 61, 56]}}]"
            ]
          },
          "execution_count": 9,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "items = []\n",
        "for i in range(len(df)):\n",
        "    id = i\n",
        "    title = df.title[i]\n",
        "    sparse_embedding = get_sparse_embedding(title)\n",
        "    items.append({\"id\": id, \"title\": title, \"sparse_embedding\": sparse_embedding})\n",
        "items[:5]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N0utF16hCv8H"
      },
      "source": [
        "Then, save them as a JSONL file `items.json` and upload to a Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M_XphPWwh_4h"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Copying file://items.json [Content-Type=application/json]...\n",
            "/ [1 files][ 39.9 KiB/ 39.9 KiB]                                                \n",
            "Operation completed over 1 objects/39.9 KiB.                                     \n"
          ]
        }
      ],
      "source": [
        "# output as a JSONL file and save to the GCS bucket\n",
        "with open(\"items.json\", \"w\") as f:\n",
        "    for item in items:\n",
        "        f.write(f\"{item}\\n\")\n",
        "! gsutil cp items.json $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "73DaRoBFffQP"
      },
      "source": [
        "## Create a sparse embedding index in Vector Search\n",
        "Next, we'll build and deploy a sparse embedding index in Vector Search. This is the same procedure that is documented in the [Vector Search quickstart](https://cloud.google.com/vertex-ai/docs/vector-search/quickstart#create_an_index)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KnTpBptndTSq"
      },
      "outputs": [],
      "source": [
        "# init the aiplatform package\n",
        "from google.cloud import aiplatform\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YW5F2nW2fq16"
      },
      "outputs": [],
      "source": [
        "# create Index (this should finish less than one minute)\n",
        "my_sparse_index = aiplatform.MatchingEngineIndex.create_tree_ah_index(\n",
        "    display_name=f\"vs-tokensearch-index-{UID}\",\n",
        "    contents_delta_uri=BUCKET_URI,\n",
        "    dimensions=768,\n",
        "    approximate_neighbors_count=10,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CCIYzxTSDBcp"
      },
      "source": [
        "To use the index, you need to create an index endpoint. It works as a server instance accepting query requests for your index."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dWYwANgoi9W7"
      },
      "outputs": [],
      "source": [
        "# create `IndexEndpoint`\n",
        "my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint.create(\n",
        "    display_name=f\"vs-hybridsearch-index-endpoint-{UID}\", public_endpoint_enabled=True\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dre2Hq2nDEWq"
      },
      "source": [
        "With the index endpoint, deploy the index by specifying a unique deployed index ID."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6eueKiLDjmgt"
      },
      "outputs": [],
      "source": [
        "# deploy the Index to the Index Endpoint (this takes upto 30 minutes)\n",
        "# Note: for Colab Enterprise, this may output timeout error, but you can ignore it\n",
        "DEPLOYED_TOKEN_INDEX_ID = f\"vs_tokensearch_deployed_{UID}\"\n",
        "my_index_endpoint.deploy_index(\n",
        "    index=my_sparse_index, deployed_index_id=DEPLOYED_TOKEN_INDEX_ID\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-HzATPOLDW8I"
      },
      "source": [
        "After waiting for the deployment, we're ready to run a test query. This may take up to 30 minutes."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "je343PRYkz8N"
      },
      "source": [
        "### Run a query with a sparse embedding index\n",
        "To run a query with a sparse embedding index, you need to create a `HybridQuery` object to encapsulate the sparse embedding of the query text, like in the following example:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ohU8qSItjvSY"
      },
      "outputs": [],
      "source": [
        "from google.cloud.aiplatform.matching_engine.matching_engine_index_endpoint import (\n",
        "    HybridQuery,\n",
        ")\n",
        "\n",
        "# create HybridQuery\n",
        "query_text = \"Kids\"\n",
        "query_emb = get_sparse_embedding(query_text)\n",
        "query = HybridQuery(\n",
        "    sparse_embedding_dimensions=query_emb[\"dimensions\"],\n",
        "    sparse_embedding_values=query_emb[\"values\"],\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QR9NZF1NDmep"
      },
      "source": [
        "This example code uses the text `Kids` for the query. Now, run a query with the `HybridQuery` object."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fHjRB4bKDppP"
      },
      "outputs": [],
      "source": [
        "# build a query request\n",
        "response = my_index_endpoint.find_neighbors(\n",
        "    deployed_index_id=DEPLOYED_TOKEN_INDEX_ID,\n",
        "    queries=[query],\n",
        "    num_neighbors=5,\n",
        ")\n",
        "\n",
        "# print results\n",
        "for idx, neighbor in enumerate(response[0]):\n",
        "    title = df.title[int(neighbor.id)]\n",
        "    print(f\"{title:<40}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "42psI25wDubo"
      },
      "source": [
        "This should provide output like the following:\n",
        "\n",
        "```\n",
        "Google Blue Kids Sunglasses\n",
        "Google Red Kids Sunglasses\n",
        "YouTube Kids Coloring Pencils\n",
        "YouTube Kids Character Sticker Sheet\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HhRUr_qgD2WU"
      },
      "source": [
        "Out of the 200 items, the result contains the item names that have the keyword `Kids`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gsf50GDiAnHy"
      },
      "source": [
        "# Example: How to use hybrid search\n",
        "\n",
        "This example combines token-based search with semantic search to create hybrid search in Vector Search.\n",
        "\n",
        "## How to create hybrid index\n",
        "\n",
        "To build a hybrid index, each item should have both `embedding` (for dense embedding) and `sparse_embedding`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sV68SZhiCynL"
      },
      "outputs": [],
      "source": [
        "# get text embedding model\n",
        "from vertexai.preview.language_models import TextEmbeddingModel\n",
        "\n",
        "model = TextEmbeddingModel.from_pretrained(\"textembedding-gecko\")\n",
        "\n",
        "\n",
        "# wrapper\n",
        "def get_dense_embedding(text):\n",
        "    return model.get_embeddings([text])[0].values\n",
        "\n",
        "\n",
        "# test it\n",
        "get_dense_embedding(\"Chrome Dino Pin\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0ZZgv0p7DMXU"
      },
      "outputs": [],
      "source": [
        "items = []\n",
        "for i in range(len(df)):\n",
        "    id = i\n",
        "    title = df.title[i]\n",
        "    dense_embedding = get_dense_embedding(title)\n",
        "    sparse_embedding = get_sparse_embedding(title)\n",
        "    items.append(\n",
        "        {\n",
        "            \"id\": id,\n",
        "            \"title\": title,\n",
        "            \"embedding\": dense_embedding,\n",
        "            \"sparse_embedding\": sparse_embedding,\n",
        "        }\n",
        "    )\n",
        "items[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dOXeX3VIELqw"
      },
      "source": [
        "The `get_dense_embedding()` function uses Vertex AI Embedding API for generating text embedding with 768 dimensions. This generates both dense and sparse embeddings in the following format:\n",
        "\n",
        "```\n",
        "{\n",
        "  'id': 0,\n",
        "  'title': 'Google Sticker',\n",
        "  'embedding':\n",
        "    [0.022880317643284798,\n",
        "    -0.03315234184265137,\n",
        "    ...\n",
        "    -0.03309667482972145,\n",
        "    0.04621824622154236],\n",
        "  'sparse_embedding': {\n",
        "    'values': [0.933008728540452, 0.359853737603667],\n",
        "    'dimensions': [191, 78]\n",
        "  }\n",
        "}\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2MbJOIbVEUZY"
      },
      "source": [
        "The rest of the process is the same as the previous example. Upload the JSONL file to the Cloud Storage bucket, create a Vector Search index with the file, and deploy the index to the index endpoint."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AehqVQdeFDeU"
      },
      "outputs": [],
      "source": [
        "# output as a JSONL file and save to the GCS bucket\n",
        "with open(\"items.json\", \"w\") as f:\n",
        "    for item in items:\n",
        "        f.write(f\"{item}\\n\")\n",
        "! gsutil cp items.json $BUCKET_URI"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5-3bb628CyzG"
      },
      "outputs": [],
      "source": [
        "# create Index (this should finish less than one minute)\n",
        "my_hybrid_index = aiplatform.MatchingEngineIndex.create_tree_ah_index(\n",
        "    display_name=f\"vs-hybridsearch-index-{UID}\",\n",
        "    contents_delta_uri=BUCKET_URI,\n",
        "    dimensions=768,\n",
        "    approximate_neighbors_count=10,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T0VRwyrHHtRn"
      },
      "outputs": [],
      "source": [
        "# deploy index\n",
        "DEPLOYED_HYBRID_INDEX_ID = f\"vs_hybridsearch_deployed_{UID}\"\n",
        "my_index_endpoint.deploy_index(\n",
        "    index=my_hybrid_index, deployed_index_id=DEPLOYED_HYBRID_INDEX_ID\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8iC0IS_2AvLV"
      },
      "source": [
        "## Run a hybrid query\n",
        "\n",
        "After deploying the hybrid index, you can run a hybrid query:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fj7UFoxTIFZR"
      },
      "outputs": [],
      "source": [
        "# create HybridQuery\n",
        "query_text = \"Kids\"\n",
        "query_dense_emb = get_dense_embedding(query_text)\n",
        "query_sparse_emb = get_sparse_embedding(query_text)\n",
        "query = HybridQuery(\n",
        "    dense_embedding=query_dense_emb,\n",
        "    sparse_embedding_dimensions=query_sparse_emb[\"dimensions\"],\n",
        "    sparse_embedding_values=query_sparse_emb[\"values\"],\n",
        "    rrf_ranking_alpha=0.5,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mfPKXIa7Ek47"
      },
      "source": [
        "For the query text `Kids`, generate both dense and sparse embeddings for the word, and encapsulate them to the `HybridQuery` object. The difference from the previous `HybridQuery` is two additional parameters: `dense_embedding` and `rrf_ranking_alpha`.\n",
        "\n",
        "This time, we'll print distances for each item:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h9bLgIdJEu51"
      },
      "outputs": [],
      "source": [
        "# run a hybrid query\n",
        "response = my_index_endpoint.find_neighbors(\n",
        "    deployed_index_id=DEPLOYED_HYBRID_INDEX_ID,\n",
        "    queries=[query],\n",
        "    num_neighbors=10,\n",
        ")\n",
        "\n",
        "# print results\n",
        "for idx, neighbor in enumerate(response[0]):\n",
        "    title = df.title[int(neighbor.id)]\n",
        "    dense_dist = neighbor.distance if neighbor.distance else 0.0\n",
        "    sparse_dist = neighbor.sparse_distance if neighbor.sparse_distance else 0.0\n",
        "    print(f\"{title:<40}: dense_dist: {dense_dist:.3f}, sparse_dist: {sparse_dist:.3f}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dVomjOksE7e2"
      },
      "source": [
        "In each `neighbor` object, there's a `distance` property that has the distance between the query and the item with the dense embedding, and a `sparse_distance` property that has the distance with the sparse embedding. These values are inverted distances, so a higher value means a shorter distance.\n",
        "\n",
        "By running a query with `HybridQuery`, you get the following result:\n",
        "\n",
        "```\n",
        "Google Blue Kids Sunglasses             : dense_dist: 0.677, sparse_dist: 0.606\n",
        "Google Red Kids Sunglasses              : dense_dist: 0.665, sparse_dist: 0.572\n",
        "YouTube Kids Coloring Pencils           : dense_dist: 0.655, sparse_dist: 0.478\n",
        "YouTube Kids Character Sticker Sheet    : dense_dist: 0.644, sparse_dist: 0.468\n",
        "Google White Classic Youth Tee          : dense_dist: 0.645, sparse_dist: 0.000\n",
        "Google Doogler Youth Tee                : dense_dist: 0.639, sparse_dist: 0.000\n",
        "Google Indigo Youth Tee                 : dense_dist: 0.637, sparse_dist: 0.000\n",
        "Google Black Classic Youth Tee          : dense_dist: 0.632, sparse_dist: 0.000\n",
        "Chrome Dino Glow-in-the-Dark Youth Tee  : dense_dist: 0.632, sparse_dist: 0.000\n",
        "Google Bike Youth Tee                   : dense_dist: 0.629, sparse_dist: 0.000\n",
        "```\n",
        "\n",
        "In addition to the token-based search results that have the `Kids` keyword, there are also semantic search results included. For example, `Google White Classic Youth Tee` is included because the embedding model knows that `Youth` and `Kids` are semantically similar.\n",
        "\n",
        "To merge the token-based and semantic search results, hybrid search uses [Reciprocal Rank Fusion (RRF)](https://plg.uwaterloo.ca/%7Egvcormac/cormacksigir09-rrf.pdf). For more information about RRF and how to specify the rrf_ranking_alpha parameter, see the later section \"What is Reciprocal Rank Fusion?\".\n",
        "\n",
        "### Reranking\n",
        "RRF provides a way to merge the ranking from semantic and token-based search results. In many production information retrieval or recommender systems, the results will be going through further precision ranking algorithms - so called reranking. With the combination of the millisecond level fast retrieval with vector search, and precision reranking on the results, you can build [multi-stage systems](https://cloud.google.com/blog/products/ai-machine-learning/scaling-deep-retrieval-tensorflow-two-towers-architecture?e=48754805) that provide higher search quality or recommendation performance.\n",
        "\n",
        "![multi-stage systems](https://cloud.google.com/vertex-ai/docs/vector-search/images/reranking.png)\n",
        "\n",
        "# Conclusion\n",
        "\n",
        "In this tutorial, we have learned the concept of hybrid search, how to build vector search index for token-based search with sparse embeddings, and how to combine it with the semantic search with dense embeddings. As we see on the example, hybrid search provides a way to mix the best part of the two approaches, resulting in higher search quality in production systems.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HRkOkkB7ruQh"
      },
      "source": [
        "# Cleaning up\n",
        "\n",
        "In case you are using your own Cloud project, not a temporary project on Qwiklab, please make sure to delete all the Indexes, Index Endpoints and Cloud Storage buckets after finishing this tutorial. Otherwise the remaining objects would incur unexpected costs.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RUy6sqgCryoU"
      },
      "outputs": [],
      "source": [
        "# wait for a confirmation\n",
        "input(\"Press Enter to delete Index Endpoint, Index and Cloud Storage bucket:\")\n",
        "\n",
        "# delete Index Endpoint\n",
        "my_index_endpoint.undeploy_all()\n",
        "my_index_endpoint.delete(force=True)\n",
        "\n",
        "# delete Indexes\n",
        "my_sparse_index.delete()\n",
        "my_hybrid_index.delete()\n",
        "\n",
        "# delete Cloud Storage bucket\n",
        "! gsutil rm -r \"{BUCKET_URI}\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v94RXQ1HrpjI"
      },
      "source": [
        "# Start using hybrid search\n",
        "The following resources can help you get started with using hybrid search in Vector Search.\n",
        "\n",
        "## Hybrid search resources\n",
        "- [Input data format and structure](https://cloud.google.com/vertex-ai/docs/vector-search/setup/format-structure): Input data format for building sparse embedding index\n",
        "- [Query public index to get nearest neighbors](https://cloud.google.com/vertex-ai/docs/vector-search/query-index-public-endpoint): How to run queries with hybrid search\n",
        "- [Reciprocal Rank Fusion outperforms Condorcet and individual Rank Learning Methods](https://plg.uwaterloo.ca/%7Egvcormac/cormacksigir09-rrf.pdf): Discussion of the RRF algorithm\n",
        "\n",
        "## Vector Search resources\n",
        "- [Overview of Vertex AI Vector Search](https://cloud.google.com/vertex-ai/docs/vector-search/overview)\n",
        "- [Vector Search quickstart](https://cloud.google.com/vertex-ai/docs/vector-search/quickstart)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d-XUnnQhrkaC"
      },
      "source": [
        "# Additional concepts\n",
        "The following sections describe TF-IDF and TfidVectorizer, Reciprical Rank Fusion, and the alpha parameter in further detail.\n",
        "\n",
        "## How do TF-IDF and `TfidfVectorizer` work?\n",
        "The `fit_transform()` function executes two important processes of the TF-IDF algorithm:\n",
        "\n",
        "- **Fit**: The vectorizer calculates the Inverse Document Frequency (IDF) for each term in the vocabulary. IDF reflects how important a term is across the entire corpus. Rare terms get higher IDF scores:\n",
        "\n",
        "```\n",
        "IDF(t) = log_e(Total number of documents / Number of documents containing term t)\n",
        "```\n",
        "\n",
        "- **Transform**:\n",
        " - **Tokenization**: Breaks the documents down into individual terms (words or phrases)\n",
        " - **Term Frequency (TF) Calculation**: Counts how often each term appears in each document with:\n",
        "\n",
        "```\n",
        "TF(t, d)** = (Number of times term t appears in document d) / (Total number of terms in document d)\n",
        "```\n",
        "\n",
        " - **TF-IDF Calculation**: Combines the TF for each term with the pre-calculated IDF to create a TF-IDF score. This score represents the importance of a term in a particular document relative to the entire corpus.\n",
        "\n",
        "```\n",
        "TF-IDF(t, d) = TF(t, d) * IDF(t)\n",
        "```\n",
        "\n",
        "The TF-IDF vectorizer tries to put higher weight to signature words in the dataset, such as \"Shirts\" or \"Dino\", compared to trivial words, such as \"The\", \"a\" or \"of\", and counts how many times those signature words are used in the specified document. Each value of a sparse embedding represents a frequency of each word based on the counts.\n",
        "\n",
        "## What is Reciprocal Rank Fusion?\n",
        "For merging the token-based and semantic search results, hybrid search uses [Reciprocal Rank Fusion (RRF)](https://plg.uwaterloo.ca/%7Egvcormac/cormacksigir09-rrf.pdf). RRF is an algorithm for combining multiple ranked lists of items into a single, unified ranking. It's a popular technique for merging search results from different sources or retrieval methods, especially in hybrid search systems and large language models.\n",
        "\n",
        "In case of the hybrid search of Vector Search, the dense distance and sparse distance are measured in different spaces and can't be directly compared to each other. Thus, RRF works effectively for merging and ranking the results from the two different spaces.\n",
        "\n",
        "Here's how RRF Works:\n",
        "\n",
        "- **Reciprocal rank**: For each item in a ranked list, calculate its reciprocal rank. This means taking the inverse of the item's position (rank) in the list. For example, the item ranked number one gets a reciprocal rank of 1/1 = 1, and the item ranked number two gets 1/2 = 0.5.\n",
        "- **Sum reciprocal ranks**: Sum the reciprocal ranks for each item across all the ranked lists. This gives a final score for each item.\n",
        "- **Sort by final score**: Sort the items by their final score in descending order. The items with the highest scores are considered the most relevant or important.\n",
        "\n",
        "In short, the items with higher ranks in both dense and sparse results will be pulled up to the top of the list. Thus, the item \"Google Blue Kids Sunglasses\" is at the top as it has higher ranks in both dense and sparse search results. Items like \"Google White Classic Youth Tee\" are low ranked as they only have ranks in the dense search result.\n",
        "\n",
        "## How the alpha parameter behaves\n",
        "The example of how to use hybrid search sets the parameter `rrf_ranking_alpha` as `0.5` when creating the `HybridQuery` object. You can specify a weight on ranking the dense and sparse search results using the following values for `rrf_ranking_alpha`:\n",
        "\n",
        "- `1`, or not specified: Hybrid search uses only dense search results and ignores sparse search results\n",
        "- `0`: Hybrid search uses only sparse search results and ignores dense search results\n",
        "- `0` to `1`: Hybrid search merges both results from dense and sparse with the weight specified by the value. `0.5` means they will be merged with the same weight\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "hybrid-search.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}