embeddings/task-type-embedding.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "wDrRkb1v_LMi",
      "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",
      "id": "pCCefZVx9pA8",
      "metadata": {
        "id": "pCCefZVx9pA8"
      },
      "source": [
        "# Using \"task type\" embeddings for improving RAG search quality"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "6yxMhQfFwuzQ",
      "metadata": {
        "id": "6yxMhQfFwuzQ"
      },
      "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/task-type-embedding.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%2Ftask-type-embedding.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/task-type-embedding.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/task-type-embedding.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/embeddings/task-type-embedding.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/embeddings/task-type-embedding.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/embeddings/task-type-embedding.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/embeddings/task-type-embedding.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/embeddings/task-type-embedding.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>            "
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7ae7abd1",
      "metadata": {
        "id": "10d04587f721"
      },
      "source": [
        "| Author |\n",
        "| --- |\n",
        "| [Kaz Sato](https://github.com/kazunori279) |"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "AD0zX6nu9-Wu",
      "metadata": {
        "id": "AD0zX6nu9-Wu"
      },
      "source": [
        "## Why my RAG can't find relevant documents?\n",
        "\n",
        "Retrieval Augmented Generation (RAG) is becoming a very popular architecture pattern of LLM grounding that retrieves and feed relevant business data to LLM and suppress hallucinations. But many RAG engineers are struggling on building the production quality retrieval engine - the search quality of RAG systems is becoming a common big challenge. In many cases, this happens when you are using \"plain vanilla\" text embeddings with vector databases for a simple similarity search. This is actually a common fault of information retrieval system design that Google has been solving for our search services for over a decade.\n",
        "\n",
        " In this tutorial we'll learn what Google has learned, and how the new \"task type\" embeddings in Vertex AI can provide a quick solution to optimize the search quality of your RAG system."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "Hvk5Y8nP-bNf",
      "metadata": {
        "id": "Hvk5Y8nP-bNf"
      },
      "source": [
        "## Questions and answers are not similar\n",
        "\n",
        "When building Retrieval Augmented Generation (RAG) systems, most designs we've seen use text embedding and vector search to conduct a [semantic similarity search](https://cloud.google.com/blog/products/ai-machine-learning/how-to-use-grounding-for-your-llms-with-text-embeddings). In many cases, this leads to degraded search quality, because of the \"question is not the answer\" problem. For example, a question like \"Why is the sky blue?\" and its answer, \"The scattering of sunlight causes the blue color\", have distinctly different meanings as separate statements. Likewise, if you ask \"What's the best birthday present for my kid?\" to a RAG system with a similarity search, it would be hard to find items like \"Nintendo Switch\" or \"Lego sets\", as those answers are not semantically similar to the question.\n",
        "\n",
        "![](https://storage.googleapis.com/github-repo/embeddings/task_type_embeddings/1.%20task_type_1.png)\n",
        "\n",
        "Technically speaking, the distribution of the question embedding space and answer embedding space have a major gap, and the simple similarity search doesn't work well. To fill the gap, you need to have an AI/ML model that learns the relationship between the query and answer.\n",
        "\n",
        "### The history of \"filling the gap\"\n",
        "\n",
        "This is actually a classic problem of semantic search for information retrieval, and Google has a long history of optimizing the semantic search quality for billions of users. Google Search started incorporating semantic search in 2015, with the introduction of [AI search innovations](https://blog.google/products/search/how-ai-powers-great-search-results/) like our deep learning ranking system RankBrain. This innovation was quickly followed with neural matching to improve the accuracy of document retrieval in Search. Neural matching contains deep learning models that learn the relationship between the query intention and relevant documents. For example, a search \"insights how to manage a green\". If a friend asked you this, you'd probably be stumped. But with neural matching, we're able to make sense of it.\n",
        "\n",
        "<center>\n",
        "  <img src=\"https://storage.googleapis.com/github-repo/embeddings/task_type_embeddings/2.%20Google%20Search.png\" width=\"300\"/><br/>\n",
        "Google Search can find relevant documents<br/>\n",
        "for an ambiguous query \"how to manage a green\"\n",
        "</center>\n",
        "<br/>\n",
        "\n",
        "### Dual encoder model\n",
        "For developers that want to incorporate sophisticated semantic search like neural matching into their RAG system the popular approach has been to [train a dual encoder model](https://cloud.google.com/blog/products/ai-machine-learning/scaling-deep-retrieval-tensorflow-two-towers-architecture) (aka two-tower) that learns the relationship between the query embeddings and answer embeddings. The following diagram depicts how a dual encoder model maps queries with relevant documents.\n",
        "\n",
        "<br/>\n",
        "<div align=\"center\">\n",
        "  <img src=\"https://storage.googleapis.com/github-repo/embeddings/task_type_embeddings/3.%20dual%20encoder.gif\">\n",
        "</div>\n",
        "\n",
        "But it is not easy to design, train and build a production retrieval system with a customized dual encoder model. It requires data science and ML engineering experience, and significant effort to train a dataset with pairs of questions and answers.\n",
        "\n",
        "### LLM-based \"Advanced RAG\"\n",
        "In the era of Large Language Models (LLM), some \"Advanced RAG\" approaches emerged to solve this problem. [HyDE](https://arxiv.org/abs/2212.10496) uses LLM reasoning capabilities to generate possible answer texts, and use the text for similarity search. Additionally, [query expansion with LLM](https://arxiv.org/abs/2305.03653) uses the LLM reasoning for expanding the original query with possible answer candidates as query texts.\n",
        "\n",
        "The downside of these LLM-based approaches is that it adds the LLM prediction latency and high cost for every query. While the vector search itself can finish within milliseconds, and can handle thousands of queries per second at low cost, the LLM reasoning will add a few seconds latency for every single query with significantly higher cost than generating embeddings.\n",
        "\n",
        "## New \"task type\" embedding\n",
        "\n",
        "Recently [Vertex AI Embeddings API](https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings) launched new text embedding models, `text-embedding-005` and `text-multilingual-embedding-002` based on [the new embedding models](https://deepmind.google/research/publications/85521/) developed by the Google DeepMind and Research team.\n",
        "\n",
        "The unique feature of these models is that they can generate optimized embeddings based on [task types](https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/task-types). With task types, you can significantly reduce the time and cost for developing dual encoder models or advanced RAG systems for achieving higher search quality, for specific tasks including question and answering, document retrieval by query, fact verification, etc.\n",
        "For example, to generate embeddings for question and answering, all you have to do is to specify a task type `QUESTION_ANSWERING` for query texts and `RETRIEVAL_DOCUMENT` for answer texts when generating their embeddings. In the embedding space optimized with the task types, the query embedding \"Why is the sky blue?\" and the answer embedding \"The scattering...\" will be placed in a much closer area, because the embedding models are trained to learn that they have the question and answer relationship. Thus you should be able to get higher search quality with vector databases to find the right answer for the query.\n",
        "\n",
        "![](https://storage.googleapis.com/github-repo/embeddings/task_type_embeddings/4.%20task_type_2.png)\n",
        "\n",
        "### Behind the scene: LLM distillation + dual encoder\n",
        "\n",
        "How do the new embedding models enable this? The novel part of the model design is that the Google DeepMind and Research team used [LLM distillation](https://arxiv.org/abs/1503.02531) (a process to train a smaller model from a larger model) to pre-train a dual encoder based embedding model.\n",
        "Here's the distillation steps. The team used an LLM for generating two kinds of data: 1) synthetic queries and 2) query-and-document pairs. First, the team generated numerous synthetic queries from the training dataset (corpus) for each task.\n",
        "\n",
        "![](https://storage.googleapis.com/github-repo/embeddings/task_type_embeddings/5.%20FRet_1.png)\n",
        "\n",
        "Then, with the generated queries, find possibly relevant documents to create query-and-document pairs. For those pairs, use the LLM to determine if the pair is considered as an appropriate query-and-document (positive pair), or not (negative pair).\n",
        "\n",
        "![](https://storage.googleapis.com/github-repo/embeddings/task_type_embeddings/6.%20FRet_2.png)\n",
        "\n",
        "With those pairs, the team trained a dual-encoder based embedding model. The result is an embedding model that inherits the reasoning capability of the LLM as a distilled model for the specific tasks. Significantly smaller, faster and inexpensive to generate embeddings, and yet, provides the capability of smart matching between queries and relevant documents.\n",
        "\n",
        "If you are interested in more details of the model design and other performance comparison results, check out [the paper](https://deepmind.google/research/publications/85521/) published by the team.\n",
        "\n",
        "### Supported task types\n",
        "\n",
        "The new embeddings models support the following task types:\n",
        "\n",
        "|Task type|Embeddings optimization criteria|\n",
        "|:---|:---|\n",
        "|`SEMANTIC_SIMILARITY`|Semantic similarity. Use this task type when retrieving similar texts from the corpus.|\n",
        "|`RETRIEVAL_QUERY`|Document search and information retrieval. Use `RETRIEVAL_QUERY` for query texts, and `RETRIEVAL_DOCUMENT` for documents to be retrieved.|\n",
        "|`QUESTION_ANSWERING`|Questions and answers applications such as RAG. Use `QUESTION_ANSWERING` for question texts, and `RETRIEVAL_DOCUMENT` for documents to be retrieved.|\n",
        "|`FACT_VERIFICATION`|Document search for fact verification. Use `FACT_VERIFICATION` for the target text, and `RETRIEVAL_DOCUMENT` for documents to be retrieved.|\n",
        "|`CODE_RETRIEVAL_QUERY`|Code search. Use `CODE_RETRIEVAL_QUERY` for query text, and `RETRIEVAL_DOCUMENT` for code blocks to be retrieved (available on embedding model `text-embedding-preview-0815` and later)|\n",
        "|`CLASSIFICATION`|Text classification. Use this task type for training a small classification model with the embedding.|\n",
        "|`CLUSTERING`|Text clustering. Use this task type for k-means or other clustering analysis.|\n",
        "\n",
        "For example, if you are building a RAG system for a question and answering use case, you may specify task type `RETRIEVAL_DOCUMENT` for generating embeddings for building with vector search, and specify `QUESTION_ANSWERING` for embeddings for question texts. Thus you should see improved search quality compared to using `SEMANTIC_SIMILARITY` for both query and document. Likewise, you may use `RETRIEVAL_QUERY` for queries for document search, and `FACT_VERIFICATION` for queries for finding documents for fact checking.\n",
        "\n",
        "Not only the document retrieval and Q&A use cases, but also Embedding models support optimized embeddings for text classification and clustering. Embeddings with task type `CLASSIFICATION` are useful for classifying texts with its semantics for use cases such as customer and product segmentation.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0xPHq0pJFdCq",
      "metadata": {
        "id": "0xPHq0pJFdCq"
      },
      "source": [
        "# Task type embeddings in Action\n",
        "\n",
        "In the following sections, we will test how task type improves search quality in Question and Answering case.\n",
        "\n",
        "## Setup\n",
        "\n",
        "Let's start setting up the SDK and environment variables.\n",
        "\n",
        "### Install Python SDK\n",
        "\n",
        "This tutorial uses Google Gen AI SDK"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "vW_vTYXEKgjt",
      "metadata": {
        "id": "vW_vTYXEKgjt"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --quiet google-genai"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "nA4IiIlSK5uO",
      "metadata": {
        "id": "nA4IiIlSK5uO"
      },
      "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,
      "id": "V2cUT0rSK8N-",
      "metadata": {
        "id": "V2cUT0rSK8N-"
      },
      "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",
      "id": "ccvdW2mgLWLE",
      "metadata": {
        "id": "ccvdW2mgLWLE"
      },
      "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,
      "id": "HPfsNnW1e5mm",
      "metadata": {
        "id": "HPfsNnW1e5mm"
      },
      "outputs": [],
      "source": [
        "# Use the environment variable if the user doesn't provide Project ID.\n",
        "import os\n",
        "\n",
        "from google import genai\n",
        "\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "\n",
        "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
        "\n",
        "client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "06489bd14f16"
      },
      "source": [
        "### Import libraries\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "07945978db0e"
      },
      "outputs": [],
      "source": [
        "from typing import Any\n",
        "\n",
        "from google.genai.types import EmbedContentConfig\n",
        "import pandas as pd\n",
        "from sklearn.metrics.pairwise import cosine_similarity\n",
        "from tqdm import tqdm  # to show the progress bar"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "KKAxo40cjnMW",
      "metadata": {
        "id": "KKAxo40cjnMW"
      },
      "source": [
        "## Try the task types\n",
        "\n",
        "Now it's ready to try the task type. Let's get started with some simple examples."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "yEeIoPNKzLad",
      "metadata": {
        "id": "yEeIoPNKzLad"
      },
      "source": [
        "### Generate embeddings with task types\n",
        "\n",
        "Here we define a wrapper function `get_embeddings()` to make it easier to call the Embeddings API.\n",
        "\n",
        "To specify task type on generating embeddings, you can use the `EmbedContentConfig` class."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "mNTyxY0dJsfc",
      "metadata": {
        "id": "mNTyxY0dJsfc"
      },
      "outputs": [],
      "source": [
        "MODEL_ID = \"text-embedding-005\"\n",
        "\n",
        "\n",
        "def get_embeddings(contents: list[str], task_type: str) -> list[list[float]]:\n",
        "    \"\"\"A wrapper to get text embeddings for a list of texts\"\"\"\n",
        "    response = client.models.embed_content(\n",
        "        model=MODEL_ID,\n",
        "        contents=contents,\n",
        "        config=EmbedContentConfig(task_type=task_type),\n",
        "    )\n",
        "    return [e.values for e in response.embeddings]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "Uh7TBAoqqxy2",
      "metadata": {
        "id": "Uh7TBAoqqxy2"
      },
      "source": [
        "### Why is the sky blue?\n",
        "\n",
        "Using this function, generate embeddings with task type `SEMANTIC_SIMILARITY` for a question and two possible answers. Then, print similarities between them using [`cosine_similarity()` function](https://scikit-learn.org/stable/modules/generated/sklearn.metrics.pairwise.cosine_similarity.html) from scikit-learn."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "-HoF2kMPkeEj",
      "metadata": {
        "id": "-HoF2kMPkeEj"
      },
      "outputs": [],
      "source": [
        "QUESTION = \"Why is the sky blue?\"\n",
        "ANSWER_1 = \"The sky is blue today\"\n",
        "ANSWER_2 = \"The scattering of sunlight causes the blue color\"\n",
        "\n",
        "q_emb = get_embeddings([QUESTION], \"SEMANTIC_SIMILARITY\")\n",
        "a1_emb = get_embeddings([ANSWER_1], \"SEMANTIC_SIMILARITY\")\n",
        "a2_emb = get_embeddings([ANSWER_2], \"SEMANTIC_SIMILARITY\")\n",
        "\n",
        "print(cosine_similarity(q_emb, a1_emb))\n",
        "print(cosine_similarity(q_emb, a2_emb))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "EtGLBZb7rIVw",
      "metadata": {
        "id": "EtGLBZb7rIVw"
      },
      "source": [
        "Since the embedding space for the task type is organized to measure the semantic similarity, naturally the first answer (which doesn't answer to the question at all) has higher similarity than the second answer (which is the correct answer).\n",
        "\n",
        "Now, let's try task type `QUESTION_ANSWERING`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "mpIea_FBlg2J",
      "metadata": {
        "id": "mpIea_FBlg2J"
      },
      "outputs": [],
      "source": [
        "q_emb = get_embeddings([QUESTION], \"QUESTION_ANSWERING\")\n",
        "a1_emb = get_embeddings([ANSWER_1], \"RETRIEVAL_DOCUMENT\")\n",
        "a2_emb = get_embeddings([ANSWER_2], \"RETRIEVAL_DOCUMENT\")\n",
        "\n",
        "print(cosine_similarity(q_emb, a1_emb))\n",
        "print(cosine_similarity(q_emb, a2_emb))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "QdzJGORirqLp",
      "metadata": {
        "id": "QdzJGORirqLp"
      },
      "source": [
        "This time, the second answer gets much higher similarity than the first one, because the embedding space is organized to capture the question-and-answering relationships.\n",
        "\n",
        "### Taylor Swift or Lego set?\n",
        "\n",
        "Try another example:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "cFywbkBOoVEj",
      "metadata": {
        "id": "cFywbkBOoVEj"
      },
      "outputs": [],
      "source": [
        "QUESTION = \"What is the best birthday present for my son?\"\n",
        "ANSWER_1 = \"My son says Taylor Swift's birthday is December 13th\"\n",
        "ANSWER_2 = \"Lego set\"\n",
        "\n",
        "q_emb = get_embeddings([QUESTION], \"SEMANTIC_SIMILARITY\")\n",
        "a1_emb = get_embeddings([ANSWER_1], \"SEMANTIC_SIMILARITY\")\n",
        "a2_emb = get_embeddings([ANSWER_2], \"SEMANTIC_SIMILARITY\")\n",
        "\n",
        "print(cosine_similarity(q_emb, a1_emb))\n",
        "print(cosine_similarity(q_emb, a2_emb))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3RaM5ZlCsIqP",
      "metadata": {
        "id": "3RaM5ZlCsIqP"
      },
      "source": [
        "The model does its job properly - choosing the first answer as similar sentence to the question. But it's not what we expect."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "PhelfltnsFbe",
      "metadata": {
        "id": "PhelfltnsFbe"
      },
      "outputs": [],
      "source": [
        "q_emb = get_embeddings([QUESTION], \"QUESTION_ANSWERING\")\n",
        "a1_emb = get_embeddings([ANSWER_1], \"RETRIEVAL_DOCUMENT\")\n",
        "a2_emb = get_embeddings([ANSWER_2], \"RETRIEVAL_DOCUMENT\")\n",
        "\n",
        "print(cosine_similarity(q_emb, a1_emb))\n",
        "print(cosine_similarity(q_emb, a2_emb))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9r_XNuwvsdus",
      "metadata": {
        "id": "9r_XNuwvsdus"
      },
      "source": [
        "With the task type `QUESTION_ANSWERING` specified, the model now understand it should find an appropriate answer to the question, not a similar sentence. It chooses Lego set over Taylor Swift's birthday."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "xKW3QCw0voGD",
      "metadata": {
        "id": "xKW3QCw0voGD"
      },
      "source": [
        "# Evaluate the search quality of task types\n",
        "\n",
        "The results above were interesting. Now, let's evaluate the search quality with the task types with a larger dataset."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "yuxVrUUdKBri",
      "metadata": {
        "id": "yuxVrUUdKBri"
      },
      "source": [
        "## Download NQ-Open dataset\n",
        "\n",
        "As a test dataset, we will use 1K rows sampled from [NQ-Open](https://huggingface.co/datasets/google-research-datasets/nq_open) dataset provided by the Google Research team. NQ-Open dataset can be easily downloaded from the Hugging Face Datasets:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "I6_uVsQxwM4M",
      "metadata": {
        "id": "I6_uVsQxwM4M"
      },
      "outputs": [],
      "source": [
        "# suppress the HF_TOKEN warning when downloading the dataset\n",
        "os.environ[\"HF_HUB_DISABLE_TELEMETRY\"] = \"1\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "SXOrnZKoJTkV",
      "metadata": {
        "id": "SXOrnZKoJTkV"
      },
      "outputs": [],
      "source": [
        "# download NQ-Open dataset\n",
        "splits = {\n",
        "    \"train\": \"nq_open/train-00000-of-00001.parquet\",\n",
        "    \"validation\": \"nq_open/validation-00000-of-00001.parquet\",\n",
        "}\n",
        "nq_open_df = pd.read_parquet(\n",
        "    \"hf://datasets/google-research-datasets/nq_open/\" + splits[\"train\"]\n",
        ")\n",
        "\n",
        "# sampling for 1K rows\n",
        "nq_open_df = nq_open_df.sample(n=1000).reset_index(drop=True)\n",
        "nq_open_df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "QYOuS6dbI3Kz",
      "metadata": {
        "id": "QYOuS6dbI3Kz"
      },
      "source": [
        "The sampled dataframe contains 1K pairs of question and answer for a wide variety topics, such as a question \"what is the solid outer layer of the earth called\" and its answer \"The crust\".\n",
        "\n",
        "In this tutorial, we will use 1K rows to measure the search quality. By increasing the rows to 10K you will see better results, but it would take about 30 minutes to finish the tutorial."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "JfArfckVWyVX",
      "metadata": {
        "id": "JfArfckVWyVX"
      },
      "source": [
        "## Calculating Q&A search quality with `SEMANTIC_SIMILARITY`"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "LkxkofhhMUJA",
      "metadata": {
        "id": "LkxkofhhMUJA"
      },
      "source": [
        "First, let's test the task type `SEMANTIC_SIMILARITY` as a baseline.\n",
        "\n",
        "Using the `get_embeddings()` function above, we define another function `generate_embeddings()` that takes a list of texts and generate embeddings with the specified task types. The embeddings will be stored on new columns `q_emb` and `a_emb`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "mMw5xmbVw5g_",
      "metadata": {
        "id": "mMw5xmbVw5g_"
      },
      "outputs": [],
      "source": [
        "def generate_embeddings(\n",
        "    df: pd.DataFrame, question_task_type: str, answer_task_type: str\n",
        ") -> pd.DataFrame:\n",
        "    \"\"\"Generate text embeddings for questions and answers\"\"\"\n",
        "\n",
        "    q_embeddings: list[list[float]] = []\n",
        "    a_embeddings: list[list[float]] = []\n",
        "\n",
        "    for i in tqdm(range(0, len(df), 5)):  # bundle 5 items for each API call\n",
        "        q_embeddings.extend(\n",
        "            get_embeddings(df.loc[i : i + 4, \"question\"].tolist(), question_task_type)\n",
        "        )\n",
        "\n",
        "        answers: list[Any] = [row[0] for row in df.loc[i : i + 4, \"answer\"]]\n",
        "        a_embeddings.extend(get_embeddings(answers, answer_task_type))\n",
        "\n",
        "    df[\"q_emb\"] = q_embeddings\n",
        "    df[\"a_emb\"] = a_embeddings\n",
        "    return df"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "Rhp_bdgyM1Yv",
      "metadata": {
        "id": "Rhp_bdgyM1Yv"
      },
      "source": [
        "It's ready to start generating the embeddings for semantic similarity task. The following code may take a few minutes to finish generating embeddings for the 1K rows."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "CvhrGyW4a8xF",
      "metadata": {
        "id": "CvhrGyW4a8xF"
      },
      "outputs": [],
      "source": [
        "# generate embeddings using SEMANTIC_SIMILARITY for both question and answer\n",
        "similarity_df = nq_open_df.copy()\n",
        "generate_embeddings(similarity_df, \"SEMANTIC_SIMILARITY\", \"SEMANTIC_SIMILARITY\")\n",
        "similarity_df"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "qFzeZmpgzVDJ",
      "metadata": {
        "id": "qFzeZmpgzVDJ"
      },
      "source": [
        "### Calculate cosine similarities between questions and answers\n",
        "\n",
        "Next, we will calculate the cosine similarity for all combinations of the questions and answers using the `cosine_similarity()` function from scikit-learn. Thus, the `calc_cosine_similarities()` function will return a two dimensional array with similarities between 1K questions x 1K answers. By printing the result for the first question, you will see 1K values that represents the cosine similarity between the first question and the 1K answers."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5sin1iKsbGsm",
      "metadata": {
        "id": "5sin1iKsbGsm"
      },
      "outputs": [],
      "source": [
        "def calc_cosine_similarities(df: pd.DataFrame) -> list[list[float]]:\n",
        "    \"\"\"Calculates similarities for all combinations of the questions and answers\"\"\"\n",
        "    q_embeddings: list[list[float]] = df[\"q_emb\"].tolist()\n",
        "    a_embeddings: list[list[float]] = df[\"a_emb\"].tolist()\n",
        "    return cosine_similarity(q_embeddings, a_embeddings)\n",
        "\n",
        "\n",
        "# print the answer similarities for the first question\n",
        "q_and_a_similarities = calc_cosine_similarities(similarity_df)\n",
        "q_and_a_similarities[:5]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "YC-cG7Mi3JZj",
      "metadata": {
        "id": "YC-cG7Mi3JZj"
      },
      "source": [
        "### Calculate MRR\n",
        "\n",
        "Finally, we will measure the search quality using [Mean Reciprocal Rank (MRR)](https://en.wikipedia.org/wiki/Mean_reciprocal_rank). It is a simple algorithm - the table from the Wikipedia page gives an example:\n",
        "\n",
        "![](https://storage.googleapis.com/github-repo/embeddings/task_type_embeddings/mrr_example.png)\n",
        "\n",
        "In this tutorial, we will use the `answer` column of the dataset as the correct response for the `question` column. Thus, if the similarity between the question and answer is ranked as 3rd in the similarity list, reciprocal rank is 1/3.\n",
        "\n",
        "To do it, we define `get_top100_similar_answers()` function that returns a top 100 list of answer index sorted by the similarity."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "mkv3rZnrlrVE",
      "metadata": {
        "id": "mkv3rZnrlrVE"
      },
      "outputs": [],
      "source": [
        "def get_top100_similar_answers(answer_similarities: list[float]) -> list[int]:\n",
        "    \"\"\"Extracts top 100 similar answers\"\"\"\n",
        "    sim_index: list[tuple[int, float]] = list(enumerate(answer_similarities))\n",
        "    sorted_sim_index: list[tuple[int, float]] = sorted(\n",
        "        sim_index, key=lambda x: x[1], reverse=True\n",
        "    )\n",
        "    sorted_index: list[int] = [index for index, _ in sorted_sim_index]\n",
        "    return sorted_index[:100]\n",
        "\n",
        "\n",
        "# print top 100 similar answers with their index\n",
        "get_top100_similar_answers(q_and_a_similarities[0])[:10]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "hkr16f5HRdgu",
      "metadata": {
        "id": "hkr16f5HRdgu"
      },
      "source": [
        "With this top 100 answer list for each question, we will add `query_ranks` column to the DataFrame that has a list like `[0 0 1 0 0...]`. In this case, the correct response is ranked at 3rd."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "oYMNvF2Oox4h",
      "metadata": {
        "id": "oYMNvF2Oox4h"
      },
      "outputs": [],
      "source": [
        "# mypy: disable-error-code=\"index\"\n",
        "\n",
        "\n",
        "def add_query_ranks(df: pd.DataFrame) -> pd.DataFrame:\n",
        "    \"\"\"Adds a 'query_ranks' column to the DataFrame\"\"\"\n",
        "\n",
        "    # Calculate similarities for all combinations of questions and answers\n",
        "    q_and_a_similarities: list[list[float]] = calc_cosine_similarities(df)\n",
        "\n",
        "    # Get query ranks for all questions\n",
        "    query_ranks_for_all_questions: list[list[int]] = []\n",
        "    for i in tqdm(range(len(q_and_a_similarities))):\n",
        "        # For each question, get ground truth and top 100 similar answers\n",
        "        ground_truth: Any = df.loc[i, \"answer\"][\n",
        "            0\n",
        "        ]  # Assuming 'answer' column contains lists\n",
        "        similar_answers: list[int] = get_top100_similar_answers(q_and_a_similarities[i])\n",
        "\n",
        "        # If an answer equals the ground truth, put \"1\" in the query ranks, otherwise \"0\"\n",
        "        query_ranks: list[int] = []\n",
        "        for index in similar_answers:\n",
        "            similar_answer: Any = df.loc[index, \"answer\"][0]\n",
        "            if similar_answer == ground_truth:\n",
        "                query_ranks.append(1)\n",
        "            else:\n",
        "                query_ranks.append(0)\n",
        "        query_ranks_for_all_questions.append(query_ranks)\n",
        "\n",
        "    # Add the query ranks to the dataframe\n",
        "    df[\"query_ranks\"] = query_ranks_for_all_questions\n",
        "    return df"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "x3sKIQsY2y7d",
      "metadata": {
        "id": "x3sKIQsY2y7d"
      },
      "outputs": [],
      "source": [
        "# add query_ranks column to the similarity_df\n",
        "add_query_ranks(similarity_df)\n",
        "similarity_df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "RsrUcuuKSKUI",
      "metadata": {
        "id": "RsrUcuuKSKUI"
      },
      "source": [
        "With this `query_ranks`, we will calculate MRR value for the dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "xCNi-TlArDw_",
      "metadata": {
        "id": "xCNi-TlArDw_"
      },
      "outputs": [],
      "source": [
        "def calc_mrr(df: pd.DataFrame) -> float:\n",
        "    \"\"\"Calculates MRR from the query ranks\"\"\"\n",
        "\n",
        "    query_ranks_for_all_questions: list[list[int]] = df[\"query_ranks\"].tolist()\n",
        "    reciprocal_ranks: list[float] = []\n",
        "\n",
        "    for query_ranks in query_ranks_for_all_questions:\n",
        "        for i, relevance in enumerate(query_ranks):\n",
        "            if relevance == 1:\n",
        "                reciprocal_ranks.append(1 / (i + 1))\n",
        "                break\n",
        "        else:  # This else is associated with the for loop, executed if no break occurred\n",
        "            reciprocal_ranks.append(0)\n",
        "\n",
        "    return sum(reciprocal_ranks) / len(query_ranks_for_all_questions)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8EKIhsd83Yro",
      "metadata": {
        "id": "8EKIhsd83Yro"
      },
      "outputs": [],
      "source": [
        "mrr = calc_mrr(similarity_df)\n",
        "mrr"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "zbjxtPeaSZ7F",
      "metadata": {
        "id": "zbjxtPeaSZ7F"
      },
      "source": [
        "You should see a value around 0.3 with this task type."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "DjP5117MufPt",
      "metadata": {
        "id": "DjP5117MufPt"
      },
      "source": [
        "## Calculating Q&A search quality with QUESTION_ANSWERING"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "PcN0-5IjSh7t",
      "metadata": {
        "id": "PcN0-5IjSh7t"
      },
      "source": [
        "Next, let's test with the task type `QUESTION_ANSWERING`.\n",
        "\n",
        "Starting with generating embeddings for the same dataset. This time, we will use task type `QUESTION_ANSWERING` for the questions and `RETRIEVAL_DOCUMENT` for the answers. This will take a few minutes to finish."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9k3sNqhG348s",
      "metadata": {
        "id": "9k3sNqhG348s"
      },
      "outputs": [],
      "source": [
        "# generate embeddings using QUESTION_ANSWERING for question and\n",
        "# RETRIEVAL_DOCUMENT for answer\n",
        "q_and_a_df = nq_open_df.copy()\n",
        "generate_embeddings(q_and_a_df, \"QUESTION_ANSWERING\", \"RETRIEVAL_DOCUMENT\")\n",
        "q_and_a_df"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "pKtOakG2Te-O",
      "metadata": {
        "id": "pKtOakG2Te-O"
      },
      "source": [
        "Then, generate `query_ranks` columns for calculating MRR."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5TFlhRBv4tgq",
      "metadata": {
        "id": "5TFlhRBv4tgq"
      },
      "outputs": [],
      "source": [
        "add_query_ranks(q_and_a_df)\n",
        "q_and_a_df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "Olqs9JdnToGt",
      "metadata": {
        "id": "Olqs9JdnToGt"
      },
      "source": [
        "Finally, calculate MRR for the task type `QUESTION_ANSWERING`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "YJA0HE68vYWU",
      "metadata": {
        "id": "YJA0HE68vYWU"
      },
      "outputs": [],
      "source": [
        "mrr = calc_mrr(q_and_a_df)\n",
        "mrr"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "_WlgppIHTt0Q",
      "metadata": {
        "id": "_WlgppIHTt0Q"
      },
      "source": [
        "You should see the MRR value around 0.4, it's about 30% increase from the task type `SEMANTIC_SIMILARITY`."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "4QEnCHzn7GMc",
      "metadata": {
        "id": "4QEnCHzn7GMc"
      },
      "source": [
        "# Cleaning up\n",
        "\n",
        "This tutorial doesn't use any persistent resources of Google Cloud. Thus no need to deleting resources for cleaning up.\n",
        "\n",
        "# Conclusion\n",
        "\n",
        "As we saw, Google has been tackling this issue to optimize the semantic search quality of its largest services for many years, and the knowledge and expertise is packaged as easy-to-use \"task type\" embedding models.\n",
        "\n",
        "In many cases, you should see around 30% - 40% increase on the MRR value by changing from `SEMANTIC_SIMILARITY` to `QUESTION_ANSWERING`. Using an appropriate task type for different requirements may improve the user experience of your gen AI applications.\n",
        "\n",
        "### When to use embedding tuning\n",
        "One caveat of task type embedding is that, the model is trained with generic corpus on the web. That means, the model may be less effective on use cases such as retrieving company proprietary document or question and answering on topics that the model doesn't know. In such cases, [Tune text embeddings](https://cloud.google.com/vertex-ai/generative-ai/docs/models/tune-embeddings) could provide better option than using the pre-trained model.\n",
        "\n",
        "### Getting Started\n",
        "To get started with task type embeddings, here's a list of resources to get started with the new embedding models:\n",
        "\n",
        "- [Choose an embeddings task type](https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/task-types)\n",
        "- [Get text embeddings](https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/get-text-embeddings)\n",
        "- [Overview of Vertex AI Vector Search](https://cloud.google.com/vertex-ai/docs/vector-search/overview)\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "task-type-embedding.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}