notebooks/official/generative_ai/tuned_text-embeddings.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-Gohq0IUAtsr"
      },
      "outputs": [],
      "source": [
        "# @title 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": "XQ7omKMxBBgp"
      },
      "source": [
        "# Getting Tuned Text-Embeddings on Vertex AI\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/generative_ai/tuned_text-embeddings.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\">Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/generative_ai/tuned_text-embeddings.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/generative_ai/tuned_text-embeddings.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-3YiUR9SCSSm"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook will walk you through the process of fine-tuning text-embeddings. By adapting a text-embedding model to your specific domain or task, you can achieve better results. See also [tuning text-embeddings](https://cloud.google.com/vertex-ai/generative-ai/docs/models/tune-embeddings)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Sy_UWYAfH4uD"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn how to tune a text-embedding model.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
        "- Vertex AI\n",
        "- Cloud Storage\n",
        "\n",
        "The steps include:\n",
        "\n",
        "- Run a text-embedding model tuning job on Vertex AI Pipelines.\n",
        "- Deploy your tuned text-embedding model to an endpoint.\n",
        "- Examine the quality metrics to assess your tuned model.\n",
        "- Get embeddings from your tuned model for downstream tasks."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fLW7At4TAV5n"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "This notebook uses a synthetic dataset (corpus, queries, and labels) from [Alphabet's annual financial performance report (10K form)](https://abc.xyz/assets/ff/7c/06d6f493f6462caf08e8502ffc33/596de1b094c32cf0592a08edfe84ae74.html).\n",
        "\n",
        "* [corpus.jsonl](https://storage.googleapis.com/cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/corpus.jsonl) (397Ki),\n",
        "* [queries.jsonl](https://storage.googleapis.com/cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/queries.jsonl) (321Ki)\n",
        "* [test.tsv](https://storage.googleapis.com/cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/test.tsv) (3.7Ki)\n",
        "* [train.tsv](https://storage.googleapis.com/cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/train.tsv) (29Ki)\n",
        "* [validation.tsv](https://storage.googleapis.com/cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/validation.tsv) (3.7Ki)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0We7bZz9EjvS"
      },
      "source": [
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing),\n",
        "and [Cloud Storage pricing](https://cloud.google.com/storage/pricing),\n",
        "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qDCccnyeJP9v"
      },
      "source": [
        "## Installation\n",
        "\n",
        "This tutorial requires you to install the latest `google-cloud-aiplatform` package release."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nn6Ze-c4IZyt"
      },
      "outputs": [],
      "source": [
        "!pip3 install --quiet --upgrade google-cloud-aiplatform"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ngQYqEkV4aBS"
      },
      "source": [
        "### Restart runtime (Colab only)\n",
        "\n",
        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VXFRhykV3jyI"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H4LtPlgQKf19"
      },
      "source": [
        "## Before you begin"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QQOMuSWFhwC2"
      },
      "source": [
        "### Set up your Google Cloud project\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
        "\n",
        "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "3. [Enable APIs](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,documentai.googleapis.com).\n",
        "\n",
        "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
        "\n",
        "See also [setting up permissions & resources for tuning text-embedding models](https://cloud.google.com/vertex-ai/generative-ai/docs/models/tune-embeddings#project-setup)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0SG0ZnxmM1Ev"
      },
      "source": [
        "### Initialize Vertex AI SDK\n",
        "\n",
        "Import and Initialize the SDK for your project and region.\n",
        "\n",
        "**If you don't know your project ID**, try the following:\n",
        "* Run `gcloud config list`.\n",
        "* Run `gcloud projects list`.\n",
        "* See also [Locate the project ID](https://support.google.com/googleapi/answer/7014113).\n",
        "\n",
        "See also [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SWbf4r0tLUvA"
      },
      "outputs": [],
      "source": [
        "# @title (Required) Set PROJECT_ID and REGION\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "REGION = \"us-central1\"  # @param {type:\"string\"}\n",
        "if not PROJECT_ID.strip():\n",
        "    raise ValueError(\"'PROJECT_ID' is required.\")\n",
        "if not REGION.strip():\n",
        "    raise ValueError(\"'REGION' is required.\")\n",
        "!gcloud config set project {PROJECT_ID}\n",
        "\n",
        "import vertexai\n",
        "from google.cloud.aiplatform import pipeline_jobs\n",
        "from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=REGION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OpT7z9MlLPxR"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XGiF1B6kMcSz"
      },
      "source": [
        "**1. Colab:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z7RZmb3VMfCP"
      },
      "outputs": [],
      "source": [
        "# @title (Required on Colab) `authenticate_user()`\n",
        "import builtins\n",
        "import os\n",
        "import sys\n",
        "\n",
        "running_in_colab = \"google.colab\" in sys.modules and hasattr(builtins, \"get_ipython\")\n",
        "if running_in_colab and not os.getenv(\"IS_TESTING\"):\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user(project_id=PROJECT_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "crt3rd6OMRLE"
      },
      "source": [
        "**2. Vertex AI Workbench**\n",
        "\n",
        "Make sure that [the Compute Engine default service account](https://cloud.google.com/compute/docs/access/service-accounts#default_service_account) running a Vertex AI workbench instance has the permission iam.serviceAccounts.actAs (most likely through [roles/iam.serviceAccountUser](https://cloud.google.com/iam/docs/understanding-roles#iam.serviceAccountUser)) at [the IAM & Admin page of the Cloud Console](https://console.cloud.google.com/iam-admin). This permission allows a workbench instance to act as the service account when interacting with other Google Cloud services."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2ymngmfhMmSW"
      },
      "source": [
        "**3. Local JupyterLab instance, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UDVsiwBbMafG"
      },
      "outputs": [],
      "source": [
        "# !gcloud auth login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N8fqs18bNNmd"
      },
      "source": [
        "## Tune Text-Embeddings\n",
        "\n",
        "(Optionally), to resume this tutorial from where you left off with your previous tuning session, set **`TUNING_JOB_ID`** accordingly. Or, clear **`TUNING_JOB_ID`** to start over a fresh tuning session.\n",
        "\n",
        "This tutorial creates a tuning job of a Vertex AI pipeline for tuning a text-embedding model within your project. See also [creating tuning jobs with parameters and defaults](https://cloud.google.com/vertex-ai/generative-ai/docs/models/tune-embeddings#create-embedding-tuning-job), and the latest [text-embedding models and eligible tasks](https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/get-text-embeddings#api_changes_to_models_released_on_or_after_august_2023)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "3VO4AFDbhnnx"
      },
      "outputs": [],
      "source": [
        "# @title (Optional) Set `TUNING_JOB_ID` to resume an existing tuning session: { run: \"auto\" }\n",
        "TUNING_JOB_ID = \"\"  # @param {type: \"string\", placeholder:\"(Optional) Enter the name of an existing Vertex Pipelines job.\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G1i__mAHNY9U"
      },
      "outputs": [],
      "source": [
        "# @title (Required) Resume an existing or start a fresh tuning session (depending on TUNING_JOB_ID)\n",
        "# @markdown ### Pipeline Parameters\n",
        "# @markdown `queries_path`: The URI of the Cloud Storage bucket storing query data, starting with `gs://`.\n",
        "queries_path = \"gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/queries.jsonl\"  # @param {type: \"string\"}\n",
        "# @markdown ---\n",
        "# @markdown `corpus_path`: The Cloud Storage URI for the corpus data, starting with `gs://`.\n",
        "corpus_path = \"gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/corpus.jsonl\"  # @param {type: \"string\"}\n",
        "# @markdown ---\n",
        "# @markdown `train_label_path`: The Cloud Storage URI of the train label data location, starting with `gs://`.\n",
        "train_label_path = \"gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/train.tsv\"  # @param{type: \"string\"}\n",
        "# @markdown ---\n",
        "# @markdown `validation_label_path`: Optional. The Cloud Storage URI of the validation label data location, starting with `gs://`. Passing an empty string will direct the pipeline to autosplit the validation dataset from the training dataset.\n",
        "validation_label_path = \"gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/validation.tsv\"  # @param{type: \"string\"}\n",
        "# @markdown ---\n",
        "# @markdown `test_label_path`: Optional. The Cloud Storage URI of the test label data location, starting with `gs://`. Passing an empty string will direct the pipeline to autosplit the test dataset from the training dataset.\n",
        "test_label_path = \"gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/test.tsv\"  # @param{type: \"string\"}\n",
        "# @markdown ---\n",
        "# @markdown `train_steps`: Optional. The number of steps to perform model tuning. Must be greater than 30.\n",
        "train_steps = 1000  # @param {type: \"number\"}\n",
        "# @markdown ---\n",
        "# @markdown `batch_size`: Optional. The training batch size.\n",
        "batch_size = 128  # @param {type: \"number\"}\n",
        "# @markdown ---\n",
        "# @markdown `base_model_version_id`: Optional. Use this to specify what text embedding model to tune. This must be a stable version, for example `text-embedding-004`. To learn more, see [Supported Models](https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/get-text-embeddings#supported-models).\n",
        "base_model_version_id = \"text-embedding-005\"  # @param [\"text-embedding-005\", \"text-embedding-004\", \"text-multilingual-embedding-002\"]\n",
        "# @markdown ---\n",
        "# @markdown `task_type`: Optional. The task type expected to be used during inference. See the [embedding API reference](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings-api#request_body) to learn more.\n",
        "task_type = \"DEFAULT\"  # @param [\"DEFAULT\", \"RETRIEVAL_QUERY\", \"RETRIEVAL_DOCUMENT\", \"SEMANTIC_SIMILARITY\", \"CLASSIFICATION\", \"CLUSTERING\", \"QUESTION_ANSWERING\", \"FACT_VERIFICATION\", \"CODE_RETRIEVAL_QUERY\"]\n",
        "# @markdown ---\n",
        "# @markdown `learning_rate_multiplier`: Optional. A multiplier to apply to the recommended learning rate. To train with the recommended learning rate, use 1.0.\n",
        "learning_rate_multiplier = 1.0  # @param {type: \"number\"}\n",
        "# @markdown ---\n",
        "# @markdown `output_dimensionality`: Optional. The desired embedding dimension of your tuned model. This is only supported for models `text-embedding-004` and `text-multilingual-embedding-002`. To produce a model with the maximum embeddings size (for all existing models, this is 768), use -1.\n",
        "output_dimensionality = 768  # @param {type: \"number\"}\n",
        "\n",
        "# Synchronously validate some edge cases that will cause async validation to fail.\n",
        "if base_model_version_id not in [\"text-embedding-005\"]:\n",
        "    if task_type in [\"CODE_RETRIEVAL_QUERY\"]:\n",
        "        raise ValueError(\n",
        "            f\"task_type '{task_type}' is not valid for model '{base_model_version_id}'.\"\n",
        "        )\n",
        "if base_model_version_id not in [\n",
        "    \"text-embedding-005\",\n",
        "    \"text-embedding-004\",\n",
        "    \"text-multilingual-embedding-002\",\n",
        "]:\n",
        "    if task_type in [\"QUESTION_ANSWERING\", \"FACT_VERIFICATION\"]:\n",
        "        raise ValueError(\n",
        "            f\"task_type '{task_type}' is not valid for model '{base_model_version_id}'.\"\n",
        "        )\n",
        "    if output_dimensionality not in [-1, 768]:\n",
        "        raise ValueError(\n",
        "            f\"Model '{base_model_version_id}' does not support the output_dimensionality parameter.\"\n",
        "        )\n",
        "\n",
        "running_interactively = not os.getenv(\"IS_TESTING\")\n",
        "base_model = TextEmbeddingModel.from_pretrained(base_model_version_id)\n",
        "if \"TUNING_JOB_ID\" in locals() and TUNING_JOB_ID:\n",
        "    filter = f'pipelineJobUserId=\"{TUNING_JOB_ID}\"'\n",
        "    tuning_job = next(iter(pipeline_jobs.PipelineJob.list(filter=filter)))\n",
        "    print(\n",
        "        f\"Got an existing tuning job '{tuning_job.name}' (state: {tuning_job.state.name}).\"\n",
        "    )\n",
        "elif running_interactively:\n",
        "    tuning_result = base_model.tune_model(\n",
        "        task_type=task_type,\n",
        "        corpus_data=corpus_path,\n",
        "        queries_data=queries_path,\n",
        "        training_data=train_label_path,\n",
        "        validation_data=validation_label_path,\n",
        "        test_data=test_label_path,\n",
        "        batch_size=batch_size,\n",
        "        train_steps=train_steps,\n",
        "        tuned_model_location=REGION,\n",
        "        learning_rate_multiplier=learning_rate_multiplier,\n",
        "        output_dimensionality=output_dimensionality,\n",
        "    )\n",
        "    tuning_job = pipeline_jobs.PipelineJob.get(tuning_result.pipeline_job_name)\n",
        "    print(\n",
        "        f\"Got a fresh tuning job '{tuning_job.name}' (state: {tuning_job.state.name}).\"\n",
        "    )\n",
        "    print(\n",
        "        f\"(OPTIONAL) Set 'TUNING_JOB_ID' to '{tuning_job.name}' when you want to resume this tuning session.\"\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jl_vjc-m8itZ"
      },
      "source": [
        "### Deploy the Tuned Text-Embedding Model\n",
        "\n",
        "You can deploy a tuned model to an endpoint, once a tuning job is completed. Deploying a model can be done in one of these ways:\n",
        "\n",
        "* Calling `deploy_tuned_model` method on a tuning result object from `tune_model` method on a base model object.\n",
        "* Calling class method `TextEmbeddingModel.deploy_tuned_model` with a tuned model resource name. You can retrieve from a tuning job the uploaded, tuned model's resource name. This resource name is a string that follows the pattern projects/{PROJECT_ID}/locations/{REGION}/models/{MODEL_ID}.\n",
        "\n",
        "**IMPORTANT**: Deployment of tuned text-embedding models, being custom trained, requires the allocation of serving resources like machines and accelerators within your project. `deploy_tuned_model`, upon its initial call, creates a tuned model deployment onto a fresh serving resource. Upon subsequent calls, it simply retrieves the tuned model deployment from the existing serving resource.\n",
        "\n",
        "See also [using tuned models](https://cloud.google.com/vertex-ai/generative-ai/docs/models/tune-embeddings#use-tuned-model).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kIJC46m7SXvW"
      },
      "outputs": [],
      "source": [
        "MACHINE_TYPE = \"a2-highgpu-1g\"  # @param {type: \"string\"}\n",
        "ACCELERATOR = \"NVIDIA_TESLA_A100\"  # @param {type: \"string\"}\n",
        "ACCELERATOR_COUNT = 1  # @param {type: \"number\"}\n",
        "\n",
        "# CAVEAT: Colab disruptions may cause 'tuning_result' to be undefined.\n",
        "if \"tuning_job\" not in locals() and running_interactively:\n",
        "    message = \"[Action Required] Run the preceding code cells to define 'tuning_job'.\"\n",
        "    raise RuntimeError(message)\n",
        "if \"tuning_job\" in locals() and running_interactively:\n",
        "    if \"tuning_result\" in locals():\n",
        "        model = tuning_result.deploy_tuned_model(\n",
        "            machine_type=MACHINE_TYPE,\n",
        "            accelerator=ACCELERATOR,\n",
        "            accelerator_count=ACCELERATOR_COUNT,\n",
        "        )\n",
        "        print(f\"Got deployed, tuned {model=}\")\n",
        "    else:\n",
        "        tuning_job.wait()\n",
        "        tasks = tuning_job.task_details\n",
        "        upload_task = next(t for t in tasks if \"uploader\" in t.task_name)\n",
        "        upload_metadata = dict(upload_task.execution.metadata)\n",
        "        tuned_model_name = upload_metadata[\"output:model_resource_name\"]\n",
        "        model = TextEmbeddingModel.deploy_tuned_model(\n",
        "            tuned_model_name=tuned_model_name,\n",
        "            machine_type=MACHINE_TYPE,\n",
        "            accelerator=ACCELERATOR,\n",
        "            accelerator_count=ACCELERATOR_COUNT,\n",
        "        )\n",
        "        print(f\"Got deployed, tuned model '{tuned_model_name}'.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GAx9zdmQ80fR"
      },
      "source": [
        "### Examine Quality Metrics from Tuning Job\n",
        "\n",
        "The pipeline assesses the performance of both the base and tuned models during tuning by calculating NDCG@10 metrics using validation and test datasets. These metrics are available in the pipeline job's `metrics` artifact."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9Q8kj3QUclGi"
      },
      "outputs": [],
      "source": [
        "import pandas as pd\n",
        "\n",
        "if \"tuning_job\" in locals() and tuning_job.done():\n",
        "    tasks = tuning_job.task_details\n",
        "    eval_task = next(t for t in tasks if \"evaluator\" in t.task_name)\n",
        "    metrics = dict(eval_task.outputs[\"metrics\"].artifacts[0].metadata)\n",
        "    metrics_df = pd.DataFrame.from_dict(\n",
        "        {\"metric\": metrics.keys(), \"value\": metrics.values()}\n",
        "    )\n",
        "    display(metrics_df.sort_values(by=\"metric\", ignore_index=True))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rmHf3P7-9FjZ"
      },
      "source": [
        "### Get Tuned Embeddings for Downstream Tasks\n",
        "\n",
        "Get embeddings from your tuned model for downstream tasks."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1NYIaUiLTJrg"
      },
      "outputs": [],
      "source": [
        "import pandas as pd\n",
        "\n",
        "if \"model\" in locals():\n",
        "    texts = [\"banana muffins?\"]  # @param {type:\"raw\"}\n",
        "    titles = [\"none\"]  # @param {type:\"raw\"}\n",
        "    embedding_inputs = [\n",
        "        TextEmbeddingInput(text=text, task_type=task_type, title=title)\n",
        "        for text, title in zip(texts, titles)\n",
        "    ]\n",
        "    tuned_embeddings = [\n",
        "        pd.Series(e.values) for e in model.get_embeddings(embedding_inputs)\n",
        "    ]\n",
        "    display(tuned_embeddings)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TAIM9cRRNNwh"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
        "\n",
        "Otherwise, you can delete the individual resources you created in this tutorial."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "tuned_text-embeddings.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}