notebooks/community/model_garden/model_garden_pytorch_nllb.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": "2bd716bf3e39"
      },
      "source": [
        "# Vertex AI Model Garden - NLLB (No Language Left Behind)\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_nllb.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_nllb.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_nllb.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",
        "    (a Python-3 CPU notebook is recommended)\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d8cd12648da4"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates deploying the pre-trained [NLLB](https://huggingface.co/facebook/nllb-200-3.3B) model on Vertex AI for online prediction.\n",
        "\n",
        "### Objective\n",
        "\n",
        "- Upload the model to [Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction).\n",
        "- Deploy the model on [Endpoint](https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints).\n",
        "- Run online predictions for translation.\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": "264c07757582"
      },
      "source": [
        "## Setup environment\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": "d73ffa0c0b83"
      },
      "source": [
        "### Colab only"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2707b02ef5df"
      },
      "outputs": [],
      "source": [
        "if \"google.colab\" in str(get_ipython()):\n",
        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
        "    from google.colab import auth as google_auth\n",
        "\n",
        "    google_auth.authenticate_user()\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": "0f826ff482a2"
      },
      "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."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8958ebc71868"
      },
      "source": [
        "Fill following variables for experiments environment:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9db30f827a65"
      },
      "outputs": [],
      "source": [
        "# Cloud project id.\n",
        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# Google Cloud Storage bucket URI.\n",
        "GCS_BUCKET = \"gs://\"  # @param {type:\"string\"}\n",
        "\n",
        "# The region you want to launch jobs in.\n",
        "REGION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "92f16e22c20b"
      },
      "source": [
        "Initialize Vertex AI API:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1680c257acfb"
      },
      "outputs": [],
      "source": [
        "from google.cloud import aiplatform\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=GCS_BUCKET)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6ca48b699d17"
      },
      "source": [
        "### Define constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "de9882ea89ea"
      },
      "outputs": [],
      "source": [
        "# The pre-built serving docker image.\n",
        "# The model artifacts are embedded within the container, except for model weights which will be downloaded during deployment.\n",
        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-transformers-serve:20231022_0832_RC00\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "10188266a5cd"
      },
      "source": [
        "### Define common functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cac4478ae098"
      },
      "outputs": [],
      "source": [
        "def deploy_model(model_id, accelerator_count):\n",
        "    \"\"\"Uploads and deploys the model to Vertex AI endpoint for prediction.\"\"\"\n",
        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_id}-endpoint\")\n",
        "    serving_env = {\n",
        "        \"MODEL_ID\": model_id,\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_id,\n",
        "        serving_container_image_uri=SERVE_DOCKER_URI,\n",
        "        serving_container_ports=[7080],\n",
        "        serving_container_predict_route=\"/predictions/transformers_serving\",\n",
        "        serving_container_health_route=\"/ping\",\n",
        "        serving_container_environment_variables=serving_env,\n",
        "    )\n",
        "\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        machine_type=\"n1-standard-8\",\n",
        "        accelerator_type=\"NVIDIA_TESLA_V100\",\n",
        "        accelerator_count=accelerator_count,\n",
        "        deploy_request_timeout=3600,\n",
        "    )\n",
        "    return model, endpoint"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d2d72ecdb8c9"
      },
      "source": [
        "## Upload and deploy models"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9448c5f545fa"
      },
      "source": [
        "This notebook and the serving can support all the 4 variations of the NLLB models listed below:\n",
        "\n",
        "\n",
        "| Model Name | Model Type | #params | checkpoint |\n",
        "| - | - | - | - |\n",
        "| NLLB-200-3.3B | Dense | 3.3B |[model](https://huggingface.co/facebook/nllb-200-3.3B) |\n",
        "| NLLB-200-1.3B | Dense | 1.3B |[model](https://huggingface.co/facebook/nllb-200-1.3B) |\n",
        "| NLLB-200-Distilled-1.3B | Dense | 1.3B | [model](https://huggingface.co/facebook/nllb-200-distilled-1.3B) |\n",
        "| NLLB-200-Distilled-600M | Dense | 600M | [model](https://huggingface.co/facebook/nllb-200-distilled-600M) |\n",
        "\n",
        "\n",
        "Among the 4 variations, the [NLLB-200-3.3B](https://huggingface.co/facebook/nllb-200-3.3B) requires 2 V100 GPUs for online prediction. The rest of the variations, however, only require a single V100 GPU. This section demonstrate how to upload the pre-trained model to Model Registry and deploys it on the Endpoint with the corresponding machine specifications.\n",
        "\n",
        "NOTE: The model deployment step will take ~25 minutes to complete."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b4b46c28d8b1"
      },
      "outputs": [],
      "source": [
        "# Please select a model from the list:\n",
        "model_id = \"facebook/nllb-200-distilled-600M\"  # @param [\"facebook/nllb-200-distilled-600M\", \"facebook/nllb-200-distilled-1.3B\", \"facebook/nllb-200-1.3B\", \"facebook/nllb-200-3.3B\"]\n",
        "\n",
        "accelerator_count = None\n",
        "# For models `NLLB-200-Distilled-600M`, `NLLB-200-Distilled-1.3B`, `NLLB-200-1.3B`,\n",
        "# we use 1 V100 GPU\n",
        "if model_id in [\n",
        "    \"facebook/nllb-200-distilled-600M\",\n",
        "    \"facebook/nllb-200-distilled-1.3B\",\n",
        "    \"facebook/nllb-200-1.3B\",\n",
        "]:\n",
        "    accelerator_count = 1\n",
        "# For model `NLLB-200-3.3B`, we require 2 V100 GPUs\n",
        "elif model_id in [\"facebook/nllb-200-3.3B\"]:\n",
        "    accelerator_count = 2\n",
        "\n",
        "\n",
        "if model_id is None or accelerator_count is None:\n",
        "    raise ValueError(\"Please select a valid model from the list.\")\n",
        "\n",
        "model, endpoint = deploy_model(model_id, accelerator_count)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "80b3fd2ace09"
      },
      "source": [
        "NOTE: The model weights will be downloaded after the deployment succeeds. Thus additional 20 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."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "68990d91bc5f"
      },
      "source": [
        "## Send prediction request"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6be655247cb1"
      },
      "outputs": [],
      "source": [
        "input_text = \"Şeful ONU spune că nu există o soluţie militară în Siria\"\n",
        "# The source and the target languages. The full list of the languages (and the\n",
        "# corresponding language codes) that are currently supported by the series of models,\n",
        "# can be found at: https://github.com/facebookresearch/flores/blob/main/flores200/README.md#languages-in-flores-200\n",
        "src_lang = \"ron_Latn\"\n",
        "tgt_lang = \"jpn_Jpan\"\n",
        "\n",
        "instances = [\n",
        "    {\"text\": input_text, \"src_lang\": src_lang, \"tgt_lang\": tgt_lang},\n",
        "]\n",
        "preds = endpoint.predict(instances=instances).predictions\n",
        "print(preds[0].get(\"translation_text\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "69453bf7230e"
      },
      "source": [
        "## Clean up"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2ccf3714dbe9"
      },
      "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_nllb.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}