notebooks/community/model_garden/model_garden_pytorch_gemma_peft_finetuning_hf.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "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",
      "metadata": {
        "id": "99c1c3fc2ca5"
      },
      "source": [
        "# Vertex AI Model Garden - Gemma Finetuning (PEFT + TGI)\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_pytorch_gemma_peft_finetuning_hf.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_pytorch_gemma_peft_finetuning_hf.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). Using Vertex AI Pipelines is the quickest way to start finetuning Gemma models, while using a Vertex AI Custom Training Job allows for a higher level of customization and control over the finetuning job. All of the examples in this notebook use parameter efficient finetuning methods [PEFT](https://github.com/huggingface/peft) to reduce training and storage costs.\n",
        "\n",
        "This notebook deploys the model with the [Text Generation Inference](https://github.com/huggingface/text-generation-inference) docker and uses [Text moderation APIs](https://cloud.google.com/natural-language/docs/moderating-text) to analyze predictions against a list of safety attributes.\n",
        "\n",
        "\n",
        "### Objective\n",
        "\n",
        "- Finetune and deploy Gemma models with a Vertex AI Custom Training Job.\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",
        "* Cloud NL APIs\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), [Cloud NL API pricing](https://cloud.google.com/natural-language/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": "B9p8QmmcD_OP"
      },
      "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 **[Optional]** Set the GCS BUCKET_URI to store the experiment artifacts, if you want to use your own bucket. **If not set, a unique GCS bucket will be created automatically on your behalf**.\n",
        "\n",
        "import json\n",
        "import os\n",
        "import sys\n",
        "import uuid\n",
        "from datetime import datetime\n",
        "from typing import Tuple\n",
        "\n",
        "from google.cloud import aiplatform, language\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",
        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\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, please change the value yourself below.\n",
        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
        "if not BUCKET_URI.strip() or BUCKET_URI == \"gs://\":\n",
        "    # Create a unique GCS bucket for this notebook if not specified\n",
        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
        "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\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",
        "\n",
        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
        "\n",
        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
        "\n",
        "# Enable Vertex AI, Cloud Compute, and Cloud Language APIs.\n",
        "! gcloud config set project $PROJECT_ID\n",
        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com language.googleapis.com\n",
        "\n",
        "# @markdown ## Access Gemma Models\n",
        "\n",
        "# @markdown Please 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",
        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate:true}\n",
        "assert HF_TOKEN, \"Please provide a read HF_TOKEN to load models from Hugging Face.\"\n",
        "\n",
        "\n",
        "# The pre-built training and serving docker images.\n",
        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20240220_0936_RC01\"\n",
        "TGI_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-hf-tgi-serve:20240220_0936_RC01\"\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user(project_id=PROJECT_ID)\n",
        "\n",
        "\n",
        "def get_job_name_with_datetime(prefix: str) -> str:\n",
        "    \"\"\"Gets the job name with date time when triggering training or deployment\n",
        "    jobs in Vertex AI.\n",
        "    \"\"\"\n",
        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
        "\n",
        "\n",
        "def moderate_text(text: str) -> language.ModerateTextResponse:\n",
        "    \"\"\"Calls Vertex AI APIs to analyze text moderations.\"\"\"\n",
        "    client = language.LanguageServiceClient()\n",
        "    document = language.Document(\n",
        "        content=text,\n",
        "        type_=language.Document.Type.PLAIN_TEXT,\n",
        "    )\n",
        "    return client.moderate_text(document=document)\n",
        "\n",
        "\n",
        "def show_text_moderation(text: str, response: language.ModerateTextResponse) -> None:\n",
        "    \"\"\"Shows text moderation results.\"\"\"\n",
        "    import pandas as pd\n",
        "\n",
        "    def confidence(category: language.ClassificationCategory) -> float:\n",
        "        return category.confidence\n",
        "\n",
        "    columns = [\"category\", \"confidence\"]\n",
        "    categories = sorted(response.moderation_categories, key=confidence, reverse=True)\n",
        "    data = ((category.name, category.confidence) for category in categories)\n",
        "    df = pd.DataFrame(columns=columns, data=data)\n",
        "\n",
        "    print(f\"Text analyzed:\\n{text}\")\n",
        "    print(df.to_markdown(index=False, tablefmt=\"presto\", floatfmt=\".0%\"))\n",
        "\n",
        "\n",
        "def deploy_model_tgi(\n",
        "    model_name: str,\n",
        "    model_id: str,\n",
        "    service_account: str,\n",
        "    machine_type: str = \"g2-standard-8\",\n",
        "    accelerator_type: str = \"NVIDIA_L4\",\n",
        "    accelerator_count: int = 1,\n",
        "    max_input_length: int = 512,\n",
        "    max_total_tokens: int = 2048,\n",
        "    max_batch_prefill_tokens: int = 2048,\n",
        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
        "    \"\"\"Deploys models with TGI on GPU in Vertex AI.\"\"\"\n",
        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
        "\n",
        "    env_vars = {\n",
        "        \"MODEL_ID\": model_id,\n",
        "        \"NUM_SHARD\": f\"{accelerator_count}\",\n",
        "        \"MAX_INPUT_LENGTH\": f\"{max_input_length}\",\n",
        "        \"MAX_TOTAL_TOKENS\": f\"{max_total_tokens}\",\n",
        "        \"MAX_BATCH_PREFILL_TOKENS\": f\"{max_batch_prefill_tokens}\",\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
        "\n",
        "    if HF_TOKEN:\n",
        "        env_vars[\"HUGGING_FACE_HUB_TOKEN\"] = HF_TOKEN\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
        "        serving_container_image_uri=TGI_DOCKER_URI,\n",
        "        serving_container_ports=[80],\n",
        "        serving_container_environment_variables=env_vars,\n",
        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
        "    )\n",
        "\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        machine_type=machine_type,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=accelerator_count,\n",
        "        deploy_request_timeout=1800,\n",
        "        service_account=service_account,\n",
        "    )\n",
        "    return model, endpoint\n",
        "\n",
        "\n",
        "def get_quota(project_id: str, region: str, resource_id: str) -> int:\n",
        "    \"\"\"Returns the quota for a resource in a region. Returns -1 if can not figure out the quota.\"\"\"\n",
        "    service_endpoint = \"aiplatform.googleapis.com\"  # noqa: F841\n",
        "    quota_list_output = !gcloud alpha services quota list --service=$service_endpoint  --consumer=projects/$project_id --filter=\"$service_endpoint/$resource_id\" --format=json\n",
        "    # Use '.s' on the command output because it is an SList type.\n",
        "    quota_data = json.loads(quota_list_output.s)\n",
        "    if len(quota_data) == 0 or \"consumerQuotaLimits\" not in quota_data[0]:\n",
        "        return -1\n",
        "    if (\n",
        "        len(quota_data[0][\"consumerQuotaLimits\"]) == 0\n",
        "        or \"quotaBuckets\" not in quota_data[0][\"consumerQuotaLimits\"][0]\n",
        "    ):\n",
        "        return -1\n",
        "    all_regions_data = quota_data[0][\"consumerQuotaLimits\"][0][\"quotaBuckets\"]\n",
        "    for region_data in all_regions_data:\n",
        "        if (\n",
        "            region_data.get(\"dimensions\")\n",
        "            and region_data[\"dimensions\"][\"region\"] == region\n",
        "        ):\n",
        "            if \"effectiveLimit\" in region_data:\n",
        "                return int(region_data[\"effectiveLimit\"])\n",
        "            else:\n",
        "                return 0\n",
        "    return -1\n",
        "\n",
        "\n",
        "def get_resource_id(accelerator_type: str, is_for_training: bool) -> str:\n",
        "    \"\"\"Returns the resource id for a given accelerator type and the use case.\n",
        "    Args:\n",
        "      accelerator_type: The accelerator type.\n",
        "      is_for_training: Whether the resource is used for training. Set false\n",
        "      for serving use case.\n",
        "    Returns:\n",
        "      The resource id.\n",
        "    \"\"\"\n",
        "    training_accelerator_map = {\n",
        "        \"NVIDIA_TESLA_V100\": \"custom_model_training_nvidia_v100_gpus\",\n",
        "        \"NVIDIA_L4\": \"custom_model_training_nvidia_l4_gpus\",\n",
        "        \"NVIDIA_TESLA_A100\": \"custom_model_training_nvidia_a100_gpus\",\n",
        "        \"NVIDIA_TESLA_T4\": \"custom_model_training_nvidia_t4_gpus\",\n",
        "        \"TPU_V5e\": \"custom_model_training_tpu_v5e\",\n",
        "        \"TPU_V3\": \"custom_model_training_tpu_v3\",\n",
        "    }\n",
        "    serving_accelerator_map = {\n",
        "        \"NVIDIA_TESLA_V100\": \"custom_model_serving_nvidia_v100_gpus\",\n",
        "        \"NVIDIA_L4\": \"custom_model_serving_nvidia_l4_gpus\",\n",
        "        \"NVIDIA_TESLA_A100\": \"custom_model_serving_nvidia_a100_gpus\",\n",
        "        \"NVIDIA_TESLA_T4\": \"custom_model_serving_nvidia_t4_gpus\",\n",
        "        \"TPU_V5e\": \"custom_model_serving_tpu_v5e\",\n",
        "    }\n",
        "    if is_for_training:\n",
        "        if accelerator_type in training_accelerator_map:\n",
        "            return training_accelerator_map[accelerator_type]\n",
        "        else:\n",
        "            raise ValueError(\n",
        "                f\"Could not find accelerator type: {accelerator_type} for training.\"\n",
        "            )\n",
        "    else:\n",
        "        if accelerator_type in serving_accelerator_map:\n",
        "            return serving_accelerator_map[accelerator_type]\n",
        "        else:\n",
        "            raise ValueError(\n",
        "                f\"Could not find accelerator type: {accelerator_type} for serving.\"\n",
        "            )\n",
        "\n",
        "\n",
        "def check_quota(\n",
        "    project_id: str,\n",
        "    region: str,\n",
        "    accelerator_type: str,\n",
        "    accelerator_count: int,\n",
        "    is_for_training: bool,\n",
        "):\n",
        "    \"\"\"Checks if the project and the region has the required quota.\"\"\"\n",
        "    resource_id = get_resource_id(accelerator_type, is_for_training)\n",
        "    quota = get_quota(project_id, region, resource_id)\n",
        "    quota_request_instruction = (\n",
        "        \"Either use \"\n",
        "        \"a different region or request additional quota. Follow \"\n",
        "        \"instructions here \"\n",
        "        \"https://cloud.google.com/docs/quotas/view-manage#requesting_higher_quota\"\n",
        "        \" to check quota in a region or request additional quota for \"\n",
        "        \"your project.\"\n",
        "    )\n",
        "    if quota == -1:\n",
        "        raise ValueError(\n",
        "            f\"\"\"Quota not found for: {resource_id} in {region}.\n",
        "            {quota_request_instruction}\"\"\"\n",
        "        )\n",
        "    if quota < accelerator_count:\n",
        "        raise ValueError(\n",
        "            f\"\"\"Quota not enough for {resource_id} in {region}:\n",
        "            {quota} < {accelerator_count}.\n",
        "            {quota_request_instruction}\"\"\"\n",
        "        )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pq4iF00YG_4T"
      },
      "source": [
        "## Finetune with Vertex AI Custom Training Jobs\n",
        "\n",
        "This section demonstrates how to finetune and deploy Gemma models with PEFT LoRA on Vertex AI Custom Training Jobs. LoRA (Low-Rank Adaptation) is one approach of PEFT (Parameter Efficient FineTuning), 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)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Yv3i0J_GG_4U"
      },
      "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\": \"A short template for vertex sample dataset.\",\n",
        "# @markdown   \"prompt_input\": \"{input_text}{output_text}\",\n",
        "# @markdown   \"prompt_no_input\": \"{input_text}{output_text}\"\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`. You can set `template` to `vertex_sample` to try out this built-in template with the dataset `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`, or build more complicated JSON templates such as [the alpaca example](https://github.com/tloen/alpaca-lora/blob/main/templates/alpaca.json). To use your own JSON template, please [upload it to Google Cloud Storage](https://cloud.google.com/storage/docs/uploading-objects) and put the `gs://` URI in the `template` field below. Leave `instruct_column_in_dataset` as `text`.\n",
        "\n",
        "\n",
        "# Hugging Face dataset name or gs:// URI to a custom JSONL dataset.\n",
        "dataset_name = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
        "\n",
        "# Name of the dataset column containing training text input.\n",
        "instruct_column_in_dataset = \"text\"  # @param {type:\"string\"}\n",
        "\n",
        "# Optional. Template name or gs:// URI to a custom template.\n",
        "template = \"\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tbWkiP3fXN6b"
      },
      "source": [
        "Execute the next cell to run the training job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ivVGS9dHXPOz"
      },
      "outputs": [],
      "source": [
        "# @title Finetune\n",
        "# @markdown Use the Vertex AI SDK to create and run the custom training jobs. It takes around 40 minutes to finetune Gemma-2B for 1000 steps on 1 NVIDIA_TESLA_V100.\n",
        "\n",
        "# @markdown **Note**: To finetune the Gemma 7B models, we recommend setting `finetuning_precision_mode` to `4bit` and using NVIDIA_L4 instead of NVIDIA_TESLA_V100.\n",
        "# @markdown Please click \"Show Code\" to see more details.\n",
        "\n",
        "# The Gemma base model.\n",
        "MODEL_ID = \"google/gemma-1.1-2b-it\"  # @param[\"google/gemma-2b\", \"google/gemma-2b-it\", \"google/gemma-7b\", \"google/gemma-7b-it\", \"google/gemma-1.1-2b-it\", \"google/gemma-1.1-7b-it\"] {isTemplate:true}\n",
        "# The accelerator to use.\n",
        "\n",
        "# fmt: off\n",
        "ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"  # @param[\"NVIDIA_TESLA_V100\", \"NVIDIA_L4\", \"NVIDIA_TESLA_A100\"]\n",
        "# fmt: on\n",
        "\n",
        "# Batch size for finetuning.\n",
        "per_device_train_batch_size = 1  # @param{type:\"integer\"}\n",
        "# Runs 10 training steps as a minimal example.\n",
        "max_steps = 1000  # @param {type:\"integer\"}\n",
        "# Precision mode for finetuning.\n",
        "finetuning_precision_mode = \"float16\"  # @param[\"4bit\", \"8bit\", \"float16\"]\n",
        "# Learning rate.\n",
        "learning_rate = 2e-4  # @param{type:\"number\"}\n",
        "# LoRA parameters.\n",
        "lora_rank = 16  # @param{type:\"integer\"}\n",
        "lora_alpha = 64  # @param{type:\"integer\"}\n",
        "lora_dropout = 0.1  # @param{type:\"number\"}\n",
        "\n",
        "# Worker pool spec.\n",
        "# Gemma only supports single-GPU finetuning.\n",
        "accelerator_count = 1\n",
        "machine_type = None\n",
        "if \"7b\" in MODEL_ID:\n",
        "    assert (\n",
        "        ACCELERATOR_TYPE != \"NVIDIA_TESLA_V100\"\n",
        "    ), \"Finetuning Gemma 7B models is not supported on NVIDIA_TESLA_V100. Try a GPU with more VRAM.\"\n",
        "    if ACCELERATOR_TYPE == \"NVIDIA_L4\":\n",
        "        assert (\n",
        "            finetuning_precision_mode == \"4bit\"\n",
        "        ), f\"{finetuning_precision_mode} finetuning of Gemma 7B models is not supported on NVIDIA_L4. Try a GPU with more VRAM or use a different precision mode.\"\n",
        "\n",
        "if ACCELERATOR_TYPE == \"NVIDIA_TESLA_V100\":\n",
        "    machine_type = \"n1-standard-8\"\n",
        "elif ACCELERATOR_TYPE == \"NVIDIA_L4\":\n",
        "    machine_type = \"g2-standard-12\"\n",
        "elif ACCELERATOR_TYPE == \"NVIDIA_TESLA_A100\":\n",
        "    machine_type = \"a2-highgpu-1g\"\n",
        "\n",
        "assert (\n",
        "    machine_type\n",
        "), f\"Cannot automatically determine machine type from {ACCELERATOR_TYPE}.\"\n",
        "\n",
        "replica_count = 1\n",
        "\n",
        "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 = get_job_name_with_datetime(\"gemma-lora-train\")\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",
        "# Create a GCS folder to store the LORA adapter.\n",
        "lora_adapter_dir = get_job_name_with_datetime(\"gemma-lora-adapter\")\n",
        "lora_output_dir = os.path.join(STAGING_BUCKET, lora_adapter_dir)\n",
        "\n",
        "# Create a GCS folder to store the merged model with the base model and the\n",
        "# finetuned LORA adapter.\n",
        "merged_model_dir = get_job_name_with_datetime(\"gemma-merged-model\")\n",
        "merged_model_output_dir = os.path.join(STAGING_BUCKET, merged_model_dir)\n",
        "\n",
        "train_job.run(\n",
        "    args=[\n",
        "        \"--task=instruct-lora\",\n",
        "        f\"--pretrained_model_id={MODEL_ID}\",\n",
        "        f\"--dataset_name={dataset_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\"--lora_rank={lora_rank}\",\n",
        "        f\"--lora_alpha={lora_alpha}\",\n",
        "        f\"--lora_dropout={lora_dropout}\",\n",
        "        f\"--max_steps={max_steps}\",\n",
        "        \"--max_seq_length=512\",\n",
        "        f\"--learning_rate={learning_rate}\",\n",
        "        f\"--precision_mode={finetuning_precision_mode}\",\n",
        "        f\"--template={template}\",\n",
        "        f\"--huggingface_access_token={HF_TOKEN}\",\n",
        "    ],\n",
        "    environment_variables={\"WANDB_DISABLED\": True},\n",
        "    replica_count=replica_count,\n",
        "    machine_type=machine_type,\n",
        "    accelerator_type=ACCELERATOR_TYPE,\n",
        "    accelerator_count=accelerator_count,\n",
        "    boot_disk_size_gb=500,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        ")\n",
        "\n",
        "print(\"LoRA adapter was saved in: \", lora_output_dir)\n",
        "print(\"Trained and merged models were saved in: \", merged_model_output_dir)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "qmHW6m8xG_4U"
      },
      "outputs": [],
      "source": [
        "# @title Deploy with TGI\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 Please click \"Show Code\" to see more details.\n",
        "\n",
        "\n",
        "print(\"Deploying models in: \", merged_model_output_dir)\n",
        "\n",
        "# Please finds 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",
        "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 larger token counts will require more GPU memory.\n",
        "max_input_length = 512\n",
        "max_total_tokens = 1024\n",
        "max_batch_prefill_tokens = 2048\n",
        "\n",
        "model, endpoint = deploy_model_tgi(\n",
        "    model_name=get_job_name_with_datetime(prefix=\"gemma-tgi-serve\"),\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_input_length=max_input_length,\n",
        "    max_total_tokens=max_total_tokens,\n",
        "    max_batch_prefill_tokens=max_batch_prefill_tokens,\n",
        ")\n",
        "print(\"endpoint_name:\", endpoint.name)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "2UYUNn60G_4U"
      },
      "outputs": [],
      "source": [
        "# @title Predict\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 Please click \"Show Code\" to see more details.\n",
        "\n",
        "# Loads an existing endpoint instance using the endpoint name:\n",
        "# - Using `endpoint_name = endpoint.name` allows us to get the\n",
        "#   endpoint name of the endpoint `endpoint` created in the cell\n",
        "#   above.\n",
        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
        "#   an existing endpoint with the ID 1234567890123456789.\n",
        "# You may uncomment the code below to load an existing endpoint.\n",
        "\n",
        "# endpoint_name = \"\"  # @param {type:\"string\"}\n",
        "# aip_endpoint_name = (\n",
        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
        "# )\n",
        "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
        "\n",
        "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",
        "# Overides 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",
        "        \"inputs\": f\"### Human: {prompt}### Assistant: \",\n",
        "        \"parameters\": {\n",
        "            \"max_new_tokens\": max_tokens,\n",
        "            \"temperature\": temperature,\n",
        "            \"top_p\": top_p,\n",
        "            \"top_k\": top_k,\n",
        "        },\n",
        "    },\n",
        "]\n",
        "response = endpoint.predict(instances=instances)\n",
        "\n",
        "for prediction in response.predictions:\n",
        "    print(prediction)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "2T_cXYJhG_4U"
      },
      "outputs": [],
      "source": [
        "# @markdown Text moderation analyzes a document against a list of safety attributes, which include \"harmful categories\" and topics that may be considered sensitive.\n",
        "# @markdown Please click \"Show Code\" to see more details.\n",
        "\n",
        "for generated_text in response.predictions:\n",
        "    # Send a request to the API.\n",
        "    response = moderate_text(generated_text)\n",
        "    # Show the results.\n",
        "    show_text_moderation(generated_text, response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "af21a3cff1e0"
      },
      "source": [
        "## Clean up resources"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "911406c1561e"
      },
      "outputs": [],
      "source": [
        "# Delete the train job.\n",
        "train_job.delete()\n",
        "\n",
        "# Undeploy model and delete endpoint.\n",
        "endpoint.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
        "model.delete()\n",
        "\n",
        "# Delete Cloud Storage bucket that was created.\n",
        "if BUCKET_URI == f\"gs://{PROJECT_ID}-tmp-{now}\":\n",
        "    ! gsutil -m rm -r $STAGING_BUCKET"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_pytorch_gemma_peft_finetuning_hf.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}