notebooks/community/model_garden/model_garden_huggingface_tei_deployment.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "20qcPG1PmFUM"
      },
      "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": "QXYOa1odnikj"
      },
      "source": [
        "# Vertex AI Model Garden - Hugging Face Text Embeddings Inference Deployment\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_huggingface_tei_deployment.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_huggingface_tei_deployment.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": "cbDI9ag4oR4C"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates deploying [nomic-ai/nomic-embed-text-v1](https://huggingface.co/nomic-ai/nomic-embed-text-v1) with [Text Embeddings Inference (TEI)](https://github.com/huggingface/text-embeddings-inference) from Hugging Face. In additional to `nomic-ai/nomic-embed-text-v1`, You can view and change the code to deploy a different Hugging Face `text-embeddings-inference` model with appropriate machine specs. **Note that some models might fail to deploy, even if they have `text-embeddings-inference` tags on the Hugging Face model card page.**\n",
        "\n",
        "\n",
        "### Objective\n",
        "\n",
        "- Download and deploy the `nomic-ai/nomic-embed-text-v1` model with TEI\n",
        "- Send prediction request to the deployed endpoint\n",
        "\n",
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/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": "hQJWRopioSKT"
      },
      "source": [
        "## Run the notebook"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "J_jmxcIZoSxU"
      },
      "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]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
        "\n",
        "REGION = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# Upgrade Vertex AI SDK.\n",
        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.64.0'\n",
        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
        "\n",
        "import importlib\n",
        "import os\n",
        "from typing import Tuple\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "\n",
        "common_util = importlib.import_module(\n",
        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
        ")\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",
        "if not REGION:\n",
        "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
        "\n",
        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
        "\n",
        "# Initialize Vertex AI API.\n",
        "print(\"Initializing Vertex AI API.\")\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
        "! gcloud config set project $PROJECT_ID\n",
        "\n",
        "models, endpoints = {}, {}\n",
        "\n",
        "# Dedicated endpoint not supported yet\n",
        "use_dedicated_endpoint = False\n",
        "HF_TOKEN = \"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "USB7dvYqvNdu"
      },
      "outputs": [],
      "source": [
        "# @title Deploy with TEI from Hugging Face\n",
        "\n",
        "# @markdown This section downloads the `nomic-ai/nomic-embed-text-v1` model from Hugging Face and deploys it to a Vertex AI Endpoint.\n",
        "# @markdown It takes ~20 minutes to complete the deployment.\n",
        "\n",
        "MODEL_ID = \"nomic-ai/nomic-embed-text-v1\"  # @param {type: \"string\", isTemplate: true}\n",
        "\n",
        "# The pre-built serving docker images for TEI.\n",
        "TEI_CPU_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-embeddings-inference-cpu.1-4\"\n",
        "TEI_GPU_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-embeddings-inference-cu122.1-4.ubuntu2204\"\n",
        "\n",
        "machine_type = \"g2-standard-8\"  # @param {type: \"string\", isTemplate: true}\n",
        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"None\"] {isTemplate: true}\n",
        "\n",
        "if accelerator_type == \"None\":\n",
        "    accelerator_type = \"\"\n",
        "\n",
        "if accelerator_type:\n",
        "    common_util.check_quota(\n",
        "        project_id=PROJECT_ID,\n",
        "        region=REGION,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=1 if accelerator_type else 0,\n",
        "        is_for_training=False,\n",
        "    )\n",
        "\n",
        "\n",
        "def deploy_model_tei(\n",
        "    model_name: str,\n",
        "    model_id: str,\n",
        "    service_account: str = \"\",\n",
        "    machine_type: str = \"g2-standard-4\",\n",
        "    accelerator_type: str = \"NVIDIA_L4\",\n",
        "    use_dedicated_endpoint: bool = False,\n",
        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
        "    \"\"\"Deploys models with TEI on Vertex AI.\"\"\"\n",
        "    endpoint = aiplatform.Endpoint.create(\n",
        "        display_name=f\"{model_name}-endpoint\",\n",
        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
        "    )\n",
        "\n",
        "    docker_uri = TEI_GPU_DOCKER_URI if accelerator_type else TEI_CPU_DOCKER_URI\n",
        "    env_vars = {\n",
        "        \"MODEL_ID\": model_id,\n",
        "        \"JSON_OUTPUT\": \"true\",\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
        "\n",
        "    # HF_TOKEN is not a compulsory field and may not be defined.\n",
        "    try:\n",
        "        if HF_TOKEN:\n",
        "            env_vars[\"HF_API_TOKEN\"] = HF_TOKEN\n",
        "    except NameError:\n",
        "        pass\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
        "        serving_container_image_uri=docker_uri,\n",
        "        serving_container_ports=[8080],\n",
        "        serving_container_environment_variables=env_vars,\n",
        "        serving_container_shared_memory_size_mb=(4 * 1024),  # 4 GB\n",
        "    )\n",
        "\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        machine_type=machine_type,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=1 if accelerator_type else 0,\n",
        "        deploy_request_timeout=1800,\n",
        "        service_account=service_account,\n",
        "    )\n",
        "    return model, endpoint\n",
        "\n",
        "\n",
        "models[\"tei\"], endpoints[\"tei\"] = deploy_model_tei(\n",
        "    model_name=common_util.get_job_name_with_datetime(prefix=MODEL_ID),\n",
        "    model_id=MODEL_ID,\n",
        "    service_account=\"\",\n",
        "    machine_type=machine_type,\n",
        "    accelerator_type=accelerator_type,\n",
        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
        ")\n",
        "\n",
        "# @markdown Click \"Show Code\" to see more details."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Aa4e1-6FvRAP"
      },
      "outputs": [],
      "source": [
        "# @title Predict\n",
        "\n",
        "# @markdown Once deployment succeeds, you can send requests to the endpoint which computes text embeddings.\n",
        "\n",
        "# @markdown Here we use a simple example: `This is a sentence`.\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",
        "text = \"This is a sentence.\"  # @param {type: \"string\"}\n",
        "\n",
        "instances = [{\"inputs\": text}]\n",
        "response = endpoints[\"tei\"].predict(\n",
        "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
        ")\n",
        "\n",
        "for prediction in response.predictions:\n",
        "    print(prediction)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tAelDidov5AW"
      },
      "source": [
        "## Clean up resources"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "8SeZCFo5v7z-"
      },
      "outputs": [],
      "source": [
        "# @title Delete the models and endpoints\n",
        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
        "# @markdown  and avoid unnecessary continuous charges that may incur.\n",
        "\n",
        "# Undeploy model and delete endpoint.\n",
        "for endpoint in endpoints.values():\n",
        "    endpoint.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
        "for model in models.values():\n",
        "    model.delete()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_huggingface_tei_deployment.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}