qodo lint

* 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

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

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

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

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

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

@@ -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:

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

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

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

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

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

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

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

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

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

@@ -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(

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

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

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

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

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

@@ -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/

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

@@ -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,

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

@@ -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,

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

@@ -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,
   )

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

@@ -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)

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

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

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

@@ -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)])

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

@@ -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()

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

@@ -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__':

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

@@ -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:

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

@@ -1,9 +1,12 @@
 """Fileutil lib to copy files between gcs and local."""
 
+import filecmp
 import fnmatch
 import os
 import pathlib
 import shutil
+import subprocess
+import time
 from typing import List, Optional, Tuple
 import uuid
 
@@ -57,6 +60,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)

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

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

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

@@ -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:

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

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

notebooks/community/model_garden/gke_model_ui_deployment_notebook.ipynb

@@ -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/"
       ]
     }
   ],

notebooks/community/model_garden/model_garden_axolotl_finetuning.ipynb

@@ -921,7 +921,9 @@
         "    is_for_training=False,\n",
         ")\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",
         "gpu_memory_utilization = 0.95\n",
         "max_model_len = 2048\n",
         "\n",

notebooks/community/model_garden/model_garden_codegemma_deployment_on_vertex.ipynb

@@ -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",

notebooks/community/model_garden/model_garden_gemma2_deployment_on_vertex.ipynb

@@ -278,7 +278,7 @@
         "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",
+        "# @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",

notebooks/community/model_garden/model_garden_gemma2_finetuning_on_vertex.ipynb

@@ -65,6 +65,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 +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",
@@ -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",
@@ -427,7 +426,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",
@@ -582,8 +581,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 +607,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 +739,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",

notebooks/community/model_garden/model_garden_gemma3_finetuning_on_vertex.ipynb

@@ -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",
@@ -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",

notebooks/community/model_garden/model_garden_gemma_deployment_on_vertex.ipynb

@@ -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",
@@ -675,7 +675,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",

notebooks/community/model_garden/model_garden_gemma_finetuning_on_vertex.ipynb

@@ -682,7 +682,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 +1240,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",

notebooks/community/model_garden/model_garden_huggingface_pytorch_inference_deployment.ipynb

@@ -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 = \"\""
       ]
     },

notebooks/community/model_garden/model_garden_huggingface_tei_deployment.ipynb

@@ -229,7 +229,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",

notebooks/community/model_garden/model_garden_huggingface_tgi_deployment.ipynb

@@ -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\"}"
       ]
     },

notebooks/community/model_garden/model_garden_huggingface_vllm_deployment.ipynb

@@ -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,7 +163,7 @@
         "\n",
         "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate: true}\n",
         "\n",
-        "# The pre-built serving docker image for TGI with vLLM.\n",
+        "# The pre-built vLLM serving docker image.\n",
         "VLLM_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/vllm-inference.cu121.0-6.ubuntu2204.py310\"\n",
         "SERVING_CONTAINER_IMAGE_URI = VLLM_DOCKER_URI\n",
         "LABEL = \"vllm\"\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",
@@ -591,7 +591,7 @@
   ],
   "metadata": {
     "colab": {
-      "name": "model_garden_huggingface_tgi_vllm_deployment.ipynb",
+      "name": "model_garden_huggingface_vllm_deployment.ipynb",
       "toc_visible": true
     },
     "kernelspec": {

notebooks/community/model_garden/model_garden_jax_paligemma_finetuning.ipynb

@@ -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",
@@ -561,6 +568,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",
@@ -618,7 +626,9 @@
         "\n",
         "# @markdown <img src=\"https://storage.googleapis.com/longcap100/91.jpeg\" width=\"400\" >\n",
         "\n",
+        "# fmt: off\n",
         "image_url = \"https://storage.googleapis.com/longcap100/91.jpeg\"  # @param {type:\"string\"}\n",
+        "# fmt: on\n",
         "language_code = \"en\"  # @param {type: \"string\"}\n",
         "\n",
         "image = common_util.download_image(image_url)\n",

notebooks/community/model_garden/model_garden_llama3_1_finetuning_with_workbench.ipynb

@@ -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",

notebooks/community/model_garden/model_garden_phi3_deployment.ipynb

@@ -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",
@@ -765,7 +765,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",

notebooks/community/model_garden/model_garden_phi4_deployment.ipynb

@@ -388,7 +388,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",
@@ -652,7 +652,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",

notebooks/community/model_garden/model_garden_pytorch_codellama.ipynb

@@ -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,19 +200,18 @@
         "\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",
-        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
+        "model_id = os.path.join(VERTEX_AI_MODEL_GARDEN_CODE_LLAMA, model_name)\n",
+        "\n",
+        "# @markdown Set use_dedicated_endpoint to True if you 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",
         "\n",
         "\n",
         "def deploy_model_vllm(\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,15 +428,49 @@
         "    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",
@@ -543,7 +542,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 +551,7 @@
         "\n",
         "# Delete models.\n",
         "for model in models.values():\n",
-        "    model.delete()\n",
-        "\n",
-        "delete_bucket = False  # @param {type:\"boolean\"}\n",
-        "if delete_bucket:\n",
-        "    ! gsutil -m rm -r $BUCKET_NAME"
+        "    model.delete()"
       ]
     }
   ],

notebooks/community/model_garden/model_garden_pytorch_flux.ipynb

@@ -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",

notebooks/community/model_garden/model_garden_pytorch_gemma_peft_finetuning_hf.ipynb

@@ -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",

notebooks/community/model_garden/model_garden_pytorch_instant_id.ipynb

@@ -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",

notebooks/community/model_garden/model_garden_pytorch_llama2_deployment.ipynb

@@ -4,11 +4,12 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
+        "cellView": "form",
         "id": "7d9bbf86da5e"
       },
       "outputs": [],
       "source": [
-        "# Copyright 2024 Google LLC\n",
+        "# Copyright 2025 Google LLC\n",
         "#\n",
         "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
         "# you may not use this file except in compliance with the License.\n",
@@ -33,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()"
       ]
     }
   ],

notebooks/community/model_garden/model_garden_pytorch_llama3_1_agent_engine.ipynb

@@ -423,7 +423,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",

notebooks/community/model_garden/model_garden_pytorch_llama3_1_finetuning.ipynb

@@ -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,9 +239,7 @@
         "\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",
@@ -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,126 @@
         "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",
+        "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",
+        "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",
+        "    )\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 +795,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 +822,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",

notebooks/community/model_garden/model_garden_pytorch_llama3_1_qwen2_5_deployment_tpu.ipynb

@@ -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",

notebooks/community/model_garden/model_garden_pytorch_llama3_1_reasoning_engine.ipynb

@@ -423,7 +423,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",

notebooks/community/model_garden/model_garden_pytorch_llama3_2_deployment.ipynb

@@ -1497,7 +1497,9 @@
         "    }\n",
         "]\n",
         "\n",
-        "model_response = endpoints[\"ref_gpu\"].predict(instances=instances)\n",
+        "model_response = endpoints[\"ref_gpu\"].predict(\n",
+        "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
+        ")\n",
         "\n",
         "print(model_response.predictions[0])\n",
         "\n",
@@ -1553,7 +1555,9 @@
         "    }\n",
         "]\n",
         "\n",
-        "model_response = endpoints[\"ref_gpu\"].predict(instances=instances)\n",
+        "model_response = endpoints[\"ref_gpu\"].predict(\n",
+        "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
+        ")\n",
         "\n",
         "print(model_response.predictions[0])\n",
         "\n",

notebooks/community/model_garden/model_garden_pytorch_llama3_3_finetuning.ipynb

@@ -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,9 +238,7 @@
         "# @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",
@@ -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,116 @@
         "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",
+        "    )\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 +776,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 +801,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",

notebooks/community/model_garden/model_garden_pytorch_llama3_finetuning.ipynb

@@ -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",
@@ -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",

notebooks/community/model_garden/model_garden_pytorch_mistral_deployment.ipynb

@@ -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()"
       ]
     }
   ],

