notebooks/community/model_garden/model_garden_gemma_evaluation.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "language": "python",
      "metadata": {
        "cellView": "form",
        "id": "B8S-yo8qTIcO"
      },
      "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": "MTRywGxLTZfU"
      },
      "source": [
        "# Vertex AI Model Garden - Gemma Evaluation\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_evaluation.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_evaluation.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": "2CXS0vZfT8_7"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates evaluating pre-trained and instruction-tuned Gemma models in Vertex AI.\n",
        "\n",
        "### Objective\n",
        "\n",
        "- Evaluate pre-trained and instruction-tuned Gemma model on any of the benchmark datasets\n",
        "- Clean up the resources\n",
        "\n",
        "| Models |\n",
        "| :- |\n",
        "| [google/gemma-2b](https://huggingface.co/google/gemma-2b)\n",
        "| [google/gemma-2b-it](https://huggingface.co/google/gemma-2b-it)\n",
        "| [google/gemma-7b](https://huggingface.co/google/gemma-7b)\n",
        "| [google/gemma-7b-it](https://huggingface.co/google/gemma-7b-it)\n",
        "| [google/gemma-1.1-2b-it](https://huggingface.co/google/gemma-1.1-2b-it)\n",
        "| [google/gemma-1.1-7b-it](https://huggingface.co/google/gemma-1.1-7b-it)\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": "HCY8PGrFUbT1"
      },
      "source": [
        "## Run the notebook"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "language": "python",
      "metadata": {
        "cellView": "form",
        "id": "81CC3tL1T_TL"
      },
      "outputs": [],
      "source": [
        "# @title Setup Google Cloud project\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. **[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 3. **[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",
        "\n",
        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
        "\n",
        "import datetime\n",
        "import importlib\n",
        "import os\n",
        "import uuid\n",
        "\n",
        "from google.cloud import aiplatform\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",
        "# The evaluation docker image.\n",
        "EVAL_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-lm-evaluation-harness:20240320_0655_RC00\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "language": "python",
      "metadata": {
        "cellView": "form",
        "id": "pNHMbjr0UjrK"
      },
      "outputs": [],
      "source": [
        "# @title Evaluate Gemma models\n",
        "\n",
        "# @markdown This section demonstrates how to evaluate the Gemma models with and without finetuned LoRA adapters using EleutherAI's [Language Model Evaluation Harness (lm-evaluation-harness)](https://github.com/EleutherAI/lm-evaluation-harness) with Vertex CustomJob. Refer the peak GPU memory usage for serving and adjust the machine type, accelerator type and accelerator count accordingly.\n",
        "\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",
        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate:true}\n",
        "\n",
        "# @markdown This example uses the dataset [HellaSwag](https://arxiv.org/abs/1905.07830). All supported tasks are listed in [this task table](https://github.com/EleutherAI/lm-evaluation-harness/blob/master/docs/task_table.md).\n",
        "# @markdown Set evaluation dataset.\n",
        "eval_dataset = \"hellaswag\"  # @param {type:\"string\"}\n",
        "\n",
        "# Worker pool spec.\n",
        "# Find Vertex AI supported accelerators and regions in:\n",
        "#  https://cloud.google.com/vertex-ai/docs/training/configure-compute\n",
        "\n",
        "\n",
        "# Setup evaluation job.\n",
        "# @markdown Set the base model id.\n",
        "base_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",
        "job_name = common_util.get_job_name_with_datetime(prefix=\"gemma-eval\")\n",
        "eval_output_dir = os.path.join(MODEL_BUCKET, job_name)\n",
        "eval_output_dir_gcsfuse = eval_output_dir.replace(\"gs://\", \"/gcs/\")\n",
        "\n",
        "# @markdown Set the accelerator type.\n",
        "accelerator_type = \"NVIDIA_L4\"  # @param[\"NVIDIA_TESLA_V100\", \"NVIDIA_L4\", \"NVIDIA_TESLA_A100\"]\n",
        "\n",
        "# @markdown To evaluate a PEFT-finetuned model, enter the PEFT output directory to the LoRA adapter below.\n",
        "# @markdown Otherwise, leave it empty.\n",
        "# @markdown See the [finetuning notebook](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_gemma_finetuning_on_vertex.ipynb) for more details.\n",
        "# @markdown Set the PEFT output directory.\n",
        "peft_output_dir = \"\"  # @param {type:\"string\"}\n",
        "peft_output_dir_gcsfuse = peft_output_dir.replace(\"gs://\", \"/gcs/\")\n",
        "\n",
        "if accelerator_type == \"NVIDIA_TESLA_A100\":\n",
        "    machine_type = \"a2-highgpu-1g\"\n",
        "    accelerator_count = 1\n",
        "elif accelerator_type == \"NVIDIA_TESLA_V100\":\n",
        "    machine_type = \"n1-standard-8\"\n",
        "    accelerator_count = 2\n",
        "elif accelerator_type == \"NVIDIA_L4\":\n",
        "    machine_type = \"g2-standard-8\"\n",
        "    accelerator_count = 1\n",
        "else:\n",
        "    print(f\"Unsupported accelerator type: {accelerator_type}\")\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",
        "# Prepare evaluation command that runs the evaluation harness.\n",
        "# Set `trust_remote_code = True` because evaluating the model requires\n",
        "# executing code from the model repository.\n",
        "# Set `use_accelerate = True` to enable evaluation across multiple GPUs.\n",
        "eval_command = [\n",
        "    \"lm_eval\",\n",
        "    \"--model\",\n",
        "    \"hf\",\n",
        "    \"--tasks\",\n",
        "    f\"{eval_dataset}\",\n",
        "    \"--output_path\",\n",
        "    f\"{eval_output_dir_gcsfuse}\",\n",
        "]\n",
        "\n",
        "if peft_output_dir_gcsfuse:\n",
        "    eval_command += [\n",
        "        \"--model_args\",\n",
        "        f\"pretrained={base_model_id},peft={peft_output_dir_gcsfuse},trust_remote_code=True,parallelize=True,device_map_option=auto\",\n",
        "    ]\n",
        "else:\n",
        "    eval_command += [\n",
        "        \"--model_args\",\n",
        "        f\"pretrained={base_model_id},trust_remote_code=True,parallelize=True,device_map_option=auto\",\n",
        "    ]\n",
        "\n",
        "# Pass evaluation arguments and launch job.\n",
        "worker_pool_specs = [\n",
        "    {\n",
        "        \"machine_spec\": {\n",
        "            \"machine_type\": machine_type,\n",
        "            \"accelerator_type\": accelerator_type,\n",
        "            \"accelerator_count\": accelerator_count,\n",
        "        },\n",
        "        \"replica_count\": replica_count,\n",
        "        \"disk_spec\": {\n",
        "            \"boot_disk_size_gb\": 500,\n",
        "        },\n",
        "        \"container_spec\": {\n",
        "            \"image_uri\": EVAL_DOCKER_URI,\n",
        "            \"env\": [\n",
        "                {\n",
        "                    \"name\": \"HF_TOKEN\",\n",
        "                    \"value\": HF_TOKEN,\n",
        "                }\n",
        "            ],\n",
        "            \"command\": eval_command,\n",
        "            \"args\": [],\n",
        "        },\n",
        "    }\n",
        "]\n",
        "\n",
        "eval_job = aiplatform.CustomJob(\n",
        "    display_name=job_name,\n",
        "    worker_pool_specs=worker_pool_specs,\n",
        "    base_output_dir=eval_output_dir,\n",
        ")\n",
        "\n",
        "eval_job.run()\n",
        "\n",
        "print(\"Evaluation results were saved in:\", eval_output_dir)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "CVBxGpwWU3kY"
      },
      "outputs": [],
      "source": [
        "# @title Fetch and print evaluation results\n",
        "import json\n",
        "\n",
        "from google.cloud import storage\n",
        "\n",
        "# Fetch evaluation results.\n",
        "storage_client = storage.Client()\n",
        "BUCKET_NAME = BUCKET_URI.split(\"gs://\")[1]\n",
        "bucket = storage_client.get_bucket(BUCKET_NAME)\n",
        "RESULT_FILE_PATH = eval_output_dir[len(BUCKET_URI) + 1 :] + \"/results.json\"\n",
        "blob = bucket.blob(RESULT_FILE_PATH)\n",
        "raw_result = blob.download_as_string()\n",
        "\n",
        "# Print evaluation results.\n",
        "result = json.loads(raw_result)\n",
        "result_formatted = json.dumps(result, indent=2)\n",
        "print(f\"Evaluation result:\\n{result_formatted}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "execution_count": null,
      "metadata": {
        "id": "unjukbcjEBOd"
      },
      "outputs": [],
      "source": [
        "## Clean up resources"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "qWN3cl_VU7pa"
      },
      "outputs": [],
      "source": [
        "# Delete evaluation job.\n",
        "\n",
        "delete_bucket = False  # @param {type:\"boolean\"}\n",
        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI\n",
        "    # Uncomment below to delete all artifacts\n",
        "    # !gsutil -m rm -r $STAGING_BUCKET $MODEL_BUCKET $EXPERIMENT_BUCKET\n",
        "\n",
        "eval_job.delete()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_gemma_evaluation.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}