notebooks/community/model_garden/model_garden_pytorch_biogpt_serve.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7d9bbf86da5e"
      },
      "outputs": [],
      "source": [
        "# Copyright 2023 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 - BioGPT\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_biogpt_serve.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_biogpt_serve.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",
        "  <td>\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_pytorch_biogpt_serve.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> (A Python-3 CPU notebook is recommended)\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3de7470326a2"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates how to serve [BioGPT](https://github.com/microsoft/BioGPT), a large language model trained on a massive dataset of biomedical text and code. BioGPT can be used for a variety of tasks, such as generating text.\n",
        "\n",
        "BioGPT has been integrated into the [Hugging Face transformers](https://huggingface.co/docs/transformers/model_doc/biogpt) library, and model checkpoints are available on the Hugging Face Hub. This means that you can use BioGPT directly with a pipeline for text generation.\n",
        "\n",
        "In this notebook, we will show you how to:\n",
        "- Serve BioGPT using Vertex AI\n",
        "- Generate text using BioGPT\n",
        "\n",
        "### Objective\n",
        "\n",
        "- To run inference with BioGPT on Vertex AI, the following models are supported\n",
        "\n",
        "| Models | LoRA |\n",
        "| :- | :- |\n",
        "| [microsoft/biogpt](https://huggingface.co/docs/transformers/model_doc/biogpt) | N |\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) and [Cloud Storage pricing](https://cloud.google.com/storage/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\n",
        "\n",
        "**NOTE**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ioensNKM8ned"
      },
      "source": [
        "### Colab only\n",
        "Run the following commands for Colab and skip this section if you are using Workbench."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2707b02ef5df"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
        "    from google.colab import auth as google_auth\n",
        "\n",
        "    google_auth.authenticate_user()\n",
        "    # Install gdown for downloading example training images.\n",
        "    ! pip3 install gdown\n",
        "\n",
        "    # Restart the notebook kernel after installs.\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bb7adab99e41"
      },
      "source": [
        "### Setup Google Cloud project\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",
        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
        "\n",
        "1. [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs.\n",
        "\n",
        "1. [Create a service account](https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console) with `Vertex AI User` and `Storage Object Admin` roles for deploying fine tuned model to Vertex AI endpoint."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6c460088b873"
      },
      "source": [
        "Fill following variables for experiments environment:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "855d6b96f291"
      },
      "outputs": [],
      "source": [
        "# Cloud project id.\n",
        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# The region you want to launch jobs in.\n",
        "REGION = \"us-central1\"  # @param {type:\"string\"}\n",
        "\n",
        "# The Cloud Storage bucket for storing experiments output.\n",
        "BUCKET_URI = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "! gcloud config set project $PROJECT_ID\n",
        "\n",
        "# The service account looks like:\n",
        "# '@.iam.gserviceaccount.com'\n",
        "# Please go to https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console\n",
        "# and create service account with `Vertex AI User` and `Storage Object Admin` roles.\n",
        "# The service account for deploying fine tuned model.\n",
        "SERVICE_ACCOUNT = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# The serving port.\n",
        "SERVE_PORT = 7080"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e828eb320337"
      },
      "source": [
        "### Initialize Vertex AI API"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "12cd25839741"
      },
      "outputs": [],
      "source": [
        "from google.cloud import aiplatform\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2cc825514deb"
      },
      "source": [
        "### Define constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0172207da484"
      },
      "outputs": [],
      "source": [
        "# The pre-built training and serving docker images.\n",
        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-biogpt-serve\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0c250872074f"
      },
      "source": [
        "### Define common functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "354da31189dc"
      },
      "outputs": [],
      "source": [
        "from datetime import datetime\n",
        "\n",
        "\n",
        "def get_job_name_with_datetime(prefix: str):\n",
        "    \"\"\"Gets the job 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(\n",
        "    model_name,\n",
        "    model_id,\n",
        "    service_account,\n",
        "    task,\n",
        "    max_length,\n",
        "    num_return_sequences,\n",
        "    machine_type=\"n1-standard-8\",\n",
        "    accelerator_type=\"NVIDIA_TESLA_V100\",\n",
        "    accelerator_count=1,\n",
        "):\n",
        "    \"\"\"Deploys trained models into Vertex AI.\"\"\"\n",
        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
        "    serving_env = {\n",
        "        \"MODEL_ID\": model_id,\n",
        "        \"TASK\": task,\n",
        "        \"MAX_LENGTH\": max_length,\n",
        "        \"NUM_RETURN_SEQUENCES\": num_return_sequences,\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
        "    # If the model_id is a GCS path, use artifact_uri to pass it to serving docker.\n",
        "    artifact_uri = model_id if model_id.startswith(\"gs://\") else None\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
        "        serving_container_image_uri=SERVE_DOCKER_URI,\n",
        "        serving_container_ports=[7080],\n",
        "        serving_container_predict_route=\"/predictions/biogpt_serving\",\n",
        "        serving_container_health_route=\"/ping\",\n",
        "        serving_container_environment_variables=serving_env,\n",
        "        artifact_uri=artifact_uri,\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",
        "    )\n",
        "    return model, endpoint"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jqmCtkGnhDmp"
      },
      "source": [
        "### Run inferences with serving images\n",
        "This section uploads the model to Model Registry and deploys it on the Endpoint.\n",
        "\n",
        "The model deployment step will take ~15 minutes to complete."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4f79683ccd37"
      },
      "outputs": [],
      "source": [
        "model_id = \"microsoft/biogpt\"  # @param [\"microsoft/biogpt\"]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bf55e38815dc"
      },
      "outputs": [],
      "source": [
        "model, endpoint = deploy_model(\n",
        "    model_name=get_job_name_with_datetime(prefix=\"biogpt-serve\"),\n",
        "    model_id=model_id,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        "    task=\"text-generation\",\n",
        "    max_length=200,\n",
        "    num_return_sequences=10,\n",
        ")\n",
        "print(\"endpoint_name:\", endpoint.name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "80b3fd2ace09"
      },
      "source": [
        "NOTE: The model weights will be downloaded after the deployment succeeds. Thus additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint.\n",
        "\n",
        "Once deployment succeeds, you can send requests to the endpoint with text prompts."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4ab04da3ec9a"
      },
      "outputs": [],
      "source": [
        "# # Loads an existing endpoint as below.\n",
        "# endpoint_name = endpoint.name\n",
        "# aip_endpoint_name = (\n",
        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
        "# )\n",
        "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
        "\n",
        "instances = [\n",
        "    {\"prompt\": \"COVID-19 is\"},\n",
        "]\n",
        "response = endpoint.predict(instances=instances)\n",
        "\n",
        "print(response.predictions[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "af21a3cff1e0"
      },
      "source": [
        "### Clean up resources"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "911406c1561e"
      },
      "outputs": [],
      "source": [
        "# Undeploy model and delete endpoint.\n",
        "endpoint.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
        "model.delete()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_pytorch_biogpt_serve.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}