notebooks/community/model_garden/model_garden_pytorch_mistral_peft_tuning.ipynb

@@ -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",

notebooks/community/model_garden/model_garden_pytorch_mixtral_deployment.ipynb

@@ -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",

notebooks/community/model_garden/model_garden_pytorch_mixtral_peft_tuning.ipynb

@@ -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",

notebooks/community/model_garden/model_garden_pytorch_prompt_guard_deployment.ipynb

@@ -277,7 +277,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",

notebooks/community/model_garden/model_garden_pytorch_qwen2_deployment.ipynb

@@ -222,7 +222,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: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).\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",
         "vllm_dtype = \"bfloat16\"\n",

notebooks/community/model_garden/model_garden_pytorch_qwq_deployment.ipynb

@@ -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",

notebooks/community/model_garden/model_garden_timesfm_deployment_on_vertex.ipynb

@@ -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",

notebooks/community/vision/image_warehouse_sdk.ipynb

@@ -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": {

notebooks/official/generative_ai/ai21labs_intro.ipynb

@@ -66,11 +66,16 @@
         "\n",
         "### AI21 Labs on Vertex AI\n",
         "\n",
-        "AI21 Labs models on Vertex AI offer fully managed and serverless models as managed APIs. To use an AI21 model on Vertex AI, send a request directly to the Vertex AI API endpoint.\n",
+        "AI21 Labs models on Vertex AI are available via fully managed and serverless models as managed APIs or self-deploy option with your own endpoint. To use an AI21 model on Vertex AI, send a request directly to the Vertex AI API endpoint.\n",
         "\n",
         "You can stream your model responses to reduce the end-user latency perception. A streamed response uses server-sent events (SSE) to incrementally stream the response.\n",
         "\n",
-        "### Available AI21 Labs models\n",
+        "### Available self-hosted AI21 Labs Models\n",
+        "\n",
+        "#### Jamba Large 1.6\n",
+        "AI21's most powerful instruction-tuned foundation model with 256K context window that's optimized for long-form input, superior accuracy, and speed.\n",
+        "\n",
+        "### Available fully managed AI21 Labs models\n",
         "\n",
         "#### Jamba 1.5 Mini\n",
         "AI21's small, powerful instruction-tuned foundation model with 256K context window that's optimized for long-form input, speed and cost efficiency.\n",
@@ -80,9 +85,9 @@
         "\n",
         "## Objective\n",
         "\n",
-        "This notebook demonstrates how to use the **Vertex AI API** to access the AI21 Jamba 1.5 Mini and Jamba 1.5 Large models on Vertex AI.\n",
+        "This notebook shows how to use Vertex AI API to use the AI21 Labs models.\n",
         "\n",
-        "For more information, see the [Use AI21 Labs](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/ai21) documentation.\n"
+        "For more information, see the [Use AI21 Labs](https://docs.ai21.com/reference/jamba-1-6-api-ref) documentation.\n"
       ]
     },
     {
@@ -191,6 +196,434 @@
         "    auth.authenticate_user()"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "AU0mJPV4hBqh"
+      },
+      "source": [
+        "### Self-Deploy AI21 Models"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "O2nAc89PhI5t"
+      },
+      "source": [
+        "#### Select one of AI21 Labs models"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "qzVPlFSFhSIw"
+      },
+      "outputs": [],
+      "source": [
+        "PUBLISHER_NAME = \"ai21\"  # @param {type:\"string\"}\n",
+        "PUBLISHER_MODEL_NAME = \"jamba-large-1.6\"  # @param [\"jamba-large-1.6\"]\n",
+        "\n",
+        "if PUBLISHER_MODEL_NAME == \"jamba-large-1.6\":\n",
+        "    available_regions = [\"us-central1\", \"us-east4\", \"europe-west4\"]"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "roMt5CG9hjhk"
+      },
+      "source": [
+        "#### Select a location and a version from the dropdown"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "20n9j-uLhl1w"
+      },
+      "outputs": [],
+      "source": [
+        "import ipywidgets as widgets\n",
+        "from IPython.display import display\n",
+        "\n",
+        "dropdown_loc = widgets.Dropdown(\n",
+        "    options=available_regions,\n",
+        "    description=\"Select a location:\",\n",
+        "    font_weight=\"bold\",\n",
+        "    style={\"description_width\": \"initial\"},\n",
+        ")\n",
+        "\n",
+        "\n",
+        "def dropdown_loc_eventhandler(change):\n",
+        "    global LOCATION\n",
+        "    if change[\"type\"] == \"change\" and change[\"name\"] == \"value\":\n",
+        "        LOCATION = change.new\n",
+        "        print(\"Selected:\", change.new)\n",
+        "\n",
+        "\n",
+        "LOCATION = dropdown_loc.value\n",
+        "dropdown_loc.observe(dropdown_loc_eventhandler, names=\"value\")\n",
+        "display(dropdown_loc)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "oLS4F4glhzXp"
+      },
+      "source": [
+        "#### Set Google Cloud project and endpoint information\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "02AKs1C9h09F"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "ENDPOINT = f\"https://{LOCATION}-aiplatform.googleapis.com\"\n",
+        "\n",
+        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
+        "    raise ValueError(\"Please set your PROJECT_ID\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tvs-Pvj5iCPV"
+      },
+      "source": [
+        "Import required libraries"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "xeBjRvYPiFh7"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "import time"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "IXy5Ch_DiHql"
+      },
+      "source": [
+        "#### Upload Model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ZcvQmwi4iKZI"
+      },
+      "outputs": [],
+      "source": [
+        "UPLOAD_MODEL_PAYLOAD = {\n",
+        "    \"model\": {\n",
+        "        \"displayName\": \"ModelGarden_LaunchPad_Model_\" + time.strftime(\"%Y%m%d-%H%M%S\"),\n",
+        "        \"baseModelSource\": {\n",
+        "            \"modelGardenSource\": {\n",
+        "                \"publicModelName\": f\"publishers/{PUBLISHER_NAME}/models/{PUBLISHER_MODEL_NAME}\",\n",
+        "            }\n",
+        "        },\n",
+        "    }\n",
+        "}\n",
+        "\n",
+        "request = json.dumps(UPLOAD_MODEL_PAYLOAD)\n",
+        "\n",
+        "! curl -X POST -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1beta1/projects/{PROJECT_ID}/locations/{LOCATION}/models:upload -d '{request}'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "mjt7mP2DiT3M"
+      },
+      "source": [
+        "#### Get Model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Sj-0k3Rqihy_"
+      },
+      "outputs": [],
+      "source": [
+        "MODEL_ID = [your-model-id]  # @param {type: \"number\"}\n",
+        "\n",
+        "! curl -X GET -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/models/{MODEL_ID}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0Ieb7Bc7ilH_"
+      },
+      "source": [
+        "#### Create the endpoint"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "MzdP-fJtimkE"
+      },
+      "outputs": [],
+      "source": [
+        "CREATE_ENDPOINT_PAYLOAD = {\n",
+        "    \"displayName\": \"ModelGarden_LaunchPad_Endpoint_\" + time.strftime(\"%Y%m%d-%H%M%S\"),\n",
+        "}\n",
+        "\n",
+        "request = json.dumps(CREATE_ENDPOINT_PAYLOAD)\n",
+        "\n",
+        "! curl -X POST -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints -d '{request}'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "qnAY0FA4irBi"
+      },
+      "source": [
+        "#### Get Endpoint"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Kg8IiVkrixrO"
+      },
+      "outputs": [],
+      "source": [
+        "ENDPOINT_ID = [your-endpoint-id]  # @param {type: \"number\"}\n",
+        "\n",
+        "! curl -X GET -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0HTovRCLi31p"
+      },
+      "source": [
+        "#### Deploy Model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ga1W0sJAi5gN"
+      },
+      "outputs": [],
+      "source": [
+        "MACHINE_TYPE = \"a3-highgpu-8g\"  # @param {type: \"string\"}\n",
+        "ACCELERATOR_TYPE = \"NVIDIA_H100_80GB\"  # @param {type: \"string\"}\n",
+        "ACCELERATOR_COUNT = 8  # @param {type: \"number\"}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "BjgAAarxi84D"
+      },
+      "outputs": [],
+      "source": [
+        "# Try CURL, if it fails, use stubby command in the next cell.\n",
+        "DEPLOY_PAYLOAD = {\n",
+        "    \"deployedModel\": {\n",
+        "        \"model\": f\"projects/{PROJECT_ID}/locations/{LOCATION}/models/{MODEL_ID}\",\n",
+        "        \"displayName\": \"ModelGarden_LaunchPad_DeployedModel_\"\n",
+        "        + time.strftime(\"%Y%m%d-%H%M%S\"),\n",
+        "        \"dedicatedResources\": {\n",
+        "            \"machineSpec\": {\n",
+        "                \"machineType\": MACHINE_TYPE,\n",
+        "                \"acceleratorType\": ACCELERATOR_TYPE,\n",
+        "                \"acceleratorCount\": ACCELERATOR_COUNT,\n",
+        "            },\n",
+        "            \"minReplicaCount\": 1,\n",
+        "            \"maxReplicaCount\": 1,\n",
+        "        },\n",
+        "    },\n",
+        "    \"trafficSplit\": {\"0\": 100},\n",
+        "}\n",
+        "\n",
+        "request = json.dumps(DEPLOY_PAYLOAD)\n",
+        "print(\"Request payload to Deploy Model:\")\n",
+        "print(json.dumps(DEPLOY_PAYLOAD, indent=2))\n",
+        "print(\"\\nResult:\")\n",
+        "! curl -X POST -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}:deployModel -d '{request}'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "lfCNC2CgjAYv"
+      },
+      "source": [
+        "#### Text generation"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ym3wTSD8jGWk"
+      },
+      "source": [
+        "##### Unary call\n",
+        "\n",
+        "Sends a POST request to the specified API endpoint to get a response from the model for a joke using the provided payload."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bJ3cSMtVjCkH"
+      },
+      "outputs": [],
+      "source": [
+        "PAYLOAD = {\n",
+        "    \"messages\": [{\"role\": \"user\", \"content\": \"Tell me a joke about whales\"}],\n",
+        "    \"max_tokens\": 100\n",
+        "}\n",
+        "\n",
+        "request = json.dumps(PAYLOAD)\n",
+        "\n",
+        "!curl -X POST \\\n",
+        "  -H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
+        "  -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}:rawPredict \\\n",
+        "  -d '{request}'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "r0eh_cH7jbGS"
+      },
+      "source": [
+        "##### Streaming call\n",
+        "\n",
+        "Sends a POST request to the specified API endpoint to stream a response from the model for a sports T-Shirt product title using provided payload."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "3F2nMlsOjVO1"
+      },
+      "outputs": [],
+      "source": [
+        "PAYLOAD = {\n",
+        "    \"messages\": [{\"role\": \"user\", \"content\": \"Write a product title for a sports T-Shirt to be published for online retail. Include these keywords: activewear, gym, dryfit.\"}],\n",
+        "    \"max_tokens\": 100\n",
+        "}\n",
+        "\n",
+        "request = json.dumps(PAYLOAD)\n",
+        "!curl -X POST \\\n",
+        "  -H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
+        "  -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}:streamRawPredict \\\n",
+        "  -d '{request}'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "IHjWW58yjUDk"
+      },
+      "source": [
+        "#### Cleaning up"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "OfvTKqj7jpBK"
+      },
+      "source": [
+        "##### Cleaning up deployment resources\n",
+        "To clean up all Google Cloud resources used in this notebook, you can delete the Google Cloud project you used for the tutorial.\n",
+        "\n",
+        "Otherwise, you can delete the individual resources you created in this tutorial:\n",
+        "\n",
+        "* Model\n",
+        "* Endpoint\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "1OjrDpV9j1st"
+      },
+      "source": [
+        "##### Undeploy model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "d9ocoA9Ej44B"
+      },
+      "outputs": [],
+      "source": [
+        "UNDEPLOY_PAYLOAD = {\"deployedModelId\": MODEL_ID}\n",
+        "\n",
+        "request = json.dumps(UNDEPLOY_PAYLOAD)\n",
+        "\n",
+        "! curl -X POST -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}:undeployModel -d '{request}'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "R6lBNu14j7zi"
+      },
+      "source": [
+        "##### Delete Endpoint"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "5cDp4Pa3j9rN"
+      },
+      "outputs": [],
+      "source": [
+        "! curl -X DELETE -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "1qsf1MKCgrF9"
+      },
+      "source": [
+        "### Fully Managed AI21 Models"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -310,7 +743,7 @@
         "id": "5ahw-uFjCAbo"
       },
       "source": [
-        "### Text generation"
+        "#### Text generation"
       ]
     },
     {
@@ -319,7 +752,7 @@
         "id": "61107099357a"
       },
       "source": [
-        "#### Unary call\n",
+        "##### Unary call\n",
         "\n",
         "Sends a POST request to the specified API endpoint to get a response from the model for a joke using the provided payload."
       ]
