{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "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": "JAPoU8Sm5E6e"
      },
      "source": [
        "# Hugging Face DLCs: Serving Gemma 2 with multiple LoRA adapters with Text Generation Inference (TGI) on Vertex AI\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/open-models/serving/vertex_ai_tgi_gemma_multi_lora_adapters_deployment.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fopen-models%2Fserving%2Fvertex_ai_tgi_gemma_multi_lora_adapters_deployment.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/open-models/serving/vertex_ai_tgi_gemma_multi_lora_adapters_deployment.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/serving/vertex_ai_tgi_gemma_multi_lora_adapters_deployment.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/serving/vertex_ai_tgi_gemma_multi_lora_adapters_deployment.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/serving/vertex_ai_tgi_gemma_multi_lora_adapters_deployment.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/serving/vertex_ai_tgi_gemma_multi_lora_adapters_deployment.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/serving/vertex_ai_tgi_gemma_multi_lora_adapters_deployment.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/serving/vertex_ai_tgi_gemma_multi_lora_adapters_deployment.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "84f0f73a0f76"
      },
      "source": [
        "| | |\n",
        "|-|-|\n",
        "| Author(s) |  [Ivan Nardini](https://github.com/inardini)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "> [**Gemma**](https://ai.google.dev/gemma) is a family of lightweight, state-of-the-art open models built from the same research and technology used to create the Gemini models, developed by Google DeepMind and other teams across Google.\n",
        "\n",
        "> [**Hugging Face DLCs**](https://github.com/huggingface/Google-Cloud-Containers) are pre-built and optimized Deep Learning Containers (DLCs) maintained by Hugging Face and Google Cloud teams to simplify environment configuration for your ML workloads.\n",
        "\n",
        "> [**Google Cloud Vertex AI**](https://cloud.google.com/vertex-ai) is a Machine Learning (ML) platform that lets you train and deploy ML models and AI applications, and customize large language models (LLMs) for use in your AI-powered applications.\n",
        "\n",
        "This notebook showcases how to deploy Gemma 2 2B from Hugging Face Hub with multiple LoRA adapters fine-tuned for different purposes such as coding, or SQL using Hugging Face's Text Generation Inference (TGI) Deep Learning Container (DLC) in combination with a [custom handler](https://huggingface.co/docs/inference-endpoints/en/guides/custom_handler#create-custom-inference-handler) on Vertex AI.\n",
        "\n",
        "By the end of this notebook, you will learn how to:\n",
        "\n",
        "- Create a custom handler and test it\n",
        "- Register any LLM from the Hugging Face Hub on Vertex AI\n",
        "- Deploy an LLM on Vertex AI\n",
        "- Send online predictions on Vertex AI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61RBz8LLbxCR"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "80q9vEu5OH59"
      },
      "source": [
        "### (Optional) Set the runtime\n",
        "\n",
        "Depending on your notebook enviroment, consider to set a GPU runtime.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install Vertex AI SDK and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --user --quiet 'torch' 'torchvision' 'torchaudio'\n",
        "%pip install --upgrade --user --quiet 'huggingface_hub[hf_transfer]' 'transformers' 'accelerate>=0.26.0'\n",
        "%pip install --upgrade --user --quiet 'google-cloud-aiplatform' 'crcmod' 'etils'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R5Xep4W9lq-Z"
      },
      "source": [
        "### Restart runtime\n",
        "\n",
        "To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which restarts the current kernel.\n",
        "\n",
        "The restart might take a minute or longer. After it's restarted, continue to the next step."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SbmM4z7FOBpM"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. In Colab or Colab Enterprise, you might see an error message that says \"Your session crashed for an unknown reason.\" This is expected. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you're running this notebook on Google Colab, run the cell below to authenticate your environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NyKGtVQjgx13"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e0vxVy5J3g41"
      },
      "source": [
        "### Set Hugging Face variables"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P5UCCYrV3pEW"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "from etils.epath import Path\n",
        "\n",
        "ROOT_PATH = Path(\".\")\n",
        "TUTORIAL_PATH = ROOT_PATH / \"deploy_gemma_with_multi_lora_adapters_tutorial\"\n",
        "\n",
        "os.environ[\"HF_HOME\"] = str(TUTORIAL_PATH)\n",
        "os.environ[\"HF_HUB_ENABLE_HF_TRANSFER\"] = \"1\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f9af3e57f89a"
      },
      "source": [
        "### Authenticate your Hugging Face account\n",
        "\n",
        "As [`google/gemma-2b`](https://huggingface.co/google/gemma-2-2b) is a gated model, you need to have a Hugging Face Hub account, and accept the Google's usage license for Gemma. Once that's done, you need to generate a new user access token with read-only access so that the weights can be downloaded from the Hub in the Hugging Face DLC for TGI.\n",
        "\n",
        "> Note that the user access token can only be generated via [the Hugging Face Hub UI](https://huggingface.co/settings/tokens/new), where you can either select read-only access to your account, or follow the recommendations and generate a fine-grained token with read-only access to [`google/gemma-2b`](https://huggingface.co/google/gemma-2-2b)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4c31c7272804"
      },
      "source": [
        "Then you can install the `huggingface_hub` that comes with a CLI that will be used for the authentication with the token generated in advance. So that then the token can be safely retrieved via `huggingface_hub.get_token`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8d836e0210fe"
      },
      "outputs": [],
      "source": [
        "from huggingface_hub import interpreter_login\n",
        "\n",
        "interpreter_login()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c71a4314c250"
      },
      "source": [
        "Read more about [Hugging Face Security](https://huggingface.co/docs/hub/en/security), specifically about [Hugging Face User Access Tokens](https://huggingface.co/docs/hub/en/security-tokens)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dy5VO78rzX8c"
      },
      "source": [
        "### Requirements"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R8Zm9y0hxU5O"
      },
      "source": [
        "#### Set Project ID and Location\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable these APIs](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,artifactregistry.googleapis.com).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "# Use the environment variable if the user doesn't provide Project ID.\n",
        "import os\n",
        "\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
        "\n",
        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "\n",
        "PROJECT_NUMBER = !gcloud projects describe {PROJECT_ID} --format=\"get(projectNumber)\"[0]\n",
        "PROJECT_NUMBER = PROJECT_NUMBER[0]\n",
        "\n",
        "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d03b04d07dc0"
      },
      "source": [
        "#### Create a bucket in Google Cloud Storage (GCS)\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as models."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0ca9b8a1bf63"
      },
      "outputs": [],
      "source": [
        "BUCKET_NAME = \"[your-bucket-name]\"  # @param {type:\"string\", isTemplate: true}\n",
        "\n",
        "if BUCKET_NAME == \"[your-bucket-name]\":\n",
        "    raise ValueError(\"A valid BUCKET_NAME needs to be specified\")\n",
        "\n",
        "BUCKET_URI = f\"gs://{BUCKET_NAME}\"\n",
        "os.environ[\"BUCKET_URI\"] = BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c1a68956ee21"
      },
      "source": [
        "Uncomment the `gcloud storage buckets create` command below if you need to create a bucket on GCS."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2ff73fc11e33"
      },
      "outputs": [],
      "source": [
        "!gcloud storage buckets create $BUCKET_URI --project $PROJECT_ID --location=$LOCATION --default-storage-class=STANDARD --uniform-bucket-level-access"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "set_service_account"
      },
      "source": [
        "#### Set Service Account and permissions\n",
        "\n",
        "You will need to have the following IAM roles set:\n",
        "\n",
        "- Vertex AI User (roles/aiplatform.user)\n",
        "- Artifact Registry Reader (roles/artifactregistry.reader)\n",
        "- Storage Object Admin (roles/storage.objectAdmin)\n",
        "\n",
        "For more information about granting roles, see [Manage access](https://cloud.google.com/iam/docs/granting-changing-revoking-access).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VX9tpdtuQI5L"
      },
      "source": [
        "> If you run following commands using Vertex AI Workbench, run directly in the terminal.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ssUJJqXJJHgC"
      },
      "outputs": [],
      "source": [
        "SERVICE_ACCOUNT = f\"{PROJECT_NUMBER}-compute@developer.gserviceaccount.com\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wqOHg5aid6HP"
      },
      "outputs": [],
      "source": [
        "for role in ['aiplatform.user', 'storage.objectAdmin', 'artifactregistry.reader']:\n",
        "\n",
        "    ! gcloud projects add-iam-policy-binding {PROJECT_ID} \\\n",
        "      --member=serviceAccount:{SERVICE_ACCOUNT} \\\n",
        "      --role=roles/{role} --condition=None"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EjebjdiNxe_D"
      },
      "source": [
        "### Initiate Vertex AI SDK\n",
        "\n",
        "Initiate Vertex AI client session."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GReulMr0xjZs"
      },
      "outputs": [],
      "source": [
        "import vertexai\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8XvMcYzm7xmm"
      },
      "source": [
        "### Import libraries\n",
        "\n",
        "Import relevant libraries."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3fXW5p7o7zlo"
      },
      "outputs": [],
      "source": [
        "import gc\n",
        "import json\n",
        "from pprint import pprint as pp\n",
        "\n",
        "from etils import epath\n",
        "from google.cloud import aiplatform\n",
        "from google.cloud.aiplatform import Endpoint, Model\n",
        "from huggingface_hub import get_token, snapshot_download\n",
        "import requests\n",
        "import torch\n",
        "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
        "from transformers.generation import GenerationConfig as TGenerationConfig\n",
        "from vertexai.generative_models import GenerationConfig, GenerativeModel"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E4rjOknk5yvq"
      },
      "source": [
        "### Helpers\n",
        "\n",
        "Define some helpers."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MOoaB5xf51BD"
      },
      "outputs": [],
      "source": [
        "def hf_model_path(hf_home_path: str, model_name: str | None = None) -> str:\n",
        "    \"\"\"Copy model files from the Hugging Face cache.\"\"\"\n",
        "    # Convert and expand the input path\n",
        "    base_path = Path(hf_home_path).expanduser()\n",
        "\n",
        "    if not base_path.exists():\n",
        "        raise FileNotFoundError(f\"Cache path does not exist: {base_path}\")\n",
        "\n",
        "    if model_name:\n",
        "        # Search in specific model directory\n",
        "        model_dir_name = \"models--\" + model_name.replace(\"/\", \"--\")\n",
        "        model_dirs = [base_path / model_dir_name]\n",
        "        if not model_dirs[0].exists():\n",
        "            raise FileNotFoundError(f\"Model directory not found: {model_dirs[0]}\")\n",
        "    else:\n",
        "        # Search all directories (original behavior)\n",
        "        model_dirs = [d for d in base_path.iterdir() if d.is_dir()]\n",
        "\n",
        "    # Find the deepest snapshot directory\n",
        "    snapshot_path = None\n",
        "    for model_dir in model_dirs:\n",
        "        # Look for 'snapshots' directory\n",
        "        snapshots_dir = model_dir / \"snapshots\"\n",
        "        if not snapshots_dir.exists():\n",
        "            continue\n",
        "\n",
        "        # Get the first snapshot folder (usually there's only one)\n",
        "        for snapshot in snapshots_dir.iterdir():\n",
        "            if snapshot.is_dir():\n",
        "                snapshot_path = snapshot\n",
        "                break\n",
        "\n",
        "        if snapshot_path:\n",
        "            break\n",
        "\n",
        "    if not snapshot_path:\n",
        "        error_msg = (\n",
        "            f\"No snapshot directory found for model: {model_name}\"\n",
        "            if model_name\n",
        "            else \"No snapshot directory found in the cache\"\n",
        "        )\n",
        "        raise FileNotFoundError(error_msg)\n",
        "\n",
        "    return str(snapshot_path)\n",
        "\n",
        "\n",
        "def get_cuda_device_names():\n",
        "    \"\"\"A function to get the list of NVIDIA GPUs\"\"\"\n",
        "    if not torch.cuda.is_available():\n",
        "        return None\n",
        "\n",
        "    return [str(i) for i in range(torch.cuda.device_count())]\n",
        "\n",
        "\n",
        "def empty_gpu_ram():\n",
        "    \"\"\"A function to empty the GPU RAM\"\"\"\n",
        "    gc.collect()\n",
        "    torch.cuda.empty_cache()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EdvJRUWRNGHE"
      },
      "source": [
        "## Prepare a custom handler to serve Lora Adapters on Vertex AI\n",
        "\n",
        "Custom Handlers are custom classes in Python that define the pre-processing, inference, and post-processing steps required to run the inference using Hugging Face Pytorch Prediction container on Vertex AI.\n",
        "\n",
        "Think of Custom Handlers as personalized instructions for Hugging Face models. They define how to prepare the input data, run the model, and handle the results. In this sense, Custom Handlers add flexibility. They let you customize how data is prepared and processed, add extra steps, and even build in custom measurements or logging. This means you can tailor the process to your exact needs when the standard setup isn't sufficient.\n",
        "\n",
        "These instructions are stored in a file named `handler.py`. If you need additional dependencies, you can list it in a `requirements.txt` file. The PyTorch container automatically finds and uses these files if they're present.\n",
        "\n",
        "Have a look at [🤗 Serve Anything with Inference Endpoints + Custom Handlers](https://huggingface.co/blog/alvarobartt/serve-anything-inference-endpoints-custom-code) to learn more."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CgPwfG07iMCS"
      },
      "source": [
        "### Test the `handler` locally\n",
        "\n",
        "Before to build the handler module, you can test its coding locally."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "npEogfIziA1O"
      },
      "source": [
        "#### Download Gemma 2 and adapters locally from Hugging Face Hub"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RZ-IsZ_niNdw"
      },
      "outputs": [],
      "source": [
        "# Download Gemma 2 and adapters\n",
        "base_model_id = \"google/gemma-2-2b-it\"\n",
        "sql_adapter_id = \"google-cloud-partnership/gemma-2-2b-it-lora-sql\"\n",
        "magicoder_adapter_id = \"google-cloud-partnership/gemma-2-2b-it-lora-magicoder\"\n",
        "\n",
        "snapshot_download(repo_id=base_model_id, token=get_token())\n",
        "snapshot_download(repo_id=sql_adapter_id, token=get_token())\n",
        "snapshot_download(repo_id=magicoder_adapter_id, token=get_token())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OVLBAtraJkQ_"
      },
      "source": [
        "#### Load Gemma model\n",
        "\n",
        "Load a pre-trained model called \"gemma\" for text generation. You first sets up the text preprocessor, which handles tasks like tokenization (breaking text into words or sub-words). Then you load the actual model itself, optimized for lower memory usage and automatic device placement (likely GPU if available)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QmHijFltikJZ"
      },
      "outputs": [],
      "source": [
        "gemma_path = hf_model_path(TUTORIAL_PATH / \"hub\", base_model_id)\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(gemma_path)\n",
        "\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    gemma_path, low_cpu_mem_usage=True, torch_dtype=torch.float16, device_map=\"auto\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "swrlVm7j99kl"
      },
      "source": [
        "#### Load Gemma's Lora adapters\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zikOj8jHFr0K"
      },
      "outputs": [],
      "source": [
        "sql_adapter_path = hf_model_path(TUTORIAL_PATH / \"hub\", sql_adapter_id)\n",
        "magicoder_adapter_path = hf_model_path(TUTORIAL_PATH / \"hub\", magicoder_adapter_id)\n",
        "\n",
        "model.load_adapter(\n",
        "    sql_adapter_path,\n",
        "    adapter_name=\"sql_adapter\",\n",
        "    is_trainable=False,\n",
        "    device_map=\"auto\",\n",
        "    offload_folder=\"/tmp/offload\",\n",
        ")\n",
        "model.load_adapter(\n",
        "    magicoder_adapter_path,\n",
        "    adapter_name=\"magicoder_adapter\",\n",
        "    is_trainable=False,\n",
        "    device_map=\"auto\",\n",
        "    offload_folder=\"/tmp/offload\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j5Qir_dtJ0-X"
      },
      "source": [
        "#### Prepare the prediction request\n",
        "\n",
        "Prepare the text generation request by creating a dictionary with user prompt and generation parameters.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6g3Mt_FVyaME"
      },
      "outputs": [],
      "source": [
        "requests = []\n",
        "user_prompts = [\n",
        "    \"I have a table called orders with columns order_id (INT), customer_id (INT), order_date (DATE), and total_amount (DECIMAL). I need to find the total revenue generated in the month of October 2023. How can I write a SQL query to achieve this?\",\n",
        "    \"# Context: You have a list of numbers called `my_numbers`.\\n# Question: How do I calculate the sum of all the numbers in `my_numbers` using a built-in function?\\n# Example `my_numbers` list:\\nmy_numbers = [1, 2, 3, 4, 5]\",\n",
        "]\n",
        "\n",
        "for prompt in user_prompts:\n",
        "    requests.append(\n",
        "        {\n",
        "            \"inputs\": prompt,\n",
        "            \"parameters\": {\"max_new_tokens\": 10, \"temperature\": 0.7, \"do_sample\": True},\n",
        "        }\n",
        "    )\n",
        "\n",
        "prediction_request = {\n",
        "    \"instances\": requests,\n",
        "}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nsyWS86Qr3CL"
      },
      "outputs": [],
      "source": [
        "pp(prediction_request, indent=3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7JUYwkwlEf9I"
      },
      "source": [
        "#### Define an LLM-based router\n",
        "\n",
        "To send the right user request to the right adapted model, define a LLM-based router which takes the user input and determines if it's a SQL query generation task (SQL) or a code generation task (CODE)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "72kz_w3-7_R7"
      },
      "outputs": [],
      "source": [
        "router_model_id = \"gemini-2.0-flash\"\n",
        "router_model = GenerativeModel(router_model_id)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0QmsqI73Ejln"
      },
      "outputs": [],
      "source": [
        "def route_prompt(prompt, router_model):\n",
        "    router_prompt = \"\"\"Analyze the following prompt and determine if it's a SQL query generation task (SQL) or a code generation task (CODE).\n",
        "\n",
        "    Guidelines:\n",
        "    - SQL: Database queries, table operations, data retrieval, SQL explanations, database schema questions\n",
        "    - CODE: Programming problems, algorithm implementations, software development tasks, code optimization\n",
        "\n",
        "    Consider these examples:\n",
        "    1. \"Given a users table with columns (id, name, age), show me all users above 25\" -> SQL\n",
        "    2. \"Write a function to reverse a linked list in Python\" -> CODE\n",
        "    3. \"How can I join these two tables and filter by date?\" -> SQL\n",
        "    4. \"Implement a binary search tree in Rust\" -> CODE\n",
        "\n",
        "    User prompt: \"{prompt}\"\n",
        "    Response:\n",
        "    \"\"\"\n",
        "\n",
        "    response_schema = {\n",
        "        \"type\": \"OBJECT\",\n",
        "        \"properties\": {\n",
        "            \"classification\": {\n",
        "                \"type\": \"string\",\n",
        "                \"enum\": [\"SQL\", \"CODE\"],\n",
        "            },\n",
        "        },\n",
        "    }\n",
        "\n",
        "    try:\n",
        "        response = router_model.generate_content(\n",
        "            router_prompt.format(prompt=prompt),\n",
        "            generation_config=GenerationConfig(\n",
        "                response_mime_type=\"application/json\", response_schema=response_schema\n",
        "            ),\n",
        "        ).text\n",
        "        result = json.loads(response)\n",
        "        return result\n",
        "    except:\n",
        "        return {\"classification\": \"CODE\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vxSTwBiLFNaw"
      },
      "outputs": [],
      "source": [
        "for instance in prediction_request[\"instances\"]:\n",
        "    prompt = instance[\"inputs\"]\n",
        "    print(prompt[:55], \"-->\", route_prompt(prompt, router_model))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DH-WU2IAJ4yY"
      },
      "source": [
        "#### Generate prediction\n",
        "\n",
        "Process a list of prediction requests, each containing a text prompt.\n",
        "\n",
        "> Note that depending on your runtime, generating predictions may require ~ 10 mins."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UyKf2v3vlOtY"
      },
      "outputs": [],
      "source": [
        "predictions = []\n",
        "\n",
        "for instance in prediction_request[\"instances\"]:\n",
        "    # Check the prompt\n",
        "    if \"inputs\" not in instance:\n",
        "        raise ValueError(\"The request body must contain the `inputs` key.\")\n",
        "\n",
        "    # Get the adapter label\n",
        "    prompt = instance[\"inputs\"]\n",
        "    prompt_classification = route_prompt(prompt, router_model)[\"classification\"]\n",
        "\n",
        "    # Set the adapter model\n",
        "    if prompt_classification == \"SQL\":\n",
        "        model.set_adapter(\"sql_adapter\")\n",
        "    else:\n",
        "        model.set_adapter(\"magicoder_adapter\")\n",
        "\n",
        "    # Prepare input\n",
        "    messages = [{\"role\": \"user\", \"content\": prompt}]\n",
        "    input_ids = tokenizer.apply_chat_template(messages, return_tensors=\"pt\").to(\n",
        "        model.device\n",
        "    )\n",
        "\n",
        "    # Generate prediction\n",
        "    input_len = input_ids.shape[-1]\n",
        "    with torch.inference_mode():\n",
        "        generation_config = instance.get(\n",
        "            \"parameters\", {\"max_new_tokens\": 10, \"temperature\": 0.7, \"do_sample\": True}\n",
        "        )\n",
        "        generation = model.generate(\n",
        "            input_ids=input_ids,\n",
        "            generation_config=TGenerationConfig(**generation_config),\n",
        "        )\n",
        "        generation = generation[0][input_len:]\n",
        "        response = tokenizer.decode(generation, skip_special_tokens=True).removeprefix(\n",
        "            \"model\\n\"\n",
        "        )\n",
        "        predictions.append(response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zzmvzmxgaiFc"
      },
      "source": [
        "Get the prediction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XlML8-o8lV9X"
      },
      "outputs": [],
      "source": [
        "for prediction in predictions:\n",
        "    print(\"------- Prediction -------\")\n",
        "    print(prediction)\n",
        "    print(\"--------------------------\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Yvh-vsUblb-W"
      },
      "source": [
        "###  Define the `handler.py` module\n",
        "\n",
        "After testing the handler code, you assemble the code in a Python module which defines the custom inference handler for Gemma.\n",
        "\n",
        "You write this handler code to a file named handler.py within the model directory together with the requirements file which contains dependencies for executing handler code."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J0uEcbJSEeKy"
      },
      "outputs": [],
      "source": [
        "serve_uri = epath.Path(BUCKET_URI) / \"serve\"\n",
        "serve_uri.mkdir(parents=True, exist_ok=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WpBfVVbtoA_m"
      },
      "outputs": [],
      "source": [
        "handler_module = '''\n",
        "from typing import Any, Dict, List\n",
        "import torch\n",
        "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
        "from transformers.generation import GenerationConfig as TGenerationConfig\n",
        "from vertexai.generative_models import GenerationConfig, GenerativeModel\n",
        "import json\n",
        "import logging\n",
        "import sys\n",
        "from huggingface_inference_toolkit.logging import logger\n",
        "\n",
        "def route_prompt(prompt, router_model):\n",
        "    router_prompt = \"\"\"Analyze the following prompt and determine if it's a SQL query generation task (SQL) or a code generation task (CODE).\n",
        "\n",
        "    Guidelines:\n",
        "    - SQL: Database queries, table operations, data retrieval, SQL explanations, database schema questions\n",
        "    - CODE: Programming problems, algorithm implementations, software development tasks, code optimization\n",
        "\n",
        "    Consider these examples:\n",
        "    1. \"Given a users table with columns (id, name, age), show me all users above 25\" -> SQL\n",
        "    2. \"Write a function to reverse a linked list in Python\" -> CODE\n",
        "    3. \"How can I join these two tables and filter by date?\" -> SQL\n",
        "    4. \"Implement a binary search tree in Rust\" -> CODE\n",
        "\n",
        "    User prompt: \"{prompt}\"\n",
        "    Response:\n",
        "    \"\"\"\n",
        "\n",
        "    response_schema = {\n",
        "    \"type\": \"OBJECT\",\n",
        "    \"properties\": {\n",
        "        \"classification\": {\n",
        "            \"type\": \"string\",\n",
        "            \"enum\": [\"SQL\", \"CODE\"],\n",
        "        },\n",
        "    }}\n",
        "\n",
        "    try:\n",
        "      response = router_model.generate_content(\n",
        "        router_prompt.format(prompt=prompt),\n",
        "        generation_config=GenerationConfig(\n",
        "        response_mime_type=\"application/json\", response_schema=response_schema\n",
        "        ),\n",
        "      ).text\n",
        "      result = json.loads(response)\n",
        "      return result\n",
        "    except:\n",
        "      return {\"classification\": \"CODE\"}\n",
        "\n",
        "class EndpointHandler:\n",
        "    def __init__(\n",
        "        self,\n",
        "        model_dir: str = \"google/gemma-2-2b-it\",\n",
        "        sql_adapter_id: str = \"/tmp/model/google-cloud-partnership/gemma-2-2b-it-lora-sql\"\n",
        "        magicoder_adapter_id: str = \"/tmp/model/google-cloud-partnership/gemma-2-2b-it-lora-magicoder\"\n",
        "        router_model_id: str = \"gemini-2.0-flash\",\n",
        "        **kwargs: Any,\n",
        "    ) -> None:\n",
        "\n",
        "        self.processor = AutoTokenizer.from_pretrained(model_dir, token=os.getenv(\"HF_TOKEN\"))\n",
        "        self.model = AutoModelForCausalLM.from_pretrained(\n",
        "                    model_dir,\n",
        "                    low_cpu_mem_usage=True,\n",
        "                    torch_dtype=torch.float16,\n",
        "                    device_map=\"auto\",\n",
        "                    token=os.getenv(\"HF_TOKEN\")\n",
        "                )\n",
        "\n",
        "        self.model.load_adapter(sql_adapter_id, adapter_name=\"sql_adapter\", is_trainable=False, device_map=\"auto\", offload_folder=\"/tmp/offload\")\n",
        "        self.model.load_adapter(magicoder_adapter_id, adapter_name=\"magicoder_adapter\", is_trainable=False, device_map=\"auto\", offload_folder=\"/tmp/offload\")\n",
        "        self.router_model = GenerativeModel(router_model_id)\n",
        "\n",
        "    def __call__(self, data: Dict[str, Any]) -> Dict[str, List[Any]]:\n",
        "        logger.info(\"Processing new request\")\n",
        "        predictions = []\n",
        "\n",
        "        for instance in data[\"instances\"]:\n",
        "            # Check the prompt\n",
        "            if \"inputs\" not in instance:\n",
        "                raise ValueError(\"The request body must contain the `inputs` key.\")\n",
        "\n",
        "            # Get the adapter label\n",
        "            logger.info(f'Getting the adapter label for the prompt: {instance[\"inputs\"]}')\n",
        "            prompt = instance[\"inputs\"]\n",
        "            prompt_classification = route_prompt(prompt, self.router_model)[\"classification\"]\n",
        "\n",
        "            # Set the adapter model\n",
        "            logger.info(f'Setting the model to {prompt_classification} adapter')\n",
        "            if prompt_classification == \"SQL\":\n",
        "                self.model.set_adapter(\"sql_adapter\")\n",
        "            else:\n",
        "                self.model.set_adapter(\"magicoder_adapter\")\n",
        "\n",
        "            # Prepare input\n",
        "            logger.info('Preparing the input for the prompt')\n",
        "            messages = [{\"role\": \"user\", \"content\": prompt}]\n",
        "            input_ids = self.tokenizer.apply_chat_template(\n",
        "                messages,\n",
        "                return_tensors=\"pt\"\n",
        "            ).to(self.model.device)\n",
        "\n",
        "            # Generate prediction\n",
        "            logger.info('Generating the prediction')\n",
        "            input_len = input_ids.shape[-1]\n",
        "            with torch.inference_mode():\n",
        "                generation_config = instance.get(\n",
        "                    \"parameters\", {\"temperature\": 0.7, \"do_sample\": True}\n",
        "                )\n",
        "                generation = self.model.generate(\n",
        "                    input_ids=input_ids,\n",
        "                    generation_config=TGenerationConfig(**generation_config),\n",
        "                )\n",
        "                generation = generation[0][input_len:]\n",
        "                response = self.tokenizer.decode(generation, skip_special_tokens=True)\n",
        "                logger.info(f'Generated response: {response[:50]}...')\n",
        "                predictions.append(response)\n",
        "\n",
        "        logger.info(f\"Successfully processed {len(predictions)} instances\")\n",
        "        return {\"predictions\": predictions}\n",
        "'''\n",
        "\n",
        "with serve_uri.joinpath(\"handler.py\").open(\"w\") as f:\n",
        "    f.write(handler_module)\n",
        "f.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UAzfd7q9REBw"
      },
      "source": [
        "### Provide requirements file\n",
        "\n",
        "Save a `requirements.txt` file with handler's dependencies."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AvOHlMKyRJaW"
      },
      "outputs": [],
      "source": [
        "requirements_file = \"\"\"\n",
        "google-cloud-aiplatform\n",
        "\"\"\"\n",
        "\n",
        "with serve_uri.joinpath(\"requirements.txt\").open(\"w\") as f:\n",
        "    f.write(requirements_file)\n",
        "f.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3zOgAT4DpO8Z"
      },
      "source": [
        "### Copy model with custom handler on Cloud Bucket\n",
        "\n",
        "Efficiently upload the model directory to Google Cloud Storage using `gsutil`.\n",
        "\n",
        "Note that `-m` enables multi-threaded uploads for faster transfer, especially for large directories. `-o GSUtil:parallel_composite_upload_threshold=150M` optimizes large file uploads by splitting them into smaller parts for parallel transfer, significantly speeding up the process for files larger than 150MB."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jZ1_2h3nTC-F"
      },
      "outputs": [],
      "source": [
        "! gsutil -o GSUtil:parallel_composite_upload_threshold=150M -m cp -r {str(gemma_path)}/* {str(serve_uri)}\n",
        "! gsutil -o GSUtil:parallel_composite_upload_threshold=150M -m cp -r {str(sql_adapter_path)}/* {str(serve_uri)}/{sql_adapter_id}\n",
        "! gsutil -o GSUtil:parallel_composite_upload_threshold=150M -m cp -r {str(magicoder_adapter_path)}/* {str(serve_uri)}/{magicoder_adapter_id}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ecK96GglLTEb"
      },
      "source": [
        "## Register Google Gemma on Vertex AI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "28d24c673821"
      },
      "source": [
        "To serve Gemma with Pytorch Inference on Vertex AI, you start importing the model on Vertex AI Model Registry, a central repository where you can manage the lifecycle of your ML models on Vertex AI.\n",
        "\n",
        "Before going into the code to upload or import a model on Vertex AI, let's quickly review the arguments provided to the `aiplatform.Model.upload` method:\n",
        "\n",
        "* **`display_name`** is the name that will be shown in the Vertex AI Model Registry.\n",
        "\n",
        "* **`serving_container_image_uri`** is the location of the Hugging Face DLC for Pytorch Inference that will be used for serving the model.\n",
        "\n",
        "* (optional) **`serving_container_ports`** is the port where the Vertex AI endpoint |will be exposed, by default 8080.\n",
        "\n",
        "For more information on the supported `aiplatform.Model.upload` arguments, check [its Python reference](https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform.Model#google_cloud_aiplatform_Model_upload)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ef487ce082f3"
      },
      "outputs": [],
      "source": [
        "model = Model.upload(\n",
        "    display_name=\"google--gemma2-tgi-multi-lora-model\",\n",
        "    artifact_uri=str(serve_uri),\n",
        "    serving_container_image_uri=\"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-pytorch-inference-cu121.2-3.transformers.4-46.ubuntu2204.py311\",\n",
        "    serving_container_ports=[8080],\n",
        "    serving_container_environment_variables={\n",
        "        \"HUGGING_FACE_HUB_TOKEN\": get_token(),\n",
        "    },\n",
        ")\n",
        "model.wait()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c427c0a87016"
      },
      "source": [
        "## Deploy Google Gemma on Vertex AI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ef4cd160dc09"
      },
      "source": [
        "After the model is registered on Vertex AI, you can deploy the model to an endpoint.\n",
        "\n",
        "You need to first deploy a model to an endpoint before that model can be used to serve online predictions. Deploying a model associates physical resources with the model so it can serve online predictions with low latency.\n",
        "\n",
        "Before going into the code to deploy a model to an endpoint, let's quickly review the arguments provided to the `aiplatform.Model.deploy` method:\n",
        "\n",
        "- **`endpoint`** is the endpoint to deploy the model to, which is optional, and by default will be set to the model display name with the `_endpoint` suffix.\n",
        "- **`machine_type`**, **`accelerator_type`** and **`accelerator_count`** are arguments that define which instance to use, and additionally, the accelerator to use and the number of accelerators, respectively. The `machine_type` and the `accelerator_type` are tied together, so you will need to select an instance that supports the accelerator that you are using and vice-versa. More information about the different instances at [Compute Engine Documentation - GPU machine types](https://cloud.google.com/compute/docs/gpus), and about the `accelerator_type` naming at [Vertex AI Documentation - MachineSpec](https://cloud.google.com/vertex-ai/docs/reference/rest/v1/MachineSpec).\n",
        "\n",
        "For more information on the supported `aiplatform.Model.deploy` arguments, you can check [its Python reference](https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform.Model#google_cloud_aiplatform_Model_deploy)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e777427015db"
      },
      "outputs": [],
      "source": [
        "deployed_model = model.deploy(\n",
        "    endpoint=Endpoint.create(display_name=\"google--gemma2-tgi-multi-lora-endpoint\"),\n",
        "    machine_type=\"g2-standard-4\",\n",
        "    accelerator_type=\"NVIDIA_L4\",\n",
        "    accelerator_count=1,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5f2380d3ea60"
      },
      "source": [
        "> Note that the model deployment on Vertex AI can take around 15 to 25 minutes; most of the time being the allocation / reservation of the resources, setting up the network and security, and such."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9b3fd1898241"
      },
      "source": [
        "## Online predictions on Vertex AI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2aa9cee03bd0"
      },
      "source": [
        "Once the model is deployed on Vertex AI, you can run the online predictions using the `aiplatform.Endpoint.predict` method, which will send the requests to the running endpoint in the `/predict` route specified within the container following Vertex AI I/O payload formatting."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "79e393e97ddb"
      },
      "source": [
        "### Via Python"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1a404fd19a08"
      },
      "source": [
        "#### Within the same session\n",
        "\n",
        "To run the online prediction via the Vertex AI SDK, you can simply use the `predict` method."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "28Hy9Ok6tz7J"
      },
      "outputs": [],
      "source": [
        "output = deployed_model.predict(instances=prediction_request[\"instances\"])\n",
        "for prediction in output.predictions:\n",
        "    print(\"------- Prediction -------\")\n",
        "    print(prediction)\n",
        "    print(\"--------------------------\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b28efb473c7e"
      },
      "source": [
        "#### From a different session\n",
        "\n",
        "To run the online prediction from a different session, you can run the following snippet."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e003ab03f30f"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "import requests\n",
        "\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\", isTemplate: true}\n",
        "if PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "\n",
        "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=LOCATION)\n",
        "\n",
        "ENDPOINT_DISPLAY_NAME = \"google--gemma2-tgi-multi-lora-endpoint\"\n",
        "\n",
        "# Iterates over all the Vertex AI Endpoints within the current project and keeps the first match (if any), otherwise set to None\n",
        "ENDPOINT_ID = next(\n",
        "    (\n",
        "        endpoint.name\n",
        "        for endpoint in aiplatform.Endpoint.list()\n",
        "        if endpoint.display_name == ENDPOINT_DISPLAY_NAME\n",
        "    ),\n",
        "    None,\n",
        ")\n",
        "assert ENDPOINT_ID, (\n",
        "    \"`ENDPOINT_ID` is not set, please make sure that the `endpoint_display_name` is correct at \"\n",
        "    f\"https://console.cloud.google.com/vertex-ai/online-prediction/endpoints?project={os.getenv('PROJECT_ID')}\"\n",
        ")\n",
        "\n",
        "# Initiate the endpoint\n",
        "endpoint = aiplatform.Endpoint(\n",
        "    f\"projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}\"\n",
        ")\n",
        "\n",
        "# Set instances\n",
        "instances = []\n",
        "user_prompts = [\n",
        "    \"I have a table called orders with columns order_id (INT), customer_id (INT), order_date (DATE), and total_amount (DECIMAL). I need to find the total revenue generated in the month of October 2023. How can I write a SQL query to achieve this?\",\n",
        "    \"# Context: You have a list of numbers called `my_numbers`.\\n# Question: How do I calculate the sum of all the numbers in `my_numbers` using a built-in function?\\n# Example `my_numbers` list:\\nmy_numbers = [1, 2, 3, 4, 5]\",\n",
        "]\n",
        "\n",
        "for prompt in user_prompts:\n",
        "    instances.append(\n",
        "        {\n",
        "            \"inputs\": prompt,\n",
        "            \"parameters\": {\"temperature\": 0.7, \"do_sample\": True},\n",
        "        }\n",
        "    )\n",
        "\n",
        "# Generate the prediction\n",
        "output = endpoint.predict(instances=instances)\n",
        "for prediction in output.predictions:\n",
        "    print(\"------- Prediction -------\")\n",
        "    print(prediction)\n",
        "    print(\"--------------------------\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7cc0b19ab641"
      },
      "source": [
        "### Via gcloud"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8fed7007c31c"
      },
      "source": [
        "You can also send the requests using the `gcloud` CLI via the `gcloud ai endpoints` command.\n",
        "\n",
        "> Note that, before proceeding, you should either replace the values or set the following environment variables in advance from the Python variables set in the example, as follows:\n",
        ">\n",
        "> ```python\n",
        "> import os\n",
        "> os.environ[\"PROJECT_ID\"] = PROJECT_ID\n",
        "> os.environ[\"LOCATION\"] = LOCATION\n",
        "> os.environ[\"ENDPOINT_NAME\"] = \"google--gemma2-tgi-multi-lora-endpoint\"\n",
        "> ```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7fb1547d409f"
      },
      "outputs": [],
      "source": [
        "%%bash\n",
        "# Get Endpoint ID\n",
        "ENDPOINT_ID=$(gcloud ai endpoints list \\\n",
        "  --project=$PROJECT_ID \\\n",
        "  --region=$LOCATION \\\n",
        "  --filter=\"display_name=$ENDPOINT_NAME\" \\\n",
        "  --format=\"value(name)\" \\\n",
        "  | cut -d'/' -f6)\n",
        "\n",
        "# Generate the prediction\n",
        "echo '{\n",
        "  \"instances\": [\n",
        "   {\n",
        "      \"inputs\":\"I have a table called orders with columns order_id (INT), customer_id (INT), order_date (DATE), and total_amount (DECIMAL). I need to find the total revenue generated in the month of October 2023. How can I write a SQL query to achieve this?\",\n",
        "      \"parameters\":{\n",
        "         \"temperature\":0.7,\n",
        "         \"do_sample\":true\n",
        "      }\n",
        "   },\n",
        "   {\n",
        "      \"inputs\":\"# Context: You have a list of numbers called `my_numbers`.\\n# Question: How do I calculate the sum of all the numbers in `my_numbers` using a built-in function?\\n# Example `my_numbers` list:\\nmy_numbers = [1, 2, 3, 4, 5]\",\n",
        "      \"parameters\":{\n",
        "         \"temperature\":0.7,\n",
        "         \"do_sample\":true\n",
        "      }\n",
        "   }]\n",
        "}' | gcloud ai endpoints predict $ENDPOINT_ID \\\n",
        "  --project=$PROJECT_ID \\\n",
        "  --region=$LOCATION \\\n",
        "  --json-request=\"-\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7e588a65759e"
      },
      "source": [
        "### Via cURL"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cf3de731bcf4"
      },
      "source": [
        "Alternatively, you can also send the requests via `cURL`.\n",
        "\n",
        "> Note that, before proceeding, you should either replace the values or set the following environment variables in advance from the Python variables set in the example, as follows:\n",
        ">\n",
        "> ```python\n",
        "> import os\n",
        "> os.environ[\"PROJECT_ID\"] = PROJECT_ID\n",
        "> os.environ[\"LOCATION\"] = LOCATION\n",
        "> os.environ[\"ENDPOINT_NAME\"] = \"google--gemma2-tgi-multi-lora-endpoint\"\n",
        "> ```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a71da4e1152b"
      },
      "outputs": [],
      "source": [
        "%%bash\n",
        "# Get Endpoint ID\n",
        "ENDPOINT_ID=$(gcloud ai endpoints list \\\n",
        "  --project=$PROJECT_ID \\\n",
        "  --region=$LOCATION \\\n",
        "  --filter=\"display_name=$ENDPOINT_NAME\" \\\n",
        "  --format=\"value(name)\" \\\n",
        "  | cut -d'/' -f6)\n",
        "\n",
        "# Generate the prediction\n",
        "curl -X POST \\\n",
        "    -H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
        "    -H \"Content-Type: application/json\" \\\n",
        "    \"https://${LOCATION}-aiplatform.googleapis.com/v1/projects/${PROJECT_ID}/locations/${LOCATION}/endpoints/${ENDPOINT_ID}:predict\" \\\n",
        "    -d @- <<EOF\n",
        "{\n",
        "    \"instances\": [\n",
        "   {\n",
        "      \"inputs\":\"I have a table called orders with columns order_id (INT), customer_id (INT), order_date (DATE), and total_amount (DECIMAL). I need to find the total revenue generated in the month of October 2023. How can I write a SQL query to achieve this?\",\n",
        "      \"parameters\":{\n",
        "         \"temperature\":0.7,\n",
        "         \"do_sample\":true\n",
        "      }\n",
        "   },\n",
        "   {\n",
        "      \"inputs\":\"# Context: You have a list of numbers called `my_numbers`.\\n# Question: How do I calculate the sum of all the numbers in `my_numbers` using a built-in function?\\n# Example `my_numbers` list:\\nmy_numbers = [1, 2, 3, 4, 5]\",\n",
        "      \"parameters\":{\n",
        "         \"temperature\":0.7,\n",
        "         \"do_sample\":true\n",
        "      }\n",
        "   }\n",
        "]\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f6f17f9aff65"
      },
      "source": [
        "## Cleaning up"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KFOXWLLTvMvJ"
      },
      "outputs": [],
      "source": [
        "delete_endpoint = False\n",
        "delete_model = False\n",
        "delete_bucket = False\n",
        "\n",
        "if delete_endpoint:\n",
        "    deployed_model.undeploy_all()\n",
        "    deployed_model.delete()\n",
        "\n",
        "if delete_endpoint:\n",
        "    delete_model.delete()\n",
        "\n",
        "if delete_bucket:\n",
        "    ! gsutil rm -r {BUCKET_URI}"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "vertex_ai_tgi_gemma_multi_lora_adapters_deployment.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
