Merge branch 'main' into ocr

small fixed
adding mistral ocr (#4069 )
2026-09-27 15:42:05 +00:00 · 2025-05-23 13:31:37 -04:00 · 2025-05-23 17:27:59 +00:00 · 2025-05-23 16:03:27 +00:00 · 2025-05-23 15:42:46 +00:00 · 2025-05-23 15:24:06 +00:00
113 changed files with 16908 additions and 3874 deletions
@@ -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
@@ -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
@@ -333,6 +333,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.
@@ -356,7 +357,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]


@@ -366,6 +369,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.
@@ -380,7 +384,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")


@@ -389,6 +395,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.
@@ -400,7 +407,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")


@@ -409,6 +418,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.
@@ -420,7 +430,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")


@@ -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:
@@ -32,6 +32,7 @@ class DockerCommandBuilder(CommandBuilder):
    super().__init__()
    self._docker_uri = [docker_uri]
    self.privilege_mode = []
+    self.entrypoint = []

    self._defaults = [
        'docker',
@@ -62,6 +63,9 @@ class DockerCommandBuilder(CommandBuilder):
  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
@@ -69,6 +73,7 @@ class DockerCommandBuilder(CommandBuilder):
        + self._mount_maps
        + self.privilege_mode
        + self._docker_uri
+        + self.entrypoint
    )


@@ -85,3 +90,6 @@ class PythonCommandBuilder(CommandBuilder):
  def build_cmd(self) -> str:
    os.environ.update(self._env_vars)
    return self._defaults
+
+  def add_entrypoint(self, entrypoint: list[str]):
+    self._defaults = entrypoint
@@ -3,6 +3,7 @@
 import copy
 import dataclasses
 import datetime
+import inspect
 import os
 import signal
 import subprocess
@@ -11,7 +12,7 @@ from absl import flags
 from absl import logging
 from absl.testing import parameterized
 import command_builder
-import frozendict
+import immutabledict
 import torch

 _DOCKER_URI = flags.DEFINE_string('docker_uri', None, 'docker image uri')
@@ -33,19 +34,19 @@ _LOCAL_OUTPUT_DIR = flags.DEFINE_string(

 _GCS_INPUT_DIR = flags.DEFINE_string(
    'gcs_input_dir',
-    'gs://peft-docker-test',
+    'gs://vmg-tuning-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',
+    'gs://vmg-tuning-docker-test/output',
    'GCS directory that stores test output.',
 )

 _GCS_TESTDATA_DIR = 'peft-train-image-test'

-_THROUGHPUT_TEST_EXCEPTIONS = frozendict.frozendict({
+_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,
@@ -101,26 +102,25 @@ class TestBase(parameterized.TestCase):
    return self.command_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
+    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 local_input_dir(self):
-    return _LOCAL_INPUT_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():
@@ -157,13 +157,44 @@ def get_test_data_path(name: str, download: bool = True) -> str:

  local_data = os.path.join(_LOCAL_INPUT_DIR.value, name)
  if not os.path.exists(local_data):
-    download_from_gcs(
-        os.path.join(_GCS_INPUT_DIR.value, name), _LOCAL_INPUT_DIR.value
-    )
+    # 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:
@@ -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)
@@ -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
@@ -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)
@@ -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}%"
@@ -11,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):
@@ -75,13 +75,15 @@ class TrainerStatsCallback(TrainerCallback):
            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: Throughput: %.2f token/s. %s, %s',
          throughput,
-          utils.gpu_stats_str(gpu_stats),
-          utils.cpu_stats_str(),
+          device_stats.gpu_stats_str(gpu_stats),
+          device_stats.cpu_stats_str(),
      )

  def on_train_begin(
@@ -95,8 +97,8 @@ class TrainerStatsCallback(TrainerCallback):
      self._start_time = time.time()
      logging.info(
          'on_train_begin: %s, %s',
-          utils.gpu_stats_str(),
-          utils.cpu_stats_str(),
+          device_stats.gpu_stats_str(),
+          device_stats.cpu_stats_str(),
      )

  def on_train_end(
@@ -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
@@ -16,6 +16,7 @@ diffusers==0.25.1
 evaluate==0.4.3
 fsspec==2024.3.1
 gcsfs==2024.3.1
+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
@@ -68,10 +68,12 @@ RUN mkdir -p ./vertex_vision_model_garden_peft/
 COPY model_oss/peft/train/vmg/configs/* ./vertex_vision_model_garden_peft/
 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/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/
@@ -37,7 +37,6 @@ class EvalConfig:
    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.
-    num_fewshot: The number of few-shot examples to use 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.
@@ -50,7 +49,6 @@ class EvalConfig:

  steps: int
  per_device_batch_size: int
-  num_fewshot: int | None
  limit: float | None
  metric_name: Sequence[str]
  tokenize_dataset: bool
@@ -99,7 +97,7 @@ def create_trainer(
  kwargs["tokenizer"] = tokenizer

  try:
-    eval_dataset = dataset_validation_util.load_dataset_with_template(
+    _, eval_dataset = dataset_validation_util.load_dataset_with_template(
        dataset_name=eval_config.dataset_path,
        split=eval_config.split,
        input_column=eval_config.column,
@@ -1,7 +1,6 @@
 """Instruct/Chat with LoRA models."""

 from collections.abc import Callable, Mapping, Sequence
-import dataclasses
 import datetime
 import json
 import os
@@ -23,6 +22,8 @@ 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
@@ -231,11 +232,6 @@ _PER_DEVICE_EVAL_BATCH_SIZE = flags.DEFINE_integer(
    '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',
@@ -255,8 +251,7 @@ _EVAL_METRIC_NAME = flags.DEFINE_list(
 _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 or path.',
+    'The Hugging Face dataset name or path to use for evaluation.',
 )

 # We set the default eval split as `test`, based on observation from
@@ -264,13 +259,13 @@ _EVAL_DATASET = 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.',
 )
@@ -278,7 +273,7 @@ _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(
@@ -576,8 +571,8 @@ def finetune_instruct(
  """Finetunes instruct."""
  logging.info(
      'on entering instruct_lora, %s,\n%s',
-      utils.gpu_stats_str(),
-      utils.cpu_stats_str(),
+      device_stats.gpu_stats_str(),
+      device_stats.cpu_stats_str(),
  )
  gradient_checkpointing_kwargs = {}
  # DDP provides limited support with the reentrant variant of gradient
@@ -594,7 +589,7 @@ def finetune_instruct(
      access_token=access_token,
  )

-  train_dataset_with_template = (
+  train_dataset, train_dataset_with_template = (
      dataset_validation_util.load_dataset_with_template(
          train_dataset,
          split=train_split,
@@ -621,18 +616,20 @@ def finetune_instruct(
          'getting tuning data stats with effective batch size %s',
          effective_batch_size,
      )
-      train_dataset_stats = utils.get_dataset_stats(
-          train_dataset_with_template,
-          tokenizer,
-          train_column,
-          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_name_or_path=pretrained_model_name_or_path,
@@ -663,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.
@@ -840,7 +839,6 @@ def main(unused_argv: Sequence[str]) -> None:
  if _EVAL_DATASET.value:
    eval_config = eval_lib.EvalConfig(
        per_device_batch_size=_PER_DEVICE_EVAL_BATCH_SIZE.value,
-        num_fewshot=_EVAL_NUM_FEWSHOT.value,
        limit=_EVAL_LIMIT.value,
        metric_name=_EVAL_METRIC_NAME.value,
        steps=_EVAL_STEPS.value,
@@ -31,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,
@@ -86,10 +86,19 @@ def main(unused_argv: Sequence[str]) -> None:
        )
    )

+  finetuned_lora_model_dir = fileutils.force_gcs_path(
+      _FINETUNED_LORA_MODEL_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_name_or_path=pretrained_model_name_or_path,
      precision_mode=_MERGE_MODEL_PRECISION_MODE.value,
-      finetuned_lora_model_dir=_FINETUNED_LORA_MODEL_DIR.value,
+      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,
  )
@@ -1,232 +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_NAME_OR_PATH = flags.DEFINE_string(
-    "pretrained_model_name_or_path",
-    None,
-    "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`.",
-)
-
-_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_TRAIN_EPOCHS = flags.DEFINE_integer(
-    "num_train_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_name_or_path: str,
-    dataset_name: str,
-    output_dir: str,
-    lora_rank: int = 8,
-    lora_alpha: int = 16,
-    lora_dropout: float = 0.1,
-    num_train_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_name_or_path for k in ("gpt", "opt", "bloom")):
-    padding_side = "left"
-  else:
-    padding_side = "right"
-
-  tokenizer = AutoTokenizer.from_pretrained(
-      pretrained_model_name_or_path, 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_name_or_path, 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_train_epochs),
-      num_training_steps=(len(train_dataloader) * num_train_epochs),
-  )
-
-  model.to(device)
-  for epoch in range(num_train_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_NAME_OR_PATH.value):
-    pretrained_model_name_or_path = (
-        dataset_validation_util.download_gcs_uri_to_local(
-            _PRETRAINED_MODEL_NAME_OR_PATH.value
-        )
-    )
-  else:
-    pretrained_model_name_or_path = _PRETRAINED_MODEL_NAME_OR_PATH.value
-  pretrained_model_path = dataset_validation_util.force_gcs_fuse_path(
-      pretrained_model_name_or_path
-  )
-  output_dir = dataset_validation_util.force_gcs_fuse_path(_OUTPUT_DIR.value)
-
-  finetune_sequence_classification(
-      pretrained_model_name_or_path=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_train_epochs=int(_NUM_TRAIN_EPOCHS.value),
-      batch_size=_BATCH_SIZE.value,
-      learning_rate=_LEARNING_RATE.value,
-  )
-
-
-if __name__ == "__main__":
-  app.run(main)
@@ -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"
+}
@@ -64,6 +64,7 @@ class TrainerThroughputTest(test_util.TestBase):
          'Mistral-7B-v0.1',
          'Mixtral-8x7B-v0.1',
          'gemma-2-9b-it',
+          'Qwen2.5-32B-Instruct',
      ],
      precision=['4bit', '8bit', 'bfloat16'],
      max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
@@ -89,6 +90,7 @@ class TrainerThroughputTest(test_util.TestBase):
          'Mistral-7B-v0.1',
          'Mixtral-8x7B-v0.1',
          'gemma-2-9b-it',
+          'Qwen2.5-32B-Instruct',
      ],
      precision=['4bit', '8bit', 'bfloat16'],
      max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
@@ -119,7 +121,11 @@ class TrainerThroughputTest(test_util.TestBase):
    self.assertEqual(self.run_cmd_and_handle_failure(), 0)

  @parameterized.product(
-      model_name=['llama3.1-8b-hf', '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],
@@ -136,9 +142,16 @@ 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.command_builder.add_env_var(
        'CUDA_VISIBLE_DEVICES', ','.join([str(x) for x in range(0, num_gpus)])
@@ -140,6 +140,70 @@ class TrainedModelQualityTest(test_util.TestBase):

    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()
@@ -9,19 +9,17 @@ environment. Otherwise, `python3` is used.

 import argparse
 from collections.abc import MutableSequence, Sequence
-import json
 import multiprocessing
-import os
 import subprocess
 import sys
 from absl import app
 from absl import flags
 from absl import logging
 from util import dataset_validation_util
-from vertex_vision_model_garden_peft.train.vmg import gcs_syncer
+from util import cluster_spec
 from vertex_vision_model_garden_peft.train.vmg import utils
 from util import constants
-from util import fileutils
+from util import gcs_syncer
 from util import hypertune_utils


@@ -41,9 +39,12 @@ _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.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.MERGE_CAUSAL_LANGUAGE_MODEL_LORA: (
+        'vertex_vision_model_garden_peft/train/vmg/merge_causal_language_model_lora.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',
 }

@@ -71,53 +72,17 @@ 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

@@ -131,45 +96,6 @@ def _append_args_to_command_in_place(
      command.append(f'--{key}={value}')


-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_training_path(path: str, node_rank: int) -> 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. Othereise, 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,
-        dataset_validation_util.force_gcs_fuse_path(path)[1:],
-    )
-    gcs_dir = fileutils.force_gcs_path(path)
-    os.makedirs(local_dir, exist_ok=True)
-  return local_dir, os.path.join(gcs_dir, f'node-{node_rank}')
-
-
 def _get_train_and_maybe_merge_cmd_and_dirs_to_sync(
    task_type: str, config_file: str, unknown: Sequence[str]
 ) -> Sequence[Sequence[str]]:
@@ -203,11 +129,11 @@ def _get_train_and_maybe_merge_cmd_and_dirs_to_sync(
      dataset_validation_util.force_gcs_fuse_path(training_args.output_dir)
  )

-  local_output_dir, gcs_output_dir = _manage_training_path(
+  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 _is_gcs_or_gcsfuse_path(gcs_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.
@@ -217,11 +143,11 @@ def _get_train_and_maybe_merge_cmd_and_dirs_to_sync(
  merge_args, unknown = merge_parser.parse_known_args(unknown)

  if merge_args.merge_base_and_lora_output_dir:
-    merge_local_dir, merge_gcs_dir = _manage_training_path(
-        merge_args.merge_base_and_lora_output_dir, node_rank
+    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 _is_gcs_or_gcsfuse_path(merge_gcs_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.
@@ -239,8 +165,10 @@ def _get_train_and_maybe_merge_cmd_and_dirs_to_sync(
  # Only the main node runs merging.
  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 = _manage_training_path(lora_dir, node_rank)
-    if _is_gcs_or_gcsfuse_path(lora_gcs_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 = [
@@ -263,46 +191,37 @@ def _get_train_and_maybe_merge_cmd_and_dirs_to_sync(
  return commands, dirs_to_sync


-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.
+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:
-    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.
+    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 GCS rsync process.
+    The bash commands to execute and the directories to sync.
  """
-  rsync_process = multiprocessing.Process(
-      target=gcs_syncer.start_gcs_rsync,
-      args=(dirs_to_sync, mp_queue, gcs_rsync_interval_secs),
-  )
-  rsync_process.start()
-  return rsync_process
+  # 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)

-
-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('training finished')
-  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
+  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:
@@ -335,6 +254,10 @@ def main(unused_argv: Sequence[str]) -> None:
    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
    cmd = launch_script_cmd(_TASK_TO_SCRIPT[task], args.config_file)
@@ -344,7 +267,7 @@ def main(unused_argv: Sequence[str]) -> None:
  rsync_process = None
  mp_queue = multiprocessing.Queue(maxsize=1)
  if dirs_to_sync:
-    rsync_process = _setup_gcs_rsync(
+    rsync_process = gcs_syncer.setup_gcs_rsync(
        dirs_to_sync, mp_queue, args.gcs_rsync_interval_secs
    )

@@ -361,7 +284,7 @@ def main(unused_argv: Sequence[str]) -> None:
        rsync_process.terminate()
      raise e
  if rsync_process is not None:
-    _cleanup_gcs_rsync(rsync_process, mp_queue)
+    gcs_syncer.cleanup_gcs_rsync(rsync_process, mp_queue)


 if __name__ == '__main__':
@@ -1,7 +1,6 @@
 """Common libraries for PEFT."""

 from collections.abc import Mapping, Sequence
-import dataclasses
 import datetime
 import gc
 import os
@@ -11,12 +10,9 @@ from absl import logging
 import accelerate
 from accelerate import DistributedType
 from accelerate import PartialState
-import numpy as np
 import peft
 from peft import PeftModel
 from peft import prepare_model_for_kbit_training
-import psutil
-import pynvml
 import torch
 import transformers
 from transformers import AutoModelForCausalLM
@@ -28,7 +24,6 @@ import trl
 from util import dataset_validation_util
 from util import constants

-
 _LLAMA_3_1_405B_MODEL_ID = "Meta-Llama-3.1-405B"
 _LOCAL_MERGED_MODEL_DIR = "/tmp/merged_model"
 _GEMMA2_MODEL = "gemma-2"
@@ -126,7 +121,7 @@ def load_model(
      "device_map": device_map,
      "torch_dtype": torch_dtype,
      "quantization_config": quantization_config,
-      "trust_remote_code": True,
+      "trust_remote_code": False,
      "token": access_token,
      "attn_implementation": attn_implementation,
  }
@@ -310,171 +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()


-@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: GpuStats | None = 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}%"
-  )
-
-
-@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}%"
-
-
 def init_partial_state(
    timeout: datetime.timedelta = datetime.timedelta(seconds=600),
 ) -> None:
@@ -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,96 @@ 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.

@@ -352,3 +445,15 @@ def get_output_video_file(video_output_file_path: str) -> str:
      file_extension, '_overlay' + file_extension
  )
  return out_local_video_file_name
+
+
+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)
@@ -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
@@ -1,17 +1,107 @@
 """Sync local directory to GCS directory using rsync."""

-from collections.abc import Sequence
 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.

@@ -57,7 +147,7 @@ def _rsync_local_to_gcs(local_dir: str, gcs_dir: str) -> None:


 def start_gcs_rsync(
-    dirs_to_sync: Sequence[tuple[str, str]],
+    dirs_to_sync: Sequence[Tuple[str, str]],
    mp_queue: multiprocessing.Queue,
    gcs_rsync_interval_secs: int,
 ) -> None:
@@ -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
@@ -57,12 +57,21 @@
      "source": [
        "# Overview\n",
        "\n",
-        "This notebook will guide you through the initial step of testing your recently deployed model with text prompts. Depending on your deployed model's inference setup, the notebook utilizes either Text Generation Inference [TGI](https://huggingface.co/docs/text-generation-inference/en/index) or [vLLM](https://developers.googleblog.com/en/inference-with-gemma-using-dataflow-and-vllm/#:~:text=model%20frameworks%20simple.-,What%20is%20vLLM%3F,-vLLM%20is%20an), two efficient serving frameworks that enhance the performance of your GPU model. Ready to see your deployed model respond? Run the cells below and start experimenting with different prompts!\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 using the Google Cloud Console. You can find an overview of AI and Machine Learning services on [GKE AI/ML](https://console.cloud.google.com/kubernetes/aiml/overview).\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",
@@ -70,33 +79,46 @@
        "\n",
        "### GPUs\n",
        "\n",
-        "GPUs let you accelerate specific workloads running on your nodes, such as machine learning and data processing. GKE provides a range of machine type options for node configuration, including machine types with NVIDIA H100, L4, and A100 GPUs.\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: vLLM or Text Generation Inference (TGI). The following sections provide a brief overview of each to give you context on the underlying technology powering your model.\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 serving throughput on GPUs. TGI includes features such as:\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",
+        "*   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 [TGI documentation](https://github.com/huggingface/text-generation-inference/blob/main/README.md)\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 known for its high throughput and efficiency, and it leverages PagedAttention. Key features include:\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 dynamic workloads.\n",
-        "* Continuous batching: Maximizes GPU utilization by batching incoming requests.\n",
-        "* High-throughput serving: Designed for production-level serving with low latency.\n",
-        "* Optimized CUDA kernels.\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 [vLLM documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/vllm/use-vllm)"
+        "To learn more, refer to the\n",
+        "[vLLM documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/vllm/use-vllm)"
      ]
    },
    {
@@ -111,9 +133,9 @@
      "source": [
        "# @title # Connect to Google Cloud Project\n",
        "# @markdown #### Run this cell to configure your Google Cloud environment for Kubernetes (GKE) operations.\n",
-        "\n",
+        "# @markdown\n",
        "# @markdown #### Actions:\n",
-        "# @markdown 1.  **Connects to Project & Region:** Retrieves and sets your Google Cloud project ID and region.\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",
@@ -121,9 +143,6 @@
        "# 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",
        "# Set up gcloud.\n",
        "! gcloud config set project \"$PROJECT_ID\"\n",
        "! gcloud services enable container.googleapis.com\n",
@@ -144,124 +163,406 @@
      "outputs": [],
      "source": [
        "# @title # Select Cluster and Deployment { vertical-output: true }\n",
-        "\n",
-        "# @markdown ## Instruction:\n",
-        "\n",
-        "# @markdown This cell provides interactive dropdown menus to select a Google Kubernetes Engine (GKE) cluster and a deployment within that cluster.\n",
-        "\n",
-        "# @markdown ***Please select a cluster and deployment before proceeding.***\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 display\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",
-        "def get_clusters(p, r):\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",
-        "        return (\n",
-        "            subprocess.run(\n",
-        "                [\n",
-        "                    \"gcloud\",\n",
-        "                    \"container\",\n",
-        "                    \"clusters\",\n",
-        "                    \"list\",\n",
-        "                    \"--project\",\n",
-        "                    p,\n",
-        "                    \"--region\",\n",
-        "                    r,\n",
-        "                    \"--format=value(name)\",\n",
-        "                ],\n",
-        "                capture_output=True,\n",
-        "                text=True,\n",
-        "                check=True,\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",
-        "            .stdout.strip()\n",
-        "            .split(\"\\n\")\n",
-        "        )\n",
-        "    except subprocess.CalledProcessError as e:\n",
-        "        print(f\"Error: {e}\")\n",
-        "        return []\n",
+        "        return None\n",
        "\n",
        "\n",
-        "def get_deployments(c, r):\n",
-        "    try:\n",
-        "        subprocess.run(\n",
-        "            [\n",
-        "                \"gcloud\",\n",
-        "                \"container\",\n",
-        "                \"clusters\",\n",
-        "                \"get-credentials\",\n",
-        "                c,\n",
-        "                \"--location\",\n",
-        "                r,\n",
-        "            ],\n",
-        "            capture_output=True,\n",
-        "            text=True,\n",
-        "            check=True,\n",
-        "        )\n",
-        "        deployments = json.loads(\n",
-        "            subprocess.run(\n",
-        "                [\"kubectl\", \"get\", \"deployments\", \"-o\", \"json\"],\n",
-        "                capture_output=True,\n",
-        "                text=True,\n",
-        "                check=True,\n",
-        "            ).stdout\n",
-        "        )\n",
-        "        return [i[\"metadata\"][\"name\"] for i in deployments[\"items\"]]\n",
-        "    except subprocess.CalledProcessError as e:\n",
-        "        print(f\"Error: {e}\")\n",
-        "        return []\n",
-        "\n",
-        "\n",
-        "def create_deployment_dropdown(cluster_name, region, on_select_deployment):\n",
-        "    deployments = get_deployments(cluster_name, region)\n",
-        "    deployments_with_prompt = [\"Select Deployment\"] + deployments\n",
-        "    deployment_dropdown = widgets.Dropdown(\n",
-        "        options=deployments_with_prompt,\n",
-        "        description=\"Deployments\",\n",
-        "        disabled=False,\n",
-        "        width=\"4000px\",\n",
-        "    )\n",
-        "    deployment_dropdown.observe(\n",
-        "        lambda c: on_select_deployment(c[\"new\"])\n",
-        "        if c[\"type\"] == \"change\" and c[\"name\"] == \"value\"\n",
-        "        else None,\n",
-        "        names=\"value\",\n",
-        "    )\n",
-        "    return deployment_dropdown\n",
-        "\n",
-        "\n",
-        "def on_deployment_select(deployment_name):\n",
+        "# --- Event Handlers ---\n",
+        "def on_deployment_select(change):\n",
+        "    \"\"\"Handles changes in the deployment selection.\"\"\"\n",
        "    global SELECTED_DEPLOYMENT\n",
-        "    SELECTED_DEPLOYMENT = deployment_name\n",
-        "    print(f\"Selected deployment: {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",
-        "        if change[\"new\"] == \"Select Cluster\":\n",
-        "            return\n",
-        "        deployment_dropdown = create_deployment_dropdown(\n",
-        "            change[\"new\"], REGION, on_deployment_select\n",
-        "        )\n",
-        "        display(deployment_dropdown)\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",
-        "clusters = get_clusters(PROJECT_ID, REGION)\n",
-        "if clusters:\n",
-        "    # @markdown Run this cell to display the Cluster dropdown menu:\n",
-        "    clusters_with_prompt = [\"Select Cluster\"] + clusters\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, description=\"Clusters\", disabled=False\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",
-        "    display(cluster_dropdown)\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",
-        "    print(f\"No clusters found in {PROJECT_ID}/{REGION}.\")"
+        "    # 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)"
      ]
    },
    {
@@ -274,91 +575,202 @@
      },
      "outputs": [],
      "source": [
-        "# @title # Chat completion for text-only models {run:\"auto\", vertical-output: true}\n",
-        "\n",
+        "# @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",
-        "\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",
-        "\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",
-        "from IPython.display import HTML\n",
+        "import ipywidgets as widgets\n",
        "\n",
        "\n",
-        "def get_deployment_pod_name(deployment):\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",
-        "        label = deployment + \"-app\"\n",
-        "        pods = json.loads(\n",
-        "            subprocess.run(\n",
-        "                [\"kubectl\", \"get\", \"pods\", \"-o\", \"json\", \"-l\", f\"app={label}\"],\n",
-        "                capture_output=True,\n",
-        "                check=True,\n",
-        "            ).stdout\n",
-        "        )\n",
-        "        return pods[\"items\"][0][\"metadata\"][\"name\"] if pods[\"items\"] else None\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",
-        "        KeyError,\n",
        "        IndexError,\n",
-        "    ):\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_vllm_label(pod_name):\n",
-        "    \"\"\"Checks if the pod has the 'ai.gke.io/inference-server=vllm' label.\"\"\"\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",
-        "        result = subprocess.run(\n",
-        "            [\"kubectl\", \"get\", \"pod\", pod_name, \"-o\", \"json\"],\n",
-        "            capture_output=True,\n",
-        "            check=True,\n",
-        "        )\n",
-        "        labels = json.loads(result.stdout)[\"metadata\"][\"labels\"]\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, KeyError, json.JSONDecodeError):\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):\n",
-        "    response = !kubectl exec -t {pod_name} -- curl -X POST http://{pod_endpoint}/generate -H \"Content-Type: application/json\" -d '{json.dumps(request)}' 2> /dev/null\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",
-        "        data = json.loads(response[0])\n",
-        "        if is_vllm:\n",
-        "            return data[\"predictions\"][0]\n",
-        "        else:\n",
-        "            return data[\"generated_text\"]\n",
-        "    except (json.JSONDecodeError, KeyError, IndexError) as e:\n",
-        "        return f\"Error: {e}, Raw: {response}\"\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",
-        "deployment_pod = get_deployment_pod_name(SELECTED_DEPLOYMENT)\n",
-        "is_vllm_inference = check_vllm_label(deployment_pod)\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",
-        "user_prompt = \"What is AI?\"  # @param {type: \"string\"}\n",
-        "temperature = 0.50  # @param {type: \"number\"}\n",
-        "max_tokens = 250  # @param {type: \"number\"}\n",
        "\n",
-        "request = {\n",
-        "    \"max_tokens\": 250 if max_tokens is None else max_tokens,\n",
-        "    \"temperature\": 0.5 if temperature is None else temperature,\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",
-        "if is_vllm_inference:\n",
-        "    request[\"prompt\"] = user_prompt\n",
-        "else:\n",
-        "    request[\"inputs\"] = user_prompt\n",
+        "        print(\n",
+        "            f\"Target: {SELECTED_DEPLOYMENT} in {SELECTED_NAMESPACE}. \\n\\nRequesting response...\"\n",
+        "        )\n",
        "\n",
-        "model_service = SELECTED_DEPLOYMENT + \"-service\"\n",
-        "output = !kubectl get endpoints {model_service}\n",
-        "pod_endpoint = output[1].split()[1]\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",
-        "# @markdown ### Response:\n",
-        "response = process_response(request, deployment_pod, pod_endpoint, is_vllm_inference)\n",
-        "HTML(\n",
-        "    '<div style=\"overflow-x: auto; font-size: 16px; line-height:'\n",
-        "    f' 1.8;\">{response}</div>'\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",
        ")"
      ]
    },
@@ -371,39 +783,67 @@
      "source": [
        "# Next Steps: Integrating the GKE Service Endpoint\n",
        "\n",
-        "After successfully deploying a model on Google Kubernetes Engine (GKE) and verifying it via a notebook, the next step is to integrate it into various applications. This involves making HTTP requests to the service's endpoint from your application code.\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 its service endpoint. Google Kubernetes Engine offers several ways to do this:\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 to your service. Set up Ingress for either an internal Load Balancer (accessible only within your VPC) or an external Load Balancer (accessible from the internet). [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 routing in Kubernetes. Similar to Ingress, Gateway API allows you to define how external and internal traffic should be directed to your services. [Explore GKE Gateway API](https://cloud.google.com/kubernetes-engine/docs/concepts/gateway-api).\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 Horizontal Pod Autoscaler (HPA). HPA automatically scales the number of Pods based on resource utilization or custom metrics, optimizing performance and cost. [See how to configure HPA](https://cloud.google.com/kubernetes-engine/docs/how-to/horizontal-pod-autoscaling).\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 Managed Service for Prometheus. Configure your model serving to expose Prometheus metrics for comprehensive insights. [Get started with Google Cloud Managed Prometheus](https://cloud.google.com/kubernetes-engine/docs/how-to/configure-automatic-application-monitoring).\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",
-        "   * Services: https://kubernetes.io/docs/concepts/services-networking/service/\n",
+        "*   #### Kubernetes Documentation:\n",
        "\n",
-        "* #### Google Cloud Documentation:\n",
-        "   * Google Kubernetes Engine (GKE): https://cloud.google.com/kubernetes-engine\n",
-        "   * Cloud Load Balancing: https://cloud.google.com/load-balancing/docs/ingress\n",
-        "   * Gateway API on GKE: https://cloud.google.com/kubernetes-engine/docs/concepts/gateway-api\n",
-        "   * Learn about GPUs in GKE: https://cloud.google.com/kubernetes-engine/docs/concepts/gpus\n",
+        "    *   Services:\n",
+        "        https://kubernetes.io/docs/concepts/services-networking/service/\n",
        "\n",
-        "* #### Python requests Library:\n",
-        "   * https://requests.readthedocs.io/en/latest/\n",
+        "*   #### Google Cloud Documentation:\n",
        "\n",
-        "* #### LangChain with Google Integrations:\n",
-        "   * The Langchain documentation is very useful: https://python.langchain.com/docs/integrations/providers/google/"
+        "    *   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/"
      ]
    }
  ],
@@ -0,0 +1,465 @@
+{
+  "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_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://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "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": "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",
+        "# @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",
+        "DEPLOYMENT_APP_LABEL = \"\"  # @param {type:\"string\", isTemplate:true}\n",
+        "\n",
+        "SERVICE = f\"{DEPLOYMENT}-service\"\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 fetch_cluster_credential(cluster, region, project_id):\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",
+        "        _run_kubectl(cred_cmd)\n",
+        "    except Exception as e:\n",
+        "        # Original code prints error and returns empty dict\n",
+        "        print(f\"Error fetching cluster credentials: {e}\")\n",
+        "        return {}\n",
+        "\n",
+        "\n",
+        "def get_deployment_pod_name(deployment, namespace, deployment_app_label):\n",
+        "    \"\"\"Finds the running pod name for a given deployment and namespace.\"\"\"\n",
+        "\n",
+        "    cmd = [\n",
+        "        \"kubectl\",\n",
+        "        \"get\",\n",
+        "        \"pods\",\n",
+        "        \"-n\",\n",
+        "        namespace,\n",
+        "        \"-o\",\n",
+        "        \"json\",\n",
+        "        \"-l\",\n",
+        "        f\"app={deployment_app_label}\",\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",
+        "\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 get_service_endpoint(service, namespace):\n",
+        "    \"\"\"Retrieve the service endpoint of the deployment\"\"\"\n",
+        "    endpoint_cmd = [\n",
+        "        \"kubectl\",\n",
+        "        \"get\",\n",
+        "        \"endpoints\",\n",
+        "        service,\n",
+        "        \"-n\",\n",
+        "        namespace,\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",
+        "            print(f\"Endpoint data incomplete for {service}.\")\n",
+        "            return None\n",
+        "        endpoint = endpoint_output[1].split()[\n",
+        "            1\n",
+        "        ]  # Assumes format: NAME ENDPOINTS AGE -> service ip:port,... age\n",
+        "        return endpoint\n",
+        "    except subprocess.CalledProcessError as e:\n",
+        "        print(f\"Error getting endpoints for {service}: {e}\")\n",
+        "        return None\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",
+        "\n",
+        "    json_data_escaped = json.dumps(request).replace(\"'\", \"'\\\\''\")\n",
+        "    curl_cmd = (\n",
+        "        f\"kubectl exec -n {namespace} -t {pod_name} -- curl -s -X POST\"\n",
+        "        f' http://{pod_endpoint}/generate -H \"Content-Type: application/json\"'\n",
+        "        f\" -d '{json_data_escaped}' 2> /dev/null\"\n",
+        "    )\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:  # vLLM format\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",
+        "    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",
+        "\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",
+        "    \"\"\"Handles the submit button click event.\"\"\"\n",
+        "    with output_area_response:\n",
+        "        clear_output()\n",
+        "\n",
+        "        fetch_cluster_credential(CLUSTER, REGION, PROJECT_ID)\n",
+        "\n",
+        "        # retrieve deployment pod\n",
+        "        pod_name = get_deployment_pod_name(DEPLOYMENT, NAMESPACE, DEPLOYMENT_APP_LABEL)\n",
+        "        if not pod_name:\n",
+        "            display(\n",
+        "                Markdown(f\"**Error:** Could not find running pod for `{DEPLOYMENT}`.\")\n",
+        "            )\n",
+        "            return\n",
+        "\n",
+        "        # build the request message\n",
+        "        is_vllm = check_inference_label(pod_name, 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",
+        "\n",
+        "        # retrieve service endpoint for the deployment\n",
+        "        endpoint = get_service_endpoint(SERVICE, NAMESPACE)\n",
+        "        if not endpoint:\n",
+        "            display(Markdown(f\"**Error getting endpoints for `{SERVICE}`:**\\n\"))\n",
+        "            return\n",
+        "\n",
+        "        # prompt test the deployment endpoint\n",
+        "        try:\n",
+        "            response = process_response(request, pod_name, endpoint, is_vllm, NAMESPACE)\n",
+        "            display(Markdown(f\"**Response:**\\n\\n{response}\"))\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_auto.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
@@ -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",
@@ -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",
@@ -118,10 +118,9 @@
      "cell_type": "code",
      "execution_count": null,
      "id": "L3dqbxovo5t6",
-      "language": "python",
      "metadata": {
        "cellView": "form",
-        "id": "L3dqbxovo5t6"
+        "id": "2b585189a670"
      },
      "outputs": [],
      "source": [
@@ -149,6 +148,8 @@
        "# 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",
@@ -224,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"
@@ -236,7 +237,7 @@
      "id": "BxlzWU2KQqmw",
      "metadata": {
        "cellView": "form",
-        "id": "BxlzWU2KQqmw"
+        "id": "798068fc0355"
      },
      "outputs": [],
      "source": [
@@ -275,7 +276,7 @@
      "cell_type": "markdown",
      "id": "JpNBJJgjWL7j",
      "metadata": {
-        "id": "JpNBJJgjWL7j"
+        "id": "10ed490e28e5"
      },
      "source": [
        "## Prefix Caching <a name=\"prefix-caching\"></a>\n",
@@ -303,7 +304,7 @@
      "cell_type": "markdown",
      "id": "9gZJ8cB27e1m",
      "metadata": {
-        "id": "9gZJ8cB27e1m"
+        "id": "30ddb93fdd7b"
      },
      "source": [
        "### Try out Prefix Caching with Hex-LLM\n",
@@ -319,7 +320,7 @@
      "id": "RpmoA2nXjdCd",
      "metadata": {
        "cellView": "form",
-        "id": "RpmoA2nXjdCd"
+        "id": "b56d82c1aa6f"
      },
      "outputs": [],
      "source": [
@@ -508,7 +509,7 @@
      "id": "5QoK8c0R9U3B",
      "metadata": {
        "cellView": "form",
-        "id": "5QoK8c0R9U3B"
+        "id": "96c5afed49b4"
      },
      "outputs": [],
      "source": [
@@ -519,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",
@@ -576,7 +575,7 @@
      "id": "29rn5ATmB2YC",
      "metadata": {
        "cellView": "form",
-        "id": "29rn5ATmB2YC"
+        "id": "9a95c9f90358"
      },
      "outputs": [],
      "source": [
@@ -649,7 +648,7 @@
      "cell_type": "markdown",
      "id": "KjbM8E9DGuuR",
      "metadata": {
-        "id": "KjbM8E9DGuuR"
+        "id": "12ad6d1ff725"
      },
      "source": [
        "#### Delete the models and endpoints"
@@ -661,7 +660,7 @@
      "id": "JpLU7GRQGuuR",
      "metadata": {
        "cellView": "form",
-        "id": "JpLU7GRQGuuR"
+        "id": "1ab4e3bb74b4"
      },
      "outputs": [],
      "source": [
@@ -687,7 +686,7 @@
      "cell_type": "markdown",
      "id": "XZ33HhYmOxCS",
      "metadata": {
-        "id": "XZ33HhYmOxCS"
+        "id": "7a8a9a1b2ddf"
      },
      "source": [
        "### Try out Prefix Caching with vLLM\n",
@@ -710,7 +709,7 @@
      "id": "E8OiHHNNE_wj",
      "metadata": {
        "cellView": "form",
-        "id": "E8OiHHNNE_wj"
+        "id": "4425cc0bdedc"
      },
      "outputs": [],
      "source": [
@@ -914,7 +913,7 @@
      "id": "zex1oXl36A70",
      "metadata": {
        "cellView": "form",
-        "id": "zex1oXl36A70"
+        "id": "bcbafec839cd"
      },
      "outputs": [],
      "source": [
@@ -922,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",
@@ -979,7 +976,7 @@
      "id": "gDOC_nfsJeUR",
      "metadata": {
        "cellView": "form",
-        "id": "gDOC_nfsJeUR"
+        "id": "e984f43422d5"
      },
      "outputs": [],
      "source": [
@@ -1052,7 +1049,7 @@
      "cell_type": "markdown",
      "id": "GdGxaTirJeUR",
      "metadata": {
-        "id": "GdGxaTirJeUR"
+        "id": "dff0d10dcc20"
      },
      "source": [
        "#### Delete the models and endpoints"
@@ -1064,7 +1061,7 @@
      "id": "OgoqXE-VJeUR",
      "metadata": {
        "cellView": "form",
-        "id": "OgoqXE-VJeUR"
+        "id": "5b8751773e7f"
      },
      "outputs": [],
      "source": [
@@ -1090,7 +1087,7 @@
      "cell_type": "markdown",
      "id": "w4Guijaw_NEs",
      "metadata": {
-        "id": "w4Guijaw_NEs"
+        "id": "863775857a46"
      },
      "source": [
        "### Best practices\n",
@@ -1105,7 +1102,7 @@
      "cell_type": "markdown",
      "id": "ml8fgoIQWSbY",
      "metadata": {
-        "id": "ml8fgoIQWSbY"
+        "id": "565cbdc3a06b"
      },
      "source": [
        "## Speculative Decoding <a name=\"spec-decoding\"></a>\n",
@@ -1147,7 +1144,7 @@
      "cell_type": "markdown",
      "id": "NmWRro8Q-Td6",
      "metadata": {
-        "id": "NmWRro8Q-Td6"
+        "id": "94eaa9050abb"
      },
      "source": [
        "### Try out Speculative Decoding with vLLM"
@@ -1159,7 +1156,7 @@
      "id": "72d1GlrYifKU",
      "metadata": {
        "cellView": "form",
-        "id": "72d1GlrYifKU"
+        "id": "5f358cc230a6"
      },
      "outputs": [],
      "source": [
@@ -1476,7 +1473,7 @@
      "id": "CNiItf5hdVFU",
      "metadata": {
        "cellView": "form",
-        "id": "CNiItf5hdVFU"
+        "id": "be3170e0e05a"
      },
      "outputs": [],
      "source": [
@@ -1502,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",
@@ -1553,7 +1548,7 @@
      "cell_type": "markdown",
      "id": "WahYGAZyq6Gl",
      "metadata": {
-        "id": "WahYGAZyq6Gl"
+        "id": "30c5d2535df3"
      },
      "source": [
        "## Clean up resources"
@@ -1563,7 +1558,7 @@
      "cell_type": "markdown",
      "id": "bV5Yjkgav9BZ",
      "metadata": {
-        "id": "bV5Yjkgav9BZ"
+        "id": "63c10917ff95"
      },
      "source": [
        "### Delete the models and endpoints"
@@ -1575,7 +1570,7 @@
      "id": "qsks36cOH9rb",
      "metadata": {
        "cellView": "form",
-        "id": "qsks36cOH9rb"
+        "id": "92892e1b1730"
      },
      "outputs": [],
      "source": [
@@ -123,8 +123,10 @@
        ")\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",
@@ -219,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",
@@ -337,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",
@@ -377,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",
@@ -312,9 +312,7 @@
        "    DEDICATED_ENDPOINT_DNS = endpoints[\n",
        "        \"my-endpoint\"\n",
        "    ].gca_resource.dedicated_endpoint_dns\n",
-        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
-        "    PROJECT_ID, REGION, endpoints[\"my-endpoint\"].name\n",
-        ")\n",
+        "ENDPOINT_RESOURCE_NAME = endpoints[\"my-endpoint\"].resource_name\n",
        "\n",
        "# @title Chat Completions Inference\n",
        "\n",
@@ -241,8 +241,7 @@
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "cellView": "form",
-        "id": "kg5MwMIfB9Uj"
+        "id": "I1u2FLa9XgVD"
      },
      "outputs": [],
      "source": [
@@ -252,9 +251,13 @@
        "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",
+        "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",
@@ -32,18 +32,18 @@
      "source": [
        "# Vertex AI Model Garden - Finetuning Tutorial\n",
        "\n",
-        "<table><tbody><tr>\n",
-        "  <td style=\"text-align: center\">\n",
-        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_finetuning_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_finetuning_tutorial.ipynb\">\n",
-        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
-        "    </a>\n",
-        "  </td>\n",
-        "</tr></tbody></table>"
+        "\u003ctable\u003e\u003ctbody\u003e\u003ctr\u003e\n",
+        "  \u003ctd style=\"text-align: center\"\u003e\n",
+        "    \u003ca 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_finetuning_tutorial.ipynb\"\u003e\n",
+        "      \u003cimg alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"\u003e\u003cbr\u003e Run in Colab Enterprise\n",
+        "    \u003c/a\u003e\n",
+        "  \u003c/td\u003e\n",
+        "  \u003ctd style=\"text-align: center\"\u003e\n",
+        "    \u003ca href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_finetuning_tutorial.ipynb\"\u003e\n",
+        "      \u003cimg alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"\u003e\u003cbr\u003e View on GitHub\n",
+        "    \u003c/a\u003e\n",
+        "  \u003c/td\u003e\n",
+        "\u003c/tr\u003e\u003c/tbody\u003e\u003c/table\u003e"
      ]
    },
    {
@@ -168,11 +168,11 @@
        "\n",
        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
-        "# @markdown 2. For finetuning, **[click here](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Frestricted_image_training_nvidia_a100_80gb_gpus)** to check if your project already has the required 8 Nvidia A100 80 GB GPUs in the us-central1 region. If yes, then run this notebook in the us-central1 region. If you do not have 8 Nvidia A100 80 GPUs or have more GPU requirements than this, then schedule your job with Nvidia H100 GPUs via Dynamic Workload Scheduler using [these instructions](https://cloud.google.com/vertex-ai/docs/training/schedule-jobs-dws). For Dynamic Workload Scheduler, check the [us-central1](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Fcustom_model_training_preemptible_nvidia_h100_gpus) or [europe-west4](https://console.cloud.google.com/iam-admin/quotas?location=europe-west4&metric=aiplatform.googleapis.com%2Fcustom_model_training_preemptible_nvidia_h100_gpus) quota for Nvidia H100 GPUs. If you do not have enough GPUs, then you can follow [these instructions](https://cloud.google.com/docs/quotas/view-manage#viewing_your_quota_console) to request quota.\n",
+        "# @markdown 2. For finetuning, **[click here](https://console.cloud.google.com/iam-admin/quotas?location=us-central1\u0026metric=aiplatform.googleapis.com%2Frestricted_image_training_nvidia_a100_80gb_gpus)** to check if your project already has the required 8 Nvidia A100 80 GB GPUs in the us-central1 region. If yes, then run this notebook in the us-central1 region. If you do not have 8 Nvidia A100 80 GPUs or have more GPU requirements than this, then schedule your job with Nvidia H100 GPUs via Dynamic Workload Scheduler using [these instructions](https://cloud.google.com/vertex-ai/docs/training/schedule-jobs-dws). For Dynamic Workload Scheduler, check the [us-central1](https://console.cloud.google.com/iam-admin/quotas?location=us-central1\u0026metric=aiplatform.googleapis.com%2Fcustom_model_training_preemptible_nvidia_h100_gpus) or [europe-west4](https://console.cloud.google.com/iam-admin/quotas?location=europe-west4\u0026metric=aiplatform.googleapis.com%2Fcustom_model_training_preemptible_nvidia_h100_gpus) quota for Nvidia H100 GPUs. If you do not have enough GPUs, then you can follow [these instructions](https://cloud.google.com/docs/quotas/view-manage#viewing_your_quota_console) to request quota.\n",
        "\n",
-        "# @markdown 3. For evaluation and deployment, **[click here](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_l4_gpus)** to check if your project already has the required 1 L4 GPU in the us-central1 region.  If yes, then run this notebook in the us-central1 region. If you need more L4 GPUs for your project, then you can follow [these instructions](https://cloud.google.com/docs/quotas/view-manage#viewing_your_quota_console) to request more. Alternatively, if you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
+        "# @markdown 3. For evaluation and deployment, **[click here](https://console.cloud.google.com/iam-admin/quotas?location=us-central1\u0026metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_l4_gpus)** to check if your project already has the required 1 L4 GPU in the us-central1 region.  If yes, then run this notebook in the us-central1 region. If you need more L4 GPUs for your project, then you can follow [these instructions](https://cloud.google.com/docs/quotas/view-manage#viewing_your_quota_console) to request more. Alternatively, if you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
        "\n",
-        "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
+        "# @markdown \u003e | 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",
@@ -188,8 +188,8 @@
        "REGION = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# Import the necessary packages.\n",
-        "! rm -rf vertex-ai-samples && git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
-        "! cd vertex-ai-samples && git reset --hard 0727e19520cf7957bceb701c248221bd3dbe4f1f\n",
+        "! rm -rf vertex-ai-samples \u0026\u0026 git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "! cd vertex-ai-samples \u0026\u0026 git reset --hard 0727e19520cf7957bceb701c248221bd3dbe4f1f\n",
        "\n",
        "import datetime\n",
        "import importlib\n",
@@ -269,13 +269,10 @@
      ]
    },
    {
-      "cell_type": "code",
-      "execution_count": null,
      "metadata": {
-        "cellView": "form",
-        "id": "36c21f10355f"
+        "id": "VUSi9jUcvdBC"
      },
-      "outputs": [],
+      "cell_type": "code",
      "source": [
        "# @title Access Llama 3.1 models\n",
        "\n",
@@ -288,7 +285,7 @@
        "base_model_id = \"meta-llama/Llama-3.1-8B-Instruct\"  # @param {type: \"string\"}\n",
        "pretrained_model_id = base_model_id\n",
        "\n",
-        "# @markdown Additionally, you must provide a Hugging Face User Access Token (read) to access the Llama 3.1 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 Additionally, you must provide a Hugging Face User Access Token (with read access) to access the Llama 3.1 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",
        "\n",
@@ -296,7 +293,9 @@
        "    assert (\n",
        "        HF_TOKEN\n",
        "    ), \"Provide a read access HF_TOKEN to load models from Hugging Face, or select a different model source. You can comment out this assert statement to skip this check.\""
-      ]
+      ],
+      "outputs": [],
+      "execution_count": null
    },
    {
      "cell_type": "markdown",
@@ -323,7 +322,7 @@
              "| coqa | 0.1158 | 0.0137 | 0.1872 | 0.0150 |\n"
            ],
            "text/plain": [
-              "<IPython.core.display.Markdown object>"
+              "\u003cIPython.core.display.Markdown object\u003e"
            ]
          },
          "metadata": {},
@@ -359,7 +358,7 @@
        "    lora_path: str = None,\n",
        "    max_num_seqs: int = 64,\n",
        "    eval_task: str = \"coqa\",\n",
-        ") -> str:\n",
+        ") -\u003e str:\n",
        "    \"\"\"Run lm-evaluation-harness to evaluate the model, and returns .\"\"\"\n",
        "\n",
        "    if accelerator_type == \"NVIDIA_L4\":\n",
@@ -607,7 +606,7 @@
              "version_minor": 0
            },
            "text/plain": [
-              "Downloading readme:   0%|          | 0.00/8.20k [00:00<?, ?B/s]"
+              "Downloading readme:   0%|          | 0.00/8.20k [00:00\u003c?, ?B/s]"
            ]
          },
          "metadata": {},
@@ -621,7 +620,7 @@
              "version_minor": 0
            },
            "text/plain": [
-              "Downloading data:   0%|          | 0.00/13.1M [00:00<?, ?B/s]"
+              "Downloading data:   0%|          | 0.00/13.1M [00:00\u003c?, ?B/s]"
            ]
          },
          "metadata": {},
@@ -635,7 +634,7 @@
              "version_minor": 0
            },
            "text/plain": [
-              "Generating train split:   0%|          | 0/15011 [00:00<?, ? examples/s]"
+              "Generating train split:   0%|          | 0/15011 [00:00\u003c?, ? examples/s]"
            ]
          },
          "metadata": {},
@@ -653,7 +652,7 @@
              "| When was Tomoaki Komorida born?                                                              | Komorida was born in Kumamoto Prefecture on July 10, 1981. After graduating from high school, he joined the J1 League club Avispa Fukuoka in 2000. Although he debuted as a midfielder in 2001, he did not play much and the club was relegated to the J2 League at the end of the 2001 season. In 2002, he moved to the J2 club Oita Trinita. He became a regular player as a defensive midfielder and the club won the championship in 2002 and was promoted in 2003. He played many matches until 2005. In September 2005, he moved to the J2 club Montedio Yamagata. In 2006, he moved to the J2 club Vissel Kobe. Although he became a regular player as a defensive midfielder, his gradually was played less during the summer. In 2007, he moved to the Japan Football League club Rosso Kumamoto (later Roasso Kumamoto) based in his local region. He played as a regular player and the club was promoted to J2 in 2008. Although he did not play as much, he still played in many matches. In 2010, he moved to Indonesia and joined Persela Lamongan. In July 2010, he returned to Japan and joined the J2 club Giravanz Kitakyushu. He played often as a defensive midfielder and center back until 2012 when he retired. | Tomoaki Komorida was born on July 10,1981.                                                                 | closed_qa      |"
            ],
            "text/plain": [
-              "<IPython.core.display.Markdown object>"
+              "\u003cIPython.core.display.Markdown object\u003e"
            ]
          },
          "metadata": {},
@@ -768,14 +767,14 @@
          "name": "stdout",
          "output_type": "stream",
          "text": [
-            "('<|begin_of_text|><|start_header_id|>system<|end_header_id|>\\n'\n",
+            "('\u003c|begin_of_text|\u003e\u003c|start_header_id|\u003esystem\u003c|end_header_id|\u003e\\n'\n",
            " '\\n'\n",
            " 'You are a helpful '\n",
-            " 'assistant.<|eot_id|><|start_header_id|>user<|end_header_id|>\\n'\n",
+            " 'assistant.\u003c|eot_id|\u003e\u003c|start_header_id|\u003euser\u003c|end_header_id|\u003e\\n'\n",
            " '\\n'\n",
-            " 'Hello, how are you?<|eot_id|><|start_header_id|>assistant<|end_header_id|>\\n'\n",
+            " 'Hello, how are you?\u003c|eot_id|\u003e\u003c|start_header_id|\u003eassistant\u003c|end_header_id|\u003e\\n'\n",
            " '\\n'\n",
-            " 'I am doing well, thank you.<|eot_id|>')\n"
+            " 'I am doing well, thank you.\u003c|eot_id|\u003e')\n"
          ]
        }
      ],
@@ -794,7 +793,7 @@
        "# @markdown to translate the Jinja template for you. For example, you can ask:\n",
        "\n",
        "# @markdown ````\n",
-        "# @markdown Translate the following Jinja template to Python, where bos_token is \"<|begin_of_text|>\":\n",
+        "# @markdown Translate the following Jinja template to Python, where bos_token is \"\u003c|begin_of_text|\u003e\":\n",
        "# @markdown ```\n",
        "# @markdown {{- bos_token }}\n",
        "# @markdown {#- This block extracts the system message, so we can slot it into the right place. #}\n",
@@ -805,14 +804,14 @@
        "# @markdown     {%- set system_message = \"\" %}\n",
        "# @markdown {%- endif %}\n",
        "# @markdown {#- System message #}\n",
-        "# @markdown {{- \"<|start_header_id|>system<|end_header_id|>\\n\\n\" }}\n",
+        "# @markdown {{- \"\u003c|start_header_id|\u003esystem\u003c|end_header_id|\u003e\\n\\n\" }}\n",
        "# @markdown {{- system_message }}\n",
-        "# @markdown {{- \"<|eot_id|>\" }}\n",
+        "# @markdown {{- \"\u003c|eot_id|\u003e\" }}\n",
        "# @markdown {%- for message in messages %}\n",
-        "# @markdown     {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\\n\\n'+ message['content'] | trim + '<|eot_id|>' }}\n",
+        "# @markdown     {{- '\u003c|start_header_id|\u003e' + message['role'] + '\u003c|end_header_id|\u003e\\n\\n'+ message['content'] | trim + '\u003c|eot_id|\u003e' }}\n",
        "# @markdown {%- endfor %}\n",
        "# @markdown {%- if add_generation_prompt %}\n",
-        "# @markdown     {{- '<|start_header_id|>assistant<|end_header_id|>\\n\\n' }}\n",
+        "# @markdown     {{- '\u003c|start_header_id|\u003eassistant\u003c|end_header_id|\u003e\\n\\n' }}\n",
        "# @markdown {%- endif %}\n",
        "# @markdown ```\n",
        "# @markdown ````\n",
@@ -821,7 +820,7 @@
        "# @markdown `bos_token` is specified in the [tokenizer_config.json](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct/blob/main/tokenizer_config.json#L2052).\n",
        "# @markdown ```\n",
        "# @markdown def render_template(messages, add_generation_prompt=False):\n",
-        "# @markdown     bos_token = \"<|begin_of_text|>\"\n",
+        "# @markdown     bos_token = \"\u003c|begin_of_text|\u003e\"\n",
        "# @markdown     output = bos_token\n",
        "# @markdown\n",
        "# @markdown     system_message = \"\"\n",
@@ -829,22 +828,22 @@
        "# @markdown         system_message = messages[0]['content'].strip()\n",
        "# @markdown         messages = messages[1:]\n",
        "# @markdown\n",
-        "# @markdown     output += \"<|start_header_id|>system<|end_header_id|>\\n\\n\"\n",
+        "# @markdown     output += \"\u003c|start_header_id|\u003esystem\u003c|end_header_id|\u003e\\n\\n\"\n",
        "# @markdown     output += system_message\n",
-        "# @markdown     output += \"<|eot_id|>\"\n",
+        "# @markdown     output += \"\u003c|eot_id|\u003e\"\n",
        "# @markdown\n",
        "# @markdown     for message in messages:\n",
-        "# @markdown         output += f\"<|start_header_id|>{message['role']}<|end_header_id|>\\n\\n{message['content'].strip()}<|eot_id|>\"\n",
+        "# @markdown         output += f\"\u003c|start_header_id|\u003e{message['role']}\u003c|end_header_id|\u003e\\n\\n{message['content'].strip()}\u003c|eot_id|\u003e\"\n",
        "# @markdown\n",
        "# @markdown     if add_generation_prompt:\n",
-        "# @markdown         output += \"<|start_header_id|>assistant<|end_header_id|>\\n\\n\"\n",
+        "# @markdown         output += \"\u003c|start_header_id|\u003eassistant\u003c|end_header_id|\u003e\\n\\n\"\n",
        "# @markdown\n",
        "# @markdown     return output\n",
        "# @markdown ```\n",
        "\n",
        "\n",
        "def render_template(messages, add_generation_prompt=False):\n",
-        "    bos_token = \"<|begin_of_text|>\"\n",
+        "    bos_token = \"\u003c|begin_of_text|\u003e\"\n",
        "    output = bos_token\n",
        "\n",
        "    system_message = \"\"\n",
@@ -852,26 +851,26 @@
        "        system_message = messages[0][\"content\"].strip()\n",
        "        messages = messages[1:]\n",
        "\n",
-        "    output += \"<|start_header_id|>system<|end_header_id|>\\n\\n\"\n",
+        "    output += \"\u003c|start_header_id|\u003esystem\u003c|end_header_id|\u003e\\n\\n\"\n",
        "    output += system_message\n",
-        "    output += \"<|eot_id|>\"\n",
+        "    output += \"\u003c|eot_id|\u003e\"\n",
        "\n",
        "    for message in messages:\n",
        "        output += (\n",
-        "            \"<|start_header_id|>\"\n",
+        "            \"\u003c|start_header_id|\u003e\"\n",
        "            + message[\"role\"]\n",
-        "            + \"<|end_header_id|>\\n\\n\"\n",
+        "            + \"\u003c|end_header_id|\u003e\\n\\n\"\n",
        "            + message[\"content\"].strip()\n",
-        "            + \"<|eot_id|>\"\n",
+        "            + \"\u003c|eot_id|\u003e\"\n",
        "        )\n",
        "\n",
        "    if add_generation_prompt:\n",
-        "        output += \"<|start_header_id|>assistant<|end_header_id|>\\n\\n\"\n",
+        "        output += \"\u003c|start_header_id|\u003eassistant\u003c|end_header_id|\u003e\\n\\n\"\n",
        "\n",
        "    return output\n",
        "\n",
        "\n",
-        "# @markdown The `<|begin_of_text|>`, `<|start_header_id|>`, `<|end_header_id|>`, and `<|eot_id|>` tokens serve specific purposes in the context of text generation models.\n",
+        "# @markdown The `\u003c|begin_of_text|\u003e`, `\u003c|start_header_id|\u003e`, `\u003c|end_header_id|\u003e`, and `\u003c|eot_id|\u003e` tokens serve specific purposes in the context of text generation models.\n",
        "# @markdown These tokens help the model delineate the boundaries of a text generation task. They provide clear markers for the start and end points, enabling the model to function effectively and produce coherent text.\n",
        "\n",
        "# @markdown Run this cell to show an example output of the template given the\n",
@@ -1030,15 +1029,15 @@
          "name": "stdout",
          "output_type": "stream",
          "text": [
-            "('<|begin_of_text|><|start_header_id|>user<|end_header_id|>\\n'\n",
+            "('\u003c|begin_of_text|\u003e\u003c|start_header_id|\u003euser\u003c|end_header_id|\u003e\\n'\n",
            " '\\n'\n",
            " 'Hello, how are you? Context: This is a test '\n",
-            " 'context.<|eot_id|><|start_header_id|>assistant<|end_header_id|>\\n'\n",
+            " 'context.\u003c|eot_id|\u003e\u003c|start_header_id|\u003eassistant\u003c|end_header_id|\u003e\\n'\n",
            " '\\n'\n",
            " \"I'm doing well, thank \"\n",
-            " 'you!<|eot_id|><|start_header_id|>user<|end_header_id|>\\n'\n",
+            " 'you!\u003c|eot_id|\u003e\u003c|start_header_id|\u003euser\u003c|end_header_id|\u003e\\n'\n",
            " '\\n'\n",
-            " 'Another question without context.<|eot_id|>')\n"
+            " 'Another question without context.\u003c|eot_id|\u003e')\n"
          ]
        }
      ],
@@ -1052,15 +1051,15 @@
        "# @markdown chat_template_string = r\"\"\"{{- bos_token }}\n",
        "# @markdown\n",
        "# @markdown {% for message in messages %}\n",
-        "# @markdown     {{- '<|start_header_id|>' + message.role + '<|end_header_id|>\\n\\n' + message.content | trim }}\n",
-        "# @markdown     {% if message.context and message.context | length > 0 %}\n",
+        "# @markdown     {{- '\u003c|start_header_id|\u003e' + message.role + '\u003c|end_header_id|\u003e\\n\\n' + message.content | trim }}\n",
+        "# @markdown     {% if message.context and message.context | length \u003e 0 %}\n",
        "# @markdown         {{- ' Context: ' + message.context }}\n",
        "# @markdown     {% endif %}\n",
-        "# @markdown     {{- '<|eot_id|>' }}\n",
+        "# @markdown     {{- '\u003c|eot_id|\u003e' }}\n",
        "# @markdown {% endfor %}\n",
        "# @markdown\n",
        "# @markdown {% if add_generation_prompt %}\n",
-        "# @markdown     {{- '<|start_header_id|>assistant<|end_header_id|>\\n\\n' }}\n",
+        "# @markdown     {{- '\u003c|start_header_id|\u003eassistant\u003c|end_header_id|\u003e\\n\\n' }}\n",
        "# @markdown {% endif %}\n",
        "# @markdown \"\"\"\n",
        "# @markdown ```\n",
@@ -1071,15 +1070,15 @@
        "chat_template_string = r\"\"\"{{- bos_token }}\n",
        "\n",
        "{% for message in messages %}\n",
-        "    {{- '<|start_header_id|>' + message.role + '<|end_header_id|>\\n\\n' + message.content | trim }}\n",
-        "    {% if message.context and message.context | length > 0 %}\n",
+        "    {{- '\u003c|start_header_id|\u003e' + message.role + '\u003c|end_header_id|\u003e\\n\\n' + message.content | trim }}\n",
+        "    {% if message.context and message.context | length \u003e 0 %}\n",
        "        {{- ' Context: ' + message.context }}\n",
        "    {% endif %}\n",
-        "    {{- '<|eot_id|>' }}\n",
+        "    {{- '\u003c|eot_id|\u003e' }}\n",
        "{% endfor %}\n",
        "\n",
        "{% if add_generation_prompt %}\n",
-        "    {{- '<|start_header_id|>assistant<|end_header_id|>\\n\\n' }}\n",
+        "    {{- '\u003c|start_header_id|\u003eassistant\u003c|end_header_id|\u003e\\n\\n' }}\n",
        "{% endif %}\n",
        "\"\"\"\n",
        "\n",
@@ -1089,19 +1088,19 @@
        "# @markdown     output = bos_token\n",
        "# @markdown\n",
        "# @markdown     for message in messages:\n",
-        "# @markdown         output += f\"<|start_header_id|>{message['role']}<|end_header_id|>\\n\\n{message['content'].strip()}\"\n",
+        "# @markdown         output += f\"\u003c|start_header_id|\u003e{message['role']}\u003c|end_header_id|\u003e\\n\\n{message['content'].strip()}\"\n",
        "# @markdown         if 'context' in message and message['context']:\n",
        "# @markdown             output += f\" Context: {message['context']}\"\n",
-        "# @markdown         output += \"<|eot_id|>\"\n",
+        "# @markdown         output += \"\u003c|eot_id|\u003e\"\n",
        "# @markdown\n",
        "# @markdown     if add_generation_prompt:\n",
-        "# @markdown         output += \"<|start_header_id|>assistant<|end_header_id|>\\n\\n\"\n",
+        "# @markdown         output += \"\u003c|start_header_id|\u003eassistant\u003c|end_header_id|\u003e\\n\\n\"\n",
        "# @markdown\n",
        "# @markdown     return output\n",
        "# @markdown ```\n",
        "\n",
        "# @markdown Run this cell to show an example output of the template given the\n",
-        "# @markdown below `messages` and `bos_token=\"<|begin_of_text|>\"`.\n",
+        "# @markdown below `messages` and `bos_token=\"\u003c|begin_of_text|\u003e\"`.\n",
        "# @markdown ```\n",
        "# @markdown messages = [\n",
        "# @markdown     {\"role\": \"user\", \"content\": \"Hello, how are you?\", \"context\": \"This is a test context.\"},\n",
@@ -1117,13 +1116,13 @@
        "    output = bos_token\n",
        "\n",
        "    for message in messages:\n",
-        "        output += f\"<|start_header_id|>{message['role']}<|end_header_id|>\\n\\n{message['content'].strip()}\"\n",
+        "        output += f\"\u003c|start_header_id|\u003e{message['role']}\u003c|end_header_id|\u003e\\n\\n{message['content'].strip()}\"\n",
        "        if \"context\" in message and message[\"context\"]:\n",
        "            output += f\" Context: {message['context']}\"\n",
-        "        output += \"<|eot_id|>\"\n",
+        "        output += \"\u003c|eot_id|\u003e\"\n",
        "\n",
        "    if add_generation_prompt:\n",
-        "        output += \"<|start_header_id|>assistant<|end_header_id|>\\n\\n\"\n",
+        "        output += \"\u003c|start_header_id|\u003eassistant\u003c|end_header_id|\u003e\\n\\n\"\n",
        "\n",
        "    return output\n",
        "\n",
@@ -1138,7 +1137,7 @@
        "    {\"role\": \"user\", \"content\": \"Another question without context.\", \"context\": \"\"},\n",
        "]\n",
        "\n",
-        "rendered_text = render_template(messages, bos_token=\"<|begin_of_text|>\")\n",
+        "rendered_text = render_template(messages, bos_token=\"\u003c|begin_of_text|\u003e\")\n",
        "pprint.pprint(rendered_text, width=80)"
      ]
    },
@@ -1173,16 +1172,16 @@
        "# @markdown template = {\n",
        "# @markdown   \"description\": \"Template used by Llama 3.1, accepting databricks dolly dataset.\",\n",
        "# @markdown   \"chat_template\": chat_template_string,\n",
-        "# @markdown   \"instruction_separator\": \"<|start_header_id|>user<|end_header_id|>\\n\\n\",\n",
-        "# @markdown   \"response_separator\": \"<|start_header_id|>assistant<|end_header_id|>\\n\\n\"\n",
+        "# @markdown   \"instruction_separator\": \"\u003c|start_header_id|\u003euser\u003c|end_header_id|\u003e\\n\\n\",\n",
+        "# @markdown   \"response_separator\": \"\u003c|start_header_id|\u003eassistant\u003c|end_header_id|\u003e\\n\\n\"\n",
        "# @markdown }\n",
        "# @markdown ```\n",
        "\n",
        "template_data = {\n",
        "    \"description\": \"Template used by Llama 3.1, accepting databricks dolly dataset.\",\n",
        "    \"chat_template\": chat_template_string,\n",
-        "    \"instruction_separator\": \"<|start_header_id|>user<|end_header_id|>\\n\\n\",\n",
-        "    \"response_separator\": \"<|start_header_id|>assistant<|end_header_id|>\\n\\n\",\n",
+        "    \"instruction_separator\": \"\u003c|start_header_id|\u003euser\u003c|end_header_id|\u003e\\n\\n\",\n",
+        "    \"response_separator\": \"\u003c|start_header_id|\u003eassistant\u003c|end_header_id|\u003e\\n\\n\",\n",
        "}\n",
        "\n",
        "template_filename = \"template.json\"\n",
@@ -1346,7 +1345,7 @@
        "\n",
        "# @markdown **Note**:\n",
        "# @markdown 1. We recommend setting `finetuning_precision_mode` to `float16`.\n",
-        "# @markdown 1. If `max_steps>0`, it takes precedence over `epochs`. One can set a small `max_steps`\n",
+        "# @markdown 1. If `max_steps\u003e0`, it takes precedence over `epochs`. One can set a small `max_steps`\n",
        "# @markdown value to quickly check the pipeline.\n",
        "\n",
        "# @markdown Acceletor type to use for training.\n",
@@ -1393,7 +1392,7 @@
        "# Set config file.\n",
        "if replica_count == 1:\n",
        "    config_file = \"vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml\"\n",
-        "elif replica_count <= 4:\n",
+        "elif replica_count \u003c= 4:\n",
        "    config_file = (\n",
        "        \"vertex_vision_model_garden_peft/\"\n",
        "        f\"llama_hsdp_{replica_count * per_node_accelerator_count}gpu.yaml\"\n",
@@ -1654,7 +1653,7 @@
        "\n",
        "\n",
        "# @markdown Expected evaluation results:\n",
-        "# @markdown  > | alias | exact_match | exact_match_stderr | f1 | f1_stderr |\n",
+        "# @markdown  \u003e | alias | exact_match | exact_match_stderr | f1 | f1_stderr |\n",
        "# @markdown | --- | --- | --- | --- | --- |\n",
        "# @markdown | coqa | 0.3213 | 0.0197 | 0.4660 | 0.0187 |\n",
        "\n",
@@ -1710,14 +1709,14 @@
        "    is_for_training=False,\n",
        ")\n",
        "\n",
-        "# Dedicated endpoint not supported yet.\n",
-        "use_dedicated_endpoint = False\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",
        "gpu_memory_utilization = 0.85\n",
        "max_model_len = 8192  # Maximum context length.\n",
        "\n",
        "# Ensure max_model_len does not exceed the limit.\n",
-        "if max_model_len > 8192:\n",
+        "if max_model_len \u003e 8192:\n",
        "    raise ValueError(\"max_model_len cannot exceed 8192\")\n",
        "\n",
        "\n",
@@ -1743,7 +1742,7 @@
        "    use_dedicated_endpoint: bool = False,\n",
        "    max_num_seqs: int = 256,\n",
        "    model_type: str = None,\n",
-        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        ") -\u003e 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",
@@ -1787,7 +1786,7 @@
        "    if enable_prefix_cache:\n",
        "        vllm_args.append(\"--enable-prefix-caching\")\n",
        "\n",
-        "    if 0 < host_prefix_kv_cache_utilization_target < 1:\n",
+        "    if 0 \u003c host_prefix_kv_cache_utilization_target \u003c 1:\n",
        "        vllm_args.append(\n",
        "            f\"--host-prefix-kv-cache-utilization-target={host_prefix_kv_cache_utilization_target}\"\n",
        "        )\n",
@@ -1846,6 +1845,7 @@
        "    top_k: int,\n",
        "    raw_response: bool,\n",
        "    lora_weight: str = \"\",\n",
+        "    use_dedicated_endpoint: bool = False,\n",
        "):\n",
        "    # Parameters for inference.\n",
        "    instance = {\n",
@@ -1901,7 +1901,7 @@
          "output_type": "stream",
          "text": [
            "Prompt:\n",
-            "<|start_header_id|>user<|end_header_id|>\\n\\nWhat was Anya looking for? Context: Anya clutched the worn teddy bear, its button eye dangling precariously. She'd lost it in the park yesterday, and the thought of never seeing Mr. Snuggles again made her tummy ache. She retraced her steps, her eyes scanning the colorful playground equipment and the sprawling green lawn.<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n",
+            "\u003c|start_header_id|\u003euser\u003c|end_header_id|\u003e\\n\\nWhat was Anya looking for? Context: Anya clutched the worn teddy bear, its button eye dangling precariously. She'd lost it in the park yesterday, and the thought of never seeing Mr. Snuggles again made her tummy ache. She retraced her steps, her eyes scanning the colorful playground equipment and the sprawling green lawn.\u003c|eot_id|\u003e\u003c|start_header_id|\u003eassistant\u003c|end_header_id|\u003e\n",
            "Output:\n",
            "Anya was looking for her teddy bear, Mr. Snuggles.\n"
          ]
@@ -1925,8 +1925,8 @@
        "\n",
        "prompt = \"What was Anya looking for? Context: Anya clutched the worn teddy bear, its button eye dangling precariously. She'd lost it in the park yesterday, and the thought of never seeing Mr. Snuggles again made her tummy ache. She retraced her steps, her eyes scanning the colorful playground equipment and the sprawling green lawn.\"  # @param {type: \"string\"}\n",
        "prompt_with_headers = (\n",
-        "    f\"<|start_header_id|>user<|end_header_id|>\\\\n\\\\n{prompt}<|eot_id|>\"\n",
-        "    \"<|start_header_id|>assistant<|end_header_id|>\\n\\n\"\n",
+        "    f\"\u003c|start_header_id|\u003euser\u003c|end_header_id|\u003e\\\\n\\\\n{prompt}\u003c|eot_id|\u003e\"\n",
+        "    \"\u003c|start_header_id|\u003eassistant\u003c|end_header_id|\u003e\\n\\n\"\n",
        ")\n",
        "\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",
@@ -1945,6 +1945,7 @@
        "    top_k=top_k,\n",
        "    raw_response=raw_response,\n",
        "    lora_weight=lora_output_dir,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
        ")\n",
        "\n",
        "# @markdown Click \"Show Code\" to see more details."
@@ -3,7 +3,6 @@
    {
      "cell_type": "code",
      "execution_count": null,
-      "language": "python",
      "metadata": {
        "cellView": "form",
        "id": "7d9bbf86da5e"
@@ -35,6 +34,11 @@
        "\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",
@@ -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.84.0'\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",
@@ -144,64 +158,26 @@
        "\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",
-        "# 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",
@@ -229,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",
@@ -271,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",
@@ -291,7 +281,6 @@
        "    model_id: str,\n",
        "    publisher: str,\n",
        "    publisher_model_id: 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",
@@ -380,7 +369,6 @@
        "        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",
@@ -396,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",
@@ -489,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",
@@ -547,6 +539,44 @@
        "    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.preview 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",
+        ")"
+      ]
+    },
+    {
+      "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",
@@ -560,7 +590,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",
@@ -591,6 +621,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",
@@ -617,12 +650,11 @@
        "    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",
@@ -630,7 +662,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."
      ]
    },
    {
@@ -713,6 +747,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",
@@ -722,11 +758,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()"
      ]
    }
  ],
@@ -36,6 +36,11 @@
        "\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",
@@ -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",
@@ -143,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 80320a9a1b818534ca785444e704f6953f2a9dd9\n",
+        "! cd vertex-ai-samples && git reset --hard c45f6a4f4d32e31a050f0e4ba52824b0caf4eda3\n",
        "\n",
        "import datetime\n",
        "import importlib\n",
@@ -155,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",
@@ -223,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",
@@ -407,9 +416,7 @@
        "# @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",
-        "# fmt: off\n",
        "training_accelerator_type = \"NVIDIA_A100_80GB\"  # @param [\"NVIDIA_A100_80GB\", \"NVIDIA_H100_80GB\"]\n",
-        "# fmt: on\n",
        "\n",
        "# The pre-built training docker image.\n",
        "if training_accelerator_type == \"NVIDIA_A100_80GB\":\n",
@@ -427,7 +434,7 @@
        "    }\n",
        "\n",
        "TRAIN_DOCKER_URI = (\n",
-        "    f\"{repo}/vertex-vision-model-garden-dockers/pytorch-peft-train:stable_20250213\"\n",
+        "    f\"{repo}/vertex-vision-model-garden-dockers/pytorch-peft-train:stable_20250409\"\n",
        ")\n",
        "\n",
        "# Worker pool spec.\n",
@@ -503,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",
@@ -582,8 +590,6 @@
        "# Wait until resource has been created.\n",
        "train_job.wait_for_resource_creation()\n",
        "\n",
-        "merged_model_output_dir = os.path.join(merged_model_output_dir, \"node-0\")\n",
-        "\n",
        "print(\"LoRA adapter will be saved in:\", lora_output_dir)\n",
        "print(\"Trained and merged models will be saved in:\", merged_model_output_dir)\n",
        "\n",
@@ -610,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,
@@ -622,18 +748,12 @@
      "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: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",
+        "# @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",
@@ -66,6 +66,10 @@
        "\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",
@@ -467,9 +471,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",
@@ -779,9 +781,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",
@@ -856,9 +856,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",
@@ -65,6 +65,7 @@
        "### Objective\n",
        "\n",
        "- Finetune and deploy Gemma 3 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 3 model.\n",
        "\n",
        "### File a bug\n",
@@ -143,7 +144,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 80320a9a1b818534ca785444e704f6953f2a9dd9\n",
+        "! cd vertex-ai-samples && git reset --hard c45f6a4f4d32e31a050f0e4ba52824b0caf4eda3\n",
        "\n",
        "import datetime\n",
        "import importlib\n",
@@ -223,7 +224,7 @@
        "\n",
        "# @markdown ## Access Gemma 3 Models\n",
        "\n",
-        "# @markdown You must provide a Hugging Face User Access Token (read) to access the Gemma 3 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 3 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",
@@ -407,9 +408,7 @@
        "# @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",
-        "# fmt: off\n",
        "training_accelerator_type = \"NVIDIA_A100_80GB\"  # @param [\"NVIDIA_A100_80GB\", \"NVIDIA_H100_80GB\"]\n",
-        "# fmt: on\n",
        "\n",
        "# The pre-built training docker image.\n",
        "if training_accelerator_type == \"NVIDIA_A100_80GB\":\n",
@@ -543,6 +542,7 @@
        "    f\"--lr_scheduler_type={lr_scheduler_type}\",\n",
        "    f\"--precision_mode={finetuning_precision_mode}\",\n",
        "    f\"--train_precision={train_precision}\",\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",
@@ -607,6 +607,113 @@
        "print(f\"Command to copy: tensorboard --logdir {base_output_dir}/logs\")"
      ]
    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "4g0woSqhvF9O"
+      },
+      "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": "635Bdo0Pt6iq"
+      },
+      "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",
+        "eval_machine_type = \"g2-standard-12\"\n",
+        "eval_accelerator_count = 1\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",
+        "gpu_memory_utilization = 0.85\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(\"gemma3-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,
@@ -619,18 +726,12 @@
      "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: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).\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",
@@ -199,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",
@@ -303,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",
@@ -325,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",
@@ -365,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",
@@ -513,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",
@@ -675,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",
@@ -909,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",
@@ -203,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",
@@ -113,10 +113,9 @@
    {
      "cell_type": "code",
      "execution_count": null,
-      "language": "python",
      "metadata": {
        "cellView": "form",
-        "id": "855d6b96f291"
+        "id": "8CQcnBfWvc-f"
      },
      "outputs": [],
      "source": [
@@ -238,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",
@@ -682,7 +681,7 @@
        "# 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",
@@ -1240,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",
@@ -316,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",
@@ -347,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",
@@ -34,6 +34,11 @@
        "\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",
@@ -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,16 +102,6 @@
      "source": [
        "# @title Setup Google Cloud project\n",
        "\n",
-        "# Used for common utilities.\n",
-        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
-        "\n",
-        "import importlib\n",
-        "# Import the necessary packages\n",
-        "import os\n",
-        "from typing import Any, Dict, Tuple\n",
-        "\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",
        "\n",
        "# @markdown 2. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
@@ -118,6 +117,27 @@
        "# @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.84.0'\n",
+        "\n",
+        "import importlib\n",
+        "# Import the necessary packages\n",
+        "import os\n",
+        "from typing import Any, Dict, 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",
+        "models, endpoints = {}, {}\n",
+        "LABEL = \"paligemma2\"\n",
+        "\n",
        "\n",
        "# Get the default cloud project id.\n",
        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
@@ -126,27 +146,116 @@
        "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",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kyMJXkfviWgl"
+      },
+      "source": [
+        "## Deploy Model to a Vertex AI Endpoint"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "toY-WPKDFesF"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "\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",
+        "MODEL_ID = os.path.join(GCS_PREFIX, MODEL_NAME)\n",
+        "\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",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=False,\n",
        ")\n",
        "\n",
-        "models, endpoints = {}, {}\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.preview 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",
+        ")"
+      ]
+    },
+    {
+      "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",
@@ -159,7 +268,6 @@
        "    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",
@@ -173,7 +281,6 @@
        "        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",
@@ -204,111 +311,15 @@
        "        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={\"NOTEBOOK_NAME\": \"model_garden_hf_paligemma2_deployment.ipynb\"},\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",
-        "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"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "kyMJXkfviWgl"
-      },
-      "source": [
-        "## Deploy Model to a Vertex AI Endpoint"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "cellView": "form",
-        "id": "toY-WPKDFesF"
-      },
-      "outputs": [],
-      "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",
-        "\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",
-        "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",
-        "MODEL_ID = os.path.join(GCS_PREFIX, 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-8\"\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",
-        "    region=REGION,\n",
-        "    accelerator_type=accelerator_type,\n",
-        "    accelerator_count=accelerator_count,\n",
-        "    is_for_training=False,\n",
-        ")\n",
-        "\n",
-        "TASK = \"paligemma_VQA\"\n",
-        "\n",
        "endpoints[\"paligemma2\"], models[\"paligemma2\"] = deploy_model(\n",
        "    model_name=MODEL_NAME,\n",
        "    model_id=MODEL_ID,\n",
@@ -316,7 +327,6 @@
        "    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",
@@ -379,6 +389,48 @@
        "\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",
@@ -138,8 +138,10 @@
        "! gcloud config set project $PROJECT_ID\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 = \"\""
      ]
    },
@@ -34,6 +34,11 @@
        "\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",
@@ -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.84.0'\n",
        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
        "\n",
        "import importlib\n",
@@ -135,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",
@@ -173,6 +173,53 @@
        "        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.preview 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",
+        ")"
+      ]
+    },
+    {
+      "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",
@@ -229,13 +276,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",
@@ -186,7 +186,7 @@
        "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",
+        "# @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\"}"
      ]
    },
@@ -236,7 +236,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",
@@ -267,6 +267,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",
@@ -30,7 +30,7 @@
        "id": "08f4AuF5eXzO"
      },
      "source": [
-        "# Vertex AI Model Garden - Hugging Face Text Generation with vLLM Container Deployment\n",
+        "# Vertex AI Model Garden - Hugging Face Deployment with vLLM Container\n",
        "\n",
        "<table><tbody><tr>\n",
        "  <td style=\"text-align: center\">\n",
@@ -39,12 +39,12 @@
        "    </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_vllm_deployment.ipynb\">\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_tgi_vllm_deployment.ipynb\">\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://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
@@ -59,13 +59,13 @@
      "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 `text-generation` model with appropriate machine specs. **Note that some models might fail to deploy, even if they have `text-generation` tags on the Hugging Face model card page.**\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 TGI\n",
-        "- Send prediction request to the deployed endpoint\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",
@@ -163,8 +163,8 @@
        "\n",
        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate: true}\n",
        "\n",
-        "# The pre-built serving docker image for TGI with vLLM.\n",
-        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/vllm-inference.cu121.0-6.ubuntu2204.py310\"\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",
@@ -384,7 +384,7 @@
        "        deploy_request_timeout=1800,\n",
        "        service_account=service_account,\n",
        "        system_labels={\n",
-        "            \"NOTEBOOK_NAME\": \"model_garden_huggingface_tgi_vllm_deployment.ipynb\",\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_huggingface_vllm_deployment.ipynb\",\n",
        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
        "        },\n",
        "    )\n",
@@ -403,6 +403,8 @@
        "    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."
@@ -591,7 +593,7 @@
  ],
  "metadata": {
    "colab": {
-      "name": "model_garden_huggingface_tgi_vllm_deployment.ipynb",
+      "name": "model_garden_huggingface_vllm_deployment.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
@@ -0,0 +1,588 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "20qcPG1PmFUM"
+      },
+      "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": "QXYOa1odnikj"
+      },
+      "source": [
+        "# Vertex AI Model Garden Integration With ADK\n",
+        "\n",
+        "\u003ctable\u003e\u003ctbody\u003e\u003ctr\u003e\n",
+        "  \u003ctd style=\"text-align: center\"\u003e\n",
+        "    \u003ca 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_integration_with_adk.ipynb\"\u003e\n",
+        "      \u003cimg alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"\u003e\u003cbr\u003e Run in Colab Enterprise\n",
+        "    \u003c/a\u003e\n",
+        "  \u003c/td\u003e\n",
+        "  \u003ctd style=\"text-align: center\"\u003e\n",
+        "    \u003ca href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_integration_with_adk.ipynb\"\u003e\n",
+        "      \u003cimg alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"\u003e\u003cbr\u003e View on GitHub\n",
+        "    \u003c/a\u003e\n",
+        "  \u003c/td\u003e\n",
+        "\u003c/tr\u003e\u003c/tbody\u003e\u003c/table\u003e"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cbDI9ag4oR4C"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates how to build an agent with a Model Garden open model deployed through Vertex endpoint and [Google Agent Development Kit](https://google.github.io/adk-docs/) (ADK). The agent can automatically call function tools like `get_weather` and `get_current_time`.\n",
+        "\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Deploy Vertex AI Model Garden OSS LLMs properly for ADK integration\n",
+        "- Test deployed endpoints\n",
+        "- Build agent web apps with deployed endpoints and ADK\n",
+        "- Deploy agent web apps to Cloud Run\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": "hQJWRopioSKT"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "J_jmxcIZoSxU",
+        "cellView": "form"
+      },
+      "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 \u003e | 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\u003e=1.84.0'\n",
+        "! pip install -qU openai google-auth requests\n",
+        "\n",
+        "# Import the necessary packages\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",
+        "if not os.path.exists(\"./vertex-ai-samples\"):\n",
+        "  ! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
+        "models, endpoints = {}, {}\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",
+        "# 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",
+        "\n",
+        "# Gets the default SERVICE_ACCOUNT.\n",
+        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
+        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
+        "\n",
+        "\n",
+        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/storage.admin\"\n",
+        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/aiplatform.user\"\n",
+        "\n",
+        "from etils import epath\n",
+        "TUTORIAL_DIR = epath.Path(\"vmg_adk_agent_tutorial\")\n",
+        "BUILD_DIR = TUTORIAL_DIR / \"build\"\n",
+        "BUILD_DIR.mkdir(exist_ok=True, parents=True)\n",
+        "\n",
+        "REPOSITORY_NAME = \"vertex-vision-model-garden-dockers\"\n",
+        "\n",
+        "!gcloud artifacts repositories create $REPOSITORY_NAME \\\n",
+        "      --repository-format=docker \\\n",
+        "      --location=$REGION \\\n",
+        "      --project=$PROJECT_ID\n",
+        "\n",
+        "! gcloud auth configure-docker $REGION-docker.pkg.dev --quiet"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## Integrate OSS LLMs With ADK"
+      ],
+      "metadata": {
+        "id": "T9UiFwbiowOV"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# @title Deploy OSS LLMs For Agents\n",
+        "# @markdown In order to use OSS LLM endpoints smoothly with ADK,\n",
+        "# @markdown these models should be deployed with:\n",
+        "\n",
+        "# @markdown - **enabling tool calls**. e.g.:\n",
+        "# @markdown If the models are deployed with *vllm*,\n",
+        "# @markdown the deployment should specify settings like `--enable-auto-tool-choice`\n",
+        "# @markdown and `--tool-call-parser=hermes`.\n",
+        "# @markdown If the models are deployed with *sglang*, the deployment should specify\n",
+        "# @markdown setting like `--tool-call-parser=qwen25`.\n",
+        "# @markdown Refer to tool calls in\n",
+        "# @markdown [vllm](https://docs.vllm.ai/en/stable/features/tool_calling.html)\n",
+        "# @markdown and [sglang](https://docs.sglang.ai/backend/function_calling.html) for more details.\n",
+        "\n",
+        "# @markdown - **disable dedicated endpoints**. The dedicated endpoints are not\n",
+        "# @markdown supported in ADK yet.\n",
+        "\n",
+        "# @markdown You can deploy models below with proper deployment settings for ADK integration.\n",
+        "\n",
+        "MODEL_ID = \"Qwen3-32B\"  # @param [\"Qwen3-32B\"] {isTemplate: true}\n",
+        "accelerator_type = \"NVIDIA_H100_80GB\"  # @param [\"NVIDIA_L4\", \"NVIDIA_A100_80GB\", \"NVIDIA_H100_80GB\"] {isTemplate: true}\n",
+        "\n",
+        "if accelerator_type == \"NVIDIA_L4\":\n",
+        "    accelerator_count = 4\n",
+        "    # Sets machine type to g2-standard-48 for 4 L4's\n",
+        "    machine_type = \"g2-standard-48\"\n",
+        "elif accelerator_type == \"NVIDIA_A100_80GB\":\n",
+        "    accelerator_count = 1\n",
+        "    # Sets machine type to a2-ultragpu-1g for 1 Nvidia A100 80 GB.\n",
+        "    machine_type = \"a2-ultragpu-1g\"\n",
+        "elif accelerator_type == \"NVIDIA_H100_80GB\":\n",
+        "    accelerator_count = 2\n",
+        "    machine_type = \"a3-highgpu-2g\"\n",
+        "\n",
+        "else:\n",
+        "    raise ValueError(\n",
+        "        \"Recommended machine settings not found for accelerator type: %s\"\n",
+        "        % accelerator_type\n",
+        "    )\n",
+        "\n",
+        "deploy_request_timeout = 1800  # 30 minutes\n",
+        "\n",
+        "from vertexai.preview import model_garden\n",
+        "publisher_model_name = f\"publishers/qwen/models/qwen3@{MODEL_ID.lower()}\"\n",
+        "model = model_garden.OpenModel(publisher_model_name)\n",
+        "\n",
+        "container_spec = model.list_deploy_options()[0].container_spec\n",
+        "updated_args = container_spec.args[:-2] + [f\"--tp={accelerator_count}\", \"--tool-call-parser=qwen25\"]\n",
+        "container_spec.args = updated_args\n",
+        "\n",
+        "print(\"The container spec are:\")\n",
+        "print(container_spec)\n",
+        "\n",
+        "print(\"Start to check quota for the deployment.\")\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",
+        "print(\"Finished to check quota for the deployment.\")\n",
+        "\n",
+        "print(\"Start to deploy models to endpoints.\")\n",
+        "endpoint = model.deploy(\n",
+        "    serving_container_spec=container_spec,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=False,\n",
+        "    spot=False,\n",
+        "    deploy_request_timeout=deploy_request_timeout,\n",
+        "    accept_eula=False,\n",
+        ")\n",
+        "print(\"Finished to deploy models to endpoints.\")\n",
+        "# @markdown After endpoints are deployed successfully, you get the endpoint\n",
+        "# @markdown resource name with the format as\n",
+        "# @markdown `projects/{PROJECT_ID}/locations/{REGION}/endpoints/{ENDPOINT_ID}`.\n",
+        "# @markdown The endpoint resource name will be used in local predictions and\n",
+        "# @markdown integration with ADK below.\n",
+        "endpoint_resource_name = endpoint.resource_name\n",
+        "print(\"The deployed endpoint resource name is:\")\n",
+        "print(endpoint_resource_name)\n",
+        "# @markdown Click \"Show Code\" to see more details.\n"
+      ],
+      "metadata": {
+        "id": "gbfFZLT5KkUV",
+        "cellView": "form"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# @title Test The Endpoint\n",
+        "# endpoint_resource_name = (\n",
+        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{ENDPOINT_NAME}\"\n",
+        "# )\n",
+        "endpoint = aiplatform.Endpoint(endpoint_resource_name)\n",
+        "\n",
+        "location = endpoint_resource_name.split('/')[3]\n",
+        "base_url = f\"https://{location}-aiplatform.googleapis.com/v1beta1/{endpoint.resource_name}\"\n",
+        "\n",
+        "# @markdown Predict locally with some requests.\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 = 100  # @param {type: \"integer\"}\n",
+        "temperature = 1.0  # @param {type: \"number\"}\n",
+        "stream = True  # @param {type: \"boolean\"}\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",
+        "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.choices[0].message.content)"
+      ],
+      "metadata": {
+        "id": "_MNVrfJomCd0",
+        "cellView": "form"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Aa4e1-6FvRAP",
+        "cellView": "form"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Build Agent Web App Dockers With VMG Endpoints\n",
+        "# @markdown The section will create required python and docker files first, and\n",
+        "# @markdown then build the dockers with cloud build.\n",
+        "\n",
+        "# @markdown 1. Create `agent.py` by loading VMG endpoints and example tool functions.\n",
+        "agent_app = '''\n",
+        "\"\"\"This is a sample agent for model garden agents.\"\"\"\n",
+        "\n",
+        "import datetime\n",
+        "import os\n",
+        "import re\n",
+        "import zoneinfo\n",
+        "\n",
+        "from google.adk.agents import LlmAgent\n",
+        "from google.adk.models.lite_llm import LiteLlm\n",
+        "import google.auth\n",
+        "\n",
+        "_MODEL_GARDEN_ENDPOINT_REGEX = r\"projects\\/.+\\/locations\\/.+\\/endpoints\\/.+\"\n",
+        "\n",
+        "\n",
+        "def get_weather(city: str) -\u003e str:\n",
+        "  \"\"\"Simulates a web search. Use it get information on weather.\n",
+        "\n",
+        "  Args:\n",
+        "      city: A string containing the location to get weather information for.\n",
+        "\n",
+        "  Returns:\n",
+        "      A string with the simulated weather information for the queried city.\n",
+        "  \"\"\"\n",
+        "  if \"sf\" in city.lower() or \"san francisco\" in city.lower():\n",
+        "    return \"It's 70 degrees and foggy.\"\n",
+        "  return \"It's 80 degrees and sunny.\"\n",
+        "\n",
+        "\n",
+        "def get_current_time(city: str) -\u003e str:\n",
+        "  \"\"\"Simulates getting the current time for a city.\n",
+        "\n",
+        "  Args:\n",
+        "      city: The name of the city to get the current time for.\n",
+        "\n",
+        "  Returns:\n",
+        "      A string with the current time information.\n",
+        "  \"\"\"\n",
+        "  if \"sf\" in city.lower() or \"san francisco\" in city.lower():\n",
+        "    tz_identifier = \"America/Los_Angeles\"\n",
+        "  else:\n",
+        "    return f\"Sorry, I don't have timezone information for city: {city}.\"\n",
+        "\n",
+        "  tz = zoneinfo.ZoneInfo(tz_identifier)\n",
+        "  now = datetime.datetime.now(tz)\n",
+        "  return (\n",
+        "      f\"The current time for city {city} is\"\n",
+        "      f\" {now.strftime('%Y-%m-%d %H:%M:%S %Z%z')}\"\n",
+        "  )\n",
+        "\n",
+        "\n",
+        "def _get_auth_headers() -\u003e dict[str, str]:\n",
+        "  \"\"\"Gets the auth headers for the model garden endpoint.\"\"\"\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",
+        "  return {\n",
+        "      \"Content-Type\": \"application/json\",\n",
+        "      \"Authorization\": f\"Bearer {creds.token}\",\n",
+        "  }\n",
+        "\n",
+        "\n",
+        "def _setup_model_garden_endpoint():\n",
+        "  \"\"\"Sets up the model garden endpoint.\"\"\"\n",
+        "  endpoint = os.environ.get(\"GOOGLE_MODEL_GARDEN_ENDPOINT\", \"\")\n",
+        "\n",
+        "  if not re.compile(_MODEL_GARDEN_ENDPOINT_REGEX).fullmatch(endpoint):\n",
+        "    raise ValueError(\n",
+        "        f\"Invalid model garden endpoint: {endpoint}. Please use the format\"\n",
+        "        \" projects/{project}/locations/{location}/endpoints/{endpoint}.\"\n",
+        "    )\n",
+        "  endpoint_parts = endpoint.split(\"/\")\n",
+        "  os.environ.setdefault(\"GOOGLE_GENAI_USE_VERTEXAI\", \"True\")\n",
+        "  os.environ[\"VERTEXAI_PROJECT\"] = endpoint_parts[1]\n",
+        "  os.environ[\"VERTEXAI_LOCATION\"] = endpoint_parts[3]\n",
+        "  os.environ[\"LITELLM_LOG\"] = \"DEBUG\"\n",
+        "  return f\"vertex_ai/openai/{endpoint_parts[5]}\"\n",
+        "\n",
+        "\n",
+        "auth_headers = _get_auth_headers()\n",
+        "model = _setup_model_garden_endpoint()\n",
+        "\n",
+        "print(\"The current model is: {model}\")\n",
+        "\n",
+        "root_agent = LlmAgent(\n",
+        "    name=\"root_agent\",\n",
+        "    model=LiteLlm(\n",
+        "        model=model,\n",
+        "        extra_headers=auth_headers,\n",
+        "    ),\n",
+        "    instruction=(\n",
+        "        \"You are a helpful AI assistant designed to provide accurate and useful\"\n",
+        "        \" information. Please output the tool callings with json format if\"\n",
+        "        \" exists.\"\n",
+        "    ),\n",
+        "    description=\"Retrieves the weather and current time using specific tools.\",\n",
+        "    tools=[get_weather, get_current_time],\n",
+        ")\n",
+        "'''\n",
+        "with BUILD_DIR.joinpath(\"agent.py\").open(\"w\") as f:\n",
+        "    f.write(agent_app)\n",
+        "\n",
+        "# @markdown 2. Create `__init__.py` to load agent.py for ADK apps.\n",
+        "initialize = '''\n",
+        "from . import agent\n",
+        "'''\n",
+        "with BUILD_DIR.joinpath(\"__init__.py\").open(\"w\") as f:\n",
+        "    f.write(initialize)\n",
+        "\n",
+        "# @markdown 3. Create `Dockerfile` to build agent app dockers.\n",
+        "dockerfile_content = '''\n",
+        "FROM python:3.11-slim\n",
+        "WORKDIR /app\n",
+        "RUN adduser --disabled-password --gecos \"\" myuser\n",
+        "RUN chown -R myuser:myuser /app\n",
+        "USER myuser\n",
+        "ENV PATH=\"/home/myuser/.local/bin:$PATH\"\n",
+        "RUN pip install \\\n",
+        "  google-adk~=0.4.0 \\\n",
+        "  google-cloud-logging~=3.11.4 \\\n",
+        "  opentelemetry-exporter-gcp-trace~=1.9.0 \\\n",
+        "  google-cloud-aiplatform[evaluation,agent-engines]~=1.88.0 \\\n",
+        "  litellm~=1.66.2\n",
+        "\n",
+        "COPY agent.py \"/app/agents/model_garden_agents/\"\n",
+        "COPY __init__.py \"/app/agents/model_garden_agents/\"\n",
+        "ENV GOOGLE_MODEL_GARDEN_ENDPOINT YOUR_ENDPOINT\n",
+        "EXPOSE 8000\n",
+        "CMD adk web --port=8000 --trace_to_cloud \"/app/agents\"\n",
+        "'''\n",
+        "with BUILD_DIR.joinpath(\"Dockerfile\").open(\"w\") as f:\n",
+        "    f.write(dockerfile_content)\n",
+        "\n",
+        "# @markdown 4. Build agent web app dockers.\n",
+        "VMG_AGENT_UI_CONTAINER_IMAGE_URI = (\n",
+        "    f\"{REGION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY_NAME}/vmg-adk-ui\"\n",
+        ")\n",
+        "! gcloud builds submit --tag $VMG_AGENT_UI_CONTAINER_IMAGE_URI --project $PROJECT_ID --machine-type e2-highcpu-32 $BUILD_DIR\n",
+        "print(\"The agent UI docker is :\")\n",
+        "print(VMG_AGENT_UI_CONTAINER_IMAGE_URI)\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# @title Deploy Agent Web App Dockers To Cloud Run\n",
+        "# @markdown After the deployment, there will be a service URL. You can click\n",
+        "# @markdown the service URL and interact with the agent web app. The agent web\n",
+        "# @markdown is a built-in development UI in [ADK](https://github.com/google/adk-python?tab=readme-ov-file)\n",
+        "# @markdown to help you test, evaluate, debug, and showcase your agent(s).\n",
+        "\n",
+        "# @markdown ![ADK WEB UI](https://raw.githubusercontent.com/google/adk-python/main/assets/adk-web-dev-ui-function-call.png)\n",
+        "\n",
+        "! gcloud run deploy vmg-agent-ui-1 \\\n",
+        "    --port 8000 \\\n",
+        "    --image=\"{VMG_AGENT_UI_CONTAINER_IMAGE_URI}\" \\\n",
+        "    --region=\"{REGION}\" \\\n",
+        "    --platform=managed \\\n",
+        "    --allow-unauthenticated \\\n",
+        "    --memory=1024Mi \\\n",
+        "    --set-env-vars=\"GOOGLE_MODEL_GARDEN_ENDPOINT={endpoint_resource_name}\"\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ],
+      "metadata": {
+        "id": "L2AGmQZmuVam",
+        "cellView": "form"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tAelDidov5AW"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "8SeZCFo5v7z-"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown  Delete the experiment resources to avoid unnecessary continuous\n",
+        "# @markdown  charges that may incur.\n",
+        "\n",
+        "delete_endpoint = False # @param {type:\"boolean\"}\n",
+        "delete_artifact_registry = False # @param {type:\"boolean\"}\n",
+        "delete_tutorial_folder = False # @param {type:\"boolean\"}\n",
+        "\n",
+        "if delete_endpoint:\n",
+        "  # Undeploy model and delete endpoint.\n",
+        "  endpoint.delete(force=True)\n",
+        "\n",
+        "if delete_artifact_registry:\n",
+        "    ! gcloud artifacts repositories delete $REPOSITORY_NAME \\\n",
+        "          --repository-format=docker \\\n",
+        "          --location=$REGION \\\n",
+        "          --project=$PROJECT_ID\n",
+        "\n",
+        "if delete_tutorial_folder:\n",
+        "    import shutil\n",
+        "    shutil.rmtree(TUTORIAL_DIR)\n"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_integration_with_adk.ipynb",
+      "toc_visible": true,
+      "provenance": [],
+      "collapsed_sections": [
+        "tAelDidov5AW"
+      ]
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
@@ -26,7 +26,6 @@
    },
    {
      "cell_type": "markdown",
-      "language": "markdown",
      "metadata": {
        "id": "VJWDivOv3OWy"
      },
@@ -35,6 +34,11 @@
        "\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",
@@ -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.84.0'\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,71 +139,38 @@
        "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",
        "# 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",
-        "# 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",
@@ -205,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",
@@ -235,6 +224,339 @@
        "    \"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.preview 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",
+        ")"
+      ]
+    },
+    {
+      "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",
@@ -277,296 +599,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",
@@ -576,10 +610,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."
@@ -706,7 +745,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",
@@ -891,6 +932,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",
@@ -900,11 +943,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()"
      ]
    }
  ],
@@ -277,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\""
      ]
    },
    {
@@ -415,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",
@@ -529,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",
@@ -545,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",
@@ -561,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",
@@ -590,6 +604,7 @@
        "    accelerator_type=accelerator_type,\n",
        "    accelerator_count=accelerator_count,\n",
        "    resolution=model_resolution,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
        ")"
      ]
    },
@@ -619,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",
@@ -628,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",
@@ -427,7 +427,7 @@
        "id": "L_q9h-SArI0c"
      },
      "source": [
-        "You must provide a Hugging Face User Access Token (read) to access the Llama 3.1 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."
+        "You must provide a Hugging Face User Access Token (with read access) to access the Llama 3.1 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."
      ]
    },
    {
@@ -1309,7 +1309,8 @@
        "    train_job.wait()\n",
        "    print(\"The training job has finished.\")\n",
        "\n",
-        "use_dedicated_endpoint = False\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",
        "if \"8b\" in base_model_id.lower():\n",
@@ -59,13 +59,14 @@
      "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",
@@ -220,14 +221,14 @@
        "\n",
        "# @markdown Select one of the three model variations.\n",
        "\n",
-        "base_model_name = \"Llama-Guard-3-11B-Vision\"  # @param [\"Llama-Guard-3-8B\", \"Llama-Guard-3-1B\", \"Llama-Guard-3-11B-Vision\"] {allow-input: true, isTemplate: true}\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:20241202_0916_RC00_maas\"\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",
@@ -238,7 +239,7 @@
        "    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",
@@ -290,7 +291,7 @@
      "source": [
        "# @title [Option 2] Deploy with customized configurations\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 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",
@@ -422,7 +423,7 @@
        "\n",
        "\n",
        "models[LABEL], endpoints[LABEL] = deploy_model_vllm(\n",
-        "    model_name=common_util.get_job_name_with_datetime(prefix=\"llama3-guard\"),\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=\"llama-guard\"),\n",
        "    model_id=model_id,\n",
        "    publisher=\"meta\",\n",
        "    publisher_model_id=\"llama-guard\",\n",
@@ -432,10 +433,10 @@
        "    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=True,\n",
+        "    enable_llama_tool_parser=False,\n",
        ")\n",
        "# @markdown Click \"Show Code\" to see more details."
      ]
@@ -686,6 +687,131 @@
        "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)"
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -340,9 +340,7 @@
        "\n",
        "if use_dedicated_endpoint:\n",
        "    DEDICATED_ENDPOINT_DNS = endpoints[\"ollama\"].gca_resource.dedicated_endpoint_dns\n",
-        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
-        "    PROJECT_ID, REGION, endpoints[\"ollama\"].name\n",
-        ")\n",
+        "ENDPOINT_RESOURCE_NAME = endpoints[\"ollama\"].resource_name\n",
        "\n",
        "# @title Chat Completions Inference\n",
        "\n",
@@ -487,7 +487,7 @@
        "    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",
+        "# @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",
@@ -674,6 +674,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",
@@ -714,6 +715,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",
@@ -765,7 +770,7 @@
        "    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",
+        "# @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",
@@ -218,15 +218,17 @@
        "# @markdown | Model Version              | Default Max Model Length | Available GPU configurations |\n",
        "# @markdown |----------------------------|------------------|-----------------------------|\n",
        "# @markdown | Phi-4                      | 16384            | 1 NVIDIA_A100 80GB a2-ultragpu-1g, 2 NVIDIA_L4 g2-standard-24  |\n",
+        "# @markdown | Phi-4-reasoning            | 32768            | 1 NVIDIA_A100 80GB a2-ultragpu-1g, 1 NVIDIA_H100 80GB a3-highgpu-1g  |\n",
+        "# @markdown | Phi-4-reasoning-plus       | 32768            | 1 NVIDIA_A100 80GB a2-ultragpu-1g, 1 NVIDIA_H100 80GB a3-highgpu-1g |\n",
        "\n",
        "# The pre-built serving docker images.\n",
-        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/vllm-inference.cu121.0-6.ubuntu2204.py310\"\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20250417_0916_RC01\"\n",
        "\n",
-        "MODEL_ID = \"Phi-4\"\n",
+        "MODEL_ID = \"Phi-4\"  # @param [\"Phi-4\", \"Phi-4-reasoning\", \"Phi-4-reasoning-plus\"] {isTemplate:true}\n",
        "model_path_prefix = \"microsoft\"\n",
        "model_id = os.path.join(model_path_prefix, MODEL_ID)\n",
        "\n",
-        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_A100_80GB\"] {isTemplate: true}\n",
+        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_A100_80GB\", \"NVIDIA_H100_80GB\"] {isTemplate: true}\n",
        "machine_type = None\n",
        "vllm_dtype = \"bfloat16\"\n",
        "accelerator_count = None\n",
@@ -234,7 +236,22 @@
        "gpu_memory_utilization = None\n",
        "enable_trust_remote_code = False\n",
        "\n",
-        "if \"Phi-4\" == MODEL_ID:\n",
+        "if \"Phi-4-reasoning\" in MODEL_ID:\n",
+        "    max_model_len = 32768\n",
+        "    if accelerator_type == \"NVIDIA_A100_80GB\":\n",
+        "        accelerator_count = 1\n",
+        "        machine_type = \"a2-ultragpu-1g\"\n",
+        "        gpu_memory_utilization = 0.85\n",
+        "    elif accelerator_type == \"NVIDIA_H100_80GB\":\n",
+        "        accelerator_count = 1\n",
+        "        machine_type = \"a3-highgpu-1g\"\n",
+        "        gpu_memory_utilization = 0.85\n",
+        "    else:\n",
+        "        raise ValueError(\n",
+        "            \"Recommended machine settings not found for accelerator type: %s\"\n",
+        "            % accelerator_type\n",
+        "        )\n",
+        "elif \"Phi-4\" == MODEL_ID:\n",
        "    max_model_len = 16384\n",
        "    if accelerator_type == \"NVIDIA_L4\":\n",
        "        accelerator_count = 2\n",
@@ -388,7 +405,7 @@
        "    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",
+        "# @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",
@@ -499,13 +516,13 @@
        "# @markdown This section uploads prebuilt Phi-4 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",
        "# @markdown Select one of the four model variations.\n",
-        "MODEL_ID = \"Phi-4\"\n",
+        "MODEL_ID = \"Phi-4\"  # @param [\"Phi-4\", \"Phi-4-reasoning\", \"Phi-4-reasoning-plus\"] {isTemplate:true}\n",
        "TPU_DEPLOYMENT_REGION = \"us-west1\"  # @param [\"us-west1\"] {isTemplate:true}\n",
        "model_path_prefix = \"microsoft\"\n",
        "model_id = os.path.join(model_path_prefix, MODEL_ID)\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",
+        "HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:phi4\"\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",
@@ -513,14 +530,21 @@
        "# @markdown | Model Version              | Default Max Model Length | Default TPU configuration |\n",
        "# @markdown |----------------------------|------------------|-----------------------------|\n",
        "# @markdown | Phi-4                      | 16384            | 4 TPU_V5e ct5lp-hightpu-4t  |\n",
+        "# @markdown | Phi-4-reasoning            | 32768            | 4 TPU_V5e ct5lp-hightpu-4t  |\n",
+        "# @markdown | Phi-4-reasoning-plus       | 32768            | 4 TPU_V5e ct5lp-hightpu-4t  |\n",
        "\n",
        "\n",
        "# Note: 1 TPU V5 chip has only one core.\n",
        "tpu_type = \"TPU_V5e\"\n",
        "\n",
-        "if \"Phi-4\" in MODEL_ID:\n",
+        "if \"Phi-4-reasoning\" in MODEL_ID:\n",
        "    tpu_count = 4\n",
-        "    tpu_topo = \"4x4\"\n",
+        "    tpu_topo = \"1x4\"\n",
+        "    max_model_len = 32768\n",
+        "    machine_type = \"ct5lp-hightpu-4t\"\n",
+        "elif \"Phi-4\" in MODEL_ID:\n",
+        "    tpu_count = 4\n",
+        "    tpu_topo = \"1x4\"\n",
        "    max_model_len = 16384\n",
        "    machine_type = \"ct5lp-hightpu-4t\"\n",
        "else:\n",
@@ -561,6 +585,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",
@@ -601,6 +626,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",
@@ -652,7 +681,7 @@
        "    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",
+        "# @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",
@@ -4,11 +4,12 @@
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
+        "cellView": "form",
        "id": "7d9bbf86da5e"
      },
      "outputs": [],
      "source": [
-        "# Copyright 2023 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",
@@ -31,26 +32,23 @@
      "source": [
        "# Vertex AI Model Garden - BLIP Image Captioning\n",
        "\n",
-        "<table align=\"left\">\n",
-        "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_blip_image_captioning.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\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>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_pytorch_blip_image_captioning.ipynb\">\n",
+        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_blip_image_captioning.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_pytorch_blip_image_captioning.ipynb\">\n",
-        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
-        "Open in Vertex AI Workbench\n",
-        "    </a>\n",
-        "    (a Python-3 CPU notebook is recommended)\n",
-        "  </td>\n",
-        "</table>"
+        "</tr></tbody></table>"
      ]
    },
    {
@@ -69,6 +67,10 @@
        "- Deploy the model on [Endpoint](https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints).\n",
        "- Run online predictions for image captioning.\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",
@@ -76,199 +78,155 @@
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
-        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [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."
+        "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"
+        "id": "G3GYgynjaKG8"
      },
      "source": [
-        "## Setup environment\n",
-        "\n",
-        "**NOTE**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "d73ffa0c0b83"
-      },
-      "source": [
-        "### Colab only"
+        "## Before you begin"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "2707b02ef5df"
+        "cellView": "form",
+        "id": "cDtpu7ZVaTFb"
      },
      "outputs": [],
      "source": [
-        "if \"google.colab\" in str(get_ipython()):\n",
-        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
-        "    from google.colab import auth as google_auth\n",
+        "# @title Setup Google Cloud project\n",
        "\n",
-        "    google_auth.authenticate_user()\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
-        "    # Restart the notebook kernel after installs.\n",
-        "    import IPython\n",
+        "# @markdown 2. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
        "\n",
-        "    app = IPython.Application.instance()\n",
-        "    app.kernel.do_shutdown(True)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0f826ff482a2"
-      },
-      "source": [
-        "### Setup Google Cloud project\n",
+        "REGION = \"\"  # @param {type:\"string\"}\n",
        "\n",
-        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\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",
-        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\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",
-        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
+        "# Import the necessary packages\n",
        "\n",
-        "1. [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "8958ebc71868"
-      },
-      "source": [
-        "Fill following variables for experiments environment:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "9db30f827a65"
-      },
-      "outputs": [],
-      "source": [
-        "# Cloud project id.\n",
-        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.84.0'\n",
        "\n",
-        "# The region you want to launch jobs in.\n",
-        "REGION = \"us-central1\"  # @param {type:\"string\"}\n",
+        "# Import the necessary packages.\n",
+        "import importlib\n",
+        "import os\n",
+        "from typing import Tuple\n",
        "\n",
-        "# The Cloud Storage bucket for storing experiments output. Fill it without the 'gs://' prefix.\n",
-        "GCS_BUCKET = \"\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "92f16e22c20b"
-      },
-      "source": [
-        "Initialize Vertex AI API:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "1680c257acfb"
-      },
-      "outputs": [],
-      "source": [
        "from google.cloud import aiplatform\n",
        "\n",
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=GCS_BUCKET)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "6ca48b699d17"
-      },
-      "source": [
-        "### Define constants"
+        "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 = \"transformers_gpu\"\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",
+        "# 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",
+        ")\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "de9882ea89ea"
+        "cellView": "form",
+        "id": "RYfNg5r8ezx1"
      },
      "outputs": [],
      "source": [
-        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "10188266a5cd"
-      },
-      "source": [
-        "### Define common functions"
+        "# @title Select the model parameters\n",
+        "\n",
+        "MODEL_ID = \"Salesforce/blip-image-captioning-base\"\n",
+        "TASK = \"image-to-text\"\n",
+        "\n",
+        "# Set the machine specifications.\n",
+        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_V100\", \"NVIDIA_TESLA_T4\"]\n",
+        "if accelerator_type == \"NVIDIA_L4\":\n",
+        "    accelerator_count = 1\n",
+        "    machine_type = \"g2-standard-8\"\n",
+        "elif accelerator_type == \"NVIDIA_TESLA_V100\":\n",
+        "    accelerator_count = 1\n",
+        "    machine_type = \"n1-standard-4\"\n",
+        "elif accelerator_type == \"NVIDIA_TESLA_T4\":\n",
+        "    accelerator_count = 1\n",
+        "    machine_type = \"n1-standard-8\"\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 `training_machine_type`, `training_accelerator_type`, and `per_node_accelerator_count` 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",
+        "# @markdown Click \"Show Code\" to see more details."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
+        "cellView": "form",
        "id": "cac4478ae098"
      },
      "outputs": [],
      "source": [
-        "import base64\n",
-        "import os\n",
-        "from datetime import datetime\n",
-        "from io import BytesIO\n",
+        "# @title Deploy with customized configs\n",
        "\n",
-        "import requests\n",
-        "from google.cloud import aiplatform\n",
-        "from PIL import Image\n",
+        "# The pre-built serving docker image. It contains serving scripts and models.\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-transformers-serve\"\n",
        "\n",
        "\n",
-        "def create_job_name(prefix):\n",
-        "    user = os.environ.get(\"USER\")\n",
-        "    now = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
-        "    job_name = f\"{prefix}-{user}-{now}\"\n",
-        "    return job_name\n",
-        "\n",
-        "\n",
-        "def download_image(url):\n",
-        "    response = requests.get(url)\n",
-        "    return Image.open(BytesIO(response.content))\n",
-        "\n",
-        "\n",
-        "def image_to_base64(image, format=\"JPEG\"):\n",
-        "    buffer = BytesIO()\n",
-        "    image.save(buffer, format=format)\n",
-        "    image_str = base64.b64encode(buffer.getvalue()).decode(\"utf-8\")\n",
-        "    return image_str\n",
-        "\n",
-        "\n",
-        "def base64_to_image(image_str):\n",
-        "    image = Image.open(BytesIO(base64.b64decode(image_str)))\n",
-        "    return image\n",
-        "\n",
-        "\n",
-        "def image_grid(imgs, rows=2, cols=2):\n",
-        "    w, h = imgs[0].size\n",
-        "    grid = Image.new(\"RGB\", size=(cols * w, rows * h))\n",
-        "    for i, img in enumerate(imgs):\n",
-        "        grid.paste(img, box=(i % cols * w, i // cols * h))\n",
-        "    return grid\n",
-        "\n",
-        "\n",
-        "def deploy_model(model_id, task):\n",
+        "def deploy_model(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    task: str,\n",
+        "    machine_type: str = \"g2-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 1,\n",
+        "    use_dedicated_endpoint: bool = True,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
        "    model_name = \"blip-image-captioning\"\n",
-        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
+        "    endpoint = aiplatform.Endpoint.create(\n",
+        "        display_name=f\"{model_name}-endpoint\",\n",
+        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
+        "    )\n",
        "    serving_env = {\n",
        "        \"MODEL_ID\": model_id,\n",
        "        \"TASK\": task,\n",
@@ -284,7 +242,7 @@
        "        serving_container_health_route=\"/ping\",\n",
        "        serving_container_environment_variables=serving_env,\n",
        "        artifact_uri=artifact_uri,\n",
-        "        model_garden_source_model_name=\"publishers/salesforce/models/blip-image-captioning-base\"\n",
+        "        model_garden_source_model_name=\"publishers/salesforce/models/blip-image-captioning-base\",\n",
        "    )\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
@@ -296,41 +254,25 @@
        "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_blip_image_captioning.ipynb\"\n",
        "        },\n",
        "    )\n",
-        "    return model, endpoint"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "d2d72ecdb8c9"
-      },
-      "source": [
-        "## Upload and deploy models"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "9448c5f545fa"
-      },
-      "source": [
-        "This section uploads the pre-trained model to Model Registry and deploys it on the Endpoint with 1 T4 GPU.\n",
+        "    return model, endpoint\n",
        "\n",
-        "The model deployment step will take ~15 minutes to complete.\n",
        "\n",
-        "Once deployed, you can send images to get descriptions."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "b4b46c28d8b1"
-      },
-      "outputs": [],
-      "source": [
-        "model, endpoint = deploy_model(\n",
-        "    model_id=\"Salesforce/blip-image-captioning-base\", task=\"image-to-text\"\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",
+        "models[LABEL], endpoints[LABEL] = deploy_model(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=\"blip-image-captioning\"),\n",
+        "    model_id=MODEL_ID,\n",
+        "    task=TASK,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
        ")"
      ]
    },
@@ -347,42 +289,49 @@
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
+        "cellView": "form",
        "id": "6be655247cb1"
      },
      "outputs": [],
      "source": [
-        "image = download_image(\"http://images.cocodataset.org/val2017/000000039769.jpg\")\n",
+        "# @title Predict\n",
+        "\n",
+        "INPUT_IMAGE = \"http://images.cocodataset.org/val2017/000000039769.jpg\"  # @param {type: \"string\"}\n",
+        "image = common_util.download_image(INPUT_IMAGE)\n",
        "display(image)\n",
        "\n",
        "instances = [\n",
-        "    {\"image\": image_to_base64(image)},\n",
+        "    {\"image\": common_util.image_to_base64(image)},\n",
        "]\n",
-        "preds = endpoint.predict(instances=instances).predictions\n",
+        "\n",
+        "preds = (\n",
+        "    endpoints[LABEL]\n",
+        "    .predict(instances=instances, use_dedicated_endpoint=use_dedicated_endpoint)\n",
+        "    .predictions\n",
+        ")\n",
        "print(preds)"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "db7ffebdb4be"
-      },
-      "source": [
-        "### Clean up resources"
-      ]
-    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
+        "cellView": "form",
        "id": "2ccf3714dbe9"
      },
      "outputs": [],
      "source": [
+        "# @title Clean up resources\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",
-        "endpoint.delete(force=True)\n",
+        "for endpoint in endpoints.values():\n",
+        "    endpoint.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
-        "model.delete()"
+        "for model in models.values():\n",
+        "    model.delete()"
      ]
    }
  ],
@@ -4,11 +4,12 @@
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
+        "cellView": "form",
        "id": "7d9bbf86da5e"
      },
      "outputs": [],
      "source": [
-        "# Copyright 2023 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",
@@ -31,26 +32,23 @@
      "source": [
        "# Vertex AI Model Garden - CLIP\n",
        "\n",
-        "<table align=\"left\">\n",
-        "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_clip.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\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>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_pytorch_clip.ipynb\">\n",
+        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_clip.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_pytorch_clip.ipynb\">\n",
-        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
-        "Open in Vertex AI Workbench\n",
-        "    </a>\n",
-        "    (a Python-3 CPU notebook is recommended)\n",
-        "  </td>\n",
-        "</table>"
+        "</tr></tbody></table>"
      ]
    },
    {
@@ -69,6 +67,10 @@
        "- Deploy the model on [Endpoint](https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints).\n",
        "- Run online predictions for image captioning.\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",
@@ -76,194 +78,79 @@
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
-        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [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": [
-        "## Setup environment\n",
-        "\n",
-        "**NOTE**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "d73ffa0c0b83"
-      },
-      "source": [
-        "### Colab only"
+        "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": "code",
      "execution_count": null,
      "metadata": {
-        "id": "2707b02ef5df"
+        "cellView": "form",
+        "id": "_53r3GqThdEP"
      },
      "outputs": [],
      "source": [
-        "if \"google.colab\" in str(get_ipython()):\n",
-        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
-        "    from google.colab import auth as google_auth\n",
+        "# @title Setup Google Cloud project\n",
        "\n",
-        "    google_auth.authenticate_user()\n",
+        "# Import the necessary packages\n",
        "\n",
-        "    # Restart the notebook kernel after installs.\n",
-        "    import IPython\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.84.0'\n",
        "\n",
-        "    app = IPython.Application.instance()\n",
-        "    app.kernel.do_shutdown(True)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0f826ff482a2"
-      },
-      "source": [
-        "### 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",
-        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\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",
-        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
        "\n",
-        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
+        "# @markdown 3. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
        "\n",
-        "1. [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "8958ebc71868"
-      },
-      "source": [
-        "Fill following variables for experiments environment:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "9db30f827a65"
-      },
-      "outputs": [],
-      "source": [
-        "# Cloud project id.\n",
-        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
+        "REGION = \"\"  # @param {type:\"string\"}\n",
        "\n",
-        "# The region you want to launch jobs in.\n",
-        "REGION = \"us-central1\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# The Cloud Storage bucket for storing experiments output. Fill it without the 'gs://' prefix.\n",
-        "GCS_BUCKET = \"\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "92f16e22c20b"
-      },
-      "source": [
-        "Initialize Vertex AI API:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "1680c257acfb"
-      },
-      "outputs": [],
-      "source": [
-        "from google.cloud import aiplatform\n",
-        "\n",
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=GCS_BUCKET)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "6ca48b699d17"
-      },
-      "source": [
-        "### Define constants"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "de9882ea89ea"
-      },
-      "outputs": [],
-      "source": [
-        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "10188266a5cd"
-      },
-      "source": [
-        "### Define common functions"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "cac4478ae098"
-      },
-      "outputs": [],
-      "source": [
-        "import base64\n",
+        "import importlib\n",
        "import os\n",
-        "from datetime import datetime\n",
-        "from io import BytesIO\n",
        "\n",
-        "import requests\n",
        "from google.cloud import aiplatform\n",
-        "from PIL import Image\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",
-        "def create_job_name(prefix):\n",
-        "    user = os.environ.get(\"USER\")\n",
-        "    now = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
-        "    job_name = f\"{prefix}-{user}-{now}\"\n",
-        "    return job_name\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",
-        "def download_image(url):\n",
-        "    response = requests.get(url)\n",
-        "    return Image.open(BytesIO(response.content))\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",
-        "def image_to_base64(image, format=\"JPEG\"):\n",
-        "    buffer = BytesIO()\n",
-        "    image.save(buffer, format=format)\n",
-        "    image_str = base64.b64encode(buffer.getvalue()).decode(\"utf-8\")\n",
-        "    return image_str\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
        "\n",
+        "# The pre-built serving docker image. It contains serving scripts and models.\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-transformers-serve\"\n",
        "\n",
-        "def base64_to_image(image_str):\n",
-        "    image = Image.open(BytesIO(base64.b64decode(image_str)))\n",
-        "    return image\n",
+        "models, endpoints = {}, {}\n",
+        "LABEL = \"endpoint\"\n",
        "\n",
+        "import vertexai\n",
        "\n",
-        "def image_grid(imgs, rows=2, cols=2):\n",
-        "    w, h = imgs[0].size\n",
-        "    grid = Image.new(\"RGB\", size=(cols * w, rows * h))\n",
-        "    for i, img in enumerate(imgs):\n",
-        "        grid.paste(img, box=(i % cols * w, i // cols * h))\n",
-        "    return grid\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")\n",
+        "\n",
+        "version_id = \"clip-vit-base-patch32\"\n",
+        "PUBLISHER_MODEL_NAME = f\"publishers/openai/models/clip-vit-base-patch32@{version_id}\"\n",
        "\n",
        "\n",
        "def deploy_model(model_id, task):\n",
@@ -284,7 +171,7 @@
        "        serving_container_health_route=\"/ping\",\n",
        "        serving_container_environment_variables=serving_env,\n",
        "        artifact_uri=artifact_uri,\n",
-        "        model_garden_source_model_name=\"publishers/openai/models/clip-vit-base-patch32\"\n",
+        "        model_garden_source_model_name=\"publishers/openai/models/clip-vit-base-patch32\",\n",
        "    )\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
@@ -293,147 +180,89 @@
        "        accelerator_count=1,\n",
        "        deploy_request_timeout=1800,\n",
        "        system_labels={\n",
-        "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_clip.ipynb\"\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_clip.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
        "        },\n",
        "    )\n",
        "    return model, endpoint"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "d2d72ecdb8c9"
-      },
-      "source": [
-        "## Upload and deploy models"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "9448c5f545fa"
-      },
-      "source": [
-        "This section uploads the pre-trained model to Model Registry and deploys it on the Endpoint with 1 T4 GPU.\n",
-        "\n",
-        "The model deployment step will take ~15 minutes to complete."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "99f3c6b404b5"
-      },
-      "source": [
-        "### Zero-shot image classification"
-      ]
-    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "b4b46c28d8b1"
+        "cellView": "form",
+        "id": "LQir_tczhdEP"
      },
      "outputs": [],
      "source": [
-        "model, endpoint = deploy_model(\n",
-        "    model_id=\"openai/clip-vit-base-patch32\", task=\"zero-shot-image-classification\"\n",
+        "accelerator_type = \"NVIDIA_TESLA_T4\"\n",
+        "machine_type = \"n1-standard-8\"\n",
+        "accelerator_count = 1\n",
+        "\n",
+        "# @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.preview 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",
+        "    accept_eula=True,  # Accept the End User License Agreement (EULA) on the model card before deploy. Otherwise, the deployment will be forbidden.\n",
        ")"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "80b3fd2ace09"
-      },
-      "source": [
-        "NOTE: The model weights will be downloaded after the deployment succeeds. Thus additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint."
-      ]
-    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "6be655247cb1"
+        "cellView": "form",
+        "id": "vBvfVTWOhdEP"
      },
      "outputs": [],
      "source": [
-        "image1 = download_image(\"http://images.cocodataset.org/val2017/000000039769.jpg\")\n",
-        "image2 = download_image(\"http://images.cocodataset.org/val2017/000000000285.jpg\")\n",
-        "grid = image_grid([image1, image2], 1, 2)\n",
+        "# @title [Option 2] Deploy with customized configs\n",
+        "\n",
+        "# @markdown This section uploads the pre-trained model to Model Registry and deploys it on the Endpoint with 1 T4 GPU.\n",
+        "# @markdown The model deployment step will take ~15 minutes to complete.\n",
+        "\n",
+        "models[LABEL], endpoints[LABEL] = deploy_model(\n",
+        "    model_id=\"openai/clip-vit-base-patch32\", task=\"zero-shot-image-classification\"\n",
+        ")\n",
+        "\n",
+        "image1 = common_util.download_image(\n",
+        "    \"http://images.cocodataset.org/val2017/000000039769.jpg\"\n",
+        ")\n",
+        "image2 = common_util.download_image(\n",
+        "    \"http://images.cocodataset.org/val2017/000000000285.jpg\"\n",
+        ")\n",
+        "grid = common_util.image_grid([image1, image2], 1, 2)\n",
        "display(grid)\n",
        "\n",
        "instances = [\n",
-        "    {\"image\": image_to_base64(image1), \"text\": \"two cats\"},\n",
-        "    {\"image\": image_to_base64(image2), \"text\": \"a bear\"},\n",
+        "    {\"image\": common_util.image_to_base64(image1), \"text\": \"two cats\"},\n",
+        "    {\"image\": common_util.image_to_base64(image2), \"text\": \"a bear\"},\n",
        "]\n",
-        "preds = endpoint.predict(instances=instances).predictions\n",
-        "print(preds)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "2ccf3714dbe9"
-      },
-      "outputs": [],
-      "source": [
-        "# Undeploy model and delete endpoint.\n",
-        "endpoint.delete(force=True)\n",
+        "preds = endpoints[LABEL].predict(instances=instances).predictions\n",
+        "print(preds)\n",
        "\n",
-        "# Delete models.\n",
-        "model.delete()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ddf9e9ec7b58"
-      },
-      "source": [
-        "### Image/text feature embedding"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "8897bbac6887"
-      },
-      "outputs": [],
-      "source": [
-        "model, endpoint = deploy_model(\n",
+        "models[LABEL], endpoints[LABEL] = deploy_model(\n",
        "    model_id=\"openai/clip-vit-base-patch32\", task=\"feature-embedding\"\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "9f854c45f8f3"
-      },
-      "source": [
-        "NOTE: The model weights will be downloaded after the deployment succeeds. Thus additional 5 minutes of waiting time is needed after the above model deployment step succeeds and before you run the next step below. Otherwise you might see a ServiceUnavailable: 503 502:Bad Gateway error when you send requests to the endpoint."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "2abd54335f36"
-      },
-      "outputs": [],
-      "source": [
+        ")\n",
+        "\n",
        "import numpy as np\n",
        "\n",
        "# Extract feature embedding of images.\n",
-        "image = download_image(\"http://images.cocodataset.org/val2017/000000039769.jpg\")\n",
+        "image = common_util.download_image(\n",
+        "    \"http://images.cocodataset.org/val2017/000000039769.jpg\"\n",
+        ")\n",
        "display(image)\n",
        "instances = [\n",
-        "    {\"image\": image_to_base64(image)},\n",
+        "    {\"image\": common_util.image_to_base64(image)},\n",
        "]\n",
-        "preds = endpoint.predict(instances=instances).predictions\n",
+        "preds = endpoints[LABEL].predict(instances=instances).predictions\n",
        "image_features = np.array(preds[0][\"image_features\"])\n",
        "print(image_features.shape)\n",
        "\n",
@@ -442,7 +271,7 @@
        "    {\"text\": \"two cats\"},\n",
        "    {\"text\": \"hello world\"},\n",
        "]\n",
-        "preds = endpoint.predict(instances=instances).predictions\n",
+        "preds = endpoints[LABEL].predict(instances=instances).predictions\n",
        "text_features = np.array(preds[0][\"text_features\"])\n",
        "print(text_features.shape)"
      ]
@@ -451,15 +280,21 @@
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
+        "cellView": "form",
        "id": "712eb9d0b336"
      },
      "outputs": [],
      "source": [
+        "# @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",
-        "endpoint.delete(force=True)\n",
+        "for endpoint in endpoints.values():\n",
+        "    endpoint.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
-        "model.delete()"
+        "for model in models.values():\n",
+        "    model.delete()"
      ]
    }
  ],
@@ -34,6 +34,11 @@
        "\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_pytorch_codellama.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",
@@ -60,6 +65,10 @@
        "\n",
        "- Deploy pre-trained Code LLaMA models with [vLLM](https://github.com/vllm-project/vllm) with best serving throughput.\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",
@@ -92,15 +101,11 @@
        "\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). 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",
+        "# @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",
@@ -112,82 +117,45 @@
        "# Import the necessary packages\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.84.0'\n",
+        "# ! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
        "\n",
-        "import datetime\n",
        "import importlib\n",
        "import os\n",
-        "import uuid\n",
        "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",
        "# 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",
-        "# 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, \"codellama\")\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\""
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")"
      ]
    },
    {
@@ -232,18 +200,17 @@
        "\n",
        "# @markdown The vLLM project is an highly optimized LLM serving framework which can increase serving throughput a lot. The higher QPS you have, the more benefits you get using vLLM.\n",
        "\n",
-        "\n",
        "# @markdown Set the model name.\n",
        "model_name = \"CodeLlama-7b-Instruct-hf\"  # @param [\"CodeLlama-7b-hf\", \"CodeLlama-7b-Python-hf\", \"CodeLlama-7b-Instruct-hf\", \"CodeLlama-13b-hf\", \"CodeLlama-13b-Python-hf\", \"CodeLlama-13b-Instruct-hf\", \"CodeLlama-34b-hf\", \"CodeLlama-34b-Python-hf\", \"CodeLlama-34b-Instruct-hf\", \"CodeLlama-70b-hf\", \"CodeLlama-70b-Python-hf\", \"CodeLlama-70b-Instruct-hf\"]\n",
-        "common_util.copy_model_artifacts(\n",
-        "    model_name, VERTEX_AI_MODEL_GARDEN_CODE_LLAMA, MODEL_BUCKET\n",
-        ")\n",
-        "model_id = os.path.join(MODEL_BUCKET, model_name)\n",
+        "version_id = model_name.lower()\n",
+        "PUBLISHER_MODEL_NAME = f\"publishers/meta/models/codellama-7b-hf@{version_id}\"\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:20241001_0916_RC00\"\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240620_1616_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",
+        "model_id = os.path.join(VERTEX_AI_MODEL_GARDEN_CODE_LLAMA, model_name)\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",
@@ -252,7 +219,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",
@@ -363,7 +329,6 @@
        "        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_pytorch_codellama.ipynb\",\n",
        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
@@ -393,7 +358,7 @@
        "        vllm_precision = \"float16\"\n",
        "    # Sets 1 L4 (24G) to deploy 7B models.\n",
        "    elif accelerator_type == \"NVIDIA_L4\":\n",
-        "        machine_type = \"g2-standard-8\"\n",
+        "        machine_type = \"g2-standard-12\"\n",
        "        accelerator_count = 1\n",
        "        vllm_precision = \"bfloat16\"\n",
        "    else:\n",
@@ -418,7 +383,7 @@
        "        vllm_precision = \"bfloat16\"\n",
        "    else:\n",
        "        raise ValueError(\n",
-        "            f\"Recommended GPU setting not found for: {accelerator_type} and {base_model_name}.\"\n",
+        "            f\"Recommended GPU setting not found for: {accelerator_type} and {model_name}.\"\n",
        "        )\n",
        "elif \"34b\" in model_name:\n",
        "    # Sets 2 A100 (40G) to deploy 34B models.\n",
@@ -438,7 +403,7 @@
        "        vllm_precision = \"bfloat16\"\n",
        "    else:\n",
        "        raise ValueError(\n",
-        "            f\"Recommended GPU setting not found for: {accelerator_type} and {base_model_name}.\"\n",
+        "            f\"Recommended GPU setting not found for: {accelerator_type} and {model_name}.\"\n",
        "        )\n",
        "elif \"70b\" in model_name:\n",
        "    # Sets 4 A100 (40G) to deploy 70B models.\n",
@@ -453,7 +418,7 @@
        "        vllm_precision = \"bfloat16\"\n",
        "    else:\n",
        "        raise ValueError(\n",
-        "            f\"Recommended GPU setting not found for: {accelerator_type} and {base_model_name}.\"\n",
+        "            f\"Recommended GPU setting not found for: {accelerator_type} and {model_name}.\"\n",
        "        )\n",
        "\n",
        "# Check quota for the selected GPU type and region.\n",
@@ -463,20 +428,55 @@
        "    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": "UBirHaZJRHzI"
+      },
+      "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.preview 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",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "OW0ykbnLRHzI"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with customized configs\n",
        "\n",
        "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
        "    model_name=common_util.get_job_name_with_datetime(prefix=\"code-llama-serve-vllm\"),\n",
        "    model_id=model_id,\n",
        "    publisher=\"meta\",\n",
        "    publisher_model_id=\"codellama-7b-hf\",\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
        "    machine_type=machine_type,\n",
        "    accelerator_type=accelerator_type,\n",
        "    accelerator_count=accelerator_count,\n",
        "    max_model_len=max_model_len,\n",
        "    dtype=vllm_precision,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
        ")"
      ]
    },
@@ -543,7 +543,6 @@
      "outputs": [],
      "source": [
        "# @title Clean up resources\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",
@@ -553,11 +552,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()"
      ]
    }
  ],
@@ -59,12 +59,16 @@
      "source": [
        "## Overview\n",
        "\n",
-        "This notebook demonstrates serving DeepSeek models with [vLLM](https://github.com/vllm-project/vllm) or [SGLang](https://github.com/sgl-project/sglang). [DeepSeek-V3](https://huggingface.co/deepseek-ai/DeepSeek-V3) is a strong Mixture-of-Experts (MoE) language model with 671B total parameters with 37B activated for each token. [DeepSeek-R1](https://huggingface.co/deepseek-ai/DeepSeek-R1) is one of the first-generation reasoning models introduced by DeepSeek and offers performance comparable to OpenAI-o1 across math, code, and reasoning tasks.\n",
+        "This notebook demonstrates serving DeepSeek models with [vLLM](https://github.com/vllm-project/vllm), [SGLang](https://github.com/sgl-project/sglang), or [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM). [DeepSeek-V3](https://huggingface.co/deepseek-ai/DeepSeek-V3) is a strong Mixture-of-Experts (MoE) language model with 671B total parameters with 37B activated for each token. [DeepSeek-R1](https://huggingface.co/deepseek-ai/DeepSeek-R1) is one of the first-generation reasoning models introduced by DeepSeek and offers performance comparable to OpenAI-o1 across math, code, and reasoning tasks.\n",
        "\n",
        "\n",
        "### Objective\n",
        "\n",
-        "- Deploy DeepSeek-V3 and DeepSeek-R1 with vLLM or SGLang on GPU using single-host and multi-host serving, and [Spot VMs](https://cloud.google.com/compute/docs/instances/spot) (Optional). Multi-host GPU serving is a preview feature.\n",
+        "- Deploy DeepSeek-V3 and DeepSeek-R1 with vLLM, SGLang, or TensorRT-LLM on GPU using single-host and multi-host serving, and [Spot VMs](https://cloud.google.com/compute/docs/instances/spot) (Optional). Multi-host GPU serving is a preview feature.\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",
@@ -106,7 +110,6 @@
    {
      "cell_type": "code",
      "execution_count": null,
-      "language": "python",
      "metadata": {
        "cellView": "form",
        "id": "YXFGIp1l-qtT"
@@ -133,6 +136,7 @@
        "\n",
        "# Import the necessary packages\n",
        "import importlib\n",
+        "import json\n",
        "import os\n",
        "import time\n",
        "from typing import Tuple\n",
@@ -179,6 +183,7 @@
        "        )\n",
        "\n",
        "\n",
+        "LABEL = \"vllm_gpu\"\n",
        "models, endpoints = {}, {}\n",
        "\n",
        "# Get the default cloud project id.\n",
@@ -188,20 +193,10 @@
        "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",
        "\n",
        "import vertexai\n",
@@ -230,7 +225,7 @@
      },
      "outputs": [],
      "source": [
-        "LABEL = \"vllm_gpu\"\n",
+        "# @title Set the model variants\n",
        "\n",
        "# @markdown It's recommended to use the region selected by the deployment button on the model card. If the deployment button is not available, it's recommended to stay with the default region of the notebook.\n",
        "\n",
@@ -247,9 +242,12 @@
        "else:\n",
        "    model_user_id = \"deepseek-v3\"\n",
        "\n",
-        "# fmt: off\n",
-        "PUBLISHER_MODEL_NAME = f\"publishers/deepseek-ai/models/{model_user_id}@{base_model_name.lower()}\"\n",
-        "# fmt: on"
+        "PUBLISHER_MODEL_NAME = (\n",
+        "    f\"publishers/deepseek-ai/models/{model_user_id}@{base_model_name.lower()}\"\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\"}"
      ]
    },
    {
@@ -269,10 +267,6 @@
        "# The pre-built serving docker image for vLLM past v0.7.3, https://github.com/vllm-project/vllm/commit/f6bb18fd9a19e5e4fb1991339638fc666d06b27a.\n",
        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20250312_0916_RC01\"\n",
        "\n",
-        "# @markdown Use a [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint) for the deployment.\n",
-        "# @markdown Only dedicated endpoints are supported at this moment.\n",
-        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
-        "\n",
        "# @markdown Choose whether to use a [Spot VM](https://cloud.google.com/compute/docs/instances/spot) for the deployment.\n",
        "is_spot = False  # @param {type:\"boolean\"}\n",
        "\n",
@@ -529,13 +523,14 @@
        "                \"requiredReplicaCount\": required_replica_count,\n",
        "                \"maxReplicaCount\": max_replica_count,\n",
        "            },\n",
-        "            \"serviceAccount\": service_account,\n",
        "            \"system_labels\": {\n",
        "                \"NOTEBOOK_NAME\": \"model_garden_pytorch_deepseek_deployment.ipynb\",\n",
        "                \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
        "            },\n",
        "        },\n",
        "    }\n",
+        "    if service_account:\n",
+        "        data[\"deployedModel\"][\"serviceAccount\"] = service_account\n",
        "    if is_spot:\n",
        "        data[\"deployedModel\"][\"dedicatedResources\"][\"spot\"] = True\n",
        "    if autoscale_by_gpu_duty_cycle_target > 0 or autoscale_by_cpu_usage_target > 0:\n",
@@ -574,7 +569,6 @@
        "    publisher=\"deepseek-ai\",\n",
        "    publisher_model_id=(\"deepseek-v3\" if \"V3\" in model_id else \"deepseek-r1\"),\n",
        "    base_model_id=hf_model_id,\n",
-        "    service_account=None,\n",
        "    machine_type=machine_type,\n",
        "    accelerator_type=accelerator_type,\n",
        "    accelerator_count=accelerator_count,\n",
@@ -681,9 +675,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 Because the DeepSeek-R1 model generates detailed reasoning steps, the output is expected to be long. We recommend using streaming for a better generation experience.\n",
        "# @title Chat Completions Inference\n",
@@ -779,15 +771,14 @@
        "hf_model_id = model_id\n",
        "\n",
        "# The pre-built serving docker images.\n",
-        "SGLANG_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/sglang-serve.cu124.0-4.ubuntu2204.py310\"\n",
-        "\n",
-        "# @markdown Use a [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint) for the deployment.\n",
-        "# @markdown Only dedicated endpoints are supported at this moment.\n",
-        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
+        "SGLANG_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/sglang-serve.cu124.0-4.ubuntu2204.py310:20250427-1800-rc0\"\n",
        "\n",
        "# @markdown Choose whether to use a [Spot VM](https://cloud.google.com/compute/docs/instances/spot) for the deployment.\n",
        "is_spot = False  # @param {type:\"boolean\"}\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_H200_141GB\"  # @param [\"NVIDIA_H200_141GB\", \"NVIDIA_H100_80GB\"] {isTemplate:true}\n",
        "accelerator_count = 8\n",
@@ -819,9 +810,11 @@
        "# @markdown 2. High throughput: This profile optimizes for high throughput on large batches of incoming requests.\n",
        "profile = \"Low latency\"  # @param [\"Low latency\", \"High throughput\"] {isTemplate:true}\n",
        "\n",
+        "# Set this value to the expected number of concurrent requests.\n",
+        "torch_compile_max_bs = 4\n",
+        "\n",
        "if profile == \"Low latency\":\n",
        "    enable_torch_compile = True\n",
-        "    enable_flashinfer_mla = True\n",
        "    disable_cuda_graph = False\n",
        "    if base_model_name not in (\"DeepSeek-V3\", \"DeepSeek-V3-0324\", \"DeepSeek-R1\"):\n",
        "        speculative_algorithm = None\n",
@@ -837,7 +830,6 @@
        "    dp_size = 1\n",
        "else:\n",
        "    enable_torch_compile = False\n",
-        "    enable_flashinfer_mla = True\n",
        "    disable_cuda_graph = False\n",
        "    speculative_algorithm = None\n",
        "    speculative_draft_model_path = \"\"\n",
@@ -882,7 +874,7 @@
        "    model_id: str,\n",
        "    publisher: str,\n",
        "    publisher_model_id: str,\n",
-        "    service_account: str,\n",
+        "    service_account: str = \"\",\n",
        "    base_model_id: str = \"\",\n",
        "    machine_type: str = \"g2-standard-8\",\n",
        "    accelerator_type: str = \"NVIDIA_L4\",\n",
@@ -893,6 +885,8 @@
        "    dtype: str | None = None,\n",
        "    enable_trust_remote_code: bool = False,\n",
        "    enable_torch_compile: bool = False,\n",
+        "    torch_compile_max_bs: int | None = None,\n",
+        "    attention_backend: str = \"\",\n",
        "    enable_flashinfer_mla: bool = False,\n",
        "    disable_cuda_graph: bool = False,\n",
        "    speculative_algorithm: str | None = None,\n",
@@ -940,6 +934,11 @@
        "\n",
        "    if enable_torch_compile:\n",
        "        sglang_args.append(\"--enable-torch-compile\")\n",
+        "        if torch_compile_max_bs:\n",
+        "            sglang_args.append(f\"--torch-compile-max-bs={torch_compile_max_bs}\")\n",
+        "\n",
+        "    if attention_backend:\n",
+        "        sglang_args.append(f\"--attention-backend={attention_backend}\")\n",
        "\n",
        "    if enable_flashinfer_mla:\n",
        "        sglang_args.append(\"--enable-flashinfer-mla\")\n",
@@ -1017,13 +1016,14 @@
        "                \"minReplicaCount\": 1,\n",
        "                \"maxReplicaCount\": 1,\n",
        "            },\n",
-        "            \"serviceAccount\": service_account,\n",
        "            \"system_labels\": {\n",
        "                \"NOTEBOOK_NAME\": \"model_garden_pytorch_deepseek_deployment.ipynb\",\n",
        "                \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
        "            },\n",
        "        },\n",
        "    }\n",
+        "    if service_account:\n",
+        "        data[\"deployedModel\"][\"serviceAccount\"] = service_account\n",
        "    if is_spot:\n",
        "        data[\"deployedModel\"][\"dedicatedResources\"][\"spot\"] = True\n",
        "    response = requests.post(url, headers=headers, json=data)\n",
@@ -1041,7 +1041,6 @@
        "    model_id=model_id,\n",
        "    publisher=\"deepseek-ai\",\n",
        "    publisher_model_id=(\"deepseek-v3\" if \"V3\" in model_id else \"deepseek-r1\"),\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
        "    base_model_id=hf_model_id,\n",
        "    machine_type=machine_type,\n",
        "    accelerator_type=accelerator_type,\n",
@@ -1049,7 +1048,8 @@
        "    multihost_gpu_node_count=multihost_gpu_node_count,\n",
        "    enable_trust_remote_code=True,\n",
        "    enable_torch_compile=enable_torch_compile,\n",
-        "    enable_flashinfer_mla=enable_flashinfer_mla,\n",
+        "    torch_compile_max_bs=torch_compile_max_bs,\n",
+        "    attention_backend=\"fa3\",\n",
        "    disable_cuda_graph=disable_cuda_graph,\n",
        "    speculative_algorithm=speculative_algorithm,\n",
        "    speculative_draft_model_path=speculative_draft_model_path,\n",
@@ -1143,9 +1143,7 @@
        "\n",
        "if use_dedicated_endpoint:\n",
        "    DEDICATED_ENDPOINT_DNS = endpoints[\"sglang_gpu\"].gca_resource.dedicated_endpoint_dns\n",
-        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
-        "    PROJECT_ID, REGION, endpoints[\"sglang_gpu\"].name\n",
-        ")\n",
+        "ENDPOINT_RESOURCE_NAME = endpoints[\"sglang_gpu\"].resource_name\n",
        "\n",
        "# @markdown Because the DeepSeek-R1 / DeepSeek-V3 model generates detailed reasoning steps, the output is expected to be long. We recommend using streaming for a better generation experience.\n",
        "# @title Chat Completions Inference\n",
@@ -1208,6 +1206,293 @@
        "# @markdown Click \"Show Code\" to see more details."
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cyg_QzpQD0DJ"
+      },
+      "source": [
+        "## Deploy DeepSeek-V3 and DeepSeek-R1 with TensorRT-LLM"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "wfoBFRv_D4mB"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "\n",
+        "# @markdown This section uploads DeepSeek models to Model Registry and deploys them to a Vertex Prediction Endpoint. It takes ~1 hour to finish.\n",
+        "\n",
+        "# @markdown It's recommended to use the region selected by the deployment button on the model card. If the deployment button is not available, it's recommended to stay with the default region of the notebook.\n",
+        "\n",
+        "# @markdown Set the model to deploy.\n",
+        "\n",
+        "base_model_name = \"DeepSeek-R1\"  # @param [\"DeepSeek-V3\", \"DeepSeek-V3-Base\", \"DeepSeek-V3-0324\", \"DeepSeek-R1\"] {isTemplate:true}\n",
+        "model_id = \"deepseek-ai/\" + base_model_name\n",
+        "hf_model_id = model_id\n",
+        "\n",
+        "# The pre-built serving docker images.\n",
+        "TRTLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/tensorrt-llm.cu128.0-18.ubuntu2404.py312:deepseek\"\n",
+        "\n",
+        "# @markdown Choose whether to use a [Spot VM](https://cloud.google.com/compute/docs/instances/spot) for the deployment.\n",
+        "is_spot = False  # @param {type:\"boolean\"}\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",
+        "trtllm_accelerator_type = \"NVIDIA_H200_141GB\"  # @param [\"NVIDIA_H200_141GB\"] {isTemplate:true}\n",
+        "accelerator_count = 8\n",
+        "if trtllm_accelerator_type == \"NVIDIA_H200_141GB\":\n",
+        "    machine_type = \"a3-ultragpu-8g\"\n",
+        "    multihost_gpu_node_count = 1\n",
+        "    resource_id = \"custom_model_serving_nvidia_h200_gpus\"\n",
+        "else:\n",
+        "    raise ValueError(\"Only NVIDIA_H200_141GB is supported for DeepSeek-R1.\")\n",
+        "\n",
+        "check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    resource_id=resource_id,\n",
+        "    accelerator_count=int(accelerator_count * multihost_gpu_node_count),\n",
+        ")\n",
+        "\n",
+        "# 18K context length. This is the maximum supported by the current version of TensorRT-LLM on DeepSeek V3/R1 models.\n",
+        "MAX_INPUT_LEN = 18000\n",
+        "MAX_MODEL_LEN = 18000\n",
+        "MAX_NUM_SEQS = 128\n",
+        "GPU_MEMORY_UTILIZATION = 0.55\n",
+        "\n",
+        "\n",
+        "def poll_operation(op_name: str) -> bool:  # noqa: F811\n",
+        "    creds, _ = auth.default()\n",
+        "    auth_req = auth.transport.requests.Request()\n",
+        "    creds.refresh(auth_req)\n",
+        "    headers = {\n",
+        "        \"Authorization\": f\"Bearer {creds.token}\",\n",
+        "    }\n",
+        "    get_resp = requests.get(\n",
+        "        f\"https://{REGION}-aiplatform.googleapis.com/ui/{op_name}\",\n",
+        "        headers=headers,\n",
+        "    )\n",
+        "    opjs = get_resp.json()\n",
+        "    if \"error\" in opjs:\n",
+        "        raise ValueError(f\"Operation failed: {opjs['error']}\")\n",
+        "    return opjs.get(\"done\", False)\n",
+        "\n",
+        "\n",
+        "def poll_and_wait(op_name: str, total_wait: int, interval: int = 60):  # noqa: F811\n",
+        "    waited = 0\n",
+        "    while not poll_operation(op_name):\n",
+        "        if waited > total_wait:\n",
+        "            raise TimeoutError(\"Operation timed out\")\n",
+        "        print(\n",
+        "            f\"\\rStill waiting for operation... Waited time in second: {waited:<6}\",\n",
+        "            end=\"\",\n",
+        "            flush=True,\n",
+        "        )\n",
+        "        waited += interval\n",
+        "        time.sleep(interval)\n",
+        "\n",
+        "\n",
+        "def deploy_model_tensorrt_llm_multihost(\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 = \"\",\n",
+        "    machine_type: str = \"g2-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 1,\n",
+        "    multihost_gpu_node_count: int = 1,\n",
+        "    gpu_memory_utilization: float | None = None,\n",
+        "    max_input_len: int | None = None,\n",
+        "    max_model_len: int | None = None,\n",
+        "    max_num_seqs: int | None = None,\n",
+        "    enable_trust_remote_code: bool = False,\n",
+        "    enable_chunked_prefill: bool = False,\n",
+        "    use_dedicated_endpoint: bool = False,\n",
+        "    is_spot: bool = True,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys trained models with TensorRT-LLM on 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",
+        "    container_args = [\n",
+        "        \"python\",\n",
+        "        \"api_server.py\",\n",
+        "        \"--host=0.0.0.0\",\n",
+        "        \"--port=7080\",\n",
+        "        f\"--model={model_id}\",\n",
+        "        f\"--tensor-parallel-size={accelerator_count * multihost_gpu_node_count}\",\n",
+        "    ]\n",
+        "\n",
+        "    if gpu_memory_utilization:\n",
+        "        container_args.append(f\"--gpu-memory-utilization={gpu_memory_utilization}\")\n",
+        "\n",
+        "    if max_input_len:\n",
+        "        container_args.append(f\"--max-input-len={max_input_len}\")\n",
+        "\n",
+        "    if max_model_len:\n",
+        "        container_args.append(f\"--max-model-len={max_model_len}\")\n",
+        "\n",
+        "    if max_num_seqs:\n",
+        "        container_args.append(f\"--max-num-seqs={max_num_seqs}\")\n",
+        "\n",
+        "    if enable_trust_remote_code:\n",
+        "        container_args.append(\"--trust-remote-code=True\")\n",
+        "\n",
+        "    if enable_chunked_prefill:\n",
+        "        container_args.append(\"--enable-chunked-prefill=True\")\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=TRTLLM_DOCKER_URI,\n",
+        "        serving_container_args=container_args,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_predict_route=\"/v1/chat/completions\",\n",
+        "        serving_container_health_route=\"/health\",\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 {int(accelerator_count * multihost_gpu_node_count)} {accelerator_type} GPU(s).\"\n",
+        "    )\n",
+        "\n",
+        "    creds, _ = auth.default()\n",
+        "    auth_req = auth.transport.requests.Request()\n",
+        "    creds.refresh(auth_req)\n",
+        "\n",
+        "    url = f\"https://{REGION}-aiplatform.googleapis.com/ui/projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint.name}:deployModel\"\n",
+        "    headers = {\n",
+        "        \"Content-Type\": \"application/json\",\n",
+        "        \"Authorization\": f\"Bearer {creds.token}\",\n",
+        "    }\n",
+        "    data = {\n",
+        "        \"deployedModel\": {\n",
+        "            \"model\": model.resource_name,\n",
+        "            \"displayName\": model_name,\n",
+        "            \"dedicatedResources\": {\n",
+        "                \"machineSpec\": {\n",
+        "                    \"machineType\": machine_type,\n",
+        "                    \"multihostGpuNodeCount\": multihost_gpu_node_count,\n",
+        "                    \"acceleratorType\": accelerator_type,\n",
+        "                    \"acceleratorCount\": accelerator_count,\n",
+        "                },\n",
+        "                \"minReplicaCount\": 1,\n",
+        "                \"maxReplicaCount\": 1,\n",
+        "            },\n",
+        "            \"system_labels\": {\n",
+        "                \"NOTEBOOK_NAME\": \"model_garden_pytorch_deepseek_deployment.ipynb\",\n",
+        "                \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
+        "            },\n",
+        "        },\n",
+        "    }\n",
+        "    if service_account:\n",
+        "        data[\"deployedModel\"][\"serviceAccount\"] = service_account\n",
+        "    if is_spot:\n",
+        "        data[\"deployedModel\"][\"dedicatedResources\"][\"spot\"] = True\n",
+        "    response = requests.post(url, headers=headers, json=data)\n",
+        "    print(f\"Deploy Model response: {response.json()}\")\n",
+        "    if response.status_code != 200 or \"name\" not in response.json():\n",
+        "        raise ValueError(f\"Failed to deploy model: {response.text}\")\n",
+        "    poll_and_wait(response.json()[\"name\"], 7200)\n",
+        "    print(\"endpoint_name:\", endpoint.name)\n",
+        "\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "models[\"trtllm_gpu\"], endpoints[\"trtllm_gpu\"] = deploy_model_tensorrt_llm_multihost(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=\"deepseek-serve\"),\n",
+        "    model_id=model_id,\n",
+        "    publisher=\"deepseek-ai\",\n",
+        "    publisher_model_id=(\"deepseek-v3\" if \"V3\" in model_id else \"deepseek-r1\"),\n",
+        "    base_model_id=hf_model_id,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=trtllm_accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    multihost_gpu_node_count=multihost_gpu_node_count,\n",
+        "    gpu_memory_utilization=GPU_MEMORY_UTILIZATION,\n",
+        "    max_input_len=MAX_INPUT_LEN,\n",
+        "    max_model_len=MAX_MODEL_LEN,\n",
+        "    max_num_seqs=MAX_NUM_SEQS,\n",
+        "    enable_trust_remote_code=True,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    is_spot=is_spot,\n",
+        ")\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "sfnEWPRGFDGD"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Raw predict\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint using `:rawPredict`. The OpenAI Client chat completions support is coming soon.\n",
+        "endpoints[\"trtllm_gpu\"] = aiplatform.Endpoint(endpoints[\"trtllm_gpu\"].resource_name)\n",
+        "\n",
+        "# @markdown Fill out some request parameters:\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",
+        "\n",
+        "# @markdown Now we can send a request.\n",
+        "\n",
+        "response = endpoints[\"trtllm_gpu\"].raw_predict(\n",
+        "    body=json.dumps(\n",
+        "        {\n",
+        "            \"model\": \"\",\n",
+        "            \"messages\": [\n",
+        "                {\n",
+        "                    \"role\": \"user\",\n",
+        "                    \"content\": user_message,\n",
+        "                }\n",
+        "            ],\n",
+        "            \"max_tokens\": max_tokens,\n",
+        "            \"temperature\": temperature,\n",
+        "        }\n",
+        "    ),\n",
+        "    headers={\"Content-Type\": \"application/json\"},\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        ")\n",
+        "print(response.json()[\"choices\"][0][\"message\"][\"content\"])\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -0,0 +1,373 @@
+{
+  "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 - Dia-1.6B\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_pytorch_dia_1_6b.ipynb\">\n",
+        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_dia_1_6b.ipynb\">\n",
+        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3de7470326a2"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates deploying the pre-trained [Dia-1.6B](https://huggingface.co/nari-labs/Dia-1.6B) model on Vertex AI for online prediction.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Upload the model to [Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction).\n",
+        "- Deploy the model on [Endpoint](https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints).\n",
+        "- Run online predictions for text-to-speech.\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": [
+        "## Run the notebook"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "ioensNKM8ned"
+      },
+      "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.84.0'\n",
+        "\n",
+        "import importlib\n",
+        "import os\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",
+        "models, endpoints = {}, {}\n",
+        "LABEL = \"text-to-speech\"\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": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "2707b02ef5df"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Set the model parameters\n",
+        "\n",
+        "MODEL_ID = \"nari-labs/Dia-1.6B\"\n",
+        "TASK = \"text-to-speech\"\n",
+        "PUBLISHER_MODEL_NAME = \"publishers/nari-labs/models/dia-1.6b@dia-1.6b\"\n",
+        "\n",
+        "machine_type = \"g2-standard-12\"\n",
+        "accelerator_type = \"NVIDIA_L4\"\n",
+        "accelerator_count = 1"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "qZ5z-4hzRzL8"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 1] Deploy with Model Garden SDK\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 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.preview 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",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "lSD2g1pYYamO"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with customized configs\n",
+        "\n",
+        "# @markdown This section uploads the Dia-1.6B model to Model Registry and deploys it on the Endpoint with NVIDIA_L4 accelerator.\n",
+        "\n",
+        "# @markdown The deployment takes ~10 minutes to finish.\n",
+        "\n",
+        "# The pre-built serving docker image. It contains serving scripts and models.\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/pytorch-dia-serve.cu125.0-1.ubuntu2204.py310\"\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_id,\n",
+        "    task,\n",
+        "    machine_type,\n",
+        "    accelerator_type,\n",
+        "    accelerator_count,\n",
+        "    use_dedicated_endpoint,\n",
+        "):\n",
+        "    \"\"\"Create a Vertex AI Endpoint and deploy the specified model to the endpoint.\"\"\"\n",
+        "\n",
+        "    model_name = model_id\n",
+        "\n",
+        "    endpoint = aiplatform.Endpoint.create(\n",
+        "        display_name=f\"{model_name}-endpoint\",\n",
+        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
+        "    )\n",
+        "    serving_env = {\n",
+        "        \"MODEL_ID\": model_id,\n",
+        "        \"TASK\": task,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=SERVE_DOCKER_URI,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_predict_route=\"/predict\",\n",
+        "        serving_container_health_route=\"/health\",\n",
+        "        serving_container_environment_variables=serving_env,\n",
+        "        model_garden_source_model_name=PUBLISHER_MODEL_NAME,\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={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_dia_1_6b.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_id=MODEL_ID,\n",
+        "    task=TASK,\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:\", endpoints[LABEL].name)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "bb7adab99e41"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts.\n",
+        "\n",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown text: \"[S1] Dia is an open weights text to dialogue model. [S2] You get full control over scripts and voices. [S1] Wow. Amazing. (laughs)\"\n",
+        "# @markdown ```\n",
+        "\n",
+        "# @markdown You may adjust the parameters below to achieve best audio quality.\n",
+        "\n",
+        "import base64\n",
+        "\n",
+        "from IPython import display\n",
+        "\n",
+        "text = \"[S1] Dia is an open weights text to dialogue model. [S2] You get full control over scripts and voices. [S1] Wow. Amazing. (laughs)\"  # @param {type: \"string\"}\n",
+        "cfg_scale = 0.3  # @param {type:\"number\"}\n",
+        "temperature = 1.3  # @param {type:\"number\"}\n",
+        "top_p = 0.95  # @param {type:\"number\"}\n",
+        "\n",
+        "instances = [{\"text\": text}]\n",
+        "parameters = {\n",
+        "    \"temperature\": temperature,\n",
+        "    \"top_p\": top_p,\n",
+        "    \"cfg_scale\": cfg_scale,\n",
+        "}\n",
+        "\n",
+        "# The default num inference steps is set to 4 in the serving container, but\n",
+        "# you can change it to your own preference for image quality in the request.\n",
+        "response = endpoints[LABEL].predict(\n",
+        "    instances=instances,\n",
+        "    parameters=parameters,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        ")\n",
+        "base64_audio = response.predictions[0][\"audio\"]\n",
+        "display.Audio(base64.b64decode(base64_audio), rate=44100)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "6c460088b873"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Clean up resources\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_pytorch_dia_1_6b.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
@@ -29,28 +29,28 @@
        "id": "2bd716bf3e39"
      },
      "source": [
-        "# Vertex AI Model Garden - Dolly V2\n",
+        "# Vertex AI Model Garden - Dolly V2 [Deprecated]\n",
        "\n",
-        "<table align=\"left\">\n",
-        "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_dolly_v2.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
-        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_dolly_v2.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "\u003ctable align=\"left\"\u003e\n",
+        "  \u003ctd\u003e\n",
+        "    \u003ca href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_dolly_v2.ipynb\"\u003e\n",
+        "      \u003cimg src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"\u003e Run in Colab\n",
+        "    \u003c/a\u003e\n",
+        "  \u003c/td\u003e\n",
+        "  \u003ctd\u003e\n",
+        "    \u003ca href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_dolly_v2.ipynb\"\u003e\n",
+        "      \u003cimg src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"\u003e\n",
        "      View on GitHub\n",
-        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_pytorch_dolly_v2.ipynb\">\n",
-        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "    \u003c/a\u003e\n",
+        "  \u003c/td\u003e\n",
+        "  \u003ctd\u003e\n",
+        "    \u003ca href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_pytorch_dolly_v2.ipynb\"\u003e\n",
+        "      \u003cimg src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"\u003e\n",
        "Open in Vertex AI Workbench\n",
-        "    </a>\n",
+        "    \u003c/a\u003e\n",
        "    (a Python-3 CPU notebook is recommended)\n",
-        "  </td>\n",
-        "</table>"
+        "  \u003c/td\u003e\n",
+        "\u003c/table\u003e"
      ]
    },
    {
@@ -61,7 +61,7 @@
      "source": [
        "## Overview\n",
        "\n",
-        "This notebook demonstrates downloading and deploying the pre-trained [Dolly-V2-3b](https://huggingface.co/databricks/dolly-v2-3b), [Dolly-V2-7b](https://huggingface.co/databricks/dolly-v2-7b) & [Dolly-V2-12b](https://huggingface.co/databricks/dolly-v2-12b) model on Vertex AI for online prediction. It also demonstrates the evaluation of popular benchmark datasets through Vertex CustomJobs using [EleutherAI’s evaluation harness](https://github.com/EleutherAI/lm-evaluation-harness).\n",
+        "This notebook demonstrates downloading and deploying the pre-trained [Dolly-V2-3b](https://huggingface.co/databricks/dolly-v2-3b), [Dolly-V2-7b](https://huggingface.co/databricks/dolly-v2-7b) \u0026 [Dolly-V2-12b](https://huggingface.co/databricks/dolly-v2-12b) model on Vertex AI for online prediction. It also demonstrates the evaluation of popular benchmark datasets through Vertex CustomJobs using [EleutherAI’s evaluation harness](https://github.com/EleutherAI/lm-evaluation-harness).\n",
        "\n",
        "\n",
        "### Objective\n",
@@ -322,7 +322,7 @@
        "    return model, endpoint\n",
        "\n",
        "\n",
-        "def get_job_name_with_datetime(prefix: str) -> str:\n",
+        "def get_job_name_with_datetime(prefix: str) -\u003e str:\n",
        "    \"\"\"Gets the job name with date time when triggering training or deployment\n",
        "    jobs in Vertex AI.\n",
        "    \"\"\"\n",
@@ -463,7 +463,7 @@
        "if \"12b\" in MODEL_ID:\n",
        "    accelerator_count = 2  # for dolly-v2-12b\n",
        "else:\n",
-        "    accelerator_count = 1  # for dolly-v2-3b & dolly-v2-7b\n",
+        "    accelerator_count = 1  # for dolly-v2-3b \u0026 dolly-v2-7b\n",
        "\n",
        "replica_count = 1"
      ]
@@ -185,10 +185,7 @@
        "elif accelerator_type == \"NVIDIA_H100_80GB\":\n",
        "    machine_type = \"a3-highgpu-2g\"\n",
        "else:\n",
-        "    raise ValueError(f\"Unsupported accelerator type: {accelerator_type}\")\n",
-        "\n",
-        "# Dedicated endpoint is not supported.\n",
-        "use_dedicated_endpoint = False"
+        "    raise ValueError(f\"Unsupported accelerator type: {accelerator_type}\")"
      ]
    },
    {
@@ -202,6 +199,8 @@
      "source": [
        "# @title [Option 1] Deploy with Model Garden SDK\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",
        "# @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.preview import model_garden\n",
        "\n",
@@ -230,11 +229,21 @@
        "\n",
        "# @markdown The deployment takes ~15 minutes to finish.\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",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/xdit-serve.cu125.0-1.ubuntu2204.py310\"\n",
        "\n",
        "\n",
-        "def deploy_model(model_id, task, machine_type, accelerator_type, accelerator_count):\n",
+        "def deploy_model(\n",
+        "    model_id,\n",
+        "    task,\n",
+        "    machine_type,\n",
+        "    accelerator_type,\n",
+        "    accelerator_count,\n",
+        "    use_dedicated_endpoint,\n",
+        "):\n",
        "    \"\"\"Create a Vertex AI Endpoint and deploy the specified model to the endpoint.\"\"\"\n",
        "    common_util.check_quota(\n",
        "        project_id=PROJECT_ID,\n",
@@ -246,7 +255,10 @@
        "\n",
        "    model_name = model_id\n",
        "\n",
-        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
+        "    endpoint = aiplatform.Endpoint.create(\n",
+        "        display_name=f\"{model_name}-endpoint\",\n",
+        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
+        "    )\n",
        "    serving_env = {\n",
        "        \"MODEL_ID\": model_id,\n",
        "        \"TASK\": task,\n",
@@ -289,6 +301,7 @@
        "    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:\", endpoints[LABEL].name)"
@@ -329,7 +342,11 @@
        "\n",
        "# The default num inference steps is set to 4 in the serving container, but\n",
        "# you can change it to your own preference for image quality in the request.\n",
-        "response = endpoints[LABEL].predict(instances=instances, parameters=parameters)\n",
+        "response = endpoints[LABEL].predict(\n",
+        "    instances=instances,\n",
+        "    parameters=parameters,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        ")\n",
        "images = [\n",
        "    common_util.base64_to_image(prediction.get(\"output\"))\n",
        "    for prediction in response.predictions\n",
@@ -656,8 +656,8 @@
        "serve_accelerator_type = \"NVIDIA_L4\"\n",
        "serve_accelerator_count = 1\n",
        "\n",
-        "# @markdown Set use_dedicated_endpoint to True if the endpoint is [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint) enabled.\n",
-        "use_dedicated_endpoint = False  # @param {type:\"boolean\"}\n",
+        "# @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",
        "common_util.check_quota(\n",
        "    project_id=PROJECT_ID,\n",
@@ -0,0 +1,366 @@
+{
+  "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 - HiDream-I1\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_pytorch_hidream_i1.ipynb\">\n",
+        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_hidream_i1.ipynb\">\n",
+        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3de7470326a2"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates deploying the pre-trained [HiDream-I1](https://huggingface.co/collections/HiDream-ai/hidream-i1-67f3e90dd509fed088a158b3) models on Vertex AI for online prediction.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Upload the model to [Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction).\n",
+        "- Deploy the model on [Endpoint](https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints).\n",
+        "- Run online predictions for text-to-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",
+        "\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": [
+        "## Run the notebook"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "ioensNKM8ned"
+      },
+      "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.84.0'\n",
+        "\n",
+        "import importlib\n",
+        "import os\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",
+        "models, endpoints = {}, {}\n",
+        "LABEL = \"text-to-image-hidream\"\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": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "2707b02ef5df"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Set the model parameters\n",
+        "\n",
+        "MODEL_ID = \"HiDream-ai/HiDream-I1-Full\"  # @param [\"HiDream-ai/HiDream-I1-Full\", \"HiDream-ai/HiDream-I1-Dev\", \"HiDream-ai/HiDream-I1-Fast\"]\n",
+        "TASK = \"text-to-image-hidream\"\n",
+        "\n",
+        "model_version = MODEL_ID.split(\"/\")[-1].lower()\n",
+        "PUBLISHER_MODEL_NAME = f\"publishers/hidream-i1/models/hidream-i1-full@{model_version}\"\n",
+        "\n",
+        "ACCELERATOR_TYPE = \"NVIDIA_A100_80GB\"  # @param [\"NVIDIA_A100_80GB\", \"NVIDIA_H100_80GB\"]\n",
+        "\n",
+        "if ACCELERATOR_TYPE == \"NVIDIA_A100_80GB\":\n",
+        "    machine_type = \"a2-ultragpu-1g\"\n",
+        "    accelerator_count = 1\n",
+        "elif ACCELERATOR_TYPE == \"NVIDIA_H100_80GB\":\n",
+        "    machine_type = \"a3-highgpu-2g\"\n",
+        "    accelerator_count = 2\n",
+        "else:\n",
+        "    raise ValueError(f\"Unsupported accelerator type: {ACCELERATOR_TYPE}\")\n",
+        "accelerator_type = ACCELERATOR_TYPE"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "lSD2g1pYYamO"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy to Vertex AI\n",
+        "\n",
+        "# @markdown This section uploads the HiDream-I1 model to Model Registry and deploys it on the Endpoint with selected accelerator type.\n",
+        "\n",
+        "# @markdown The deployment takes ~25 minutes to finish.\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",
+        "# The pre-built serving docker image. It contains serving scripts and models.\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/pytorch-inference.cu125.0-4.ubuntu2204.py310\"\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_id,\n",
+        "    task,\n",
+        "    machine_type,\n",
+        "    accelerator_type,\n",
+        "    accelerator_count,\n",
+        "    use_dedicated_endpoint,\n",
+        "):\n",
+        "    \"\"\"Create a Vertex AI Endpoint and deploy the specified model to the endpoint.\"\"\"\n",
+        "\n",
+        "    model_name = model_id\n",
+        "\n",
+        "    endpoint = aiplatform.Endpoint.create(\n",
+        "        display_name=f\"{model_name}-endpoint\",\n",
+        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
+        "    )\n",
+        "    serving_env = {\n",
+        "        \"MODEL_ID\": model_id,\n",
+        "        \"TASK\": task,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=SERVE_DOCKER_URI,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_predict_route=\"/predict\",\n",
+        "        serving_container_health_route=\"/health\",\n",
+        "        serving_container_environment_variables=serving_env,\n",
+        "        model_garden_source_model_name=PUBLISHER_MODEL_NAME,\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={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_hidream_i1.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_id=MODEL_ID,\n",
+        "    task=TASK,\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:\", endpoints[LABEL].name)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "bb7adab99e41"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts.\n",
+        "\n",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown text: \"A cat holding a sign that says hello world\"\n",
+        "# @markdown ```\n",
+        "\n",
+        "# @markdown Recommended parameters:\n",
+        "# @markdown - HiDream-I1-Full: num_inference_steps=50, guidance_scale=5.0\n",
+        "# @markdown - HiDream-I1-Dev: num_inference_steps=28, guidance_scale=0.0\n",
+        "# @markdown - HiDream-I1-Fast: num_inference_steps=16, guidance_scale=0.0\n",
+        "\n",
+        "# @markdown You may adjust the parameters below to achieve best image quality.\n",
+        "\n",
+        "text = \"A cat holding a sign that says hello world\"  # @param {type: \"string\"}\n",
+        "height = 1024  # @param {type:\"number\"}\n",
+        "width = 1024  # @param {type:\"number\"}\n",
+        "num_inference_steps = 50  # @param {type:\"number\"}\n",
+        "guidance_scale = 5.0  # @param {type:\"number\"}\n",
+        "\n",
+        "instances = [{\"text\": text}]\n",
+        "parameters = {\n",
+        "    \"height\": height,\n",
+        "    \"width\": width,\n",
+        "    \"num_inference_steps\": num_inference_steps,\n",
+        "    \"guidance_scale\": guidance_scale,\n",
+        "}\n",
+        "\n",
+        "# The default num inference steps is set to 4 in the serving container, but\n",
+        "# you can change it to your own preference for image quality in the request.\n",
+        "response = endpoints[LABEL].predict(\n",
+        "    instances=instances,\n",
+        "    parameters=parameters,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        ")\n",
+        "images = [\n",
+        "    common_util.base64_to_image(prediction.get(\"output\"))\n",
+        "    for prediction in response.predictions\n",
+        "]\n",
+        "common_util.image_grid(images, rows=1)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "6c460088b873"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Clean up resources\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_pytorch_hidream_i1.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
@@ -100,6 +100,9 @@
      "source": [
        "# @title Setup Google Cloud project\n",
        "\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.84.0'\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",
@@ -116,9 +119,6 @@
        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
        "\n",
        "# Import the necessary packages\n",
-        "\n",
-        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.84.0'\n",
-        "\n",
        "import importlib\n",
        "import os\n",
        "from typing import Tuple\n",
@@ -133,6 +133,7 @@
        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
        ")\n",
        "\n",
+        "LABEL = \"instant_id_gpu\"\n",
        "models, endpoints = {}, {}\n",
        "LABEL = \"vllm_gpu\"\n",
        "\n",
@@ -155,8 +156,6 @@
        "    project=PROJECT_ID,\n",
        "    location=REGION,\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 Click \"Show code\" to see more details."
      ]
@@ -184,15 +183,24 @@
        "model_id = \"instantx/instantid\"\n",
        "task = \"instant-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",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "accelerator_type = \"NVIDIA_L4\"\n",
        "accelerator_count = 1\n",
-        "machine_type = \"g2-standard-8\"\n",
+        "machine_type = \"g2-standard-12\"\n",
        "\n",
        "\n",
-        "def deploy_model(model_id, task, accelerator_type, machine_type, accelerator_count=1):\n",
+        "def deploy_model(\n",
+        "    model_id: str,\n",
+        "    task: str,\n",
+        "    accelerator_type: str,\n",
+        "    machine_type: str,\n",
+        "    accelerator_count: int,\n",
+        "):\n",
        "    \"\"\"Create a Vertex AI Endpoint and deploy the specified model to the endpoint.\"\"\"\n",
        "    model_name = model_id\n",
        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-{task}-endpoint\")\n",
@@ -218,7 +226,10 @@
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=accelerator_count,\n",
        "        deploy_request_timeout=1800,\n",
-        "        system_labels={\"NOTEBOOK_NAME\": \"model_garden_pytorch_instant_id.ipynb\"},\n",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_instant_id.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
+        "        },\n",
        "    )\n",
        "    print(\"To load this existing endpoint from a different session:\")\n",
        "    print(\n",
@@ -246,7 +257,7 @@
        "    accelerator_type=accelerator_type,\n",
        "    accelerator_count=accelerator_count,\n",
        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
-        "    accept_eula=False,  # Change to True if you have accepted the EULA on the model card.\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",
        "# @title [Option 2] Deploy with customized configs\n",
@@ -278,6 +289,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",
@@ -330,6 +342,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",
@@ -375,13 +391,14 @@
        "    return model, endpoint\n",
        "\n",
        "\n",
-        "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model(\n",
+        "models[LABEL], endpoints[LABEL] = deploy_model(\n",
        "    model_id=model_id,\n",
        "    task=task,\n",
        "    accelerator_type=accelerator_type,\n",
        "    machine_type=machine_type,\n",
+        "    accelerator_count=accelerator_count,\n",
        ")\n",
-        "print(\"endpoint_name:\", endpoints[\"vllm_gpu\"].name)\n",
+        "print(\"endpoint_name:\", endpoints[LABEL].name)\n",
        "\n",
        "# @markdown Click \"Show code\" to see more details."
      ]
@@ -445,7 +462,7 @@
        "    },\n",
        "]\n",
        "\n",
-        "response = endpoints[\"vllm_gpu\"].predict(instances=instances)\n",
+        "response = endpoints[LABEL].predict(instances=instances)\n",
        "images = [common_util.base64_to_image(image) for image in response.predictions]\n",
        "display(common_util.image_grid(images, rows=1))\n",
        "\n",
@@ -491,7 +508,7 @@
        "    },\n",
        "]\n",
        "\n",
-        "response = endpoints[\"vllm_gpu\"].predict(instances=instances)\n",
+        "response = endpoints[LABEL].predict(instances=instances)\n",
        "images = [common_util.base64_to_image(image) for image in response.predictions]\n",
        "display(common_util.image_grid(images, rows=1))\n",
        "\n",
@@ -178,7 +178,10 @@
        "# The machine and accelerator specs for model deployment.\n",
        "accelerator_type = \"NVIDIA_L4\"\n",
        "machine_type = \"g2-standard-12\"\n",
-        "accelerator_count = 1"
+        "accelerator_count = 1\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\"}"
      ]
    },
    {
@@ -200,6 +203,7 @@
        "    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",
        ")"
      ]
@@ -226,9 +230,13 @@
        "    machine_type: str,\n",
        "    accelerator_type: str,\n",
        "    accelerator_count: int,\n",
+        "    use_dedicated_endpoint: bool = False,\n",
        "):\n",
        "    model_name = \"instruct-pix2pix\"\n",
-        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
+        "    endpoint = aiplatform.Endpoint.create(\n",
+        "        display_name=f\"{model_name}-endpoint\",\n",
+        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
+        "    )\n",
        "    serving_env = {\n",
        "        \"MODEL_ID\": model_id,\n",
        "        \"TASK\": task,\n",
@@ -272,6 +280,7 @@
        "    machine_type=machine_type,\n",
        "    accelerator_type=accelerator_type,\n",
        "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
        ")"
      ]
    },
@@ -302,7 +311,9 @@
        "        \"num_inference_steps\": num_inference_steps,\n",
        "    },\n",
        "]\n",
-        "response = endpoints[LABEL].predict(instances=instances)\n",
+        "response = endpoints[LABEL].predict(\n",
+        "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
+        ")\n",
        "images = [common_util.base64_to_image(image) for image in response.predictions]\n",
        "common_util.image_grid([init_image, images[0]], rows=1, cols=2)"
      ]
@@ -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,6 +34,11 @@
        "\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_pytorch_llama2_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",
@@ -61,6 +67,11 @@
        "- Download and deploy prebuilt LLaMA2 models\n",
        "- Deploy LLaMA2 with [vLLM](https://github.com/vllm-project/vllm) to improve serving throughput\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",
@@ -93,85 +104,63 @@
        "\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 specified region (`REGION`). 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",
-        "# @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",
+        "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, 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.84.0'\n",
        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
        "\n",
        "# Import the necessary packages\n",
        "import importlib\n",
        "import os\n",
-        "import uuid\n",
-        "from datetime import datetime\n",
+        "import re\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",
        "# Get the default region for launching jobs.\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 below.\n",
-        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
-        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\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",
-        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
-        "else:\n",
-        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\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",
-        "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, \"llama2\")\n",
-        "\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, 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",
-        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\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",
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")\n",
        "\n",
        "# @markdown # Access LLaMA2 models on Vertex AI for GPU based serving\n",
        "# @markdown The original models from Meta are converted into the Hugging Face format for serving in Vertex AI.\n",
@@ -182,102 +171,18 @@
        "# @markdown 4. Paste the URI in the `VERTEX_AI_MODEL_GARDEN_LLAMA2` field below.\n",
        "# @markdown 5. The LLaMA2 models will be copied into `BUCKET_URI`.\n",
        "\n",
-        "\n",
        "VERTEX_AI_MODEL_GARDEN_LLAMA2 = \"\"  # @param {type:\"string\", isTemplate:true}\n",
        "assert (\n",
        "    VERTEX_AI_MODEL_GARDEN_LLAMA2\n",
-        "), \"Model artifact path is required. Click the agreement of LLaMA2 in Vertex AI Model Garden, and get the GCS path of LLaMA2 model artifacts.\"\n",
-        "print(\n",
-        "    \"Copying LLaMA2 model artifacts from\",\n",
-        "    VERTEX_AI_MODEL_GARDEN_LLAMA2,\n",
-        "    \"to \",\n",
-        "    MODEL_BUCKET,\n",
-        ")\n",
+        "), \"Click the agreement of Llama2 in Vertex AI Model Garden, and get the GCS path of Llama 3.2 model artifacts.\"\n",
+        "parsed_gcs_url = re.search(\"gs://.*?(?=[ ]|$)\", VERTEX_AI_MODEL_GARDEN_LLAMA2)\n",
+        "if parsed_gcs_url:\n",
+        "    VERTEX_AI_MODEL_GARDEN_LLAMA2 = parsed_gcs_url.group()\n",
+        "assert VERTEX_AI_MODEL_GARDEN_LLAMA2.startswith(\n",
+        "    \"gs://\"\n",
+        "), \"VERTEX_AI_MODEL_GARDEN_LLAMA2 is expected to be a GCS URI and must start with `gs://`.\"\n",
        "\n",
-        "! gsutil -m cp -R $VERTEX_AI_MODEL_GARDEN_LLAMA2/* $MODEL_BUCKET\n",
-        "base_model_path_prefix = MODEL_BUCKET\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:20240222_0916_RC00\"\n",
-        "\n",
-        "\n",
-        "def get_job_name_with_datetime(prefix: str) -> str:\n",
-        "    \"\"\"Gets the job name with date time when triggering training or deployment\n",
-        "    jobs in Vertex AI.\n",
-        "    \"\"\"\n",
-        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
-        "\n",
-        "\n",
-        "def deploy_model(\n",
-        "    model_name: str,\n",
-        "    model_id: str,\n",
-        "    base_model_id: str,\n",
-        "    service_account: str,\n",
-        "    machine_type: str = \"g2-standard-8\",\n",
-        "    accelerator_type: str = \"NVIDIA_L4\",\n",
-        "    accelerator_count: int = 1,\n",
-        "    max_model_len: int = 4096,\n",
-        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
-        "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\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",
-        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
-        "\n",
-        "    vllm_args = [\n",
-        "        \"--host=0.0.0.0\",\n",
-        "        \"--port=7080\",\n",
-        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
-        "        \"--swap-space=16\",\n",
-        "        \"--gpu-memory-utilization=0.95\",\n",
-        "        f\"--max-model-len={max_model_len}\",\n",
-        "        \"--max-num-batched-tokens=4096\",\n",
-        "        \"--disable-log-stats\",\n",
-        "    ]\n",
-        "\n",
-        "    env_vars = {\n",
-        "        \"MODEL_ID\": base_model_id,\n",
-        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
-        "    }\n",
-        "    model = aiplatform.Model.upload(\n",
-        "        display_name=model_name,\n",
-        "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
-        "        serving_container_command=[\"python\", \"-m\", \"vllm.entrypoints.api_server\"],\n",
-        "        serving_container_args=vllm_args,\n",
-        "        serving_container_ports=[7080],\n",
-        "        serving_container_predict_route=\"/generate\",\n",
-        "        serving_container_health_route=\"/ping\",\n",
-        "        serving_container_environment_variables=env_vars,\n",
-        "        artifact_uri=model_id,\n",
-        "        model_garden_source_model_name=\"publishers/meta/models/llama2\"\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_pytorch_llama2_deployment.ipynb\"\n",
-        "        },\n",
-        "    )\n",
-        "    print(\"endpoint_name:\", endpoint.name)\n",
-        "\n",
-        "    print(\"To load this existing endpoint from a different session:\")\n",
-        "    print(\"from google.cloud import aiplatform\")\n",
-        "    print(\n",
-        "        f'endpoint = aiplatform.Endpoint(\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint.name}\")'\n",
-        "    )\n",
-        "    return model, endpoint"
+        "# @markdown Click \"Show Code\" to see more details."
      ]
    },
    {
@@ -294,39 +199,26 @@
      "execution_count": null,
      "metadata": {
        "cellView": "form",
-        "id": "E8OiHHNNE_wj"
+        "id": "CwlOMvd2lDCH"
      },
      "outputs": [],
      "source": [
-        "# @title Deploy\n",
+        "# @title Select the model variants\n",
        "\n",
-        "# @markdown This section uploads prebuilt LLaMA2 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",
-        "\n",
-        "# @markdown V100 GPUs are used for demonstration. Note that V100 serving generally offers better throughput and latency performance than L4 serving, while L4 serving is generally more cost efficient than V100 serving. The serving efficiency of V100 and L4 GPUs is inferior to that of A100 GPUs, but V100 and L4 GPUs are nevertheless good serving solutions if you do not have A100 quota.\n",
-        "\n",
-        "# @markdown Set the model to deploy.\n",
-        "\n",
-        "base_model_name = \"llama2-7b-chat-hf\"  # @param [\"llama2-7b-hf\", \"llama2-7b-chat-hf\", \"llama2-13b-hf\", \"llama2-13b-chat-hf\", \"llama2-70b-hf\", \"llama2-70b-chat-hf\"] {isTemplate:true}\n",
-        "model_id = os.path.join(MODEL_BUCKET, base_model_name)\n",
-        "if base_model_name == \"llama2-7b-chat-hf\":\n",
-        "    hf_model_id = \"meta-llama/Llama-2-7b-chat-hf\"\n",
-        "elif base_model_name == \"llama2-7b-hf\":\n",
-        "    hf_model_id = \"meta-llama/Llama-2-7b-hf\"\n",
-        "elif base_model_name == \"llama2-13b-hf\":\n",
-        "    hf_model_id = \"meta-llama/Llama-2-13b-hf\"\n",
-        "elif base_model_name == \"llama2-13b-chat-hf\":\n",
-        "    hf_model_id = \"meta-llama/Llama-2-13b-chat-hf\"\n",
-        "elif base_model_name == \"llama2-70b-hf\":\n",
-        "    hf_model_id = \"meta-llama/Llama-2-70b-hf\"\n",
-        "elif base_model_name == \"llama2-70b-chat-hf\":\n",
-        "    hf_model_id = \"meta-llama/Llama-2-70b-chat-hf\"\n",
-        "else:\n",
-        "    raise ValueError(f\"Unsupported base model name: {base_model_name}\")\n",
+        "base_model_name = \"llama2-7b-hf\"  # @param [\"llama2-7b-hf\", \"llama2-7b-chat-hf\", \"llama2-13b-hf\", \"llama2-13b-chat-hf\", \"llama2-70b-hf\", \"llama2-70b-chat-hf\"] {isTemplate:true}\n",
+        "model_id = os.path.join(VERTEX_AI_MODEL_GARDEN_LLAMA2, base_model_name)\n",
+        "version_id = \"llama-2-\" + base_model_name.split(\"-\")[1]\n",
+        "PUBLISHER_MODEL_NAME = f\"publishers/meta/models/llama2@{version_id}\"\n",
+        "hf_model_id = \"meta-llama/Llama-2-\" + base_model_name.split(\"-\", 1)[1]\n",
        "\n",
        "# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
        "\n",
-        "accelerator_type = \"NVIDIA_TESLA_V100\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_V100\", \"NVIDIA_TESLA_A100\", \"NVIDIA_H100_80GB\"]\n",
+        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_T4\", \"NVIDIA_TESLA_V100\", \"NVIDIA_TESLA_A100\", \"4 NVIDIA_H100_80GB\", \"8 NVIDIA_H100_80GB\"]\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",
+        "# Find Vertex AI prediction supported accelerators and regions in\n",
+        "# https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
        "\n",
        "# @markdown If A100 is not available, you may serve LLaMA2 13B models with multiple V100s\n",
        "# @markdown or L4s. Keep in mind that the efficiency of serving with multiple\n",
@@ -340,27 +232,32 @@
        "        accelerator_count = 1\n",
        "    elif accelerator_type == \"NVIDIA_L4\":\n",
        "        # L4 serving is more cost efficient than V100 serving.\n",
-        "        machine_type = \"g2-standard-8\"\n",
+        "        machine_type = \"g2-standard-12\"\n",
+        "        accelerator_count = 1\n",
+        "    elif accelerator_type == \"NVIDIA_TESLA_T4\":\n",
+        "        machine_type = \"n1-standard-16\"\n",
+        "        accelerator_count = 1\n",
+        "    elif accelerator_type == \"NVIDIA_TESLA_A100\":\n",
+        "        machine_type = \"a2-highgpu-1g\"\n",
+        "        accelerator_count = 1\n",
+        "    elif accelerator_type == \"NVIDIA_A100_80GB\":\n",
+        "        machine_type = \"a2-ultragpu-1g\"\n",
        "        accelerator_count = 1\n",
        "    else:\n",
        "        raise ValueError(\n",
        "            f\"Recommended GPU setting not found for: {accelerator_type} and {base_model_name}.\"\n",
        "        )\n",
        "elif \"13b\" in base_model_name:\n",
-        "    if accelerator_type == \"NVIDIA_TESLA_V100\":\n",
-        "        # V100 serving has better throughput and latency performance than L4 serving.\n",
-        "        machine_type = \"n1-standard-8\"\n",
-        "        accelerator_count = 2\n",
-        "    elif accelerator_type == \"NVIDIA_L4\":\n",
-        "        # L4 serving is more cost efficient than V100 serving.\n",
+        "    if accelerator_type == \"NVIDIA_L4\":\n",
+        "        # L4 serving is more cost efficient than A100 serving.\n",
        "        machine_type = \"g2-standard-24\"\n",
        "        accelerator_count = 2\n",
-        "    elif accelerator_type == \"NVIDIA_TESLA_A100\":\n",
-        "        machine_type = \"a2-highgpu-1g\"\n",
+        "    elif accelerator_type == \"NVIDIA_A100_80GB\":\n",
+        "        machine_type = \"a2-ultragpu-1g\"\n",
        "        accelerator_count = 1\n",
        "    else:\n",
        "        raise ValueError(\n",
-        "            f\"Recommended GPU setting not found for: {accelerator_type} and {base_model_name}.\"\n",
+        "            f\"{base_model_name} supports only NVIDIA_L4 and NVIDIA_A100_80GB\"\n",
        "        )\n",
        "elif \"70b\" in base_model_name:\n",
        "    # If you do not have access to 4 A100 (40G) GPUs, you may serve LLaMA 2 70B\n",
@@ -372,32 +269,227 @@
        "    if accelerator_type == \"NVIDIA_L4\":\n",
        "        machine_type = \"g2-standard-96\"\n",
        "        accelerator_count = 8\n",
-        "    elif accelerator_type == \"NVIDIA_TESLA_A100\":\n",
-        "        machine_type = \"a2-highgpu-4g\"\n",
-        "        accelerator_count = 4\n",
-        "    elif accelerator_type == \"NVIDIA_H100_80GB\":\n",
+        "    elif accelerator_type == \"4 NVIDIA_H100_80GB\":\n",
+        "        accelerator_type = \"NVIDIA_H100_80GB\"\n",
        "        machine_type = \"a3-highgpu-4g\"\n",
        "        accelerator_count = 4\n",
+        "    elif accelerator_type == \"8 NVIDIA_H100_80GB\":\n",
+        "        accelerator_type = \"NVIDIA_H100_80GB\"\n",
+        "        machine_type = \"a3-highgpu-8g\"\n",
+        "        accelerator_count = 8\n",
        "    else:\n",
        "        raise ValueError(\n",
-        "            f\"Recommended GPU setting not found for: {accelerator_type} and {base_model_name}.\"\n",
+        "            f\"{base_model_name} supports only NVIDIA_L4 and NVIDIA_H100_80GB\"\n",
        "        )\n",
        "else:\n",
        "    raise ValueError(\n",
        "        f\"Recommended GPU setting not found for: {accelerator_type} and {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",
+        ")\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "B9cqWafKlU5U"
+      },
+      "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.preview 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",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "E8OiHHNNE_wj"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with customized configs\n",
+        "\n",
+        "# @markdown This section uploads prebuilt LLaMA2 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",
+        "# @markdown V100 GPUs are used for demonstration. Note that V100 serving generally offers better throughput and latency performance than L4 serving, while L4 serving is generally more cost efficient than V100 serving. The serving efficiency of V100 and L4 GPUs is inferior to that of A100 GPUs, but V100 and L4 GPUs are nevertheless good serving solutions if you do not have A100 quota.\n",
+        "\n",
+        "# @markdown Set the model to deploy.\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:20240715_0916_RC00\"\n",
+        "\n",
        "# Note that a larger max_model_len will require more GPU memory.\n",
-        "max_model_len = 2048\n",
-        "model, endpoint = deploy_model(\n",
-        "    model_name=get_job_name_with_datetime(prefix=\"llama2-serve\"),\n",
+        "if accelerator_type in [\"NVIDIA_TESLA_T4\", \"NVIDIA_TESLA_V100\"]:\n",
+        "    max_model_len = 1024\n",
+        "elif accelerator_type in [\"NVIDIA_L4\"]:\n",
+        "    max_model_len = 2048\n",
+        "else:\n",
+        "    max_model_len = 4096\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_pytorch_llama2_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[LABEL], endpoints[LABEL] = deploy_model_vllm(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=\"llama2-serve\"),\n",
        "    model_id=model_id,\n",
+        "    publisher=\"meta\",\n",
+        "    publisher_model_id=\"llama2\",\n",
        "    base_model_id=hf_model_id,\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
        "    machine_type=machine_type,\n",
        "    accelerator_type=accelerator_type,\n",
        "    accelerator_count=accelerator_count,\n",
        "    max_model_len=max_model_len,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
        ")\n",
+        "\n",
        "# @markdown Click \"Show Code\" to see more details."
      ]
    },
@@ -412,7 +504,7 @@
      "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://github.com/vllm-project/vllm/blob/2e8e49fce3775e7704d413b2f02da6d7c99525c9/vllm/sampling_params.py#L23-L64).\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",
@@ -422,16 +514,30 @@
        "# @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.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",
-        "# Overides parameters for inferences.\n",
-        "# If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`,\n",
-        "# you can reduce the max length, such as set max_tokens as 20.\n",
+        "# Overrides parameters for inferences.\n",
        "instances = [\n",
        "    {\n",
        "        \"prompt\": prompt,\n",
@@ -442,12 +548,12 @@
        "        \"raw_response\": raw_response,\n",
        "    },\n",
        "]\n",
-        "response = endpoint.predict(instances=instances)\n",
+        "response = endpoints[LABEL].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."
+        "    print(prediction)"
      ]
    },
    {
@@ -461,17 +567,15 @@
      "source": [
        "# @title Clean up resources\n",
        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
-        "# @markdown  and avoid unnecessary continouous charges that may incur.\n",
+        "# @markdown  and avoid unnecessary continuous charges that may incur.\n",
        "\n",
        "# Undeploy model and delete endpoint.\n",
-        "endpoint.delete(force=True)\n",
+        "for endpoint in endpoints.values():\n",
+        "    endpoint.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
-        "model.delete()\n",
-        "\n",
-        "delete_bucket = False  # @param {type:\"boolean\"}\n",
-        "if delete_bucket:\n",
-        "    ! gsutil -m rm -r $BUCKET_URI"
+        "for model in models.values():\n",
+        "    model.delete()"
      ]
    }
  ],
@@ -380,9 +380,7 @@
        "models[\"vllm_fast\"], endpoints[\"vllm_fast\"] = fast_deploy(\n",
        "    \"meta\", \"llama3_1\", \"llama-3.1-8b-instruct\"\n",
        ")\n",
-        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
-        "    PROJECT_ID, REGION, endpoints[\"vllm_fast\"].name\n",
-        ")\n",
+        "ENDPOINT_RESOURCE_NAME = endpoints[\"vllm_fast\"].resource_name\n",
        "BASE_URL = (\n",
        "    f\"https://{REGION}-aiplatform.googleapis.com/v1beta1/{ENDPOINT_RESOURCE_NAME}\"\n",
        ")\n",
@@ -423,7 +421,7 @@
        "# The pre-built serving docker images.\n",
        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241210_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",
+        "# @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",
        "# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
        "if \"8b\" in base_model_name.lower():\n",
@@ -111,7 +111,6 @@
    {
      "cell_type": "code",
      "execution_count": null,
-      "language": "python",
      "metadata": {
        "cellView": "form",
        "id": "YXFGIp1l-qtT"
@@ -778,10 +777,8 @@
        "REGION = FAST_DEPLOYMENT_REGION\n",
        "\n",
        "if use_dedicated_endpoint:\n",
-        "DEDICATED_ENDPOINT_DNS = endpoints[\"vllm_fast\"].gca_resource.dedicated_endpoint_dns\n",
-        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
-        "    PROJECT_ID, REGION, endpoints[\"vllm_fast\"].name\n",
-        ")\n",
+        "    DEDICATED_ENDPOINT_DNS = endpoints[\"vllm_fast\"].gca_resource.dedicated_endpoint_dns\n",
+        "    ENDPOINT_RESOURCE_NAME = endpoints[\"vllm_fast\"].resource_name\n",
        "\n",
        "# @title Chat Completions Inference\n",
        "\n",
@@ -1070,8 +1067,12 @@
        "        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",
+        "        if \"Llama-4\" not in model_id:\n",
+        "            vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "            vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "        else:\n",
+        "            vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "            vllm_args.append(\"--tool-call-parser=llama3_json\")\n",
        "\n",
        "    env_vars = {\n",
        "        \"MODEL_ID\": base_model_id,\n",
@@ -1263,9 +1264,7 @@
        "\n",
        "if use_dedicated_endpoint:\n",
        "    DEDICATED_ENDPOINT_DNS = endpoints[LABEL].gca_resource.dedicated_endpoint_dns\n",
-        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
-        "    PROJECT_ID, REGION, endpoints[LABEL].name\n",
-        ")\n",
+        "ENDPOINT_RESOURCE_NAME = endpoints[LABEL].resource_name\n",
        "\n",
        "# @title Chat Completions Inference\n",
        "\n",
@@ -1564,9 +1563,7 @@
      "source": [
        "# @title Chat completion\n",
        "\n",
-        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
-        "    PROJECT_ID, REGION, endpoints[\"optimized_vllm_gpu\"].name\n",
-        ")\n",
+        "ENDPOINT_RESOURCE_NAME = endpoints[\"optimized_vllm_gpu\"].resource_name\n",
        "\n",
        "# @title Chat Completions Inference\n",
        "\n",
@@ -62,6 +62,7 @@
        "### Objective\n",
        "\n",
        "- Finetune Llama 3.1 models with Vertex AI Custom Training Jobs.\n",
+        "- Evaluate the finetuned model using [lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness).\n",
        "- Deploy finetuned Llama 3.1 models on Vertex AI Prediction.\n",
        "- Send prediction requests to your finetuned Llama 3.1 models.\n",
        "\n",
@@ -140,7 +141,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 80320a9a1b818534ca785444e704f6953f2a9dd9\n",
+        "! cd vertex-ai-samples && git reset --hard c45f6a4f4d32e31a050f0e4ba52824b0caf4eda3\n",
        "\n",
        "import datetime\n",
        "import importlib\n",
@@ -238,14 +239,12 @@
        "\n",
        "# @markdown **Only select and fill one of the following sections.**\n",
        "# It is recommended to use \"Google Cloud\" for 405B model since it can be downloaded faster.\n",
-        "# fmt: off\n",
        "LOAD_MODEL_FROM = \"Google Cloud\"  # @param [\"Hugging Face\", \"Google Cloud\"] {isTemplate:true}\n",
-        "# fmt: on\n",
        "\n",
        "# @markdown ---\n",
        "\n",
        "# @markdown ### Access Llama 3.1 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 Llama 3.1 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 Llama 3.1 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",
        "# @markdown *--- Or ---*\n",
@@ -304,13 +303,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",
@@ -450,9 +449,7 @@
        "# @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 Acceletor type to use for training.\n",
-        "# fmt: off\n",
        "training_accelerator_type = \"NVIDIA_A100_80GB\"  # @param [\"NVIDIA_A100_80GB\", \"NVIDIA_H100_80GB\"]\n",
-        "# fmt: on\n",
        "\n",
        "# The pre-built training docker image.\n",
        "if training_accelerator_type == \"NVIDIA_A100_80GB\":\n",
@@ -474,7 +471,7 @@
        "    }\n",
        "\n",
        "TRAIN_DOCKER_URI = (\n",
-        "    f\"{repo}/vertex-vision-model-garden-dockers/pytorch-peft-train:stable_20250213\"\n",
+        "    f\"{repo}/vertex-vision-model-garden-dockers/pytorch-peft-train:stable_20250409\"\n",
        ")\n",
        "\n",
        "# Worker pool spec.\n",
@@ -666,6 +663,129 @@
        "print(f\"Command to copy: tensorboard --logdir {base_output_dir}/logs\")"
      ]
    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "1f-gJ_dldAuQ"
+      },
+      "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": "7j5p83ps88N8"
+      },
+      "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",
+        "if \"8b\" in base_model_id.lower():\n",
+        "    eval_machine_type = \"g2-standard-24\"\n",
+        "    eval_accelerator_type = \"NVIDIA_L4\"\n",
+        "    eval_accelerator_count = 2\n",
+        "    dws_kwargs = {}\n",
+        "    is_dynamic_workload_scheduler = False\n",
+        "elif \"70b\" in base_model_id.lower():\n",
+        "    eval_machine_type = \"a2-ultragpu-4g\"\n",
+        "    eval_accelerator_type = \"NVIDIA_A100_80GB\"\n",
+        "    eval_accelerator_count = 4\n",
+        "    dws_kwargs = {\n",
+        "        \"max_wait_duration\": 1800,  # 30 minutes\n",
+        "        \"scheduling_strategy\": gca_custom_job_compat.Scheduling.Strategy.FLEX_START,\n",
+        "    }\n",
+        "    is_dynamic_workload_scheduler = True\n",
+        "elif \"405b\" in base_model_id.lower():\n",
+        "    print(\n",
+        "        \"405B model is not supported for evaluation. We will skip the evaluation job.\"\n",
+        "    )\n",
+        "    RUN_EVALUATION = False\n",
+        "else:\n",
+        "    raise ValueError(f\"Unsupported model ID or GCS path: {base_model_id}.\")\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 final checkpoint.\n",
+        "evaluation_checkpoint_dir = \"\"  # @param {type:\"string\"}\n",
+        "if not evaluation_checkpoint_dir:\n",
+        "    evaluation_checkpoint_dir = final_checkpoint\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 \"-Instruct\" in pretrained_model_id else False\n",
+        "gpu_memory_utilization = 0.8\n",
+        "max_model_len = 8192  # Maximum context length.\n",
+        "eval_output_dir = os.path.join(base_output_dir, \"lm_eval\")\n",
+        "\n",
+        "if RUN_EVALUATION:\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",
+        "    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_model_id}\",\n",
+        "        f\"--model_args={model_args}\",\n",
+        "        f'--lora_path={evaluation_checkpoint_dir.rstrip(\"/\")}',\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",
+        "    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",
+        "        is_dynamic_workload_scheduler=is_dynamic_workload_scheduler,\n",
+        "    )\n",
+        "    # Pass evaluation arguments and launch job.\n",
+        "    lm_eval_job = aiplatform.CustomContainerTrainingJob(\n",
+        "        display_name=common_util.get_job_name_with_datetime(\"llama3_1-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=boot_disk_size_gb,\n",
+        "        service_account=SERVICE_ACCOUNT,\n",
+        "        base_output_dir=base_output_dir,\n",
+        "        **dws_kwargs,\n",
+        "    )\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
    {
      "cell_type": "code",
      "execution_count": null,
@@ -678,11 +798,6 @@
        "# @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",
        "# 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:20250116_0916_RC00\"\n",
        "\n",
@@ -710,8 +825,8 @@
        "    is_for_training=False,\n",
        ")\n",
        "\n",
-        "# Dedicated endpoint not supported yet.\n",
-        "use_dedicated_endpoint = False\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",
        "gpu_memory_utilization = 0.95\n",
        "max_model_len = 8192  # Maximum context length.\n",
@@ -209,7 +209,7 @@
        "# @markdown 1. Open the [Llama 3.1 model](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct) and [Qwen2.5](https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct) from [Hugging Face](https://huggingface.co/).\n",
        "# @markdown 2. Review and accept the agreement.\n",
        "# @markdown 3. After accepting the agreement, Llama 3.1 and Qwen2.5 models will be available for serving.\n",
-        "# @markdown 4. You must provide a Hugging Face User Access Token (read) to access the 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 4. You must provide a Hugging Face User Access Token (with read access) to access the 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\""
@@ -257,8 +257,8 @@
        "# The pre-built serving docker image for vLLM on TPU\n",
        "vLLM_TPU_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241107_0917_tpu_experimental_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).\n",
-        "use_dedicated_endpoint = False  # @param {type:\"boolean\"}\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",
        "common_util.check_quota(\n",
        "    project_id=PROJECT_ID,\n",
@@ -380,9 +380,7 @@
        "models[\"vllm_fast\"], endpoints[\"vllm_fast\"] = fast_deploy(\n",
        "    \"meta\", \"llama3_1\", \"llama-3.1-8b-instruct\"\n",
        ")\n",
-        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
-        "    PROJECT_ID, REGION, endpoints[\"vllm_fast\"].name\n",
-        ")\n",
+        "ENDPOINT_RESOURCE_NAME = endpoints[\"vllm_fast\"].resource_name\n",
        "BASE_URL = (\n",
        "    f\"https://{REGION}-aiplatform.googleapis.com/v1beta1/{ENDPOINT_RESOURCE_NAME}\"\n",
        ")\n",
@@ -423,7 +421,7 @@
        "# The pre-built serving docker images.\n",
        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241210_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",
+        "# @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",
        "# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
        "if \"8b\" in base_model_name.lower():\n",
@@ -65,8 +65,7 @@
        "\n",
        "- Deploy Llama 3.2 1B and 3B with [Hex-LLM](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-hex-llm) on TPU.\n",
        "- Deploy Llama 3.2 1B and 3B Instruct with the Fast Deployment feature.\n",
-        "- Deploy Llama 3.2 11B-vision and 90B-vision with [vLLM](https://github.com/vllm-project/vllm) on GPU with limited functions supported.\n",
-        "- Deploy Llama 3.2 11B-vision and 90B-vision with reference server on GPU with full functions supported.\n",
+        "- Deploy Llama 3.2 11B-vision and 90B-vision with [vLLM](https://github.com/vllm-project/vllm) on GPU.\n",
        "\n",
        "### File a bug\n",
        "\n",
@@ -112,7 +111,6 @@
    {
      "cell_type": "code",
      "execution_count": null,
-      "language": "python",
      "metadata": {
        "cellView": "form",
        "id": "YXFGIp1l-qtT"
@@ -159,7 +157,6 @@
        "    \"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",
@@ -780,9 +777,7 @@
        "\n",
        "if use_dedicated_endpoint:\n",
        "    DEDICATED_ENDPOINT_DNS = endpoints[\"vllm_fast\"].gca_resource.dedicated_endpoint_dns\n",
-        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
-        "    PROJECT_ID, REGION, endpoints[\"vllm_fast\"].name\n",
-        ")\n",
+        "ENDPOINT_RESOURCE_NAME = endpoints[\"vllm_fast\"].resource_name\n",
        "\n",
        "# @title Chat Completions Inference\n",
        "\n",
@@ -865,7 +860,7 @@
      "outputs": [],
      "source": [
        "# @title Select the model variants\n",
-        "\n",
+        "LABEL = \"vllm_gpu\"\n",
        "base_model_name = \"Llama-3.2-11B-Vision-Instruct\"  # @param [\"Llama-3.2-11B-Vision\", \"Llama-3.2-11B-Vision-Instruct\", \"Llama-3.2-90B-Vision\", \"Llama-3.2-90B-Vision-Instruct\"] {isTemplate:true}\n",
        "hf_model_id = \"meta-llama/\" + base_model_name\n",
        "model_id = os.path.join(VERTEX_AI_MODEL_GARDEN_LLAMA_3_2, base_model_name)\n",
@@ -919,11 +914,17 @@
        "# @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.preview import model_garden\n",
        "\n",
+        "if REGION == \"us-central1\":\n",
+        "    fast_tryout_enabled = True\n",
+        "else:\n",
+        "    fast_tryout_enabled = False\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",
+        "    fast_tryout_enabled=fast_tryout_enabled,\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",
        ")"
@@ -1046,8 +1047,12 @@
        "        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",
+        "        if \"Llama-4\" not in model_id:\n",
+        "            vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "            vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "        else:\n",
+        "            vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "            vllm_args.append(\"--tool-call-parser=llama3_json\")\n",
        "\n",
        "    env_vars = {\n",
        "        \"MODEL_ID\": base_model_id,\n",
@@ -1178,9 +1183,7 @@
        "\n",
        "if use_dedicated_endpoint:\n",
        "    DEDICATED_ENDPOINT_DNS = endpoints[LABEL].gca_resource.dedicated_endpoint_dns\n",
-        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
-        "    PROJECT_ID, REGION, endpoints[LABEL].name\n",
-        ")\n",
+        "ENDPOINT_RESOURCE_NAME = endpoints[LABEL].resource_name\n",
        "\n",
        "# @title Chat Completions Inference\n",
        "\n",
@@ -1255,9 +1258,7 @@
        "\n",
        "if use_dedicated_endpoint:\n",
        "    DEDICATED_ENDPOINT_DNS = endpoints[LABEL].gca_resource.dedicated_endpoint_dns\n",
-        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
-        "    PROJECT_ID, REGION, endpoints[LABEL].name\n",
-        ")\n",
+        "ENDPOINT_RESOURCE_NAME = endpoints[LABEL].resource_name\n",
        "\n",
        "# @title Chat Completions Inference\n",
        "\n",
@@ -1314,252 +1315,6 @@
        "# @markdown Click \"Show Code\" to see more details."
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "4AoiFsQWhYtM"
-      },
-      "source": [
-        "## Deploy prebuilt Llama 3.2 11B-Vision and 90B-Vision with reference server"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "cellView": "form",
-        "id": "wXkUpLOXhYtM"
-      },
-      "outputs": [],
-      "source": [
-        "# @title Deploy\n",
-        "\n",
-        "# @markdown This section uploads prebuilt Llama 3.2 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",
-        "# @markdown The reference server can be used with full input formats supported.\n",
-        "\n",
-        "# @markdown Set the model to deploy.\n",
-        "\n",
-        "base_model_name = \"Llama-3.2-11B-Vision-Instruct-meta\"  # @param [\"Llama-3.2-11B-Vision-meta\", \"Llama-3.2-11B-Vision-Instruct-meta\", \"Llama-3.2-90B-Vision-meta\", \"Llama-3.2-90B-Vision-Instruct-meta\"] {isTemplate:true}\n",
-        "model_id = os.path.join(VERTEX_AI_MODEL_GARDEN_LLAMA_3_2, base_model_name)\n",
-        "hf_model_id = \"meta-llama/\" + base_model_name\n",
-        "\n",
-        "# The pre-built serving docker images.\n",
-        "REF_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-serve:llama_ref_impl\"\n",
-        "\n",
-        "# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
-        "if \"11B\" in base_model_name:\n",
-        "    accelerator_type = \"NVIDIA_TESLA_A100\"\n",
-        "    machine_type = \"a2-highgpu-1g\"\n",
-        "    accelerator_count = 1\n",
-        "elif \"90B\" in base_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: {base_model_name}.\")\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_ref(\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",
-        "    max_model_len: int = 4096,\n",
-        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
-        "    \"\"\"Deploys models with reference server 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",
-        "    ref_args = [\n",
-        "        \"torchrun\",\n",
-        "        f\"--nproc_per_node={accelerator_count}\",\n",
-        "        \"server.py\",\n",
-        "        \"--port=8080\",\n",
-        "        f\"--model={model_id}\",\n",
-        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
-        "        f\"--max-model-len={max_model_len}\",\n",
-        "    ]\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=REF_DOCKER_URI,\n",
-        "        serving_container_args=ref_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_pytorch_llama3_2_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",
-        "max_model_len = 4096\n",
-        "\n",
-        "models[\"ref_gpu\"], endpoints[\"ref_gpu\"] = deploy_model_ref(\n",
-        "    model_name=common_util.get_job_name_with_datetime(prefix=\"llama3_2-serve-ref\"),\n",
-        "    model_id=model_id,\n",
-        "    publisher=\"meta\",\n",
-        "    publisher_model_id=\"llama3-2\",\n",
-        "    base_model_id=hf_model_id,\n",
-        "    machine_type=machine_type,\n",
-        "    accelerator_type=accelerator_type,\n",
-        "    accelerator_count=accelerator_count,\n",
-        "    max_model_len=max_model_len,\n",
-        ")\n",
-        "# @markdown Click \"Show Code\" to see more details."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "cellView": "form",
-        "id": "Gl72VUchVmK5"
-      },
-      "outputs": [],
-      "source": [
-        "# @title Chat completion for text-only models\n",
-        "\n",
-        "if use_dedicated_endpoint:\n",
-        "    DEDICATED_ENDPOINT_DNS = endpoints[\"ref_gpu\"].gca_resource.dedicated_endpoint_dns\n",
-        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
-        "    PROJECT_ID, REGION, endpoints[\"ref_gpu\"].name\n",
-        ")\n",
-        "\n",
-        "# @markdown Once deployment succeeds, you can send requests to the endpoint using the Vertex SDK.\n",
-        "\n",
-        "# @markdown Next fill out some request parameters:\n",
-        "\n",
-        "user_message = \"Tell me about large language models\"  # @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",
-        "top_p = 0.9  # @param {type: \"number\"}\n",
-        "\n",
-        "# @markdown Now we can send a request.\n",
-        "\n",
-        "instances = [\n",
-        "    {\n",
-        "        \"messages\": [{\"role\": \"user\", \"content\": user_message}],\n",
-        "        \"temperature\": temperature,\n",
-        "        \"max_tokens\": max_tokens,\n",
-        "        \"top_p\": top_p,\n",
-        "    }\n",
-        "]\n",
-        "\n",
-        "model_response = endpoints[\"ref_gpu\"].predict(instances=instances)\n",
-        "\n",
-        "print(model_response.predictions[0])\n",
-        "\n",
-        "# @markdown Click \"Show Code\" to see more details."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "cellView": "form",
-        "id": "pY_hkPI2hYtM"
-      },
-      "outputs": [],
-      "source": [
-        "# @title Chat completion for vision models\n",
-        "\n",
-        "if use_dedicated_endpoint:\n",
-        "    DEDICATED_ENDPOINT_DNS = endpoints[\"ref_gpu\"].gca_resource.dedicated_endpoint_dns\n",
-        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
-        "    PROJECT_ID, REGION, endpoints[\"ref_gpu\"].name\n",
-        ")\n",
-        "\n",
-        "# @markdown Once deployment succeeds, you can send requests to the endpoint using the Vertex SDK.\n",
-        "\n",
-        "# @markdown Next fill out some request parameters:\n",
-        "\n",
-        "user_image1 = \"https://upload.wikimedia.org/wikipedia/commons/thumb/8/84/Male_and_female_chicken_sitting_together.jpg/440px-Male_and_female_chicken_sitting_together.jpg\"  # @param {type: \"string\"}\n",
-        "user_image2 = \"https://upload.wikimedia.org/wikipedia/commons/thumb/9/91/Joseph_Crawhall_-_Spanish_Cock_And_Snail.jpg/337px-Joseph_Crawhall_-_Spanish_Cock_And_Snail.jpg\"  # @param {type: \"string\"}\n",
-        "user_message = \"How many chickens are there in these two images?\"  # @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",
-        "top_p = 0.9  # @param {type: \"number\"}\n",
-        "\n",
-        "# @markdown Now we can send a request.\n",
-        "\n",
-        "instances = [\n",
-        "    {\n",
-        "        \"messages\": [\n",
-        "            {\n",
-        "                \"role\": \"user\",\n",
-        "                \"content\": [\n",
-        "                    {\"type\": \"image_url\", \"image_url\": {\"url\": user_image1}},\n",
-        "                    {\"type\": \"image_url\", \"image_url\": {\"url\": user_image2}},\n",
-        "                    {\"type\": \"text\", \"text\": user_message},\n",
-        "                ],\n",
-        "            },\n",
-        "        ],\n",
-        "        \"temperature\": temperature,\n",
-        "        \"max_tokens\": max_tokens,\n",
-        "        \"top_p\": top_p,\n",
-        "    }\n",
-        "]\n",
-        "\n",
-        "model_response = endpoints[\"ref_gpu\"].predict(instances=instances)\n",
-        "\n",
-        "print(model_response.predictions[0])\n",
-        "\n",
-        "# @markdown Click \"Show Code\" to see more details."
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -59,12 +59,12 @@
      "source": [
        "## Overview\n",
        "\n",
-        "This notebook demonstrates downloading, deploying, and serving prebuilt Llama 3.3 model with [vLLM](https://github.com/vllm-project/vllm).\n",
+        "This notebook demonstrates downloading, deploying, and serving prebuilt Llama 3.3 model with [vLLM](https://github.com/vllm-project/vllm) or [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM).\n",
        "\n",
        "\n",
        "### Objective\n",
        "\n",
-        "- Deploy Llama 3.3 70B Instruct with vLLM on GPU, optionally with dynamic LoRA adapters.\n",
+        "- Deploy Llama 3.3 70B Instruct with vLLM (optionally with dynamic LoRA adapters) or TensorRT-LLM on GPU.\n",
        "\n",
        "### File a bug\n",
        "\n",
@@ -136,8 +136,10 @@
        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
        "\n",
        "import importlib\n",
+        "import json\n",
        "import os\n",
        "import re\n",
+        "import time\n",
        "from typing import Tuple\n",
        "\n",
        "import requests\n",
@@ -217,9 +219,7 @@
        "\n",
        "# @markdown Set the model to deploy.\n",
        "\n",
-        "# fmt: off\n",
        "base_model_name = \"Llama-3.3-70B-Instruct\"  # @param [\"Llama-3.3-70B-Instruct\"] {isTemplate:true}\n",
-        "# fmt: on\n",
        "\n",
        "model_id = os.path.join(VERTEX_AI_MODEL_GARDEN_LLAMA_3_3, base_model_name)\n",
        "ENABLE_DYNAMIC_LORA = True  # @param {type:\"boolean\", isTemplate:true}\n",
@@ -401,8 +401,12 @@
        "        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",
+        "        if \"Llama-4\" not in model_id:\n",
+        "            vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "            vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "        else:\n",
+        "            vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "            vllm_args.append(\"--tool-call-parser=llama3_json\")\n",
        "\n",
        "    env_vars = {\n",
        "        \"MODEL_ID\": base_model_id,\n",
@@ -591,9 +595,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",
@@ -655,6 +657,284 @@
        "# @markdown Click \"Show Code\" to see more details."
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "fQL54XQV1Fq2"
+      },
+      "source": [
+        "## Deploy Llama 3.3 70B Instruct FP8 with TensorRT-LLM"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "Apv7MP2-1I1B"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "\n",
+        "# @markdown This section uploads Llama 3.3 70B Instruct FP8 model to Model Registry and deploys it to a Vertex Prediction Endpoint. It takes ~1 hour to finish.\n",
+        "\n",
+        "# @markdown The [Llama 3.3 70B Instruct FP8 model](https://huggingface.co/nvidia/Llama-3.3-70B-Instruct-FP8) is provided by Nvidia.\n",
+        "\n",
+        "# @markdown It's recommended to use the region selected by the deployment button on the model card. If the deployment button is not available, it's recommended to stay with the default region of the notebook.\n",
+        "\n",
+        "# @markdown Set the model to deploy.\n",
+        "\n",
+        "base_model_name = \"Llama-3.3-70B-Instruct-FP8\"  # @param [\"Llama-3.3-70B-Instruct-FP8\"] {isTemplate:true}\n",
+        "model_id = os.path.join(VERTEX_AI_MODEL_GARDEN_LLAMA_3_3, base_model_name)\n",
+        "\n",
+        "# The pre-built serving docker images.\n",
+        "TRTLLM_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/tensorrt-llm.cu128.0-18.ubuntu2404.py312:20250514-1800-rc0\"\n",
+        "\n",
+        "# @markdown Choose whether to use a [Spot VM](https://cloud.google.com/compute/docs/instances/spot) for the deployment.\n",
+        "is_spot = False  # @param {type:\"boolean\"}\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",
+        "accelerator_type = \"NVIDIA_H100_80GB\"\n",
+        "accelerator_count = 8\n",
+        "machine_type = \"a3-highgpu-8g\"\n",
+        "multihost_gpu_node_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",
+        ")\n",
+        "\n",
+        "# 128K context length.\n",
+        "MAX_MODEL_LEN = 131072\n",
+        "\n",
+        "\n",
+        "def poll_operation(op_name: str) -> bool:  # noqa: F811\n",
+        "    creds, _ = auth.default()\n",
+        "    auth_req = auth.transport.requests.Request()\n",
+        "    creds.refresh(auth_req)\n",
+        "    headers = {\n",
+        "        \"Authorization\": f\"Bearer {creds.token}\",\n",
+        "    }\n",
+        "    get_resp = requests.get(\n",
+        "        f\"https://{REGION}-aiplatform.googleapis.com/ui/{op_name}\",\n",
+        "        headers=headers,\n",
+        "    )\n",
+        "    opjs = get_resp.json()\n",
+        "    if \"error\" in opjs:\n",
+        "        raise ValueError(f\"Operation failed: {opjs['error']}\")\n",
+        "    return opjs.get(\"done\", False)\n",
+        "\n",
+        "\n",
+        "def poll_and_wait(op_name: str, total_wait: int, interval: int = 60):  # noqa: F811\n",
+        "    waited = 0\n",
+        "    while not poll_operation(op_name):\n",
+        "        if waited > total_wait:\n",
+        "            raise TimeoutError(\"Operation timed out\")\n",
+        "        print(\n",
+        "            f\"\\rStill waiting for operation... Waited time in second: {waited:<6}\",\n",
+        "            end=\"\",\n",
+        "            flush=True,\n",
+        "        )\n",
+        "        waited += interval\n",
+        "        time.sleep(interval)\n",
+        "\n",
+        "\n",
+        "def deploy_model_tensorrt_llm_multihost(\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 = \"\",\n",
+        "    machine_type: str = \"g2-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 1,\n",
+        "    multihost_gpu_node_count: int = 1,\n",
+        "    gpu_memory_utilization: float | None = None,\n",
+        "    max_input_len: int | None = None,\n",
+        "    max_model_len: int | None = None,\n",
+        "    max_num_seqs: int | None = None,\n",
+        "    enable_trust_remote_code: bool = False,\n",
+        "    enable_chunked_prefill: bool = False,\n",
+        "    use_dedicated_endpoint: bool = False,\n",
+        "    is_spot: bool = True,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys trained models with TensorRT-LLM on 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",
+        "    container_args = [\n",
+        "        \"python\",\n",
+        "        \"api_server.py\",\n",
+        "        \"--host=0.0.0.0\",\n",
+        "        \"--port=7080\",\n",
+        "        f\"--model={model_id}\",\n",
+        "        f\"--tensor-parallel-size={accelerator_count * multihost_gpu_node_count}\",\n",
+        "    ]\n",
+        "\n",
+        "    if gpu_memory_utilization:\n",
+        "        container_args.append(f\"--gpu-memory-utilization={gpu_memory_utilization}\")\n",
+        "\n",
+        "    if max_input_len:\n",
+        "        container_args.append(f\"--max-input-len={max_input_len}\")\n",
+        "\n",
+        "    if max_model_len:\n",
+        "        container_args.append(f\"--max-model-len={max_model_len}\")\n",
+        "\n",
+        "    if max_num_seqs:\n",
+        "        container_args.append(f\"--max-num-seqs={max_num_seqs}\")\n",
+        "\n",
+        "    if enable_trust_remote_code:\n",
+        "        container_args.append(\"--trust-remote-code=True\")\n",
+        "\n",
+        "    if enable_chunked_prefill:\n",
+        "        container_args.append(\"--enable-chunked-prefill=True\")\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=TRTLLM_DOCKER_URI,\n",
+        "        serving_container_args=container_args,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_predict_route=\"/v1/chat/completions\",\n",
+        "        serving_container_health_route=\"/health\",\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 {int(accelerator_count * multihost_gpu_node_count)} {accelerator_type} GPU(s).\"\n",
+        "    )\n",
+        "\n",
+        "    creds, _ = auth.default()\n",
+        "    auth_req = auth.transport.requests.Request()\n",
+        "    creds.refresh(auth_req)\n",
+        "\n",
+        "    url = f\"https://{REGION}-aiplatform.googleapis.com/ui/projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint.name}:deployModel\"\n",
+        "    headers = {\n",
+        "        \"Content-Type\": \"application/json\",\n",
+        "        \"Authorization\": f\"Bearer {creds.token}\",\n",
+        "    }\n",
+        "    data = {\n",
+        "        \"deployedModel\": {\n",
+        "            \"model\": model.resource_name,\n",
+        "            \"displayName\": model_name,\n",
+        "            \"dedicatedResources\": {\n",
+        "                \"machineSpec\": {\n",
+        "                    \"machineType\": machine_type,\n",
+        "                    \"multihostGpuNodeCount\": multihost_gpu_node_count,\n",
+        "                    \"acceleratorType\": accelerator_type,\n",
+        "                    \"acceleratorCount\": accelerator_count,\n",
+        "                },\n",
+        "                \"minReplicaCount\": 1,\n",
+        "                \"maxReplicaCount\": 1,\n",
+        "            },\n",
+        "            \"system_labels\": {\n",
+        "                \"NOTEBOOK_NAME\": \"model_garden_pytorch_llama3_3_deployment.ipynb\",\n",
+        "                \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
+        "            },\n",
+        "        },\n",
+        "    }\n",
+        "    if service_account:\n",
+        "        data[\"deployedModel\"][\"serviceAccount\"] = service_account\n",
+        "    if is_spot:\n",
+        "        data[\"deployedModel\"][\"dedicatedResources\"][\"spot\"] = True\n",
+        "    response = requests.post(url, headers=headers, json=data)\n",
+        "    print(f\"Deploy Model response: {response.json()}\")\n",
+        "    if response.status_code != 200 or \"name\" not in response.json():\n",
+        "        raise ValueError(f\"Failed to deploy model: {response.text}\")\n",
+        "    poll_and_wait(response.json()[\"name\"], 7200)\n",
+        "    print(\"endpoint_name:\", endpoint.name)\n",
+        "\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "models[\"trtllm_gpu\"], endpoints[\"trtllm_gpu\"] = deploy_model_tensorrt_llm_multihost(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=\"llama3-3-serve\"),\n",
+        "    model_id=model_id,\n",
+        "    publisher=\"meta-llama\",\n",
+        "    publisher_model_id=\"llama-3.3-70b-instruct\",\n",
+        "    base_model_id=\"meta-llama/Llama-3.3-70B-Instruct\",\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    multihost_gpu_node_count=multihost_gpu_node_count,\n",
+        "    max_model_len=MAX_MODEL_LEN,\n",
+        "    enable_chunked_prefill=True,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    is_spot=is_spot,\n",
+        ")\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "tcLecC90K6Br"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Raw predict\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint using `:rawPredict`. The OpenAI Client chat completions support is coming soon.\n",
+        "endpoints[\"trtllm_gpu\"] = aiplatform.Endpoint(endpoints[\"trtllm_gpu\"].resource_name)\n",
+        "\n",
+        "# @markdown Fill out some request parameters:\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",
+        "\n",
+        "# @markdown Now we can send a request.\n",
+        "\n",
+        "response = endpoints[\"trtllm_gpu\"].raw_predict(\n",
+        "    body=json.dumps(\n",
+        "        {\n",
+        "            \"model\": \"\",\n",
+        "            \"messages\": [\n",
+        "                {\n",
+        "                    \"role\": \"user\",\n",
+        "                    \"content\": user_message,\n",
+        "                }\n",
+        "            ],\n",
+        "            \"max_tokens\": max_tokens,\n",
+        "            \"temperature\": temperature,\n",
+        "        }\n",
+        "    ),\n",
+        "    headers={\"Content-Type\": \"application/json\"},\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        ")\n",
+        "print(response.json()[\"choices\"][0][\"message\"][\"content\"])\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -62,6 +62,7 @@
        "### Objective\n",
        "\n",
        "- Finetune Llama 3.3 models with Vertex AI Custom Training Jobs.\n",
+        "- Evaluate the finetuned model using [lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness).\n",
        "- Deploy finetuned Llama 3.3 models on Vertex AI Prediction.\n",
        "- Send prediction requests to your finetuned Llama 3.3 models.\n",
        "\n",
@@ -99,17 +100,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.34.2\n",
        "! pip install --upgrade --quiet transformers==4.47.1\n",
-        "! pip install --upgrade --quiet datasets==2.20.0\n",
-        "\n",
-        "# Load local tensorboard.\n",
-        "%load_ext tensorboard"
+        "! pip install --upgrade --quiet datasets==2.20.0"
      ]
    },
    {
@@ -146,7 +141,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",
@@ -243,14 +238,12 @@
        "# @markdown Alternatively, you can also load the original Llama 3.3 models for finetuning and serving from Vertex AI after accepting the agreement.\n",
        "\n",
        "# @markdown **Only select and fill one of the following sections.**\n",
-        "# fmt: off\n",
        "LOAD_MODEL_FROM = \"Google Cloud\"  # @param [\"Hugging Face\", \"Google Cloud\"] {isTemplate:true}\n",
-        "# fmt: on\n",
        "\n",
        "# @markdown ---\n",
        "\n",
        "# @markdown ### Access Llama 3.3 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 Llama 3.3 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 Llama 3.3 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",
        "# @markdown *--- Or ---*\n",
@@ -300,7 +293,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",
@@ -309,13 +302,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",
@@ -341,23 +334,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\"}"
      ]
    },
    {
@@ -392,8 +387,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",
@@ -409,32 +402,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",
+        "    )"
      ]
    },
    {
@@ -476,18 +467,17 @@
        "    }\n",
        "\n",
        "TRAIN_DOCKER_URI = (\n",
-        "    f\"{repo}/vertex-vision-model-garden-dockers/pytorch-peft-train:stable_20250320\"\n",
+        "    f\"{repo}/vertex-vision-model-garden-dockers/pytorch-peft-train:stable_20250409\"\n",
        ")\n",
        "\n",
        "# Worker pool spec.\n",
+        "boot_disk_size_gb = 500\n",
        "if training_accelerator_type == \"NVIDIA_A100_80GB\":\n",
        "    per_node_accelerator_count = 8\n",
        "    training_machine_type = \"a2-ultragpu-8g\"\n",
-        "    boot_disk_size_gb = 500\n",
        "elif training_accelerator_type == \"NVIDIA_H100_80GB\":\n",
        "    per_node_accelerator_count = 8\n",
        "    training_machine_type = \"a3-highgpu-8g\"\n",
-        "    boot_disk_size_gb = 2000\n",
        "else:\n",
        "    raise ValueError(\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",
@@ -572,10 +562,10 @@
        "labels[\"versioned-mg-tune\"] = f\"{labels['mg-tune']}-{versioned_model_id}\"\n",
        "\n",
        "eval_args = [\n",
-        "    f\"--eval_dataset={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_metric_name=loss\",\n",
        "]\n",
@@ -585,9 +575,9 @@
        "    \"--task=instruct-lora\",\n",
        "    \"--input_masking=True\",\n",
        "    f\"--pretrained_model_name_or_path={pretrained_model_id}\",\n",
-        "    f\"--train_dataset={train_dataset_name}\",\n",
-        "    f\"--train_split={train_split_name}\",\n",
-        "    f\"--train_column={instruct_column_in_dataset}\",\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\"--per_device_train_batch_size={per_device_train_batch_size}\",\n",
        "    f\"--gradient_accumulation_steps={gradient_accumulation_steps}\",\n",
@@ -663,6 +653,117 @@
        "print(f\"Command to copy: tensorboard --logdir {base_output_dir}/logs\")"
      ]
    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "maTaH3s_dWw9"
+      },
+      "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": "i2nSJldtrZeA"
+      },
+      "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_machine_type = \"a3-highgpu-8g\"\n",
+        "eval_accelerator_type = \"NVIDIA_H100_80GB\"\n",
+        "eval_accelerator_count = 8\n",
+        "dws_kwargs = {\n",
+        "    \"max_wait_duration\": 1800,  # 30 minutes\n",
+        "    \"scheduling_strategy\": gca_custom_job_compat.Scheduling.Strategy.FLEX_START,\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 final checkpoint.\n",
+        "evaluation_checkpoint_dir = \"\"  # @param {type:\"string\"}\n",
+        "if not evaluation_checkpoint_dir:\n",
+        "    evaluation_checkpoint_dir = final_checkpoint\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 \"-Instruct\" in pretrained_model_id else False\n",
+        "gpu_memory_utilization = 0.9\n",
+        "max_model_len = 8192  # Maximum context length.\n",
+        "eval_output_dir = os.path.join(base_output_dir, \"lm_eval\")\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",
+        "\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_model_id}\",\n",
+        "    f\"--model_args={model_args}\",\n",
+        "    f'--lora_path={evaluation_checkpoint_dir.rstrip(\"/\")}',\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 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",
+        "        is_dynamic_workload_scheduler=True,\n",
+        "    )\n",
+        "\n",
+        "    # Pass evaluation arguments and launch job.\n",
+        "    lm_eval_job = aiplatform.CustomContainerTrainingJob(\n",
+        "        display_name=common_util.get_job_name_with_datetime(\"llama3_3-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=boot_disk_size_gb,\n",
+        "        service_account=SERVICE_ACCOUNT,\n",
+        "        base_output_dir=base_output_dir,\n",
+        "        **dws_kwargs,\n",
+        "    )\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
    {
      "cell_type": "code",
      "execution_count": null,
@@ -676,11 +777,6 @@
        "# @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",
        "# 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:20250116_0916_RC00\"\n",
        "\n",
@@ -706,8 +802,8 @@
        "    is_for_training=False,\n",
        ")\n",
        "\n",
-        "# Dedicated endpoint not supported yet.\n",
-        "use_dedicated_endpoint = False\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",
        "gpu_memory_utilization = 0.95\n",
        "max_model_len = 8192  # Maximum context length.\n",
@@ -544,9 +544,7 @@
      "source": [
        "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",
@@ -135,8 +135,6 @@
        ")\n",
        "\n",
        "models, endpoints = {}, {}\n",
-        "# Dedicated endpoint not supported yet\n",
-        "use_dedicated_endpoint = False\n",
        "\n",
        "# Get the default cloud project id.\n",
        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
@@ -224,7 +222,7 @@
        "# @markdown ---\n",
        "\n",
        "# @markdown ### Access Llama 3 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 Llama 3 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 Llama 3 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",
@@ -554,6 +552,9 @@
        "accelerator_type = \"NVIDIA_H100_80GB\"  # @param [\"NVIDIA_L4\", \"NVIDIA_H100_80GB\"]\n",
        "machine_type = None\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 in [here](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute).\n",
        "if \"8b\" in MODEL_ID.lower():\n",
        "    if accelerator_type == \"NVIDIA_L4\":\n",
@@ -59,11 +59,12 @@
      "source": [
        "## Overview\n",
        "\n",
-        "This notebook demonstrates deploying Llama 4 models on GPU.\n",
+        "This notebook demonstrates deploying Llama 4 models on GPUs and TPUs.\n",
        "\n",
        "### Objective\n",
        "\n",
        "- Deploy Llama 4 models with [vLLM](https://github.com/vllm-project/vllm) on GPU.\n",
+        "- Deploy Llama 4 models with [Hex-LLM](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-hex-llm) on TPU.\n",
        "\n",
        "\n",
        "### File a bug\n",
@@ -226,7 +227,7 @@
        "PUBLISHER_MODEL_NAME = \"publishers/meta/models/llama4\"\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:20250405_1205_RC01\"\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20250417_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",
@@ -374,8 +375,12 @@
        "        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",
+        "        if \"Llama-4\" not in model_id:\n",
+        "            vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "            vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "        else:\n",
+        "            vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "            vllm_args.append(\"--tool-call-parser=llama3_json\")\n",
        "\n",
        "    env_vars = {\n",
        "        \"MODEL_ID\": base_model_id,\n",
@@ -480,6 +485,7 @@
        "    gpu_memory_utilization=gpu_memory_utilization,\n",
        "    max_model_len=max_model_len,\n",
        "    enforce_eager=False,\n",
+        "    enable_llama_tool_parser=True,\n",
        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
        ")\n",
        "# @markdown Click \"Show Code\" to see more details."
@@ -563,9 +569,7 @@
        "\n",
        "if use_dedicated_endpoint:\n",
        "    DEDICATED_ENDPOINT_DNS = endpoints[LABEL].gca_resource.dedicated_endpoint_dns\n",
-        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
-        "    PROJECT_ID, REGION, endpoints[LABEL].name\n",
-        ")\n",
+        "ENDPOINT_RESOURCE_NAME = endpoints[LABEL].resource_name\n",
        "\n",
        "# @title Chat Completions Inference\n",
        "\n",
@@ -640,9 +644,426 @@
        "\n",
        "if use_dedicated_endpoint:\n",
        "    DEDICATED_ENDPOINT_DNS = endpoints[LABEL].gca_resource.dedicated_endpoint_dns\n",
-        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
-        "    PROJECT_ID, REGION, endpoints[LABEL].name\n",
+        "ENDPOINT_RESOURCE_NAME = endpoints[LABEL].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": "U4tRt3ZtClzH"
+      },
+      "source": [
+        "## Deploy Llama 4 models with Hex-LLM"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "pIoNTUDsCr8m"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "\n",
+        "# @markdown This section uploads prebuilt Llama 4 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",
+        "# @markdown Select one of the model variations.\n",
+        "base_model_name = \"Llama-4-Scout-17B-16E-Instruct\"  # @param [\"Llama-4-Scout-17B-16E\", \"Llama-4-Scout-17B-16E-Instruct\"] {isTemplate:true}\n",
+        "TPU_DEPLOYMENT_REGION = \"us-west1\"  # @param [\"us-west1\"] {isTemplate:true}\n",
+        "model_path_prefix = \"meta\"\n",
+        "hf_model_id = \"meta-llama/\" + base_model_name\n",
+        "model_id = os.path.join(VERTEX_AI_MODEL_GARDEN_LLAMA_4, base_model_name)\n",
+        "\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:llama4\"\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",
+        "\n",
+        "# @markdown | Model Version              | Default Max Model Length | Default TPU configuration |\n",
+        "# @markdown |----------------------------|------------------|-----------------------------|\n",
+        "# @markdown | Llama-4-Scout-17B-16E                     | 8192            | 4 x 4 TPU_V5e ct5lp-hightpu-4t  |\n",
+        "# @markdown | Llama-4-Scout-17B-16E-Instruct                     | 8192            | 4 x 4 TPU_V5e ct5lp-hightpu-4t  |\n",
+        "\n",
+        "# Note: 1 TPU V5 chip has only one core.\n",
+        "tpu_type = \"TPU_V5e\"\n",
+        "\n",
+        "if \"Llama-4\" in model_id:\n",
+        "    tpu_count = 16\n",
+        "    tpu_topo = \"4x4\"\n",
+        "    num_hosts = 4\n",
+        "    max_model_len = 8192\n",
+        "    machine_type = \"ct5lp-hightpu-4t\"\n",
+        "else:\n",
+        "    raise ValueError(f\"Unsupported MODEL_ID: {model_id}\")\n",
+        "\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=TPU_DEPLOYMENT_REGION,\n",
+        "    accelerator_type=tpu_type,\n",
+        "    accelerator_count=tpu_count,\n",
+        "    is_for_training=False,\n",
+        ")\n",
+        "\n",
+        "# Server parameters.\n",
+        "tensor_parallel_size = tpu_count\n",
+        "\n",
+        "# 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",
+        "hbm_utilization_factor = 0.85\n",
+        "\n",
+        "max_running_seqs = 256\n",
+        "decode_seqs_padding = 128\n",
+        "\n",
+        "# Endpoint configurations.\n",
+        "min_replica_count = 1\n",
+        "max_replica_count = 1\n",
+        "\n",
+        "\n",
+        "def deploy_model_hexllm(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    publisher: str,\n",
+        "    publisher_model_id: 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",
+        "    machine_type: str = \"ct5lp-hightpu-1t\",\n",
+        "    tpu_topology: str = \"1x1\",\n",
+        "    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",
+        "    min_replica_count: int = 1,\n",
+        "    max_replica_count: int = 1,\n",
+        "    use_dedicated_endpoint: bool = False,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys models with Hex-LLM on TPU in Vertex AI.\"\"\"\n",
+        "    if endpoint_id:\n",
+        "        aip_endpoint_name = (\n",
+        "            f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_id}\"\n",
+        "        )\n",
+        "        endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "    else:\n",
+        "        endpoint = aiplatform.Endpoint.create(\n",
+        "            display_name=f\"{model_name}-endpoint\",\n",
+        "            location=TPU_DEPLOYMENT_REGION,\n",
+        "            dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
+        "        )\n",
+        "\n",
+        "    if not base_model_id:\n",
+        "        base_model_id = model_id\n",
+        "\n",
+        "    if not tensor_parallel_size:\n",
+        "        tensor_parallel_size = int(machine_type[-2])\n",
+        "\n",
+        "    num_hosts = int(tpu_topology.split(\"x\")[0])\n",
+        "\n",
+        "    # Learn more about the supported arguments and environment variables at https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-hex-llm#config-server.\n",
+        "    hexllm_args = [\n",
+        "        \"--host=0.0.0.0\",\n",
+        "        \"--port=7080\",\n",
+        "        f\"--model={model_id}\",\n",
+        "        f\"--data_parallel_size={data_parallel_size}\",\n",
+        "        f\"--tensor_parallel_size={tensor_parallel_size}\",\n",
+        "        f\"--num_hosts={num_hosts}\",\n",
+        "        f\"--hbm_utilization_factor={hbm_utilization_factor}\",\n",
+        "        f\"--max_running_seqs={max_running_seqs}\",\n",
+        "        f\"--max_model_len={max_model_len}\",\n",
+        "    ]\n",
+        "\n",
+        "    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",
+        "        hexllm_args.append(\"--enable_prefix_cache_hbm\")\n",
+        "\n",
+        "    env_vars = {\n",
+        "        \"MODEL_ID\": base_model_id,\n",
+        "        \"HEX_LLM_LOG_LEVEL\": \"info\",\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "\n",
+        "    # HF_TOKEN is not a compulsory field and may not be defined.\n",
+        "    try:\n",
+        "        if HF_TOKEN:\n",
+        "            env_vars.update({\"HF_TOKEN\": HF_TOKEN})\n",
+        "    except:\n",
+        "        pass\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=HEXLLM_DOCKER_URI,\n",
+        "        serving_container_command=[\"python\", \"-m\", \"hex_llm.server.api_server\"],\n",
+        "        serving_container_args=hexllm_args,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_predict_route=\"/generate\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        serving_container_environment_variables=env_vars,\n",
+        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
+        "        serving_container_deployment_timeout=7200,\n",
+        "        location=TPU_DEPLOYMENT_REGION,\n",
+        "        model_garden_source_model_name=(\n",
+        "            f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
+        "        ),\n",
+        "    )\n",
+        "\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        tpu_topology=tpu_topology if num_hosts > 1 else None,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=service_account,\n",
+        "        min_replica_count=min_replica_count,\n",
+        "        max_replica_count=max_replica_count,\n",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_llama4_deployment.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
+        "        },\n",
+        "    )\n",
+        "    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[\"hexllm_tpu\"], endpoints[\"hexllm_tpu\"] = deploy_model_hexllm(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=\"llama4-serve-hexllm\"),\n",
+        "    model_id=model_id,\n",
+        "    publisher=\"meta\",\n",
+        "    publisher_model_id=\"llama4\",\n",
+        "    tensor_parallel_size=tensor_parallel_size,\n",
+        "    machine_type=machine_type,\n",
+        "    tpu_topology=tpu_topo,\n",
+        "    hbm_utilization_factor=hbm_utilization_factor,\n",
+        "    max_running_seqs=max_running_seqs,\n",
+        "    decode_seqs_padding=decode_seqs_padding,\n",
+        "    max_model_len=max_model_len,\n",
+        "    min_replica_count=min_replica_count,\n",
+        "    max_replica_count=max_replica_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "jp77ifJ6CvZ4"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict\n",
+        "\n",
+        "# @markdown **Note that the first few prompts will take longer to execute.**\n",
+        "# @markdown **Note that the framework supports an 8k context window, which can process a combination of both image and text inputs up to this limit (in tokens).**\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts based on your `template`. Note that the first few prompts will take longer to execute.\n",
+        "\n",
+        "# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for more details.\n",
+        "\n",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown > What is a car?\n",
+        "# @markdown > A car is a four-wheeled vehicle designed for the transportation of passengers and their belongings.\n",
+        "# @markdown ```\n",
+        "\n",
+        "# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for more details.\n",
+        "\n",
+        "# Loads an existing endpoint instance using the endpoint name:\n",
+        "# - Using `endpoint_name = endpoint.name` allows us to get the endpoint\n",
+        "#   name of the endpoint `endpoint` created in the cell 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",
+        "# endpoint_name = endpoint_without_peft.name\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 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",
+        "top_p = 1.0  # @param {type: \"number\"}\n",
+        "top_k = 1  # @param {type: \"integer\"}\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",
+        "    },\n",
+        "]\n",
+        "response = endpoints[\"hexllm_tpu\"].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": "B6WA_XX-I2NL"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Chat completion with text-only requests\n",
+        "\n",
+        "if use_dedicated_endpoint:\n",
+        "    DEDICATED_ENDPOINT_DNS = endpoints[\"hexllm_tpu\"].gca_resource.dedicated_endpoint_dns\n",
+        "ENDPOINT_RESOURCE_NAME = endpoints[\"hexllm_tpu\"].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": "mAbnF4uEJCGk"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Chat completion for multi-modal requests\n",
+        "\n",
+        "if use_dedicated_endpoint:\n",
+        "    DEDICATED_ENDPOINT_DNS = endpoints[\"hexllm_tpu\"].gca_resource.dedicated_endpoint_dns\n",
+        "ENDPOINT_RESOURCE_NAME = endpoints[\"hexllm_tpu\"].resource_name\n",
        "\n",
        "# @title Chat Completions Inference\n",
        "\n",
@@ -157,7 +157,7 @@
        "    location=REGION,\n",
        ")\n",
        "\n",
-        "# @markdown You must provide a Hugging Face User Access Token (read) to access models from Hugging Face. 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 models from Hugging Face. 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\"}\n",
        "assert HF_TOKEN, \"Provide a read HF_TOKEN to load models from Hugging Face.\"\n",
@@ -34,6 +34,11 @@
        "\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_pytorch_mistral_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",
@@ -64,9 +69,13 @@
        "    - [mistralai/Mistral-7B-Instruct-v0.2](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2): improved instruction fine-tuned version of Mistral-7B-Instruct-v0.1 supporting 32k context length\n",
        "    - [mistralai/Mistral-7B-v0.3](https://huggingface.co/mistralai/Mistral-7B-v0.3): Mistral-7B-v0.2 with extended vocabulary of 32768 and supports function calling\n",
        "    - [mistralai/Mistral-7B-Instruct-v0.3](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.3): instruction fine-tuned version of the Mistral-7B-v0.3 generative text model\n",
-        "    - [mistralai/Mistral-Nemo-Base-2407](https://huggingface.co/mistralai/Mistral-Nemo-Base-2407):  pretrained generative text model of 12B parameters \n",
+        "    - [mistralai/Mistral-Nemo-Base-2407](https://huggingface.co/mistralai/Mistral-Nemo-Base-2407):  pretrained generative text model of 12B parameters\n",
        "    - [mistralai/Mistral-Nemo-Instruct-2407](https://huggingface.co/mistralai/Mistral-Nemo-Instruct-2407): instruct fine-tuned version of the Mistral-Nemo-Base-2407\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 +83,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) 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."
      ]
    },
    {
@@ -83,7 +92,7 @@
        "id": "Wi-aGs442yve"
      },
      "source": [
-        "## Run the notebook"
+        "## Before you begin"
      ]
    },
    {
@@ -99,94 +108,182 @@
        "\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 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.64.0'\n",
-        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.84.0'\n",
        "\n",
        "# Import the necessary packages.\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",
+        "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",
        "# 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",
-        "# 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, \"mistral\")\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\""
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cJJe9QRBAEUv"
+      },
+      "source": [
+        "## Deploy prebuilt Mistral model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "UMooIiHxAML0"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Select the model variants\n",
+        "\n",
+        "prebuilt_model_id = \"mistralai/Mistral-7B-Instruct-v0.3\"  # @param [\"mistralai/Mistral-7B-v0.1\", \"mistralai/Mistral-7B-Instruct-v0.1\", \"mistralai/Mistral-7B-Instruct-v0.2\", \"mistralai/Mistral-7B-v0.3\", \"mistralai/Mistral-7B-Instruct-v0.3\", \"mistralai/Mistral-Nemo-Base-2407\", \"mistralai/Mistral-Nemo-Instruct-2407\"]\n",
+        "model_id = f\"gs://vertex-model-garden-public-us/{prebuilt_model_id}\"\n",
+        "\n",
+        "version_id = prebuilt_model_id.split(\"/\")[1].lower()\n",
+        "PUBLISHER_MODEL_NAME = f\"publishers/mistral-ai/models/mistral@{version_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",
+        "# Find Vertex AI prediction supported accelerators and regions in\n",
+        "# https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
+        "\n",
+        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_V100\", \"NVIDIA_TESLA_T4\", \"NVIDIA_TESLA_A100\"]\n",
+        "\n",
+        "if prebuilt_model_id in [\n",
+        "    \"mistralai/Mistral-7B-v0.3\",\n",
+        "    \"mistralai/Mistral-7B-Instruct-v0.3\",\n",
+        "]:\n",
+        "    if accelerator_type == \"NVIDIA_L4\":\n",
+        "        machine_type = \"g2-standard-24\"\n",
+        "        accelerator_count = 2\n",
+        "    else:\n",
+        "        raise ValueError(f\"{prebuilt_model_id} supports only NVIDIA_L4\")\n",
+        "elif prebuilt_model_id == \"mistralai/Mistral-7B-v0.1\":\n",
+        "    if accelerator_type == \"NVIDIA_L4\":\n",
+        "        machine_type = \"g2-standard-12\"\n",
+        "        accelerator_count = 1\n",
+        "    elif accelerator_type == \"NVIDIA_TESLA_V100\":\n",
+        "        machine_type = \"n1-standard-16\"\n",
+        "        accelerator_count = 2\n",
+        "    elif accelerator_type == \"NVIDIA_TESLA_T4\":\n",
+        "        machine_type = \"n1-standard-32\"\n",
+        "        accelerator_count = 4\n",
+        "    else:\n",
+        "        raise ValueError(\n",
+        "            f\"{prebuilt_model_id} supports only NVIDIA_L4, NVIDIA_TESLA_V100, NVIDIA_TESLA_T4\"\n",
+        "        )\n",
+        "elif prebuilt_model_id in [\n",
+        "    \"mistralai/Mistral-7B-Instruct-v0.1\",\n",
+        "    \"mistralai/Mistral-7B-Instruct-v0.2\",\n",
+        "]:\n",
+        "    if accelerator_type == \"NVIDIA_L4\":\n",
+        "        machine_type = \"g2-standard-12\"\n",
+        "        accelerator_count = 1\n",
+        "    elif accelerator_type == \"NVIDIA_TESLA_V100\":\n",
+        "        machine_type = \"n1-standard-16\"\n",
+        "        accelerator_count = 2\n",
+        "    elif accelerator_type == \"NVIDIA_TESLA_T4\":\n",
+        "        machine_type = \"n1-standard-16\"\n",
+        "        accelerator_count = 2\n",
+        "    else:\n",
+        "        raise ValueError(\n",
+        "            f\"{prebuilt_model_id} supports only NVIDIA_L4, NVIDIA_TESLA_V100, NVIDIA_TESLA_T4\"\n",
+        "        )\n",
+        "elif prebuilt_model_id in [\n",
+        "    \"mistralai/Mistral-Nemo-Base-2407\",\n",
+        "    \"mistralai/Mistral-Nemo-Instruct-2407\",\n",
+        "]:\n",
+        "    if accelerator_type == \"NVIDIA_L4\":\n",
+        "        machine_type = \"g2-standard-48\"\n",
+        "        accelerator_count = 4\n",
+        "    else:\n",
+        "        raise ValueError(f\"{prebuilt_model_id} supports only NVIDIA_L4\")\n",
+        "else:\n",
+        "    raise ValueError(f\"{prebuilt_model_id} is not supported.\")\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details.\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": "W86V0UYUHkeH"
+      },
+      "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.preview 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",
+        ")"
      ]
    },
    {
@@ -198,49 +295,12 @@
      },
      "outputs": [],
      "source": [
-        "# @title Deploy\n",
+        "# @title [Option 2] Deploy with customized configs\n",
        "\n",
        "# @markdown This section deploys the prebuilt Mistral model with [vLLM](https://github.com/vllm-project/vllm) on a Vertex endpoint. It takes 15 minutes to 1 hour to finish depending on the model and the accelerator.\n",
        "\n",
-        "# @markdown Set the model to deploy.\n",
-        "\n",
-        "prebuilt_model_id = \"mistralai/Mistral-7B-Instruct-v0.3\"  # @param [\"mistralai/Mistral-7B-v0.1\", \"mistralai/Mistral-7B-Instruct-v0.1\", \"mistralai/Mistral-7B-Instruct-v0.2\", \"mistralai/Mistral-7B-v0.3\", \"mistralai/Mistral-7B-Instruct-v0.3\", \"mistralai/Mistral-Nemo-Base-2407\", \"mistralai/Mistral-Nemo-Instruct-2407\"]\n",
-        "model_id = f\"gs://vertex-model-garden-public-us/{prebuilt_model_id}\"\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:20240721_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",
-        "# Find Vertex AI prediction supported accelerators and regions in\n",
-        "# https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
-        "\n",
-        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_V100\", \"NVIDIA_TESLA_T4\", \"NVIDIA_TESLA_A100\"]\n",
-        "\n",
-        "if accelerator_type == \"NVIDIA_L4\":\n",
-        "    if \"Mistral-Nemo\" in prebuilt_model_id:\n",
-        "        machine_type = \"g2-standard-48\"\n",
-        "        accelerator_count = 4\n",
-        "    else:\n",
-        "        machine_type = \"g2-standard-8\"\n",
-        "        accelerator_count = 1\n",
-        "elif accelerator_type == \"NVIDIA_TESLA_V100\":\n",
-        "    machine_type = \"n1-standard-16\"\n",
-        "    accelerator_count = 2\n",
-        "elif accelerator_type == \"NVIDIA_TESLA_T4\":\n",
-        "    machine_type = \"n1-standard-16\"\n",
-        "    accelerator_count = 2\n",
-        "elif accelerator_type == \"NVIDIA_TESLA_A100\":\n",
-        "    machine_type = \"a2-highgpu-1g\"\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",
-        ")\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240910_0916_RC00\"\n",
        "\n",
        "# Larger setting of `max-model-len` can lead to higher requirements on\n",
        "# `gpu-memory-utilization` and GPU configuration. Larger setting of\n",
@@ -263,7 +323,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",
@@ -374,7 +433,6 @@
        "        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_pytorch_mistral_deployment.ipynb\",\n",
        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
@@ -385,12 +443,11 @@
        "    return model, endpoint\n",
        "\n",
        "\n",
-        "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
+        "models[LABEL], endpoints[LABEL] = deploy_model_vllm(\n",
        "    model_name=common_util.get_job_name_with_datetime(prefix=\"mistral-serve-vllm\"),\n",
        "    model_id=model_id,\n",
        "    publisher=\"mistral-ai\",\n",
        "    publisher_model_id=\"mistral\",\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
        "    base_model_id=prebuilt_model_id,\n",
        "    machine_type=machine_type,\n",
        "    accelerator_type=accelerator_type,\n",
@@ -399,7 +456,9 @@
        "    gpu_memory_utilization=gpu_memory_utilization,\n",
        "    dtype=dtype,\n",
        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
-        ")"
+        ")\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
      ]
    },
    {
@@ -407,7 +466,7 @@
      "execution_count": null,
      "metadata": {
        "cellView": "form",
-        "id": "NYN1Z49SJ-MM"
+        "id": "KZAcN3F2NiFv"
      },
      "outputs": [],
      "source": [
@@ -457,7 +516,7 @@
        "        \"raw_response\": raw_response,\n",
        "    },\n",
        "]\n",
-        "response = endpoints[\"vllm_gpu\"].predict(\n",
+        "response = endpoints[LABEL].predict(\n",
        "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
        ")\n",
        "\n",
@@ -479,7 +538,9 @@
        "#     headers={\"Content-Type\": \"application/json\"},\n",
        "#     use_dedicated_endpoint=use_dedicated_endpoint,\n",
        "# )\n",
-        "# print(response.json())"
+        "# print(response.json())\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
      ]
    },
    {
@@ -501,11 +562,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()"
      ]
    }
  ],
@@ -633,8 +633,8 @@
        "    is_for_training=False,\n",
        ")\n",
        "\n",
-        "# Dedicated endpoint not supported yet.\n",
-        "use_dedicated_endpoint = False\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",
        "gpu_memory_utilization = 0.85\n",
        "max_model_len = 8192  # Maximum context length.\n",
@@ -230,7 +230,7 @@
        "# 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:20240721_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",
+        "# @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",
        "# Find Vertex AI prediction supported accelerators and regions in\n",
        "# https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
@@ -739,9 +739,7 @@
        "    DEDICATED_ENDPOINT_DNS = endpoints[\n",
        "        \"optimized_vllm_gpu\"\n",
        "    ].gca_resource.dedicated_endpoint_dns\n",
-        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
-        "    PROJECT_ID, REGION, endpoints[\"optimized_vllm_gpu\"].name\n",
-        ")\n",
+        "ENDPOINT_RESOURCE_NAME = endpoints[\"optimized_vllm_gpu\"].resource_name\n",
        "\n",
        "# @title Chat Completions Inference\n",
        "\n",
@@ -638,8 +638,8 @@
        "    is_for_training=False,\n",
        ")\n",
        "\n",
-        "# Dedicated endpoint not supported yet.\n",
-        "use_dedicated_endpoint = False\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",
        "gpu_memory_utilization = 0.85\n",
        "max_model_len = 8192  # Maximum context length.\n",
@@ -90,107 +90,11 @@
      },
      "outputs": [],
      "source": [
-        "# @title Setup Google Cloud project\n",
+        "# @title Request for quota\n",
        "\n",
-        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "# @markdown By default, the quota for A100_80GB and H100 deployment `Custom model serving per region` is 0. You need to request quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
        "\n",
-        "# @markdown 2. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
-        "\n",
-        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# @markdown 3. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
-        "\n",
-        "REGION = \"\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# @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, europe-west4, us-west1, asia-southeast1 |\n",
-        "\n",
-        "# Import the necessary packages\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 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",
-        "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",
-        "# 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, \"Prompt-Guard-86M\")\n",
-        "\n",
-        "\n",
-        "# Initialize Vertex AI API.\n",
-        "print(\"Initializing Vertex AI API.\")\n",
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
-        "\n",
-        "# Gets the default SERVICE_ACCOUNT.\n",
-        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
-        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
-        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
-        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
-        "\n",
-        "\n",
-        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
-        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
-        "\n",
-        "! gcloud config set project $PROJECT_ID\n",
-        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/storage.admin\"\n",
-        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/aiplatform.user\"\n",
-        "\n",
-        "models, endpoints = {}, {}"
+        "# @markdown For better chance to get resources, we recommend to request A100_80GB quota in the regions `us-central1, us-east1`, and request H100 quota in the regions `us-central1, us-west1`."
      ]
    },
    {
@@ -202,57 +106,87 @@
      },
      "outputs": [],
      "source": [
-        "# @title Access Prompt Guard\n",
+        "# @title Setup Google Cloud project\n",
        "\n",
-        "# @markdown For GPU based serving, choose between accessing the Prompt Guard model on [Hugging Face](https://huggingface.co/)\n",
-        "# @markdown or Vertex AI as described below.\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 If you already obtained access to Prompt Guard on [Hugging Face](https://huggingface.co/), you can load the model from there.\n",
-        "# @markdown Alternatively, you can also load the original Prompt Guard model for serving from Vertex AI after accepting the agreement.\n",
+        "# @markdown 2. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
        "\n",
-        "# @markdown **Only select and fill one of the following sections.**\n",
-        "# fmt: off\n",
-        "LOAD_MODEL_FROM = \"Hugging Face\"  # @param [\"Hugging Face\", \"Google Cloud\"] {isTemplate:true}\n",
-        "# fmt: on\n",
+        "REGION = \"\"  # @param {type:\"string\"}\n",
        "\n",
-        "# @markdown ---\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 ### Access Prompt Guard on Hugging Face for GPU based serving\n",
-        "# @markdown You must provide a Hugging Face User Access Token (read) to access Prompt Guard. 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 > | 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",
-        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate:true}\n",
+        "# Import the necessary packages\n",
        "\n",
-        "# @markdown *--- Or ---*\n",
-        "# @markdown ### Access Prompt Guard on Vertex AI for GPU based serving\n",
-        "# @markdown The original model from Meta is converted into the Hugging Face format for serving in Vertex AI.\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.84.0'\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
+        "import importlib\n",
+        "import os\n",
+        "import re\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 = \"pytorch_inference_gpu\"\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",
+        "# 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",
+        ")\n",
+        "\n",
+        "# @markdown # Access Llama Prompt 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",
        "# @markdown Accept the model agreement to access the models:\n",
-        "# @markdown 1. Open the [Prompt Guard model card](https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/prompt-guard) from [Vertex AI Model Garden](https://cloud.google.com/model-garden).\n",
+        "# @markdown 1. Open the [Llama Prompt Guard model card](https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/prompt-guard) from [Vertex AI Model Garden](https://cloud.google.com/model-garden).\n",
        "# @markdown 2. Review and accept the agreement in the pop-up window on the model card page. If you have previously accepted the model agreement, there will not be a pop-up window on the model card page and this step is not needed.\n",
-        "# @markdown 3. After accepting the agreement of Prompt Guard, a `gs://` URI containing the Prompt Guard model artifacts will be shared.\n",
+        "# @markdown 3. After accepting the agreement, a `gs://` URI containing Llama Prompt Guard models will be shared.\n",
        "# @markdown 4. Paste the URI in the `VERTEX_AI_MODEL_GARDEN_PROMPT_GUARD` field below.\n",
+        "# @markdown 5. The Llama Prompt Guard models will be copied into `BUCKET_URI`.\n",
+        "\n",
        "\n",
        "VERTEX_AI_MODEL_GARDEN_PROMPT_GUARD = \"\"  # @param {type:\"string\", isTemplate:true}\n",
-        "if LOAD_MODEL_FROM == \"Hugging Face\":\n",
-        "    assert (\n",
-        "        HF_TOKEN\n",
-        "    ), \"Provide a read HF_TOKEN to load models from Hugging Face, or select a different model source.\"\n",
-        "else:\n",
-        "    assert (\n",
-        "        VERTEX_AI_MODEL_GARDEN_PROMPT_GUARD\n",
-        "    ), \"Click the agreement of Prompt Guard in Vertex AI Model Garden via https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/prompt-guard, and get the GCS path of Prompt Guard model artifacts.\"\n",
-        "    parsed_gcs_url = re.search(\"gs://.*?(?=[ ]|$)\", VERTEX_AI_MODEL_GARDEN_PROMPT_GUARD)\n",
-        "    if parsed_gcs_url:\n",
-        "        VERTEX_AI_MODEL_GARDEN_PROMPT_GUARD = parsed_gcs_url.group()\n",
-        "    assert VERTEX_AI_MODEL_GARDEN_PROMPT_GUARD.startswith(\n",
-        "        \"gs://\"\n",
-        "    ), \"VERTEX_AI_MODEL_GARDEN_PROMPT_GUARD is expected to be a GCS URI and must start with `gs://`.\"\n",
-        "    print(\n",
-        "        \"Copying Prompt Guard model artifacts from\",\n",
-        "        VERTEX_AI_MODEL_GARDEN_PROMPT_GUARD,\n",
-        "        \"to \",\n",
-        "        MODEL_BUCKET,\n",
-        "    )\n",
-        "    ! gsutil -m cp -R $VERTEX_AI_MODEL_GARDEN_PROMPT_GUARD/* $MODEL_BUCKET\n",
+        "assert (\n",
+        "    VERTEX_AI_MODEL_GARDEN_PROMPT_GUARD\n",
+        "), \"Click the agreement in Vertex AI Model Garden at https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/prompt-guard, and get the GCS path of Llama Prompt Guard model artifacts.\"\n",
+        "parsed_gcs_url = re.search(\"gs://.*?(?=[ ]|$)\", VERTEX_AI_MODEL_GARDEN_PROMPT_GUARD)\n",
+        "if parsed_gcs_url:\n",
+        "    VERTEX_AI_MODEL_GARDEN_PROMPT_GUARD = parsed_gcs_url.group()\n",
+        "assert VERTEX_AI_MODEL_GARDEN_PROMPT_GUARD.startswith(\n",
+        "    \"gs://\"\n",
+        "), \"VERTEX_AI_MODEL_GARDEN_PROMPT_GUARD is expected to be a GCS URI and must start with `gs://`.\"\n",
        "\n",
        "# @markdown ---\n",
        "# @markdown Click \"Show Code\" to see more details."
@@ -277,7 +211,7 @@
      "outputs": [],
      "source": [
        "# @title Deploy\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",
        "# @markdown This section uploads Prompt Guard to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish.\n",
@@ -285,14 +219,15 @@
        "# @markdown NVIDIA_L4 GPUs are used for demonstration. The serving efficiency of L4 GPUs is inferior to that of A100 GPUs, but L4 GPUs are nevertheless good serving solutions if you do not have A100 quota.\n",
        "\n",
        "# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
-        "\n",
-        "if LOAD_MODEL_FROM == \"Hugging Face\":\n",
-        "    model_id = \"meta-llama/Prompt-Guard-86M\"\n",
-        "else:\n",
-        "    model_id = MODEL_BUCKET\n",
+        "base_model_name = \"Llama-Prompt-Guard-2-22M\"  # @param [\"Prompt-Guard-86M\", \"Llama-Prompt-Guard-2-22M\", \"Llama-Prompt-Guard-2-86M\"] {allow-input: true, isTemplate: true}\n",
+        "model_id = os.path.join(VERTEX_AI_MODEL_GARDEN_PROMPT_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/prompt-guard@{version_id}\"\n",
        "\n",
        "# The pre-built serving docker image.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-pytorch-inference-cu121.2-2.transformers.4-41.ubuntu2204.py311\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-pytorch-inference-cu121.2-3.transformers.4-46.ubuntu2204.py311\"\n",
+        "SERVICE_ACCOUNT = None\n",
        "\n",
        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_A100_80GB\"]\n",
        "\n",
@@ -387,9 +322,9 @@
        "    return model, endpoint\n",
        "\n",
        "\n",
-        "models[\"pytorch_inference_gpu\"], endpoints[\"pytorch_inference_gpu\"] = deploy_model(\n",
+        "models[LABEL], endpoints[LABEL] = deploy_model(\n",
        "    model_name=common_util.get_job_name_with_datetime(prefix=\"prompt-guard-serve\"),\n",
-        "    model_id=model_id,\n",
+        "    model_id=model_id if is_gcs_path(model_id) else hf_model_id,\n",
        "    publisher=\"meta\",\n",
        "    publisher_model_id=\"prompt-guard\",\n",
        "    task=task,\n",
@@ -62,6 +62,10 @@
        "- Download and deploy prebuilt Qwen2/Qwen2.5 models\n",
        "- Deploy Qwen2/Qwen2.5 with [vLLM](https://github.com/vllm-project/vllm) to improve serving throughput\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",
@@ -94,15 +98,11 @@
        "\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). 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",
+        "# @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",
@@ -112,84 +112,47 @@
        "# @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",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.84.0'\n",
        "\n",
        "# Import the necessary packages\n",
-        "import datetime\n",
+        "\n",
        "import importlib\n",
        "import os\n",
-        "import uuid\n",
        "from typing import Tuple\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "\n",
-        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
-        "\n",
        "models, endpoints = {}, {}\n",
+        "LABEL = \"vllm_gpu\"\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",
-        "# 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, \"qwen2\")\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",
+        ")"
      ]
    },
    {
@@ -211,20 +174,18 @@
      },
      "outputs": [],
      "source": [
-        "# @title Deploy\n",
-        "\n",
-        "# @markdown This section uploads prebuilt Qwen2/Qwen2.5 models to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 to 30 minutes to finish depending on the size of the model.\n",
-        "\n",
        "MODEL_ID = \"Qwen2.5-0.5B-Instruct\"  # @param [\"Qwen2-0.5B-Instruct\", \"Qwen2-1.5B-Instruct\", \"Qwen2-7B-Instruct\", \"Qwen2.5-0.5B-Instruct\", \"Qwen2.5-1.5B-Instruct\", \"Qwen2.5-7B-Instruct\", \"Qwen2.5-14B-Instruct\", \"Qwen2.5-32B-Instruct\"] {isTemplate: true}\n",
        "model_path_prefix = \"Qwen\"\n",
        "model_id = os.path.join(model_path_prefix, MODEL_ID)\n",
        "\n",
-        "# The pre-built serving docker image for vLLM.\n",
+        "PUBLISHER_MODEL_NAME = f\"publishers/qwen/models/qwen2@{MODEL_ID}\"\n",
+        "\n",
        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241008_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",
+        "# @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",
-        "accelerator_type = \"NVIDIA_A100_80GB\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_A100\", \"NVIDIA_A100_80GB\"] {isTemplate: true}\n",
+        "\n",
+        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_A100\", \"NVIDIA_A100_80GB\"] {isTemplate: true}\n",
        "vllm_dtype = \"bfloat16\"\n",
        "max_model_len = 32768  # Maximum context length.\n",
        "gpu_memory_utilization = 0.85\n",
@@ -298,6 +259,45 @@
        "else:\n",
        "    raise ValueError(\"Invalid model id: %s\" % MODEL_ID)\n",
        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "Y8_2LKmZPHQG"
+      },
+      "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.preview 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",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "GxUAAhJKPHQG"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with customized configs\n",
+        "\n",
+        "# @markdown This section uploads prebuilt Qwen2/Qwen2.5 models to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 to 30 minutes to finish depending on the size of the model.\n",
        "\n",
        "common_util.check_quota(\n",
        "    project_id=PROJECT_ID,\n",
@@ -313,7 +313,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",
@@ -424,7 +423,6 @@
        "        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_pytorch_qwen2_deployment.ipynb\",\n",
        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
@@ -440,7 +438,6 @@
        "    model_id=model_id,\n",
        "    publisher=\"qwen\",\n",
        "    publisher_model_id=\"qwen2\",\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
        "    machine_type=machine_type,\n",
        "    accelerator_type=accelerator_type,\n",
        "    accelerator_count=accelerator_count,\n",
@@ -448,9 +445,7 @@
        "    gpu_memory_utilization=gpu_memory_utilization,\n",
        "    dtype=vllm_dtype,\n",
        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
-        ")\n",
-        "\n",
-        "# @markdown Click \"Show Code\" to see more details."
+        ")"
      ]
    },
    {
@@ -536,16 +531,14 @@
      },
      "outputs": [],
      "source": [
-        "# @title Deploy\n",
-        "\n",
-        "# @markdown This section uploads prebuilt Qwen2 & Qwen2.5 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",
        "# @markdown Select one of the four model variations.\n",
        "MODEL_ID = \"Qwen2.5-0.5B-Instruct\"  # @param [\"Qwen2-0.5B-Instruct\", \"Qwen2-1.5B-Instruct\", \"Qwen2-7B-Instruct\", \"Qwen2.5-0.5B-Instruct\", \"Qwen2.5-1.5B-Instruct\", \"Qwen2.5-7B-Instruct\", \"Qwen2.5-14B-Instruct\", \"Qwen2.5-32B-Instruct-AWQ\"] {isTemplate: true}\n",
        "TPU_DEPLOYMENT_REGION = \"us-west1\"  # @param [\"us-west1\"] {isTemplate:true}\n",
        "model_path_prefix = \"Qwen\"\n",
        "model_id = os.path.join(model_path_prefix, MODEL_ID)\n",
        "\n",
+        "PUBLISHER_MODEL_NAME = f\"publishers/qwen/models/qwen2@{MODEL_ID}\"\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",
@@ -584,7 +577,21 @@
        "\n",
        "# Endpoint configurations.\n",
        "min_replica_count = 1\n",
-        "max_replica_count = 1\n",
+        "max_replica_count = 1"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "ZOqzKtyrPHQH"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "\n",
+        "# @markdown This section uploads prebuilt Qwen2 & Qwen2.5 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",
        "\n",
        "def deploy_model_hexllm(\n",
@@ -592,7 +599,6 @@
        "    model_id: str,\n",
        "    publisher: str,\n",
        "    publisher_model_id: 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",
@@ -681,7 +687,6 @@
        "        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",
@@ -697,7 +702,6 @@
        "    model_id=model_id,\n",
        "    publisher=\"qwen\",\n",
        "    publisher_model_id=\"qwen2\",\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
        "    tensor_parallel_size=tpu_count,\n",
        "    machine_type=machine_type,\n",
        "    tpu_topology=tpu_topo,\n",
@@ -802,11 +806,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()"
      ]
    }
  ],
@@ -0,0 +1,782 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "SgQ6t5bqZVlH"
+      },
+      "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 - Qwen3 (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_pytorch_qwen3_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_pytorch_qwen3_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",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3de7470326a2"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates serving Qwen3 models with [SGLang](https://github.com/sgl-project/sglang). [Qwen3](https://huggingface.co/collections/Qwen/qwen3-67dd247413f0e2e4f653967f) is the latest generation of large language models in Qwen series, offering a comprehensive suite of dense and mixture-of-experts (MoE) models. Built upon extensive training, Qwen3 delivers groundbreaking advancements in reasoning, instruction-following, agent capabilities, and multilingual support\n",
+        "\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Deploy Qwen3 with SGLang on GPU using single-host serving, and [Spot VMs](https://cloud.google.com/compute/docs/instances/spot) (Optional).\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": "ax7zWynUDcjk"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Request for quota\n",
+        "\n",
+        "# @markdown To deploy the largest variant of the Qwen3 models, you need 1 host of 8 x H100 machine. Check that you have sufficient quota:\n",
+        "# @markdown - For Spot VM quota, check [`CustomModelServingPreemptibleH100GPUsPerProjectPerRegion`](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_preemptible_nvidia_h100_gpus).\n",
+        "# @markdown - For regular VM quota, check [`CustomModelServingH100GPUsPerProjectPerRegion`](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
+        "#\n",
+        "# @markdown If you don't have sufficient quota, request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "YXFGIp1l-qtT"
+      },
+      "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.84.0'\n",
+        "\n",
+        "# Import the necessary packages\n",
+        "import importlib\n",
+        "import os\n",
+        "import time\n",
+        "from typing import Tuple\n",
+        "\n",
+        "import requests\n",
+        "from google import auth\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",
+        "\n",
+        "def check_quota(\n",
+        "    project_id: str,\n",
+        "    region: str,\n",
+        "    resource_id: str,\n",
+        "    accelerator_count: int,\n",
+        "):\n",
+        "    \"\"\"Checks if the project and the region has the required quota.\"\"\"\n",
+        "    quota = common_util.get_quota(project_id, region, resource_id)\n",
+        "    quota_request_instruction = (\n",
+        "        \"Either use \"\n",
+        "        \"a different region or request additional quota. Follow \"\n",
+        "        \"instructions here \"\n",
+        "        \"https://cloud.google.com/docs/quotas/view-manage#requesting_higher_quota\"\n",
+        "        \" to check quota in a region or request additional quota for \"\n",
+        "        \"your project.\"\n",
+        "    )\n",
+        "    if quota == -1:\n",
+        "        raise ValueError(\n",
+        "            f\"Quota not found for: {resource_id} in {region}.\"\n",
+        "            f\" {quota_request_instruction}\"\n",
+        "        )\n",
+        "    if quota < accelerator_count:\n",
+        "        raise ValueError(\n",
+        "            f\"Quota not enough for {resource_id} in {region}: {quota} <\"\n",
+        "            f\" {accelerator_count}. {quota_request_instruction}\"\n",
+        "        )\n",
+        "\n",
+        "\n",
+        "LABEL = \"sglang_gpu\"\n",
+        "models, endpoints = {}, {}\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "if not REGION:\n",
+        "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "\n",
+        "# 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",
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3zJJDmldn7rw"
+      },
+      "source": [
+        "## Deploy Qwen3 with SGLang"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "_3Swj3pxn7rw"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Select the model variants\n",
+        "\n",
+        "# @markdown Set the model to deploy.\n",
+        "\n",
+        "base_model_name = \"Qwen3-235B-A22B\"  # @param [\"Qwen3-235B-A22B\", \"Qwen3-235B-A22B-FP8\", \"Qwen3-30B-A3B\", \"Qwen3-30B-A3B-Base\", \"Qwen/Qwen3-30B-A3B-FP8\", \"Qwen3-32B\", \"Qwen/Qwen3-32B-FP8\", \"Qwen3-14B\", \"Qwen3-14B-Base\", \"Qwen/Qwen3-14B-FP8\", \"Qwen3-8B\", \"Qwen3-8B-Base\", \"Qwen/Qwen3-8B-FP8\", \"Qwen3-4B\", \"Qwen3-4B-Base\", \"Qwen/Qwen3-4B-FP8\", \"Qwen3-1.7B\", \"Qwen3-1.7B-Base\", \"Qwen/Qwen3-1.7B-FP8\", \"Qwen3-0.6B\", \"Qwen3-0.6B-Base\", \"Qwen/Qwen3-0.6B-FP8\"] {isTemplate:true}\n",
+        "model_id = \"Qwen/\" + base_model_name\n",
+        "hf_model_id = model_id\n",
+        "\n",
+        "# The pre-built serving docker images.\n",
+        "SGLANG_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/sglang-serve.cu124.0-4.ubuntu2204.py310:20250428-1803-rc0\"\n",
+        "\n",
+        "# @markdown Choose whether to use a [Spot VM](https://cloud.google.com/compute/docs/instances/spot) for the deployment.\n",
+        "is_spot = False  # @param {type:\"boolean\"}\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_H100_80GB\"  # @param [\"NVIDIA_H100_80GB\", \"NVIDIA_L4\"] {isTemplate:true}\n",
+        "\n",
+        "PUBLISHER_MODEL_NAME = f\"publishers/qwen/models/qwen3@{base_model_name.lower()}\"\n",
+        "\n",
+        "if accelerator_type == \"NVIDIA_H100_80GB\":\n",
+        "    if is_spot:\n",
+        "        resource_id = \"custom_model_serving_preemptible_nvidia_h100_gpus\"\n",
+        "    else:\n",
+        "        resource_id = \"custom_model_serving_nvidia_h100_gpus\"\n",
+        "    if base_model_name in [\"Qwen3-235B-A22B\"]:\n",
+        "        machine_type = \"a3-highgpu-8g\"\n",
+        "        accelerator_count = 8\n",
+        "        model_id = \"gs://vertex-model-garden-restricted-us/qwen3/Qwen3-235B-A22B\"\n",
+        "    elif base_model_name in [\"Qwen3-235B-A22B-FP8\"]:\n",
+        "        machine_type = \"a3-highgpu-4g\"\n",
+        "        accelerator_count = 4\n",
+        "    else:\n",
+        "        machine_type = \"a3-highgpu-1g\"\n",
+        "        accelerator_count = 1\n",
+        "elif accelerator_type == \"NVIDIA_L4\":\n",
+        "    if is_spot:\n",
+        "        resource_id = \"custom_model_serving_preemptible_nvidia_l4_gpus\"\n",
+        "    else:\n",
+        "        resource_id = \"custom_model_serving_nvidia_l4_gpus\"\n",
+        "    if base_model_name in [\"Qwen3-235B-A22B\", \"Qwen3-235B-A22B-FP8\"]:\n",
+        "        raise ValueError(\"L4s are insufficient to server Qwen3-235B-A22B.\")\n",
+        "    elif base_model_name in [\"Qwen3-30B-A3B\", \"Qwen3-30B-A3B-Base\", \"Qwen3-32B\"]:\n",
+        "        machine_type = \"g2-standard-48\"\n",
+        "        accelerator_count = 4\n",
+        "    elif base_model_name in [\n",
+        "        \"Qwen3-14B\",\n",
+        "        \"Qwen3-14B-Base\",\n",
+        "        \"Qwen3-30B-A3B-FP8\",\n",
+        "        \"Qwen3-32B-FP8\",\n",
+        "    ]:\n",
+        "        machine_type = \"g2-standard-24\"\n",
+        "        accelerator_count = 2\n",
+        "    elif base_model_name in [\n",
+        "        \"Qwen3-8B\",\n",
+        "        \"Qwen3-8B-Base\",\n",
+        "        \"Qwen3-4B\",\n",
+        "        \"Qwen3-4B-Base\",\n",
+        "        \"Qwen3-1.7B\",\n",
+        "        \"Qwen3-1.7B-Base\",\n",
+        "        \"Qwen3-0.6B\",\n",
+        "        \"Qwen3-0.6B-Base\",\n",
+        "        \"Qwen/Qwen3-14B-FP8\",\n",
+        "        \"Qwen/Qwen3-8B-FP8\",\n",
+        "        \"Qwen/Qwen3-4B-FP8\",\n",
+        "        \"Qwen/Qwen3-1.7B-FP8\",\n",
+        "        \"Qwen/Qwen3-0.6B-FP8\",\n",
+        "    ]:\n",
+        "        machine_type = \"g2-standard-12\"\n",
+        "        accelerator_count = 1\n",
+        "    else:\n",
+        "        raise ValueError(f\"Recommended GPU setting not found for: {base_model_name}.\")\n",
+        "else:\n",
+        "    raise ValueError(f\"Recommended GPU setting not found for: {base_model_name}.\")\n",
+        "\n",
+        "check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    resource_id=resource_id,\n",
+        "    accelerator_count=accelerator_count,\n",
+        ")\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "omW0LaC8wWz5"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 1] Deploy with Model Garden SDK\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",
+        "deploy_request_timeout = 1800  # 30 minutes\n",
+        "from vertexai.preview 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",
+        "    spot=is_spot,\n",
+        "    deploy_request_timeout=deploy_request_timeout,\n",
+        "    accept_eula=False,\n",
+        ")\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "3m-tDxgawYhU"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with customized configs\n",
+        "\n",
+        "# @markdown This section uploads Qwen3 models to Model Registry and deploys them to a Vertex Prediction Endpoint. It takes ~1 hour to finish.\n",
+        "\n",
+        "# @markdown It's recommended to use the region selected by the deployment button on the model card. If the deployment button is not available, it's recommended to stay with the default region of the notebook.\n",
+        "\n",
+        "\n",
+        "def poll_operation(op_name: str) -> bool:  # noqa: F811\n",
+        "    creds, _ = auth.default()\n",
+        "    auth_req = auth.transport.requests.Request()\n",
+        "    creds.refresh(auth_req)\n",
+        "    headers = {\n",
+        "        \"Authorization\": f\"Bearer {creds.token}\",\n",
+        "    }\n",
+        "    get_resp = requests.get(\n",
+        "        f\"https://{REGION}-aiplatform.googleapis.com/ui/{op_name}\",\n",
+        "        headers=headers,\n",
+        "    )\n",
+        "    opjs = get_resp.json()\n",
+        "    if \"error\" in opjs:\n",
+        "        raise ValueError(f\"Operation failed: {opjs['error']}\")\n",
+        "    return opjs.get(\"done\", False)\n",
+        "\n",
+        "\n",
+        "def poll_and_wait(op_name: str, total_wait: int, interval: int = 60):  # noqa: F811\n",
+        "    waited = 0\n",
+        "    while not poll_operation(op_name):\n",
+        "        if waited > total_wait:\n",
+        "            raise TimeoutError(\"Operation timed out\")\n",
+        "        print(\n",
+        "            f\"\\rStill waiting for operation... Waited time in second: {waited:<6}\",\n",
+        "            end=\"\",\n",
+        "            flush=True,\n",
+        "        )\n",
+        "        waited += interval\n",
+        "        time.sleep(interval)\n",
+        "\n",
+        "\n",
+        "def deploy_model_sglang_multihost(\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 = \"\",\n",
+        "    machine_type: str = \"g2-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 1,\n",
+        "    multihost_gpu_node_count: int = 1,\n",
+        "    gpu_memory_utilization: float | None = None,\n",
+        "    context_length: int | None = None,\n",
+        "    dtype: str | None = None,\n",
+        "    enable_trust_remote_code: bool = False,\n",
+        "    enable_torch_compile: bool = False,\n",
+        "    torch_compile_max_bs: int | None = None,\n",
+        "    attention_backend: str = \"\",\n",
+        "    enable_flashinfer_mla: bool = False,\n",
+        "    disable_cuda_graph: bool = False,\n",
+        "    speculative_algorithm: str | None = None,\n",
+        "    speculative_draft_model_path: str = \"\",\n",
+        "    speculative_num_steps: int = 3,\n",
+        "    speculative_eagle_topk: int = 1,\n",
+        "    speculative_num_draft_tokens: int = 4,\n",
+        "    enable_jit_deepgemm: bool = False,\n",
+        "    enable_dp_attention: bool = False,\n",
+        "    dp_size: int = 1,\n",
+        "    use_dedicated_endpoint: bool = False,\n",
+        "    max_num_seqs: int | None = None,\n",
+        "    is_spot: bool = True,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys trained models with SGLang 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.sglang.ai/backend/server_arguments.html for a list of possible arguments with descriptions.\n",
+        "    sglang_args = [\n",
+        "        f\"--model={model_id}\",\n",
+        "        f\"--tp={accelerator_count * multihost_gpu_node_count}\",\n",
+        "        f\"--dp={dp_size}\",\n",
+        "    ]\n",
+        "\n",
+        "    if context_length:\n",
+        "        sglang_args.append(f\"--context-length={context_length}\")\n",
+        "\n",
+        "    if gpu_memory_utilization:\n",
+        "        sglang_args.append(f\"--mem-fraction-static={gpu_memory_utilization}\")\n",
+        "\n",
+        "    if max_num_seqs:\n",
+        "        sglang_args.append(f\"--max-running-requests={max_num_seqs}\")\n",
+        "\n",
+        "    if dtype:\n",
+        "        sglang_args.append(f\"--dtype={dtype}\")\n",
+        "\n",
+        "    if enable_trust_remote_code:\n",
+        "        sglang_args.append(\"--trust-remote-code\")\n",
+        "\n",
+        "    if enable_torch_compile:\n",
+        "        sglang_args.append(\"--enable-torch-compile\")\n",
+        "        if torch_compile_max_bs:\n",
+        "            sglang_args.append(f\"--torch-compile-max-bs={torch_compile_max_bs}\")\n",
+        "\n",
+        "    if attention_backend:\n",
+        "        sglang_args.append(f\"--attention-backend={attention_backend}\")\n",
+        "\n",
+        "    if enable_flashinfer_mla:\n",
+        "        sglang_args.append(\"--enable-flashinfer-mla\")\n",
+        "\n",
+        "    if disable_cuda_graph:\n",
+        "        sglang_args.append(\"--disable-cuda-graph\")\n",
+        "\n",
+        "    if speculative_algorithm:\n",
+        "        sglang_args.append(f\"--speculative-algorithm={speculative_algorithm}\")\n",
+        "        sglang_args.append(\n",
+        "            f\"--speculative-draft-model-path={speculative_draft_model_path}\"\n",
+        "        )\n",
+        "        sglang_args.append(f\"--speculative-num-steps={speculative_num_steps}\")\n",
+        "        sglang_args.append(f\"--speculative-eagle-topk={speculative_eagle_topk}\")\n",
+        "        sglang_args.append(\n",
+        "            f\"--speculative-num-draft-tokens={speculative_num_draft_tokens}\"\n",
+        "        )\n",
+        "\n",
+        "    if enable_dp_attention:\n",
+        "        sglang_args.append(\"--enable-dp-attention\")\n",
+        "\n",
+        "    env_vars = {\n",
+        "        \"MODEL_ID\": base_model_id,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "\n",
+        "    if enable_jit_deepgemm:\n",
+        "        env_vars[\"SGL_ENABLE_JIT_DEEPGEMM\"] = \"1\"\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=SGLANG_DOCKER_URI,\n",
+        "        serving_container_args=sglang_args,\n",
+        "        serving_container_ports=[30000],\n",
+        "        serving_container_predict_route=\"/vertex_generate\",\n",
+        "        serving_container_health_route=\"/health\",\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 {int(accelerator_count * multihost_gpu_node_count)} {accelerator_type} GPU(s).\"\n",
+        "    )\n",
+        "\n",
+        "    creds, _ = auth.default()\n",
+        "    auth_req = auth.transport.requests.Request()\n",
+        "    creds.refresh(auth_req)\n",
+        "\n",
+        "    url = f\"https://{REGION}-aiplatform.googleapis.com/ui/projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint.name}:deployModel\"\n",
+        "    headers = {\n",
+        "        \"Content-Type\": \"application/json\",\n",
+        "        \"Authorization\": f\"Bearer {creds.token}\",\n",
+        "    }\n",
+        "    data = {\n",
+        "        \"deployedModel\": {\n",
+        "            \"model\": model.resource_name,\n",
+        "            \"displayName\": model_name,\n",
+        "            \"dedicatedResources\": {\n",
+        "                \"machineSpec\": {\n",
+        "                    \"machineType\": machine_type,\n",
+        "                    \"multihostGpuNodeCount\": multihost_gpu_node_count,\n",
+        "                    \"acceleratorType\": accelerator_type,\n",
+        "                    \"acceleratorCount\": accelerator_count,\n",
+        "                },\n",
+        "                \"minReplicaCount\": 1,\n",
+        "                \"maxReplicaCount\": 1,\n",
+        "            },\n",
+        "            \"system_labels\": {\n",
+        "                \"NOTEBOOK_NAME\": \"model_garden_pytorch_qwen3_deployment.ipynb\",\n",
+        "                \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
+        "            },\n",
+        "        },\n",
+        "    }\n",
+        "    if service_account:\n",
+        "        data[\"deployedModel\"][\"serviceAccount\"] = service_account\n",
+        "    if is_spot:\n",
+        "        data[\"deployedModel\"][\"dedicatedResources\"][\"spot\"] = True\n",
+        "    response = requests.post(url, headers=headers, json=data)\n",
+        "    print(f\"Deploy Model response: {response.json()}\")\n",
+        "    if response.status_code != 200 or \"name\" not in response.json():\n",
+        "        raise ValueError(f\"Failed to deploy model: {response.text}\")\n",
+        "    poll_and_wait(response.json()[\"name\"], 7200)\n",
+        "    print(\"endpoint_name:\", endpoint.name)\n",
+        "\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "models[LABEL], endpoints[LABEL] = deploy_model_sglang_multihost(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=\"qwen3-serve\"),\n",
+        "    model_id=model_id,\n",
+        "    publisher=\"qwen\",\n",
+        "    publisher_model_id=\"qwen3\",\n",
+        "    base_model_id=hf_model_id,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    is_spot=is_spot,\n",
+        ")\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "AGVPzwHkn7rw"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Raw predict\n",
+        "\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. Sampling parameters supported by SGLang can be found [here](https://docs.sglang.ai/backend/sampling_params.html).\n",
+        "\n",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown User: What is the best way to diagnose and fix a flickering light in my house?\n",
+        "# @markdown Assistant: Okay, the user is asking about the best way to diagnose and fix a flickering light in their house. Let me start by breaking down the steps. First, I need to consider the possible causes of flickering. Common issues could be a faulty light bulb, a blown fuse, a short circuit, or a problem with the circuit itself.\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",
+        "# @markdown A chat template formatted prompt for Qwen3 models are shown below as an example.\n",
+        "prompt = \"<|im_start|>user What is the best way to diagnose and fix a flickering light in my house?<|im_end|><|im_start|>assistant\"  # @param {type: \"string\"}\n",
+        "\n",
+        "# @markdown By default, Qwen3 has thinking capabilities enabled, similar to QwQ-32B. This means the model will use its reasoning abilities to enhance the quality of generated responses.\n",
+        "# @markdown The model will generate think content wrapped in a \\<think>...\\</think> block, followed by the final response.\n",
+        "# @markdown `max_new_tokens` may need to be increased to accommodate the additional think content.\n",
+        "enable_thinking = True  # @param {type:\"boolean\"}\n",
+        "if not enable_thinking:\n",
+        "    prompt += \"<think></think>\"\n",
+        "\n",
+        "\n",
+        "max_new_tokens = 1024  # @param {type:\"integer\"}\n",
+        "\n",
+        "# @markdown Sampling parameters such as `temperature` and `top_p` are automatically set according to [Qwen3's best practices](https://huggingface.co/Qwen/Qwen3-30B-A3B#best-practices).\n",
+        "if enable_thinking:\n",
+        "    temperature = 0.6\n",
+        "    top_p = 0.95\n",
+        "    top_k = 20\n",
+        "    min_p = 0\n",
+        "else:\n",
+        "    temperature = 0.7\n",
+        "    top_p = 0.8\n",
+        "    top_k = 20\n",
+        "    min_p = 0\n",
+        "\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_new_tokens`.\n",
+        "\n",
+        "\n",
+        "# Overrides parameters for inferences.\n",
+        "instances = [{\"text\": prompt}]\n",
+        "parameters = {\n",
+        "    \"sampling_params\": {\n",
+        "        \"max_new_tokens\": max_new_tokens,\n",
+        "        \"temperature\": temperature,\n",
+        "        \"top_p\": top_p,\n",
+        "        \"top_k\": top_k,\n",
+        "        \"min_p\": min_p,\n",
+        "    }\n",
+        "}\n",
+        "response = endpoints[\"sglang_gpu\"].predict(\n",
+        "    instances=instances,\n",
+        "    parameters=parameters,\n",
+        "    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": "ZauMzfXJzAKZ"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Chat completion\n",
+        "\n",
+        "if use_dedicated_endpoint:\n",
+        "    DEDICATED_ENDPOINT_DNS = endpoints[\"sglang_gpu\"].gca_resource.dedicated_endpoint_dns\n",
+        "ENDPOINT_RESOURCE_NAME = endpoints[\"sglang_gpu\"].resource_name\n",
+        "\n",
+        "# @markdown Because the Qwen3 models generate detailed reasoning steps, the output is expected to be long. We recommend using streaming for a better generation experience.\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": "JETd33jIDcjm"
+      },
+      "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_pytorch_qwen3_deployment.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
@@ -220,7 +220,7 @@
        "# The pre-built serving docker image for vLLM.\n",
        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/vllm-inference.cu121.0-6.ubuntu2204.py310\"\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",
        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_H100_80GB\"] {isTemplate: true}\n",
@@ -34,6 +34,11 @@
        "\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_pytorch_sam.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",
@@ -67,6 +72,10 @@
        "- Deploy the model on [Endpoint](https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints).\n",
        "- Run online predictions for image captioning.\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",
@@ -99,42 +108,32 @@
        "\n",
        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
-        "# @markdown 2. For finetuning, **[click here](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Frestricted_image_training_nvidia_a100_80gb_gpus)** to check if your project already has the required 8 Nvidia A100 80 GB GPUs in the us-central1 region. If yes, then run this notebook in the us-central1 region. If you do not have 8 Nvidia A100 80 GPUs or have more GPU requirements than this, then schedule your job with Nvidia H100 GPUs via Dynamic Workload Scheduler using [these instructions](https://cloud.google.com/vertex-ai/docs/training/schedule-jobs-dws). For Dynamic Workload Scheduler, check the [us-central1](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Fcustom_model_training_preemptible_nvidia_h100_gpus) or [europe-west4](https://console.cloud.google.com/iam-admin/quotas?location=europe-west4&metric=aiplatform.googleapis.com%2Fcustom_model_training_preemptible_nvidia_h100_gpus) quota for Nvidia H100 GPUs. If you do not have enough GPUs, then you can follow [these instructions](https://cloud.google.com/docs/quotas/view-manage#viewing_your_quota_console) to request quota.\n",
-        "\n",
-        "# @markdown 3. For serving, **[click here](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_l4_gpus)** to check if your project already has the required 1 L4 GPU in the us-central1 region.  If yes, then run this notebook in the us-central1 region. If you need more L4 GPUs for your project, then you can follow [these instructions](https://cloud.google.com/docs/quotas/view-manage#viewing_your_quota_console) to request more. Alternatively, if you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
-        "\n",
-        "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
-        "# @markdown | ----------- | ----------- | ----------- |\n",
-        "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
-        "# @markdown | a3-highgpu-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",
-        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# @markdown 5. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
+        "# @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",
+        "# Upgrade Vertex AI SDK.\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.84.0'\n",
        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
        "\n",
-        "import datetime\n",
        "import importlib\n",
        "import os\n",
-        "import uuid\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import pycocotools.mask as mask_util\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",
        "models, endpoints = {}, {}\n",
+        "LABEL = \"sam_model\"\n",
        "\n",
        "\n",
        "# Get the default cloud project id.\n",
@@ -142,60 +141,23 @@
        "\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",
-        "# 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, \"sam\")\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",
+        "base_model_name = \"sam-vit-large\"\n",
+        "PUBLISHER_MODEL_NAME = f\"publishers/meta/models/segment-anything@{base_model_name}\""
      ]
    },
    {
@@ -209,29 +171,28 @@
      "source": [
        "# @title Deploy\n",
        "\n",
-        "# @markdown This section deploys a pre-trained `sam-vit-large` model on Model Registry by using 1 L4 Machine.\n",
-        "\n",
-        "# @markdown The model deploy step will take around 20 minutes to complete.\n",
-        "\n",
        "# The pre-built serving docker image.\n",
        "# The model artifacts are embedded within the container, except for model weights which will be downloaded during deployment.\n",
        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/sam-serve:public-image-20240121\"\n",
        "\n",
        "# @markdown Set the accelerator type.\n",
-        "serve_accelerator_type = \"NVIDIA_L4\"  # @param[\"NVIDIA_TESLA_V100\", \"NVIDIA_L4\"]\n",
+        "accelerator_type = \"NVIDIA_L4\"  # @param[\"NVIDIA_TESLA_V100\", \"NVIDIA_L4\"]\n",
        "\n",
-        "if serve_accelerator_type == \"NVIDIA_TESLA_V100\":\n",
-        "    serve_machine_type = \"n1-standard-8\"\n",
-        "    serve_accelerator_count = 1\n",
-        "elif serve_accelerator_type == \"NVIDIA_L4\":\n",
-        "    serve_machine_type = \"g2-standard-12\"\n",
-        "    serve_accelerator_count = 1\n",
+        "if accelerator_type == \"NVIDIA_TESLA_V100\":\n",
+        "    machine_type = \"n1-standard-8\"\n",
+        "    accelerator_count = 1\n",
+        "elif accelerator_type == \"NVIDIA_L4\":\n",
+        "    machine_type = \"g2-standard-12\"\n",
+        "    accelerator_count = 1\n",
        "else:\n",
        "    print(f\"Unsupported accelerator type: {accelerator_type}\")\n",
        "\n",
        "MODEL_ID = \"facebook/sam-vit-large\"\n",
        "task = \"mask-generation\"\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",
        "def deploy_model(\n",
        "    task, display_name, model_id, machine_type, accelerator_type, accelerator_count\n",
@@ -255,20 +216,62 @@
        "    )\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
-        "        machine_type=serve_machine_type,\n",
-        "        accelerator_type=serve_accelerator_type,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=1,\n",
        "        deploy_request_timeout=1800,\n",
-        "        system_labels={\"NOTEBOOK_NAME\": \"model_garden_pytorch_sam.ipynb\"},\n",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_sam.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
+        "        },\n",
        "    )\n",
-        "    return model, endpoint\n",
+        "    return model, endpoint"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "HF8tcQUFG7jB"
+      },
+      "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.preview 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",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "lyZSkLIwG7jB"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with customized configs\n",
+        "\n",
+        "# @markdown This section deploys a pre-trained `sam-vit-large` model on Model Registry by using 1 L4 Machine.\n",
+        "\n",
+        "# @markdown The model deploy step will take around 20 minutes to complete.\n",
        "\n",
        "common_util.check_quota(\n",
        "    project_id=PROJECT_ID,\n",
        "    region=REGION,\n",
-        "    accelerator_type=serve_accelerator_type,\n",
-        "    accelerator_count=serve_accelerator_count,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
        "    is_for_training=False,\n",
        ")\n",
        "\n",
@@ -276,8 +279,8 @@
        "    task=task,\n",
        "    display_name=common_util.get_job_name_with_datetime(prefix=task),\n",
        "    model_id=\"facebook/sam-vit-large\",\n",
-        "    machine_type=serve_machine_type,\n",
-        "    accelerator_type=serve_accelerator_type,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
        "    accelerator_count=1,\n",
        ")"
      ]
@@ -373,11 +376,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()"
      ]
    }
  ],
@@ -29,28 +29,28 @@
        "id": "2bd716bf3e39"
      },
      "source": [
-        "# Vertex AI Model Garden - ViT-GPT2 Image Captioning\n",
+        "# Vertex AI Model Garden - ViT-GPT2 Image Captioning [Deprecated]\n",
        "\n",
-        "<table align=\"left\">\n",
-        "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_vit_gpt2_image_captioning.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
-        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_vit_gpt2_image_captioning.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "\u003ctable align=\"left\"\u003e\n",
+        "  \u003ctd\u003e\n",
+        "    \u003ca href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_vit_gpt2_image_captioning.ipynb\"\u003e\n",
+        "      \u003cimg src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"\u003e Run in Colab\n",
+        "    \u003c/a\u003e\n",
+        "  \u003c/td\u003e\n",
+        "  \u003ctd\u003e\n",
+        "    \u003ca href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_vit_gpt2_image_captioning.ipynb\"\u003e\n",
+        "      \u003cimg src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"\u003e\n",
        "      View on GitHub\n",
-        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_pytorch_vit_gpt2_image_captioning.ipynb\">\n",
-        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "    \u003c/a\u003e\n",
+        "  \u003c/td\u003e\n",
+        "  \u003ctd\u003e\n",
+        "    \u003ca href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_pytorch_vit_gpt2_image_captioning.ipynb\"\u003e\n",
+        "      \u003cimg src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"\u003e\n",
        "Open in Vertex AI Workbench\n",
-        "    </a>\n",
+        "    \u003c/a\u003e\n",
        "    (a Python-3 CPU notebook is recommended)\n",
-        "  </td>\n",
-        "</table>"
+        "  \u003c/td\u003e\n",
+        "\u003c/table\u003e"
      ]
    },
    {
@@ -312,7 +312,7 @@
        "print(\"Creating endpoint.\")\n",
        "\n",
        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/jax-timesfm-serve:20240828_1036_RC00\"\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",
@@ -31,18 +31,18 @@
      },
      "source": [
        "# Vertex AI Model Garden - WeatherNext Forecasting\n",
-        "<table><tbody><tr>\n",
-        "  <td style=\"text-align: center\">\n",
-        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_weather_prediction_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_weather_prediction_on_vertex.ipynb\">\n",
-        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
-        "    </a>\n",
-        "  </td>\n",
-        "</tr></tbody></table>"
+        "\u003ctable\u003e\u003ctbody\u003e\u003ctr\u003e\n",
+        "  \u003ctd style=\"text-align: center\"\u003e\n",
+        "    \u003ca 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_weather_prediction_on_vertex.ipynb\"\u003e\n",
+        "      \u003cimg alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"\u003e\u003cbr\u003e Run in Colab Enterprise\n",
+        "    \u003c/a\u003e\n",
+        "  \u003c/td\u003e\n",
+        "  \u003ctd style=\"text-align: center\"\u003e\n",
+        "    \u003ca href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_weather_prediction_on_vertex.ipynb\"\u003e\n",
+        "      \u003cimg alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"\u003e\u003cbr\u003e View on GitHub\n",
+        "    \u003c/a\u003e\n",
+        "  \u003c/td\u003e\n",
+        "\u003c/tr\u003e\u003c/tbody\u003e\u003c/table\u003e"
      ]
    },
    {
@@ -81,7 +81,7 @@
        "\n",
        "### Request For TPU Quota\n",
        "\n",
-        "By default, the quota for TPU training [Custom model training TPU v5e cores per region](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Fcustom_model_training_tpu_v5e) is 0. TPU quota is only available in `us-west1`, `us-west4`, `us-central1`. You can request for higher TPU quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota). It is suggested to request at least 4 v5e to run this notebook."
+        "By default, the quota for TPU training [Custom model training TPU v5e cores per region](https://console.cloud.google.com/iam-admin/quotas?location=us-central1\u0026metric=aiplatform.googleapis.com%2Fcustom_model_training_tpu_v5e) is 0. TPU quota is only available in `us-west1`, `us-west4`, `us-central1`. You can request for higher TPU quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota). It is suggested to request at least 4 v5e to run this notebook."
      ]
    },
    {
@@ -112,12 +112,12 @@
        "import importlib\n",
        "import os\n",
        "import uuid\n",
-        "from typing import Tuple\n",
+        "from typing import Tuple, List\n",
        "import glob\n",
        "from google.cloud import aiplatform, storage\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\u003e=1.64.0'\n",
        "if not os.path.exists(\"./vertex-ai-samples\"):\n",
        "  ! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
        "! pip3 uninstall --quiet -y xarray\n",
@@ -181,7 +181,7 @@
        "\n",
        "\n",
        "# Utility functions for vertex jobs.\n",
-        "def get_job_name_with_datetime(prefix: str) -> str:\n",
+        "def get_job_name_with_datetime(prefix: str) -\u003e str:\n",
        "    \"\"\"Gets the job name with date time when triggering training or deployment\n",
        "    jobs in Vertex AI.\n",
        "    \"\"\"\n",
@@ -229,16 +229,67 @@
        "\n",
        "print(xarray.backends.list_engines())\n",
        "\n",
+        "def get_existing_demo_step(\n",
+        "    num_forecast_steps: int,\n",
+        "    model_type: str = \"gen_small\",\n",
+        "    ) -\u003e str:\n",
+        "  # The demo data only supports some steps, and we can only run the predictions\n",
+        "  # if num_forecast_steps is smaller than or equal to the maximum supported setps.\n",
+        "  if model_type == \"gen_small\":\n",
+        "    # The max supported steps are obtained from gs://dm_graphcast/gencast/dataset.\n",
+        "    supported_demo_steps = [1, 4, 12, 20, 30]\n",
+        "  elif model_type == \"graph_small\":\n",
+        "    # The max supported steps are obtained from gs://dm_graphcast/graphcast/dataset.\n",
+        "    supported_demo_steps = [1, 4, 12, 20, 40]\n",
+        "  elif model_type == \"graph_operational\":\n",
+        "    # The max supported steps are obtained from gs://dm_graphcast/graphcast/dataset.\n",
+        "    supported_demo_steps = [1, 4, 12]\n",
+        "  else:\n",
+        "    raise ValueError(\"Invalid model_type.\")\n",
+        "\n",
+        "  # Find the proper demo data for forecasting.\n",
+        "  found_supported_step = supported_demo_steps[-1]\n",
+        "  for i, supported_step in enumerate(supported_demo_steps):\n",
+        "    if num_forecast_steps \u003c= supported_step:\n",
+        "      found_supported_step = supported_step\n",
+        "      break\n",
+        "  if num_forecast_steps \u003e found_supported_step:\n",
+        "    raise ValueError(f\"Supported demo steps for {model_type} in gs://dm_graphcast are {supported_demo_steps}. {num_forecast_steps} is too large, and could not find proper demo data.\")\n",
+        "  return found_supported_step\n",
+        "\n",
+        "def get_suggested_machines(\n",
+        "    num_forecast_steps: int,\n",
+        "    model_type: str = \"gen_small\",) -\u003e Tuple[str, str, int]:\n",
+        "    if model_type == \"gen_small\":\n",
+        "      if num_forecast_steps \u003c= 16:\n",
+        "        machine_type = \"ct5lp-hightpu-4t\"\n",
+        "        tpu_topology = \"2x2\"\n",
+        "        accelerator_count = 4\n",
+        "      else:\n",
+        "        machine_type = \"ct5lp-hightpu-8t\"\n",
+        "        tpu_topology = \"2x4\"\n",
+        "        accelerator_count = 8\n",
+        "    else:\n",
+        "      if num_forecast_steps \u003c= 16:\n",
+        "        machine_type = \"ct5lp-hightpu-1t\"\n",
+        "        tpu_topology = \"1x1\"\n",
+        "        accelerator_count = 1\n",
+        "      else:\n",
+        "        machine_type = \"ct5lp-hightpu-4t\"\n",
+        "        tpu_topology = \"2x2\"\n",
+        "        accelerator_count = 4\n",
+        "    return machine_type, tpu_topology, accelerator_count\n",
+        "\n",
        "def select(\n",
        "    data: xarray.Dataset,\n",
        "    variable: str,\n",
        "    level: Optional[int] = None,\n",
        "    max_steps: Optional[int] = None\n",
-        "    ) -> xarray.Dataset:\n",
+        "    ) -\u003e xarray.Dataset:\n",
        "  data = data[variable]\n",
        "  if \"batch\" in data.dims:\n",
        "    data = data.isel(batch=0)\n",
-        "  if max_steps is not None and \"time\" in data.sizes and max_steps < data.sizes[\"time\"]:\n",
+        "  if max_steps is not None and \"time\" in data.sizes and max_steps \u003c data.sizes[\"time\"]:\n",
        "    data = data.isel(time=range(0, max_steps))\n",
        "  if level is not None and \"level\" in data.coords:\n",
        "    data = data.sel(level=level)\n",
@@ -248,7 +299,7 @@
        "    data: xarray.Dataset,\n",
        "    center: Optional[float] = None,\n",
        "    robust: bool = False,\n",
-        "    ) -> tuple[xarray.Dataset, matplotlib.colors.Normalize, str]:\n",
+        "    ) -\u003e tuple[xarray.Dataset, matplotlib.colors.Normalize, str]:\n",
        "  vmin = np.nanpercentile(data, (2 if robust else 0))\n",
        "  vmax = np.nanpercentile(data, (98 if robust else 100))\n",
        "  if center is not None:\n",
@@ -264,7 +315,7 @@
        "    plot_size: float = 5,\n",
        "    robust: bool = False,\n",
        "    cols: int = 4\n",
-        "    ) -> tuple[xarray.Dataset, matplotlib.colors.Normalize, str]:\n",
+        "    ) -\u003e tuple[xarray.Dataset, matplotlib.colors.Normalize, str]:\n",
        "\n",
        "  first_data = next(iter(data.values()))[0]\n",
        "  max_steps = first_data.sizes.get(\"time\", 1)\n",
@@ -326,45 +377,40 @@
      "cell_type": "code",
      "source": [
        "# @title Configure Models\n",
+        "# @markdown You can config WeatherNext models with *data_type*, *model_type* and *num_forecast_steps*.\n",
        "\n",
        "output_dir = f\"{BUCKET_URI}/science\"\n",
        "accelerator_type = \"TPU_V5e\"\n",
-        "SCIENCE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/science-serve.tpu.0-1.debian12.py310:20250331_0715_RC03\"\n",
        "\n",
-        "# @markdown All demo data are from the public gcs bucket gs://dm_graphcast.\n",
-        "# @markdown You can prepare your own data similarly for forecasting.\n",
-        "# @markdown   - The demo data for gen_small is from the date 2019-03-29 with resolution 1.0.\n",
-        "# @markdown   - The demo data for graph_small is from the date 2022-01-01 with resolution 1.0.\n",
-        "# @markdown   - The demo data for graph_operational is from the date 2022-01-01 with resolution 0.25.\n",
-        "model_type = \"gen_small\" # @param [\"gen_small\", \"graph_small\", \"graph_operational\"]\n",
+        "# @markdown All demo data are from the public gcs bucket `gs://dm_graphcast`.\n",
+        "# @markdown You can prepare your own data similarly to demo data for forecasting.\n",
+        "# @markdown The demo data for gen_small is from the date 2019-03-29 with resolution 1.0.\n",
+        "# @markdown The demo data for graph_small is from the date 2022-01-01 with resolution 1.0.\n",
+        "# @markdown The demo data for graph_operational is from the date 2022-01-01 with resolution 0.25.\n",
        "\n",
+        "model_type = \"graph_operational\" # @param [\"gen_small\", \"graph_small\", \"graph_operational\"]\n",
+        "\n",
+        "\n",
+        "num_forecast_steps = 10 # @param {type:\"integer\"}\n",
        "# @markdown *num_forecast_steps* will specific the number of forecast steps, which will indicate the forcasting time combined with model leading time.\n",
        "# @markdown Assuming num_forecast_steps=4, and the leading time is 6 hours, then the results will contain forecasts with 6 hours, 12 hours, 18 hours and 24 hours.\n",
        "# @markdown WeatherNext Gen and Graph models support leading time as 12 hours and 6 hours separately.\n",
-        "# @markdown num_forecast_steps will be truncated to the maximum of allowed values if it is beyond.\n",
-        "num_forecast_steps = 4 # @param {type:\"integer\"}\n",
+        "# @markdown num_forecast_steps will be truncated to the maximum of allowed values if it is beyond. Maximum of num_forecast_steps for Weather Gen and Graph models are 30 and 40 separately.\n",
+        "machine_type, tpu_topology, accelerator_count = get_suggested_machines(num_forecast_steps, model_type)\n",
        "\n",
+        "SCIENCE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/science-serve.tpu.0-1.debian12.py310:20250331_0715_RC03\"\n",
        "data_storage_dir = \"gs://dm_graphcast\"\n",
-        "\n",
+        "existing_demo_step = get_existing_demo_step(num_forecast_steps, model_type)\n",
        "if model_type == \"gen_small\":\n",
-        "  input_file = f\"{data_storage_dir}/gencast/dataset/source-era5_date-2019-03-29_res-1.0_levels-13_steps-04.nc\"\n",
-        "  machine_type = \"ct5lp-hightpu-4t\"\n",
-        "  tpu_topology = \"2x2\"\n",
-        "  accelerator_count = 4\n",
+        "  input_file = f\"{data_storage_dir}/gencast/dataset/source-era5_date-2019-03-29_res-1.0_levels-13_steps-{existing_demo_step:02d}.nc\"\n",
        "  # num_ensemble_samples (WeatherNext Gen models only) specified the number of ensembling samples per step.\n",
        "  num_ensemble_samples = 8\n",
        "  parameters = {\"num_forecast_steps\": num_forecast_steps, \"num_ensemble_samples\": num_ensemble_samples}\n",
        "elif model_type == \"graph_small\":\n",
-        "  input_file = f\"{data_storage_dir}/graphcast/dataset/source-era5_date-2022-01-01_res-1.0_levels-13_steps-12.nc\"\n",
-        "  machine_type = \"ct5lp-hightpu-1t\"\n",
-        "  tpu_topology = \"1x1\"\n",
-        "  accelerator_count = 1\n",
+        "  input_file = f\"{data_storage_dir}/graphcast/dataset/source-era5_date-2022-01-01_res-1.0_levels-13_steps-{existing_demo_step:02d}.nc\"\n",
        "  parameters = {\"num_forecast_steps\": num_forecast_steps}\n",
        "elif model_type == \"graph_operational\":\n",
-        "  input_file = f\"{data_storage_dir}/graphcast/dataset/source-hres_date-2022-01-01_res-0.25_levels-13_steps-12.nc\"\n",
-        "  machine_type = \"ct5lp-hightpu-1t\"\n",
-        "  tpu_topology = \"1x1\"\n",
-        "  accelerator_count = 1\n",
+        "  input_file = f\"{data_storage_dir}/graphcast/dataset/source-hres_date-2022-01-01_res-0.25_levels-13_steps-{existing_demo_step:02d}.nc\"\n",
        "  parameters = {\"num_forecast_steps\": num_forecast_steps}\n",
        "else:\n",
        "  raise ValueError(\"Invalid example_type.\")\n",
@@ -381,6 +427,7 @@
        "\n",
        "print(f\"machine_type is {machine_type}.\")\n",
        "print(f\"tpu_topology is {tpu_topology}.\")\n",
+        "print(f\"SCIENCE_DOCKER_URI is {SCIENCE_DOCKER_URI}.\")\n",
        "print(f\"The prediction instances: {instances}\")\n"
      ],
      "metadata": {
@@ -398,7 +445,6 @@
        "# @markdown It usually takes a couple of minutes to finish.\n",
        "# @markdown Click on the generated link in the output to see your run in the Cloud Console.\n",
        "\n",
-        "\n",
        "print(\"Check if there are enough quota.\")\n",
        "common_util.check_quota(\n",
        "    project_id=PROJECT_ID,\n",
@@ -529,8 +575,7 @@
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "911406c1561e",
-        "cellView": "form"
+        "id": "911406c1561e"
      },
      "outputs": [],
      "source": [
@@ -538,7 +583,8 @@
        "\n",
        "delete_bucket = True  # @param {type:\"boolean\"}\n",
        "if delete_bucket:\n",
-        "    ! gsutil -m rm -r $BUCKET_URI\n"
+        "    ! gsutil -m rm -r $BUCKET_URI\n",
+        "\n"
      ]
    }
  ],
@@ -8,7 +8,7 @@
      },
      "outputs": [],
      "source": [
-        "# Copyright 2023 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",
@@ -140,7 +140,8 @@
      "outputs": [],
      "source": [
        "!gsutil cp gs://visionai-artifacts/visionai-0.0.6-py3-none-any.whl .\n",
-        "!pip install visionai-0.0.6-py3-none-any.whl --force-reinstall"
+        "!pip install visionai-0.0.6-py3-none-any.whl --force-reinstall\n",
+        "!pip install ipywidgets requests"
      ]
    },
    {
@@ -206,7 +207,6 @@
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
-        "\n",
        "# Set the project id\n",
        "! gcloud config set project {PROJECT_ID}"
      ]
@@ -785,7 +785,7 @@
      "source": [
        "MAX_RESULTS = 10  # @param {type: \"integer\"} Set to 0 to allow all results.\n",
        "QUERY = \"multiple purple tulips\"  # @param {type: \"string\"}\n",
-        "\n",
+        "print(\"endpoint_name:\", endpoint_name)\n",
        "results = warehouse_client.search_index_endpoint(\n",
        "    visionai_v1.SearchIndexEndpointRequest(\n",
        "        index_endpoint=endpoint_name,\n",
@@ -877,6 +877,158 @@
        "RenderImages(image_uris=uris)"
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "c44759617e80"
+      },
+      "source": [
+        "### Adding metadata filters"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ec46febeebfa"
+      },
+      "outputs": [],
+      "source": [
+        "IMAGE_GCS_FILE = \"gs://cloud-samples-data/ai-platform/flowers/roses/14312910041_b747240d56_n.jpg\"  # @param {type: \"string\"} example: gs://iwh_fishfood/sample-image.jpg\n",
+        "MAX_RESULTS = 10  # @param {type: \"integer\"} Set to 0 to allow all results.\n",
+        "IMAGE_FILE = \"/tmp/sample-image.jpg\"\n",
+        "!gsutil cp $IMAGE_GCS_FILE $IMAGE_FILE\n",
+        "\n",
+        "with open(IMAGE_FILE, \"rb\") as f:\n",
+        "    image_content = f.read()\n",
+        "grid = GridspecLayout(1, 1)\n",
+        "grid[0, 0] = ipywidgets.Image(value=image_content, width=200)\n",
+        "\n",
+        "print(\"Query image:\")\n",
+        "display(grid)\n",
+        "\n",
+        "aspect_ratios = [\"1.3\", \"1.4\"]  # @#param {type: \"list\", itemType: \"string\"}\n",
+        "aspect_ratio_criteria = visionai_v1.types.StringArray(txt_values=aspect_ratios)\n",
+        "aspect_ratio_filter = visionai_v1.Criteria(\n",
+        "    field=\"aspect-ratio\", text_array=aspect_ratio_criteria\n",
+        ")\n",
+        "\n",
+        "# Define creator filter criteria\n",
+        "creator = [\"Saige Fuentes\"]  # @#param {type: \"list\", itemType: \"string\"}\n",
+        "creator_criteria = visionai_v1.types.StringArray(txt_values=creator)\n",
+        "creator_filter = visionai_v1.Criteria(field=\"creator\", text_array=creator_criteria)\n",
+        "\n",
+        "criteria = [aspect_ratio_filter, creator_filter]\n",
+        "\n",
+        "results = warehouse_client.search_index_endpoint(\n",
+        "    visionai_v1.SearchIndexEndpointRequest(\n",
+        "        index_endpoint=endpoint_name,\n",
+        "        image_query=visionai_v1.ImageQuery(\n",
+        "            input_image=image_content,\n",
+        "        ),\n",
+        "        criteria=criteria,\n",
+        "    ),\n",
+        ")\n",
+        "\n",
+        "results_cnt = 0\n",
+        "asset_names = []\n",
+        "for r in results:\n",
+        "    asset_names.append(r.asset)\n",
+        "    results_cnt += 1\n",
+        "    if results_cnt >= MAX_RESULTS:\n",
+        "        break\n",
+        "\n",
+        "uris = list(\n",
+        "    map(\n",
+        "        lambda asset_name: warehouse_client.generate_retrieval_url(\n",
+        "            visionai_v1.GenerateRetrievalUrlRequest(\n",
+        "                name=asset_name,\n",
+        "            )\n",
+        "        ).signed_uri,\n",
+        "        asset_names,\n",
+        "    )\n",
+        ")\n",
+        "\n",
+        "print(\"Filtered search results:\")\n",
+        "RenderImages(image_uris=uris)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "c44759617e80"
+      },
+      "source": [
+        "### Adding metadata filters"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ec702c5248f9"
+      },
+      "outputs": [],
+      "source": [
+        "IMAGE_GCS_FILE = \"gs://cloud-samples-data/ai-platform/flowers/roses/14312910041_b747240d56_n.jpg\"  # @param {type: \"string\"} example: gs://iwh_fishfood/sample-image.jpg\n",
+        "MAX_RESULTS = 10  # @param {type: \"integer\"} Set to 0 to allow all results.\n",
+        "IMAGE_FILE = \"/tmp/sample-image.jpg\"\n",
+        "!gsutil cp $IMAGE_GCS_FILE $IMAGE_FILE\n",
+        "\n",
+        "with open(IMAGE_FILE, \"rb\") as f:\n",
+        "    image_content = f.read()\n",
+        "grid = GridspecLayout(1, 1)\n",
+        "grid[0, 0] = ipywidgets.Image(value=image_content, width=200)\n",
+        "\n",
+        "print(\"Query image:\")\n",
+        "display(grid)\n",
+        "\n",
+        "aspect_ratios = [\"1.3\", \"1.4\"]  # # @param {type: \"list\", itemType: \"string\"}\n",
+        "aspect_ratio_criteria = visionai_v1.types.StringArray(txt_values=aspect_ratios)\n",
+        "aspect_ratio_filter = visionai_v1.Criteria(\n",
+        "    field=\"aspect-ratio\", text_array=aspect_ratio_criteria\n",
+        ")\n",
+        "\n",
+        "# Define creator filter criteria\n",
+        "creator = [\"Saige Fuentes\"]  # # @param {type: \"list\", itemType: \"string\"}\n",
+        "creator_criteria = visionai_v1.types.StringArray(txt_values=creator)\n",
+        "creator_filter = visionai_v1.Criteria(field=\"creator\", text_array=creator_criteria)\n",
+        "\n",
+        "criteria = [aspect_ratio_filter, creator_filter]\n",
+        "\n",
+        "results = warehouse_client.search_index_endpoint(\n",
+        "    visionai_v1.SearchIndexEndpointRequest(\n",
+        "        index_endpoint=endpoint_name,\n",
+        "        image_query=visionai_v1.ImageQuery(\n",
+        "            input_image=image_content,\n",
+        "        ),\n",
+        "        criteria=criteria,\n",
+        "    ),\n",
+        ")\n",
+        "\n",
+        "results_cnt = 0\n",
+        "asset_names = []\n",
+        "for r in results:\n",
+        "    asset_names.append(r.asset)\n",
+        "    results_cnt += 1\n",
+        "    if results_cnt >= MAX_RESULTS:\n",
+        "        break\n",
+        "\n",
+        "uris = list(\n",
+        "    map(\n",
+        "        lambda asset_name: warehouse_client.generate_retrieval_url(\n",
+        "            visionai_v1.GenerateRetrievalUrlRequest(\n",
+        "                name=asset_name,\n",
+        "            )\n",
+        "        ).signed_uri,\n",
+        "        asset_names,\n",
+        "    )\n",
+        ")\n",
+        "\n",
+        "print(\"Filtered search results:\")\n",
+        "RenderImages(image_uris=uris)"
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -57,8 +57,12 @@
 /training/tpuv5e_gemma_peft_finetuning_and_serving.ipynb @brianchunkang
 /training/tpuv5e_llama2_pytorch_finetuning_and_serving.ipynb @brianchunkang @chiefkarlin
 /prediction/get_started_with_psc_private_endpoint.ipynb @tianjiaoliu
+/prediction/get_started_with_dedicated_endpoint.ipynb @tianjiaoliu
 /ray_on_vertex_ai/spark_on_ray_on_vertex_ai.ipynb @ravi-dalal
 /generative_ai/mistralai_intro.ipynb @sujituk
-/generative_ai/ai21labs_intro.ipynb @sujituk
+/generative_ai/ai21labs_intro.ipynb @miri-bar
 /forecasting/starry_net_pipeline.ipynb @tsteve
 /generative_ai/nvidia_nim_vertexai.ipynb @sujituk
+/generative_ai/cambai_intro.ipynb @arnavmehta7
+/generative_ai/qodo_intro.ipynb @talshef
+
@@ -409,7 +409,7 @@
        "df = bq_client.list_rows(table).to_dataframe()\n",
        "\n",
        "# Drop unusable rows\n",
-        "df = df.replace(to_replace=NA_VALUES, value=np.NaN).dropna()\n",
+        "df = df.replace(to_replace=NA_VALUES, value=np.nan).dropna()\n",
        "\n",
        "# Convert categorical columns to numeric\n",
        "df[\"island\"], _ = pd.factorize(df[\"island\"])\n",
@@ -153,53 +153,13 @@
      },
      "outputs": [],
      "source": [
-        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
+        "%pip install --upgrade --quiet google-cloud-aiplatform \\\n",
        "                                 google-cloud \\\n",
        "                                 google-cloud-storage \\\n",
        "                                 pillow \\\n",
        "                                 numpy"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ff555b32bab8"
-      },
-      "source": [
-        "### Restart runtime (Colab only)\n",
-        "\n",
-        "To use the newly installed packages, you must restart the runtime on Google Colab."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "f09b4dff629a"
-      },
-      "outputs": [],
-      "source": [
-        "import sys\n",
-        "\n",
-        "if \"google.colab\" in sys.modules:\n",
-        "\n",
-        "    import IPython\n",
-        "\n",
-        "    app = IPython.Application.instance()\n",
-        "    app.kernel.do_shutdown(True)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "54c5ef8a8f43"
-      },
-      "source": [
-        "<div class=\"alert alert-block alert-warning\">\n",
-        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
-        "</div>\n"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
denisj3030andGitHub	49dfbf7f14	Merge branch 'main' into ocr	2025-05-23 13:31:37 -04:00
denisj3030	f0d75e59e1	small fixed	2025-05-23 17:27:59 +00:00
denisj3030andGitHub	ca8cb4480a	adding mistral ocr (#4069 ) * adding mistral ocr * adding mistral ocr * fix missing ] * fix codestral ver	2025-05-23 16:03:27 +00:00
denisj3030	9dea76a52a	fix codestral ver	2025-05-23 15:42:46 +00:00
denisj3030	2e56b498ec	fix missing ]	2025-05-23 15:24:06 +00:00
denisj3030	d66ae4cb78	adding mistral ocr	2025-05-23 14:10:50 +00:00
denisj3030	5001259780	adding mistral ocr	2025-05-23 14:02:11 +00:00
Rayan DasoriyaandCopybara-Service	5e4bdf4a2c	Add different variable name for trtllm accelerator type PiperOrigin-RevId: 762039341	2025-05-22 11:04:55 -07:00
denisj3030andGitHub	011422636f	adding Anthropic v4 (#4065 ) * adding v4 * adding v4 lint fix * adding v4 lint fix * adding v4 lint fix	2025-05-22 17:43:28 +00:00
Vertex MG TeamandCopybara-Service	dace15a300	Publish Qwen3 Axolotl Notebook. PiperOrigin-RevId: 761753984	2025-05-21 18:51:56 -07:00
Changyu ZhuandCopybara-Service	f3134943a4	Update SGLang docker version in Mode Garden DeepSeek deployment notebook. PiperOrigin-RevId: 761742957	2025-05-21 18:06:23 -07:00
skarukasandGitHub	5eb0a7e114	Update embedding notebooks to reference gemini-embedding-001 (#4058 ) * Update embedding notebooks to reference gemini-embedding-001 * Update embedding notebooks to reference gemini-embedding-001	2025-05-21 12:52:44 +00:00
Rayan DasoriyaandCopybara-Service	48ead2727f	Fix llama notebooks PiperOrigin-RevId: 761112151	2025-05-20 09:22:11 -07:00
Vertex MG TeamandCopybara-Service	dbb226a7a4	Update the vllm demo notebook to use the vanilla container image. PiperOrigin-RevId: 760757673	2025-05-19 13:53:14 -07:00
Changyu ZhuandCopybara-Service	04352e92cc	Add TensorRT-LLM deployment sample to Deepseek and Llama3.3 notebooks. PiperOrigin-RevId: 759818622	2025-05-16 17:44:44 -07:00
Genquan DuanandCopybara-Service	23975591bc	Add notebook examples for smooth integration of vmg oss llms + adk. PiperOrigin-RevId: 759741078	2025-05-16 13:46:02 -07:00
Vertex MG TeamandCopybara-Service	9027adebc2	[VMG Tutorial] Create serving notebook tutorial for hexLLM Llama 3 deep dive. PiperOrigin-RevId: 759632399	2025-05-16 08:52:42 -07:00
Vertex MG TeamandCopybara-Service	ce05c8af80	Fix indentation error in the Fast deploy Chat Completion PiperOrigin-RevId: 759431771	2025-05-15 21:18:07 -07:00
Vertex MG TeamandCopybara-Service	cfb870a323	Updated Llama 4 notebook PiperOrigin-RevId: 759146983	2025-05-15 07:53:14 -07:00
TJ-Liu GitHub TJ Liu	22709379dd	feat: Add prediction dedicated endpoint colab sample (#3942 ) * feat: Add prediction dedicated endpoint colab sample * Update get_started_with_dedicated_endpoint.ipynb --------- Co-authored-by: TJ(Tianjiao) Liu <tianjiaoliu@google.com>	2025-05-15 01:14:40 +00:00
Vertex MG TeamandCopybara-Service	4aafcfb40f	Add qwen3 specific instruction to runtime creation. PiperOrigin-RevId: 758737197	2025-05-14 10:21:56 -07:00
Rayan DasoriyaandCopybara-Service	d53aa0c816	Fix gemma2 notebook PiperOrigin-RevId: 758283299	2025-05-13 10:43:57 -07:00
Rayan DasoriyaandCopybara-Service	21976bf94f	Clarify the HF token requirement PiperOrigin-RevId: 758277776	2025-05-13 10:33:15 -07:00
Vertex MG TeamandCopybara-Service	434e8ac8fc	[Fix] Add `NVIDIA_A100_80GB` option to vertex finetuning. PiperOrigin-RevId: 757854822	2025-05-12 11:55:28 -07:00
Vertex MG TeamandCopybara-Service	2376532e3d	Fix A100_80GB machine type. PiperOrigin-RevId: 757734855	2025-05-12 06:17:50 -07:00
Vertex MG TeamandCopybara-Service	8b2fbe3f34	Make runtime connection instructions clearer for axolotl notebook. PiperOrigin-RevId: 757720267	2025-05-12 05:29:51 -07:00
Vertex MG TeamandCopybara-Service	f991ae44ab	Fix local finetuning for axolotl. PiperOrigin-RevId: 757655772	2025-05-12 02:01:51 -07:00
Dustin LuongandCopybara-Service	7eb8b76d19	Remove reference server deployment and inference from llama3.2 notebook. PiperOrigin-RevId: 757455670	2025-05-11 11:01:47 -07:00
Vertex MG TeamandCopybara-Service	7da9c0b644	Remove a3-ultragpu from Axolotl notebook. PiperOrigin-RevId: 756910136	2025-05-09 14:35:55 -07:00
Genquan DuanandCopybara-Service	46a75498d2	Update notebook to add more machine type suggestions, and support allowlisted real-time forecasting for graph_operational model. PiperOrigin-RevId: 756860759	2025-05-09 12:17:35 -07:00
Vertex MG TeamandCopybara-Service	c2208eb454	Fix hf cache dir for local finetuning and use gcsfuse for local training and merging of model. PiperOrigin-RevId: 756772796	2025-05-09 08:15:32 -07:00
Vertex MG TeamandCopybara-Service	f9c43d4a8a	Add a3-ultragpu with dws to Axolotl notebook. PiperOrigin-RevId: 756625405	2025-05-08 23:39:48 -07:00
Vertex MG TeamandCopybara-Service	32ae5b4af0	Add gpu type check in axolotl local run. PiperOrigin-RevId: 756598801	2025-05-08 22:04:50 -07:00
Vertex MG TeamandCopybara-Service	2345f89625	Refactoring and minor fixes PiperOrigin-RevId: 756570419	2025-05-08 20:25:13 -07:00
Vertex MG TeamandCopybara-Service	46a80b731b	Add Dia-1.6B serving notebook PiperOrigin-RevId: 756499377	2025-05-08 16:13:17 -07:00
Vertex MG TeamandCopybara-Service	2781808a96	Add n1+T4 as one deployment option for e5 model. PiperOrigin-RevId: 756479268	2025-05-08 15:16:27 -07:00
Vertex MG TeamandCopybara-Service	5af1bc523d	Fix Axolotl installation. PiperOrigin-RevId: 756401696	2025-05-08 11:54:00 -07:00
Rayan DasoriyaandCopybara-Service	9742d29e51	Add eval harness support for Qwen2.5 finetuning notebook PiperOrigin-RevId: 756350741	2025-05-08 09:50:30 -07:00
Vertex MG TeamandCopybara-Service	30fda48397	change default gpu type as empty to avoid unintended gpu runtime creation PiperOrigin-RevId: 756199290	2025-05-08 01:17:09 -07:00
Vertex MG TeamandCopybara-Service	6888728be5	Fix dataset vars definition issue and make runtime creation session optional in Axolotl notebook. PiperOrigin-RevId: 756138907	2025-05-07 21:40:57 -07:00
Vertex MG TeamandCopybara-Service	8c7fbc6210	Phi-4 reasoning variants PiperOrigin-RevId: 756047269	2025-05-07 16:17:35 -07:00
Dustin LuongandCopybara-Service	4c4e224d31	Update Qwen3 deployment notebook with FP8 support. PiperOrigin-RevId: 756041267	2025-05-07 15:59:11 -07:00
Vertex MG TeamandCopybara-Service	ac98f72005	Fix runtime creation codes in Axolotl notebook to disallow L4x16 combination. PiperOrigin-RevId: 755968781	2025-05-07 12:51:45 -07:00
Vertex MG TeamandCopybara-Service	ddebceb70c	Support Workbench in Gemma 2 tuning notebook PiperOrigin-RevId: 755884581	2025-05-07 09:24:14 -07:00
Vertex MG TeamandCopybara-Service	d8e5c3b461	Update Axolotl finetuning notebook to support Qwen3. PiperOrigin-RevId: 755577906	2025-05-06 16:52:10 -07:00
Eric DongandGitHub	0bfe70afcd	refactor: remove runtime reboot (#4026 ) * refactor: remove runtime reboot * Remove spaces * Use %pip instead * Update model path * Update BQ path * Downgrade numpy for backfoward compability	2025-05-05 21:12:55 +00:00
Ravi DalalandGitHub	7488a5dc27	upgraded spark on ray on vertex ai notebook to 2.42.0 version (#4027 ) * upgraded spark on ray on vertex ai notebook to 2.42.0 version * upgraded spark on ray on vertex ai notebook to 2.42.0 version * upgraded spark on ray on vertex ai notebook to 2.42.0 version * upgraded spark on ray on vertex ai notebook to 2.42.0 version	2025-05-05 19:05:30 +00:00
Vertex MG TeamandCopybara-Service	4914987998	Update Axolotl docker version for the finetuning notebook PiperOrigin-RevId: 754984828	2025-05-05 10:20:32 -07:00
Vertex MG TeamandCopybara-Service	abd30065d0	Add support for dedicated endpoint in paligemma finetuning PiperOrigin-RevId: 753860258	2025-05-01 22:26:25 -07:00
Dustin LuongandCopybara-Service	68960f1221	Add SDK deploy option to Qwen3 deployment notebook. PiperOrigin-RevId: 753636431	2025-05-01 10:17:08 -07:00
Vertex MG TeamandCopybara-Service	836d1ebbd1	No public description PiperOrigin-RevId: 753635866	2025-05-01 10:15:36 -07:00
Vertex MG TeamandCopybara-Service	f96c830a66	chore: Add back `faster_deploy_enabled` in deploy SDK. PiperOrigin-RevId: 753454732	2025-04-30 23:17:36 -07:00
Dustin LuongandCopybara-Service	ecd1fc28d2	Set model id to GCS bucket for Qwen3 235B model for better stability. PiperOrigin-RevId: 753430192	2025-04-30 21:33:20 -07:00
Vertex MG TeamandCopybara-Service	da0874682f	Set dedicate endpoint as default for model_garden_finetuning_tutorial PiperOrigin-RevId: 753294164	2025-04-30 13:38:52 -07:00
Dustin LuongandCopybara-Service	e851c1ad99	Upload Qwen3 deployment notebook. PiperOrigin-RevId: 753265446	2025-04-30 12:19:21 -07:00
Vertex MG TeamandCopybara-Service	93fd7088ba	Add deployment options for the smaller Qwen 2.5 model versions too in the finetuning notebook. PiperOrigin-RevId: 753246638	2025-04-30 11:30:02 -07:00
Vertex MG TeamandCopybara-Service	cab440a05b	Fixes for Llama Prompt Guard variant to deployment notebook. PiperOrigin-RevId: 752852958	2025-04-29 13:08:40 -07:00
Vertex MG TeamandCopybara-Service	f03f88a0f7	Add Llama Prompt Guard variant to deployment notebook. PiperOrigin-RevId: 752802325	2025-04-29 10:58:25 -07:00
Vertex MG TeamandCopybara-Service	841fbf9f53	Add Llama Guard variant to deployment notebook. PiperOrigin-RevId: 752801901	2025-04-29 10:56:47 -07:00
Vertex MG TeamandCopybara-Service	bed0c09ad9	Add Model Garden deploy SDK PiperOrigin-RevId: 752589085	2025-04-28 22:43:24 -07:00
Vertex MG TeamandCopybara-Service	2497476009	Remove 0.5B and 1.5B Qwen 2.5 models from the notebook PiperOrigin-RevId: 752453505	2025-04-28 14:42:09 -07:00
Eric DongandGitHub	5e2841384f	refactor: Use raw string in regex (#4016 ) * refactor: Use raw string in regex * np.NaN was removed in the NumPy 2.0 * Move default python version to 3.10 * Move back default python version to 3.9 * Remove version mistmatched tes notebooks	2025-04-28 19:54:44 +00:00
Rayan Dasoriya GitHub Rayan Dasoriya	49de587b2c	Add optional support for dedicated endpoint (#4014 ) Co-authored-by: Rayan Dasoriya <dasoriya@google.com>	2025-04-28 12:03:33 +00:00
Vertex MG TeamandCopybara-Service	feea47a206	Enable the use of dedicated Endpoints in instructpix2pix notebook PiperOrigin-RevId: 751641417	2025-04-25 19:45:08 -07:00
Vertex MG TeamandCopybara-Service	f4b1b277cf	Add HiDream-I1 serving notebook PiperOrigin-RevId: 751521086	2025-04-25 12:38:34 -07:00
Vertex MG TeamandCopybara-Service	e4608c983b	chore: Clean up endpoint resource names. PiperOrigin-RevId: 751415593	2025-04-25 07:58:12 -07:00
Vertex MG TeamandCopybara-Service	f147b50332	Support VPC-SC. PiperOrigin-RevId: 751028021	2025-04-24 09:45:20 -07:00
Changyu ZhuandCopybara-Service	b1b16f718b	Fix missing positional argument in DeepSeek deployment notebook PiperOrigin-RevId: 750742323	2025-04-23 15:19:37 -07:00
Vertex MG TeamandCopybara-Service	871eb25dc3	Deprecate the colab notebooks for a few models. PiperOrigin-RevId: 750692428	2025-04-23 12:54:53 -07:00
denisj3030andGitHub	793515bac2	Launchpad qodo (#3994 ) * launchpad notebooks * Update ai21labs_intro.ipynb * Update ai21labs_intro.ipynb * launchpad * audio play * lint fixes * Update ai21labs_intro.ipynb * Update cambai_intro.ipynb * Update cambai_intro.ipynb * Update cambai_intro.ipynb * qodo lint	2025-04-22 15:25:46 +00:00
Vertex MG TeamandCopybara-Service	4496842a86	Support VPC-SC. Add Model Garden deploy SDK PiperOrigin-RevId: 750197006	2025-04-22 08:12:30 -07:00
Vertex MG TeamandCopybara-Service	0687238a97	Refactoring PiperOrigin-RevId: 750127462	2025-04-22 03:47:19 -07:00
Vertex MG TeamandCopybara-Service	cd9120bc15	Support VPC-SC. Add Model Garden deploy SDK PiperOrigin-RevId: 749803092	2025-04-21 08:20:30 -07:00
Vertex MG TeamandCopybara-Service	77388979d7	Add Qwen 2.5 PEFT finetuning notebook PiperOrigin-RevId: 748978717	2025-04-18 02:28:04 -07:00
Vertex MG TeamandCopybara-Service	d9198306a8	Support dedicated endpoints PiperOrigin-RevId: 748922184	2025-04-17 21:48:54 -07:00
Vertex MG TeamandCopybara-Service	093e32658b	Fix check_quota fn for eval in finetuning notebooks PiperOrigin-RevId: 748701270	2025-04-17 09:09:24 -07:00
Vertex MG TeamandCopybara-Service	e6821c94c9	Add a link to file bugs PiperOrigin-RevId: 748650287	2025-04-17 05:41:07 -07:00
Vertex MG TeamandCopybara-Service	2fbf8c4379	Workbench support and refactoring PiperOrigin-RevId: 748167158	2025-04-15 23:09:41 -07:00
Vertex MG TeamandCopybara-Service	257c478ef3	Support VPC-SC. Add Model Garden deploy SDK PiperOrigin-RevId: 748076170	2025-04-15 17:06:01 -07:00
Vertex MG TeamandCopybara-Service	9ce7e61434	Support VPC-SC. Add Model Garden deploy SDK PiperOrigin-RevId: 748075941	2025-04-15 17:04:26 -07:00
denisj3030andGitHub	42fa7ac1a3	lint fixed (#3993 ) * launchpad notebooks * Update ai21labs_intro.ipynb * Update ai21labs_intro.ipynb * launchpad * audio play * lint fixes * Update ai21labs_intro.ipynb * Update cambai_intro.ipynb * Update cambai_intro.ipynb * Update cambai_intro.ipynb	2025-04-15 17:15:08 +00:00
Yehia Elshater GitHub Yehia Elshater	85d43c76ca	Add metadata filtering to image warehouse SDK notebook (#3959 ) * Add metadata filtering to image warehouse SDK notebook. * Add metadata filtering to image warehouse SDK notebook and installing dependencies * linter fix * Linter fix * Fix formatting in image_warehouse_sdk notebook * Fix formatting in image_warehouse_sdk notebook --------- Co-authored-by: Yehia Elshater <elshater@google.com>	2025-04-15 12:31:27 +00:00
Rayan Dasoriya GitHub Rayan Dasoriya	9f4d837e54	Update peft code for stable_20250409 (#3991 ) Co-authored-by: Rayan Dasoriya <dasoriya@google.com>	2025-04-15 12:29:40 +00:00
Minwoo ParkandCopybara-Service	27db486275	Improve paligemma notebook documentation. PiperOrigin-RevId: 747492281	2025-04-14 11:09:52 -07:00
Vertex MG TeamandCopybara-Service	2df60d7862	Fix dedicated endpoint codellama PiperOrigin-RevId: 747448256	2025-04-14 09:22:18 -07:00
Vertex MG TeamandCopybara-Service	7850587517	Support VPC-SC and workbench PiperOrigin-RevId: 747376863	2025-04-14 05:41:22 -07:00
Vertex MG TeamandCopybara-Service	b924278b03	Add lm eval harness to the finetuning notebook PiperOrigin-RevId: 747343434	2025-04-14 03:44:40 -07:00
Vertex MG TeamandCopybara-Service	adb16ca8b4	Update cell ordering in Llama 4 MaaS notebook. PiperOrigin-RevId: 746284188	2025-04-10 20:47:06 -07:00
Vertex MG TeamandCopybara-Service	42383514ff	Add linter change for Llama 4 MaaS notebook. PiperOrigin-RevId: 746223674	2025-04-10 16:48:23 -07:00
Vertex MG TeamandCopybara-Service	2609db529e	Add Llama 4 MaaS Notebook PiperOrigin-RevId: 746217133	2025-04-10 16:27:22 -07:00
Vertex MG TeamandCopybara-Service	632385c6bc	Support VPC-SC and workbench PiperOrigin-RevId: 746058497	2025-04-10 09:21:44 -07:00
Vertex MG TeamandCopybara-Service	4c4519f679	Support VPC-SC. Add Model Garden deploy SDK PiperOrigin-RevId: 746057566	2025-04-10 09:18:53 -07:00
Vertex MG TeamandCopybara-Service	98cded3a79	Fix chat completion in model_garden_pytorch_llama3_2_deployment.ipynb PiperOrigin-RevId: 745642276	2025-04-09 10:26:41 -07:00
Vertex MG TeamandCopybara-Service	8c37e3995b	Set dedicate endpoint as default for some model garden samples PiperOrigin-RevId: 745421163	2025-04-08 21:47:01 -07:00
Vertex MG TeamandCopybara-Service	1281bce438	see b/380317852 for details PiperOrigin-RevId: 745345688	2025-04-08 17:02:32 -07:00
talshefandGitHub	aaf0fd2ea8	Qodo intro notebook (#3968 ) * update qodo notebook * update qodo notebook * update qodo notebook * update qodo notebook * update qodo notebook * update qodo notebook * update qodo notebook	2025-04-08 21:18:10 +00:00
Vertex MG TeamandCopybara-Service	8e249b61d9	Update Hugging Face vLLM deploy notebook title and documents. PiperOrigin-RevId: 745208445	2025-04-08 10:48:41 -07:00
Vertex MG TeamandCopybara-Service	250418bab9	Set dedicate endpoint as default for some model garden samples PiperOrigin-RevId: 745192594	2025-04-08 10:10:59 -07:00
Vertex MG TeamandCopybara-Service	70ec8b4ae2	Support VPC-SC. Add Model Garden deploy SDK PiperOrigin-RevId: 745132382	2025-04-08 07:18:19 -07:00
Vertex MG TeamandCopybara-Service	e48f3f05be	Fix gemma3 finetuning. PiperOrigin-RevId: 744956022	2025-04-07 20:45:44 -07:00
denisj3030andGitHub	a60c0376c4	launchpad notebooks (#3964 ) * launchpad notebooks * Update ai21labs_intro.ipynb * Update ai21labs_intro.ipynb * launchpad * audio play	2025-04-07 19:23:47 +00:00
denisj3030andGitHub	4bb5ffa1ec	Update CODEOWNERS (#3970 )	2025-04-07 18:53:44 +00:00