model_garden/notebooks/community/model_garden/model_garden_gemma_evaluation.ipynb

In [ ]:

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Vertex AI Model Garden - Gemma Evaluation

Run in Colab Enterprise

View on GitHub

Overview

This notebook demonstrates evaluating pre-trained and instruction-tuned Gemma models in Vertex AI.

Objective

• Evaluate pre-trained and instruction-tuned Gemma model on any of the benchmark datasets
• Clean up the resources

Models
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

Costs

This tutorial uses billable components of Google Cloud:

• Vertex AI
• Cloud Storage

Learn about Vertex AI pricing, Cloud Storage pricing, and use the Pricing Calculator to generate a cost estimate based on your projected usage.

Run the notebook

In [ ]:

# @title Setup Google Cloud project
# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).

# @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.

BUCKET_URI = "gs://"  # @param {type:"string"}

# @markdown 3. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.

REGION = ""  # @param {type:"string"}

# Import the necessary packages

! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git

import datetime
import importlib
import os
import uuid

from google.cloud import aiplatform

common_util = importlib.import_module(
    "vertex-ai-samples.notebooks.community.model_garden.docker_source_codes.notebook_util.common_util"
)

models, endpoints = {}, {}

# Get the default cloud project id.
PROJECT_ID = os.environ["GOOGLE_CLOUD_PROJECT"]

# Get the default region for launching jobs.
if not REGION:
    if not os.environ.get("GOOGLE_CLOUD_REGION"):
        raise ValueError(
            "REGION must be set. See"
            " https://cloud.google.com/vertex-ai/docs/general/locations for"
            " available cloud locations."
        )
    REGION = os.environ["GOOGLE_CLOUD_REGION"]

# Enable the Vertex AI API and Compute Engine API, if not already.
print("Enabling Vertex AI API and Compute Engine API.")
! gcloud services enable aiplatform.googleapis.com compute.googleapis.com

# Cloud Storage bucket for storing the experiment artifacts.
# A unique GCS bucket will be created for the purpose of this notebook. If you
# prefer using your own GCS bucket, change the value yourself below.
now = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
BUCKET_NAME = "/".join(BUCKET_URI.split("/")[:3])

if BUCKET_URI is None or BUCKET_URI.strip() == "" or BUCKET_URI == "gs://":
    BUCKET_URI = f"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}"
    BUCKET_NAME = "/".join(BUCKET_URI.split("/")[:3])
    ! gsutil mb -l {REGION} {BUCKET_URI}
else:
    assert BUCKET_URI.startswith("gs://"), "BUCKET_URI must start with `gs://`."
    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep "Location constraint:" | sed "s/Location constraint://"
    bucket_region = shell_output[0].strip().lower()
    if bucket_region != REGION:
        raise ValueError(
            "Bucket region %s is different from notebook region %s"
            % (bucket_region, REGION)
        )
print(f"Using this GCS Bucket: {BUCKET_URI}")

STAGING_BUCKET = os.path.join(BUCKET_URI, "temporal")
MODEL_BUCKET = os.path.join(BUCKET_URI, "gemma")


# Initialize Vertex AI API.
print("Initializing Vertex AI API.")
aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)

# Gets the default SERVICE_ACCOUNT.
shell_output = ! gcloud projects describe $PROJECT_ID
project_number = shell_output[-1].split(":")[1].strip().replace("'", "")
SERVICE_ACCOUNT = f"{project_number}-compute@developer.gserviceaccount.com"
print("Using this default Service Account:", SERVICE_ACCOUNT)


# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket
! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME

! gcloud config set project $PROJECT_ID
! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role="roles/storage.admin"
! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role="roles/aiplatform.user"

In [ ]:

# @title Evaluate Gemma models

# @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.

# @markdown You must provide a Hugging Face User Access Token (with read access) 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.
HF_TOKEN = ""  # @param {type:"string", isTemplate:true}

# @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).
# @markdown Set evaluation dataset.
eval_dataset = "hellaswag"  # @param {type:"string"}

# Worker pool spec.
# Find Vertex AI supported accelerators and regions in:
#  https://cloud.google.com/vertex-ai/docs/training/configure-compute


# Setup evaluation job.
# @markdown Set the base model id.
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}
job_name = common_util.get_job_name_with_datetime(prefix="gemma-eval")
eval_output_dir = os.path.join(MODEL_BUCKET, job_name)
eval_output_dir_gcsfuse = eval_output_dir.replace("gs://", "/gcs/")