@@ -354,7 +787,7 @@
         "id": "e6f52fae9379"
       },
       "source": [
-        "#### Streaming call\n",
+        "##### Streaming call\n",
         "\n",
         "Sends a POST request to the specified API endpoint to stream a response from the model for a sports T-Shirt product title using provided payload."
       ]

notebooks/official/generative_ai/cambai_intro.ipynb

@@ -0,0 +1,815 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "9A9NkTRTfo2I"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2024 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5HQofEWl0LZp"
+      },
+      "source": [
+        "# Getting Started with CAMB.AI Models\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/generative_ai/cambai_intro.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\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%2Fofficial/generative_ai/cambai_intro.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/generative_ai/cambai_intro.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/generative_ai/cambai_intro.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8fK_rdvvx1iZ"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "### Camb AI on Vertex AI\n",
+        "\n",
+        "You can deploy the Camb AI models in your own endpoint.\n",
+        "\n",
+        "### Available Camb AI models\n",
+        "\n",
+        "#### MARS7\n",
+        "(Multilingual AutoRegressive Speech 7) is the latest generation in CAMB.AI’s MARS series of speech synthesis models. MARS7 creates hyper-realistic, prosodic, multilingual text-to-speech (TTS) outputs, featuring optional voice cloning and fine-grained emotional control. Deployed on Google Cloud’s Vertex AI Launchpad, it brings near real-time latency, with an architectural focus on parameter efficiency and global context understanding.\n",
+        "\n",
+        "## Objective\n",
+        "This notebook shows how to use **Vertex AI API** to deploy the Camb AI models.\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "nwYvaaW25jYS"
+      },
+      "source": [
+        "## Get Started\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5FO67G8c0LZq"
+      },
+      "source": [
+        "### Install required packages\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 1,
+      "metadata": {
+        "id": "cvwSgvqQ0LZq"
+      },
+      "outputs": [],
+      "source": [
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 2,
+      "metadata": {
+        "id": "ZpqsCL0f0LZq"
+      },
+      "outputs": [],
+      "source": [
+        "! pip3 install -U -q httpx"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Terpy8VeVIN6"
+      },
+      "outputs": [],
+      "source": [
+        "! apt-get install jq -y"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "b9f4c57a43f6"
+      },
+      "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": 4,
+      "metadata": {
+        "id": "3b9119a60525"
+      },
+      "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": "e767418763cd"
+      },
+      "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": {
+        "id": "6a5bea26f60f"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
+        "\n",
+        "Authenticate your environment on Google Colab.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 1,
+      "metadata": {
+        "id": "c97be6a73155"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2fxZn4SAbxdl"
+      },
+      "source": [
+        "### Select one of Camb AI models"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 2,
+      "metadata": {
+        "id": "Y8X70FTSbx7U"
+      },
+      "outputs": [],
+      "source": [
+        "PUBLISHER_NAME = \"cambai\"  # @param {type:\"string\"}\n",
+        "PUBLISHER_MODEL_NAME = \"mars7\"\n",
+        "available_regions = [\"us-central1\"]"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bpuX3sKtexlK"
+      },
+      "source": [
+        "### Select a location and a version from the dropdown"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 3,
+      "metadata": {
+        "id": "dHl8xW45ex_O"
+      },
+      "outputs": [],
+      "source": [
+        "import ipywidgets as widgets\n",
+        "from IPython.display import display\n",
+        "\n",
+        "dropdown_loc = widgets.Dropdown(\n",
+        "    options=available_regions,\n",
+        "    description=\"Select a location:\",\n",
+        "    font_weight=\"bold\",\n",
+        "    style={\"description_width\": \"initial\"},\n",
+        ")\n",
+        "\n",
+        "\n",
+        "def dropdown_loc_eventhandler(change):\n",
+        "    global LOCATION\n",
+        "    if change[\"type\"] == \"change\" and change[\"name\"] == \"value\":\n",
+        "        LOCATION = change.new\n",
+        "        print(\"Selected:\", change.new)\n",
+        "\n",
+        "\n",
+        "LOCATION = dropdown_loc.value\n",
+        "dropdown_loc.observe(dropdown_loc_eventhandler, names=\"value\")\n",
+        "display(dropdown_loc)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "JZQFZJOY0LZr"
+      },
+      "source": [
+        "### Set Google Cloud project information\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 11,
+      "metadata": {
+        "id": "Q7jkRBhB0LZr"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "ENDPOINT = f\"https://{LOCATION}-aiplatform.googleapis.com\"\n",
+        "\n",
+        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
+        "    raise ValueError(\"Please set your PROJECT_ID\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4NAstKRFBt4N"
+      },
+      "source": [
+        "### Import required libraries"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 5,
+      "metadata": {
+        "id": "QZEFLE6a6bqy"
+      },
+      "outputs": [],
+      "source": [
+        "import base64\n",
+        "import json\n",
+        "import time"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "_FI94lP20LZr"
+      },
+      "source": [
+        "## Using Vertex AI API"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "qjsDpa8jlTRu"
+      },
+      "source": [
+        "### Upload Model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "y1R2BRsBlu-k"
+      },
+      "outputs": [],
+      "source": [
+        "UPLOAD_MODEL_PAYLOAD = {\n",
+        "    \"model\": {\n",
+        "        \"displayName\": \"ModelGarden_LaunchPad_Model_\" + time.strftime(\"%Y%m%d-%H%M%S\"),\n",
+        "        \"baseModelSource\": {\n",
+        "            \"modelGardenSource\": {\n",
+        "                \"publicModelName\": f\"publishers/{PUBLISHER_NAME}/models/{PUBLISHER_MODEL_NAME}\",\n",
+        "            }\n",
+        "        },\n",
+        "    }\n",
+        "}\n",
+        "\n",
+        "request = json.dumps(UPLOAD_MODEL_PAYLOAD)\n",
+        "\n",
+        "! curl -X POST -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1beta1/projects/{PROJECT_ID}/locations/{LOCATION}/models:upload -d '{request}'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "V2j0nVGwlf9b"
+      },
+      "source": [
+        "### Get Model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bxwM0GXTmQhh"
+      },
+      "outputs": [],
+      "source": [
+        "MODEL_ID = [your - model - id]  # @param {type: \"number\"}\n",
+        "\n",
+        "! curl -X GET -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/models/{MODEL_ID}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3q3ygq8VlZAp"
+      },
+      "source": [
+        "### Create Endpoint"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "O1ChDOt7mPBQ"
+      },
+      "outputs": [],
+      "source": [
+        "CREATE_ENDPOINT_PAYLOAD = {\n",
+        "    \"displayName\": \"ModelGarden_LaunchPad_Endpoint_\" + time.strftime(\"%Y%m%d-%H%M%S\"),\n",
+        "}\n",
+        "\n",
+        "request = json.dumps(CREATE_ENDPOINT_PAYLOAD)\n",
+        "\n",
+        "! curl -X POST -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints -d '{request}'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "GuMZCdhmlpCE"
+      },
+      "source": [
+        "### Get Endpoint"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tHq_cLT6mPp_"
+      },
+      "outputs": [],
+      "source": [
+        "ENDPOINT_ID = [your - endpoint - id]  # @param {type: \"number\"}\n",
+        "\n",
+        "! curl -X GET -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "G0amEPXolbP7"
+      },
+      "source": [
+        "### Deploy Model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 19,
+      "metadata": {
+        "id": "Ucj-Xa-fpGrg"
+      },
+      "outputs": [],
+      "source": [
+        "MACHINE_TYPE = \"g2-standard-4\"  # @param {type: \"string\"}\n",
+        "ACCELERATOR_TYPE = \"NVIDIA_L4\"  # @param {type: \"string\"}\n",
+        "ACCELERATOR_COUNT = 1  # @param {type: \"number\"}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "VGTyCQQhlrAR"
+      },
+      "outputs": [],
+      "source": [
+        "DEPLOY_PAYLOAD = {\n",
+        "    \"deployedModel\": {\n",
+        "        \"model\": f\"projects/{PROJECT_ID}/locations/{LOCATION}/models/{MODEL_ID}\",\n",
+        "        \"displayName\": \"ModelGarden_LaunchPad_DeployedModel_\"\n",
+        "        + time.strftime(\"%Y%m%d-%H%M%S\"),\n",
+        "        \"dedicatedResources\": {\n",
+        "            \"machineSpec\": {\n",
+        "                \"machineType\": MACHINE_TYPE,\n",
+        "                \"acceleratorType\": ACCELERATOR_TYPE,\n",
+        "                \"acceleratorCount\": ACCELERATOR_COUNT,\n",
+        "            },\n",
+        "            \"minReplicaCount\": 1,\n",
+        "            \"maxReplicaCount\": 1,\n",
+        "        },\n",
+        "    },\n",
+        "    \"trafficSplit\": {\"0\": 100},\n",
+        "}\n",
+        "\n",
+        "request = json.dumps(DEPLOY_PAYLOAD)\n",
+        "print(\"Request payload to Deploy Model:\")\n",
+        "print(json.dumps(DEPLOY_PAYLOAD, indent=2))\n",
+        "print(\"\\nResult:\")\n",
+        "! curl -X POST -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}:deployModel -d '{request}'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5ahw-uFjCAbo"
+      },
+      "source": [
+        "### Generate audio"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "61107099357a"
+      },
+      "source": [
+        "#### Unary call\n",
+        "\n",
+        "Sends a POST request to the specified API endpoint to get a response from the model for a joke using the provided payload."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "EOkEkRBiDUpH"
+      },
+      "outputs": [],
+      "source": [
+        "!curl -o ref.flac \"https://storage.googleapis.com/cambai-prod-public/public/ref.flac\"\n",
+        "print(\"Downloaded ref.flac\")\n",
+        "\n",
+        "text = \"The quick brown fox jumps over the lazy dog\"\n",
+        "audio_ref_path = \"ref.flac\"\n",
+        "ref_text = None\n",
+        "language = \"en-us\"\n",
+        "\n",
+        "# Prepare the JSON payload in a separate file\n",
+        "json_payload_file = \"payload.json\"\n",
+        "\n",
+        "with open(audio_ref_path, \"rb\") as audio_file:\n",
+        "    encoded_string = base64.b64encode(audio_file.read()).decode(\"utf-8\")\n",
+        "\n",
+        "payload = {\n",
+        "    \"text\": text,\n",
+        "    \"audio_ref\": encoded_string,\n",
+        "    \"ref_text\": ref_text,\n",
+        "    \"language\": language,\n",
+        "}\n",
+        "\n",
+        "with open(json_payload_file, \"w\") as f:\n",
+        "    json.dump(payload, f, indent=2)\n",
+        "\n",
+        "save_output_cmd = f\"\"\"\n",
+        "curl -X POST {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}:rawPredict \\\\\n",
+        "        -H \"Content-Type: application/json\" \\\\\n",
+        "        -H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\\\n",
+        "        -d @{json_payload_file} > response.json\n",
+        "\n",
+        "# save audio\n",
+        "cat response.json | jq -r '.predictions[0]' | base64 -d > output.flac\n",
+        "\"\"\"\n",
+        "print(save_output_cmd)\n",
+        "!{save_output_cmd}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import soundfile as sf\n",
+        "import tempfile\n",
+        "from IPython.display import Audio\n",
+        "\n",
+        "def display_flac_in_notebook(path_to_flac):\n",
+        "    data, samplerate = sf.read(path_to_flac)\n",
+        "    with tempfile.NamedTemporaryFile(suffix=\".wav\", delete=False) as tmp:\n",
+        "        sf.write(tmp.name, data, samplerate, format='WAV')\n",
+        "        return Audio(filename=tmp.name)\n",
+        "\n",
+        "display_flac_in_notebook(\"output.flac\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "x2r2YeT70LZt"
+      },
+      "source": [
+        "## Using Vertex AI SDK for *Python*"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "kNons5WV0LZt"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "5RaYIOoR0LZt"
+      },
+      "outputs": [],
+      "source": [
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4a4U1cTF0LZt"
+      },
+      "source": [
+        "### Upload Model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "082M23Nb0LaA"
+      },
+      "outputs": [],
+      "source": [
+        "model = aiplatform.Model.upload(\n",
+        "    display_name=\"ModelGarden_LaunchPad_Endpoint_\" + time.strftime(\"%Y%m%d-%H%M%S\"),\n",
+        "    model_garden_source_model_name=f\"publishers/{PUBLISHER_NAME}/models/{PUBLISHER_MODEL_NAME}\",\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-r3_PPU00LaA"
+      },
+      "source": [
+        "### Create Endpoint"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "LvRKh7dl0LaB"
+      },
+      "outputs": [],
+      "source": [
+        "my_endpoint = aiplatform.Endpoint.create(\n",
+        "    display_name=\"ModelGarden_LaunchPad_Endpoint_\" + time.strftime(\"%Y%m%d-%H%M%S\")\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "IW9tGPMc0LaB"
+      },
+      "source": [
+        "### Deploy Model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "PMIsYtTr0LaB"
+      },
+      "outputs": [],
+      "source": [
+        "MACHINE_TYPE = \"g2-standard-4\"  # @param {type: \"string\"}\n",
+        "ACCELERATOR_TYPE = \"NVIDIA_L4\"  # @param {type: \"string\"}\n",
+        "ACCELERATOR_COUNT = 1  # @param {type: \"number\"}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "AghmHYof0LaB"
+      },
+      "outputs": [],
+      "source": [
+        "model.deploy(\n",
+        "    endpoint=my_endpoint,\n",
+        "    deployed_model_display_name=\"ModelGarden_LaunchPad_DeployedModel_\"\n",
+        "    + time.strftime(\"%Y%m%d-%H%M%S\"),\n",
+        "    traffic_split={\"0\": 100},\n",
+        "    machine_type=MACHINE_TYPE,\n",
+        "    accelerator_type=ACCELERATOR_TYPE,\n",
+        "    accelerator_count=ACCELERATOR_COUNT,\n",
+        "    min_replica_count=1,\n",
+        "    max_replica_count=1,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "optmPVWs0LaB"
+      },
+      "source": [
+        "### Generate audio"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Stp-MvTy0LaB"
+      },
+      "source": [
+        "#### Unary call"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "rz1n3gfW0LaB"
+      },
+      "outputs": [],
+      "source": [
+        "def predict_custom_trained_model_sample(\n",
+        "    instances: dict, target_file: str = \"output.flac\"\n",
+        "):\n",
+        "    # If you created your endpoint using a different method you can set the value by uncommenting the next line and making sure ENDPOINT_ID is set\n",
+        "    # my_endpoint = aiplatform.Endpoint(f'{ENDPOINT_ID}')\n",
+        "    DATA = {\"instances\": [instances]}\n",
+        "    response = my_endpoint.raw_predict(\n",
+        "        body=json.dumps(DATA).encode(\"utf-8\"),\n",
+        "        headers={\"Content-Type\": \"application/json\"},\n",
+        "    )\n",
+        "    with open(target_file, \"wb\") as f:\n",
+        "        audio_bytes = base64.b64decode(json.loads(response.content)[\"predictions\"][0])\n",
+        "        f.write(audio_bytes)\n",
+        "\n",
+        "\n",
+        "with open(\"ref.flac\", \"rb\") as f:\n",
+        "    audio_ref_bytes = base64.b64encode(f.read()).decode(\"utf-8\")\n",
+        "\n",
+        "predict_custom_trained_model_sample(\n",
+        "    instances={\n",
+        "        \"text\": \"The quick brown fox jumps over the lazy dog\",\n",
+        "        \"audio_ref\": audio_ref_bytes,\n",
+        "        \"ref_text\": None,\n",
+        "        \"language\": \"en-us\",\n",
+        "    },\n",
+        "    target_file=\"output.flac\",\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "UJyNltvFJeF9"
+      },
+      "source": [
+        "## Cleaning up"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dCdn4H-jJlNP"
+      },
+      "source": [
+        "##### Cleaning up deployment resources\n",
+        "To clean up all Google Cloud resources used in this notebook, you can delete the Google Cloud project you used for the tutorial.\n",
+        "\n",
+        "Otherwise, you can delete the individual resources you created in this tutorial:\n",
+        "\n",
+        "* Model\n",
+        "* Endpoint\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "00wtIF8lJpKS"
+      },
+      "source": [
+        "##### Undeploy model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ovQSxWqfJmMs"
+      },
+      "outputs": [],
+      "source": [
+        "UNDEPLOY_PAYLOAD = {\"deployedModelId\": MODEL_ID}\n",
+        "\n",
+        "request = json.dumps(UNDEPLOY_PAYLOAD)\n",
+        "\n",
+        "! curl -X POST -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}:undeployModel -d '{request}'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "nQ8ERh4uJvr4"
+      },
+      "source": [
+        "##### Delete Endpoint"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "zj37lITNJy-u"
+      },
+      "outputs": [],
+      "source": [
+        "! curl -X DELETE -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "cambai_intro.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/official/generative_ai/qodo_intro.ipynb

@@ -0,0 +1,984 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "9A9NkTRTfo2I"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2024 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "b9f8293b0643"
+      },
+      "source": [
+        "# Getting Started with qodo Models\n",
+        "\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/generative_ai/Qodo_intro.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\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%2Fnotebook_template.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/generative_ai/Qodo_intro.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/generative_ai/Qodo_intro.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8fK_rdvvx1iZ"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates how to deploy and use Qodo's state-of-the-art code embedding models on Google Cloud's Vertex AI platform. You'll learn how to set up, deploy, and make predictions with these specialized embedding models that enhance code retrieval and search capabilities.\n",
+        "In this notebook, you will:\n",
+        "\n",
+        "Set up your Google Cloud environment and initialize the Vertex AI SDK\n",
+        "Upload a Qodo model to your Vertex AI Model Registry\n",
+        "Create a Vertex AI endpoint for model deployment\n",
+        "Deploy the Qodo model to your endpoint with appropriate compute resources\n",
+        "Make predictions using the deployed model.\n",
+        "\n",
+        "### Qodo on Vertex AI\n",
+        "\n",
+        "You can deploy the Qodo models in your own endpoint.\n",
+        "\n",
+        "\n",
+        "\n",
+        "### Available Qodo models\n",
+        "\n",
+        "#### Qodo-Embed-1-7B\n",
+        "Qodo-Embed-1-7B is a state-of-the-art code embedding model for efficient code & text retrieval, enhancing the search accuracy of RAG methods.\n",
+        "\n",
+        "\n",
+        "\n",
+        "## Objective\n",
+        "\n",
+        "This notebook shows how to use **Vertex AI API** to deploy the qodo models.\n",
+        "\n",
+        "For more information, see the [qodo website](https://www.qodo.ai/blog/qodo-embed-1-code-embedding-code-retreival/).\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "nwYvaaW25jYS"
+      },
+      "source": [
+        "## Get Started\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "d10e8895d2d4"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python or other required packages\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "08dd6d2ac629"
+      },
+      "outputs": [],
+      "source": [
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "754611260f53"
+      },
+      "outputs": [],
+      "source": [
+        "! pip3 install -U -q httpx"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "b9f4c57a43f6"
+      },
+      "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": "3b9119a60525"
+      },
+      "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": "6a5bea26f60f"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
+        "\n",
+        "Authenticate your environment on Google Colab.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "c97be6a73155"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Y8X70FTSbx7U"
+      },
+      "outputs": [],
+      "source": [
+        "PUBLISHER_NAME = \"qodo\"  # @param {type:\"string\"}\n",
+        "PUBLISHER_MODEL_NAME = \"qodo-embed-1-7b-v1\"  # @param [\"publisher-model-name-1\", \"publisher-model-name-2\", \"test-marketplace-publisher-model-e2e-01\"]\n",
+        "\n",
+        "available_regions = [\"us-central1\", \"europe-west4\"]"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bpuX3sKtexlK"
+      },
+      "source": [
+        "### Select a location and a version from the dropdown"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "dHl8xW45ex_O"
+      },
+      "outputs": [],
+      "source": [
+        "import ipywidgets as widgets\n",
+        "from IPython.display import display\n",
+        "\n",
+        "dropdown_loc = widgets.Dropdown(\n",
+        "    options=available_regions,\n",
+        "    description=\"Select a location:\",\n",
+        "    font_weight=\"bold\",\n",
+        "    style={\"description_width\": \"initial\"},\n",
+        ")\n",
+        "\n",
+        "\n",
+        "def dropdown_loc_eventhandler(change):\n",
+        "    global LOCATION\n",
+        "    if change[\"type\"] == \"change\" and change[\"name\"] == \"value\":\n",
+        "        LOCATION = change.new\n",
+        "        print(\"Selected:\", change.new)\n",
+        "\n",
+        "\n",
+        "LOCATION = dropdown_loc.value\n",
+        "dropdown_loc.observe(dropdown_loc_eventhandler, names=\"value\")\n",
+        "display(dropdown_loc)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4f872cd812d0"
+      },
+      "source": [
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ce2e2765bc2d"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "ENDPOINT = f\"https://{LOCATION}-aiplatform.googleapis.com\"\n",
+        "\n",
+        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
+        "    raise ValueError(\"Please set your PROJECT_ID\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4NAstKRFBt4N"
+      },
+      "source": [
+        "### Import required libraries"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "QZEFLE6a6bqy"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "import time"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4fa6f083d253"
+      },
+      "source": [
+        "## Using Vertex AI API"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "qjsDpa8jlTRu"
+      },
+      "source": [
+        "### Upload Model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "y1R2BRsBlu-k"
+      },
+      "outputs": [],
+      "source": [
+        "UPLOAD_MODEL_PAYLOAD = {\n",
+        "    \"model\": {\n",
+        "        \"displayName\": \"ModelGarden_LaunchPad_Model_\" + time.strftime(\"%Y%m%d-%H%M%S\"),\n",
+        "        \"baseModelSource\": {\n",
+        "            \"modelGardenSource\": {\n",
+        "                \"publicModelName\": f\"publishers/{PUBLISHER_NAME}/models/{PUBLISHER_MODEL_NAME}\",\n",
+        "            }\n",
+        "        },\n",
+        "    }\n",
+        "}\n",
+        "\n",
+        "request = json.dumps(UPLOAD_MODEL_PAYLOAD)\n",
+        "\n",
+        "! curl -X POST -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1beta1/projects/{PROJECT_ID}/locations/{LOCATION}/models:upload -d '{request}'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6afd1e782a3f"
+      },
+      "source": [
+        "## Extract the Model ID\n",
+        "\n",
+        "After uploading your model to Vertex AI, you'll need to extract the model ID from the response for use in subsequent steps.\n",
+        "\n",
+        "The response from the upload command will look similar to this:\n",
+        "\n",
+        "```json\n",
+        "{\n",
+        "  \"name\": \"projects/123456789/locations/us-central1/models/9876543210/operations/1122334455\",\n",
+        "  \"metadata\": {\n",
+        "    \"@type\": \"type.googleapis.com/google.cloud.aiplatform.v1beta1.UploadModelOperationMetadata\",\n",
+        "    \"genericMetadata\": {\n",
+        "      \"createTime\": \"2025-04-07T16:47:27.076450Z\",\n",
+        "      \"updateTime\": \"2025-04-07T16:47:27.076450Z\"\n",
+        "    }\n",
+        "  }\n",
+        "}\n",
+        "```\n",
+        "\n",
+        "Your **model ID** is the number between `models/` and `/operations` in the \"name\" field.\n",
+        "\n",
+        "In the example above, the model ID would be `9876543210`."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "76cba0adc39c"
+      },
+      "source": "Extract the model ID from the response"
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "V2j0nVGwlf9b"
+      },
+      "source": [
+        "### Verify Your Model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bxwM0GXTmQhh"
+      },
+      "outputs": [],
+      "source": [
+        "MODEL_ID = \"[extracted_model_id]\"  # @param {type: \"number\"}\n",
+        "\n",
+        "! curl -X GET -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/models/{MODEL_ID}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3q3ygq8VlZAp"
+      },
+      "source": [
+        "### Create an Endpoint and Extract the Endpoint ID"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "O1ChDOt7mPBQ"
+      },
+      "outputs": [],
+      "source": [
+        "CREATE_ENDPOINT_PAYLOAD = {\n",
+        "    \"displayName\": \"ModelGarden_LaunchPad_Endpoint_\" + time.strftime(\"%Y%m%d-%H%M%S\"),\n",
+        "}\n",
+        "\n",
+        "request = json.dumps(CREATE_ENDPOINT_PAYLOAD)\n",
+        "\n",
+        "! curl -X POST -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints -d '{request}'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "400ec43e8a5c"
+      },
+      "source": [
+        "### Extracting the Endpoint ID\n",
+        "\n",
+        " After running the command above, you'll see a JSON response similar to:\n",
+        "```json\n",
+        "{\n",
+        "  \"name\": \"projects/PROJECT_NUMBER/locations/LOCATION/endpoints/ENDPOINT_ID/operations/OPERATION_ID\",\n",
+        "  \"metadata\": {\n",
+        "    \"@type\": \"type.googleapis.com/google.cloud.aiplatform.v1.CreateEndpointOperationMetadata\",\n",
+        "    \"genericMetadata\": {\n",
+        "      \"createTime\": \"2025-04-07T16:55:27.076450Z\",\n",
+        "      \"updateTime\": \"2025-04-07T16:55:27.076450Z\"\n",
+        "    }\n",
+        "  }\n",
+        "}\n",
+        "```\n",
+        "\n",
+        "Your endpoint ID is the number that appears after \"endpoints/\" and before \"/operations\" in the \"name\" field.\n",
+        "\n",
+        "For example, if the \"name\" field shows:\n",
+        "\"projects/123456789/locations/us-central1/endpoints/9876543210/operations/1122334455\"\n",
+        "\n",
+        "Then your endpoint ID is: 9876543210"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "GuMZCdhmlpCE"
+      },
+      "source": [
+        "### Verify Your Endpoint"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tHq_cLT6mPp_"
+      },
+      "outputs": [],
+      "source": [
+        "ENDPOINT_ID = \"[extracted_endpoint_id]\"  # @param {type: \"number\"}\n",
+        "\n",
+        "! curl -X GET -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "G0amEPXolbP7"
+      },
+      "source": [
+        "### Deploy Model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Ucj-Xa-fpGrg"
+      },
+      "outputs": [],
+      "source": [
+        "import ipywidgets as widgets\n",
+        "from IPython.display import display\n",
+        "\n",
+        "# Initial data\n",
+        "PUBLISHER_NAME = \"qodo\"  # @param {type:\"string\"}\n",
+        "PUBLISHER_MODEL_NAME = \"qodo-embed-1-7b-v1\"\n",
+        "available_regions = [\"us-central1\", \"europe-west4\"]\n",
+        "compatible_machines = [\n",
+        "    \"a2-highgpu-1g\",\n",
+        "    \"a2-highgpu-4g\",\n",
+        "    \"a2-ultragpu-1g\",\n",
+        "    \"a2-ultragpu-2g\",\n",
+        "    \"a3-highgpu-2g\",\n",
+        "]\n",
+        "\n",
+        "# Machine type to accelerator mapping (you can customize this based on your needs)\n",
+        "machine_config = {\n",
+        "    \"a2-highgpu-1g\": {\"type\": \"NVIDIA_A100\", \"count\": 1},\n",
+        "    \"a2-highgpu-4g\": {\"type\": \"NVIDIA_A100\", \"count\": 4},\n",
+        "    \"a2-ultragpu-1g\": {\"type\": \"NVIDIA_A100_80GB\", \"count\": 1},\n",
+        "    \"a2-ultragpu-2g\": {\"type\": \"NVIDIA_A100_80GB\", \"count\": 2},\n",
+        "    \"a3-highgpu-2g\": {\"type\": \"NVIDIA_H100\", \"count\": 2},\n",
+        "}\n",
+        "\n",
+        "# Create widgets\n",
+        "\n",
+        "\n",
+        "dropdown_machine = widgets.Dropdown(\n",
+        "    options=compatible_machines,\n",
+        "    description=\"Machine type:\",\n",
+        "    font_weight=\"bold\",\n",
+        "    style={\"description_width\": \"initial\"},\n",
+        ")\n",
+        "\n",
+        "label_accelerator_type = widgets.HTML(\n",
+        "    value=f\"<b>Accelerator type:</b> {machine_config[compatible_machines[0]]['type']}\"\n",
+        ")\n",
+        "\n",
+        "label_accelerator_count = widgets.HTML(\n",
+        "    value=f\"<b>Accelerator count:</b> {machine_config[compatible_machines[0]]['count']}\"\n",
+        ")\n",
+        "\n",
+        "# Event handlers\n",
+        "\n",
+        "\n",
+        "def dropdown_machine_eventhandler(change):\n",
+        "    global MACHINE_TYPE, ACCELERATOR_TYPE, ACCELERATOR_COUNT\n",
+        "    if change[\"type\"] == \"change\" and change[\"name\"] == \"value\":\n",
+        "        MACHINE_TYPE = change.new\n",
+        "        machine_info = machine_config.get(change.new, {\"type\": \"Unknown\", \"count\": 0})\n",
+        "        ACCELERATOR_TYPE = machine_info[\"type\"]\n",
+        "        ACCELERATOR_COUNT = machine_info[\"count\"]\n",
+        "\n",
+        "        # Update the displayed information\n",
+        "        label_accelerator_type.value = f\"<b>Accelerator type:</b> {ACCELERATOR_TYPE}\"\n",
+        "        label_accelerator_count.value = f\"<b>Accelerator count:</b> {ACCELERATOR_COUNT}\"\n",
+        "\n",
+        "\n",
+        "# Initialize global variables\n",
+        "MACHINE_TYPE = dropdown_machine.value\n",
+        "ACCELERATOR_TYPE = machine_config[MACHINE_TYPE][\"type\"]\n",
+        "ACCELERATOR_COUNT = machine_config[MACHINE_TYPE][\"count\"]\n",
+        "\n",
+        "# Set up observers\n",
+        "dropdown_machine.observe(dropdown_machine_eventhandler, names=\"value\")\n",
+        "\n",
+        "# Display widgets\n",
+        "display(\n",
+        "    widgets.VBox([dropdown_machine, label_accelerator_type, label_accelerator_count])\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "VGTyCQQhlrAR"
+      },
+      "outputs": [],
+      "source": [
+        "DEPLOY_PAYLOAD = {\n",
+        "    \"deployedModel\": {\n",
+        "        \"model\": f\"projects/{PROJECT_ID}/locations/{LOCATION}/models/{MODEL_ID}\",\n",
+        "        \"displayName\": \"ModelGarden_LaunchPad_DeployedModel_\"\n",
+        "        + time.strftime(\"%Y%m%d-%H%M%S\"),\n",
+        "        \"dedicatedResources\": {\n",
+        "            \"machineSpec\": {\n",
+        "                \"machineType\": MACHINE_TYPE,\n",
+        "                \"acceleratorType\": ACCELERATOR_TYPE,\n",
+        "                \"acceleratorCount\": ACCELERATOR_COUNT,\n",
+        "            },\n",
+        "            \"minReplicaCount\": 1,\n",
+        "            \"maxReplicaCount\": 1,\n",
+        "        },\n",
+        "    },\n",
+        "    \"trafficSplit\": {\"0\": 100},\n",
+        "}\n",
+        "\n",
+        "request = json.dumps(DEPLOY_PAYLOAD)\n",
+        "\n",
+        "! curl -X POST -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}:deployModel -d '{request}'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "21399a078364"
+      },
+      "source": [
+        "### Wait for Deployment to Complete\n",
+        "\n",
+        "Deployment can take several minutes. This cell will check the status of the operation."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "d95cc714a451"
+      },
+      "source": [
+        "### Extracting the operation ID\n",
+        "\n",
+        " After running the command above, you'll see a JSON response similar to:\n",
+        "```json\n",
+        "{\n",
+        "  \"name\": \"projects/513257720056/locations/us-central1/endpoints/3978337634014461952/operations/2501704616106786816\",\n",
+        "  \"metadata\": {\n",
+        "    \"@type\": \"type.googleapis.com/google.cloud.aiplatform.v1.DeployModelOperationMetadata\",\n",
+        "    \"genericMetadata\": {\n",
+        "      \"createTime\": \"2025-04-07T17:10:55.383719Z\",\n",
+        "      \"updateTime\": \"2025-04-07T17:10:55.383719Z\"\n",
+        "    }\n",
+        "  }\n",
+        "}\n",
+        "```\n",
+        "\n",
+        "Your operation ID is the number that appears after \"operation/\" in the \"name\" field.\n",
+        "\n",
+        "For example, if the \"name\" field shows:\n",
+        "\"projects/123456789/locations/us-central1/endpoints/9876543210/operations/1122334455\"\n",
+        "\n",
+        "Then your endpoint ID is: 1122334455"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "7dc1861dbde8"
+      },
+      "source": [
+        "###  Check Operation Status\n",
+        "\n",
+        "Run this cell to check the current status of the deployment operation. You may need to run this cell multiple times until the operation is complete (`\"done\": true`).\n",
+        "\n",
+        "**Note:** Model deployment typically takes 5-20 minutes to complete."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "346c2fc14a21"
+      },
+      "outputs": [],
+      "source": [
+        "OPERATION_ID = \"[extracted_operation_id]\"  # @param {type: \"number\"}\n",
+        "# Check operation status\n",
+        "print(\"Checking deployment status...\")\n",
+        "!curl -X GET -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/operations/{OPERATION_ID}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5ahw-uFjCAbo"
+      },
+      "source": [
+        "### Prediction"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "7cb14ea3b257"
+      },
+      "source": [
+        "Sends a POST request to the specified API endpoint to get a response from the model for a joke using the provided payload."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "4zFz260B50oi"
+      },
+      "outputs": [],
+      "source": [
+        "PAYLOAD = {\n",
+        "   \"instances\":[\n",
+        "      {\n",
+        "         \"input\":[\n",
+        "            \"def hello_world(): \n",
+        " print('hello_world')\"\n",
+        "         ]\n",
+        "      }\n",
+        "   ]\n",
+        "}\n",
+        "\n",
+        "\n",
+        "request = json.dumps(PAYLOAD)\n",
+        "\n",
+        "!curl -X POST \\\n",
+        "  -H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
+        "  -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}:rawPredict \\\n",
+        "  -d '{request}'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "90b684e180b8"
+      },
+      "source": [
+        "## Using Vertex AI SDK for *Python*"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "141fe83f051b"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "cc8edb9d1d3b"
+      },
+      "outputs": [],
+      "source": [
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6175fddd280b"
+      },
+      "source": [
+        "### Upload Model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "3cf31f867d7f"
+      },
+      "outputs": [],
+      "source": [
+        "model = aiplatform.Model.upload(\n",
+        "    display_name=\"ModelGarden_LaunchPad_Endpoint_\" + time.strftime(\"%Y%m%d-%H%M%S\"),\n",
+        "    model_garden_source_model_name=f\"publishers/{PUBLISHER_NAME}/models/{PUBLISHER_MODEL_NAME}\",\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e61cc3ae9860"
+      },
+      "source": [
+        "### Create Endpoint"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "db0a82502964"
+      },
+      "outputs": [],
+      "source": [
+        "my_endpoint = aiplatform.Endpoint.create(\n",
+        "    display_name=\"ModelGarden_LaunchPad_Endpoint_\" + time.strftime(\"%Y%m%d-%H%M%S\")\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "050f96c8c2b8"
+      },
+      "source": [
+        "### Deploy Model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "16da63ba97a7"
+      },
+      "outputs": [],
+      "source": [
+        "import ipywidgets as widgets\n",
+        "from IPython.display import display\n",
+        "\n",
+        "# Initial data\n",
+        "PUBLISHER_NAME = \"qodo\"  # @param {type:\"string\"}\n",
+        "PUBLISHER_MODEL_NAME = \"qodo-embed-1-7b-v1\"\n",
+        "available_regions = [\"us-central1\", \"europe-west4\"]\n",
+        "compatible_machines = [\n",
+        "    \"a2-highgpu-1g\",\n",
+        "    \"a2-highgpu-4g\",\n",
+        "    \"a2-ultragpu-1g\",\n",
+        "    \"a2-ultragpu-2g\",\n",
+        "    \"a3-highgpu-2g\",\n",
+        "]\n",
+        "\n",
+        "# Machine type to accelerator mapping (you can customize this based on your needs)\n",
+        "machine_config = {\n",
+        "    \"a2-highgpu-1g\": {\"type\": \"NVIDIA_A100\", \"count\": 1},\n",
+        "    \"a2-highgpu-4g\": {\"type\": \"NVIDIA_A100\", \"count\": 4},\n",
+        "    \"a2-ultragpu-1g\": {\"type\": \"NVIDIA_A100_80GB\", \"count\": 1},\n",
+        "    \"a2-ultragpu-2g\": {\"type\": \"NVIDIA_A100_80GB\", \"count\": 2},\n",
+        "    \"a3-highgpu-2g\": {\"type\": \"NVIDIA_H100\", \"count\": 2},\n",
+        "}\n",
+        "\n",
+        "# Create widgets\n",
+        "\n",
+        "\n",
+        "dropdown_machine = widgets.Dropdown(\n",
+        "    options=compatible_machines,\n",
+        "    description=\"Machine type:\",\n",
+        "    font_weight=\"bold\",\n",
+        "    style={\"description_width\": \"initial\"},\n",
+        ")\n",
+        "\n",
+        "label_accelerator_type = widgets.HTML(\n",
+        "    value=f\"<b>Accelerator type:</b> {machine_config[compatible_machines[0]]['type']}\"\n",
+        ")\n",
+        "\n",
+        "label_accelerator_count = widgets.HTML(\n",
+        "    value=f\"<b>Accelerator count:</b> {machine_config[compatible_machines[0]]['count']}\"\n",
+        ")\n",
+        "\n",
+        "# Event handlers\n",
+        "\n",
+        "\n",
+        "def dropdown_machine_eventhandler(change):\n",
+        "    global MACHINE_TYPE, ACCELERATOR_TYPE, ACCELERATOR_COUNT\n",
+        "    if change[\"type\"] == \"change\" and change[\"name\"] == \"value\":\n",
+        "        MACHINE_TYPE = change.new\n",
+        "        machine_info = machine_config.get(change.new, {\"type\": \"Unknown\", \"count\": 0})\n",
+        "        ACCELERATOR_TYPE = machine_info[\"type\"]\n",
+        "        ACCELERATOR_COUNT = machine_info[\"count\"]\n",
+        "\n",
+        "        # Update the displayed information\n",
+        "        label_accelerator_type.value = f\"<b>Accelerator type:</b> {ACCELERATOR_TYPE}\"\n",
+        "        label_accelerator_count.value = f\"<b>Accelerator count:</b> {ACCELERATOR_COUNT}\"\n",
+        "\n",
+        "\n",
+        "# Initialize global variables\n",
+        "MACHINE_TYPE = dropdown_machine.value\n",
+        "ACCELERATOR_TYPE = machine_config[MACHINE_TYPE][\"type\"]\n",
+        "ACCELERATOR_COUNT = machine_config[MACHINE_TYPE][\"count\"]\n",
+        "\n",
+        "# Set up observers\n",
+        "dropdown_machine.observe(dropdown_machine_eventhandler, names=\"value\")\n",
+        "\n",
+        "# Display widgets\n",
+        "display(\n",
+        "    widgets.VBox([dropdown_machine, label_accelerator_type, label_accelerator_count])\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "a4afefe566f6"
+      },
+      "outputs": [],
+      "source": [
+        "model.deploy(\n",
+        "    endpoint=my_endpoint,\n",
+        "    deployed_model_display_name=\"ModelGarden_LaunchPad_DeployedModel_\"\n",
+        "    + time.strftime(\"%Y%m%d-%H%M%S\"),\n",
+        "    traffic_split={\"0\": 100},\n",
+        "    machine_type=MACHINE_TYPE,\n",
+        "    accelerator_type=ACCELERATOR_TYPE,\n",
+        "    accelerator_count=ACCELERATOR_COUNT,\n",
+        "    min_replica_count=1,\n",
+        "    max_replica_count=1,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e793c9b3e13b"
+      },
+      "source": [
+        "### Prediction"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "8cee72080193"
+      },
+      "outputs": [],
+      "source": [
+        "PAYLOAD = {\"instances\": [{\"input\": [\"def hello_world(): \\n  print('hello_world')\"]}]}\n",
+        "\n",
+        "request = json.dumps(PAYLOAD)\n",
+        "\n",
+        "response = my_endpoint.raw_predict(\n",
+        "    body=request, headers={\"Content-Type\": \"application/json\"}\n",
+        ")\n",
+        "data = response.json()\n",
+        "embedding = data[\"predictions\"][0][\"data\"][0][\"embedding\"]\n",
+        "print(embedding)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "d45c572a4b7d"
+      },
+      "source": [
+        "## Cleaning up\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ec834cd282d4"
+      },
+      "outputs": [],
+      "source": [
+        "# Cleaning up\n",
+        "\n",
+        "print(\"Starting cleanup process...\\n\")\n",
+        "\n",
+        "# First undeploy the model from the endpoint\n",
+        "print(\"Undeploying model from endpoint...\")\n",
+        "try:\n",
+        "    my_endpoint.undeploy_all()\n",
+        "    print(\"✓ Model successfully undeployed from endpoint\")\n",
+        "except Exception as e:\n",
+        "    print(f\"Error undeploying model: {e}\")\n",
+        "\n",
+        "# Delete the endpoint\n",
+        "print(\"\\nDeleting endpoint...\")\n",
+        "try:\n",
+        "    my_endpoint.delete()\n",
+        "    print(\"✓ Endpoint successfully deleted\")\n",
+        "except Exception as e:\n",
+        "    print(f\"Error deleting endpoint: {e}\")\n",
+        "\n",
+        "# Delete the model\n",
+        "print(\"\\nDeleting model...\")\n",
+        "try:\n",
+        "    model.delete()\n",
+        "    print(\"✓ Model successfully deleted\")\n",
+        "except Exception as e:\n",
+        "    print(f\"Error deleting model: {e}\")\n",
+        "\n",
+        "print(\"\\nCleanup complete! All resources have been removed.\")\n"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "qodo_intro.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}