Formatting and refactoring of Detectron2 Notebook notebook

PiperOrigin-RevId: 772946944

.cloud-build/execute_changed_notebooks_helper.py

@@ -238,7 +238,7 @@ def _get_notebook_python_version(notebook_path: str) -> str:
 
             # Look for the python version specification pattern
             re_match = re.search(
-                "python version = (\d+\.\d+)", markdown, flags=re.IGNORECASE
+                r"python version = (\d+\.\d+)", markdown, flags=re.IGNORECASE
             )
             if re_match:
                 # get the version number

.cloud-build/test_notebook_vm.txt

@@ -1,5 +1,3 @@
-notebooks/official/vizier/gapic-vizier-multi-objective-optimization.ipynb
-notebooks/official/pipelines/lightweight_functions_component_io_kfp.ipynb
 notebooks/official/ml_metadata/sdk-metric-parameter-tracking-for-locally-trained-models.ipynb
 notebooks/official/custom/custom-tabular-bq-managed-dataset.ipynb
 .cloud-build/tests/python_version_test.ipynb

.github/workflows/linter/requirements.txt

@@ -4,7 +4,7 @@ jupyter
 nbconvert
 black==25.1.0
 pyupgrade==3.19.1
-isort==6.0.0
-flake8==7.1.1
+isort==6.0.1
+flake8==7.2.0
 nbqa==1.9.1
 

community-content/vertex_model_garden/model_oss/notebook_util/common_util.py

@@ -1,13 +1,16 @@
 """Common util functions for notebook."""
 
 import base64
+from collections.abc import Sequence
 import datetime
 import io
 import json
 import os
 import subprocess
-from typing import Any, Dict, Sequence
+import time
+from typing import Any
 
+from google import auth
 from google.cloud import storage
 import matplotlib.pyplot as plt
 import numpy as np
@@ -281,7 +284,7 @@ def decode_image(
   return image
 
 
-def get_label_map(label_map_yaml_filepath: str) -> Dict[int, str]:
+def get_label_map(label_map_yaml_filepath: str) -> dict[int, str]:
   """Returns class id to label mapping given a filepath to the label map.
 
   Args:
@@ -331,6 +334,7 @@ def vqa_predict(
     image: Any,
     language_code: str = "en",
     new_width: int = 1000,
+    use_dedicated_endpoint: bool = False,
 ) -> Sequence[str]:
   """Predicts the answer to a question about an image using an Endpoint."""
   # Resize and convert image to base64 string.
@@ -354,7 +358,9 @@ def vqa_predict(
         "image": resized_image_base64,
     })
 
-  response = endpoint.predict(instances=instances)
+  response = endpoint.predict(
+      instances=instances, use_dedicated_endpoint=use_dedicated_endpoint
+  )
   return [pred.get("response") for pred in response.predictions]
 
 
@@ -364,6 +370,7 @@ def caption_predict(
     image: Any,
     caption_prompt: bool = False,
     new_width: int = 1000,
+    use_dedicated_endpoint: bool = False,
 ) -> str:
   """Predicts a caption for a given image using an Endpoint."""
   # Resize and convert image to base64 string.
@@ -378,7 +385,9 @@ def caption_predict(
     instance["prompt"] = caption_prompt_format.format(language_code)
 
   instances = [instance]
-  response = endpoint.predict(instances=instances)
+  response = endpoint.predict(
+      instances=instances, use_dedicated_endpoint=use_dedicated_endpoint
+  )
   return response.predictions[0].get("response")
 
 
@@ -387,6 +396,7 @@ def ocr_predict(
     ocr_prompt: str,
     image: Any,
     new_width: int = 1000,
+    use_dedicated_endpoint: bool = False,
 ) -> str:
   """Extracts text from a given image using an Endpoint."""
   # Resize and convert image to base64 string.
@@ -398,7 +408,9 @@ def ocr_predict(
     instance["prompt"] = ocr_prompt
   instances = [instance]
 
-  response = endpoint.predict(instances=instances)
+  response = endpoint.predict(
+      instances=instances, use_dedicated_endpoint=use_dedicated_endpoint
+  )
   return response.predictions[0].get("response")
 
 
@@ -407,6 +419,7 @@ def detect_predict(
     detect_prompt: str,
     image: Any,
     new_width: int = 1000,
+    use_dedicated_endpoint: bool = False,
 ) -> str:
   """Predicts the answer to a question about an image using an Endpoint."""
   # Resize and convert image to base64 string.
@@ -418,7 +431,9 @@ def detect_predict(
     instance["prompt"] = detect_prompt
   instances = [instance]
 
-  response = endpoint.predict(instances=instances)
+  response = endpoint.predict(
+      instances=instances, use_dedicated_endpoint=use_dedicated_endpoint
+  )
   return response.predictions[0].get("response")
 
 
@@ -495,6 +510,17 @@ def get_quota(project_id: str, region: str, resource_id: str) -> int:
   ):
     return -1
   all_regions_data = quota_data[0]["consumerQuotaLimits"][0]["quotaBuckets"]
+
+  # If the quota data does not have dimensions, it is global quota. However,
+  # global quota may be overridden by regional quota. So we need to check the
+  # global quota first.
+  global_quota = -1
+  if (
+      all_regions_data
+      and "dimensions" not in all_regions_data[0]
+      and "effectiveLimit" in all_regions_data[0]
+  ):
+    global_quota = int(all_regions_data[0]["effectiveLimit"])
   for region_data in all_regions_data:
     if (
         region_data.get("dimensions")
@@ -504,12 +530,13 @@ def get_quota(project_id: str, region: str, resource_id: str) -> int:
         return int(region_data["effectiveLimit"])
       else:
         return 0
-  return -1
+  return global_quota
 
 
 def get_resource_id(
     accelerator_type: str,
     is_for_training: bool,
+    is_spot: bool = False,
     is_restricted_image: bool = False,
     is_dynamic_workload_scheduler: bool = False,
 ) -> str:
@@ -519,6 +546,7 @@ def get_resource_id(
     accelerator_type: The accelerator type.
     is_for_training: Whether the resource is used for training. Set false for
       serving use case.
+    is_spot: Whether the resource is used with Spot.
     is_restricted_image: Whether the image is hosted in `vertex-ai-restricted`.
     is_dynamic_workload_scheduler: Whether the resource is used with Dynamic
       Workload Scheduler.
@@ -533,7 +561,10 @@ def get_resource_id(
       "NVIDIA_TESLA_A100": "nvidia_a100_gpus",
       "NVIDIA_A100_80GB": "nvidia_a100_80gb_gpus",
       "NVIDIA_H100_80GB": "nvidia_h100_gpus",
+      "NVIDIA_H100_MEGA_80GB": "nvidia_h100_mega_gpus",
+      "NVIDIA_H200_141GB": "nvidia_h200_gpus",
       "NVIDIA_TESLA_T4": "nvidia_t4_gpus",
+      "TPU_V6e": "tpu_v6e",
       "TPU_V5e": "tpu_v5e",
       "TPU_V3": "tpu_v3",
   }
@@ -548,6 +579,10 @@ def get_resource_id(
   restricted_image_training_accelerator_map = {
       "NVIDIA_A100_80GB": "restricted_image_training_nvidia_a100_80gb_gpus",
   }
+  spot_serving_accelerator_map = {
+      key: f"custom_model_serving_preemptible_{accelerator_suffix_map[key]}"
+      for key in accelerator_suffix_map
+  }
   serving_accelerator_map = {
       key: f"custom_model_serving_{accelerator_suffix_map[key]}"
       for key in accelerator_suffix_map
@@ -576,8 +611,11 @@ def get_resource_id(
   else:
     if is_dynamic_workload_scheduler:
       raise ValueError("Dynamic Workload Scheduler does not work for serving.")
-    if accelerator_type in serving_accelerator_map:
-      return serving_accelerator_map[accelerator_type]
+    accelerator_map = (
+        spot_serving_accelerator_map if is_spot else serving_accelerator_map
+    )
+    if accelerator_type in accelerator_map:
+      return accelerator_map[accelerator_type]
     else:
       raise ValueError(
           f"Could not find accelerator type: {accelerator_type} for serving."
@@ -590,13 +628,28 @@ def check_quota(
     accelerator_type: str,
     accelerator_count: int,
     is_for_training: bool,
+    is_spot: bool = False,
     is_restricted_image: bool = False,
     is_dynamic_workload_scheduler: bool = False,
-):
-  """Checks if the project and the region has the required quota."""
+) -> None:
+  """Checks if the project and the region has the required quota.
+
+  Args:
+    project_id: The project id.
+    region: The region.
+    accelerator_type: The accelerator type.
+    accelerator_count: The number of accelerators to check quota for.
+    is_for_training: Whether the resource is used for training. Set false for
+      serving use case.
+    is_spot: Whether the resource is used with Spot.
+    is_restricted_image: Whether the image is hosted in `vertex-ai-restricted`.
+    is_dynamic_workload_scheduler: Whether the resource is used with Dynamic
+      Workload Scheduler.
+  """
   resource_id = get_resource_id(
       accelerator_type,
       is_for_training=is_for_training,
+      is_spot=is_spot,
       is_restricted_image=is_restricted_image,
       is_dynamic_workload_scheduler=is_dynamic_workload_scheduler,
   )
@@ -620,3 +673,75 @@ def check_quota(
         f" {accelerator_count}. {quota_request_instruction}"
     )
 
+
+def get_deploy_source() -> str:
+  """Gets deploy_source string based on running environment."""
+  vertex_product = os.environ.get("VERTEX_PRODUCT", "")
+  match vertex_product:
+    case "COLAB_ENTERPRISE":
+      return "notebook_colab_enterprise"
+    case "WORKBENCH_INSTANCE":
+      return "notebook_workbench"
+    case _:
+      # Legacy workbench, legacy colab, or other custom environments.
+      return "notebook_environment_unspecified"
+
+
+def _is_operation_done(op_name: str, region: str) -> bool:
+  """Checks if the operation is done.
+
+  Args:
+    op_name: The name of the operation to poll.
+    region: The region of the operation.
+
+  Returns:
+    True if the operation is done, False otherwise.
+
+  Raises:
+    ValueError: If the operation failed.
+  """
+  creds, _ = auth.default()
+  auth_req = auth.transport.requests.Request()
+  creds.refresh(auth_req)
+  headers = {
+      "Authorization": f"Bearer {creds.token}",
+  }
+  url = f"https://{region}-aiplatform.googleapis.com/ui/{op_name}"
+  response = requests.get(url, headers=headers)
+  operation_data = response.json()
+  if "error" in operation_data:
+    raise ValueError(f"Operation failed: {operation_data['error']}")
+  return operation_data.get("done", False)
+
+
+def poll_and_wait(
+    op_name: str, region: str, total_wait: int, interval: int = 60
+) -> None:
+  """Polls the operation and waits for it to complete.
+
+  Args:
+    op_name: The name of the operation to poll.
+    region: The region of the operation.
+    total_wait: The total wait time in seconds.
+    interval: The interval between each poll in seconds.
+
+  Raises:
+    TimeoutError: If the operation times out.
+  """
+  start_time = time.time()
+  while True:
+    if _is_operation_done(op_name, region):
+      break
+    time_elapsed = time.time() - start_time
+    if time_elapsed > total_wait:
+      raise TimeoutError(
+          f"Operation timed out after {int(time_elapsed)} seconds."
+      )
+    print(
+        "\rStill waiting for operation... Elapsed time in seconds:"
+        f" {int(time_elapsed):<6}",
+        end="",
+        flush=True,
+    )
+    time.sleep(interval)
+

community-content/vertex_model_garden/model_oss/notebook_util/dataset_validation_util.py

@@ -7,7 +7,7 @@ import json
 import multiprocessing
 import os
 import subprocess
-from typing import Any, Callable, Dict, Union
+from typing import Any, Callable, Dict, Tuple, Union
 from absl import logging
 import accelerate
 import datasets
@@ -70,7 +70,9 @@ def force_gcs_fuse_path(gcs_uri: str) -> str:
 
 
 def download_gcs_uri_to_local(
-    gcs_uri: str, destination_dir: str = LOCAL_BASE_MODEL_DIR
+    gcs_uri: str,
+    destination_dir: str = LOCAL_BASE_MODEL_DIR,
+    check_path_exists: bool = True,
 ) -> str:
   """Downloads GCS URI to local.
 
@@ -81,6 +83,7 @@ def download_gcs_uri_to_local(
   Args:
     gcs_uri: GCS URI to download.
     destination_dir: Local directory directory.
+    check_path_exists: Whether to check if the path exists.
 
   Returns:
     Local path to target folder/file.
@@ -89,7 +92,7 @@ def download_gcs_uri_to_local(
       destination_dir,
       os.path.basename(os.path.normpath(gcs_uri)),
   )
-  if os.path.exists(target):
+  if check_path_exists and os.path.exists(target):
     logging.info("File %s already exists.", target)
     return target
   if accelerate.PartialState().is_local_main_process:
@@ -415,13 +418,42 @@ def get_filtered_dataset(
   return filtered_dataset
 
 
+def format_dataset(
+    dataset: datasets.Dataset,
+    input_column: str,
+    template: str = None,
+    tokenizer: transformers.PreTrainedTokenizer | None = None,
+) -> datasets.Dataset:
+  """Takes a raw dataset and formats it using a template and tokenizer.
+
+  Args:
+    dataset: The raw (unprocessed) dataset to format.
+    input_column: The input column in the dataset to be used or updaded by the
+      template. If it does not exist, the template's `prompt_no_input` will be
+      used, and the input_column will be created.
+    template: Name of the JSON template file under `templates/` or GCS path to
+      the template file.
+    tokenizer: The tokenizer to use for chat_template templates.
+
+  Returns:
+    A dataset compatible with the template.
+  """
+  return dataset.map(
+      _format_template_fn(
+          template,
+          input_column=input_column,
+          tokenizer=tokenizer,
+      )
+  )
+
+
 def load_dataset_with_template(
     dataset_name: str,
     split: str,
     input_column: str,
     template: str = None,
     tokenizer: transformers.PreTrainedTokenizer | None = None,
-) -> Any:
+) -> Tuple[Any, Any]:
   """Loads dataset with templates.
 
   Args:
@@ -435,19 +467,15 @@ def load_dataset_with_template(
     tokenizer: The tokenizer to use for chat_template templates.
 
   Returns:
-    A dataset compatible with the template.
+    The raw dataset and the dataset compatible with the template.
   """
-  dataset = _get_dataset(dataset_name, split=split)
+  raw = _get_dataset(dataset_name, split=split)
   if template:
-    dataset = dataset.map(
-        _format_template_fn(
-            template,
-            input_column=input_column,
-            tokenizer=tokenizer,
-        )
-    )
+    templated = format_dataset(raw, input_column, template, tokenizer)
+  else:
+    templated = None
 
-  return dataset
+  return raw, templated
 
 
 def validate_dataset_with_template(
@@ -521,12 +549,11 @@ def validate_dataset_with_template(
         f" https://github.com/GoogleCloudPlatform/{_VERTEX_AI_SAMPLES_GITHUB_REPO_NAME}/tree/main/{_VERTEX_AI_SAMPLES_GITHUB_TEMPLATE_DIR}."
     )
 
-  dataset = _get_dataset(dataset_name, split, num_proc).map(
-      _format_template_fn(
-          template_path,
-          input_column=input_column,
-          tokenizer=tokenizer,
-      )
+  dataset = format_dataset(
+      _get_dataset(dataset_name, split, num_proc),
+      input_column,
+      template_path,
+      tokenizer,
   )
 
   if tokenizer is not None:
@@ -541,4 +568,3 @@ def validate_dataset_with_template(
           os.path.basename(dataset_name), os.path.basename(template)
       )
   )
-

community-content/vertex_model_garden/model_oss/peft/train/test_utils/command_builder.py

@@ -0,0 +1,95 @@
+"""Class that bundles docker related flags."""
+
+import getpass
+import os
+import pwd
+
+
+class CommandBuilder:
+  """Base class for building commands."""
+
+  def __init__(self):
+    self._defaults = []
+    self._env_vars = {}
+
+  def add_env_var(self, var: str, val: str) -> None:
+    """Add environment variable to the command.
+
+    Args:
+      var: environment variable name.
+      val: environment variable value.
+    """
+    self._env_vars[var] = val
+
+  def add_mount_map(self, host_path, docker_path):
+    pass
+
+
+class DockerCommandBuilder(CommandBuilder):
+  """Bundle docker related flags."""
+
+  def __init__(self, docker_uri: str, shm_size: str = '128gb'):
+    super().__init__()
+    self._docker_uri = [docker_uri]
+    self.privilege_mode = []
+    self.entrypoint = []
+
+    self._defaults = [
+        'docker',
+        'run',
+        '--gpus=all',
+        '--net=host',
+        '--rm',
+        f'--shm-size={shm_size}',
+    ]
+
+    self._mount_maps = []
+    user = getpass.getuser()
+    # username ends with `_google_com` is managed by ldap and does not have a
+    # corresponding entry in /etc/passwd or /etc/group file. We cannot enable
+    # non-root docker user with below method.
+    if not user.endswith('_google_com'):
+      uid = os.getuid()
+      gid = pwd.getpwuid(uid).pw_gid
+      self._defaults += [
+          f'--user={uid}:{gid}',
+          '--volume=/etc/group:/etc/group:ro',
+          '--volume=/etc/passwd:/etc/passwd:ro',
+      ]
+
+  def add_mount_map(self, host_path, docker_path):
+    self._mount_maps.append(f'--volume={host_path}:{docker_path}')
+
+  def add_privilege_mode(self):
+    self.privilege_mode = ['--privileged']
+
+  def add_entrypoint(self, entrypoint: list[str]):
+    self.entrypoint = entrypoint
+
+  def build_cmd(self) -> str:
+    return (
+        self._defaults
+        + [f'--env={var}={val}' for var, val in self._env_vars.items()]
+        + self._mount_maps
+        + self.privilege_mode
+        + self._docker_uri
+        + self.entrypoint
+    )
+
+
+class PythonCommandBuilder(CommandBuilder):
+  """Bundle Python test command related flags."""
+
+  def __init__(self):
+    super().__init__()
+    self._defaults = [
+        'python3',
+        './vertex_vision_model_garden_peft/train/vmg/train_entrypoint.py',
+    ]
+
+  def build_cmd(self) -> str:
+    os.environ.update(self._env_vars)
+    return self._defaults
+
+  def add_entrypoint(self, entrypoint: list[str]):
+    self._defaults = entrypoint

community-content/vertex_model_garden/model_oss/peft/train/test_utils/test_util.py

@@ -0,0 +1,471 @@
+"""Test util class."""
+
+import copy
+import dataclasses
+import datetime
+import inspect
+import os
+import signal
+import subprocess
+import sys
+from absl import flags
+from absl import logging
+from absl.testing import parameterized
+import command_builder
+import immutabledict
+import torch
+
+_DOCKER_URI = flags.DEFINE_string('docker_uri', None, 'docker image uri')
+
+_DRY_RUN = flags.DEFINE_bool('dry_run', False, 'dry-run the commands')
+
+_LOCAL_INPUT_DIR = flags.DEFINE_string(
+    'local_input_dir',
+    os.path.expanduser('~/test_input'),
+    'local directory for storing input data.',
+)
+
+_LOCAL_OUTPUT_DIR = flags.DEFINE_string(
+    'local_output_dir',
+    '/tmp',
+    'local directory for storing test output.',
+)
+
+
+_GCS_INPUT_DIR = flags.DEFINE_string(
+    'gcs_input_dir',
+    'gs://vmg-tuning-docker-test',
+    'GCS directory that stores model checkpoint, dataset and etc.',
+)
+
+_GCS_OUTPUT_DIR = flags.DEFINE_string(
+    'gcs_output_dir',
+    'gs://vmg-tuning-docker-test/output',
+    'GCS directory that stores test output.',
+)
+
+_GCS_TESTDATA_DIR = 'peft-train-image-test'
+
+_THROUGHPUT_TEST_EXCEPTIONS = immutabledict.immutabledict({
+    ('bm_deepspeed_zero3_8gpu_gemma-2-9b-it_4bit.txt', '12.0'): float('inf'),
+    ('bm_fsdp_8gpu_llama3.1-70b-hf_4bit.txt', '20.0'): float('inf'),
+    ('bm_deepspeed_zero2_8gpu_gemma-2-2b-it_bfloat16.txt', '12.0'): 20.0,
+    ('bm_deepspeed_zero3_8gpu_gemma-2-2b-it_4bit.txt', '4.0'): 20.0,
+    ('bm_deepspeed_zero3_8gpu_gemma-2-27b-it_4bit.txt', '4.0'): 20.0,
+})
+
+
+@dataclasses.dataclass
+class BenchmarkStats:
+  """Class to store the benchmark result.
+
+  Attributes:
+    peak_mem: peak memory in GB.
+    throughput: throughput in tokens/sec.
+  """
+
+  peak_mem: float
+  throughput: float
+
+
+class TestBase(parameterized.TestCase):
+  """Test base class that defines how to run commands."""
+
+  def setUp(self):
+    super().setUp()
+
+    # Create a copy of the environment variables
+    self.old_env_var = copy.deepcopy(os.environ)
+    if _DOCKER_URI.value:
+      self.command_builder = command_builder.DockerCommandBuilder(
+          _DOCKER_URI.value
+      )
+    else:
+      self.command_builder = command_builder.PythonCommandBuilder()
+
+    self.command_builder.add_mount_map(
+        os.path.expanduser('~'), os.path.expanduser('~')
+    )
+    self.command_builder.add_mount_map(
+        self.local_input_dir(), self.local_input_dir()
+    )
+
+    self.task_cmd_builder = None
+
+  def tearDown(self):
+    super().tearDown()
+    # Restore the original environment variables
+    os.environ.clear()
+    os.environ.update(self.old_env_var)
+
+  def cmd(self):
+    return self.command_builder.build_cmd() + self.task_cmd_builder.build_cmd()
+
+  def run_cmd(self) -> int:
+    return run_cmd(self.cmd(), output_file=None)
+
+  def gcs_output_dir(self):
+    return _GCS_OUTPUT_DIR.value
+
+  def local_input_dir(self):
+    """Returns local input dir in host/docker."""
+    return _LOCAL_INPUT_DIR.value
+
+  def local_output_dir(self):
+    """Returns local output dir in host/docker."""
+    return _LOCAL_OUTPUT_DIR.value
+
+  def get_testcase_name(self):
+    """Returns the function name at the calling site."""
+    # https://docs.python.org/3/library/inspect.html#inspect.FrameInfo
+    cur_frame = inspect.currentframe()
+    # https://stackoverflow.com/a/17366561
+    return cur_frame.f_back.f_code.co_name
+
+
+def get_timestamp():
+  return datetime.datetime.now(datetime.timezone.utc).strftime(
+      '%Y%m%d_%H%M%S%Z'
+  )
+
+
+def download_from_gcs(gcs_uri: str, local_dir: str):
+  if not os.path.exists(local_dir):
+    os.mkdir(local_dir)
+  subprocess.check_output([
+      'gcloud',
+      'storage',
+      'cp',
+      '-r',
+      gcs_uri,
+      local_dir,
+  ])
+
+
+def get_test_data_path(name: str, download: bool = True) -> str:
+  """Gets test data path.
+
+  Args:
+    name: name of the test data
+    download: if True, then download data from GCS and returns its local path.
+
+  Returns:
+    test data path.
+  """
+  if not download:
+    return os.path.join(_GCS_INPUT_DIR.value, name)
+
+  local_data = os.path.join(_LOCAL_INPUT_DIR.value, name)
+  if not os.path.exists(local_data):
+    # If `name` is a file in sub-folders, then create the sub-folders under
+    # `_LOCAL_INPUT_DIR`.
+    local_data_dir = os.path.dirname(local_data)
+    if not os.path.exists(local_data_dir):
+      os.makedirs(local_data_dir)
+
+    download_from_gcs(os.path.join(_GCS_INPUT_DIR.value, name), local_data_dir)
+
+  return local_data
+
+
+def run_cmd(cmd: list[str], output_file: str = None) -> int:
+  """Runs the command and returns the return code.
+
+  Args:
+    cmd: The command to run.
+    output_file: The file to write the output to.
+
+  Returns:
+    The return code of the command.
+  """
+  logging.info('running command: \n%s', ' \\\n'.join(cmd))
+  if _DRY_RUN.value:
+    return 0
+  stdout = sys.stdout if output_file is None else open(output_file, 'w')
+  p = subprocess.Popen(cmd, stdout=stdout, stderr=sys.stderr)
+  try:
+    unused_output, unused_error = p.communicate()
+    return_code = p.returncode
+  except KeyboardInterrupt:
+    p.send_signal(signal.SIGINT)
+    return_code = 0
+  finally:
+    if output_file is not None:
+      stdout.close()
+  return return_code
+
+
+def get_pretrained_model_name_or_path(model_id: str) -> str:
+  # If `model_id` contains `/`, it is assumed to be HF model or model from GCS.
+  if '/' in model_id:
+    return model_id
+
+  return get_test_data_path(model_id, download=True)
+
+
+def is_gpu_h100():
+  """Checks if the GPU is H100."""
+  return 'H100' in torch.cuda.get_device_name()
+
+
+def is_gpu_a100():
+  """Checks if the GPU is A100."""
+  return 'A100' in torch.cuda.get_device_name()
+
+
+def _get_formatted_string(max_seq_length: int) -> str:
+  """Returns the formatted string for max_seq_length.
+
+  Args:
+    max_seq_length: max sequence length to get the formatted string.
+
+  Returns:
+    formatted string for max_seq_length.
+  """
+  return f'{max_seq_length/1024.0:.1f}'
+
+
+def get_benchmark_results(
+    benchmark_file_path: str, max_seq_length: int
+) -> BenchmarkStats:
+  """Gets benchmark results from the benchmark file.
+
+  Args:
+    benchmark_file_path: path to the benchmark file.
+    max_seq_length: max sequence length to get the benchmark results.
+
+  Returns:
+    peak_mem: peak memory in GB.
+    throughput: throughput in tokens/sec.
+  """
+  formatted_max_seq_length = _get_formatted_string(max_seq_length)
+  peak_mem, throughput = None, None
+  with open(benchmark_file_path, 'r') as f:
+    for line in f:
+      if line.startswith(formatted_max_seq_length):
+        metrics = line.split('|')
+        try:
+          peak_mem = float(metrics[1].strip())
+        except ValueError:
+          pass
+        try:
+          throughput = float(metrics[2].strip())
+        except ValueError:
+          pass
+        break
+    else:
+      logging.error(
+          'No metrics found for max_seq_length %s in %s',
+          formatted_max_seq_length,
+          benchmark_file_path,
+      )
+  return BenchmarkStats(peak_mem, throughput)
+
+
+def print_benchmark_file(file_path: str) -> None:
+  """Prints the contents of the file.
+
+  Args:
+    file_path: path to the file.
+  """
+  with open(file_path, 'r') as f:
+    for line in f:
+      logging.info(line.strip())
+
+
+def print_benchmark_results(
+    benchmark_file_path: str, benchmark_type: str
+) -> None:
+  """Prints the benchmark results.
+
+  Args:
+    benchmark_file_path: path to the benchmark file.
+    benchmark_type: type of the benchmark.
+  """
+  benchmark_filename = os.path.basename(benchmark_file_path)
+  logging.info('--------------------------------------------------------------')
+  logging.info('%s benchmark for %s', benchmark_type, benchmark_filename)
+  logging.info('--------------------------------------------------------------')
+  print_benchmark_file(benchmark_file_path)
+
+
+def _calculate_percent_change(
+    actual_value: float, expected_value: float
+) -> float:
+  """Calculates the percent change between the actual and expected values.
+
+  Args:
+    actual_value: actual value to compare.
+    expected_value: expected value to compare.
+
+  Returns:
+    percent change between the actual and expected values.
+  """
+  return ((actual_value - expected_value) / expected_value) * 100.0
+
+
+def compare_benchmark_results(
+    expected_benchmark_file_path: str,
+    actual_benchmark_file_path: str,
+    allowed_threshold: float,
+    max_seq_length: int,
+) -> bool:
+  """Compares if the benchmark results are the similar.
+
+  Args:
+    expected_benchmark_file_path: path to the expected benchmark file.
+    actual_benchmark_file_path: path to the actual benchmark file.
+    allowed_threshold: allowed percent range of the benchmark results.
+    max_seq_length: max sequence length to get the benchmark results.
+
+  Returns:
+    True if the benchmark results are the similar, False otherwise.
+  """
+  benchmark_filename = os.path.basename(expected_benchmark_file_path)
+  expected_results = get_benchmark_results(
+      expected_benchmark_file_path, max_seq_length
+  )
+  expected_peak_mem, expected_throughput = (
+      expected_results.peak_mem,
+      expected_results.throughput,
+  )
+  actual_results = get_benchmark_results(
+      actual_benchmark_file_path, max_seq_length
+  )
+  actual_peak_mem, actual_throughput = (
+      actual_results.peak_mem,
+      actual_results.throughput,
+  )
+  formatted_max_seq_length = _get_formatted_string(max_seq_length)
+
+  # Case 1: both peak mem and throughput are None(ideally due to OOM)
+  if expected_peak_mem is None and actual_peak_mem is None:
+    logging.info(
+        'Both peak mem and throughput are None for max_seq_length %d.',
+        max_seq_length,
+    )
+    return True
+
+  check_oom_exception = _THROUGHPUT_TEST_EXCEPTIONS.get(
+      (benchmark_filename, formatted_max_seq_length), 0.0
+  ) == float('inf')
+  # Case 2: When something strated to fail recently, or something which failed
+  # before but is working now.
+  if expected_peak_mem is None and actual_peak_mem is not None:
+    if check_oom_exception:
+      return True
+    logging.error(
+        'One of the failing benchmarks in %s is passing now for max_seq_length'
+        ' %d. The expected peak mem and throughput are None, but the actual'
+        ' peak mem is %f and actual throughput is %f',
+        benchmark_filename,
+        max_seq_length,
+        actual_peak_mem,
+        actual_throughput,
+    )
+    return False
+  if actual_peak_mem is None and expected_peak_mem is not None:
+    if check_oom_exception:
+      return True
+    logging.error(
+        'One of the passing benchmarks in %s is failing now for max_seq_length'
+        ' %d. The actual peak mem and throughput are None, but the expected'
+        ' peak mem is %f and expected throughput is %f',
+        benchmark_filename,
+        max_seq_length,
+        expected_peak_mem,
+        expected_throughput,
+    )
+    return False
+  # Case 3: When both actual peak mem and throughput lies within the range
+  # of their respective expected values.
+  mem_percent_change = _calculate_percent_change(
+      actual_peak_mem, expected_peak_mem
+  )
+  throughput_percent_change = _calculate_percent_change(
+      actual_throughput, expected_throughput
+  )
+  allowed_threshold = _THROUGHPUT_TEST_EXCEPTIONS.get(
+      (benchmark_filename, formatted_max_seq_length), allowed_threshold
+  )
+
+  if abs(mem_percent_change) > allowed_threshold:
+    logging.error(
+        'The peak memory is changing by more than %f%% for max_seq_length %d.'
+        ' Expected: %f, Actual: %f',
+        allowed_threshold,
+        max_seq_length,
+        expected_peak_mem,
+        actual_peak_mem,
+    )
+    return False
+  if abs(throughput_percent_change) > allowed_threshold:
+    logging.error(
+        'The throughput is changing by more than %f%% for max_seq_length %d.'
+        ' Expected throughput: %f, Actual throughput: %f',
+        allowed_threshold,
+        max_seq_length,
+        expected_throughput,
+        actual_throughput,
+    )
+    return False
+
+  return True
+
+
+def check_benchmark_results(
+    actual_benchmark_file_path: str,
+    model_family: str,
+    allowed_threshold: float,
+    max_seq_length: int,
+) -> bool:
+  """Checks the benchmark result between the actual and expected benchmark files.
+
+  Args:
+    actual_benchmark_file_path: path to the actual benchmark file.
+    model_family: family of the model.
+    allowed_threshold: allowed range of the benchmark results in percent.
+    max_seq_length: max sequence length to get the benchmark results.
+
+  Returns:
+    True if the benchmark results are the similar, False otherwise.
+  """
+  benchmark_filename = os.path.basename(actual_benchmark_file_path)
+  get_test_data_path(_GCS_TESTDATA_DIR)
+  expected_benchmark_file_path = os.path.join(
+      _LOCAL_INPUT_DIR.value,
+      _GCS_TESTDATA_DIR,
+      model_family,
+      benchmark_filename,
+  )
+  print_benchmark_results(expected_benchmark_file_path, 'Expected')
+  print_benchmark_results(actual_benchmark_file_path, 'Actual')
+
+  return compare_benchmark_results(
+      expected_benchmark_file_path,
+      actual_benchmark_file_path,
+      allowed_threshold,
+      max_seq_length,
+  )
+
+
+def list_gcs_directories(bucket: str, directory: str) -> list[str]:
+  """Lists GCS files."""
+  output = subprocess.check_output([
+      'gcloud',
+      'storage',
+      'ls',
+      f'gs://{bucket}/{directory}',
+  ])
+  return output.decode('utf-8').splitlines()
+
+
+def delete_gcs_object(gcs_directory: str):
+  """Deletes GCS object."""
+  subprocess.check_output([
+      'gcloud',
+      'storage',
+      'rm',
+      '-r',
+      f'{gcs_directory}',
+  ])

community-content/vertex_model_garden/model_oss/peft/train/tests/docker_command_builder.py

@@ -1,45 +0,0 @@
-"""Class that bundles docker related flags."""
-
-import getpass
-import os
-import pwd
-
-
-class DockerCommandBuilder:
-  """Bundle docker related flags."""
-
-  def __init__(self, docker_uri, shm_size='128gb'):
-    self._docker_uri = [docker_uri]
-
-    self._defaults = [
-        'docker',
-        'run',
-        '--gpus=all',
-        '--net=host',
-        '--rm',
-        f'--shm-size={shm_size}',
-    ]
-
-    user = getpass.getuser()
-    # username ends with `_google_com` is managed by ldap and does not have a
-    # corresponding entry in /etc/passwd or /etc/group file. We cannot enable
-    # non-root docker user with below method.
-    if not user.endswith('_google_com'):
-      uid = os.getuid()
-      gid = pwd.getpwuid(uid).pw_gid
-      self._defaults += [
-          f'--user={uid}:{gid}',
-          '--volume=/etc/group:/etc/group:ro',
-          '--volume=/etc/passwd:/etc/passwd:ro',
-      ]
-    self._env_vars = []
-    self._mount_maps = []
-
-  def add_env_var(self, var, val):
-    self._env_vars.append(f'--env={var}={val}')
-
-  def add_mount_map(self, host_path, docker_path):
-    self._mount_maps.append(f'--volume={host_path}:{docker_path}')
-
-  def build_cmd(self) -> str:
-    return self._defaults + self._env_vars + self._mount_maps + self._docker_uri

community-content/vertex_model_garden/model_oss/peft/train/tests/test_instruct_lora_trained_model_quality.py

@@ -1,152 +0,0 @@
-# pylint: disable=missing-function-docstring
-# pylint: disable=missing-class-docstring
-"""Tests to make sure trained model achieves decent quality.
-
-Right now, the metric is loss decreasing and we'll eyeball the TB graphs.
-"""
-
-import os
-
-from absl.testing import absltest
-from absl.testing import parameterized
-import instruct_lora_command_builder as task_cmd_builder
-import test_util
-
-
-class TrainedModelQualityTest(test_util.TestBase):
-
-  _TEST_OUTPUT_DIR = os.path.expanduser('~/output')
-
-  @classmethod
-  def setUpClass(cls):
-    super().setUpClass()
-    cls.test_suite_output_dir = os.path.join(
-        cls._TEST_OUTPUT_DIR,
-        os.path.splitext(os.path.basename(__file__))[0],
-        cls.__class__.__name__,
-    )
-
-  def setUp(self):
-    super().setUp()
-
-    self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
-    self.task_cmd_builder.task = 'instruct-lora'
-    self.task_cmd_builder.eval_tasks = 'builtin_eval'
-    self.task_cmd_builder.eval_metric_name = 'loss'
-    self.task_cmd_builder.per_device_batch_size = 1
-    self.task_cmd_builder.gradient_accumulation_steps = 8
-    self.task_cmd_builder.lora_rank = 16
-    self.task_cmd_builder.lora_alpha = 32
-    self.task_cmd_builder.lora_dropout = 0.05
-    self.task_cmd_builder.learning_rate = 5e-5
-    self.task_cmd_builder.num_epochs = 2.0
-    self.task_cmd_builder.warmup_ratio = 0.01
-    self.task_cmd_builder.max_steps = -1
-    self.task_cmd_builder.save_steps = 10
-    self.task_cmd_builder.eval_steps = 10
-    self.task_cmd_builder.max_seq_length = 4096
-    self.task_cmd_builder.load_precision = '4bit'
-    self.task_cmd_builder.gradient_checkpointing = True
-    self.task_cmd_builder.completion_only = True
-    self.task_cmd_builder.attn_implementation = 'flash_attention_2'
-    self.task_cmd_builder.report_to = 'tensorboard'
-
-  def setup_output_dir(self, testcase_name: str):
-    testcase_output_dir = os.path.join(
-        self.test_suite_output_dir, testcase_name
-    )
-    self.task_cmd_builder.ckpt_dir = os.path.join(
-        testcase_output_dir, 'adapter'
-    )
-    self.task_cmd_builder.logging_dir = os.path.join(
-        testcase_output_dir, 'logs'
-    )
-    self.task_cmd_builder.merged_model_dir = os.path.join(
-        testcase_output_dir, 'merged'
-    )
-
-  @parameterized.named_parameters(
-      ('llama3-8b', 'llama3-8b-hf'),
-      ('llama3.1-8b', 'llama3.1-8b-hf'),
-  )
-  def test_8b_model_deepspeed(self, model_name):
-    self.setup_output_dir(f'test_deepspeed_{model_name}')
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(model_name)
-    )
-    self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
-        'peft_train_sample.jsonl'
-    )
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'input_text'
-    self.task_cmd_builder.template = 'llama3-text-bison'
-    self.task_cmd_builder.eval_dataset = test_util.get_test_data_path(
-        'peft_eval_sample.jsonl'
-    )
-    self.task_cmd_builder.eval_split_name = 'train'
-    self.task_cmd_builder.eval_instruct_column = (
-        self.task_cmd_builder.instruct_column
-    )
-    self.task_cmd_builder.eval_template = self.task_cmd_builder.template
-
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
-
-    self.assertEqual(self.run_cmd(), 0)
-
-  @parameterized.named_parameters(
-      ('llama3-70b', 'llama3-70b-hf'),
-      ('llama3.1-70b', 'llama3.1-70b-hf'),
-  )
-  def test_70b_model_deepspeed(self, model_name):
-    self.setup_output_dir(f'test_deepspeed_{model_name}')
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(model_name)
-    )
-    self.task_cmd_builder.config_file = (
-        'vertex_vision_model_garden_peft/deepspeed_zero2_8gpu.yaml'
-    )
-    self.task_cmd_builder.train_dataset = 'timdettmers/openassistant-guanaco'
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'text'
-    self.task_cmd_builder.template = 'openassistant-guanaco'
-    self.task_cmd_builder.eval_dataset = self.task_cmd_builder.train_dataset
-    self.task_cmd_builder.eval_split_name = 'test'
-    self.task_cmd_builder.eval_instruct_column = (
-        self.task_cmd_builder.instruct_column
-    )
-    self.task_cmd_builder.eval_template = self.task_cmd_builder.template
-
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
-
-    self.assertEqual(self.run_cmd(), 0)
-
-  @parameterized.named_parameters(
-      ('llama3-70b', 'llama3-70b-hf'),
-      ('llama3.1-70b', 'llama3.1-70b-hf'),
-  )
-  def test_70b_model_fsdp(self, model_name):
-    self.setup_output_dir(f'test_fsdp_{model_name}')
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(model_name)
-    )
-    self.task_cmd_builder.config_file = (
-        'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
-    )
-    self.task_cmd_builder.train_dataset = 'timdettmers/openassistant-guanaco'
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'text'
-    self.task_cmd_builder.template = 'openassistant-guanaco'
-    self.task_cmd_builder.eval_dataset = self.task_cmd_builder.train_dataset
-    self.task_cmd_builder.eval_split_name = 'test'
-    self.task_cmd_builder.eval_instruct_column = (
-        self.task_cmd_builder.instruct_column
-    )
-    self.task_cmd_builder.eval_template = self.task_cmd_builder.template
-
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
-
-    self.assertEqual(self.run_cmd(), 0)
-
-
-if __name__ == '__main__':
-  absltest.main()

community-content/vertex_model_garden/model_oss/peft/train/tests/test_quantize_model.py

@@ -1,49 +0,0 @@
-# pylint: disable=missing-function-docstring
-# pylint: disable=missing-class-docstring
-"""Tests quantize model task in PEFT docker."""
-
-import os
-import time
-
-from absl.testing import absltest
-import quantize_model_command_builder as task_cmd_builder
-import test_util
-
-
-class QuantizeModelTest(test_util.TestBase):
-
-  def setUp(self):
-    super().setUp()
-
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '')
-    self.docker_builder.add_mount_map(
-        os.path.expanduser('~'), os.path.expanduser('~')
-    )
-
-    self.task_cmd_builder = task_cmd_builder.QuantizeModelCommandBuilder()
-    self.task_cmd_builder.task = 'quantize-model'
-    self.task_cmd_builder.pretrained_model_id = (
-        'gs://vertex-model-garden-public-us/llama3/llama3-8b-hf'
-    )
-    self.task_cmd_builder.quantization_method = 'awq'
-    self.task_cmd_builder.quantization_precision_mode = '4bit'
-    self.task_cmd_builder.quantization_dataset_name = 'pileval'
-    self.task_cmd_builder.text_column_in_quantization_dataset = 'text'
-    self.task_cmd_builder.quantization_output_dir = '~/llama3-8b-hf-quantized'
-    self.task_cmd_builder.device_map = None
-    self.task_cmd_builder.max_memory = None
-    self.task_cmd_builder.group_size = 128
-    self.task_cmd_builder.desc_act = False
-    self.task_cmd_builder.damp_percent = 0.1
-    self.task_cmd_builder.cache_examples_on_gpu = False
-    self.task_cmd_builder.awq_version = 'GEMM'
-
-  def test_llama3_8b_model_awq_quantization(self):
-    start_time = time.time()
-    self.assertEqual(self.run_cmd(), 0)
-    end_time = time.time()
-    self.assertLess(end_time - start_time, 1.5 * 60 * 60)
-
-
-if __name__ == '__main__':
-  absltest.main()

community-content/vertex_model_garden/model_oss/peft/train/tests/test_util.py

@@ -1,134 +0,0 @@
-"""Test util class."""
-
-import datetime
-import os
-import signal
-import subprocess
-import sys
-
-from absl import flags
-from absl import logging
-from absl.testing import parameterized
-import docker_command_builder as docker_cmd_builder
-
-_DOCKER_URI = flags.DEFINE_string(
-    'docker_uri', None, 'docker image uri', required=True
-)
-
-_DRY_RUN = flags.DEFINE_bool('dry_run', False, 'dry-run the commands')
-
-_LOCAL_INPUT_DIR = flags.DEFINE_string(
-    'local_input_dir',
-    os.path.expanduser('~/test_input'),
-    'local directory for storing input data.',
-)
-
-_LOCAL_OUTPUT_DIR = flags.DEFINE_string(
-    'local_output_dir',
-    '/tmp',
-    'local directory for storing test output.',
-)
-
-
-_GCS_INPUT_DIR = flags.DEFINE_string(
-    'gcs_input_dir',
-    'gs://peft-docker-test',
-    'GCS directory that stores model checkpoint, dataset and etc.',
-)
-
-_GCS_OUTPUT_DIR = flags.DEFINE_string(
-    'gcs_output_dir',
-    'gs://peft-docker-test/output',
-    'GCS directory that stores test output.',
-)
-
-
-class TestBase(parameterized.TestCase):
-  """Test base class that defines how to run commands."""
-
-  def setUp(self):
-    super().setUp()
-
-    self.docker_builder = docker_cmd_builder.DockerCommandBuilder(
-        _DOCKER_URI.value
-    )
-    self.docker_builder.add_mount_map(
-        os.path.expanduser('~'), os.path.expanduser('~')
-    )
-    self.docker_builder.add_mount_map(
-        self.local_input_dir(), self.local_input_dir()
-    )
-
-    self.task_cmd_builder = None
-
-  def cmd(self):
-    return self.docker_builder.build_cmd() + self.task_cmd_builder.build_cmd()
-
-  def run_cmd(self) -> int:
-    logging.info('running command: \n%s', ' \\\n'.join(self.cmd()))
-    if _DRY_RUN.value:
-      return 0
-
-    p = subprocess.Popen(self.cmd(), stdout=sys.stdout, stderr=sys.stderr)
-    try:
-      unused_output, unused_error = p.communicate()
-      return p.returncode
-    except KeyboardInterrupt:
-      p.send_signal(signal.SIGINT)
-      return 0
-
-  def gcs_output_dir(self):
-    return _GCS_OUTPUT_DIR.value
-
-  def local_output_dir(self):
-    return _LOCAL_OUTPUT_DIR.value
-
-  def local_input_dir(self):
-    return _LOCAL_INPUT_DIR.value
-
-
-def get_timestamp():
-  return datetime.datetime.now(datetime.timezone.utc).strftime(
-      '%Y%m%d_%H%M%S%Z'
-  )
-
-
-def get_test_data_path(name: str, download: bool = True) -> str:
-  """Gets test data path.
-
-  Args:
-    name: name of the test data
-    download: if True, then download data from GCS and returns its local path.
-
-  Returns:
-    test data path.
-  """
-
-  def _download_from_gcs(name):
-    if not os.path.exists(_LOCAL_INPUT_DIR.value):
-      os.mkdir(_LOCAL_INPUT_DIR.value)
-    subprocess.check_output([
-        'gsutil',
-        '-m',
-        'cp',
-        '-r',
-        os.path.join(_GCS_INPUT_DIR.value, name),
-        _LOCAL_INPUT_DIR.value,
-    ])
-
-  if not download:
-    return os.path.join(_GCS_INPUT_DIR.value, name)
-
-  local_data = os.path.join(_LOCAL_INPUT_DIR.value, name)
-  if not os.path.exists(local_data):
-    _download_from_gcs(name)
-
-  return local_data
-
-
-def get_pretrained_model_id(model_id: str) -> str:
-  # If `model_id` contains `/`, it is assumed to be HF model or model from GCS.
-  if '/' in model_id:
-    return model_id
-
-  return get_test_data_path(model_id, download=True)

community-content/vertex_model_garden/model_oss/peft/train/tests/tools/command_builder_gen.py

@@ -1,109 +0,0 @@
-"""Tools to generate CommandBuilder class.
-
-See go/vmg-oss-peft-tests#commandbuilder-class-generation for details.
-"""
-
-import argparse
-import dataclasses
-from typing import List
-
-_DO_NOT_MODIFY_WARNING = """
-# DO NOT MODIFY: this file is auto-generated
-# See go/vmg-oss-peft-tests#command-builder-genpy
-"""
-
-_GETTER_TMPL = """
-    @property
-    def {}(self):
-        return self._{}
-"""
-
-_SETTER_TMPL = """
-    @{}.setter
-    def {}(self, val: {}):
-        self._{} = val
-"""
-
-_INIT_NAME = """
-    def __init__(self):"""
-
-_INIT_FIELDS = """
-        self._{} = None"""
-
-_BUILD_CMD = r"""
-    def build_cmd(self) -> str:
-        cmd = []
-        for k, v in self.__dict__.items():
-            if v is not None:
-                cmd.append(f'--{k[1:]}={v}')
-        return cmd
-"""
-
-
-@dataclasses.dataclass
-class FlagInfo:
-  api_name: str
-  impl_name: str
-  arg_type: str
-
-
-def get_flag_info(line: str) -> FlagInfo:
-  api_name, impl_name, arg_type = [x.strip() for x in line.split(',')]
-  return FlagInfo(api_name, impl_name, arg_type)
-
-
-def gen_getter(info: FlagInfo) -> str:
-  return _GETTER_TMPL.format(info.api_name, info.impl_name)
-
-
-def gen_setter(info: FlagInfo) -> str:
-  return _SETTER_TMPL.format(
-      info.api_name, info.api_name, info.arg_type, info.impl_name
-  )
-
-
-def gen_init(infos: List[FlagInfo]) -> str:
-  fields = [_INIT_FIELDS.format(i.impl_name) for i in infos]
-  return ''.join([_INIT_NAME] + fields)
-
-
-def main():
-  parser = argparse.ArgumentParser()
-  parser.add_argument(
-      '--flags_def', required=True, help='file path contain flags definition.'
-  )
-  parser.add_argument(
-      '--generated_file',
-      required=True,
-      help='file path to the generated command builder.',
-  )
-  parser.add_argument(
-      '--class_name',
-      required=True,
-      help='class name for command build',
-  )
-  args = parser.parse_args()
-
-  flags_info = []
-  with open(args.flags_def, 'r') as flags_f:
-    for line in flags_f:
-      if not line.startswith('#'):
-        flags_info.append(get_flag_info(line))
-
-  with open(args.generated_file, 'w') as gen_f:
-    # Disables pylint messages.
-    # See https://stackoverflow.com/a/43510297
-    print('# pylint: disable=W,C,R', file=gen_f)
-    print(_DO_NOT_MODIFY_WARNING, file=gen_f)
-    print(f'class {args.class_name}:', file=gen_f)
-    print(gen_init(flags_info), file=gen_f)
-    for info in flags_info:
-      print(gen_getter(info), file=gen_f)
-      print(gen_setter(info), file=gen_f)
-    print(_BUILD_CMD, file=gen_f)
-
-  print(f'file generated at {args.generated_file}')
-
-
-if __name__ == '__main__':
-  main()

community-content/vertex_model_garden/model_oss/peft/train/tests/tools/instruct_lora_flags.txt

@@ -1,50 +0,0 @@
-# api_name, impl_name, value_type
-
-
-#                    eval related and etc.
-config_file, config_file, str
-task, task, str
-pretrained_model_id, pretrained_model_id, str
-train_dataset, dataset_name, str
-train_split_name, train_split_name, str
-template, template, str
-instruct_column, instruct_column_in_dataset, str
-ckpt_dir, output_dir, str
-merged_model_dir, merge_base_and_lora_output_dir, str
-logging_dir, logging_output_dir, str
-per_device_batch_size, per_device_train_batch_size, int
-gradient_accumulation_steps, gradient_accumulation_steps, int
-lora_rank, lora_rank, int
-lora_alpha, lora_alpha, int
-lora_dropout, lora_dropout, float
-max_steps, max_steps, int
-num_epochs, num_epochs, float
-max_seq_length, max_seq_length, int
-learning_rate, learning_rate, float
-lr_scheduler_type, lr_scheduler_type, str
-load_precision, precision_mode, str
-train_precision, train_precision, str
-gradient_checkpointing, enable_gradient_checkpointing, bool
-example_packing, use_example_packing, bool
-attn_implementation, attn_implementation, str
-optimizer, optimizer, str
-warmup_ratio, warmup_ratio, float
-report_to, report_to, str
-save_steps, save_steps, int
-logging_steps, logging_steps, int
-huggingface_access_token, huggingface_access_token, str
-eval_dataset, eval_dataset_path, str
-eval_instruct_column, eval_column, str
-eval_template, eval_template, str
-eval_split_name, eval_split, str
-eval_steps, eval_steps, int
-eval_tasks, eval_tasks, str
-eval_metric_name, eval_metric_name, str
-completion_only, completion_only, bool
-max_grad_norm, max_grad_norm, float
-logger_level, logger_level, str
-benchmark_out_file, benchmark_out_file, str
-tuning_data_stats_file, tuning_data_stats_file, str
-enable_peft, enable_peft, bool
-merge_model_precision_mode, merge_model_precision_mode, str
-target_modules, target_modules, str

community-content/vertex_model_garden/model_oss/peft/train/tests/tools/quantize_model_flags.txt

@@ -1,15 +0,0 @@
-# api_name, impl_name, value_type
-task, task, str
-pretrained_model_id, pretrained_model_id, str
-quantization_method, quantization_method, str
-quantization_precision_mode, quantization_precision_mode, str
-quantization_dataset_name, quantization_dataset_name, str
-text_column_in_quantization_dataset, text_column_in_quantization_dataset, str
-quantization_output_dir, quantization_output_dir, str
-device_map, device_map, str
-max_memory, max_memory, str
-group_size, group_size, int
-desc_act, desc_act, bool
-damp_percent, damp_percent, float
-cache_examples_on_gpu, cache_examples_on_gpu, bool
-awq_version, awq_version, str

community-content/vertex_model_garden/model_oss/peft/train/tests/tools/validate_dataset_with_template_flags.txt

@@ -1,9 +0,0 @@
-# api_name, impl_name, value_type
-task, task, str
-template, template, str
-dataset_name, dataset_name, str
-train_split_name, train_split_name, str
-instruct_column_in_dataset, instruct_column_in_dataset, str
-use_multiprocessing, use_multiprocessing, bool
-validate_k_rows_of_dataset, validate_k_rows_of_dataset, int
-validate_percentage_of_dataset, validate_percentage_of_dataset, int

community-content/vertex_model_garden/model_oss/peft/train/util/cluster_spec.py

@@ -0,0 +1,79 @@
+"""Get cluster info from environment variables."""
+
+import dataclasses
+import json
+import os
+
+from absl import logging
+
+
+@dataclasses.dataclass
+class ClusterInfo:
+  """Contains information about the cluster.
+
+  Attributes:
+    primary_node_addr: The address of the primary node.
+    primary_node_port: The port of the primary node.
+    node_rank: The rank of the node.
+    num_nodes: The number of nodes in the cluster.
+  """
+
+  primary_node_addr: str | None = None
+  primary_node_port: str | None = None
+  node_rank: int = 0
+  num_nodes: int = 1
+
+  # Allows unpacking operation like
+  # primary_node_addr, primary_node_port, _, _ = ClusterInfo()
+  # See https://stackoverflow.com/a/70753113
+  def __iter__(self):
+    return iter(dataclasses.astuple(self))
+
+
+def get_cluster_spec() -> ClusterInfo:
+  """Parses CLUSTER_SPEC environment variable and returns the cluster info.
+
+  Returns:
+    A ClusterInfo object.
+  """
+  cluster_spec = os.getenv('CLUSTER_SPEC', None)
+
+  # If CLUSTER_SPEC is not set, use individual vars to construct cluster info.
+  if not cluster_spec:
+    cluster_info = ClusterInfo(
+        primary_node_addr=os.getenv('MASTER_ADDR', None),
+        primary_node_port=os.getenv('MASTER_PORT', None),
+        node_rank=int(os.getenv('RANK', '0')),
+        num_nodes=int(os.getenv('NNODES', '1')),
+    )
+    return cluster_info
+
+  cluster_data = json.loads(cluster_spec)
+  # Get primary node info
+  primary_node = cluster_data['cluster']['workerpool0'][0]
+  logging.info('primary node: %s', primary_node)
+  primary_node_addr, primary_node_port = primary_node.split(':')
+  logging.info('primary node address: %s', primary_node_addr)
+  logging.info('primary node port: %s', primary_node_port)
+
+  # Determine node rank of this machine
+  workerpool = cluster_data['task']['type']
+  if workerpool == 'workerpool0':
+    node_rank = 0
+  elif workerpool == 'workerpool1':
+    # Add 1 for the primary node, since `index` is the index of workerpool1.
+    node_rank = cluster_data['task']['index'] + 1
+  else:
+    raise ValueError(
+        'Only workerpool0 and workerpool1 are supported. Unknown workerpool:'
+        f' {workerpool}'
+    )
+  logging.info('node rank: %s', node_rank)
+
+  # Calculate total nodes.
+  num_nodes = 1  # For the primary node.
+  if 'workerpool1' in cluster_data['cluster']:
+    num_nodes += len(cluster_data['cluster']['workerpool1'])
+  logging.info('num nodes: %s', num_nodes)
+
+  return ClusterInfo(primary_node_addr, primary_node_port, node_rank, num_nodes)

community-content/vertex_model_garden/model_oss/peft/train/util/common_util.py

@@ -0,0 +1,24 @@
+"""Utility functions."""
+
+import logging
+import subprocess
+import sys
+import time
+
+
+def run_cmd(cmd: list[str]) -> float:
+  """Runs the command and logs the output.
+
+  Args:
+    cmd: The command to run.
+
+  Returns:
+    The time it took to run the command.
+  """
+  cmd_str = ' \\\n'.join(cmd)
+  logging.info('launching cmd: \n%s', cmd_str)
+  start_time = time.time()
+  subprocess.run(cmd, stdout=sys.stdout, stderr=sys.stdout, check=True)
+  elapsed_time = round(time.time() - start_time, 2)
+  logging.info('Command %s finished in %0.2f seconds.', cmd_str, elapsed_time)
+  return elapsed_time

community-content/vertex_model_garden/model_oss/peft/train/util/dataset_stats.py

@@ -0,0 +1,197 @@
+"""Calculate dataset statistics like token, example and character counts."""
+
+from collections.abc import Mapping, Sequence
+import dataclasses
+import json
+from typing import Any
+import datasets
+import numpy as np
+import transformers
+from util import dataset_validation_util
+
+_MAX_NUM_DATASET_SAMPLES = 6
+
+
+@dataclasses.dataclass
+class SupervisedTuningDatasetBucket:
+  """Represents a histogram bucket for tuning dataset distribution stats."""
+
+  count: float = 0
+  left: float = 0
+  right: float = 0
+
+
+@dataclasses.dataclass
+class SupervisedTuningDatasetDistribution:
+  """Represents a histogram with summary statistics for tuning dataset distribution stats."""
+
+  sum: int = 0
+  billable_sum: int = 0
+  min: float = 0
+  max: float = 0
+  mean: float = 0
+  median: float = 0
+  p5: float = 0
+  p95: float = 0
+  buckets: list[SupervisedTuningDatasetBucket] = dataclasses.field(
+      default_factory=list
+  )
+
+
+# Represents detailed tuning dataset statistics.
+@dataclasses.dataclass
+class SupervisedTuningDataStats:
+  """Represents detailed tuning dataset stats."""
+
+  tuning_dataset_example_count: int = 0
+  total_tuning_character_count: int = 0
+  total_billable_token_count: int = 0
+  tuning_step_count: int = 0
+  # Represents a histogram and some summary statistics of the number of input
+  # tokens across examples.
+  user_input_token_distribution: SupervisedTuningDatasetDistribution | None = (
+      None
+  )
+  # Represents a histogram and some summary statistics for the number of output
+  # tokens across examples.
+  user_output_token_distribution: SupervisedTuningDatasetDistribution | None = (
+      None
+  )
+  # Represents the number of "messages" (a single-turn conversation will have a
+  # single message) across examples.
+  user_message_per_example_distribution: (
+      SupervisedTuningDatasetDistribution | None
+  ) = None
+  user_dataset_examples: list[str] = dataclasses.field(default_factory=list)
+
+
+def get_dataset_stats(
+    *,
+    raw: Any,
+    templated: Any,
+    template: str,
+    tokenizer: transformers.PreTrainedTokenizer,
+    column: str,
+    effective_batch_size: int,
+) -> Mapping[str, Any]:
+  """Calculates dataset statistics for managed fine-tuning, e.g., total number of tokens."""
+  tokenized_dataset = templated.map(lambda x: tokenizer(x[column]))
+  inputs = tokenized_dataset["input_ids"]
+  tuning_dataset_example_count = int(len(inputs))
+  total_billable_token_count = int(np.sum([len(ex) for ex in inputs]))
+  total_tuning_character_count = int(
+      np.sum([len(ex[column]) for ex in templated])
+  )
+  tuning_step_count = (
+      tuning_dataset_example_count + effective_batch_size - 1
+  ) // effective_batch_size
+
+  # Assume that data is represented as ChatCompletions or Vertex Text-Bison
+  # formats to extract per-example input/output tokens.
+  user_inputs = []
+  user_outputs = []
+  user_input_messages_counts = []
+
+  for ex in raw:
+    if "messages" in ex:
+      messages = ex["messages"]
+      if messages:
+        # For ChatCompletions assume the last turn (i.e. the instruction
+        # response) is the expected output.
+        user_inputs.append({**ex, "messages": messages[:-1]})
+        user_outputs.append({**ex, "messages": messages[-1:]})
+        # Exclude everything but the last message for the number of input
+        # messages.
+        user_input_messages_counts.append(len(messages[:-1]))
+    elif "input_text" in ex:
+      # For Vertex Text-Bison, the `output_text` field is the expected output.
+      user_inputs.append({**ex, "output_text": ""})
+      user_outputs.append(
+          {**ex, "input_text": ex["output_text"], "output_text": ""}
+      )
+      # Vertex Text-Bison goes from input -> output; i.e. there is only a single
+      # input "message".
+      user_input_messages_counts.append(1)
+
+  def calc_histogram(
+      counts: Sequence[int],
+  ) -> SupervisedTuningDatasetDistribution:
+    mean = np.mean(counts)
+    median = np.median(counts).item()
+    max_count = np.max(counts).item()
+    min_count = np.min(counts).item()
+    count_sum = np.sum(counts).item()
+    p5 = np.percentile(counts, 0.05).item()
+    p95 = np.percentile(counts, 0.95).item()
+    hist, bin_edges = np.histogram(counts, bins=10)
+
+    return SupervisedTuningDatasetDistribution(
+        sum=count_sum,
+        billable_sum=count_sum,
+        min=min_count,
+        max=max_count,
+        mean=mean,
+        median=median,
+        p5=p5,
+        p95=p95,
+        buckets=[
+            SupervisedTuningDatasetBucket(
+                count=hist[i].item(),
+                left=bin_edges[i].item(),
+                right=bin_edges[i + 1].item(),
+            )
+            for i in range(len(hist))
+        ],
+    )
+
+  # Tokenize input and output messages separately to generate separate summary
+  # statistics about them.
+  user_input_token_distribution = None
+  if user_inputs:
+    user_input_dataset = dataset_validation_util.format_dataset(
+        datasets.Dataset.from_list(user_inputs), column, template, tokenizer
+    )
+    user_input_tokenized_dataset = user_input_dataset.map(
+        lambda x: tokenizer(x[column])
+    )
+    user_input_tokens = user_input_tokenized_dataset["input_ids"]
+    user_input_token_counts = np.array([len(ex) for ex in user_input_tokens])
+    user_input_token_distribution = calc_histogram(user_input_token_counts)
+
+  user_output_token_distribution = None
+  if user_outputs:
+    user_output_dataset = dataset_validation_util.format_dataset(
+        datasets.Dataset.from_list(user_outputs), column, template, tokenizer
+    )
+    user_output_tokenized_dataset = user_output_dataset.map(
+        lambda x: tokenizer(x[column])
+    )
+    user_output_tokens = user_output_tokenized_dataset["input_ids"]
+    user_output_token_counts = np.array([len(ex) for ex in user_output_tokens])
+    user_output_token_distribution = calc_histogram(user_output_token_counts)
+
+  user_messages_per_example_distribution = None
+  if user_input_messages_counts:
+    user_input_messages_counts = np.array(user_input_messages_counts)
+    user_messages_per_example_distribution = calc_histogram(
+        user_input_messages_counts
+    )
+
+  user_dataset_examples = [
+      json.dumps(ex)
+      for ex in raw.shuffle().select(
+          range(min(len(raw), _MAX_NUM_DATASET_SAMPLES))
+      )
+  ]
+
+  dataset_stats = SupervisedTuningDataStats(
+      tuning_dataset_example_count=tuning_dataset_example_count,
+      total_tuning_character_count=total_tuning_character_count,
+      total_billable_token_count=total_billable_token_count,
+      tuning_step_count=tuning_step_count,
+      user_input_token_distribution=user_input_token_distribution,
+      user_output_token_distribution=user_output_token_distribution,
+      user_message_per_example_distribution=user_messages_per_example_distribution,
+      user_dataset_examples=user_dataset_examples,
+  )
+  return dataclasses.asdict(dataset_stats)

community-content/vertex_model_garden/model_oss/peft/train/util/device_stats.py

@@ -0,0 +1,140 @@
+"""Util functions for reporting device (GPU, CPU) stats."""
+
+import dataclasses
+
+import psutil
+import pynvml
+import torch
+
+
+@dataclasses.dataclass
+class GpuStats:
+  """Holds information about GPU usage stats.
+
+  For memory related, see
+  https://pytorch.org/docs/stable/notes/cuda.html#cuda-memory-management
+  """
+
+  # device id
+  device_id: int
+  # memory reserved.
+  reserved: float
+  # memory occupied.
+  occupied: float
+  # memory reserved, but not used.
+  unused: float
+  # nvidia-smi usually reports more memory usages than pytorch (for driver,
+  # kernel and etc). `smi_diff` tracks this difference.
+  smi_diff: float
+  # Gpu utilization.
+  util: float
+
+  # Allows unpacking operation like
+  # device_id, reserved, occupied, unused, smi_diff, util = GpuStats(...)
+  # See https://stackoverflow.com/a/70753113
+  def __iter__(self):
+    return iter(dataclasses.astuple(self))
+
+
+def gpu_stats() -> GpuStats:
+  """Reports GPU memory usage and utilization."""
+  # See https://pytorch.org/docs/stable/notes/cuda.html#memory-management
+  bytes_per_gb = 1024.0**3
+  device = torch.cuda.current_device()
+  occupied = torch.cuda.memory_allocated(device) / bytes_per_gb
+  reserved = torch.cuda.memory_reserved(device) / bytes_per_gb
+  unused = reserved - occupied
+
+  def smi_mem(device):
+    try:
+      pynvml.nvmlInit()
+      handle = pynvml.nvmlDeviceGetHandleByIndex(device)
+      info = pynvml.nvmlDeviceGetMemoryInfo(handle)
+      return info.used / bytes_per_gb
+    except pynvml.NVMLError:
+      return 0.0
+
+  mem_used_smi = smi_mem(device)
+  smi_diff = mem_used_smi - reserved
+
+  util = torch.cuda.utilization(device)
+  return GpuStats(device, reserved, occupied, unused, smi_diff, util)
+
+
+def gpu_stats_str(stats: GpuStats | None = None) -> str:
+  if stats is None:
+    stats = gpu_stats()
+  device, reserved, occupied, unused, smi_diff, util = stats
+  return (
+      f"GPU ({device=}) memory: {reserved:.2f}({occupied=:.2f}, {unused=:.2f}),"
+      f" {smi_diff=:.2f} GB. Utilization: {util:.2f}%"
+  )
+
+
+@dataclasses.dataclass
+class CpuStats:
+  """Holds information about CPU usage stats."""
+
+  # Total CPU virtual memory i.e. virtual memory allocated + unallocated.
+  total_virtual_mem: float
+  # CPU virtual memory available for use.
+  unallocated_virtual_mem: float
+  # CPU virtual memory already used.
+  allocated_virtual_mem: float
+  # Total CPU swap memory i.e. swap memory allocated + unallocated.
+  total_swap_mem: float
+  # CPU swap memory available for use.
+  unallocated_swap_mem: float
+  # CPU swap memory already used.
+  allocated_swap_mem: float
+  # CPU utilization percentage.
+  utilization: float
+
+
+def cpu_stats() -> CpuStats:
+  """Reports CPU memory usage and utilization."""
+
+  # https://psutil.readthedocs.io/en/latest/#memory
+  gb = 1024.0**3
+  vmem = psutil.virtual_memory()
+  vmem_total = vmem.total / gb
+  vmem_available = vmem.available / gb
+  vmem_used = vmem_total - vmem_available
+  smem = psutil.swap_memory()
+  swap_total = smem.total / gb
+  swap_free = smem.free / gb
+  swap_used = smem.used / gb
+  # https://psutil.readthedocs.io/en/latest/#psutil.cpu_percent
+  cpu_util = psutil.cpu_percent(interval=1e-6)
+  return CpuStats(
+      total_virtual_mem=vmem_total,
+      unallocated_virtual_mem=vmem_available,
+      allocated_virtual_mem=vmem_used,
+      total_swap_mem=swap_total,
+      unallocated_swap_mem=swap_free,
+      allocated_swap_mem=swap_used,
+      utilization=cpu_util,
+  )
+
+
+def cpu_stats_str(stats: CpuStats | None = None) -> str:
+  """Returns a string representation of the CPU stats."""
+
+  if stats is None:
+    stats = cpu_stats()
+  total, occupied, unused = (
+      stats.total_virtual_mem,
+      stats.allocated_virtual_mem,
+      stats.unallocated_virtual_mem,
+  )
+  virtual_mem = (
+      f"CPU virtual memory: {total:.2f}({occupied=:.2f}, {unused=:.2f}) GB"
+  )
+  total, occupied, unused = (
+      stats.total_swap_mem,
+      stats.allocated_swap_mem,
+      stats.unallocated_swap_mem,
+  )
+  swap_mem = f"CPU swap memory: {total:.2f}({occupied=:.2f}, {unused=:.2f}) GB"
+  percent = stats.utilization
+  return f"{virtual_mem} {swap_mem} CPU Utilization: {percent:.2f}%"

community-content/vertex_model_garden/model_oss/peft/train/vmg/callbacks.py

@@ -1,5 +1,7 @@
 """Different trainer callbacks for PEFT Trainer."""
 
+from collections.abc import MutableMapping
+import math
 import time
 
 from absl import logging
@@ -9,7 +11,7 @@ from transformers.trainer_callback import TrainerCallback
 from transformers.trainer_callback import TrainerControl
 from transformers.trainer_callback import TrainerState
 
-from vertex_vision_model_garden_peft.train.vmg import utils
+from util import device_stats
 
 
 class TrainerStatsCallback(TrainerCallback):
@@ -25,6 +27,30 @@ class TrainerStatsCallback(TrainerCallback):
     self._peak_mem = 0.0
     self._avg_throughput = 0.0
 
+  def on_log(
+      self,
+      args: TrainingArguments,
+      state: TrainerState,
+      control: TrainerControl,
+      logs: MutableMapping[str, float] | None = None,
+      **kwargs,
+  ) -> None:
+    """Calculates perplexity from train loss.
+
+    Args:
+      args: Arguments passed to the trainer.
+      state: State of the trainer.
+      control: Control of the trainer.
+      logs: A dict of logs from the training loop.
+      **kwargs: Additional keyword arguments, not used in this callback.
+    """
+    del kwargs  # Unused.
+    if self._partial_state.is_main_process:
+      train_loss = logs.get('loss') if logs is not None else None
+      if train_loss is not None:
+        perplexity = round(float(math.exp(train_loss)), 4)
+        logs['perplexity'] = perplexity
+
   def on_step_end(
       self,
       args: TrainingArguments,
@@ -35,21 +61,29 @@ class TrainerStatsCallback(TrainerCallback):
     if self._partial_state.is_main_process:
       if state.global_step == 1:
         self._prev_time = time.time()
-        delta_t = float('nan')
+        self._prev_num_token = state.num_input_tokens_seen
+        throughput = 0.0
       else:
         cur_time = time.time()
-        delta_t = cur_time - self._prev_time
+        cur_num_token = state.num_input_tokens_seen
+        throughput = (cur_num_token - self._prev_num_token) / (
+            cur_time - self._prev_time
+        )
         self._prev_time = cur_time
-        self._avg_throughput += (delta_t - self._avg_throughput) / (
+        self._prev_num_token = cur_num_token
+        self._avg_throughput += (throughput - self._avg_throughput) / (
             state.global_step - 1
         )
 
-      gpu_stats = utils.gpu_stats()
-      self._peak_mem = max(gpu_stats.total_mem, self._peak_mem)
+      gpu_stats = device_stats.gpu_stats()
+      self._peak_mem = max(
+          gpu_stats.reserved + gpu_stats.smi_diff, self._peak_mem
+      )
       logging.info(
-          'on_step_end: %s, throughput: %.2f s/it',
-          utils.gpu_stats_str(gpu_stats),
-          delta_t,
+          'on_step_end: Throughput: %.2f token/s. %s, %s',
+          throughput,
+          device_stats.gpu_stats_str(gpu_stats),
+          device_stats.cpu_stats_str(),
       )
 
   def on_train_begin(
@@ -61,7 +95,11 @@ class TrainerStatsCallback(TrainerCallback):
   ):
     if self._partial_state.is_main_process:
       self._start_time = time.time()
-      logging.info('on_train_begin: %s', utils.gpu_stats_str())
+      logging.info(
+          'on_train_begin: %s, %s',
+          device_stats.gpu_stats_str(),
+          device_stats.cpu_stats_str(),
+      )
 
   def on_train_end(
       self,
@@ -72,15 +110,17 @@ class TrainerStatsCallback(TrainerCallback):
   ):
     if self._partial_state.is_main_process:
       train_time = time.time() - self._start_time
+      throughput = state.num_input_tokens_seen / train_time
       logging.info(
-          'training time %.2f s, throughput: %.2f s/it, peak_mem: %.2f GB',
+          'training time %.2f s, throughput (including overhead, e.g., ckpt'
+          ' saving): %.2f token/s, peak_mem: %.2f GB',
           train_time,
-          self._avg_throughput,
+          throughput,
           self._peak_mem,
       )
       if self._filename:
         with open(self._filename, 'a') as out_f:
           out_f.write(
-              f'{self._max_seq_length/1024.0:.1f}k | {self._peak_mem:.2f} |'
+              f'{self._max_seq_length/1024.0:.1f} | {self._peak_mem:.2f} |'
               f' {self._avg_throughput:.2f}\n'
           )

community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/gemma2_fsdp_8gpu.yaml

@@ -0,0 +1,28 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: FSDP
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+fsdp_config:
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: Gemma2DecoderLayer
+  fsdp_backward_prefetch: NO_PREFETCH
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_forward_prefetch: false
+  fsdp_offload_params: true
+  fsdp_sharding_strategy: FULL_SHARD
+  fsdp_state_dict_type: SHARDED_STATE_DICT
+  fsdp_sync_module_states: true
+  fsdp_use_orig_params: false
+  fsdp_activation_checkpointing: false
+main_training_function: main
+mixed_precision: bf16
+machine_rank: 0
+num_machines: 1
+num_processes: 8
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/qwen2_fsdp_8gpu.yaml

@@ -0,0 +1,28 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: FSDP
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+fsdp_config:
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: Qwen2DecoderLayer
+  fsdp_backward_prefetch: NO_PREFETCH
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_forward_prefetch: false
+  fsdp_offload_params: true
+  fsdp_sharding_strategy: FULL_SHARD
+  fsdp_state_dict_type: SHARDED_STATE_DICT
+  fsdp_sync_module_states: true
+  fsdp_use_orig_params: false
+  fsdp_activation_checkpointing: false
+main_training_function: main
+mixed_precision: bf16
+machine_rank: 0
+num_machines: 1
+num_processes: 8
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

community-content/vertex_model_garden/model_oss/peft/train/vmg/dockerfile/merge_env.yaml

@@ -31,7 +31,7 @@ dependencies:
 - pip:
   - --extra-index-url https://download.pytorch.org/whl/cu121
   - absl-py==2.1.0
-  - accelerate==0.33.0  # Needed for fp8
+  - accelerate==0.34.2  # Needed for fp8
   - datasets==2.19.2
   - fbgemm-gpu==0.8.0+cu121  # Needed for fp8
   - kfp==2.5.0
@@ -39,5 +39,5 @@ dependencies:
   - protobuf==3.20.3
   - pynvml==11.5.3
   - torch==2.4.0+cu121  # Needed for fp8
-  - transformers==4.43.1
-  - trl==0.9.6
+  - transformers==4.47.1
+  - trl==0.11.2

community-content/vertex_model_garden/model_oss/peft/train/vmg/dockerfile/requirements.txt

@@ -5,23 +5,28 @@
 --extra-index-url https://huggingface.github.io/autogptq-index/whl/cu118/
 
 # keep sorted
-accelerate==0.31.0
+accelerate==0.34.2
 auto_gptq==0.7.1+cu118
-autoawq==0.2.5
+autoawq==0.2.8
 bitsandbytes==0.43.2
 cloudml-hypertune==0.1.0.dev6
-datasets==2.19.2
-deepspeed==0.14.4
+datasets==2.20.0
+deepspeed==0.15.2
 diffusers==0.25.1
+evaluate==0.4.3
 fsspec==2024.3.1
 gcsfs==2024.3.1
-lm_eval==0.4.3
+immutabledict==4.2.1
 ninja==1.11.1 # Needed to avoid `ninja 1.11.1.1 is not supported on this platform` error
+nltk==3.9.1
 optimum==1.17.1
 peft==0.12.0
 pynvml==11.5.3
+rouge_score==0.1.2
 torch==2.2.2+cu118
 torchvision==0.17.2+cu118
-transformers==4.43.1
-trl==0.9.6
+transformers==4.47.1
+trl==0.11.2
 wandb==0.17.1
+ydata-profiling==4.7.0 # Upgrade the version from 4.6.0 to 4.7.0 to fix the old `pydantic` package error.
+psutil==6.0.0

community-content/vertex_model_garden/model_oss/peft/train/vmg/dockerfile/train.Dockerfile

@@ -23,6 +23,14 @@ RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples
 ENV PIP_ROOT_USER_ACTION=ignore
 RUN pip install --upgrade pip
 
+# Remove packages that are not needed and are causing conflicts.
+# dataproc_jupyter_plugin was installed as a part of pytorch-cu121.2-2.py310
+# container which we don't need. It depends on ibis-framework and bigframes.
+# The package and its dependencies request lower versions of pyarrow/pydantic
+# than deepspeed/datasets. So, dataproc_jupyter_plugin conflicts with
+# deepspeed/datasets.
+RUN pip uninstall -y dataproc_jupyter_plugin ibis-framework bigframes
+
 # Prefer to install with requirement file as much as possible for reasons
 # described in b/355034754.
 COPY model_oss/peft/train/vmg/dockerfile/requirements.txt /tmp/requirements.txt
@@ -58,15 +66,14 @@ WORKDIR /diffusers/examples
 
 RUN mkdir -p ./vertex_vision_model_garden_peft/
 COPY model_oss/peft/train/vmg/configs/* ./vertex_vision_model_garden_peft/
-# custom `lm_eval` task.
-ARG LM_EVAL_DIR=$(python -c 'import site; print(site.getsitepackages()[0])')/lm_eval
-RUN mkdir -p $LM_EVAL_DIR/tasks/vertex && \
-  mv ./vertex_vision_model_garden_peft/custom_loglikelihood.yaml $LM_EVAL_DIR/tasks/vertex/
 COPY model_oss/peft/train/vmg/*.py ./vertex_vision_model_garden_peft/train/vmg/
 COPY model_oss/peft/train/vmg/templates /diffusers/examples/util/templates
-COPY model_oss/util /diffusers/examples/util
+COPY model_oss/peft/train/util/*.py /diffusers/examples/util/
+COPY model_oss/util/* /diffusers/examples/util/
 COPY model_oss/notebook_util/dataset_validation_util.py /diffusers/examples/util
-COPY model_oss/peft/train/tests/*.py ./vertex_vision_model_garden_peft/tests/
+COPY model_oss/peft/train/vmg/tests/*.py ./vertex_vision_model_garden_peft/tests/
+COPY model_oss/peft/train/test_utils/test_util.py ./vertex_vision_model_garden_peft/tests/
+COPY model_oss/peft/train/test_utils/command_builder.py ./vertex_vision_model_garden_peft/tests/
 
 RUN chmod a+rwX -R /diffusers/examples/
 ENV PYTHONPATH /diffusers/examples/

community-content/vertex_model_garden/model_oss/peft/train/vmg/eval_lib.py

@@ -1,183 +1,284 @@
 """Library for running evaluations during training."""
 
+from collections.abc import Callable, Mapping, MutableMapping, Sequence
 import dataclasses
-from typing import Any, Optional, Type
+import string
+from typing import Type
 
 from absl import logging
-import datasets
-from lm_eval import evaluator
-from lm_eval import tasks
-from lm_eval import utils
-from lm_eval.api import model as lm_model
-from lm_eval.api import registry
-from lm_eval.models import huggingface
-from peft import peft_model
+import evaluate
+import numpy as np
+import torch
 import transformers
-from transformers import trainer
 
 from util import dataset_validation_util
 from util import constants
 
 
-_DESCRIPTION_EVALUATION = "evaluation"
-_BUILTIN_EVAL_TASK = "builtin_eval"
+_STRING_TRANSLATOR = str.maketrans("", "", string.punctuation)
+
+_GREATER_IS_BETTER_MAP = {
+    "loss": False,
+    "perplexity": False,
+    "bleu": True,
+    "google_bleu": True,
+    "rouge1": True,
+    "rouge2": True,
+    "rougeL": True,
+    "rougeLsum": True,
+}
 
 
 @dataclasses.dataclass(frozen=True)
 class EvalConfig:
+  """Configuration for running evaluations during training.
+
+  Attributes:
+    steps: The number of steps to run evaluation.
+    tasks: The list of tasks to run evaluation on.
+    per_device_batch_size: The per device batch size for evaluation.
+    limit: The maximum number of examples to evaluate.
+    metric_name: The name of the metric to compute.
+    tokenize_dataset: Whether to tokenize the dataset.
+    dataset_path: The path to the dataset.
+    split: The split of the dataset to evaluate.
+    template: The template to use for the dataset.
+    column: The column name of the dataset.
+    metric_for_best_model: The metric to use for loading the best model.
+  """
+
   steps: int
-  tasks: list[str]
   per_device_batch_size: int
-  num_fewshot: Optional[int]
-  limit: Optional[float]
-  metric_name: str
+  limit: float | None
+  metric_name: Sequence[str]
   tokenize_dataset: bool
   dataset_path: str = ""
   split: str = "test"
   template: str = ""
-  column: str = constants.DEFAULT_INSTRUCT_COLUMN_IN_DATASET
-
-
-class PeftCausalLMModel(huggingface.HFLM):
-  """PeftCausalLMModel that supports loading an in-memory model."""
-
-  AUTO_MODEL_CLASS = transformers.AutoModelForCausalLM
-
-  def __init__(
-      self,
-      model: peft_model.PeftModelForCausalLM,
-      tokenizer: transformers.PreTrainedTokenizerBase,
-      batch_size_per_gpu: int,
-  ):
-    lm_model.LM.__init__(self)
-    self._model = model
-    self.tokenizer = tokenizer
-    self.vocab_size = tokenizer.vocab_size
-    tokenizer.pad_token_id = tokenizer.eos_token_id
-    self._config = model.config
-    self.batch_size_per_gpu = batch_size_per_gpu
-    self._device = model.device
-    self._max_length = None  # Will be automatically determined from config.
-    self._add_special_tokens = (
-        None  # Will be automatically determined from AUTO_MODEL_CLASS.
-    )
+  column: str = constants.DEFAULT_TRAIN_COLUMN
+  metric_for_best_model: str | None = None
 
 
 def create_trainer(
     cls: Type[transformers.Trainer],
-    eval_config: Optional[EvalConfig],
-    tokenizer: Optional[transformers.PreTrainedTokenizerBase],
-    args: trainer.TrainingArguments,
+    eval_config: EvalConfig | None,
+    tokenizer: transformers.PreTrainedTokenizerBase | None,
+    args: transformers.TrainingArguments,
     **kwargs,
 ) -> transformers.Trainer:
-  """Creates a trainer. If eval config is provided, injects evaluation loop."""
+  """Creates a trainer. If eval config is provided, injects evaluation loop.
+
+  Args:
+    cls: The trainer class.
+    eval_config: The evaluation config.
+    tokenizer: The tokenizer.
+    args: The training arguments.
+    **kwargs: The keyword arguments.
+
+  Returns:
+    A trainer.
+  """
   if not eval_config:
     return cls(args=args, **kwargs)
 
   args.eval_strategy = "steps"
   args.eval_steps = eval_config.steps
   args.per_device_eval_batch_size = eval_config.per_device_batch_size
+  args.metric_for_best_model = eval_config.metric_for_best_model
+  args.greater_is_better = _GREATER_IS_BETTER_MAP.get(
+      eval_config.metric_for_best_model, None
+  )
+  args.save_strategy = (
+      transformers.trainer_utils.SaveStrategy.STEPS
+      if eval_config.metric_for_best_model is None
+      else transformers.trainer_utils.SaveStrategy.BEST
+  )
+
   kwargs["tokenizer"] = tokenizer
 
-  if eval_config.tasks == [_BUILTIN_EVAL_TASK]:
-    try:
-      eval_dataset = dataset_validation_util.load_dataset_with_template(
-          dataset_name=eval_config.dataset_path,
-          split=eval_config.split,
+  try:
+    _, eval_dataset = dataset_validation_util.load_dataset_with_template(
+        dataset_name=eval_config.dataset_path,
+        split=eval_config.split,
+        input_column=eval_config.column,
+        template=eval_config.template,
+        tokenizer=tokenizer,
+    )
+    if eval_config.limit is not None:
+      if eval_config.limit >= 1:
+        limit = int(eval_config.limit)
+      else:
+        limit = int(eval_config.limit * len(eval_dataset))
+      eval_dataset = eval_dataset.select(range(limit))
+    if tokenizer is not None:
+      eval_dataset = dataset_validation_util.get_filtered_dataset(
+          dataset=eval_dataset,
           input_column=eval_config.column,
-          template=eval_config.template,
+          max_seq_length=kwargs["max_seq_length"],
           tokenizer=tokenizer,
       )
-      if eval_config.limit is not None:
-        if eval_config.limit >= 1:
-          limit = int(eval_config.limit)
-        else:
-          limit = int(eval_config.limit * len(eval_dataset))
-        eval_dataset = eval_dataset.select(range(limit))
-      if eval_config.tokenize_dataset:
-        eval_dataset = eval_dataset.map(
-            lambda samples: tokenizer(samples[eval_config.column])
-        )
-      kwargs["eval_dataset"] = eval_dataset
-    except (OSError, ValueError, IndexError) as e:
-      logging.warning(
-          "Failed to load eval dataset %s. Evaluation will be skipped.\n%s",
-          eval_config.dataset_path,
-          e,
+    if eval_config.tokenize_dataset:
+      eval_dataset = eval_dataset.map(
+          lambda samples: tokenizer(samples[eval_config.column])
       )
-      del args.evaluation_strategy
-      del args.eval_steps
-      del args.per_device_eval_batch_size
-    return cls(args=args, **kwargs)
+    kwargs["eval_dataset"] = eval_dataset
+  except (OSError, ValueError, IndexError) as e:
+    logging.warning(
+        "Failed to load eval dataset %s. Evaluation will be skipped.\n%s",
+        eval_config.dataset_path,
+        e,
+    )
+    del args.evaluation_strategy
+    del args.eval_steps
+    del args.per_device_eval_batch_size
+  return cls(args=args, **kwargs)
 
-  class LMEvalTrainer(cls):
-    """Trainer with lm_eval injected as the eval library."""
 
-    def __init__(self, **kwargs):
-      super().__init__(**kwargs)
-      task_names = utils.pattern_match(eval_config.tasks, registry.ALL_TASKS)
-      logging.info("Selected Eval Tasks: %s", task_names)
-      task_args = {}
-      if eval_config.num_fewshot is not None:
-        task_args["num_fewshot"] = eval_config.num_fewshot
-      if eval_config.dataset_path:
-        task_args["dataset_path"] = "json"
-        task_args["dataset_kwargs"] = {
-            "data_files": {"test": eval_config.dataset_path},
-        }
-      self._eval_task_dict = tasks.get_task_dict(task_names, **task_args)
+def _cleanup_text(text: str) -> str:
+  """Cleans up the prediction and references text.
 
-    def evaluation_loop(
-        self,
-        dataloader: trainer.DataLoader,
-        description: str,
-        prediction_loss_only: Optional[bool] = None,
-        ignore_keys: Optional[list[str]] = None,
-        metric_key_prefix: str = "eval",
-    ) -> trainer.EvalLoopOutput:
-      """Custom evaluation loop that invokes lm_eval."""
-      if description.lower() != _DESCRIPTION_EVALUATION:
-        return super().evaluation_loop(
-            dataloader,
-            description,
-            prediction_loss_only,
-            ignore_keys,
-            metric_key_prefix,
-        )
+  Args:
+    text: The text to clean up.
 
-      model = self._wrap_model(self.model, training=False)
-      lm = PeftCausalLMModel(
-          model,
-          self.tokenizer or self.data_collator.tokenizer,
-          eval_config.per_device_batch_size,
-      )
-      results: dict[str, Any] = evaluator.evaluate(
-          lm=lm,
-          task_dict=self._eval_task_dict,
-          limit=eval_config.limit,
-      )["results"]
-      metric_name = eval_config.metric_name
-      # Compute average value if there are multiple tasks.
-      metric_values: list[float] = []
-      for result in results.values():
-        for key, value in result.items():
-          if key.split(",")[0] == metric_name:
-            metric_values.append(value)
-      if not metric_values:
-        raise ValueError(
-            f"Metric {metric_name} not found in eval response: {results}"
-        )
-      metric_average = sum(metric_values) / len(metric_values)
-      logging.info("%s value: %f\n%s", metric_name, metric_average, results)
-      return trainer.EvalLoopOutput(
-          # Only metrics field is set. Other fields are dummy values.
-          predictions=None,
-          label_ids=None,
-          metrics={f"{metric_key_prefix}_{metric_name}": metric_average},
-          num_samples=0,
-      )
+  Returns:
+    Cleaned up text.
+  """
+  text = text.translate(_STRING_TRANSLATOR)
+  text = text.strip()
+  text = " ".join(text.split())
+  return text.lower()
 
-  # Use empty eval dataset as a placeholder.
-  return LMEvalTrainer(
-      args=args, eval_dataset=datasets.Dataset.from_dict({"test": []}), **kwargs
+
+def create_compute_metrics(
+    tokenizer: transformers.PreTrainedTokenizerBase,
+    eval_metrics: Mapping[str, evaluate.EvaluationModule],
+) -> Callable[[transformers.EvalPrediction], MutableMapping[str, float]]:
+  """Creates a compute_metrics function using Hugging Face evaluate library.
+
+  Args:
+    tokenizer: The tokenizer for decoding predictions.
+    eval_metrics: The eval metrics to compute.
+
+  Returns:
+    Function that computes comprehensive metrics.
+  """
+
+  def _preprocess_data(
+      predictions: np.ndarray, labels: np.ndarray
+  ) -> tuple[Sequence[str], Sequence[str]]:
+    """Preprocesses predictions and lavels before evaluation.
+
+    Args:
+      predictions: The predictions to preprocess.
+      labels: The labels to preprocess.
+
+    Returns:
+      A tuple (preprocessed predictions, labels).
+    """
+    # Handle padding and special tokens.
+    predictions = np.where(
+        predictions != -100, predictions, tokenizer.pad_token_id
+    )
+    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
+
+    # Decode to text.
+    pred_texts = tokenizer.batch_decode(predictions, skip_special_tokens=True)
+    label_texts = tokenizer.batch_decode(labels, skip_special_tokens=True)
+
+    # Clean up text.
+    cleaned_pred_texts = [_cleanup_text(text) for text in pred_texts]
+    cleaned_label_texts = [_cleanup_text(text) for text in label_texts]
+
+    return cleaned_pred_texts, cleaned_label_texts
+
+  def _compute_metrics_with_tokenizer(
+      eval_pred: transformers.EvalPrediction,
+  ) -> MutableMapping[str, float]:
+    """Computes metrics using Hugging Face evaluate library.
+
+    Args:
+      eval_pred: The evaluation prediction.
+
+    Returns:
+      A dictionary of metrics.
+    """
+    predictions, perplexities = eval_pred.predictions
+    labels = eval_pred.label_ids
+
+    pred_texts, label_texts = _preprocess_data(predictions, labels)
+
+    metrics = {}
+
+    for eval_metric, computed_eval_metric in eval_metrics.items():
+      match eval_metric:
+        case "perplexity":
+          # We don't use the perplexity from HF Evaluate since it loads the
+          # model again. This causes an increase in the GPU utilization and
+          # hence an OOM. Due to this, we compute the perplexity ourselves
+          # using the eval_loss over the unmasked tokens in
+          # preprocess_logits_for_metrics fn.
+          metrics[eval_metric] = np.mean(perplexities)
+        case "bleu" | "google_bleu":
+          num_valid_labels = len(list(filter(None, label_texts)))
+          if num_valid_labels:
+            eval_score = computed_eval_metric.compute(
+                predictions=pred_texts,
+                references=[[text] for text in label_texts],
+            )
+            metrics[eval_metric] = eval_score[eval_metric]
+          else:
+            metrics[eval_metric] = 0.0
+        case "rouge1" | "rouge2" | "rougeL" | "rougeLsum":
+          rouge_scores = computed_eval_metric.compute(
+              predictions=pred_texts,
+              references=label_texts,
+              use_stemmer=True,
+          )
+          metrics[eval_metric] = rouge_scores[eval_metric]
+
+    pred_lengths = [len(pred.split()) for pred in pred_texts]
+    label_lengths = [len(label.split()) for label in label_texts]
+
+    metrics["gen_len"] = np.mean(pred_lengths)
+    metrics["ref_len"] = np.mean(label_lengths)
+    metrics["length_ratio"] = np.mean(
+        [len(p) / len(r) if r else 0 for p, r in zip(pred_texts, label_texts)]
+    )
+
+    # Round all metrics to 4 decimal places.
+    return {k: round(float(v), 4) for k, v in metrics.items()}
+
+  return _compute_metrics_with_tokenizer
+
+
+def preprocess_logits_for_metrics(
+    logits: torch.Tensor, labels: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+  """Preprocesses the logits before caching them for eval metric calculation.
+
+  Args:
+    logits: Logits predicted by the model.
+    labels: Ground truth labels.
+
+  Returns:
+    A tuple (pred_ids, perplexities).
+  """
+  # Calculate prediction IDs.
+  pred_ids = logits.argmax(dim=-1)
+
+  # This step shifts the logits and labels to align them correctly, where we are
+  # predicting the next token in a sequence. The last logit doesn't have a
+  # corresponding label, and the first label doesn't have a preceding logit to
+  # predict it. This calculation of perplexity is inspired from
+  # https://github.com/huggingface/evaluate/blob/main/metrics/perplexity/perplexity.py.
+  shift_logits = logits[..., :-1, :].contiguous()
+  shift_labels = labels[..., 1:].contiguous()
+  attn_mask = shift_labels != -100
+  loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
+
+  perplexities = torch.exp(
+      (loss_fct(shift_logits.transpose(1, 2), shift_labels) * attn_mask).sum(1)
+      / attn_mask.sum(1)
   )
+
+  return (pred_ids, perplexities)

community-content/vertex_model_garden/model_oss/peft/train/vmg/instruct_lora.py

@@ -1,10 +1,10 @@
 """Instruct/Chat with LoRA models."""
 
-import dataclasses
+from collections.abc import Callable, Mapping, Sequence
 import datetime
 import json
 import os
-from typing import Any, Dict, Optional, Sequence
+from typing import Any
 import warnings
 
 from absl import app
@@ -13,17 +13,17 @@ from absl import logging
 from accelerate import DistributedType
 from accelerate import PartialState
 import bitsandbytes as bnb
-import hypertune
+import evaluate
 from peft import get_peft_model
 from peft import LoraConfig
 import torch
-from transformers import AutoModelForCausalLM
-from transformers import TrainingArguments
-from trl import DataCollatorForCompletionOnlyLM
-from trl import SFTTrainer
+import transformers
+import trl
 import wandb
 
 from util import dataset_validation_util
+from util import dataset_stats
+from util import device_stats
 from vertex_vision_model_garden_peft.train.vmg import callbacks
 from vertex_vision_model_garden_peft.train.vmg import eval_lib
 from vertex_vision_model_garden_peft.train.vmg import utils
@@ -31,14 +31,15 @@ from util import constants
 from util import fileutils
 
 
-_PRETRAINED_MODEL_ID = flags.DEFINE_string(
-    'pretrained_model_id',
+_PRETRAINED_MODEL_NAME_OR_PATH = flags.DEFINE_string(
+    'pretrained_model_name_or_path',
     None,
-    'The pretrained model id. Supported models can be causal language modeling'
-    ' models from https://github.com/huggingface/peft/tree/main. Note, there'
-    ' might be different paddings for different models. This tool assumes the'
-    ' pretrained_model_id contains model name, and then choose proper padding'
-    ' methods. e.g. it must contain `llama` for `Llama2 models`.',
+    'The pretrained model name or path. Supported models can be causal language'
+    ' modeling models from https://github.com/huggingface/peft/tree/main. Note,'
+    ' there might be different paddings for different models. This tool assumes'
+    ' the pretrained_model_name_or_path contains model name, and then choose'
+    ' proper padding methods. e.g. it must contain `llama` for `Llama2'
+    ' models`.',
     required=True,
 )
 
@@ -48,10 +49,10 @@ _HUGGINGFACE_ACCESS_TOKEN = flags.DEFINE_string(
     'The access token for loading huggingface gated models.',
 )
 
-_DATASET_NAME = flags.DEFINE_string(
-    'dataset_name',
+_TRAIN_DATASET = flags.DEFINE_string(
+    'train_dataset',
     None,
-    'The dataset name in huggingface.',
+    'The training dataset name in huggingface or path.',
 )
 
 _OUTPUT_DIR = flags.DEFINE_string(
@@ -119,8 +120,8 @@ _WEIGHT_DECAY = flags.DEFINE_float(
     'The weight decay in the learning rate scheduler.',
 )
 
-_NUM_EPOCHS = flags.DEFINE_float(
-    'num_epochs',
+_NUM_TRAIN_EPOCHS = flags.DEFINE_float(
+    'num_train_epochs',
     None,
     'The number of training epochs. Only used for'
     ' "sequence-classification-lora" with an integer value and for'
@@ -130,8 +131,8 @@ _NUM_EPOCHS = flags.DEFINE_float(
 _MAX_STEPS = flags.DEFINE_integer(
     'max_steps',
     None,
-    'Total number of training steps. Overrides num_epochs if set. Only used for'
-    ' "instruct-lora."',
+    'Total number of training steps. Overrides num_train_epochs if set. Only'
+    ' used for "instruct-lora."',
 )
 
 _MAX_SEQ_LENGTH = flags.DEFINE_integer(
@@ -146,9 +147,9 @@ _LEARNING_RATE = flags.DEFINE_float(
     'The learning rate after the potential warmup period.',
 )
 
-_INSTRUCT_COLUMN_IN_DATASET = flags.DEFINE_string(
-    'instruct_column_in_dataset',
-    constants.DEFAULT_INSTRUCT_COLUMN_IN_DATASET,
+_TRAIN_COLUMN = flags.DEFINE_string(
+    'train_column',
+    constants.DEFAULT_TRAIN_COLUMN,
     'The instruct column in dataset.',
 )
 
@@ -170,8 +171,8 @@ _GRADIENT_ACCUMULATION_STEPS = flags.DEFINE_integer(
     'The gradient accumulation steps.',
 )
 
-_ENABLE_GRADIENT_CHECKPOINTING = flags.DEFINE_boolean(
-    'enable_gradient_checkpointing',
+_GRADIENT_CHECKPOINTING = flags.DEFINE_boolean(
+    'gradient_checkpointing',
     False,
     'Whether to enable gradient checkpointing.',
 )
@@ -181,8 +182,8 @@ _ENABLE_PEFT = flags.DEFINE_boolean(
     True,
     'Whether to enable peft.',
 )
-_TEMPLATE = flags.DEFINE_string(
-    'template',
+_TRAIN_TEMPLATE = flags.DEFINE_string(
+    'train_template',
     None,
     'Template for formatting language model training data. Must be a filename'
     ' under `templates` folder, without `.json` extension, e.g. `alpaca`, or a'
@@ -219,32 +220,18 @@ _EVAL_STEPS = flags.DEFINE_integer(
     'The number of training steps between evaluations.',
 )
 
-_TRAIN_SPLIT_NAME = flags.DEFINE_string(
-    'train_split_name',
+_TRAIN_SPLIT = flags.DEFINE_string(
+    'train_split',
     'train',
     'The train split name.',
 )
 
-_EVAL_TASKS = flags.DEFINE_list(
-    'eval_tasks',
-    None,
-    'List of eval task names (can have wildcards) as in'
-    ' https://github.com/EleutherAI/lm-evaluation-harness. Will not run'
-    ' evaluation if not set. Runs the built-in trainer evaluation loop if set'
-    ' to `builtin_eval`.',
-)
-
-_EVAL_PER_DEVICE_BATCH_SIZE = flags.DEFINE_integer(
-    'eval_per_device_batch_size',
+_PER_DEVICE_EVAL_BATCH_SIZE = flags.DEFINE_integer(
+    'per_device_eval_batch_size',
     1,
     'The per device batch size for model evaluation.',
 )
 
-_EVAL_NUM_FEWSHOT = flags.DEFINE_integer(
-    'eval_num_fewshot',
-    None,
-    'Run N-shot language model evaluation. Not implemented in `builtin_eval`.',
-)
 
 _EVAL_LIMIT = flags.DEFINE_float(
     'eval_limit',
@@ -253,17 +240,18 @@ _EVAL_LIMIT = flags.DEFINE_float(
     ' total number of examples.',
 )
 
-_EVAL_METRIC_NAME = flags.DEFINE_string(
+_EVAL_METRIC_NAME = flags.DEFINE_list(
     'eval_metric_name',
-    'acc',
-    'The metric name to aggregate during model evaluation.',
+    ['loss'],
+    'A comma-separated list of metric names to aggregate during model'
+    ' evaluation. The supported metrics are: '
+    + ', '.join(constants.SUPPORTED_EVAL_METRICS),
 )
 
-_EVAL_DATASET_PATH = flags.DEFINE_string(
-    'eval_dataset_path',
+_EVAL_DATASET = flags.DEFINE_string(
+    'eval_dataset',
     None,
-    'Overrides the default evaluation dataset path. In `builtin_eval` mode,'
-    ' this can be any Hugging Face dataset name.',
+    'The Hugging Face dataset name or path to use for evaluation.',
 )
 
 # We set the default eval split as `test`, based on observation from
@@ -271,13 +259,13 @@ _EVAL_DATASET_PATH = flags.DEFINE_string(
 _EVAL_SPLIT = flags.DEFINE_string(
     'eval_split',
     'test',
-    'Eval split name in the eval dataset for `builtin_eval`.',
+    'Eval split name in the eval dataset.',
 )
 
 _EVAL_TEMPLATE = flags.DEFINE_string(
     'eval_template',
     None,
-    'Template for formatting language model evaluation data for `builtin_eval`.'
+    'Template for formatting language model evaluation data.'
     ' Must be a filename under `templates` folder, without `.json` extension,'
     ' e.g. `alpaca`, or a Cloud Storage URI to a JSON file.',
 )
@@ -285,7 +273,14 @@ _EVAL_TEMPLATE = flags.DEFINE_string(
 _EVAL_COLUMN = flags.DEFINE_string(
     'eval_column',
     None,
-    'Eval column name in the eval dataset for `builtin_eval`.',
+    'Eval column name in the eval dataset.',
+)
+
+_METRIC_FOR_BEST_MODEL = flags.DEFINE_string(
+    'metric_for_best_model',
+    None,
+    'If set, the best model is saved at the end of training based on the'
+    ' metric',
 )
 
 _TRAIN_PRECISION = flags.DEFINE_enum(
@@ -299,15 +294,15 @@ _TRAIN_PRECISION = flags.DEFINE_enum(
     'Precision to train the model.',
 )
 
-_USE_EXAMPLE_PACKING = flags.DEFINE_boolean(
-    'use_example_packing',
+_EXAMPLE_PACKING = flags.DEFINE_boolean(
+    'example_packing',
     False,
     'Enables example packing during training, which uses '
     '`ConstantLengthDataset` under the hood.',
 )
 
-_COMPLETION_ONLY = flags.DEFINE_boolean(
-    'completion_only',
+_INPUT_MASKING = flags.DEFINE_boolean(
+    'input_masking',
     False,
     'If set, it uses DataCollatorForCompletionOnlyLM to train the model on the'
     ' generated prompts only, i.e., masking out the input',
@@ -356,28 +351,53 @@ _TARGET_MODULES = flags.DEFINE_list(
     'target_modules', None, 'The names of the modules to apply LoRA adapter to.'
 )
 
+_MAX_GPU_MEMORY_FRACTION = flags.DEFINE_float(
+    'max_gpu_memory_fraction',
+    '0.9',
+    'Maximum GPU memory a caching allocator is allowed to use per GPU.',
+)
+
 
 @flags.multi_flags_validator(
     [
-        _COMPLETION_ONLY.name,
-        _USE_EXAMPLE_PACKING.name,
+        _INPUT_MASKING.name,
+        _EXAMPLE_PACKING.name,
     ],
-    message=(
-        '`use_example_packing=True` does not work with `completion_only=True`'
-    ),
+    message='`example_packing=True` does not work with `input_masking=True`',
 )
-def check_example_packing(flags_dict: Dict[str, Any]) -> bool:
+def check_example_packing(flags_dict: Mapping[str, Any]) -> bool:
   """Check to make sure example packing is enabled properly.
 
   Args:
     flags_dict: Dictionary containing flags to check.
 
   Returns:
-    If `use_example_packing` is set properly.
+    If `example_packing` is set properly.
+  """
+  if flags_dict[_INPUT_MASKING.name] and flags_dict[_EXAMPLE_PACKING.name]:
+    return False
+  return True
+
+
+@flags.multi_flags_validator(
+    [
+        _INPUT_MASKING.name,
+        _TRAIN_TEMPLATE.name,
+    ],
+    message='`train_template` should be provided if using `input_masking=True`',
+)
+def check_input_masking(flags_dict: Mapping[str, Any]) -> bool:
+  """Check to make sure input_masking is enabled properly.
+
+  Args:
+    flags_dict: Dictionary containing flags to check.
+
+  Returns:
+    If `input_masking` is set properly
   """
   if (
-      flags_dict[_COMPLETION_ONLY.name]
-      and flags_dict[_USE_EXAMPLE_PACKING.name]
+      flags_dict[_INPUT_MASKING.name]
+      and flags_dict[_TRAIN_TEMPLATE.name] is None
   ):
     return False
   return True
@@ -385,22 +405,65 @@ def check_example_packing(flags_dict: Dict[str, Any]) -> bool:
 
 @flags.multi_flags_validator(
     [
-        _COMPLETION_ONLY.name,
-        _TEMPLATE.name,
+        _EVAL_DATASET.name,
+        _EVAL_METRIC_NAME.name,
     ],
-    message='`template` should be provided if using `completion_only=True`',
+    message=(
+        '`eval_metric_name` should be a valid metric name and present when'
+        ' eval_dataset is provided.'
+    ),
 )
-def check_completion_only(flags_dict: Dict[str, Any]) -> bool:
-  """Check to make sure completion_only is enabled properly.
+def _validate_eval_metrics(flags_dict: Mapping[str, Any]) -> bool:
+  """Validates the eval metric name.
 
   Args:
     flags_dict: Dictionary containing flags to check.
 
   Returns:
-    If `completion_only` is set properly
+    If the eval metrics are valid.
   """
-  if flags_dict[_COMPLETION_ONLY.name] and flags_dict[_TEMPLATE.name] is None:
-    return False
+  if flags_dict[_EVAL_DATASET.name] is None:
+    return True
+  eval_metrics = flags_dict[_EVAL_METRIC_NAME.name]
+  for eval_metric in eval_metrics:
+    if eval_metric not in constants.SUPPORTED_EVAL_METRICS:
+      raise flags.ValidationError(f'Invalid eval metric: {eval_metric}')
+  if 'perplexity' in eval_metrics and 'loss' not in eval_metrics:
+    _EVAL_METRIC_NAME.value.append('loss')
+    logging.warning(
+        'Adding `loss` to eval_metric_name because `perplexity` is present.'
+    )
+  return True
+
+
+@flags.multi_flags_validator(
+    [
+        _METRIC_FOR_BEST_MODEL.name,
+        _EVAL_METRIC_NAME.name,
+    ],
+    message='`metric_for_best_model` should be in `eval_metric_name`.',
+)
+def _validate_metric_for_best_model(flags_dict: Mapping[str, Any]) -> bool:
+  """Validates the metric for best model.
+
+  Args:
+    flags_dict: Dictionary containing flags to check.
+
+  Returns:
+    If the metric for best model is valid.
+  """
+  if flags_dict[_METRIC_FOR_BEST_MODEL.name] is None:
+    return True
+
+  metric_for_best_model = flags_dict[_METRIC_FOR_BEST_MODEL.name]
+  eval_metric_name = flags_dict[_EVAL_METRIC_NAME.name]
+
+  if metric_for_best_model not in eval_metric_name:
+    raise flags.ValidationError(
+        'Invalid metric for picking the best model:'
+        f' {metric_for_best_model}. The metric should be one'
+        f' of the {eval_metric_name}.'
+    )
   return True
 
 
@@ -411,10 +474,43 @@ def check_completion_only(flags_dict: Dict[str, Any]) -> bool:
 # https://github.com/huggingface/notebooks/blob/main/sagemaker/28_train_llms_with_qlora/scripts/run_clm.py.
 
 
+def _calculate_hf_eval_metrics(
+    tokenizer: transformers.PreTrainedTokenizerBase,
+    eval_config: eval_lib.EvalConfig | None,
+) -> tuple[
+    Callable[[transformers.EvalPrediction], Mapping[str, float]], torch.Tensor
+]:
+  """Calculates the HF evaluation metrics.
+
+  Args:
+    tokenizer: The tokenizer to use for evaluation.
+    eval_config: The evaluation config to use.
+
+  Returns:
+    The compute metrics and preprocess logits for metrics.
+  """
+  if eval_config is None:
+    return None, None
+  hf_eval_metrics = {}
+  for metric in eval_config.metric_name:
+    if metric in constants.SUPPORTED_HF_EVAL_METRICS:
+      if metric in constants.ROUGE_VARIANTS:
+        hf_eval_metrics[metric] = evaluate.load('rouge')
+      else:
+        hf_eval_metrics[metric] = evaluate.load(metric)
+
+  if not hf_eval_metrics:
+    return None, None
+  return (
+      eval_lib.create_compute_metrics(tokenizer, hf_eval_metrics),
+      eval_lib.preprocess_logits_for_metrics,
+  )
+
+
 # Copied from https://github.com/artidoro/qlora/blob/main/qlora.py.
 def find_all_linear_names(
-    model: AutoModelForCausalLM, precision_mode: str
-) -> list[str]:
+    model: transformers.AutoModelForCausalLM, precision_mode: str
+) -> Sequence[str]:
   """Finds all linear module names."""
   if precision_mode == constants.PRECISION_MODE_4:
     cls = bnb.nn.Linear4bit
@@ -433,47 +529,51 @@ def find_all_linear_names(
 
 
 def finetune_instruct(
-    pretrained_model_id: str,
-    dataset_name: str,
+    pretrained_model_name_or_path: str,
+    train_dataset: str,
     output_dir: str,
     logging_output_dir: str,
     lora_rank: int = 64,
     lora_alpha: int = 16,
     lora_dropout: float = 0.1,
     warmup_ratio: int = 0.03,
-    num_epochs: Optional[float] = None,
-    max_steps: Optional[int] = None,
+    num_train_epochs: float | None = None,
+    max_steps: int | None = None,
     warmup_steps: int = 10,
     max_seq_length: int = 512,
     learning_rate: float = 2e-4,
     precision_mode: str = None,
-    instruct_column_in_dataset: str = constants.DEFAULT_INSTRUCT_COLUMN_IN_DATASET,
+    train_column: str = constants.DEFAULT_TRAIN_COLUMN,
     per_device_train_batch_size: int = 4,
     gradient_accumulation_steps: int = 4,
     optim: str = 'paged_adamw_32bit',
     weight_decay: float = 0.001,
-    enable_gradient_checkpointing: bool = False,
+    gradient_checkpointing: bool = False,
     enable_peft: bool = True,
-    template: str = None,
+    train_template: str = None,
     lr_scheduler_type: str = 'constant',
     save_steps: int = 10,
     logging_steps: int = 10,
-    train_split_name: str = 'train',
-    eval_config: Optional[eval_lib.EvalConfig] = None,
+    train_split: str = 'train',
+    eval_config: eval_lib.EvalConfig | None = None,
     report_to: str = constants.REPORT_TO_NONE,
-    access_token: Optional[str] = None,
+    access_token: str | None = None,
     train_precision: str = constants.PRECISION_MODE_16B,
-    use_example_packing: bool = False,
-    attn_implementation: Optional[str] = None,
+    example_packing: bool = False,
+    attn_implementation: str | None = None,
     max_grad_norm: float = 0.3,
-    completion_only: bool = False,
+    input_masking: bool = False,
     logger_level: str = 'passive',
-    benchmark_out_file: Optional[str] = None,
-    tuning_data_stats_file: Optional[str] = None,
-    target_modules: Optional[str] = None,
+    benchmark_out_file: str | None = None,
+    tuning_data_stats_file: str | None = None,
+    target_modules: str | None = None,
 ) -> None:
   """Finetunes instruct."""
-  logging.info('on entering instruct_lora, %s', utils.gpu_stats_str())
+  logging.info(
+      'on entering instruct_lora, %s,\n%s',
+      device_stats.gpu_stats_str(),
+      device_stats.cpu_stats_str(),
+  )
   gradient_checkpointing_kwargs = {}
   # DDP provides limited support with the reentrant variant of gradient
   # checkpoint [1]. Below is an indirect way of checking whether DDP will be
@@ -483,17 +583,25 @@ def finetune_instruct(
   if PartialState().distributed_type == DistributedType.MULTI_GPU:
     gradient_checkpointing_kwargs['use_reentrant'] = False
 
-  tokenizer = utils.load_tokenizer(
-      pretrained_model_id,
+  tokenizer = dataset_validation_util.load_tokenizer(
+      pretrained_model_name_or_path,
       'right',
       access_token=access_token,
   )
 
-  train_dataset = dataset_validation_util.load_dataset_with_template(
-      dataset_name,
-      split=train_split_name,
-      input_column=instruct_column_in_dataset,
-      template=template,
+  train_dataset, train_dataset_with_template = (
+      dataset_validation_util.load_dataset_with_template(
+          train_dataset,
+          split=train_split,
+          input_column=train_column,
+          template=train_template,
+          tokenizer=tokenizer,
+      )
+  )
+  train_dataset_with_template = dataset_validation_util.get_filtered_dataset(
+      dataset=train_dataset_with_template,
+      input_column=train_column,
+      max_seq_length=max_seq_length,
       tokenizer=tokenizer,
   )
 
@@ -508,24 +616,26 @@ def finetune_instruct(
           'getting tuning data stats with effective batch size %s',
           effective_batch_size,
       )
-      train_dataset_stats = utils.get_dataset_stats(
-          train_dataset,
-          tokenizer,
-          instruct_column_in_dataset,
-          effective_batch_size,
+      train_dataset_stats = dataset_stats.get_dataset_stats(
+          raw=train_dataset,
+          templated=train_dataset_with_template,
+          template=train_template,
+          tokenizer=tokenizer,
+          column=train_column,
+          effective_batch_size=effective_batch_size,
       )
       logging.info('stats: %s', train_dataset_stats)
       tuning_data_stats_file = dataset_validation_util.force_gcs_fuse_path(
           tuning_data_stats_file
       )
       with open(tuning_data_stats_file, 'w') as out_f:
-        json.dump(dataclasses.asdict(train_dataset_stats), out_f)
+        json.dump(train_dataset_stats, out_f)
 
   model = utils.load_model(
-      pretrained_model_id=pretrained_model_id,
+      pretrained_model_name_or_path=pretrained_model_name_or_path,
       tokenizer=tokenizer,
       precision_mode=precision_mode,
-      enable_gradient_checkpointing=enable_gradient_checkpointing,
+      gradient_checkpointing=gradient_checkpointing,
       access_token=access_token,
       attn_implementation=attn_implementation,
       train_precision=train_precision,
@@ -550,7 +660,9 @@ def finetune_instruct(
     # `get_peft_model`, which may revert other changes we did before. That's why
     # we are calling `get_peft_model` explicitly here.
     model = get_peft_model(model, peft_config)
-
+    adapter_for_eval_dir = os.path.join(output_dir, 'adapter_for_eval')
+    logging.info('saving adapter for evaluation to %s...', adapter_for_eval_dir)
+    peft_config.save_pretrained(adapter_for_eval_dir)
     # This is to work-around mix-precision training. This issue is not fixed as
     # of transformers==4.41.2.
     # See b/332760883#comment30 for more details.
@@ -568,14 +680,13 @@ def finetune_instruct(
   # b/357970482#comment3
   accelerator_config = {'use_configured_state': True}
 
-  training_arguments = TrainingArguments(
+  training_arguments = transformers.TrainingArguments(
       report_to=report_to,
       output_dir=output_dir,
       per_device_train_batch_size=per_device_train_batch_size,
       gradient_accumulation_steps=gradient_accumulation_steps,
       optim=optim,
       save_steps=save_steps,
-      save_strategy='steps',
       save_total_limit=3,
       logging_dir=os.path.join(logging_output_dir, 'logs'),
       logging_steps=logging_steps,
@@ -583,21 +694,24 @@ def finetune_instruct(
       fp16=(train_precision == constants.PRECISION_MODE_16),
       bf16=(train_precision == constants.PRECISION_MODE_16B),
       max_grad_norm=max_grad_norm,
-      num_train_epochs=num_epochs if num_epochs else -1,
+      num_train_epochs=num_train_epochs if num_train_epochs else -1,
       max_steps=max_steps if max_steps else -1,
       warmup_ratio=warmup_ratio,
       warmup_steps=warmup_steps,
       group_by_length=False,
       lr_scheduler_type=lr_scheduler_type,
-      gradient_checkpointing=enable_gradient_checkpointing,
+      gradient_checkpointing=gradient_checkpointing,
       gradient_checkpointing_kwargs=gradient_checkpointing_kwargs,
       weight_decay=weight_decay,
       log_level=logger_level,
       accelerator_config=accelerator_config,
+      include_num_input_tokens_seen=True,
   )
   trainer_kwargs = {}
-  if completion_only and template:
-    template_json = dataset_validation_util.get_template(template_path=template)
+  if input_masking and train_template:
+    template_json = dataset_validation_util.get_template(
+        template_path=train_template
+    )
     instruction_sep = dataset_validation_util.get_instruction_separator(
         template_json
     )
@@ -608,7 +722,7 @@ def finetune_instruct(
           ' `DataCollatorForCompletionOnlyLM`'
       )
 
-    trainer_kwargs['data_collator'] = DataCollatorForCompletionOnlyLM(
+    trainer_kwargs['data_collator'] = trl.DataCollatorForCompletionOnlyLM(
         instruction_template=instruction_sep,
         response_template=response_sep,
         tokenizer=tokenizer,
@@ -618,17 +732,23 @@ def finetune_instruct(
   trainer_stats_callback = callbacks.TrainerStatsCallback(
       max_seq_length, benchmark_out_file
   )
+  compute_metrics, preprocess_logits = _calculate_hf_eval_metrics(
+      tokenizer, eval_config
+  )
+
   trainer = eval_lib.create_trainer(
-      cls=SFTTrainer,
+      cls=trl.SFTTrainer,
       eval_config=eval_config,
       model=model,
-      train_dataset=train_dataset,
-      dataset_text_field=instruct_column_in_dataset,
+      train_dataset=train_dataset_with_template,
+      dataset_text_field=train_column,
       max_seq_length=max_seq_length,
       tokenizer=tokenizer,
       args=training_arguments,
-      packing=use_example_packing,
+      packing=example_packing,
       callbacks=[trainer_stats_callback],
+      compute_metrics=compute_metrics,
+      preprocess_logits_for_metrics=preprocess_logits,
       **trainer_kwargs,
   )
 
@@ -658,7 +778,7 @@ def finetune_instruct(
     # This method saves the sharded weights like `accelerator.save_state`, see
     # https://huggingface.co/docs/accelerate/en/usage_guides/fsdp#saving-and-loading
     trainer.save_model(output_dir)
-    model = trainer.model.cpu()  # Avoids GPU OOM
+    model = trainer.model
     state_dict = trainer.accelerator.get_state_dict(model)
     # To aggregate the weights from all the devices, we need to use
     # `state_dict=state_dict`.
@@ -668,7 +788,6 @@ def finetune_instruct(
         is_main_process=PartialState().is_main_process,
         save_embedding_layers=False,  # Only pad token is added. See go/lora-adapter-pad-token #pylint: disable=line-too-long
     )
-    model.cuda()  # Move back to GPU to do eval.
   else:
     trainer.model.save_pretrained(
         final_checkpoint,
@@ -683,15 +802,6 @@ def finetune_instruct(
     # https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/trainer_pt_utils.py#L1001 #pylint: disable=line-too-long
     trainer.log_metrics('eval', metrics)
     trainer.save_metrics('eval', metrics)
-    if PartialState().is_main_process:
-      hp_metric = metrics[f'eval_{eval_config.metric_name}']
-      hpt = hypertune.HyperTune()
-      hpt.report_hyperparameter_tuning_metric(
-          hyperparameter_metric_tag=constants.HP_METRIC_TAG,
-          metric_value=hp_metric,
-      )
-      logging.info('Send HP metric: %f to hyperparameter tuning.', hp_metric)
-    PartialState().wait_for_everyone()
 
   if not enable_peft:
     tokenizer.save_pretrained(
@@ -705,38 +815,41 @@ def main(unused_argv: Sequence[str]) -> None:
       timeout=datetime.timedelta(seconds=_NCCL_TIMEOUT.value)
   )
 
+  torch.cuda.set_per_process_memory_fraction(
+      _MAX_GPU_MEMORY_FRACTION.value, device=PartialState().local_process_index
+  )
+
   utils.print_library_versions()
   warnings.simplefilter(_WARNINGS_FILTER.value)
 
-  pretrained_model_id = fileutils.force_gcs_path(_PRETRAINED_MODEL_ID.value)
-  if dataset_validation_util.is_gcs_path(pretrained_model_id):
-    pretrained_model_id = dataset_validation_util.download_gcs_uri_to_local(
-        pretrained_model_id
-    )
-
-  output_dir = utils.GcsOrLocalDirectory(
-      _OUTPUT_DIR.value, check_empty=True, upload_from_all_nodes=True
+  pretrained_model_name_or_path = fileutils.force_gcs_path(
+      _PRETRAINED_MODEL_NAME_OR_PATH.value
   )
+  if dataset_validation_util.is_gcs_path(pretrained_model_name_or_path):
+    pretrained_model_name_or_path = (
+        dataset_validation_util.download_gcs_uri_to_local(
+            pretrained_model_name_or_path
+        )
+    )
 
   # GCS Fuse does not sync flushed files if not closed. See b/361771727.
   logging_output_dir = fileutils.force_gcs_path(_LOGGING_OUTPUT_DIR.value)
 
   # Creates evaluation config.
-  if _EVAL_TASKS.value:
+  if _EVAL_DATASET.value:
     eval_config = eval_lib.EvalConfig(
-        tasks=_EVAL_TASKS.value,
-        per_device_batch_size=_EVAL_PER_DEVICE_BATCH_SIZE.value,
-        num_fewshot=_EVAL_NUM_FEWSHOT.value,
+        per_device_batch_size=_PER_DEVICE_EVAL_BATCH_SIZE.value,
         limit=_EVAL_LIMIT.value,
         metric_name=_EVAL_METRIC_NAME.value,
         steps=_EVAL_STEPS.value,
         dataset_path=dataset_validation_util.force_gcs_fuse_path(
-            _EVAL_DATASET_PATH.value
+            _EVAL_DATASET.value
         ),
         split=_EVAL_SPLIT.value,
         template=_EVAL_TEMPLATE.value,
         column=_EVAL_COLUMN.value,
         tokenize_dataset=False,
+        metric_for_best_model=_METRIC_FOR_BEST_MODEL.value,
     )
   else:
     eval_config = None
@@ -745,40 +858,40 @@ def main(unused_argv: Sequence[str]) -> None:
     wandb.login()
 
   finetune_instruct(
-      pretrained_model_id=pretrained_model_id,
-      dataset_name=_DATASET_NAME.value,
-      output_dir=output_dir.local_dir,
+      pretrained_model_name_or_path=pretrained_model_name_or_path,
+      train_dataset=_TRAIN_DATASET.value,
+      output_dir=_OUTPUT_DIR.value,
       logging_output_dir=logging_output_dir,
       precision_mode=_PRECISION_MODE.value,
       lora_rank=_LORA_RANK.value,
       lora_alpha=_LORA_ALPHA.value,
       lora_dropout=_LORA_DROPOUT.value,
       warmup_ratio=_WARMUP_RATIO.value,
-      num_epochs=_NUM_EPOCHS.value,
+      num_train_epochs=_NUM_TRAIN_EPOCHS.value,
       warmup_steps=_WARMUP_STEPS.value,
       max_steps=_MAX_STEPS.value,
       max_seq_length=_MAX_SEQ_LENGTH.value,
       learning_rate=_LEARNING_RATE.value,
-      instruct_column_in_dataset=_INSTRUCT_COLUMN_IN_DATASET.value,
+      train_column=_TRAIN_COLUMN.value,
       per_device_train_batch_size=_PER_DEVICE_TRAIN_BATCH_SIZE.value,
       optim=_OPTIMIZER.value,
       weight_decay=_WEIGHT_DECAY.value,
       gradient_accumulation_steps=_GRADIENT_ACCUMULATION_STEPS.value,
-      enable_gradient_checkpointing=_ENABLE_GRADIENT_CHECKPOINTING.value,
+      gradient_checkpointing=_GRADIENT_CHECKPOINTING.value,
       enable_peft=_ENABLE_PEFT.value,
-      template=_TEMPLATE.value,
+      train_template=_TRAIN_TEMPLATE.value,
       lr_scheduler_type=_LR_SCHEDULER_TYPE.value,
       save_steps=_SAVE_STEPS.value,
       logging_steps=_LOGGING_STEPS.value,
-      train_split_name=_TRAIN_SPLIT_NAME.value,
+      train_split=_TRAIN_SPLIT.value,
       eval_config=eval_config,
       report_to=_REPORT_TO.value,
       access_token=_HUGGINGFACE_ACCESS_TOKEN.value,
       train_precision=_TRAIN_PRECISION.value,
-      use_example_packing=_USE_EXAMPLE_PACKING.value,
+      example_packing=_EXAMPLE_PACKING.value,
       attn_implementation=_ATTN_IMPLEMENTATION.value,
       max_grad_norm=_MAX_GRAD_NORM.value,
-      completion_only=_COMPLETION_ONLY.value,
+      input_masking=_INPUT_MASKING.value,
       logger_level=_LOGGER_LEVEL.value,
       benchmark_out_file=_BENCHMARK_OUT_FILE.value,
       tuning_data_stats_file=_TUNING_DATA_STATS_FILE.value,
@@ -787,8 +900,6 @@ def main(unused_argv: Sequence[str]) -> None:
   # Frees the model from GPU.
   utils.force_gc()
 
-  output_dir.upload_to_gcs(skip_if_exists=True)
-
 
 if __name__ == '__main__':
   app.run(main)

community-content/vertex_model_garden/model_oss/peft/train/vmg/merge_causal_language_model_lora.py

@@ -1,6 +1,7 @@
 """Script to merge PEFT adapter with base model."""
 
-from typing import Any, Dict, Sequence
+from collections.abc import Mapping, Sequence
+from typing import Any
 
 from absl import app
 from absl import flags
@@ -11,14 +12,14 @@ from util import constants
 from util import fileutils
 
 
-_PRETRAINED_MODEL_ID = flags.DEFINE_string(
-    'pretrained_model_id',
+_PRETRAINED_MODEL_NAME_OR_PATH = flags.DEFINE_string(
+    'pretrained_model_name_or_path',
     None,
     'The pretrained model id. Supported models can be causal language modeling'
     ' models from https://github.com/huggingface/peft/tree/main. Note, there'
     ' might be different paddings for different models. This tool assumes the'
-    ' pretrained_model_id contains model name, and then choose proper padding'
-    ' methods. e.g. it must contain `llama` for `Llama2 models`.',
+    ' pretrained_model_name_or_path contains model name, and then choose proper'
+    ' padding methods. e.g. it must contain `llama` for `Llama2 models`.',
     required=True,
 )
 
@@ -30,7 +31,7 @@ _MERGE_BASE_AND_LORA_OUTPUT_DIR = flags.DEFINE_string(
 
 _MERGE_MODEL_PRECISION_MODE = flags.DEFINE_enum(
     'merge_model_precision_mode',
-    constants.PRECISION_MODE_16,
+    constants.PRECISION_MODE_16B,
     [
         constants.PRECISION_MODE_4,
         constants.PRECISION_MODE_8,
@@ -48,20 +49,6 @@ _FINETUNED_LORA_MODEL_DIR = flags.DEFINE_string(
     'The directory storing finetuned LoRA model weights.',
 )
 
-_RESTRICT_MODEL_UPLOAD_DOCKER_URI = flags.DEFINE_string(
-    'restrict_model_upload_docker_uri',
-    '',
-    'If set, mark output model as only uploadable to Model Registry with the'
-    ' specified Docker URI.',
-)
-
-_EXECUTOR_INPUT = flags.DEFINE_string(
-    'executor_input',
-    '',
-    'For internal use. Kubeflow pipeline context when running trainer as part'
-    ' of an internal pipeline.',
-)
-
 _HUGGINGFACE_ACCESS_TOKEN = flags.DEFINE_string(
     'huggingface_access_token',
     None,
@@ -71,12 +58,12 @@ _HUGGINGFACE_ACCESS_TOKEN = flags.DEFINE_string(
 
 @flags.multi_flags_validator(
     [
-        _PRETRAINED_MODEL_ID.name,
+        _PRETRAINED_MODEL_NAME_OR_PATH.name,
         _FINETUNED_LORA_MODEL_DIR.name,
         _MERGE_BASE_AND_LORA_OUTPUT_DIR.name,
     ],
 )
-def check_merge_lora_model_flags(flags_dict: Dict[str, Any]) -> bool:
+def check_merge_lora_model_flags(flags_dict: Mapping[str, Any]) -> bool:
   """Check if required flags are set on merge model LoRA task.
 
   Args:
@@ -89,44 +76,33 @@ def check_merge_lora_model_flags(flags_dict: Dict[str, Any]) -> bool:
 
 
 def main(unused_argv: Sequence[str]) -> None:
-  pretrained_model_id = fileutils.force_gcs_path(_PRETRAINED_MODEL_ID.value)
-  if dataset_validation_util.is_gcs_path(pretrained_model_id):
-    pretrained_model_id = dataset_validation_util.download_gcs_uri_to_local(
-        pretrained_model_id
+  pretrained_model_name_or_path = fileutils.force_gcs_path(
+      _PRETRAINED_MODEL_NAME_OR_PATH.value
+  )
+  if dataset_validation_util.is_gcs_path(pretrained_model_name_or_path):
+    pretrained_model_name_or_path = (
+        dataset_validation_util.download_gcs_uri_to_local(
+            pretrained_model_name_or_path
+        )
     )
 
-  finetuned_lora_model_dir = utils.GcsOrLocalDirectory(
+  finetuned_lora_model_dir = fileutils.force_gcs_path(
       _FINETUNED_LORA_MODEL_DIR.value
   )
-
-  merge_base_and_lora_output_dir = utils.GcsOrLocalDirectory(
-      _MERGE_BASE_AND_LORA_OUTPUT_DIR.value
-  )
-
+  if dataset_validation_util.is_gcs_path(finetuned_lora_model_dir):
+    finetuned_lora_model_dir = (
+        dataset_validation_util.download_gcs_uri_to_local(
+            finetuned_lora_model_dir
+        )
+    )
   utils.merge_causal_language_model_with_lora(
-      pretrained_model_id=pretrained_model_id,
+      pretrained_model_name_or_path=pretrained_model_name_or_path,
       precision_mode=_MERGE_MODEL_PRECISION_MODE.value,
-      finetuned_lora_model_dir=finetuned_lora_model_dir.local_dir,
-      merged_model_output_dir=merge_base_and_lora_output_dir.local_dir,
+      finetuned_lora_model_dir=finetuned_lora_model_dir,
+      merged_model_output_dir=_MERGE_BASE_AND_LORA_OUTPUT_DIR.value,
       access_token=_HUGGINGFACE_ACCESS_TOKEN.value,
   )
 
-  if _RESTRICT_MODEL_UPLOAD_DOCKER_URI.value:
-    utils.write_first_party_model_metadata(
-        merge_base_and_lora_output_dir.local_dir,
-        _RESTRICT_MODEL_UPLOAD_DOCKER_URI.value,
-    )
-
-  if _EXECUTOR_INPUT.value:
-    utils.write_kfp_outputs(
-        _EXECUTOR_INPUT.value,
-        {
-            'saved_model': _MERGE_BASE_AND_LORA_OUTPUT_DIR.value,
-        },
-    )
-
-  merge_base_and_lora_output_dir.upload_to_gcs(skip_if_exists=True)
-
 
 if __name__ == '__main__':
   app.run(main)

community-content/vertex_model_garden/model_oss/peft/train/vmg/quantize_model.py

@@ -1,349 +0,0 @@
-"""Quantizes the model."""
-
-import json
-import os
-from typing import Any, Dict, List, Sequence, Union
-
-from absl import app
-from absl import flags
-from absl import logging
-from auto_gptq import AutoGPTQForCausalLM
-from auto_gptq import BaseQuantizeConfig
-from awq import AutoAWQForCausalLM
-from optimum.gptq.data import get_dataset
-from transformers import AutoTokenizer
-
-from util import dataset_validation_util
-from vertex_vision_model_garden_peft.train.vmg import utils
-from util import constants
-
-_PRETRAINED_MODEL_ID = flags.DEFINE_string(
-    'pretrained_model_id',
-    None,
-    'The pretrained model id. Supported models can be causal language modeling'
-    ' models from https://github.com/huggingface/peft/tree/main. Note, there'
-    ' might be different paddings for different models. This tool assumes the'
-    ' pretrained_model_id contains model name, and then choose proper padding'
-    ' methods. e.g. it must contain `llama` for `Llama2 models`.',
-)
-
-_QUANTIZATION_METHOD = flags.DEFINE_enum(
-    'quantization_method',
-    None,
-    [constants.GPTQ, constants.AWQ],
-    'The quantization method. Choose from ["gtpq", "awq"].',
-)
-
-_QUANTIZATION_PRECISION_MODE = flags.DEFINE_enum(
-    'quantization_precision_mode',
-    constants.PRECISION_MODE_4,
-    [
-        constants.PRECISION_MODE_8,
-        constants.PRECISION_MODE_4,
-        constants.PRECISION_MODE_3,
-        constants.PRECISION_MODE_2,
-    ],
-    'Quantization precision mode.',
-)
-
-_QUANTIZATION_DATASET_NAME = flags.DEFINE_string(
-    'quantization_dataset_name',
-    None,
-    'The dataset used for quantization. You can provide your own dataset in a'
-    ' list of string or just use the original datasets used in GPTQ paper'
-    ' ["wikitext2","c4","c4-new","ptb","ptb-new"] for GPTQ quantization. Using'
-    " a dataset more appropriate to the model's training can improve"
-    ' quantisation accuracy. Note that the GPTQ dataset is not the same as the'
-    ' dataset used to train the model.',
-)
-
-_TEXT_COLUMN_IN_QUANTIZATION_DATASET = flags.DEFINE_string(
-    'text_column_in_quantization_dataset',
-    constants.DEFAULT_TEXT_COLUMN_IN_QUANTIZATION_DATASET,
-    'The text column in quantization dataset.',
-)
-
-_QUANTIZATION_OUTPUT_DIR = flags.DEFINE_string(
-    'quantization_output_dir',
-    None,
-    'The directory to store the quantized model.',
-)
-
-_QUANTIZATION_DEVICE_MAP = flags.DEFINE_string(
-    'device_map', None, 'The device map.'
-)
-
-_QUANTIZATION_MAX_MEMORY = flags.DEFINE_string(
-    'max_memory', None, 'The maximum memory.'
-)
-
-_GROUP_SIZE = flags.DEFINE_integer(
-    'group_size',
-    None,
-    'The group size to use for quantization. Recommended value is 128 and -1'
-    ' uses per-column quantization. Higher numbers use less VRAM, but have'
-    ' lower quantisation accuracy. "None" is the lowest possible value.',
-)
-
-_DESC_ACT = flags.DEFINE_boolean(
-    'desc_act',
-    False,
-    'Whether to quantize columns in order of decreasing activation size.'
-    ' Setting it to False can significantly speed up inference but the'
-    ' perplexity may become slightly worse. Also known as act-order.',
-)
-
-_DAMP_PERCENT = flags.DEFINE_float(
-    'damp_percent',
-    0.1,
-    'The percent of the average Hessian diagonal to use for dampening.',
-)
-
-_CACHE_EXAMPLES_ON_GPU = flags.DEFINE_boolean(
-    'cache_examples_on_gpu',
-    True,
-    'Whether to cache the examples on GPU. Disabling will reduce VRAM usage,'
-    ' but increase quantization time.',
-)
-
-_AWQ_VERSION = flags.DEFINE_enum(
-    'awq_version',
-    constants.GEMM,
-    [constants.GEMM, constants.GEMV],
-    'The version of the AWQ to use. It determines how matrix multiplication'
-    ' runs under the hood. GEMV is 20% faster than GEMM, only at batch size 1'
-    ' (not good for large contexts). GEMM is much faster than FP16 at batch'
-    ' sizes below 8 (good with large contexts).',
-)
-
-
-@flags.multi_flags_validator(
-    [
-        _PRETRAINED_MODEL_ID.name,
-        _QUANTIZATION_METHOD.name,
-        _QUANTIZATION_PRECISION_MODE.name,
-        _QUANTIZATION_DATASET_NAME.name,
-        _QUANTIZATION_OUTPUT_DIR.name,
-    ],
-)
-def check_quantization_flags(flags_dict: Dict[str, Any]) -> bool:
-  """Check if required flags are set on quantization task.
-
-  Args:
-    flags_dict: Dictionary containing task and flags to check.
-
-  Returns:
-    If required flags are not None.
-  """
-  required_flags = [
-      _QUANTIZATION_METHOD.name,
-      _PRETRAINED_MODEL_ID.name,
-      _QUANTIZATION_PRECISION_MODE.name,
-      _QUANTIZATION_DATASET_NAME.name,
-      _QUANTIZATION_OUTPUT_DIR.name,
-  ]
-
-  return all(map(lambda x: flags_dict[x] is not None, required_flags))
-
-
-def quantize_model(
-    quantization_method: str,
-    pretrained_model_id: str,
-    quantization_output_dir: str,
-    quantization_precision_mode: str = None,
-    quantization_dataset_name: Union[List[str]] = None,
-    text_column_in_quantization_dataset: str = constants.DEFAULT_TEXT_COLUMN_IN_QUANTIZATION_DATASET,
-    group_size: int = None,
-    desc_act: bool = True,
-    damp_percent: float = 0.1,
-    awq_version: str = 'GEMM',
-    device_map: str = None,
-    max_memory: Dict[Any, str] = None,
-    cache_examples_on_gpu: bool = True,
-) -> None:
-  """Quantizes the model using `quantization_method`."""
-  if quantization_method == constants.GPTQ:
-    gptq_quantize_model(
-        pretrained_model_id=pretrained_model_id,
-        gptq_output_dir=quantization_output_dir,
-        gptq_precision_mode=quantization_precision_mode,
-        gptq_dataset_name=quantization_dataset_name,
-        group_size=group_size,
-        desc_act=desc_act,
-        damp_percent=damp_percent,
-        cache_examples_on_gpu=cache_examples_on_gpu,
-    )
-  elif quantization_method == constants.AWQ:
-    awq_quantize_model(
-        pretrained_model_id=pretrained_model_id,
-        quantization_output_dir=quantization_output_dir,
-        quantization_precision_mode=quantization_precision_mode,
-        quantization_dataset_name=quantization_dataset_name,
-        text_column_in_quantization_dataset=text_column_in_quantization_dataset,
-        group_size=group_size,
-        awq_version=awq_version,
-        device_map=device_map,
-        max_memory=max_memory,
-    )
-
-
-def awq_quantize_model(
-    pretrained_model_id: str,
-    quantization_output_dir: str,
-    quantization_precision_mode: str = None,
-    quantization_dataset_name: Union[List[str]] = None,
-    text_column_in_quantization_dataset: str = constants.DEFAULT_TEXT_COLUMN_IN_QUANTIZATION_DATASET,
-    group_size: int = None,
-    awq_version: str = 'GEMM',
-    device_map: str = None,
-    max_memory: Dict[Any, str] = None,
-) -> None:
-  """Quantizes the model using AWQ."""
-  if quantization_precision_mode != constants.PRECISION_MODE_4:
-    raise ValueError(
-        f'Invalid precision mode: {quantization_precision_mode} for AWQ. 4bit'
-        ' quantization must be used.'
-    )
-  else:
-    bits = 4
-  if not group_size:
-    group_size = 128
-  if not device_map:
-    device_map = 'cpu'
-  if dataset_validation_util.is_gcs_path(quantization_dataset_name):
-    logging.info('Using custom dataset: %s', quantization_dataset_name)
-    with open(
-        dataset_validation_util.force_gcs_fuse_path(quantization_dataset_name),
-        'r',
-    ) as f:
-      quantization_dataset = [line.rstrip('\n') for line in f]
-  else:
-    quantization_dataset = quantization_dataset_name
-  quant_config = {
-      'zero_point': True,
-      'q_group_size': group_size,
-      'w_bit': bits,
-      'version': awq_version,
-  }
-  logging.info('Quantization config: %s', quant_config)
-  model = AutoAWQForCausalLM.from_pretrained(
-      pretrained_model_id,
-      trust_remote_code=True,
-      device_map=device_map,
-      max_memory=max_memory,
-      low_cpu_mem_usage=True,
-  )
-  tokenizer = AutoTokenizer.from_pretrained(
-      pretrained_model_id, trust_remote_code=True
-  )
-  model.quantize(
-      tokenizer,
-      quant_config=quant_config,
-      calib_data=quantization_dataset,
-      text_column=text_column_in_quantization_dataset,
-  )
-  model.save_quantized(quantization_output_dir)
-  tokenizer.save_pretrained(quantization_output_dir)
-
-
-def gptq_quantize_model(
-    pretrained_model_id: str,
-    gptq_output_dir: str,
-    gptq_precision_mode: str = None,
-    gptq_dataset_name: Union[List[str]] = None,
-    group_size: int = -1,
-    desc_act: bool = False,
-    damp_percent: float = 0.1,
-    cache_examples_on_gpu: bool = True,
-) -> None:
-  """Quantizes the model using GPTQ."""
-  logging.info(
-      'PYTORCH_CUDA_ALLOC_CONF: %s',
-      os.environ.get('PYTORCH_CUDA_ALLOC_CONF', ''),
-  )
-  if dataset_validation_util.is_gcs_path(gptq_dataset_name):
-    logging.info('Using custom dataset: %s', gptq_dataset_name)
-    with open(
-        dataset_validation_util.force_gcs_fuse_path(gptq_dataset_name), 'r'
-    ) as f:
-      gptq_dataset = [line.rstrip('\n') for line in f]
-  else:
-    gptq_dataset = gptq_dataset_name
-  if gptq_precision_mode == constants.PRECISION_MODE_8:
-    bits = 8
-  elif gptq_precision_mode == constants.PRECISION_MODE_4:
-    bits = 4
-  elif gptq_precision_mode == constants.PRECISION_MODE_3:
-    bits = 3
-  elif gptq_precision_mode == constants.PRECISION_MODE_2:
-    bits = 2
-  else:
-    raise ValueError(f'Invalid precision mode: {gptq_precision_mode} for GPTQ.')
-  if not group_size:
-    group_size = -1
-
-  tokenizer = AutoTokenizer.from_pretrained(pretrained_model_id)
-  gptq_dataset = get_dataset(gptq_dataset, tokenizer)
-
-  quantization_config = BaseQuantizeConfig(
-      bits=bits,
-      group_size=group_size,
-      damp_percent=damp_percent,
-      desc_act=desc_act,
-  )
-
-  logging.info('Quantization config: %s', quantization_config.to_dict())
-
-  model = AutoGPTQForCausalLM.from_pretrained(
-      pretrained_model_id,
-      quantization_config,
-      low_cpu_mem_usage=True,
-      torch_dtype='auto',
-      trust_remote_code=True,
-  )
-  model.quantize(
-      examples=gptq_dataset,
-      cache_examples_on_gpu=cache_examples_on_gpu,
-  )
-
-  if utils.should_add_pad_token(pretrained_model_id):
-    tokenizer.add_special_tokens({'pad_token': '[PAD]'})
-    model.resize_token_embeddings(len(tokenizer))
-  model.save_pretrained(gptq_output_dir)
-  tokenizer.save_pretrained(gptq_output_dir)
-
-
-def main(unused_argv: Sequence[str]) -> None:
-  pretrained_model_id = _PRETRAINED_MODEL_ID.value
-  if dataset_validation_util.is_gcs_path(pretrained_model_id):
-    pretrained_model_id = dataset_validation_util.download_gcs_uri_to_local(
-        pretrained_model_id
-    )
-  pretrained_model_id = dataset_validation_util.force_gcs_fuse_path(
-      pretrained_model_id
-  )
-
-  if _QUANTIZATION_MAX_MEMORY.value:
-    max_memory = json.loads(_QUANTIZATION_MAX_MEMORY.value)
-  else:
-    max_memory = None
-
-  quantize_model(
-      quantization_method=_QUANTIZATION_METHOD.value,
-      pretrained_model_id=pretrained_model_id,
-      quantization_output_dir=_QUANTIZATION_OUTPUT_DIR.value,
-      quantization_precision_mode=_QUANTIZATION_PRECISION_MODE.value,
-      quantization_dataset_name=_QUANTIZATION_DATASET_NAME.value,
-      text_column_in_quantization_dataset=_TEXT_COLUMN_IN_QUANTIZATION_DATASET.value,
-      group_size=_GROUP_SIZE.value,
-      desc_act=_DESC_ACT.value,
-      damp_percent=_DAMP_PERCENT.value,
-      awq_version=_AWQ_VERSION.value,
-      device_map=_QUANTIZATION_DEVICE_MAP.value,
-      max_memory=max_memory,
-      cache_examples_on_gpu=_CACHE_EXAMPLES_ON_GPU.value,
-  )
-
-
-if __name__ == '__main__':
-  app.run(main)

community-content/vertex_model_garden/model_oss/peft/train/vmg/sequence_classification_lora.py

@@ -1,229 +0,0 @@
-"""Sequence classification with LoRA models."""
-
-from typing import Sequence
-
-from absl import app
-from absl import flags
-from datasets import load_dataset
-import evaluate
-from peft import get_peft_model
-from peft import LoraConfig
-import torch
-from torch.optim import AdamW
-from torch.utils.data import DataLoader
-from tqdm import tqdm
-from transformers import AutoModelForSequenceClassification
-from transformers import AutoTokenizer
-from transformers import get_linear_schedule_with_warmup
-
-from util import dataset_validation_util
-
-
-_PRETRAINED_MODEL_ID = flags.DEFINE_string(
-    "pretrained_model_id",
-    None,
-    "The pretrained model id. Supported models can be causal language modeling"
-    " models from https://github.com/huggingface/peft/tree/main. Note, there"
-    " might be different paddings for different models. This tool assumes the"
-    " pretrained_model_id contains model name, and then choose proper padding"
-    " methods. e.g. it must contain `llama` for `Llama2 models`.",
-)
-
-_OUTPUT_DIR = flags.DEFINE_string(
-    "output_dir",
-    None,
-    "The output directory.",
-)
-
-_DATASET_NAME = flags.DEFINE_string(
-    "dataset_name",
-    None,
-    "The dataset name in huggingface.",
-)
-
-_LORA_RANK = flags.DEFINE_integer(
-    "lora_rank",
-    16,
-    "The rank of the update matrices, expressed in int. Lower rank results in"
-    " smaller update matrices with fewer trainable parameters, referring to"
-    " https://huggingface.co/docs/peft/conceptual_guides/lora.",
-)
-
-_LORA_ALPHA = flags.DEFINE_integer(
-    "lora_alpha",
-    32,
-    "LoRA scaling factor, referring to"
-    " https://huggingface.co/docs/peft/conceptual_guides/lora.",
-)
-
-_LORA_DROPOUT = flags.DEFINE_float(
-    "lora_dropout",
-    0.05,
-    "dropout probability of the LoRA layers, referring to"
-    " https://huggingface.co/docs/peft/task_guides/token-classification-lora.",
-)
-
-_NUM_EPOCHS = flags.DEFINE_integer(
-    "num_epochs",
-    None,
-    "The number of training epochs.",
-)
-
-_BATCH_SIZE = flags.DEFINE_integer(
-    "batch_size",
-    32,
-    "The batch size.",
-)
-
-_LEARNING_RATE = flags.DEFINE_float(
-    "learning_rate",
-    2e-4,
-    "The learning rate after the potential warmup period.",
-)
-
-
-def finetune_sequence_classification(
-    pretrained_model_id: str,
-    dataset_name: str,
-    output_dir: str,
-    lora_rank: int = 8,
-    lora_alpha: int = 16,
-    lora_dropout: float = 0.1,
-    num_epochs: int = 20,
-    batch_size: int = 32,
-    learning_rate: float = 3e-4,
-) -> None:
-  """Finetunes sequence classification."""
-  task = "mrpc"
-  device = "cuda"
-
-  peft_config = LoraConfig(
-      task_type="SEQ_CLS",
-      inference_mode=False,
-      r=lora_rank,
-      lora_alpha=lora_alpha,
-      lora_dropout=lora_dropout,
-  )
-  if any(k in pretrained_model_id for k in ("gpt", "opt", "bloom")):
-    padding_side = "left"
-  else:
-    padding_side = "right"
-
-  tokenizer = AutoTokenizer.from_pretrained(
-      pretrained_model_id, padding_side=padding_side
-  )
-  if getattr(tokenizer, "pad_token_id") is None:
-    tokenizer.pad_token_id = tokenizer.eos_token_id
-
-  datasets = load_dataset(dataset_name, task)
-  metric = evaluate.load(dataset_name, task)
-
-  def tokenize_function(examples):
-    # max_length=None => use the model max length (it's actually the default)
-    outputs = tokenizer(
-        examples["sentence1"],
-        examples["sentence2"],
-        truncation=True,
-        max_length=None,
-    )
-    return outputs
-
-  tokenized_datasets = datasets.map(
-      tokenize_function,
-      batched=True,
-      remove_columns=["idx", "sentence1", "sentence2"],
-  )
-
-  # We also rename the 'label' column to 'labels' which is the expected name for
-  # labels by the models of the transformers library.
-  tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
-
-  def collate_fn(examples):
-    return tokenizer.pad(examples, padding="longest", return_tensors="pt")
-
-  # Instantiate dataloaders.
-  train_dataloader = DataLoader(
-      tokenized_datasets["train"],
-      shuffle=True,
-      collate_fn=collate_fn,
-      batch_size=batch_size,
-  )
-  eval_dataloader = DataLoader(
-      tokenized_datasets["validation"],
-      shuffle=False,
-      collate_fn=collate_fn,
-      batch_size=batch_size,
-  )
-
-  model = AutoModelForSequenceClassification.from_pretrained(
-      pretrained_model_id, return_dict=True
-  )
-  model = get_peft_model(model, peft_config)
-  model.print_trainable_parameters()
-
-  optimizer = AdamW(params=model.parameters(), lr=learning_rate)
-
-  # Instantiate scheduler
-  lr_scheduler = get_linear_schedule_with_warmup(
-      optimizer=optimizer,
-      num_warmup_steps=0.06 * (len(train_dataloader) * num_epochs),
-      num_training_steps=(len(train_dataloader) * num_epochs),
-  )
-
-  model.to(device)
-  for epoch in range(num_epochs):
-    model.train()
-    for _, batch in enumerate(tqdm(train_dataloader)):
-      batch.to(device)
-      outputs = model(**batch)
-      loss = outputs.loss
-      loss.backward()
-      optimizer.step()
-      lr_scheduler.step()
-      optimizer.zero_grad()
-
-    model.eval()
-    for _, batch in enumerate(tqdm(eval_dataloader)):
-      batch.to(device)
-      with torch.no_grad():
-        outputs = model(**batch)
-      predictions = outputs.logits.argmax(dim=-1)
-      references = batch["labels"]
-      metric.add_batch(
-          predictions=predictions,
-          references=references,
-      )
-
-    eval_metric = metric.compute()
-    print(f"epoch {epoch}:", eval_metric)
-
-  model.save_pretrained(output_dir)
-
-
-def main(unused_argv: Sequence[str]) -> None:
-  if dataset_validation_util.is_gcs_path(_PRETRAINED_MODEL_ID.value):
-    pretrained_model_id = dataset_validation_util.download_gcs_uri_to_local(
-        _PRETRAINED_MODEL_ID.value
-    )
-  else:
-    pretrained_model_id = _PRETRAINED_MODEL_ID.value
-  pretrained_model_path = dataset_validation_util.force_gcs_fuse_path(
-      pretrained_model_id
-  )
-  output_dir = dataset_validation_util.force_gcs_fuse_path(_OUTPUT_DIR.value)
-
-  finetune_sequence_classification(
-      pretrained_model_id=pretrained_model_path,
-      dataset_name=_DATASET_NAME.value,
-      output_dir=output_dir,
-      lora_rank=_LORA_RANK.value,
-      lora_alpha=_LORA_ALPHA.value,
-      lora_dropout=_LORA_DROPOUT.value,
-      num_epochs=int(_NUM_EPOCHS.value),
-      batch_size=_BATCH_SIZE.value,
-      learning_rate=_LEARNING_RATE.value,
-  )
-
-
-if __name__ == "__main__":
-  app.run(main)

community-content/vertex_model_garden/model_oss/peft/train/vmg/templates/llama3.json

@@ -1,7 +1,7 @@
 {
   "description": "Chat template used by Llama 3.",
   "source": "https://huggingface.co/meta-llama/Meta-Llama-3-70B-Instruct/blob/a5a71a7527eac1d651bb145436c72026887fb68e/tokenizer_config.json#L2053",
-  "chat_template": "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %}",
+  "chat_template": "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '\n\n<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '\n\n<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %}",
   "instruction_separator": "<|start_header_id|>user<|end_header_id|>\n\n",
   "response_separator": "<|start_header_id|>assistant<|end_header_id|>\n\n"
 }

community-content/vertex_model_garden/model_oss/peft/train/vmg/templates/openai-chat.json

@@ -0,0 +1,7 @@
+{
+  "description": "Template used by openai chat.",
+  "source": "https://platform.openai.com/docs/api-reference/fine-tuning/chat-input",
+  "chat_template": "{% set loop_messages = messages %}{% set content = '' %}{% for message in loop_messages %}{% set content = content ~ '\n\n<|start_header_id|>' ~ message.role ~ '<|end_header_id|>\n\n' %}{% if message.content is string %}{% set content = content ~ message.content|trim ~ '<|eot_id|>' %}{% else %}{% set content = content ~ message.content|join(' ', attribute='text')|trim ~ '<|eot_id|>' %}{% endif %}{% if loop.index0 == 0 %}{% set content = bos_token ~ content %}{% endif %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %}",
+  "instruction_separator": "<|start_header_id|>user<|end_header_id|>\n\n",
+  "response_separator": "<|start_header_id|>assistant<|end_header_id|>\n\n"
+}

community-content/vertex_model_garden/model_oss/peft/train/vmg/templates/qwen2_5.json

@@ -0,0 +1,7 @@
+{
+  "description": "Chat template used by Qwen 2.5.",
+  "source": "https://huggingface.co/Qwen/Qwen2.5-72B-Instruct/blob/main/tokenizer_config.json#L198",
+  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- 'You are Qwen, created by Alibaba Cloud. You are a helpful assistant.' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else %}\n        {{- '<|im_start|>system\\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n",
+  "instruction_separator": "<|im_start|>user\n",
+  "response_separator": "<|im_start|>assistant\n"
+}

community-content/vertex_model_garden/model_oss/peft/train/vmg/templates/vertex-chat.json

@@ -1,7 +1,7 @@
 {
   "description": "Template used for chat based models.",
   "source": "https://huggingface.co/meta-llama/Meta-Llama-3-70B-Instruct/blob/a5a71a7527eac1d651bb145436c72026887fb68e/tokenizer_config.json#L2053",
-  "chat_template": "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>model<|end_header_id|>\n\n' }}{% endif %}",
+  "chat_template": "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '\n\n<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '\n\n<|start_header_id|>model<|end_header_id|>\n\n' }}{% endif %}",
   "instruction_separator": "<|start_header_id|>user<|end_header_id|>\n\n",
   "response_separator": "<|start_header_id|>model<|end_header_id|>\n\n"
 }

community-content/vertex_model_garden/model_oss/peft/train/vmg/tests/instruct_lora_command_builder.py

@@ -3,16 +3,18 @@
 # DO NOT MODIFY: this file is auto-generated
 # See go/vmg-oss-peft-tests#command-builder-genpy
 
+
 class InstructLoraCommandBuilder:
 
   def __init__(self):
     self._config_file = None
     self._task = None
-    self._pretrained_model_id = None
-    self._dataset_name = None
-    self._train_split_name = None
-    self._template = None
-    self._instruct_column_in_dataset = None
+    self._gcs_rsync_interval_secs = None
+    self._pretrained_model_name_or_path = None
+    self._train_dataset = None
+    self._train_split = None
+    self._train_template = None
+    self._train_column = None
     self._output_dir = None
     self._merge_base_and_lora_output_dir = None
     self._logging_output_dir = None
@@ -22,14 +24,14 @@ class InstructLoraCommandBuilder:
     self._lora_alpha = None
     self._lora_dropout = None
     self._max_steps = None
-    self._num_epochs = None
+    self._num_train_epochs = None
     self._max_seq_length = None
     self._learning_rate = None
     self._lr_scheduler_type = None
     self._precision_mode = None
     self._train_precision = None
-    self._enable_gradient_checkpointing = None
-    self._use_example_packing = None
+    self._gradient_checkpointing = None
+    self._example_packing = None
     self._attn_implementation = None
     self._optimizer = None
     self._warmup_ratio = None
@@ -37,14 +39,14 @@ class InstructLoraCommandBuilder:
     self._save_steps = None
     self._logging_steps = None
     self._huggingface_access_token = None
-    self._eval_dataset_path = None
+    self._eval_dataset = None
     self._eval_column = None
     self._eval_template = None
     self._eval_split = None
     self._eval_steps = None
-    self._eval_tasks = None
     self._eval_metric_name = None
-    self._completion_only = None
+    self._metric_for_best_model = None
+    self._input_masking = None
     self._max_grad_norm = None
     self._logger_level = None
     self._benchmark_out_file = None
@@ -52,6 +54,7 @@ class InstructLoraCommandBuilder:
     self._enable_peft = None
     self._merge_model_precision_mode = None
     self._target_modules = None
+    self._unnamed_args = None
 
   @property
   def config_file(self):
@@ -70,44 +73,52 @@ class InstructLoraCommandBuilder:
     self._task = val
 
   @property
-  def pretrained_model_id(self):
-    return self._pretrained_model_id
+  def gcs_rsync_interval_secs(self):
+    return self._gcs_rsync_interval_secs
 
-  @pretrained_model_id.setter
-  def pretrained_model_id(self, val: str):
-    self._pretrained_model_id = val
+  @gcs_rsync_interval_secs.setter
+  def gcs_rsync_interval_secs(self, val: str):
+    self._gcs_rsync_interval_secs = val
+
+  @property
+  def pretrained_model_name_or_path(self):
+    return self._pretrained_model_name_or_path
+
+  @pretrained_model_name_or_path.setter
+  def pretrained_model_name_or_path(self, val: str):
+    self._pretrained_model_name_or_path = val
 
   @property
   def train_dataset(self):
-    return self._dataset_name
+    return self._train_dataset
 
   @train_dataset.setter
   def train_dataset(self, val: str):
-    self._dataset_name = val
+    self._train_dataset = val
 
   @property
-  def train_split_name(self):
-    return self._train_split_name
+  def train_split(self):
+    return self._train_split
 
-  @train_split_name.setter
-  def train_split_name(self, val: str):
-    self._train_split_name = val
+  @train_split.setter
+  def train_split(self, val: str):
+    self._train_split = val
 
   @property
-  def template(self):
-    return self._template
+  def train_template(self):
+    return self._train_template
 
-  @template.setter
-  def template(self, val: str):
-    self._template = val
+  @train_template.setter
+  def train_template(self, val: str):
+    self._train_template = val
 
   @property
-  def instruct_column(self):
-    return self._instruct_column_in_dataset
+  def train_column(self):
+    return self._train_column
 
-  @instruct_column.setter
-  def instruct_column(self, val: str):
-    self._instruct_column_in_dataset = val
+  @train_column.setter
+  def train_column(self, val: str):
+    self._train_column = val
 
   @property
   def ckpt_dir(self):
@@ -182,12 +193,12 @@ class InstructLoraCommandBuilder:
     self._max_steps = val
 
   @property
-  def num_epochs(self):
-    return self._num_epochs
+  def num_train_epochs(self):
+    return self._num_train_epochs
 
-  @num_epochs.setter
-  def num_epochs(self, val: float):
-    self._num_epochs = val
+  @num_train_epochs.setter
+  def num_train_epochs(self, val: float):
+    self._num_train_epochs = val
 
   @property
   def max_seq_length(self):
@@ -231,19 +242,19 @@ class InstructLoraCommandBuilder:
 
   @property
   def gradient_checkpointing(self):
-    return self._enable_gradient_checkpointing
+    return self._gradient_checkpointing
 
   @gradient_checkpointing.setter
   def gradient_checkpointing(self, val: bool):
-    self._enable_gradient_checkpointing = val
+    self._gradient_checkpointing = val
 
   @property
   def example_packing(self):
-    return self._use_example_packing
+    return self._example_packing
 
   @example_packing.setter
   def example_packing(self, val: bool):
-    self._use_example_packing = val
+    self._example_packing = val
 
   @property
   def attn_implementation(self):
@@ -303,18 +314,18 @@ class InstructLoraCommandBuilder:
 
   @property
   def eval_dataset(self):
-    return self._eval_dataset_path
+    return self._eval_dataset
 
   @eval_dataset.setter
   def eval_dataset(self, val: str):
-    self._eval_dataset_path = val
+    self._eval_dataset = val
 
   @property
-  def eval_instruct_column(self):
+  def eval_column(self):
     return self._eval_column
 
-  @eval_instruct_column.setter
-  def eval_instruct_column(self, val: str):
+  @eval_column.setter
+  def eval_column(self, val: str):
     self._eval_column = val
 
   @property
@@ -326,11 +337,11 @@ class InstructLoraCommandBuilder:
     self._eval_template = val
 
   @property
-  def eval_split_name(self):
+  def eval_split(self):
     return self._eval_split
 
-  @eval_split_name.setter
-  def eval_split_name(self, val: str):
+  @eval_split.setter
+  def eval_split(self, val: str):
     self._eval_split = val
 
   @property
@@ -341,14 +352,6 @@ class InstructLoraCommandBuilder:
   def eval_steps(self, val: int):
     self._eval_steps = val
 
-  @property
-  def eval_tasks(self):
-    return self._eval_tasks
-
-  @eval_tasks.setter
-  def eval_tasks(self, val: str):
-    self._eval_tasks = val
-
   @property
   def eval_metric_name(self):
     return self._eval_metric_name
@@ -358,12 +361,20 @@ class InstructLoraCommandBuilder:
     self._eval_metric_name = val
 
   @property
-  def completion_only(self):
-    return self._completion_only
+  def metric_for_best_model(self):
+    return self._metric_for_best_model
 
-  @completion_only.setter
-  def completion_only(self, val: bool):
-    self._completion_only = val
+  @metric_for_best_model.setter
+  def metric_for_best_model(self, val: str):
+    self._metric_for_best_model = val
+
+  @property
+  def input_masking(self):
+    return self._input_masking
+
+  @input_masking.setter
+  def input_masking(self, val: bool):
+    self._input_masking = val
 
   @property
   def max_grad_norm(self):
@@ -421,9 +432,22 @@ class InstructLoraCommandBuilder:
   def target_modules(self, val: str):
     self._target_modules = val
 
-  def build_cmd(self) -> str:
+  @property
+  def unnamed_args(self):
+    return self._unnamed_args
+
+  @unnamed_args.setter
+  def unnamed_args(self, val: list):
+    self._unnamed_args = val
+
+  def build_cmd(self) -> list[str]:
     cmd = []
+    args = ''
     for k, v in self.__dict__.items():
+      if k == '_unnamed_args' and v is not None:
+        args += ' '.join(v)
+        continue
       if v is not None:
         cmd.append(f'--{k[1:]}={v}')
+    cmd.append(f'{args}')
     return cmd

community-content/vertex_model_garden/model_oss/peft/train/vmg/tests/quantize_model_command_builder.py

@@ -8,7 +8,7 @@ class QuantizeModelCommandBuilder:
 
   def __init__(self):
     self._task = None
-    self._pretrained_model_id = None
+    self._pretrained_model_name_or_path = None
     self._quantization_method = None
     self._quantization_precision_mode = None
     self._quantization_dataset_name = None
@@ -31,12 +31,12 @@ class QuantizeModelCommandBuilder:
     self._task = val
 
   @property
-  def pretrained_model_id(self):
-    return self._pretrained_model_id
+  def pretrained_model_name_or_path(self):
+    return self._pretrained_model_name_or_path
 
-  @pretrained_model_id.setter
-  def pretrained_model_id(self, val: str):
-    self._pretrained_model_id = val
+  @pretrained_model_name_or_path.setter
+  def pretrained_model_name_or_path(self, val: str):
+    self._pretrained_model_name_or_path = val
 
   @property
   def quantization_method(self):

community-content/vertex_model_garden/model_oss/peft/train/vmg/tests/run_tests.py

@@ -0,0 +1,61 @@
+"""Run the tests from docker command line."""
+
+import subprocess
+import sys
+from typing import Sequence
+
+from absl import app
+from absl import flags
+
+
+_ALLOWED_TEST_FILE_PATHS = (
+    "test_instruct_lora_adapters",
+    "test_instruct_lora_features",
+    "test_instruct_lora_throughput",
+    "test_instruct_lora_trained_model_quality",
+    "test_validate_dataset_with_template",
+)
+
+_TEST_FILE_PATH = flags.DEFINE_multi_enum(
+    "test_file_path",
+    None,
+    _ALLOWED_TEST_FILE_PATHS + ("all",),
+    "The test file path.",
+    required=True,
+)
+
+_IS_AUTOMATED_TEST = flags.DEFINE_bool(
+    "is_automated_test",
+    True,
+    "Whether the test is an automated test.",
+)
+
+
+def main(argv: Sequence[str]) -> None:
+  if len(argv) > 1:
+    raise app.UsageError("Too many command-line arguments.")
+  test_file_path = _TEST_FILE_PATH.value
+  if "all" in test_file_path:
+    test_file_path = _ALLOWED_TEST_FILE_PATHS
+
+  for test_file in test_file_path:
+    cmd = [
+        "python3",
+        f"vertex_vision_model_garden_peft/tests/{test_file}.py",
+    ]
+    if (
+        test_file == "test_instruct_lora_throughput"
+        and _IS_AUTOMATED_TEST.value
+    ):
+      subprocess.run(
+          cmd + ["--", "-k", "peft_train_image_automated_test"],
+          stdout=sys.stdout,
+          stderr=sys.stdout,
+          check=True,
+      )
+    else:
+      subprocess.run(cmd, stdout=sys.stdout, stderr=sys.stdout, check=True)
+
+
+if __name__ == "__main__":
+  app.run(main)

community-content/vertex_model_garden/model_oss/peft/train/vmg/tests/test_instruct_lora_adapters.py

@@ -31,7 +31,7 @@ class AdapterTest(test_util.TestBase):
   def setUp(self):
     super().setUp()
 
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
 
     self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
     self.task_cmd_builder.task = 'instruct-lora'
@@ -39,9 +39,9 @@ class AdapterTest(test_util.TestBase):
     self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
         'peft_train_sample.jsonl'
     )
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'input_text'
-    self.task_cmd_builder.template = 'llama3-text-bison'
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'input_text'
+    self.task_cmd_builder.train_template = 'llama3-text-bison'
     self.task_cmd_builder.gradient_accumulation_steps = 1
     self.task_cmd_builder.lora_rank = 16
     self.task_cmd_builder.lora_alpha = 32
@@ -87,8 +87,8 @@ class AdapterTest(test_util.TestBase):
   def test_llama_adapters(self, model_name):
     test_function_name = inspect.stack()[0][3]
     self.setup_output_dir(f'{test_function_name}-{model_name}')
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(model_name)
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(model_name)
     )
     start_time = time.time()
     self.assertEqual(self.run_cmd(), 0)

community-content/vertex_model_garden/model_oss/peft/train/vmg/tests/test_instruct_lora_features.py

@@ -4,19 +4,81 @@
 
 import os
 import time
-
+import unittest
 from absl.testing import absltest
 from absl.testing import parameterized
 import instruct_lora_command_builder as task_cmd_builder
 import test_util
 
 
+class EvalConfigTest(test_util.TestBase):
+
+  def setUp(self):
+    super().setUp()
+
+    self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path('llama3.1-8b-hf')
+    )
+    self.task_cmd_builder.config_file = (
+        'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
+    )
+    self.task_cmd_builder.task = 'instruct-lora'
+    self.task_cmd_builder.per_device_batch_size = 1
+    self.task_cmd_builder.gradient_accumulation_steps = 1
+    self.task_cmd_builder.lora_rank = 16
+    self.task_cmd_builder.lora_alpha = 32
+    self.task_cmd_builder.lora_dropout = 0.05
+    self.task_cmd_builder.learning_rate = 5e-5
+    self.task_cmd_builder.warmup_ratio = 0.01
+    self.task_cmd_builder.max_steps = 10
+    self.task_cmd_builder.save_steps = 1000
+    self.task_cmd_builder.logging_steps = 1
+    self.task_cmd_builder.gradient_checkpointing = True
+    self.task_cmd_builder.attn_implementation = 'flash_attention_2'
+    self.task_cmd_builder.example_packing = True
+    self.task_cmd_builder.train_dataset = 'mlabonne/guanaco-llama2'
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'text'
+    self.task_cmd_builder.train_template = 'openassistant-guanaco'
+    self.task_cmd_builder.ckpt_dir = '/tmp/adapter'
+    self.task_cmd_builder.logging_dir = '/tmp/logs'
+    self.task_cmd_builder.eval_steps = 10
+    self.task_cmd_builder.eval_dataset = 'mlabonne/guanaco-llama2'
+    self.task_cmd_builder.eval_split = 'test'
+    self.task_cmd_builder.eval_column = 'text'
+    self.task_cmd_builder.eval_template = 'openassistant-guanaco'
+
+  @parameterized.named_parameters(
+      ('all_eval_metric', 'loss,perplexity,bleu,google_bleu,rouge1', 0),
+      ('invalid_metric', 'invalid_metric', 1),
+      ('only_loss', 'loss', 0),
+      ('perplexity_without_loss', 'perplexity,bleu', 0),
+      ('unsupported_eval_metric', 'f1', 1),
+  )
+  def test_hf_eval_metrics(self, eval_metric_name, expected_return_code):
+    self.task_cmd_builder.eval_metric_name = eval_metric_name
+    self.assertEqual(self.run_cmd(), expected_return_code)
+
+  @parameterized.named_parameters(
+      ('valid_best_model_metric', 'loss,perplexity', 'perplexity', 0),
+      ('only_loss', None, 'loss', 0),
+      ('invalid_best_model_metric', 'loss', 'invalid_metric', 1),
+  )
+  def test_metric_for_best_model(
+      self, eval_metric_name, metric_for_best_model, expected_return_code
+  ):
+    self.task_cmd_builder.eval_metric_name = eval_metric_name
+    self.task_cmd_builder.metric_for_best_model = metric_for_best_model
+    self.assertEqual(self.run_cmd(), expected_return_code)
+
+
 class GcsUploadDownloadTest(test_util.TestBase):
 
   def setUp(self):
     super().setUp()
 
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
 
     self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
     self.task_cmd_builder.task = 'instruct-lora'
@@ -24,9 +86,9 @@ class GcsUploadDownloadTest(test_util.TestBase):
     self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
         'peft_train_sample.jsonl'
     )
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'input_text'
-    self.task_cmd_builder.template = 'llama3-text-bison'
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'input_text'
+    self.task_cmd_builder.train_template = 'llama3-text-bison'
     self.task_cmd_builder.gradient_accumulation_steps = 1
     self.task_cmd_builder.lora_rank = 16
     self.task_cmd_builder.lora_alpha = 32
@@ -45,9 +107,11 @@ class GcsUploadDownloadTest(test_util.TestBase):
       ),
       ('llama2_7b_hf', 'NousResearch/Llama-2-7b-hf'),
   )
-  def test_model_download_single_process(self, pretrained_model_id):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(pretrained_model_id)
+  def test_model_download_single_process(self, pretrained_model_name_or_path):
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(
+            pretrained_model_name_or_path
+        )
     )
     start_time = time.time()
     self.assertEqual(self.run_cmd(), 0)
@@ -61,24 +125,26 @@ class GcsUploadDownloadTest(test_util.TestBase):
       ),
       ('llama2_7b_hf', 'NousResearch/Llama-2-7b-hf'),
   )
-  def test_model_download_multi_process(self, pretrained_model_id):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(pretrained_model_id)
+  def test_model_download_multi_process(self, pretrained_model_name_or_path):
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(
+            pretrained_model_name_or_path
+        )
     )
     self.task_cmd_builder.config_file = (
         'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
     )
 
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
 
     start_time = time.time()
     self.assertEqual(self.run_cmd(), 0)
     end_time = time.time()
     self.assertLess(end_time - start_time, 5 * 60.0)
 
-  def test_70b_model_download(self):
-    self.task_cmd_builder.pretrained_model_id = (
-        'gs://vertex-model-garden-public-us/llama3/llama3-70b-hf'
+  def test_8b_model_download(self):
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        'gs://vertex-model-garden-public-us/llama3/llama3-8b-hf'
     )
     start_time = time.time()
     self.assertEqual(self.run_cmd(), 0)
@@ -90,8 +156,8 @@ class GcsUploadDownloadTest(test_util.TestBase):
       ('merged-and-upload-to-gcs', 'gs://vmg-test-ttl-1y/tests/merged'),
   )
   def test_model_merge(self, output_dir):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path('llama3.1-8b-hf')
     )
 
     ckpt_dir = os.path.join(
@@ -107,9 +173,13 @@ class GcsUploadDownloadTest(test_util.TestBase):
     end_time = time.time()
     self.assertLess(end_time - start_time, 5 * 60.0)
 
+  @unittest.skipIf(
+      not test_util.is_gpu_h100(),
+      'Skipping because this test is only for H100',
+  )
   def test_model_fp8_conversion(self):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path('llama3.1-8b-hf')
     )
 
     ckpt_dir = f'/tmp/output/output-{test_util.get_timestamp()}'
@@ -128,17 +198,17 @@ class GcsUploadDownloadTest(test_util.TestBase):
       ('merged-and-upload-to-gcs', 'gs://vmg-test-ttl-1y/tests/merged'),
   )
   def test_model_merge_and_upload_deepspeed(self, merged_model_dir):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path('llama3.1-8b-hf')
     )
     self.task_cmd_builder.config_file = (
-        'vertex_vision_model_garden_peft/deepspeed_zero3_8gpu.yaml'
+        'vertex_vision_model_garden_peft/deepspeed_zero2_8gpu.yaml'
     )
     self.task_cmd_builder.merged_model_dir = os.path.join(
         merged_model_dir, f'merged-{test_util.get_timestamp()}'
     )
 
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
     start_time = time.time()
     self.assertEqual(self.run_cmd(), 0)
     end_time = time.time()
@@ -149,8 +219,8 @@ class GcsUploadDownloadTest(test_util.TestBase):
       ('save-multiple-times', 1),
   )
   def test_llama3_8b_save_and_merge_8_gpus_fsdp(self, save_steps):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path('llama3.1-8b-hf')
     )
     self.task_cmd_builder.save_steps = save_steps
     self.task_cmd_builder.max_steps = 3
@@ -159,7 +229,7 @@ class GcsUploadDownloadTest(test_util.TestBase):
     )
     self.task_cmd_builder.merged_model_dir = '/tmp/merged'
 
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
     start_time = time.time()
     self.assertEqual(self.run_cmd(), 0)
     end_time = time.time()
@@ -171,7 +241,7 @@ class TemplateAndDataStatsTest(test_util.TestBase):
   def setUp(self):
     super().setUp()
 
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
 
     self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
     self.task_cmd_builder.task = 'instruct-lora'
@@ -187,35 +257,42 @@ class TemplateAndDataStatsTest(test_util.TestBase):
     self.task_cmd_builder.attn_implementation = 'flash_attention_2'
     self.task_cmd_builder.ckpt_dir = '/tmp'
 
-  def test_openai_chat_template(self):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+  @parameterized.named_parameters(
+      ('multi-chat-string-content', 'openai-multi-chat-example-data.jsonl'),
+      (
+          'multi-chat-array-content',
+          'openai-multi-chat-example-data-array-content.jsonl',
+      ),
+  )
+  def test_openai_chat_template(self, example_dataset):
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path('llama3.1-8b-hf')
     )
     self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
-        'openai-multi-chat-example-data.jsonl'
+        example_dataset
     )
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'messages'
-    self.task_cmd_builder.template = 'llama3'
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'messages'
+    self.task_cmd_builder.train_template = 'openai-chat'
 
     self.assertEqual(self.run_cmd(), 0)
 
   def test_openai_completion_template(self):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path('llama3.1-8b-hf')
     )
     self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
         'openai-completion-example-data.jsonl'
     )
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'prompt'
-    self.task_cmd_builder.template = 'openai-completion'
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'prompt'
+    self.task_cmd_builder.train_template = 'openai-completion'
 
     self.assertEqual(self.run_cmd(), 0)
 
   def test_data_stats_chat_template(self):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path('llama3.1-8b-hf')
     )
     self.task_cmd_builder.config_file = (
         'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
@@ -223,17 +300,17 @@ class TemplateAndDataStatsTest(test_util.TestBase):
     self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
         'openai-multi-chat-example-data.jsonl'
     )
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'messages'
-    self.task_cmd_builder.template = 'llama3'
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'messages'
+    self.task_cmd_builder.train_template = 'llama3'
     self.task_cmd_builder.tuning_data_stats_file = '/tmp/data-stats.json'
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
 
     self.assertEqual(self.run_cmd(), 0)
 
   def test_data_stats_completion_template(self):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path('llama3.1-8b-hf')
     )
     self.task_cmd_builder.config_file = (
         'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
@@ -241,12 +318,12 @@ class TemplateAndDataStatsTest(test_util.TestBase):
     self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
         'openai-completion-example-data.jsonl'
     )
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'prompt'
-    self.task_cmd_builder.template = 'openai-completion'
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'prompt'
+    self.task_cmd_builder.train_template = 'openai-completion'
     self.task_cmd_builder.tuning_data_stats_file = '/tmp/data-stats.json'
 
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
     self.assertEqual(self.run_cmd(), 0)
 
 
@@ -255,7 +332,7 @@ class TargetModulesTest(test_util.TestBase):
   def setUp(self):
     super().setUp()
 
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
 
     self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
     self.task_cmd_builder.task = 'instruct-lora'
@@ -263,9 +340,9 @@ class TargetModulesTest(test_util.TestBase):
     self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
         'peft_train_sample.jsonl'
     )
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'input_text'
-    self.task_cmd_builder.template = 'llama3-text-bison'
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'input_text'
+    self.task_cmd_builder.train_template = 'llama3-text-bison'
     self.task_cmd_builder.gradient_accumulation_steps = 1
     self.task_cmd_builder.lora_rank = 16
     self.task_cmd_builder.lora_alpha = 32
@@ -278,8 +355,8 @@ class TargetModulesTest(test_util.TestBase):
     self.task_cmd_builder.ckpt_dir = '/tmp'
 
   def test_target_modules(self):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path('llama3.1-8b-hf')
     )
     self.task_cmd_builder.target_modules = 'q_proj, v_proj, k_proj'
 

community-content/vertex_model_garden/model_oss/peft/train/vmg/tests/test_instruct_lora_throughput.py

@@ -43,9 +43,9 @@ class TrainerThroughputTest(test_util.TestBase):
     self.task_cmd_builder.attn_implementation = 'flash_attention_2'
     self.task_cmd_builder.example_packing = True
     self.task_cmd_builder.train_dataset = 'mlabonne/guanaco-llama2-1k'
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'text'
-    self.task_cmd_builder.template = 'openassistant-guanaco'
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'text'
+    self.task_cmd_builder.train_template = 'openassistant-guanaco'
     self.task_cmd_builder.ckpt_dir = '/tmp/adapter'
     self.task_cmd_builder.logging_dir = '/tmp/logs'
 
@@ -54,23 +54,24 @@ class TrainerThroughputTest(test_util.TestBase):
     if ret != 0:
       with open(self.task_cmd_builder.benchmark_out_file, 'a') as f:
         max_seq_length = self.task_cmd_builder.max_seq_length
-        f.write(f'{max_seq_length/1024.0:.1f}k | failed | n/a\n')
+        f.write(f'{max_seq_length/1024.0:.1f} | failed | n/a\n')
     return ret
 
   @parameterized.product(
       model_name=[
-          'llama3-70b-hf',
+          'llama3.1-8b-hf',
           'llama3.1-70b-hf',
           'Mistral-7B-v0.1',
           'Mixtral-8x7B-v0.1',
-          'Gemma2-9b-it',
+          'gemma-2-9b-it',
+          'Qwen2.5-32B-Instruct',
       ],
       precision=['4bit', '8bit', 'bfloat16'],
       max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
   )
   def test_model_single_gpu(self, model_name, precision, max_seq_length):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(model_name)
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(model_name)
     )
     self.task_cmd_builder.max_seq_length = max_seq_length
     self.task_cmd_builder.load_precision = precision
@@ -78,29 +79,30 @@ class TrainerThroughputTest(test_util.TestBase):
         self.test_suite_output_dir, f'bm_{model_name}_{precision}.txt'
     )
 
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
 
     self.assertEqual(self.run_cmd_and_handle_failure(), 0)
 
   @parameterized.product(
       model_name=[
-          'llama3-70b-hf',
+          'llama3.1-8b-hf',
           'llama3.1-70b-hf',
           'Mistral-7B-v0.1',
           'Mixtral-8x7B-v0.1',
-          'Gemma2-9b-it',
+          'gemma-2-9b-it',
+          'Qwen2.5-32B-Instruct',
       ],
       precision=['4bit', '8bit', 'bfloat16'],
       max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
       num_gpus=[8],
-      config=['deepspeed_zero2', 'deepspeed_zero3'],
+      config=['deepspeed_zero2'],
   )
   def test_model_multi_gpu_deepspeed(
       self, model_name, precision, max_seq_length, num_gpus, config
   ):
     self.assertTrue(num_gpus == 4 or num_gpus == 8)
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(model_name)
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(model_name)
     )
     self.task_cmd_builder.max_seq_length = max_seq_length
     self.task_cmd_builder.load_precision = precision
@@ -112,14 +114,18 @@ class TrainerThroughputTest(test_util.TestBase):
     self.task_cmd_builder.config_file = (
         f'vertex_vision_model_garden_peft/{config}_{num_gpus}gpu.yaml'
     )
-    self.docker_builder.add_env_var(
+    self.command_builder.add_env_var(
         'CUDA_VISIBLE_DEVICES', ','.join([str(x) for x in range(0, num_gpus)])
     )
 
     self.assertEqual(self.run_cmd_and_handle_failure(), 0)
 
   @parameterized.product(
-      model_name=['llama3.1-70b-hf'],
+      model_name=[
+          'llama3.1-8b-hf',
+          'llama3.1-70b-hf',
+          'Qwen2.5-32B-Instruct',
+      ],
       precision=['4bit', '8bit', 'bfloat16'],
       max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
       num_gpus=[8],
@@ -127,8 +133,8 @@ class TrainerThroughputTest(test_util.TestBase):
   def test_model_multi_gpu_fsdp_lora(
       self, model_name, precision, max_seq_length, num_gpus
   ):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(model_name)
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(model_name)
     )
     self.task_cmd_builder.max_seq_length = max_seq_length
     self.task_cmd_builder.load_precision = precision
@@ -136,18 +142,108 @@ class TrainerThroughputTest(test_util.TestBase):
         self.test_suite_output_dir,
         f'bm_fsdp_{num_gpus}gpu_{model_name}_{precision}.txt',
     )
-    self.task_cmd_builder.config_file = (
-        'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
-    )
+    if 'llama' in model_name.lower():
+      self.task_cmd_builder.config_file = (
+          'vertex_vision_model_garden_peft/llama2_fsdp_8gpu.yaml'
+      )
+    elif 'qwen' in model_name.lower():
+      self.task_cmd_builder.config_file = (
+          'vertex_vision_model_garden_peft/qwen2_fsdp_8gpu.yaml'
+      )
+    else:
+      self.fail(f'Unsupported model: {model_name}')
 
-    self.docker_builder.add_env_var(
+    self.command_builder.add_env_var(
         'CUDA_VISIBLE_DEVICES', ','.join([str(x) for x in range(0, num_gpus)])
     )
 
     self.assertEqual(self.run_cmd_and_handle_failure(), 0)
 
   @parameterized.product(
-      model_name=['llama3.1-70b-hf'],
+      model_name=['llama3.1-8b-hf', 'llama3.1-70b-hf'],
+      precision=['4bit', 'bfloat16'],
+      max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
+      config=['deepspeed_zero2', 'fsdp'],
+  )
+  def test_peft_train_image_automated_test_llama(
+      self, model_name, precision, max_seq_length, config
+  ):
+    num_gpus = 8
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(model_name)
+    )
+    self.task_cmd_builder.max_seq_length = max_seq_length
+    self.task_cmd_builder.load_precision = precision
+    benchmark_out_file = os.path.join(
+        self.test_suite_output_dir,
+        f'bm_{config}_{num_gpus}gpu_{model_name}_{precision}.txt',
+    )
+    self.task_cmd_builder.benchmark_out_file = benchmark_out_file
+    if config == 'fsdp':
+      self.task_cmd_builder.config_file = (
+          f'vertex_vision_model_garden_peft/llama_{config}_{num_gpus}gpu.yaml'
+      )
+    else:
+      self.task_cmd_builder.config_file = (
+          f'vertex_vision_model_garden_peft/{config}_{num_gpus}gpu.yaml'
+      )
+
+    self.command_builder.add_env_var(
+        'CUDA_VISIBLE_DEVICES', ','.join([str(x) for x in range(0, num_gpus)])
+    )
+    self.run_cmd_and_handle_failure()
+    if test_util.is_gpu_h100():
+      self.assertEqual(
+          test_util.check_benchmark_results(
+              benchmark_out_file, 'llama', 10.0, max_seq_length
+          ),
+          True,
+      )
+
+  @parameterized.product(
+      model_name=['gemma-2-2b-it', 'gemma-2-9b-it', 'gemma-2-27b-it'],
+      precision=['4bit', 'bfloat16'],
+      max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
+      config=['deepspeed_zero2', 'deepspeed_zero3', 'fsdp'],
+  )
+  def test_peft_train_image_automated_test_gemma(
+      self, model_name, precision, max_seq_length, config
+  ):
+    num_gpus = 8
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(model_name)
+    )
+    self.task_cmd_builder.attn_implementation = 'eager'
+    self.task_cmd_builder.max_seq_length = max_seq_length
+    self.task_cmd_builder.load_precision = precision
+    benchmark_out_file = os.path.join(
+        self.test_suite_output_dir,
+        f'bm_{config}_{num_gpus}gpu_{model_name}_{precision}.txt',
+    )
+    self.task_cmd_builder.benchmark_out_file = benchmark_out_file
+    if config == 'fsdp':
+      self.task_cmd_builder.config_file = (
+          f'vertex_vision_model_garden_peft/gemma2_{config}_{num_gpus}gpu.yaml'
+      )
+    else:
+      self.task_cmd_builder.config_file = (
+          f'vertex_vision_model_garden_peft/{config}_{num_gpus}gpu.yaml'
+      )
+
+    self.command_builder.add_env_var(
+        'CUDA_VISIBLE_DEVICES', ','.join([str(x) for x in range(0, num_gpus)])
+    )
+    self.run_cmd_and_handle_failure()
+    if test_util.is_gpu_h100():
+      self.assertEqual(
+          test_util.check_benchmark_results(
+              benchmark_out_file, 'gemma', 10.0, max_seq_length
+          ),
+          True,
+      )
+
+  @parameterized.product(
+      model_name=['llama3.1-8b-hf', 'llama3.1-70b-hf'],
       precision=['bfloat16'],
       max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
       num_gpus=[8],
@@ -155,8 +251,8 @@ class TrainerThroughputTest(test_util.TestBase):
   def test_model_multi_gpu_fsdp_full_finetuning(
       self, model_name, precision, max_seq_length, num_gpus
   ):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(model_name)
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(model_name)
     )
     self.task_cmd_builder.max_seq_length = max_seq_length
     self.task_cmd_builder.load_precision = precision
@@ -169,7 +265,7 @@ class TrainerThroughputTest(test_util.TestBase):
     )
     self.task_cmd_builder.enable_peft = False
 
-    self.docker_builder.add_env_var(
+    self.command_builder.add_env_var(
         'CUDA_VISIBLE_DEVICES', ','.join([str(x) for x in range(0, num_gpus)])
     )
 

community-content/vertex_model_garden/model_oss/peft/train/vmg/tests/test_instruct_lora_trained_model_quality.py

@@ -0,0 +1,209 @@
+# pylint: disable=missing-function-docstring
+# pylint: disable=missing-class-docstring
+"""Tests to make sure trained model achieves decent quality.
+
+Right now, the metric is loss decreasing and we'll eyeball the TB graphs.
+"""
+
+import os
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import instruct_lora_command_builder as task_cmd_builder
+import test_util
+
+
+class TrainedModelQualityTest(test_util.TestBase):
+
+  _TEST_OUTPUT_DIR = os.path.expanduser('~/output')
+
+  @classmethod
+  def setUpClass(cls):
+    super().setUpClass()
+    cls.test_suite_output_dir = os.path.join(
+        cls._TEST_OUTPUT_DIR,
+        os.path.splitext(os.path.basename(__file__))[0],
+        cls.__class__.__name__,
+    )
+
+  def setUp(self):
+    super().setUp()
+
+    self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
+    self.task_cmd_builder.task = 'instruct-lora'
+    self.task_cmd_builder.eval_metric_name = 'loss'
+    self.task_cmd_builder.per_device_batch_size = 1
+    self.task_cmd_builder.gradient_accumulation_steps = 8
+    self.task_cmd_builder.lora_rank = 16
+    self.task_cmd_builder.lora_alpha = 32
+    self.task_cmd_builder.lora_dropout = 0.05
+    self.task_cmd_builder.learning_rate = 5e-5
+    self.task_cmd_builder.num_train_epochs = 2.0
+    self.task_cmd_builder.warmup_ratio = 0.01
+    self.task_cmd_builder.max_steps = -1
+    self.task_cmd_builder.save_steps = 10
+    self.task_cmd_builder.eval_steps = 10
+    self.task_cmd_builder.max_seq_length = 4096
+    self.task_cmd_builder.load_precision = '4bit'
+    self.task_cmd_builder.gradient_checkpointing = True
+    self.task_cmd_builder.input_masking = True
+    self.task_cmd_builder.attn_implementation = 'flash_attention_2'
+    self.task_cmd_builder.report_to = 'tensorboard'
+
+  def setup_output_dir(self, testcase_name: str):
+    testcase_output_dir = os.path.join(
+        self.test_suite_output_dir, testcase_name
+    )
+    self.task_cmd_builder.ckpt_dir = os.path.join(
+        testcase_output_dir, 'adapter'
+    )
+    self.task_cmd_builder.logging_dir = os.path.join(
+        testcase_output_dir, 'logs'
+    )
+    self.task_cmd_builder.merged_model_dir = os.path.join(
+        testcase_output_dir, 'merged'
+    )
+
+  @parameterized.named_parameters(
+      ('llama3-8b', 'llama3-8b-hf'),
+      ('llama3.1-8b', 'llama3.1-8b-hf'),
+  )
+  def test_8b_model_deepspeed(self, model_name):
+    self.setup_output_dir(f'test_deepspeed_{model_name}')
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(model_name)
+    )
+    self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
+        'peft_train_sample.jsonl'
+    )
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'input_text'
+    self.task_cmd_builder.train_template = 'llama3-text-bison'
+    self.task_cmd_builder.eval_dataset = test_util.get_test_data_path(
+        'peft_eval_sample.jsonl'
+    )
+    self.task_cmd_builder.eval_split = 'train'
+    self.task_cmd_builder.eval_column = self.task_cmd_builder.train_column
+    self.task_cmd_builder.eval_template = self.task_cmd_builder.train_template
+
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
+
+    self.assertEqual(self.run_cmd(), 0)
+
+  @parameterized.named_parameters(
+      ('llama3-70b', 'llama3-70b-hf'),
+      ('llama3.1-70b', 'llama3.1-70b-hf'),
+  )
+  def test_70b_model_deepspeed(self, model_name):
+    self.setup_output_dir(f'test_deepspeed_{model_name}')
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(model_name)
+    )
+    self.task_cmd_builder.config_file = (
+        'vertex_vision_model_garden_peft/deepspeed_zero2_8gpu.yaml'
+    )
+    self.task_cmd_builder.train_dataset = 'timdettmers/openassistant-guanaco'
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'text'
+    self.task_cmd_builder.train_template = 'openassistant-guanaco'
+    self.task_cmd_builder.eval_dataset = self.task_cmd_builder.train_dataset
+    self.task_cmd_builder.eval_split = 'test'
+    self.task_cmd_builder.eval_column = self.task_cmd_builder.train_column
+    self.task_cmd_builder.eval_template = self.task_cmd_builder.train_template
+
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+
+    self.assertEqual(self.run_cmd(), 0)
+
+  @parameterized.named_parameters(
+      ('llama3-70b', 'llama3-70b-hf'),
+      ('llama3.1-70b', 'llama3.1-70b-hf'),
+  )
+  def test_70b_model_fsdp(self, model_name):
+    self.setup_output_dir(f'test_fsdp_{model_name}')
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(model_name)
+    )
+    self.task_cmd_builder.config_file = (
+        'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
+    )
+    self.task_cmd_builder.train_dataset = 'timdettmers/openassistant-guanaco'
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'text'
+    self.task_cmd_builder.train_template = 'openassistant-guanaco'
+    self.task_cmd_builder.eval_dataset = self.task_cmd_builder.train_dataset
+    self.task_cmd_builder.eval_split = 'test'
+    self.task_cmd_builder.eval_column = self.task_cmd_builder.train_column
+    self.task_cmd_builder.eval_template = self.task_cmd_builder.train_template
+
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+
+    self.assertEqual(self.run_cmd(), 0)
+
+  @parameterized.named_parameters(
+      ('Qwen2.5-32B-Instruct', 'Qwen2.5-32B-Instruct'),
+  )
+  def test_qwen_model_deepspeed(self, model_name):
+    self.setup_output_dir(f'test_deepspeed_{model_name}')
+    self.task_cmd_builder.pretrained_model_name_or_path = model_name
+    self.task_cmd_builder.config_file = (
+        'vertex_vision_model_garden_peft/deepspeed_zero2_8gpu.yaml'
+    )
+    self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
+        'llama-tuning-test/opposite-examples-train.jsonl'
+    )
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'messages'
+    self.task_cmd_builder.train_template = 'qwen2_5'
+    self.task_cmd_builder.eval_dataset = test_util.get_test_data_path(
+        'llama-tuning-test/opposite-examples-eval.jsonl'
+    )
+    self.task_cmd_builder.eval_split = 'train'
+    self.task_cmd_builder.eval_column = self.task_cmd_builder.train_column
+    self.task_cmd_builder.eval_template = self.task_cmd_builder.train_template
+
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    # Note(lavrai): The following parameters are needed for the opposite-word
+    # dataset to converge properly.
+    self.task_cmd_builder.gradient_accumulation_steps = 1
+    self.task_cmd_builder.num_train_epochs = 10.0
+    self.task_cmd_builder.logging_steps = 1
+
+    self.assertEqual(self.run_cmd(), 0)
+
+  @parameterized.named_parameters(
+      ('Qwen2.5-32B-Instruct', 'Qwen2.5-32B-Instruct'),
+  )
+  def test_qwen_model_fsdp(self, model_name):
+    self.setup_output_dir(f'test_fsdp_{model_name}')
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(model_name)
+    )
+    self.task_cmd_builder.config_file = (
+        'vertex_vision_model_garden_peft/qwen2_fsdp_8gpu.yaml'
+    )
+    self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
+        'llama-tuning-test/opposite-examples-train.jsonl'
+    )
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'messages'
+    self.task_cmd_builder.train_template = 'qwen2_5'
+    self.task_cmd_builder.eval_dataset = test_util.get_test_data_path(
+        'llama-tuning-test/opposite-examples-eval.jsonl'
+    )
+    self.task_cmd_builder.eval_split = 'train'
+    self.task_cmd_builder.eval_column = self.task_cmd_builder.train_column
+    self.task_cmd_builder.eval_template = self.task_cmd_builder.train_template
+
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    # Note(lavrai): The following parameters are needed for the opposite-word
+    # dataset to converge properly.
+    self.task_cmd_builder.gradient_accumulation_steps = 1
+    self.task_cmd_builder.num_train_epochs = 10.0
+    self.task_cmd_builder.logging_steps = 1
+
+    self.assertEqual(self.run_cmd(), 0)
+
+
+if __name__ == '__main__':
+  absltest.main()

community-content/vertex_model_garden/model_oss/peft/train/vmg/tests/test_validate_dataset_with_template.py

@@ -31,7 +31,7 @@ class ValidateDatasetWithTemplateTest(test_util.TestBase):
       dict(
           testcase_name="out_of_range_rows",
           validate_top_k_rows=100000,
-          expected_result=1,
+          expected_result=0,
       ),
   )
   def test_validate_dataset_with_template_top_k_rows(
@@ -40,8 +40,8 @@ class ValidateDatasetWithTemplateTest(test_util.TestBase):
       expected_result,
   ):
     self.task_cmd_builder.dataset_name = "timdettmers/openassistant-guanaco"
-    self.task_cmd_builder.train_split_name = "train"
-    self.task_cmd_builder.instruct_column_in_dataset = "text"
+    self.task_cmd_builder.train_split = "train"
+    self.task_cmd_builder.train_column = "text"
     self.task_cmd_builder.template = (
         "gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
     )
@@ -79,8 +79,8 @@ class ValidateDatasetWithTemplateTest(test_util.TestBase):
       expected_result,
   ):
     self.task_cmd_builder.dataset_name = "timdettmers/openassistant-guanaco"
-    self.task_cmd_builder.train_split_name = "train"
-    self.task_cmd_builder.instruct_column_in_dataset = "text"
+    self.task_cmd_builder.train_split = "train"
+    self.task_cmd_builder.train_column = "text"
     self.task_cmd_builder.template = (
         "gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
     )
@@ -92,6 +92,32 @@ class ValidateDatasetWithTemplateTest(test_util.TestBase):
     result = self.run_cmd()
     self.assertEqual(result, expected_result)
 
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="small_max_seq_length",
+          max_seq_length=10,
+      ),
+      dict(
+          testcase_name="large_max_seq_length",
+          max_seq_length=1024,
+      ),
+  )
+  def test_validate_dataset_with_template_max_seq_length(
+      self,
+      max_seq_length,
+  ):
+    self.task_cmd_builder.dataset_name = "timdettmers/openassistant-guanaco"
+    self.task_cmd_builder.train_split = "train"
+    self.task_cmd_builder.train_column = "text"
+    self.task_cmd_builder.template = (
+        "gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
+    )
+    self.task_cmd_builder.max_seq_length = max_seq_length
+    self.task_cmd_builder.validate_k_rows_of_dataset = None
+    self.task_cmd_builder.use_multiprocessing = True
+    result = self.run_cmd()
+    self.assertEqual(result, 0)
+
   @parameterized.named_parameters(
       dict(
           testcase_name="invalid_default_input_column",
@@ -204,8 +230,8 @@ class ValidateDatasetWithTemplateTest(test_util.TestBase):
       expected_result,
   ):
     self.task_cmd_builder.dataset_name = dataset_name
-    self.task_cmd_builder.train_split_name = split
-    self.task_cmd_builder.instruct_column_in_dataset = input_column
+    self.task_cmd_builder.train_split = split
+    self.task_cmd_builder.train_column = input_column
     self.task_cmd_builder.template = template
     self.task_cmd_builder.validate_percentage_of_dataset = (
         validate_percentage_of_dataset
@@ -311,8 +337,8 @@ class ValidateDatasetWithTemplateTest(test_util.TestBase):
       expected_result,
   ):
     self.task_cmd_builder.dataset_name = dataset_name
-    self.task_cmd_builder.train_split_name = "train"
-    self.task_cmd_builder.instruct_column_in_dataset = "text"
+    self.task_cmd_builder.train_split = "train"
+    self.task_cmd_builder.train_column = "text"
     self.task_cmd_builder.template = template
     self.task_cmd_builder.validate_percentage_of_dataset = (
         validate_percentage_of_dataset

community-content/vertex_model_garden/model_oss/peft/train/vmg/tests/validate_dataset_with_template_command_builder.py

@@ -10,8 +10,9 @@ class ValidateDatasetWithTemplateCommandBuilder:
     self._task = None
     self._template = None
     self._dataset_name = None
-    self._train_split_name = None
-    self._instruct_column_in_dataset = None
+    self._train_split = None
+    self._train_column = None
+    self._max_seq_length = None
     self._use_multiprocessing = None
     self._validate_k_rows_of_dataset = None
     self._validate_percentage_of_dataset = None
@@ -41,20 +42,28 @@ class ValidateDatasetWithTemplateCommandBuilder:
     self._dataset_name = val
 
   @property
-  def train_split_name(self):
-    return self._train_split_name
+  def train_split(self):
+    return self._train_split
 
-  @train_split_name.setter
-  def train_split_name(self, val: str):
-    self._train_split_name = val
+  @train_split.setter
+  def train_split(self, val: str):
+    self._train_split = val
 
   @property
-  def instruct_column_in_dataset(self):
-    return self._instruct_column_in_dataset
+  def train_column(self):
+    return self._train_column
 
-  @instruct_column_in_dataset.setter
-  def instruct_column_in_dataset(self, val: str):
-    self._instruct_column_in_dataset = val
+  @train_column.setter
+  def train_column(self, val: str):
+    self._train_column = val
+
+  @property
+  def max_seq_length(self):
+    return self._max_seq_length
+
+  @max_seq_length.setter
+  def max_seq_length(self, val: int):
+    self._max_seq_length = val
 
   @property
   def use_multiprocessing(self):

community-content/vertex_model_garden/model_oss/peft/train/vmg/train_entrypoint.py

@@ -8,16 +8,18 @@ environment. Otherwise, `python3` is used.
 """
 
 import argparse
-import json
-import os
+from collections.abc import MutableSequence, Sequence
+import multiprocessing
 import subprocess
-from typing import List, Optional, Sequence
+import sys
 from absl import app
 from absl import flags
 from absl import logging
 from util import dataset_validation_util
+from util import cluster_spec
 from vertex_vision_model_garden_peft.train.vmg import utils
 from util import constants
+from util import gcs_syncer
 from util import hypertune_utils
 
 
@@ -37,20 +39,21 @@ _TASK_TO_SCRIPT = {
     constants.INSTRUCT_LORA: (
         'vertex_vision_model_garden_peft/train/vmg/instruct_lora.py'
     ),
-    constants.MERGE_CAUSAL_LANGUAGE_MODEL_LORA: 'vertex_vision_model_garden_peft/train/vmg/merge_causal_language_model_lora.py',
-    constants.QUANTIZE_MODEL: (
-        'vertex_vision_model_garden_peft/train/vmg/quantize_model.py'
+    constants.MERGE_CAUSAL_LANGUAGE_MODEL_LORA: (
+        'vertex_vision_model_garden_peft/train/vmg/merge_causal_language_model_lora.py'
     ),
-    constants.SEQUENCE_CLASSIFICATION_LORA: 'vertex_vision_model_garden_peft/train/vmg/sequence_classification_lora.py',
-    constants.VALIDATE_DATASET_WITH_TEMPLATE: 'vertex_vision_model_garden_peft/train/vmg/validate_dataset_with_template.py',
+    constants.VALIDATE_DATASET_WITH_TEMPLATE: (
+        'vertex_vision_model_garden_peft/train/vmg/validate_dataset_with_template.py'
+    ),
+    constants.RUN_TESTS: 'vertex_vision_model_garden_peft/tests/run_tests.py',
 }
 
 
 def launch_script_cmd(
     script: str,
-    config_file: Optional[str],
+    config_file: str | None,
     accelerate_args: argparse.Namespace = argparse.Namespace(),
-) -> List[str]:
+) -> MutableSequence[str]:
   """Returns the command to launch the script."""
   if config_file:
     cmd = [
@@ -69,59 +72,23 @@ def launch_script_cmd(
 
 def _get_accelerate_args() -> argparse.Namespace:
   """Returns the accelerate args."""
-  # For the format of the cluster spec, see
-  # https://cloud.google.com/vertex-ai/docs/training/distributed-training#cluster-spec-format # pylint: disable=line-too-long
-  cluster_spec = os.getenv('CLUSTER_SPEC', default=None)
-  if not cluster_spec:
-    return argparse.Namespace()
-  logging.info('CLUSTER_SPEC: %s', cluster_spec)
-
-  cluster_data = json.loads(cluster_spec)
-  if (
-      'workerpool1' not in cluster_data['cluster']
-      or not cluster_data['cluster']['workerpool1']
-  ):
-    return argparse.Namespace()
-
-  # Get primary node info
-  primary_node = cluster_data['cluster']['workerpool0'][0]
-  logging.info('primary node: %s', primary_node)
-  primary_node_addr, primary_node_port = primary_node.split(':')
-  logging.info('primary node address: %s', primary_node_addr)
-  logging.info('primary node port: %s', primary_node_port)
-
-  # Determine node rank of this machine
-  workerpool = cluster_data['task']['type']
-  if workerpool == 'workerpool0':
-    node_rank = 0
-  elif workerpool == 'workerpool1':
-    # Add 1 for the primary node, since `index` is the index of workerpool1.
-    node_rank = cluster_data['task']['index'] + 1
-  else:
-    raise ValueError(
-        'Only workerpool0 and workerpool1 are supported. Unknown workerpool:'
-        f' {workerpool}'
-    )
-  logging.info('node rank: %s', node_rank)
-
-  # Calculate total nodes
-  num_worker_nodes = len(cluster_data['cluster']['workerpool1'])
-  num_nodes = num_worker_nodes + 1  # Add 1 for the primary node
-  logging.info('num nodes: %s', num_nodes)
-
+  primary_node_addr, primary_node_port, node_rank, num_nodes = (
+      cluster_spec.get_cluster_spec()
+  )
   accelerate_args = argparse.Namespace()
-  accelerate_args.machine_rank = node_rank
-  accelerate_args.num_machines = num_nodes
-  accelerate_args.main_process_ip = primary_node_addr
-  accelerate_args.main_process_port = primary_node_port
-  accelerate_args.max_restarts = 0
-  accelerate_args.monitor_interval = 120
+  if num_nodes > 1:
+    accelerate_args.machine_rank = node_rank
+    accelerate_args.num_machines = num_nodes
+    accelerate_args.main_process_ip = primary_node_addr
+    accelerate_args.main_process_port = primary_node_port
+    accelerate_args.max_restarts = 0
+    accelerate_args.monitor_interval = 120
 
   return accelerate_args
 
 
 def _append_args_to_command_in_place(
-    args: argparse.Namespace, command: List[str]
+    args: argparse.Namespace, command: MutableSequence[str]
 ):
   for key, value in vars(args).items():
     # If not specified, skip.
@@ -129,15 +96,25 @@ def _append_args_to_command_in_place(
       command.append(f'--{key}={value}')
 
 
-def _get_train_cmd_and_maybe_merge_cmd(
-    task: str, config_file: str, unknown: Sequence[str]
+def _get_train_and_maybe_merge_cmd_and_dirs_to_sync(
+    task_type: str, config_file: str, unknown: Sequence[str]
 ) -> Sequence[Sequence[str]]:
-  """Returns the training command and maybe the merge command if applicable."""
+  """Returns the training and merge command(if applicable) and dirs to sync.
 
+  Args:
+    task_type: The task type.
+    config_file: The accelerate config file path.
+    unknown: The unknown args which are not recognised by the parser.
+
+  Returns:
+    The bash commands to execute and the directories to sync.
+  """
+  dirs_to_sync = []
   # Only populated when multi-node is used.
   accelerate_args = _get_accelerate_args()
+  node_rank = getattr(accelerate_args, 'machine_rank', 0)
   training_cmd = launch_script_cmd(
-      _TASK_TO_SCRIPT[task],
+      _TASK_TO_SCRIPT[task_type],
       config_file,
       accelerate_args=accelerate_args,
   )
@@ -152,17 +129,30 @@ def _get_train_cmd_and_maybe_merge_cmd(
       dataset_validation_util.force_gcs_fuse_path(training_args.output_dir)
   )
 
+  local_output_dir, gcs_output_dir = gcs_syncer.manage_sync_path(
+      training_args.output_dir, node_rank
+  )
+  training_args.output_dir = local_output_dir
+  if gcs_syncer.is_gcs_or_gcsfuse_path(gcs_output_dir):
+    dirs_to_sync.append((local_output_dir, gcs_output_dir))
+
   # Merge only flags.
   merge_parser = argparse.ArgumentParser()
   merge_parser.add_argument('--merge_model_precision_mode')
-  merge_parser.add_argument('--executor_input')
-  merge_parser.add_argument('--restrict_model_upload_docker_uri')
   merge_parser.add_argument('--merge_base_and_lora_output_dir')
   merge_args, unknown = merge_parser.parse_known_args(unknown)
 
+  if merge_args.merge_base_and_lora_output_dir:
+    merge_local_dir, merge_gcs_dir = gcs_syncer.manage_sync_path(
+        merge_args.merge_base_and_lora_output_dir, None
+    )
+    merge_args.merge_base_and_lora_output_dir = merge_local_dir
+    if gcs_syncer.is_gcs_or_gcsfuse_path(merge_gcs_dir):
+      dirs_to_sync.append((merge_local_dir, merge_gcs_dir))
+
   # Common flags shared by merging and training.
   common_parser = argparse.ArgumentParser()
-  common_parser.add_argument('--pretrained_model_id', required=True)
+  common_parser.add_argument('--pretrained_model_name_or_path', required=True)
   common_parser.add_argument('--huggingface_access_token')
   common_args, remaining = common_parser.parse_known_args(unknown)
 
@@ -173,17 +163,19 @@ def _get_train_cmd_and_maybe_merge_cmd(
   commands = [training_cmd]
 
   # Only the main node runs merging.
-  if (
-      merge_args.merge_base_and_lora_output_dir
-      and getattr(accelerate_args, 'machine_rank', 0) == 0
-  ):
+  if merge_args.merge_base_and_lora_output_dir and node_rank == 0:
     lora_dir = utils.get_final_checkpoint_path(training_args.output_dir)
+    lora_local_dir, lora_gcs_dir = gcs_syncer.manage_sync_path(
+        lora_dir, node_rank
+    )
+    if gcs_syncer.is_gcs_or_gcsfuse_path(lora_gcs_dir):
+      dirs_to_sync.append((lora_local_dir, lora_gcs_dir))
 
     merge_cmd = [
         'WORLD_SIZE=1',  # To ignore other nodes in multi-node setting.
         'python3',
         _TASK_TO_SCRIPT[constants.MERGE_CAUSAL_LANGUAGE_MODEL_LORA],
-        f'--finetuned_lora_model_dir={lora_dir}',
+        f'--finetuned_lora_model_dir={lora_local_dir}',
     ]
     _append_args_to_command_in_place(merge_args, merge_cmd)
     _append_args_to_command_in_place(common_args, merge_cmd)
@@ -196,16 +188,51 @@ def _get_train_cmd_and_maybe_merge_cmd(
     ]
     commands.append(conda_run_cmd)
 
-  return commands
+  return commands, dirs_to_sync
+
+
+def _get_merge_cmd_and_dirs_to_sync(
+    task_type: str, config_file: str, unknown: Sequence[str]
+) -> Sequence[Sequence[str]]:
+  """Returns the merge command and dirs to sync.
+
+  Args:
+    task_type: The task type.
+    config_file: The accelerate config file path.
+    unknown: The unknown args which are not recognised by the parser.
+
+  Returns:
+    The bash commands to execute and the directories to sync.
+  """
+  # Merge only flags.
+  merge_parser = argparse.ArgumentParser()
+  merge_parser.add_argument('--merge_base_and_lora_output_dir')
+  merge_args, unknown = merge_parser.parse_known_args(unknown)
+
+  dirs_to_sync = []
+  if merge_args.merge_base_and_lora_output_dir:
+    merge_local_dir, merge_gcs_dir = gcs_syncer.manage_sync_path(
+        merge_args.merge_base_and_lora_output_dir, None
+    )
+    merge_args.merge_base_and_lora_output_dir = merge_local_dir
+    if gcs_syncer.is_gcs_or_gcsfuse_path(merge_gcs_dir):
+      dirs_to_sync.append((merge_local_dir, merge_gcs_dir))
+
+  cmd = launch_script_cmd(_TASK_TO_SCRIPT[task_type], config_file)
+  _append_args_to_command_in_place(merge_args, cmd)
+  cmd.extend(unknown)
+  return [cmd], dirs_to_sync
 
 
 def main(unused_argv: Sequence[str]) -> None:
   parser = argparse.ArgumentParser()
   parser.add_argument('--config_file')
   parser.add_argument('--task')
+  parser.add_argument('--gcs_rsync_interval_secs', type=int, default=60)
   args, unknown = parser.parse_known_args()
 
   task = args.task
+  dirs_to_sync = None
 
   if task in _TEXT_TO_IMAGE_TASKS_SCRIPTS:
     # Setup accelerate config before running trainer.
@@ -224,8 +251,12 @@ def main(unused_argv: Sequence[str]) -> None:
     ] + list(map(dataset_validation_util.force_gcs_fuse_path, unknown))
     commands = [config_gen_cmd, task_cmd]
   elif task in [constants.INSTRUCT_LORA]:
-    commands = _get_train_cmd_and_maybe_merge_cmd(
-        task=task, config_file=args.config_file, unknown=unknown
+    commands, dirs_to_sync = _get_train_and_maybe_merge_cmd_and_dirs_to_sync(
+        task_type=task, config_file=args.config_file, unknown=unknown
+    )
+  elif task in [constants.MERGE_CAUSAL_LANGUAGE_MODEL_LORA]:
+    commands, dirs_to_sync = _get_merge_cmd_and_dirs_to_sync(
+        task_type=task, config_file=args.config_file, unknown=unknown
     )
   else:
     assert task in _TASK_TO_SCRIPT
@@ -233,10 +264,29 @@ def main(unused_argv: Sequence[str]) -> None:
     cmd.extend(unknown)
     commands = [cmd]
 
+  rsync_process = None
+  mp_queue = multiprocessing.Queue(maxsize=1)
+  if dirs_to_sync:
+    rsync_process = gcs_syncer.setup_gcs_rsync(
+        dirs_to_sync, mp_queue, args.gcs_rsync_interval_secs
+    )
+
   for cmd in commands:
     logging.info('launching task=%s with cmd: \n%s', task, ' \\\n'.join(cmd))
-    subprocess.run(cmd, check=True)
+    # Both absl logging and python's logging module writes to stderr by default.
+    # Redirect output to stdout on purpose, such that log entries do not get
+    # marked as `Error` in Cloud's Log Explorer.
+    try:
+      subprocess.run(cmd, stdout=sys.stdout, stderr=sys.stdout, check=True)
+    except subprocess.CalledProcessError as e:
+      if rsync_process is not None and rsync_process.is_alive():
+        logging.info('Terminating GCS rsync process.')
+        rsync_process.terminate()
+      raise e
+  if rsync_process is not None:
+    gcs_syncer.cleanup_gcs_rsync(rsync_process, mp_queue)
 
 
 if __name__ == '__main__':
+  logging.get_absl_handler().python_handler.stream = sys.stdout
   app.run(main, flags_parser=lambda _args: flags.FLAGS(_args, known_only=True))

community-content/vertex_model_garden/model_oss/peft/train/vmg/utils.py

@@ -1,232 +1,44 @@
 """Common libraries for PEFT."""
 
-import dataclasses
+from collections.abc import Mapping, Sequence
 import datetime
 import gc
-import multiprocessing as mp
 import os
-import subprocess
-from typing import Any, Dict, Optional, Sequence
+from typing import Any
 
 from absl import logging
 import accelerate
 from accelerate import DistributedType
 from accelerate import PartialState
-from google.protobuf import json_format
-from kfp.pipeline_spec import pipeline_spec_pb2
-import numpy as np
 import peft
 from peft import PeftModel
 from peft import prepare_model_for_kbit_training
-import pynvml
 import torch
 import transformers
 from transformers import AutoModelForCausalLM
 from transformers import AutoTokenizer
 from transformers import BitsAndBytesConfig
 from transformers import FbgemmFp8Config
-from transformers.integrations import is_deepspeed_zero3_enabled
 import trl
 
 from util import dataset_validation_util
 from util import constants
-from util import fileutils
 
-
-_MODELS_REQUIRING_PAD_TOKEN = ("llama", "falcon", "mistral", "mixtral")
-_MODELS_REQUIRING_EOS_TOEKN = ("gemma-2b", "gemma-7b")
 _LLAMA_3_1_405B_MODEL_ID = "Meta-Llama-3.1-405B"
 _LOCAL_MERGED_MODEL_DIR = "/tmp/merged_model"
-
-
-
-class GcsOrLocalDirectory(os.PathLike):
-  """A class to represent a directory with upload support if GCS path is given.
-
-  This class is used to represent a directory. It can be used for a temporary
-  local directory and for uploading files to the GCS directory later if the
-  given path is a GCS directory. If the given path is a local directory, a call
-  to gcs_dir attribute will raise an error. This class has multi-node and
-  multi-process support with accelerate.
-
-  Attributes:
-    local_dir: The local directory to store the files.
-    gcs_dir: The path to the GCS directory.
-  """
-
-  def __init__(
-      self,
-      path: str,
-      check_empty: bool = False,
-      upload_from_all_nodes: bool = False,
-  ):
-    """Initializes the GcsOrLocalDirectory.
-
-    Args:
-      path: The path to the directory.
-      check_empty: If True, check if the GCS directory is empty. No-op for local
-        directory.
-      upload_from_all_nodes: If True, upload the local directory to GCS from all
-        nodes.
-    """
-    if len(path) > 1:
-      path = path.rstrip("/")
-
-    self._upload_from_all_nodes = upload_from_all_nodes
-
-    if path.startswith(constants.GCS_URI_PREFIX) or path.startswith(
-        constants.GCSFUSE_URI_PREFIX
-    ):
-      self._is_gcs_path = True
-      self._local_dir = _get_local_dir_from_gcs_dir(path)
-      self._gcs_dir = fileutils.force_gcs_path(path)
-      os.makedirs(self.local_dir, exist_ok=True)
-
-      with PartialState().main_process_first():
-        if (
-            check_empty
-            and PartialState().is_main_process
-            and not _is_gcs_dir_empty(self._gcs_dir)
-        ):
-          raise ValueError(f"{self._gcs_dir} needs to be empty.")
-    else:
-      self._is_gcs_path = False
-      self._local_dir = path
-      self._gcs_dir = path
-
-  def __fspath__(self) -> str:
-    return self.local_dir
-
-  @property
-  def local_dir(self) -> str:
-    return self._local_dir
-
-  @property
-  def gcs_dir(self) -> str:
-    """Returns the GCS directory path.
-
-    Returns:
-      The GCS directory path.
-
-    Raises:
-      ValueError: If the path is not a GCS path.
-    """
-    if not self._is_gcs_path:
-      raise ValueError(f"{self._gcs_dir} is not a GCS path.")
-    return self._gcs_dir
-
-  def upload_to_gcs(
-      self,
-      skip_if_exists: bool = True,
-      force_upload: bool = False,
-  ):
-    """Uploads the local directory to GCS."""
-    if not self._is_gcs_path:
-      logging.info(
-          "Not uploading to GCS since %s is not a GCS path.", self.local_dir
-      )
-      return
-
-    if not os.listdir(self.local_dir):
-      logging.info("Not uploading to GCS since %s is empty.", self.local_dir)
-      return
-
-    target = os.path.dirname(self.gcs_dir) + "/"
-    # Avoid race condition uploading the same file from multiple processes.
-    with PartialState().main_process_first():
-      if not PartialState().is_local_main_process:
-        # Non local main processes don't upload.
-        pass
-      elif self._upload_from_all_nodes or PartialState().is_main_process:
-        logging.info("Uploading %s to %s...", self.local_dir, target)
-        cmd = [
-            "gsutil",
-            "-m",
-            "cp",
-            "-r",
-        ]
-        if skip_if_exists:
-          cmd.append("-n")
-        if force_upload:
-          cmd.append("-f")
-        cmd.extend([self.local_dir, target])
-        subprocess.check_output(cmd)
-        logging.info("%s uploaded.", self.local_dir)
-
-
-def _get_local_dir_from_gcs_dir(path: str) -> str:
-  return os.path.join(
-      constants.LOCAL_OUTPUT_DIR,
-      dataset_validation_util.force_gcs_fuse_path(path)[1:],
-  )
-
-
-def _is_gcs_dir_empty(path: str) -> bool:
-  """Checks if a GCS directory is empty.
-
-  Args:
-    path: The GCS directory path.
-
-  Returns:
-    True if the directory is empty.
-
-  Raises:
-    subprocess.CalledProcessError: If the gsutil command failure reason is not
-      because the dir is empty.
-  """
-  path = path.rstrip("/") + "/"
-  try:
-    subprocess.check_output(["gsutil", "ls", path], stderr=subprocess.STDOUT)
-  except subprocess.CalledProcessError as e:
-    if (
-        str(e.output, encoding="utf-8")
-        == "CommandException: One or more URLs matched no objects.\n"
-    ):
-      return True
-    else:
-      logging.info(str(e.output, encoding="utf-8"))
-      raise
-  else:
-    return False
-
-
-def load_tokenizer(
-    pretrained_model_id: str,
-    padding_side: Optional[str] = None,
-    access_token: Optional[str] = None,
-) -> AutoTokenizer:
-  """Loads tokenizer based on `pretrained_model_id`."""
-  tokenizer_kwargs = {}
-  if should_add_eos_token(pretrained_model_id):
-    tokenizer_kwargs["add_eos_token"] = True
-  if padding_side:
-    tokenizer_kwargs["padding_side"] = padding_side
-
-  with PartialState().local_main_process_first():
-    tokenizer = AutoTokenizer.from_pretrained(
-        pretrained_model_id,
-        trust_remote_code=False,
-        use_fast=True,
-        token=access_token,
-        **tokenizer_kwargs,
-    )
-
-  if should_add_pad_token(pretrained_model_id):
-    tokenizer.add_special_tokens({"pad_token": "[PAD]"})
-
-  return tokenizer
+_GEMMA2_MODEL = "gemma-2"
 
 
 def load_model(
-    pretrained_model_id: str,
+    pretrained_model_name_or_path: str,
     tokenizer: AutoTokenizer,
     precision_mode: str = None,
-    enable_gradient_checkpointing: bool = False,
-    gradient_checkpointing_kwargs: Optional[Dict[str, Any]] = None,
-    access_token: Optional[str] = None,
-    attn_implementation: Optional[str] = None,
-    train_precision: Optional[str] = None,
-    device_map: Optional[str] = None,
+    gradient_checkpointing: bool = False,
+    gradient_checkpointing_kwargs: Mapping[str, Any] | None = None,
+    access_token: str | None = None,
+    attn_implementation: str | None = None,
+    train_precision: str | None = None,
+    device_map: str | None = None,
     is_training: bool = True,
 ) -> AutoModelForCausalLM:
   """Loads models from the local dir if specified or from huggingface."""
@@ -304,25 +116,37 @@ def load_model(
     raise ValueError(f"Invalid precision mode: {precision_mode}")
   logging.info("using torch_type=%s", torch_dtype)
 
+  model_kwargs = {
+      "use_cache": not gradient_checkpointing,
+      "device_map": device_map,
+      "torch_dtype": torch_dtype,
+      "quantization_config": quantization_config,
+      "trust_remote_code": False,
+      "token": access_token,
+      "attn_implementation": attn_implementation,
+  }
+  if _GEMMA2_MODEL in pretrained_model_name_or_path:
+    # The cache_implementation for Gemma 2 is set to hybrid by default. This
+    # param is only supported by Gemma 2. The default 'hybrid' value causes an
+    # issue when use_cache is set to False. So we have to use 'None' in such
+    # cases.
+    # https://github.com/huggingface/transformers/commit/238b13478df209ab534f2195a397dc64a3930883
+    model_kwargs["cache_implementation"] = (
+        None if gradient_checkpointing else "hybrid"
+    )
+
   model = AutoModelForCausalLM.from_pretrained(
-      pretrained_model_id,
-      use_cache=not enable_gradient_checkpointing,
-      device_map=device_map,
-      torch_dtype=torch_dtype,
-      quantization_config=quantization_config,
-      trust_remote_code=True,
-      token=access_token,
-      attn_implementation=attn_implementation,
+      pretrained_model_name_or_path, **model_kwargs
   )
 
   if precision_mode in (constants.PRECISION_MODE_4, constants.PRECISION_MODE_8):
     model = prepare_model_for_kbit_training(
         model,
-        use_gradient_checkpointing=enable_gradient_checkpointing,
+        use_gradient_checkpointing=gradient_checkpointing,
         gradient_checkpointing_kwargs=gradient_checkpointing_kwargs,
     )
 
-  if enable_gradient_checkpointing:
+  if gradient_checkpointing:
     model.gradient_checkpointing_enable(
         gradient_checkpointing_kwargs=gradient_checkpointing_kwargs
     )
@@ -345,8 +169,10 @@ def load_model(
     # https://stackoverflow.com/a/77408076
     model.config.use_cache = False
 
-  if should_add_pad_token(pretrained_model_id):
-    model.resize_token_embeddings(len(tokenizer))
+  if dataset_validation_util.should_add_pad_token(
+      pretrained_model_name_or_path
+  ):
+    model.resize_token_embeddings(len(tokenizer), mean_resizing=False)
     if is_training:
       # The following is needed since we added a new token that needs to be
       # learned.
@@ -357,22 +183,24 @@ def load_model(
 
 
 def _merge_causal_language_model_with_lora_internal(
-    pretrained_model_id: str,
+    pretrained_model_name_or_path: str,
     merge_precision_mode: str,
     finetuned_lora_model_dir: str,
     merged_model_output_dir: str,
-    access_token: Optional[str] = None,
+    access_token: str | None = None,
 ) -> None:
   """Internal function to merges the base model with the lora adapter."""
   logging.info("loading tokenizer...")
-  tokenizer = load_tokenizer(pretrained_model_id)
+  tokenizer = dataset_validation_util.load_tokenizer(
+      pretrained_model_name_or_path
+  )
 
   # Note: merging peft adapter requires loading model in 16 bits, so merging
   # is done on CPU on purpose in case one GPU cannot hold the base model.
-  logging.info("loading model %s...", pretrained_model_id)
+  logging.info("loading model %s...", pretrained_model_name_or_path)
   device_map = "cpu"
   model = load_model(
-      pretrained_model_id=pretrained_model_id,
+      pretrained_model_name_or_path=pretrained_model_name_or_path,
       tokenizer=tokenizer,
       precision_mode=merge_precision_mode,
       access_token=access_token,
@@ -402,42 +230,12 @@ def _merge_causal_language_model_with_lora_internal(
   )
 
 
-def merge_causal_language_model_with_lora_fsdp(
-    pretrained_model_id: str,
-    merge_precision_mode: str,
-    finetuned_lora_model_dir: str,
-    merged_model_output_dir: str,
-    access_token: Optional[str] = None,
-) -> None:
-  """Merges the base model with the lora adapter for FSDP.
-
-  Only the main process should call this function.
-
-  Args:
-    pretrained_model_id: Predefined base model name or path to directory
-      containing model checkpoints.
-    merge_precision_mode: Precision mode for saving model weights.
-    finetuned_lora_model_dir: Path to directory containing PEFT-finetuned model
-      weights.
-    merged_model_output_dir: Path to directory to save the merged model.
-    access_token: Access token for accessing the model.
-  """
-  assert PartialState().is_main_process
-  _merge_causal_language_model_with_lora_internal(
-      pretrained_model_id=pretrained_model_id,
-      merge_precision_mode=merge_precision_mode,
-      finetuned_lora_model_dir=finetuned_lora_model_dir,
-      merged_model_output_dir=merged_model_output_dir,
-      access_token=access_token,
-  )
-
-
 def merge_causal_language_model_with_lora(
-    pretrained_model_id: str,
+    pretrained_model_name_or_path: str,
     precision_mode: str,
     finetuned_lora_model_dir: str,
     merged_model_output_dir: str,
-    access_token: Optional[str] = None,
+    access_token: str | None = None,
 ) -> None:
   """Merges the base model with the lora adapter."""
 
@@ -452,39 +250,13 @@ def merge_causal_language_model_with_lora(
 
   if PartialState().is_main_process:
     logging.info("Starting merging job...")
-    # When deepspeed Zero3 is enabled, users are not allowed to specify
-    # `device_map` when loading the model (even on CPU).
-    #
-    # To work-around this, we kick off another process (from the
-    # is_main_process) and set up the environment to avoid using Deepspeed when
-    # doing the merging.
-    if is_deepspeed_zero3_enabled():
-      ctx = mp.get_context("spawn")
-      os.environ["ACCELERATE_USE_DEEPSPEED"] = "false"
-      merge_job = ctx.Process(
-          target=_merge_causal_language_model_with_lora_internal,
-          args=(
-              pretrained_model_id,
-              merge_precision_mode,
-              finetuned_lora_model_dir,
-              local_merged_model_dir,
-          ),
-          kwargs={
-              "access_token": access_token,
-          },
-      )
-      merge_job.start()
-      merge_job.join()
-      os.environ["ACCELERATE_USE_DEEPSPEED"] = "true"
-    else:
-      _merge_causal_language_model_with_lora_internal(
-          pretrained_model_id=pretrained_model_id,
-          merge_precision_mode=merge_precision_mode,
-          finetuned_lora_model_dir=finetuned_lora_model_dir,
-          merged_model_output_dir=local_merged_model_dir,
-          access_token=access_token,
-      )
-
+    _merge_causal_language_model_with_lora_internal(
+        pretrained_model_name_or_path=pretrained_model_name_or_path,
+        merge_precision_mode=merge_precision_mode,
+        finetuned_lora_model_dir=finetuned_lora_model_dir,
+        merged_model_output_dir=local_merged_model_dir,
+        access_token=access_token,
+    )
     logging.info("merging job is done")
 
   # Wait for all processes to sync here.
@@ -492,7 +264,7 @@ def merge_causal_language_model_with_lora(
 
   if precision_mode == constants.PRECISION_MODE_FP8:
     convert_model_to_fp8(
-        pretrained_model_name_or_path=pretrained_model_id,
+        pretrained_model_name_or_path=pretrained_model_name_or_path,
         merged_model_output_dir=local_merged_model_dir,
         quantized_model_output_dir=merged_model_output_dir,
         access_token=access_token,
@@ -503,7 +275,7 @@ def convert_model_to_fp8(
     pretrained_model_name_or_path: str,
     merged_model_output_dir: str,
     quantized_model_output_dir: str,
-    access_token: Optional[str] = None,
+    access_token: str | None = None,
 ) -> None:
   """Converts the model to fp8.
 
@@ -533,166 +305,12 @@ def convert_model_to_fp8(
   PartialState().wait_for_everyone()
 
 
-@dataclasses.dataclass
-class TuningDataStats:
-  tuning_dataset_example_count: int
-  total_billable_token_count: int
-  tuning_step_count: int
-
-
-def get_dataset_stats(
-    dataset: Any,
-    tokenizer: transformers.PreTrainedTokenizer,
-    column: str,
-    effective_batch_size: int,
-) -> TuningDataStats:
-  """Calculates dataset statistics, e.g., total number of tokens."""
-  tokenized_dataset = dataset.map(lambda x: tokenizer(x[column]))
-  inputs = tokenized_dataset["input_ids"]
-  tuning_dataset_example_count = int(len(inputs))
-  total_billable_token_count = int(np.sum([len(ex) for ex in inputs]))
-  tuning_step_count = (
-      tuning_dataset_example_count + effective_batch_size - 1
-  ) // effective_batch_size
-  return TuningDataStats(
-      tuning_dataset_example_count,
-      total_billable_token_count,
-      tuning_step_count,
-  )
-
-
 def force_gc():
   """Collects garbage immediately to release unused CPU/GPU resources."""
   gc.collect()
   torch.cuda.empty_cache()
 
 
-def should_add_pad_token(model_id: str) -> bool:
-  """Returns whether the model requires adding a special pad token."""
-  return any(s.lower() in model_id.lower() for s in _MODELS_REQUIRING_PAD_TOKEN)
-
-
-def should_add_eos_token(model_id: str) -> bool:
-  """Returns whether the model requires adding a special eos token."""
-  return any(m in model_id for m in _MODELS_REQUIRING_EOS_TOEKN)
-
-
-def write_kfp_outputs(
-    executor_input: str, output_artifacts: Dict[str, str]
-) -> None:
-  """Writes KFP outputs given a dict of output artifact names and URIs."""
-  # Only the main process writes to avoid race condition.
-  if PartialState().is_main_process:
-    executor_input = json_format.Parse(
-        executor_input, pipeline_spec_pb2.ExecutorInput()
-    )
-    outputs = executor_input.outputs
-    # set all artifacts
-    for name, uri in output_artifacts.items():
-      artifact_list = outputs.artifacts.get(name)
-      if not artifact_list or not artifact_list.artifacts:
-        raise ValueError(f"Artifact name={name} does not exist.")
-      artifact_list.artifacts[0].uri = uri
-
-    # write output file
-    executor_output = pipeline_spec_pb2.ExecutorOutput(
-        artifacts=outputs.artifacts
-    )
-    os.makedirs(os.path.dirname(outputs.output_file), exist_ok=True)
-    with open(outputs.output_file, "w") as f:
-      f.write(json_format.MessageToJson(executor_output, indent=None))
-
-  # Wait for the main process to finish before moving on to the next task.
-  PartialState().wait_for_everyone()
-
-
-def upload_local_dir_to_gcs(local_dir: str, gcs_path: str):
-  """Uploads local dir to GCS."""
-
-  if PartialState().is_main_process:
-    logging.info("uploading %s to %s...", local_dir, gcs_path)
-    subprocess.check_output([
-        "gsutil",
-        "-m",
-        "cp",
-        "-r",
-        local_dir,
-        gcs_path,
-    ])
-    logging.info("%s uploaded.", local_dir)
-
-  PartialState().wait_for_everyone()
-
-
-def write_first_party_model_metadata(output_dir: str, docker_uri: str) -> None:
-  """Multi-process friendly version of fileutils.write_first_party_model_metadata."""
-  if PartialState().is_main_process:
-    fileutils.write_first_party_model_metadata(output_dir, docker_uri)
-  PartialState().wait_for_everyone()
-
-
-@dataclasses.dataclass
-class GpuStats:
-  """Holds information about GPU usage stats.
-
-  For memory related, see
-  https://pytorch.org/docs/stable/notes/cuda.html#cuda-memory-management
-  """
-
-  # total memory
-  total_mem: float
-  # memory occupied.
-  occupied: float
-  # memory reserved, but not used.
-  unused: float
-  # nvidia-smi usually reports more memory usages than pytorch (for driver,
-  # kernel and etc). `smi_diff` tracks this difference.
-  smi_diff: float
-  # Gpu utilization.
-  util: float
-
-  # Allows unpacking operation like
-  # total_mem, occupied, unused, smi_diff, util = GpuStats(...)
-  # See https://stackoverflow.com/a/70753113
-  def __iter__(self):
-    return iter(dataclasses.astuple(self))
-
-
-def gpu_stats() -> GpuStats:
-  """Reports GPU memory usage and utilization."""
-  # See https://pytorch.org/docs/stable/notes/cuda.html#memory-management
-  bytes_per_gb = 1024.0**3
-  device = torch.cuda.current_device()
-  occupied = torch.cuda.memory_allocated(device) / bytes_per_gb
-  reserved = torch.cuda.memory_reserved(device) / bytes_per_gb
-  unused = reserved - occupied
-
-  def smi_mem(device):
-    try:
-      pynvml.nvmlInit()
-      handle = pynvml.nvmlDeviceGetHandleByIndex(device)
-      info = pynvml.nvmlDeviceGetMemoryInfo(handle)
-      return info.used / bytes_per_gb
-    except pynvml.NVMLError:
-      return 0.0
-
-  mem_used_smi = smi_mem(device)
-  smi_diff = mem_used_smi - reserved
-
-  util = torch.cuda.utilization(device)
-  return GpuStats(mem_used_smi, occupied, unused, smi_diff, util)
-
-
-def gpu_stats_str(stats: Optional[GpuStats] = None) -> str:
-  if stats is None:
-    stats = gpu_stats()
-  total, occupied, unused, smi_diff, util = stats
-  return (
-      f"GPU memory: {total:.2f}({occupied=:.2f}, {unused=:.2f},"
-      f" {smi_diff=:.2f}) GB. Utilization: {util:.2f}%"
-  )
-
-
 def init_partial_state(
     timeout: datetime.timedelta = datetime.timedelta(seconds=600),
 ) -> None:
@@ -722,7 +340,7 @@ def get_final_checkpoint_path(output_dir: str) -> str:
 
 def _maybe_get_modules_to_not_convert_by_model_id(
     pretrained_model_name_or_path: str,
-) -> Optional[Sequence[str]]:
+) -> Sequence[str] | None:
   """Returns the modules to not convert for the model."""
   if _LLAMA_3_1_405B_MODEL_ID in pretrained_model_name_or_path:
     return _get_llama_3_1_405b_modules_to_not_convert()

community-content/vertex_model_garden/model_oss/peft/train/vmg/validate_dataset_with_template.py

@@ -17,15 +17,15 @@ _DATASET_NAME = flags.DEFINE_string(
     required=True,
 )
 
-_TRAIN_SPLIT_NAME = flags.DEFINE_string(
-    'train_split_name',
+_TRAIN_SPLIT = flags.DEFINE_string(
+    'train_split',
     'train',
     'The train split name.',
 )
 
-_INSTRUCT_COLUMN_IN_DATASET = flags.DEFINE_string(
-    'instruct_column_in_dataset',
-    constants.DEFAULT_INSTRUCT_COLUMN_IN_DATASET,
+_TRAIN_COLUMN = flags.DEFINE_string(
+    'train_column',
+    constants.DEFAULT_TRAIN_COLUMN,
     'The instruct column in dataset.',
 )
 
@@ -38,6 +38,12 @@ _TEMPLATE = flags.DEFINE_string(
     required=True,
 )
 
+_MAX_SEQ_LENGTH = flags.DEFINE_integer(
+    'max_seq_length',
+    None,
+    'The maximum sequence length.',
+)
+
 _VALIDATE_PERCENTAGE_OF_DATASET = flags.DEFINE_integer(
     'validate_percentage_of_dataset',
     None,
@@ -64,9 +70,10 @@ def main(unused_argv: Sequence[str]) -> None:
 
   dataset_validation_util.validate_dataset_with_template(
       dataset_name=_DATASET_NAME.value,
-      split=_TRAIN_SPLIT_NAME.value,
-      input_column=_INSTRUCT_COLUMN_IN_DATASET.value,
+      split=_TRAIN_SPLIT.value,
+      input_column=_TRAIN_COLUMN.value,
       template=_TEMPLATE.value,
+      max_seq_length=_MAX_SEQ_LENGTH.value,
       use_multiprocessing=_USE_MULTIPROCESSING.value,
       validate_percentage_of_dataset=_VALIDATE_PERCENTAGE_OF_DATASET.value,
       validate_k_rows_of_dataset=_VALIDATE_K_ROWS_OF_DATASET.value,

community-content/vertex_model_garden/model_oss/util/constants.py

@@ -90,12 +90,11 @@ TEXT_TO_IMAGE_DREAMBOOTH_LORA = 'text-to-image-dreambooth-lora'
 TEXT_TO_IMAGE_DREAMBOOTH_LORA_SDXL = 'text-to-image-dreambooth-lora-sdxl'
 SEQUENCE_CLASSIFICATION_LORA = 'sequence-classification-lora'
 MERGE_CAUSAL_LANGUAGE_MODEL_LORA = 'merge-causal-language-model-lora'
-QUANTIZE_MODEL = 'quantize-model'
 INSTRUCT_LORA = 'instruct-lora'
 VALIDATE_DATASET_WITH_TEMPLATE = 'validate-dataset-with-template'
+RUN_TESTS = 'test'
 DEFAULT_TEXT_COLUMN_IN_DATASET = 'quote'
-DEFAULT_TEXT_COLUMN_IN_QUANTIZATION_DATASET = 'text'
-DEFAULT_INSTRUCT_COLUMN_IN_DATASET = 'text'
+DEFAULT_TRAIN_COLUMN = 'text'
 
 FINAL_CHECKPOINT_DIRNAME = 'checkpoint-final'
 
@@ -113,13 +112,17 @@ PRECISION_MODE_16 = 'float16'
 PRECISION_MODE_16B = 'bfloat16'
 PRECISION_MODE_32 = 'float32'
 
-# Quantization modes.
-GPTQ = 'gptq'
-AWQ = 'awq'
+ROUGE_VARIANTS = ('rouge1', 'rouge2', 'rougeL', 'rougeLsum')
 
-# AWQ versions.
-GEMM = 'GEMM'
-GEMV = 'GEMV'
+# Supported HF evaluation metrics.
+SUPPORTED_HF_EVAL_METRICS = (
+    'perplexity',
+    'bleu',
+    'google_bleu',
+) + ROUGE_VARIANTS
+
+# Supported evaluation metrics.
+SUPPORTED_EVAL_METRICS = ('loss',) + SUPPORTED_HF_EVAL_METRICS
 
 # Environment variable keys.
 PRIVATE_BUCKET_ENV_KEY = 'AIP_PRIVATE_BUCKET_NAME'

community-content/vertex_model_garden/model_oss/util/fileutils.py

@@ -1,9 +1,12 @@
 """Fileutil lib to copy files between gcs and local."""
 
+import filecmp
 import fnmatch
 import os
 import pathlib
 import shutil
+import subprocess
+import time
 from typing import List, Optional, Tuple
 import uuid
 
@@ -57,6 +60,113 @@ def force_gcs_path(uri: str) -> str:
     return uri
 
 
+def is_file_available(
+    file_path: str, retry_interval_secs: int = 60, timeout_secs: int = 3600
+) -> bool:
+  """Checks and waits for a file to be available in GCS.
+
+  Args:
+    file_path: The file path to check.
+    retry_interval_secs: The interval in seconds to check the file.
+    timeout_secs: The timeout in seconds to wait for the file.
+
+  Returns:
+    True if the file is available, False otherwise.
+  """
+  start_time = time.time()
+  while True:
+    try:
+      file_check_cmd = ['gcloud', 'storage', 'ls', file_path]
+      result = subprocess.run(
+          file_check_cmd, capture_output=True, text=True, check=True
+      )
+      if file_path in result.stdout:
+        logging.info('File %s exists.', file_path)
+        return True
+    except subprocess.CalledProcessError as e:
+      elapsed_time = time.time() - start_time
+      if elapsed_time > timeout_secs:
+        logging.info(
+            "Timeout: File '%s' not found after %d seconds. Error: %s",
+            file_path,
+            elapsed_time,
+            e,
+        )
+        return False
+
+      logging.info(
+          "File '%s' not found yet. Checking again in %d seconds. Error: %s",
+          file_path,
+          retry_interval_secs,
+          e,
+      )
+      time.sleep(retry_interval_secs)
+
+
+def compare_dirs(
+    local_dir: str,
+    gcsfuse_dir: str,
+    retry_interval_secs: int = 30,
+    timeout_secs: int = 3600,
+) -> bool:
+  """Compares two directories and returns True if they are the same.
+
+  Args:
+    local_dir: The local directory.
+    gcsfuse_dir: The gcsfuse directory.
+    retry_interval_secs: The interval in seconds to check the directories.
+    timeout_secs: The timeout in seconds to wait for the directories.
+
+  Returns:
+    True if the directories are the same, False otherwise.
+  """
+  start_time = time.time()
+  while True:
+    if os.path.exists(local_dir) and os.path.exists(gcsfuse_dir):
+      comparison = filecmp.dircmp(local_dir, gcsfuse_dir)
+      if (
+          not comparison.left_only
+          and not comparison.right_only
+          and not comparison.diff_files
+      ):
+        return True
+    elapsed_time = time.time() - start_time
+    if elapsed_time > timeout_secs:
+      logging.info(
+          "Timeout: Directories '%s' and '%s' do not match after %d seconds.",
+          local_dir,
+          gcsfuse_dir,
+          elapsed_time,
+      )
+      return False
+
+    logging.info(
+        "Directories '%s' and '%s' do not match yet. Checking again in %d"
+        ' seconds.',
+        local_dir,
+        gcsfuse_dir,
+        retry_interval_secs,
+    )
+    time.sleep(retry_interval_secs)
+
+
+def download_gcs_file_to_memory(gcs_uri: str) -> bytes:
+  """Downloads a gcs file to in memory.
+
+  Args:
+    gcs_uri: A string of GCS uri.
+
+  Returns:
+    The content of the gcs file in byte format.
+  """
+  bucket = gcs_uri.split('/')[2]
+  file_path = gcs_uri[len(constants.GCS_URI_PREFIX + bucket + '/') :]
+  client = _get_gcs_client()
+  bucket = client.bucket(bucket)
+  blob = bucket.blob(file_path)
+  return blob.download_as_bytes()
+
+
 def download_gcs_file_to_local_dir(gcs_uri: str, local_dir: str):
   """Download a gcs file to a local dir.
 
@@ -337,22 +447,13 @@ def get_output_video_file(video_output_file_path: str) -> str:
   return out_local_video_file_name
 
 
-def write_first_party_model_metadata(
-    output_path: str, required_container_uri: str
-) -> None:
-  """Write Vertex internal model metadata for first party artifacts."""
-  model_metadata_fname = 'model_metadata.jsonl'
-  if len(required_container_uri) > 126:
-    raise ValueError(f'Docker URI exceeds 126 chars: {required_container_uri}')
-  payload = '\n{}{}'.format(  # serialized proto
-      chr(len(required_container_uri)),
-      required_container_uri,
-  )
-  os.makedirs(output_path, exist_ok=True)
-  output_dirs = [output_path]
-  if output_path.startswith('/gcs'):
-    # include all parent dirs, except "/", "/gcs"
-    output_dirs.extend([str(p) for p in pathlib.Path(output_path).parents][:-2])
-  for output_dir in output_dirs:
-    with open(os.path.join(output_dir, model_metadata_fname), 'w') as f:
-      f.write(payload)
+def delete_local_file(local_file_path: str) -> None:
+  """Deletes a local file."""
+  if os.path.exists(local_file_path):
+    os.remove(local_file_path)
+
+
+def delete_local_dir(local_dir: str) -> None:
+  """Deletes a local directory recursively."""
+  if os.path.exists(local_dir):
+    shutil.rmtree(local_dir)

community-content/vertex_model_garden/model_oss/util/gcs_download_launcher.sh

@@ -0,0 +1,119 @@
+#!/bin/bash
+#
+# This launcher downloads model files from GCS to local model directory before
+# launching the actual command.
+#
+# If GCS URI is passed as an environment variable, set GCS_URI_ENV_KEY to the
+# environment variable name.
+# If GCS URI is passed as an argument, set GCS_URI_ARG_KEY to the argument name.
+# The argument must be in the format of '--$GCS_URI_ARG_KEY=gs://*'. Do not
+# separate argument name and value with spaces.
+# This script will also try reading from AIP_STORAGE_URI or AIP_STORAGE_DIR.
+# Note that AIP_STORAGE_DIR is expected to be a local path, so it bypasses the
+# download process.
+#
+# Input priority: AIP_STORAGE_DIR > AIP_STORAGE_URI > GCS_URI_ENV_KEY > GCS_URI_ARG_KEY.
+# Will output the local model directory to GCS_URI_ENV_KEY and GCS_URI_ARG_KEY
+# if they are set. Both will be updated if both set.
+#
+# Requires google-cloud-sdk as a dependency (for gcloud storage CLI).
+
+set -e
+
+readonly LOCAL_MODEL_DIR=${LOCAL_MODEL_DIR:-"/tmp/model_dir"}
+readonly LOCAL_ARGS_FILE=${LOCAL_ARGS_FILE:-"/tmp/args.txt"}
+
+update_model_id() {
+  if [[ ! -z "$GCS_URI_ENV_KEY" ]]; then
+    echo "Updating env var $GCS_URI_ENV_KEY to $AIP_STORAGE_DIR."
+    export "$GCS_URI_ENV_KEY"="$AIP_STORAGE_DIR"
+  fi
+
+  if [[ ! -z "$GCS_URI_ARG_KEY" ]]; then
+    echo "Updating args $GCS_URI_ARG_KEY to $AIP_STORAGE_DIR."
+    updated=0
+    for (( i=1; i <= $#; i++)); do
+      arg="${!i}"
+      if [[ "$arg" == "--$GCS_URI_ARG_KEY="* ]]; then
+        echo "Found $arg, updating to $AIP_STORAGE_DIR."
+        set -- "${@:1:(($i-1))}" "--$GCS_URI_ARG_KEY=$AIP_STORAGE_DIR" "${@:$(($i+1))}";
+        updated=1
+        break
+      fi
+    done
+    if [[ $updated -eq 0 ]]; then
+      echo "Appending args $GCS_URI_ARG_KEY to $AIP_STORAGE_DIR."
+      set -- "$@" "--$GCS_URI_ARG_KEY=$AIP_STORAGE_DIR";
+    fi
+  fi
+  echo "$*" > "$LOCAL_ARGS_FILE"
+}
+
+maybe_download_model() {
+  if [[ -z "$GCS_URI_ENV_KEY" ]] && [[ -z "$GCS_URI_ARG_KEY" ]]; then
+    echo "Internal error: Required GCS_URI_ENV_KEY or GCS_URI_ARG_KEY."
+    exit 1
+  fi
+
+  echo "$*" > "$LOCAL_ARGS_FILE"
+  gcs_uri=""
+  if [[ ! -z "$AIP_STORAGE_DIR" ]]; then
+    # AIP_STORAGE_DIR is expected to be a local path.
+    echo "AIP_STORAGE_DIR set, proceeding to run the launcher."
+    update_model_id "$@"
+    return
+  elif [[ $AIP_STORAGE_URI == gs://* ]]; then
+    # Check AIP_STORAGE_URI environment variable.
+    echo "AIP_STORAGE_URI set and starts with 'gs://', proceeding to download from GCS."
+    gcs_uri="$AIP_STORAGE_URI"
+  elif [[ ! -z "$GCS_URI_ENV_KEY" ]] && [[ ${!GCS_URI_ENV_KEY} == gs://* ]]; then
+    # Check custom environment variable.
+    echo "Custom environment variable ${GCS_URI_ENV_KEY} set and starts with 'gs://', proceeding to download from GCS."
+    gcs_uri="${!GCS_URI_ENV_KEY}"
+  elif [[ ! -z "$GCS_URI_ARG_KEY" ]]; then
+    # Check custom args.
+    for arg in "$@"; do
+      if [[ "$arg" == "--$GCS_URI_ARG_KEY=gs://"* ]]; then
+        gcs_uri="${arg#*=}"
+        echo "Custom args ${GCS_URI_ARG_KEY} set and starts with 'gs://', proceeding to download from GCS."
+        break
+      elif [[ "$arg" == "--$GCS_URI_ARG_KEY" ]]; then
+        echo "Found $GCS_URI_ARG_KEY, but it's not in the format of '--$GCS_URI_ARG_KEY=gs://*'."
+        echo "Ensure the value of $GCS_URI_ARG_KEY is within the same arg, separated by '='."
+        exit 1
+      fi
+    done
+  fi
+
+  if [[ -z "$gcs_uri" ]]; then
+    echo "No GCS URI found, proceeding to run the launcher."
+    return
+  fi
+
+  # Remove trailing '/' if any.
+  gcs_uri="${gcs_uri%%/}"
+  export AIP_STORAGE_DIR="$LOCAL_MODEL_DIR/${gcs_uri##gs://}"
+
+  # Create the target directory.
+  mkdir -p "$AIP_STORAGE_DIR"
+  echo "Downloading model from ${gcs_uri} to ${AIP_STORAGE_DIR}."
+
+  # Use gcloud storage CLI to copy the content from GCS to the target directory.
+  if gcloud storage cp -r "$gcs_uri/*" "$AIP_STORAGE_DIR"; then
+    echo "Model downloaded successfully to ${AIP_STORAGE_DIR}."
+    update_model_id "$@"
+  else
+    echo "Failed to download model from GCS."
+    exit 1
+  fi
+}
+
+run_local_command() {
+  command=$(cat "$LOCAL_ARGS_FILE")
+  rm -f "$LOCAL_ARGS_FILE"
+  echo "Launch command: $command"
+  eval "$command"
+}
+
+maybe_download_model "$@"
+run_local_command

community-content/vertex_model_garden/model_oss/util/gcs_syncer.py

@@ -0,0 +1,176 @@
+"""Sync local directory to GCS directory using rsync."""
+
+import multiprocessing
+import os
+import subprocess
+import time
+from typing import Optional, Sequence, Tuple
+
+from absl import logging
+
+from util import constants
+from util import fileutils
+
+_GCS_COMMAND_RETRIES = 3
+_RSYNC_RETRY_INTERVAL_SECS = 30
+
+
+def is_gcs_or_gcsfuse_path(path: str) -> bool:
+  """Returns if the path is a GCS or gcsfuse path.
+
+  Args:
+    path: The path to check.
+
+  Returns:
+    True if the path is a GCS or gcsfuse path.
+  """
+  return path.startswith(
+      (constants.GCS_URI_PREFIX, constants.GCSFUSE_URI_PREFIX)
+  )
+
+
+def manage_sync_path(
+    path: str, node_rank: Optional[int] = None
+) -> Tuple[str, str]:
+  """Returns local dir and GCS location for the given path if the given path is a GCS or gcsfuse path.
+
+  It will also create a local directory if it does not exist. Otherwise, it
+  returns the same path.
+
+  Args:
+    path: The local or GCS path to manage.
+    node_rank: The node rank to be appended to the GCS path.
+
+  Returns:
+    The local and GCS paths.
+  """
+  local_dir = path
+  gcs_dir = path
+  if is_gcs_or_gcsfuse_path(path):
+    local_dir = os.path.join(
+        constants.LOCAL_OUTPUT_DIR,
+        fileutils.force_gcs_fuse_path(path)[1:],
+    )
+    gcs_dir = fileutils.force_gcs_path(path)
+  if not os.path.exists(local_dir):
+    os.makedirs(local_dir, exist_ok=True)
+
+  if node_rank is None:
+    return local_dir, gcs_dir
+  return local_dir, os.path.join(gcs_dir, f"node-{node_rank}")
+
+
+def setup_gcs_rsync(
+    dirs_to_sync: Sequence[Tuple[str, str]],
+    mp_queue: multiprocessing.Queue,
+    gcs_rsync_interval_secs: int,
+) -> multiprocessing.Process:
+  """Sets up the GCS rsync process.
+
+  Args:
+    dirs_to_sync: The absolute directory paths which will be synced to GCS.
+    mp_queue: The multiprocessing queue to check if the training is finished.
+    gcs_rsync_interval_secs: Integer, interval in seconds to run gcs rsync.
+
+  Returns:
+    The GCS rsync process.
+  """
+  rsync_process = multiprocessing.Process(
+      target=start_gcs_rsync,
+      args=(dirs_to_sync, mp_queue, gcs_rsync_interval_secs),
+  )
+  rsync_process.start()
+  return rsync_process
+
+
+def cleanup_gcs_rsync(
+    rsync_process: multiprocessing.Process, mp_queue: multiprocessing.Queue
+) -> None:
+  """Cleans up the GCS rsync process.
+
+  Args:
+    rsync_process: The GCS rsync process.
+    mp_queue: The multiprocessing queue.
+  """
+  mp_queue.put("finish rsync process")
+  rsync_process.join()
+  if rsync_process.exitcode == 0:
+    logging.info("Artifacts have been uploaded to GCS.")
+  else:
+    logging.error(
+        "GCS rsync process failed with exit code %d.", rsync_process.exitcode
+    )
+
+
+def _rsync_local_to_gcs(local_dir: str, gcs_dir: str) -> None:
+  """Syncs the local directory to GCS.
+
+  Args:
+    local_dir: The local directory to sync.
+    gcs_dir: The GCS directory to sync to.
+  """
+  if not os.listdir(local_dir):
+    logging.info("Not rsyncing to GCS since %s is empty.", local_dir)
+    return
+
+  logging.info("Rsyncing %s <--> %s...", local_dir, gcs_dir)
+  cmd = [
+      "gcloud",
+      "storage",
+      "rsync",
+      "-r",
+      "--delete-unmatched-destination-objects",
+  ]
+  cmd.extend([local_dir, gcs_dir])
+
+  attempt = 0
+  while attempt < _GCS_COMMAND_RETRIES:
+    try:
+      subprocess.check_output(cmd)
+      break
+    except subprocess.CalledProcessError as e:
+      attempt += 1
+      if attempt < _GCS_COMMAND_RETRIES:
+        logging.exception(
+            "Attempt %d: Command failed: %s. Retrying in %d seconds...",
+            attempt,
+            e,
+            _RSYNC_RETRY_INTERVAL_SECS,
+        )
+        time.sleep(_RSYNC_RETRY_INTERVAL_SECS)
+      else:
+        logging.exception(
+            "Command failed after %d attempts: %s.", e, _GCS_COMMAND_RETRIES
+        )
+
+  logging.info("%s rsynced to %s.", local_dir, gcs_dir)
+
+
+def start_gcs_rsync(
+    dirs_to_sync: Sequence[Tuple[str, str]],
+    mp_queue: multiprocessing.Queue,
+    gcs_rsync_interval_secs: int,
+) -> None:
+  """Starts a rsync process to sync local directories to GCS directories.
+
+  Args:
+    dirs_to_sync: A list of tuples, where each tuple contains local directory
+      which will be synced to GCS. For example: [('/tmp/local_dir_1',
+      'gs://bucket/gcs_dir_1'), ('/tmp/local_dir_2', 'gs://bucket/gcs_dir_2')]
+    mp_queue: The multiprocessing queue to check if the training is finished.
+    gcs_rsync_interval_secs: Integer, interval in seconds to run gcs rsync.
+  """
+  while True:
+    for local_dir, gcs_dir in dirs_to_sync:
+      _rsync_local_to_gcs(local_dir, gcs_dir)
+    if not mp_queue.empty():
+      break
+    time.sleep(gcs_rsync_interval_secs)
+
+  # Sync up the directory one more time to avoid a race condition.
+  # There can be a case when we are doing an rsync and receive a signal that
+  # the training has been done. The final checkpoint will be skipped in such
+  # case. So we do a final sync to make sure that the all directories
+  # are synced.
+  for local_dir, gcs_dir in dirs_to_sync:
+    _rsync_local_to_gcs(local_dir, gcs_dir)

community-content/vertex_model_garden/model_oss/util/pytorch_startup_prober.sh

@@ -0,0 +1,35 @@
+#!/bin/bash
+
+# !/bin/bash
+# The Startup prober built to check whether models listed in local disk are
+# loaded in memory and are ready to serve traffic. The script returns 0 if
+# succeed. Any other returned value are consider as an error. More detail could be
+# found from [shell script Exit codes](http://shellscript.sh/exitcodes.html).
+#
+# TorchServe: The Management API listens on port 8081 and is only accessible
+# from localhost by default.
+
+if [[ -z "${MNG_PORT}" ]]; then
+  MNG_PORT=7081  # We default the management_port to 7081.
+else
+  MNG_PORT="${MNG_PORT}"
+fi
+
+check_model_availability(){
+  local MODEL_NAME=$1
+  # Returns whether "READY" is found in the model status.
+  # Reference: https://pytorch.org/serve/management_api.html#describe-model.
+  curl -s "http://localhost:${MNG_PORT}/models/${MODEL_NAME}" | grep "READY" -q
+}
+
+main(){
+  check_model_availability "$MODEL"  # Assume Dockerfile sets MODEL environment parameter.
+  local available=$?
+  if [[ $available -gt 0 ]]
+  then
+    echo "Warning: Model(${MODEL}) is not yet available."
+    return 1
+  fi
+  return 0
+}
+main

notebooks/community/model_garden/gke_model_ui_deployment_notebook.ipynb

@@ -0,0 +1,862 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "Pr9TgOcV9vAXeqGiyTaTI5kS",
+      "metadata": {
+        "cellView": "form",
+        "id": "Pr9TgOcV9vAXeqGiyTaTI5kS"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2025 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",
+      "id": "M1CpgYundFwz",
+      "metadata": {
+        "id": "M1CpgYundFwz"
+      },
+      "source": [
+        "# Get started with your deployed model on GKE\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%2Fgke_model_ui_deployment_notebook.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/gke_model_ui_deployment_notebook.ipynb\">\n",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "t2jj2XOgkS4F",
+      "metadata": {
+        "id": "t2jj2XOgkS4F"
+      },
+      "source": [
+        "# Overview\n",
+        "\n",
+        "This notebook will guide you through the initial step of testing your recently\n",
+        "deployed model with text prompts. Depending on your deployed model's inference\n",
+        "setup, the notebook utilizes either Text Generation Inference\n",
+        "[TGI](https://huggingface.co/docs/text-generation-inference/en/index) or\n",
+        "[vLLM](https://developers.googleblog.com/en/inference-with-gemma-using-dataflow-and-vllm/#:~:text=model%20frameworks%20simple.-,What%20is%20vLLM%3F,-vLLM%20is%20an),\n",
+        "two efficient serving frameworks that enhance the performance of your GPU model.\n",
+        "Ready to see your deployed model respond? Run the cells below and start\n",
+        "experimenting with different prompts!\n",
+        "\n",
+        "### Prerequisites\n",
+        "\n",
+        "Before proceeding with this notebook, ensure you have already deployed a model\n",
+        "using the Google Cloud Console. You can find an overview of AI and Machine\n",
+        "Learning services on\n",
+        "[GKE AI/ML](https://console.cloud.google.com/kubernetes/aiml/overview).\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "Enable prompt-based testing of the AI model deployed on GKE\n",
+        "\n",
+        "### GPUs\n",
+        "\n",
+        "GPUs let you accelerate specific workloads running on your nodes, such as\n",
+        "machine learning and data processing. GKE provides a range of machine type\n",
+        "options for node configuration, including machine types with NVIDIA H100, L4,\n",
+        "and A100 GPUs.\n",
+        "\n",
+        "### Understanding the Inference Frameworks\n",
+        "\n",
+        "Your model is running on one of two popular and efficient serving frameworks:\n",
+        "vLLM or Text Generation Inference (TGI). The following sections provide a brief\n",
+        "overview of each to give you context on the underlying technology powering your\n",
+        "model.\n",
+        "\n",
+        "#### TGI\n",
+        "\n",
+        "TGI is a highly optimized open-source LLM serving framework that can increase\n",
+        "serving throughput on GPUs. TGI includes features such as:\n",
+        "\n",
+        "*   Optimized transformer implementation with PagedAttention\n",
+        "*   Continuous batching to improve the overall serving throughput\n",
+        "*   Tensor parallelism and distributed serving on multiple GPUs\n",
+        "\n",
+        "To learn more, refer to the\n",
+        "[TGI documentation](https://github.com/huggingface/text-generation-inference/blob/main/README.md)\n",
+        "\n",
+        "#### vLLM\n",
+        "\n",
+        "vLLM is another fast and easy-to-use library for LLM inference and serving. It's\n",
+        "known for its high throughput and efficiency, and it leverages PagedAttention.\n",
+        "Key features include:\n",
+        "\n",
+        "*   PagedAttention: Efficient memory management for handling long sequences and\n",
+        "    dynamic workloads.\n",
+        "*   Continuous batching: Maximizes GPU utilization by batching incoming\n",
+        "    requests.\n",
+        "*   High-throughput serving: Designed for production-level serving with low\n",
+        "    latency.\n",
+        "*   Optimized CUDA kernels.\n",
+        "\n",
+        "To learn more, refer to the\n",
+        "[vLLM documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/vllm/use-vllm)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "XMf-T58TkDy1",
+      "metadata": {
+        "cellView": "form",
+        "id": "XMf-T58TkDy1"
+      },
+      "outputs": [],
+      "source": [
+        "# @title # Connect to Google Cloud Project\n",
+        "# @markdown #### Run this cell to configure your Google Cloud environment for Kubernetes (GKE) operations.\n",
+        "# @markdown\n",
+        "# @markdown #### Actions:\n",
+        "# @markdown 1.  **Connects to Project:** Retrieves and sets your Google Cloud project ID.\n",
+        "# @markdown 3.  **Installs `kubectl`:** Installs the Kubernetes command-line tool.\n",
+        "\n",
+        "import os\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Set up gcloud.\n",
+        "! gcloud config set project \"$PROJECT_ID\"\n",
+        "! gcloud services enable container.googleapis.com\n",
+        "\n",
+        "# Add kubectl to the set of available tools.\n",
+        "! mkdir -p /tools/google-cloud-sdk/.install\n",
+        "! gcloud components install kubectl --quiet"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "1oG8ymQenHyD",
+      "metadata": {
+        "cellView": "form",
+        "id": "1oG8ymQenHyD"
+      },
+      "outputs": [],
+      "source": [
+        "# @title # Select Cluster and Deployment { vertical-output: true }\n",
+        "# @markdown **Instructions:**\n",
+        "# @markdown\n",
+        "# @markdown Run this cell using the ▶ button. Then, use the interactive widgets that appear below:\n",
+        "# @markdown 1. **Select Cluster:** From the first dropdown, choose the GKE cluster where your model deployment is running. Note: the list only contains autopilot clusters.\n",
+        "# @markdown 2.   **Select Namespace:** After selecting a cluster, choose the Kubernetes *Namespace* where your deployment resides within that cluster.\n",
+        "# @markdown 3. **Select Deployment:** After selecting a cluster, this dropdown will populate with the names of deployments found.\n",
+        "\n",
+        "import json\n",
+        "import subprocess\n",
+        "\n",
+        "import ipywidgets as widgets\n",
+        "from IPython.display import Markdown, clear_output, display\n",
+        "\n",
+        "# --- Globals and Configuration ---\n",
+        "DEFAULT_NAMESPACE = \"default\"\n",
+        "SELECTED_DEPLOYMENT = None\n",
+        "SELECTED_NAMESPACE = DEFAULT_NAMESPACE\n",
+        "deployment_dropdown = None\n",
+        "namespace_dropdown = None\n",
+        "cluster_dropdown = None\n",
+        "output_area = widgets.Output()\n",
+        "\n",
+        "\n",
+        "# --- Data Fetching Functions ---\n",
+        "def get_clusters(project_id):\n",
+        "    \"\"\"Fetches autopilot GKE clusters for a given project.\"\"\"\n",
+        "    # Note: Uses broad exception handling as per original code.\n",
+        "    try:\n",
+        "        cmd = f\"gcloud container clusters list --filter=autopilot.enabled=true --format=json --project={project_id}\"\n",
+        "        result = subprocess.run(\n",
+        "            cmd, shell=True, capture_output=True, text=True, check=True, timeout=60\n",
+        "        )\n",
+        "        clusters_data = json.loads(result.stdout)\n",
+        "        # Create a map of cluster name to its region/location\n",
+        "        return {c[\"name\"]: c[\"location\"] for c in clusters_data}\n",
+        "    except Exception as e:\n",
+        "        # Original code prints error and returns empty dict\n",
+        "        print(f\"Error getting clusters: {e}\")\n",
+        "        return {}\n",
+        "\n",
+        "\n",
+        "# Fetch clusters immediately using PROJECT_ID assumed to be globally defined\n",
+        "# Note: This relies on PROJECT_ID being set *before* this cell runs.\n",
+        "try:\n",
+        "    CLUSTER_REGION_MAP = get_clusters(PROJECT_ID)\n",
+        "except NameError:\n",
+        "    print(\n",
+        "        \"Error: PROJECT_ID variable is not defined. Please define it in a previous cell.\"\n",
+        "    )\n",
+        "    CLUSTER_REGION_MAP = {}  # Define as empty to prevent errors later\n",
+        "\n",
+        "\n",
+        "def get_deployments(cluster, region, namespace):\n",
+        "    \"\"\"Fetches deployments from a specific namespace in a cluster.\"\"\"\n",
+        "    # Note: Uses PROJECT_ID as a global variable as per original code.\n",
+        "    # Note: Uses broad exception handling as per original code.\n",
+        "    target_namespace = namespace if namespace else DEFAULT_NAMESPACE\n",
+        "    try:\n",
+        "        # Ensure credentials for the target cluster\n",
+        "        cred_cmd = [\n",
+        "            \"gcloud\",\n",
+        "            \"container\",\n",
+        "            \"clusters\",\n",
+        "            \"get-credentials\",\n",
+        "            cluster,\n",
+        "            f\"--location={region}\",\n",
+        "            f\"--project={PROJECT_ID}\",\n",
+        "        ]\n",
+        "        subprocess.run(cred_cmd, capture_output=True, text=True, check=True, timeout=60)\n",
+        "\n",
+        "        # Fetch deployments using kubectl\n",
+        "        kubectl_cmd = [\n",
+        "            \"kubectl\",\n",
+        "            \"get\",\n",
+        "            \"deployments\",\n",
+        "            f\"--namespace={target_namespace}\",\n",
+        "            \"-o\",\n",
+        "            \"json\",\n",
+        "        ]\n",
+        "        result = subprocess.run(\n",
+        "            kubectl_cmd, capture_output=True, text=True, check=True, timeout=60\n",
+        "        )\n",
+        "        deployments_data = json.loads(result.stdout)\n",
+        "        # Extract deployment names\n",
+        "        return [item[\"metadata\"][\"name\"] for item in deployments_data.get(\"items\", [])]\n",
+        "    except Exception as e:\n",
+        "        # Original code prints error and returns empty list\n",
+        "        print(f\"Error fetching deployments from namespace '{target_namespace}': {e}\")\n",
+        "        return []\n",
+        "\n",
+        "\n",
+        "def get_namespaces(cluster, region, project_id):\n",
+        "    \"\"\"Fetches namespaces for a given cluster.\"\"\"\n",
+        "    # Note: Uses broad exception handling as per original code.\n",
+        "    try:\n",
+        "        # Ensure credentials for the target cluster\n",
+        "        cred_cmd = [\n",
+        "            \"gcloud\",\n",
+        "            \"container\",\n",
+        "            \"clusters\",\n",
+        "            \"get-credentials\",\n",
+        "            cluster,\n",
+        "            f\"--location={region}\",\n",
+        "            f\"--project={project_id}\",\n",
+        "        ]\n",
+        "        subprocess.run(cred_cmd, capture_output=True, text=True, check=True, timeout=60)\n",
+        "\n",
+        "        # Fetch namespaces using kubectl\n",
+        "        kubectl_cmd = [\"kubectl\", \"get\", \"namespaces\", \"-o\", \"json\"]\n",
+        "        result = subprocess.run(\n",
+        "            kubectl_cmd, capture_output=True, text=True, check=True, timeout=60\n",
+        "        )\n",
+        "        namespaces_data = json.loads(result.stdout)\n",
+        "        # Extract namespace names\n",
+        "        all_ns = [item[\"metadata\"][\"name\"] for item in namespaces_data.get(\"items\", [])]\n",
+        "        return all_ns\n",
+        "    except Exception as e:\n",
+        "        # Original code displays error in output_area and returns None\n",
+        "        with output_area:\n",
+        "            # Clear previous output before showing error\n",
+        "            clear_output(wait=True)\n",
+        "            display(\n",
+        "                Markdown(\n",
+        "                    f\"<font color='red'>Error processing namespaces for **{cluster}**: {e}</font>\"\n",
+        "                )\n",
+        "            )\n",
+        "        return None\n",
+        "\n",
+        "\n",
+        "# --- Event Handlers ---\n",
+        "def on_deployment_select(change):\n",
+        "    \"\"\"Handles changes in the deployment selection.\"\"\"\n",
+        "    global SELECTED_DEPLOYMENT\n",
+        "    if change[\"type\"] == \"change\" and change[\"name\"] == \"value\":\n",
+        "        SELECTED_DEPLOYMENT = change[\"new\"]\n",
+        "        with output_area:\n",
+        "            clear_output(wait=True)\n",
+        "            current_cluster = cluster_dropdown.value\n",
+        "\n",
+        "            # Display context message\n",
+        "            if current_cluster != \"Select Cluster\":\n",
+        "                # Use SELECTED_NAMESPACE global which should be set by on_namespace_change\n",
+        "                # or default if namespace hasn't been selected yet.\n",
+        "                ns_context = SELECTED_NAMESPACE or DEFAULT_NAMESPACE\n",
+        "                ns_info = f\"Cluster: **{current_cluster}**, Namespace: **{ns_context}**\"\n",
+        "                display(Markdown(ns_info))\n",
+        "\n",
+        "            # Display selection message if a valid deployment is chosen\n",
+        "            if (\n",
+        "                SELECTED_DEPLOYMENT\n",
+        "                and SELECTED_DEPLOYMENT != \"Select Deployment\"\n",
+        "                and SELECTED_DEPLOYMENT != \"Loading...\"\n",
+        "            ):\n",
+        "                mes = f\"\"\"Selected deployment: **{SELECTED_DEPLOYMENT}**\"\"\"\n",
+        "                display(Markdown(mes))\n",
+        "\n",
+        "\n",
+        "def update_deployment_dropdown(cluster_name, namespace_to_use):\n",
+        "    \"\"\"Updates the deployment list based on cluster/namespace change.\"\"\"\n",
+        "    global deployment_dropdown, SELECTED_DEPLOYMENT\n",
+        "    target_namespace = namespace_to_use if namespace_to_use else DEFAULT_NAMESPACE\n",
+        "\n",
+        "    # Reset selection before fetching/updating\n",
+        "    SELECTED_DEPLOYMENT = None\n",
+        "    deployment_dropdown.disabled = True  # Disable while loading/updating\n",
+        "    deployment_dropdown.options = [\"Loading...\"]\n",
+        "    deployment_dropdown.value = \"Loading...\"\n",
+        "\n",
+        "    # Clear output area and show loading context\n",
+        "    with output_area:\n",
+        "        clear_output(wait=True)\n",
+        "        display(Markdown(f\"Cluster: **{cluster_name}**\"))\n",
+        "        if namespace_to_use:\n",
+        "            display(Markdown(f\"Namespace: **{namespace_to_use}**\"))\n",
+        "        display(Markdown(\"Fetching deployments...\"))\n",
+        "\n",
+        "    # Fetch deployments (assuming CLUSTER_REGION_MAP and PROJECT_ID are available)\n",
+        "    region = CLUSTER_REGION_MAP.get(cluster_name)\n",
+        "    if not region:\n",
+        "        with output_area:\n",
+        "            clear_output(wait=True)\n",
+        "            display(\n",
+        "                Markdown(\n",
+        "                    f\"<font color='red'>Error: Region not found for cluster {cluster_name}.</font>\"\n",
+        "                )\n",
+        "            )\n",
+        "        deployment_dropdown.options = [\"Error loading\"]\n",
+        "        deployment_dropdown.value = \"Error loading\"\n",
+        "        return  # Stop if region is missing\n",
+        "\n",
+        "    deployments = get_deployments(cluster_name, region, target_namespace)\n",
+        "\n",
+        "    # Update dropdown options\n",
+        "    new_options = [\"Select Deployment\"] + deployments\n",
+        "    deployment_dropdown.options = new_options\n",
+        "\n",
+        "    # Set final state based on results\n",
+        "    if deployments:\n",
+        "        deployment_dropdown.value = \"Select Deployment\"\n",
+        "        deployment_dropdown.disabled = False\n",
+        "        status_message = f\"Found {len(deployments)} deployment(s) in namespace **{target_namespace}**.\"\n",
+        "    else:\n",
+        "        deployment_dropdown.value = \"Select Deployment\"  # Keep prompt\n",
+        "        deployment_dropdown.disabled = True  # No valid options to select\n",
+        "        # Check if get_deployments printed an error or if it just returned empty\n",
+        "        if not output_area.outputs:  # If no error printed by get_deployments\n",
+        "            status_message = (\n",
+        "                f\"No deployments found in namespace **{target_namespace}**.\"\n",
+        "            )\n",
+        "        else:\n",
+        "            status_message = None  # Error likely already shown\n",
+        "\n",
+        "    # Update output area with final status\n",
+        "    with output_area:\n",
+        "        clear_output(wait=True)\n",
+        "        display(Markdown(f\"Cluster: **{cluster_name}**\"))\n",
+        "        if namespace_to_use:\n",
+        "            display(Markdown(f\"Namespace: **{namespace_to_use}**\"))\n",
+        "        if status_message:\n",
+        "            display(Markdown(status_message))\n",
+        "\n",
+        "\n",
+        "def update_namespace_dropdown(cluster_name):\n",
+        "    \"\"\"Updates the namespace list based on cluster change.\"\"\"\n",
+        "    global namespace_dropdown, SELECTED_NAMESPACE\n",
+        "    global deployment_dropdown, SELECTED_DEPLOYMENT  # Need to reset deployment too\n",
+        "\n",
+        "    # Reset namespace state and dependent deployment dropdown\n",
+        "    SELECTED_NAMESPACE = None  # Reset selection\n",
+        "    SELECTED_DEPLOYMENT = None\n",
+        "    namespace_dropdown.disabled = True\n",
+        "    namespace_dropdown.options = [\"Loading...\"]\n",
+        "    namespace_dropdown.value = \"Loading...\"\n",
+        "    deployment_dropdown.options = [\"Select Deployment\"]\n",
+        "    deployment_dropdown.value = \"Select Deployment\"\n",
+        "    deployment_dropdown.disabled = True\n",
+        "\n",
+        "    # Clear output area and show loading context\n",
+        "    with output_area:\n",
+        "        clear_output(wait=True)\n",
+        "        display(Markdown(f\"Cluster: **{cluster_name}**\"))\n",
+        "        display(Markdown(\"Fetching namespaces...\"))\n",
+        "\n",
+        "    # Fetch namespaces (assuming CLUSTER_REGION_MAP and PROJECT_ID are available)\n",
+        "    region = CLUSTER_REGION_MAP.get(cluster_name)\n",
+        "    if not region:\n",
+        "        with output_area:\n",
+        "            clear_output(wait=True)\n",
+        "            display(\n",
+        "                Markdown(\n",
+        "                    f\"<font color='red'>Error: Region not found for cluster {cluster_name}.</font>\"\n",
+        "                )\n",
+        "            )\n",
+        "        namespace_dropdown.options = [\"Error loading\"]\n",
+        "        namespace_dropdown.value = \"Error loading\"\n",
+        "        return  # Stop if region is missing\n",
+        "\n",
+        "    # Assuming PROJECT_ID is globally available\n",
+        "    namespaces = get_namespaces(cluster_name, region, PROJECT_ID)\n",
+        "\n",
+        "    # Update dropdown options based on fetch result\n",
+        "    if namespaces is not None:  # Success (get_namespaces returns None on error)\n",
+        "        new_options = [\"Select Namespace\"] + namespaces  # Use \"Select Namespace\" prompt\n",
+        "        namespace_dropdown.options = new_options\n",
+        "        namespace_dropdown.value = \"Select Namespace\"\n",
+        "        namespace_dropdown.disabled = False\n",
+        "        status_message = (\n",
+        "            f\"Found {len(namespaces)} namespace(s). Select one to list deployments.\"\n",
+        "        )\n",
+        "    else:  # Error occurred during fetch\n",
+        "        namespace_dropdown.options = [\"Error loading\"]  # Keep error state\n",
+        "        namespace_dropdown.value = \"Error loading\"\n",
+        "        namespace_dropdown.disabled = True\n",
+        "        status_message = None  # Error already displayed by get_namespaces\n",
+        "\n",
+        "    # Update output area with final status\n",
+        "    with output_area:\n",
+        "        clear_output(wait=True)\n",
+        "        display(Markdown(f\"Cluster: **{cluster_name}**\"))\n",
+        "        if status_message:\n",
+        "            display(Markdown(status_message))\n",
+        "\n",
+        "\n",
+        "def on_cluster_change(change):\n",
+        "    \"\"\"Handles cluster selection changes.\"\"\"\n",
+        "    # Globals not strictly needed here as it calls update_namespace_dropdown which uses them\n",
+        "    if change[\"type\"] == \"change\" and change[\"name\"] == \"value\":\n",
+        "        cluster = change[\"new\"]\n",
+        "\n",
+        "        # Clear output area for new selection process\n",
+        "        with output_area:\n",
+        "            clear_output(wait=True)\n",
+        "\n",
+        "        if cluster == \"Select Cluster\":\n",
+        "            # Reset namespace dropdown\n",
+        "            namespace_dropdown.options = [\"Select Namespace\"]  # Correct prompt\n",
+        "            namespace_dropdown.value = \"Select Namespace\"\n",
+        "            namespace_dropdown.disabled = True\n",
+        "            # Reset deployment dropdown\n",
+        "            deployment_dropdown.options = [\"Select Deployment\"]\n",
+        "            deployment_dropdown.value = \"Select Deployment\"\n",
+        "            deployment_dropdown.disabled = True\n",
+        "            # Clear globals\n",
+        "            global SELECTED_NAMESPACE, SELECTED_DEPLOYMENT\n",
+        "            SELECTED_NAMESPACE = None\n",
+        "            SELECTED_DEPLOYMENT = None\n",
+        "        else:\n",
+        "            # Trigger update for the namespace dropdown\n",
+        "            update_namespace_dropdown(cluster)\n",
+        "\n",
+        "\n",
+        "def on_namespace_change(change):\n",
+        "    \"\"\"Handles namespace selection: fetches deployments.\"\"\"\n",
+        "    global SELECTED_NAMESPACE, cluster_dropdown, deployment_dropdown  # Added deployment_dropdown\n",
+        "    if change[\"type\"] == \"change\" and change[\"name\"] == \"value\":\n",
+        "        new_namespace = change[\"new\"]\n",
+        "\n",
+        "        # Get current cluster value\n",
+        "        current_cluster = cluster_dropdown.value\n",
+        "\n",
+        "        # Handle placeholder/loading/error values or if cluster isn't selected\n",
+        "        if (\n",
+        "            new_namespace in [\"Select Namespace\", \"Loading...\", \"Error loading\"]\n",
+        "            or current_cluster == \"Select Cluster\"\n",
+        "        ):\n",
+        "            SELECTED_NAMESPACE = None\n",
+        "            # Reset deployment dropdown state\n",
+        "            deployment_dropdown.options = [\"Select Deployment\"]\n",
+        "            deployment_dropdown.value = \"Select Deployment\"\n",
+        "            deployment_dropdown.disabled = True\n",
+        "            global SELECTED_DEPLOYMENT\n",
+        "            SELECTED_DEPLOYMENT = None\n",
+        "            # Clear output area for clean state\n",
+        "            with output_area:\n",
+        "                clear_output(wait=True)\n",
+        "                if current_cluster != \"Select Cluster\":  # Keep cluster context\n",
+        "                    display(Markdown(f\"Cluster: **{current_cluster}**\"))\n",
+        "                if new_namespace == \"Select Namespace\":\n",
+        "                    display(Markdown(\"Select a namespace to list deployments.\"))\n",
+        "            return  # Don't proceed to fetch deployments\n",
+        "\n",
+        "        # Valid namespace selected\n",
+        "        SELECTED_NAMESPACE = new_namespace\n",
+        "\n",
+        "        # Trigger update for the deployment dropdown\n",
+        "        if current_cluster != \"Select Cluster\":\n",
+        "            update_deployment_dropdown(current_cluster, SELECTED_NAMESPACE)\n",
+        "\n",
+        "\n",
+        "# --- Main Widget Setup ---\n",
+        "if CLUSTER_REGION_MAP:\n",
+        "    clusters_with_prompt = [\"Select Cluster\"] + sorted(list(CLUSTER_REGION_MAP.keys()))\n",
+        "    cluster_dropdown = widgets.Dropdown(\n",
+        "        options=clusters_with_prompt,\n",
+        "        value=\"Select Cluster\",  # Set initial value\n",
+        "        description=\"Cluster:\",\n",
+        "        style={\"description_width\": \"initial\"},\n",
+        "        layout=widgets.Layout(width=\"auto\"),  # Auto width\n",
+        "    )\n",
+        "\n",
+        "    namespace_dropdown = widgets.Dropdown(\n",
+        "        options=[\"Select Namespace\"],  # Correct initial prompt\n",
+        "        value=\"Select Namespace\",\n",
+        "        description=\"Namespace:\",\n",
+        "        disabled=True,  # Initially disabled\n",
+        "        style={\"description_width\": \"initial\"},\n",
+        "        layout=widgets.Layout(width=\"auto\"),\n",
+        "    )\n",
+        "\n",
+        "    deployment_dropdown = widgets.Dropdown(\n",
+        "        options=[\"Select Deployment\"],\n",
+        "        value=\"Select Deployment\",\n",
+        "        description=\"Deployment:\",\n",
+        "        disabled=True,  # Initially disabled\n",
+        "        style={\"description_width\": \"initial\"},\n",
+        "        layout=widgets.Layout(width=\"auto\"),\n",
+        "    )\n",
+        "\n",
+        "    # Observe changes\n",
+        "    cluster_dropdown.observe(on_cluster_change, names=\"value\")\n",
+        "    namespace_dropdown.observe(on_namespace_change, names=\"value\")\n",
+        "    deployment_dropdown.observe(on_deployment_select, names=\"value\")\n",
+        "\n",
+        "    # Display initial status and widgets\n",
+        "    print(\n",
+        "        f\"Found {len(CLUSTER_REGION_MAP)} Autopilot Cluster(s) in Project '{PROJECT_ID}'.\\n\"\n",
+        "    )\n",
+        "    display(cluster_dropdown, namespace_dropdown, deployment_dropdown, output_area)\n",
+        "\n",
+        "else:\n",
+        "    # Handle case where PROJECT_ID might be missing or no clusters found\n",
+        "    if \"PROJECT_ID\" not in globals() or not PROJECT_ID:\n",
+        "        error_message = \"Error: PROJECT_ID variable is not defined or empty. Please define it in a previous cell.\"\n",
+        "    else:\n",
+        "        error_message = f\"Error: No Autopilot clusters found or accessible in project '{PROJECT_ID}'. Check Project ID, permissions, and ensure Autopilot clusters exist.\"\n",
+        "    print(error_message)\n",
+        "    # Display error message using a widget for better integration in notebook\n",
+        "    display(widgets.HTML(f\"<font color='red'>{error_message}</font>\"))\n",
+        "    # Keep output_area widget displayed even on error for potential messages from retries etc.\n",
+        "    display(output_area)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "IKGTaN84p8rX",
+      "metadata": {
+        "cellView": "form",
+        "id": "IKGTaN84p8rX"
+      },
+      "outputs": [],
+      "source": [
+        "# @title # Chat completion for text-only models { vertical-output: true}\n",
+        "# @markdown You may send prompts to the model server for prediction.\n",
+        "# @markdown\n",
+        "# @markdown * **user_prompt (string):** This is the text prompt you provide to the language model. It's the question or instruction e (e.g., \"Explain neural networks\").\n",
+        "# @markdown * **temperature (number):** This  parameter controls the randomness of the model's output. It influences how the model selects the next token in the sequence it generates. Typical values range from 0.2 to 1.0.\n",
+        "# @markdown * **max_tokens (number):** This parameter refers to the maximum number of tokens (words or sub-word units) that the model is allowed to generate in its response.\n",
+        "\n",
+        "import ipywidgets as widgets\n",
+        "\n",
+        "\n",
+        "def _run_kubectl(cmd):\n",
+        "    \"\"\"Executes a kubectl command and returns its stdout.\"\"\"\n",
+        "    result = subprocess.run(cmd, capture_output=True, text=True, check=True, timeout=60)\n",
+        "    return result.stdout.strip()\n",
+        "\n",
+        "\n",
+        "def get_deployment_pod_name(deployment, namespace):\n",
+        "    \"\"\"Finds the running pod name for a given deployment and namespace.\"\"\"\n",
+        "    cmd = [\n",
+        "        \"kubectl\",\n",
+        "        \"get\",\n",
+        "        \"pods\",\n",
+        "        \"-n\",\n",
+        "        namespace,\n",
+        "        \"-o\",\n",
+        "        \"json\",\n",
+        "        \"-l\",\n",
+        "        f\"app={deployment}-app\",\n",
+        "        \"--field-selector=status.phase=Running\",\n",
+        "    ]\n",
+        "    try:\n",
+        "        pods_json = _run_kubectl(cmd)\n",
+        "        pods = json.loads(pods_json)\n",
+        "        if pods.get(\"items\"):\n",
+        "            return pods[\"items\"][0][\"metadata\"][\"name\"]\n",
+        "        print(f\"No running pods found for {deployment} in {namespace}.\")\n",
+        "        return None\n",
+        "    except (\n",
+        "        subprocess.CalledProcessError,\n",
+        "        json.JSONDecodeError,\n",
+        "        IndexError,\n",
+        "        KeyError,\n",
+        "    ) as e:\n",
+        "        print(f\"Error getting pod name for {deployment} in {namespace}: {e}\")\n",
+        "        return None\n",
+        "\n",
+        "\n",
+        "def check_inference_label(pod_name, namespace):\n",
+        "    \"\"\"Checks if the specified pod has the vLLM inference server label.\"\"\"\n",
+        "    cmd = [\"kubectl\", \"get\", \"pod\", pod_name, \"-n\", namespace, \"-o\", \"json\"]\n",
+        "    try:\n",
+        "        pod_json = _run_kubectl(cmd)\n",
+        "        labels = json.loads(pod_json).get(\"metadata\", {}).get(\"labels\", {})\n",
+        "        return labels.get(\"ai.gke.io/inference-server\") == \"vllm\"\n",
+        "    except (subprocess.CalledProcessError, json.JSONDecodeError, KeyError) as e:\n",
+        "        print(f\"Error checking labels for pod {pod_name} in {namespace}: {e}\")\n",
+        "        return False\n",
+        "\n",
+        "\n",
+        "def process_response(request, pod_name, pod_endpoint, is_vllm_inference, namespace):\n",
+        "    \"\"\"Sends a request to the pod and processes the response.\"\"\"\n",
+        "    json_data_escaped = json.dumps(request).replace(\"'\", \"'\\\\''\")\n",
+        "    curl_cmd = f\"kubectl exec -n {namespace} -t {pod_name} -- curl -s -X POST http://{pod_endpoint}/generate -H \\\"Content-Type: application/json\\\" -d '{json_data_escaped}' 2> /dev/null\"\n",
+        "    try:\n",
+        "        response_raw = _run_kubectl([\"bash\", \"-c\", curl_cmd])\n",
+        "        if not response_raw:\n",
+        "            return f\"Error: Empty response from pod {pod_name}.\"\n",
+        "        first_line = response_raw.splitlines()[0]\n",
+        "        data = json.loads(first_line)\n",
+        "\n",
+        "        if is_vllm_inference:\n",
+        "            predictions = data.get(\"predictions\")\n",
+        "            if isinstance(predictions, (list, tuple)) and predictions:\n",
+        "                return predictions[0]\n",
+        "            return f\"Error: Unexpected vLLM format. Raw: {first_line}\"\n",
+        "        else:  # TGI format\n",
+        "            generated_text = data.get(\"generated_text\")\n",
+        "            if generated_text is not None:\n",
+        "                return generated_text\n",
+        "            return f\"Error: Unexpected TGI format. Raw: {first_line}\"\n",
+        "\n",
+        "    except json.JSONDecodeError as e:\n",
+        "        raw_response = (\n",
+        "            response_raw.splitlines()[0]\n",
+        "            if \"response_raw\" in locals() and response_raw\n",
+        "            else \"N/A\"\n",
+        "        )\n",
+        "        return f\"Error decoding JSON: {e}. Raw: {raw_response}\"\n",
+        "    except (subprocess.CalledProcessError, IndexError, KeyError, TypeError) as e:\n",
+        "        raw_response = (\n",
+        "            response_raw.splitlines()[0]\n",
+        "            if \"response_raw\" in locals() and response_raw\n",
+        "            else \"N/A\"\n",
+        "        )\n",
+        "        return f\"Error processing response: {e}. Raw: {raw_response}\"\n",
+        "    except Exception as e:\n",
+        "        return f\"Unexpected error during response processing: {e}\"\n",
+        "\n",
+        "\n",
+        "# --- Widgets Setup ---\n",
+        "user_prompt_widget = widgets.Textarea(\n",
+        "    value=\"What is AI?\",\n",
+        "    description=\"User Prompt:\",\n",
+        "    layout=widgets.Layout(width=\"95%\", height=\"100px\"),\n",
+        ")\n",
+        "temperature_widget = widgets.FloatSlider(\n",
+        "    value=0.50, min=0.0, max=1.0, step=0.01, description=\"Temperature:\"\n",
+        ")\n",
+        "max_tokens_widget = widgets.IntSlider(\n",
+        "    value=250, min=1, max=2048, step=1, description=\"Max Tokens:\"\n",
+        ")\n",
+        "submit_button = widgets.Button(description=\"Submit\")\n",
+        "output_area_response = widgets.Output()\n",
+        "\n",
+        "\n",
+        "# --- Submit Button Logic ---\n",
+        "def on_submit_clicked(b):\n",
+        "    \"\"\"Handles the submit button click event.\"\"\"\n",
+        "    with output_area_response:\n",
+        "        clear_output()\n",
+        "        if (\n",
+        "            \"SELECTED_DEPLOYMENT\" not in globals()\n",
+        "            or \"SELECTED_NAMESPACE\" not in globals()\n",
+        "        ):\n",
+        "            display(\n",
+        "                Markdown(\n",
+        "                    \"**Error:** `SELECTED_DEPLOYMENT` or `SELECTED_NAMESPACE` not defined.\"\n",
+        "                )\n",
+        "            )\n",
+        "            return\n",
+        "\n",
+        "        print(\n",
+        "            f\"Target: {SELECTED_DEPLOYMENT} in {SELECTED_NAMESPACE}. \\n\\nRequesting response...\"\n",
+        "        )\n",
+        "\n",
+        "        pod_name = get_deployment_pod_name(SELECTED_DEPLOYMENT, SELECTED_NAMESPACE)\n",
+        "        if not pod_name:\n",
+        "            display(\n",
+        "                Markdown(\n",
+        "                    f\"**Error:** Could not find running pod for `{SELECTED_DEPLOYMENT}`.\"\n",
+        "                )\n",
+        "            )\n",
+        "            return\n",
+        "\n",
+        "        is_vllm = check_inference_label(pod_name, SELECTED_NAMESPACE)\n",
+        "        request = {\n",
+        "            \"max_tokens\": max_tokens_widget.value,\n",
+        "            \"temperature\": temperature_widget.value,\n",
+        "            \"prompt\" if is_vllm else \"inputs\": user_prompt_widget.value,\n",
+        "        }\n",
+        "        service = f\"{SELECTED_DEPLOYMENT}-service\"\n",
+        "        endpoint_cmd = [\n",
+        "            \"kubectl\",\n",
+        "            \"get\",\n",
+        "            \"endpoints\",\n",
+        "            service,\n",
+        "            \"-n\",\n",
+        "            SELECTED_NAMESPACE,\n",
+        "        ]\n",
+        "\n",
+        "        try:\n",
+        "            endpoint_output = _run_kubectl(endpoint_cmd).splitlines()\n",
+        "            if len(endpoint_output) < 2 or len(endpoint_output[1].split()) < 2:\n",
+        "                display(\n",
+        "                    Markdown(\n",
+        "                        f\"**Error:** Endpoint data incomplete for service `{service}`.\"\n",
+        "                    )\n",
+        "                )\n",
+        "                print(\"kubectl output:\\n\", \"\\n\".join(endpoint_output))\n",
+        "                return\n",
+        "            endpoint = endpoint_output[1].split()[\n",
+        "                1\n",
+        "            ]  # Assumes format: NAME ENDPOINTS AGE -> service ip:port,... age\n",
+        "            response = process_response(\n",
+        "                request, pod_name, endpoint, is_vllm, SELECTED_NAMESPACE\n",
+        "            )\n",
+        "            display(Markdown(f\"**Response:**\\n\\n{response}\"))\n",
+        "\n",
+        "        except subprocess.CalledProcessError as e:\n",
+        "            display(\n",
+        "                Markdown(\n",
+        "                    f\"**Error getting endpoints for `{service}`:**\\n```\\n{e.stderr}\\n```\"\n",
+        "                )\n",
+        "            )\n",
+        "        except Exception as e:\n",
+        "            display(Markdown(f\"**Unexpected Error:**\\n```\\n{e}\\n```\"))\n",
+        "\n",
+        "\n",
+        "# --- Display Widgets ---\n",
+        "submit_button.on_click(on_submit_clicked)\n",
+        "display(\n",
+        "    user_prompt_widget,\n",
+        "    temperature_widget,\n",
+        "    max_tokens_widget,\n",
+        "    submit_button,\n",
+        "    output_area_response,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "5b6ZM2K3fux0",
+      "metadata": {
+        "id": "5b6ZM2K3fux0"
+      },
+      "source": [
+        "# Next Steps: Integrating the GKE Service Endpoint\n",
+        "\n",
+        "After successfully deploying a model on Google Kubernetes Engine (GKE) and\n",
+        "verifying it via a notebook, the next step is to integrate it into various\n",
+        "applications. This involves making HTTP requests to the service's endpoint from\n",
+        "your application code.\n",
+        "\n",
+        "### Exposing the Service\n",
+        "\n",
+        "To make your deployed model accessible to applications, you'll need to expose\n",
+        "its service endpoint. Google Kubernetes Engine offers several ways to do this:\n",
+        "\n",
+        "1.  **Ingress:** Configure an Ingress resource to route external HTTP(S) traffic\n",
+        "    to your service. Set up Ingress for either an internal Load Balancer\n",
+        "    (accessible only within your VPC) or an external Load Balancer (accessible\n",
+        "    from the internet).\n",
+        "    [Learn more about GKE Ingress](https://cloud.google.com/kubernetes-engine/docs/concepts/ingress).\n",
+        "2.  **Gateway API:** A more modern and feature-rich API for managing traffic\n",
+        "    routing in Kubernetes. Similar to Ingress, Gateway API allows you to define\n",
+        "    how external and internal traffic should be directed to your services.\n",
+        "    [Explore GKE Gateway API](https://cloud.google.com/kubernetes-engine/docs/concepts/gateway-api).\n",
+        "\n",
+        "### Setting Up Autoscaling\n",
+        "\n",
+        "Ensure your model serving can handle varying traffic by configuring the\n",
+        "Horizontal Pod Autoscaler (HPA). HPA automatically scales the number of Pods\n",
+        "based on resource utilization or custom metrics, optimizing performance and\n",
+        "cost.\n",
+        "[See how to configure HPA](https://cloud.google.com/kubernetes-engine/docs/how-to/horizontal-pod-autoscaling).\n",
+        "\n",
+        "### Setting Up Monitoring\n",
+        "\n",
+        "Monitor the health and performance of your deployed model using Google Cloud\n",
+        "Managed Service for Prometheus. Configure your model serving to expose\n",
+        "Prometheus metrics for comprehensive insights.\n",
+        "[Get started with Google Cloud Managed Prometheus](https://cloud.google.com/kubernetes-engine/docs/how-to/configure-automatic-application-monitoring).\n",
+        "\n",
+        "### Additional Resources:\n",
+        "\n",
+        "*   #### Kubernetes Documentation:\n",
+        "\n",
+        "    *   Services:\n",
+        "        https://kubernetes.io/docs/concepts/services-networking/service/\n",
+        "\n",
+        "*   #### Google Cloud Documentation:\n",
+        "\n",
+        "    *   Google Kubernetes Engine (GKE):\n",
+        "        https://cloud.google.com/kubernetes-engine\n",
+        "    *   Cloud Load Balancing:\n",
+        "        https://cloud.google.com/load-balancing/docs/ingress\n",
+        "    *   Gateway API on GKE:\n",
+        "        https://cloud.google.com/kubernetes-engine/docs/concepts/gateway-api\n",
+        "    *   Learn about GPUs in GKE:\n",
+        "        https://cloud.google.com/kubernetes-engine/docs/concepts/gpus\n",
+        "\n",
+        "*   #### Python requests Library:\n",
+        "\n",
+        "    *   https://requests.readthedocs.io/en/latest/\n",
+        "\n",
+        "*   #### LangChain with Google Integrations:\n",
+        "\n",
+        "    *   The Langchain documentation is very useful:\n",
+        "        https://python.langchain.com/docs/integrations/providers/google/"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "gke_model_ui_deployment_notebook.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/gke_model_ui_deployment_notebook_auto.ipynb

@@ -0,0 +1,481 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "Pr9TgOcV9vAXeqGiyTaTI5kS"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2025 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": "M1CpgYundFwz"
+      },
+      "source": [
+        "# Get started with your deployed model on GKE\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%2Fgke_model_ui_deployment_notebook_auto.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/gke_model_ui_deployment_notebook_auto.ipynb\">\n",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "t2jj2XOgkS4F"
+      },
+      "source": [
+        "# Overview\n",
+        "\n",
+        "This notebook will guide you through the initial step of testing your recently\n",
+        "deployed model with text prompts. Depending on your deployed model's inference\n",
+        "setup, the notebook utilizes either Text Generation Inference\n",
+        "[TGI](https://huggingface.co/docs/text-generation-inference/en/index) or\n",
+        "[vLLM](https://developers.googleblog.com/en/inference-with-gemma-using-dataflow-and-vllm/#:~:text=model%20frameworks%20simple.-,What%20is%20vLLM%3F,-vLLM%20is%20an),\n",
+        "two efficient serving frameworks that enhance the performance of your GPU model.\n",
+        "Ready to see your deployed model respond? Run the cells below and start\n",
+        "experimenting with different prompts!\n",
+        "\n",
+        "### Prerequisites\n",
+        "\n",
+        "Before proceeding with this notebook, ensure you have already deployed a model\n",
+        "using the Google Cloud Console. You can find an overview of AI and Machine\n",
+        "Learning services on\n",
+        "[GKE AI/ML](https://console.cloud.google.com/kubernetes/aiml/overview).\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "Enable prompt-based testing of the AI model deployed on GKE\n",
+        "\n",
+        "### GPUs\n",
+        "\n",
+        "GPUs let you accelerate specific workloads running on your nodes, such as\n",
+        "machine learning and data processing. GKE provides a range of machine type\n",
+        "options for node configuration, including machine types with NVIDIA H100, L4,\n",
+        "and A100 GPUs.\n",
+        "\n",
+        "### Understanding the Inference Frameworks\n",
+        "\n",
+        "Your model is running on one of two popular and efficient serving frameworks:\n",
+        "vLLM or Text Generation Inference (TGI). The following sections provide a brief\n",
+        "overview of each to give you context on the underlying technology powering your\n",
+        "model.\n",
+        "\n",
+        "#### TGI\n",
+        "\n",
+        "TGI is a highly optimized open-source LLM serving framework that can increase\n",
+        "serving throughput on GPUs. TGI includes features such as:\n",
+        "\n",
+        "*   Optimized transformer implementation with PagedAttention\n",
+        "*   Continuous batching to improve the overall serving throughput\n",
+        "*   Tensor parallelism and distributed serving on multiple GPUs\n",
+        "\n",
+        "To learn more, refer to the\n",
+        "[TGI documentation](https://github.com/huggingface/text-generation-inference/blob/main/README.md)\n",
+        "\n",
+        "#### vLLM\n",
+        "\n",
+        "vLLM is another fast and easy-to-use library for LLM inference and serving. It's\n",
+        "known for its high throughput and efficiency, and it leverages PagedAttention.\n",
+        "Key features include:\n",
+        "\n",
+        "*   PagedAttention: Efficient memory management for handling long sequences and\n",
+        "    dynamic workloads.\n",
+        "*   Continuous batching: Maximizes GPU utilization by batching incoming\n",
+        "    requests.\n",
+        "*   High-throughput serving: Designed for production-level serving with low\n",
+        "    latency.\n",
+        "*   Optimized CUDA kernels.\n",
+        "\n",
+        "To learn more, refer to the\n",
+        "[vLLM documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/vllm/use-vllm)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "XMf-T58TkDy1"
+      },
+      "outputs": [],
+      "source": [
+        "# @title # Connect to Google Cloud Project\n",
+        "# @markdown #### Run this cell to configure your Google Cloud environment for Kubernetes (GKE) operations.\n",
+        "# @markdown\n",
+        "# @markdown #### Actions:\n",
+        "# @markdown 1.  **Connects to Project:** Retrieves and sets your Google Cloud project ID.\n",
+        "# @markdown 3.  **Installs `kubectl`:** Installs the Kubernetes command-line tool.\n",
+        "\n",
+        "import os\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Set up gcloud.\n",
+        "! gcloud config set project \"$PROJECT_ID\"\n",
+        "! gcloud services enable container.googleapis.com\n",
+        "\n",
+        "# Add kubectl to the set of available tools.\n",
+        "! mkdir -p /tools/google-cloud-sdk/.install\n",
+        "! gcloud components install kubectl --quiet"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "IKGTaN84p8rX"
+      },
+      "outputs": [],
+      "source": [
+        "# @title # Chat completion for text-only models {vertical-output: true}\n",
+        "# @markdown Run cell to prompt the model server for prediction.\n",
+        "# @markdown\n",
+        "# @markdown * **user_prompt (string):** This is the text prompt you provide to the language model. It's the question or instruction e (e.g., \"Explain neural networks\").\n",
+        "# @markdown * **temperature (number):** This  parameter controls the randomness of the model's output. It influences how the model selects the next token in the sequence it generates. Typical values range from 0.2 to 1.0.\n",
+        "# @markdown * **max_tokens (number):** This parameter refers to the maximum number of tokens (words or sub-word units) that the model is allowed to generate in its response.\n",
+        "# @markdown\n",
+        "\n",
+        "import json\n",
+        "import subprocess\n",
+        "\n",
+        "import ipywidgets as widgets\n",
+        "from IPython.display import Markdown, clear_output, display\n",
+        "\n",
+        "CLUSTER = \"\"  # @param {type:\"string\", isTemplate:true}\n",
+        "REGION = \"\"  # @param {type:\"string\", isTemplate:true}\n",
+        "NAMESPACE = \"\"  # @param {type:\"string\", isTemplate:true}\n",
+        "DEPLOYMENT = \"\"  # @param {type:\"string\", isTemplate:true}\n",
+        "POD_PORT = \"\"  # @param {type:\"string\", isTemplate:true}\n",
+        "\n",
+        "\n",
+        "def _run_kubectl(cmd, timeout=60):\n",
+        "    \"\"\"Executes a kubectl command.\"\"\"\n",
+        "    try:\n",
+        "        result = subprocess.run(\n",
+        "            cmd, capture_output=True, text=True, check=True, timeout=timeout\n",
+        "        )\n",
+        "        return result.stdout.strip()\n",
+        "    except subprocess.CalledProcessError as e:\n",
+        "        raise RuntimeError(\n",
+        "            f\"Kubectl command failed: {' '.join(e.cmd)}\\nStderr: {e.stderr}\"\n",
+        "        ) from e\n",
+        "    except subprocess.TimeoutExpired as e:\n",
+        "        raise RuntimeError(f\"Kubectl command timed out: {' '.join(e.cmd)}\") from e\n",
+        "\n",
+        "\n",
+        "def fetch_cluster_credentials(cluster, region, project_id):\n",
+        "    \"\"\"Ensures credentials for the target GKE cluster.\"\"\"\n",
+        "    cred_cmd = [\n",
+        "        \"gcloud\",\n",
+        "        \"container\",\n",
+        "        \"clusters\",\n",
+        "        \"get-credentials\",\n",
+        "        cluster,\n",
+        "        f\"--location={region}\",\n",
+        "        f\"--project={project_id}\",\n",
+        "    ]\n",
+        "    _run_kubectl(cred_cmd)\n",
+        "\n",
+        "\n",
+        "def get_deployment_selector_labels(deployment_name, namespace):\n",
+        "    \"\"\"Retrieves the selector labels for a given Kubernetes deployment.\"\"\"\n",
+        "    cmd = [\n",
+        "        \"kubectl\",\n",
+        "        \"get\",\n",
+        "        \"deployment\",\n",
+        "        deployment_name,\n",
+        "        \"-n\",\n",
+        "        namespace,\n",
+        "        \"-o\",\n",
+        "        \"json\",\n",
+        "    ]\n",
+        "    deployment_json = _run_kubectl(cmd)\n",
+        "    deployment_data = json.loads(deployment_json)\n",
+        "\n",
+        "    selector_labels = (\n",
+        "        deployment_data.get(\"spec\", {}).get(\"selector\", {}).get(\"matchLabels\")\n",
+        "    )\n",
+        "    if not selector_labels:\n",
+        "        raise RuntimeError(\n",
+        "            f\"No selector labels found for deployment '{deployment_name}' in\"\n",
+        "            f\" namespace '{namespace}'.\"\n",
+        "        )\n",
+        "    return selector_labels\n",
+        "\n",
+        "\n",
+        "def get_running_pod_name(deployment_name, namespace):\n",
+        "    \"\"\"Retrieves the name of a running pod associated with a deployment.\"\"\"\n",
+        "    selector_labels = get_deployment_selector_labels(deployment_name, namespace)\n",
+        "    label_selector_str = \",\".join(f\"{k}={v}\" for k, v in selector_labels.items())\n",
+        "\n",
+        "    cmd = [\n",
+        "        \"kubectl\",\n",
+        "        \"get\",\n",
+        "        \"pods\",\n",
+        "        \"-n\",\n",
+        "        namespace,\n",
+        "        \"-o\",\n",
+        "        \"json\",\n",
+        "        \"-l\",\n",
+        "        label_selector_str,\n",
+        "        \"--field-selector=status.phase=Running\",\n",
+        "    ]\n",
+        "    pods_json = _run_kubectl(cmd)\n",
+        "    pods_data = json.loads(pods_json)\n",
+        "\n",
+        "    if not pods_data.get(\"items\"):\n",
+        "        raise RuntimeError(\n",
+        "            f\"No running pods found for deployment '{deployment_name}' in namespace\"\n",
+        "            f\" '{namespace}' with selector '{label_selector_str}'.\"\n",
+        "        )\n",
+        "    return pods_data[\"items\"][0][\"metadata\"][\"name\"]\n",
+        "\n",
+        "\n",
+        "def check_vllm_inference_label(pod_name, namespace):\n",
+        "    \"\"\"Checks if the specified pod has the vLLM inference server label.\"\"\"\n",
+        "    cmd = [\"kubectl\", \"get\", \"pod\", pod_name, \"-n\", namespace, \"-o\", \"json\"]\n",
+        "    pod_json = _run_kubectl(cmd)\n",
+        "    labels = json.loads(pod_json).get(\"metadata\", {}).get(\"labels\", {})\n",
+        "    return labels.get(\"ai.gke.io/inference-server\") == \"vllm\"\n",
+        "\n",
+        "\n",
+        "def send_inference_request(\n",
+        "    request_payload, pod_name, pod_port, is_vllm_inference, namespace\n",
+        "):\n",
+        "    \"\"\"Sends an inference request to the specified pod and returns the model's response.\"\"\"\n",
+        "    json_data_escaped = json.dumps(request_payload).replace(\"'\", \"'\\\\''\")\n",
+        "    curl_cmd = (\n",
+        "        f\"kubectl exec -n {namespace} -t {pod_name} -- curl -s -X POST\"\n",
+        "        f' http://localhost:{pod_port}/generate -H \"Content-Type:'\n",
+        "        ' application/json\"'\n",
+        "        f\" -d '{json_data_escaped}' 2\u003e /dev/null\"\n",
+        "    )\n",
+        "\n",
+        "    response_raw = _run_kubectl([\"bash\", \"-c\", curl_cmd])\n",
+        "\n",
+        "    if not response_raw:\n",
+        "        raise RuntimeError(f\"Empty response received from pod '{pod_name}'.\")\n",
+        "\n",
+        "    try:\n",
+        "        first_line = response_raw.splitlines()[0]\n",
+        "        data = json.loads(first_line)\n",
+        "    except json.JSONDecodeError as e:\n",
+        "        raise RuntimeError(\n",
+        "            f\"Failed to decode JSON response from pod: {e}. Raw: {response_raw}\"\n",
+        "        ) from e\n",
+        "    except IndexError:\n",
+        "        raise RuntimeError(\n",
+        "            f\"Unexpected empty response line from pod. Raw: {response_raw}\"\n",
+        "        )\n",
+        "\n",
+        "    if is_vllm_inference:\n",
+        "        predictions = data.get(\"predictions\")\n",
+        "        if isinstance(predictions, list) and predictions:\n",
+        "            return predictions[0]\n",
+        "        raise RuntimeError(f\"Unexpected vLLM response format. Raw data: {data}\")\n",
+        "    else:  # TGI format\n",
+        "        generated_text = data.get(\"generated_text\")\n",
+        "        if generated_text is not None:\n",
+        "            return generated_text\n",
+        "        raise RuntimeError(f\"Unexpected TGI response format. Raw data: {data}\")\n",
+        "\n",
+        "\n",
+        "# --- Main Execution Logic ---\n",
+        "\n",
+        "\n",
+        "def execute_chat_completion(\n",
+        "    deployment_name, namespace, pod_port, user_prompt, temperature, max_tokens\n",
+        "):\n",
+        "    \"\"\"Executes the full chat completion process: fetches credentials, finds a pod,\n",
+        "\n",
+        "    determines inference type, sends a request, and returns the response.\n",
+        "    \"\"\"\n",
+        "    display(Markdown(\"Establishing cluster credentials...\"))\n",
+        "    fetch_cluster_credentials(CLUSTER, REGION, PROJECT_ID)\n",
+        "\n",
+        "    display(Markdown(\"Retrieving pod information...\"))\n",
+        "    pod_name = get_running_pod_name(deployment_name, namespace)\n",
+        "    display(Markdown(f\"Successfully identified pod: `{pod_name}`\"))\n",
+        "\n",
+        "    is_vllm = check_vllm_inference_label(pod_name, namespace)\n",
+        "\n",
+        "    request_payload = {\n",
+        "        \"max_tokens\": max_tokens,\n",
+        "        \"temperature\": temperature,\n",
+        "        \"prompt\" if is_vllm else \"inputs\": user_prompt,\n",
+        "    }\n",
+        "    display(Markdown(\"Sending inference request...\"))\n",
+        "    response = send_inference_request(\n",
+        "        request_payload, pod_name, pod_port, is_vllm, namespace\n",
+        "    )\n",
+        "\n",
+        "    return response\n",
+        "\n",
+        "\n",
+        "# --- Widgets Setup ---\n",
+        "user_prompt_widget = widgets.Textarea(\n",
+        "    value=\"What is AI?\",\n",
+        "    description=\"User Prompt:\",\n",
+        "    layout=widgets.Layout(width=\"95%\", height=\"100px\"),\n",
+        ")\n",
+        "\n",
+        "temperature_widget = widgets.FloatSlider(\n",
+        "    value=0.50, min=0.0, max=1.0, step=0.01, description=\"Temperature:\"\n",
+        ")\n",
+        "\n",
+        "max_tokens_widget = widgets.IntSlider(\n",
+        "    value=250, min=1, max=2048, step=1, description=\"Max Tokens:\"\n",
+        ")\n",
+        "\n",
+        "submit_button = widgets.Button(description=\"Submit\")\n",
+        "output_area_response = widgets.Output()\n",
+        "\n",
+        "\n",
+        "# --- Submit Button Logic ---\n",
+        "def on_submit_clicked(b):\n",
+        "    with output_area_response:\n",
+        "        clear_output()\n",
+        "        display(Markdown(\"Loading...\"))\n",
+        "\n",
+        "        try:\n",
+        "            model_response = execute_chat_completion(\n",
+        "                DEPLOYMENT,\n",
+        "                NAMESPACE,\n",
+        "                POD_PORT,\n",
+        "                user_prompt_widget.value,\n",
+        "                temperature_widget.value,\n",
+        "                max_tokens_widget.value,\n",
+        "            )\n",
+        "            clear_output()\n",
+        "            display(Markdown(f\"**Response:**\\n\\n{model_response}\"))\n",
+        "        except Exception as e:\n",
+        "            clear_output()\n",
+        "            display(Markdown(f\"**An error occurred:**\\n```\\n{e}\\n```\"))\n",
+        "\n",
+        "\n",
+        "# --- Display Widgets ---\n",
+        "submit_button.on_click(on_submit_clicked)\n",
+        "display(\n",
+        "    user_prompt_widget,\n",
+        "    temperature_widget,\n",
+        "    max_tokens_widget,\n",
+        "    submit_button,\n",
+        "    output_area_response,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5b6ZM2K3fux0"
+      },
+      "source": [
+        "# Next Steps: Integrating the GKE Service Endpoint\n",
+        "\n",
+        "After successfully deploying a model on Google Kubernetes Engine (GKE) and\n",
+        "verifying it via a notebook, the next step is to integrate it into various\n",
+        "applications. This involves making HTTP requests to the service's endpoint from\n",
+        "your application code.\n",
+        "\n",
+        "### Exposing the Service\n",
+        "\n",
+        "To make your deployed model accessible to applications, you'll need to expose\n",
+        "its service endpoint. Google Kubernetes Engine offers several ways to do this:\n",
+        "\n",
+        "1.  **Ingress:** Configure an Ingress resource to route external HTTP(S) traffic\n",
+        "    to your service. Set up Ingress for either an internal Load Balancer\n",
+        "    (accessible only within your VPC) or an external Load Balancer (accessible\n",
+        "    from the internet).\n",
+        "    [Learn more about GKE Ingress](https://cloud.google.com/kubernetes-engine/docs/concepts/ingress).\n",
+        "2.  **Gateway API:** A more modern and feature-rich API for managing traffic\n",
+        "    routing in Kubernetes. Similar to Ingress, Gateway API allows you to define\n",
+        "    how external and internal traffic should be directed to your services.\n",
+        "    [Explore GKE Gateway API](https://cloud.google.com/kubernetes-engine/docs/concepts/gateway-api).\n",
+        "\n",
+        "### Setting Up Autoscaling\n",
+        "\n",
+        "Ensure your model serving can handle varying traffic by configuring the\n",
+        "Horizontal Pod Autoscaler (HPA). HPA automatically scales the number of Pods\n",
+        "based on resource utilization or custom metrics, optimizing performance and\n",
+        "cost.\n",
+        "[See how to configure HPA](https://cloud.google.com/kubernetes-engine/docs/how-to/horizontal-pod-autoscaling).\n",
+        "\n",
+        "### Setting Up Monitoring\n",
+        "\n",
+        "Monitor the health and performance of your deployed model using Google Cloud\n",
+        "Managed Service for Prometheus. Configure your model serving to expose\n",
+        "Prometheus metrics for comprehensive insights.\n",
+        "[Get started with Google Cloud Managed Prometheus](https://cloud.google.com/kubernetes-engine/docs/how-to/configure-automatic-application-monitoring).\n",
+        "\n",
+        "### Additional Resources:\n",
+        "\n",
+        "*   #### Kubernetes Documentation:\n",
+        "\n",
+        "    *   Services:\n",
+        "        https://kubernetes.io/docs/concepts/services-networking/service/\n",
+        "\n",
+        "*   #### Google Cloud Documentation:\n",
+        "\n",
+        "    *   Google Kubernetes Engine (GKE):\n",
+        "        https://cloud.google.com/kubernetes-engine\n",
+        "    *   Cloud Load Balancing:\n",
+        "        https://cloud.google.com/load-balancing/docs/ingress\n",
+        "    *   Gateway API on GKE:\n",
+        "        https://cloud.google.com/kubernetes-engine/docs/concepts/gateway-api\n",
+        "    *   Learn about GPUs in GKE:\n",
+        "        https://cloud.google.com/kubernetes-engine/docs/concepts/gpus\n",
+        "\n",
+        "*   #### Python requests Library:\n",
+        "\n",
+        "    *   https://requests.readthedocs.io/en/latest/\n",
+        "\n",
+        "*   #### LangChain with Google Integrations:\n",
+        "\n",
+        "    *   The Langchain documentation is very useful:\n",
+        "        https://python.langchain.com/docs/integrations/providers/google/"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "gke_model_ui_deployment_notebook_auto.ipynb",
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_advanced_features.ipynb

@@ -6,7 +6,7 @@
       "id": "DZ1j6RRg-Td6",
       "metadata": {
         "cellView": "form",
-        "id": "DZ1j6RRg-Td6"
+        "id": "f705f4be70e9"
       },
       "outputs": [],
       "source": [
@@ -29,7 +29,7 @@
       "cell_type": "markdown",
       "id": "99c1c3fc2ca5",
       "metadata": {
-        "id": "99c1c3fc2ca5"
+        "id": "778cc1227be8"
       },
       "source": [
         "# Vertex AI Model Garden - Advanced Features\n",
@@ -42,7 +42,7 @@
         "  </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_advanced_features.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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"
@@ -52,7 +52,7 @@
       "cell_type": "markdown",
       "id": "f9-tJ6RfDLIs",
       "metadata": {
-        "id": "f9-tJ6RfDLIs"
+        "id": "0779b48f654e"
       },
       "source": [
         "## Overview\n",
@@ -90,7 +90,7 @@
       "cell_type": "markdown",
       "id": "47GcOrZjosOx",
       "metadata": {
-        "id": "47GcOrZjosOx"
+        "id": "69453bf7230e"
       },
       "source": [
         "## Before you begin"
@@ -100,7 +100,7 @@
       "cell_type": "markdown",
       "id": "1D_pWejJPHP3",
       "metadata": {
-        "id": "1D_pWejJPHP3"
+        "id": "bf3706e69f61"
       },
       "source": [
         "### Request for quota\n",
@@ -114,26 +114,18 @@
         "The quota for A100_80GB deployment `Custom model serving Nvidia A100 80GB GPUs per region` is 0. You need to request at least 4 for 70B model and 1 for 8B model following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota)."
       ]
     },
-    {
-      "cell_type": "markdown",
-      "id": "bdbcf9e7",
-      "metadata": {
-        "id": "jAELn17_1GfV"
-      },
-      "source": [
-        "### Setup Google Cloud project"
-      ]
-    },
     {
       "cell_type": "code",
       "execution_count": null,
       "id": "L3dqbxovo5t6",
       "metadata": {
         "cellView": "form",
-        "id": "L3dqbxovo5t6"
+        "id": "2b585189a670"
       },
       "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. **[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",
@@ -144,18 +136,20 @@
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
-        "# @markdown 4. 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",
+        "# @markdown 4. 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). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\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-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# Install and import the necessary packages\n",
         "! pip install -q openai google-auth requests\n",
         "\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.93.1'\n",
+        "\n",
         "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
         "\n",
         "import datetime\n",
@@ -231,7 +225,7 @@
       "cell_type": "markdown",
       "id": "SeGqxuMfRBS5",
       "metadata": {
-        "id": "SeGqxuMfRBS5"
+        "id": "4782dd003acb"
       },
       "source": [
         "### Access Llama 3.1, 3.2, and 3.3 models on Vertex AI for serving"
@@ -243,7 +237,7 @@
       "id": "BxlzWU2KQqmw",
       "metadata": {
         "cellView": "form",
-        "id": "BxlzWU2KQqmw"
+        "id": "798068fc0355"
       },
       "outputs": [],
       "source": [
@@ -282,7 +276,7 @@
       "cell_type": "markdown",
       "id": "JpNBJJgjWL7j",
       "metadata": {
-        "id": "JpNBJJgjWL7j"
+        "id": "10ed490e28e5"
       },
       "source": [
         "## Prefix Caching <a name=\"prefix-caching\"></a>\n",
@@ -310,7 +304,7 @@
       "cell_type": "markdown",
       "id": "9gZJ8cB27e1m",
       "metadata": {
-        "id": "9gZJ8cB27e1m"
+        "id": "30ddb93fdd7b"
       },
       "source": [
         "### Try out Prefix Caching with Hex-LLM\n",
@@ -326,7 +320,7 @@
       "id": "RpmoA2nXjdCd",
       "metadata": {
         "cellView": "form",
-        "id": "RpmoA2nXjdCd"
+        "id": "b56d82c1aa6f"
       },
       "outputs": [],
       "source": [
@@ -515,7 +509,7 @@
       "id": "5QoK8c0R9U3B",
       "metadata": {
         "cellView": "form",
-        "id": "5QoK8c0R9U3B"
+        "id": "96c5afed49b4"
       },
       "outputs": [],
       "source": [
@@ -526,9 +520,7 @@
         "\n",
         "if use_dedicated_endpoint:\n",
         "    DEDICATED_ENDPOINT_DNS = endpoints[\"hexllm_tpu\"].gca_resource.dedicated_endpoint_dns\n",
-        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
-        "    PROJECT_ID, REGION, endpoints[\"hexllm_tpu\"].name\n",
-        ")\n",
+        "ENDPOINT_RESOURCE_NAME = endpoints[\"hexllm_tpu\"].resource_name\n",
         "\n",
         "# @markdown Once deployment succeeds, you can send requests to the endpoint using the OpenAI SDK.\n",
         "\n",
@@ -583,7 +575,7 @@
       "id": "29rn5ATmB2YC",
       "metadata": {
         "cellView": "form",
-        "id": "29rn5ATmB2YC"
+        "id": "9a95c9f90358"
       },
       "outputs": [],
       "source": [
@@ -656,7 +648,7 @@
       "cell_type": "markdown",
       "id": "KjbM8E9DGuuR",
       "metadata": {
-        "id": "KjbM8E9DGuuR"
+        "id": "12ad6d1ff725"
       },
       "source": [
         "#### Delete the models and endpoints"
@@ -668,7 +660,7 @@
       "id": "JpLU7GRQGuuR",
       "metadata": {
         "cellView": "form",
-        "id": "JpLU7GRQGuuR"
+        "id": "1ab4e3bb74b4"
       },
       "outputs": [],
       "source": [
@@ -694,7 +686,7 @@
       "cell_type": "markdown",
       "id": "XZ33HhYmOxCS",
       "metadata": {
-        "id": "XZ33HhYmOxCS"
+        "id": "7a8a9a1b2ddf"
       },
       "source": [
         "### Try out Prefix Caching with vLLM\n",
@@ -717,7 +709,7 @@
       "id": "E8OiHHNNE_wj",
       "metadata": {
         "cellView": "form",
-        "id": "E8OiHHNNE_wj"
+        "id": "4425cc0bdedc"
       },
       "outputs": [],
       "source": [
@@ -921,7 +913,7 @@
       "id": "zex1oXl36A70",
       "metadata": {
         "cellView": "form",
-        "id": "zex1oXl36A70"
+        "id": "bcbafec839cd"
       },
       "outputs": [],
       "source": [
@@ -929,9 +921,7 @@
         "\n",
         "if use_dedicated_endpoint:\n",
         "    DEDICATED_ENDPOINT_DNS = endpoints[\"vllm_gpu\"].gca_resource.dedicated_endpoint_dns\n",
-        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
-        "    PROJECT_ID, REGION, endpoints[\"vllm_gpu\"].name\n",
-        ")\n",
+        "ENDPOINT_RESOURCE_NAME = endpoints[\"vllm_gpu\"].resource_name\n",
         "\n",
         "# @markdown Once deployment succeeds, you can send requests to the endpoint using the OpenAI SDK.\n",
         "\n",
@@ -986,7 +976,7 @@
       "id": "gDOC_nfsJeUR",
       "metadata": {
         "cellView": "form",
-        "id": "gDOC_nfsJeUR"
+        "id": "e984f43422d5"
       },
       "outputs": [],
       "source": [
@@ -1059,7 +1049,7 @@
       "cell_type": "markdown",
       "id": "GdGxaTirJeUR",
       "metadata": {
-        "id": "GdGxaTirJeUR"
+        "id": "dff0d10dcc20"
       },
       "source": [
         "#### Delete the models and endpoints"
@@ -1071,7 +1061,7 @@
       "id": "OgoqXE-VJeUR",
       "metadata": {
         "cellView": "form",
-        "id": "OgoqXE-VJeUR"
+        "id": "5b8751773e7f"
       },
       "outputs": [],
       "source": [
@@ -1097,7 +1087,7 @@
       "cell_type": "markdown",
       "id": "w4Guijaw_NEs",
       "metadata": {
-        "id": "w4Guijaw_NEs"
+        "id": "863775857a46"
       },
       "source": [
         "### Best practices\n",
@@ -1112,7 +1102,7 @@
       "cell_type": "markdown",
       "id": "ml8fgoIQWSbY",
       "metadata": {
-        "id": "ml8fgoIQWSbY"
+        "id": "565cbdc3a06b"
       },
       "source": [
         "## Speculative Decoding <a name=\"spec-decoding\"></a>\n",
@@ -1154,7 +1144,7 @@
       "cell_type": "markdown",
       "id": "NmWRro8Q-Td6",
       "metadata": {
-        "id": "NmWRro8Q-Td6"
+        "id": "94eaa9050abb"
       },
       "source": [
         "### Try out Speculative Decoding with vLLM"
@@ -1166,7 +1156,7 @@
       "id": "72d1GlrYifKU",
       "metadata": {
         "cellView": "form",
-        "id": "72d1GlrYifKU"
+        "id": "5f358cc230a6"
       },
       "outputs": [],
       "source": [
@@ -1483,7 +1473,7 @@
       "id": "CNiItf5hdVFU",
       "metadata": {
         "cellView": "form",
-        "id": "CNiItf5hdVFU"
+        "id": "be3170e0e05a"
       },
       "outputs": [],
       "source": [
@@ -1509,9 +1499,7 @@
         "    DEDICATED_ENDPOINT_DNS = endpoints[\n",
         "        \"vllm_gpu_spec\"\n",
         "    ].gca_resource.dedicated_endpoint_dns\n",
-        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
-        "    PROJECT_ID, REGION, endpoints[\"vllm_gpu_spec\"].name\n",
-        ")\n",
+        "ENDPOINT_RESOURCE_NAME = endpoints[\"vllm_gpu_spec\"].resource_name\n",
         "\n",
         "BASE_URL = (\n",
         "    f\"https://{REGION}-aiplatform.googleapis.com/v1beta1/{ENDPOINT_RESOURCE_NAME}\"\n",
@@ -1560,7 +1548,7 @@
       "cell_type": "markdown",
       "id": "WahYGAZyq6Gl",
       "metadata": {
-        "id": "WahYGAZyq6Gl"
+        "id": "30c5d2535df3"
       },
       "source": [
         "## Clean up resources"
@@ -1570,7 +1558,7 @@
       "cell_type": "markdown",
       "id": "bV5Yjkgav9BZ",
       "metadata": {
-        "id": "bV5Yjkgav9BZ"
+        "id": "63c10917ff95"
       },
       "source": [
         "### Delete the models and endpoints"
@@ -1582,7 +1570,7 @@
       "id": "qsks36cOH9rb",
       "metadata": {
         "cellView": "form",
-        "id": "qsks36cOH9rb"
+        "id": "92892e1b1730"
       },
       "outputs": [],
       "source": [

notebooks/community/model_garden/model_garden_camp_zipnerf.ipynb

@@ -38,7 +38,7 @@
         "  </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_camp_zipnerf.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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"

notebooks/community/model_garden/model_garden_camp_zipnerf_gradio.ipynb

@@ -38,7 +38,7 @@
         "  </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_camp_zipnerf_gradio.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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"

notebooks/community/model_garden/model_garden_codegemma_deployment_on_vertex.ipynb

@@ -40,7 +40,7 @@
         "  </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_codegemma_deployment_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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"
@@ -123,14 +123,22 @@
         ")\n",
         "\n",
         "models, endpoints = {}, {}\n",
-        "# Dedicated endpoint not supported yet\n",
-        "use_dedicated_endpoint = False\n",
+        "\n",
+        "# @markdown Set `use_dedicated_endpoint` to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
+        "\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",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -213,7 +221,7 @@
         "# @markdown *--- Or ---*\n",
         "\n",
         "# @markdown #### Access CodeGemma models on HuggingFace\n",
-        "# @markdown You must provide a Hugging Face User Access Token (read) to access the CodeGemma 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",
+        "# @markdown You must provide a Hugging Face User Access Token (with read access) to access the CodeGemma 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",
         "if LOAD_MODEL_FROM == \"Hugging Face\":\n",
         "    assert (\n",
@@ -322,7 +330,7 @@
         "    model_id: str,\n",
         "    publisher: str,\n",
         "    publisher_model_id: str,\n",
-        "    service_account: str,\n",
+        "    service_account: str = None,\n",
         "    base_model_id: str = None,\n",
         "    data_parallel_size: int = 1,\n",
         "    tensor_parallel_size: int = 1,\n",
@@ -331,6 +339,7 @@
         "    disagg_topology: str = None,\n",
         "    hbm_utilization_factor: float = 0.6,\n",
         "    max_running_seqs: int = 256,\n",
+        "    decode_seqs_padding: int = None,\n",
         "    max_model_len: int = 4096,\n",
         "    enable_prefix_cache_hbm: bool = False,\n",
         "    endpoint_id: str = \"\",\n",
@@ -371,6 +380,10 @@
         "        f\"--max_running_seqs={max_running_seqs}\",\n",
         "        f\"--max_model_len={max_model_len}\",\n",
         "    ]\n",
+        "\n",
+        "    if decode_seqs_padding is not None:\n",
+        "        hexllm_args.append(f\"--decode_seqs_padding={decode_seqs_padding}\")\n",
+        "\n",
         "    if disagg_topology:\n",
         "        hexllm_args.append(f\"--disagg_topo={disagg_topology}\")\n",
         "    if enable_prefix_cache_hbm and not disagg_topology:\n",
@@ -416,6 +429,7 @@
         "        max_replica_count=max_replica_count,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_codegemma_deployment_on_vertex.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    return model, endpoint\n",
@@ -586,6 +600,7 @@
         "    use_dedicated_endpoint: bool = False,\n",
         "    max_num_seqs: int = 256,\n",
         "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -638,6 +653,10 @@
         "    if model_type:\n",
         "        vllm_args.append(f\"--model-type={model_type}\")\n",
         "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -676,6 +695,7 @@
         "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_codegemma_deployment_on_vertex.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",

notebooks/community/model_garden/model_garden_deployment_tutorial.ipynb

@@ -0,0 +1,424 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "GyJPhQ2_Om3X"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2025 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": "L1BRnURwUb87"
+      },
+      "source": [
+        "# Vertex AI Model Garden - Deployment Tutorial\n",
+        "\n",
+        "<table><tbody><tr>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/instances\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_deployment_tutorial.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_deployment_tutorial.ipynb\">\n",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "K4W8U0MQUb87"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "You can deploy open models (including Hugging Face models) by using [Google Gen AI SDK or Google Cloud CLI](https://cloud.google.com/vertex-ai/generative-ai/docs/model-garden/use-models#deploy_a_model).\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": "x7Z25I3xUb87"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "WnEYdVEKUb87"
+      },
+      "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. **[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",
+        "# @markdown 3. 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). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\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-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
+        "\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform==1.93.1'\n",
+        "\n",
+        "# Import the necessary packages\n",
+        "import os\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
+        "    ! pip install --upgrade tensorflow\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\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",
+        "LABEL = \"my-endpoint\"\n",
+        "\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",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2NRInIKZ3nSJ"
+      },
+      "source": [
+        "## Deploy"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "xpIv3skP3nSK"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Choose the model to deploy\n",
+        "from vertexai import model_garden\n",
+        "\n",
+        "# @markdown List all deployable models and then get the ID of the model to deploy.\n",
+        "\n",
+        "# @markdown You can also use the Hugging Face model ID.\n",
+        "list_hf_models = False  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# @markdown You can also filter by model name.\n",
+        "model_filter = \"gemma\"  # @param {type:\"string\"}\n",
+        "\n",
+        "model_garden.list_deployable_models(\n",
+        "    list_hf_models=list_hf_models, model_filter=model_filter\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "uaLkN2mta5T2"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown Select a model from the list above.\n",
+        "MODEL_ID = \"google/gemma3@gemma-3-1b-it\"  # @param [\"google/gemma3@gemma-3-1b-it\", \"google/gemma-3-1b-it\"] {isTemplate:true}\n",
+        "\n",
+        "# @markdown 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}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "fR9wl0-Na5T2"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "Fw9W4TrOUb87"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy with Model Garden SDK\n",
+        "\n",
+        "# @markdown Deploy with Gen AI model-centric SDK. This section uploads the prebuilt model to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model. See [use open models with Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-open-models) for documentation on other use cases.\n",
+        "from vertexai import model_garden\n",
+        "\n",
+        "model = model_garden.OpenModel(MODEL_ID)\n",
+        "endpoints[LABEL] = model.deploy(\n",
+        "    hugging_face_access_token=HF_TOKEN,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    accept_eula=True,  # Accept the End User License Agreement (EULA) on the model card before deploy. Otherwise, the deployment will be forbidden.\n",
+        ")\n",
+        "\n",
+        "endpoint = endpoints[LABEL]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "rKjAOh7FUb87"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Raw predict\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. Sampling parameters supported by vLLM can be found [here](https://docs.vllm.ai/en/latest/dev/sampling_params.html).\n",
+        "\n",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown Human: What is a car?\n",
+        "# @markdown Assistant:  A car, or a motor car, is a road-connected human-transportation system used to move people or goods from one place to another. The term also encompasses a wide range of vehicles, including motorboats, trains, and aircrafts. Cars typically have four wheels, a cabin for passengers, and an engine or motor. They have been around since the early 19th century and are now one of the most popular forms of transportation, used for daily commuting, shopping, and other purposes.\n",
+        "# @markdown ```\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",
+        "prompt = \"What is a car?\"  # @param {type: \"string\"}\n",
+        "# @markdown If you encounter an issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, by lowering `max_tokens`.\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  # @param {type:\"integer\"}\n",
+        "# @markdown Set `raw_response` to `True` to obtain the raw model output. Set `raw_response` to `False` to apply additional formatting in the structure of `\"Prompt:\\n{prompt.strip()}\\nOutput:\\n{output}\"`.\n",
+        "raw_response = False  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# Overrides parameters for inferences.\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",
+        "        \"raw_response\": raw_response,\n",
+        "    },\n",
+        "]\n",
+        "response = endpoints[\"my-endpoint\"].predict(\n",
+        "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
+        ")\n",
+        "\n",
+        "for prediction in response.predictions:\n",
+        "    print(prediction)\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "wxe8GIwpUb87"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Chat completion\n",
+        "\n",
+        "if use_dedicated_endpoint:\n",
+        "    DEDICATED_ENDPOINT_DNS = endpoints[\n",
+        "        \"my-endpoint\"\n",
+        "    ].gca_resource.dedicated_endpoint_dns\n",
+        "ENDPOINT_RESOURCE_NAME = endpoints[\"my-endpoint\"].resource_name\n",
+        "\n",
+        "# @title Chat Completions Inference\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint using the OpenAI SDK.\n",
+        "\n",
+        "# @markdown First you will need to install the SDK and some auth-related dependencies.\n",
+        "\n",
+        "! pip install -qU openai google-auth requests\n",
+        "\n",
+        "# @markdown Next fill out some request parameters:\n",
+        "\n",
+        "user_message = \"How is your day going?\"  # @param {type: \"string\"}\n",
+        "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
+        "max_tokens = 50  # @param {type: \"integer\"}\n",
+        "temperature = 1.0  # @param {type: \"number\"}\n",
+        "stream = False  # @param {type: \"boolean\"}\n",
+        "\n",
+        "# @markdown Now we can send a request.\n",
+        "\n",
+        "import google.auth\n",
+        "import openai\n",
+        "\n",
+        "creds, project = google.auth.default()\n",
+        "auth_req = google.auth.transport.requests.Request()\n",
+        "creds.refresh(auth_req)\n",
+        "\n",
+        "BASE_URL = (\n",
+        "    f\"https://{REGION}-aiplatform.googleapis.com/v1beta1/{ENDPOINT_RESOURCE_NAME}\"\n",
+        ")\n",
+        "try:\n",
+        "    if use_dedicated_endpoint:\n",
+        "        BASE_URL = f\"https://{DEDICATED_ENDPOINT_DNS}/v1beta1/{ENDPOINT_RESOURCE_NAME}\"\n",
+        "except NameError:\n",
+        "    pass\n",
+        "\n",
+        "client = openai.OpenAI(base_url=BASE_URL, api_key=creds.token)\n",
+        "\n",
+        "model_response = client.chat.completions.create(\n",
+        "    model=\"\",\n",
+        "    messages=[{\"role\": \"user\", \"content\": user_message}],\n",
+        "    temperature=temperature,\n",
+        "    max_tokens=max_tokens,\n",
+        "    stream=stream,\n",
+        ")\n",
+        "\n",
+        "if stream:\n",
+        "    usage = None\n",
+        "    contents = []\n",
+        "    for chunk in model_response:\n",
+        "        if chunk.usage is not None:\n",
+        "            usage = chunk.usage\n",
+        "            continue\n",
+        "        print(chunk.choices[0].delta.content, end=\"\")\n",
+        "        contents.append(chunk.choices[0].delta.content)\n",
+        "    print(f\"\\n\\n{usage}\")\n",
+        "else:\n",
+        "    print(model_response)\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cqHoUEEMXNAH"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "5qKldg25Ub87"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Delete the models and endpoints\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()"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_deployment_tutorial.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_e5.ipynb

@@ -4,11 +4,12 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
+        "cellView": "form",
         "id": "7d9bbf86da5e"
       },
       "outputs": [],
       "source": [
-        "# Copyright 2024 Google LLC\n",
+        "# Copyright 2025 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",
@@ -33,13 +34,18 @@
         "\n",
         "<table><tbody><tr>\n",
         "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/instances\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_e5.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_e5.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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"
@@ -67,6 +73,10 @@
         "- Run inference on the deployed Vertex AI Endpoint\n",
         "\n",
         "\n",
+        "### File a bug\n",
+        "\n",
+        "File a bug on [GitHub](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues/new) if you encounter any issue with the notebook.\n",
+        "\n",
         "### Costs\n",
         "\n",
         "This tutorial uses billable components of Google Cloud:\n",
@@ -74,7 +84,7 @@
         "* 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), [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."
+        "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."
       ]
     },
     {
@@ -99,133 +109,61 @@
         "\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",
+        "# @markdown 2. **[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",
+        "# @markdown 3. 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). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\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-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
+        "\n",
+        "# Import the necessary packages\n",
+        "\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform==1.93.1'\n",
+        "! pip3 install --quiet torchvision\n",
+        "\n",
+        "import importlib\n",
         "import os\n",
-        "import uuid\n",
-        "from datetime import datetime\n",
         "from typing import Tuple\n",
         "\n",
         "from google.cloud import aiplatform\n",
+        "from torch import Tensor\n",
+        "\n",
+        "if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
+        "    ! pip install --upgrade tensorflow\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\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",
         "\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, change the value yourself below.\n",
-        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
-        "BUCKET_URI = \"gs://\"  # @param {type: \"string\"}\n",
-        "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
-        "\n",
-        "# @markdown Click \"Show code\" to see more details.\n",
-        "\n",
-        "# Create a unique GCS bucket for this notebook, if not specified by the user.\n",
-        "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\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",
-        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
-        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
-        "else:\n",
-        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\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",
-        "\n",
-        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
+        "if not REGION:\n",
+        "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
         "\n",
         "! gcloud config set project $PROJECT_ID\n",
-        "! gcloud services enable language.googleapis.com\n",
+        "import vertexai\n",
         "\n",
-        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
-        "\n",
-        "# Gets the default BUCKET_URI and SERVICE_ACCOUNT if they were not specified by the user.\n",
-        "\n",
-        "SERVICE_ACCOUNT = None\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",
-        "\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",
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
-        "\n",
-        "\n",
-        "def create_name_with_datetime(prefix: str) -> str:\n",
-        "    \"\"\"Creates a 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 deploy_model_tei(\n",
-        "    model_name: str,\n",
-        "    model_id: str,\n",
-        "    service_account: str,\n",
-        "    docker_uri: str,\n",
-        "    machine_type: str = \"g2-standard-8\",\n",
-        "    accelerator_type: str = \"NVIDIA_L4\",\n",
-        "    accelerator_count: int = 1,\n",
-        "    max_model_len: int = 512,\n",
-        "    gpu_memory_utilization: float = 0.9,\n",
-        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
-        "    \"\"\"Deploys E5 models with TEI on Vertex AI.\n",
-        "\n",
-        "    Args:\n",
-        "        model_name: Display name of the model.\n",
-        "        model_id: Model ID or path to model weights.\n",
-        "        service_account: Service account for model uploading and deployment.\n",
-        "        machine_type: Deployment machine type.\n",
-        "        accelerator_type: Deployment accelerator type.\n",
-        "        accelerator_count: Number of accelerators to use.\n",
-        "        max_model_len: Maximum model length.\n",
-        "        gpu_memory_utilization: Fraction of GPU memory to be used for the model\n",
-        "            executor.\n",
-        "\n",
-        "    Returns:\n",
-        "        Model instance and endpoint instance.\n",
-        "    \"\"\"\n",
-        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
-        "\n",
-        "    tei_args = [\n",
-        "        f\"--model-id={model_id}\",\n",
-        "    ]\n",
-        "    serving_env = {\n",
-        "        \"MODEL_ID\": model_id,\n",
-        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
-        "    }\n",
-        "    model = aiplatform.Model.upload(\n",
-        "        display_name=model_name,\n",
-        "        serving_container_image_uri=docker_uri,\n",
-        "        serving_container_args=tei_args,\n",
-        "        serving_container_ports=[7080],\n",
-        "        serving_container_environment_variables=serving_env,\n",
-        "        model_garden_source_model_name=\"publishers/intfloat/models/e5\"\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",
-        "        system_labels={\n",
-        "            \"NOTEBOOK_NAME\": \"model_garden_e5.ipynb\"\n",
-        "        },\n",
-        "    )\n",
-        "    return model, endpoint"
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")"
       ]
     },
     {
@@ -242,19 +180,25 @@
       "execution_count": null,
       "metadata": {
         "cellView": "form",
-        "id": "kg5MwMIfB9Uj"
+        "id": "I1u2FLa9XgVD"
       },
       "outputs": [],
       "source": [
-        "# @title Deploy\n",
+        "# @title Select the model variants\n",
         "# @markdown This section uploads a prebuilt model to Model Registry and deploys it on the Endpoint. The model deployment step will take ~15 minutes to complete.\n",
         "\n",
         "prebuilt_model_id = \"intfloat/e5-small-v2\"  # @param [\"intfloat/multilingual-e5-large-instruct\", \"intfloat/multilingual-e5-large\", \"intfloat/e5-large-v2\", \"intfloat/multilingual-e5-small\", \"intfloat/e5-base-v2\", \"intfloat/e5-small-v2\"]\n",
         "\n",
         "# @markdown Specify a processor for the TEI docker image. E5 models can be run on either GPU or CPU.\n",
-        "processor = \"NVIDIA_L4\"  # @param[\"NVIDIA_TESLA_V100\", \"NVIDIA_L4\", \"NVIDIA_TESLA_A100\", \"CPU\"]\n",
+        "# Find Vertex AI prediction supported accelerators and regions in\n",
+        "# https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
+        "processor = \"NVIDIA_L4\"  # @param[\"NVIDIA_TESLA_T4\", \"NVIDIA_TESLA_V100\", \"NVIDIA_L4\", \"NVIDIA_TESLA_A100\", \"CPU\"]\n",
         "\n",
-        "if processor == \"NVIDIA_TESLA_V100\":\n",
+        "if processor == \"NVIDIA_TESLA_T4\":\n",
+        "    accelerator_type = \"NVIDIA_TESLA_T4\"\n",
+        "    machine_type = \"n1-highmem-16\"\n",
+        "    accelerator_count = 1\n",
+        "elif processor == \"NVIDIA_TESLA_V100\":\n",
         "    accelerator_type = \"NVIDIA_TESLA_V100\"\n",
         "    machine_type = \"n1-highmem-16\"\n",
         "    accelerator_count = 2\n",
@@ -280,24 +224,108 @@
         "else:\n",
         "    TEI_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-embeddings-inference-cu122.1-2.ubuntu2204\"\n",
         "\n",
-        "# @markdown Click \"Show code\" to see more details.\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
         "\n",
-        "# Finds Vertex AI prediction supported accelerators and regions in\n",
-        "# https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
+        "# @markdown Click \"Show code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "6dY6_ppyObQy"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy model using custom configuration\n",
+        "# @markdown This section uploads prebuilt E5 models to Model Registry and deploys it to a Vertex AI Endpoint. It might take ~15 minutes to 1 hour to finish depending on the size of the model.\n",
         "\n",
-        "model, endpoint = deploy_model_tei(\n",
-        "    model_name=create_name_with_datetime(prefix=\"e5-serve-tei\"),\n",
+        "\n",
+        "def deploy_model_tei(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    docker_uri: str,\n",
+        "    machine_type: str = \"g2-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 1,\n",
+        "    max_model_len: int = 512,\n",
+        "    gpu_memory_utilization: float = 0.9,\n",
+        "    use_dedicated_endpoint: bool = True,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys E5 models with TEI on Vertex AI.\n",
+        "\n",
+        "    Args:\n",
+        "        model_name: Display name of the model.\n",
+        "        model_id: Model ID or path to model weights.\n",
+        "        machine_type: Deployment machine type.\n",
+        "        accelerator_type: Deployment accelerator type.\n",
+        "        accelerator_count: Number of accelerators to use.\n",
+        "        max_model_len: Maximum model length.\n",
+        "        gpu_memory_utilization: Fraction of GPU memory to be used for the model\n",
+        "            executor.\n",
+        "        use_dedicated_endpoint: A dedicated endpoint is an endpoint for online\n",
+        "            prediction,provide a secure connection for private communication\n",
+        "            between on-premises and Google Cloud.\n",
+        "\n",
+        "    Returns:\n",
+        "        Model instance and endpoint instance.\n",
+        "    \"\"\"\n",
+        "    endpoint = aiplatform.Endpoint.create(\n",
+        "        display_name=f\"{model_name}-endpoint\",\n",
+        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
+        "    )\n",
+        "\n",
+        "    tei_args = [\n",
+        "        f\"--model-id={model_id}\",\n",
+        "    ]\n",
+        "    serving_env = {\n",
+        "        \"MODEL_ID\": model_id,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=docker_uri,\n",
+        "        serving_container_args=tei_args,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_environment_variables=serving_env,\n",
+        "        model_garden_source_model_name=\"publishers/intfloat/models/e5\",\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",
+        "        system_labels={\"NOTEBOOK_NAME\": \"model_garden_e5.ipynb\"},\n",
+        "    )\n",
+        "    return model, endpoint\n",
+        "\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",
+        "\n",
+        "LABEL = \"tei\"\n",
+        "models[LABEL], endpoints[LABEL] = deploy_model_tei(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=\"e5-serve-tei\"),\n",
         "    model_id=prebuilt_model_id,\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
         "    docker_uri=TEI_DOCKER_URI,\n",
         "    machine_type=machine_type,\n",
         "    accelerator_type=accelerator_type,\n",
         "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
         ")\n",
         "\n",
-        "print(\"endpoint_name:\", endpoint.name)\n",
-        "print(\"model_name:\", model.display_name)\n",
-        "print(\"model_id:\", model.resource_name)"
+        "model = models[LABEL]\n",
+        "endpoint = endpoints[LABEL]"
       ]
     },
     {
@@ -355,7 +383,6 @@
         "# )\n",
         "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
         "\n",
-        "from torch import Tensor\n",
         "\n",
         "# Each input text should start with \"query: \" or \"passage: \".\n",
         "# For tasks other than retrieval, you can simply use the \"query: \" prefix.\n",
@@ -369,7 +396,9 @@
         "        ],\n",
         "    },\n",
         "]\n",
-        "response = endpoint.predict(instances=instances)\n",
+        "response = endpoint.predict(\n",
+        "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
+        ")\n",
         "\n",
         "for prediction in response.predictions:\n",
         "    embeddings = Tensor(prediction)\n",
@@ -397,9 +426,6 @@
         "# @markdown Instruct: Given a web search query, retrieve relevant passages that answer the query\n",
         "# @markdown Query: how much protein should a female eat\n",
         "# @markdown Instruct: Given a web search query, retrieve relevant passages that answer the query\n",
-        "# @markdown Query: 南瓜的家常做法\n",
-        "# @markdown As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.\n",
-        "# @markdown 1.清炒南瓜丝 原料:嫩南瓜半个 调料:葱、盐、白糖、鸡精 做法: 1、南瓜用刀薄薄的削去表面一层皮,用勺子刮去瓤 2、擦成细丝(没有擦菜板就用刀慢慢切成细丝) 3、锅烧热放油,入葱花煸出香味 4、入南瓜丝快速翻炒一分钟左右,放盐、一点白糖和鸡精调味出锅 2.香葱炒南瓜 原料:南瓜1只 调料:香葱、蒜末、橄榄油、盐 做法: 1、将南瓜去皮,切成片 2、油锅8成热后,将蒜末放入爆香 3、爆香后,将南瓜片放入,翻炒 4、在翻炒的同时,可以不时地往锅里加水,但不要太多 5、放入盐,炒匀 6、南瓜差不多软和绵了之后,就可以关火 7、撒入香葱,即可出锅\n",
         "# @markdown ```\n",
         "\n",
         "# @markdown API reference link to HuggingFace : [Text Embeddings Inference API](https://huggingface.github.io/text-embeddings-inference/#/).\n",
@@ -428,8 +454,6 @@
         "# )\n",
         "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
         "\n",
-        "from torch import Tensor\n",
-        "\n",
         "\n",
         "def get_detailed_instruct(task_description: str, query: str) -> str:\n",
         "    return f\"Instruct: {task_description}\\nQuery: {query}\"\n",
@@ -448,7 +472,9 @@
         "\n",
         "instances = [{\"inputs\": queries + documents}]\n",
         "\n",
-        "response = endpoint.predict(instances=instances)\n",
+        "response = endpoint.predict(\n",
+        "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
+        ")\n",
         "\n",
         "for prediction in response.predictions:\n",
         "    embeddings = Tensor(prediction)\n",
@@ -479,12 +505,13 @@
         "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
         "# @markdown  and avoid unnecessary continuous charges that may incur.\n",
         "\n",
-        "endpoint.delete(force=True)\n",
-        "model.delete()\n",
+        "# Undeploy model and delete endpoint.\n",
+        "for endpoint in endpoints.values():\n",
+        "    endpoint.delete(force=True)\n",
         "\n",
-        "delete_bucket = False  # @param {type:\"boolean\"}\n",
-        "if delete_bucket:\n",
-        "    ! gsutil -m rm -r $BUCKET_URI"
+        "# Delete models.\n",
+        "for model in models.values():\n",
+        "    model.delete()"
       ]
     }
   ],

notebooks/community/model_garden/model_garden_gemma2_deployment_on_vertex.ipynb

@@ -3,7 +3,6 @@
     {
       "cell_type": "code",
       "execution_count": null,
-      "language": "python",
       "metadata": {
         "cellView": "form",
         "id": "7d9bbf86da5e"
@@ -35,13 +34,18 @@
         "\n",
         "<table><tbody><tr>\n",
         "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/instances\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_gemma2_deployment_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_gemma2_deployment_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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"
@@ -65,6 +69,10 @@
         "- Deploy Gemma 2 with Hex-LLM on TPU\n",
         "- Deploy Gemma with [TGI](https://github.com/huggingface/text-generation-inference) on GPU\n",
         "\n",
+        "### File a bug\n",
+        "\n",
+        "File a bug on [GitHub](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues/new) if you encounter any issue with the notebook.\n",
+        "\n",
         "### Costs\n",
         "\n",
         "This tutorial uses billable components of Google Cloud:\n",
@@ -111,32 +119,38 @@
         "\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",
+        "# @markdown 2. **[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",
+        "# @markdown 3. 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). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\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-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# Upgrade Vertex AI SDK.\n",
-        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.64.0'\n",
-        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform==1.93.1'\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",
         "\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
+        "    ! pip install --upgrade tensorflow\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\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",
+        "LABEL = \"tgi\"\n",
         "models, endpoints = {}, {}\n",
         "\n",
         "# Get the default cloud project id.\n",
@@ -146,56 +160,24 @@
         "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, \"gemma2\")\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",
+        "aiplatform.init(project=PROJECT_ID, location=REGION)\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",
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")\n",
+        "\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
         "# @markdown ## Access Gemma 2 Models\n",
         "\n",
-        "# @markdown You must provide a Hugging Face User Access Token (read) to access the Gemma 2 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",
+        "# @markdown You must provide a Hugging Face User Access Token (with read access) to access the Gemma 2 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",
         "assert (\n",
@@ -223,21 +205,22 @@
       "execution_count": null,
       "metadata": {
         "cellView": "form",
-        "id": "E8OiHHNNE_wj"
+        "id": "B7bg9nM0S0Mp"
       },
       "outputs": [],
       "source": [
-        "# @title Deploy\n",
-        "# @markdown Set the model ID. Model weights can be loaded from HuggingFace or from a GCS bucket.\n",
-        "\n",
-        "# The pre-built serving docker images.\n",
-        "HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:20241210_2323_RC00\"\n",
+        "# @title Select the model variants\n",
         "\n",
         "# @markdown Select one of the four model variations.\n",
         "MODEL_ID = \"gemma-2-2b-it\"  # @param [\"gemma-2-2b\", \"gemma-2-2b-it\", \"gemma-2-9b\", \"gemma-2-9b-it\", \"gemma-2-27b\", \"gemma-2-27b-it\"] {allow-input: true, isTemplate: true}\n",
+        "version_id = f\"publishers/google/models/gemma2/@{MODEL_ID}\"\n",
+        "\n",
         "TPU_DEPLOYMENT_REGION = \"us-west1\"  # @param [\"us-west1\"] {isTemplate:true}\n",
         "model_id = os.path.join(model_path_prefix, MODEL_ID)\n",
         "\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
+        "\n",
         "# @markdown Find Vertex AI prediction TPUv5e machine types in\n",
         "# @markdown https://cloud.google.com/vertex-ai/docs/predictions/use-tpu#deploy_a_model.\n",
         "if \"2b\" in model_id:\n",
@@ -265,16 +248,29 @@
         "    accelerator_type=accelerator_type,\n",
         "    accelerator_count=accelerator_count,\n",
         "    is_for_training=False,\n",
-        ")\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "E8OiHHNNE_wj"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy Gemma2 models with Hex-LLM on TPU\n",
+        "# @markdown Set the model ID. Model weights can be loaded from HuggingFace or from a GCS bucket.\n",
+        "\n",
+        "# The pre-built serving docker images.\n",
+        "HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:20241210_2323_RC00\"\n",
         "\n",
         "# Server parameters.\n",
         "tensor_parallel_size = accelerator_count\n",
         "hbm_utilization_factor = 0.6  # Fraction of HBM memory allocated for KV cache after model loading. A larger value improves throughput but gives higher risk of TPU out-of-memory errors with long prompts.\n",
         "max_running_seqs = 256  # Maximum number of running sequences in a continuous batch.\n",
         "\n",
-        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
-        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
-        "\n",
         "# Endpoint configurations.\n",
         "min_replica_count = 1\n",
         "max_replica_count = 1\n",
@@ -285,7 +281,6 @@
         "    model_id: str,\n",
         "    publisher: str,\n",
         "    publisher_model_id: str,\n",
-        "    service_account: str,\n",
         "    base_model_id: str = None,\n",
         "    data_parallel_size: int = 1,\n",
         "    tensor_parallel_size: int = 1,\n",
@@ -374,11 +369,11 @@
         "        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",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_gemma2_deployment_on_vertex.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    return model, endpoint\n",
@@ -389,7 +384,6 @@
         "    model_id=model_id,\n",
         "    publisher=\"google\",\n",
         "    publisher_model_id=\"gemma2\",\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
         "    machine_type=machine_type,\n",
         "    tensor_parallel_size=tensor_parallel_size,\n",
         "    hbm_utilization_factor=hbm_utilization_factor,\n",
@@ -482,23 +476,28 @@
       "execution_count": null,
       "metadata": {
         "cellView": "form",
-        "id": "TBNJYZMlBNwZ"
+        "id": "eYst8GHqcGco"
       },
       "outputs": [],
       "source": [
-        "# @title Deploy\n",
+        "# @title Select the model variants\n",
         "\n",
         "# The pre-built serving docker image.\n",
         "TGI_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu121.2-1.ubuntu2204.py310\"\n",
         "\n",
         "MODEL_ID = \"gemma-2-2b\"  # @param [\"gemma-2-2b\", \"gemma-2-2b-it\", \"gemma-2-9b\", \"gemma-2-9b-it\", \"gemma-2-27b\", \"gemma-2-27b-it\"] {allow-input: true, isTemplate: true}\n",
         "model_id = os.path.join(model_path_prefix, MODEL_ID)\n",
+        "PUBLISHER_MODEL_NAME = f\"publishers/google/models/gemma2@{MODEL_ID}\"\n",
         "\n",
         "# @markdown Finds Vertex AI prediction supported accelerators and regions in\n",
         "# @markdown https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
         "\n",
         "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\"] {isTemplate: true}\n",
         "\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
+        "\n",
+        "\n",
         "if \"2b\" in MODEL_ID:\n",
         "    if accelerator_type == \"NVIDIA_L4\":\n",
         "        # Sets 1 L4 (24G) to deploy Gemma 2 2B models.\n",
@@ -540,6 +539,46 @@
         "    is_for_training=False,\n",
         ")\n",
         "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "nlRmOQmZhjvp"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 1] Deploy with Model Garden SDK\n",
+        "\n",
+        "# @markdown Deploy with Gen AI model-centric SDK. This section uploads the prebuilt model to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model. See [use open models with Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-open-models) for documentation on other use cases.\n",
+        "from vertexai import model_garden\n",
+        "\n",
+        "model = model_garden.OpenModel(PUBLISHER_MODEL_NAME)\n",
+        "endpoints[LABEL] = model.deploy(\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    accept_eula=True,  # Accept the End User License Agreement (EULA) on the model card before deploy. Otherwise, the deployment will be forbidden.\n",
+        ")\n",
+        "\n",
+        "endpoint = endpoints[LABEL]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "TBNJYZMlBNwZ"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy Gemma models with TGI on GPU\n",
+        "\n",
         "# Note that larger token counts will require more GPU memory. For example, if you'd\n",
         "# like to increase the `max_total_tokens` and `max_batch_prefill_tokens` to 8192,\n",
         "# you may need 1 L4 for 2b model, 4 L4s for the 9b model, and 8 L4s for the 27b model.\n",
@@ -553,7 +592,7 @@
         "    model_id: str,\n",
         "    publisher: str,\n",
         "    publisher_model_id: str,\n",
-        "    service_account: str,\n",
+        "    service_account: str = None,\n",
         "    machine_type: str = \"g2-standard-8\",\n",
         "    accelerator_type: str = \"NVIDIA_L4\",\n",
         "    accelerator_count: int = 1,\n",
@@ -584,6 +623,9 @@
         "    except NameError:\n",
         "        pass\n",
         "\n",
+        "    if service_account:\n",
+        "        env_vars[\"SERVICE_ACCOUNT\"] = service_account\n",
+        "\n",
         "    model = aiplatform.Model.upload(\n",
         "        display_name=model_name,\n",
         "        serving_container_image_uri=TGI_DOCKER_URI,\n",
@@ -604,17 +646,17 @@
         "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_gemma2_deployment_on_vertex.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    return model, endpoint\n",
         "\n",
         "\n",
-        "models[\"tgi\"], endpoints[\"tgi\"] = deploy_model_tgi(\n",
+        "models[LABEL], endpoints[LABEL] = deploy_model_tgi(\n",
         "    model_name=common_util.get_job_name_with_datetime(prefix=MODEL_ID),\n",
         "    model_id=model_id,\n",
         "    publisher=\"google\",\n",
         "    publisher_model_id=\"gemma2\",\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
         "    machine_type=machine_type,\n",
         "    accelerator_type=accelerator_type,\n",
         "    accelerator_count=accelerator_count,\n",
@@ -622,7 +664,9 @@
         "    max_total_tokens=max_total_tokens,\n",
         "    max_batch_prefill_tokens=max_batch_prefill_tokens,\n",
         "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
-        ")"
+        ")\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
       ]
     },
     {
@@ -705,6 +749,8 @@
       },
       "outputs": [],
       "source": [
+        "# @title Delete the models and endpoints\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",
@@ -714,11 +760,7 @@
         "\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"
+        "    model.delete()"
       ]
     }
   ],

notebooks/community/model_garden/model_garden_gemma2_finetuning_on_vertex.ipynb

@@ -10,7 +10,7 @@
       },
       "outputs": [],
       "source": [
-        "# Copyright 2024 Google LLC\n",
+        "# Copyright 2025 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",
@@ -36,13 +36,18 @@
         "\n",
         "<table><tbody><tr>\n",
         "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/instances\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_gemma2_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_gemma2_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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"
@@ -65,6 +70,7 @@
         "### Objective\n",
         "\n",
         "- Finetune and deploy Gemma 2 models with Vertex AI Custom Training Jobs.\n",
+        "- Evaluate the finetuned model using [lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness).\n",
         "- Send prediction requests to your finetuned Gemma 2 model.\n",
         "\n",
         "### File a bug\n",
@@ -101,17 +107,11 @@
       "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",
+        "# @markdown 1. Install packages to validate dataset with template.\n",
         "! pip install --upgrade --quiet gcsfs==2024.3.1\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"
+        "! pip install --upgrade --quiet accelerate==0.34.2\n",
+        "! pip install --upgrade --quiet transformers==4.47.1\n",
+        "! pip install --upgrade --quiet datasets==2.20.0"
       ]
     },
     {
@@ -135,7 +135,9 @@
         "# @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",
+        "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, 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",
@@ -147,7 +149,7 @@
         "\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",
+        "! cd vertex-ai-samples && git reset --hard c45f6a4f4d32e31a050f0e4ba52824b0caf4eda3\n",
         "\n",
         "import datetime\n",
         "import importlib\n",
@@ -159,6 +161,9 @@
         "from google.cloud.aiplatform.compat.types import \\\n",
         "    custom_job as gca_custom_job_compat\n",
         "\n",
+        "if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
+        "    ! pip install --upgrade tensorflow\n",
+        "\n",
         "common_util = importlib.import_module(\n",
         "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
         ")\n",
@@ -170,6 +175,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -221,7 +232,7 @@
         "\n",
         "# @markdown ## Access Gemma 2 Models\n",
         "\n",
-        "# @markdown You must provide a Hugging Face User Access Token (read) to access the Gemma 2 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",
+        "# @markdown You must provide a Hugging Face User Access Token (with read access) to access the Gemma 2 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",
         "assert HF_TOKEN, \"Provide a read HF_TOKEN to load models from Hugging Face.\"\n",
@@ -252,7 +263,7 @@
         "# @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",
+        "# @markdown You can set `train_dataset` to any existing [Hugging Face dataset](https://huggingface.co/datasets) name, and set `train_column` 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 `train_column` to `text` in this notebook.\n",
         "\n",
         "# @markdown ### (Optional) Prepare a custom JSONL dataset for finetuning\n",
         "\n",
@@ -261,13 +272,13 @@
         "# @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",
+        "# @markdown The JSON object has a key `text`, which should match `train_column`; 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",
+        "# @markdown Optionally update the `train_column` 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",
@@ -293,23 +304,25 @@
         "# @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",
+        "# @markdown 1. Set `train_dataset` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`\n",
+        "# @markdown 1. Set `train_split` to `train`\n",
+        "# @markdown 1. Set `eval_dataset` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_eval_sample.jsonl`\n",
+        "# @markdown 1. Set `eval_split` to `train` (**NOT** `test`)\n",
+        "# @markdown 1. Set `train_column` 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",
+        "train_dataset = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
+        "train_split = \"train\"  # @param {type:\"string\"}\n",
+        "eval_dataset = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
+        "eval_split = \"test\"  # @param {type:\"string\"}\n",
         "\n",
         "# Name of the dataset column containing training text input.\n",
-        "instruct_column_in_dataset = \"text\"  # @param {type:\"string\"}"
+        "train_column = \"text\"  # @param {type:\"string\"}\n",
+        "# Maximum sequence length.\n",
+        "max_seq_length = 4096  # @param{type:\"integer\"}"
       ]
     },
     {
@@ -341,8 +354,6 @@
         "\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",
@@ -358,32 +369,30 @@
         "    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",
+        "tokenizer = dataset_validation_util.load_tokenizer(tokenizer_path, None, access_token)\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",
+        "    dataset_name=train_dataset,\n",
+        "    split=train_split,\n",
+        "    input_column=train_column,\n",
         "    template=template,\n",
+        "    max_seq_length=max_seq_length,\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",
-        ")"
+        "# Validate the eval dataset if it exists.\n",
+        "if eval_dataset:\n",
+        "    dataset_validation_util.validate_dataset_with_template(\n",
+        "        dataset_name=eval_dataset,\n",
+        "        split=eval_split,\n",
+        "        input_column=train_column,\n",
+        "        template=template,\n",
+        "        max_seq_length=max_seq_length,\n",
+        "        use_multiprocessing=False,\n",
+        "        tokenizer=tokenizer,\n",
+        "    )"
       ]
     },
     {
@@ -407,10 +416,10 @@
         "# @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",
         "# @markdown Accelerator type to use for training.\n",
-        "accelerator_type = \"NVIDIA_A100_80GB\"  # @param [\"NVIDIA_A100_80GB\", \"NVIDIA_H100_80GB\"]\n",
+        "training_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",
+        "if training_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",
@@ -425,30 +434,28 @@
         "    }\n",
         "\n",
         "TRAIN_DOCKER_URI = (\n",
-        "    f\"{repo}/vertex-vision-model-garden-dockers/pytorch-peft-train:stable_20240909\"\n",
+        "    f\"{repo}/vertex-vision-model-garden-dockers/pytorch-peft-train:stable_20250409\"\n",
         ")\n",
         "\n",
         "# Worker pool spec.\n",
-        "if accelerator_type == \"NVIDIA_A100_80GB\":\n",
+        "if training_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",
+        "    training_machine_type = \"a2-ultragpu-8g\"\n",
+        "elif training_accelerator_type == \"NVIDIA_H100_80GB\":\n",
         "    per_node_accelerator_count = 8\n",
-        "    machine_type = \"a3-highgpu-8g\"\n",
+        "    training_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",
+        "        f\"Recommended machine settings not found for: {training_accelerator_type}. To use another accelerator type, edit this code block to pass in an appropriate `training_machine_type`, `training_accelerator_type`, and `per_node_accelerator_count` 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",
+        "# @markdown Setting a positive `max_steps` here will override `num_train_epochs`.\n",
         "max_steps = -1  # @param{type:\"integer\"}\n",
-        "num_epochs = 1.0  # @param{type:\"number\"}\n",
+        "num_train_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",
@@ -460,7 +467,7 @@
         "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",
+        "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",
@@ -476,12 +483,17 @@
         "# Train precision of the model.\n",
         "train_precision = \"bfloat16\"\n",
         "\n",
+        "# @markdown Evaluation metrics to compute. Supported eval metrics: loss, perplexity, bleu, google_bleu, rouge1, rouge2, rougeL, rougeLsum.\n",
+        "eval_metric_name = \"loss,perplexity,bleu\"  # @param{type:\"string\"}\n",
+        "# @markdown Metric to use for best model selection. This will save the best checkpoint based on the eval metric.\n",
+        "metric_for_best_model = \"perplexity\"  # @param{type:\"string\"}\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_type=training_accelerator_type,\n",
         "    accelerator_count=per_node_accelerator_count * replica_count,\n",
         "    is_for_training=True,\n",
         "    is_restricted_image=is_restricted_image,\n",
@@ -498,8 +510,9 @@
         "merged_model_output_dir = os.path.join(base_output_dir, \"merged-model\")\n",
         "\n",
         "# Add labels for the finetuning job.\n",
+        "\n",
         "labels = {\n",
-        "    \"mg-source\": \"notebook\",\n",
+        "    \"mg-source\": common_util.get_deploy_source(),\n",
         "    \"mg-notebook-name\": \"model_garden_gemma2_finetuning_on_vertex.ipynb\".split(\".\")[0],\n",
         "}\n",
         "\n",
@@ -508,23 +521,23 @@
         "labels[\"versioned-mg-tune\"] = f\"{labels['mg-tune']}-{versioned_model_id}\"\n",
         "\n",
         "eval_args = [\n",
-        "    f\"--eval_dataset_path={eval_dataset_name}\",\n",
-        "    f\"--eval_column={instruct_column_in_dataset}\",\n",
+        "    f\"--eval_dataset={eval_dataset}\",\n",
+        "    f\"--eval_column={train_column}\",\n",
         "    f\"--eval_template={template}\",\n",
-        "    f\"--eval_split={eval_split_name}\",\n",
+        "    f\"--eval_split={eval_split}\",\n",
         "    f\"--eval_steps={save_steps}\",\n",
-        "    \"--eval_tasks=builtin_eval\",\n",
-        "    \"--eval_metric_name=loss\",\n",
+        "    f\"--eval_metric_name={eval_metric_name}\",\n",
+        "    f\"--metric_for_best_model={metric_for_best_model}\",\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",
+        "    \"--input_masking=True\",\n",
+        "    f\"--pretrained_model_name_or_path={pretrained_model_id}\",\n",
+        "    f\"--train_dataset={train_dataset}\",\n",
+        "    f\"--train_split={train_split}\",\n",
+        "    f\"--train_column={train_column}\",\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",
@@ -538,8 +551,9 @@
         "    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\"--merge_model_precision_mode={train_precision}\",\n",
+        "    f\"--gradient_checkpointing={gradient_checkpointing}\",\n",
+        "    f\"--num_train_epochs={num_train_epochs}\",\n",
         "    f\"--attn_implementation={attn_implementation}\",\n",
         "    f\"--optimizer={optimizer}\",\n",
         "    f\"--warmup_ratio={warmup_ratio}\",\n",
@@ -547,7 +561,7 @@
         "    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\"--train_template={template}\",\n",
         "    f\"--huggingface_access_token={HF_TOKEN}\",\n",
         "] + eval_args\n",
         "\n",
@@ -563,8 +577,8 @@
         "train_job.run(\n",
         "    args=train_job_args,\n",
         "    replica_count=replica_count,\n",
-        "    machine_type=machine_type,\n",
-        "    accelerator_type=accelerator_type,\n",
+        "    machine_type=training_machine_type,\n",
+        "    accelerator_type=training_accelerator_type,\n",
         "    accelerator_count=per_node_accelerator_count,\n",
         "    boot_disk_size_gb=500,\n",
         "    service_account=SERVICE_ACCOUNT,\n",
@@ -602,6 +616,126 @@
         "print(f\"Command to copy: tensorboard --logdir {base_output_dir}/logs\")"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "KdtcMGHgtrVC"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Select Evaluation Checkpoint\n",
+        "\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",
+        "# @markdown The following checkpoints are available for evaluation:\n",
+        "! gcloud storage ls \"{lora_output_dir}/node-0\" | grep \"checkpoint-\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "1LBADPr6tTqy"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Run Evaluation Job\n",
+        "# @markdown This section runs the evaluation using [lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness) on the finetuned model. The evaluation takes approximately 20 mins to finish.\n",
+        "\n",
+        "# The pre-built evaluation docker image for LM Evaluation Harness.\n",
+        "LM_EVAL_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-lm-evaluation-harness:20250410_1035_RC00\"\n",
+        "\n",
+        "# @markdown Set `RUN_EVALUATION` to False to skip the evaluation job.\n",
+        "RUN_EVALUATION = True  # @param {type:\"boolean\"}\n",
+        "\n",
+        "eval_accelerator_type = \"NVIDIA_L4\"\n",
+        "gpu_memory_utilization = 0.85\n",
+        "\n",
+        "if \"2b\" in base_model_id:\n",
+        "    eval_machine_type = \"g2-standard-12\"\n",
+        "    eval_accelerator_count = 1\n",
+        "elif \"9b\" in base_model_id:\n",
+        "    eval_machine_type = \"g2-standard-48\"\n",
+        "    eval_accelerator_count = 4\n",
+        "elif \"27b\" in base_model_id:\n",
+        "    eval_machine_type = \"g2-standard-96\"\n",
+        "    eval_accelerator_count = 8\n",
+        "    gpu_memory_utilization = 0.8\n",
+        "else:\n",
+        "    raise ValueError(\n",
+        "        \"Recommended machine settings not found for model: %s\" % base_model_id\n",
+        "    )\n",
+        "\n",
+        "# @markdown Set `evaluation_checkpoint_dir` to an intermediate checkpoint from the above training job. If not set, the evaluation job will use the merged model.\n",
+        "evaluation_checkpoint_dir = \"\"  # @param {type:\"string\"}\n",
+        "if evaluation_checkpoint_dir:\n",
+        "    pretrained = pretrained_model_id\n",
+        "else:\n",
+        "    pretrained = merged_model_output_dir\n",
+        "\n",
+        "# @markdown Evaluation tasks to run.\n",
+        "eval_tasks = \"coqa\"  # @param {type:\"string\"}\n",
+        "# @markdown Model to use for evaluation.\n",
+        "model = \"vllm\"  # @param {type:\"string\"}\n",
+        "# @markdown Batch size for evaluation.\n",
+        "batch_size = \"auto\"  # @param {type:\"string\"}\n",
+        "apply_chat_template = True if \"-it\" in pretrained_model_id else False\n",
+        "max_model_len = 4096  # Maximum context length.\n",
+        "\n",
+        "model_args = f\"tensor_parallel_size={eval_accelerator_count},max_model_len={max_model_len},gpu_memory_utilization={gpu_memory_utilization},enforce_eager=True\"\n",
+        "eval_output_dir = os.path.join(base_output_dir, \"lm_eval\")\n",
+        "\n",
+        "lm_eval_job_args = [\n",
+        "    \"--task=lm_eval\",\n",
+        "    f\"--model={model}\",\n",
+        "    f\"--eval_tasks={eval_tasks}\",\n",
+        "    f\"--pretrained_model_name_or_path={pretrained}\",\n",
+        "    f\"--model_args={model_args}\",\n",
+        "    f\"--output_dir={eval_output_dir}\",\n",
+        "    f\"--apply_chat_template={apply_chat_template}\",\n",
+        "    f\"--batch_size={batch_size}\",\n",
+        "    f\"--huggingface_access_token={HF_TOKEN}\",\n",
+        "]\n",
+        "\n",
+        "if evaluation_checkpoint_dir:\n",
+        "    lm_eval_job_args.append(f\"--lora_path={evaluation_checkpoint_dir}\")\n",
+        "\n",
+        "if RUN_EVALUATION:\n",
+        "    common_util.check_quota(\n",
+        "        project_id=PROJECT_ID,\n",
+        "        region=REGION,\n",
+        "        accelerator_type=eval_accelerator_type,\n",
+        "        accelerator_count=eval_accelerator_count,\n",
+        "        is_for_training=True,\n",
+        "    )\n",
+        "    lm_eval_job = aiplatform.CustomContainerTrainingJob(\n",
+        "        display_name=common_util.get_job_name_with_datetime(\"gemma2-lm-eval\"),\n",
+        "        container_uri=LM_EVAL_DOCKER_URI,\n",
+        "        labels=labels,\n",
+        "    )\n",
+        "\n",
+        "    print(\"Running evaluation job with args:\")\n",
+        "    print(\" \\\\\\n\".join(lm_eval_job_args))\n",
+        "    lm_eval_job.run(\n",
+        "        args=lm_eval_job_args,\n",
+        "        replica_count=1,\n",
+        "        machine_type=eval_machine_type,\n",
+        "        accelerator_type=eval_accelerator_type,\n",
+        "        accelerator_count=eval_accelerator_count,\n",
+        "        boot_disk_size_gb=500,\n",
+        "        service_account=SERVICE_ACCOUNT,\n",
+        "        base_output_dir=base_output_dir,\n",
+        "    )\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -614,50 +748,47 @@
       "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",
-        "\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",
         "# 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",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20250116_0916_RC00\"\n",
+        "\n",
+        "# @markdown Set `use_dedicated_endpoint` to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
         "\n",
         "# Find Vertex AI prediction supported accelerators and regions [here](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute).\n",
         "# @markdown Accelerator type to use for serving.\n",
-        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\"] {isTemplate: true}\n",
+        "serving_accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\"] {isTemplate: true}\n",
         "\n",
         "if \"2b\" in base_model_id:\n",
-        "    if accelerator_type == \"NVIDIA_L4\":\n",
+        "    if serving_accelerator_type == \"NVIDIA_L4\":\n",
         "        # Sets 1 L4 (24G) to deploy Gemma 2 2B models.\n",
-        "        machine_type = \"g2-standard-12\"\n",
-        "        accelerator_count = 1\n",
+        "        serving_machine_type = \"g2-standard-12\"\n",
+        "        serving_accelerator_count = 1\n",
         "    else:\n",
         "        raise ValueError(\n",
         "            \"Recommended machine settings not found for accelerator type: %s\"\n",
-        "            % accelerator_type\n",
+        "            % serving_accelerator_type\n",
         "        )\n",
         "elif \"9b\" in base_model_id:\n",
-        "    if accelerator_type == \"NVIDIA_L4\":\n",
+        "    if serving_accelerator_type == \"NVIDIA_L4\":\n",
         "        # Sets 2 L4 (24G) to deploy Gemma 2 9B models.\n",
-        "        machine_type = \"g2-standard-24\"\n",
-        "        accelerator_count = 2\n",
+        "        serving_machine_type = \"g2-standard-24\"\n",
+        "        serving_accelerator_count = 2\n",
         "    else:\n",
         "        raise ValueError(\n",
         "            \"Recommended machine settings not found for accelerator type: %s\"\n",
-        "            % accelerator_type\n",
+        "            % serving_accelerator_type\n",
         "        )\n",
         "elif \"27b\" in base_model_id:\n",
-        "    if accelerator_type == \"NVIDIA_L4\":\n",
+        "    if serving_accelerator_type == \"NVIDIA_L4\":\n",
         "        # Sets 4 L4 (24G) to deploy Gemma 2 27B models.\n",
-        "        machine_type = \"g2-standard-48\"\n",
-        "        accelerator_count = 4\n",
+        "        serving_machine_type = \"g2-standard-48\"\n",
+        "        serving_accelerator_count = 4\n",
         "    else:\n",
         "        raise ValueError(\n",
         "            \"Recommended machine settings not found for accelerator type: %s\"\n",
-        "            % accelerator_type\n",
+        "            % serving_accelerator_type\n",
         "        )\n",
         "else:\n",
         "    raise ValueError(\n",
@@ -667,21 +798,32 @@
         "common_util.check_quota(\n",
         "    project_id=PROJECT_ID,\n",
         "    region=REGION,\n",
-        "    accelerator_type=accelerator_type,\n",
-        "    accelerator_count=accelerator_count,\n",
+        "    accelerator_type=serving_accelerator_type,\n",
+        "    accelerator_count=serving_accelerator_count,\n",
         "    is_for_training=False,\n",
         ")\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",
         "gpu_memory_utilization = 0.85\n",
         "max_model_len = 4096  # Maximum context length.\n",
         "\n",
         "\n",
+        "def get_deploy_source() -> str:\n",
+        "    \"\"\"Gets deploy_source string based on running environment.\"\"\"\n",
+        "    vertex_product = os.environ.get(\"VERTEX_PRODUCT\", \"\")\n",
+        "    if vertex_product == \"COLAB_ENTERPRISE\":\n",
+        "        return \"notebook_colab_enterprise\"\n",
+        "    elif vertex_product == \"WORKBENCH_INSTANCE\":\n",
+        "        return \"notebook_workbench\"\n",
+        "    else:\n",
+        "        # Legacy workbench, legacy colab, or other custom environments.\n",
+        "        return \"notebook_environment_unspecified\"\n",
+        "\n",
+        "\n",
         "def deploy_model_vllm(\n",
         "    model_name: str,\n",
         "    model_id: str,\n",
+        "    publisher: str,\n",
+        "    publisher_model_id: str,\n",
         "    service_account: str,\n",
         "    base_model_id: str = None,\n",
         "    machine_type: str = \"g2-standard-8\",\n",
@@ -693,11 +835,15 @@
         "    enable_trust_remote_code: bool = False,\n",
         "    enforce_eager: bool = False,\n",
         "    enable_lora: bool = False,\n",
+        "    enable_chunked_prefill: bool = False,\n",
+        "    enable_prefix_cache: bool = False,\n",
+        "    host_prefix_kv_cache_utilization_target: float = 0.0,\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",
+        "    enable_llama_tool_parser: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -736,9 +882,24 @@
         "    if enable_lora:\n",
         "        vllm_args.append(\"--enable-lora\")\n",
         "\n",
+        "    if enable_chunked_prefill:\n",
+        "        vllm_args.append(\"--enable-chunked-prefill\")\n",
+        "\n",
+        "    if enable_prefix_cache:\n",
+        "        vllm_args.append(\"--enable-prefix-caching\")\n",
+        "\n",
+        "    if 0 < host_prefix_kv_cache_utilization_target < 1:\n",
+        "        vllm_args.append(\n",
+        "            f\"--host-prefix-kv-cache-utilization-target={host_prefix_kv_cache_utilization_target}\"\n",
+        "        )\n",
+        "\n",
         "    if model_type:\n",
         "        vllm_args.append(f\"--model-type={model_type}\")\n",
         "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -761,6 +922,9 @@
         "        serving_container_environment_variables=env_vars,\n",
         "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
         "        serving_container_deployment_timeout=7200,\n",
+        "        model_garden_source_model_name=(\n",
+        "            f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
+        "        ),\n",
         "    )\n",
         "    print(\n",
         "        f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
@@ -772,6 +936,10 @@
         "        accelerator_count=accelerator_count,\n",
         "        deploy_request_timeout=1800,\n",
         "        service_account=service_account,\n",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_gemma2_finetuning_on_vertex.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": get_deploy_source(),\n",
+        "        },\n",
         "    )\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",
         "\n",
@@ -781,10 +949,12 @@
         "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
         "    model_name=common_util.get_job_name_with_datetime(prefix=\"gemma2-vllm-serve\"),\n",
         "    model_id=merged_model_output_dir,\n",
+        "    publisher=\"google\",\n",
+        "    publisher_model_id=\"gemma2\",\n",
         "    service_account=SERVICE_ACCOUNT,\n",
-        "    machine_type=machine_type,\n",
-        "    accelerator_type=accelerator_type,\n",
-        "    accelerator_count=accelerator_count,\n",
+        "    machine_type=serving_machine_type,\n",
+        "    accelerator_type=serving_accelerator_type,\n",
+        "    accelerator_count=serving_accelerator_count,\n",
         "    gpu_memory_utilization=gpu_memory_utilization,\n",
         "    max_model_len=max_model_len,\n",
         "    use_dedicated_endpoint=use_dedicated_endpoint,\n",

notebooks/community/model_garden/model_garden_gemma3_deployment_on_vertex.ipynb

@@ -0,0 +1,978 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "7d9bbf86da5e"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2025 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 3 (Deployment)\n",
+        "\n",
+        "<table><tbody><tr>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/instances\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_gemma3_deployment_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_gemma3_deployment_on_vertex.ipynb\">\n",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.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 deploying Gemma 3 models on GPU using [vLLM](https://github.com/vllm-project/vllm).\n",
+        "\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Deploy Gemma 3 with vLLM on GPU\n",
+        "\n",
+        "### File a bug\n",
+        "\n",
+        "File a bug on [GitHub](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues/new) if you encounter any issue with the notebook.\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": "2707b02ef5df"
+      },
+      "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. **[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",
+        "# @markdown 3. 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). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\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-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
+        "\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform==1.93.1'\n",
+        "\n",
+        "# Import the necessary packages\n",
+        "import importlib\n",
+        "import os\n",
+        "from typing import Tuple\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
+        "    ! pip install --upgrade tensorflow\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\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",
+        "# Initialize models and endpoints as a dict\n",
+        "models, endpoints = {}, {}\n",
+        "\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",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION)\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",
+        "! gcloud config set project $PROJECT_ID\n",
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "GFfqpQm8BNwZ"
+      },
+      "source": [
+        "## Deploy Gemma 3 1B models with vLLM on GPU"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "W1PYsnz2JXHm"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown Set the model to deploy.\n",
+        "\n",
+        "base_model_name = \"gemma-3-1b-it\"  # @param [\"gemma-3-1b-pt\", \"gemma-3-1b-it\"] {isTemplate:true}\n",
+        "hf_model_id = \"google/\" + base_model_name\n",
+        "PUBLISHER_MODEL_NAME = f\"publishers/google/models/gemma3@{base_model_name}\"\n",
+        "model_id = f\"gs://vertex-model-garden-restricted-us/gemma3/{base_model_name}\"\n",
+        "\n",
+        "# The pre-built serving docker image.\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20250312_0916_RC01\"\n",
+        "\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
+        "accelerator_type = \"NVIDIA_L4\"\n",
+        "machine_type = \"g2-standard-12\"\n",
+        "accelerator_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=False,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "d2SO5S5fMKE2"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 1] Deploy with Model Garden SDK\n",
+        "\n",
+        "LABEL = \"sdk-deploy-1b\"\n",
+        "# @markdown Deploy with Gen AI model-centric SDK. This section uploads the prebuilt model to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model. See [use open models with Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-open-models) for documentation on other use cases.\n",
+        "from vertexai import model_garden\n",
+        "\n",
+        "model = model_garden.OpenModel(PUBLISHER_MODEL_NAME)\n",
+        "endpoints[LABEL] = model.deploy(\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    accept_eula=True,  # Accept the End User License Agreement (EULA) on the model card before deploy. Otherwise, the deployment will be forbidden.\n",
+        ")\n",
+        "\n",
+        "endpoint = endpoints[LABEL]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "TBNJYZMlBNwZ"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with customized configs\n",
+        "\n",
+        "# @markdown This section uploads Gemma 3 1B models to Model Registry and deploys them to a Vertex Prediction Endpoint. It takes 15 minutes to 30 minutes to finish.\n",
+        "\n",
+        "gpu_memory_utilization = 0.95\n",
+        "max_model_len = 32768\n",
+        "\n",
+        "\n",
+        "def deploy_model_vllm(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    publisher: str,\n",
+        "    publisher_model_id: 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",
+        "    enable_chunked_prefill: bool = False,\n",
+        "    enable_prefix_cache: bool = False,\n",
+        "    host_prefix_kv_cache_utilization_target: float = 0.0,\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",
+        "    enable_llama_tool_parser: bool = False,\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 enable_chunked_prefill:\n",
+        "        vllm_args.append(\"--enable-chunked-prefill\")\n",
+        "\n",
+        "    if enable_prefix_cache:\n",
+        "        vllm_args.append(\"--enable-prefix-caching\")\n",
+        "\n",
+        "    if 0 < host_prefix_kv_cache_utilization_target < 1:\n",
+        "        vllm_args.append(\n",
+        "            f\"--host-prefix-kv-cache-utilization-target={host_prefix_kv_cache_utilization_target}\"\n",
+        "        )\n",
+        "\n",
+        "    if model_type:\n",
+        "        vllm_args.append(f\"--model-type={model_type}\")\n",
+        "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\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",
+        "        model_garden_source_model_name=(\n",
+        "            f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
+        "        ),\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",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_gemma3_deployment_on_vertex.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
+        "        },\n",
+        "    )\n",
+        "    print(\"endpoint_name:\", endpoint.name)\n",
+        "\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "LABEL = \"custom-deploy-1b\"\n",
+        "\n",
+        "models[LABEL], endpoints[LABEL] = deploy_model_vllm(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=\"gemma3-serve\"),\n",
+        "    model_id=model_id,\n",
+        "    publisher=\"google\",\n",
+        "    publisher_model_id=\"gemma3\",\n",
+        "    base_model_id=hf_model_id,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    gpu_memory_utilization=gpu_memory_utilization,\n",
+        "    max_model_len=max_model_len,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        ")\n",
+        "\n",
+        "model = models[LABEL]\n",
+        "endpoint = endpoints[LABEL]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "rDHsCOqvFYBi"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Raw predict\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. Sampling parameters supported by vLLM can be found [here](https://docs.vllm.ai/en/latest/dev/sampling_params.html).\n",
+        "\n",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown Human: What is a car?\n",
+        "# @markdown Assistant:  A car, or a motor car, is a road-connected human-transportation system used to move people or goods from one place to another. The term also encompasses a wide range of vehicles, including motorboats, trains, and aircrafts. Cars typically have four wheels, a cabin for passengers, and an engine or motor. They have been around since the early 19th century and are now one of the most popular forms of transportation, used for daily commuting, shopping, and other purposes.\n",
+        "# @markdown ```\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",
+        "prompt = \"What is a car?\"  # @param {type: \"string\"}\n",
+        "# @markdown If you encounter an issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, by lowering `max_tokens`.\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  # @param {type:\"integer\"}\n",
+        "# @markdown Set `raw_response` to `True` to obtain the raw model output. Set `raw_response` to `False` to apply additional formatting in the structure of `\"Prompt:\\n{prompt.strip()}\\nOutput:\\n{output}\"`.\n",
+        "raw_response = False  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# Overrides parameters for inferences.\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",
+        "        \"raw_response\": raw_response,\n",
+        "    },\n",
+        "]\n",
+        "response = endpoint.predict(\n",
+        "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
+        ")\n",
+        "\n",
+        "for prediction in response.predictions:\n",
+        "    print(prediction)\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "LSG9ITWTbTb7"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Chat completion\n",
+        "\n",
+        "if use_dedicated_endpoint:\n",
+        "    DEDICATED_ENDPOINT_DNS = endpoint.gca_resource.dedicated_endpoint_dns\n",
+        "ENDPOINT_RESOURCE_NAME = endpoint.resource_name\n",
+        "\n",
+        "# @title Chat Completions Inference\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint using the OpenAI SDK.\n",
+        "\n",
+        "# @markdown First you will need to install the SDK and some auth-related dependencies.\n",
+        "\n",
+        "! pip install -qU openai google-auth requests\n",
+        "\n",
+        "# @markdown Next fill out some request parameters:\n",
+        "\n",
+        "user_message = \"How is your day going?\"  # @param {type: \"string\"}\n",
+        "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
+        "max_tokens = 50  # @param {type: \"integer\"}\n",
+        "temperature = 1.0  # @param {type: \"number\"}\n",
+        "stream = False  # @param {type: \"boolean\"}\n",
+        "\n",
+        "# @markdown Now we can send a request.\n",
+        "\n",
+        "import google.auth\n",
+        "import openai\n",
+        "\n",
+        "creds, project = google.auth.default()\n",
+        "auth_req = google.auth.transport.requests.Request()\n",
+        "creds.refresh(auth_req)\n",
+        "\n",
+        "BASE_URL = (\n",
+        "    f\"https://{REGION}-aiplatform.googleapis.com/v1beta1/{ENDPOINT_RESOURCE_NAME}\"\n",
+        ")\n",
+        "try:\n",
+        "    if use_dedicated_endpoint:\n",
+        "        BASE_URL = f\"https://{DEDICATED_ENDPOINT_DNS}/v1beta1/{ENDPOINT_RESOURCE_NAME}\"\n",
+        "except NameError:\n",
+        "    pass\n",
+        "\n",
+        "client = openai.OpenAI(base_url=BASE_URL, api_key=creds.token)\n",
+        "\n",
+        "model_response = client.chat.completions.create(\n",
+        "    model=\"\",\n",
+        "    messages=[{\"role\": \"user\", \"content\": user_message}],\n",
+        "    temperature=temperature,\n",
+        "    max_tokens=max_tokens,\n",
+        "    stream=stream,\n",
+        ")\n",
+        "\n",
+        "if stream:\n",
+        "    usage = None\n",
+        "    contents = []\n",
+        "    for chunk in model_response:\n",
+        "        if chunk.usage is not None:\n",
+        "            usage = chunk.usage\n",
+        "            continue\n",
+        "        print(chunk.choices[0].delta.content, end=\"\")\n",
+        "        contents.append(chunk.choices[0].delta.content)\n",
+        "    print(f\"\\n\\n{usage}\")\n",
+        "else:\n",
+        "    print(model_response)\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "_0IYHNsYJO55"
+      },
+      "source": [
+        "## Deploy Gemma 3 4B, 12B and 27B multimodal models with vLLM on GPU"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "zpKbaFK_Jeny"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown Set the model to deploy.\n",
+        "\n",
+        "base_model_name = \"gemma-3-4b-it\"  # @param [\"gemma-3-4b-pt\", \"gemma-3-4b-it\", \"gemma-3-12b-pt\", \"gemma-3-12b-it\", \"gemma-3-27b-pt\", \"gemma-3-27b-it\"] {isTemplate:true}\n",
+        "hf_model_id = \"google/\" + base_model_name\n",
+        "PUBLISHER_MODEL_NAME = f\"publishers/google/models/gemma3@{base_model_name}\"\n",
+        "model_id = f\"gs://vertex-model-garden-restricted-us/gemma3/{base_model_name}\"\n",
+        "\n",
+        "# The pre-built serving docker image.\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20250312_0916_RC01\"\n",
+        "\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
+        "if \"4b\" in model_id or \"12b\" in model_id:\n",
+        "    accelerator_type = \"NVIDIA_H100_80GB\"\n",
+        "    machine_type = \"a3-highgpu-2g\"\n",
+        "    accelerator_count = 2\n",
+        "elif \"27b\" in model_id:\n",
+        "    accelerator_type = \"NVIDIA_H100_80GB\"\n",
+        "    machine_type = \"a3-highgpu-4g\"\n",
+        "    accelerator_count = 4\n",
+        "else:\n",
+        "    raise ValueError(\n",
+        "        f\"Recommended machine settings not found for model: {base_model_name}.\"\n",
+        "    )\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",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "CA_dsYyjLy1b"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 1] Deploy with Model Garden SDK\n",
+        "\n",
+        "LABEL = \"sdk-deploy\"\n",
+        "# @markdown Deploy with Gen AI model-centric SDK. This section uploads the prebuilt model to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model. See [use open models with Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-open-models) for documentation on other use cases.\n",
+        "from vertexai import model_garden\n",
+        "\n",
+        "model = model_garden.OpenModel(PUBLISHER_MODEL_NAME)\n",
+        "endpoints[LABEL] = model.deploy(\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    accept_eula=True,  # Accept the End User License Agreement (EULA) on the model card before deploy. Otherwise, the deployment will be forbidden.\n",
+        ")\n",
+        "\n",
+        "endpoint = endpoints[LABEL]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "plirtyNxJO55"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with customized configs\n",
+        "\n",
+        "# @markdown This section uploads Gemma 3 multimodal models to Model Registry and deploys them to a Vertex Prediction Endpoint. It takes 15 minutes to 1 hour to finish.\n",
+        "\n",
+        "gpu_memory_utilization = 0.95\n",
+        "max_model_len = 131072\n",
+        "\n",
+        "LABEL = \"multimodal-deploy\"\n",
+        "\n",
+        "\n",
+        "def deploy_model_vllm(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    publisher: str,\n",
+        "    publisher_model_id: 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",
+        "    enable_chunked_prefill: bool = False,\n",
+        "    enable_prefix_cache: bool = False,\n",
+        "    host_prefix_kv_cache_utilization_target: float = 0.0,\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",
+        "    enable_llama_tool_parser: bool = False,\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 enable_chunked_prefill:\n",
+        "        vllm_args.append(\"--enable-chunked-prefill\")\n",
+        "\n",
+        "    if enable_prefix_cache:\n",
+        "        vllm_args.append(\"--enable-prefix-caching\")\n",
+        "\n",
+        "    if 0 < host_prefix_kv_cache_utilization_target < 1:\n",
+        "        vllm_args.append(\n",
+        "            f\"--host-prefix-kv-cache-utilization-target={host_prefix_kv_cache_utilization_target}\"\n",
+        "        )\n",
+        "\n",
+        "    if model_type:\n",
+        "        vllm_args.append(f\"--model-type={model_type}\")\n",
+        "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\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",
+        "        model_garden_source_model_name=(\n",
+        "            f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
+        "        ),\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",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_gemma3_deployment_on_vertex.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
+        "        },\n",
+        "    )\n",
+        "    print(\"endpoint_name:\", endpoint.name)\n",
+        "\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "models[LABEL], endpoints[LABEL] = deploy_model_vllm(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=\"gemma3-serve\"),\n",
+        "    model_id=model_id,\n",
+        "    publisher=\"google\",\n",
+        "    publisher_model_id=\"gemma3\",\n",
+        "    base_model_id=hf_model_id,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    gpu_memory_utilization=gpu_memory_utilization,\n",
+        "    max_model_len=max_model_len,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        ")\n",
+        "\n",
+        "model = models[LABEL]\n",
+        "endpoint = endpoints[LABEL]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "eMBKLoUPJO55"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Chat completion with text-only requests\n",
+        "\n",
+        "if use_dedicated_endpoint:\n",
+        "    DEDICATED_ENDPOINT_DNS = endpoint.gca_resource.dedicated_endpoint_dns\n",
+        "ENDPOINT_RESOURCE_NAME = endpoint.resource_name\n",
+        "\n",
+        "# @title Chat Completions Inference\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint using the OpenAI SDK.\n",
+        "\n",
+        "# @markdown First you will need to install the SDK and some auth-related dependencies.\n",
+        "\n",
+        "! pip install -qU openai google-auth requests\n",
+        "\n",
+        "# @markdown Next fill out some request parameters:\n",
+        "\n",
+        "user_message = \"How is your day going?\"  # @param {type: \"string\"}\n",
+        "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
+        "max_tokens = 50  # @param {type: \"integer\"}\n",
+        "temperature = 1.0  # @param {type: \"number\"}\n",
+        "stream = False  # @param {type: \"boolean\"}\n",
+        "\n",
+        "# @markdown Now we can send a request.\n",
+        "\n",
+        "import google.auth\n",
+        "import openai\n",
+        "\n",
+        "creds, project = google.auth.default()\n",
+        "auth_req = google.auth.transport.requests.Request()\n",
+        "creds.refresh(auth_req)\n",
+        "\n",
+        "BASE_URL = (\n",
+        "    f\"https://{REGION}-aiplatform.googleapis.com/v1beta1/{ENDPOINT_RESOURCE_NAME}\"\n",
+        ")\n",
+        "try:\n",
+        "    if use_dedicated_endpoint:\n",
+        "        BASE_URL = f\"https://{DEDICATED_ENDPOINT_DNS}/v1beta1/{ENDPOINT_RESOURCE_NAME}\"\n",
+        "except NameError:\n",
+        "    pass\n",
+        "\n",
+        "client = openai.OpenAI(base_url=BASE_URL, api_key=creds.token)\n",
+        "\n",
+        "model_response = client.chat.completions.create(\n",
+        "    model=\"\",\n",
+        "    messages=[{\"role\": \"user\", \"content\": user_message}],\n",
+        "    temperature=temperature,\n",
+        "    max_tokens=max_tokens,\n",
+        "    stream=stream,\n",
+        ")\n",
+        "\n",
+        "if stream:\n",
+        "    usage = None\n",
+        "    contents = []\n",
+        "    for chunk in model_response:\n",
+        "        if chunk.usage is not None:\n",
+        "            usage = chunk.usage\n",
+        "            continue\n",
+        "        print(chunk.choices[0].delta.content, end=\"\")\n",
+        "        contents.append(chunk.choices[0].delta.content)\n",
+        "    print(f\"\\n\\n{usage}\")\n",
+        "else:\n",
+        "    print(model_response)\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "b3c2f6559a67"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Chat completion with multimodal requests\n",
+        "\n",
+        "if use_dedicated_endpoint:\n",
+        "    DEDICATED_ENDPOINT_DNS = endpoint.gca_resource.dedicated_endpoint_dns\n",
+        "ENDPOINT_RESOURCE_NAME = endpoint.resource_name\n",
+        "\n",
+        "# @title Chat Completions Inference\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint using the OpenAI SDK.\n",
+        "\n",
+        "# @markdown First you will need to install the SDK and some auth-related dependencies.\n",
+        "\n",
+        "! pip install -qU openai google-auth requests\n",
+        "\n",
+        "# @markdown Next fill out some request parameters:\n",
+        "\n",
+        "user_image = \"https://upload.wikimedia.org/wikipedia/commons/thumb/c/cb/The_Blue_Marble_%28remastered%29.jpg/580px-The_Blue_Marble_%28remastered%29.jpg\"  # @param {type: \"string\"}\n",
+        "user_message = \"What is in the image?\"  # @param {type: \"string\"}\n",
+        "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
+        "max_tokens = 50  # @param {type: \"integer\"}\n",
+        "temperature = 1.0  # @param {type: \"number\"}\n",
+        "\n",
+        "# @markdown Now we can send a request.\n",
+        "\n",
+        "import google.auth\n",
+        "import openai\n",
+        "\n",
+        "creds, project = google.auth.default()\n",
+        "auth_req = google.auth.transport.requests.Request()\n",
+        "creds.refresh(auth_req)\n",
+        "\n",
+        "BASE_URL = (\n",
+        "    f\"https://{REGION}-aiplatform.googleapis.com/v1beta1/{ENDPOINT_RESOURCE_NAME}\"\n",
+        ")\n",
+        "try:\n",
+        "    if use_dedicated_endpoint:\n",
+        "        BASE_URL = f\"https://{DEDICATED_ENDPOINT_DNS}/v1beta1/{ENDPOINT_RESOURCE_NAME}\"\n",
+        "except NameError:\n",
+        "    pass\n",
+        "\n",
+        "client = openai.OpenAI(base_url=BASE_URL, api_key=creds.token)\n",
+        "\n",
+        "model_response = client.chat.completions.create(\n",
+        "    model=\"\",\n",
+        "    messages=[\n",
+        "        {\n",
+        "            \"role\": \"user\",\n",
+        "            \"content\": [\n",
+        "                {\"type\": \"image_url\", \"image_url\": {\"url\": user_image}},\n",
+        "                {\"type\": \"text\", \"text\": user_message},\n",
+        "            ],\n",
+        "        }\n",
+        "    ],\n",
+        "    temperature=temperature,\n",
+        "    max_tokens=max_tokens,\n",
+        ")\n",
+        "print(model_response)\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "af21a3cff1e0"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "911406c1561e"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Delete the models and endpoints\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()"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_gemma3_deployment_on_vertex.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_gemma_deployment_on_gke.ipynb

@@ -40,7 +40,7 @@
         "  </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_deployment_on_gke.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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"

notebooks/community/model_garden/model_garden_gemma_deployment_on_vertex.ipynb

@@ -40,7 +40,7 @@
         "  </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_deployment_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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"
@@ -131,6 +131,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -193,7 +199,7 @@
         "# @markdown ---\n",
         "\n",
         "# @markdown ### Access Gemma models on Hugging Face\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",
+        "# @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.\n",
         "\n",
         "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate:true}\n",
         "if LOAD_MODEL_FROM == \"Hugging Face\":\n",
@@ -297,7 +303,7 @@
         "hbm_utilization_factor = 0.6  # A larger value improves throughput but gives higher risk of TPU out-of-memory errors with long prompts.\n",
         "max_running_seqs = 256\n",
         "\n",
-        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
+        "# @markdown Set `use_dedicated_endpoint` to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
         "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
         "\n",
         "# Endpoint configurations.\n",
@@ -310,7 +316,7 @@
         "    model_id: str,\n",
         "    publisher: str,\n",
         "    publisher_model_id: str,\n",
-        "    service_account: str,\n",
+        "    service_account: str = None,\n",
         "    base_model_id: str = None,\n",
         "    data_parallel_size: int = 1,\n",
         "    tensor_parallel_size: int = 1,\n",
@@ -319,6 +325,7 @@
         "    disagg_topology: str = None,\n",
         "    hbm_utilization_factor: float = 0.6,\n",
         "    max_running_seqs: int = 256,\n",
+        "    decode_seqs_padding: int = None,\n",
         "    max_model_len: int = 4096,\n",
         "    enable_prefix_cache_hbm: bool = False,\n",
         "    endpoint_id: str = \"\",\n",
@@ -359,6 +366,10 @@
         "        f\"--max_running_seqs={max_running_seqs}\",\n",
         "        f\"--max_model_len={max_model_len}\",\n",
         "    ]\n",
+        "\n",
+        "    if decode_seqs_padding is not None:\n",
+        "        hexllm_args.append(f\"--decode_seqs_padding={decode_seqs_padding}\")\n",
+        "\n",
         "    if disagg_topology:\n",
         "        hexllm_args.append(f\"--disagg_topo={disagg_topology}\")\n",
         "    if enable_prefix_cache_hbm and not disagg_topology:\n",
@@ -404,6 +415,7 @@
         "        max_replica_count=max_replica_count,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_gemma_deployment_on_vertex.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    return model, endpoint\n",
@@ -506,9 +518,7 @@
         "\n",
         "if use_dedicated_endpoint:\n",
         "    DEDICATED_ENDPOINT_DNS = endpoints[\"hexllm_tpu\"].gca_resource.dedicated_endpoint_dns\n",
-        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
-        "    PROJECT_ID, REGION, endpoints[\"hexllm_tpu\"].name\n",
-        ")\n",
+        "ENDPOINT_RESOURCE_NAME = endpoints[\"hexllm_tpu\"].resource_name\n",
         "\n",
         "# @title Chat Completions Inference\n",
         "\n",
@@ -524,6 +534,7 @@
         "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
         "max_tokens = 50  # @param {type: \"integer\"}\n",
         "temperature = 1.0  # @param {type: \"number\"}\n",
+        "stream = False  # @param {type: \"boolean\"}\n",
         "\n",
         "# @markdown Now we can send a request.\n",
         "\n",
@@ -550,8 +561,21 @@
         "    messages=[{\"role\": \"user\", \"content\": user_message}],\n",
         "    temperature=temperature,\n",
         "    max_tokens=max_tokens,\n",
+        "    stream=stream,\n",
         ")\n",
-        "print(model_response)\n",
+        "\n",
+        "if stream:\n",
+        "    usage = None\n",
+        "    contents = []\n",
+        "    for chunk in model_response:\n",
+        "        if chunk.usage is not None:\n",
+        "            usage = chunk.usage\n",
+        "            continue\n",
+        "        print(chunk.choices[0].delta.content, end=\"\")\n",
+        "        contents.append(chunk.choices[0].delta.content)\n",
+        "    print(f\"\\n\\n{usage}\")\n",
+        "else:\n",
+        "    print(model_response)\n",
         "\n",
         "REGION = _region\n",
         "\n",
@@ -654,7 +678,7 @@
         "# Note that a larger max_model_len will require more GPU memory.\n",
         "max_model_len = 2048\n",
         "\n",
-        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
+        "# @markdown Set `use_dedicated_endpoint` to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
         "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
         "\n",
         "\n",
@@ -682,6 +706,7 @@
         "    use_dedicated_endpoint: bool = False,\n",
         "    max_num_seqs: int = 256,\n",
         "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -734,6 +759,10 @@
         "    if model_type:\n",
         "        vllm_args.append(f\"--model-type={model_type}\")\n",
         "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -772,6 +801,7 @@
         "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_gemma_deployment_on_vertex.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",
@@ -882,9 +912,7 @@
         "\n",
         "if use_dedicated_endpoint:\n",
         "    DEDICATED_ENDPOINT_DNS = endpoints[\"vllm_gpu\"].gca_resource.dedicated_endpoint_dns\n",
-        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
-        "    PROJECT_ID, REGION, endpoints[\"vllm_gpu\"].name\n",
-        ")\n",
+        "ENDPOINT_RESOURCE_NAME = endpoints[\"vllm_gpu\"].resource_name\n",
         "\n",
         "# @title Chat Completions Inference\n",
         "\n",
@@ -900,6 +928,7 @@
         "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
         "max_tokens = 50  # @param {type: \"integer\"}\n",
         "temperature = 1.0  # @param {type: \"number\"}\n",
+        "stream = False  # @param {type: \"boolean\"}\n",
         "\n",
         "# @markdown Now we can send a request.\n",
         "\n",
@@ -926,8 +955,21 @@
         "    messages=[{\"role\": \"user\", \"content\": user_message}],\n",
         "    temperature=temperature,\n",
         "    max_tokens=max_tokens,\n",
+        "    stream=stream,\n",
         ")\n",
-        "print(model_response)\n",
+        "\n",
+        "if stream:\n",
+        "    usage = None\n",
+        "    contents = []\n",
+        "    for chunk in model_response:\n",
+        "        if chunk.usage is not None:\n",
+        "            usage = chunk.usage\n",
+        "            continue\n",
+        "        print(chunk.choices[0].delta.content, end=\"\")\n",
+        "        contents.append(chunk.choices[0].delta.content)\n",
+        "    print(f\"\\n\\n{usage}\")\n",
+        "else:\n",
+        "    print(model_response)\n",
         "\n",
         "# @markdown Click \"Show Code\" to see more details."
       ]

notebooks/community/model_garden/model_garden_gemma_evaluation.ipynb

@@ -41,7 +41,7 @@
         "  </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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"
@@ -133,6 +133,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -197,7 +203,7 @@
         "\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",
+        "# @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.\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",

notebooks/community/model_garden/model_garden_gemma_finetuning_on_vertex.ipynb

@@ -42,7 +42,7 @@
         "  </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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"
@@ -113,10 +113,9 @@
     {
       "cell_type": "code",
       "execution_count": null,
-      "language": "python",
       "metadata": {
         "cellView": "form",
-        "id": "855d6b96f291"
+        "id": "8CQcnBfWvc-f"
       },
       "outputs": [],
       "source": [
@@ -133,7 +132,7 @@
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, 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",
@@ -168,6 +167,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -232,7 +237,7 @@
         "# @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",
+        "# @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.\n",
         "\n",
         "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate:true}\n",
         "if LOAD_MODEL_FROM == \"Hugging Face\":\n",
@@ -676,10 +681,22 @@
         "# 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",
-        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
+        "# @markdown Set `use_dedicated_endpoint` to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
         "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
         "\n",
         "\n",
+        "def get_deploy_source() -> str:\n",
+        "    \"\"\"Gets deploy_source string based on running environment.\"\"\"\n",
+        "    vertex_product = os.environ.get(\"VERTEX_PRODUCT\", \"\")\n",
+        "    if vertex_product == \"COLAB_ENTERPRISE\":\n",
+        "        return \"notebook_colab_enterprise\"\n",
+        "    elif vertex_product == \"WORKBENCH_INSTANCE\":\n",
+        "        return \"notebook_workbench\"\n",
+        "    else:\n",
+        "        # Legacy workbench, legacy colab, or other custom environments.\n",
+        "        return \"notebook_environment_unspecified\"\n",
+        "\n",
+        "\n",
         "def deploy_model_vllm(\n",
         "    model_name: str,\n",
         "    model_id: str,\n",
@@ -704,6 +721,7 @@
         "    use_dedicated_endpoint: bool = False,\n",
         "    max_num_seqs: int = 256,\n",
         "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -756,6 +774,10 @@
         "    if model_type:\n",
         "        vllm_args.append(f\"--model-type={model_type}\")\n",
         "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -794,6 +816,7 @@
         "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_gemma_finetuning_on_vertex.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",
@@ -1216,7 +1239,7 @@
         "# 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:20241210_2323_RC00\"\n",
         "\n",
-        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
+        "# @markdown Set `use_dedicated_endpoint` to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
         "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
         "\n",
         "\n",
@@ -1225,7 +1248,7 @@
         "    model_id: str,\n",
         "    publisher: str,\n",
         "    publisher_model_id: str,\n",
-        "    service_account: str,\n",
+        "    service_account: str = None,\n",
         "    base_model_id: str = None,\n",
         "    data_parallel_size: int = 1,\n",
         "    tensor_parallel_size: int = 1,\n",
@@ -1319,6 +1342,7 @@
         "        max_replica_count=max_replica_count,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_gemma_finetuning_on_vertex.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    return model, endpoint\n",

notebooks/community/model_garden/model_garden_gradio_streaming_chat_completions.ipynb

@@ -40,7 +40,7 @@
         "  </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_gradio_streaming_chat_completions.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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"
@@ -208,6 +208,7 @@
         "    use_dedicated_endpoint: bool = False,\n",
         "    max_num_seqs: int = 256,\n",
         "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -260,6 +261,10 @@
         "    if model_type:\n",
         "        vllm_args.append(f\"--model-type={model_type}\")\n",
         "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -298,6 +303,7 @@
         "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_gradio_streaming_chat_completions.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",
@@ -310,7 +316,7 @@
         "    model_id: str,\n",
         "    publisher: str,\n",
         "    publisher_model_id: str,\n",
-        "    service_account: str,\n",
+        "    service_account: str = None,\n",
         "    machine_type: str = \"g2-standard-8\",\n",
         "    accelerator_type: str = \"NVIDIA_L4\",\n",
         "    accelerator_count: int = 1,\n",
@@ -341,6 +347,9 @@
         "    except NameError:\n",
         "        pass\n",
         "\n",
+        "    if service_account:\n",
+        "        env_vars[\"SERVICE_ACCOUNT\"] = service_account\n",
+        "\n",
         "    model = aiplatform.Model.upload(\n",
         "        display_name=model_name,\n",
         "        serving_container_image_uri=TGI_DOCKER_URI,\n",
@@ -361,6 +370,7 @@
         "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_gradio_streaming_chat_completions.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    return model, endpoint\n",

notebooks/community/model_garden/model_garden_hf_paligemma2_deployment.ipynb

@@ -34,13 +34,18 @@
         "\n",
         "<table><tbody><tr>\n",
         "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/instances\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_hf_paligemma2_deployment.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_hf_paligemma2_deployment.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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"
@@ -63,6 +68,10 @@
         "- Make predictions to the endpoint including:\n",
         "  - Answering questions about a given image.\n",
         "\n",
+        "### File a bug\n",
+        "\n",
+        "File a bug on [GitHub](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues/new) if you encounter any issue with the notebook.\n",
+        "\n",
         "### Costs\n",
         "\n",
         "This tutorial uses billable components of Google Cloud:\n",
@@ -93,8 +102,23 @@
       "source": [
         "# @title Setup Google Cloud project\n",
         "\n",
-        "# Used for common utilities.\n",
-        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\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]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
+        "\n",
+        "REGION = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# @markdown 3. 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). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\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-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
+        "\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform==1.93.1'\n",
         "\n",
         "import importlib\n",
         "# Import the necessary packages\n",
@@ -103,20 +127,16 @@
         "\n",
         "from google.cloud import aiplatform\n",
         "\n",
-        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
+        "    ! pip install --upgrade tensorflow\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
         "\n",
-        "# @markdown 2. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
+        "common_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
+        ")\n",
         "\n",
-        "REGION = \"\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# @markdown 3. 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-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "models, endpoints = {}, {}\n",
+        "LABEL = \"paligemma2\"\n",
         "\n",
         "\n",
         "# Get the default cloud project id.\n",
@@ -126,130 +146,17 @@
         "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",
         "# Initialize Vertex AI API.\n",
         "print(\"Initializing Vertex AI API.\")\n",
         "aiplatform.init(project=PROJECT_ID, location=REGION)\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",
         "! gcloud config set project $PROJECT_ID\n",
+        "import vertexai\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",
-        "# The pre-built serving docker images.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-one-serve:20250205_0822_RC00\"\n",
-        "\n",
-        "\n",
-        "def deploy_model(\n",
-        "    model_id: str = None,\n",
-        "    task: str = None,\n",
-        "    machine_type: str = \"g2-standard-8\",\n",
-        "    accelerator_type: str = \"NVIDIA_L4\",\n",
-        "    accelerator_count: int = 1,\n",
-        "    service_account: str = None,\n",
-        "    serving_port: int = 7080,\n",
-        "    serving_route: str = \"/predict\",\n",
-        "    serving_docker_uri: str = SERVE_DOCKER_URI,\n",
-        ") -> Tuple[aiplatform.Endpoint, aiplatform.Model]:\n",
-        "    \"\"\"Deploys a model to a real-time prediction endpoint.\n",
-        "\n",
-        "    Args:\n",
-        "        model_id: The model ID.\n",
-        "        task: The task to perform.\n",
-        "        machine_type: The machine type.\n",
-        "        accelerator_type: The accelerator type.\n",
-        "        accelerator_count: The accelerator count.\n",
-        "        service_account: The service account.\n",
-        "        serving_port: The serving port.\n",
-        "        serving_route: The serving route.\n",
-        "        hf_token: HuggingFace token for model access.\n",
-        "\n",
-        "    Returns:\n",
-        "        A tuple containing the created endpoint and deployed model objects.\n",
-        "    \"\"\"\n",
-        "\n",
-        "    endpoint = aiplatform.Endpoint.create(\n",
-        "        display_name=common_util.get_job_name_with_datetime(prefix=\"paligemma-2\")\n",
-        "    )\n",
-        "    serving_env = {\n",
-        "        \"MODEL_ID\": model_id,\n",
-        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
-        "        \"TASK\": task,\n",
-        "    }\n",
-        "    model = aiplatform.Model.upload(\n",
-        "        display_name=task,\n",
-        "        serving_container_image_uri=serving_docker_uri,\n",
-        "        serving_container_ports=[serving_port],\n",
-        "        serving_container_predict_route=serving_route,\n",
-        "        serving_container_health_route=\"/ping\",\n",
-        "        serving_container_environment_variables=serving_env,\n",
-        "        model_garden_source_model_name=\"publishers/google/models/paligemma\",\n",
-        "    )\n",
-        "    model.deploy(\n",
-        "        endpoint=endpoint,\n",
-        "        machine_type=machine_type,\n",
-        "        accelerator_type=accelerator_type,\n",
-        "        accelerator_count=accelerator_count,\n",
-        "        service_account=service_account,\n",
-        "        sync=False,\n",
-        "        system_labels={\n",
-        "            \"NOTEBOOK_NAME\": (\"model_garden_hf_paligemma2_deployment_templated.ipynb\")\n",
-        "        },\n",
-        "    )\n",
-        "    return endpoint, model\n",
-        "\n",
-        "\n",
-        "def vqa_predict(\n",
-        "    endpoint: aiplatform.Endpoint,\n",
-        "    image_url: str,\n",
-        "    text_prompt: str,\n",
-        "    parameters: Dict[str, Any] = None,\n",
-        ") -> str:\n",
-        "    \"\"\"Predicts the answer to a question about an image using an Endpoint,\n",
-        "\n",
-        "       and passes parameters in the payload.\n",
-        "\n",
-        "    Args:\n",
-        "        endpoint: The deployed Vertex AI endpoint.\n",
-        "        image_url: URL of the image to ask about.\n",
-        "        text_prompt: The text prompt question.\n",
-        "        parameters: Additional parameters for the prediction request.\n",
-        "\n",
-        "    Returns:\n",
-        "        The predicted answer string or None if no prediction.\n",
-        "    \"\"\"\n",
-        "\n",
-        "    instances = []\n",
-        "    if text_prompt:\n",
-        "        instances.append(\n",
-        "            {\n",
-        "                \"text_prompt\": text_prompt,\n",
-        "                \"image_url\": image_url,\n",
-        "            }\n",
-        "        )\n",
-        "\n",
-        "    # Construct the prediction payload\n",
-        "    payload = {\"instances\": instances}\n",
-        "    if parameters:\n",
-        "        payload[\"parameters\"] = parameters\n",
-        "\n",
-        "    response = endpoint.predict(instances=instances, parameters=parameters)\n",
-        "    answer = None\n",
-        "    if response.predictions:\n",
-        "        answer = response.predictions[0][\"text\"].split(\"\\n\")[1]\n",
-        "    return answer"
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")"
       ]
     },
     {
@@ -272,24 +179,30 @@
       "source": [
         "# @title Deploy\n",
         "\n",
-        "# @markdown This section uploads the prebuilt PaliGemma 2 models to Model Registry and deploys it to a Vertex AI Endpoint. It takes approximately 15 minutes to finish.\n",
+        "MODEL_NAME = \"paligemma2-3b-pt-224\"  # @param [\"paligemma2-3b-pt-224\", \"paligemma2-3b-mix-224\", \"paligemma2-3b-ft-docci-448\", \"paligemma2-3b-mix-448\", \"paligemma2-3b-pt-448\", \"paligemma2-3b-pt-896\", \"paligemma2-10b-mix-224\", \"paligemma2-10b-pt-224\", \"paligemma2-10b-ft-docci-448\", \"paligemma2-10b-mix-448\", \"paligemma2-10b-pt-448\", \"paligemma2-10b-pt-896\", \"paligemma2-28b-mix-224\", \"paligemma2-28b-pt-224\", \"paligemma2-28b-mix-448\", \"paligemma2-28b-pt-448\", \"paligemma2-28b-pt-896\"]\n",
+        "GCS_PREFIX = \"gs://vertex-model-garden-restricted-us/paligemma2\"\n",
         "\n",
-        "# @markdown Select the desired resolution and precision of prebuilt model to deploy, leaving the optional `custom_paligemma_model_uri` as is. Higher resolution and precision_type can result in better inference results, but may require additional GPU.\n",
+        "MODEL_ID = os.path.join(GCS_PREFIX, MODEL_NAME)\n",
         "\n",
-        "\n",
-        "MODEL_ID = \"paligemma2-3b-pt-224\"  # @param [\"paligemma2-3b-pt-224\", \"paligemma2-3b-pt-448\", \"paligemma2-10b-ft-docci-448\"]\n",
-        "\n",
-        "GSC_PREFIX = \"gs://vertex-model-garden-public-us/paligemma2/\"\n",
-        "MODEL_ID = GSC_PREFIX + MODEL_ID\n",
-        "\n",
-        "TASK = \"paligemma_VQA\"  # @param [\"paligemma_VQA\"]\n",
-        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\"]\n",
-        "accelerator_count = 1  # @param [1]\n",
-        "machine_type = \"g2-standard-8\"  # @param  [\"g2-standard-8\"]\n",
+        "PUBLISHER_MODEL_NAME = f\"publishers/google/models/paligemma@{MODEL_NAME}\"\n",
         "\n",
         "\n",
         "# @markdown If you want to use other accelerator types not listed above, then check other Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute. You may need to manually set the `machine_type`, `accelerator_type`, and `accelerator_count` in the code by clicking `Show code` first.\n",
         "\n",
+        "if \"3b\" in MODEL_NAME:\n",
+        "    accelerator_type = \"NVIDIA_L4\"\n",
+        "    machine_type = \"g2-standard-16\"\n",
+        "    accelerator_count = 1\n",
+        "elif \"10b\" in MODEL_NAME:\n",
+        "    accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+        "    machine_type = \"a2-highgpu-1g\"\n",
+        "    accelerator_count = 1\n",
+        "elif \"28b\" in MODEL_NAME:\n",
+        "    accelerator_type = \"NVIDIA_H100_80GB\"\n",
+        "    machine_type = \"a3-highgpu-8g\"\n",
+        "    accelerator_count = 8\n",
+        "else:\n",
+        "    raise ValueError(f\"Recommended GPU setting not found for: {MODEL_NAME}.\")\n",
         "\n",
         "common_util.check_quota(\n",
         "    project_id=PROJECT_ID,\n",
@@ -299,13 +212,123 @@
         "    is_for_training=False,\n",
         ")\n",
         "\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "pe_qbTCA6nKf"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 1] Deploy with Model Garden SDK\n",
+        "\n",
+        "# @markdown Deploy with Gen AI model-centric SDK. This section uploads the prebuilt model to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model. See [use open models with Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-open-models) for documentation on other use cases.\n",
+        "from vertexai import model_garden\n",
+        "\n",
+        "model = model_garden.OpenModel(PUBLISHER_MODEL_NAME)\n",
+        "endpoints[LABEL] = model.deploy(\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    accept_eula=True,  # Accept the End User License Agreement (EULA) on the model card before deploy. Otherwise, the deployment will be forbidden.\n",
+        ")\n",
+        "\n",
+        "endpoint = endpoints[LABEL]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "jbeLl-9C6nKf"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with customized configs\n",
+        "\n",
+        "# @markdown This section uploads the prebuilt PaliGemma 2 models to Model Registry and deploys it to a Vertex AI Endpoint. It takes approximately 15 minutes to finish.\n",
+        "\n",
+        "# @markdown Select the desired resolution and precision of prebuilt model to deploy, leaving the optional `custom_paligemma_model_uri` as is. Higher resolution and precision_type can result in better inference results, but may require additional GPU.\n",
+        "\n",
+        "TASK = \"paligemma_VQA\"\n",
+        "\n",
+        "# The pre-built serving docker images.\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-one-serve:20250205_0822_RC00\"\n",
+        "\n",
+        "\n",
+        "def deploy_model(\n",
+        "    model_name: str = None,\n",
+        "    model_id: str = None,\n",
+        "    task: str = None,\n",
+        "    machine_type: str = \"g2-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 1,\n",
+        "    serving_port: int = 7080,\n",
+        "    serving_route: str = \"/predict\",\n",
+        "    serving_docker_uri: str = SERVE_DOCKER_URI,\n",
+        ") -> Tuple[aiplatform.Endpoint, aiplatform.Model]:\n",
+        "    \"\"\"Deploys a model to a real-time prediction endpoint.\n",
+        "\n",
+        "    Args:\n",
+        "        model_name: The base name of the model.\n",
+        "        model_id: The model ID.\n",
+        "        task: The task to perform.\n",
+        "        machine_type: The machine type.\n",
+        "        accelerator_type: The accelerator type.\n",
+        "        accelerator_count: The accelerator count.\n",
+        "        serving_port: The serving port.\n",
+        "        serving_route: The serving route.\n",
+        "        hf_token: HuggingFace token for model access.\n",
+        "\n",
+        "    Returns:\n",
+        "        A tuple containing the created endpoint and deployed model objects.\n",
+        "    \"\"\"\n",
+        "\n",
+        "    endpoint = aiplatform.Endpoint.create(\n",
+        "        display_name=common_util.get_job_name_with_datetime(prefix=model_name)\n",
+        "    )\n",
+        "    serving_env = {\n",
+        "        \"MODEL_ID\": model_id,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "        \"TASK\": task,\n",
+        "    }\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=task,\n",
+        "        serving_container_image_uri=serving_docker_uri,\n",
+        "        serving_container_ports=[serving_port],\n",
+        "        serving_container_predict_route=serving_route,\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        serving_container_environment_variables=serving_env,\n",
+        "        model_garden_source_model_name=\"publishers/google/models/paligemma\",\n",
+        "    )\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        sync=False,\n",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_hf_paligemma2_deployment.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
+        "        },\n",
+        "    )\n",
+        "    return endpoint, model\n",
+        "\n",
+        "\n",
         "endpoints[\"paligemma2\"], models[\"paligemma2\"] = deploy_model(\n",
+        "    model_name=MODEL_NAME,\n",
         "    model_id=MODEL_ID,\n",
         "    task=TASK,\n",
         "    machine_type=machine_type,\n",
         "    accelerator_type=accelerator_type,\n",
         "    accelerator_count=accelerator_count,\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
         "    serving_port=7080,\n",
         "    serving_route=\"/predict\",\n",
         "    serving_docker_uri=SERVE_DOCKER_URI,\n",
@@ -364,10 +387,52 @@
         "image_url = \"https://images.pexels.com/photos/1006293/pexels-photo-1006293.jpeg\"  # @param {type:\"string\"}\n",
         "\n",
         "# @markdown You may leave question prompts empty and they will be ignored.\n",
-        "question_prompt = \"What animal is shown in the picture?\"  # @param {type: \"string\"}\n",
+        "question_prompt = \"What is shown in the picture?\"  # @param {type: \"string\"}\n",
         "\n",
         "# @markdown The question prompt can be non-English languages.\n",
         "\n",
+        "\n",
+        "def vqa_predict(\n",
+        "    endpoint: aiplatform.Endpoint,\n",
+        "    image_url: str,\n",
+        "    text_prompt: str,\n",
+        "    parameters: Dict[str, Any] = None,\n",
+        ") -> str:\n",
+        "    \"\"\"Predicts the answer to a question about an image using an Endpoint,\n",
+        "\n",
+        "       and passes parameters in the payload.\n",
+        "\n",
+        "    Args:\n",
+        "        endpoint: The deployed Vertex AI endpoint.\n",
+        "        image_url: URL of the image to ask about.\n",
+        "        text_prompt: The text prompt question.\n",
+        "        parameters: Additional parameters for the prediction request.\n",
+        "\n",
+        "    Returns:\n",
+        "        The predicted answer string or None if no prediction.\n",
+        "    \"\"\"\n",
+        "\n",
+        "    instances = []\n",
+        "    if text_prompt:\n",
+        "        instances.append(\n",
+        "            {\n",
+        "                \"text_prompt\": text_prompt,\n",
+        "                \"image_url\": image_url,\n",
+        "            }\n",
+        "        )\n",
+        "\n",
+        "    # Construct the prediction payload\n",
+        "    payload = {\"instances\": instances}\n",
+        "    if parameters:\n",
+        "        payload[\"parameters\"] = parameters\n",
+        "\n",
+        "    response = endpoint.predict(instances=instances, parameters=parameters)\n",
+        "    answer = None\n",
+        "    if response.predictions:\n",
+        "        answer = response.predictions[0][\"text\"].split(\"\\n\")[1]\n",
+        "    return answer\n",
+        "\n",
+        "\n",
         "# Using max_new_tokens along with other parameters\n",
         "parameters_with_tokens = {\"max_new_tokens\": 50}\n",
         "predictions_with_tokens = vqa_predict(\n",

notebooks/community/model_garden/model_garden_huggingface_local_inference.ipynb

@@ -40,7 +40,7 @@
         "  </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_huggingface_local_inference.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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"

notebooks/community/model_garden/model_garden_huggingface_pytorch_inference_deployment.ipynb

@@ -40,7 +40,7 @@
         "  </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_huggingface_pytorch_inference_deployment.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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"
@@ -62,6 +62,10 @@
         "- Download and deploy the `distilbert/distilbert-base-uncased-finetuned-sst-2-english` model with Hugging Face Pytorch Inference\n",
         "- Send prediction request to the deployed endpoint\n",
         "\n",
+        "### File a bug\n",
+        "\n",
+        "File a bug on [GitHub](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues/new) if you encounter any issue with the notebook.\n",
+        "\n",
         "### Costs\n",
         "\n",
         "This tutorial uses billable components of Google Cloud:\n",
@@ -97,6 +101,14 @@
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
+        "# @markdown 3. 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). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\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-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "# Upgrade Vertex AI SDK.\n",
         "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.64.0'\n",
         "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
@@ -111,6 +123,7 @@
         "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
         ")\n",
         "\n",
+        "\n",
         "# Get the default cloud project id.\n",
         "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
         "\n",
@@ -118,20 +131,17 @@
         "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",
         "# Initialize Vertex AI API.\n",
         "print(\"Initializing Vertex AI API.\")\n",
         "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
+        "\n",
         "! gcloud config set project $PROJECT_ID\n",
         "\n",
-        "models, endpoints = {}, {}\n",
-        "\n",
         "HF_TOKEN = \"\"\n",
-        "# Dedicated endpoint not supported yet\n",
-        "use_dedicated_endpoint = False\n",
+        "\n",
+        "# @markdown Set `use_dedicated_endpoint` to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
+        "\n",
         "SERVICE_ACCOUNT = \"\""
       ]
     },
@@ -172,6 +182,7 @@
         "        is_for_training=False,\n",
         "    )\n",
         "\n",
+        "models, endpoints = {}, {}\n",
         "\n",
         "GCS_PREFIX = \"gs://\"\n",
         "\n",

notebooks/community/model_garden/model_garden_huggingface_tei_deployment.ipynb

@@ -34,13 +34,18 @@
         "\n",
         "<table><tbody><tr>\n",
         "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/instances\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_huggingface_tei_deployment.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_huggingface_tei_deployment.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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"
@@ -62,6 +67,10 @@
         "- Download and deploy the `nomic-ai/nomic-embed-text-v1` model with TEI\n",
         "- Send prediction request to the deployed endpoint\n",
         "\n",
+        "### File a bug\n",
+        "\n",
+        "File a bug on [GitHub](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues/new) if you encounter any issue with the notebook.\n",
+        "\n",
         "### Costs\n",
         "\n",
         "This tutorial uses billable components of Google Cloud:\n",
@@ -97,8 +106,7 @@
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
-        "# Upgrade Vertex AI SDK.\n",
-        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.64.0'\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform==1.93.1'\n",
         "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
         "\n",
         "import importlib\n",
@@ -116,6 +124,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -129,24 +143,16 @@
         "\n",
         "models, endpoints = {}, {}\n",
         "\n",
-        "HF_TOKEN = \"\""
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "cellView": "form",
-        "id": "USB7dvYqvNdu"
-      },
-      "outputs": [],
-      "source": [
-        "# @title Deploy with TEI from Hugging Face\n",
+        "import vertexai\n",
         "\n",
-        "# @markdown This section downloads the `nomic-ai/nomic-embed-text-v1` model from Hugging Face and deploys it to a Vertex AI Endpoint.\n",
-        "# @markdown It takes ~20 minutes to complete the deployment.\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")\n",
         "\n",
-        "MODEL_ID = \"nomic-ai/nomic-embed-text-v1\"  # @param {type: \"string\", isTemplate: true}\n",
+        "HF_TOKEN = \"\"\n",
+        "\n",
+        "HUGGING_FACE_MODEL_ID = \"nomic-ai/nomic-embed-text-v1\"  # @param {type: \"string\", isTemplate: true}\n",
         "\n",
         "# The pre-built serving docker images for TEI.\n",
         "TEI_CPU_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-embeddings-inference-cpu.1-4\"\n",
@@ -167,6 +173,55 @@
         "        is_for_training=False,\n",
         "    )\n",
         "\n",
+        "LABEL = \"tei\"\n",
+        "\n",
+        "# @markdown Set `use_dedicated_endpoint` to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "obbeTtMJ5C8j"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 1] Deploy with Model Garden SDK\n",
+        "\n",
+        "accelerator_count = 1\n",
+        "\n",
+        "# @markdown Deploy with Gen AI model-centric SDK. This section uploads the prebuilt model to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model. See [use open models with Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-open-models) for documentation on other use cases.\n",
+        "from vertexai import model_garden\n",
+        "\n",
+        "model = model_garden.OpenModel(HUGGING_FACE_MODEL_ID)\n",
+        "endpoints[LABEL] = model.deploy(\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    hugging_face_access_token=HF_TOKEN,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    accept_eula=True,  # Accept the End User License Agreement (EULA) on the model card before deploy. Otherwise, the deployment will be forbidden.\n",
+        ")\n",
+        "\n",
+        "endpoint = endpoints[LABEL]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "USB7dvYqvNdu"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with customized configs\n",
+        "\n",
+        "# @markdown This section downloads the `nomic-ai/nomic-embed-text-v1` model from Hugging Face and deploys it to a Vertex AI Endpoint.\n",
+        "# @markdown It takes ~20 minutes to complete the deployment.\n",
+        "\n",
         "\n",
         "def deploy_model_tei(\n",
         "    model_name: str,\n",
@@ -223,13 +278,9 @@
         "    return model, endpoint\n",
         "\n",
         "\n",
-        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
-        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
-        "\n",
-        "\n",
         "models[\"tei\"], endpoints[\"tei\"] = deploy_model_tei(\n",
-        "    model_name=common_util.get_job_name_with_datetime(prefix=MODEL_ID),\n",
-        "    model_id=MODEL_ID,\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=HUGGING_FACE_MODEL_ID),\n",
+        "    model_id=HUGGING_FACE_MODEL_ID,\n",
         "    publisher=\"hf-nomic-ai\",\n",
         "    publisher_model_id=\"nomic-embed-text-v1\",\n",
         "    service_account=\"\",\n",

notebooks/community/model_garden/model_garden_huggingface_tgi_deployment.ipynb

@@ -34,13 +34,18 @@
         "\n",
         "<table><tbody><tr>\n",
         "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/instances\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_huggingface_tgi_deployment.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_huggingface_tgi_deployment.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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"
@@ -97,13 +102,22 @@
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
-        "# @markdown 3. You must agree to the license on the [model card](https://huggingface.co/google/gemma-2-2b-it) before accessing the Gemma 2 models.\n",
+        "# @markdown 3. 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). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
         "\n",
-        "# @markdown 4. 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",
+        "# @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-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform==1.93.1'\n",
+        "\n",
+        "# @markdown 5. You must agree to the license on the [model card](https://huggingface.co/google/gemma-2-2b-it) before accessing the Gemma 2 models.\n",
+        "\n",
+        "# @markdown 6. 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",
         "# Upgrade Vertex AI SDK.\n",
         "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.64.0'\n",
-        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
         "\n",
         "import importlib\n",
         "import os\n",
@@ -111,6 +125,10 @@
         "\n",
         "from google.cloud import aiplatform\n",
         "\n",
+        "if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
+        "    ! pip install --upgrade tensorflow\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
         "common_util = importlib.import_module(\n",
         "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
         ")\n",
@@ -121,6 +139,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -134,32 +158,18 @@
         "\n",
         "models, endpoints = {}, {}\n",
         "\n",
-        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate: true}"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "cellView": "form",
-        "id": "USB7dvYqvNdu"
-      },
-      "outputs": [],
-      "source": [
-        "# @title Deploy with TGI from Hugging Face\n",
+        "import vertexai\n",
         "\n",
-        "# @markdown This section downloads the `Gemma-2-2b-it` model from Hugging Face and deploys it to a Vertex AI Endpoint.\n",
-        "# @markdown It takes ~20 minutes to complete the deployment.\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")\n",
         "\n",
-        "MODEL_ID = \"google/gemma-2-2b-it\"  # @param {type: \"string\", isTemplate: true}\n",
+        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate: true}\n",
         "\n",
-        "# The pre-built serving docker image for TGI.\n",
-        "TGI_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu124.2-4.ubuntu2204.py311\"\n",
+        "HUGGING_FACE_MODEL_ID = \"google/gemma-2-2b-it\"  # @param {type: \"string\", isTemplate: true}\n",
         "\n",
-        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
-        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
-        "\n",
-        "machine_type = \"g2-standard-8\"  # @param {type: \"string\", isTemplate: true}\n",
+        "machine_type = \"g2-standard-12\"  # @param {type: \"string\", isTemplate: true}\n",
         "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\"] {isTemplate: true}\n",
         "accelerator_count = 1  # @param {type: \"integer\", isTemplate: true}\n",
         "\n",
@@ -171,13 +181,64 @@
         "    is_for_training=False,\n",
         ")\n",
         "\n",
+        "# The pre-built serving docker image for TGI.\n",
+        "TGI_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu124.2-4.ubuntu2204.py311\"\n",
+        "SERVING_CONTAINER_IMAGE_URI = TGI_DOCKER_URI\n",
+        "LABEL = \"tgi\"\n",
+        "\n",
+        "# @markdown Set `use_dedicated_endpoint` to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "LSkIcKP_FrkO"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 1] Deploy with Model Garden SDK\n",
+        "\n",
+        "# @markdown Deploy with Gen AI model-centric SDK. This section uploads the prebuilt model to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model. See [use open models with Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-open-models) for documentation on other use cases.\n",
+        "from vertexai import model_garden\n",
+        "\n",
+        "model = model_garden.OpenModel(HUGGING_FACE_MODEL_ID)\n",
+        "endpoints[LABEL] = model.deploy(\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    hugging_face_access_token=HF_TOKEN,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    accept_eula=True,  # Accept the End User License Agreement (EULA) on the model card before deploy. Otherwise, the deployment will be forbidden.\n",
+        ")\n",
+        "\n",
+        "endpoint = endpoints[LABEL]\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "USB7dvYqvNdu"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with customized configs\n",
+        "\n",
+        "# @markdown This section downloads the `Gemma-2-2b-it` model from Hugging Face and deploys it to a Vertex AI Endpoint.\n",
+        "# @markdown It takes ~20 minutes to complete the deployment.\n",
+        "\n",
         "\n",
         "def deploy_model_tgi(\n",
         "    model_name: str,\n",
         "    model_id: str,\n",
         "    publisher: str,\n",
         "    publisher_model_id: str,\n",
-        "    service_account: str,\n",
+        "    service_account: str = None,\n",
         "    machine_type: str = \"g2-standard-8\",\n",
         "    accelerator_type: str = \"NVIDIA_L4\",\n",
         "    accelerator_count: int = 1,\n",
@@ -208,6 +269,9 @@
         "    except NameError:\n",
         "        pass\n",
         "\n",
+        "    if service_account:\n",
+        "        env_vars[\"SERVICE_ACCOUNT\"] = service_account\n",
+        "\n",
         "    model = aiplatform.Model.upload(\n",
         "        display_name=model_name,\n",
         "        serving_container_image_uri=TGI_DOCKER_URI,\n",
@@ -228,14 +292,15 @@
         "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_huggingface_tgi_deployment.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    return model, endpoint\n",
         "\n",
         "\n",
         "models[\"tgi\"], endpoints[\"tgi\"] = deploy_model_tgi(\n",
-        "    model_name=common_util.get_job_name_with_datetime(prefix=MODEL_ID),\n",
-        "    model_id=MODEL_ID,\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=HUGGING_FACE_MODEL_ID),\n",
+        "    model_id=HUGGING_FACE_MODEL_ID,\n",
         "    publisher=\"google\",\n",
         "    publisher_model_id=\"gemma2\",\n",
         "    service_account=\"\",\n",

notebooks/community/model_garden/model_garden_huggingface_vllm_deployment.ipynb

@@ -0,0 +1,608 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "qZCBJ69YYS0S"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2025 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": "08f4AuF5eXzO"
+      },
+      "source": [
+        "# Vertex AI Model Garden - Hugging Face Deployment with vLLM Container\n",
+        "\n",
+        "<table><tbody><tr>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/instances\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_huggingface_vllm_deployment.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_huggingface_vllm_deployment.ipynb\">\n",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "j0yLwcdReXzO"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates deploying [qwen/qwq-32b](https://huggingface.co/Qwen/QwQ-32B) model with vLLM container from Hugging Face. In additional to `qwen/qwq-32b`, You can view and change the code to deploy a different Hugging Face model with appropriate machine specs.\n",
+        "\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Download and deploy the `qwen/qwq-32b` model with vLLM container.\n",
+        "- Send prediction request to the deployed endpoint.\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/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": "wkhp_C5leXzO"
+      },
+      "source": [
+        "## Setup"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "D8wjhOsKeXzO"
+      },
+      "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. **[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",
+        "# @markdown 3. 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). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\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-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
+        "\n",
+        "# @markdown 4. 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",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform==1.93.1'\n",
+        "\n",
+        "import importlib\n",
+        "import os\n",
+        "from typing import Tuple\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
+        "    ! pip install --upgrade tensorflow\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\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",
+        "\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",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\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",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "\n",
+        "models, endpoints = {}, {}\n",
+        "\n",
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")\n",
+        "\n",
+        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate: true}\n",
+        "\n",
+        "# The pre-built vLLM serving docker image.\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20250506_0916_RC01\"\n",
+        "SERVING_CONTAINER_IMAGE_URI = VLLM_DOCKER_URI\n",
+        "LABEL = \"vllm\"\n",
+        "\n",
+        "MACHINE_SPEC_MAP = {\n",
+        "    \"2 NVIDIA_H100_80GB a3-highgpu-2g\": {\n",
+        "        \"machine_type\": \"a3-highgpu-2g\",\n",
+        "        \"accelerator_type\": \"NVIDIA_H100_80GB\",\n",
+        "        \"accelerator_count\": 2,\n",
+        "    },\n",
+        "    \"4 NVIDIA_L4 g2-standard-48\": {\n",
+        "        \"machine_type\": \"g2-standard-48\",\n",
+        "        \"accelerator_type\": \"NVIDIA_L4\",\n",
+        "        \"accelerator_count\": 4,\n",
+        "    },\n",
+        "}\n",
+        "\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9j0VbjSNwGJM"
+      },
+      "source": [
+        "## Deploy and predict"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "e8-PBYQ7wGJM"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 1] Deploy with Model Garden SDK\n",
+        "\n",
+        "# @markdown Deploy with Gen AI model-centric SDK. This section uploads the prebuilt model to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model. See [use open models with Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-open-models) for documentation on other use cases.\n",
+        "\n",
+        "HUGGING_FACE_MODEL_ID = \"qwen/qwq-32b\"  # @param {type: \"string\", isTemplate: true}\n",
+        "DEPLOY_CONFIG = \"2 NVIDIA_H100_80GB a3-highgpu-2g\"  # @param [\"2 NVIDIA_H100_80GB a3-highgpu-2g\", \"4 NVIDIA_L4 g2-standard-48\",] {isTemplate: true}\n",
+        "\n",
+        "machine_type = MACHINE_SPEC_MAP[DEPLOY_CONFIG][\"machine_type\"]\n",
+        "accelerator_type = MACHINE_SPEC_MAP[DEPLOY_CONFIG][\"accelerator_type\"]\n",
+        "accelerator_count = MACHINE_SPEC_MAP[DEPLOY_CONFIG][\"accelerator_count\"]\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",
+        "from vertexai import model_garden\n",
+        "\n",
+        "model = model_garden.OpenModel(HUGGING_FACE_MODEL_ID)\n",
+        "endpoints[LABEL] = model.deploy(\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    hugging_face_access_token=HF_TOKEN,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    accept_eula=True,  # Accept the End User License Agreement (EULA) on the model card before deploy. Otherwise, the deployment will be forbidden.\n",
+        ")\n",
+        "\n",
+        "endpoint = endpoints[LABEL]\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "qbNI0WkVeXzO"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with customized configs\n",
+        "\n",
+        "# @markdown This section downloads the `qwen/qwq-32b` model from Hugging Face and deploys it to a Vertex AI Endpoint.\n",
+        "# @markdown It takes ~20 minutes to complete the deployment.\n",
+        "\n",
+        "HUGGING_FACE_MODEL_ID = \"qwen/qwq-32b\"  # @param {type: \"string\", isTemplate: true}\n",
+        "DEPLOY_CONFIG = \"2 NVIDIA_H100_80GB a3-highgpu-2g\"  # @param [\"2 NVIDIA_H100_80GB a3-highgpu-2g\", \"4 NVIDIA_L4 g2-standard-48\"] {isTemplate: true}\n",
+        "\n",
+        "machine_type = MACHINE_SPEC_MAP[DEPLOY_CONFIG][\"machine_type\"]\n",
+        "accelerator_type = MACHINE_SPEC_MAP[DEPLOY_CONFIG][\"accelerator_type\"]\n",
+        "accelerator_count = MACHINE_SPEC_MAP[DEPLOY_CONFIG][\"accelerator_count\"]\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",
+        "\n",
+        "def deploy_model_vllm(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    publisher: str,\n",
+        "    publisher_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",
+        "    enable_chunked_prefill: bool = False,\n",
+        "    enable_prefix_cache: bool = False,\n",
+        "    host_prefix_kv_cache_utilization_target: float = 0.0,\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",
+        "    enable_llama_tool_parser: bool = False,\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 enable_chunked_prefill:\n",
+        "        vllm_args.append(\"--enable-chunked-prefill\")\n",
+        "\n",
+        "    if enable_prefix_cache:\n",
+        "        vllm_args.append(\"--enable-prefix-caching\")\n",
+        "\n",
+        "    if 0 < host_prefix_kv_cache_utilization_target < 1:\n",
+        "        vllm_args.append(\n",
+        "            f\"--host-prefix-kv-cache-utilization-target={host_prefix_kv_cache_utilization_target}\"\n",
+        "        )\n",
+        "\n",
+        "    if model_type:\n",
+        "        vllm_args.append(f\"--model-type={model_type}\")\n",
+        "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\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",
+        "        model_garden_source_model_name=(\n",
+        "            f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
+        "        ),\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",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_huggingface_vllm_deployment.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
+        "        },\n",
+        "    )\n",
+        "    print(\"endpoint_name:\", endpoint.name)\n",
+        "\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "models[\"vllm\"], endpoints[\"vllm\"] = deploy_model_vllm(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=HUGGING_FACE_MODEL_ID),\n",
+        "    model_id=HUGGING_FACE_MODEL_ID,\n",
+        "    publisher=HUGGING_FACE_MODEL_ID.split(\"/\")[0].lower(),\n",
+        "    publisher_model_id=HUGGING_FACE_MODEL_ID.split(\"/\")[1].lower(),\n",
+        "    service_account=\"\",\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    enforce_eager=True,\n",
+        "    max_num_seqs=5,\n",
+        ")\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "K0B98U88eXzO"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. Sampling parameters supported by vLLM can be found [here](https://docs.vllm.ai/en/latest/dev/sampling_params.html).\n",
+        "\n",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown Human: What is a car?\n",
+        "# @markdown Assistant:  A car, or a motor car, is a road-connected human-transportation system used to move people or goods from one place to another. The term also encompasses a wide range of vehicles, including motorboats, trains, and aircrafts. Cars typically have four wheels, a cabin for passengers, and an engine or motor. They have been around since the early 19th century and are now one of the most popular forms of transportation, used for daily commuting, shopping, and other purposes.\n",
+        "# @markdown ```\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",
+        "prompt = \"What is a car?\"  # @param {type: \"string\"}\n",
+        "# @markdown If you encounter an issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, by lowering `max_tokens`.\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  # @param {type:\"integer\"}\n",
+        "# @markdown Set `raw_response` to `True` to obtain the raw model output. Set `raw_response` to `False` to apply additional formatting in the structure of `\"Prompt:\\n{prompt.strip()}\\nOutput:\\n{output}\"`.\n",
+        "raw_response = False  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# Overrides parameters for inferences.\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",
+        "        \"raw_response\": raw_response,\n",
+        "    },\n",
+        "]\n",
+        "response = endpoints[\"vllm\"].predict(\n",
+        "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
+        ")\n",
+        "\n",
+        "for prediction in response.predictions:\n",
+        "    print(prediction)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "iK3kcE3CeXzO"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict: Streaming Chat Completions\n",
+        "\n",
+        "# @markdown You can also send chat completions requests in streaming mode, if supported by the model.\n",
+        "\n",
+        "# @markdown Here we use an example from the [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) to show the chat completions output:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown User: How would the Future of AI in 10 Years look?\n",
+        "# @markdown 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 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",
+        "import json\n",
+        "\n",
+        "import requests\n",
+        "\n",
+        "# Initialize request argument(s).\n",
+        "prompt = \"How would the Future of AI in 10 Years look?\"  # @param {type: \"string\"}\n",
+        "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
+        "max_tokens = 128  # @param {type:\"integer\"}\n",
+        "temperature = 1.0  # @param {type:\"number\"}\n",
+        "top_p = 0.9  # @param {type:\"number\"}\n",
+        "\n",
+        "# Send streaming chat completions request.\n",
+        "endpoint = endpoints[\"vllm\"]\n",
+        "url = f\"https://{endpoint.location}-aiplatform.googleapis.com/v1beta1/{endpoint.resource_name}/chat/completions\"\n",
+        "if use_dedicated_endpoint:\n",
+        "    url = f\"https://{endpoint.gca_resource.dedicated_endpoint_dns}/v1beta1/{endpoint.resource_name}/chat/completions\"\n",
+        "\n",
+        "payload = {\n",
+        "    \"messages\": [{\"role\": \"user\", \"content\": prompt}],\n",
+        "    \"max_tokens\": max_tokens,\n",
+        "    \"temperature\": temperature,\n",
+        "    \"top_p\": top_p,\n",
+        "    \"stream\": True,\n",
+        "}\n",
+        "access_token = ! gcloud auth print-access-token\n",
+        "access_token = access_token[0]\n",
+        "\n",
+        "response = requests.post(\n",
+        "    url,\n",
+        "    headers={\"Authorization\": f\"Bearer {access_token}\"},\n",
+        "    json=payload,\n",
+        "    stream=True,\n",
+        ")\n",
+        "\n",
+        "if not response.ok:\n",
+        "    raise ValueError(response.text)\n",
+        "\n",
+        "for chunk in response.iter_lines(chunk_size=8192, decode_unicode=False):\n",
+        "    if chunk:\n",
+        "        chunk = chunk.decode(\"utf-8\").removeprefix(\"data:\").strip()\n",
+        "        if chunk == \"[DONE]\":\n",
+        "            break\n",
+        "        data = json.loads(chunk)\n",
+        "        if type(data) is not dict or \"error\" in data:\n",
+        "            raise ValueError(data)\n",
+        "        delta = data[\"choices\"][0][\"delta\"].get(\"content\")\n",
+        "        if delta:\n",
+        "            print(delta, end=\"\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "UowpdPJOeXzO"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "JgvATV8aeXzO"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Delete the models and endpoints\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()"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_huggingface_vllm_deployment.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_jax_dito.ipynb

@@ -40,7 +40,7 @@
         "\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_jax_dito.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",

notebooks/community/model_garden/model_garden_jax_fvlm.ipynb

@@ -40,7 +40,7 @@
         "\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_jax_fvlm.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",

notebooks/community/model_garden/model_garden_jax_owl_vit_v2.ipynb

@@ -40,7 +40,7 @@
         "  </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_jax_owl_vit_v2.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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"
@@ -118,14 +118,14 @@
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
-        "# @markdown 4. 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",
+        "# @markdown 4. 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). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\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-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "import base64\n",
         "import datetime\n",
@@ -155,6 +155,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",

notebooks/community/model_garden/model_garden_jax_paligemma_deployment.ipynb

@@ -26,7 +26,6 @@
     },
     {
       "cell_type": "markdown",
-      "language": "markdown",
       "metadata": {
         "id": "VJWDivOv3OWy"
       },
@@ -35,13 +34,18 @@
         "\n",
         "<table><tbody><tr>\n",
         "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/instances\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_jax_paligemma_deployment.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_jax_paligemma_deployment.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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"
@@ -68,6 +72,10 @@
         "  - Detecting objects.\n",
         "- Create a playground website to use with the PaliGemma Vertex AI Endpoint.\n",
         "\n",
+        "### File a bug\n",
+        "\n",
+        "File a bug on [GitHub](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues/new) if you encounter any issue with the notebook.\n",
+        "\n",
         "### Costs\n",
         "\n",
         "This tutorial uses billable components of Google Cloud:\n",
@@ -90,7 +98,6 @@
     {
       "cell_type": "code",
       "execution_count": null,
-      "language": "python",
       "metadata": {
         "cellView": "form",
         "id": "QvQjsmIJ6Y3f"
@@ -100,25 +107,29 @@
         "# @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",
+        "# @markdown 2. **[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",
-        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "# @markdown 3. 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). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\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-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
+        "\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform==1.93.1'\n",
         "\n",
         "# Import the necessary packages\n",
         "! pip install -q gradio==4.21.0\n",
-        "import datetime\n",
         "import enum\n",
         "import importlib\n",
         "import io\n",
         "import os\n",
         "import re\n",
-        "import uuid\n",
         "from typing import Sequence, Tuple\n",
         "\n",
         "import gradio as gr\n",
@@ -128,12 +139,19 @@
         "from google.cloud import aiplatform\n",
         "from PIL import Image\n",
         "\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
+        "    ! pip install --upgrade tensorflow\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\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",
+        "LABEL = \"endpoint\"\n",
         "models, endpoints = {}, {}\n",
         "\n",
+        "\n",
         "# Get the default cloud project id.\n",
         "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
         "\n",
@@ -141,52 +159,18 @@
         "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, \"paligemma\")\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",
+        "aiplatform.init(project=PROJECT_ID, location=REGION)\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",
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")\n",
         "\n",
         "# @markdown ### Access PaliGemma models on Vertex AI for GPU based serving\n",
         "# @markdown Accept the model agreement to access the models:\n",
@@ -199,17 +183,28 @@
         "VERTEX_AI_MODEL_GARDEN_PALIGEMMA = \"gs://\"  # @param {type:\"string\", isTemplate:true}\n",
         "assert (\n",
         "    VERTEX_AI_MODEL_GARDEN_PALIGEMMA\n",
-        "), \"Click the agreement of PaliGemma in Vertex AI Model Garden, and get the GCS path of PaliGemma model artifacts.\"\n",
-        "print(\n",
-        "    \"Copying PaliGemma model artifacts from\",\n",
-        "    VERTEX_AI_MODEL_GARDEN_PALIGEMMA,\n",
-        "    \"to \",\n",
-        "    MODEL_BUCKET,\n",
-        ")\n",
-        "\n",
-        "! gsutil -m cp -R $VERTEX_AI_MODEL_GARDEN_PALIGEMMA/* $MODEL_BUCKET\n",
-        "\n",
-        "model_path_prefix = MODEL_BUCKET\n",
+        "), \"Click the agreement of PaliGemma in Vertex AI Model Garden, and get the GCS path of PaliGemma model artifacts.\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kyMJXkfviWgl"
+      },
+      "source": [
+        "## Deploy PaliGemma to a Vertex AI Endpoint"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "JThvioAxy8-a"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Select the model variants\n",
         "\n",
         "pretrained_filename_lookup = {\n",
         "    \"paligemma-224-float32\": \"pt_224.npz\",\n",
@@ -229,6 +224,341 @@
         "    \"paligemma-mix-448-bfloat16\": \"mix_448.bf16.npz\",\n",
         "}\n",
         "\n",
+        "# @markdown Select the desired resolution and precision of prebuilt model to deploy, leaving the optional `custom_paligemma_model_uri` as is. Higher resolution and precision_type can result in better inference results, but may require additional GPU.\n",
+        "\n",
+        "# @markdown You can also serve a finetuned PaliGemma model by setting `resolution` and `precision_type` to the resolution and precision type of the original base model and then setting `custom_paligemma_model_uri` to the GCS URI containing the model.\n",
+        "\n",
+        "# @markdown **Note**: You cannot use accelerator type `NVIDIA_TESLA_V100` to serve prebuilt or finetuned PaliGemma models with resolution `896` and precision_type `float32`.\n",
+        "\n",
+        "model_variant = \"mix\"  # @param [\"mix\", \"pt\"]\n",
+        "resolution = 224  # @param [224, 448, 896]\n",
+        "precision_type = \"float32\"  # @param [\"float32\", \"float16\", \"bfloat16\"]\n",
+        "custom_paligemma_model_uri = \"gs://\"  # @param {type: \"string\"}\n",
+        "\n",
+        "if model_variant == \"mix\":\n",
+        "    model_name_prefix = \"paligemma-mix\"\n",
+        "else:\n",
+        "    model_name_prefix = \"paligemma\"\n",
+        "\n",
+        "\n",
+        "if custom_paligemma_model_uri == \"gs://\" or not custom_paligemma_model_uri:\n",
+        "    model_name = f\"{model_name_prefix}-{resolution}-{precision_type}\"\n",
+        "    checkpoint_filename = pretrained_filename_lookup[model_name]\n",
+        "    checkpoint_path = os.path.join(\n",
+        "        VERTEX_AI_MODEL_GARDEN_PALIGEMMA, checkpoint_filename\n",
+        "    )\n",
+        "    PUBLISHER_MODEL_NAME = f\"publishers/google/models/paligemma@{model_name}\"\n",
+        "else:\n",
+        "    model_name = f\"{model_name_prefix}-{resolution}-{precision_type}-custom\"\n",
+        "    checkpoint_path = custom_paligemma_model_uri\n",
+        "\n",
+        "# @markdown If you want to use other accelerator types not listed below, then check other Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute. You may need to manually set the `machine_type`, `accelerator_type`, and `accelerator_count` in the code by clicking `Show code` first.\n",
+        "# @markdown Select the accelerator type to use to deploy the model:\n",
+        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_V100\"]\n",
+        "if accelerator_type == \"NVIDIA_L4\":\n",
+        "    machine_type = \"g2-standard-16\"\n",
+        "    accelerator_count = 1\n",
+        "elif accelerator_type == \"NVIDIA_TESLA_V100\":\n",
+        "    if resolution == 896 and precision_type == \"float32\":\n",
+        "        raise ValueError(\n",
+        "            \"NVIDIA_TESLA_V100 is not sufficient. Multi-gpu is not supported for PaLIGemma.\"\n",
+        "        )\n",
+        "    else:\n",
+        "        machine_type = \"n1-highmem-8\"\n",
+        "        accelerator_count = 1\n",
+        "else:\n",
+        "    raise ValueError(\n",
+        "        f\"Recommended machine settings not found for: {accelerator_type}. To use another another accelerator, edit this code block to pass in an appropriate `machine_type`, `accelerator_type`, and `accelerator_count` to the deploy_model function by clicking `Show Code` and then modifying the code.\"\n",
+        "    )\n",
+        "\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @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",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "144IKkHrzrMs"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 1] Deploy with Model Garden SDK\n",
+        "\n",
+        "# @markdown Kindly note that the deployment using custom_paligemma_model_uri is not supported.\n",
+        "\n",
+        "# @markdown Deploy with Gen AI model-centric SDK. This section uploads the prebuilt model to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model. See [use open models with Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-open-models) for documentation on other use cases.\n",
+        "from vertexai import model_garden\n",
+        "\n",
+        "model = model_garden.OpenModel(PUBLISHER_MODEL_NAME)\n",
+        "endpoints[LABEL] = model.deploy(\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    accept_eula=True,  # Accept the End User License Agreement (EULA) on the model card before deploy. Otherwise, the deployment will be forbidden.\n",
+        ")\n",
+        "\n",
+        "endpoint = endpoints[LABEL]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "toY-WPKDFesF"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with custom configs\n",
+        "\n",
+        "# @markdown This section uploads the prebuilt PaliGemma model to Model Registry and deploys it to a Vertex AI Endpoint. It takes approximately 15 minutes to finish.\n",
+        "\n",
+        "# The pre-built serving docker image.\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/jax-paligemma-serve-gpu:20240807_0916_RC00\"\n",
+        "\n",
+        "\n",
+        "def deploy_model(\n",
+        "    model_name: str,\n",
+        "    checkpoint_path: str,\n",
+        "    machine_type: str = \"g2-standard-32\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 1,\n",
+        "    resolution: int = 224,\n",
+        "    use_dedicated_endpoint: bool = False,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Create a Vertex AI Endpoint and deploy the specified model to the endpoint.\"\"\"\n",
+        "    model_name_with_time = common_util.get_job_name_with_datetime(model_name)\n",
+        "    endpoint = aiplatform.Endpoint.create(\n",
+        "        display_name=f\"{model_name_with_time}-endpoint\",\n",
+        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
+        "    )\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name_with_time,\n",
+        "        serving_container_image_uri=SERVE_DOCKER_URI,\n",
+        "        serving_container_ports=[8080],\n",
+        "        serving_container_predict_route=\"/predict\",\n",
+        "        serving_container_health_route=\"/health\",\n",
+        "        serving_container_environment_variables={\n",
+        "            \"CKPT_PATH\": checkpoint_path,\n",
+        "            \"RESOLUTION\": resolution,\n",
+        "            \"MODEL_ID\": \"google/\" + model_name,\n",
+        "            \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "        },\n",
+        "        model_garden_source_model_name=\"publishers/google/models/paligemma\",\n",
+        "    )\n",
+        "    print(\n",
+        "        f\"Deploying {model_name_with_time} 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",
+        "        enable_access_logging=True,\n",
+        "        min_replica_count=1,\n",
+        "        sync=True,\n",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_jax_paligemma_deployment.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
+        "        },\n",
+        "    )\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "models[LABEL], endpoints[LABEL] = deploy_model(\n",
+        "    model_name=model_name,\n",
+        "    checkpoint_path=checkpoint_path,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    resolution=resolution,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "tOtYOhZa3lsx"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Optional] Loading an existing Endpoint\n",
+        "# @markdown If you've already deployed an Endpoint, you can load it by filling in the Endpoint's ID below.\n",
+        "# @markdown You can view deployed Endpoints at [Vertex Online Prediction](https://console.cloud.google.com/vertex-ai/online-prediction/endpoints).\n",
+        "endpoint_id = \"\"  # @param {type: \"string\"}\n",
+        "\n",
+        "if endpoint_id:\n",
+        "    endpoints[LABEL] = aiplatform.Endpoint(\n",
+        "        endpoint_name=endpoint_id,\n",
+        "        project=PROJECT_ID,\n",
+        "        location=REGION,\n",
+        "    )"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "MlP2Y7XE4SS5"
+      },
+      "source": [
+        "### Predict\n",
+        "\n",
+        "The following sections will use images from [pexels.com](https://www.pexels.com/) for demoing purposes. All the images have the following license: https://www.pexels.com/license/.\n",
+        "\n",
+        "Images will be resized to a width of 1000 pixels by default since requests made to a Vertex Endpoint are limited to 1.500MB."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "xnZw8wNyQhmN"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Visual Question Answering\n",
+        "\n",
+        "# @markdown This section uses the deployed PaliGemma model to answer questions about a given image.\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with images and questions.\n",
+        "# @markdown ![](https://images.pexels.com/photos/4012966/pexels-photo-4012966.jpeg?w=1260&h=750)\n",
+        "image_url = \"https://images.pexels.com/photos/4012966/pexels-photo-4012966.jpeg\"  # @param {type:\"string\"}\n",
+        "\n",
+        "image = common_util.download_image(image_url)\n",
+        "display(image)\n",
+        "\n",
+        "# @markdown You may leave question prompts empty and they will be ignored.\n",
+        "question_prompt_1 = \"Which of laptop, book, pencil, clock, flower are in the image?\"  # @param {type: \"string\"}\n",
+        "question_prompt_2 = \"Do the book and the cup have the same color?\"  # @param {type: \"string\"}\n",
+        "question_prompt_3 = \"Is there a person in the image?\"  # @param {type: \"string\"}\n",
+        "question_prompt_4 = \"How many laptop are in the image?\"  # @param {type: \"string\"}\n",
+        "question_prompt_5 = \"桌子是什么颜色的?\"  # @param {type: \"string\"}\n",
+        "\n",
+        "# @markdown The question prompt can be non-English languages.\n",
+        "questions_list = [\n",
+        "    question_prompt_1,\n",
+        "    question_prompt_2,\n",
+        "    question_prompt_3,\n",
+        "    question_prompt_4,\n",
+        "    question_prompt_5,\n",
+        "]\n",
+        "questions = [question for question in questions_list if question]\n",
+        "\n",
+        "answers = common_util.vqa_predict(\n",
+        "    endpoints[\"endpoint\"],\n",
+        "    questions,\n",
+        "    image,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        ")\n",
+        "\n",
+        "for question, answer in zip(questions, answers):\n",
+        "    print(f\"Question: {question}\")\n",
+        "    print(f\"Answer: {answer}\")\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "mF1MxC1ouzqj"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Image Captioning\n",
+        "# @markdown This section uses the deployed PaliGemma model to caption and describe an image in a chosen language.\n",
+        "\n",
+        "caption_prompt = True\n",
+        "\n",
+        "# @markdown <img src=\"https://storage.googleapis.com/longcap100/91.jpeg\" width=\"400\" >\n",
+        "\n",
+        "image_url = \"https://storage.googleapis.com/longcap100/91.jpeg\"  # @param {type:\"string\"}\n",
+        "\n",
+        "language_code = \"en\"  # @param {type: \"string\"}\n",
+        "\n",
+        "image = common_util.download_image(image_url)\n",
+        "display(image)\n",
+        "\n",
+        "# Make a prediction.\n",
+        "image_base64 = common_util.image_to_base64(image)\n",
+        "\n",
+        "caption = common_util.caption_predict(\n",
+        "    endpoints[\"endpoint\"],\n",
+        "    language_code,\n",
+        "    image,\n",
+        "    caption_prompt,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        ")\n",
+        "\n",
+        "print(\"Caption: \", caption)\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "TtkXMZTIegLq"
+      },
+      "outputs": [],
+      "source": [
+        "# @title OCR\n",
+        "# @markdown This section uses the deployed PaliGemma model to extract text from an image, starting from the top left.\n",
+        "ocr_prompt = \"ocr\"\n",
+        "\n",
+        "# @markdown ![](https://images.pexels.com/photos/8919535/pexels-photo-8919535.jpeg?auto=compress&cs=tinysrgb&w=630&h=375&dpr=2)\n",
+        "image_url = \"https://images.pexels.com/photos/8919535/pexels-photo-8919535.jpeg?auto=compress&cs=tinysrgb&w=1260&h=750&dpr=2\"  # @param {type:\"string\"}\n",
+        "\n",
+        "image = common_util.download_image(image_url)\n",
+        "display(image)\n",
+        "text_found = common_util.ocr_predict(\n",
+        "    endpoints[\"endpoint\"],\n",
+        "    ocr_prompt,\n",
+        "    image,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        ")\n",
+        "\n",
+        "print(f\"Text found: {text_found}\")\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "JlLr3nu-YEon"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Object Detection\n",
+        "# @markdown This section uses the deployed PaliGemma model to output bounding boxes for specified object image in a given image.\n",
+        "# @markdown The text output will be parsed into bounding boxes and overlaid on the original image.\n",
+        "\n",
+        "# @markdown Specify what object to detect. To specify multiple objects, enter them as a semicolon separated list as shown below.\n",
+        "objects = \"plant ; pineapple ; glasses\"  # @param {type:\"string\"}\n",
+        "detect_promt = f\"detect {objects}\"\n",
+        "\n",
         "\n",
         "def parse_detections(txt):\n",
         "    \"\"\"Parses bounding boxes from a detection string.\"\"\"\n",
@@ -271,296 +601,8 @@
         "    buf = io.BytesIO()\n",
         "    fig.savefig(buf)\n",
         "    buf.seek(0)\n",
-        "    return Image.open(buf)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "kyMJXkfviWgl"
-      },
-      "source": [
-        "## Deploy PaliGemma to a Vertex AI Endpoint"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "cellView": "form",
-        "id": "toY-WPKDFesF"
-      },
-      "outputs": [],
-      "source": [
-        "# @title Deploy\n",
+        "    return Image.open(buf)\n",
         "\n",
-        "# @markdown This section uploads the prebuilt PaliGemma model to Model Registry and deploys it to a Vertex AI Endpoint. It takes approximately 15 minutes to finish.\n",
-        "\n",
-        "# @markdown Select the desired resolution and precision of prebuilt model to deploy, leaving the optional `custom_paligemma_model_uri` as is. Higher resolution and precision_type can result in better inference results, but may require additional GPU.\n",
-        "\n",
-        "# @markdown You can also serve a finetuned PaliGemma model by setting `resolution` and `precision_type` to the resolution and precision type of the original base model and then setting `custom_paligemma_model_uri` to the GCS URI containing the model.\n",
-        "\n",
-        "# @markdown **Note**: You cannot use accelerator type `NVIDIA_TESLA_V100` to serve prebuilt or finetuned PaliGemma models with resolution `896` and precision_type `float32`.\n",
-        "\n",
-        "model_variant = \"mix\"  # @param [\"mix\", \"pt\"]\n",
-        "resolution = 224  # @param [224, 448, 896]\n",
-        "precision_type = \"float32\"  # @param [\"float32\", \"float16\", \"bfloat16\"]\n",
-        "custom_paligemma_model_uri = \"gs://\"  # @param {type: \"string\"}\n",
-        "\n",
-        "if model_variant == \"mix\":\n",
-        "    model_name_prefix = \"paligemma-mix\"\n",
-        "else:\n",
-        "    model_name_prefix = \"paligemma\"\n",
-        "\n",
-        "if custom_paligemma_model_uri == \"gs://\" or not custom_paligemma_model_uri:\n",
-        "    print(\"Deploying prebuilt PaliGemma model.\")\n",
-        "    model_name = f\"{model_name_prefix}-{resolution}-{precision_type}\"\n",
-        "    checkpoint_filename = pretrained_filename_lookup[model_name]\n",
-        "    checkpoint_path = os.path.join(model_path_prefix, checkpoint_filename)\n",
-        "else:\n",
-        "    print(\"Deploying custom PaliGemma model.\")\n",
-        "    model_name = f\"{model_name_prefix}-{resolution}-{precision_type}-custom\"\n",
-        "    checkpoint_path = custom_paligemma_model_uri\n",
-        "\n",
-        "# The pre-built serving docker image.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/jax-paligemma-serve-gpu:20240807_0916_RC00\"\n",
-        "\n",
-        "# @markdown If you want to use other accelerator types not listed below, then check other Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute. You may need to manually set the `machine_type`, `accelerator_type`, and `accelerator_count` in the code by clicking `Show code` first.\n",
-        "# @markdown Select the accelerator type to use to deploy the model:\n",
-        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_V100\"]\n",
-        "if accelerator_type == \"NVIDIA_L4\":\n",
-        "    machine_type = \"g2-standard-16\"\n",
-        "    accelerator_count = 1\n",
-        "elif accelerator_type == \"NVIDIA_TESLA_V100\":\n",
-        "    if resolution == 896 and precision_type == \"float32\":\n",
-        "        raise ValueError(\n",
-        "            \"NVIDIA_TESLA_V100 is not sufficient. Multi-gpu is not supported for PaLIGemma.\"\n",
-        "        )\n",
-        "    else:\n",
-        "        machine_type = \"n1-highmem-8\"\n",
-        "        accelerator_count = 1\n",
-        "else:\n",
-        "    raise ValueError(\n",
-        "        f\"Recommended machine settings not found for: {accelerator_type}. To use another another accelerator, edit this code block to pass in an appropriate `machine_type`, `accelerator_type`, and `accelerator_count` to the deploy_model function by clicking `Show Code` and then modifying the code.\"\n",
-        "    )\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",
-        "\n",
-        "def deploy_model(\n",
-        "    model_name: str,\n",
-        "    checkpoint_path: str,\n",
-        "    machine_type: str = \"g2-standard-32\",\n",
-        "    accelerator_type: str = \"NVIDIA_L4\",\n",
-        "    accelerator_count: int = 1,\n",
-        "    resolution: int = 224,\n",
-        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
-        "    \"\"\"Create a Vertex AI Endpoint and deploy the specified model to the endpoint.\"\"\"\n",
-        "    model_name_with_time = common_util.get_job_name_with_datetime(model_name)\n",
-        "    endpoint = aiplatform.Endpoint.create(\n",
-        "        display_name=f\"{model_name_with_time}-endpoint\"\n",
-        "    )\n",
-        "    model = aiplatform.Model.upload(\n",
-        "        display_name=model_name_with_time,\n",
-        "        serving_container_image_uri=SERVE_DOCKER_URI,\n",
-        "        serving_container_ports=[8080],\n",
-        "        serving_container_predict_route=\"/predict\",\n",
-        "        serving_container_health_route=\"/health\",\n",
-        "        serving_container_environment_variables={\n",
-        "            \"CKPT_PATH\": checkpoint_path,\n",
-        "            \"RESOLUTION\": resolution,\n",
-        "            \"MODEL_ID\": \"google/\" + model_name,\n",
-        "        },\n",
-        "        model_garden_source_model_name=\"publishers/google/models/paligemma\",\n",
-        "    )\n",
-        "    print(\n",
-        "        f\"Deploying {model_name_with_time} 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",
-        "        enable_access_logging=True,\n",
-        "        min_replica_count=1,\n",
-        "        sync=True,\n",
-        "        system_labels={\"NOTEBOOK_NAME\": \"model_garden_jax_paligemma_deployment.ipynb\"},\n",
-        "    )\n",
-        "    return model, endpoint\n",
-        "\n",
-        "\n",
-        "models[\"model\"], endpoints[\"endpoint\"] = deploy_model(\n",
-        "    model_name=model_name,\n",
-        "    checkpoint_path=checkpoint_path,\n",
-        "    machine_type=machine_type,\n",
-        "    accelerator_type=accelerator_type,\n",
-        "    accelerator_count=accelerator_count,\n",
-        "    resolution=resolution,\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "cellView": "form",
-        "id": "tOtYOhZa3lsx"
-      },
-      "outputs": [],
-      "source": [
-        "# @title [Optional] Loading an existing Endpoint\n",
-        "# @markdown If you've already deployed an Endpoint, you can load it by filling in the Endpoint's ID below.\n",
-        "# @markdown You can view deployed Endpoints at [Vertex Online Prediction](https://console.cloud.google.com/vertex-ai/online-prediction/endpoints).\n",
-        "endpoint_id = \"\"  # @param {type: \"string\"}\n",
-        "\n",
-        "if endpoint_id:\n",
-        "    endpoint = aiplatform.Endpoint(\n",
-        "        endpoint_name=endpoint_id,\n",
-        "        project=PROJECT_ID,\n",
-        "        location=REGION,\n",
-        "    )"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "MlP2Y7XE4SS5"
-      },
-      "source": [
-        "### Predict\n",
-        "\n",
-        "The following sections will use images from [pexels.com](https://www.pexels.com/) for demoing purposes. All the images have the following license: https://www.pexels.com/license/.\n",
-        "\n",
-        "Images will be resized to a width of 1000 pixels by default since requests made to a Vertex Endpoint are limited to 1.500MB."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "language": "python",
-      "metadata": {
-        "cellView": "form",
-        "id": "xnZw8wNyQhmN"
-      },
-      "outputs": [],
-      "source": [
-        "# @title Visual Question Answering\n",
-        "\n",
-        "# @markdown This section uses the deployed PaliGemma model to answer questions about a given image.\n",
-        "\n",
-        "# @markdown Once deployment succeeds, you can send requests to the endpoint with images and questions.\n",
-        "# @markdown ![](https://images.pexels.com/photos/4012966/pexels-photo-4012966.jpeg?w=1260&h=750)\n",
-        "image_url = \"https://images.pexels.com/photos/4012966/pexels-photo-4012966.jpeg\"  # @param {type:\"string\"}\n",
-        "\n",
-        "image = common_util.download_image(image_url)\n",
-        "display(image)\n",
-        "\n",
-        "# @markdown You may leave question prompts empty and they will be ignored.\n",
-        "question_prompt_1 = \"Which of laptop, book, pencil, clock, flower are in the image?\"  # @param {type: \"string\"}\n",
-        "question_prompt_2 = \"Do the book and the cup have the same color?\"  # @param {type: \"string\"}\n",
-        "question_prompt_3 = \"Is there a person in the image?\"  # @param {type: \"string\"}\n",
-        "question_prompt_4 = \"How many laptop are in the image?\"  # @param {type: \"string\"}\n",
-        "question_prompt_5 = \"桌子是什么颜色的?\"  # @param {type: \"string\"}\n",
-        "\n",
-        "# @markdown The question prompt can be non-English languages.\n",
-        "questions_list = [\n",
-        "    question_prompt_1,\n",
-        "    question_prompt_2,\n",
-        "    question_prompt_3,\n",
-        "    question_prompt_4,\n",
-        "    question_prompt_5,\n",
-        "]\n",
-        "questions = [question for question in questions_list if question]\n",
-        "\n",
-        "answers = common_util.vqa_predict(endpoints[\"endpoint\"], questions, image)\n",
-        "\n",
-        "for question, answer in zip(questions, answers):\n",
-        "    print(f\"Question: {question}\")\n",
-        "    print(f\"Answer: {answer}\")\n",
-        "# @markdown Click \"Show Code\" to see more details."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "cellView": "form",
-        "id": "mF1MxC1ouzqj"
-      },
-      "outputs": [],
-      "source": [
-        "# @title Image Captioning\n",
-        "# @markdown This section uses the deployed PaliGemma model to caption and describe an image in a chosen language.\n",
-        "\n",
-        "caption_prompt = True\n",
-        "\n",
-        "# @markdown <img src=\"https://storage.googleapis.com/longcap100/91.jpeg\" width=\"400\" >\n",
-        "\n",
-        "image_url = \"https://storage.googleapis.com/longcap100/91.jpeg\"  # @param {type:\"string\"}\n",
-        "language_code = \"en\"  # @param {type: \"string\"}\n",
-        "\n",
-        "image = common_util.download_image(image_url)\n",
-        "display(image)\n",
-        "\n",
-        "# Make a prediction.\n",
-        "image_base64 = common_util.image_to_base64(image)\n",
-        "\n",
-        "caption = common_util.caption_predict(\n",
-        "    endpoints[\"endpoint\"], language_code, image, caption_prompt\n",
-        ")\n",
-        "\n",
-        "print(\"Caption: \", caption)\n",
-        "# @markdown Click \"Show Code\" to see more details."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "cellView": "form",
-        "id": "TtkXMZTIegLq"
-      },
-      "outputs": [],
-      "source": [
-        "# @title OCR\n",
-        "# @markdown This section uses the deployed PaliGemma model to extract text from an image, starting from the top left.\n",
-        "ocr_prompt = \"ocr\"\n",
-        "\n",
-        "# @markdown ![](https://images.pexels.com/photos/8919535/pexels-photo-8919535.jpeg?auto=compress&cs=tinysrgb&w=630&h=375&dpr=2)\n",
-        "image_url = \"https://images.pexels.com/photos/8919535/pexels-photo-8919535.jpeg?auto=compress&cs=tinysrgb&w=1260&h=750&dpr=2\"  # @param {type:\"string\"}\n",
-        "\n",
-        "image = common_util.download_image(image_url)\n",
-        "display(image)\n",
-        "text_found = common_util.ocr_predict(endpoints[\"endpoint\"], ocr_prompt, image)\n",
-        "\n",
-        "print(f\"Text found: {text_found}\")\n",
-        "# @markdown Click \"Show Code\" to see more details."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "cellView": "form",
-        "id": "JlLr3nu-YEon"
-      },
-      "outputs": [],
-      "source": [
-        "# @title Object Detection\n",
-        "# @markdown This section uses the deployed PaliGemma model to output bounding boxes for specified object image in a given image.\n",
-        "# @markdown The text output will be parsed into bounding boxes and overlaid on the original image.\n",
-        "\n",
-        "# @markdown Specify what object to detect. To specify multiple objects, enter them as a semicolon separated list as shown below.\n",
-        "objects = \"plant ; pineapple ; glasses\"  # @param {type:\"string\"}\n",
-        "detect_promt = f\"detect {objects}\"\n",
         "\n",
         "# @markdown ![](https://images.pexels.com/photos/1006293/pexels-photo-1006293.jpeg?auto=compress&cs=tinysrgb&w=630&h=375&dpr=2)\n",
         "image_url = \"https://images.pexels.com/photos/1006293/pexels-photo-1006293.jpeg?auto=compress&cs=tinysrgb&w=1260&h=750&dpr=2\"  # @param {type:\"string\"}\n",
@@ -570,10 +612,15 @@
         "\n",
         "# Make a prediction.\n",
         "detection_response = common_util.detect_predict(\n",
-        "    endpoints[\"endpoint\"], detect_promt, image\n",
+        "    endpoints[\"endpoint\"],\n",
+        "    detect_promt,\n",
+        "    image,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
         ")\n",
         "\n",
         "print(\"Output: \", detection_response)\n",
+        "\n",
+        "\n",
         "bboxes = parse_detections(detection_response)\n",
         "plot_bounding_boxes(image, bboxes)\n",
         "# @markdown Click \"Show Code\" to see more details."
@@ -700,7 +747,9 @@
         "    resolution = int(resolution)\n",
         "    model, endpoint = deploy_model(\n",
         "        model_name=model_choice,\n",
-        "        checkpoint_path=os.path.join(model_path_prefix, checkpoint_filename),\n",
+        "        checkpoint_path=os.path.join(\n",
+        "            VERTEX_AI_MODEL_GARDEN_PALIGEMMA, checkpoint_filename\n",
+        "        ),\n",
         "        machine_type=\"g2-standard-16\",\n",
         "        accelerator_type=\"NVIDIA_L4\",\n",
         "        accelerator_count=1,\n",
@@ -885,6 +934,8 @@
       },
       "outputs": [],
       "source": [
+        "# @title Delete the models and endpoints\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",
@@ -894,11 +945,7 @@
         "\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"
+        "    model.delete()"
       ]
     }
   ],

notebooks/community/model_garden/model_garden_jax_paligemma_finetuning.ipynb

@@ -43,7 +43,7 @@
         "  </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_jax_paligemma_finetuning.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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"
@@ -108,7 +108,7 @@
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, 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",
@@ -142,6 +142,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -271,9 +277,17 @@
         "\n",
         "dataset_gcs_uri = \"gs://longcap100/data_train90.jsonl\"  # @param {type: \"string\"}\n",
         "\n",
-        "# @markdown [Optional] You can optionally specify the image fields in the JSONL file to use the\n",
-        "# @markdown filename and fill in the `dataset_image_dir` with the location where the images are stored.\n",
-        "dataset_image_dir = \"\"  # @param {type:\"string\"}"
+        "# @markdown [Optional] You can specify the `image` fields in the JSONL file to\n",
+        "# @markdown contain only filenames. In this case, you must also provide the\n",
+        "# @markdown image storage location in `dataset_image_dir`. If the JSONL file\n",
+        "# @markdown already contains full paths to the images, leave\n",
+        "# @markdown `dataset_image_dir` blank. Note that the `SERVICE_ACCOUNT` defined\n",
+        "# @markdown above must have read access to the images.\n",
+        "dataset_image_dir = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Set defaults for the example dataset.\n",
+        "if dataset_gcs_uri == \"gs://longcap100/data_train90.jsonl\" and not dataset_image_dir:\n",
+        "    dataset_image_dir = \"gs://longcap100\""
       ]
     },
     {
@@ -409,8 +423,7 @@
         "if learning_rate:\n",
         "    train_args.append(f\"--config.lr={learning_rate}\")\n",
         "\n",
-        "if dataset_image_dir:\n",
-        "    train_args.append(f\"--config.input.data.fopen_keys.image={dataset_image_dir}\")\n",
+        "train_args.append(f\"--config.input.data.fopen_keys.image={dataset_image_dir}\")\n",
         "train_job.run(\n",
         "    args=train_args,\n",
         "    replica_count=replica_count,\n",
@@ -523,6 +536,10 @@
         "    raise ValueError(\n",
         "        f\"Recommended machine settings not found for: {accelerator_type}. To use another another accelerator, edit this code block to pass in an appropriate `machine_type`, `accelerator_type`, and `accelerator_count` to the deploy_model function by clicking `Show Code` and then modifying the code.\"\n",
         "    )\n",
+        "\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
+        "\n",
         "common_util.check_quota(\n",
         "    project_id=PROJECT_ID,\n",
         "    region=REGION,\n",
@@ -539,11 +556,13 @@
         "    accelerator_type: str = \"NVIDIA_L4\",\n",
         "    accelerator_count: int = 1,\n",
         "    resolution: int = 224,\n",
+        "    use_dedicated_endpoint: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Create a Vertex AI Endpoint and deploy the specified model to the endpoint.\"\"\"\n",
         "    model_name_with_time = common_util.get_job_name_with_datetime(model_name)\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
-        "        display_name=f\"{model_name_with_time}-endpoint\"\n",
+        "        display_name=f\"{model_name_with_time}-endpoint\",\n",
+        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
         "    )\n",
         "    model = aiplatform.Model.upload(\n",
         "        display_name=model_name_with_time,\n",
@@ -555,6 +574,7 @@
         "            \"CKPT_PATH\": checkpoint_path,\n",
         "            \"RESOLUTION\": resolution,\n",
         "            \"MODEL_ID\": model_name,\n",
+        "            \"DEPLOY_SOURCE\": \"notebook\",\n",
         "        },\n",
         "        model_garden_source_model_name=\"publishers/google/models/paligemma\",\n",
         "    )\n",
@@ -584,6 +604,7 @@
         "    accelerator_type=accelerator_type,\n",
         "    accelerator_count=accelerator_count,\n",
         "    resolution=model_resolution,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
         ")"
       ]
     },
@@ -613,6 +634,7 @@
         "# @markdown <img src=\"https://storage.googleapis.com/longcap100/91.jpeg\" width=\"400\" >\n",
         "\n",
         "image_url = \"https://storage.googleapis.com/longcap100/91.jpeg\"  # @param {type:\"string\"}\n",
+        "\n",
         "language_code = \"en\"  # @param {type: \"string\"}\n",
         "\n",
         "image = common_util.download_image(image_url)\n",
@@ -622,7 +644,11 @@
         "image_base64 = common_util.image_to_base64(image)\n",
         "\n",
         "caption = common_util.caption_predict(\n",
-        "    endpoints[\"endpoint\"], language_code, image, caption_prompt\n",
+        "    endpoints[\"endpoint\"],\n",
+        "    language_code,\n",
+        "    image,\n",
+        "    caption_prompt,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
         ")\n",
         "\n",
         "print(\"Caption: \", caption)\n",

notebooks/community/model_garden/model_garden_jax_stable_diffusion_xl.ipynb

@@ -39,7 +39,7 @@
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_jax_stable_diffusion_xl.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",

notebooks/community/model_garden/model_garden_jax_vision_transformer.ipynb

@@ -40,7 +40,7 @@
         "\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_jax_vision_transformer.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",

notebooks/community/model_garden/model_garden_keras_stable_diffusion.ipynb

@@ -40,7 +40,7 @@
         "\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_keras_stable_diffusion.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",

notebooks/community/model_garden/model_garden_keras_yolov8.ipynb

@@ -40,7 +40,7 @@
         "  </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_keras_yolov8.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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"
@@ -108,14 +108,14 @@
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
-        "# @markdown 4. 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",
+        "# @markdown 4. 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). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\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-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# Configs for all notebooks.\n",
         "! pip3 install --quiet keras-cv==0.9.0\n",
@@ -154,6 +154,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",

notebooks/community/model_garden/model_garden_llama3_2_deployment_on_gke.ipynb

@@ -40,7 +40,7 @@
         "  </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_llama3_2_deployment_on_gke.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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"

notebooks/community/model_garden/model_garden_llama3_2_evaluation.ipynb

@@ -40,7 +40,7 @@
         "  </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_llama3_2_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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"
@@ -136,6 +136,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",

notebooks/community/model_garden/model_garden_llama_guard_deployment.ipynb

@@ -34,13 +34,18 @@
         "\n",
         "<table><tbody><tr>\n",
         "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/instances\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_llama_guard_deployment.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_llama_guard_deployment.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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"
@@ -54,13 +59,18 @@
       "source": [
         "## Overview\n",
         "\n",
-        "This notebook demonstrates downloading and deploying prebuilt [Llama Guard models](https://huggingface.co/meta-llama) with [vLLM](https://github.com/vllm-project/vllm) on GPU, and demonstrates using the Llama Guard model to safeguard LLM inputs and outputs with the Vertex Llama 3 API service.\n",
+        "This notebook demonstrates downloading and deploying [Llama Guard models](https://huggingface.co/meta-llama) with [vLLM](https://github.com/vllm-project/vllm) on GPU, and demonstrates using the Llama Guard model to safeguard LLM inputs and outputs with the Vertex Llama API service.\n",
         "\n",
         "### Objective\n",
         "\n",
-        "- Download and deploy prebuilt Llama Guard models with [vLLM](https://github.com/vllm-project/vllm) on GPU\n",
+        "- Download and deploy Llama Guard models with [vLLM](https://github.com/vllm-project/vllm) on GPU\n",
         "- Use the Llama Guard models to safeguard LLM inputs and outputs with the Vertex Llama 3.1 API service\n",
         "- Use the Llama Guard models to safeguard LLM vision inputs and outputs with the Vertex Llama 3.2 API service\n",
+        "- Use the Llama Guard models to safeguard LLM vision inputs and outputs with the Vertex Llama 4 API service\n",
+        "\n",
+        "### File a bug\n",
+        "\n",
+        "File a bug on [GitHub](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues/new) if you encounter any issue with the notebook.\n",
         "\n",
         "### Costs\n",
         "\n",
@@ -110,44 +120,44 @@
         "\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",
+        "# @markdown 2. **[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",
-        "# @markdown 4. 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",
+        "# @markdown 3. 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). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\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-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# Import the necessary packages\n",
         "\n",
         "# Upgrade Vertex AI SDK.\n",
-        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.64.0'\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform==1.93.1'\n",
         "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
         "\n",
-        "import datetime\n",
         "import importlib\n",
         "import os\n",
         "import re\n",
-        "import uuid\n",
         "from typing import Tuple\n",
         "\n",
         "from google.cloud import aiplatform\n",
         "\n",
+        "if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
+        "    ! pip install --upgrade tensorflow\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\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",
+        "LABEL = \"vllm_gpu\"\n",
         "models, endpoints = {}, {}\n",
         "\n",
+        "\n",
         "# Get the default cloud project id.\n",
         "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
         "\n",
@@ -155,52 +165,17 @@
         "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, \"llama_guard\")\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",
+        "aiplatform.init(project=PROJECT_ID, location=REGION)\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",
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")\n",
         "\n",
         "# @markdown # Access Llama Guard models on Vertex AI\n",
         "# @markdown The original models from Meta are converted into the Hugging Face format for serving in Vertex AI.\n",
@@ -221,15 +196,7 @@
         "    VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD = parsed_gcs_url.group()\n",
         "assert VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD.startswith(\n",
         "    \"gs://\"\n",
-        "), \"VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD is expected to be a GCS URI and must start with `gs://`.\"\n",
-        "print(\n",
-        "    \"Copying Llama Guard model artifacts from\",\n",
-        "    VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD,\n",
-        "    \"to \",\n",
-        "    MODEL_BUCKET,\n",
-        ")\n",
-        "\n",
-        "! gsutil -m cp -R $VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD/* $MODEL_BUCKET"
+        "), \"VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD is expected to be a GCS URI and must start with `gs://`.\""
       ]
     },
     {
@@ -238,7 +205,7 @@
         "id": "z-XybZjtgF9M"
       },
       "source": [
-        "## Deploy Llama Guard with vLLM on GPU"
+        "## Deploy Llama Guard"
       ]
     },
     {
@@ -246,30 +213,33 @@
       "execution_count": null,
       "metadata": {
         "cellView": "form",
-        "id": "E8OiHHNNE_wj"
+        "id": "kRiRTAMxxUoq"
       },
       "outputs": [],
       "source": [
-        "# @title Deploy\n",
+        "# @title Select the model variants\n",
         "\n",
-        "# @markdown This section uploads prebuilt Llama Guard models to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model.\n",
+        "# @markdown Select one of the three model variations.\n",
         "\n",
-        "# @markdown Set the model to deploy.\n",
-        "\n",
-        "base_model_name = \"Llama-Guard-3-8B\"  # @param [\"Llama-Guard-3-8B\", \"Llama-Guard-3-1B\", \"Llama-Guard-3-11B-Vision\"] {allow-input: true, isTemplate: true}\n",
-        "model_id = os.path.join(MODEL_BUCKET, base_model_name)\n",
+        "base_model_name = \"Llama-Guard-4-12B\"  # @param [\"Llama-Guard-4-12B\", \"Llama-Guard-3-8B\", \"Llama-Guard-3-1B\", \"Llama-Guard-3-11B-Vision\"] {allow-input: true, isTemplate: true}\n",
+        "model_id = os.path.join(VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD, base_model_name)\n",
         "hf_model_id = \"meta-llama/\" + base_model_name\n",
+        "version_id = base_model_name.lower()\n",
+        "PUBLISHER_MODEL_NAME = f\"publishers/meta/models/llama-guard@{version_id}\"\n",
         "\n",
         "# The pre-built serving docker images.\n",
-        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240925_0916_RC01\"\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20250429_0916_RC01\"\n",
+        "\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
         "\n",
         "# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
         "if \"3-1B\" in base_model_name or \"3-8B\" in base_model_name:\n",
         "    accelerator_type = \"NVIDIA_L4\"\n",
-        "    machine_type = \"g2-standard-8\"\n",
+        "    machine_type = \"g2-standard-12\"\n",
         "    accelerator_count = 1\n",
         "    max_num_seqs = 256\n",
-        "elif \"3-11B\" in base_model_name:\n",
+        "elif \"3-11B\" in base_model_name or \"4-12B\" in base_model_name:\n",
         "    accelerator_type = \"NVIDIA_TESLA_A100\"\n",
         "    machine_type = \"a2-highgpu-1g\"\n",
         "    accelerator_count = 1\n",
@@ -283,8 +253,48 @@
         "    accelerator_type=accelerator_type,\n",
         "    accelerator_count=accelerator_count,\n",
         "    is_for_training=False,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "2DiRl36FzauJ"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 1] Deploy with Model Garden SDK\n",
+        "\n",
+        "# @markdown Deploy with Gen AI model-centric SDK. This section uploads the prebuilt model to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model. See [use open models with Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-open-models) for documentation on other use cases.\n",
+        "from vertexai import model_garden\n",
+        "\n",
+        "model = model_garden.OpenModel(PUBLISHER_MODEL_NAME)\n",
+        "endpoints[LABEL] = model.deploy(\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    accept_eula=True,  # Accept the End User License Agreement (EULA) on the model card before deploy. Otherwise, the deployment will be forbidden.\n",
         ")\n",
         "\n",
+        "endpoint = endpoints[LABEL]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "E8OiHHNNE_wj"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with customized configurations\n",
+        "\n",
+        "# @markdown This section uploads Llama Guard models to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model.\n",
+        "\n",
         "gpu_memory_utilization = 0.9\n",
         "max_model_len = 4096\n",
         "\n",
@@ -294,7 +304,6 @@
         "    model_id: str,\n",
         "    publisher: str,\n",
         "    publisher_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",
@@ -313,6 +322,7 @@
         "    use_dedicated_endpoint: bool = False,\n",
         "    max_num_seqs: int = 256,\n",
         "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -365,6 +375,10 @@
         "    if model_type:\n",
         "        vllm_args.append(f\"--model-type={model_type}\")\n",
         "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -400,9 +414,9 @@
         "        accelerator_type=accelerator_type,\n",
         "        accelerator_count=accelerator_count,\n",
         "        deploy_request_timeout=1800,\n",
-        "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_llama_guard_deployment.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",
@@ -410,24 +424,21 @@
         "    return model, endpoint\n",
         "\n",
         "\n",
-        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
-        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
-        "\n",
-        "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
-        "    model_name=common_util.get_job_name_with_datetime(prefix=\"llama3-guard\"),\n",
+        "models[LABEL], endpoints[LABEL] = deploy_model_vllm(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=\"llama-guard\"),\n",
         "    model_id=model_id,\n",
         "    publisher=\"meta\",\n",
-        "    publisher=\"llama-guard\",\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
+        "    publisher_model_id=\"llama-guard\",\n",
         "    base_model_id=hf_model_id,\n",
         "    machine_type=machine_type,\n",
         "    accelerator_type=accelerator_type,\n",
         "    accelerator_count=accelerator_count,\n",
         "    gpu_memory_utilization=gpu_memory_utilization,\n",
         "    max_model_len=max_model_len,\n",
-        "    enforce_eager=True,\n",
+        "    enforce_eager=False,\n",
         "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
         "    max_num_seqs=max_num_seqs,\n",
+        "    enable_llama_tool_parser=False,\n",
         ")\n",
         "# @markdown Click \"Show Code\" to see more details."
       ]
@@ -549,8 +560,8 @@
         "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
         ")\n",
         "\n",
-        "prediction = response.predictions[0]\n",
-        "print(\"Llama Guard prediction:\", prediction[0][\"message\"][\"content\"])"
+        "prediction = response.predictions[\"choices\"][0][\"message\"][\"content\"]\n",
+        "print(\"Llama Guard prediction:\", prediction)"
       ]
     },
     {
@@ -616,29 +627,6 @@
       "source": [
         "# @markdown Define input message in conversation and get output message from model.\n",
         "\n",
-        "\n",
-        "from urllib.request import urlopen\n",
-        "\n",
-        "from google.cloud import storage\n",
-        "\n",
-        "IMAGE_GCS_PATH = \"input_image.jpg\"\n",
-        "\n",
-        "\n",
-        "def image_url_to_gcs(url):\n",
-        "    # Uploads image to GCS and returns GCS URL.\n",
-        "    storage_client = storage.Client()\n",
-        "    bucket_name = BUCKET_URI.replace(\"gs://\", \"\")\n",
-        "    bucket = storage_client.get_bucket(bucket_name)\n",
-        "    blob = bucket.blob(IMAGE_GCS_PATH)\n",
-        "\n",
-        "    with urlopen(url) as response:\n",
-        "        image_data = response.read()\n",
-        "        image_type = response.info().get_content_type()\n",
-        "        blob.upload_from_string(image_data, content_type=image_type)\n",
-        "\n",
-        "    return os.path.join(BUCKET_URI, IMAGE_GCS_PATH)\n",
-        "\n",
-        "\n",
         "user_image = \"https://upload.wikimedia.org/wikipedia/commons/thumb/c/cb/The_Blue_Marble_%28remastered%29.jpg/580px-The_Blue_Marble_%28remastered%29.jpg\"  # @param {type: \"string\"}\n",
         "user_message = \"What is in the image?\"  # @param {type: \"string\"}\n",
         "\n",
@@ -654,23 +642,12 @@
         "        ],\n",
         "    }\n",
         "]\n",
-        "messages_gcs = [\n",
-        "    {\n",
-        "        \"role\": \"user\",\n",
-        "        \"content\": [\n",
-        "            {\n",
-        "                \"type\": \"image_url\",\n",
-        "                \"image_url\": {\"url\": image_url_to_gcs(user_image)},\n",
-        "            },\n",
-        "            {\"type\": \"text\", \"text\": user_message},\n",
-        "        ],\n",
-        "    }\n",
-        "]\n",
-        "print(\"Conversation [turn 1]:\", messages_gcs)\n",
+        "\n",
+        "print(\"Conversation [turn 1]:\", messages)\n",
         "\n",
         "response = client.chat.completions.create(\n",
         "    model=LLAMA3_90B_VISION_INSTRUCT,\n",
-        "    messages=messages_gcs,\n",
+        "    messages=messages,\n",
         ")\n",
         "print(\"Response:\", response)\n",
         "\n",
@@ -680,13 +657,8 @@
         "        \"content\": response.choices[0].message.content,\n",
         "    }\n",
         ")\n",
-        "messages_gcs.append(\n",
-        "    {\n",
-        "        \"role\": response.choices[0].message.role,\n",
-        "        \"content\": response.choices[0].message.content,\n",
-        "    }\n",
-        ")\n",
-        "print(\"Conversation [turn 2]:\", messages_gcs)"
+        "\n",
+        "print(\"Conversation [turn 2]:\", messages)"
       ]
     },
     {
@@ -713,8 +685,133 @@
         "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
         ")\n",
         "\n",
-        "prediction = response.predictions[0]\n",
-        "print(\"Llama Guard prediction:\", prediction[0][\"message\"][\"content\"])"
+        "prediction = response.predictions[\"choices\"][0][\"message\"][\"content\"]\n",
+        "print(\"Llama Guard prediction:\", prediction)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "h_LgDtVyO13s"
+      },
+      "source": [
+        "## Use the Llama Guard models to safeguard LLM vision inputs and outputs with the Vertex Llama 4 API service\n",
+        "\n",
+        "We use [meta-llama/Llama-Guard-4-12B](https://huggingface.co/meta-llama/Llama-Guard-4-12B) to safeguard input and output conversations with the [Llama 4 model API service on Vertex](https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama-4-maverick-17b-128e-instruct-maas)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "4hgRrEuqO13s"
+      },
+      "outputs": [],
+      "source": [
+        "!pip install --upgrade --quiet openai"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "zw5BkBd4O13s"
+      },
+      "outputs": [],
+      "source": [
+        "import google.auth\n",
+        "import openai\n",
+        "\n",
+        "# @markdown Set up the Llama 4 model API service.\n",
+        "\n",
+        "# Programmatically get an access token\n",
+        "creds, _ = google.auth.default(\n",
+        "    scopes=[\"https://www.googleapis.com/auth/cloud-platform\"]\n",
+        ")\n",
+        "auth_req = google.auth.transport.requests.Request()\n",
+        "creds.refresh(auth_req)\n",
+        "# Note: the credential lives for 1 hour by default (https://cloud.google.com/docs/authentication/token-types#at-lifetime); after expiration, it must be refreshed.\n",
+        "\n",
+        "client = openai.OpenAI(\n",
+        "    base_url=f\"https://us-east5-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/{REGION}/endpoints/openapi\",\n",
+        "    api_key=creds.token,\n",
+        ")\n",
+        "LLAMA4_MODEL_ID = \"meta/llama-4-scout-17b-16e-instruct-maas\"  # @param [\"meta/llama-4-scout-17b-16e-instruct-maas\", \"meta/llama-4-maverick-17b-128e-instruct-maas\"]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "3xX8VqWFO13s"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown Define input message in conversation and get output message from model.\n",
+        "\n",
+        "user_image = \"https://upload.wikimedia.org/wikipedia/commons/thumb/c/cb/The_Blue_Marble_%28remastered%29.jpg/580px-The_Blue_Marble_%28remastered%29.jpg\"  # @param {type: \"string\"}\n",
+        "user_message = \"What is in the image?\"  # @param {type: \"string\"}\n",
+        "\n",
+        "messages = [\n",
+        "    {\n",
+        "        \"role\": \"user\",\n",
+        "        \"content\": [\n",
+        "            {\n",
+        "                \"type\": \"image_url\",\n",
+        "                \"image_url\": {\"url\": user_image},\n",
+        "            },\n",
+        "            {\"type\": \"text\", \"text\": user_message},\n",
+        "        ],\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "print(\"Conversation [turn 1]:\", messages)\n",
+        "\n",
+        "response = client.chat.completions.create(\n",
+        "    model=LLAMA4_MODEL_ID,\n",
+        "    messages=messages,\n",
+        ")\n",
+        "print(\"Response:\", response)\n",
+        "\n",
+        "messages.append(\n",
+        "    {\n",
+        "        \"role\": response.choices[0].message.role,\n",
+        "        \"content\": response.choices[0].message.content,\n",
+        "    }\n",
+        ")\n",
+        "\n",
+        "print(\"Conversation [turn 2]:\", messages)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "6zhDnfAcO13s"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown Use Llama Guard to classify the conversation: safe versus unsafe.\n",
+        "# @markdown Classification is performed on the last turn of the conversation.\n",
+        "# @markdown If the content is safe, the model will return `safe`. If the content is unsafe, the model will return `unsafe` and additionally the list of offending categories as a comma-separated list in a new line.\n",
+        "# @markdown Set `\"@requestFormat\": \"chatCompletions\"` to use the OpenAI chat completions format.\n",
+        "\n",
+        "instances = [\n",
+        "    {\n",
+        "        \"messages\": messages,\n",
+        "        \"@requestFormat\": \"chatCompletions\",\n",
+        "    },\n",
+        "]\n",
+        "response = endpoints[\"vllm_gpu\"].predict(\n",
+        "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
+        ")\n",
+        "\n",
+        "prediction = response.predictions[\"choices\"][0][\"message\"][\"content\"]\n",
+        "print(\"Llama Guard prediction:\", prediction)"
       ]
     },
     {
@@ -745,11 +842,7 @@
         "\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"
+        "    model.delete()"
       ]
     }
   ],

notebooks/community/model_garden/model_garden_mammut.ipynb

@@ -39,7 +39,7 @@
         "  </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_mammut.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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"

notebooks/community/model_garden/model_garden_mediapipe_face_stylizer.ipynb

@@ -40,7 +40,7 @@
         "  </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_mediapipe_face_stylizer.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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"
@@ -110,7 +110,7 @@
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, 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",
@@ -140,6 +140,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",

notebooks/community/model_garden/model_garden_mediapipe_gesture_recognition.ipynb

@@ -40,7 +40,7 @@
         "\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_mediapipe_gesture_recognition.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",

notebooks/community/model_garden/model_garden_mediapipe_image_classification.ipynb

@@ -40,7 +40,7 @@
         "\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_mediapipe_image_classification.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",

notebooks/community/model_garden/model_garden_mediapipe_image_generation.ipynb

@@ -41,7 +41,7 @@
         "\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_mediapipe_image_generation.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",

notebooks/community/model_garden/model_garden_mediapipe_object_detection.ipynb

@@ -40,7 +40,7 @@
         "  </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_mediapipe_object_detection.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",
+        "      <img alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "</tr></tbody></table>"
@@ -104,7 +104,7 @@
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, 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",
@@ -136,6 +136,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",