notebooks/community/model_garden/model_garden_gemma_finetuning_on_vertex.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "language": "python",
      "metadata": {
        "cellView": "form",
        "id": "7d9bbf86da5e"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "language": "markdown",
      "metadata": {
        "id": "99c1c3fc2ca5"
      },
      "source": [
        "# Vertex AI Model Garden - Gemma Finetuning\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_gemma_finetuning_on_vertex.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_gemma_finetuning_on_vertex.ipynb\">\n",
        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</tr></tbody></table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3de7470326a2"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates finetuning and deploying Gemma models with [Vertex AI Custom Training Job](https://cloud.google.com/vertex-ai/docs/training/create-custom-job). All of the examples in this notebook use parameter efficient finetuning methods [PEFT (LoRA)](https://github.com/huggingface/peft) to reduce training and storage costs. LoRA (Low-Rank Adaptation) is one approach of Parameter Efficient FineTuning (PEFT), where pretrained model weights are frozen and rank decomposition matrices representing the change in model weights are trained during finetuning. Read more about LoRA in the following publication: [Hu, E.J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L. and Chen, W., 2021. Lora: Low-rank adaptation of large language models. *arXiv preprint arXiv:2106.09685*](https://arxiv.org/abs/2106.09685).\n",
        "\n",
        "\n",
        "After tuning, we can deploy models on Vertex with GPU or TPU.\n",
        "\n",
        "\n",
        "### Objective\n",
        "\n",
        "- Finetune and deploy Gemma models with Vertex AI Custom Training Jobs.\n",
        "- Send prediction requests to your finetuned Gemma model.\n",
        "\n",
        "\n",
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "264c07757582"
      },
      "source": [
        "## Before you begin"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "BNAlJh_pGxbL"
      },
      "outputs": [],
      "source": [
        "# @title Install Python Packages for Finetuning\n",
        "\n",
        "# @markdown 1. Install google-cloud-aiplatform package and restart the session if instructed.\n",
        "! pip install --upgrade --quiet 'google-cloud-aiplatform>=1.66.0'\n",
        "\n",
        "# @markdown 2. Install packages to validate dataset with template.\n",
        "! pip install --upgrade --quiet accelerate==0.31.0\n",
        "! pip install --upgrade --quiet transformers==4.43.1\n",
        "! pip install --upgrade --quiet datasets==2.19.2\n",
        "\n",
        "# Load local tensorboard.\n",
        "%load_ext tensorboard"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "language": "python",
      "metadata": {
        "cellView": "form",
        "id": "855d6b96f291"
      },
      "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. For finetuning, **[click here](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Frestricted_image_training_nvidia_a100_80gb_gpus)** to check if your project already has the required 8 Nvidia A100 80 GB GPUs in the us-central1 region. If yes, then run this notebook in the us-central1 region. If you do not have 8 Nvidia A100 80 GPUs or have more GPU requirements than this, then schedule your job with Nvidia H100 GPUs via Dynamic Workload Scheduler using [these instructions](https://cloud.google.com/vertex-ai/docs/training/schedule-jobs-dws). For Dynamic Workload Scheduler, check the [us-central1](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Fcustom_model_training_preemptible_nvidia_h100_gpus) or [europe-west4](https://console.cloud.google.com/iam-admin/quotas?location=europe-west4&metric=aiplatform.googleapis.com%2Fcustom_model_training_preemptible_nvidia_h100_gpus) quota for Nvidia H100 GPUs. If you do not have enough GPUs, then you can follow [these instructions](https://cloud.google.com/docs/quotas/view-manage#viewing_your_quota_console) to request quota.\n",
        "\n",
        "# @markdown 3. For serving, **[click here](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_l4_gpus)** to check if your project already has the required 1 L4 GPU in the us-central1 region.  If yes, then run this notebook in the us-central1 region. If you need more L4 GPUs for your project, then you can follow [these instructions](https://cloud.google.com/docs/quotas/view-manage#viewing_your_quota_console) to request more. Alternatively, if you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
        "\n",
        "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
        "# @markdown | ----------- | ----------- | ----------- |\n",
        "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-west1, europe-west4, asia-southeast1 |\n",
        "\n",
        "# @markdown 4. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
        "\n",
        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
        "\n",
        "# @markdown 5. **[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",
        "# Import the necessary packages\n",
        "! rm -rf vertex-ai-samples && git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
        "! cd vertex-ai-samples && git reset --hard 0727e19520cf7957bceb701c248221bd3dbe4f1f\n",
        "\n",
        "import datetime\n",
        "import importlib\n",
        "import os\n",
        "import uuid\n",
        "from typing import Tuple\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "from google.cloud.aiplatform.compat.types import \\\n",
        "    custom_job as gca_custom_job_compat\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",
        "models, endpoints = {}, {}\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",
        "# Cloud Storage bucket for storing the experiment artifacts.\n",
        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
        "# prefer using your own GCS bucket, change the value yourself below.\n",
        "now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "\n",
        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
        "else:\n",
        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
        "    bucket_region = shell_output[0].strip().lower()\n",
        "    if bucket_region != REGION:\n",
        "        raise ValueError(\n",
        "            \"Bucket region %s is different from notebook region %s\"\n",
        "            % (bucket_region, REGION)\n",
        "        )\n",
        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
        "\n",
        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"gemma\")\n",
        "\n",
        "\n",
        "# Initialize Vertex AI API.\n",
        "print(\"Initializing Vertex AI API.\")\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
        "\n",
        "# Gets the default SERVICE_ACCOUNT.\n",
        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
        "\n",
        "\n",
        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
        "! gcloud config set project $PROJECT_ID\n",
        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/storage.admin\"\n",
        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/aiplatform.user\"\n",
        "\n",
        "# @markdown ## Access Gemma Models\n",
        "# @markdown For GPU based finetuning and serving, choose between accessing Gemma models on [Hugging Face](https://huggingface.co/)\n",
        "# @markdown or Vertex AI as described below.\n",
        "\n",
        "# @markdown If you already obtained access to Gemma models on [Hugging Face](https://huggingface.co/), you can load models from there.\n",
        "# @markdown Alternatively, you can also load the original Gemma models for finetuning and serving from Vertex AI after accepting the agreement.\n",
        "\n",
        "# @markdown For TPU based finetuning and serving with KerasNLP, choose the Kaggle option.\n",
        "\n",
        "# @markdown **Select and fill one of the three following sections.**\n",
        "LOAD_MODEL_FROM = \"Hugging Face\"  # @param [\"Hugging Face\", \"Google Cloud\", \"Kaggle\"] {isTemplate:true}\n",
        "\n",
        "# @markdown ---\n",
        "\n",
        "# @markdown ### Access Gemma models on Hugging Face for GPU based finetuning and serving\n",
        "# @markdown You must provide a Hugging Face User Access Token (read) to access the Gemma models. You can follow the [Hugging Face documentation](https://huggingface.co/docs/hub/en/security-tokens) to create a **read** access token and put it in the `HF_TOKEN` field below.\n",
        "\n",
        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate:true}\n",
        "if LOAD_MODEL_FROM == \"Hugging Face\":\n",
        "    assert (\n",
        "        HF_TOKEN\n",
        "    ), \"Provide a read HF_TOKEN to load models from Hugging Face, or select a different model source.\"\n",
        "\n",
        "# @markdown *--- Or ---*\n",
        "# @markdown ### Access Gemma models on Vertex AI for GPU based finetuning and serving\n",
        "# @markdown Accept the model agreement to access the models:\n",
        "# @markdown 1. Open the [Gemma model card](https://console.cloud.google.com/vertex-ai/publishers/google/model-garden/335) from [Vertex AI Model Garden](https://cloud.google.com/model-garden).\n",
        "# @markdown 2. Review the agreement on the model card page.\n",
        "# @markdown 3. After accepting the agreement of Gemma, a `https://` link containing Gemma pretrained and finetuned models will be shared.\n",
        "# @markdown 4. Paste the link in the `VERTEX_MODEL_GARDEN_GEMMA` field below.\n",
        "# @markdown **Note:** This will unzip and copy the Gemma model artifacts to your Cloud Storage bucket, which will take around 1 hour.\n",
        "\n",
        "VERTEX_AI_MODEL_GARDEN_GEMMA = \"\"  # @param {type:\"string\", isTemplate:true}\n",
        "\n",
        "if LOAD_MODEL_FROM == \"Google Cloud\":\n",
        "    assert (\n",
        "        VERTEX_AI_MODEL_GARDEN_GEMMA\n",
        "    ), \"Accept the agreement of Gemma in Vertex AI Model Garden and get the URL to Gemma model artifacts, or select a different model source.\"\n",
        "\n",
        "    # Only use the last part in case a full command is pasted.\n",
        "    signed_url = VERTEX_AI_MODEL_GARDEN_GEMMA.split(\" \")[-1].strip('\"')\n",
        "\n",
        "    ! mkdir -p ./gemma\n",
        "    ! curl -X GET \"{signed_url}\" | tar -xzvf - -C ./gemma/\n",
        "    ! gsutil -m cp -R ./gemma/* {MODEL_BUCKET}\n",
        "\n",
        "    model_path_prefix = MODEL_BUCKET\n",
        "    HF_TOKEN = \"\"\n",
        "else:\n",
        "    model_path_prefix = \"google/\"\n",
        "\n",
        "conversion_job = None\n",
        "\n",
        "# @markdown *--- Or ---*\n",
        "# @markdown ### Access Gemma models from Kaggle for TPU based finetuning and serving\n",
        "# @markdown Kaggle credentials are required for KerasNLP training and Hex-LLM deployment with TPUs.\n",
        "# @markdown Generate the Kaggle username and key by following [these instructions](https://github.com/Kaggle/kaggle-api?tab=readme-ov-file#api-credentials).\n",
        "# @markdown You will need to review and accept the model license.\n",
        "KAGGLE_USERNAME = \"\"  # @param {type:\"string\", isTemplate:true}\n",
        "KAGGLE_KEY = \"\"  # @param {type:\"string\", isTemplate:true}\n",
        "if LOAD_MODEL_FROM == \"Kaggle\":\n",
        "    assert (\n",
        "        KAGGLE_USERNAME and KAGGLE_KEY\n",
        "    ), \"Provide Kaggle credentials to load models from Kaggle, or select a different model source.\"\n",
        "# @markdown ---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cb56d402e84a"
      },
      "source": [
        "## Finetune with HuggingFace PEFT and Deploy with vLLM on GPUs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "KwAW99YZHTdy"
      },
      "outputs": [],
      "source": [
        "# @title Set dataset\n",
        "\n",
        "# @markdown Use the Vertex AI SDK to create and run the custom training jobs.\n",
        "\n",
        "# @markdown This notebook uses [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) dataset as an example.\n",
        "# @markdown You can set `dataset_name` to any existing [Hugging Face dataset](https://huggingface.co/datasets) name, and set `instruct_column_in_dataset` to the name of the dataset column containing training data. The [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) has only one column `text`, and therefore we set `instruct_column_in_dataset` to `text` in this notebook.\n",
        "\n",
        "# @markdown ### (Optional) Prepare a custom JSONL dataset for finetuning\n",
        "\n",
        "# @markdown You can prepare a JSONL file where each line is a valid JSON string as your custom training dataset. For example, here is one line from the [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) dataset:\n",
        "# @markdown ```\n",
        "# @markdown {\"text\": \"### Human: Hola### Assistant: \\u00a1Hola! \\u00bfEn qu\\u00e9 puedo ayudarte hoy?\"}\n",
        "# @markdown ```\n",
        "\n",
        "# @markdown The JSON object has a key `text`, which should match `instruct_column_in_dataset`; The value should be one training data point, i.e. a string. After you prepared your JSONL file, you can either upload it to [Hugging Face datasets](https://huggingface.co/datasets) or [Google Cloud Storage](https://cloud.google.com/storage).\n",
        "\n",
        "# @markdown - To upload a JSONL dataset to [Hugging Face datasets](https://huggingface.co/datasets), follow the instructions on [Uploading Datasets](https://huggingface.co/docs/hub/en/datasets-adding). Then, set `dataset_name` to the name of your newly created dataset on Hugging Face.\n",
        "\n",
        "# @markdown - To upload a JSONL dataset to [Google Cloud Storage](https://cloud.google.com/storage), follow the instructions on [Upload objects from a filesystem](https://cloud.google.com/storage/docs/uploading-objects). Then, set `dataset_name` to the `gs://` URI to your JSONL file. For example: `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`.\n",
        "\n",
        "# @markdown Optionally update the `instruct_column_in_dataset` field below if your JSON objects use a key other than the default `text`.\n",
        "\n",
        "# @markdown ### (Optional) Format your data with custom JSON template\n",
        "\n",
        "# @markdown Sometimes, your dataset might have multiple text columns and you want to construct the training data with a template. You can prepare a JSON template in the following format:\n",
        "\n",
        "# @markdown ```\n",
        "# @markdown {\n",
        "# @markdown   \"description\": \"Template that accepts text-bison format.\",\n",
        "# @markdown   \"source\": \"https://cloud.google.com/vertex-ai/generative-ai/docs/models/tune-text-models-supervised#dataset-format\",\n",
        "# @markdown   \"prompt_input\": \"\\n\\n<|start_header_id|>user<|end_header_id|>\\n\\n{input_text}<|eot_id|>\\n\\n<|start_header_id|>assistant<|end_header_id|>\\n\\n{output_text}<|eot_id|>\",\n",
        "# @markdown   \"instruction_separator\": \"<|start_header_id|>user<|end_header_id|>\\n\\n\",\n",
        "# @markdown   \"response_separator\": \"<|start_header_id|>assistant<|end_header_id|>\\n\\n\"\n",
        "# @markdown }\n",
        "# @markdown ```\n",
        "\n",
        "# @markdown As an example, the template above can be used to format the following training data (this line comes from `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`):\n",
        "\n",
        "# @markdown ```\n",
        "# @markdown {\"input_text\":\"TRANSCRIPT: \\nREASON FOR EVALUATION:,\\n\\n LABEL:\",\"output_text\":\"Chiropractic\"}\n",
        "# @markdown ```\n",
        "\n",
        "# @markdown This example template simply concatenates `input_text` with `output_text` with some special tokens in between.\n",
        "# @markdown\n",
        "# @markdown To try such custom dataset, you can make the following changes:\n",
        "# @markdown 1. Set `template` to `llama3-text-bison`\n",
        "# @markdown 1. Set `train_dataset_name` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`\n",
        "# @markdown 1. Set `train_split_name` to `train`\n",
        "# @markdown 1. Set `eval_dataset_name` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_eval_sample.jsonl`\n",
        "# @markdown 1. Set `eval_split_name` to `train` (**NOT** `test`)\n",
        "# @markdown 1. Set `instruct_column_in_dataset` as `input_text`.\n",
        "\n",
        "# Template name or gs:// URI to a custom template.\n",
        "template = \"openassistant-guanaco\"  # @param {type:\"string\"}\n",
        "\n",
        "# Hugging Face dataset name or gs:// URI to a custom JSONL dataset.\n",
        "train_dataset_name = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
        "train_split_name = \"train\"  # @param {type:\"string\"}\n",
        "eval_dataset_name = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
        "eval_split_name = \"test\"  # @param {type:\"string\"}\n",
        "\n",
        "# Name of the dataset column containing training text input.\n",
        "instruct_column_in_dataset = \"text\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "SdiyOeyFGxbM"
      },
      "outputs": [],
      "source": [
        "# @title Set model\n",
        "\n",
        "# @markdown Select a model variant of Gemma 2.\n",
        "base_model_id = \"gemma-2b\"  # @param[\"gemma-2b\", \"gemma-2b-it\", \"gemma-7b\", \"gemma-7b-it\", \"gemma-1.1-2b-it\", \"gemma-1.1-7b-it\"] {isTemplate:true}\n",
        "pretrained_model_id = os.path.join(model_path_prefix, base_model_id)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "R5PcRc0MGxbM"
      },
      "outputs": [],
      "source": [
        "# @title Validate Dataset with Template\n",
        "\n",
        "# @markdown This section validates the train and eval datasets with the template before starting the fine tuning process.\n",
        "\n",
        "import transformers\n",
        "\n",
        "dataset_validation_util = importlib.import_module(\n",
        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.dataset_validation_util\"\n",
        ")\n",
        "\n",
        "if dataset_validation_util.is_gcs_path(pretrained_model_id):\n",
        "    # Download tokenizer.\n",
        "    ! mkdir tokenizer\n",
        "    ! gsutil cp {pretrained_model_id}/tokenizer.json ./tokenizer\n",
        "    ! gsutil cp {pretrained_model_id}/config.json ./tokenizer\n",
        "    tokenizer_path = \"./tokenizer\"\n",
        "    access_token = \"\"\n",
        "else:\n",
        "    tokenizer_path = pretrained_model_id\n",
        "    access_token = HF_TOKEN\n",
        "\n",
        "tokenizer = transformers.AutoTokenizer.from_pretrained(\n",
        "    tokenizer_path,\n",
        "    trust_remote_code=False,\n",
        "    use_fast=True,\n",
        "    token=access_token,\n",
        ")\n",
        "\n",
        "# Validate the train dataset.\n",
        "dataset_validation_util.validate_dataset_with_template(\n",
        "    dataset_name=train_dataset_name,\n",
        "    split=train_split_name,\n",
        "    input_column=instruct_column_in_dataset,\n",
        "    template=template,\n",
        "    use_multiprocessing=False,\n",
        "    tokenizer=tokenizer,\n",
        ")\n",
        "\n",
        "# Validate the eval dataset.\n",
        "dataset_validation_util.validate_dataset_with_template(\n",
        "    dataset_name=eval_dataset_name,\n",
        "    split=eval_split_name,\n",
        "    input_column=instruct_column_in_dataset,\n",
        "    template=template,\n",
        "    use_multiprocessing=False,\n",
        "    tokenizer=tokenizer,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ivVGS9dHXPOz"
      },
      "outputs": [],
      "source": [
        "# @title Finetune\n",
        "# @markdown This section demonstrates how to finetune the Gemma model and merge the finetuned LoRA adapter with the base model on Vertex AI. It uses the Vertex AI SDK to create and run the custom training jobs.\n",
        "\n",
        "# @markdown The training job takes approximately between 10 to 20 mins to set-up. Once done, the training job is expected to take around 20 mins with the default configuration. To find the training time, throughput, and memory usage of your training job, you can go to the training logs and check the log line of the last training epoch.\n",
        "\n",
        "# @markdown **Note**:\n",
        "# @markdown 1. We recommend setting `finetuning_precision_mode` to `4bit` because it enables using fewer hardware resources for finetuning.\n",
        "# @markdown 1. If `max_steps > 0`, it takes precedence over `epochs`. One can set a small `max_steps` value to quickly check the pipeline.\n",
        "\n",
        "if LOAD_MODEL_FROM == \"Kaggle\":\n",
        "    print(\n",
        "        \"Skipped: Cannot load model from Kaggle, which is only supported in the KerasNLP section.\"\n",
        "    )\n",
        "else:\n",
        "    # @markdown Accelerator type to use for training.\n",
        "    accelerator_type = \"NVIDIA_A100_80GB\"  # @param [\"NVIDIA_A100_80GB\", \"NVIDIA_H100_80GB\"]\n",
        "\n",
        "    # The pre-built training docker image.\n",
        "    if accelerator_type == \"NVIDIA_A100_80GB\":\n",
        "        repo = \"us-docker.pkg.dev/vertex-ai-restricted\"\n",
        "        is_restricted_image = True\n",
        "        is_dynamic_workload_scheduler = False\n",
        "        dws_kwargs = {}\n",
        "    else:\n",
        "        repo = \"us-docker.pkg.dev/vertex-ai\"\n",
        "        is_restricted_image = False\n",
        "        is_dynamic_workload_scheduler = True\n",
        "        dws_kwargs = {\n",
        "            \"max_wait_duration\": 1800,  # 30 minutes\n",
        "            \"scheduling_strategy\": gca_custom_job_compat.Scheduling.Strategy.FLEX_START,\n",
        "        }\n",
        "\n",
        "    TRAIN_DOCKER_URI = (\n",
        "        f\"{repo}/vertex-vision-model-garden-dockers/pytorch-peft-train:stable_20240909\"\n",
        "    )\n",
        "\n",
        "    # Worker pool spec.\n",
        "    if accelerator_type == \"NVIDIA_A100_80GB\":\n",
        "        per_node_accelerator_count = 8\n",
        "        machine_type = \"a2-ultragpu-8g\"\n",
        "    elif accelerator_type == \"NVIDIA_H100_80GB\":\n",
        "        per_node_accelerator_count = 8\n",
        "        machine_type = \"a3-highgpu-8g\"\n",
        "    else:\n",
        "        raise ValueError(\n",
        "            f\"Recommended machine settings not found for: {accelerator_type}. To use another accelerator type, edit this code block to pass in an appropriate `machine_type`, `accelerator_type`, and `per_node_accelerator_count` to the deploy_model_vllm function by clicking `Show Code` and then modifying the code.\"\n",
        "        )\n",
        "\n",
        "    # @markdown Batch size for finetuning.\n",
        "    per_device_train_batch_size = 1  # @param{type:\"integer\"}\n",
        "    # @markdown Number of updates steps to accumulate the gradients for, before performing a backward/update pass.\n",
        "    gradient_accumulation_steps = 4  # @param{type:\"integer\"}\n",
        "    # @markdown Maximum sequence length.\n",
        "    max_seq_length = 4096  # @param{type:\"integer\"}\n",
        "    # @markdown Setting a positive `max_steps` here will override `num_epochs`.\n",
        "    max_steps = -1  # @param{type:\"integer\"}\n",
        "    num_epochs = 1.0  # @param{type:\"number\"}\n",
        "    # @markdown Precision mode for finetuning.\n",
        "    finetuning_precision_mode = \"4bit\"  # @param [\"4bit\", \"8bit\", \"float16\"]\n",
        "    # @markdown Learning rate.\n",
        "    learning_rate = 5e-5  # @param{type:\"number\"}\n",
        "    # @markdown The scheduler type to use.\n",
        "    lr_scheduler_type = \"cosine\"  # @param{type:\"string\"}\n",
        "    # @markdown LoRA parameters.\n",
        "    lora_rank = 16  # @param{type:\"integer\"}\n",
        "    lora_alpha = 32  # @param{type:\"integer\"}\n",
        "    lora_dropout = 0.05  # @param{type:\"number\"}\n",
        "    # Activates gradient checkpointing for the current model (may be referred to as activation checkpointing or checkpoint activations in other frameworks).\n",
        "    enable_gradient_checkpointing = True\n",
        "    # Attention implementation to use in the model.\n",
        "    attn_implementation = \"eager\"\n",
        "    # The optimizer for which to schedule the learning rate.\n",
        "    optimizer = \"paged_adamw_32bit\"\n",
        "    # Define the proportion of training to be dedicated to a linear warmup where learning rate gradually increases.\n",
        "    warmup_ratio = \"0.01\"\n",
        "    # The list or string of integrations to report the results and logs to.\n",
        "    report_to = \"tensorboard\"\n",
        "    # Number of updates steps before two checkpoint saves.\n",
        "    save_steps = 10\n",
        "    # Number of update steps between two logs.\n",
        "    logging_steps = save_steps\n",
        "    # Train precision of the model.\n",
        "    train_precision = \"bfloat16\"\n",
        "\n",
        "    replica_count = 1\n",
        "\n",
        "    common_util.check_quota(\n",
        "        project_id=PROJECT_ID,\n",
        "        region=REGION,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=per_node_accelerator_count * replica_count,\n",
        "        is_for_training=True,\n",
        "        is_restricted_image=is_restricted_image,\n",
        "        is_dynamic_workload_scheduler=is_dynamic_workload_scheduler,\n",
        "    )\n",
        "\n",
        "    job_name = common_util.get_job_name_with_datetime(\"gemma-lora-train\")\n",
        "\n",
        "    base_output_dir = os.path.join(STAGING_BUCKET, job_name)\n",
        "    # Create a GCS folder to store the LORA adapter.\n",
        "    lora_output_dir = os.path.join(base_output_dir, \"adapter\")\n",
        "    # Create a GCS folder to store the merged model with the base model and the\n",
        "    # finetuned LORA adapter.\n",
        "    merged_model_output_dir = os.path.join(base_output_dir, \"merged-model\")\n",
        "\n",
        "    eval_args = [\n",
        "        f\"--eval_dataset_path={eval_dataset_name}\",\n",
        "        f\"--eval_column={instruct_column_in_dataset}\",\n",
        "        f\"--eval_template={template}\",\n",
        "        f\"--eval_split={eval_split_name}\",\n",
        "        f\"--eval_steps={save_steps}\",\n",
        "        \"--eval_tasks=builtin_eval\",\n",
        "        \"--eval_metric_name=loss\",\n",
        "    ]\n",
        "\n",
        "    train_job_args = [\n",
        "        \"--config_file=vertex_vision_model_garden_peft/deepspeed_zero2_8gpu.yaml\",\n",
        "        \"--task=instruct-lora\",\n",
        "        \"--completion_only=True\",\n",
        "        f\"--pretrained_model_id={pretrained_model_id}\",\n",
        "        f\"--dataset_name={train_dataset_name}\",\n",
        "        f\"--train_split_name={train_split_name}\",\n",
        "        f\"--instruct_column_in_dataset={instruct_column_in_dataset}\",\n",
        "        f\"--output_dir={lora_output_dir}\",\n",
        "        f\"--merge_base_and_lora_output_dir={merged_model_output_dir}\",\n",
        "        f\"--per_device_train_batch_size={per_device_train_batch_size}\",\n",
        "        f\"--gradient_accumulation_steps={gradient_accumulation_steps}\",\n",
        "        f\"--lora_rank={lora_rank}\",\n",
        "        f\"--lora_alpha={lora_alpha}\",\n",
        "        f\"--lora_dropout={lora_dropout}\",\n",
        "        f\"--max_steps={max_steps}\",\n",
        "        f\"--max_seq_length={max_seq_length}\",\n",
        "        f\"--learning_rate={learning_rate}\",\n",
        "        f\"--lr_scheduler_type={lr_scheduler_type}\",\n",
        "        f\"--precision_mode={finetuning_precision_mode}\",\n",
        "        f\"--train_precision={train_precision}\",\n",
        "        f\"--enable_gradient_checkpointing={enable_gradient_checkpointing}\",\n",
        "        f\"--num_epochs={num_epochs}\",\n",
        "        f\"--attn_implementation={attn_implementation}\",\n",
        "        f\"--optimizer={optimizer}\",\n",
        "        f\"--warmup_ratio={warmup_ratio}\",\n",
        "        f\"--report_to={report_to}\",\n",
        "        f\"--logging_output_dir={base_output_dir}\",\n",
        "        f\"--save_steps={save_steps}\",\n",
        "        f\"--logging_steps={logging_steps}\",\n",
        "        f\"--template={template}\",\n",
        "        f\"--huggingface_access_token={HF_TOKEN}\",\n",
        "    ] + eval_args\n",
        "\n",
        "    # Pass training arguments and launch job.\n",
        "    train_job = aiplatform.CustomContainerTrainingJob(\n",
        "        display_name=job_name,\n",
        "        container_uri=TRAIN_DOCKER_URI,\n",
        "    )\n",
        "\n",
        "    print(\"Running training job with args:\")\n",
        "    print(\" \\\\\\n\".join(train_job_args))\n",
        "    train_job.run(\n",
        "        args=train_job_args,\n",
        "        replica_count=replica_count,\n",
        "        machine_type=machine_type,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=per_node_accelerator_count,\n",
        "        boot_disk_size_gb=500,\n",
        "        service_account=SERVICE_ACCOUNT,\n",
        "        base_output_dir=base_output_dir,\n",
        "        sync=False,  # Non-blocking call to run.\n",
        "        **dws_kwargs,\n",
        "    )\n",
        "\n",
        "    # Wait until resource has been created.\n",
        "    train_job.wait_for_resource_creation()\n",
        "\n",
        "    print(\"LoRA adapter will be saved in:\", lora_output_dir)\n",
        "    print(\"Trained and merged models will be saved in:\", merged_model_output_dir)\n",
        "\n",
        "# @markdown Click \"Show Code\" to see more details."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "FvJV3FwFGxbM"
      },
      "outputs": [],
      "source": [
        "# @title Run TensorBoard\n",
        "# @markdown This section shows how to launch TensorBoard in a [Cloud Shell](https://cloud.google.com/shell/docs).\n",
        "# @markdown 1. Click the Cloud Shell icon(![terminal](https://github.com/google/material-design-icons/blob/master/png/action/terminal/materialicons/24dp/1x/baseline_terminal_black_24dp.png?raw=true)) on the top right to open the Cloud Shell.\n",
        "# @markdown 2. Copy the `tensorboard` command shown below by running this cell.\n",
        "# @markdown 3. Paste and run the command in the Cloud Shell to launch TensorBoard.\n",
        "# @markdown 4. Once the command runs (You may have to click `Authorize` if prompted), click the link starting with `http://localhost`.\n",
        "\n",
        "# @markdown Note: You may need to wait around 10 minutes after the job starts in order for the TensorBoard logs to be written to the GCS bucket.\n",
        "print(f\"Command to copy: tensorboard --logdir {base_output_dir}/logs\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "qmHW6m8xG_4U"
      },
      "outputs": [],
      "source": [
        "# @title Deploy\n",
        "\n",
        "# @markdown This section uploads the model to Model Registry and deploys it on the Endpoint. It takes 15 minutes to 1 hour to finish.\n",
        "\n",
        "# The pre-built serving docker image for vLLM.\n",
        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240815_1634_RC00\"\n",
        "\n",
        "\n",
        "def deploy_model_vllm(\n",
        "    model_name: str,\n",
        "    model_id: str,\n",
        "    service_account: str,\n",
        "    base_model_id: str = None,\n",
        "    machine_type: str = \"g2-standard-8\",\n",
        "    accelerator_type: str = \"NVIDIA_L4\",\n",
        "    accelerator_count: int = 1,\n",
        "    gpu_memory_utilization: float = 0.9,\n",
        "    max_model_len: int = 4096,\n",
        "    dtype: str = \"auto\",\n",
        "    enable_trust_remote_code: bool = False,\n",
        "    enforce_eager: bool = False,\n",
        "    enable_lora: bool = False,\n",
        "    max_loras: int = 1,\n",
        "    max_cpu_loras: int = 8,\n",
        "    use_dedicated_endpoint: bool = False,\n",
        "    max_num_seqs: int = 256,\n",
        "    model_type: str = None,\n",
        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
        "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
        "    endpoint = aiplatform.Endpoint.create(\n",
        "        display_name=f\"{model_name}-endpoint\",\n",
        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
        "    )\n",
        "\n",
        "    if not base_model_id:\n",
        "        base_model_id = model_id\n",
        "\n",
        "    # See https://docs.vllm.ai/en/latest/models/engine_args.html for a list of possible arguments with descriptions.\n",
        "    vllm_args = [\n",
        "        \"python\",\n",
        "        \"-m\",\n",
        "        \"vllm.entrypoints.api_server\",\n",
        "        \"--host=0.0.0.0\",\n",
        "        \"--port=8080\",\n",
        "        f\"--model={model_id}\",\n",
        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
        "        \"--swap-space=16\",\n",
        "        f\"--gpu-memory-utilization={gpu_memory_utilization}\",\n",
        "        f\"--max-model-len={max_model_len}\",\n",
        "        f\"--dtype={dtype}\",\n",
        "        f\"--max-loras={max_loras}\",\n",
        "        f\"--max-cpu-loras={max_cpu_loras}\",\n",
        "        f\"--max-num-seqs={max_num_seqs}\",\n",
        "        \"--disable-log-stats\",\n",
        "    ]\n",
        "\n",
        "    if enable_trust_remote_code:\n",
        "        vllm_args.append(\"--trust-remote-code\")\n",
        "\n",
        "    if enforce_eager:\n",
        "        vllm_args.append(\"--enforce-eager\")\n",
        "\n",
        "    if enable_lora:\n",
        "        vllm_args.append(\"--enable-lora\")\n",
        "\n",
        "    if model_type:\n",
        "        vllm_args.append(f\"--model-type={model_type}\")\n",
        "\n",
        "    env_vars = {\n",
        "        \"MODEL_ID\": base_model_id,\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_TOKEN\"] = HF_TOKEN\n",
        "    except NameError:\n",
        "        pass\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
        "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
        "        serving_container_args=vllm_args,\n",
        "        serving_container_ports=[8080],\n",
        "        serving_container_predict_route=\"/generate\",\n",
        "        serving_container_health_route=\"/ping\",\n",
        "        serving_container_environment_variables=env_vars,\n",
        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
        "        serving_container_deployment_timeout=7200,\n",
        "    )\n",
        "    print(\n",
        "        f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\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",
        "    print(\"endpoint_name:\", endpoint.name)\n",
        "\n",
        "    return model, endpoint\n",
        "\n",
        "\n",
        "if LOAD_MODEL_FROM == \"Kaggle\":\n",
        "    print(\n",
        "        \"Skipped: Cannot load model from Kaggle, which is only supported in the KerasNLP section.\"\n",
        "    )\n",
        "else:\n",
        "    if train_job.end_time is None:\n",
        "        print(\"Waiting for the training job to finish...\")\n",
        "        train_job.wait()\n",
        "        print(\"The training job has finished.\")\n",
        "\n",
        "    print(\"Deploying models in: \", merged_model_output_dir)\n",
        "\n",
        "    # Find Vertex AI prediction supported accelerators and regions in [here](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute).\n",
        "    # Sets 1 L4 (24G) to deploy Gemma models.\n",
        "    machine_type = \"g2-standard-12\"\n",
        "    accelerator_type = \"NVIDIA_L4\"\n",
        "    accelerator_count = 1\n",
        "\n",
        "    # @markdown Set use_dedicated_endpoint to True if the endpoint is [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint) enabled.\n",
        "    use_dedicated_endpoint = False  # @param {type:\"boolean\"}\n",
        "\n",
        "    common_util.check_quota(\n",
        "        project_id=PROJECT_ID,\n",
        "        region=REGION,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=accelerator_count,\n",
        "        is_for_training=False,\n",
        "    )\n",
        "\n",
        "    # Note that a larger max_model_len will require more GPU memory.\n",
        "    max_model_len = 2048\n",
        "\n",
        "    models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
        "        model_name=common_util.get_job_name_with_datetime(prefix=\"gemma-vllm-serve\"),\n",
        "        base_model_id=f\"google/{base_model_id}\",\n",
        "        model_id=merged_model_output_dir,\n",
        "        service_account=SERVICE_ACCOUNT,\n",
        "        machine_type=machine_type,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=accelerator_count,\n",
        "        max_model_len=max_model_len,\n",
        "        use_dedicated_endpoint=use_dedicated_endpoint,\n",
        "    )\n",
        "    print(\"endpoint_name:\", endpoints[\"vllm_gpu\"].name)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "2UYUNn60G_4U"
      },
      "outputs": [],
      "source": [
        "# @title Predict\n",
        "\n",
        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts.\n",
        "\n",
        "# @markdown Here we use an example from the [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) to show the finetuning outcome:\n",
        "\n",
        "# @markdown ```\n",
        "# @markdown ### Human: How would the Future of AI in 10 Years look?### Assistant: Predicting the future is always a challenging task, but here are some possible ways that AI could evolve over the next 10 years: Continued advancements in deep learning: Deep learning has been one of the main drivers of recent AI breakthroughs, and we can expect continued advancements in this area. This may include improvements to existing algorithms, as well as the development of new architectures that are better suited to specific types of data and tasks. Increased use of AI in healthcare: AI has the potential to revolutionize healthcare, by improving the accuracy of diagnoses, developing new treatments, and personalizing patient care. We can expect to see continued investment in this area, with more healthcare providers and researchers using AI to improve patient outcomes. Greater automation in the workplace: Automation is already transforming many industries, and AI is likely to play an increasingly important role in this process. We can expect to see more jobs being automated, as well as the development of new types of jobs that require a combination of human and machine skills. More natural and intuitive interactions with technology: As AI becomes more advanced, we can expect to see more natural and intuitive ways of interacting with technology. This may include voice and gesture recognition, as well as more sophisticated chatbots and virtual assistants. Increased focus on ethical considerations: As AI becomes more powerful, there will be a growing need to consider its ethical implications. This may include issues such as bias in AI algorithms, the impact of automation on employment, and the use of AI in surveillance and policing. Overall, the future of AI in 10 years is likely to be shaped by a combination of technological advancements, societal changes, and ethical considerations. While there are many exciting possibilities for AI in the future, it will be important to carefully consider its potential impact on society and to work towards ensuring that its benefits are shared fairly and equitably.\n",
        "# @markdown ```\n",
        "\n",
        "# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for 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",
        "if LOAD_MODEL_FROM == \"Kaggle\":\n",
        "    print(\n",
        "        \"Skipped: Cannot load model from Kaggle, which is only supported in the KerasNLP section.\"\n",
        "    )\n",
        "else:\n",
        "    prompt = \"How would the Future of AI in 10 Years look?\"  # @param {type: \"string\"}\n",
        "    max_tokens = 128  # @param {type:\"integer\"}\n",
        "    temperature = 1.0  # @param {type:\"number\"}\n",
        "    top_p = 0.9  # @param {type:\"number\"}\n",
        "    top_k = 1  # @param {type:\"integer\"}\n",
        "\n",
        "    # Overrides max_tokens and top_k parameters during inferences.\n",
        "    # If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`,\n",
        "    # you can reduce the max length, such as set max_tokens as 20.\n",
        "    instances = [\n",
        "        {\n",
        "            \"prompt\": f\"### Human: {prompt}### Assistant: \",\n",
        "            \"max_tokens\": max_tokens,\n",
        "            \"temperature\": temperature,\n",
        "            \"top_p\": top_p,\n",
        "            \"top_k\": top_k,\n",
        "        },\n",
        "    ]\n",
        "    response = endpoints[\"vllm_gpu\"].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": "qujA6HkYyL5Z"
      },
      "source": [
        "## Finetune with KerasNLP PEFT and Deploy with HexLLM on TPUs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "xyfeOl1HyL5Z"
      },
      "outputs": [],
      "source": [
        "# @title Set Dataset\n",
        "\n",
        "# @markdown For the training dataset, you can\n",
        "# @markdown  - either use a pre-built TensorFlow\n",
        "# @markdown    dataset such as this [imdb_reviews](https://www.tensorflow.org/datasets/catalog/imdb_reviews) dataset\n",
        "# @markdown  - or your own dataset in [JSONL format](https://jsonlines.org/examples/) such as this [databricks-dolly-15K](https://huggingface.co/datasets/databricks/databricks-dolly-15k/blob/main/databricks-dolly-15k.jsonl) JSONL file. [See license information here.](https://huggingface.co/datasets/databricks/databricks-dolly-15k#licenseattribution)\n",
        "\n",
        "# @markdown Whether you use a TensorFlow dataset or a JSONL dataset, each data-item\n",
        "# @markdown will be in the form of a dictionary containing multiple key-value pairs. For example,\n",
        "# @markdown [the `imdb_reviews` data-item dictionary](https://www.tensorflow.org/datasets/catalog/imdb_reviews#imdb_reviewsplain_text_default_config)\n",
        "# @markdown contains keys `text` and `label` and [the `databricks-dolly-15k` data-item dictionary](https://huggingface.co/datasets/databricks/databricks-dolly-15k) contains\n",
        "# @markdown keys `instruction`, `context`, `response`, and  `category`.\n",
        "# @markdown However, the Gemma model only takes a single string as a training example. To arbitrarily select and combine\n",
        "# @markdown multiple key-values into a single training string, you can set a `template` variable in the next section.\n",
        "# @markdown For example, for the [`databricks-dolly-15k`](https://huggingface.co/datasets/databricks/databricks-dolly-15k), you can set the\n",
        "# @markdown `template` as `Instruction: {instruction} Response: {response}`\n",
        "# @markdown which will then automatically fill each `instruction` and `response` key-values into this string template\n",
        "# @markdown and generate a single training string which will look like:\n",
        "\n",
        "# @markdown ```\n",
        "# @markdown  Instruction: Why can camels survive for long without water? Response: Camels use the fat in their humps to keep them filled with energy and hydration for long periods of time.\n",
        "# @markdown  ```\n",
        "\n",
        "# @markdown And for the [imdb_reviews](https://www.tensorflow.org/datasets/catalog/imdb_reviews#imdb_reviewsplain_text_default_config) you can set the `template` as `{text}`\n",
        "# @markdown which will then select each `text` key-value as a single training string which will look like:\n",
        "\n",
        "# @markdown ```\n",
        "# @markdown I have been known to fall asleep during films, but this is usually due to a combination of things including, really tired, being warm and comfortable on the sette and having just eaten a lot. However on this occasion I fell asleep because the film was rubbish. The plot development was constant. Constantly slow and boring. Things seemed to happen, but with no explanation of what was causing them or why. I admit, I may have missed part of the film, but i watched the majority of it and everything just seemed to happen of its own accord without any real concern for anything else. I cant recommend this film at all.\n",
        "# @markdown ```\n",
        "\n",
        "# @markdown Set a template suitable for the selected dataset whether TensorFlow Dataset or JSONL format. The following value is\n",
        "# @markdown set for the `databricks-dolly-15k` dataset.\n",
        "template = \"Instruction: {instruction} Response: {response}\"  # @param {type:\"string\"}\n",
        "\n",
        "# @markdown ### Fill only one of the sections below:\n",
        "# @markdown ---\n",
        "\n",
        "# @markdown For example `imdb_reviews`.\n",
        "\n",
        "# The TensorFlow dataset name.\n",
        "tfds_dataset_name = \"\"  # @param {type:\"string\"}\n",
        "# The dataset split to use.\n",
        "tfds_dataset_split = \"train\"  # @param {type:\"string\"}\n",
        "# @markdown ---\n",
        "\n",
        "# @markdown or\n",
        "\n",
        "# @markdown ---\n",
        "\n",
        "# @markdown `jsonl_dataset_file` can be your Cloud Storage path\n",
        "# @markdown such as `<BUCKET_URI>/<path-to-your-jsonl-file>` or a link to an online JSONL file\n",
        "# @markdown in which case the code here will download and then copy the file to `BUCKET_URI`.\n",
        "# @markdown If you want to upload a JSONL dataset to [Google Cloud Storage](https://cloud.google.com/storage) by yourself, then follow the instructions on [Upload objects from a filesystem](https://cloud.google.com/storage/docs/uploading-objects). Then, set `jsonl_dataset_file` to the `gs://` URI to your JSONL file such as `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`.\n",
        "\n",
        "jsonl_dataset_file = \"https://huggingface.co/datasets/databricks/databricks-dolly-15k/resolve/main/databricks-dolly-15k.jsonl\"  # @param {type:\"string\"}\n",
        "# @markdown ---\n",
        "\n",
        "# @markdown Click `Show code` to see more details.\n",
        "\n",
        "assert (\n",
        "    tfds_dataset_name or jsonl_dataset_file\n",
        "), \"Fill in either `tfds_dataset_name` or `jsonl_dataset_file`.\"\n",
        "assert not (\n",
        "    tfds_dataset_name and jsonl_dataset_file\n",
        "), \"Fill in only one of `tfds_dataset_name` or `jsonl_dataset_file`.\"\n",
        "\n",
        "# Download the JSONL dataset.\n",
        "jsonl_dataset_uri_gcsfuse = \"\"\n",
        "if jsonl_dataset_file:\n",
        "    if jsonl_dataset_file.startswith(\"gs://\"):\n",
        "        # Using cloud storage location.\n",
        "        jsonl_dataset_uri = jsonl_dataset_file\n",
        "    else:\n",
        "        # Download the file and copy to cloud storage.\n",
        "        !wget -O dataset.jsonl $jsonl_dataset_file\n",
        "        jsonl_dataset_uri = f\"{BUCKET_URI}/dataset.jsonl\"\n",
        "        print(\"Copying dataset.jsonl to \", jsonl_dataset_uri)\n",
        "        !gsutil cp dataset.jsonl $jsonl_dataset_uri\n",
        "        print(\"JSONL url copied to: \", jsonl_dataset_uri)\n",
        "    jsonl_dataset_uri_gcsfuse = jsonl_dataset_uri.replace(\"gs://\", \"/gcs/\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "s1bR_nk_yL5a"
      },
      "outputs": [],
      "source": [
        "# @title Finetune\n",
        "# @markdown Use the Vertex AI SDK to create and run the custom training jobs.\n",
        "# @markdown The training job uses TPU V3 8 cores and takes around 10 mins to\n",
        "# @markdown finish once it starts running.\n",
        "# @markdown Click `View backing custom job` link in the output of this cell to follow training job progress.\n",
        "# @markdown **Note that to make the training run faster, only a subset of dataset (2000 examples) is used here during fine tuning and the fine tuning runs for just one epoch. To improve the performance of the model, use more training samples, fine tune for more epochs and experiment with increasing the LoRA rank.**\n",
        "# @markdown Click `Show code` to see more details.\n",
        "\n",
        "# The pre-built training docker images for KerasNLP training.\n",
        "KERAS_TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/jax-keras-train-tpu:20240422_0939_RC00\"\n",
        "KERAS_MODEL_CONVERSION_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/jax-keras-model-conversion:20240422_0949_RC00\"\n",
        "\n",
        "if LOAD_MODEL_FROM != \"Kaggle\":\n",
        "    print(\"Skipped: Expect to load model from Kaggle, got\", LOAD_MODEL_FROM)\n",
        "else:\n",
        "    # The Gemma base model.\n",
        "    KAGGLE_MODEL_ID = \"gemma_2b_en\"  # @param[\"gemma_2b_en\", \"gemma_instruct_2b_en\", \"gemma_7b_en\", \"gemma_instruct_7b_en\", \"gemma_1.1_instruct_2b_en\", \"gemma_1.1_instruct_7b_en\"] {isTemplate:true}\n",
        "    # @markdown Set `num_train_subset_samples` as `-1` to use all the training samples.\n",
        "    num_train_subset_samples = 2000  # @param {type:\"integer\"}\n",
        "    # Number of train epochs.\n",
        "    num_epochs = 1  # @param{type:\"integer\"}\n",
        "    # Learning rate.\n",
        "    learning_rate = 5e-5  # @param{type:\"number\"}\n",
        "    # Weight decay.\n",
        "    weight_decay = 0.01  # @param{type:\"number\"}\n",
        "    # Input sequence length. It determines the memory required by the model.\n",
        "    input_sequence_length = 512  # @param{type:\"integer\"}\n",
        "    # LoRA rank.\n",
        "    lora_rank = 4  # @param{type:\"integer\"}\n",
        "    # Batch size for training.\n",
        "    train_batch_size = 2  # @param{type:\"integer\"}\n",
        "    # The KerasNLP checkpoint filename.\n",
        "    # Note: Do not add folder name here.\n",
        "    checkpoint_filename = \"fine_tuned.weights.h5\"  # @param{type:\"string\"}\n",
        "\n",
        "    # Worker pool spec.\n",
        "    machine_type = \"cloud-tpu\"\n",
        "    # NOTE: The models have been test only with 8 cores.\n",
        "    accelerator_type = \"TPU_V3\"\n",
        "    # Number of TPU cores.\n",
        "    accelerator_count = 8\n",
        "    # Set model parallelism related parameters for 8 cores.\n",
        "    model_parallel_batch_dim = 1\n",
        "    model_parallel_model_dim = 8\n",
        "\n",
        "    replica_count = 1\n",
        "\n",
        "    common_util.check_quota(\n",
        "        project_id=PROJECT_ID,\n",
        "        region=REGION,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=accelerator_count,\n",
        "        is_for_training=True,\n",
        "    )\n",
        "\n",
        "    # Setup training job.\n",
        "    job_name = common_util.get_job_name_with_datetime(\"gemma-keras-lora-train\")\n",
        "\n",
        "    # Pass training arguments and launch job.\n",
        "    train_job = aiplatform.CustomContainerTrainingJob(\n",
        "        display_name=job_name,\n",
        "        container_uri=KERAS_TRAIN_DOCKER_URI,\n",
        "    )\n",
        "\n",
        "    # Create a GCS folder to save the finetuned model.\n",
        "    output_folder = os.path.join(BUCKET_URI, job_name)\n",
        "    output_folder_gcsfuse = output_folder.replace(\"gs://\", \"/gcs/\")\n",
        "\n",
        "    train_job.run(\n",
        "        args=[\n",
        "            f\"--model_type={KAGGLE_MODEL_ID}\",\n",
        "            f\"--num_epochs={num_epochs}\",\n",
        "            f\"--learning_rate={learning_rate}\",\n",
        "            f\"--weight_decay={weight_decay}\",\n",
        "            f\"--input_sequence_length={input_sequence_length}\",\n",
        "            f\"--lora_rank={lora_rank}\",\n",
        "            f\"--model_parallel_batch_dim={model_parallel_batch_dim}\",\n",
        "            f\"--model_parallel_model_dim={model_parallel_model_dim}\",\n",
        "            f\"--tfds_dataset_name={tfds_dataset_name}\",\n",
        "            f\"--tfds_dataset_split={tfds_dataset_split}\",\n",
        "            f\"--jsonl_dataset_file={jsonl_dataset_uri_gcsfuse}\",\n",
        "            f\"--template={template}\",\n",
        "            f\"--train_batch_size={train_batch_size}\",\n",
        "            f\"--num_train_subset_samples={num_train_subset_samples}\",\n",
        "            f\"--output_folder={output_folder_gcsfuse}\",\n",
        "            f\"--checkpoint_filename={checkpoint_filename}\",\n",
        "        ],\n",
        "        environment_variables={\n",
        "            \"KAGGLE_USERNAME\": KAGGLE_USERNAME,\n",
        "            \"KAGGLE_KEY\": KAGGLE_KEY,\n",
        "        },\n",
        "        replica_count=replica_count,\n",
        "        machine_type=machine_type,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=accelerator_count,\n",
        "        service_account=SERVICE_ACCOUNT,\n",
        "    )\n",
        "\n",
        "    print(\"Trained model is saved in: \", output_folder)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "h2xEEcgsAQsb"
      },
      "outputs": [],
      "source": [
        "# @title Convert model\n",
        "# @markdown Convert the KerasNLP model checkpoint to Hex-LLM format.\n",
        "# @markdown  Use the Vertex AI SDK to create and run the custom job.\n",
        "# @markdown Click `View backing custom job` link in the output of this cell to follow job progress.\n",
        "# @markdown  The jobs takes around 6 mins to finish.\n",
        "# @markdown  Click `Show code` to see more details.\n",
        "\n",
        "if LOAD_MODEL_FROM != \"Kaggle\":\n",
        "    print(\"Skipped: Expect to load model from Kaggle, got\", LOAD_MODEL_FROM)\n",
        "else:\n",
        "    model_type_to_size = {\n",
        "        \"gemma_2b_en\": \"2b\",\n",
        "        \"gemma_instruct_2b_en\": \"2b\",\n",
        "        \"gemma_7b_en\": \"7b\",\n",
        "        \"gemma_instruct_7b_en\": \"7b\",\n",
        "        \"gemma_1.1_instruct_2b_en\": \"2b\",\n",
        "        \"gemma_1.1_instruct_7b_en\": \"7b\",\n",
        "    }\n",
        "    size = model_type_to_size[KAGGLE_MODEL_ID]\n",
        "    # NOTE: The Hex-LLM serving code looks for model type tag in\n",
        "    # the checkpoint filename.\n",
        "    if \"_2b_\" in KAGGLE_MODEL_ID:\n",
        "        model_type_to_file_suffix = \"_gemma-2b.ckpt\"\n",
        "    else:\n",
        "        assert (\n",
        "            \"_7b_\" in KAGGLE_MODEL_ID\n",
        "        ), \"KAGGLE_MODEL_ID should contain `_2b_` or `_7b_`.\"\n",
        "        model_type_to_file_suffix = \"_gemma-7b.ckpt\"\n",
        "    hexllm_checkpoint_file = \"finetuned_hexllm\" + model_type_to_file_suffix\n",
        "\n",
        "    # Worker pool spec.\n",
        "    machine_type = \"n1-highmem-16\"\n",
        "\n",
        "    replica_count = 1\n",
        "\n",
        "    # Setup training job.\n",
        "    job_name = common_util.get_job_name_with_datetime(\"gemma-keras-model-conversion\")\n",
        "\n",
        "    # Pass training arguments and launch job.\n",
        "    conversion_job = aiplatform.CustomContainerTrainingJob(\n",
        "        display_name=job_name,\n",
        "        container_uri=KERAS_MODEL_CONVERSION_DOCKER_URI,\n",
        "    )\n",
        "\n",
        "    conversion_job.run(\n",
        "        args=[\n",
        "            f\"--weights_file={output_folder_gcsfuse}/{checkpoint_filename}\",\n",
        "            f\"--size={size}\",\n",
        "            f\"--output_file={output_folder_gcsfuse}/{hexllm_checkpoint_file}\",\n",
        "        ],\n",
        "        environment_variables={\n",
        "            \"KAGGLE_USERNAME\": KAGGLE_USERNAME,\n",
        "            \"KAGGLE_KEY\": KAGGLE_KEY,\n",
        "        },\n",
        "        replica_count=replica_count,\n",
        "        machine_type=machine_type,\n",
        "        service_account=SERVICE_ACCOUNT,\n",
        "    )\n",
        "\n",
        "    print(\n",
        "        \"Converted Hexllm checkpoint is saved in: \",\n",
        "        output_folder + \"/\" + hexllm_checkpoint_file,\n",
        "    )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "VJkzxF9tyL5a"
      },
      "outputs": [],
      "source": [
        "# @title Deploy\n",
        "# @markdown This section uploads the model to Model Registry and deploys it on the Endpoint. It takes 15 minutes to 1 hour to finish.\n",
        "# @markdown **Hex-LLM** is a **H**igh-**E**fficiency **L**arge **L**anguage **M**odel (LLM) TPU serving solution built with **XLA**, which is being developed by Google Cloud. This notebook uses TPU v5e machines. Click `Show code` to see more details.\n",
        "\n",
        "# @markdown Region to deploy the model on TPU.\n",
        "TPU_DEPLOYMENT_REGION = \"us-west1\"  # @param [\"us-west1\"] {isTemplate:true}\n",
        "\n",
        "# The pre-built serving docker image for Hex-LLM.\n",
        "HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:deploy\"\n",
        "\n",
        "\n",
        "def deploy_model_hexllm(\n",
        "    model_name: str,\n",
        "    model_id: str,\n",
        "    service_account: str,\n",
        "    base_model_id: str = None,\n",
        "    tensor_parallel_size: int = 1,\n",
        "    machine_type: str = \"ct5lp-hightpu-1t\",\n",
        "    tpu_topology: str = \"1x1\",\n",
        "    hbm_utilization_factor: float = 0.6,\n",
        "    max_running_seqs: int = 256,\n",
        "    max_model_len: int = 4096,\n",
        "    endpoint_id: str = \"\",\n",
        "    min_replica_count: int = 1,\n",
        "    max_replica_count: int = 1,\n",
        "    use_dedicated_endpoint: bool = False,\n",
        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
        "    \"\"\"Deploys models with Hex-LLM on TPU in Vertex AI.\"\"\"\n",
        "    if endpoint_id:\n",
        "        aip_endpoint_name = (\n",
        "            f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_id}\"\n",
        "        )\n",
        "        endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
        "    else:\n",
        "        endpoint = aiplatform.Endpoint.create(\n",
        "            display_name=f\"{model_name}-endpoint\",\n",
        "            location=TPU_DEPLOYMENT_REGION,\n",
        "            dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
        "        )\n",
        "\n",
        "    if not base_model_id:\n",
        "        base_model_id = model_id\n",
        "\n",
        "    if not tensor_parallel_size:\n",
        "        tensor_parallel_size = int(machine_type[-2])\n",
        "\n",
        "    num_hosts = int(tpu_topology.split(\"x\")[0])\n",
        "\n",
        "    hexllm_args = [\n",
        "        \"--host=0.0.0.0\",\n",
        "        \"--port=7080\",\n",
        "        f\"--model={model_id}\",\n",
        "        f\"--tensor_parallel_size={tensor_parallel_size}\",\n",
        "        f\"--num_hosts={num_hosts}\",\n",
        "        f\"--hbm_utilization_factor={hbm_utilization_factor}\",\n",
        "        f\"--max_running_seqs={max_running_seqs}\",\n",
        "        f\"--max_model_len={max_model_len}\",\n",
        "    ]\n",
        "\n",
        "    env_vars = {\n",
        "        \"MODEL_ID\": base_model_id,\n",
        "        \"HEX_LLM_LOG_LEVEL\": \"info\",\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.update({\"HF_TOKEN\": HF_TOKEN})\n",
        "    except:\n",
        "        pass\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
        "        serving_container_image_uri=HEXLLM_DOCKER_URI,\n",
        "        serving_container_command=[\"python\", \"-m\", \"hex_llm.server.api_server\"],\n",
        "        serving_container_args=hexllm_args,\n",
        "        serving_container_ports=[7080],\n",
        "        serving_container_predict_route=\"/generate\",\n",
        "        serving_container_health_route=\"/ping\",\n",
        "        serving_container_environment_variables=env_vars,\n",
        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
        "        serving_container_deployment_timeout=7200,\n",
        "        location=TPU_DEPLOYMENT_REGION,\n",
        "    )\n",
        "\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        machine_type=machine_type,\n",
        "        tpu_topology=tpu_topology if num_hosts > 1 else None,\n",
        "        deploy_request_timeout=1800,\n",
        "        service_account=service_account,\n",
        "        min_replica_count=min_replica_count,\n",
        "        max_replica_count=max_replica_count,\n",
        "    )\n",
        "    return model, endpoint\n",
        "\n",
        "\n",
        "if LOAD_MODEL_FROM != \"Kaggle\":\n",
        "    print(\"Skipped: Expect to load model from Kaggle, got\", LOAD_MODEL_FROM)\n",
        "else:\n",
        "    if \"2b\" in KAGGLE_MODEL_ID:\n",
        "        # Sets ct5lp-hightpu-1t (1 TPU chip) to deploy Gemma 2B models.\n",
        "        machine_type = \"ct5lp-hightpu-1t\"\n",
        "        accelerator_type = \"TPU_V5e\"\n",
        "        # Note: 1 TPU-V5e chip has only 1 core.\n",
        "        accelerator_count = 1\n",
        "    else:\n",
        "        # Sets ct5lp-hightpu-4t (4 TPU chips) to deploy Gemma 7B models.\n",
        "        machine_type = \"ct5lp-hightpu-4t\"\n",
        "        accelerator_type = \"TPU_V5e\"\n",
        "        # Note: 1 TPU-V5e chip has only 1 core.\n",
        "        accelerator_count = 4\n",
        "\n",
        "    common_util.check_quota(\n",
        "        project_id=PROJECT_ID,\n",
        "        region=REGION,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=accelerator_count,\n",
        "        is_for_training=False,\n",
        "    )\n",
        "\n",
        "    # Note that a larger max_num_batched_tokens will require more TPU memory.\n",
        "    max_num_batched_tokens = 11264\n",
        "    # Multiple of tokens for padding alignment. A higher value can reduce\n",
        "    # re-compilation but can also increase the waste in computation.\n",
        "    tokens_pad_multiple = 1024\n",
        "    # Multiple of sequences for padding alignment. A higher value can reduce\n",
        "    # re-compilation but can also increase the waste in computation.\n",
        "    seqs_pad_multiple = 32\n",
        "\n",
        "    print(\"Using model from: \", output_folder)\n",
        "    models[\"hexllm_tpu\"], endpoints[\"hexllm_tpu\"] = deploy_model_hexllm(\n",
        "        model_name=common_util.get_job_name_with_datetime(prefix=\"gemma-serve-hexllm\"),\n",
        "        base_model_id=f\"google/{KAGGLE_MODEL_ID}\",\n",
        "        model_id=output_folder,\n",
        "        service_account=SERVICE_ACCOUNT,\n",
        "        machine_type=machine_type,\n",
        "        max_num_batched_tokens=max_num_batched_tokens,\n",
        "        tokens_pad_multiple=tokens_pad_multiple,\n",
        "        seqs_pad_multiple=seqs_pad_multiple,\n",
        "        use_dedicated_endpoint=use_dedicated_endpoint,\n",
        "    )\n",
        "    print(\"endpoint_name:\", endpoints[\"hexllm_tpu\"].name)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "language": "python",
      "metadata": {
        "cellView": "form",
        "id": "lgJPw3e9yL5a"
      },
      "outputs": [],
      "source": [
        "# @title Predict\n",
        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts based on your `template`.\n",
        "\n",
        "# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for 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",
        "# If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`,\n",
        "# you can reduce the max length, such as set max_tokens as 20.\n",
        "\n",
        "# @markdown **Note:** The following input corresponds to the default `template` set for the `databricks-dolly-15k` which uses `instruction` and `response` keys.\n",
        "# @markdown   If you modify the `template` or use another dataset, then modify the `prompt` accordingly. For example for the `imdb_reviews` dataset  where `template = \"{text}\"`, set `prompt = \"Inception is \"`.\n",
        "\n",
        "if LOAD_MODEL_FROM != \"Kaggle\":\n",
        "    print(\"Skipped: Expect to load model from Kaggle, got\", LOAD_MODEL_FROM)\n",
        "else:\n",
        "    prompt = \"Instruction: What should I do on a trip to Europe? Response: \"  # @param {type: \"string\"}\n",
        "    # Overrides max_tokens and top_k parameters during inferences.\n",
        "    # If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`,\n",
        "    # you can reduce the max length, such as set max_tokens as 20.\n",
        "    max_tokens = 50  # @param {type:\"integer\"}\n",
        "    temperature = 1.0  # @param {type:\"number\"}\n",
        "    top_p = 1.0  # @param {type:\"number\"}\n",
        "    top_k = 1.0  # @param {type:\"number\"}\n",
        "\n",
        "    # @markdown **Note that the first few prompts will take longer to execute.**\n",
        "\n",
        "    instances = [\n",
        "        {\n",
        "            \"prompt\": prompt,\n",
        "            \"max_tokens\": max_tokens,\n",
        "            \"temperature\": temperature,\n",
        "            \"top_p\": top_p,\n",
        "            \"top_k\": top_k,\n",
        "        },\n",
        "    ]\n",
        "    prediction_response = endpoints[\"hexllm_tpu\"].predict(\n",
        "        instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
        "    )\n",
        "\n",
        "    for prediction in prediction_response.predictions:\n",
        "        print(prediction)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "af21a3cff1e0"
      },
      "source": [
        "## Clean up resources"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "911406c1561e"
      },
      "outputs": [],
      "source": [
        "# Delete the train job.\n",
        "train_job.delete()\n",
        "\n",
        "# Delete the conversion job.\n",
        "if conversion_job:\n",
        "    conversion_job.delete()\n",
        "\n",
        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
        "# @markdown  and avoid unnecessary continuous charges that may incur.\n",
        "\n",
        "# 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()\n",
        "\n",
        "delete_bucket = False  # @param {type:\"boolean\"}\n",
        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_NAME"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_gemma_finetuning_on_vertex.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}