fix: official standard

* fix: install db-dtypes

* fix: TF import issue

* fix: TF related import

* debug: import issue

* debug: TF issue

* debug: try again

* fix: install issue

* fix: install within job

* fix: install within job

community-content/CODEOWNERS

@@ -12,4 +12,5 @@
 /prediction_featurestore_integration @googleapis/vertex-prediction-team
 /vertex_vision_model_garden/model_oss/util @weigary
 /vertex_vision_model_garden/model_oss/diffusers @weigary
+/vertex_vision_model_garden/model_oss/keras @dstnluong-google
 /vertex_vision_model_garden/model_oss/transformers @dstnluong-google

community-content/vertex_vision_model_garden/model_oss/diffusers/handler.py

@@ -4,9 +4,10 @@
 # pylint: disable=logging-fstring-interpolation
 
 import base64
+import io
 import logging
 import os
-from typing import Any, List, Tuple
+from typing import Any, List, Sequence, Tuple
 
 from diffusers import ControlNetModel
 from diffusers import DiffusionPipeline
@@ -20,6 +21,7 @@ from diffusers import StableDiffusionPipeline
 from diffusers import StableDiffusionUpscalePipeline
 from diffusers import TextToVideoZeroPipeline
 from diffusers import UniPCMultistepScheduler
+import imageio
 import numpy as np
 from PIL import Image
 import torch
@@ -43,6 +45,13 @@ TEXT_TO_VIDEO_ZERO_SHOT = "text-to-video-zero-shot"
 TEXT_TO_VIDEO = "text-to-video"
 
 
+def frames_to_video_bytes(frames: Sequence[np.ndarray], fps: int) -> bytes:
+  images = [Image.fromarray(array) for array in frames]
+  io_obj = io.BytesIO()
+  imageio.mimsave(io_obj, images, format=".mp4", fps=fps)
+  return io_obj.getvalue()
+
+
 class DiffusersHandler(BaseHandler):
   """Custom handler for TIMM models."""
 
@@ -214,7 +223,7 @@ class DiffusersHandler(BaseHandler):
         numpy_arrays = self.pipeline(prompt=prompt).images
         numpy_arrays = [(i * 255).astype("uint8") for i in numpy_arrays]
         videos.append(
-            video_format_converter.frames_to_video_bytes(numpy_arrays, fps=4)
+            frames_to_video_bytes(numpy_arrays, fps=4)
         )
       return videos
     elif self.task == TEXT_TO_VIDEO:
@@ -224,7 +233,7 @@ class DiffusersHandler(BaseHandler):
       # Therefore we need to split the output into different videos.
       predicted_images = np.array_split(predicted_images, len(prompts), axis=2)
       videos = [
-          video_format_converter.frames_to_video_bytes(images, fps=8)
+          frames_to_video_bytes(images, fps=8)
           for images in predicted_images
       ]
       return videos

community-content/vertex_vision_model_garden/model_oss/keras/dockerfile/serve.Dockerfile

@@ -0,0 +1,118 @@
+# Dockerfile for basic serving dockers with Keras.
+#
+# To build:
+# docker build -f model_oss/keras/dockerfile/serve.Dockerfile . -t ${YOUR_IMAGE_TAG}
+#
+# To push to gcr:
+# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+
+FROM tensorflow/tensorflow:2.12.0-gpu
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+# This is added to fix docker build error related to Nvidia key update.
+RUN rm -f /etc/apt/sources.list.d/cuda.list
+RUN curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add -
+
+# Install basic libs.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        cmake \
+        curl \
+        wget \
+        sudo \
+        gnupg \
+        libsm6 \
+        libxext6 \
+        libxrender-dev \
+        lsb-release \
+        ca-certificates \
+        build-essential \
+        git \
+        vim \
+        screen \
+        libtcmalloc-minimal4
+
+
+# Install google cloud SDK.
+RUN wget -q https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-sdk-359.0.0-linux-x86_64.tar.gz
+RUN tar xzf google-cloud-sdk-359.0.0-linux-x86_64.tar.gz
+RUN ./google-cloud-sdk/install.sh -q
+# Make sure gsutil will use the default service account.
+RUN echo '[GoogleCompute]\nservice_account = default' > /etc/boto.cfg
+
+
+# Install required libs.
+RUN pip install --upgrade pip
+RUN pip install cloud-tpu-client==0.10
+RUN pip install pyyaml==5.4.1
+RUN pip install fsspec==2021.10.1
+RUN pip install gcsfs==2021.10.1
+RUN pip install tensorflow-text==2.11.0
+RUN pip install pyglove==0.1.0
+RUN pip install cloudml-hypertune==0.1.0.dev6
+RUN pip install pylint==2.17.2
+RUN pip install keras-cv==0.4.0
+RUN pip install tensorflow-datasets==4.8.3
+RUN pip install protobuf==3.20.3
+RUN pip install Pillow==9.5.0
+RUN pip install flask==2.3.2
+RUN pip install waitress==2.1.2
+
+# Installs Reduction Server NCCL plugin.
+RUN echo "deb https://packages.cloud.google.com/apt google-fast-socket main" | tee /etc/apt/sources.list.d/google-fast-socket.list \
+&&  curl -s -L https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - \
+&&  apt update && apt install -y google-reduction-server
+
+# Downloading gcloud package
+RUN curl https://dl.google.com/dl/cloudsdk/release/google-cloud-sdk.tar.gz > /tmp/google-cloud-sdk.tar.gz
+
+# Installing the package
+RUN mkdir -p /usr/local/gcloud \
+  && tar -C /usr/local/gcloud -xvf /tmp/google-cloud-sdk.tar.gz \
+  && /usr/local/gcloud/google-cloud-sdk/install.sh
+
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Adding the package path to local
+ENV PATH $PATH:/usr/local/gcloud/google-cloud-sdk/bin
+
+
+ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=cpp
+
+# Lower the memory fragmentation, and speed up the training.
+# https://github.com/tensorflow/tensorflow/issues/44176#issuecomment-783768033
+ENV LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4
+
+# Enable userspace DNS cache
+ENV GCS_RESOLVE_REFRESH_SECS=60
+ENV GCS_REQUEST_CONNECTION_TIMEOUT_SECS=300
+ENV GCS_METADATA_REQUEST_TIMEOUT_SECS=300
+ENV GCS_READ_REQUEST_TIMEOUT_SECS=300
+ENV GCS_WRITE_REQUEST_TIMEOUT_SECS=600
+# Each opened GCS file takes GCS_READ_CACHE_BLOCK_SIZE_MB of RAM, reduce the
+# value from the default 64MB to 8MB to decrease memory footprint.
+ENV GCS_READ_CACHE_BLOCK_SIZE_MB=8
+
+EXPOSE 8501
+
+WORKDIR /usr/local/lib/python3.8/dist-packages/official/vision
+
+COPY model_oss/keras /automl_vision/keras
+COPY model_oss/util /automl_vision/util
+
+WORKDIR /automl_vision
+
+ENV PYTHONPATH "${PYTHONPATH}:/automl_vision/util"
+ENV MODEL_PATH ""
+ENV IMAGE_WIDTH "512"
+ENV IMAGE_HEIGHT "512"
+
+COPY model_oss/keras/serve.py ./app.py
+
+# Run pylint to validate code.
+COPY .pylintrc /automl_vision/.pylintrc
+RUN find . -type f -name "*.py" | xargs pylint --rcfile=./.pylintrc --errors-only
+
+ENTRYPOINT ["flask","run"]
+CMD ["--host=0.0.0.0", "--port=8501"]

community-content/vertex_vision_model_garden/model_oss/keras/dockerfile/train.Dockerfile

@@ -0,0 +1,111 @@
+# Dockerfile for basic training dockers with Keras.
+#
+# To build:
+# docker build -f model_oss/keras/dockerfile/train.Dockerfile . -t ${YOUR_IMAGE_TAG}
+#
+# To push to gcr:
+# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+
+FROM tensorflow/tensorflow:2.12.0-gpu
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+# This is added to fix docker build error related to Nvidia key update.
+RUN rm -f /etc/apt/sources.list.d/cuda.list
+RUN curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add -
+
+# Install basic libs.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        cmake \
+        curl \
+        wget \
+        sudo \
+        gnupg \
+        libsm6 \
+        libxext6 \
+        libxrender-dev \
+        lsb-release \
+        ca-certificates \
+        build-essential \
+        git \
+        vim \
+        screen \
+        libtcmalloc-minimal4
+
+
+# Install google cloud SDK.
+RUN wget -q https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-sdk-359.0.0-linux-x86_64.tar.gz
+RUN tar xzf google-cloud-sdk-359.0.0-linux-x86_64.tar.gz
+RUN ./google-cloud-sdk/install.sh -q
+# Make sure gsutil will use the default service account.
+RUN echo '[GoogleCompute]\nservice_account = default' > /etc/boto.cfg
+
+
+# Install required libs.
+RUN pip install --upgrade pip
+RUN pip install cloud-tpu-client==0.10
+RUN pip install pyyaml==5.4.1
+RUN pip install fsspec==2021.10.1
+RUN pip install gcsfs==2021.10.1
+RUN pip install tensorflow-text==2.11.0
+RUN pip install pyglove==0.1.0
+RUN pip install cloudml-hypertune==0.1.0.dev6
+RUN pip install pylint==2.17.2
+RUN pip install keras-cv==0.4.0
+RUN pip install tensorflow-datasets==4.8.3
+RUN pip install tensorflow-estimator==2.12.0
+RUN pip install tensorflow-gcs-config==2.12.0
+RUN pip install tensorflow-hub==0.13.0
+RUN pip install tensorflow-io-gcs-filesystem==0.32.0
+RUN pip install tensorflow-metadata==1.13.1
+RUN pip install tensorflow-probability==0.19.0
+RUN pip install tensorboard==2.12.2
+RUN pip install tensorboard-data-server==0.7.0
+RUN pip install tensorboard-plugin-wit==1.8.1
+RUN pip install protobuf==3.20.3
+RUN pip install pandas==1.5.3
+RUN pip install pandas-datareader==0.10.0
+RUN pip install pandas-gbq==0.17.9
+RUN pip install pycocotools==2.0.6
+
+# Installs Reduction Server NCCL plugin.
+RUN echo "deb https://packages.cloud.google.com/apt google-fast-socket main" | tee /etc/apt/sources.list.d/google-fast-socket.list \
+&&  curl -s -L https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - \
+&&  apt update && apt install -y google-reduction-server
+
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=cpp
+
+# Lower the memory fragmentation, and speed up the training.
+# https://github.com/tensorflow/tensorflow/issues/44176#issuecomment-783768033
+ENV LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4
+
+# Enable userspace DNS cache
+ENV GCS_RESOLVE_REFRESH_SECS=60
+ENV GCS_REQUEST_CONNECTION_TIMEOUT_SECS=300
+ENV GCS_METADATA_REQUEST_TIMEOUT_SECS=300
+ENV GCS_READ_REQUEST_TIMEOUT_SECS=300
+ENV GCS_WRITE_REQUEST_TIMEOUT_SECS=600
+# Each opened GCS file takes GCS_READ_CACHE_BLOCK_SIZE_MB of RAM, reduce the
+# value from the default 64MB to 8MB to decrease memory footprint.
+ENV GCS_READ_CACHE_BLOCK_SIZE_MB=8
+
+WORKDIR /usr/local/lib/python3.8/dist-packages/official/vision
+
+COPY model_oss/keras /automl_vision/keras
+COPY model_oss/util /automl_vision/util
+
+WORKDIR /automl_vision
+
+# Keras stable diffusion training codes set width and height as RESOLUTION.
+ENV RESOLUTION "512"
+
+ENV PYTHONPATH "${PYTHONPATH}:/automl_vision/util"
+
+# Run pylint to validate code.
+COPY .pylintrc /automl_vision/.pylintrc
+RUN find . -type f -name "*.py" | xargs pylint --rcfile=./.pylintrc --errors-only
+
+ENTRYPOINT ["python3","keras/train.py"]

community-content/vertex_vision_model_garden/model_oss/keras/serve.py

@@ -0,0 +1,184 @@
+r"""Servers Keras Stable Diffusion models.
+
+python serve.py --model_path=<model path in gcs>
+
+curl -d \
+'{"prompt":"Hello Kitty"}' \
+-H "Content-Type: application/json" \
+-X POST http://localhost:8501/predict
+"""
+
+import base64
+import io
+import json
+import os
+from typing import List, Tuple
+
+from absl import app
+# The docker builds could not find flask and waitress.
+# pylint: disable=import-error
+from flask import Flask
+from flask import request
+from flask import Response
+import keras_cv
+from PIL import Image
+from waitress import serve
+
+from util import constants
+from util import fileutils
+
+
+flask_app = Flask(__name__)
+
+stable_diffusion_model = None
+
+
+model_path = os.environ.get('MODEL_PATH', '')
+if model_path.startswith(constants.GCS_URI_PREFIX):
+  print('Downloading models from gcs to local.')
+  os.makedirs(constants.LOCAL_MODEL_DIR, exist_ok=True)
+  fileutils.download_gcs_dir_to_local(
+      os.path.dirname(model_path), constants.LOCAL_MODEL_DIR
+  )
+  model_path = os.path.join(
+      constants.LOCAL_MODEL_DIR, os.path.basename(model_path)
+  )
+
+image_width = int(os.environ.get('IMAGE_WIDTH', 512))
+image_height = int(os.environ.get('IMAGE_HEIGHT', 512))
+
+print('image_width=', image_width, 'image_height=', image_height)
+print('Create Keras stable diffusion models.')
+stable_diffusion_model = keras_cv.models.StableDiffusion(
+    img_width=image_width,
+    img_height=image_height,
+    jit_compile=True,
+)
+
+if model_path:
+  # We just reload the weights of the fine-tuned diffusion model.
+  print('Initialize finetuned models from: ', model_path)
+  stable_diffusion_model.diffusion_model.load_weights(model_path)
+
+
+def error(message: str) -> str:
+  """Returns a JSON representing an error response."""
+  return json.dumps({
+      'success': False,
+      'error': message,
+  })
+
+
+def check_key_in_json(content: str, keys: List[str]) -> str:
+  for key in keys:
+    if key not in content:
+      return error('No {} in request {}.'.format(key, content))
+  return None
+
+
+def validate_json_key(json_key_string: str) -> Tuple[str, bool]:
+  try:
+    json_key = json.loads(json_key_string)
+  except (ValueError, TypeError):
+    return (error('Invalid key found in request'), False)
+  return (json_key, True)
+
+
+# The health check route is required for docker deployment in google cloud.
+@flask_app.route('/ping')
+def ping() -> Response:
+  """Health checks."""
+  return Response(status=200)
+
+
+# The return should be `Response` for docker deployment in google cloud.
+@flask_app.route('/predict', methods=['GET', 'POST'])
+def predict_model() -> Response:
+  """Predictions."""
+  if request.method == 'POST':
+    contents = request.get_json(force=True)
+
+    print('The input contents are:', contents)
+    batch_size = 1
+    num_steps = 25
+    seed = 1234
+    if 'parameters' in contents:
+      parameters = contents['parameters']
+      if 'batch_size' in parameters:
+        batch_size = int(parameters['batch_size'])
+      if 'num_steps' in parameters:
+        num_steps = int(parameters['num_steps'])
+      if 'seed' in parameters:
+        seed = int(parameters['seed'])
+    print('batch_size=', batch_size, 'num_steps=', num_steps, 'seed=', seed)
+    if batch_size < 1:
+      return Response(
+          response=error('The batch size must be a positive integar.'),
+          status=200,
+          mimetype='text/plain',
+      )
+    if num_steps < 1:
+      return Response(
+          response=error('The num steps must be a positive integar.'),
+          status=200,
+          mimetype='text/plain',
+      )
+    predictions = []
+    for content in contents['instances']:
+      print('Processing:', content)
+      prompt = content['prompt']
+      generated_image_array = stable_diffusion_model.text_to_image(
+          prompt=prompt,
+          batch_size=batch_size,
+          num_steps=num_steps,
+          seed=seed,
+      )
+
+      generated_image_bytes_array = []
+      for i in range(batch_size):
+        generated_image = Image.fromarray(generated_image_array[i])
+        # Converts the image to a base64-encoded string.
+        buffered_image = io.BytesIO()
+        generated_image.save(buffered_image, format='JPEG')
+        generated_image_bytes = base64.b64encode(
+            buffered_image.getvalue()
+        ).decode('utf-8')
+        generated_image_bytes_array.append(generated_image_bytes)
+      prediction = {
+          'prompt': prompt,
+          'predicted_image': generated_image_bytes_array,
+      }
+      predictions.append(prediction)
+
+    return Response(
+        response=json.dumps({
+            'success': True,
+            'predictions': predictions,
+        }),
+        status=200,
+        mimetype='text/plain',
+    )
+  else:
+    return Response(
+        response=json.dumps({
+            'success': True,
+            'isalive': stable_diffusion_model is not None,
+        }),
+        status=200,
+        mimetype='text/plain',
+    )
+
+
+def serve_main(unused_argv):
+  """The main function to serve Keras models."""
+  del unused_argv
+  # This is used when running locally only. When deploying to Google App
+  # Engine, a webserver process such as Gunicorn will serve the app.
+  # # Debug deployment.
+  # flask_app.run(host='0.0.0.0', port=8501, debug=True)
+  # Prod deployment.
+  serve(flask_app, host='0.0.0.0', port=8501)
+
+
+if __name__ == '__main__':
+  app.run(serve_main)

community-content/vertex_vision_model_garden/model_oss/keras/train.py

@@ -0,0 +1,363 @@
+"""Train Keras Stable Diffusion.
+
+Most the codes below are from
+https://keras.io/examples/generative/finetune_stable_diffusion/.
+"""
+import os
+
+from absl import app
+from absl import flags
+from absl import logging
+import keras_cv
+# pylint: disable=g-importing-member
+from keras_cv.models.stable_diffusion.clip_tokenizer import SimpleTokenizer
+from keras_cv.models.stable_diffusion.diffusion_model import DiffusionModel
+from keras_cv.models.stable_diffusion.image_encoder import ImageEncoder
+from keras_cv.models.stable_diffusion.noise_scheduler import NoiseScheduler
+from keras_cv.models.stable_diffusion.text_encoder import TextEncoder
+import numpy as np
+# The docker builds could not find pandas.
+# pylint: disable=import-error
+import pandas as pd
+import tensorflow as tf
+from tensorflow import keras
+import tensorflow.experimental.numpy as tnp
+
+from util import constants
+from util import fileutils
+
+_INPUT_CSV_PATH = flags.DEFINE_string(
+    'input_csv_path',
+    None,
+    'The input csv path.',
+    required=True,
+)
+
+_USE_MP = flags.DEFINE_bool(
+    'use_mp',
+    True,
+    'Enable mixed-precision training if the underlying GPU has tensor cores.',
+)
+
+_EPOCHS = flags.DEFINE_integer('epochs', 1, 'The number of epochs.')
+
+_OUTPUT_MODEL_DIR = flags.DEFINE_string(
+    'output_model_dir',
+    None,
+    'The output model dir.',
+    required=True,
+)
+
+# These hyperparameters defaults come from this tutorial by Hugging Face:
+# https://huggingface.co/docs/diffusers/training/text2image
+_LEARNING_RATE = flags.DEFINE_float(
+    'learning_rate', 1e-5, 'The learning rate parameter for AdamW optimizer.'
+)
+
+_BETA_1 = flags.DEFINE_float(
+    'beta_1', 0.9, 'The beta_1 parameter for AdamW optimizer.'
+)
+
+_BETA_2 = flags.DEFINE_float(
+    'beta_2', 0.999, 'The beta_2 parameter for AdamW optimizer.'
+)
+
+_WEIGHT_DECAY = flags.DEFINE_float(
+    'weight_decay', 1e-2, 'The weight decay parameter for AdamW optimizer.'
+)
+
+_EPSILON = flags.DEFINE_float(
+    'epsilon', 1e-08, 'The epsilon parameter for AdamW optimizer.'
+)
+
+RESOLUTION = int(os.environ.get('RESOLUTION', 512))
+
+# The padding token and maximum prompt length are specific to the text encoder.
+# If you're using a different text encoder be sure to change them accordingly.
+PADDING_TOKEN = 49407
+MAX_PROMPT_LENGTH = 77
+
+AUTO = tf.data.AUTOTUNE
+POS_IDS = tf.convert_to_tensor([list(range(MAX_PROMPT_LENGTH))], dtype=tf.int32)
+
+
+augmenter = keras.Sequential(
+    layers=[
+        keras_cv.layers.CenterCrop(RESOLUTION, RESOLUTION),
+        keras_cv.layers.RandomFlip(),
+        tf.keras.layers.Rescaling(scale=1.0 / 127.5, offset=-1),
+    ]
+)
+text_encoder = TextEncoder(MAX_PROMPT_LENGTH)
+
+
+def process_image(image_path, tokenized_text):
+  image = tf.io.read_file(image_path)
+  image = tf.io.decode_png(image, 3)
+  image = tf.image.resize(image, (RESOLUTION, RESOLUTION))
+  return image, tokenized_text
+
+
+def apply_augmentation(image_batch, token_batch):
+  return augmenter(image_batch), token_batch
+
+
+def run_text_encoder(image_batch, token_batch):
+  return (
+      image_batch,
+      token_batch,
+      text_encoder([token_batch, POS_IDS], training=False),
+  )
+
+
+def prepare_dict(image_batch, token_batch, encoded_text_batch):
+  return {
+      'images': image_batch,
+      'tokens': token_batch,
+      'encoded_text': encoded_text_batch,
+  }
+
+
+def prepare_dataset(image_paths, tokenized_texts, batch_size=1):
+  dataset = tf.data.Dataset.from_tensor_slices((image_paths, tokenized_texts))
+  dataset = dataset.shuffle(batch_size * 10)
+  dataset = dataset.map(process_image, num_parallel_calls=AUTO).batch(
+      batch_size
+  )
+  dataset = dataset.map(apply_augmentation, num_parallel_calls=AUTO)
+  dataset = dataset.map(run_text_encoder, num_parallel_calls=AUTO)
+  dataset = dataset.map(prepare_dict, num_parallel_calls=AUTO)
+  return dataset.prefetch(AUTO)
+
+
+def prepare_training_dataset(dataset_csv):
+  """Prepares training datasets."""
+  if dataset_csv.startswith(constants.GCS_URI_PREFIX):
+    if not os.path.exists(constants.LOCAL_DATA_DIR):
+      os.makedirs(constants.LOCAL_DATA_DIR)
+    logging.info(
+        'Start to download data from %s to %s.',
+        os.path.dirname(dataset_csv),
+        constants.LOCAL_DATA_DIR,
+    )
+    fileutils.download_gcs_dir_to_local(
+        os.path.dirname(dataset_csv), constants.LOCAL_DATA_DIR
+    )
+    data_frame = pd.read_csv(
+        os.path.join(constants.LOCAL_DATA_DIR, os.path.basename(dataset_csv))
+    )
+    data_frame['image_path'] = data_frame['image_path'].apply(
+        lambda x: os.path.join(constants.LOCAL_DATA_DIR, x)
+    )
+  else:
+    # Keeps the following codes for experiments with
+    # https://keras.io/examples/generative/finetune_stable_diffusion/.
+    data_path = tf.keras.utils.get_file(origin=dataset_csv, untar=True)
+    data_frame = pd.read_csv(os.path.join(data_path, 'data.csv'))
+    data_frame['image_path'] = data_frame['image_path'].apply(
+        lambda x: os.path.join(data_path, x)
+    )
+  data_frame.head()
+
+  # Load the tokenizer.
+  tokenizer = SimpleTokenizer()
+
+  #  Method to tokenize and pad the tokens.
+  def process_text(caption):
+    tokens = tokenizer.encode(caption)
+    tokens = tokens + [PADDING_TOKEN] * (MAX_PROMPT_LENGTH - len(tokens))
+    return np.array(tokens)
+
+  # Collate the tokenized captions into an array.
+  tokenized_texts = np.empty((len(data_frame), MAX_PROMPT_LENGTH))
+
+  all_captions = list(data_frame['caption'].values)
+  for i, caption in enumerate(all_captions):
+    tokenized_texts[i] = process_text(caption)
+
+  # Prepare the dataset.
+  training_dataset = prepare_dataset(
+      np.array(data_frame['image_path']), tokenized_texts, batch_size=4
+  )
+
+  return training_dataset
+
+
+class Trainer(tf.keras.Model):
+  """The trainer for Keras Stable Diffusion."""
+
+  # Reference:
+  # https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py
+
+  def __init__(
+      self,
+      diffusion_model,
+      vae,
+      noise_scheduler,
+      use_mixed_precision=False,
+      max_grad_norm=1.0,
+      **kwargs,
+  ):
+    super().__init__(**kwargs)
+
+    self.diffusion_model = diffusion_model
+    self.vae = vae
+    self.noise_scheduler = noise_scheduler
+    self.max_grad_norm = max_grad_norm
+
+    self.use_mixed_precision = use_mixed_precision
+    self.vae.trainable = False
+
+  def train_step(self, inputs):
+    images = inputs['images']
+    encoded_text = inputs['encoded_text']
+    batch_size = tf.shape(images)[0]
+
+    with tf.GradientTape() as tape:
+      # Project image into the latent space and sample from it.
+      latents = self.sample_from_encoder_outputs(
+          self.vae(images, training=False)
+      )
+      # Know more about the magic number here:
+      # https://keras.io/examples/generative/fine_tune_via_textual_inversion/
+      latents = latents * 0.18215
+
+      # Sample noise that we'll add to the latents.
+      noise = tf.random.normal(tf.shape(latents))
+
+      # Sample a random timestep for each image.
+      timesteps = tnp.random.randint(
+          0, self.noise_scheduler.train_timesteps, (batch_size,)
+      )
+
+      # Add noise to the latents according to the noise magnitude at each
+      # timestep (this is the forward diffusion process).
+      noisy_latents = self.noise_scheduler.add_noise(
+          tf.cast(latents, noise.dtype), noise, timesteps
+      )
+
+      # Get the target for loss depending on the prediction type
+      # just the sampled noise for now.
+      target = noise  # noise_schedule.predict_epsilon == True
+
+      # Predict the noise residual and compute loss.
+      # pylint: disable=unnecessary-lambda
+      timestep_embedding = tf.map_fn(
+          lambda t: self.get_timestep_embedding(t), timesteps, dtype=tf.float32
+      )
+      timestep_embedding = tf.squeeze(timestep_embedding, 1)
+      model_pred = self.diffusion_model(
+          [noisy_latents, timestep_embedding, encoded_text], training=True
+      )
+      loss = self.compiled_loss(target, model_pred)
+      if self.use_mixed_precision:
+        loss = self.optimizer.get_scaled_loss(loss)
+
+    # Update parameters of the diffusion model.
+    trainable_vars = self.diffusion_model.trainable_variables
+    gradients = tape.gradient(loss, trainable_vars)
+    if self.use_mixed_precision:
+      gradients = self.optimizer.get_unscaled_gradients(gradients)
+    gradients = [tf.clip_by_norm(g, self.max_grad_norm) for g in gradients]
+    self.optimizer.apply_gradients(zip(gradients, trainable_vars))
+
+    return {m.name: m.result() for m in self.metrics}
+
+  def get_timestep_embedding(self, timestep, dim=320, max_period=10000):
+    half = dim // 2
+    log_max_preiod = tf.math.log(tf.cast(max_period, tf.float32))
+    # The docker builds could not support unary `-`.
+    # pylint: disable=invalid-unary-operand-type
+    freqs = tf.math.exp(
+        -log_max_preiod * tf.range(0, half, dtype=tf.float32) / half
+    )
+    args = tf.convert_to_tensor([timestep], dtype=tf.float32) * freqs
+    embedding = tf.concat([tf.math.cos(args), tf.math.sin(args)], 0)
+    embedding = tf.reshape(embedding, [1, -1])
+    return embedding
+
+  def sample_from_encoder_outputs(self, outputs):
+    mean, logvar = tf.split(outputs, 2, axis=-1)
+    logvar = tf.clip_by_value(logvar, -30.0, 20.0)
+    std = tf.exp(0.5 * logvar)
+    sample = tf.random.normal(tf.shape(mean), dtype=mean.dtype)
+    return mean + std * sample
+
+  def save_weights(
+      self, filepath, overwrite=True, save_format=None, options=None
+  ):
+    # Overriding this method will allow us to use the `ModelCheckpoint`
+    # callback directly with this trainer class. In this case, it will
+    # only checkpoint the `diffusion_model` since that's what we're training
+    # during fine-tuning.
+    self.diffusion_model.save_weights(
+        filepath=filepath,
+        overwrite=overwrite,
+        save_format=save_format,
+        options=options,
+    )
+
+
+def main(_) -> None:
+  # _INPUT_CSV_PATH and _OUTPUT_MODEL_DIR should have the format as
+  # gs://<bucket_name>/<object_name>.
+  if _INPUT_CSV_PATH.value:
+    if not _INPUT_CSV_PATH.value.startswith(constants.GCS_URI_PREFIX):
+      raise ValueError('The input csv path should be a gcs path like gs://<>')
+  if _OUTPUT_MODEL_DIR.value:
+    if not _OUTPUT_MODEL_DIR.value.startswith(constants.GCS_URI_PREFIX):
+      raise ValueError('The output model dir should be a gcs path like gs://<>')
+
+  if _USE_MP.value:
+    keras.mixed_precision.set_global_policy('mixed_float16')
+
+  image_encoder = ImageEncoder(RESOLUTION, RESOLUTION)
+  diffusion_ft_trainer = Trainer(
+      diffusion_model=DiffusionModel(RESOLUTION, RESOLUTION, MAX_PROMPT_LENGTH),
+      # Remove the top layer from the encoder, which cuts off the variance and
+      # only returns the mean.
+      vae=tf.keras.Model(
+          image_encoder.input,
+          image_encoder.layers[-2].output,
+      ),
+      noise_scheduler=NoiseScheduler(),
+      use_mixed_precision=_USE_MP.value,
+  )
+
+  optimizer = tf.keras.optimizers.experimental.AdamW(
+      learning_rate=_LEARNING_RATE.value,
+      weight_decay=_WEIGHT_DECAY.value,
+      beta_1=_BETA_1.value,
+      beta_2=_BETA_2.value,
+      epsilon=_EPSILON.value,
+  )
+  diffusion_ft_trainer.compile(optimizer=optimizer, loss='mse')
+
+  training_dataset = prepare_training_dataset(_INPUT_CSV_PATH.value)
+
+  # Note: gcsfuse does not work for Keras. We saves the trained models locally
+  # first, and then copy to gcs storages.
+  if not os.path.exists(constants.LOCAL_MODEL_DIR):
+    os.makedirs(constants.LOCAL_MODEL_DIR)
+  # The default saved model is in HDF5.
+  ckpt_path = os.path.join(constants.LOCAL_MODEL_DIR, 'saved_model.h5')
+  ckpt_callback = tf.keras.callbacks.ModelCheckpoint(
+      ckpt_path,
+      save_weights_only=True,
+      monitor='loss',
+      mode='min',
+  )
+  diffusion_ft_trainer.fit(
+      training_dataset, epochs=_EPOCHS.value, callbacks=[ckpt_callback]
+  )
+
+  # Copies the files in constants.LOCAL_MODEL_DIR to output_model_dir.
+  fileutils.upload_local_dir_to_gcs(
+      constants.LOCAL_MODEL_DIR, _OUTPUT_MODEL_DIR.value
+  )
+
+  return
+
+
+if __name__ == '__main__':
+  app.run(main)

