notebooks/community/model_garden/model_garden_e5.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7d9bbf86da5e"
      },
      "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": "99c1c3fc2ca5"
      },
      "source": [
        "# Vertex AI Model Garden - E5 Text Embedding Models\n",
        "\n",
        "<table><tbody><tr>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_e5.ipynb\">\n",
        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_e5.ipynb\">\n",
        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</tr></tbody></table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3de7470326a2"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates deploying E5 text embedding models in Vertex AI.\n",
        "\n",
        "### Objective\n",
        "\n",
        "- Deploy prebuilt E5 models with Hugging Face [Text Embeddings Inference](https://github.com/huggingface/text-embeddings-inference) (TEI) docker image on a Vertex AI Endpoint\n",
        "    - [intfloat/multilingual-e5-large-instruct](https://huggingface.co/intfloat/multilingual-e5-large-instruct): 560M params, instruction-tuned\n",
        "    - [intfloat/multilingual-e5-large](https://huggingface.co/intfloat/multilingual-e5-large): 560M params\n",
        "    - [intfloat/e5-large-v2](https://huggingface.co/intfloat/e5-large-v2): 335M params\n",
        "    - [intfloat/multilingual-e5-small](https://huggingface.co/intfloat/multilingual-e5-small): 118M params\n",
        "    - [intfloat/e5-base-v2](https://huggingface.co/intfloat/e5-base-v2): 109M params\n",
        "    - [intfloat/e5-small-v2](https://huggingface.co/intfloat/e5-small-v2): 33M params\n",
        "- Run inference on the deployed Vertex AI Endpoint\n",
        "\n",
        "\n",
        "### 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), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), [Cloud NL API pricing](https://cloud.google.com/natural-language/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "264c07757582"
      },
      "source": [
        "## Before you begin"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "LyEVDkHhAUHF"
      },
      "outputs": [],
      "source": [
        "# @title Setup Google Cloud project\n",
        "\n",
        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
        "\n",
        "import os\n",
        "import uuid\n",
        "from datetime import datetime\n",
        "from typing import Tuple\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "\n",
        "# Get the default cloud project id.\n",
        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
        "\n",
        "# Get the default region for launching jobs.\n",
        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
        "\n",
        "# Cloud Storage bucket for storing the experiment artifacts.\n",
        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
        "# prefer using your own GCS bucket, change the value yourself below.\n",
        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "BUCKET_URI = \"gs://\"  # @param {type: \"string\"}\n",
        "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
        "\n",
        "# @markdown Click \"Show code\" to see more details.\n",
        "\n",
        "# Create a unique GCS bucket for this notebook, if not specified by the user.\n",
        "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "else:\n",
        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
        "    bucket_region = shell_output[0].strip().lower()\n",
        "    if bucket_region != REGION:\n",
        "        raise ValueError(\n",
        "            \"Bucket region %s is different from notebook region %s\"\n",
        "            % (bucket_region, REGION)\n",
        "        )\n",
        "\n",
        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
        "\n",
        "\n",
        "! gcloud config set project $PROJECT_ID\n",
        "! gcloud services enable language.googleapis.com\n",
        "\n",
        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
        "\n",
        "# Gets the default BUCKET_URI and SERVICE_ACCOUNT if they were not specified by the user.\n",
        "\n",
        "SERVICE_ACCOUNT = None\n",
        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
        "\n",
        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
        "\n",
        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
        "\n",
        "\n",
        "def create_name_with_datetime(prefix: str) -> str:\n",
        "    \"\"\"Creates a name with date time when triggering training or deployment\n",
        "    jobs in Vertex AI.\n",
        "    \"\"\"\n",
        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
        "\n",
        "\n",
        "def deploy_model_tei(\n",
        "    model_name: str,\n",
        "    model_id: str,\n",
        "    service_account: str,\n",
        "    docker_uri: str,\n",
        "    machine_type: str = \"g2-standard-8\",\n",
        "    accelerator_type: str = \"NVIDIA_L4\",\n",
        "    accelerator_count: int = 1,\n",
        "    max_model_len: int = 512,\n",
        "    gpu_memory_utilization: float = 0.9,\n",
        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
        "    \"\"\"Deploys E5 models with TEI on Vertex AI.\n",
        "\n",
        "    Args:\n",
        "        model_name: Display name of the model.\n",
        "        model_id: Model ID or path to model weights.\n",
        "        service_account: Service account for model uploading and deployment.\n",
        "        machine_type: Deployment machine type.\n",
        "        accelerator_type: Deployment accelerator type.\n",
        "        accelerator_count: Number of accelerators to use.\n",
        "        max_model_len: Maximum model length.\n",
        "        gpu_memory_utilization: Fraction of GPU memory to be used for the model\n",
        "            executor.\n",
        "\n",
        "    Returns:\n",
        "        Model instance and endpoint instance.\n",
        "    \"\"\"\n",
        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
        "\n",
        "    tei_args = [\n",
        "        f\"--model-id={model_id}\",\n",
        "    ]\n",
        "    serving_env = {\n",
        "        \"MODEL_ID\": model_id,\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
        "        serving_container_image_uri=docker_uri,\n",
        "        serving_container_args=tei_args,\n",
        "        serving_container_ports=[7080],\n",
        "        serving_container_environment_variables=serving_env,\n",
        "        model_garden_source_model_name=\"publishers/intfloat/models/e5\"\n",
        "    )\n",
        "\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        machine_type=machine_type,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=accelerator_count,\n",
        "        deploy_request_timeout=1800,\n",
        "        service_account=service_account,\n",
        "        system_labels={\n",
        "            \"NOTEBOOK_NAME\": \"model_garden_e5.ipynb\"\n",
        "        },\n",
        "    )\n",
        "    return model, endpoint"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2xeBQF0iVwSr"
      },
      "source": [
        "## Deploy and Predict"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "kg5MwMIfB9Uj"
      },
      "outputs": [],
      "source": [
        "# @title Deploy\n",
        "# @markdown This section uploads a prebuilt model to Model Registry and deploys it on the Endpoint. The model deployment step will take ~15 minutes to complete.\n",
        "\n",
        "prebuilt_model_id = \"intfloat/e5-small-v2\"  # @param [\"intfloat/multilingual-e5-large-instruct\", \"intfloat/multilingual-e5-large\", \"intfloat/e5-large-v2\", \"intfloat/multilingual-e5-small\", \"intfloat/e5-base-v2\", \"intfloat/e5-small-v2\"]\n",
        "\n",
        "# @markdown Specify a processor for the TEI docker image. E5 models can be run on either GPU or CPU.\n",
        "processor = \"NVIDIA_L4\"  # @param[\"NVIDIA_TESLA_V100\", \"NVIDIA_L4\", \"NVIDIA_TESLA_A100\", \"CPU\"]\n",
        "\n",
        "if processor == \"NVIDIA_TESLA_V100\":\n",
        "    accelerator_type = \"NVIDIA_TESLA_V100\"\n",
        "    machine_type = \"n1-highmem-16\"\n",
        "    accelerator_count = 2\n",
        "elif processor == \"NVIDIA_L4\":\n",
        "    accelerator_type = \"NVIDIA_L4\"\n",
        "    machine_type = \"g2-standard-8\"\n",
        "    accelerator_count = 1\n",
        "elif processor == \"NVIDIA_TESLA_A100\":\n",
        "    accelerator_type = \"NVIDIA_TESLA_A100\"\n",
        "    machine_type = \"a2-highgpu-1g\"\n",
        "    accelerator_count = 1\n",
        "elif processor == \"CPU\":\n",
        "    accelerator_type = None\n",
        "    machine_type = None\n",
        "    accelerator_count = None\n",
        "else:\n",
        "    raise ValueError(f\"Unsupported processor: {processor}\")\n",
        "\n",
        "\n",
        "# The pre-built serving docker images with TEI.\n",
        "if processor == \"CPU\":\n",
        "    TEI_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-embeddings-inference-cpu.1-2\"\n",
        "else:\n",
        "    TEI_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-embeddings-inference-cu122.1-2.ubuntu2204\"\n",
        "\n",
        "# @markdown Click \"Show code\" to see more details.\n",
        "\n",
        "# Finds Vertex AI prediction supported accelerators and regions in\n",
        "# https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
        "\n",
        "model, endpoint = deploy_model_tei(\n",
        "    model_name=create_name_with_datetime(prefix=\"e5-serve-tei\"),\n",
        "    model_id=prebuilt_model_id,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        "    docker_uri=TEI_DOCKER_URI,\n",
        "    machine_type=machine_type,\n",
        "    accelerator_type=accelerator_type,\n",
        "    accelerator_count=accelerator_count,\n",
        ")\n",
        "\n",
        "print(\"endpoint_name:\", endpoint.name)\n",
        "print(\"model_name:\", model.display_name)\n",
        "print(\"model_id:\", model.resource_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cUGriDwUQyx4"
      },
      "source": [
        "### Predict\n",
        "\n",
        "Once deployment succeeds, you can send requests to the endpoint with text prompts."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "UKfBeMJCN2wl"
      },
      "outputs": [],
      "source": [
        "# @title Run sample prompt\n",
        "# @markdown Below is an example to encode queries and passages from the [MS-MARCO passage ranking](https://github.com/microsoft/MSMARCO-Passage-Ranking) dataset.\n",
        "\n",
        "# @markdown Example:\n",
        "\n",
        "# @markdown ```\n",
        "# @markdown query: how much protein should a female eat\n",
        "# @markdown query: summit define\n",
        "# @markdown passage: As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.\n",
        "# @markdown passage: Definition of summit for English Language Learners. : 1  the highest point of a mountain : the top of a mountain. : 2  the highest level. : 3  a meeting or series of meetings between the leaders of two or more governments.\n",
        "# @markdown ```\n",
        "\n",
        "# @markdown NOTE: Inputs are not limited to 2 queries and 2 passages. To add more inputs, you may modify the code directly.\n",
        "\n",
        "query1 = \"how much protein should a female eat?\"  # @param {type: \"string\"}\n",
        "query2 = \"summit define\"  # @param {type: \"string\"}\n",
        "passage1 = \"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.\"  # @param {type: \"string\"}\n",
        "passage2 = \"Definition of summit for English Language Learners. : 1  the highest point of a mountain : the top of a mountain. : 2  the highest level. : 3  a meeting or series of meetings between the leaders of two or more governments.\"  # @param {type: \"string\"}\n",
        "\n",
        "# @markdown Click \"Show code\" to see more details.\n",
        "\n",
        "# Loads an existing endpoint instance using the endpoint name:\n",
        "# - Using `endpoint_name = endpoint.name` allows us to get the\n",
        "#   endpoint name of the endpoint `endpoint` created in the cell\n",
        "#   above.\n",
        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
        "#   an existing endpoint with the ID 1234567890123456789.\n",
        "# You may uncomment the code below to load an existing endpoint.\n",
        "\n",
        "# endpoint_name = \"\"  # @param {type:\"string\"}\n",
        "# aip_endpoint_name = (\n",
        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
        "# )\n",
        "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
        "\n",
        "from torch import Tensor\n",
        "\n",
        "# Each input text should start with \"query: \" or \"passage: \".\n",
        "# For tasks other than retrieval, you can simply use the \"query: \" prefix.\n",
        "instances = [\n",
        "    {\n",
        "        \"inputs\": [\n",
        "            f\"query: {query1}\",\n",
        "            f\"query: {query2}\",\n",
        "            f\"passage: {passage1}\",\n",
        "            f\"passage: {passage2}\",\n",
        "        ],\n",
        "    },\n",
        "]\n",
        "response = endpoint.predict(instances=instances)\n",
        "\n",
        "for prediction in response.predictions:\n",
        "    embeddings = Tensor(prediction)\n",
        "    scores = (embeddings[:2] @ embeddings[2:].T) * 100\n",
        "    print(scores.tolist())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "3hHR7xxyOklz"
      },
      "outputs": [],
      "source": [
        "# @title Run sample prompt for instruction-tuned models\n",
        "# @markdown For instruction-tuned models (e.g. intfloat/multilingual-e5-large-instruct), the task definition should be a one-sentence instruction that describes the task. This is a way to customize text embeddings for different scenarios through natural language instructions.\n",
        "\n",
        "# @markdown Below is an example to encode queries and passages from the [MS-MARCO passage ranking](https://github.com/microsoft/MSMARCO-Passage-Ranking) dataset.\n",
        "\n",
        "# @markdown Example:\n",
        "\n",
        "# @markdown ```\n",
        "# @markdown Instruct: Given a web search query, retrieve relevant passages that answer the query\n",
        "# @markdown Query: how much protein should a female eat\n",
        "# @markdown Instruct: Given a web search query, retrieve relevant passages that answer the query\n",
        "# @markdown Query: 南瓜的家常做法\n",
        "# @markdown As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.\n",
        "# @markdown 1.清炒南瓜丝 原料:嫩南瓜半个 调料:葱、盐、白糖、鸡精 做法: 1、南瓜用刀薄薄的削去表面一层皮,用勺子刮去瓤 2、擦成细丝(没有擦菜板就用刀慢慢切成细丝) 3、锅烧热放油,入葱花煸出香味 4、入南瓜丝快速翻炒一分钟左右,放盐、一点白糖和鸡精调味出锅 2.香葱炒南瓜 原料:南瓜1只 调料:香葱、蒜末、橄榄油、盐 做法: 1、将南瓜去皮,切成片 2、油锅8成热后,将蒜末放入爆香 3、爆香后,将南瓜片放入,翻炒 4、在翻炒的同时,可以不时地往锅里加水,但不要太多 5、放入盐,炒匀 6、南瓜差不多软和绵了之后,就可以关火 7、撒入香葱,即可出锅\n",
        "# @markdown ```\n",
        "\n",
        "# @markdown API reference link to HuggingFace : [Text Embeddings Inference API](https://huggingface.github.io/text-embeddings-inference/#/).\n",
        "\n",
        "# @markdown NOTE: Inputs are not limited to 1 instruction, 2 queries, and 2 documents. To add more inputs, you may modify the code directly.\n",
        "\n",
        "instruction = \"Given a web search query, retrieve relevant passages that answer the query\"  # @param {type: \"string\"}\n",
        "query1 = \"how much protein should a female eat\"  # @param {type: \"string\"}\n",
        "query2 = \"南瓜的家常做法\"  # @param {type: \"string\"}\n",
        "document1 = \"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.\"  # @param {type: \"string\"}\n",
        "document2 = \"1.清炒南瓜丝 原料:嫩南瓜半个 调料:葱、盐、白糖、鸡精 做法: 1、南瓜用刀薄薄的削去表面一层皮,用勺子刮去瓤 2、擦成细丝(没有擦菜板就用刀慢慢切成细丝) 3、锅烧热放油,入葱花煸出香味 4、入南瓜丝快速翻炒一分钟左右,放盐、一点白糖和鸡精调味出锅 2.香葱炒南瓜 原料:南瓜1只 调料:香葱、蒜末、橄榄油、盐 做法: 1、将南瓜去皮,切成片 2、油锅8成热后,将蒜末放入爆香 3、爆香后,将南瓜片放入,翻炒 4、在翻炒的同时,可以不时地往锅里加水,但不要太多 5、放入盐,炒匀 6、南瓜差不多软和绵了之后,就可以关火 7、撒入香葱,即可出锅\"  # @param {type: \"string\"}\n",
        "\n",
        "# @markdown Click \"Show code\" to see more details.\n",
        "\n",
        "# Loads an existing endpoint instance using the endpoint name:\n",
        "# - Using `endpoint_name = endpoint.name` allows us to get the\n",
        "#   endpoint name of the endpoint `endpoint` created in the cell\n",
        "#   above.\n",
        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
        "#   an existing endpoint with the ID 1234567890123456789.\n",
        "# You may uncomment the code below to load an existing endpoint.\n",
        "\n",
        "# endpoint_name = \"\"  # @param {type:\"string\"}\n",
        "# aip_endpoint_name = (\n",
        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
        "# )\n",
        "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
        "\n",
        "from torch import Tensor\n",
        "\n",
        "\n",
        "def get_detailed_instruct(task_description: str, query: str) -> str:\n",
        "    return f\"Instruct: {task_description}\\nQuery: {query}\"\n",
        "\n",
        "\n",
        "# Each query must come with a one-sentence instruction that describes the task\n",
        "queries = [\n",
        "    get_detailed_instruct(instruction, query1),\n",
        "    get_detailed_instruct(instruction, query2),\n",
        "]\n",
        "# No need to add instruction for retrieval documents\n",
        "documents = [\n",
        "    document1,\n",
        "    document2,\n",
        "]\n",
        "\n",
        "instances = [{\"inputs\": queries + documents}]\n",
        "\n",
        "response = endpoint.predict(instances=instances)\n",
        "\n",
        "for prediction in response.predictions:\n",
        "    embeddings = Tensor(prediction)\n",
        "    scores = (embeddings[:2] @ embeddings[2:].T) * 100\n",
        "    print(scores.tolist())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Nun3w71JYbss"
      },
      "source": [
        "## Clean up resources"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "EOL0Qt_0YT5D"
      },
      "outputs": [],
      "source": [
        "# @title Delete the models and endpoints\n",
        "\n",
        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
        "# @markdown  and avoid unnecessary continuous charges that may incur.\n",
        "\n",
        "endpoint.delete(force=True)\n",
        "model.delete()\n",
        "\n",
        "delete_bucket = False  # @param {type:\"boolean\"}\n",
        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_e5.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}