# @markdown Set the accelerator type.
accelerator_type = "NVIDIA_L4"  # @param["NVIDIA_TESLA_V100", "NVIDIA_L4", "NVIDIA_TESLA_A100"]

# @markdown To evaluate a PEFT-finetuned model, enter the PEFT output directory to the LoRA adapter below.
# @markdown Otherwise, leave it empty.
# @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.
# @markdown Set the PEFT output directory.
peft_output_dir = ""  # @param {type:"string"}
peft_output_dir_gcsfuse = peft_output_dir.replace("gs://", "/gcs/")

if accelerator_type == "NVIDIA_TESLA_A100":
    machine_type = "a2-highgpu-1g"
    accelerator_count = 1
elif accelerator_type == "NVIDIA_TESLA_V100":
    machine_type = "n1-standard-8"
    accelerator_count = 2
elif accelerator_type == "NVIDIA_L4":
    machine_type = "g2-standard-8"
    accelerator_count = 1
else:
    print(f"Unsupported accelerator type: {accelerator_type}")

replica_count = 1

common_util.check_quota(
    project_id=PROJECT_ID,
    region=REGION,
    accelerator_type=accelerator_type,
    accelerator_count=accelerator_count,
    is_for_training=True,
)

# Prepare evaluation command that runs the evaluation harness.
# Set `trust_remote_code = True` because evaluating the model requires
# executing code from the model repository.
# Set `use_accelerate = True` to enable evaluation across multiple GPUs.
eval_command = [
    "lm_eval",
    "--model",
    "hf",
    "--tasks",
    f"{eval_dataset}",
    "--output_path",
    f"{eval_output_dir_gcsfuse}",
]

if peft_output_dir_gcsfuse:
    eval_command += [
        "--model_args",
        f"pretrained={base_model_id},peft={peft_output_dir_gcsfuse},trust_remote_code=True,parallelize=True",
    ]
else:
    eval_command += [
        "--model_args",
        f"pretrained={base_model_id},trust_remote_code=True,parallelize=True",
    ]


# The evaluation docker image.
EVAL_DOCKER_URI = "us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-lm-evaluation-harness:20241016_0934_RC00"

# Pass evaluation arguments and launch job.
worker_pool_specs = [
    {
        "machine_spec": {
            "machine_type": machine_type,
            "accelerator_type": accelerator_type,
            "accelerator_count": accelerator_count,
        },
        "replica_count": replica_count,
        "disk_spec": {
            "boot_disk_size_gb": 500,
        },
        "container_spec": {
            "image_uri": EVAL_DOCKER_URI,
            "env": [
                {
                    "name": "HF_TOKEN",
                    "value": HF_TOKEN,
                }
            ],
            "command": eval_command,
            "args": [],
        },
    }
]

eval_job = aiplatform.CustomJob(
    display_name=job_name,
    worker_pool_specs=worker_pool_specs,
    base_output_dir=eval_output_dir,
)

eval_job.run()

print("Evaluation results were saved in:", eval_output_dir)

In [ ]:

# @title Fetch and print evaluation results
import json
import re

from google.cloud import storage

# Fetch evaluation results.
storage_client = storage.Client()
BUCKET_NAME = BUCKET_URI.split("gs://")[1]
bucket = storage_client.get_bucket(BUCKET_NAME)

blobs = [b.name for b in bucket.list_blobs()]

result_file_path = None
for file_path in filter(re.compile(".*/*.json").match, blobs):
    result_file_path = file_path
    print(f"Found result file: {file_path}")

if result_file_path is None:
    raise ValueError("No result file found.")

blob = bucket.blob(result_file_path)
raw_result = blob.download_as_string()

# Print evaluation results.
result = json.loads(raw_result)
result_formatted = json.dumps(result, indent=2)
print(f"Evaluation result:\n{result_formatted}")

Clean up resources

In [ ]:

# Delete evaluation job.

delete_bucket = False  # @param {type:"boolean"}
if delete_bucket:
    ! gsutil -m rm -r $BUCKET_URI
    # Uncomment below to delete all artifacts
    # !gsutil -m rm -r $STAGING_BUCKET $MODEL_BUCKET $EXPERIMENT_BUCKET

eval_job.delete()