community-content/vertex_vision_model_garden/model_oss/video_util/video_format_converter.py

@@ -1,14 +0,0 @@
-"""Video format converter util lib."""
-
-import io
-from typing import Sequence
-import imageio
-import numpy as np
-from PIL import Image
-
-
-def frames_to_video_bytes(frames: Sequence[np.ndarray], fps: int) -> bytes:
-  images = [Image.fromarray(array) for array in frames]
-  io_obj = io.BytesIO()
-  imageio.mimsave(io_obj, images, format=".mp4", fps=fps)
-  return io_obj.getvalue()

notebooks/community/CODEOWNERS

@@ -44,7 +44,9 @@
 /notebooks/community/feature_store/get_started_vertex_feature_store.ipynb @junkourata
 /notebooks/community/model_garden/model_garden_huggingface_local_inference.ipynb @dstnluong-google
 /notebooks/community/model_garden/model_garden_mediapipe_image_classification.ipynb @schmidt-sebastian
+/notebooks/community/model_garden/model_garden_mediapipe_gesture_recognition.ipynb @schmidt-sebastian
 /notebooks/community/model_garden/model_garden_mediapipe_object_detection.ipynb @schmidt-sebastian
+/notebooks/community/model_garden/model_garden_mediapipe_text_classification.ipynb @schmidt-sebastian
 /notebooks/community/model_garden/model_garden_proprietary_image_classification.ipynb @weigary
 /notebooks/community/model_garden/model_garden_proprietary_image_object_detection.ipynb @weigary
 /notebooks/community/model_garden/model_garden_tfvision_image_classification.ipynb @genquan9
@@ -73,4 +75,9 @@
 /notebooks/community/bigquery_ml_inference/bq_ml_with_vision_translation_nlp.ipynb @deaconsmith
 /notebooks/community/model_garden/model_garden_keras_stable_diffusion.ipynb @genquan9
 /notebooks/community/model_garden/model_garden_pytorch_sam.ipynb @huguensjean
+/notebooks/community/model_garden/model_garden_pytorch_pic2word.ipynb @jismailyan
 /notebooks/community/model_garden/model_garden_pytorch_peft.ipynb @genquan9
+/notebooks/community/model_garden/model_garden_pytorch_openllama_peft.ipynb @genquan9
+/notebooks/community/model_garden/model_garden_pytorch_falcon_instruct_peft.ipynb @genquan9
+/notebooks/community/model_garden/model_garden_movinet_clip_classification.ipynb @KCFindstr
+

notebooks/community/model_garden/model_garden_jax_vision_transformer.ipynb

@@ -44,7 +44,7 @@
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
-        "  <td>                                                                                               <td>\n",
+        "  <td>\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_jax_vision_transformer.ipynb\">\n",
         "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
         "Open in Vertex AI Workbench\n",

notebooks/community/model_garden/model_garden_keras_stable_diffusion.ipynb

@@ -44,7 +44,7 @@
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
-        "  <td>                                                                                               <td>\n",
+        "  <td>\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_keras_stable_diffusion.ipynb\">\n",
         "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
         "Open in Vertex AI Workbench\n",
@@ -282,18 +282,21 @@
         "from io import BytesIO\n",
         "\n",
         "import matplotlib.pyplot as plt\n",
+        "from google.cloud import storage\n",
         "from PIL import Image\n",
         "\n",
+        "GCS_URI_PREFIX = \"gs://\"\n",
+        "\n",
         "# Training constants.\n",
         "TRAINING_JOB_PREFIX = \"train\"\n",
-        "TRAIN_CONTAINER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/keras-train:latest\"\n",
+        "TRAIN_CONTAINER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/keras-train:latest\"\n",
         "TRAIN_MACHINE_TYPE = \"a2-highgpu-1g\"\n",
         "TRAIN_ACCELERATOR_TYPE = \"NVIDIA_TESLA_A100\"\n",
         "TRAIN_NUM_GPU = 1\n",
         "RESOLUTION = 512\n",
         "\n",
         "# Prediction constants.\n",
-        "PREDICTION_CONTAINER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/keras-serve:latest\"\n",
+        "PREDICTION_CONTAINER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/keras-serve:latest\"\n",
         "PREDICTION_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
         "PREDICTION_MACHINE_TYPE = \"n1-standard-8\"\n",
         "DEPLOY_JOB_PREFIX = \"deploy\"\n",
@@ -317,6 +320,21 @@
         "    return gcs_path\n",
         "\n",
         "\n",
+        "def download_gcs_file_to_local(gcs_uri: str, local_path: str):\n",
+        "    \"\"\"Download a gcs file to a local path.\n",
+        "\n",
+        "    Args:\n",
+        "      gcs_uri: A string of file path on GCS.\n",
+        "      local_path: A string of local file path.\n",
+        "    \"\"\"\n",
+        "    if not gcs_uri.startswith(GCS_URI_PREFIX):\n",
+        "        raise ValueError(f\"{gcs_uri} is not a GCS path starting with {GCS_URI_PREFIX}.\")\n",
+        "    client = storage.Client()\n",
+        "    os.makedirs(os.path.dirname(local_path), exist_ok=True)\n",
+        "    with open(local_path, \"wb\") as f:\n",
+        "        client.download_blob_to_file(gcs_uri, f)\n",
+        "\n",
+        "\n",
         "def deploy_model(model_path, service_account):\n",
         "\n",
         "    deploy_model_name = get_job_name_with_datetime(DEPLOY_JOB_PREFIX)\n",
@@ -420,7 +438,11 @@
         "from keras_cv.models import StableDiffusion\n",
         "\n",
         "model = StableDiffusion(img_height=RESOLUTION, img_width=RESOLUTION, jit_compile=True)\n",
-        "if model_path:\n",
+        "if model_path.startswith(GCS_URI_PREFIX):\n",
+        "    local_model_path = \"/tmp/saved_model.h5\"\n",
+        "    download_gcs_file_to_local(model_path, local_model_path)\n",
+        "    model.diffusion_model.load_weights(local_model_path)\n",
+        "elif model_path:\n",
         "    model.diffusion_model.load_weights(model_path)"
       ]
     },
@@ -568,7 +590,7 @@
       },
       "source": [
         "## Finetune models\n",
-        "This section shows how to finetune Keras Stable diffusion models with trainig dockers.\n",
+        "This section shows how to finetune Keras Stable diffusion models with training dockers.\n",
         "\n",
         "If you would like to use finetuned models, please go to the section `Run inferences`."
       ]

notebooks/community/model_garden/model_garden_mediapipe_gesture_recognition.ipynb

@@ -0,0 +1,641 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ur8xi4C7S06n"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2023 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": "TirJ-SGQseby"
+      },
+      "source": [
+        "# Vertex AI Model Garden MediaPipe with gesture recognition\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_mediapipe_gesture_recognition.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_mediapipe_gesture_recognition.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>                                                                                               <td>\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_mediapipe_gesture_recognition.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "Open in Vertex AI Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dwGLvtIeECLK"
+      },
+      "source": [
+        "**_NOTE_**: This notebook has been tested in the following environment:\n",
+        "\n",
+        "* Python version = 3.9"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tvgnzT1CKxrO"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates how to use [MediaPipe Model Maker](https://developers.google.com/mediapipe/solutions/model_maker) to train an on-device gesture recognition model in Vertex AI Model Garden.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "* Train new models\n",
+        "  * Convert input data to training formats\n",
+        "  * Create [custom jobs](https://cloud.google.com/vertex-ai/docs/training/create-custom-job) to train new models\n",
+        "  * Export models\n",
+        "\n",
+        "* Cleanup resources\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI\n",
+        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
+        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
+        "Calculator](https://cloud.google.com/products/calculator/)\n",
+        "to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "KEukV6uRk_S3"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "z__i0w0lCAsW"
+      },
+      "source": [
+        "### Colab only\n",
+        "Run the following commands to install dependencies and to authenticate with Google Cloud if running on Colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Jvqs-ehKlaYh"
+      },
+      "outputs": [],
+      "source": [
+        "! pip3 install --upgrade pip\n",
+        "\n",
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
+        "\n",
+        "    # Automatically restart kernel after installs\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)\n",
+        "\n",
+        "    from google.colab import auth as google_auth\n",
+        "\n",
+        "    google_auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "WReHDGG5g0XY"
+      },
+      "source": [
+        "#### Set your project ID\n",
+        "\n",
+        "**If you don't know your project ID**, see the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "oM1iC_MfAts1"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Set the project id\n",
+        "! gcloud config set project {PROJECT_ID}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "region"
+      },
+      "source": [
+        "#### Region\n",
+        "\n",
+        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tTy1gX11kCJY"
+      },
+      "outputs": [],
+      "source": [
+        "REGION = \"us-central1\"  # @param {type: \"string\"}\n",
+        "REGION_PREFIX = REGION.split(\"-\")[0]\n",
+        "assert REGION_PREFIX in (\n",
+        "    \"us\",\n",
+        "    \"europe\",\n",
+        "    \"asia\",\n",
+        "), f'{REGION} is not supported. It must be prefixed by \"us\", \"asia\", or \"europe\".'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zgPO1eR3CYjk"
+      },
+      "source": [
+        "### Create a Cloud Storage bucket\n",
+        "\n",
+        "Create a storage bucket to store intermediate artifacts such as datasets."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "MzGDU7TWdts_"
+      },
+      "outputs": [],
+      "source": [
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-EcIXiGsCePi"
+      },
+      "source": [
+        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NIq7R4HZCfIc"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "960505627ddf"
+      },
+      "source": [
+        "### Import libraries"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "PyQmSRbKA8r-"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "import os\n",
+        "from datetime import datetime\n",
+        "\n",
+        "import tensorflow\n",
+        "from google.cloud import aiplatform"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "init_aip:mbsdk,all"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
+        "\n",
+        "Initialize the Vertex AI SDK for Python for your project."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9wExiMUxFk91"
+      },
+      "outputs": [],
+      "source": [
+        "now = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temp/%s\" % now)\n",
+        "\n",
+        "EVALUATION_RESULT_OUTPUT_DIRECTORY = os.path.join(STAGING_BUCKET, \"evaluation\")\n",
+        "EVALUATION_RESULT_OUTPUT_FILE = os.path.join(\n",
+        "    EVALUATION_RESULT_OUTPUT_DIRECTORY, \"evaluation.json\"\n",
+        ")\n",
+        "\n",
+        "EXPORTED_MODEL_OUTPUT_DIRECTORY = os.path.join(STAGING_BUCKET, \"model\")\n",
+        "EXPORTED_MODEL_OUTPUT_FILE = os.path.join(\n",
+        "    EXPORTED_MODEL_OUTPUT_DIRECTORY, \"gesture_recognizer.task\"\n",
+        ")\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "n6IFz75WGCam"
+      },
+      "source": [
+        "### Define training machine specs"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "riG_qUokg0XZ"
+      },
+      "outputs": [],
+      "source": [
+        "TRAINING_JOB_DISPLAY_NAME = \"mediapipe_gesture_recognizer_%s\" % now\n",
+        "TRAINING_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/mediapipe-train\"\n",
+        "TRAINING_MACHINE_TYPE = \"n1-highmem-16\"\n",
+        "TRAINING_ACCELARATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
+        "TRAINING_ACCELERATOR_COUNT = 2"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-rsdAcBV-vlf"
+      },
+      "source": [
+        "## Train your customized models"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zgPO1eR3CYjk"
+      },
+      "source": [
+        "### Prepare input data for training\n",
+        "\n",
+        "Finetuning a model for gesture recognition requires a dataset with a directory structure following the pattern `<dataset_path>/<label_name>/<img_name>.*` (e.g. `my_custom_dataset/thumbs_up/img12.jpg`). In addition, one of the label names must be none. The none label represents any gesture that isn't classified as one of the other gestures.\n",
+        "\n",
+        "This example uses a rock paper scissors dataset sample which is available on Cloud Storage.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "IndQ_m6ddUEM"
+      },
+      "outputs": [],
+      "source": [
+        "training_data_path = (\n",
+        "    \"gs://mediapipe-tasks/gesture_recognizer/rps_data_sample\"  # @param {type:\"string\"}\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ci4IV6vdXRMD"
+      },
+      "source": [
+        "When Model Maker loads the dataset, it runs the pre-packaged hand detection model from MediaPipe Hands to detect the hand landmarks from the images. Any images without detected hands are ommitted from the dataset. The resulting dataset will contain the extracted hand landmark positions from each image, rather than images themselves.\n",
+        "\n",
+        "You can configure a few options that determine how the dataset is loaded:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "aNHLSyFtXP7I"
+      },
+      "outputs": [],
+      "source": [
+        "# A boolean controlling whether to shuffle the dataset. Defaults to true.\n",
+        "shuffle = True  # @param {type:\"boolean\"}\n",
+        "# A float between 0 and 1 controlling the confidence threshold for hand detection\n",
+        "min_detection_confidence = 0.6  # @param {type:\"number\"}\n",
+        "# Configures how to split the dataset between training, validation and test data. Must sum to up 1.\n",
+        "split_ratio = \"0.8,0.1,0.1\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aaff6f5be7f6"
+      },
+      "source": [
+        "### Set fine-tuning options\n",
+        "\n",
+        "You can customize the model using the by specifying ModelOptions and HParams. The ModelOptions contain parameters related to the model itself, while the HParams  contains parameters related to training and saving the model.\n",
+        "\n",
+        "The ModelOptions contain these customizable parameter that affects accuracy:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bDxsEaoGcibW"
+      },
+      "outputs": [],
+      "source": [
+        "# The fraction of the input units to drop. Used in dropout layer.\n",
+        "dropout_rate: float = 0.05  # @param {type:\"number\"}\n",
+        "# A list of hidden layer widths for the gesture model. Each element\n",
+        "# in the list will create a new hidden layer with the specified width.\n",
+        "# The hidden layers are separated with BatchNorm, Dropout, and ReLU.\n",
+        "layer_widths: str = \"\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "fk0TTZbDdJPX"
+      },
+      "source": [
+        "HParams has the following list of customizable parameters which affect model accuracy:\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "um_XKbmpTaHx"
+      },
+      "outputs": [],
+      "source": [
+        "# The learning rate to use for gradient descent training.\n",
+        "learning_rate: float = 0.001  # @param {type:\"number\"}\n",
+        "# Batch size for training.\n",
+        "batch_size: int = 2  # @param {type:\"number\"}\n",
+        "# Number of training iterations over the dataset.\n",
+        "epochs: int = 10  # @param {type:\"slider\", min:0, max:100, step:1}\n",
+        "# An optional integer that indicates the number of training steps per\n",
+        "# epoch. If set to 0, the training pipeline calculates the default\n",
+        "# steps per epoch as the training dataset size divided by batch size.\n",
+        "steps_per_epoch: int = 0  # @param {type:\"number\"}\n",
+        "# Whether to shuffle the dataset before training\n",
+        "shuffle: bool = False  # @param {type:\"boolean\"}\n",
+        "# Learning rate decay to use for gradient descent training.\n",
+        "lr_decay: float = 0.99  # @param {type:\"number\"}\n",
+        "# Gamma parameter for focal loss. Defaults to 2\n",
+        "gamma: float = 2  # @param {type:\"number\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "HwcCjwlBTQIz"
+      },
+      "source": [
+        "### Run fine-tuning\n",
+        "With your training dataset and fine-tuning options prepared, you are ready to start the fine-tuning process. This process is resource intensive and can take a few minutes to complete. On Vertex AI with GPU processing, the example fine-tuning below takes between 1-2 minutes to train on approximately 500 images.\n",
+        "\n",
+        "To begin the fine-tuning process, use the following code:\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "aec22792ee84"
+      },
+      "outputs": [],
+      "source": [
+        "model_export_path = EXPORTED_MODEL_OUTPUT_DIRECTORY\n",
+        "evaluation_result_path = EVALUATION_RESULT_OUTPUT_DIRECTORY\n",
+        "\n",
+        "model_options = {\"dropout_rate\": dropout_rate}\n",
+        "if layer_widths:\n",
+        "    model_options[\"layer_widths\"] = layer_widths\n",
+        "\n",
+        "hparams = {\n",
+        "    \"learning_rate\": learning_rate,\n",
+        "    \"batch_size\": batch_size,\n",
+        "    \"epochs\": epochs,\n",
+        "    \"shuffle\": shuffle,\n",
+        "    \"lr_decay\": lr_decay,\n",
+        "    \"gamma\": gamma,\n",
+        "}\n",
+        "if steps_per_epoch:\n",
+        "    hparams[\"steps_per_epoch\"] = steps_per_epoch\n",
+        "\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": TRAINING_MACHINE_TYPE,\n",
+        "            \"accelerator_type\": TRAINING_ACCELARATOR_TYPE,\n",
+        "            \"accelerator_count\": TRAINING_ACCELERATOR_COUNT,\n",
+        "        },\n",
+        "        \"replica_count\": 1,\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": TRAINING_CONTAINER,\n",
+        "            \"command\": [],\n",
+        "            \"args\": [\n",
+        "                \"--task_name=gesture_recognizer\",\n",
+        "                \"--training_data_path=%s\" % training_data_path,\n",
+        "                \"--model_export_path=%s\" % model_export_path,\n",
+        "                \"--evaluation_result_path=%s\" % evaluation_result_path,\n",
+        "                \"--split_ratio=%s\" % split_ratio,\n",
+        "                \"--model_options=%s\" % json.dumps(model_options),\n",
+        "                \"--hparams=%s\" % json.dumps(hparams),\n",
+        "            ],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "training_job = aiplatform.CustomJob(\n",
+        "    display_name=TRAINING_JOB_DISPLAY_NAME,\n",
+        "    project=PROJECT_ID,\n",
+        "    worker_pool_specs=worker_pool_specs,\n",
+        "    staging_bucket=STAGING_BUCKET,\n",
+        ")\n",
+        "\n",
+        "training_job.run()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "rXMF2tnV_WS0"
+      },
+      "source": [
+        "## Evaluate and export model"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "mV-Djz-frBni"
+      },
+      "source": [
+        "### Evaluate performance\n",
+        "\n",
+        "After fine-tuning the model, we evaluate the training result on a test dataset, which is typically a portion of your original dataset not used during training. Accuracy levels between 0.8 and 0.9 are generally considered very good, but your use case requirements may differ. You should also consider how fast the model can produce an inference. Higher accuracy frequently comes at the cost of longer inference times.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "09Rz1AYspK19"
+      },
+      "outputs": [],
+      "source": [
+        "def get_evaluation_result(evaluation_result_path):\n",
+        "    try:\n",
+        "        with tensorflow.io.gfile.GFile(evaluation_result_path, \"r\") as input_file:\n",
+        "            evalutation_result = json.loads(input_file.read())\n",
+        "        return evalutation_result[\"accuracy\"], evalutation_result[\"loss\"]\n",
+        "    except:\n",
+        "        print(\n",
+        "            \"Evaluation result not found. Your test dataset is likely \"\n",
+        "            + \"empty. You can adjust the size of your test dataset or adjust \"\n",
+        "            + \"how you split your dataset.\"\n",
+        "        )\n",
+        "        return None\n",
+        "\n",
+        "\n",
+        "evaluation_result = get_evaluation_result(EVALUATION_RESULT_OUTPUT_FILE)\n",
+        "\n",
+        "if evaluation_result is not None:\n",
+        "    print(\"Accuracy:\", evaluation_result[0])\n",
+        "    print(\"Loss:\", evaluation_result[1])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "g0BGaofgsMsy"
+      },
+      "source": [
+        "### Export model\n",
+        "After finetuning and evaluating the model, you can save the Tensorflow Lite model, try it out in the [Gesture Recognizer](https://mediapipe-studio.webapps.google.com/demo/gesture_recognizer) demo in MediaPipe Studio or integrate it with your on-device application by following the [Gesture recognizer task guide](https://developers.google.com/mediapipe/solutions/vision/gesture_recognizer). The exported model contains the generates required model metadata, as well as a classification label file."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NYuQowyZEtxK"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "\n",
+        "def copy_model(model_source, model_dest):\n",
+        "    ! gsutil cp {model_source} {model_dest}\n",
+        "\n",
+        "copy_model(EXPORTED_MODEL_OUTPUT_FILE, \"gesture_recognizer.task\")\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    from google.colab import files\n",
+        "\n",
+        "    files.download(\"gesture_recognizer.task\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kkH2nrpdp4sp"
+      },
+      "source": [
+        "## Clean up"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Ax6vQVZhp9pR"
+      },
+      "outputs": [],
+      "source": [
+        "# Delete training data and jobs.\n",
+        "if training_job.list(filter=f'display_name=\"{TRAINING_JOB_DISPLAY_NAME}\"'):\n",
+        "    training_job.delete()\n",
+        "\n",
+        "!gsutil rm -r {STAGING_BUCKET}"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_mediapipe_gesture_recognition.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_mediapipe_image_classification.ipynb

@@ -29,7 +29,7 @@
         "id": "TirJ-SGQseby"
       },
       "source": [
-        "# Vertex AI Model Garden MediaPipe With Image Classification\n",
+        "# Vertex AI Model Garden MediaPipe with image classification\n",
         "\n",
         "<table align=\"left\">\n",
         "  <td>\n",
@@ -44,7 +44,7 @@
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
-        "  <td>                                                                                               <td>\n",
+        "  <td>\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_mediapipe_image_classification.ipynb\">\n",
         "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
         "Open in Vertex AI Workbench\n",
@@ -145,24 +145,132 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "BF1j6f9HApxa"
+        "id": "WReHDGG5g0XY"
       },
       "source": [
-        "### Set up your Google Cloud project\n",
+        "#### Set your project ID\n",
         "\n",
-        "**The following steps are required, regardless of your notebook environment.**\n",
+        "**If you don't know your project ID**, see the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "oM1iC_MfAts1"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
         "\n",
-        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "# Set the project id\n",
+        "! gcloud config set project {PROJECT_ID}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "region"
+      },
+      "source": [
+        "#### Region\n",
         "\n",
-        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tTy1gX11kCJY"
+      },
+      "outputs": [],
+      "source": [
+        "REGION = \"us-central1\"  # @param {type: \"string\"}\n",
+        "REGION_PREFIX = REGION.split(\"-\")[0]\n",
+        "assert REGION_PREFIX in (\n",
+        "    \"us\",\n",
+        "    \"europe\",\n",
+        "    \"asia\",\n",
+        "), f'{REGION} is not supported. It must be prefixed by \"us\", \"asia\", or \"europe\".'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zgPO1eR3CYjk"
+      },
+      "source": [
+        "### Create a Cloud Storage bucket\n",
         "\n",
-        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
-        "1. If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
+        "Create a storage bucket to store intermediate artifacts such as datasets."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "MzGDU7TWdts_"
+      },
+      "outputs": [],
+      "source": [
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-EcIXiGsCePi"
+      },
+      "source": [
+        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NIq7R4HZCfIc"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "960505627ddf"
+      },
+      "source": [
+        "### Import libraries"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "PyQmSRbKA8r-"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "import os\n",
+        "from datetime import datetime\n",
         "\n",
-        "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
-        "Cloud SDK uses the right project for all the commands in this notebook.\n",
+        "import tensorflow\n",
+        "from google.cloud import aiplatform"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "init_aip:mbsdk,all"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
         "\n",
-        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
+        "Initialize the Vertex AI SDK for Python for your project."
       ]
     },
     {
@@ -173,13 +281,6 @@
       },
       "outputs": [],
       "source": [
-        "import json\n",
-        "import os\n",
-        "from datetime import datetime\n",
-        "\n",
-        "import tensorflow\n",
-        "from google.cloud import aiplatform\n",
-        "\n",
         "now = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
         "\n",
         "# The project and bucket are for experiments below.\n",
@@ -280,14 +381,14 @@
       "source": [
         "### Set fine-tuning options\n",
         "\n",
-        "There are a few required settings to run fine-tuning aside from your training dataset: output directory for the model and the model architecture. Use HParams object export_dir parameter to specify a model output directory. Use the SupportedModels class to specify the model architecture. The image classifier solution supports the following model architectures:\n",
+        "You can pick between different model architectures to further customize your training:\n",
         "\n",
         "*   MobileNet-V2\n",
         "*   EfficientNet-Lite0\n",
         "*   EfficientNet-Lite2\n",
         "*   EfficientNet-Lite4\n",
         "\n",
-        "To set the required parameters, use the following code:"
+        "To set the model architecture and other training parameters, adjust the following values:"
       ]
     },
     {

notebooks/community/model_garden/model_garden_mediapipe_object_detection.ipynb

@@ -29,7 +29,7 @@
         "id": "TirJ-SGQseby"
       },
       "source": [
-        "# Vertex AI Model Garden MediaPipe With Object Detection\n",
+        "# Vertex AI Model Garden MediaPipe with object detection\n",
         "\n",
         "<table align=\"left\">\n",
         "  <td>\n",
@@ -44,7 +44,7 @@
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
-        "  <td>                                                                                               <td>\n",
+        "  <td>\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_mediapipe_object_detection.ipynb\">\n",
         "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
         "Open in Vertex AI Workbench\n",
@@ -112,7 +112,7 @@
         "id": "z__i0w0lCAsW"
       },
       "source": [
-        "### Colab Only\n",
+        "### Colab only\n",
         "Run the following commands to install dependencies and to authenticate with Google Cloud if running on Colab."
       ]
     },
@@ -145,24 +145,132 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "BF1j6f9HApxa"
+        "id": "WReHDGG5g0XY"
       },
       "source": [
-        "### Set up your Google Cloud project\n",
+        "#### Set your project ID\n",
         "\n",
-        "**The following steps are required, regardless of your notebook environment.**\n",
+        "**If you don't know your project ID**, see the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "oM1iC_MfAts1"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
         "\n",
-        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "# Set the project id\n",
+        "! gcloud config set project {PROJECT_ID}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "region"
+      },
+      "source": [
+        "#### Region\n",
         "\n",
-        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tTy1gX11kCJY"
+      },
+      "outputs": [],
+      "source": [
+        "REGION = \"us-central1\"  # @param {type: \"string\"}\n",
+        "REGION_PREFIX = REGION.split(\"-\")[0]\n",
+        "assert REGION_PREFIX in (\n",
+        "    \"us\",\n",
+        "    \"europe\",\n",
+        "    \"asia\",\n",
+        "), f'{REGION} is not supported. It must be prefixed by \"us\", \"asia\", or \"europe\".'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zgPO1eR3CYjk"
+      },
+      "source": [
+        "### Create a Cloud Storage bucket\n",
         "\n",
-        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
-        "1. If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
+        "Create a storage bucket to store intermediate artifacts such as datasets."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "MzGDU7TWdts_"
+      },
+      "outputs": [],
+      "source": [
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-EcIXiGsCePi"
+      },
+      "source": [
+        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NIq7R4HZCfIc"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "960505627ddf"
+      },
+      "source": [
+        "### Import libraries"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "PyQmSRbKA8r-"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "import os\n",
+        "from datetime import datetime\n",
         "\n",
-        "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
-        "Cloud SDK uses the right project for all the commands in this notebook.\n",
+        "import tensorflow\n",
+        "from google.cloud import aiplatform"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "init_aip:mbsdk,all"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
         "\n",
-        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
+        "Initialize the Vertex AI SDK for Python for your project."
       ]
     },
     {
@@ -173,32 +281,8 @@
       },
       "outputs": [],
       "source": [
-        "import json\n",
-        "import os\n",
-        "from datetime import datetime\n",
-        "\n",
-        "import tensorflow\n",
-        "from google.cloud import aiplatform\n",
-        "\n",
         "now = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
         "\n",
-        "# The project and bucket are for experiments below.\n",
-        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
-        "# The form for BUCKET_URI is gs://<bucket-name>.\\n\",\n",
-        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# You can choose a region from https://cloud.google.com/about/locations.\n",
-        "# Only regions prefixed by \"us\", \"asia\", or \"europe\" are supported.\n",
-        "REGION = \"us-central1\"  # @param {type:\"string\"}\n",
-        "REGION_PREFIX = REGION.split(\"-\")[0]\n",
-        "assert REGION_PREFIX in (\n",
-        "    \"us\",\n",
-        "    \"europe\",\n",
-        "    \"asia\",\n",
-        "), f'{REGION} is not supported. It must be prefixed by \"us\", \"asia\", or \"europe\".'\n",
-        "\n",
-        "! gcloud config set project $PROJECT_ID\n",
-        "\n",
         "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temp/%s\" % now)\n",
         "\n",
         "EVALUATION_RESULT_OUTPUT_DIRECTORY = os.path.join(STAGING_BUCKET, \"evaluation\")\n",
@@ -359,12 +443,12 @@
       "source": [
         "### Set fine-tuning options\n",
         "\n",
-        "There are a few required settings to run fine-tuning aside from your training dataset: output directory for the model, and the model architecture. Use HParams to specify the export_dir parameter for the output directory. Use the SupportedModels class to specify the model architecture. The object detector solution supports the following model architectures:\n",
+        "You can pick between different model architectures to further customize your training:\n",
         "\n",
         "*   MobileNet-V2\n",
         "*   MobileNet-MultiHW-AVG\n",
         "\n",
-        "To set the parameters, use the following code:"
+        "To set the model architecture and other training parameters, adjust the following values:"
       ]
     },
     {

notebooks/community/model_garden/model_garden_mediapipe_text_classification.ipynb

@@ -0,0 +1,616 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ur8xi4C7S06n"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2023 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": "TirJ-SGQseby"
+      },
+      "source": [
+        "# Vertex AI Model Garden MediaPipe with text classification\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_mediapipe_text_classification.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_mediapipe_text_classification.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_mediapipe_image_classification.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "Open in Vertex AI Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dwGLvtIeECLK"
+      },
+      "source": [
+        "**_NOTE_**: This notebook has been tested in the following environment:\n",
+        "\n",
+        "* Python version = 3.9"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tvgnzT1CKxrO"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates how to use [MediaPipe Model Maker](https://developers.google.com/mediapipe/solutions/model_maker) to train an on-device text classification model in Vertex AI Model Garden.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "* Train new models\n",
+        "  * Convert input data to training formats\n",
+        "  * Create [custom jobs](https://cloud.google.com/vertex-ai/docs/training/create-custom-job) to train new models\n",
+        "  * Export models\n",
+        "\n",
+        "* Cleanup resources\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI\n",
+        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
+        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
+        "Calculator](https://cloud.google.com/products/calculator/)\n",
+        "to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "KEukV6uRk_S3"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "z__i0w0lCAsW"
+      },
+      "source": [
+        "### Colab only\n",
+        "Run the following commands to install dependencies and to authenticate with Google Cloud if running on Colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Jvqs-ehKlaYh"
+      },
+      "outputs": [],
+      "source": [
+        "! pip3 install --upgrade pip\n",
+        "\n",
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
+        "\n",
+        "    # Automatically restart kernel after installs\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)\n",
+        "\n",
+        "    from google.colab import auth as google_auth\n",
+        "\n",
+        "    google_auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "WReHDGG5g0XY"
+      },
+      "source": [
+        "#### Set your project ID\n",
+        "\n",
+        "**If you don't know your project ID**, see the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "oM1iC_MfAts1"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Set the project id\n",
+        "! gcloud config set project {PROJECT_ID}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "region"
+      },
+      "source": [
+        "#### Region\n",
+        "\n",
+        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tTy1gX11kCJY"
+      },
+      "outputs": [],
+      "source": [
+        "REGION = \"us-central1\"  # @param {type: \"string\"}\n",
+        "REGION_PREFIX = REGION.split(\"-\")[0]\n",
+        "assert REGION_PREFIX in (\n",
+        "    \"us\",\n",
+        "    \"europe\",\n",
+        "    \"asia\",\n",
+        "), f'{REGION} is not supported. It must be prefixed by \"us\", \"asia\", or \"europe\".'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zgPO1eR3CYjk"
+      },
+      "source": [
+        "### Create a Cloud Storage bucket\n",
+        "\n",
+        "Create a storage bucket to store intermediate artifacts such as datasets."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "MzGDU7TWdts_"
+      },
+      "outputs": [],
+      "source": [
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-EcIXiGsCePi"
+      },
+      "source": [
+        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NIq7R4HZCfIc"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "960505627ddf"
+      },
+      "source": [
+        "### Import libraries"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "PyQmSRbKA8r-"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "import os\n",
+        "from datetime import datetime\n",
+        "\n",
+        "from google.cloud import aiplatform"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "init_aip:mbsdk,all"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
+        "\n",
+        "Initialize the Vertex AI SDK for Python for your project."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9wExiMUxFk91"
+      },
+      "outputs": [],
+      "source": [
+        "now = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temp/%s\" % now)\n",
+        "\n",
+        "EVALUATION_RESULT_OUTPUT_DIRECTORY = os.path.join(STAGING_BUCKET, \"evaluation\")\n",
+        "EVALUATION_RESULT_OUTPUT_FILE = os.path.join(\n",
+        "    EVALUATION_RESULT_OUTPUT_DIRECTORY, \"evaluation.json\"\n",
+        ")\n",
+        "\n",
+        "EXPORTED_MODEL_OUTPUT_DIRECTORY = os.path.join(STAGING_BUCKET, \"model\")\n",
+        "EXPORTED_MODEL_OUTPUT_FILE = os.path.join(\n",
+        "    EXPORTED_MODEL_OUTPUT_DIRECTORY, \"model.tflite\"\n",
+        ")\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "n6IFz75WGCam"
+      },
+      "source": [
+        "### Define training machine specs"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "riG_qUokg0XZ"
+      },
+      "outputs": [],
+      "source": [
+        "TRAINING_JOB_DISPLAY_NAME = \"mediapipe_text_classifier_%s\" % now\n",
+        "TRAINING_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/mediapipe-train\"\n",
+        "TRAINING_MACHINE_TYPE = \"n1-highmem-16\"\n",
+        "TRAINING_ACCELARATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
+        "TRAINING_ACCELERATOR_COUNT = 2"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-rsdAcBV-vlf"
+      },
+      "source": [
+        "## Train your customized models"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zgPO1eR3CYjk"
+      },
+      "source": [
+        "### Get the Dataset\n",
+        "\n",
+        "The following code block uses the [SST-2](https://nlp.stanford.edu/sentiment/index.html) (Stanford Sentiment Treebank) dataset which contains 67,349 movie reviews for training and 872 movie reviews for testing. The dataset has two classes: positive and negative movie reviews. Positive reviews are labeled with 1 and negative reviews with 0.\n",
+        "\n",
+        "The SST-2 dataset is stored as a TSV file. The only difference between the TSV and CSV formats is that TSV uses a tab `\\t` character as its delimiter and CSV uses a comma `,`.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "IndQ_m6ddUEM"
+      },
+      "outputs": [],
+      "source": [
+        "training_data_path = (\n",
+        "    \"gs://mediapipe-tasks/text_classifier/SST-2/train.tsv\"  # @param {type:\"string\"}\n",
+        ")\n",
+        "validation_data_path = (\n",
+        "    \"gs://mediapipe-tasks/text_classifier/SST-2/dev.tsv\"  # @param {type:\"string\"}\n",
+        ")\n",
+        "\n",
+        "# The delimiter used in the dataset.\n",
+        "delimiter = \"\\t\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Character used to quote fields that contain special characters\n",
+        "# like the `delimiter`.\n",
+        "quotechar = \"\\t\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Sequence of keys for the CSV columns (represented as a comma\n",
+        "# separated list). If empty, the first row of the CSV file is used\n",
+        "# as the keys\n",
+        "fieldnames = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Column name for the input text.\n",
+        "text_column = \"sentence\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Column name for the labels.\n",
+        "label_column = \"label\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aaff6f5be7f6"
+      },
+      "source": [
+        "### Set fine-tuning options\n",
+        "\n",
+        "You can pick between different model architectures to further customize your training:\n",
+        "\n",
+        "*   Average Word Embedding Model\n",
+        "*   BERT-classifier\n",
+        "\n",
+        "To set the model architecture and other training parameters, adjust the following values:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "um_XKbmpTaHx"
+      },
+      "outputs": [],
+      "source": [
+        "model_architecture = (\n",
+        "    \"average_word_embedding\"  # @param [\"average_word_embedding\", \"mobilebert\"]\n",
+        ")\n",
+        "\n",
+        "# The learning rate to use for gradient descent-based\n",
+        "# optimizers. Defaults to 3e-5 for the BERT-based classifier\n",
+        "# and 0 for the average word-embedding classifier because\n",
+        "# it does not need such an optimizer.\n",
+        "learning_rate: float = 0.0  # @param {type:\"number\"}\n",
+        "\n",
+        "# Batch size for training. Defaults to 32 for the average\n",
+        "# word-embedding classifier and 48 for the BERT-based\n",
+        "# classifier.\n",
+        "batch_size: int = 48  # @param {type:\"number\"}\n",
+        "\n",
+        "# Number of training iterations over the dataset. Defaults\n",
+        "# to 10 for the average word-embedding classifier and 3\n",
+        "# for the BERT-based classifier.\n",
+        "epochs: int = 10  # @param {type:\"slider\", min:0, max:100, step:1}\n",
+        "\n",
+        "# An integer that indicates the number of training steps per\n",
+        "# epoch. If set to 0, the training pipeline calculates the\n",
+        "# default steps per epoch as the training dataset size\n",
+        "# divided by batch size.\n",
+        "steps_per_epoch: int = 0  # @param {type:\"number\"}\n",
+        "\n",
+        "# Controls whether the dataset is shuffled before training.\n",
+        "shuffle: bool = False  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# Length of the sequence to feed into the model.\n",
+        "seq_len: int = 256  # @param {type:\"number\"}\n",
+        "\n",
+        "# Whether to convert all uppercase characters to lowercase\n",
+        "# during preprocessing.\n",
+        "do_lower_case: bool = True  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# The rate for dropout.\n",
+        "dropout_rate: float = 0.2  # @param {type:\"number\"}\n",
+        "\n",
+        "# Dimension of the word embedding. Only used for the Average Word\n",
+        "# Embedding Model.\n",
+        "wordvec_dim: int = 16  # @param {type:\"number\"}\n",
+        "\n",
+        "# Number of words to generate the vocabulary from data.\n",
+        "# Only used for the Average Word Embedding Model.\n",
+        "vocab_size: int = 10000  # @param {type:\"number\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "HwcCjwlBTQIz"
+      },
+      "source": [
+        "### Run fine-tuning\n",
+        "With your training dataset and fine-tuning options prepared, you are ready to start the fine-tuning process. This process is resource intensive and can take a few minutes to a few hours depending on the model archtiecture and your available compute resources. On Vertex AI with GPU processing, the example fine-tuning below takes between 2-3 minutes to train an Average Word Embedding Model on the SST-2 dataset.\n",
+        "\n",
+        "To begin the fine-tuning process, use the following code:\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "aec22792ee84"
+      },
+      "outputs": [],
+      "source": [
+        "model_export_path = EXPORTED_MODEL_OUTPUT_DIRECTORY\n",
+        "evaluation_result_path = EVALUATION_RESULT_OUTPUT_DIRECTORY\n",
+        "\n",
+        "preprocessing_params = {\n",
+        "    \"text_column\": text_column,\n",
+        "    \"label_column\": label_column,\n",
+        "    \"delimiter\": delimiter,\n",
+        "    \"quotechar\": quotechar,\n",
+        "}\n",
+        "if fieldnames:\n",
+        "    preprocessing_params[\"fieldnames\"] = [\n",
+        "        fieldname.strip() for fieldname in fieldnames.split(\",\")\n",
+        "    ]\n",
+        "\n",
+        "hparams = {\n",
+        "    \"learning_rate\": learning_rate,\n",
+        "    \"batch_size\": batch_size,\n",
+        "    \"epochs\": epochs,\n",
+        "    \"shuffle\": shuffle,\n",
+        "}\n",
+        "if steps_per_epoch:\n",
+        "    hparams[\"steps_per_epoch\"] = steps_per_epoch\n",
+        "\n",
+        "model_options = {\n",
+        "    \"dropout_rate\": dropout_rate,\n",
+        "    \"wordvec_dim\": wordvec_dim,\n",
+        "    \"do_lower_case\": do_lower_case,\n",
+        "    \"vocab_size\": vocab_size,\n",
+        "    \"dropout_rate\": dropout_rate,\n",
+        "}\n",
+        "\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": TRAINING_MACHINE_TYPE,\n",
+        "            \"accelerator_type\": TRAINING_ACCELARATOR_TYPE,\n",
+        "            \"accelerator_count\": TRAINING_ACCELERATOR_COUNT,\n",
+        "        },\n",
+        "        \"replica_count\": 1,\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": TRAINING_CONTAINER,\n",
+        "            \"command\": [],\n",
+        "            \"args\": [\n",
+        "                \"--task_name=text_classifier\",\n",
+        "                \"--training_data_path=%s\" % training_data_path,\n",
+        "                \"--validation_data_path=%s\" % validation_data_path,\n",
+        "                \"--evaluation_result_path=%s\" % evaluation_result_path,\n",
+        "                \"--model_export_path=%s\" % model_export_path,\n",
+        "                \"--model_architecture=%s\" % model_architecture,\n",
+        "                \"--preprocessing_params=%s\" % json.dumps(preprocessing_params),\n",
+        "                \"--hparams=%s\" % json.dumps(hparams),\n",
+        "                \"--model_options=%s\" % json.dumps(model_options),\n",
+        "            ],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "training_job = aiplatform.CustomJob(\n",
+        "    display_name=TRAINING_JOB_DISPLAY_NAME,\n",
+        "    project=PROJECT_ID,\n",
+        "    worker_pool_specs=worker_pool_specs,\n",
+        "    staging_bucket=STAGING_BUCKET,\n",
+        ")\n",
+        "\n",
+        "training_job.run()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "rXMF2tnV_WS0"
+      },
+      "source": [
+        "## Export model"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "g0BGaofgsMsy"
+      },
+      "source": [
+        "After finetuning, you can save the Tensorflow Lite model, try it out in the [Text Classification](https://mediapipe-studio.webapps.google.com/demo/text_classifier) demo in MediaPipe Studio or integrate it with your on-device application by following the [Text classification task guide](https://developers.google.com/mediapipe/solutions/text/text_classifier). The exported model contains the generates required model metadata, as well as a classification label file."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NYuQowyZEtxK"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "\n",
+        "def copy_model(model_source, model_dest):\n",
+        "    ! gsutil cp {model_source} {model_dest}\n",
+        "\n",
+        "copy_model(EXPORTED_MODEL_OUTPUT_FILE, \"text_classification_model.tflite\")\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    from google.colab import files\n",
+        "\n",
+        "    files.download(\"text_classification_model.tflite\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kkH2nrpdp4sp"
+      },
+      "source": [
+        "## Clean up"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Ax6vQVZhp9pR"
+      },
+      "outputs": [],
+      "source": [
+        "# Delete training data and jobs.\n",
+        "if training_job.list(filter=f'display_name=\"{TRAINING_JOB_DISPLAY_NAME}\"'):\n",
+        "    training_job.delete()\n",
+        "\n",
+        "!gsutil rm -r {STAGING_BUCKET}"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_mediapipe_text_classification.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_movinet_clip_classification.ipynb

@@ -0,0 +1,873 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ur8xi4C7S06n"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2023 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": "TirJ-SGQseby"
+      },
+      "source": [
+        "# Vertex AI Model Garden MoViNet video clip classification\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_movinet_clip_classification.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_movinet_clip_classification.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>                                                                                               <td>\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_movinet_clip_classification.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "Open in Vertex AI Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dwGLvtIeECLK"
+      },
+      "source": [
+        "**_NOTE_**: This notebook has been tested in the following environment:\n",
+        "\n",
+        "* Python version = 3.9"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tvgnzT1CKxrO"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates how to use [MoViNet](https://github.com/tensorflow/models/tree/master/official/projects/movinet) in Vertex AI Model Garden.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "* Train new models\n",
+        "  * Convert input data to training formats\n",
+        "  * Create [hyperparameter tuning jobs](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview) to train new models\n",
+        "  * Find and export best models\n",
+        "\n",
+        "* Test trained models\n",
+        "  * Upload models to the [Vertex AI Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction)\n",
+        "  * Run batch predictions\n",
+        "\n",
+        "* Clean up resources\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI\n",
+        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
+        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
+        "Calculator](https://cloud.google.com/products/calculator/)\n",
+        "to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "KEukV6uRk_S3"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "z__i0w0lCAsW"
+      },
+      "source": [
+        "### Colab Only\n",
+        "Run the following commands for Colab or skip this section if you use Workbench."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Jvqs-ehKlaYh"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
+        "\n",
+        "    # Automatically restart kernel after installs\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)\n",
+        "\n",
+        "    from google.colab import auth as google_auth\n",
+        "\n",
+        "    google_auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "BF1j6f9HApxa"
+      },
+      "source": [
+        "### Set up your Google Cloud project\n",
+        "\n",
+        "**The following steps are required, regardless of your notebook environment.**\n",
+        "\n",
+        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "\n",
+        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
+        "1. If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
+        "\n",
+        "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
+        "Cloud SDK uses the right project for all the commands in this notebook.\n",
+        "\n",
+        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9wExiMUxFk91"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "# The GCP project ID for experiments.\n",
+        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Bucket URI with gs:// prefix.\n",
+        "BUCKET_URI = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# You can choose a region from https://cloud.google.com/about/locations.\n",
+        "# Only regions prefixed by \"us\", \"asia\", or \"europe\" are supported.\n",
+        "REGION = \"us-central1\"  # @param {type:\"string\"}\n",
+        "REGION_PREFIX = REGION.split(\"-\")[0]\n",
+        "assert REGION_PREFIX in (\n",
+        "    \"us\",\n",
+        "    \"europe\",\n",
+        "    \"asia\",\n",
+        "), f'{REGION} is not supported. It must be prefixed by \"us\", \"asia\", or \"europe\".'\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "CHECKPOINT_BUCKET = os.path.join(BUCKET_URI, \"ckpt\")\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
+        "\n",
+        "# Download config files.\n",
+        "CONFIG_DIR = os.path.join(BUCKET_URI, \"config\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "n6IFz75WGCam"
+      },
+      "source": [
+        "### Define constants"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "riG_qUokg0XZ"
+      },
+      "outputs": [],
+      "source": [
+        "OBJECTIVE = \"vcn\"\n",
+        "\n",
+        "# Data converter constants.\n",
+        "DATA_CONVERTER_JOB_PREFIX = \"data_converter\"\n",
+        "DATA_CONVERTER_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/data-converter\"\n",
+        "DATA_CONVERTER_MACHINE_TYPE = \"n1-highmem-8\"\n",
+        "\n",
+        "# Training constants.\n",
+        "TRAINING_JOB_PREFIX = \"train\"\n",
+        "TRAIN_CONTAINER_URI = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/movinet-train\"\n",
+        "TRAIN_MACHINE_TYPE = \"n1-highmem-16\"\n",
+        "TRAIN_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
+        "TRAIN_NUM_GPU = 2\n",
+        "\n",
+        "# Evaluation constants.\n",
+        "EVALUATION_METRIC = \"accuracy\"\n",
+        "\n",
+        "# Export constants.\n",
+        "EXPORT_JOB_PREFIX = \"export\"\n",
+        "EXPORT_CONTAINER_URI = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/movinet-model-export\"\n",
+        "EXPORT_MACHINE_TYPE = \"n1-highmem-8\"\n",
+        "\n",
+        "# Prediction constants.\n",
+        "# You can adjust accelerator types and machine types to get faster predictions.\n",
+        "UPLOAD_JOB_PREFIX = \"upload\"\n",
+        "PREDICTION_CONTAINER_URI = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/movinet-serve\"\n",
+        "PREDICTION_PORT = 8501\n",
+        "PREDICTION_ACCELERATOR_COUNT = 1\n",
+        "PREDICTION_ACCELERATOR_TYPE = \"NVIDIA_TESLA_T4\"\n",
+        "PREDICTION_MACHINE_TYPE = \"n1-standard-4\"\n",
+        "PREDICTION_JOB_PREFIX = \"predict\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ZZFPe_GezXg8"
+      },
+      "source": [
+        "### Define common helper functions"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "XcYUGwr-AJGY"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "from datetime import datetime\n",
+        "\n",
+        "import numpy as np\n",
+        "import tensorflow as tf\n",
+        "import yaml\n",
+        "\n",
+        "\n",
+        "def get_job_name_with_datetime(prefix: str):\n",
+        "    \"\"\"Returns a timestamped job name with the given prefix.\"\"\"\n",
+        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
+        "\n",
+        "\n",
+        "def print_response_instance(json_str: str, label_map: dict[int, str]):\n",
+        "    \"\"\"Prints summary of a prediction JSON result from the model response.\"\"\"\n",
+        "    json_obj = json.loads(json_str)\n",
+        "    if \"prediction\" not in json_obj:\n",
+        "        print(\"Error:\", json_str)\n",
+        "        return\n",
+        "    instance = json_obj[\"instance\"]\n",
+        "    prediction = json_obj[\"prediction\"]\n",
+        "    gcs_uri = instance[\"content\"]\n",
+        "    time_start = instance.get(\"timeSegmentStart\", \"0.0s\")\n",
+        "    time_end = instance.get(\"timeSegmentEnd\", \"Infinity\")\n",
+        "    max_idx = np.argmax(prediction)\n",
+        "    print(f\"{gcs_uri} {time_start}-{time_end}:\", label_map[max_idx])\n",
+        "\n",
+        "\n",
+        "def get_label_map(label_map_yaml_filepath: str) -> tuple[dict[int, str], int]:\n",
+        "    \"\"\"Reads label map from a YAML file and returns the label map with the number of classes.\"\"\"\n",
+        "    with tf.io.gfile.GFile(label_map_yaml_filepath, \"rb\") as input_file:\n",
+        "        label_map = yaml.safe_load(input_file.read())[\"label_map\"]\n",
+        "    num_classes = max(label_map.keys()) + 1\n",
+        "    return label_map, num_classes\n",
+        "\n",
+        "\n",
+        "def get_best_trial(model_dir, max_trial_count, evaluation_metric):\n",
+        "    \"\"\"Finds the best trial directory and eval results from a hyperparameter tuning job.\"\"\"\n",
+        "    best_trial_dir = \"\"\n",
+        "    best_trial_evaluation_results = {}\n",
+        "    best_performance = -1\n",
+        "\n",
+        "    for i in range(max_trial_count):\n",
+        "        current_trial = i + 1\n",
+        "        current_trial_dir = os.path.join(model_dir, \"trial_\" + str(current_trial))\n",
+        "        current_trial_best_ckpt_dir = os.path.join(current_trial_dir, \"best_ckpt\")\n",
+        "        current_trial_best_ckpt_evaluation_filepath = os.path.join(\n",
+        "            current_trial_best_ckpt_dir, \"info.json\"\n",
+        "        )\n",
+        "        with tf.io.gfile.GFile(current_trial_best_ckpt_evaluation_filepath, \"rb\") as f:\n",
+        "            eval_metric_results = json.load(f)\n",
+        "            current_performance = eval_metric_results[evaluation_metric]\n",
+        "            if current_performance > best_performance:\n",
+        "                best_performance = current_performance\n",
+        "                best_trial_dir = current_trial_dir\n",
+        "                best_trial_evaluation_results = eval_metric_results\n",
+        "    return best_trial_dir, best_trial_evaluation_results\n",
+        "\n",
+        "\n",
+        "def find_checkpoint_in_dir(checkpoint_dir: str):\n",
+        "    \"\"\"Finds a checkpoint path relative to the directory.\"\"\"\n",
+        "    for root, dirs, files in tf.io.gfile.walk(checkpoint_dir):\n",
+        "        for file in files:\n",
+        "            if file.endswith(\".index\"):\n",
+        "                return os.path.join(root, os.path.splitext(file)[0])\n",
+        "\n",
+        "\n",
+        "def upload_checkpoint_to_gcs(checkpoint_url: str):\n",
+        "    \"\"\"Uploads a compressed .tar.gz checkpoint at the given URL to Cloud Storage.\"\"\"\n",
+        "    filename = os.path.basename(checkpoint_url)\n",
+        "    checkpoint_name = filename.replace(\".tar.gz\", \"\")\n",
+        "    print(\"Download checkpoint from\", checkpoint_url, \"and store to\", CHECKPOINT_BUCKET)\n",
+        "    ! wget $checkpoint_url -O $filename\n",
+        "    ! mkdir -p $checkpoint_name\n",
+        "    ! tar -xvzf $filename -C $checkpoint_name\n",
+        "\n",
+        "    checkpoint_path = find_checkpoint_in_dir(checkpoint_name)\n",
+        "    checkpoint_path = os.path.relpath(checkpoint_path, checkpoint_name)\n",
+        "\n",
+        "    ! gsutil cp -r $checkpoint_name $CHECKPOINT_BUCKET/\n",
+        "    checkpoint_uri = os.path.join(CHECKPOINT_BUCKET, checkpoint_name, checkpoint_path)\n",
+        "    print(\"Checkpoint uploaded to\", checkpoint_uri)\n",
+        "    return checkpoint_uri\n",
+        "\n",
+        "\n",
+        "def upload_config_to_gcs(url: str):\n",
+        "    \"\"\"Uploads a config file at the given URL to Cloud Storage.\"\"\"\n",
+        "    filename = os.path.basename(url)\n",
+        "    destination = os.path.join(CONFIG_DIR, filename)\n",
+        "    print(\"Copy\", url, \"to\", destination)\n",
+        "    ! wget \"$url\" -O \"$filename\"\n",
+        "    ! gsutil cp \"$filename\" \"$destination\"\n",
+        "    return destination"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "RB_xY9ipr7ZU"
+      },
+      "source": [
+        "## Train new models\n",
+        "This section shows how to train new models.\n",
+        "1. Convert input data to training formats\n",
+        "2. Create hyperparameter tuning jobs to train new models\n",
+        "3. Find and export best models\n",
+        "\n",
+        "If you already trained models, please go to the section `Test Trained models`.\n",
+        "\n",
+        "Please select a model:\n",
+        "* `model_id`: MoViNet model variant ID, one of `a0`, `a1`, `a2`, `a3`, `a4`, `a5`. The model with a larger number requires more resources to train, and is expected to have a higher accuracy and latency. Here, we use `a0` for demonstration purpose.\n",
+        "* `model_mode`: MoViNet model type, either `base` or `stream`. The base model has a slightly higher accuracy, while the streaming model is optimized for streaming and faster CPU inference. See [official MoViNet docs](https://github.com/tensorflow/models/tree/master/official/projects/movinet) for more information.\n",
+        "\n",
+        "**Note**: The prediction container only supports base model (non-streaming) for now. If you train a streaming model, you need to download the model and refer to the [MoViNet official guide](https://github.com/tensorflow/models/blob/master/official/projects/movinet/movinet_streaming_model_training_and_inference.ipynb) for running predictions locally."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "3Ry1mw6AHLTy"
+      },
+      "outputs": [],
+      "source": [
+        "model_id = \"a0\"  # @param [\"a0\", \"a1\", \"a2\", \"a3\", \"a4\", \"a5\"]\n",
+        "model_mode = \"base\"  # @param [\"base\", \"stream\"]\n",
+        "is_stream = model_mode == \"stream\"\n",
+        "model_name = f\"movinet_{model_id}_{model_mode}\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zgPO1eR3CYjk"
+      },
+      "source": [
+        "### Prepare input data for training\n",
+        "\n",
+        "Prepare data in the format as described [here](https://cloud.google.com/vertex-ai/docs/video-data/classification/prepare-data), and then convert them to the training formats by running the cell below:\n",
+        "\n",
+        "* `input_file_path`: The input file path to the prepared data.\n",
+        "* `input_file_type`: The input file type, such as `csv` or `jsonl`.\n",
+        "* `split_ratio`: Three comma separated floats indicating the proportion of data to split into train/validation/test. They must add up to 1.\n",
+        "* `num_shard`: Three comma separated integers indicating the shards for train/validation/test.\n",
+        "* `output_dir`: The output directory, which will contain converted train/test/validation data.\n",
+        "* `output_fps`: The sampling rate of the video; Frames per second."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "IndQ_m6ddUEM"
+      },
+      "outputs": [],
+      "source": [
+        "# This job will convert input data as training format, with given split ratios\n",
+        "# and number of shards on train/test/validation.\n",
+        "\n",
+        "data_converter_job_name = get_job_name_with_datetime(\n",
+        "    DATA_CONVERTER_JOB_PREFIX + \"_\" + OBJECTIVE\n",
+        ")\n",
+        "\n",
+        "input_file_path = \"\"  # @param {type:\"string\"}\n",
+        "input_file_type = \"csv\"  # @param [\"csv\", \"jsonl\"]\n",
+        "output_fps = 5  # @param {type:\"integer\"}\n",
+        "split_ratio = \"0.8,0.1,0.1\"\n",
+        "num_shard = \"10,10,10\"\n",
+        "data_converter_output_dir = os.path.join(BUCKET_URI, data_converter_job_name)\n",
+        "\n",
+        "\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": DATA_CONVERTER_MACHINE_TYPE,\n",
+        "        },\n",
+        "        \"replica_count\": 1,\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": DATA_CONVERTER_CONTAINER,\n",
+        "            \"command\": [],\n",
+        "            \"args\": [\n",
+        "                \"--input_file_path=%s\" % input_file_path,\n",
+        "                \"--input_file_type=%s\" % input_file_type,\n",
+        "                \"--objective=%s\" % OBJECTIVE,\n",
+        "                \"--num_shard=%s\" % num_shard,\n",
+        "                \"--split_ratio=%s\" % split_ratio,\n",
+        "                \"--output_dir=%s\" % data_converter_output_dir,\n",
+        "                \"--output_fps=%d\" % output_fps,\n",
+        "            ],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "data_converter_custom_job = aiplatform.CustomJob(\n",
+        "    display_name=data_converter_job_name,\n",
+        "    project=PROJECT_ID,\n",
+        "    worker_pool_specs=worker_pool_specs,\n",
+        "    staging_bucket=STAGING_BUCKET,\n",
+        ")\n",
+        "\n",
+        "data_converter_custom_job.run()\n",
+        "\n",
+        "input_train_data_path = os.path.join(data_converter_output_dir, \"train.tfrecord*\")\n",
+        "input_validation_data_path = os.path.join(data_converter_output_dir, \"val.tfrecord*\")\n",
+        "label_map_path = os.path.join(data_converter_output_dir, \"label_map.yaml\")\n",
+        "print(\"input_train_data_path for training: \", input_train_data_path)\n",
+        "print(\"input_validation_data_path for training: \", input_validation_data_path)\n",
+        "print(\"label_map_path for prediction: \", label_map_path)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aaff6f5be7f6"
+      },
+      "source": [
+        "### Create a Vertex AI custom job with hyperparameter tuning\n",
+        "\n",
+        "You use the Vertex AI SDK to create and run the [hyperparameter tuning job](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview) with Vertex AI Model Garden training docker images.\n",
+        "\n",
+        "#### Define the following specifications\n",
+        "\n",
+        "* `worker_pool_specs`: A list of dictionaries specifying the machine type and docker image. This example defines a single node cluster with one `n1-standard-4` machine with 2 `NVIDIA_TESLA_V100` GPUs.\n",
+        "\n",
+        "  **Note**: We recommend using 8 GPUs for MoViNet-A2 and larger. Since loading video data requires a lot of GPU memory, it is recommended to experiment with a small batch size first.\n",
+        "* `parameter_spec`: Dictionary specifying the parameters to optimize. The dictionary key is the string assigned to the command line argument for each hyperparameter in your training application code, and the dictionary value is the parameter specification. The parameter specification includes the type, min/max values, and scale for the hyperparameter.\n",
+        "* `metric_spec`: Dictionary specifying the metric to optimize. The dictionary key is the `hyperparameter_metric_tag` that you set in your training application code, and the value is the optimization goal."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "um_XKbmpTaHx"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud.aiplatform import hyperparameter_tuning as hpt\n",
+        "\n",
+        "# Input train and validation datasets can be found from the section above\n",
+        "# `Prepare input data for training`.\n",
+        "# Or, set prepared datasets paths if already exist.\n",
+        "# input_train_data_path = \"\"\n",
+        "# input_validation_data_path = \"\"\n",
+        "# label_map_path = \"\"\n",
+        "\n",
+        "train_job_name = get_job_name_with_datetime(f\"{TRAINING_JOB_PREFIX}_{model_name}\")\n",
+        "model_dir = os.path.join(BUCKET_URI, train_job_name)\n",
+        "label_map, num_classes = get_label_map(label_map_path)\n",
+        "\n",
+        "# Uploads pretained checkpoint to GCS bucket.\n",
+        "init_checkpoint = f\"https://storage.googleapis.com/tf_model_garden/vision/movinet/{model_name}_with_backbone.tar.gz\"\n",
+        "init_checkpoint = upload_checkpoint_to_gcs(init_checkpoint)\n",
+        "\n",
+        "# Uploads config file according to model_id and streaming options.\n",
+        "config_file = f\"{model_id}_stream\" if is_stream else model_id\n",
+        "config_file = f\"https://raw.githubusercontent.com/tensorflow/models/master/official/projects/movinet/configs/yaml/movinet_{config_file}_gpu.yaml\"\n",
+        "config_file = upload_config_to_gcs(config_file)\n",
+        "\n",
+        "# The parameters here are mainly for demonstration purpose. Please update them\n",
+        "# for better performance.\n",
+        "trainer_args = {\n",
+        "    \"experiment\": \"movinet_kinetics600\",\n",
+        "    \"config_file\": config_file,\n",
+        "    \"input_train_data_path\": input_train_data_path,\n",
+        "    \"input_validation_data_path\": input_validation_data_path,\n",
+        "    \"init_checkpoint\": init_checkpoint,\n",
+        "    \"model_dir\": model_dir,\n",
+        "    \"num_classes\": num_classes,\n",
+        "    \"global_batch_size\": 4,\n",
+        "    \"prefetch_buffer_size\": 8,\n",
+        "    \"shuffle_buffer_size\": 32,\n",
+        "    \"train_steps\": 2000,\n",
+        "}\n",
+        "\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": TRAIN_MACHINE_TYPE,\n",
+        "            \"accelerator_type\": TRAIN_ACCELERATOR_TYPE,\n",
+        "            # Each training job uses TRAIN_NUM_GPU GPUs.\n",
+        "            \"accelerator_count\": TRAIN_NUM_GPU,\n",
+        "        },\n",
+        "        \"replica_count\": 1,\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": TRAIN_CONTAINER_URI,\n",
+        "            \"args\": [\n",
+        "                \"--mode=train_and_eval\",\n",
+        "                \"--params_override=runtime.num_gpus=%d\" % TRAIN_NUM_GPU,\n",
+        "            ]\n",
+        "            + [\"--{}={}\".format(k, v) for k, v in trainer_args.items()],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "metric_spec = {\"model_performance\": \"maximize\"}\n",
+        "\n",
+        "# These learning rates might not be optimal for your selected model type; To\n",
+        "# tune learning rates, try hpt.DoubleParameterSpec with more trials.\n",
+        "LEARNING_RATES = [1e-3, 3e-3]\n",
+        "MAX_TRIAL_COUNT = len(LEARNING_RATES)\n",
+        "parameter_spec = {\n",
+        "    \"learning_rate\": hpt.DiscreteParameterSpec(values=LEARNING_RATES, scale=\"linear\"),\n",
+        "}\n",
+        "\n",
+        "print(worker_pool_specs, metric_spec, parameter_spec)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "HwcCjwlBTQIz"
+      },
+      "source": [
+        "#### Run the hyperparameter tuning job\n",
+        "* `max_trial_count`: Sets an upper bound on the number of trials the service will run. The recommended practice is to start with a smaller number of trials and get a sense of how impactful your chosen hyperparameters are before scaling up.\n",
+        "\n",
+        "* `parallel_trial_count`:  If you use parallel trials, the service provisions multiple training processing clusters. The worker pool spec that you specify when creating the job is used for each individual training cluster.  Increasing the number of parallel trials reduces the amount of time the hyperparameter tuning job takes to run; however, it can reduce the effectiveness of the job overall. This is because the default tuning strategy uses results of previous trials to inform the assignment of values in subsequent trials.\n",
+        "\n",
+        "* `search_algorithm`: The available search algorithms are grid, random, or default (None). The default option applies Bayesian optimization to search the space of possible hyperparameter values and is the recommended algorithm.\n",
+        "\n",
+        "Click on the generated link in the output to see your run in the Cloud Console."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "aec22792ee84"
+      },
+      "outputs": [],
+      "source": [
+        "train_custom_job = aiplatform.CustomJob(\n",
+        "    display_name=train_job_name,\n",
+        "    project=PROJECT_ID,\n",
+        "    worker_pool_specs=worker_pool_specs,\n",
+        "    staging_bucket=STAGING_BUCKET,\n",
+        ")\n",
+        "\n",
+        "train_hpt_job = aiplatform.HyperparameterTuningJob(\n",
+        "    display_name=train_job_name,\n",
+        "    custom_job=train_custom_job,\n",
+        "    metric_spec=metric_spec,\n",
+        "    parameter_spec=parameter_spec,\n",
+        "    max_trial_count=MAX_TRIAL_COUNT,\n",
+        "    parallel_trial_count=MAX_TRIAL_COUNT,\n",
+        "    project=PROJECT_ID,\n",
+        "    search_algorithm=None,\n",
+        ")\n",
+        "\n",
+        "train_hpt_job.run()\n",
+        "\n",
+        "print(\"model_dir is:\", model_dir)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "vugUfJEC2HrK"
+      },
+      "source": [
+        "### Export model in Tensorflow SavedModel format"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "09Rz1AYspK19"
+      },
+      "outputs": [],
+      "source": [
+        "# This job will export models from TF checkpoints to TF saved model format.\n",
+        "# model_dir is from the section above.\n",
+        "best_trial_dir, best_trial_evaluation_results = get_best_trial(\n",
+        "    model_dir, MAX_TRIAL_COUNT, EVALUATION_METRIC\n",
+        ")\n",
+        "best_checkpoint_path = find_checkpoint_in_dir(f\"{best_trial_dir}/best_ckpt/\")\n",
+        "print(\"best_trial_dir: \", best_trial_dir)\n",
+        "print(\"best_trial_evaluation_results: \", best_trial_evaluation_results)\n",
+        "print(\"best_checkpoint: \", best_checkpoint_path)\n",
+        "\n",
+        "container_args = {\n",
+        "    \"export_path\": f\"{model_dir}/best_model\",\n",
+        "    \"model_id\": model_id,\n",
+        "    \"num_classes\": num_classes,\n",
+        "    \"causal\": is_stream,\n",
+        "    \"checkpoint_path\": best_checkpoint_path,\n",
+        "    \"assert_checkpoint_objects_matched\": False,\n",
+        "}\n",
+        "\n",
+        "if is_stream:\n",
+        "    container_args.update(\n",
+        "        {\n",
+        "            \"conv_type\": \"2plus1d\",\n",
+        "            \"se_type\": \"2plus3d\",\n",
+        "            \"activation\": \"hard_swish\",\n",
+        "            \"gating_activation\": \"hard_sigmoid\",\n",
+        "            \"use_positional_encoding\": model_id in {\"a3\", \"a4\", \"a5\"},\n",
+        "        }\n",
+        "    )\n",
+        "\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": EXPORT_MACHINE_TYPE,\n",
+        "        },\n",
+        "        \"replica_count\": 1,\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": EXPORT_CONTAINER_URI,\n",
+        "            \"args\": [\"--{}={}\".format(k, v) for k, v in container_args.items()],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "model_export_job_name = get_job_name_with_datetime(EXPORT_JOB_PREFIX + \"_\" + OBJECTIVE)\n",
+        "model_export_custom_job = aiplatform.CustomJob(\n",
+        "    display_name=model_export_job_name,\n",
+        "    project=PROJECT_ID,\n",
+        "    worker_pool_specs=worker_pool_specs,\n",
+        "    staging_bucket=STAGING_BUCKET,\n",
+        ")\n",
+        "\n",
+        "model_export_custom_job.run()\n",
+        "\n",
+        "print(\"best model is saved to: \", container_args[\"export_path\"])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "g0BGaofgsMsy"
+      },
+      "source": [
+        "## Test trained models\n",
+        "This section shows the way to test with trained models.\n",
+        "1. Upload and deploy models to the [Vertex AI Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction)\n",
+        "2. Run batch predictions\n",
+        "\n",
+        "**Note:** The prediction container only works with the base model. If you trained a streaming model, download the model from the exported path and refer to the [MoViNet official guide](https://github.com/tensorflow/models/blob/master/official/projects/movinet/movinet_streaming_model_training_and_inference.ipynb) for running predictions locally."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "gdlca3BOypXU"
+      },
+      "source": [
+        "### Upload model to Vertex AI Model Registry"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NYuQowyZEtxK"
+      },
+      "outputs": [],
+      "source": [
+        "upload_job_name = get_job_name_with_datetime(f\"{UPLOAD_JOB_PREFIX}_{model_name}\")\n",
+        "\n",
+        "serving_env = {\n",
+        "    \"MODEL_PATH\": container_args[\"export_path\"],\n",
+        "    \"BATCH_SIZE\": 1,  # Select a larger batch size to accelerate GPU prediction.\n",
+        "    \"NUM_FRAMES\": 32,\n",
+        "    \"FPS\": output_fps,\n",
+        "    \"OVERLAP_FRAMES\": 24,\n",
+        "    \"OBJECTIVE\": OBJECTIVE,\n",
+        "}\n",
+        "\n",
+        "model = aiplatform.Model.upload(\n",
+        "    display_name=model_name,\n",
+        "    serving_container_image_uri=PREDICTION_CONTAINER_URI,\n",
+        "    serving_container_ports=[PREDICTION_PORT],\n",
+        "    serving_container_predict_route=\"/predict\",\n",
+        "    serving_container_health_route=\"/ping\",\n",
+        "    serving_container_environment_variables=serving_env,\n",
+        ")\n",
+        "\n",
+        "model.wait()\n",
+        "\n",
+        "print(\"The uploaded model name is: \", model_name)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9SZsKGeS3x6S"
+      },
+      "source": [
+        "### Run batch predictions\n",
+        "\n",
+        "We will now run batch predictions with the trained MoViNet clip classification model with [Vertex AI Batch Prediction](https://cloud.google.com/vertex-ai/docs/predictions/get-batch-predictions).\n",
+        "\n",
+        "Please prepare an input JSONL file where each line follows [this format](https://cloud.google.com/vertex-ai/docs/video-data/classification/get-predictions?hl=en#input_data_requirements) and store it in a Cloud Storage bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "vbIW9me1F2RY"
+      },
+      "outputs": [],
+      "source": [
+        "# Path to the prediction input JSONL file.\n",
+        "test_jsonl_path = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "predict_job_name = get_job_name_with_datetime(f\"{PREDICTION_JOB_PREFIX}_{model_name}\")\n",
+        "predict_destination_prefix = os.path.join(STAGING_BUCKET, predict_job_name)\n",
+        "\n",
+        "batch_prediction_job = model.batch_predict(\n",
+        "    job_display_name=predict_job_name,\n",
+        "    gcs_source=test_jsonl_path,\n",
+        "    gcs_destination_prefix=predict_destination_prefix,\n",
+        "    machine_type=PREDICTION_MACHINE_TYPE,\n",
+        "    accelerator_count=PREDICTION_ACCELERATOR_COUNT,\n",
+        "    accelerator_type=PREDICTION_ACCELERATOR_TYPE,\n",
+        "    max_replica_count=1,\n",
+        ")\n",
+        "\n",
+        "batch_prediction_job.wait()\n",
+        "\n",
+        "print(batch_prediction_job.display_name)\n",
+        "print(batch_prediction_job.resource_name)\n",
+        "print(batch_prediction_job.state)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ik-XPjfx9OCE"
+      },
+      "source": [
+        "You can then read the prediction response JSONL files in the output directory:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tdkW9e5B9OU1"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "\n",
+        "# The label map file was generated from the section above (`Prepare input data for training`).\n",
+        "for file in tf.io.gfile.glob(os.path.join(predict_destination_prefix, \"*/*\")):\n",
+        "    with tf.io.gfile.GFile(file, \"r\") as f:\n",
+        "        for line in f:\n",
+        "            print_response_instance(line, label_map)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kkH2nrpdp4sp"
+      },
+      "source": [
+        "## Clean up"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Ax6vQVZhp9pR"
+      },
+      "outputs": [],
+      "source": [
+        "# Delete the trained model.\n",
+        "model.delete()\n",
+        "# Delete custom and hpt jobs.\n",
+        "if data_converter_custom_job.list(filter=f'display_name=\"{data_converter_job_name}\"'):\n",
+        "    data_converter_custom_job.delete()\n",
+        "if train_hpt_job.list(filter=f'display_name=\"{train_job_name}\"'):\n",
+        "    train_hpt_job.delete()\n",
+        "if model_export_custom_job.list(filter=f'display_name=\"{model_export_job_name}\"'):\n",
+        "    model_export_custom_job.delete()\n",
+        "if batch_prediction_job.list(filter=f'display_name=\"{predict_job_name}\"'):\n",
+        "    batch_prediction_job.delete()"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_movinet_clip_classification.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_pipeline_templates_t5x.ipynb

@@ -319,10 +319,8 @@
       },
       "outputs": [],
       "source": [
-        "IS_COLAB = False\n",
         "# from google.colab import auth\n",
-        "# auth.authenticate_user()\n",
-        "# IS_COLAB=True"
+        "# auth.authenticate_user()"
       ]
     },
     {
@@ -407,6 +405,9 @@
       },
       "outputs": [],
       "source": [
+        "import sys\n",
+        "\n",
+        "IS_COLAB = \"google.colab\" in sys.modules\n",
         "if (\n",
         "    SERVICE_ACCOUNT == \"\"\n",
         "    or SERVICE_ACCOUNT is None\n",
@@ -469,6 +470,7 @@
       "outputs": [],
       "source": [
         "import json\n",
+        "import os\n",
         "\n",
         "import google.cloud.aiplatform as aiplatform\n",
         "from kfp.registry import RegistryClient"
@@ -496,6 +498,33 @@
         "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2d242773d707"
+      },
+      "source": [
+        "### Enable Artifact Registry API\n",
+        "You must enable the Artifact Registry API service for your project.\n",
+        "\n",
+        "<a href=\"https://cloud.google.com/artifact-registry/docs/enable-service\">Learn more about Enabling service</a>."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "162b5e8883c2"
+      },
+      "outputs": [],
+      "source": [
+        "! gcloud services enable artifactregistry.googleapis.com\n",
+        "\n",
+        "if os.getenv(\"IS_TESTING\"):\n",
+        "    ! sudo apt-get update --yes && sudo apt-get --only-upgrade --yes install google-cloud-sdk-cloud-run-proxy google-cloud-sdk-harbourbridge google-cloud-sdk-cbt google-cloud-sdk-gke-gcloud-auth-plugin google-cloud-sdk-kpt google-cloud-sdk-local-extract google-cloud-sdk-minikube google-cloud-sdk-app-engine-java google-cloud-sdk-app-engine-go google-cloud-sdk-app-engine-python google-cloud-sdk-spanner-emulator google-cloud-sdk-bigtable-emulator google-cloud-sdk-nomos google-cloud-sdk-package-go-module google-cloud-sdk-firestore-emulator kubectl google-cloud-sdk-datastore-emulator google-cloud-sdk-app-engine-python-extras google-cloud-sdk-cloud-build-local google-cloud-sdk-kubectl-oidc google-cloud-sdk-anthos-auth google-cloud-sdk-app-engine-grpc google-cloud-sdk-pubsub-emulator google-cloud-sdk-datalab google-cloud-sdk-skaffold google-cloud-sdk google-cloud-sdk-terraform-tools google-cloud-sdk-config-connector\n",
+        "    ! gcloud components update --quiet"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -504,9 +533,7 @@
       "source": [
         "## Create repo in Artifact Registry\n",
         "\n",
-        "First, you create your own (user-defined) repository in the `Artifact Registry`. You use this repository to upload and retreive your pipeline templates.\n",
-        "\n",
-        "The name of your repo is `quickstart-kfp-repo`"
+        "First, you create your own (user-defined) repository in the `Artifact Registry`. You use this repository to upload and retreive your pipeline templates."
       ]
     },
     {
@@ -517,7 +544,7 @@
       },
       "outputs": [],
       "source": [
-        "REPO_NAME = \"quickstart-kfp-repo\"\n",
+        "REPO_NAME = \"my-docker-repo-unique\"\n",
         "\n",
         "! gcloud artifacts repositories create {REPO_NAME} --location={REGION} --repository-format=KFP"
       ]
@@ -822,7 +849,7 @@
       },
       "outputs": [],
       "source": [
-        "delete_bucket = True\n",
+        "delete_bucket = False\n",
         "\n",
         "endpoint.undeploy_all()\n",
         "endpoint.delete()\n",
@@ -831,7 +858,9 @@
         "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
         "    ! gsutil rm -r $BUCKET_URI\n",
         "\n",
-        "! rm -rf custom custom.tar.gz"
+        "! rm -rf custom custom.tar.gz\n",
+        "\n",
+        "! gcloud artifacts repositories delete $REPO_NAME --project {PROJECT_ID} --location {REGION} --quiet"
       ]
     }
   ],

notebooks/community/model_garden/model_garden_proprietary_image_classification.ipynb

@@ -24,7 +24,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "TirJ-SGQseby"
@@ -55,7 +54,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "dwGLvtIeECLK"
@@ -67,7 +65,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "tvgnzT1CKxrO"
@@ -107,7 +104,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "KEukV6uRk_S3"
@@ -138,7 +134,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "BF1j6f9HApxa"
@@ -196,7 +191,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "n6IFz75WGCam"
@@ -229,16 +223,13 @@
         "# Evaluation constants.\n",
         "EVALUATION_METRIC = \"accuracy\"\n",
         "\n",
-        "# Prediction constants.\n",
-        "# The example in this notebook uses optimized tensorflow runtime dockers.\n",
-        "# You can adjust accelerator types and machine types to get faster predictions.\n",
-        "PREDICTION_ACCELERATOR_TYPE = \"NVIDIA_TESLA_T4\"\n",
-        "PREDICTION_MACHINE_TYPE = \"n1-standard-4\"\n",
+        "# Prediction constant.\n",
+        "# Model does not support dedicated deployment resources.\n",
+        "# An n1-standard-4 machine with 1 P100 GPU will be used.\n",
         "DEPLOY_JOB_PREFIX = \"deploy\""
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "ZZFPe_GezXg8"
@@ -301,7 +292,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "Q149N3V6Uynm"
@@ -330,7 +320,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "8yfBZ1_8VZvq"
@@ -363,7 +352,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "RB_xY9ipr7ZU"
@@ -389,7 +377,6 @@
         "- `model_type`: The type of model for deployment.\n",
         "  - `EFFICIENTNET`: A model that is available in Vertex Model Garden image classification training with customizable hyperparameters. Best tailored to be used within Google Cloud, and cannot be exported externally.\n",
         "  - `MAXVIT`: A model that is available in Vertex Model Garden image classification training with customizable hyperparameters. Best tailored to be used within Google Cloud, and cannot be exported externally.\n",
-        "  - `COCA`: A model that is available in Vertex Model Garden image classification training with customizable hyperparameters. Best tailored to be used within Google Cloud, and cannot be exported externally.\n",
         "- `checkpoint_name`: Optional. The field is reserved for Model Garden model training, based on the provided pre-trained model checkpoint.\n",
         "- `trainer_config`: Optional. The field is usually used together with the Model Garden model training when passing the customized configs for the trainer.\n",
         "\n",
@@ -457,7 +444,9 @@
         "METRIC_SPEC_VALUE = \"maximize\"\n",
         "SEARCH_ALGORITHM = \"random\"\n",
         "MEASUREMENT_SELECTION = \"best\"\n",
-        "MODEL_TYPE = \"COCA\"  # @param {type:\"string\"} one of the values [\"COCA\", \"MAXVIT\", \"EFFICIENTNET\"]\n",
+        "MODEL_TYPE = (\n",
+        "    \"MAXVIT\"  # @param {type:\"string\"} one of the values [\"MAXVIT\", \"EFFICIENTNET\"]\n",
+        ")\n",
         "\n",
         "job = aiplatform.AutoMLImageTrainingJob(\n",
         "    display_name=get_job_name_with_datetime(TRAINING_JOB_PREFIX),\n",
@@ -478,7 +467,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "HwcCjwlBTQIz"
@@ -523,7 +511,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "g0BGaofgsMsy"
@@ -550,9 +537,9 @@
         "\n",
         "endpoint = model.deploy(\n",
         "    deployed_model_display_name=deploy_model_name,\n",
-        "    machine_type=PREDICTION_MACHINE_TYPE,\n",
+        "    machine_type=\"\",\n",
         "    traffic_split={\"0\": 100},\n",
-        "    accelerator_type=PREDICTION_ACCELERATOR_TYPE,\n",
+        "    accelerator_type=None,\n",
         "    accelerator_count=1,\n",
         "    min_replica_count=1,\n",
         "    max_replica_count=1,\n",
@@ -589,7 +576,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "kkH2nrpdp4sp"

notebooks/community/model_garden/model_garden_proprietary_image_object_detection.ipynb

@@ -24,7 +24,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "TirJ-SGQseby"
@@ -55,7 +54,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "dwGLvtIeECLK"
@@ -67,7 +65,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "tvgnzT1CKxrO"
@@ -107,7 +104,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "KEukV6uRk_S3"
@@ -138,7 +134,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "BF1j6f9HApxa"
@@ -181,11 +176,11 @@
         "# You can choose a region from https://cloud.google.com/about/locations.\n",
         "# Only regions prefixed by \"us\", \"europe\", or \"asia\" are supported.\n",
         "REGION = \"us-central1\"  # @param {type:\"string\"}\n",
-        "REGION_PREFIX = REGION.split('-')[0]\n",
+        "REGION_PREFIX = REGION.split(\"-\")[0]\n",
         "assert REGION_PREFIX in (\n",
         "    \"us\",\n",
         "    \"europe\",\n",
-        "    \"asia\"\n",
+        "    \"asia\",\n",
         "), f'{REGION} is not supported. It must be prefixed by \"us\", \"europe\", or \"asia\".'\n",
         "\n",
         "! gcloud config set project $PROJECT_ID\n",
@@ -196,7 +191,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "n6IFz75WGCam"
@@ -227,14 +221,11 @@
         "EVALUATION_METRIC = \"AP50\"\n",
         "\n",
         "# Prediction constants.\n",
-        "# You can adjust accelerator types and machine types to get faster predictions.\n",
-        "PREDICTION_ACCELERATOR_TYPE = \"NVIDIA_TESLA_T4\"\n",
-        "PREDICTION_MACHINE_TYPE = \"n1-standard-4\"\n",
+        "# An n1-standard-4 machine with 1 P100 GPU will be used.\n",
         "DEPLOY_JOB_PREFIX = \"deploy\""
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "ZZFPe_GezXg8"
@@ -357,7 +348,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "nZLVI9TtUuif"
@@ -386,7 +376,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "SZEdBfNZUxQn"
@@ -419,7 +408,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "RB_xY9ipr7ZU"
@@ -550,7 +538,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "HwcCjwlBTQIz"
@@ -595,7 +582,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "g0BGaofgsMsy"
@@ -622,9 +608,9 @@
         "\n",
         "endpoint = model.deploy(\n",
         "    deployed_model_display_name=deploy_model_name,\n",
-        "    machine_type=PREDICTION_MACHINE_TYPE,\n",
+        "    machine_type=\"\",\n",
         "    traffic_split={\"0\": 100},\n",
-        "    accelerator_type=PREDICTION_ACCELERATOR_TYPE,\n",
+        "    accelerator_type=None,\n",
         "    accelerator_count=1,\n",
         "    min_replica_count=1,\n",
         "    max_replica_count=1,\n",
@@ -661,7 +647,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "frcGP5HFX1XN"

notebooks/community/model_garden/model_garden_pytorch_bart_large_cnn.ipynb

@@ -208,7 +208,7 @@
       "source": [
         "# The pre-built serving docker image.\n",
         "# The model artifacts are embedded within the container, except for model weights which will be downloaded during deployment.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_blip2.ipynb

@@ -205,7 +205,7 @@
       "source": [
         "# The pre-built serving docker image.\n",
         "# The model artifacts are embedded within the container, except for model weights which will be downloaded during deployment.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_blip_image_captioning.ipynb

@@ -204,7 +204,7 @@
       "outputs": [],
       "source": [
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_blip_vqa.ipynb

@@ -204,7 +204,7 @@
       "outputs": [],
       "source": [
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_clip.ipynb

@@ -204,7 +204,7 @@
       "outputs": [],
       "source": [
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_controlnet.ipynb

@@ -259,10 +259,10 @@
       "outputs": [],
       "source": [
         "# The pre-built training docker image. It contains training scripts and models.\n",
-        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-diffusers-train:latest\"\n",
+        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-train:latest\"\n",
         "\n",
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_detectron2.ipynb

@@ -44,7 +44,7 @@
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
-        "  <td>                                                                                               <td>\n",
+        "  <td>\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_pytorch_detectron2.ipynb\">\n",
         "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
         "Open in Vertex AI Workbench\n",

notebooks/community/model_garden/model_garden_pytorch_dolly_v2.ipynb

@@ -208,7 +208,7 @@
       "source": [
         "# The pre-built serving docker image.\n",
         "# The model artifacts are embedded within the container, except for model weights which will be downloaded during deployment.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-dolly-v2-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-dolly-v2-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_falcon_instruct_peft.ipynb

@@ -0,0 +1,655 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7d9bbf86da5e"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2023 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "99c1c3fc2ca5"
+      },
+      "source": [
+        "# Vertex AI Model Garden - Falcon Instruct (PEFT)\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_falcon_instruct_peft.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_falcon_instruct_peft.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_pytorch_falcon_instruct_peft.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "Open in Vertex AI Workbench\n",
+        "    </a> (A Python-3 CPU notebook is recommended)\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3de7470326a2"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates deploying prebuilt Falcon Instruct models, and also finetuning and deploying Falcon Instruct models with performance efficient finetuning libraries ([PEFT](https://github.com/huggingface/peft)) in Vertex AI.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Deploy prebuilt Falcon Instruct models\n",
+        "- Finetune and deploy Falcon Instruct models with PEFT, supporting\n",
+        "\n",
+        "| Models | LoRA |\n",
+        "| :- | :- |\n",
+        "| [tiiuae/falcon-7b-instruct](https://huggingface.co/tiiuae/falcon-7b-instruct) | Y |\n",
+        "| [tiiuae/falcon-40b-instruct](https://huggingface.co/tiiuae/falcon-40b-instruct) | Y |\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "264c07757582"
+      },
+      "source": [
+        "## Before you begin\n",
+        "\n",
+        "**NOTE**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ioensNKM8ned"
+      },
+      "source": [
+        "### Colab only\n",
+        "Run the following commands for Colab and skip this section if you are using Workbench."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "2707b02ef5df"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
+        "    from google.colab import auth as google_auth\n",
+        "\n",
+        "    google_auth.authenticate_user()\n",
+        "    # Install gdown for downloading example training images.\n",
+        "    ! pip3 install gdown\n",
+        "\n",
+        "    # Restart the notebook kernel after installs.\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bb7adab99e41"
+      },
+      "source": [
+        "### Setup Google Cloud project\n",
+        "\n",
+        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "\n",
+        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
+        "\n",
+        "1. [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs.\n",
+        "\n",
+        "1. [Create a service account](https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console) with `Vertex AI User` and `Storage Object Admin` roles for deploying fine tuned model to Vertex AI endpoint."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6c460088b873"
+      },
+      "source": [
+        "Fill following variables for experiments environment:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "855d6b96f291"
+      },
+      "outputs": [],
+      "source": [
+        "# Cloud project id.\n",
+        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# The region you want to launch jobs in.\n",
+        "REGION = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# The Cloud Storage bucket for storing experiments output.\n",
+        "BUCKET_URI = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "\n",
+        "import os\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "EXPERIMENT_BUCKET = os.path.join(BUCKET_URI, \"peft\")\n",
+        "DATA_BUCKET = os.path.join(EXPERIMENT_BUCKET, \"data\")\n",
+        "MODEL_BUCKET = os.path.join(EXPERIMENT_BUCKET, \"model\")\n",
+        "\n",
+        "# The service account looks like:\n",
+        "# '@.iam.gserviceaccount.com'\n",
+        "# Please go to https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console\n",
+        "# and create service account with `Vertex AI User` and `Storage Object Admin` roles.\n",
+        "# The service account for deploying fine tuned model.\n",
+        "SERVICE_ACCOUNT = \"\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e828eb320337"
+      },
+      "source": [
+        "### Initialize Vertex AI API"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "12cd25839741"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2cc825514deb"
+      },
+      "source": [
+        "### Define constants"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "b42bd4fa2b2d"
+      },
+      "outputs": [],
+      "source": [
+        "# The pre-built training and serving docker images.\n",
+        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-peft-train\"\n",
+        "PREDICTION_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-peft-serve\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0c250872074f"
+      },
+      "source": [
+        "### Define common functions"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "354da31189dc"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "from datetime import datetime\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "\n",
+        "def get_job_name_with_datetime(prefix: 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,\n",
+        "    base_model_id,\n",
+        "    finetuned_lora_model_path,\n",
+        "    service_account,\n",
+        "    task,\n",
+        "    machine_type=\"n1-standard-8\",\n",
+        "    accelerator_type=\"NVIDIA_TESLA_V100\",\n",
+        "):\n",
+        "    \"\"\"Deploys trained models into Vertex AI.\"\"\"\n",
+        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
+        "    serving_env = {\n",
+        "        \"BASE_MODEL_ID\": base_model_id,\n",
+        "        \"TASK\": task,\n",
+        "    }\n",
+        "    if finetuned_lora_model_path:\n",
+        "        serving_env[\"FINETUNED_LORA_MODEL_PATH\"] = finetuned_lora_model_path\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=PREDICTION_DOCKER_URI,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_predict_route=\"/predictions/peft_serving\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        serving_container_environment_variables=serving_env,\n",
+        "    )\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=1,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=service_account,\n",
+        "    )\n",
+        "    return model, endpoint"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8neJc8CnDDpu"
+      },
+      "source": [
+        "## Deploy prebuilt Falcon Instruct models\n",
+        "\n",
+        "This section deploys prebuilt Falcon Instruct models on the Endpoint. The model deployment step will take ~15 minutes to complete.\n",
+        "\n",
+        "The peak GPU memory usages for [tiiuae/falcon-7b-instruct](https://huggingface.co/tiiuae/falcon-7b-instruct), and [tiiuae/falcon-40b-instruct](https://huggingface.co/tiiuae/falcon-40b-instruct) are ~15.5G and ~38.2G separately with the default settings. We use V100 in deployments as an example. Please use A100 (40G) or A100 (80G) to get better inferences."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2MjaORIIFDVu"
+      },
+      "source": [
+        "Set the prebuilt model id."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "E8OiHHNNE_wj"
+      },
+      "outputs": [],
+      "source": [
+        "prebuilt_model_id = \"tiiuae/falcon-7b-instruct\"  # @param [\"tiiuae/falcon-7b-instruct\", \"tiiuae/falcon-40b-instruct\"]"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dHFW7yvjaVFV"
+      },
+      "source": [
+        "We use the PEFT serving images to deploy prebuilt Falcon Instruct models, by setting finetuning LoRA model paths as empty."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Uak1pyEeExYM"
+      },
+      "outputs": [],
+      "source": [
+        "# Sets V100 to deploy tiiuae/falcon-7b-instruct as an example.\n",
+        "machine_type = \"n1-standard-8\"\n",
+        "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
+        "\n",
+        "# Sets A100 (40G) to deploy tiiuae/falcon-7b-instruct or tiiuae/falcon-40b-instruct.\n",
+        "# machine_type = \"a2-highgpu-1g\"\n",
+        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+        "\n",
+        "# Sets A100 (80G) to deploy falcon-40b-instruct models for faster inferences.\n",
+        "# machine_type = \"a2-ultragpu-1g\"\n",
+        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+        "\n",
+        "model_without_peft, endpoint_without_peft = deploy_model(\n",
+        "    model_name=get_job_name_with_datetime(prefix=\"falcon-instruct-serve\"),\n",
+        "    base_model_id=prebuilt_model_id,\n",
+        "    finetuned_lora_model_path=\"\",  # This will avoid override finetuning models.\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    task=\"instruct-lora\",\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        ")\n",
+        "print(\"endpoint_name:\", endpoint_without_peft.name)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "sGKIjgmDFRW2"
+      },
+      "source": [
+        "NOTE: The prebuilt model weights will be downloaded on the fly from the orginal location after the deployment succeeds. Thus additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint.\n",
+        "\n",
+        "Once deployment succeeds, you can send requests to the endpoint with text prompts.\n",
+        "\n",
+        "Example:\n",
+        "\n",
+        "```\n",
+        "Human: What is a car?\n",
+        "Assistant:  A car, or a motor car, is a road-connected human-transportation system used to move people or goods from one place to another. The term also encompasses a wide range of vehicles, including motorboats, trains, and aircrafts. Cars typically have four wheels, a cabin for passengers, and an engine or motor. They have been around since the early 19th century and are now one of the most popular forms of transportation, used for daily commuting, shopping, and other purposes.\n",
+        "```"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "rDHsCOqvFYBi"
+      },
+      "outputs": [],
+      "source": [
+        "# # Loads an existing endpoint as below.\n",
+        "# endpoint_name = endpoint_without_peft.name\n",
+        "# aip_endpoint_name = (\n",
+        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
+        "# )\n",
+        "# endpoint_without_peft = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "instances = [\n",
+        "    {\"prompt\": \"What is a car?\"},\n",
+        "]\n",
+        "response = endpoint_without_peft.predict(instances=instances)\n",
+        "\n",
+        "for prediction in response.predictions[0]:\n",
+        "    print(prediction[\"generated_text\"])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e70e3519ff8b"
+      },
+      "source": [
+        "## Finetune and deploy Falcon Instruct models with PEFT\n",
+        "\n",
+        "This section demonstrates how to finetune and dpeloy Falcon Instruct models with PEFT LoRA."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5qCrm_kJH5cz"
+      },
+      "source": [
+        "Set the base model id."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "N3UBLiYrM3sU"
+      },
+      "outputs": [],
+      "source": [
+        "base_model_id = \"tiiuae/falcon-7b-instruct\"  # @param [\"tiiuae/falcon-7b-instruct\", \"tiiuae/falcon-40b-instruct\"]"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "iWGwJHqI7LMs"
+      },
+      "source": [
+        "### Finetune"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "KKEYoRfiHDVv"
+      },
+      "source": [
+        "Use the Vertex AI SDK to create and run the custom training jobs with Vertex AI Model Garden training images.\n",
+        "\n",
+        "This example uses the dataset [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco).\n",
+        "\n",
+        "The peak GPU memory usages are ~11G and ~34G for finetuning LoRA models for [tiiuae/falcon-7b-instruct](https://huggingface.co/tiiuae/falcon-7b-instruct), and [tiiuae/falcon-40b-instruct](https://huggingface.co/tiiuae/falcon-40b-instruct) separately with default training parameters and the example dataset. Falcon-7b-instruct can be finetuned on 1 P100/V100, and falcon-40b-instruct can be finetuned on 1 A100 (40G)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "65467b361315"
+      },
+      "outputs": [],
+      "source": [
+        "dataset_name = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Worker pool spec.\n",
+        "# Uses V100 to finetune falcon-70b-instruct.\n",
+        "machine_type = \"n1-standard-8\"\n",
+        "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
+        "# Uses A100 to finetune falcon-40b-instruct.\n",
+        "# machine_type = \"a2-highgpu-1g\"\n",
+        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+        "replica_count = 1\n",
+        "accelerator_count = 1\n",
+        "\n",
+        "# Setup training job.\n",
+        "job_name = get_job_name_with_datetime(\"falcon-instruct-lora-train\")\n",
+        "train_job = aiplatform.CustomContainerTrainingJob(\n",
+        "    display_name=job_name,\n",
+        "    container_uri=TRAIN_DOCKER_URI,\n",
+        ")\n",
+        "output_dir = os.path.join(MODEL_BUCKET, job_name)\n",
+        "output_dir_gcsfuse = output_dir.replace(\"gs://\", \"/gcs/\")\n",
+        "\n",
+        "# Pass training arguments and launch job.\n",
+        "max_steps = 10\n",
+        "train_job.run(\n",
+        "    args=[\n",
+        "        \"--task=instruct-lora\",\n",
+        "        f\"--pretrained_model_id={base_model_id}\",\n",
+        "        f\"--dataset_name={dataset_name}\",\n",
+        "        f\"--output_dir={output_dir_gcsfuse}\",\n",
+        "        \"--lora_rank=64\",\n",
+        "        \"--lora_alpha=16\",\n",
+        "        \"--lora_dropout=0.1\",\n",
+        "        \"--warmup_ratio=0.03\",\n",
+        "        f\"--max_steps={max_steps}\",\n",
+        "        \"--max_seq_length=512\",\n",
+        "        \"--learning_rate=2e-4\",\n",
+        "    ],\n",
+        "    replica_count=replica_count,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    boot_disk_size_gb=500,\n",
+        ")\n",
+        "\n",
+        "print(\"Trained models were saved in: \", output_dir)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "jqmCtkGnhDmp"
+      },
+      "source": [
+        "### Deploy\n",
+        "This section uploads the model to Model Registry and deploys it on the Endpoint.\n",
+        "\n",
+        "The model deployment step will take ~15 minutes to complete.\n",
+        "\n",
+        "The peak GPU memory usages for [tiiuae/falcon-7b-instruct](https://huggingface.co/tiiuae/falcon-7b-instruct), and [tiiuae/falcon-40b-instruct](https://huggingface.co/tiiuae/falcon-40b-instruct) with LoRA weights are ~15.5G and ~38.2G separately with the default settings. We use V100 in deployments as an example. Please use A100 (40G) or A100 (80G) to get better inferences."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bf55e38815dc"
+      },
+      "outputs": [],
+      "source": [
+        "# Sets V100 to deploy tiiuae/falcon-7b-instruct as an example.\n",
+        "machine_type = \"n1-standard-8\"\n",
+        "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
+        "\n",
+        "# Sets A100 (40G) to deploy tiiuae/falcon-7b-instruct or tiiuae/falcon-40b-instruct.\n",
+        "# machine_type = \"a2-highgpu-1g\"\n",
+        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+        "\n",
+        "# Sets A100 (80G) to deploy falcon-40b-instruct models for faster inferences.\n",
+        "# machine_type = \"a2-ultragpu-1g\"\n",
+        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+        "\n",
+        "model_with_peft, endpoint_with_peft = deploy_model(\n",
+        "    model_name=get_job_name_with_datetime(prefix=\"falcon-instruct-peft-serve\"),\n",
+        "    base_model_id=base_model_id,\n",
+        "    finetuned_lora_model_path=os.path.join(output_dir, \"checkpoint-\" + str(max_steps)),\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    task=\"instruct-lora\",\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        ")\n",
+        "print(\"endpoint_name:\", endpoint_with_peft.name)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "80b3fd2ace09"
+      },
+      "source": [
+        "NOTE: After the deployment succeeds, the base model weights will be downloaded one the fly from the original location and LoRA model weights will be downloaded from the GCS bucket used in training above. Thus additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint.\n",
+        "\n",
+        "Once deployment succeeds, you can send requests to the endpoint with text prompts.\n",
+        "\n",
+        "Example:\n",
+        "\n",
+        "```\n",
+        "Human: What is a car?\n",
+        "Assistant:  A car, or a motor car, is a road-connected human-transportation system used to move people or goods from one place to another. The term also encompasses a wide range of vehicles, including motorboats, trains, and aircrafts. Cars typically have four wheels, a cabin for passengers, and an engine or motor. They have been around since the early 19th century and are now one of the most popular forms of transportation, used for daily commuting, shopping, and other purposes.\n",
+        "```"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "4ab04da3ec9a"
+      },
+      "outputs": [],
+      "source": [
+        "# # Loads an existing endpoint as below.\n",
+        "# endpoint_name = endpoint_with_peft.name\n",
+        "# aip_endpoint_name = (\n",
+        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
+        "# )\n",
+        "# endpoint_with_peft = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "instances = [\n",
+        "    {\"prompt\": \"What is a car?\"},\n",
+        "]\n",
+        "response = endpoint_with_peft.predict(instances=instances)\n",
+        "\n",
+        "for prediction in response.predictions[0]:\n",
+        "    print(prediction[\"generated_text\"])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "af21a3cff1e0"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "911406c1561e"
+      },
+      "outputs": [],
+      "source": [
+        "# Delete custom train jobs.\n",
+        "train_job.delete()\n",
+        "\n",
+        "# Undeploy model and delete endpoint.\n",
+        "endpoint_without_peft.delete(force=True)\n",
+        "endpoint_with_peft.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "model_without_peft.delete()\n",
+        "model_with_peft.delete()"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_pytorch_falcon_instruct_peft.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_pytorch_instructpix2pix.ipynb

@@ -204,7 +204,7 @@
       "outputs": [],
       "source": [
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_layoutml_document_qa.ipynb

@@ -204,7 +204,7 @@
       "outputs": [],
       "source": [
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_openllama_peft.ipynb

@@ -0,0 +1,610 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7d9bbf86da5e"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2023 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "99c1c3fc2ca5"
+      },
+      "source": [
+        "# Vertex AI Model Garden - OpenLLaMA (PEFT)\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_openllama_peft.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_openllama_peft.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_pytorch_openllama_peft.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "Open in Vertex AI Workbench\n",
+        "    </a> (A Python-3 CPU notebook is recommended)\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3de7470326a2"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates deploying prebuilt OpenLLaMA, and also finetuning and deploying OpenLLaMA with performance efficient finetuning libraries ([PEFT](https://github.com/huggingface/peft)) in Vertex AI.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Deploy prebuilt OpenLLaMA\n",
+        "- Finetune and deploy OpenLLaMA with PEFT, supporting\n",
+        "\n",
+        "| Models | LoRA |\n",
+        "| :- | :- |\n",
+        "| [openlm-research/open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b) | Y |\n",
+        "| [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b) | Y |\n",
+        "| [openlm-research/open_llama_13b](https://huggingface.co/openlm-research/open_llama_13b) | Y |\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "264c07757582"
+      },
+      "source": [
+        "## Before you begin\n",
+        "\n",
+        "**NOTE**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ioensNKM8ned"
+      },
+      "source": [
+        "### Colab only\n",
+        "Run the following commands for Colab and skip this section if you are using Workbench."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "2707b02ef5df"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
+        "    from google.colab import auth as google_auth\n",
+        "\n",
+        "    google_auth.authenticate_user()\n",
+        "    # Install gdown for downloading example training images.\n",
+        "    ! pip3 install gdown\n",
+        "\n",
+        "    # Restart the notebook kernel after installs.\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bb7adab99e41"
+      },
+      "source": [
+        "### Setup Google Cloud project\n",
+        "\n",
+        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "\n",
+        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
+        "\n",
+        "1. [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs.\n",
+        "\n",
+        "1. [Create a service account](https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console) with `Vertex AI User` and `Storage Object Admin` roles for deploying fine tuned model to Vertex AI endpoint."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6c460088b873"
+      },
+      "source": [
+        "Fill following variables for experiments environment:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "855d6b96f291"
+      },
+      "outputs": [],
+      "source": [
+        "# Cloud project id.\n",
+        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# The region you want to launch jobs in.\n",
+        "REGION = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# The Cloud Storage bucket for storing experiments output.\n",
+        "BUCKET_URI = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "\n",
+        "import os\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "EXPERIMENT_BUCKET = os.path.join(BUCKET_URI, \"peft\")\n",
+        "DATA_BUCKET = os.path.join(EXPERIMENT_BUCKET, \"data\")\n",
+        "MODEL_BUCKET = os.path.join(EXPERIMENT_BUCKET, \"model\")\n",
+        "\n",
+        "# The service account looks like:\n",
+        "# '@.iam.gserviceaccount.com'\n",
+        "# Please go to https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console\n",
+        "# and create service account with `Vertex AI User` and `Storage Object Admin` roles.\n",
+        "# The service account for deploying fine tuned model.\n",
+        "SERVICE_ACCOUNT = \"\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e828eb320337"
+      },
+      "source": [
+        "### Initialize Vertex AI API"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "12cd25839741"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2cc825514deb"
+      },
+      "source": [
+        "### Define constants"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "b42bd4fa2b2d"
+      },
+      "outputs": [],
+      "source": [
+        "# The pre-built training and serving docker images.\n",
+        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-peft-train\"\n",
+        "PREDICTION_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-peft-serve\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0c250872074f"
+      },
+      "source": [
+        "### Define common functions"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "354da31189dc"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "from datetime import datetime\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "\n",
+        "def get_job_name_with_datetime(prefix: 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,\n",
+        "    base_model_id,\n",
+        "    finetuned_lora_model_path,\n",
+        "    service_account,\n",
+        "    task,\n",
+        "    machine_type=\"n1-standard-8\",\n",
+        "    accelerator_type=\"NVIDIA_TESLA_V100\",\n",
+        "):\n",
+        "    \"\"\"Deploys trained models into Vertex AI.\"\"\"\n",
+        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
+        "    serving_env = {\n",
+        "        \"BASE_MODEL_ID\": base_model_id,\n",
+        "        \"TASK\": task,\n",
+        "    }\n",
+        "    if finetuned_lora_model_path:\n",
+        "        serving_env[\"FINETUNED_LORA_MODEL_PATH\"] = finetuned_lora_model_path\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=PREDICTION_DOCKER_URI,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_predict_route=\"/predictions/peft_serving\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        serving_container_environment_variables=serving_env,\n",
+        "    )\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=1,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=service_account,\n",
+        "    )\n",
+        "    return model, endpoint"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8neJc8CnDDpu"
+      },
+      "source": [
+        "## Deploy Prebuilt OpenLLaMA\n",
+        "\n",
+        "This section deploys prebuilt OpenLLaMA models on the Endpoint. The model deployment step will take ~15 minutes to complete.\n",
+        "\n",
+        "The peak GPU memory usages for [openlm-research/open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b), [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b), and [openlm-research/open_llama_13b](https://huggingface.co/openlm-research/open_llama_13b) are ~5.3G, ~8.7G and ~15.2G separately with the default settings. We use V100 in deployments for simplicity."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2MjaORIIFDVu"
+      },
+      "source": [
+        "Set the prebuilt model id."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "E8OiHHNNE_wj"
+      },
+      "outputs": [],
+      "source": [
+        "prebuilt_model_id = \"openlm-research/open_llama_3b\"  # @param [\"openlm-research/open_llama_3b\", \"openlm-research/open_llama_7b\", \"openlm-research/open_llama_13b\"]"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dHFW7yvjaVFV"
+      },
+      "source": [
+        "We use the PEFT serving images to deploy prebuilt OpenLLaMA models, by setting finetuning LoRA model paths as empty."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Uak1pyEeExYM"
+      },
+      "outputs": [],
+      "source": [
+        "model_without_peft, endpoint_without_peft = deploy_model(\n",
+        "    model_name=get_job_name_with_datetime(prefix=\"openllama-serve\"),\n",
+        "    base_model_id=prebuilt_model_id,\n",
+        "    finetuned_lora_model_path=\"\",  # This will avoid override finetuning models.\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    task=\"causal-language-modeling-lora\",\n",
+        ")\n",
+        "print(\"endpoint_name:\", endpoint_without_peft.name)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "sGKIjgmDFRW2"
+      },
+      "source": [
+        "NOTE: The prebuilt model weights will be downloaded on the fly from the orginal location after the deployment succeeds. Thus additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint.\n",
+        "\n",
+        "Once deployment succeeds, you can send requests to the endpoint with text prompts."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "rDHsCOqvFYBi"
+      },
+      "outputs": [],
+      "source": [
+        "# # Loads an existing endpoint as below.\n",
+        "# endpoint_name = endpoint_without_peft.name\n",
+        "# aip_endpoint_name = (\n",
+        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
+        "# )\n",
+        "# endpoint_without_peft = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "instances = [\n",
+        "    {\"prompt\": \"Hi, Google.\"},\n",
+        "]\n",
+        "response = endpoint_without_peft.predict(instances=instances)\n",
+        "\n",
+        "for prediction in response.predictions[0]:\n",
+        "    print(prediction[\"generated_text\"])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e70e3519ff8b"
+      },
+      "source": [
+        "## Finetune and deploy OpenLLaMA with PEFT\n",
+        "\n",
+        "This section demonstrates how to finetune and dpeloy OpenLLaMA with PEFT LoRA."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5qCrm_kJH5cz"
+      },
+      "source": [
+        "Set the base model id."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "N3UBLiYrM3sU"
+      },
+      "outputs": [],
+      "source": [
+        "base_model_id = \"openlm-research/open_llama_3b\"  # @param [\"openlm-research/open_llama_3b\", \"openlm-research/open_llama_7b\", \"openlm-research/open_llama_13b\"]"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "iWGwJHqI7LMs"
+      },
+      "source": [
+        "### Finetune"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "KKEYoRfiHDVv"
+      },
+      "source": [
+        "Use the Vertex AI SDK to create and run the custom training jobs with Vertex AI Model Garden training images.\n",
+        "\n",
+        "This example uses the dataset [Abirate/english_quotes](https://huggingface.co/datasets/Abirate/english_quotes).\n",
+        "\n",
+        "In order to make the finetuning efficiently, we enabled quantization (8bits) when loading pretrained models for finetuning LoRA models. The peak GPU memory usages are ~7G, ~10G and ~16G for finetuning LoRA models for [openlm-research/open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b), [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b), and [openlm-research/open_llama_13b](https://huggingface.co/openlm-research/open_llama_13b) separately with default training parameters and the example dataset. In theory, open_llama_3b and open_llama_7b can be finetuned on 1 V100, and open_llama_13b can be finetuned on 1 A100 (40G). We choose to use 1 A100 (40G) by default to support all these models in this notebook for simplicity."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "65467b361315"
+      },
+      "outputs": [],
+      "source": [
+        "dataset_name = \"Abirate/english_quotes\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Worker pool spec.\n",
+        "# machine_type = \"n1-standard-8\"\n",
+        "# accelerator_type = \"NVIDIA_TESLA_V100\"\n",
+        "machine_type = \"a2-highgpu-1g\"\n",
+        "accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+        "replica_count = 1\n",
+        "accelerator_count = 1\n",
+        "\n",
+        "# Setup training job.\n",
+        "job_name = get_job_name_with_datetime(\"openllama-lora-train\")\n",
+        "train_job = aiplatform.CustomContainerTrainingJob(\n",
+        "    display_name=job_name,\n",
+        "    container_uri=TRAIN_DOCKER_URI,\n",
+        ")\n",
+        "output_dir = os.path.join(MODEL_BUCKET, job_name)\n",
+        "output_dir_gcsfuse = output_dir.replace(\"gs://\", \"/gcs/\")\n",
+        "\n",
+        "# Pass training arguments and launch job.\n",
+        "train_job.run(\n",
+        "    args=[\n",
+        "        \"--task=causal-language-modeling-lora\",\n",
+        "        f\"--pretrained_model_id={base_model_id}\",\n",
+        "        f\"--dataset_name={dataset_name}\",\n",
+        "        f\"--output_dir={output_dir_gcsfuse}\",\n",
+        "        \"--lora_rank=16\",\n",
+        "        \"--lora_alpha=32\",\n",
+        "        \"--lora_dropout=0.05\",\n",
+        "        \"--warmup_steps=10\",\n",
+        "        \"--max_steps=10\",\n",
+        "        \"--learning_rate=2e-4\",\n",
+        "    ],\n",
+        "    replica_count=replica_count,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    boot_disk_size_gb=500,\n",
+        ")\n",
+        "\n",
+        "print(\"Trained models were saved in: \", output_dir)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "jqmCtkGnhDmp"
+      },
+      "source": [
+        "### Deploy\n",
+        "This section uploads the model to Model Registry and deploys it on the Endpoint.\n",
+        "\n",
+        "The model deployment step will take ~15 minutes to complete.\n",
+        "\n",
+        "The peak GPU memory usages for [openlm-research/open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b), [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b), and [openlm-research/open_llama_13b](https://huggingface.co/openlm-research/open_llama_13b) with LoRA weights are ~5.3G, ~8.7G and ~15.2G separately with the default settings. We use V100 in deployments for simplicity."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bf55e38815dc"
+      },
+      "outputs": [],
+      "source": [
+        "model_with_peft, endpoint_with_peft = deploy_model(\n",
+        "    model_name=get_job_name_with_datetime(prefix=\"openllama-peft-serve\"),\n",
+        "    base_model_id=base_model_id,\n",
+        "    finetuned_lora_model_path=output_dir,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    task=\"causal-language-modeling-lora\",\n",
+        ")\n",
+        "print(\"endpoint_name:\", endpoint_with_peft.name)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "80b3fd2ace09"
+      },
+      "source": [
+        "NOTE: After the deployment succeeds, the base model weights will be downloaded one the fly from the original location and LoRA model weights will be downloaded from the GCS bucket used in training above. Thus additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint.\n",
+        "\n",
+        "Once deployment succeeds, you can send requests to the endpoint with text prompts."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "4ab04da3ec9a"
+      },
+      "outputs": [],
+      "source": [
+        "# # Loads an existing endpoint as below.\n",
+        "# endpoint_name = endpoint_with_peft.name\n",
+        "# aip_endpoint_name = (\n",
+        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
+        "# )\n",
+        "# endpoint_with_peft = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "instances = [\n",
+        "    {\"prompt\": \"Hi, Google.\"},\n",
+        "]\n",
+        "response = endpoint_with_peft.predict(instances=instances)\n",
+        "\n",
+        "for prediction in response.predictions[0]:\n",
+        "    print(prediction[\"generated_text\"])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "af21a3cff1e0"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "911406c1561e"
+      },
+      "outputs": [],
+      "source": [
+        "# Delete custom train jobs.\n",
+        "train_job.delete()\n",
+        "\n",
+        "# Undeploy model and delete endpoint.\n",
+        "endpoint_without_peft.delete(force=True)\n",
+        "endpoint_with_peft.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "model_without_peft.delete()\n",
+        "model_with_peft.delete()"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_pytorch_openllama_peft.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_pytorch_owlvit.ipynb

@@ -204,7 +204,7 @@
       "outputs": [],
       "source": [
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_peft.ipynb

@@ -649,14 +649,13 @@
         "# # If deploy finetuned falcon-40b-instruct models, please set\n",
         "# machine_type = \"a2-highgpu-1g\",\n",
         "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
-        "machine_type = \"n1-standard-8\",\n",
+        "machine_type = \"n1-standard-8\"\n",
         "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
         "\n",
-        "accelerator_type =\n",
         "model, endpoint = deploy_model(\n",
         "    model_name=get_job_name_with_datetime(prefix=\"falcon-peft-serve\"),\n",
         "    base_model_id=base_model_id,\n",
-        "    finetuned_lora_model_path=os.path.join(output_dir, \"checkpoint-\"+str(max_steps)),\n",
+        "    finetuned_lora_model_path=os.path.join(output_dir, \"checkpoint-\" + str(max_steps)),\n",
         "    service_account=SERVICE_ACCOUNT,\n",
         "    task=\"instruct-lora\",\n",
         "    machine_type=machine_type,\n",

notebooks/community/model_garden/model_garden_pytorch_pic2word.ipynb

@@ -0,0 +1,633 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "ad41232f-4ac9-4607-8ee6-9b1d98d0f5c3",
+      "metadata": {
+        "id": "7d9bbf86da5e"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2023 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."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "30e2cb63",
+      "metadata": {
+        "id": "99c1c3fc2ca5"
+      },
+      "source": [
+        "# Vertex AI Model Garden - Pic2Word\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pic2word.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_jpic2word.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>                                                                                               <td>\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_pic2word.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "Open in Vertex AI Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "5d5af88c-a3c8-46fa-b90a-1b6737a91534",
+      "metadata": {
+        "id": "7e3e5205fbfd"
+      },
+      "source": [
+        "## Overview"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "87b71198-5c24-4e8c-aab1-44e080aa7fe7",
+      "metadata": {
+        "id": "523084ce4894"
+      },
+      "source": [
+        "This notebook demonstrates how to use the [Pic2Word](https://github.com/google-research/composed_image_retrieval) model in Vertex AI Model Garden. "
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "0ba32ec8-e06f-474f-b753-20aaf74ebd12",
+      "metadata": {
+        "id": "f9cbeb1704e1"
+      },
+      "source": [
+        "## Objective"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "95694bc4-0626-4371-8b2a-6b89967742db",
+      "metadata": {
+        "id": "da71cd41e69e"
+      },
+      "source": [
+        "Following the notebook you will conduct experiments using the pre-built docker image on Vertex AI.\n",
+        "\n",
+        "- Deploy pretrained Pic2Word models in Google Cloud Vertex AI\n",
+        "\n",
+        "This tutorial uses the following Google Cloud ML services and resources:\n",
+        "\n",
+        "- Vertex AI Model Registry\n",
+        "- Vertex AI Online Prediction"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "b1893e0c-859b-4d42-b1d8-d133855d8767",
+      "metadata": {
+        "id": "8264c04e0f34"
+      },
+      "source": [
+        "## Dataset"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "a099439c-3122-429b-8215-caaf9aee69c5",
+      "metadata": {
+        "id": "52e4f9df87b2"
+      },
+      "source": [
+        "We use the [COCO](https://cocodataset.org/#home) validation set (5,000 images) for evaluation."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "56d2a9f8-2d99-407e-a378-57f99754185b",
+      "metadata": {
+        "id": "169ef74e43a6"
+      },
+      "source": [
+        "## Costs"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "bb62de0d-71ac-45ee-82ac-5f7465f2f0eb",
+      "metadata": {
+        "id": "61a64432ac87"
+      },
+      "source": [
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "- Vertex AI\n",
+        "- Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing?_ga=2.46650789.-341051769.1686949237) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing?_ga=2.46650789.-341051769.1686949237), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/?_ga=2.247379078.-341051769.1686949237) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "ea27403e",
+      "metadata": {
+        "id": "586c9147ba8a"
+      },
+      "source": [
+        "# Installation\n",
+        "\n",
+        "Install the following packages required to execute this notebook."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "f9c685d9",
+      "metadata": {
+        "id": "68f1241e68a7"
+      },
+      "outputs": [],
+      "source": [
+        "if \"google.colab\" in str(get_ipython()):\n",
+        "    # Configs for colab notebooks.\n",
+        "    ! pip3 install --upgrade --quiet google-cloud-aiplatform\n",
+        "\n",
+        "    # Automatically restart kernel after installs\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)\n",
+        "\n",
+        "    from google.colab import auth as google_auth\n",
+        "\n",
+        "    google_auth.authenticate_user()"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "e75fe500-5dc1-4ab3-8829-3eeada5be7f7",
+      "metadata": {
+        "id": "d0c7bca32b78"
+      },
+      "source": [
+        "## Setup environment"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "e8b0ef65",
+      "metadata": {
+        "id": "BF1j6f9HApxa"
+      },
+      "source": [
+        "## Before you begin\n",
+        "\n",
+        "### Set up your Google Cloud project\n",
+        "\n",
+        "**The following steps are required, regardless of your notebook environment.**\n",
+        "\n",
+        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "\n",
+        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
+        "\n",
+        "1. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
+        "\n",
+        "1. [Create a service account](https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console) with **Vertex AI User** and **Storage Object Admin** roles for deploying fine tuned model to Vertex AI endpoint.\n",
+        "\n",
+        "1. Enter your project ID in the cell below. Then run the cell to make sure the Cloud SDK uses the right project for all the commands in this notebook.\n",
+        "\n",
+        "\n",
+        "**Note**: Jupyter runs lines prefixed with ! as shell commands, and it interpolates Python variables prefixed with $ into these commands."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "09490824",
+      "metadata": {
+        "id": "WReHDGG5g0XY"
+      },
+      "source": [
+        "#### Set your project ID\n",
+        "\n",
+        "**If you don't know your project ID**, try the following:\n",
+        "* Run `gcloud config list`.\n",
+        "* Run `gcloud projects list`.\n",
+        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "6ddc1d11",
+      "metadata": {
+        "id": "oM1iC_MfAts1"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"your-project-id\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Set the project id\n",
+        "! gcloud config set project {PROJECT_ID}"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "9ed7d668",
+      "metadata": {
+        "id": "region"
+      },
+      "source": [
+        "#### Region\n",
+        "\n",
+        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "90d0108f",
+      "metadata": {
+        "id": "twgKk-LsLmX3"
+      },
+      "outputs": [],
+      "source": [
+        "REGION = \"us-central1\"  # @param {type: \"string\"}"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "ca43938d",
+      "metadata": {
+        "id": "409f1effab6c"
+      },
+      "source": [
+        "### Buckets\n",
+        "\n",
+        "You can create a storage bucket to store model input and output images."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "889d10ac",
+      "metadata": {
+        "id": "bcdadd216c34"
+      },
+      "outputs": [],
+      "source": [
+        "# The form for BUCKET_URI is gs://.\n",
+        "BUCKET_URI = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "import os\n",
+        "\n",
+        "INPUT_BUCKET = os.path.join(BUCKET_URI, \"/input\")\n",
+        "OUTPUT_BUCKET = os.path.join(BUCKET_URI, \"/output\")"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "15c0e671",
+      "metadata": {
+        "id": "sBCra4QMA2wR"
+      },
+      "source": [
+        "### Authenticate your Google Cloud account\n",
+        "\n",
+        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "623a03cc",
+      "metadata": {
+        "id": "74ccc9e52986"
+      },
+      "source": [
+        "**1. Vertex AI Workbench**\n",
+        "* Do nothing as you are already authenticated."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "fbee4160",
+      "metadata": {
+        "id": "de775a3773ba"
+      },
+      "source": [
+        "**2. Local JupyterLab instance, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "3eabaf6b",
+      "metadata": {
+        "id": "254614fa0c46"
+      },
+      "outputs": [],
+      "source": [
+        "# ! gcloud auth login"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "90a615fe",
+      "metadata": {
+        "id": "ef21552ccea8"
+      },
+      "source": [
+        "**3. Colab, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "3d4e50e1",
+      "metadata": {
+        "id": "603adbbf0532"
+      },
+      "outputs": [],
+      "source": [
+        "# from google.colab import auth\n",
+        "# auth.authenticate_user()"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "b0b40233",
+      "metadata": {
+        "id": "f6b2ccc891ed"
+      },
+      "source": [
+        "**4. Service account or other**\n",
+        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "a8019fcc-26ce-4536-b394-aa1fb18794ee",
+      "metadata": {
+        "id": "4226467373ef"
+      },
+      "source": [
+        "If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk) and [gsutil](https://cloud.google.com/storage/docs/gsutil_install)."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "315598c1-3ff4-4bc8-9b82-7e2fe55d4678",
+      "metadata": {
+        "id": "8ba36d3a171d"
+      },
+      "source": [
+        "### Setup variables"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "6d26ddb2-04ad-4973-8965-9b25eb24993f",
+      "metadata": {
+        "id": "25ca675e59e7"
+      },
+      "outputs": [],
+      "source": [
+        "# Prediction constants.\n",
+        "PREDICTION_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pic2word_serve:latest\"\n",
+        "PREDICTION_ACCELERATOR_TYPE = \"NVIDIA_TESLA_T4\"\n",
+        "PREDICTION_MACHINE_TYPE = \"n1-standard-8\"\n",
+        "\n",
+        "# The serving port.\n",
+        "SERVE_PORT = 7080\n",
+        "\n",
+        "# The service account looks like:\n",
+        "# '@.iam.gserviceaccount.com'\n",
+        "# Please go to https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console\n",
+        "# and create service account with `Vertex AI User` and `Storage Object Admin` roles.\n",
+        "SERVICE_ACCOUNT = \"\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "0daec6e5-7642-458f-bd93-60a15431b0a1",
+      "metadata": {
+        "id": "71dd15118703"
+      },
+      "source": [
+        "## Deploy model for online prediction"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "b286ce6b-cb00-441b-8249-93497de17974",
+      "metadata": {
+        "id": "5c891f6352ad"
+      },
+      "source": [
+        "This section uploads the model to Vertex Model Registry and deploys it on an Endpoint resource. This will take around 15 minutes to finish."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "01f686b2-9227-496b-9c4d-a63becb59fb8",
+      "metadata": {
+        "id": "e6abf4ee450a"
+      },
+      "source": [
+        "### Upload and deploy model to Vertex AI"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "8885d5be-5007-4778-9562-976102fefe73",
+      "metadata": {
+        "id": "645d8d1df8d4"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "# Init common setup.\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)\n",
+        "\n",
+        "# Upload model.\n",
+        "serving_env = {}\n",
+        "model = aiplatform.Model.upload(\n",
+        "    display_name=\"pic2word-model\",\n",
+        "    serving_container_image_uri=PREDICTION_DOCKER_URI,\n",
+        "    serving_container_ports=[SERVE_PORT],\n",
+        "    serving_container_predict_route=\"/predictions/pic2word\",\n",
+        "    serving_container_health_route=\"/ping\",\n",
+        "    serving_container_environment_variables=serving_env,\n",
+        ")\n",
+        "# Or reuse a pre-uploaded model.\n",
+        "# model = aiplatform.Model('projects/123456789/locations/us-central1/models/123456789@1')\n",
+        "\n",
+        "# Create an endpoint.\n",
+        "endpoint = aiplatform.Endpoint.create(display_name=\"pytorch-pic2word-endpoint\")\n",
+        "# Or reuse a pre-created endpoint.\n",
+        "# endpoint = aiplatform.Endpoint('projects/123456789/locations/us-central1/endpoints/123456789')\n",
+        "\n",
+        "# Deploy model to endpoint.\n",
+        "model.deploy(\n",
+        "    endpoint=endpoint,\n",
+        "    machine_type=PREDICTION_MACHINE_TYPE,\n",
+        "    accelerator_type=PREDICTION_ACCELERATOR_TYPE,\n",
+        "    accelerator_count=1,\n",
+        "    traffic_percentage=100,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        ")"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "f0bd36bb-f988-479a-be40-18c54980ca1f",
+      "metadata": {
+        "id": "2bb47a2855a2"
+      },
+      "source": [
+        "You can manage your uploaded models in the [Model Registry](https://pantheon.corp.google.com/vertex-ai/models) and your endpoints in the [Endpoints](https://pantheon.corp.google.com/vertex-ai/endpoints)."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "95edbcdd-8246-444a-aef2-a9cadbc03000",
+      "metadata": {
+        "id": "f14ff5e7a476"
+      },
+      "source": [
+        "## Send a prediction request to the endpoint"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "3a83d6bb",
+      "metadata": {
+        "id": "c41c7d51d654"
+      },
+      "source": [
+        "Set the query text and upload your desired query image to the specified INPUT_BUCKET. Use an asterisk in the query text to specify which token gets replaced by the image token. "
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "f7c6fc99-42cd-4544-86df-b14f49189bfd",
+      "metadata": {
+        "id": "58d4f8e84e02"
+      },
+      "outputs": [],
+      "source": [
+        "payload = json.dumps(\n",
+        "    {\n",
+        "        \"query\": \"a bunch of *\",\n",
+        "        \"image_path\": INPUT_BUCKET,\n",
+        "        \"image_file_name\": \"cat.png\",\n",
+        "        \"output_storage_dir\": OUTPUT_BUCKET,\n",
+        "    }\n",
+        ")\n",
+        "response = endpoint.predict(payload).predictions\n",
+        "print(response)"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "8e71ee23-967c-4871-802f-ccb39958386e",
+      "metadata": {
+        "id": "608b936179df"
+      },
+      "source": [
+        "## Cleaning Up\n",
+        "\n",
+        "To clean up all Google Cloud resources used in this project, 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:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "d9be2a18-db11-4108-9622-647b33be2594",
+      "metadata": {
+        "id": "7d17b385141d"
+      },
+      "outputs": [],
+      "source": [
+        "# Delete endpoint resource.\n",
+        "endpoint.delete(force=True)\n",
+        "\n",
+        "# Delete model resource.\n",
+        "model.delete()\n",
+        "\n",
+        "# Delete Cloud Storage objects that were created.\n",
+        "delete_bucket = False\n",
+        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "    ! gsutil -m rm -r $BUCKET_URI"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_pytorch_pic2word.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_pytorch_sam.ipynb

@@ -231,7 +231,7 @@
       "source": [
         "# The pre-built serving docker image.\n",
         "# The model artifacts are embedded within the container, except for model weights which will be downloaded during deployment.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/sam-serve\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/sam-serve\"\n",
         "\n",
         "# The serving port.\n",
         "SERVE_PORT = 7080\n",

notebooks/community/model_garden/model_garden_pytorch_stable_diffusion.ipynb

@@ -259,10 +259,10 @@
       "outputs": [],
       "source": [
         "# The pre-built training docker image. It contains training scripts and models.\n",
-        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-diffusers-train:latest\"\n",
+        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-train:latest\"\n",
         "\n",
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_stable_diffusion_2_1.ipynb

@@ -249,10 +249,10 @@
       "outputs": [],
       "source": [
         "# The pre-built training docker image. It contains training scripts and models.\n",
-        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-diffusers-train:latest\"\n",
+        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-train:latest\"\n",
         "\n",
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_stable_diffusion_inpainting.ipynb

@@ -252,10 +252,10 @@
       "outputs": [],
       "source": [
         "# The pre-built training docker image. It contains training scripts and models.\n",
-        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-diffusers-train:latest\"\n",
+        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-train:latest\"\n",
         "\n",
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_stable_diffusion_upscaler.ipynb

@@ -243,7 +243,7 @@
       "outputs": [],
       "source": [
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_text_to_video.ipynb

@@ -215,7 +215,7 @@
       "outputs": [],
       "source": [
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_text_to_video_zero_shot.ipynb

@@ -216,7 +216,7 @@
       "outputs": [],
       "source": [
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_timm.ipynb

@@ -262,6 +262,8 @@
       },
       "outputs": [],
       "source": [
+        "import urllib\n",
+        "\n",
         "import timm\n",
         "import torch\n",
         "from PIL import Image\n",

notebooks/community/model_garden/model_garden_pytorch_vilt_vqa.ipynb

@@ -204,7 +204,7 @@
       "outputs": [],
       "source": [
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_vit_gpt2_image_captioning.ipynb

@@ -43,7 +43,7 @@
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
-        "  <td>                                                                                               <td>\n",
+        "  <td>\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_pytorch_vit_gpt2_image_captioning.ipynb\">\n",
         "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
         "Open in Vertex AI Workbench\n",
@@ -204,7 +204,7 @@
       "outputs": [],
       "source": [
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_tfvision_image_classification.ipynb

@@ -44,7 +44,7 @@
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
-        "  <td>                                                                                               <td>\n",
+        "  <td>\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_tfvision_image_classification.ipynb\">\n",
         "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
         "Open in Vertex AI Workbench\n",

notebooks/community/model_garden/model_garden_tfvision_image_object_detection.ipynb

@@ -24,7 +24,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "TirJ-SGQseby"
@@ -45,7 +44,7 @@
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
-        "  <td>                                                                                               <td>\n",
+        "  <td>\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_tfvision_image_object_detection.ipynb\">\n",
         "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
         "Open in Vertex AI Workbench\n",
@@ -55,7 +54,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "dwGLvtIeECLK"
@@ -67,7 +65,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "tvgnzT1CKxrO"
@@ -106,7 +103,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "KEukV6uRk_S3"
@@ -116,7 +112,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "z__i0w0lCAsW"
@@ -149,7 +144,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "BF1j6f9HApxa"
@@ -217,11 +211,13 @@
         "! gsutil cp coco_spinenet143_gpu_multiworker_mirrored.yaml $CONFIG_DIR/\n",
         "\n",
         "! wget https://raw.githubusercontent.com/tensorflow/models/master/official/projects/yolo/configs/experiments/yolov4/detection/scaled_yolov4_1280_gpu.yaml\n",
-        "! gsutil cp scaled_yolov4_1280_gpu.yaml $CONFIG_DIR/"
+        "! gsutil cp scaled_yolov4_1280_gpu.yaml $CONFIG_DIR/\n",
+        "\n",
+        "! wget https://github.com/tensorflow/models/blob/master/official/projects/yolo/configs/experiments/yolov7/detection/yolov7_gpu.yaml\n",
+        "! gsutil cp yolov7_gpu.yaml $CONFIG_DIR/"
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "n6IFz75WGCam"
@@ -262,6 +258,7 @@
         "    CONFIG_DIR, \"coco_spinenet143_gpu_multiworker_mirrored.yaml\"\n",
         ")\n",
         "TRAIN_YOLOV4_CONFIG = os.path.join(CONFIG_DIR, \"scaled_yolov4_1280_gpu.yaml\")\n",
+        "TRAIN_YOLOV7_CONFIG = os.path.join(CONFIG_DIR, \"yolov7_gpu.yaml\")\n",
         "\n",
         "# Evaluation constants.\n",
         "EVALUATION_METRIC = \"AP50\"\n",
@@ -285,7 +282,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "ZZFPe_GezXg8"
@@ -510,7 +506,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "RB_xY9ipr7ZU"
@@ -526,7 +521,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "zgPO1eR3CYjk"
@@ -603,7 +597,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "SA8DVTn7j69v"
@@ -615,7 +608,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "aaff6f5be7f6"
@@ -649,7 +641,7 @@
         "\n",
         "# Refer to https://github.com/tensorflow/models/blob/master/official/vision/MODEL_GARDEN.md\n",
         "# for more model details.\n",
-        "experiment = \"retinanet_spinenet96\"  # @param ['retinanet_spinenet49', \"retinanet_spinenet96\", 'retinanet_spinenet143', 'scaled_yolo_v4']\n",
+        "experiment = \"retinanet_spinenet96\"  # @param ['retinanet_spinenet49', \"retinanet_spinenet96\", 'retinanet_spinenet143', 'scaled_yolo_v4', 'yolov7']\n",
         "\n",
         "train_job_name = get_job_name_with_datetime(TRAINING_JOB_PREFIX + \"_\" + OBJECTIVE)\n",
         "model_dir = os.path.join(BUCKET_URI, train_job_name)\n",
@@ -706,6 +698,15 @@
         "            \"input_size\": \"1280,1280\",\n",
         "        },\n",
         "    ),\n",
+        "    # yolov7 experiment args.\n",
+        "    \"yolov7\": dict(\n",
+        "        common_args,\n",
+        "        **{\n",
+        "            \"experiment\": \"coco_yolov7\",\n",
+        "            \"config_file\": TRAIN_YOLOV7_CONFIG,\n",
+        "            \"input_size\": \"640,640\",\n",
+        "        },\n",
+        "    ),\n",
         "}\n",
         "experiment_container_args = experiment_container_args_dict[experiment]\n",
         "\n",
@@ -715,6 +716,8 @@
         "    experiment_container_args[\"init_checkpoint\"] = upload_checkpoint_to_gcs(\n",
         "        init_checkpoint\n",
         "    )\n",
+        "if \"yolov7\" in experiment:\n",
+        "    TRAIN_CONTAINER_URI = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/tfvision-oss-v2\"\n",
         "\n",
         "params_override = \"runtime.num_gpus=%s\" % TRAIN_NUM_GPU\n",
         "eval_params_override = \"runtime.num_gpus=1,runtime.distribution_strategy=mirrored\"\n",
@@ -768,7 +771,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "HwcCjwlBTQIz"
@@ -817,7 +819,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "mV-Djz-frBni"
@@ -878,7 +879,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "g0BGaofgsMsy"
@@ -904,11 +904,12 @@
         "\n",
         "upload_job_name = get_job_name_with_datetime(UPLOAD_JOB_PREFIX + \"_\" + OBJECTIVE)\n",
         "\n",
-        "if 'yolo' in experiment:\n",
+        "if \"yolo\" in experiment:\n",
         "    SERVING_CONTAINER_ARGS = [\"--allow_precompilation\"]\n",
         "else:\n",
         "    SERVING_CONTAINER_ARGS = [\"--allow_precompilation\", \"--allow_compression\"]\n",
         "\n",
+        "\n",
         "model = aiplatform.Model.upload(\n",
         "    display_name=upload_job_name,\n",
         "    artifact_uri=trained_model_dir,\n",
@@ -982,7 +983,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "kkH2nrpdp4sp"

notebooks/community/model_garden/model_garden_tfvision_image_segmentation.ipynb

@@ -44,7 +44,7 @@
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
-        "  <td>                                                                                               <td>\n",
+        "  <td>\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_tfvision_image_segmentation.ipynb\">\n",
         "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
         "Open in Vertex AI Workbench\n",

notebooks/official/automl/automl_forecasting_bqml_arima_plus_comparison.ipynb

@@ -88,8 +88,8 @@
         "- Make a batch prediction with the BigQuery ML model.\n",
         "- Create a Vertex AI `Dataset` resource.\n",
         "- Train the Vertex AI Forecasting model.\n",
-        "- View the Model evaluation.\n",
-        "- Make a batch prediction with the Model.\n"
+        "- View the Vertex AI Model Evaluation results.\n",
+        "- Make a batch prediction with the Vertex AI Forecasting model.\n"
       ]
     },
     {

notebooks/official/automl/sdk_automl_tabular_forecasting_batch.ipynb

@@ -620,10 +620,15 @@
         ")\n",
         "# Must be the same region as batch_predict_bq_input_uri\n",
         "client = bigquery.Client(project=PROJECT_ID)\n",
-        "bq_dataset = bigquery.Dataset(batch_predict_bq_output_dataset_path)\n",
+        "bq_dataset_id = bigquery.Dataset(batch_predict_bq_output_dataset_path)\n",
         "dataset_region = \"US\"  # @param {type : \"string\"}\n",
-        "bq_dataset.location = dataset_region\n",
-        "bq_dataset = client.create_dataset(bq_dataset)\n",
+        "bq_dataset_id.location = dataset_region\n",
+        "# delete any existing dataset\n",
+        "try:\n",
+        "    client.delete_dataset(bq_dataset_id, delete_contents=True)\n",
+        "except Exception as e:\n",
+        "    print(e)\n",
+        "bq_dataset = client.create_dataset(bq_dataset_id)\n",
         "print(\n",
         "    \"Created bigquery dataset {} in {}\".format(\n",
         "        batch_predict_bq_output_dataset_path, dataset_region\n",
@@ -857,6 +862,12 @@
         "# Delete batch prediction job\n",
         "batch_prediction_job.delete()\n",
         "\n",
+        "# Delete the dataset\n",
+        "try:\n",
+        "    client.delete_dataset(bq_dataset_id)\n",
+        "except Exception as e:\n",
+        "    print(e)\n",
+        "\n",
         "# Set this to true only if you'd like to delete your bucket\n",
         "delete_bucket = False\n",
         "\n",

notebooks/official/custom/SDK_Custom_Container_Prediction.ipynb

@@ -166,7 +166,7 @@
         "! pip3 install -U  \"uvicorn[standard]>=0.12.0,<0.14.0\" fastapi~=0.63 -q\n",
         "\n",
         "# Vertex SDK for Python\n",
-        "! pip3 install -U  google-cloud-aiplatform -q"
+        "! pip3 install --upgrade --quiet  google-cloud-aiplatform"
       ]
     },
     {

notebooks/official/experiments/comparing_pipeline_runs.ipynb

@@ -714,7 +714,12 @@
         "):\n",
         "\n",
         "    custom_trainer(\n",
-        "        train_uri, label_uri, max_depth, learning_rate, boost_rounds, model_uri\n",
+        "        train_uri=train_uri,\n",
+        "        label_uri=label_uri,\n",
+        "        max_depth=max_depth,\n",
+        "        learning_rate=learning_rate,\n",
+        "        boost_rounds=boost_rounds,\n",
+        "        model_uri=model_uri,\n",
         "    )"
       ]
     },

notebooks/official/model_evaluation/custom_tabular_classification_model_evaluation.ipynb

@@ -29,7 +29,7 @@
         "id": "JAPoU8Sm5E6e"
       },
       "source": [
-        "# Vertex AI Pipelines: Evaluating BatchPrediction results from a Custom Tabular classification model\n",
+        "# Vertex AI Pipelines: Evaluating BatchPrediction results from a custom tabular classification model\n",
         "\n",
         "<table align=\"left\">\n",
         "\n",
@@ -151,15 +151,16 @@
       "outputs": [],
       "source": [
         "# Install the latest versions of the following packages\n",
-        "! pip3 install --upgrade google-cloud-aiplatform \\\n",
-        "                            google-cloud-pipeline-components==1.0.26 \\\n",
-        "                            matplotlib \\\n",
-        "                            pyarrow -q\n",
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
+        "                                 google-cloud-pipeline-components==1.0.26 \\\n",
+        "                                 matplotlib \\\n",
+        "                                 pyarrow \n",
         "# Install the specified versions of the following packages\n",
-        "! pip3 install scikit-learn==1.0 \\\n",
-        "                  pandas \\\n",
-        "                  joblib==1.2.0 \\\n",
-        "                  numpy==1.23.3 -q"
+        "! pip3 install --quiet scikit-learn==1.0 \\\n",
+        "                       pandas \\\n",
+        "                       joblib==1.2.0 \\\n",
+        "                       numpy==1.23.3 \\\n",
+        "                       db-dtypes"
       ]
     },
     {
@@ -401,12 +402,25 @@
       },
       "outputs": [],
       "source": [
-        "if SERVICE_ACCOUNT == \"[your-service-account]\":\n",
-        "    shell_output = ! gcloud projects list --filter=\"PROJECT_ID:'{PROJECT_ID}'\" --format='value(PROJECT_NUMBER)'\n",
-        "    PROJECT_NUMBER = shell_output[0]\n",
-        "    SERVICE_ACCOUNT = f\"{PROJECT_NUMBER}-compute@developer.gserviceaccount.com\"\n",
+        "import sys\n",
         "\n",
-        "print(\"Service Account:\", SERVICE_ACCOUNT)"
+        "IS_COLAB = \"google.colab\" in sys.modules\n",
+        "if (\n",
+        "    SERVICE_ACCOUNT == \"\"\n",
+        "    or SERVICE_ACCOUNT is None\n",
+        "    or SERVICE_ACCOUNT == \"[your-service-account]\"\n",
+        "):\n",
+        "    # Get your service account from gcloud\n",
+        "    if not IS_COLAB:\n",
+        "        shell_output = !gcloud auth list 2>/dev/null\n",
+        "        SERVICE_ACCOUNT = shell_output[2].replace(\"*\", \"\").strip()\n",
+        "\n",
+        "    else:  # IS_COLAB:\n",
+        "        shell_output = ! gcloud projects describe  $PROJECT_ID\n",
+        "        project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "        SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "\n",
+        "    print(\"Service Account:\", SERVICE_ACCOUNT)"
       ]
     },
     {
@@ -1035,6 +1049,34 @@
         "RUN pip install -r requirements.txt"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "OrpUIkAIs_uQ"
+      },
+      "source": [
+        "#### Create a private Docker repository\n",
+        "\n",
+        "Your first step is to create your own Docker repository in Google Artifact Registry."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "0amu4063tDnG"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "\n",
+        "! gcloud services enable artifactregistry.googleapis.com\n",
+        "\n",
+        "if os.getenv(\"IS_TESTING\"):\n",
+        "    ! sudo apt-get update --yes && sudo apt-get --only-upgrade --yes install google-cloud-sdk-cloud-run-proxy google-cloud-sdk-harbourbridge google-cloud-sdk-cbt google-cloud-sdk-gke-gcloud-auth-plugin google-cloud-sdk-kpt google-cloud-sdk-local-extract google-cloud-sdk-minikube google-cloud-sdk-app-engine-java google-cloud-sdk-app-engine-go google-cloud-sdk-app-engine-python google-cloud-sdk-spanner-emulator google-cloud-sdk-bigtable-emulator google-cloud-sdk-nomos google-cloud-sdk-package-go-module google-cloud-sdk-firestore-emulator kubectl google-cloud-sdk-datastore-emulator google-cloud-sdk-app-engine-python-extras google-cloud-sdk-cloud-build-local google-cloud-sdk-kubectl-oidc google-cloud-sdk-anthos-auth google-cloud-sdk-app-engine-grpc google-cloud-sdk-pubsub-emulator google-cloud-sdk-datalab google-cloud-sdk-skaffold google-cloud-sdk google-cloud-sdk-terraform-tools google-cloud-sdk-config-connector\n",
+        "    ! gcloud components update --quiet"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -1159,6 +1201,8 @@
       "outputs": [],
       "source": [
         "# define the evaluation pipeline\n",
+        "\n",
+        "\n",
         "@kfp.dsl.pipeline(name=\"custom-tabular-classification-evaluation-pipeline\")\n",
         "def evaluation_custom_tabular_feature_attribution_pipeline(\n",
         "    project: str,\n",
@@ -1246,6 +1290,58 @@
         "    )"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "c7b87e9196f3"
+      },
+      "source": [
+        "### Optional: Workaround to import a BigQuery table for predictions_bigquery_source"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "fc2df1b02f1a"
+      },
+      "outputs": [],
+      "source": [
+        "\"\"\"\n",
+        "# Set constants for BigQuery Table\n",
+        "BIGQUERY_PROJECT_ID = \"your-project-id\"\n",
+        "BIGQUERY_DATASET_ID = \"your-dataset-id\"\n",
+        "BIGQUERY_PREDICTION_RESULTS_TABLE_ID = \"your-table-id\"\n",
+        "\n",
+        "# Import the BigQuery table using the importer to obtain a BQTable artifact\n",
+        "bq_table_uri = f\"bq://{BIGQUERY_PROJECT_ID}.{BIGQUERY_DATASET_ID}.{BIGQUERY_PREDICTION_RESULTS_TABLE_ID}\"\n",
+        "bq_table = kfp.v2.dsl.importer(\n",
+        "    artifact_uri=bq_table_uri,\n",
+        "    artifact_class=artifact_types.BQTable,\n",
+        "    metadata={\n",
+        "        \"projectId\": BIGQUERY_PROJECT_ID,\n",
+        "        \"datasetId\": BIGQUERY_DATASET_ID,\n",
+        "        \"tableId\": BIGQUERY_PREDICTION_RESULTS_TABLE_ID,\n",
+        "    },\n",
+        ").output\n",
+        "\n",
+        "# Run the evaluation based on prediction type\n",
+        "eval_task = ModelEvaluationClassificationOp(\n",
+        "    project=project,\n",
+        "    location=location,\n",
+        "    root_dir=root_dir,\n",
+        "    class_labels=evaluation_class_names,\n",
+        "    prediction_label_column=evaluation_prediction_label_column,\n",
+        "    prediction_score_column=evaluation_prediction_score_column,\n",
+        "    target_field_name=target_field_name,\n",
+        "    ground_truth_format=batch_predict_instances_format,\n",
+        "    ground_truth_bigquery_source=bq_table_uri,\n",
+        "    predictions_format=batch_predict_predictions_format,\n",
+        "    predictions_bigquery_source=bq_table,\n",
+        ")\n",
+        "\"\"\""
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {

notebooks/official/prediction/custom_batch_prediction_feature_filter.ipynb

@@ -154,10 +154,11 @@
       "outputs": [],
       "source": [
         "# Install the packages\n",
-        "! pip3 install --upgrade google-cloud-aiplatform \\\n",
-        "                                    google-cloud-storage \\\n",
-        "                                    google-cloud-bigquery \\\n",
-        "                                    pyarrow -q"
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
+        "                                 google-cloud-storage \\\n",
+        "                                 google-cloud-bigquery \\\n",
+        "                                 pyarrow \\\n",
+        "                                 db-dtypes\n"
       ]
     },
     {
@@ -352,7 +353,7 @@
       },
       "outputs": [],
       "source": [
-        "BUCKET_URI = \"gs://your-bucket-name-unique\"  # @param {type:\"string\"}"
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
       ]
     },
     {
@@ -372,7 +373,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
+        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
       ]
     },
     {
@@ -1016,7 +1017,7 @@
         "    display_name=JOB_NAME,\n",
         "    script_path=\"task.py\",\n",
         "    container_uri=TRAIN_IMAGE,\n",
-        "    requirements=[\"google-cloud-bigquery>=2.20.0\", \"db-dtypes\"],\n",
+        "    requirements=[\"google-cloud-bigquery>=2.20.0\", \"db-dtypes\", \"protobuf==3.20.3\"],\n",
         "    model_serving_container_image_uri=DEPLOY_IMAGE,\n",
         ")\n",
         "\n",
@@ -1439,12 +1440,7 @@
         "\n",
         "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
         "\n",
-        "Otherwise, you can delete the individual resources you created in this notebook:\n",
-        "\n",
-        "- Training Job\n",
-        "- Model\n",
-        "- Cloud Storage Bucket\n",
-        "- BigQuery Dataset"
+        "Otherwise, you can delete the individual resources you created in this notebook."
       ]
     },
     {
@@ -1461,6 +1457,9 @@
         "# Delete the training job\n",
         "job.delete()\n",
         "\n",
+        "# Delete the dataset\n",
+        "dataset.delete()\n",
+        "\n",
         "# Delete the model\n",
         "model.delete()\n",
         "\n",

notebooks/official/tabnet/tabnet_vertex_tutorial.ipynb

@@ -149,7 +149,9 @@
       "outputs": [],
       "source": [
         "! pip3 install --quiet --upgrade tensorflow\n",
-        "! pip3 install --quiet --upgrade google-cloud-aiplatform tensorboard-plugin-profile\n",
+        "! pip3 install --quiet --upgrade google-cloud-aiplatform \\\n",
+        "                                 tensorboard-plugin-profile \\\n",
+        "                                 google-cloud-pipeline-components\n",
         "! gcloud components update --quiet"
       ]
     },

notebooks/official/tabular_workflows/wide_and_deep_on_vertex_pipelines.ipynb

@@ -141,7 +141,8 @@
       },
       "outputs": [],
       "source": [
-        "! pip3 install --upgrade --quiet google-cloud-aiplatform google-cloud-pipeline-components"
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
+        "                                 \"google-cloud-pipeline-components<2.0\""
       ]
     },
     {

notebooks/official/tensorboard/tensorboard_custom_training_with_custom_container.ipynb

@@ -147,39 +147,6 @@
         "Use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "rAPBk_OCNN_h"
-      },
-      "source": [
-        "### Set up your local development environment\n",
-        "\n",
-        "**If you are using Colab or Vertex AI Workbench**, your environment already meets all the requirements to run this notebook. You can skip this step.\n",
-        "\n",
-        "Otherwise, make sure your environment meets this notebook's requirements. You need the following:\n",
-        "\n",
-        "- The Cloud Storage SDK\n",
-        "- Git\n",
-        "- Python 3\n",
-        "- virtualenv\n",
-        "- Jupyter notebook running in a virtual environment with Python 3\n",
-        "\n",
-        "The Cloud Storage guide to [Setting up a Python development environment](https://cloud.google.com/python/setup) and the [Jupyter installation guide](https://jupyter.org/install) provide detailed instructions for meeting these requirements. The following steps provide a condensed set of instructions:\n",
-        "\n",
-        "1. [Install and initialize the SDK](https://cloud.google.com/sdk/docs/).\n",
-        "\n",
-        "2. [Install Python 3](https://cloud.google.com/python/setup#installing_python).\n",
-        "\n",
-        "3. [Install virtualenv](https://cloud.google.com/python/setup#installing_and_using_virtualenv) and create a virtual environment that uses Python 3.  Activate the virtual environment.\n",
-        "\n",
-        "4. To install Jupyter, run `pip3 install jupyter` on the command-line in a terminal shell.\n",
-        "\n",
-        "5. To launch Jupyter, run `jupyter notebook` on the command-line in a terminal shell.\n",
-        "\n",
-        "6. Open this notebook in the Jupyter Notebook Dashboard."
-      ]
-    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -199,136 +166,102 @@
       },
       "outputs": [],
       "source": [
-        "import os\n",
-        "\n",
-        "# The Vertex AI Workbench Notebook product has specific requirements\n",
-        "IS_WORKBENCH_NOTEBOOK = os.getenv(\"DL_ANACONDA_HOME\")\n",
-        "\n",
-        "# Vertex AI Notebook requires dependencies to be installed with '--user'\n",
-        "USER_FLAG = \"\"\n",
-        "if IS_WORKBENCH_NOTEBOOK:\n",
-        "    USER_FLAG = \"--user\"\n",
-        "\n",
-        "! pip3 install google-cloud-aiplatform {USER_FLAG} -q\n"
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform "
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "b24902cde81b"
+        "id": "58707a750154"
       },
       "source": [
-        "### Restart the kernel\n",
-        "\n",
-        "Once you've installed the additional packages, you need to restart the notebook kernel so it can find the packages."
+        "### Colab only: Uncomment the following cell to restart the kernel."
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "c61d171395d7"
+        "id": "f200f10a1da3"
       },
       "outputs": [],
       "source": [
-        "import os\n",
+        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
+        "# import IPython\n",
         "\n",
-        "if not os.getenv(\"IS_TESTING\"):\n",
-        "    # Automatically restart kernel after installs\n",
-        "    import IPython\n",
-        "\n",
-        "    app = IPython.Application.instance()\n",
-        "    app.kernel.do_shutdown(True)"
+        "# app = IPython.Application.instance()\n",
+        "# app.kernel.do_shutdown(True)"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "1Dunp1YrhPYo"
+        "id": "BF1j6f9HApxa"
       },
       "source": [
         "## Before you begin\n",
+        "\n",
         "### Set up your Google Cloud project\n",
         "\n",
         "**The following steps are required, regardless of your notebook environment.**\n",
         "\n",
         "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
         "\n",
-        "2. [Make sure that billing is enabled for your project.](https://cloud.google.com/billing/docs/how-to/modify-project)\n",
+        "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
         "\n",
-        "3. [Enable the Vertex AI, Cloud Storage, Cloud Build, and Artifact Registry APIs.](https://console.cloud.google.com/flows/enableapi?apiid=ml.googleapis.com,storage-component.googleapis.com,cloudbuild.googleapis.com,artifactregistry.googleapis.com)\n",
+        "3. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
         "\n",
-        "4. If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
-        "\n",
-        "\n",
-        "5. Enter your project ID in the cell below. Then run the  cell to make sure the\n",
-        "Cloud SDK uses the right project for all the commands in this notebook.\n",
-        "\n",
-        "Note: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$`."
+        "4. If you are running this notebook locally, install the [Cloud SDK](https://cloud.google.com/sdk)."
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "lMqIOZfqhXyD"
+        "id": "WReHDGG5g0XY"
       },
       "source": [
-        "### Set your project ID"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "siN0RvWyZYv9"
-      },
-      "source": [
-        "**If you don't know your project ID**, you may be able to get your project ID using `gcloud`."
+        "#### Set your project ID\n",
+        "\n",
+        "**If you don't know your project ID**, try the following:\n",
+        "* Run `gcloud config list`.\n",
+        "* Run `gcloud projects list`.\n",
+        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "cde8e0876d62"
+        "id": "oM1iC_MfAts1"
       },
       "outputs": [],
       "source": [
-        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "lU-1yD4jZc-6"
-      },
-      "outputs": [],
-      "source": [
-        "if PROJECT_ID == \"\" or PROJECT_ID is None or PROJECT_ID == \"[your-project-id]\":\n",
-        "    # Get your GCP project id from gcloud\n",
-        "    shell_output = ! gcloud config list --format 'value(core.project)' 2>/dev/null\n",
-        "    PROJECT_ID = shell_output[0]\n",
-        "    print(\"Project ID:\", PROJECT_ID)"
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Set the project id\n",
+        "! gcloud config set project {PROJECT_ID}"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "XNVqyIvYaPO-"
+        "id": "region"
       },
       "source": [
-        "Otherwise, set your project ID here."
+        "#### Set the region\n",
+        "\n",
+        "**Optional**: Update the 'REGION' variable to specify the region that you want to use. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "s0F0U7WZhfAZ"
+        "id": "nsN5NJKSu-GU"
       },
       "outputs": [],
       "source": [
-        "! gcloud config set project $PROJECT_ID"
+        "REGION = \"us-central1\"  # @param {type: \"string\"}"
       ]
     },
     {
@@ -365,182 +298,103 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "Ee3vBgvdhgTb"
-      },
-      "source": [
-        "#### Region\n",
-        "\n",
-        "You can also change the `REGION` variable, which is used for operations\n",
-        "throughout the rest of this notebook.  Below are regions supported for Vertex AI. It is recommended that you choose the region closest to you.\n",
-        "\n",
-        "- Americas: `us-central1`\n",
-        "- Europe: `europe-west4`\n",
-        "- Asia Pacific: `asia-east1`\n",
-        "\n",
-        "You may not use a multi-regional bucket for training with Vertex AI. Not all regions provide support for all Vertex AI services.\n",
-        "\n",
-        "Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "flcwBahRhi8b"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"[your-region]\"  # @param {type: \"string\"}\n",
-        "\n",
-        "if REGION == \"[your-region]\":\n",
-        "    REGION = \"us-central1\""
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "KuNRbXkIijp6"
+        "id": "sBCra4QMA2wR"
       },
       "source": [
         "### Authenticate your Google Cloud account\n",
         "\n",
-        "**If you are using Vertex AI Workbench Notebooks**, your environment is already\n",
-        "authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "f40aa139740f"
-      },
-      "source": [
-        "**If you are using Colab**, run the cell below and follow the instructions\n",
-        "when prompted to authenticate your account via oAuth.\n",
+        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
         "\n",
-        "**Otherwise**, follow these steps:\n",
+        "**1. Vertex AI Workbench**\n",
+        "* Do nothing as you are already authenticated.\n",
         "\n",
-        "1. In the Cloud Console, go to the [**Create service account key**\n",
-        "   page](https://console.cloud.google.com/apis/credentials/serviceaccountkey).\n",
-        "\n",
-        "2. Click **Create service account**.\n",
-        "\n",
-        "3. In the **Service account name** field, enter a name, and\n",
-        "   click **Create**.\n",
-        "\n",
-        "4. In the **Grant this service account access to project** section, click the **Role** drop-down list. Type \"Vertex AI\"\n",
-        "into the filter box, and select\n",
-        "   **Vertex AI Administrator**. Type \"Storage Object Admin\" into the filter box, and select **Storage Object Admin**.\n",
-        "\n",
-        "5. Click *Create*. A JSON file that contains your key downloads to your\n",
-        "local environment.\n",
-        "\n",
-        "6. Enter the path to your service account key as the\n",
-        "`GOOGLE_APPLICATION_CREDENTIALS` variable in the cell below and run the cell."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "P9vQxUzfirCV"
-      },
-      "source": [
-        "# The Google Cloud Notebook product has specific requirements\n",
-        "import os\n",
-        "import sys\n",
-        "\n",
-        "IS_GOOGLE_CLOUD_NOTEBOOK = os.path.exists(\"/opt/deeplearning/metadata/env_version\")\n",
-        "\n",
-        "IS_COLAB = \"google.colab\" in sys.modules\n",
-        "# If on Google Cloud Notebooks, then don't execute this code\n",
-        "if not IS_GOOGLE_CLOUD_NOTEBOOK:\n",
-        "    if IS_COLAB:\n",
-        "        from google.colab import auth as google_auth\n",
-        "\n",
-        "        google_auth.authenticate_user()\n",
-        "\n",
-        "    # If you are running this notebook locally, replace the string below with the\n",
-        "    # path to your service account key and run this cell to authenticate your GCP\n",
-        "    # account.\n",
-        "    elif not os.getenv(\"IS_TESTING\"):\n",
-        "        %env GOOGLE_APPLICATION_CREDENTIALS"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "OG3dstAuVtgz"
-      },
-      "source": [
-        "### Create Cloud Storage bucket"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "TKVkz5uzV07p"
-      },
-      "source": [
-        "A Cloud Storage buckets will be used store your training code output (including TensorBoard logs). The bucket must be regional that is, not multi-region or dual-region, and the following resources must be in same region:\n",
-        "\n",
-        "* the Cloud Storage bucket\n",
-        "* the Vertex AI training job\n",
-        "* the Vertex AI TensorBoard instance"
+        "**2. Local JupyterLab instance, uncomment and run:**"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "_BzfwD8EV7Fw"
+        "id": "254614fa0c46"
       },
       "outputs": [],
       "source": [
-        "BUCKET_URI = \"gs://[your-bucket-name]\"  # @param {type:\"string\"}\n",
-        "\n",
-        "if BUCKET_URI == \"\" or BUCKET_URI is None or BUCKET_URI == \"gs://[your-bucket-name]\":\n",
-        "    BUCKET_URI = \"gs://\" + PROJECT_ID + \"aip-\" + UUID"
+        "# ! gcloud auth login"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "iVe3Y5UGWDoC"
+        "id": "ef21552ccea8"
       },
       "source": [
-        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket. The created bucket will be deleted in the cleaning up section in the end. "
+        "**3. Colab, uncomment and run:**"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "M0krsPd9WJYl"
+        "id": "603adbbf0532"
+      },
+      "outputs": [],
+      "source": [
+        "# from google.colab import auth\n",
+        "# auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f6b2ccc891ed"
+      },
+      "source": [
+        "**4. Service account or other**\n",
+        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zgPO1eR3CYjk"
+      },
+      "source": [
+        "### Create a Cloud Storage bucket\n",
+        "\n",
+        "Create a storage bucket to store intermediate artifacts such as datasets."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "MzGDU7TWdts_"
+      },
+      "outputs": [],
+      "source": [
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-EcIXiGsCePi"
+      },
+      "source": [
+        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NIq7R4HZCfIc"
       },
       "outputs": [],
       "source": [
         "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "oXo9TNzQWQBS"
-      },
-      "source": [
-        "Finally, validate access to your Cloud Storage bucket by examining its contents:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "BD-n41FzWY0G"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil ls -al {BUCKET_URI}"
-      ]
-    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -558,6 +412,8 @@
       },
       "outputs": [],
       "source": [
+        "import os\n",
+        "\n",
         "import google.cloud.aiplatform as aiplatform"
       ]
     },
@@ -603,7 +459,11 @@
       },
       "outputs": [],
       "source": [
-        "!gcloud services enable artifactregistry.googleapis.com"
+        "! gcloud services enable artifactregistry.googleapis.com\n",
+        "\n",
+        "if os.getenv(\"IS_TESTING\"):\n",
+        "    ! sudo apt-get update --yes && sudo apt-get --only-upgrade --yes install google-cloud-sdk-cloud-run-proxy google-cloud-sdk-harbourbridge google-cloud-sdk-cbt google-cloud-sdk-gke-gcloud-auth-plugin google-cloud-sdk-kpt google-cloud-sdk-local-extract google-cloud-sdk-minikube google-cloud-sdk-app-engine-java google-cloud-sdk-app-engine-go google-cloud-sdk-app-engine-python google-cloud-sdk-spanner-emulator google-cloud-sdk-bigtable-emulator google-cloud-sdk-nomos google-cloud-sdk-package-go-module google-cloud-sdk-firestore-emulator kubectl google-cloud-sdk-datastore-emulator google-cloud-sdk-app-engine-python-extras google-cloud-sdk-cloud-build-local google-cloud-sdk-kubectl-oidc google-cloud-sdk-anthos-auth google-cloud-sdk-app-engine-grpc google-cloud-sdk-pubsub-emulator google-cloud-sdk-datalab google-cloud-sdk-skaffold google-cloud-sdk google-cloud-sdk-terraform-tools google-cloud-sdk-config-connector\n",
+        "    ! gcloud components update --quiet\n"
       ]
     },
     {
@@ -626,14 +486,7 @@
       },
       "outputs": [],
       "source": [
-        "DOCKER_REPOSITORY = \"[your-docker-repository-name]\"  # @param {type:\"string\"}\n",
-        "\n",
-        "if (\n",
-        "    DOCKER_REPOSITORY == \"\"\n",
-        "    or DOCKER_REPOSITORY is None\n",
-        "    or DOCKER_REPOSITORY == \"[your-docker-repository-name]\"\n",
-        "):\n",
-        "    DOCKER_REPOSITORY = \"tb-docker-repo-\" + PROJECT_ID + \"-\" + UUID\n",
+        "DOCKER_REPOSITORY = \"my-docker-repo-unique\"  # @param {type:\"string\"}\n",
         "\n",
         "print(\"Docker repository to create:\", DOCKER_REPOSITORY)"
       ]
@@ -812,6 +665,7 @@
         "\n",
         "# Specifies base image and tag\n",
         "FROM us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-8:latest\n",
+        "RUN pip install tensorflow-datasets\n",
         "WORKDIR /root\n",
         "\n",
         "\n",
@@ -1016,6 +870,8 @@
       },
       "outputs": [],
       "source": [
+        "import os\n",
+        "\n",
         "# Delete docker repository.\n",
         "! gcloud artifacts repositories delete $DOCKER_REPOSITORY --project {PROJECT_ID} --location {REGION} --quiet\n",
         "\n",

notebooks/official/tensorboard/tensorboard_custom_training_with_prebuilt_container.ipynb

@@ -24,6 +24,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "eBEO2w9My9py"
@@ -50,10 +51,11 @@
         "    </a>\n",
         "  </td>\n",
         "</table>\n",
-        "<br/><br/><br/>"
+        "<br/><br/><br/>\n"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "VL7XCFV7yCBU"
@@ -95,6 +97,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "UE8vLw7SlpwE"
@@ -118,6 +121,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "fbF2lF8rlp3I"
@@ -129,6 +133,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "yajq2millpnu"
@@ -145,39 +150,7 @@
       ]
     },
     {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "lAcDYZfslpeF"
-      },
-      "source": [
-        "### Set up your local development environment\n",
-        "\n",
-        "**If you are using Colab or Vertex AI Workbench**, your environment already meets all the requirements to run this notebook. You can skip this step.\n",
-        "\n",
-        "Otherwise, make sure your environment meets this notebook's requirements. You need the following:\n",
-        "\n",
-        "- The Cloud Storage SDK\n",
-        "- Git\n",
-        "- Python 3\n",
-        "- virtualenv\n",
-        "- Jupyter notebook running in a virtual environment with Python 3\n",
-        "\n",
-        "The Cloud Storage guide to [Setting up a Python development environment](https://cloud.google.com/python/setup) and the [Jupyter installation guide](https://jupyter.org/install) provide detailed instructions for meeting these requirements. The following steps provide a condensed set of instructions:\n",
-        "\n",
-        "1. [Install and initialize the SDK](https://cloud.google.com/sdk/docs/).\n",
-        "\n",
-        "2. [Install Python 3](https://cloud.google.com/python/setup#installing_python).\n",
-        "\n",
-        "3. [Install virtualenv](https://cloud.google.com/python/setup#installing_and_using_virtualenv) and create a virtual environment that uses Python 3.  Activate the virtual environment.\n",
-        "\n",
-        "4. To install Jupyter, run `pip3 install jupyter` on the command-line in a terminal shell.\n",
-        "\n",
-        "5. To launch Jupyter, run `jupyter notebook` on the command-line in a terminal shell.\n",
-        "\n",
-        "6. Open this notebook in the Jupyter Notebook Dashboard."
-      ]
-    },
-    {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "s3moH5AexXpk"
@@ -196,343 +169,221 @@
       },
       "outputs": [],
       "source": [
-        "import os\n",
-        "\n",
-        "# The Vertex AI Workbench Notebook product has specific requirements\n",
-        "IS_WORKBENCH_NOTEBOOK = os.getenv(\"DL_ANACONDA_HOME\")\n",
-        "\n",
-        "# Vertex AI Notebook requires dependencies to be installed with '--user'\n",
-        "USER_FLAG = \"\"\n",
-        "if IS_WORKBENCH_NOTEBOOK:\n",
-        "    USER_FLAG = \"--user\"\n",
-        "\n",
-        "! pip3 install google-cloud-aiplatform {USER_FLAG} -q\n",
-        "\n",
-        "# Automatically restart kernel after installs\n",
-        "\n",
-        "if not os.getenv(\"IS_TESTING\"):\n",
-        "    # Automatically restart kernel after installs\n",
-        "    import IPython\n",
-        "\n",
-        "    app = IPython.Application.instance()\n",
-        "    app.kernel.do_shutdown(True)"
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform "
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
-        "id": "cczxMMkYK9a4"
+        "id": "restart"
       },
       "source": [
-        "## Before you begin"
+        "### Colab only: Uncomment the following cell to restart the kernel.\n"
       ]
     },
     {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bzPxhxS5lugp"
+      },
+      "outputs": [],
+      "source": [
+        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
+        "# import IPython\n",
+        "\n",
+        "# app = IPython.Application.instance()\n",
+        "# app.kernel.do_shutdown(True)"
+      ]
+    },
+    {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
-        "id": "64BZ1jL5GEi0"
+        "id": "d2qpIurSjmpT"
       },
       "source": [
+        "## Before you begin\n",
+        "\n",
         "### Set up your Google Cloud project\n",
         "\n",
         "**The following steps are required, regardless of your notebook environment.**\n",
         "\n",
         "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
         "\n",
-        "2. [Make sure that billing is enabled for your project.](https://cloud.google.com/billing/docs/how-to/modify-project)\n",
+        "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
         "\n",
-        "3. [Enable the Vertex AI APIs and Cloud Storage.](https://console.cloud.google.com/flows/enableapi?apiid=ml.googleapis.com,storage-component.googleapis.com)\n",
+        "3. [Enable the following APIs: Vertex AI API, Cloud Resource Manager API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,cloudresourcemanager.googleapis.com).\n",
         "\n",
-        "4. If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
-        "\n",
-        "\n",
-        "5. Enter your project ID in the cell below. Then run the  cell to make sure the\n",
-        "Cloud SDK uses the right project for all the commands in this notebook.\n",
-        "\n",
-        "Note: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$`."
+        "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
-        "id": "UIl_kn1pGH_T"
+        "id": "project_id"
       },
       "source": [
-        "### Set your project ID"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "Maw6BYbPA0kn"
-      },
-      "source": [
-        "**If you don't know your project ID**, you may be able to get your project ID using `gcloud`.\n"
+        "#### Set your project ID\n",
+        "\n",
+        "**If you don't know your project ID**, try the following:\n",
+        "* Run `gcloud config list`.\n",
+        "* Run `gcloud projects list`.\n",
+        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "sw4J6j5tBZWp"
+        "id": "wsePm9c4jmpT"
       },
       "outputs": [],
       "source": [
-        "PROJECT_ID = \"\"\n",
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
         "\n",
-        "import os\n",
-        "\n",
-        "# Get your Google Cloud project ID from gcloud\n",
-        "if not os.getenv(\"IS_TESTING\"):\n",
-        "    shell_output = !gcloud config list --format 'value(core.project)' 2>/dev/null\n",
-        "    PROJECT_ID = shell_output[0]\n",
-        "    print(\"Project ID: \", PROJECT_ID)"
+        "# Set the project id\n",
+        "! gcloud config set project {PROJECT_ID}"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
-        "id": "EJj3d9GxBd3b"
+        "id": "a54f9d7c1876"
       },
       "source": [
-        "Otherwise, set your project ID here."
+        "#### Region\n",
+        "\n",
+        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "75C2px4XtS5l"
+        "id": "3aaadaaf9b30"
       },
       "outputs": [],
       "source": [
-        "if PROJECT_ID == \"\" or PROJECT_ID is None:\n",
-        "    PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}"
+        "REGION = \"us-central1\"  # @param {type: \"string\"}"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
-        "id": "Z36ycwA4IjYC"
+        "id": "5c0404984792"
       },
       "source": [
-        "### Timestamp"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "MngMavafIrQa"
-      },
-      "source": [
-        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a timestamp for each instance session, and append it onto the name of resources you create in this tutorial."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "uddA3D7yIn-L"
-      },
-      "outputs": [],
-      "source": [
-        "from datetime import datetime\n",
+        "### Authenticate your Google Cloud account\n",
         "\n",
-        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "6yPUOueCF3pI"
-      },
-      "source": [
-        "### Set your region"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "gzEibeHGF1jb"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"[your-region]\"  # @param {type: \"string\"}\n",
+        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
         "\n",
-        "if REGION == \"[your-region]\":\n",
-        "    REGION = \"us-central1\""
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "jetPysDMtV1-"
-      },
-      "source": [
-        "### Login to your Google Cloud account"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "IK8eOR8tt6TI"
-      },
-      "outputs": [],
-      "source": [
-        "# The Google Cloud Notebook product has specific requirements\n",
-        "import os\n",
-        "import sys\n",
+        "**1. Vertex AI Workbench**\n",
+        "* Do nothing as you are already authenticated.\n",
         "\n",
-        "IS_GOOGLE_CLOUD_NOTEBOOK = os.path.exists(\"/opt/deeplearning/metadata/env_version\")\n",
+        "**2. Local JupyterLab instance, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Nt8cEM2GjmpU"
+      },
+      "outputs": [],
+      "source": [
+        "# ! gcloud auth login"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "XUSL_JcpjmpU"
+      },
+      "source": [
+        "**3. Colab, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "_2zemfGvjmpU"
+      },
+      "outputs": [],
+      "source": [
+        "# from google.colab import auth\n",
+        "# auth.authenticate_user()"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "TCPJ38n7jmpU"
+      },
+      "source": [
+        "**4. Service account or other**\n",
+        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bucket:custom"
+      },
+      "source": [
+        "### Create a Cloud Storage bucket\n",
         "\n",
-        "IS_COLAB = \"google.colab\" in sys.modules\n",
-        "# If on Google Cloud Notebooks, then don't execute this code\n",
-        "if not IS_GOOGLE_CLOUD_NOTEBOOK:\n",
-        "    if IS_COLAB:\n",
-        "        from google.colab import auth as google_auth\n",
+        "Create a storage bucket to store intermediate artifacts such as datasets.\n",
         "\n",
-        "        google_auth.authenticate_user()\n",
-        "\n",
-        "    # If you are running this notebook locally, replace the string below with the\n",
-        "    # path to your service account key and run this cell to authenticate your GCP\n",
-        "    # account.\n",
-        "    elif not os.getenv(\"IS_TESTING\"):\n",
-        "        %env GOOGLE_APPLICATION_CREDENTIALS ''"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "tU1OEt8ibhLS"
-      },
-      "source": [
-        "###Create Cloud Storage bucket\n",
-        "A Cloud Storage bucket will be used to a) store your training code distribution (details below), and b) the outputs (including TensorBoard logs) your training code generates. The bucket must be regional that is, not multi-region or dual-region, and the following resources must be in same region:\n",
-        "\n",
-        "* the Cloud Storage bucket\n",
-        "* the Vertex AI training job\n",
-        "* the Vertex AI TensorBoard instance"
+        "When you submit a training job using the Cloud SDK, you upload a Python package\n",
+        "containing your training code to a Cloud Storage bucket. Vertex AI runs\n",
+        "the code from this package. In this tutorial, Vertex AI also saves the\n",
+        "trained model that results from your job in the same bucket. Using this model artifact, you can then\n",
+        "create Vertex AI Model resource and use for prediction."
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "HYu2_qM9b3Cn"
+        "id": "bucket"
       },
       "outputs": [],
       "source": [
-        "BUCKET_URI = \"gs://[your-bucket-name]\"  # @param {type:\"string\"}\n",
-        "\n",
-        "if BUCKET_URI == \"\" or BUCKET_URI is None or BUCKET_URI == \"gs://[your-bucket-name]\":\n",
-        "    BUCKET_URI = \"gs://\" + PROJECT_ID + \"aip-\" + TIMESTAMP"
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
-        "id": "qo1KHfJ2b83V"
+        "id": "create_bucket"
       },
       "source": [
-        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket. The created bucket will be deleted in the cleaning up section in the end. "
+        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "imVC-rXxb8A1"
+        "id": "Oz8J0vmSlugt"
       },
       "outputs": [],
       "source": [
-        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "1VU7ukLOcCa1"
-      },
-      "source": [
-        "Finally, validate access to your Cloud Storage bucket by examining its contents:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "_wil_Y9lcHhZ"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil ls -al {BUCKET_URI}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "lx8UKHsccMDe"
-      },
-      "source": [
-        "Set up the GCS paths for traing code and outputs."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "phyaIjuHcPdH"
-      },
-      "outputs": [],
-      "source": [
-        "GCS_BUCKET_TRAINING = BUCKET_URI + \"/training/\"\n",
-        "GCS_BUCKET_OUTPUT = BUCKET_URI + \"/output/\""
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "KIPcg_Xhwvsn"
-      },
-      "source": [
-        "### Import aiplatform"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "p4w8c1pHw2Yt"
-      },
-      "outputs": [],
-      "source": [
-        "import google.cloud.aiplatform as aiplatform"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "cVmlv9sRbCSs"
-      },
-      "source": [
-        "### Initialize Vertex AI SDK for Python\n",
-        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "JmEsq1fda_1N"
-      },
-      "outputs": [],
-      "source": [
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
+        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "hwXxa4Qgnh4Y"
@@ -542,6 +393,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "7qXFUiHLoFRw"
@@ -569,6 +421,9 @@
       },
       "outputs": [],
       "source": [
+        "import sys\n",
+        "\n",
+        "IS_COLAB = \"google.colab\" in sys.modules\n",
         "if (\n",
         "    SERVICE_ACCOUNT == \"\"\n",
         "    or SERVICE_ACCOUNT is None\n",
@@ -616,6 +471,50 @@
       ]
     },
     {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "KIPcg_Xhwvsn"
+      },
+      "source": [
+        "### Import libraries"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "p4w8c1pHw2Yt"
+      },
+      "outputs": [],
+      "source": [
+        "import google.cloud.aiplatform as aiplatform"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cVmlv9sRbCSs"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
+        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "JmEsq1fda_1N"
+      },
+      "outputs": [],
+      "source": [
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
+      ]
+    },
+    {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "dR2mOCllvlqN"
@@ -664,6 +563,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "vjdLIqYyDZFS"
@@ -790,6 +690,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "L067Jw_QFcZ3"
@@ -814,6 +715,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "a7Gt9WPFG9V7"
@@ -830,10 +732,12 @@
       },
       "outputs": [],
       "source": [
+        "GCS_BUCKET_TRAINING = f\"{BUCKET_URI}/data/\"\n",
         "! gsutil cp dist/hello-custom-training-3.0.tar.gz {GCS_BUCKET_TRAINING}"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "IaQjIPvuKLwW"
@@ -843,6 +747,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "svUGBOow_Obj"
@@ -866,7 +771,7 @@
         "    or TENSORBOARD_NAME is None\n",
         "    or TENSORBOARD_NAME == \"[your-tensorboard-name]\"\n",
         "):\n",
-        "    TENSORBOARD_NAME = PROJECT_ID + \"-tb-\" + TIMESTAMP\n",
+        "    TENSORBOARD_NAME = PROJECT_ID + \"-tb\"\n",
         "\n",
         "tensorboard = aiplatform.Tensorboard.create(\n",
         "    display_name=TENSORBOARD_NAME, project=PROJECT_ID, location=REGION\n",
@@ -876,6 +781,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "mudxBDal_a_k"
@@ -892,8 +798,9 @@
       },
       "outputs": [],
       "source": [
-        "JOB_NAME = \"tensorboard-example-job-{}\".format(TIMESTAMP)\n",
-        "BASE_OUTPUT_DIR = \"{}{}\".format(GCS_BUCKET_OUTPUT, JOB_NAME)\n",
+        "JOB_NAME = \"tensorboard-example-job\"\n",
+        "GCS_BUCKET_OUTPUT = BUCKET_URI\n",
+        "BASE_OUTPUT_DIR = \"{}/{}\".format(GCS_BUCKET_OUTPUT, JOB_NAME)\n",
         "\n",
         "job = aiplatform.CustomPythonPackageTrainingJob(\n",
         "    display_name=JOB_NAME,\n",
@@ -914,6 +821,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "VfMsn_RnEtnj"
@@ -923,6 +831,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "TFEriiywMZga"

notebooks/official/tensorboard/tensorboard_profiler_custom_training_with_prebuilt_container.ipynb

@@ -24,6 +24,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "l2mMvIUG9meX"
@@ -91,6 +92,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "zfXf0r-K81Y-"
@@ -102,6 +104,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "I3KFLvpq87rs"
@@ -122,6 +125,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "ze4-nDLfK4pw"
@@ -144,6 +148,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "aUw6ibN-n5Za"
@@ -168,6 +173,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "LgFWLeJfoGQu"
@@ -189,6 +195,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "8ckyxpX_oSzD"
@@ -217,6 +224,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "mSQjVQmMosMl"
@@ -239,6 +247,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "IfJRIMBpo5Pg"
@@ -246,25 +255,11 @@
       "source": [
         "### Authenticate your Google Cloud account\n",
         "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "acFN0s3So9-Y"
-      },
-      "source": [
+        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
+        "\n",
         "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "dQ_mNwuapE5T"
-      },
-      "source": [
+        "* Do nothing as you are already authenticated.\n",
+        "\n",
         "**2. Local JupyterLab instance, uncomment and run:**"
       ]
     },
@@ -280,6 +275,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "h-MuVI_ypJfw"
@@ -301,6 +297,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "3ivZkPUjpaFz"
@@ -322,6 +319,35 @@
         "SERVICE_ACCOUNT = \"[your-service-account]\"  # @param {type:\"string\"}"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "autoset_service_account"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "IS_COLAB = \"google.colab\" in sys.modules\n",
+        "if (\n",
+        "    SERVICE_ACCOUNT == \"\"\n",
+        "    or SERVICE_ACCOUNT is None\n",
+        "    or SERVICE_ACCOUNT == \"[your-service-account]\"\n",
+        "):\n",
+        "    # Get your service account from gcloud\n",
+        "    if not IS_COLAB:\n",
+        "        shell_output = !gcloud auth list 2>/dev/null\n",
+        "        SERVICE_ACCOUNT = shell_output[2].replace(\"*\", \"\").strip()\n",
+        "\n",
+        "    if IS_COLAB:\n",
+        "        shell_output = ! gcloud projects describe  $PROJECT_ID\n",
+        "        project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "        SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "\n",
+        "    print(\"Service Account:\", SERVICE_ACCOUNT)"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -347,6 +373,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "OKtKGmr9pfr6"
@@ -365,10 +392,11 @@
       },
       "outputs": [],
       "source": [
-        "BUCKET_URI = \"gs://your-bucket-name-unique\"  # @param {type:\"string\"}"
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "GOaOsIjxp0oB"
@@ -385,10 +413,11 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
+        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "ankcS-vtp7Wv"
@@ -405,10 +434,13 @@
       },
       "outputs": [],
       "source": [
+        "import os\n",
+        "\n",
         "from google.cloud import aiplatform"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "OMrAJ8RGqBQu"
@@ -431,6 +463,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "-ayTbNdi62_t"
@@ -442,6 +475,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "9c3QrDTZdaxk"
@@ -462,6 +496,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "vJrWKK0mY7H7"
@@ -507,6 +542,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "syw3GabNGgJz"
@@ -573,6 +609,7 @@
         "\n",
         "REQUIRED_PACKAGES = [\n",
         "    'google-cloud-aiplatform[cloud_profiler]>=1.20.0',\n",
+        "    'protobuf==3.20.2',\n",
         "]\n",
         "\n",
         "setup(\n",
@@ -587,6 +624,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "hyAwgsoQmaYI"
@@ -706,6 +744,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "ihYFahRAr6sj"
@@ -732,6 +771,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "k4e6OYmimqTR"
@@ -779,6 +819,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "51hKGTbU32Eg"
@@ -810,6 +851,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "JkEe2Nb_85UD"
@@ -823,6 +865,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "TpV-iwP9qw9c"
@@ -849,7 +892,7 @@
         "job.delete()\n",
         "tensorboard.delete()\n",
         "\n",
-        "if delete_bucket and \"BUCKET_URI\" in globals():\n",
+        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
         "    ! gsutil -m rm -r $BUCKET_URI"
       ]
     }

notebooks/official/training/multi_node_ddp_gloo_vertex_training_with_custom_container.ipynb

@@ -980,7 +980,11 @@
       },
       "outputs": [],
       "source": [
-        "! gcloud services enable artifactregistry.googleapis.com"
+        "! gcloud services enable artifactregistry.googleapis.com\n",
+        "\n",
+        "if os.getenv(\"IS_TESTING\"):\n",
+        "    ! sudo apt-get update --yes && sudo apt-get --only-upgrade --yes install google-cloud-sdk-cloud-run-proxy google-cloud-sdk-harbourbridge google-cloud-sdk-cbt google-cloud-sdk-gke-gcloud-auth-plugin google-cloud-sdk-kpt google-cloud-sdk-local-extract google-cloud-sdk-minikube google-cloud-sdk-app-engine-java google-cloud-sdk-app-engine-go google-cloud-sdk-app-engine-python google-cloud-sdk-spanner-emulator google-cloud-sdk-bigtable-emulator google-cloud-sdk-nomos google-cloud-sdk-package-go-module google-cloud-sdk-firestore-emulator kubectl google-cloud-sdk-datastore-emulator google-cloud-sdk-app-engine-python-extras google-cloud-sdk-cloud-build-local google-cloud-sdk-kubectl-oidc google-cloud-sdk-anthos-auth google-cloud-sdk-app-engine-grpc google-cloud-sdk-pubsub-emulator google-cloud-sdk-datalab google-cloud-sdk-skaffold google-cloud-sdk google-cloud-sdk-terraform-tools google-cloud-sdk-config-connector\n",
+        "    ! gcloud components update --quiet"
       ]
     },
     {

notebooks/official/training/pytorch_gcs_data_training.ipynb

@@ -156,7 +156,7 @@
         "\n",
         "google-cloud-aiplatform\n",
         "google-cloud-storage\n",
-        "torch"
+        "torch==1.12"
       ]
     },
     {
@@ -239,7 +239,7 @@
       },
       "outputs": [],
       "source": [
-        "PROJECT_ID = \"your-project-id\"  # @param\n",
+        "PROJECT_ID = \"[your-project-id]\"  # @param\n",
         "\n",
         "# Set the project id\n",
         "! gcloud config set project {PROJECT_ID}"
@@ -358,7 +358,7 @@
       },
       "outputs": [],
       "source": [
-        "BUCKET_URI = \"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}\n",
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}\n",
         "BUCKET_PREFIX = \"pytorch-on-gcs\"  # @param"
       ]
     },
@@ -379,7 +379,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
+        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
       ]
     },
     {

notebooks/official/workbench/ads_targetting/training-multi-class-classification-model-for-ads-targeting-usecase.ipynb

@@ -789,7 +789,7 @@
       "source": [
         "import numpy as np\n",
         "\n",
-        "X_train = np.asarray(X_train, dtype=np.float)\n",
+        "X_train = np.asarray(X_train, dtype=np.float32)\n",
         "\n",
         "history = model.fit(X_train, y_train_categ, epochs=50, verbose=1)"
       ]
@@ -811,7 +811,7 @@
       },
       "outputs": [],
       "source": [
-        "X_test = np.asarray(X_test, dtype=np.float)\n",
+        "X_test = np.asarray(X_test, dtype=np.float32)\n",
         "\n",
         "test_results = model.evaluate(X_test, y_test_categ, verbose=1)\n",
         "print(f\"Test results - Loss: {test_results}\")"