fix: conflict

community-content/Train_tabular_models_with_many_frameworks_and_import_to_Vertex_AI_using_Pipelines/Train_tabular_classification_model_using_XGBoost_and_import_to_Vertex_AI/pipeline.py

@@ -9,7 +9,7 @@ binarize_column_using_Pandas_on_CSV_data_op = components.load_component_from_url
 split_rows_into_subsets_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/dataset_manipulation/Split_rows_into_subsets/in_CSV/component.yaml")
 train_XGBoost_model_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Train/component.yaml")
 xgboost_predict_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Predict/component.yaml")
-upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
+upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/d5c9918850a6cc70004c4269dae066cfe2e664eb/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
 deploy_model_to_endpoint_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Deploy_to_endpoint/component.yaml")
 
 # %% Pipeline definition

community-content/Train_tabular_models_with_many_frameworks_and_import_to_Vertex_AI_using_Pipelines/Train_tabular_classification_model_using_all_frameworks_and_import_to_Vertex_AI/pipeline.py

@@ -23,7 +23,7 @@ upload_PyTorch_model_archive_to_Google_Cloud_Vertex_AI_op = components.load_comp
 # XGBoost
 train_XGBoost_model_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Train/component.yaml")
 xgboost_predict_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Predict/component.yaml")
-upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
+upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/d5c9918850a6cc70004c4269dae066cfe2e664eb/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
 
 # Scikit-learn
 #train_linear_regression_model_using_scikit_learn_from_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/1f5cf6e06409b704064b2086c0a705e4e6b4fcde/community-content/pipeline_components/ML_frameworks/Scikit_learn/Train_linear_regression_model/from_CSV/component.yaml")

community-content/Train_tabular_models_with_many_frameworks_and_import_to_Vertex_AI_using_Pipelines/Train_tabular_regression_model_using_XGBoost_and_import_to_Vertex_AI/pipeline.py

@@ -8,7 +8,7 @@ fill_all_missing_values_using_Pandas_on_CSV_data_op = components.load_component_
 split_rows_into_subsets_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/dataset_manipulation/Split_rows_into_subsets/in_CSV/component.yaml")
 train_XGBoost_model_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Train/component.yaml")
 xgboost_predict_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Predict/component.yaml")
-upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
+upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/d5c9918850a6cc70004c4269dae066cfe2e664eb/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
 deploy_model_to_endpoint_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Deploy_to_endpoint/component.yaml")
 
 # %% Pipeline definition

community-content/Train_tabular_models_with_many_frameworks_and_import_to_Vertex_AI_using_Pipelines/Train_tabular_regression_model_using_all_frameworks_and_import_to_Vertex_AI/pipeline.py

@@ -22,7 +22,7 @@ upload_PyTorch_model_archive_to_Google_Cloud_Vertex_AI_op = components.load_comp
 # XGBoost
 train_XGBoost_model_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Train/component.yaml")
 xgboost_predict_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Predict/component.yaml")
-upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
+upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/d5c9918850a6cc70004c4269dae066cfe2e664eb/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
 
 # Scikit-learn
 train_linear_regression_model_using_scikit_learn_from_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/1f5cf6e06409b704064b2086c0a705e4e6b4fcde/community-content/pipeline_components/ML_frameworks/Scikit_learn/Train_linear_regression_model/from_CSV/component.yaml")

community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml

@@ -64,8 +64,8 @@ implementation:
           labels["component-source"] = "github-com-ark-kun-pipeline-components"
 
           # The serving container decides the model type based on the model file extension.
-          # So we need to rename the mode file (e.g. /tmp/inputs/model/data) to *.pkl
-          _, renamed_model_path = tempfile.mkstemp(suffix=".pkl")
+          # So we need to rename the mode file (e.g. /tmp/inputs/model/data) to *.bst
+          _, renamed_model_path = tempfile.mkstemp(suffix=".bst")
           shutil.copyfile(src=model_path, dst=renamed_model_path)
 
           model = aiplatform.Model.upload_xgboost_model_file(

community-content/pipeline_components/image_ml_model_training/load_image_classification_model/component.yaml

@@ -87,7 +87,7 @@ outputs:
 - {name: image_size_path, type: HeightWidth}
 implementation:
   container:
-    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.1
+    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.2
     # command is a list of strings (command-line arguments).
     # The YAML language has two syntaxes for lists and you can use either of them.
     # Here we use the "flow syntax" - comma-separated strings inside square brackets.
@@ -109,4 +109,4 @@ implementation:
       {inputValue: l2_regularization_penalty},
       --image-size-path,
       {outputPath: image_size_path},
-    ]
+    ]

community-content/pipeline_components/image_ml_model_training/preprocess_image_data/component.yaml

@@ -34,7 +34,7 @@ outputs:
     path for the validation data,'}
 implementation:
   container:
-    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.1
+    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.2
     # command is a list of strings (command-line arguments).
     # The YAML language has two syntaxes for lists and you can use either of them.
     # Here we use the "flow syntax" - comma-separated strings inside square brackets.

community-content/pipeline_components/image_ml_model_training/train_image_classification_model/component.yaml

@@ -55,7 +55,7 @@ outputs:
     for the saved model,'}
 implementation:
   container:
-    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.1
+    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.2
     # command is a list of strings (command-line arguments).
     # The YAML language has two syntaxes for lists and you can use either of them.
     # Here we use the "flow syntax" - comma-separated strings inside square brackets.

community-content/pipeline_components/image_ml_model_training/transcode_tfrecord_image_dataset_from_csv/component.yaml

@@ -20,7 +20,7 @@ outputs:
     path for the TFRecord image data}
 implementation:
   container:
-    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.1
+    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.2
     # command is a list of strings (command-line arguments).
     # The YAML language has two syntaxes for lists and you can use either of them.
     # Here we use the "flow syntax" - comma-separated strings inside square brackets.

community-content/pipeline_components/image_ml_model_training/transcode_tfrecord_image_dataset_from_jsonl/component.yaml

@@ -22,7 +22,7 @@ outputs:
     path for the TFRecord image data}
 implementation:
   container:
-    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.1
+    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.2
     # command is a list of strings (command-line arguments).
     # The YAML language has two syntaxes for lists and you can use either of them.
     # Here we use the "flow syntax" - comma-separated strings inside square brackets.

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

@@ -0,0 +1,59 @@
+# Dockerfile for serving dockers with timm.
+#
+# To build:
+# docker build -f model_oss/timm/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 pytorch/torchserve:0.7.0-gpu
+
+USER root
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        curl \
+        wget \
+        vim
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+ENV infer_port=7080
+ENV mng_port=7081
+ENV model_name="timm_serving"
+
+# Install timm.
+RUN python3 -m pip install --upgrade pip
+RUN python3 -m pip install timm==0.6.12
+RUN python3 -m pip install google-cloud-storage==2.9.0
+
+# Copy model artifacts.
+COPY model_oss/timm/handler.py /home/model-server/handler.py
+
+# Create torchserve configuration file.
+RUN echo \
+    "default_response_timeout=1200\n" \
+    "service_envelope=json\n" \
+    "inference_address=http://0.0.0.0:${infer_port}\n"  \
+    "management_address=http://0.0.0.0:${mng_port}" >> /home/model-server/config.properties
+
+# Expose ports.
+EXPOSE ${infer_port}
+EXPOSE ${mng_port}
+
+# Archive eager mode model artifacts and dependencies.
+# Do not set --model-file and --serialized-file because model and checkpoint will be dynamically loaded in handler.py.
+RUN torch-model-archiver \
+--model-name=${model_name} \
+--version=1.0 \
+--handler=/home/model-server/handler.py \
+--runtime=python3 \
+--export-path=/home/model-server/model-store \
+--archive-format=default \
+--force
+
+# Run Torchserve HTTP serve to respond to prediction requests.
+CMD ["torchserve", "--start", \
+     "--ts-config", "/home/model-server/config.properties", \
+     "--models", "${model_name}=${model_name}.mar", \
+     "--model-store", "/home/model-server/model-store"]

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

@@ -0,0 +1,47 @@
+# Dockerfile for basic training dockers with timm.
+#
+# To build:
+# docker build -f model_oss/timm/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}
+
+# Base on pytorch-cuda image.
+FROM pytorch/pytorch:1.13.0-cuda11.6-cudnn8-runtime
+
+# Install tools.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        curl \
+        wget \
+        vim
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Download timm source code with pinned version.
+RUN wget -q https://github.com/rwightman/pytorch-image-models/archive/refs/tags/v0.6.12.tar.gz
+RUN tar xzf v0.6.12.tar.gz
+
+# Install libraries.
+RUN pip install cloudml-hypertune==0.1.0.dev6
+
+# Switch to timm repo.
+WORKDIR /workspace/pytorch-image-models-0.6.12
+
+# NOTE: use 'sed' to modify the timm source code to
+#       make timm CheckpointSaver can work with gcsfuse.
+RUN sed -i "1 i\import shutil" timm/utils/checkpoint_saver.py
+RUN sed -i "s#os.link#shutil.copyfile#g" timm/utils/checkpoint_saver.py
+RUN sed -i "s#os.unlink#os.remove#g" timm/utils/checkpoint_saver.py
+
+# NOTE: use 'sed' to modify the timm source code to
+#       add hp training support to timm trainer.
+RUN sed -i "693 a\        if saver is not None: hpt = hypertune.HyperTune(); hpt.report_hyperparameter_tuning_metric(hyperparameter_metric_tag='top1_accuracy', metric_value=best_metric, global_step=best_epoch)" train.py
+RUN sed -i "1 i\import hypertune" train.py
+
+# Install timm from source code.
+RUN pip install -e .
+
+# https://pytorch.org/docs/stable/elastic/run.html
+ENTRYPOINT ["torchrun"]

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

@@ -0,0 +1,97 @@
+"""Custom handler for TIMM models."""
+
+import logging
+import os
+from typing import Any
+
+from google.cloud import storage
+import timm
+import torch
+from ts.torch_handler.base_handler import load_label_mapping
+from ts.torch_handler.image_classifier import ImageClassifier
+
+
+GCS_PREFIX = "gs://"
+DOWNLOAD_DIR = "/tmp/download"
+
+
+def download_gcs_file(gcs_uri: str, local_dir: str) -> str:
+  """Download a GCS file to a local directory.
+
+  Arguments:
+    gcs_uri: A string of file path on GCS.
+    local_dir: A string of local directory path.
+
+  Returns:
+    Local path to downloaded file.
+  """
+  if not gcs_uri.startswith(GCS_PREFIX):
+    raise ValueError(f"{gcs_uri} is not a GCS path starting with gs://.")
+
+  file_name = os.path.basename(gcs_uri)
+  local_file_path = os.path.join(local_dir, file_name)
+  os.makedirs(local_dir, exist_ok=True)
+  client = storage.Client()
+  with open(local_file_path, "wb") as f:
+    client.download_blob_to_file(gcs_uri, f)
+  return local_file_path
+
+
+class TimmHandler(ImageClassifier):
+  """Custom handler for TIMM models."""
+
+  def initialize(self, context: Any):
+    """Custom initialize."""
+
+    properties = context.system_properties
+    self.map_location = (
+        "cuda"
+        if torch.cuda.is_available() and properties.get("gpu_id") is not None
+        else "cpu"
+    )
+    self.device = torch.device(
+        self.map_location + ":" + str(properties.get("gpu_id"))
+        if torch.cuda.is_available() and properties.get("gpu_id") is not None
+        else self.map_location
+    )
+    self.manifest = context.manifest
+
+    # Load timm model by model name.
+    self.model_name = os.environ["MODEL_NAME"]
+    # Whether to use timm pretrained weights, MODEL_PT_PATH overrides this.
+    timm_pretrained = True if os.environ.get("TIMM_PRETRAINED") else False
+    # Load custom checkpoint, it overrides TIMM_PRETRAINED model.
+    self.model_pt_path = os.environ.get("MODEL_PT_PATH")
+    if self.model_pt_path and self.model_pt_path.startswith(GCS_PREFIX):
+      self.model_pt_path = download_gcs_file(self.model_pt_path, DOWNLOAD_DIR)
+
+    if self.model_pt_path and self.model_pt_path.endswith(".pt"):
+      logging.info(
+          "Load model with .pt in jit mode, not working for all timm models"
+          " yet."
+      )
+      self.model = self._load_torchscript_model(self.model_pt_path)
+    else:
+      logging.info("Load model with .pth in eager mode.")
+      self.model = timm.create_model(
+          self.model_name, pretrained=timm_pretrained
+      )
+      if self.model_pt_path and (
+          self.model_pt_path.endswith(".pth")
+          or self.model_pt_path.endswith(".pth.tar")
+      ):
+        checkpoint = torch.load(self.model_pt_path, map_location=self.device)
+        state_dict = checkpoint["state_dict"]
+        self.model.load_state_dict(state_dict)
+      self.model.to(self.device)
+    self.model.eval()
+
+    mapping_file_path = os.environ.get("INDEX_TO_NAME_FILE")
+    if mapping_file_path:
+      if mapping_file_path.startswith(GCS_PREFIX):
+        mapping_file_path = download_gcs_file(mapping_file_path, DOWNLOAD_DIR)
+      self.mapping = load_label_mapping(mapping_file_path)
+
+    self.initialized = True
+
+  # NOTE: Preprocess and postprocess are implemented by ImageClassifier.

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

@@ -0,0 +1,73 @@
+# Dockerfile for vLLM serving.
+#
+# To build:
+# docker build -f model_oss/vllm/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}
+
+# The base image is required by vllm
+# https://vllm.readthedocs.io/en/latest/getting_started/installation.html
+FROM nvcr.io/nvidia/pytorch:22.12-py3
+
+USER root
+
+# Install tools.
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt-get update
+RUN apt-get install -y --no-install-recommends apt-utils
+RUN apt-get install -y --no-install-recommends curl
+RUN apt-get install -y --no-install-recommends wget
+RUN apt-get install -y --no-install-recommends git
+RUN apt-get install -y --no-install-recommends jq
+RUN apt-get install -y --no-install-recommends gnupg
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Install libraries.
+ENV PIP_ROOT_USER_ACTION=ignore
+RUN python3 -m pip install --upgrade pip
+RUN pip install google-cloud-storage==2.7.0
+RUN pip install absl-py==1.4.0
+
+# Install pytorch
+RUN pip install --upgrade torch==2.0.1
+
+# Install vllm deps.
+RUN pip install xformers==0.0.20
+RUN pip install ninja==1.11.1
+RUN pip install psutil==5.9.5
+RUN pip install ray==2.6.2
+RUN pip install sentencepiece==0.1.99
+RUN pip install fastapi==0.100.1
+RUN pip install uvicorn==0.23.2
+RUN pip install pydantic==1.10.12
+
+# Install transformers from source.
+WORKDIR /workspace
+RUN git clone https://github.com/huggingface/transformers.git
+WORKDIR transformers
+# Pin the commit to add-code-llama at 08/25/2023
+RUN git reset --hard 015f8e110d270a0ad42de4ae5b98198d69eb1964
+RUN pip install -e .
+WORKDIR /workspace
+
+# Install vllm from source.
+RUN git clone https://github.com/vllm-project/vllm.git
+WORKDIR vllm
+# Pin the version to a fixed git commit on 08/16/2023.
+RUN git reset --hard d1744376ae9fdbfa6a2dc763e1c67309e138fa3d
+# Apply a patch to vllm source:
+# 1) For models on Huggingface hub: if the model has multiple bin files, each
+#    bin file is downloaded separately and gets deleted after loading to GPU
+# 2) For models on GCS bucket: each model bin files is download separately
+#    and gets deleted after loading to GPU.
+# 3) Support code-llama model loading.
+COPY model_oss/vllm/vllm.patch /tmp/vllm.patch
+RUN git apply /tmp/vllm.patch
+RUN pip install -e .
+
+# Expose port 7080 for host serving.
+EXPOSE 7080

community-content/vertex_vision_model_garden/model_oss/vllm/vllm.patch

@@ -0,0 +1,311 @@
+diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py
+index 99fe593..e11246b 100644
+--- a/vllm/engine/arg_utils.py
++++ b/vllm/engine/arg_utils.py
+@@ -1,12 +1,43 @@
+ import argparse
+ import dataclasses
+ from dataclasses import dataclass
++import os
+ from typing import Optional, Tuple
+ 
++from google.cloud import storage
+ from vllm.config import (CacheConfig, ModelConfig, ParallelConfig,
+                          SchedulerConfig)
+ 
+ 
++GCS_PREFIX = "gs://"
++
++
++def is_gcs_path(input_path: str) -> bool:
++    return input_path.startswith(GCS_PREFIX)
++
++
++def download_gcs_dir_to_local(gcs_dir: str, local_dir: str):
++    if os.path.isdir(local_dir):
++        return
++    # gs://bucket_name/dir
++    bucket_name = gcs_dir.split('/')[2]
++    prefix = gcs_dir[len(GCS_PREFIX + bucket_name) :].strip('/')
++    client = storage.Client()
++    blobs = client.list_blobs(bucket_name, prefix=prefix)
++    for blob in blobs:
++        if blob.name[-1] == '/':
++            continue
++        file_path = blob.name[len(prefix) :].strip('/')
++        local_file_path = os.path.join(local_dir, file_path)
++        os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
++        if file_path.endswith(".bin"):
++            with open(local_file_path, 'w') as f:
++                f.write(f'{GCS_PREFIX}{bucket_name}/{prefix}/{file_path}')
++        else:
++            print(f"==> Download {gcs_dir}/{file_path} to {local_file_path}")
++            blob.download_to_filename(local_file_path)
++
++
+ @dataclass
+ class EngineArgs:
+     """Arguments for vLLM engine."""
+@@ -143,6 +174,19 @@ class EngineArgs:
+     def create_engine_configs(
+         self,
+     ) -> Tuple[ModelConfig, CacheConfig, ParallelConfig, SchedulerConfig]:
++        # Preprocess GCS paths.
++        if is_gcs_path(self.tokenizer) and self.tokenizer != self.model:
++            local_dir = "/tmp/gcs_tokenizer"
++            download_gcs_dir_to_local(self.tokenizer, local_dir)
++            self.tokenizer = local_dir
++        if is_gcs_path(self.model):
++            # Download GCS model without bin files.
++            local_dir = "/tmp/gcs_model"
++            download_gcs_dir_to_local(self.model, local_dir)
++            if self.tokenizer == self.model:
++                self.tokenizer = local_dir
++            self.model = local_dir
++
+         # Initialize the configs.
+         model_config = ModelConfig(self.model, self.tokenizer,
+                                    self.tokenizer_mode, self.trust_remote_code,
+diff --git a/vllm/entrypoints/api_server.py b/vllm/entrypoints/api_server.py
+index 58ea2e2..350e209 100644
+--- a/vllm/entrypoints/api_server.py
++++ b/vllm/entrypoints/api_server.py
+@@ -15,6 +15,10 @@ TIMEOUT_KEEP_ALIVE = 5  # seconds.
+ TIMEOUT_TO_PREVENT_DEADLOCK = 1  # seconds.
+ app = FastAPI()
+ 
++# Required by Vertex deployment.
++@app.get("/ping")
++async def ping() -> Response:
++    return Response(status_code=200)
+ 
+ @app.post("/generate")
+ async def generate(request: Request) -> Response:
+@@ -26,6 +30,9 @@ async def generate(request: Request) -> Response:
+     - other fields: the sampling parameters (See `SamplingParams` for details).
+     """
+     request_dict = await request.json()
++    is_on_vertex = "instances" in request_dict
++    if is_on_vertex:
++        request_dict = request_dict["instances"][0]
+     prompt = request_dict.pop("prompt")
+     stream = request_dict.pop("stream", False)
+     sampling_params = SamplingParams(**request_dict)
+@@ -63,7 +70,10 @@ async def generate(request: Request) -> Response:
+     assert final_output is not None
+     prompt = final_output.prompt
+     text_outputs = [prompt + output.text for output in final_output.outputs]
+-    ret = {"text": text_outputs}
++    if is_on_vertex:
++        ret = {"predictions": text_outputs}
++    else:
++        ret = {"text": text_outputs}
+     return JSONResponse(ret)
+ 
+ 
+diff --git a/vllm/model_executor/models/llama.py b/vllm/model_executor/models/llama.py
+index 93ab499..eca1b89 100644
+--- a/vllm/model_executor/models/llama.py
++++ b/vllm/model_executor/models/llama.py
+@@ -85,6 +85,7 @@ class LlamaAttention(nn.Module):
+         hidden_size: int,
+         num_heads: int,
+         num_kv_heads: int,
++        rope_theta: float = 10000,
+     ):
+         super().__init__()
+         self.hidden_size = hidden_size
+@@ -99,6 +100,7 @@ class LlamaAttention(nn.Module):
+         self.q_size = self.num_heads * self.head_dim
+         self.kv_size = self.num_kv_heads * self.head_dim
+         self.scaling = self.head_dim**-0.5
++        self.rope_theta = rope_theta
+ 
+         self.qkv_proj = ColumnParallelLinear(
+             hidden_size,
+@@ -118,6 +120,7 @@ class LlamaAttention(nn.Module):
+         self.attn = PagedAttentionWithRoPE(self.num_heads,
+                                            self.head_dim,
+                                            self.scaling,
++                                           base=self.rope_theta,
+                                            rotary_dim=self.head_dim,
+                                            num_kv_heads=self.num_kv_heads)
+ 
+@@ -143,10 +146,15 @@ class LlamaDecoderLayer(nn.Module):
+     def __init__(self, config: LlamaConfig):
+         super().__init__()
+         self.hidden_size = config.hidden_size
++        try:
++            rope_theta = config.rope_theta
++        except AttributeError:
++            rope_theta = 10000
+         self.self_attn = LlamaAttention(
+             hidden_size=self.hidden_size,
+             num_heads=config.num_attention_heads,
+             num_kv_heads=config.num_key_value_heads,
++            rope_theta=rope_theta,
+         )
+         self.mlp = LlamaMLP(
+             hidden_size=self.hidden_size,
+diff --git a/vllm/model_executor/weight_utils.py b/vllm/model_executor/weight_utils.py
+index a9d899a..57f39b5 100644
+--- a/vllm/model_executor/weight_utils.py
++++ b/vllm/model_executor/weight_utils.py
+@@ -3,13 +3,17 @@ import filelock
+ import glob
+ import json
+ import os
++import time
+ from typing import Iterator, List, Optional, Tuple
+ 
+-from huggingface_hub import snapshot_download
++from google.cloud import storage
++from huggingface_hub import hf_hub_download, snapshot_download
+ import numpy as np
+ import torch
+ from tqdm.auto import tqdm
+ 
++HF_PREFIX = "hf://"
++
+ 
+ class Disabledtqdm(tqdm):
+ 
+@@ -22,60 +26,90 @@ def hf_model_weights_iterator(
+     cache_dir: Optional[str] = None,
+     use_np_cache: bool = False,
+ ) -> Iterator[Tuple[str, torch.Tensor]]:
++    if use_np_cache:
++        raise ValueError("Do not support use_np_cache for lazy download.")
++
+     # Prepare file lock directory to prevent multiple processes from
+     # downloading the same model weights at the same time.
+     lock_dir = cache_dir if cache_dir is not None else "/tmp"
+     lock_file_name = model_name_or_path.replace("/", "-") + ".lock"
+     lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name))
+ 
+-    # Download model weights from huggingface.
+-    is_local = os.path.isdir(model_name_or_path)
+-    if not is_local:
+-        with lock:
+-            hf_folder = snapshot_download(model_name_or_path,
+-                                          allow_patterns="*.bin",
+-                                          cache_dir=cache_dir,
+-                                          tqdm_class=Disabledtqdm)
+-    else:
+-        hf_folder = model_name_or_path
+-
+-    hf_bin_files = [
+-        x for x in glob.glob(os.path.join(hf_folder, "*.bin"))
+-        if not x.endswith("training_args.bin")
+-    ]
+-
+-    if use_np_cache:
+-        # Convert the model weights from torch tensors to numpy arrays for
+-        # faster loading.
+-        np_folder = os.path.join(hf_folder, "np")
+-        os.makedirs(np_folder, exist_ok=True)
+-        weight_names_file = os.path.join(np_folder, "weight_names.json")
+-        with lock:
+-            if not os.path.exists(weight_names_file):
+-                weight_names = []
+-                for bin_file in hf_bin_files:
+-                    state = torch.load(bin_file, map_location="cpu")
+-                    for name, param in state.items():
+-                        param_path = os.path.join(np_folder, name)
+-                        with open(param_path, "wb") as f:
+-                            np.save(f, param.cpu().detach().numpy())
+-                        weight_names.append(name)
+-                with open(weight_names_file, "w") as f:
+-                    json.dump(weight_names, f)
+-
+-        with open(weight_names_file, "r") as f:
+-            weight_names = json.load(f)
+-
+-        for name in weight_names:
+-            param_path = os.path.join(np_folder, name)
+-            with open(param_path, "rb") as f:
+-                param = np.load(f)
+-            yield name, torch.from_numpy(param)
++    bin_files = []
++    if not os.path.isdir(model_name_or_path):
++        try:
++            with lock:
++                index_file = hf_hub_download(repo_id=model_name_or_path,
++                                             filename="pytorch_model.bin.index.json",
++                                             cache_dir=cache_dir)
++        except:
++            print("==> The model is in HF hub with 1 bin file, download it directly.", flush=True)
++            with lock:
++                hf_folder = snapshot_download(repo_id=model_name_or_path,
++                                              allow_patterns="*.bin",
++                                              cache_dir=cache_dir,
++                                              tqdm_class=Disabledtqdm)
++            bin_files = [x for x in glob.glob(os.path.join(hf_folder, "*.bin"))]
++        else:
++            print("==> The model is in HF hub with multiple bin file, do not download it now.", flush=True)
++            with open(index_file, "r") as f:
++                index = json.loads(f.read())
++            bin_filenames = set(index["weight_map"].values())
++            bin_files = [f"{HF_PREFIX}{model_name_or_path}/{bin_filename}" for bin_filename in bin_filenames]
+     else:
+-        for bin_file in hf_bin_files:
+-            state = torch.load(bin_file, map_location="cpu")
+-            for name, param in state.items():
+-                yield name, param
++        print("==> The model is in local disk.", flush=True)
++        bin_files = [x for x in glob.glob(os.path.join(model_name_or_path, "*.bin"))]
++
++    if "training_args.bin" in bin_files:
++        bin_files.remove("training_args.bin")
++    bin_files.sort()
++    print(f"==> Fetched bin files: {bin_files}", flush=True)
++
++    model_dir = "/tmp/model"
++    os.makedirs(model_dir, exist_ok=True)
++    for bin_file in bin_files:
++        delete_download = False
++
++        if os.path.exists(bin_file):
++            if open(bin_file, "rb").read(2) == b"gs":
++                gcs_path = open(bin_file).read()
++                bin_filename = gcs_path.split("/")[-1]
++                local_file = os.path.join(model_dir, bin_filename)
++                with lock:
++                    if not os.path.exists(local_file):
++                        client = storage.Client()
++                        with open(local_file, 'wb') as f:
++                            print(f"==> Download {gcs_path} to {bin_file}", flush=True)
++                            client.download_blob_to_file(gcs_path, f)
++                bin_file = local_file
++                delete_download = True
++        else:
++            assert bin_file.startswith(HF_PREFIX)
++            bin_filename = os.path.basename(bin_file)
++            local_file = os.path.join(model_dir, bin_filename)
++            with lock:
++                if not os.path.exists(local_file):
++                    print(f"==> Download {model_name_or_path}/{bin_filename} to {local_file}", flush=True)
++                    hf_hub_download(repo_id=model_name_or_path,
++                                    filename=bin_filename,
++                                    local_dir=model_dir,
++                                    local_dir_use_symlinks=False,
++                                    force_download=True)
++            bin_file = local_file
++            delete_download = True
++
++        torch.distributed.barrier()
++        print(f"==> Load {bin_file} to memory.", flush=True)
++        state = torch.load(bin_file, map_location="cpu")
++        for name, param in state.items():
++            yield name, param
++        torch.distributed.barrier()
++
++        if delete_download:
++            with lock:
++                if os.path.exists(bin_file):
++                    print(f"==> Delete {bin_file}", flush=True)
++                    os.remove(bin_file)
+ 
+ 
+ def load_tensor_parallel_weights(
+ 

notebooks/community/CODEOWNERS

@@ -43,9 +43,10 @@
 /notebooks/community/pipelines/google_cloud_pipeline_components_ready_to_go_text_classification_pipeline.ipynb @Narwhalprime
 /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_face_stylizer.pynb @schmidt-sebastian
+/notebooks/community/model_garden/model_garden_mediapipe_gesture_recognition.ipynb @schmidt-sebastian
 /notebooks/community/model_garden/model_garden_mediapipe_image_classification.ipynb @schmidt-sebastian
 /notebooks/community/model_garden/model_garden_mediapipe_image_generation.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
@@ -75,6 +76,7 @@
 /notebooks/community/model_garden/model_garden_jax_fvlm.ipynb @lavraicse
 /notebooks/community/model_garden/model_garden_pytorch_text_to_video_zero_shot.ipynb @bingatgoogle
 /notebooks/community/model_garden/model_garden_pytorch_text_to_video.ipynb @KCFindstr
+/notebooks/community/generative_ai/text_embedding_api_cloud_next_new_models.ipynb @xqr-g
 /notebooks/community/generative_ai/text_embedding_api_semantic_search_with_scann.ipynb @henrytansetiawan
 /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/generative_ai/text_embedding_api_cloud_next_new_models.ipynb

@@ -0,0 +1,363 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/xqr-g/vertex-ai-samples/blob/main/notebooks/community/generative_ai/text_embedding_api_cloud_next_new_models.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "KSP1duKDeaDR"
+      },
+      "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": "JAPoU8Sm5E6e"
+      },
+      "source": [
+        "# Cloud Next Embedding models\n",
+        "\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/generative_ai/text_embedding_api_cloud_next_new_models.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/generative_ai/text_embedding_api_cloud_next_new_models.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/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/generative_ai/text_embedding_api_cloud_next_new_models.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": "24743cf4a1e1"
+      },
+      "source": [
+        "**_NOTE_**: This notebook has been tested in the following environment:\n",
+        "\n",
+        "* Python version = 3.10"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tvgnzT1CKxrO"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This colab is used as a code example for how to call our newly released text embedding models (textembedding-gecko@latest and textembedding-gecko-multilingual@latest).\n",
+        "\n",
+        "Learn more about [text embedding api](https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-text-embeddings).\n",
+        "\n",
+        "This tutorial uses the following Google Cloud ML services and resources:\n",
+        "- Vertex LLM SDK\n",
+        "\n",
+        "The steps performed include:\n",
+        "- Installation and imports\n",
+        "- Generate embeddings\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aed92deeb4a0"
+      },
+      "source": [
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing),\n",
+        "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n",
+        "to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "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",
+        "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 API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
+        "\n",
+        "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "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."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "74ccc9e52986"
+      },
+      "source": [
+        "**1. Vertex AI Workbench**\n",
+        "* Do nothing as you are already authenticated."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "de775a3773ba"
+      },
+      "source": [
+        "**2. Local JupyterLab instance, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "254614fa0c46"
+      },
+      "outputs": [],
+      "source": [
+        "# ! gcloud auth login"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ef21552ccea8"
+      },
+      "source": [
+        "**3. Colab, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "603adbbf0532"
+      },
+      "outputs": [],
+      "source": [
+        "# from google.colab import auth\n",
+        "# auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "FyyMdUeAJIVv"
+      },
+      "source": [
+        "## Installation\n",
+        "\n",
+        "Install the following packages required to execute this notebook.\n",
+        "\n",
+        "**Remember to restart the runtime after installation.**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "snBUuUamoJPz"
+      },
+      "outputs": [],
+      "source": [
+        "!pip install git+https://github.com/googleapis/python-aiplatform.git"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "WX3CHZitmSJM"
+      },
+      "source": [
+        "### Please restart the runtime."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dae340cb-0583-4e7e-a562-6817ee4d7f6d"
+      },
+      "source": [
+        "### Imports libraries"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "412d00f1-08db-4880-8ced-52a9583757b8"
+      },
+      "outputs": [],
+      "source": [
+        "import vertexai\n",
+        "from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "MyMXIZoRlUcR"
+      },
+      "source": [
+        "#### Set your project ID and initiate Vertex AI\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": "3EdtdqnoldX4"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
+        "REGION = \"us-central1\"\n",
+        "\n",
+        "# Set the project id\n",
+        "! gcloud config set project {PROJECT_ID}\n",
+        "\n",
+        "# Initiate Vertex AI\n",
+        "vertexai.init(project=PROJECT_ID, location=REGION)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f50f22f3-ec85-463e-b6fe-5c8e6b80b07b"
+      },
+      "source": [
+        "## Generate embeddings"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "hQZoBXNGjizH"
+      },
+      "outputs": [],
+      "source": [
+        "# Set the model name.\n",
+        "MODEL_NAME = \"textembedding-gecko@latest\"  # @param [\"textembedding-gecko@latest\", \"textembedding-gecko-multilingual@latest\"]\n",
+        "\n",
+        "# Set the task_type, text and optional title as the model inputs.\n",
+        "TASK_TYPE = \"RETRIEVAL_DOCUMENT\"  # @param [\"RETRIEVAL_QUERY\", \"RETRIEVAL_DOCUMENT\", \"SEMANTIC_SIMILARITY\", \"CLASSIFICATION\", \"CLUSTERING\"]\n",
+        "TITLE = \"Google\"  # @param {type:\"string\"}\n",
+        "TEXT = \"Embed text.\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Verify the input is valid.\n",
+        "if not MODEL_NAME:\n",
+        "    raise ValueError(\"Please set MODEL_NAME.\")\n",
+        "if not TASK_TYPE:\n",
+        "    raise ValueError(\"Please set TASK_TYPE.\")\n",
+        "if not TEXT:\n",
+        "    raise ValueError(\"Please set TEXT.\")\n",
+        "if TITLE and TASK_TYPE != \"RETRIEVAL_DOCUMENT\":\n",
+        "    raise ValueError(\"Title can only be provided if the task_type is RETRIEVAL_DOCUMENT\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "BNPapKXviHlE"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "768\n"
+          ]
+        }
+      ],
+      "source": [
+        "def text_embedding(\n",
+        "  model_name: str, task_type: str, text: str, title: str = \"\") -> list:\n",
+        "    \"\"\"Generate text embedding with a Large Language Model.\"\"\"\n",
+        "    model = TextEmbeddingModel.from_pretrained(model_name)\n",
+        "\n",
+        "    text_embedding_input = TextEmbeddingInput(\n",
+        "        task_type=task_type, title=title, text=text)\n",
+        "    embeddings = model.get_embeddings([text_embedding_input])\n",
+        "    return embeddings[0].values\n",
+        "\n",
+        "embedding = text_embedding(\n",
+        "    model_name=MODEL_NAME, task_type=TASK_TYPE, text=TEXT, title=TITLE)\n",
+        "print(len(embedding))"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "text_embedding_api_cloud_next_new_models.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_jax_fvlm.ipynb

@@ -33,19 +33,19 @@
         "\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_jax_f_vlm.ipynb\">\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_jax_fvlm.ipynb\">\n",
         "      <img src=\"https://cloud.google.com/ml-engine/images/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_jax_f_vlm.ipynb\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_jax_fvlm.ipynb\">\n",
         "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
         "      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_jax_f_vlm.ipynb\">\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_fvlm.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",
@@ -72,7 +72,7 @@
       "source": [
         "## Overview\n",
         "\n",
-        "This notebook demonstrates serving a [JAX F-VLM model](https://github.com/google-research/google-research/tree/master/fvlm) for [open-vocabulary object detection](https://arxiv.org/abs/2209.15639) task and deploying them on Vertex AI for online prediction."
+        "This notebook demonstrates training and serving a [JAX F-VLM model](https://github.com/google-research/google-research/tree/master/fvlm) for [open-vocabulary object detection and instance segmentation](https://arxiv.org/abs/2209.15639) tasks and deploying them on Vertex AI for online prediction."
       ]
     },
     {
@@ -85,9 +85,11 @@
         "\n",
         "In this tutorial, you learn how to:\n",
         "\n",
+        "- Prepare a training dataset.\n",
+        "- Train a new JAX F-VLM model.\n",
         "- Upload the model to [Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction).\n",
         "- Deploy the model on [Endpoint](https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints).\n",
-        "- Run online predictions for image classification.\n",
+        "- Run online predictions for open-vocabulary image object detection and instance segmentation.\n",
         "\n",
         "This tutorial uses the following Google Cloud ML services and resources:\n",
         "\n",
@@ -150,22 +152,32 @@
         "# Install the packages.\n",
         "! pip3 install --upgrade google-cloud-aiplatform\n",
         "# Get F-VLM repository by using svn to avoid downloading entire google-research repository.\n",
-        "! apt install subversion\n",
+        "! sudo apt-get install -y subversion\n",
         "! rm -rf ./fvlm\n",
-        "! svn export -r 59152 https://github.com/google-research/google-research/trunk/fvlm\n",
+        "! svn export -r 60422 https://github.com/google-research/google-research/trunk/fvlm\n",
         "# Note: The following libraries are pinned down versions of:\n",
         "# https://github.com/google-research/google-research/blob/master/fvlm/requirements.txt\n",
-        "! pip3 install tensorflow==2.12.0\n",
-        "! pip3 install numpy==1.23.5\n",
-        "! pip3 install jax==0.4.14\n",
-        "! pip3 install jaxlib==0.4.14+cuda11.cudnn86\n",
-        "! pip3 install flax==0.7.1\n",
-        "! pip3 install torch==2.0.1+cu118\n",
-        "! pip3 install torchvision==0.15.2+cu118\n",
-        "! pip3 install opencv-python==4.7.0.72\n",
-        "! pip3 install tqdm==4.65.0\n",
-        "! pip3 install git+https://github.com/openai/CLIP.git@a1d071733d7111c9c014f024669f959182114e33\n",
-        "! pip3 install Pillow==9.5.0"
+        "! pip install tensorflow==2.12.0\n",
+        "! pip install tensorflow-datasets==4.9.2\n",
+        "! pip install numpy==1.23.5\n",
+        "! pip install torch==2.0.1+cu118 --extra-index-url https://download.pytorch.org/whl/cu118\n",
+        "! pip install torchvision==0.15.2+cu118 --extra-index-url https://download.pytorch.org/whl/cu118\n",
+        "! pip install opencv-python==4.7.0.72\n",
+        "! pip install tqdm==4.65.0\n",
+        "! pip install git+https://github.com/openai/CLIP.git@a1d071733d7111c9c014f024669f959182114e33\n",
+        "! pip install Pillow==9.1.1\n",
+        "! pip install orbax-checkpoint==0.3.3\n",
+        "! pip install gin-config==0.5.0\n",
+        "! pip install pycocotools==2.0.6\n",
+        "! pip install contextlib2==21.6.0\n",
+        "! pip install ml-collections==0.1.1\n",
+        "! pip install chex==0.1.7\n",
+        "! pip install optax==0.1.5\n",
+        "# Dependencies already included. Use no-deps to not update numpy.\n",
+        "! pip install --no-deps flax==0.7.2\n",
+        "! pip install --no-deps clu==0.0.9\n",
+        "! pip install jax[cuda11_cudnn86]==0.4.9 \\\n",
+        "--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html"
       ]
     },
     {
@@ -232,7 +244,9 @@
         "\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"
+        "1. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
+        "\n",
+        "1. Check your accelerator [quota](https://console.cloud.google.com/iam-admin/quotas). This notebook uses TPU V3 8 cores, you can filter and request `Custom model training TPU V3 cores per region` quota for 8 cores for the `Vertex AI API` service in your region.\n"
       ]
     },
     {
@@ -241,7 +255,7 @@
         "id": "WReHDGG5g0XY"
       },
       "source": [
-        "#### Set your project ID\n",
+        "### Set your project ID\n",
         "\n",
         "**If you don't know your project ID**, try the following:\n",
         "* Run `gcloud config list`.\n",
@@ -269,7 +283,7 @@
         "id": "region"
       },
       "source": [
-        "#### Region\n",
+        "### 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)."
       ]
@@ -417,24 +431,30 @@
       "outputs": [],
       "source": [
         "import base64\n",
+        "import datetime\n",
         "import functools\n",
         "import os\n",
         "import sys\n",
         "from io import BytesIO\n",
         "\n",
+        "import gin\n",
         "import jax\n",
+        "import matplotlib.pyplot as plt\n",
         "import numpy as np\n",
         "import tensorflow as tf\n",
         "import tqdm\n",
+        "import yaml\n",
         "from PIL import Image\n",
         "\n",
-        "sys.path.append(\"./fvlm\")\n",
-        "import inputs\n",
-        "import jax_clip\n",
-        "import utils\n",
+        "sys.path.append(os.path.join(os.getcwd(), \"fvlm\"))\n",
+        "# Work in fvlm directory in order to create gin config later.\n",
+        "os.chdir(os.path.join(os.getcwd(), \"fvlm\"))\n",
+        "from demo_utils import input_utils as inputs\n",
+        "from demo_utils import vis_utils\n",
         "from google.cloud import aiplatform\n",
         "from google.protobuf import json_format\n",
-        "from google.protobuf.struct_pb2 import Value"
+        "from google.protobuf.struct_pb2 import Value\n",
+        "from utils import clip_utils"
       ]
     },
     {
@@ -477,39 +497,62 @@
       },
       "outputs": [],
       "source": [
+        "REGION_PREFIX = REGION.split(\"-\")[0]\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",
+        "OBJECTIVE = \"iod\"\n",
+        "\n",
+        "# Category embedding constants.\n",
+        "GCS_CATEGORY_EMBEDDING_PATH = os.path.join(BUCKET_URI, \"category_embedding.npy\")\n",
+        "# The F-VLM model to use. Choose between 'resnet_50', 'resnet_50x4', or 'resnet_50x16'.\n",
+        "MODEL = \"resnet_50\"\n",
+        "# An upper bound on the number of classes.\n",
+        "MAX_NUM_CLS = 91\n",
+        "\n",
+        "# The local path to the F-VLM folder.\n",
+        "F_VLM_FOLDER = \".\"  # Current directory.\n",
+        "# Train config template path.\n",
+        "LOCAL_TRAIN_CONFIG_TEMPLATE_PATH = f\"{F_VLM_FOLDER}/configs/fvlm_train_and_eval.gin\"\n",
+        "# Final train config path.\n",
+        "GCS_TRAIN_CONFIG_PATH = f\"{BUCKET_URI}/fvlm_train_and_eval.gin\"\n",
+        "# Base output path for training artifacts.\n",
+        "GCS_TRAIN_BASE_PATH = f\"{BUCKET_URI}/train\"\n",
+        "# Training container image.\n",
+        "TRAIN_DOCKER_URI = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/jax-f-vlm-train\"\n",
+        "\n",
+        "# The train/eval batch sizes should be divisible by the number of GPUs used.\n",
+        "TRAIN_BATCH_SIZE = 8\n",
+        "EVAL_BATCH_SIZE = 1\n",
+        "# Whether to predict mask.\n",
+        "# Set False for object detection task and True for segmentation task.\n",
+        "INCLUDE_MASK = False\n",
+        "# Number of TPU cores to be used.\n",
+        "NUM_CORES = 8\n",
+        "# Total train steps.\n",
+        "TRAIN_STEPS = 1000\n",
+        "# Total evaluation steps.\n",
+        "EVAL_STEPS = 500\n",
+        "\n",
+        "# The pre-built TF SavedModel conversion docker image.\n",
+        "MODEL_CONVERSION_DOCKER_URI = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/jax-f-vlm-model-conversion\"\n",
+        "\n",
         "# The pre-built prediction docker image.\n",
         "OPTIMIZED_TF_RUNTIME_IMAGE_URI = (\n",
         "    \"us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.nightly:latest\"\n",
         ")\n",
-        "# The local path to the F-VLM folder.\n",
-        "F_VLM_FOLDER = \"./fvlm\"\n",
-        "# The F-VLM model to use. Choose between 'resnet_50', 'resnet_50x4', or 'resnet_50x16'.\n",
-        "MODEL = \"resnet_50\"\n",
-        "# The list of object categories to detect. For example: \"person, car, oven\".\n",
-        "CATEGORIES = [\n",
-        "    \"kiwi\",\n",
-        "    \"orange\",\n",
-        "    \"lemon\",\n",
-        "    \"blackberry\",\n",
-        "    \"pine cone\",\n",
-        "    \"red orange\",\n",
-        "    \"table\",\n",
-        "    \"spoon\",\n",
-        "    \"pine needles\",\n",
-        "    \"seed\",\n",
-        "]\n",
-        "# An upper bound on the number of classes.\n",
-        "MAX_NUM_CLS = 91\n",
         "# The max number of boxes to draw on the output image.\n",
         "MAX_BOXES_TO_DRAW = 25\n",
         "# The minimum score required to draw a detected object.\n",
         "MIN_SCORE_THRESH = 0.2  # @param {type:\"slider\", min:0, max:0.9, step:0.05}\n",
         "# The local path to the output image.\n",
         "OUTPUT_IMAGE_PATH = \"./output.jpg\"\n",
-        "# The original F-VLM SavedModel folder which takes image and text embeddings as inputs.\n",
-        "SAVED_MODEL_DIR = f'{F_VLM_FOLDER}/checkpoints/{MODEL.replace(\"resnet_\",\"r\")}'\n",
+        "# The F-VLM SavedModel folder which takes image and text embeddings as inputs.\n",
+        "GCS_SAVED_MODEL_DIR = f\"{BUCKET_URI}/saved_model\"\n",
         "# The converted SavedModel folder which takes jpeg bytes and text-embeddings bytes as inputs.\n",
-        "CONVERTED_SAVED_MODEL_DIR = \"./converted_saved_model\"\n",
+        "LOCAL_CONVERTED_SAVED_MODEL_DIR = f\"{F_VLM_FOLDER}/converted_saved_model\"\n",
         "# The Cloud Storage location for the converted SavedModel.\n",
         "GCS_CONVERTED_SAVED_MODEL_DIR = f\"{BUCKET_URI}/fvlm_saved_model\""
       ]
@@ -525,7 +568,9 @@
         "This section defines functions for:\n",
         "\n",
         "- Loading and converting input image into the required prediction format.\n",
-        "- Visualization of detection outputs."
+        "- Visualization of detection outputs.\n",
+        "- Getting GCS Fuse path.\n",
+        "- Getting a job name with current time."
       ]
     },
     {
@@ -545,7 +590,7 @@
         "\n",
         "def generate_text_embeddings(categories):\n",
         "    \"\"\"Generates text embeddings in numpy format from object categories.\"\"\"\n",
-        "    clip_text_fn = jax_clip.get_clip_text_fn(MODEL)\n",
+        "    clip_text_fn = clip_utils.get_clip_text_fn(MODEL)\n",
         "    class_clip_features = []\n",
         "    print(\"Computing custom category text embeddings.\")\n",
         "    for cls_name in tqdm.tqdm(categories, total=len(categories)):\n",
@@ -585,7 +630,7 @@
         "\n",
         "\n",
         "def generate_prediction_output_image(\n",
-        "    input_image_path, prediction_output, output_image_path\n",
+        "    input_image_path, prediction_output, output_image_path, categories\n",
         "):\n",
         "    \"\"\"Generates prediction output image with detected objects and bounding boxes.\"\"\"\n",
         "    # Generate tensors from prediction outputs.\n",
@@ -607,21 +652,20 @@
         "    labels = np_data.pop(\"labels\")\n",
         "    # Generate visualization.\n",
         "    print(\"Preparing visualization.\")\n",
-        "    categories = CATEGORIES\n",
         "    id_mapping = {(i + 1): c for i, c in enumerate(categories)}\n",
         "    id_mapping[0] = \"background\"\n",
         "    for k in range(len(categories) + 2, MAX_NUM_CLS):\n",
         "        id_mapping[k] = \"empty\"\n",
         "    category_index = inputs.get_category_index(id_mapping)\n",
         "    maskrcnn_visualizer_fn = functools.partial(\n",
-        "        utils.visualize_boxes_and_labels_on_image_array,\n",
+        "        vis_utils.visualize_boxes_and_labels_on_image_array,\n",
         "        category_index=category_index,\n",
         "        use_normalized_coordinates=False,\n",
         "        max_boxes_to_draw=MAX_BOXES_TO_DRAW,\n",
         "        min_score_thresh=MIN_SCORE_THRESH,\n",
         "        skip_labels=False,\n",
         "    )\n",
-        "    vis_image = utils.visualize_instance_segmentations(\n",
+        "    vis_image = vis_utils.visualize_instance_segmentations(\n",
         "        prediction_output_tensor,\n",
         "        image_embeddings,\n",
         "        labels[\"image_info\"],\n",
@@ -629,7 +673,20 @@
         "    )\n",
         "    pil_vis_image = Image.fromarray(vis_image, mode=\"RGB\")\n",
         "    pil_vis_image.save(output_image_path)\n",
-        "    print(\"Completed saving the output image at: \", output_image_path)"
+        "    print(\"Completed saving the output image at: \", output_image_path)\n",
+        "\n",
+        "\n",
+        "def gcs_fuse_path(path: str) -> str:\n",
+        "    \"\"\"Try to convert path to gcsfuse path if it starts with gs:// else do not modify it.\"\"\"\n",
+        "    path = path.strip()\n",
+        "    if path.startswith(\"gs://\"):\n",
+        "        return \"/gcs/\" + path[5:]\n",
+        "    return path\n",
+        "\n",
+        "\n",
+        "def get_job_name_with_datetime(prefix: str):\n",
+        "    \"\"\"Gets a job name by adding current time to prefix.\"\"\"\n",
+        "    return prefix + datetime.datetime.now().strftime(\"_%Y%m%d_%H%M%S\")"
       ]
     },
     {
@@ -638,9 +695,326 @@
         "id": "ayNrua2txk0B"
       },
       "source": [
-        "# Convert F-VLM SavedModel to support smaller input size\n",
+        "## Train new models\n",
         "\n",
-        "The F-VLM SavedModel takes image embeddings and text embeddings as input. But you can not send these inputs directly for Vertex AI Online Prediction because there is a limit of 1.5 MB on the prediction request size. So you will first convert the SavedModel format to take jpeg bytes and text-embeddings bytes as an input instead. This modified input format will meet the 1.5 MB limit requirement."
+        "This section shows how to train new models:\n",
+        "\n",
+        "1. Convert input data to tfrecord format.\n",
+        "2. Prepare category embeddings.\n",
+        "3. Create Custom training job to train a new JAX model.\n",
+        "\n",
+        "If you want to use the provided pretrained saved model, you can skip this section and go to the `Convert F-VLM SavedModel to support smaller input size` section."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2d45a55bf581"
+      },
+      "source": [
+        "### Prepare input data for training\n",
+        "\n",
+        "Prepare data in the format as described [here](https://cloud.google.com/vertex-ai/docs/image-data/object-detection/prepare-data?_ga=2.18235642.-2108754545.1692303912), and then convert them to the training formats as below:\n",
+        "\n",
+        "- `input_file_path`: The input file path for preparing data.\n",
+        "- `input_file_type`: The input file type, such as csv or jsonl.\n",
+        "- `split_ratio`: The proportion of data to split into train/validation/test.\n",
+        "- `num_shard`: The number of shards for train/validation/test.\n",
+        "- `data_converter_output_dir`: The output directory, which will contain prepared train/test/validation data."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "f6ecd1749020"
+      },
+      "outputs": [],
+      "source": [
+        "# This job will convert input data as training format, with given split ratios\n",
+        "# and number of shards on train/validation/test.\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', 'coco_json']\n",
+        "split_ratio = \"0.8,0.2,0.0\"\n",
+        "num_shard = \"10,10,10\"\n",
+        "data_converter_output_dir = os.path.join(BUCKET_URI, data_converter_job_name)\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",
+        "            ],\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 = gcs_fuse_path(\n",
+        "    os.path.join(data_converter_output_dir, \"train.tfrecord*\")\n",
+        ")\n",
+        "input_validation_data_path = gcs_fuse_path(\n",
+        "    os.path.join(data_converter_output_dir, \"val.tfrecord*\")\n",
+        ")\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": "c88e7507727a"
+      },
+      "source": [
+        "### Prepare category embeddings\n",
+        "\n",
+        "Get the list of categories from input dataset. Alternatively, you can choose your own categories and override the `CATEGORIES` list."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "0e77c3010012"
+      },
+      "outputs": [],
+      "source": [
+        "with tf.io.gfile.GFile(label_map_path) as f:\n",
+        "    label_map = yaml.safe_load(f)[\"label_map\"]\n",
+        "\n",
+        "CATEGORIES = [label_map[i] for i in range(1, len(label_map) + 1)]\n",
+        "\n",
+        "print(f\"List of categories: {CATEGORIES}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "24c293e29742"
+      },
+      "source": [
+        "Generate text embeddings for the labels and save it as an `npy` file."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "3d003f59bf70"
+      },
+      "outputs": [],
+      "source": [
+        "text_embeddings = generate_text_embeddings(categories=CATEGORIES)\n",
+        "with tf.io.gfile.GFile(GCS_CATEGORY_EMBEDDING_PATH, \"w\") as f:\n",
+        "    np.save(f, text_embeddings)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dd1b3a5d5171"
+      },
+      "source": [
+        "### Create Custom training job and train a new model"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8a30f2fa7d34"
+      },
+      "source": [
+        "Set training hyperparameters. To see the full list of customizable hyperparameters, see the [config template file](https://github.com/google-research/google-research/blob/e1f9fae637db06ba885217518cbbf7f4fa4b9d7b/fvlm/configs/fvlm_train_and_eval.gin)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "b2f5dd9951ff"
+      },
+      "outputs": [],
+      "source": [
+        "assert TRAIN_BATCH_SIZE % NUM_CORES == 0, \"NUM_CORES must divide TRAIN_BATCH_SIZE.\"",
+        "\n",
+        "config_overrides = [\n",
+        "    \"evaluate.host_evaluator = None\",  # For evaluation without COCO annotation json file.\n",
+        "    \"evaluate.eval_metrics = {'coco_metric': @COCODetectionMetric}\",  # For evaluation without COCO annotation json file.\n",
+        "    f\"TRAIN_FILE_PATTERN = '{gcs_fuse_path(input_train_data_path)}'\",\n",
+        "    f\"EVAL_FILE_PATTERN = '{gcs_fuse_path(input_validation_data_path)}'\",\n",
+        "    f\"TRAIN_BS = {TRAIN_BATCH_SIZE}\",\n",
+        "    f\"EVAL_BS = {EVAL_BATCH_SIZE}\",\n",
+        "    f\"TRAIN_STEPS = {TRAIN_STEPS}\",\n",
+        "    f\"EVAL_STEPS = {EVAL_STEPS}\",\n",
+        "    f\"EMBED_PATH = '{GCS_CATEGORY_EMBEDDING_PATH}'\",\n",
+        "    f\"CATG_PAD_SIZE = {MAX_NUM_CLS}\",\n",
+        "    f\"INCLUDE_MASK = {INCLUDE_MASK}\",\n",
+        "]\n",
+        "print(config_overrides)\n",
+        "\n",
+        "gin.parse_config_files_and_bindings(\n",
+        "    [LOCAL_TRAIN_CONFIG_TEMPLATE_PATH], config_overrides, finalize_config=False\n",
+        ")\n",
+        "config = gin.config_str()\n",
+        "\n",
+        "with tf.io.gfile.GFile(GCS_TRAIN_CONFIG_PATH, \"w\") as f:\n",
+        "    f.write(config)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "290fd279e7f0"
+      },
+      "source": [
+        "Create and run the training job with the model-garden JAX F-VLM training docker using the Vertex AI SDK."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "003068435a36"
+      },
+      "outputs": [],
+      "source": [
+        "JOB_NAME = get_job_name_with_datetime(prefix=\"jax_fvlm\")\n",
+        "\n",
+        "docker_args_list = [\n",
+        "    f\"--output_dir={gcs_fuse_path(GCS_TRAIN_BASE_PATH)}\",\n",
+        "    f\"--config_path={gcs_fuse_path(GCS_TRAIN_CONFIG_PATH)}\",\n",
+        "    \"--mode=train_and_eval\",\n",
+        "]\n",
+        "print(docker_args_list)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9de7e802b53c"
+      },
+      "outputs": [],
+      "source": [
+        "# Click on the generated link in the output under \"View backing custom job:\" to see your run in the Cloud Console.\n",
+        "# The job will run for appoximately 15 minutes in the current settings.\n",
+        "job = aiplatform.CustomContainerTrainingJob(\n",
+        "    display_name=JOB_NAME,\n",
+        "    container_uri=TRAIN_DOCKER_URI,\n",
+        ")\n",
+        "model = job.run(\n",
+        "    args=docker_args_list,\n",
+        "    base_output_dir=f\"{GCS_TRAIN_BASE_PATH}\",\n",
+        "    replica_count=1,\n",
+        "    machine_type=\"cloud-tpu\",\n",
+        "    accelerator_type=\"TPU_V3\",\n",
+        "    accelerator_count=NUM_CORES,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "798d2268f869"
+      },
+      "source": [
+        "## Convert JAX checkpoint to TensorFlow SavedModel for inference\n",
+        "\n",
+        "Convert the previously fine-tuned JAX F-VLM model to a TF SavedModel for online prediction."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "f31f3e19c6ee"
+      },
+      "outputs": [],
+      "source": [
+        "JOB_NAME = get_job_name_with_datetime(prefix=\"jax_model_conversion\")\n",
+        "\n",
+        "docker_args_list = [\n",
+        "    f\"--input_dir={gcs_fuse_path(GCS_TRAIN_BASE_PATH)}\",\n",
+        "    f\"--output_dir={gcs_fuse_path(GCS_SAVED_MODEL_DIR)}\",\n",
+        "    f\"--max_num_classes={MAX_NUM_CLS}\",\n",
+        "    f\"--model_name={MODEL}\",\n",
+        "    f\"--include_mask={INCLUDE_MASK}\",\n",
+        "]\n",
+        "print(docker_args_list)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "d0a4457779b1"
+      },
+      "outputs": [],
+      "source": [
+        "# Create and run the model conversion job.\n",
+        "# Click on the generated link in the output under \"View backing custom job:\" to see your run in the Cloud Console.\n",
+        "container_uri = MODEL_CONVERSION_DOCKER_URI\n",
+        "job = aiplatform.CustomContainerTrainingJob(\n",
+        "    display_name=JOB_NAME,\n",
+        "    container_uri=container_uri,\n",
+        ")\n",
+        "model_conversion_workdir = os.path.join(BUCKET_URI, JOB_NAME)\n",
+        "model = job.run(\n",
+        "    args=docker_args_list,\n",
+        "    base_output_dir=f\"{model_conversion_workdir}\",\n",
+        "    replica_count=1,\n",
+        "    machine_type=\"n1-highmem-8\",\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ayNrua2txk0B"
+      },
+      "source": [
+        "## Convert F-VLM SavedModel to support smaller input size\n",
+        "\n",
+        "The F-VLM SavedModel takes image embeddings and text embeddings as input. But you can not send these inputs directly for Vertex AI Online Prediction because there is a limit of 1.5 MB on the prediction request size. So you will first convert the SavedModel format to take jpeg bytes and text-embeddings bytes as an input instead. This modified input format will meet the 1.5 MB limit requirement.\n",
+        "\n",
+        "If you did not train a new model, uncomment and run the cell below to use the provided saved model."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "12e3d9c0221a"
+      },
+      "outputs": [],
+      "source": [
+        "# # Path to the downloaded pre-trained saved model.\n",
+        "# GCS_SAVED_MODEL_DIR = f'{F_VLM_FOLDER}/checkpoints/{MODEL.replace(\"resnet_\",\"r\")}'"
       ]
     },
     {
@@ -688,15 +1062,15 @@
         "    return serve_fn\n",
         "\n",
         "\n",
-        "! rm -rf {CONVERTED_SAVED_MODEL_DIR}\n",
-        "model = tf.saved_model.load(SAVED_MODEL_DIR)\n",
+        "! rm -rf {LOCAL_CONVERTED_SAVED_MODEL_DIR}\n",
+        "model = tf.saved_model.load(GCS_SAVED_MODEL_DIR)\n",
         "signatures = {\n",
         "    \"serving_default\": get_serve_fn(model=model).get_concrete_function(\n",
         "        tf.TensorSpec(shape=[None], dtype=tf.string), tf.TensorSpec([None], tf.string)\n",
         "    )\n",
         "}\n",
-        "tf.saved_model.save(model, CONVERTED_SAVED_MODEL_DIR, signatures=signatures)\n",
-        "print(\"Saved the converted SavedModel to directory: \", CONVERTED_SAVED_MODEL_DIR)"
+        "tf.saved_model.save(model, LOCAL_CONVERTED_SAVED_MODEL_DIR, signatures=signatures)\n",
+        "print(\"Saved the converted SavedModel to directory: \", LOCAL_CONVERTED_SAVED_MODEL_DIR)"
       ]
     },
     {
@@ -717,7 +1091,7 @@
       "outputs": [],
       "source": [
         "! gsutil -m rm -R -f {GCS_CONVERTED_SAVED_MODEL_DIR}\n",
-        "! gsutil -m cp -R {CONVERTED_SAVED_MODEL_DIR} {GCS_CONVERTED_SAVED_MODEL_DIR}\n",
+        "! gsutil -m cp -R {LOCAL_CONVERTED_SAVED_MODEL_DIR} {GCS_CONVERTED_SAVED_MODEL_DIR}\n",
         "! gsutil ls {GCS_CONVERTED_SAVED_MODEL_DIR}"
       ]
     },
@@ -757,7 +1131,7 @@
         ")\n",
         "\n",
         "jax_fvlm_endpoint = jax_fvlm_model.deploy(\n",
-        "    deployed_model_display_name=\"jax_vlm_deployed\",\n",
+        "    deployed_model_display_name=\"jax_fvlm_deployed\",\n",
         "    traffic_split={\"0\": 100},\n",
         "    machine_type=\"n1-highmem-16\",\n",
         "    min_replica_count=1,\n",
@@ -795,6 +1169,37 @@
       "metadata": {
         "id": "B1Q7AbmJ4QxZ"
       },
+      "source": [
+        "If you did not train a new model and are using the provided saved model, uncomment and define your own categories in the cell below."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "50109d914484"
+      },
+      "outputs": [],
+      "source": [
+        "# CATEGORIES = [\n",
+        "#     \"kiwi\",\n",
+        "#     \"orange\",\n",
+        "#     \"lemon\",\n",
+        "#     \"blackberry\",\n",
+        "#     \"pine cone\",\n",
+        "#     \"red orange\",\n",
+        "#     \"table\",\n",
+        "#     \"spoon\",\n",
+        "#     \"pine needles\",\n",
+        "#     \"seed\",\n",
+        "# ]"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "1ad9b7e69dc9"
+      },
       "source": [
         "Prepare jpeg bytes and text-embeddings bytes inputs."
       ]
@@ -810,6 +1215,7 @@
         "image_jpeg_bytes_inputs = get_jpeg_bytes(\n",
         "    local_image_path=DEMO_IMAGE_PATH, new_width=1024\n",
         ")\n",
+        "\n",
         "text_embeddings = generate_text_embeddings(categories=CATEGORIES)\n",
         "text_embeddings_bytes_inputs = convert_numpy_array_to_byte_string_via_tf_tensor(\n",
         "    text_embeddings\n",
@@ -871,7 +1277,11 @@
         "    input_image_path=DEMO_IMAGE_PATH,\n",
         "    prediction_output=prediction_output,\n",
         "    output_image_path=OUTPUT_IMAGE_PATH,\n",
-        ")"
+        "    categories=CATEGORIES,\n",
+        ")\n",
+        "\n",
+        "img = Image.open(OUTPUT_IMAGE_PATH)\n",
+        "plt.imshow(img)"
       ]
     },
     {

notebooks/community/model_garden/model_garden_keras_yolov8.ipynb

@@ -303,22 +303,23 @@
         "\n",
         "RESOLUTION = 512\n",
         "\n",
+        "\n",
         "def get_job_name_with_datetime(prefix: str):\n",
-        "  \"\"\"Generates a job name with date time when triggering training or deployment\n",
+        "    \"\"\"Generates a job name with date time when triggering training or deployment\n",
         "    jobs in Vertex AI.\n",
-        "  \"\"\"\n",
+        "    \"\"\"\n",
         "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
         "\n",
         "\n",
         "def load_img(path):\n",
-        "    \"\"\"Reads image from path and return PIL.Image instance.\"\"\"\"\n",
+        "    \"\"\"Reads image from path and return PIL.Image instance.\"\"\"\n",
         "    img = tf.io.read_file(path)\n",
         "    img = tf.image.decode_jpeg(img, channels=3)\n",
         "    return Image.fromarray(np.uint8(img)).convert(\"RGB\")\n",
         "\n",
         "\n",
         "def decode_image(image_str_tensor: tf.string) -> tf.float32:\n",
-        "    \"\"\"Converts and resizes image bytes to image tensor.\"\"\"\"\n",
+        "    \"\"\"Converts and resizes image bytes to image tensor.\"\"\"\n",
         "    image = tf.io.decode_image(image_str_tensor, 3, expand_animations=False)\n",
         "    image = tf.image.resize(image, (RESOLUTION, RESOLUTION))\n",
         "    return image\n",

notebooks/community/model_garden/model_garden_mediapipe_face_stylizer.ipynb

@@ -0,0 +1,580 @@
+{
+  "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 Face Stylizer\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_face_stylizer.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_face_stylizer.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_face_stylizer.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\n",
+        "\n",
+        "**_NOTE_**: The checkpoint and the dataset linked in this Colab are not owned or distributed by Google, and are made available by third parties. Please review the terms and conditions made available by the third parties before using the checkpoint and data."
+      ]
+    },
+    {
+      "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 customize an on-device face stylizer model in Vertex AI Model Garden.\n",
+        "\n",
+        "The MediaPipe face stylizer solution provides several models you can use immediately to transform the face to the styles including (cartoon, oil painting, etc.) in your application. However, if you need to transfer the face to an unseen style not covered by the provided models, you can customize the pretrained model with your own data and MediaPipe Model Maker. This model modification tool fine-tune a portion of the model using data you provide. This method is faster than training a new model from scatch and can produce a model adapt to your specific application.\n",
+        "\n",
+        "The following sections show you how to use Model Maker to retrain a pre-built model for face stylization with your own data on Vertex AI, which you can then use with the MediaPipe Face Stylizer.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "* Customize a Face Stylizer model\n",
+        "  * Convert input data to training formats\n",
+        "  * Create [custom jobs](https://cloud.google.com/vertex-ai/docs/training/create-custom-job) to customize new models\n",
+        "  * Export customized 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",
+        "\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_face_stylizer_%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_ACCELERATOR_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": "LmZ4efOd-sak"
+      },
+      "source": [
+        "### Prepare input data for training\n",
+        "\n",
+        "Retraining the face stylizer model requires user to provide a single stylized face image. The stylized face is expected to be forward facing with visible left right eyes and mouth. The face should only have minor rotation, i.e. less than 30 degress around the yaw, pitch, and roll axes.\n",
+        "\n",
+        "You can upload an image to Google Cloud Storage or use our [provided example](https://storage.googleapis.com/mediapipe-assets/face_stylizer_style_color_sketch.jpg)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "IndQ_m6ddUEM"
+      },
+      "outputs": [],
+      "source": [
+        "training_data_path = \"gs://mediapipe-assets/face_stylizer_style_color_sketch.jpg\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aaff6f5be7f6"
+      },
+      "source": [
+        "# Retrain model\n",
+        "\n",
+        "Once you have provided an input image, you can begin retraining the face stylizer model to adapt to the new style. This type of model modification is called transfer learning. The instructions below use the data prepared in the previous section to retrain a face stylizer model to apply cartoon style to the raw human face.\n",
+        "\n",
+        "**_NOTE_**: For this type of model, the retraining process causes the model to forget any style it can apply before. Once the retraining is complete, the new model can only apply the new style defined by the new stylized image.\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kySYKb3-hnr6"
+      },
+      "source": [
+        "## Set retraining options\n",
+        "There are a few required settings to run a retraining aside from your training dataset:\n",
+        "\n",
+        "* **Swap layers:** The `swap_layers` parameter is used to determine how to mix the latent code layers between the learned style and the raw face images. The latent code is represented as a tensor of shape [1, 12, 512]. The second dimension of the latent code tensor is called the layer. The face stylizer mixes the learned style and raw face images by generating a weighted sum of the two latent codes on the swap layers. The swap layers are therefore integers within [1, 12]. The more layers are set, the more style will be applied to the output image. Although there is no explicit mapping between the style semantics and the layer index, the shallow layers, e.g. 8, 9, represent the global features of the face, while the deep layers, e.g. 10, 11, represent the fine-grained features of the face. The output stylized image is sensitive to the setting of swap layers. By default, it is set to [8, 9, 10, 11].\n",
+        "* **Learning rate and epochs:** Use `learning_rate` and epochs` to specify the these two hyperparameters. learning_rate is set to 4e-4 by default. epochs defines the number of iterations to fine-tune the BlazeStyleGAN model and are set to 100 by default. The lower the learning rate is, the greater the epochs is expected to retrain the model to converge.\n",
+        "* **Batch size:** The `batch_size` is used to define the number of latent code samples we sample around the latent code extracted by the encoder with the input image. The batch of latent codes are used to fine-tune the decoder. The greater the batch size usually yield to better performance. It is also limited by the hardware memory. For A100 GPU, the maximum batch size is 8. For P100 and T4 GPU, the maximum batch size is 2.\n",
+        "\n",
+        "Further more advanced parameters that you can configure are `alpha`, `perception_loss_weight`, `adv_loss_weight`, `beta_1` and `beta_2`."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "um_XKbmpTaHx"
+      },
+      "outputs": [],
+      "source": [
+        "# The layers of feature to be interpolated between encoding features and\n",
+        "# StyleGAN input features.\n",
+        "swap_layers: str = \"[8, 9, 10, 11]\"  # @param {type:\"string\"}\n",
+        "# The learning rate to use for gradient descent training.\n",
+        "learning_rate: float = 0.0001  # @param {type:\"number\"}\n",
+        "# Number of training iterations over the dataset.\n",
+        "epochs: int = 100  # @param {type:\"slider\", min:0, max:100, step:1}\n",
+        "# Batch size for training.\n",
+        "batch_size: int = 2  # @param {type:\"number\"}\n",
+        "\n",
+        "\n",
+        "# Other supported options\n",
+        "\n",
+        "# Weighting coefficient of style latent for swapping layer interpolation.\n",
+        "# Its valid range is [0, 1]. The greater weight means stronger style is\n",
+        "# applied to the output image. Expect to set it to a small value,\n",
+        "# i.e. < 0.1.\n",
+        "alpha: float = 0.1  # @param {type:\"number\"}\n",
+        "\n",
+        "# Weighting coefficients of image perception quality loss. It contains three\n",
+        "# coefficients, l1, content, and style which control the difference between the\n",
+        "# generated image and raw input image, the content difference between generated\n",
+        "# face and raw input face, and the how similar the style between the generated\n",
+        "# image and raw input image. Users can increase the style weight to enforce\n",
+        "# stronger style or the content weight to reserve more raw input face details.\n",
+        "# Weight for L1 loss.\n",
+        "perception_loss_l1: float = 0.5  # @param {type:\"number\"}\n",
+        "# Weight for content loss.\n",
+        "perception_loss_content: float = 4.0  # @param {type:\"number\"}\n",
+        "# Weight for stlye loss.\n",
+        "perception_loss_style: float = 1.0  # @param {type:\"number\"}\n",
+        "\n",
+        "# Weighting coeffcieint of adversarial loss versus image perceptual quality loss.\n",
+        "# This hyperparameter is used to control the realism of the generated image. It\n",
+        "# expects a small value, i.e. < 0.2.\n",
+        "adv_loss_weight: float = 0.2  # @param {type:\"number\"}\n",
+        "# beta_1 used in tf.keras.optimizers.Adam.\n",
+        "beta_1: float = 0.0  # @param {type:\"number\"}\n",
+        "# beta_2 used in tf.keras.optimizers.Adam.\n",
+        "beta_2: float = 0.99  # @param {type:\"number\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "HwcCjwlBTQIz"
+      },
+      "source": [
+        "### Run retraining\n",
+        "With your training dataset and retraining options prepared, you are ready to start the retraining process. This process requires running on GPU and can take a few minutes to a few hours depending on your available compute resources. On Vertex AI with GPU processing, the example retraining below takes about 2 minutes.\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",
+        "\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": TRAINING_MACHINE_TYPE,\n",
+        "            \"accelerator_type\": TRAINING_ACCELERATOR_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=face_stylizer\",\n",
+        "                \"--training_data_path=%s\" % training_data_path,\n",
+        "                \"--model_export_path=%s\" % model_export_path,\n",
+        "                \"--evaluation_result_path=%s\" % model_export_path,\n",
+        "                \"--hparams=%s\"\n",
+        "                % json.dumps(\n",
+        "                    {\n",
+        "                        \"learning_rate\": learning_rate,\n",
+        "                        \"batch_size\": batch_size,\n",
+        "                        \"epochs\": epochs,\n",
+        "                        \"beta_1\": beta_1,\n",
+        "                        \"beta_2\": beta_2,\n",
+        "                    }\n",
+        "                ),\n",
+        "                \"--model_options=%s\"\n",
+        "                % json.dumps(\n",
+        "                    {\n",
+        "                        \"swap_layers\": json.loads(swap_layers),\n",
+        "                        \"alpha\": alpha,\n",
+        "                        \"perception_loss_l1\": perception_loss_l1,\n",
+        "                        \"perception_loss_content\": perception_loss_content,\n",
+        "                        \"perception_loss_style\": perception_loss_style,\n",
+        "                        \"adv_loss_weight\": adv_loss_weight,\n",
+        "                    }\n",
+        "                ),\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": "g0BGaofgsMsy"
+      },
+      "source": [
+        "## Export model\n",
+        "After retraining the model, you can save the Tensorflow Lite model and integrate it with your on-device application by following the [Face stylization task guide](https://developers.google.com/mediapipe/solutions/vision/face_stylizer)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NYuQowyZEtxK"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "def copy_model(model_source, model_dest):\n",
+        "    ! gsutil cp {model_source} {model_dest}\n",
+        "\n",
+        "copy_model(EXPORTED_MODEL_OUTPUT_FILE, \"face_stylizer.task\")\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    from google.colab import files\n",
+        "\n",
+        "    files.download(\"face_stylizer.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_face_stylizer.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

@@ -285,23 +285,6 @@
       "source": [
         "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",

notebooks/community/model_garden/model_garden_mediapipe_image_generation.ipynb

@@ -299,7 +299,7 @@
         "id": "n6IFz75WGCam"
       },
       "source": [
-        "### Define training machine specs"
+        "### Define training and serving constants"
       ]
     },
     {
@@ -314,7 +314,13 @@
         "TRAINING_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/mediapipe-stable-diffusion-train\"\n",
         "TRAINING_MACHINE_TYPE = \"a2-highgpu-1g\"\n",
         "TRAINING_ACCELERATOR_TYPE = \"NVIDIA_TESLA_A100\"\n",
-        "TRAINING_ACCELERATOR_COUNT = 1"
+        "TRAINING_ACCELERATOR_COUNT = 1\n",
+        "\n",
+        "PREDICTION_CONTAINER_URI = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-serve\"\n",
+        "PREDICTION_PORT = 7080\n",
+        "PREDICTION_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
+        "PREDICTION_MACHINE_TYPE = \"n1-standard-8\"\n",
+        "UPLOAD_MODEL_NAME = \"mediapipe_stable_diffusion_model_%s\" % now"
       ]
     },
     {
@@ -359,6 +365,31 @@
       "source": [
         "### Set training options\n",
         "\n",
+        "The Image Generator comes with a set of pre-defined HParams settings that work best for specific situations. You should select a template that best matches your use case.\n",
+        "\n",
+        "If applicable, you can also use one of our pre-trained models for these templates. These can be used directly and without training:\n",
+        "* [Object - Berry Bowl](https://storage.googleapis.com/mediapipe-tasks/image_generator/object/pytorch_lora_weights.bin)\n",
+        "* [Face](https://storage.googleapis.com/mediapipe-tasks/image_generator/face/pytorch_lora_weights.bin)\n",
+        "* [Style](https://storage.googleapis.com/mediapipe-tasks/image_generator/style/pytorch_lora_weights.bin)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bj31hD4W02Ui"
+      },
+      "outputs": [],
+      "source": [
+        "template = \"\"  # @param [\"\", \"Face\", \"Object\", \"Style\"]"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "lYF-4ydq0k0u"
+      },
+      "source": [
         "To set custom training parameters, adjust the following values:"
       ]
     },
@@ -373,8 +404,8 @@
         "# Parameters about training configuration\n",
         "# The learning rate to use for gradient descent training.\n",
         "learning_rate: float = 0.00001  # @param {type:\"number\"}\n",
-        "#  Number of training steps.\n",
-        "num_train_steps: int = 1000  # @param {type:\"integer\"}\n",
+        "# Number of training steps. If set to 0, uses the default value.\n",
+        "num_train_steps: int = 0  # @param {type:\"integer\"}\n",
         "# Save the checkpoint in every n steps.\n",
         "save_checkpoints_every_n: int = 100  # @param {type:\"integer\"}\n",
         "# Batch size for training.\n",
@@ -383,9 +414,14 @@
         "# Dataset-related parameters\n",
         "# Whether to use random horizontal flip on data.\n",
         "random_flip: bool = False  # @param {type:\"boolean\"}\n",
+        "# Whether to use random largest square crop.\n",
+        "random_crop: bool = False  # @param {type:\"boolean\"}\n",
+        "# Whether to distort the color of the image (jittering order is random).\n",
+        "random_color_jitter: bool = False  # @param {type:\"boolean\"}\n",
         "\n",
         "# Hyperparameters for LoRA tuning\n",
-        "lora_rank: int = 4  # @param {type:\"integer\"}"
+        "# The rank in the low-rank matrices. If set to 0, uses the default value.\n",
+        "lora_rank: int = 0  # @param {type:\"integer\"}"
       ]
     },
     {
@@ -425,7 +461,7 @@
         "id": "nvhSeCt1GbUy"
       },
       "source": [
-        "## Customize image generation\n",
+        "### Customize image generation\n",
         "\n",
         "You may customize the image generation."
       ]
@@ -493,6 +529,8 @@
         "                        \"save_checkpoints_every_n\": save_checkpoints_every_n,\n",
         "                        \"batch_size\": batch_size,\n",
         "                        \"random_flip\": random_flip,\n",
+        "                        \"random_crop\": random_crop,\n",
+        "                        \"random_color_jitter\": random_color_jitter,\n",
         "                        \"lora_rank\": lora_rank,\n",
         "                    }\n",
         "                ),\n",
@@ -525,7 +563,7 @@
         "id": "Jdm1dOgsMH9O"
       },
       "source": [
-        "# Download generated images\n",
+        "## Download generated images\n",
         "\n",
         "You can download and preview the generated images."
       ]
@@ -545,9 +583,10 @@
         "\n",
         "def copy_image(images_source, images_dest):\n",
         "    os.makedirs(images_dest, exist_ok=True)\n",
-        "    ! gsutil cp -r {images_source} {images_dest}\n",
+        "    ! gsutil cp -r {images_source}/* {images_dest}\n",
         "\n",
-        "local_image_path = \"/images\"\n",
+        "\n",
+        "local_image_path = \"./images/\"\n",
         "copy_image(IMAGE_EXPORT_PATH, local_image_path)\n",
         "\n",
         "for filename in os.listdir(local_image_path):\n",
@@ -581,12 +620,13 @@
         "\n",
         "def copy_model(model_source, model_dest):\n",
         "    os.makedirs(model_dest, exist_ok=True)\n",
-        "    ! gsutil cp -r {model_source} {model_dest}\n",
+        "    ! gsutil -m cp -r {model_source}/* {model_dest}\n",
+        "\n",
         "\n",
         "local_model_path = \"/models\"\n",
         "copy_model(MODEL_EXPORT_PATH, local_model_path)\n",
         "\n",
-        "! tar czf models.tar.gz {local_model_path}\n",
+        "! tar czf models.tar.gz {local_model_path}/*\n",
         "\n",
         "if \"google.colab\" in sys.modules:\n",
         "    from google.colab import files\n",
@@ -594,6 +634,132 @@
         "    files.download(\"models.tar.gz\")"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-gqs4Kc4u6U3"
+      },
+      "source": [
+        "## Use trained models on Vertex AI\n",
+        "\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. Get [online predictions](https://cloud.google.com/vertex-ai/docs/predictions/get-online-predictions) from the deployed model"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Y_HV9vMbvF_6"
+      },
+      "source": [
+        "### Upload model to Vertex AI Model Registry"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "OL6wdXf-1rta"
+      },
+      "outputs": [],
+      "source": [
+        "serving_env = {\n",
+        "    \"TASK\": \"text-to-image-lora\",\n",
+        "    \"BASE_MODEL_ID\": \"runwayml/stable-diffusion-v1-5\",\n",
+        "    \"FINETUNED_LORA_MODEL_PATH\": MODEL_EXPORT_PATH,\n",
+        "}\n",
+        "\n",
+        "model = aiplatform.Model.upload(\n",
+        "    display_name=UPLOAD_MODEL_NAME,\n",
+        "    serving_container_image_uri=PREDICTION_CONTAINER_URI,\n",
+        "    serving_container_ports=[PREDICTION_PORT],\n",
+        "    serving_container_predict_route=\"/predictions/peft_serving\",\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: \", UPLOAD_MODEL_NAME)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "NjXg3QHw2GCT"
+      },
+      "source": [
+        "### Deploy the uploaded model\n",
+        "\n",
+        "You will deploy models in Google Cloud Vertex AI. The default setting will use 1 V100 GPU for deployment.\n",
+        "\n",
+        "Please create a Service Account for serving with dockers if you do not have one yet.\n",
+        "\n",
+        "The model deployment will take around 1 minute to finish."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "BLTA7nRcw57D"
+      },
+      "outputs": [],
+      "source": [
+        "# 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\"}\n",
+        "\n",
+        "endpoint = aiplatform.Endpoint.create(display_name=f\"{UPLOAD_MODEL_NAME}-endpoint\")\n",
+        "model.deploy(\n",
+        "    endpoint=endpoint,\n",
+        "    machine_type=PREDICTION_MACHINE_TYPE,\n",
+        "    accelerator_type=PREDICTION_ACCELERATOR_TYPE,\n",
+        "    accelerator_count=1,\n",
+        "    deploy_request_timeout=1800,\n",
+        "    service_account=service_account,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tMfwrH32_uIM"
+      },
+      "source": [
+        "The docker container still needs to download and load the model after the endpoint is created. Therefore, we recommend waiting for 3 extra minutes before proceeding to the next cell.\n",
+        "\n",
+        "Once deployed, you can send a batch of text prompts to the endpoint to generate images."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "wr7wg9JQ_tjM"
+      },
+      "outputs": [],
+      "source": [
+        "import base64\n",
+        "from io import BytesIO\n",
+        "\n",
+        "import matplotlib.pyplot as plt\n",
+        "from PIL import Image\n",
+        "\n",
+        "instances = [\n",
+        "    {\"prompt\": \"Two monadikos teapots on a table\"},\n",
+        "    {\"prompt\": \"Two monadikos teapots on the floor\"},\n",
+        "]\n",
+        "response = endpoint.predict(instances=instances)\n",
+        "\n",
+        "plt.figure()\n",
+        "_, grid = plt.subplots(1, len(instances))\n",
+        "for cell, prediction in zip(grid, response.predictions):\n",
+        "    image = Image.open(BytesIO(base64.b64decode(prediction)))\n",
+        "    cell.imshow(image)"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -613,7 +779,10 @@
       "outputs": [],
       "source": [
         "if training_job.list(filter=f'display_name=\"{TRAINING_JOB_DISPLAY_NAME}\"'):\n",
-        "    training_job.delete()"
+        "    training_job.delete()\n",
+        "# Undeploys models and deletes endpoints.\n",
+        "endpoint.delete(force=True)\n",
+        "model.delete()"
       ]
     },
     {

notebooks/community/model_garden/model_garden_movinet_action_recognition.ipynb

@@ -170,6 +170,8 @@
         "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",
+        "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 running batch predictions with the fine tuned model.\n",
+        "\n",
         "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
       ]
     },
@@ -248,9 +250,9 @@
         "# Training constants.\n",
         "TRAINING_JOB_PREFIX = \"train\"\n",
         "TRAIN_CONTAINER_URI = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/movinet-train\"\n",
-        "TRAIN_MACHINE_TYPE = \"n1-highmem-16\"\n",
+        "TRAIN_MACHINE_TYPE = \"n1-highmem-32\"\n",
         "TRAIN_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
-        "TRAIN_NUM_GPU = 2\n",
+        "TRAIN_NUM_GPU = 8\n",
         "\n",
         "# Evaluation constants.\n",
         "EVALUATION_METRIC = \"accuracy\"\n",
@@ -262,7 +264,6 @@
         "\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/vertex-vision-model-garden-dockers/movinet-serve\"\n",
         "PREDICTION_PORT = 8501\n",
         "PREDICTION_ACCELERATOR_COUNT = 1\n",
@@ -291,12 +292,13 @@
       "source": [
         "import json\n",
         "from datetime import datetime\n",
+        "from typing import Any\n",
         "\n",
         "import tensorflow as tf\n",
         "import yaml\n",
         "\n",
         "\n",
-        "def get_job_name_with_datetime(prefix: str):\n",
+        "def get_job_name_with_datetime(prefix: str) -> 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",
@@ -309,7 +311,9 @@
         "    return label_map, num_classes\n",
         "\n",
         "\n",
-        "def get_best_trial(model_dir, max_trial_count, evaluation_metric):\n",
+        "def get_best_trial(\n",
+        "    model_di: str, max_trial_count: int, evaluation_metric: str\n",
+        ") -> tuple[str, Any]:\n",
         "    \"\"\"Finds the best trial directory and eval results from a hyperparameter tuning job.\"\"\"\n",
         "    best_trial_dir = \"\"\n",
         "    best_trial_evaluation_results = {}\n",
@@ -351,7 +355,7 @@
         "        print(f\"At {time}, detected {label} with {confidence} confidence.\")\n",
         "\n",
         "\n",
-        "def find_checkpoint_in_dir(checkpoint_dir: str):\n",
+        "def find_checkpoint_in_dir(checkpoint_dir: str) -> 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",
@@ -359,7 +363,7 @@
         "                return os.path.join(root, os.path.splitext(file)[0])\n",
         "\n",
         "\n",
-        "def upload_checkpoint_to_gcs(checkpoint_url: str):\n",
+        "def upload_checkpoint_to_gcs(checkpoint_url: str) -> 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",
@@ -377,7 +381,7 @@
         "    return checkpoint_uri\n",
         "\n",
         "\n",
-        "def upload_config_to_gcs(url: str):\n",
+        "def upload_config_to_gcs(url: str) -> 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",
@@ -403,7 +407,7 @@
         "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_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 `a3` for demonstration purpose. **`a0`, `a1`, and `a2` are not recommended for now as we are currently investigating some inference issues with them.**\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."
@@ -417,14 +421,14 @@
       },
       "outputs": [],
       "source": [
-        "model_id = \"a0\"  # @param [\"a0\", \"a1\", \"a2\", \"a3\", \"a4\", \"a5\"]\n",
+        "model_id = \"a3\"  # @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}\"\n",
         "image_size = IMAGE_SIZES[model_id]\n",
         "\n",
         "if is_stream:\n",
-        "    train_container_args = {\n",
+        "    export_container_args = {\n",
         "        \"conv_type\": \"2plus1d\",\n",
         "        \"se_type\": \"2plus3d\",\n",
         "        \"activation\": \"hard_swish\",\n",
@@ -432,7 +436,7 @@
         "        \"use_positional_encoding\": model_id in {\"a3\", \"a4\", \"a5\"},\n",
         "    }\n",
         "else:\n",
-        "    train_container_args = {\n",
+        "    export_container_args = {\n",
         "        \"conv_type\": \"3d\",\n",
         "        \"se_type\": \"3d\",\n",
         "        \"activation\": \"swish\",\n",
@@ -454,13 +458,12 @@
         "\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.\n",
         "* `num_frames`: Number of frame to sample around keyframe inputs.\n",
         "* `min_duration_sec`: Minimum duration in seconds for sampling video clips around keyframe inputs. This is for validation purpose - an error will be thrown if there is not enough context around a keyframe.\n",
         "* `pos_neg_ratio`: Sampling ratio between positive and negative segments. For example, a pos_neg_ratio of 0.5 samples 1 negative instance every 2 positive instances.\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",
         "\n",
         "**Note**: For JSONL input, please use `aiplatform.googleapis.com/ml_use` instead of `ml_use` as the JSON key for ML use in `dataItemResourceLabels`. This is to be consistent with other objectives."
       ]
@@ -546,7 +549,7 @@
         "\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",
+        "* `worker_pool_specs`: A list of dictionaries specifying the machine type and docker image. This example defines a single node cluster with one `n1-highmem-32` machine with 8 `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",
@@ -593,7 +596,7 @@
         "    \"init_checkpoint\": init_checkpoint,\n",
         "    \"model_dir\": model_dir,\n",
         "    \"num_classes\": num_classes,\n",
-        "    \"global_batch_size\": 8,\n",
+        "    \"global_batch_size\": 16,\n",
         "    \"prefetch_buffer_size\": 16,\n",
         "    \"shuffle_buffer_size\": 32,\n",
         "    \"train_steps\": 2000,\n",
@@ -670,7 +673,7 @@
         "    metric_spec=metric_spec,\n",
         "    parameter_spec=parameter_spec,\n",
         "    max_trial_count=MAX_TRIAL_COUNT,\n",
-        "    parallel_trial_count=MAX_TRIAL_COUNT,\n",
+        "    parallel_trial_count=1,\n",
         "    project=PROJECT_ID,\n",
         "    search_algorithm=None,\n",
         ")\n",
@@ -715,7 +718,7 @@
         "    \"causal\": is_stream,\n",
         "    \"checkpoint_path\": best_checkpoint_path,\n",
         "    \"assert_checkpoint_objects_matched\": False,\n",
-        "    **train_container_args,\n",
+        "    **export_container_args,\n",
         "}\n",
         "\n",
         "worker_pool_specs = [\n",
@@ -766,7 +769,17 @@
         "id": "gdlca3BOypXU"
       },
       "source": [
-        "### Upload model to Vertex AI Model Registry"
+        "### Upload model to Vertex AI Model Registry\n",
+        "\n",
+        "The following cell uploads the trained model to Vertex AI Model Registry. Skip it if you want to run batch predictions on an already uploaded model instead.\n",
+        "\n",
+        "#### Configurable environment variables\n",
+        "\n",
+        "* `MODEL_PATH`: Cloud Storage URI to the MoViNet model.\n",
+        "* `BATCH_SIZE`: Batch size for inference. Use a larger value to accelerate GPU prediction.\n",
+        "* `NUM_FRAMES`: Number of frames for a single prediction with the model.\n",
+        "* `FPS`: Video sampling frame per second.\n",
+        "* `OVERLAP_FRAMES`: Allowed overlapping frames between consecutive prediction windows. Set a smaller value for faster inference but less accurate."
       ]
     },
     {
@@ -777,11 +790,9 @@
       },
       "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",
+        "    \"BATCH_SIZE\": 1,\n",
         "    \"NUM_FRAMES\": 32,\n",
         "    \"FPS\": output_fps,\n",
         "    \"OVERLAP_FRAMES\": 24,\n",
@@ -804,6 +815,26 @@
         "print(\"The uploaded model name is: \", model_name)"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "c2b47e629a01"
+      },
+      "source": [
+        "Alternatively, uncomment the following cell to use an already uploaded model. Replace the model name string with that of the existing model."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "aa503a565b8f"
+      },
+      "outputs": [],
+      "source": [
+        "# model = aiplatform.Model(\"projects/123456789/locations/us-central1/models/12345678901234567890\")"
+      ]
+    },
     {
       "attachments": {},
       "cell_type": "markdown",
@@ -815,7 +846,7 @@
         "\n",
         "We will now run batch predictions with the trained MoViNet action recognition 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/action-recognition/get-predictions?hl=en#input_data_requirements) and store it in a Cloud Storage bucket.\n",
+        "Please prepare an input JSONL file where each line follows [this format](https://cloud.google.com/vertex-ai/docs/video-data/action-recognition/get-predictions?hl=en#input_data_requirements) and store it in a Cloud Storage bucket. The service account should have read access to the buckets containing the trained model and the input data. See [Service accounts overview](https://cloud.google.com/iam/docs/service-account-overview) for more information.\n",
         "\n",
         "The [Vertex AI Batch Prediction](https://cloud.google.com/vertex-ai/docs/predictions/get-batch-predictions) has a default timeout of 10 minutes. Therefore, please make sure the input video clip is around 5 minutes at 5~10 FPS or you may experience a timeout error. To use this model at a larger scale beyond this notebook demontration, you can try one of the following:\n",
         "\n",
@@ -834,6 +865,8 @@
       "source": [
         "# Path to the prediction input JSONL file.\n",
         "test_jsonl_path = \"\"  # @param {type:\"string\"}\n",
+        "# Full service account name with the suffix `gserviceaccount.com`.\n",
+        "batch_predict_service_account = \"\"  # @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",
@@ -846,6 +879,7 @@
         "    accelerator_count=PREDICTION_ACCELERATOR_COUNT,\n",
         "    accelerator_type=PREDICTION_ACCELERATOR_TYPE,\n",
         "    max_replica_count=1,\n",
+        "    service_account=batch_predict_service_account,\n",
         ")\n",
         "\n",
         "batch_prediction_job.wait()\n",
@@ -873,8 +907,6 @@
       },
       "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",

notebooks/community/model_garden/model_garden_movinet_clip_classification.ipynb

@@ -164,6 +164,8 @@
         "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",
+        "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 running batch predictions with the fine tuned model.\n",
+        "\n",
         "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
       ]
     },
@@ -247,7 +249,6 @@
         "\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/vertex-vision-model-garden-dockers/movinet-serve\"\n",
         "PREDICTION_PORT = 8501\n",
         "PREDICTION_ACCELERATOR_COUNT = 1\n",
@@ -275,13 +276,14 @@
       "source": [
         "import json\n",
         "from datetime import datetime\n",
+        "from typing import Any\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",
+        "def get_job_name_with_datetime(prefix: str) -> 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",
@@ -309,7 +311,9 @@
         "    return label_map, num_classes\n",
         "\n",
         "\n",
-        "def get_best_trial(model_dir, max_trial_count, evaluation_metric):\n",
+        "def get_best_trial(\n",
+        "    model_di: str, max_trial_count: int, evaluation_metric: str\n",
+        ") -> tuple[str, Any]:\n",
         "    \"\"\"Finds the best trial directory and eval results from a hyperparameter tuning job.\"\"\"\n",
         "    best_trial_dir = \"\"\n",
         "    best_trial_evaluation_results = {}\n",
@@ -332,7 +336,7 @@
         "    return best_trial_dir, best_trial_evaluation_results\n",
         "\n",
         "\n",
-        "def find_checkpoint_in_dir(checkpoint_dir: str):\n",
+        "def find_checkpoint_in_dir(checkpoint_dir: str) -> 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",
@@ -340,7 +344,7 @@
         "                return os.path.join(root, os.path.splitext(file)[0])\n",
         "\n",
         "\n",
-        "def upload_checkpoint_to_gcs(checkpoint_url: str):\n",
+        "def upload_checkpoint_to_gcs(checkpoint_url: str) -> 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",
@@ -358,7 +362,7 @@
         "    return checkpoint_uri\n",
         "\n",
         "\n",
-        "def upload_config_to_gcs(url: str):\n",
+        "def upload_config_to_gcs(url: str) -> 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",
@@ -400,7 +404,24 @@
         "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}\""
+        "model_name = f\"movinet_{model_id}_{model_mode}\"\n",
+        "\n",
+        "if is_stream:\n",
+        "    export_container_args = {\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",
+        "else:\n",
+        "    export_container_args = {\n",
+        "        \"conv_type\": \"3d\",\n",
+        "        \"se_type\": \"3d\",\n",
+        "        \"activation\": \"swish\",\n",
+        "        \"gating_activation\": \"sigmoid\",\n",
+        "        \"use_positional_encoding\": False,\n",
+        "    }"
       ]
     },
     {
@@ -415,10 +436,9 @@
         "\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",
+        "* `output_fps`: The sampling rate of the video; Frames per second.\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."
+        "* `num_shard`: Three comma separated integers indicating the shards for train/validation/test."
       ]
     },
     {
@@ -495,7 +515,7 @@
         "\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",
+        "* `worker_pool_specs`: A list of dictionaries specifying the machine type and docker image. This example defines a single node cluster with one `n1-highmem-16` 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",
@@ -662,19 +682,9 @@
         "    \"causal\": is_stream,\n",
         "    \"checkpoint_path\": best_checkpoint_path,\n",
         "    \"assert_checkpoint_objects_matched\": False,\n",
+        "    **export_container_args,\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",
@@ -721,7 +731,17 @@
         "id": "gdlca3BOypXU"
       },
       "source": [
-        "### Upload model to Vertex AI Model Registry"
+        "### Upload model to Vertex AI Model Registry\n",
+        "\n",
+        "The following cell uploads the trained model to Vertex AI Model Registry. Skip it if you want to run batch predictions on an already uploaded model instead.\n",
+        "\n",
+        "#### Configurable environment variables\n",
+        "\n",
+        "* `MODEL_PATH`: Cloud Storage URI to the MoViNet model.\n",
+        "* `BATCH_SIZE`: Batch size for inference. Use a larger value to accelerate GPU prediction.\n",
+        "* `NUM_FRAMES`: Number of frames for a single prediction with the model.\n",
+        "* `FPS`: Video sampling frame per second.\n",
+        "* `OVERLAP_FRAMES`: Allowed overlapping frames between consecutive prediction windows. Set a smaller value for faster inference but less accurate."
       ]
     },
     {
@@ -732,8 +752,6 @@
       },
       "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",
@@ -757,6 +775,26 @@
         "print(\"The uploaded model name is: \", model_name)"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "c2b47e629a01"
+      },
+      "source": [
+        "Alternatively, uncomment the following cell to use an already uploaded model. Replace the model name string with that of the existing model."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "aa503a565b8f"
+      },
+      "outputs": [],
+      "source": [
+        "# model = aiplatform.Model(\"projects/123456789/locations/us-central1/models/12345678901234567890\")"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -767,7 +805,7 @@
         "\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."
+        "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. The service account should have read access to the buckets containing the trained model and the input data. See [Service accounts overview](https://cloud.google.com/iam/docs/service-account-overview) for more information."
       ]
     },
     {
@@ -780,6 +818,8 @@
       "source": [
         "# Path to the prediction input JSONL file.\n",
         "test_jsonl_path = \"\"  # @param {type:\"string\"}\n",
+        "# Full service account name with the suffix `gserviceaccount.com`.\n",
+        "batch_predict_service_account = \"\"  # @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",
@@ -792,6 +832,7 @@
         "    accelerator_count=PREDICTION_ACCELERATOR_COUNT,\n",
         "    accelerator_type=PREDICTION_ACCELERATOR_TYPE,\n",
         "    max_replica_count=1,\n",
+        "    service_account=batch_predict_service_account,\n",
         ")\n",
         "\n",
         "batch_prediction_job.wait()\n",
@@ -818,8 +859,6 @@
       },
       "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",

notebooks/community/model_garden/model_garden_pytorch_codellama.ipynb

@@ -164,15 +164,16 @@
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
         "# The Cloud Storage bucket for storing experiments output.\n",
-        "BUCKET_URI = \"\"  # @param {type:\"string\"}\n",
+        "# Start with gs:// prefix, e.g. gs://foo_bucket.\n",
+        "GCS_BUCKET = \"\"  # @param {type:\"string\"}\n",
         "\n",
         "! gcloud config set project $PROJECT_ID\n",
         "! gcloud services enable language.googleapis.com\n",
         "\n",
         "import os\n",
         "\n",
-        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"staging\")\n",
-        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"code-llama\")\n",
+        "STAGING_BUCKET = os.path.join(GCS_BUCKET, \"staging\")\n",
+        "MODEL_BUCKET = os.path.join(GCS_BUCKET, \"code-llama\")\n",
         "\n",
         "# The service account looks like:\n",
         "# '@.iam.gserviceaccount.com'\n",
@@ -345,7 +346,7 @@
       "outputs": [],
       "source": [
         "model_name = \"CodeLlama-7b-Instruct-hf\"  # @param [\"CodeLlama-7b-hf\", \"CodeLlama-7b-Python-hf\", \"CodeLlama-7b-Instruct-hf\", \"CodeLlama-13b-hf\", \"CodeLlama-13b-Python-hf\", \"CodeLlama-13b-Instruct-hf\", \"CodeLlama-34b-hf\", \"CodeLlama-34b-Python-hf\", \"CodeLlama-34b-Instruct-hf\"]\n",
-        "model_id = \"gs://\" + os.path.join(MODEL_BUCKET, model_name)\n",
+        "model_id = os.path.join(MODEL_BUCKET, model_name)\n",
         "print(model_id)"
       ]
     },
@@ -383,7 +384,7 @@
         "# accelerator_count = 1\n",
         "\n",
         "# Sets 2 A100 (40G) to deploy 34B models.\n",
-        "# machine_type = \"a2-highgpu-4g\"\n",
+        "# machine_type = \"a2-highgpu-2g\"\n",
         "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
         "# accelerator_count = 2\n",
         "\n",

notebooks/community/model_garden/model_garden_pytorch_falcon_instruct_peft.ipynb

@@ -47,7 +47,7 @@
         "    <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",
+        "    </a> (A Python-3 GPU notebook is recommended)\n",
         "  </td>\n",
         "</table>"
       ]
@@ -60,13 +60,14 @@
       "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",
+        "This notebook demonstrates running inferences locally with prebuilt Falcon Instruct models, deploying prebuilt Falcon Instruct models, finetuning and deploying Falcon Instruct models with performance efficient finetuning libraries ([PEFT](https://github.com/huggingface/peft)), and evaluating PEFT-finetuned Falcon Instruct models in Vertex AI.\n",
         "\n",
         "### Objective\n",
         "\n",
         "- Run inferences locally on prebuilt Falcon Instruct models\n",
         "- Deploy prebuilt Falcon Instruct models\n",
-        "- Finetune and deploy Falcon Instruct models with PEFT, supporting\n",
+        "- Finetune and deploy Falcon Instruct models with PEFT\n",
+        "- Evaluate PEFT-finetuned Falcon Instruct models\n",
         "\n",
         "| Models | LoRA |\n",
         "| :- | :- |\n",
@@ -91,7 +92,9 @@
       "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."
+        "**NOTE**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands.\n",
+        "\n",
+        "Running inferences locally with Falcon Instruct models requires a GPU."
       ]
     },
     {
@@ -157,7 +160,7 @@
         "id": "6c460088b873"
       },
       "source": [
-        "Fill following variables for experiments environment:"
+        "Set the following variables for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the specified region (`REGION`). Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\")."
       ]
     },
     {
@@ -174,8 +177,8 @@
         "# 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",
+        "# The Cloud Storage bucket for storing experiments output with gs:// prefix.\n",
+        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
         "\n",
         "! gcloud config set project $PROJECT_ID\n",
         "\n",
@@ -233,9 +236,14 @@
       },
       "outputs": [],
       "source": [
-        "# The pre-built training and serving docker images.\n",
-        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train\"\n",
-        "PREDICTION_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-serve\""
+        "# The pre-built training, serving and evaluation docker images.\n",
+        "TRAIN_DOCKER_URI = (\n",
+        "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train\"\n",
+        ")\n",
+        "PREDICTION_DOCKER_URI = (\n",
+        "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-serve\"\n",
+        ")\n",
+        "EVAL_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-lm-evaluation-harness\""
       ]
     },
     {
@@ -261,7 +269,7 @@
         "from google.cloud import aiplatform\n",
         "\n",
         "\n",
-        "def get_job_name_with_datetime(prefix: str):\n",
+        "def get_job_name_with_datetime(prefix: str) -> str:\n",
         "    \"\"\"Gets the job name with date time when triggering training or deployment\n",
         "    jobs in Vertex AI.\n",
         "    \"\"\"\n",
@@ -269,14 +277,15 @@
         "\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",
+        "    model_name: str,\n",
+        "    base_model_id: str,\n",
+        "    finetuned_lora_model_path: str,\n",
+        "    service_account: str,\n",
+        "    task: str,\n",
+        "    machine_type: str = \"n1-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_TESLA_V100\",\n",
+        "    accelerator_count: int = 1,\n",
+        ") -> tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
         "    serving_env = {\n",
@@ -297,7 +306,7 @@
         "        endpoint=endpoint,\n",
         "        machine_type=machine_type,\n",
         "        accelerator_type=accelerator_type,\n",
-        "        accelerator_count=1,\n",
+        "        accelerator_count=accelerator_count,\n",
         "        deploy_request_timeout=1800,\n",
         "        service_account=service_account,\n",
         "    )\n",
@@ -360,7 +369,7 @@
         "\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."
+        "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 ~84G separately with the default settings. Please adjust the machine type, accelerator type and accelerator count accordingly. We use V100 in deployments as an example. Please use A100 (40G) or A100 (80G) to get better inferences."
       ]
     },
     {
@@ -400,17 +409,28 @@
       },
       "outputs": [],
       "source": [
+        "# Find Vertex AI supported accelerators and regions in:\n",
+        "#  https://cloud.google.com/vertex-ai/docs/predictions/configure-compute\n",
+        "\n",
         "# Sets V100 to deploy tiiuae/falcon-7b-instruct as an example.\n",
+        "# If A100 is not available, you may deploy tiiuae/falcon-40b-instruct with\n",
+        "#  multiple V100s. Please keep in mind that the efficiency of serving with\n",
+        "#  multiple V100s is inferior to that of serving with A100s.\n",
         "machine_type = \"n1-standard-8\"\n",
         "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
+        "accelerator_count = 1  # for tiiuae/falcon-7b-instruct\n",
+        "# accelerator_count = 8  # for tiiuae/falcon-40b-instruct\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",
+        "# accelerator_count = 1  # for tiiuae/falcon-7b-instruct\n",
+        "# accelerator_count = 4  # for tiiuae/falcon-40b-instruct\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",
+        "# accelerator_count = 2\n",
         "\n",
         "model_without_peft, endpoint_without_peft = deploy_model(\n",
         "    model_name=get_job_name_with_datetime(prefix=\"falcon-instruct-serve\"),\n",
@@ -420,6 +440,7 @@
         "    task=\"instruct-lora\",\n",
         "    machine_type=machine_type,\n",
         "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
         ")\n",
         "print(\"endpoint_name:\", endpoint_without_peft.name)"
       ]
@@ -430,7 +451,7 @@
         "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",
+        "NOTE: The prebuilt model weights will be downloaded on the fly from the original location after the deployment succeeds. Thus, an additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you can 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",
@@ -473,7 +494,7 @@
       "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."
+        "This section demonstrates how to finetune and deploy Falcon Instruct models with PEFT LoRA."
       ]
     },
     {
@@ -529,7 +550,10 @@
         "dataset_name = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
         "\n",
         "# Worker pool spec.\n",
-        "# Uses V100 to finetune falcon-70b-instruct.\n",
+        "# Find Vertex AI supported accelerators and regions in:\n",
+        "#  https://cloud.google.com/vertex-ai/docs/training/configure-compute\n",
+        "\n",
+        "# Uses V100 to finetune falcon-7b-instruct.\n",
         "machine_type = \"n1-standard-8\"\n",
         "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
         "# Uses A100 to finetune falcon-40b-instruct.\n",
@@ -584,7 +608,7 @@
         "\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."
+        "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 ~84G separately with the default settings. Please adjust the machine type, accelerator type and accelerator count accordingly. We use V100 in deployments as an example. Please use A100 (40G) or A100 (80G) to get better inferences."
       ]
     },
     {
@@ -595,17 +619,28 @@
       },
       "outputs": [],
       "source": [
+        "# Find Vertex AI supported accelerators and regions in:\n",
+        "#  https://cloud.google.com/vertex-ai/docs/predictions/configure-compute\n",
+        "\n",
         "# Sets V100 to deploy tiiuae/falcon-7b-instruct as an example.\n",
+        "# If A100 is not available, you may deploy tiiuae/falcon-40b-instruct with\n",
+        "#  multiple V100s. Please keep in mind that the efficiency of serving with\n",
+        "#  multiple V100s is inferior to that of serving with A100s.\n",
         "machine_type = \"n1-standard-8\"\n",
         "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
+        "accelerator_count = 1  # for tiiuae/falcon-7b-instruct\n",
+        "# accelerator_count = 8  # for tiiuae/falcon-40b-instruct\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",
+        "# accelerator_count = 1  # for tiiuae/falcon-7b-instruct\n",
+        "# accelerator_count = 4  # for tiiuae/falcon-40b-instruct\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",
+        "# accelerator_count = 2\n",
         "\n",
         "model_with_peft, endpoint_with_peft = deploy_model(\n",
         "    model_name=get_job_name_with_datetime(prefix=\"falcon-instruct-peft-serve\"),\n",
@@ -615,6 +650,7 @@
         "    task=\"instruct-lora\",\n",
         "    machine_type=machine_type,\n",
         "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
         ")\n",
         "print(\"endpoint_name:\", endpoint_with_peft.name)"
       ]
@@ -625,7 +661,7 @@
         "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",
+        "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, an additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you can 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",
@@ -660,6 +696,166 @@
         "    print(prediction[\"generated_text\"])"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4ffffac5b20a"
+      },
+      "source": [
+        "## Evaluate PEFT-finetuned Falcon Instruct models\n",
+        "\n",
+        "This section demonstrates how to evaluate the Falcon Instruct models fintuned with PEFT LoRA using EleutherAI's [Language Model Evaluation Harness (lm-evaluation-harness)](https://github.com/EleutherAI/lm-evaluation-harness) with Vertex CustomJob. Please reference the peak GPU memory usgaes for serving and adjust the machine type, accelerator type and accelerator count accordingly.\n",
+        "\n",
+        "This example uses the dataset [HellaSwag](https://allenai.org/data/hellaswag). All supported tasks are listed in [this task table](https://github.com/EleutherAI/lm-evaluation-harness/blob/master/docs/task_table.md)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "078b4d178624"
+      },
+      "outputs": [],
+      "source": [
+        "eval_dataset = \"hellaswag\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Worker pool spec.\n",
+        "# Find Vertex AI supported accelerators and regions in:\n",
+        "#  https://cloud.google.com/vertex-ai/docs/training/configure-compute\n",
+        "\n",
+        "# Sets V100 to deploy tiiuae/falcon-7b-instruct as an example.\n",
+        "# If A100 is not available, you may deploy tiiuae/falcon-40b-instruct with\n",
+        "#  multiple V100s. Please keep in mind that the efficiency of serving with\n",
+        "#  multiple V100s is inferior to that of serving with A100s.\n",
+        "machine_type = \"n1-standard-8\"\n",
+        "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
+        "accelerator_count = 1  # for tiiuae/falcon-7b-instruct\n",
+        "# accelerator_count = 8  # for tiiuae/falcon-40b-instruct\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",
+        "# accelerator_count = 1  # for tiiuae/falcon-7b-instruct\n",
+        "# accelerator_count = 4  # for tiiuae/falcon-40b-instruct\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",
+        "# accelerator_count = 2\n",
+        "\n",
+        "replica_count = 1\n",
+        "\n",
+        "# Setup evaluation job.\n",
+        "job_name = get_job_name_with_datetime(prefix=\"falcon-instruct-peft-eval\")\n",
+        "eval_output_dir = os.path.join(MODEL_BUCKET, job_name)\n",
+        "eval_output_dir_gcsfuse = eval_output_dir.replace(\"gs://\", \"/gcs/\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "c5e32112815e"
+      },
+      "outputs": [],
+      "source": [
+        "# Prepare evaluation script that runs the evaluation harness.\n",
+        "# We set `trust_remote_code = True` because evaluating the model requires\n",
+        "#  executing code from the model repository.\n",
+        "# We set `use_accelerate = True` to enable evaluation across multiple GPUs.\n",
+        "script_path = \"./eval_script.py\"  # @param {type:\"string\"}\n",
+        "\n",
+        "eval_command = f\"\"\"import subprocess\n",
+        "\n",
+        "\n",
+        "subprocess.call([\n",
+        "    'python',\n",
+        "    'main.py',\n",
+        "    '--model',\n",
+        "    'hf-causal-experimental',\n",
+        "    '--model_args',\n",
+        "    'pretrained={base_model_id},peft={output_dir_gcsfuse},trust_remote_code=True,use_accelerate=True,device_map_option=auto',\n",
+        "    '--tasks',\n",
+        "    '{eval_dataset}',\n",
+        "    '--output_path',\n",
+        "    '{eval_output_dir_gcsfuse}',\n",
+        "])\n",
+        "\"\"\"\n",
+        "\n",
+        "with open(script_path, \"w\") as fp:\n",
+        "    fp.write(eval_command)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "afd9305771ba"
+      },
+      "source": [
+        "### Submit evaluation CustomJob"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "602413896b58"
+      },
+      "outputs": [],
+      "source": [
+        "# Pass evaluation arguments and launch job.\n",
+        "eval_job = aiplatform.CustomJob.from_local_script(\n",
+        "    display_name=job_name,\n",
+        "    script_path=script_path,\n",
+        "    container_uri=EVAL_DOCKER_URI,\n",
+        "    replica_count=replica_count,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    base_output_dir=eval_output_dir,\n",
+        "    boot_disk_size_gb=500,\n",
+        ")\n",
+        "\n",
+        "eval_job.run()\n",
+        "\n",
+        "print(\"Evaluation results were saved in:\", eval_output_dir)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "de47e182a37e"
+      },
+      "source": [
+        "### Fetch and print evaluation results"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "1f15ed6d375a"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "\n",
+        "from google.cloud import storage\n",
+        "\n",
+        "# Fetch evaluation results.\n",
+        "storage_client = storage.Client()\n",
+        "BUCKET_NAME = BUCKET_URI.split(\"gs://\")[1]\n",
+        "bucket = storage_client.get_bucket(BUCKET_NAME)\n",
+        "RESULT_FILE_PATH = eval_output_dir[len(BUCKET_URI) + 1 :]\n",
+        "blob = bucket.blob(RESULT_FILE_PATH)\n",
+        "raw_result = blob.download_as_string()\n",
+        "\n",
+        "# Print evaluation results.\n",
+        "result = json.loads(raw_result)\n",
+        "result_formatted = json.dumps(result, indent=2)\n",
+        "print(f\"Evaluation result:\\n{result_formatted}\")"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -677,10 +873,11 @@
       },
       "outputs": [],
       "source": [
-        "# Delete custom train jobs.\n",
+        "# Delete custom train and evaluation jobs.\n",
         "train_job.delete()\n",
+        "eval_job.delete()\n",
         "\n",
-        "# Undeploy model and delete endpoint.\n",
+        "# Undeploy models and delete endpoints.\n",
         "endpoint_without_peft.delete(force=True)\n",
         "endpoint_with_peft.delete(force=True)\n",
         "\n",

notebooks/community/model_garden/model_garden_pytorch_openllama_peft.ipynb

@@ -47,7 +47,7 @@
         "    <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",
+        "    </a> (A Python-3 GPU notebook is recommended)\n",
         "  </td>\n",
         "</table>"
       ]
@@ -60,14 +60,15 @@
       "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",
+        "This notebook demonstrates running local inference with prebuilt OpenLLaMA, deploying prebuilt OpenLLaMA, deploying prebuilt OpenLLaMA with [vLLM](https://github.com/vllm-project/vllm), finetuning and deploying OpenLLaMA with performance efficient finetuning libraries ([PEFT](https://github.com/huggingface/peft)), and evaluating PEFT-finetuned OpenLLaMA in Vertex AI.\n",
         "\n",
         "### Objective\n",
         "\n",
         "- Run local inference with prebuilt OpenLLaMA\n",
         "- Deploy prebuilt OpenLLaMA\n",
-        "- Finetune and deploy OpenLLaMA with PEFT, supporting\n",
-        "- Deploy OpenLLaMA with [vLLM](https://github.com/vllm-project/vllm) to improve serving throughput\n",
+        "- Deploy prebuilt OpenLLaMA with [vLLM](https://github.com/vllm-project/vllm) to improve serving throughput\n",
+        "- Finetune and deploy OpenLLaMA with PEFT\n",
+        "- Evaluate finetuned OpenLLaMA with PEFT\n",
         "\n",
         "| Models | LoRA |\n",
         "| :- | :- |\n",
@@ -93,7 +94,9 @@
       "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."
+        "**NOTE**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands.\n",
+        "\n",
+        "Running local inference with OpenLLaMA requires a GPU."
       ]
     },
     {
@@ -159,7 +162,7 @@
         "id": "6c460088b873"
       },
       "source": [
-        "Fill following variables for experiments environment:"
+        "Set the following variables for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the specified region (`REGION`). Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\")."
       ]
     },
     {
@@ -177,6 +180,7 @@
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
         "# The Cloud Storage bucket for storing experiments output.\n",
+        "# Start with gs:// prefix, e.g. gs://foo_bucket.\n",
         "BUCKET_URI = \"\"  # @param {type:\"string\"}\n",
         "\n",
         "! gcloud config set project $PROJECT_ID\n",
@@ -236,11 +240,17 @@
       "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\"\n",
+        "TRAIN_DOCKER_URI = (\n",
+        "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train\"\n",
+        ")\n",
+        "PREDICTION_DOCKER_URI = (\n",
+        "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-serve\"\n",
+        ")\n",
         "VLLM_DOCKER_URI = (\n",
         "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve\"\n",
-        ")"
+        ")\n",
+        "\n",
+        "EVAL_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-lm-evaluation-harness\""
       ]
     },
     {
@@ -260,13 +270,12 @@
       },
       "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",
+        "def get_job_name_with_datetime(prefix: str) -> str:\n",
         "    \"\"\"Gets the job name with date time when triggering training or deployment\n",
         "    jobs in Vertex AI.\n",
         "    \"\"\"\n",
@@ -274,14 +283,14 @@
         "\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",
+        "    model_name: str,\n",
+        "    base_model_id: str,\n",
+        "    finetuned_lora_model_path: str,\n",
+        "    service_account: str,\n",
+        "    task: str,\n",
+        "    machine_type: str = \"n1-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_TESLA_V100\",\n",
+        ") -> tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
         "    serving_env = {\n",
@@ -310,13 +319,13 @@
         "\n",
         "\n",
         "def deploy_model_vllm(\n",
-        "    model_name,\n",
-        "    model_id,\n",
-        "    service_account,\n",
-        "    machine_type=\"n1-standard-8\",\n",
-        "    accelerator_type=\"NVIDIA_TESLA_V100\",\n",
-        "    accelerator_count=1,\n",
-        "):\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    service_account: str,\n",
+        "    machine_type: str = \"n1-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_TESLA_V100\",\n",
+        "    accelerator_count: int = 1,\n",
+        ") -> tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
         "\n",
@@ -453,7 +462,7 @@
         "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",
+        "NOTE: The prebuilt model weights will be downloaded on the fly from the original location after the deployment succeeds. Thus, an additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you can 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."
       ]
@@ -484,27 +493,47 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "8neJc8CnDDpu"
+        "id": "V7VOhhHGpUrj"
       },
       "source": [
         "## Deploy Prebuilt OpenLLaMA with vLLM\n",
         "\n",
         "This section deploys prebuilt OpenLLaMA models with [vLLM](https://github.com/vllm-project/vllm) on the Endpoint. The model deployment step will take ~15 minutes to complete.\n",
         "\n",
-        "vLLM is an highly optimized LLM serving framework which can increase serving throughput a lot. The higher QPS you have, the more benefits you get using vLLM."
+        "vLLM is a highly optimized LLM serving framework that can significantly increase serving throughput. The higher QPS you have, the more improvement you get using vLLM."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4GTNnnuYqrW_"
+      },
+      "source": [
+        "Set the prebuilt model id."
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "ee610586eddb"
+        "id": "kLsRoc4Kqrkx"
+      },
+      "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": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YI0vaDi6p2fi"
       },
       "outputs": [],
       "source": [
         "model_without_peft_vllm, endpoint_without_peft_vllm = deploy_model_vllm(\n",
         "    model_name=get_job_name_with_datetime(prefix=\"openllama-serve-vllm\"),\n",
-        "    model_id=\"openlm-research/open_llama_13b\",\n",
+        "    model_id=prebuilt_model_id,\n",
         "    service_account=SERVICE_ACCOUNT,\n",
         "    machine_type=\"n1-highmem-8\",\n",
         "    accelerator_type=\"NVIDIA_TESLA_V100\",\n",
@@ -515,10 +544,10 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "sGKIjgmDFRW2"
+        "id": "dWYmYWoqqBuZ"
       },
       "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",
+        "NOTE: The prebuilt model weights will be downloaded on the fly from the original location after the deployment succeeds. Thus, an additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you can 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."
       ]
@@ -527,7 +556,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "63bb8fb78d0b"
+        "id": "fjO4z3qAp3pK"
       },
       "outputs": [],
       "source": [
@@ -590,7 +619,7 @@
         "\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. open_llama_3b and open_llama_7b can be finetuned on 1 V100, and open_llama_13b can be finetuned on 1 A100 (40G)."
+        "In order to make the finetuning efficiently, we enabled quantization for loading pretrained models for finetuning LoRA models. Precision options include `\"4bit\"`, `\"8bit\"`, `\"float16\"` (default) and `\"float32\"`, and the precision can be set via `\"--precision_mode\"`. 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. open_llama_3b and open_llama_7b can be finetuned on 1 V100, and open_llama_13b can be finetuned on 1 A100 (40G)."
       ]
     },
     {
@@ -607,7 +636,7 @@
         "# Finetunes open_llama_3b and open_llama_7b with 1 V100.\n",
         "machine_type = \"n1-standard-8\"\n",
         "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
-        "# Finetunes and open_llama_13b with 1 A100 (40G).\n",
+        "# Finetunes open_llama_13b with 1 A100 (40G).\n",
         "# machine_type = \"a2-highgpu-1g\"\n",
         "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
         "replica_count = 1\n",
@@ -684,7 +713,7 @@
         "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",
+        "NOTE: After the deployment succeeds, the base model weights will be downloaded on the fly from the original location and LoRA model weights will be downloaded from the GCS bucket used in training above. Thus, an additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you can 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."
       ]
@@ -712,6 +741,146 @@
         "    print(prediction[\"generated_text\"])"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "JmuUk3l1DoEo"
+      },
+      "source": [
+        "## Evaluate PEFT-finetuned OpenLLaMA\n",
+        "\n",
+        "This section demonstrates how to evaluate the OpenLLaMA model fintuned with PEFT LoRA using EleutherAI's [Language Model Evaluation Harness (lm-evaluation-harness)](https://github.com/EleutherAI/lm-evaluation-harness) with Vertex CustomJob.\n",
+        "\n",
+        "This example uses the dataset [HellaSwag](https://allenai.org/data/hellaswag). All supported tasks are listed in [this task table](https://github.com/EleutherAI/lm-evaluation-harness/blob/master/docs/task_table.md)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "gM4SXaquDoEo"
+      },
+      "outputs": [],
+      "source": [
+        "eval_dataset = \"hellaswag\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Worker pool spec.\n",
+        "# Evaluates open_llama_3b and open_llama_7b with 1 V100.\n",
+        "machine_type = \"n1-standard-8\"\n",
+        "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
+        "# Evaluates open_llama_13b with 1 A100 (40G).\n",
+        "# machine_type = \"a2-highgpu-1g\"\n",
+        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+        "replica_count = 1\n",
+        "accelerator_count = 1\n",
+        "\n",
+        "# Setup evaluation job.\n",
+        "job_name = get_job_name_with_datetime(prefix=\"openllama-peft-eval\")\n",
+        "eval_output_dir = os.path.join(MODEL_BUCKET, job_name)\n",
+        "eval_output_dir_gcsfuse = eval_output_dir.replace(\"gs://\", \"/gcs/\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Yt4Tth8hDoEo"
+      },
+      "outputs": [],
+      "source": [
+        "# Prepare evaluation script that runs the evaluation harness.\n",
+        "script_path = \"./eval_script.py\"  # @param {type:\"string\"}\n",
+        "\n",
+        "eval_command = f\"\"\"import subprocess\n",
+        "\n",
+        "\n",
+        "subprocess.call([\n",
+        "    'python',\n",
+        "    'main.py',\n",
+        "    '--model',\n",
+        "    'hf-causal-experimental',\n",
+        "    '--model_args',\n",
+        "    'pretrained={base_model_id},peft={output_dir_gcsfuse}',\n",
+        "    '--tasks',\n",
+        "    '{eval_dataset}',\n",
+        "    '--output_path',\n",
+        "    '{eval_output_dir_gcsfuse}',\n",
+        "])\n",
+        "\"\"\"\n",
+        "\n",
+        "with open(script_path, \"w\") as fp:\n",
+        "    fp.write(eval_command)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ItWB0WS__CX-"
+      },
+      "source": [
+        "### Submit evaluation CustomJob"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "BbyIk99bDoEo"
+      },
+      "outputs": [],
+      "source": [
+        "# Pass evaluation arguments and launch job.\n",
+        "eval_job = aiplatform.CustomJob.from_local_script(\n",
+        "    display_name=job_name,\n",
+        "    script_path=script_path,\n",
+        "    container_uri=EVAL_DOCKER_URI,\n",
+        "    replica_count=replica_count,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    base_output_dir=eval_output_dir,\n",
+        ")\n",
+        "\n",
+        "eval_job.run()\n",
+        "\n",
+        "print(\"Evaluation results were saved in:\", eval_output_dir)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kN0lE2iu_NXN"
+      },
+      "source": [
+        "### Fetch and print evaluation results"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "927oRxoADoEp"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "\n",
+        "from google.cloud import storage\n",
+        "\n",
+        "# Fetch evaluation results.\n",
+        "storage_client = storage.Client()\n",
+        "BUCKET_NAME = BUCKET_URI.split(\"gs://\")[1]\n",
+        "bucket = storage_client.get_bucket(BUCKET_NAME)\n",
+        "RESULT_FILE_PATH = eval_output_dir[len(BUCKET_URI) + 1 :]\n",
+        "blob = bucket.blob(RESULT_FILE_PATH)\n",
+        "raw_result = blob.download_as_string()\n",
+        "\n",
+        "# Print evaluation results.\n",
+        "result = json.loads(raw_result)\n",
+        "result_formatted = json.dumps(result, indent=2)\n",
+        "print(f\"Evaluation result:\\n{result_formatted}\")"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -729,18 +898,19 @@
       },
       "outputs": [],
       "source": [
-        "# Delete custom train jobs.\n",
+        "# Delete custom train and evaluation jobs.\n",
         "train_job.delete()\n",
+        "eval_job.delete()\n",
         "\n",
-        "# Undeploy model and delete endpoint.\n",
+        "# Undeploy models and delete endpoints.\n",
         "endpoint_without_peft.delete(force=True)\n",
-        "endpoint_with_peft.delete(force=True)\n",
         "endpoint_without_peft_vllm.delete(force=True)\n",
+        "endpoint_with_peft.delete(force=True)\n",
         "\n",
         "# Delete models.\n",
         "model_without_peft.delete()\n",
-        "model_with_peft.delete()\n",
-        "model_without_peft_vllm.delete()"
+        "model_without_peft_vllm.delete()\n",
+        "model_with_peft.delete()"
       ]
     }
   ],

notebooks/official/custom/get_started_with_vertex_endpoint_and_shared_vm.ipynb

@@ -1215,20 +1215,6 @@
         "    time.sleep(30)"
       ]
     },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "86a659bf60f0"
-      },
-      "outputs": [],
-      "source": [
-        "! curl -X GET \\\n",
-        " -H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
-        " -H \"Content-Type: application/json\" \\\n",
-        "https://{REGION}-aiplatform.googleapis.com/v1/projects/759209241365/locations/us-central1/endpoints/2259566763823857664"
-      ]
-    },
     {
       "cell_type": "markdown",
       "metadata": {

notebooks/official/explainable_ai/sdk_automl_tabular_binary_classification_batch_explain.ipynb

@@ -149,8 +149,6 @@
       },
       "outputs": [],
       "source": [
-        "import os\n",
-        "\n",
         "! pip3 install --upgrade --quiet google-cloud-aiplatform\n",
         "! pip3 install --upgrade --quiet google-cloud-storage\n",
         "! pip3 install --upgrade --quiet tensorflow"
@@ -784,6 +782,8 @@
       },
       "outputs": [],
       "source": [
+        "import os\n",
+        "\n",
         "# Set this to true only if you'd like to delete your bucket\n",
         "delete_bucket = False\n",
         "\n",

notebooks/official/explainable_ai/sdk_custom_image_classification_online_explain.ipynb

@@ -356,6 +356,8 @@
       },
       "outputs": [],
       "source": [
+        "import os\n",
+        "\n",
         "import google.cloud.aiplatform as aip"
       ]
     },

notebooks/official/explainable_ai/sdk_custom_tabular_regression_online_explain_get_metadata.ipynb

@@ -409,6 +409,8 @@
       },
       "outputs": [],
       "source": [
+        "import os\n",
+        "\n",
         "if os.getenv(\"IS_TESTING_TRAIN_GPU\"):\n",
         "    TRAIN_GPU, TRAIN_NGPU = (\n",
         "        aip.gapic.AcceleratorType.NVIDIA_TESLA_K80,\n",

notebooks/official/generative_ai/rlhf_tune_llm.ipynb

@@ -34,18 +34,18 @@
         "\n",
         "<table align=\"left\">\n",
         "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/gen_ai/rlhf_tune_llm.ipynb\">\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/generative_ai/rlhf_tune_llm.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/official/gen_ai/rlhf_tune_llm.ipynb\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/generative_ai/rlhf_tune_llm.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/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/gen_ai/rlhf_tune_llm.ipynb\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/generative_ai/rlhf_tune_llm.ipynb\">\n",
         "       <img src=\"https://www.gstatic.com/cloud/images/navigation/vertex-ai.svg\" alt=\"Vertex AI logo\">Open in Vertex AI Workbench\n",
         "    </a>\n",
         "</table>"
@@ -61,7 +61,9 @@
         "\n",
         "This tutorial demonstrates how to use Vertex AI RLHF to tune a large-language model (LLM). This workflow improves a model's accuracy by fine-tuning a base model with a training dataset.\n",
         "\n",
-        "*Preview releases are covered by the Pre-GA Offerings Terms of the Google Cloud Platform Terms of Service. They are not intended for production use or covered by any SLA, support obligation, or deprecation policy and might be subject to backward-incompatible changes.*"
+        "*Preview releases are covered by the Pre-GA Offerings Terms of the Google Cloud Platform Terms of Service. They are not intended for production use or covered by any SLA, support obligation, or deprecation policy and might be subject to backward-incompatible changes.*\n",
+        "\n",
+        "Learn more about [Overview of Generative AI support on Vertex AI](https://cloud.google.com/vertex-ai/docs/generative-ai/learn/overview)."
       ]
     },
     {
@@ -103,7 +105,7 @@
         "\n",
         "| large_model_reference | supported accelerator_type | supported region |\n",
         "|---|---|---|\n",
-        "|  `BISON` | `TPU_V3`, `NVIDIA_TESLA_A100` | europe-west4, us-central1 |\n",
+        "|  `text-bison@001` | `TPU_V3`, `NVIDIA_TESLA_A100` | europe-west4, us-central1 |\n",
         "|  `t5-small` | `TPU_V3`, `NVIDIA_TESLA_A100` | europe-west4, us-central1 |\n",
         "|  `t5-large` | `TPU_V3` | europe-west4 |\n",
         "|  `t5-xl` | `TPU_V3` | europe-west4 |\n",
@@ -617,10 +619,10 @@
         "    pipeline_root=f\"{BUCKET_URI}/rlhf\",\n",
         "    template_path=RLHF_PIPELINE_PKG_PATH,\n",
         "    parameter_values={\n",
-        "        \"preference_dataset\": \"gs://vertex-ai/generative-ai/rlhf/text_small/summarize_from_feedback_tfds/comparisons/train/*.jsonl\",\n",
-        "        \"prompt_dataset\": \"gs://vertex-ai/generative-ai/rlhf/text_small/reddit_tfds/train/*.jsonl\",\n",
-        "        \"eval_dataset\": \"gs://vertex-ai/generative-ai/rlhf/text_small/reddit_tfds/val/*.jsonl\",\n",
-        "        \"large_model_reference\": \"BISON\",  # GA value will be text-bison@001\n",
+        "        \"preference_dataset\": \"gs://cloud-samples-data/vertex-ai/llm/rlhf/text_small/summarize_from_feedback_tfds/comparisons/train/*.jsonl\",\n",
+        "        \"prompt_dataset\": \"gs://cloud-samples-data/vertex-ai/llm/rlhf/text_small/reddit_tfds/train/*.jsonl\",\n",
+        "        \"eval_dataset\": \"gs://cloud-samples-data/vertex-ai/llm/rlhf/text_small/reddit_tfds/val/*.jsonl\",\n",
+        "        \"large_model_reference\": \"text-bison@001\",\n",
         "        \"model_display_name\": \"my_rlhf_tutorial_model\",  # Optional. If omitted, a default model_display_name will be created.\n",
         "        \"reward_model_train_steps\": 100,  # Please remember to read \"A Note on choosing train_steps\" section.\n",
         "        \"reinforcement_learning_train_steps\": 100,  # Please remember to read \"A Note on choosing train_steps\" section.\n",

notebooks/official/pipelines/Train_tabular_models_with_many_frameworks_and_import_to_Vertex_AI_using_Pipelines/Train_tabular_classification_model_using_XGBoost_and_import_to_Vertex_AI/pipeline.py

@@ -9,7 +9,7 @@ binarize_column_using_Pandas_on_CSV_data_op = components.load_component_from_url
 split_rows_into_subsets_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/dataset_manipulation/Split_rows_into_subsets/in_CSV/component.yaml")
 train_XGBoost_model_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Train/component.yaml")
 xgboost_predict_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Predict/component.yaml")
-upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
+upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/d5c9918850a6cc70004c4269dae066cfe2e664eb/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
 deploy_model_to_endpoint_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Deploy_to_endpoint/component.yaml")
 
 # %% Pipeline definition

notebooks/official/pipelines/Train_tabular_models_with_many_frameworks_and_import_to_Vertex_AI_using_Pipelines/Train_tabular_classification_model_using_all_frameworks_and_import_to_Vertex_AI/pipeline.py

@@ -23,7 +23,7 @@ upload_PyTorch_model_archive_to_Google_Cloud_Vertex_AI_op = components.load_comp
 # XGBoost
 train_XGBoost_model_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Train/component.yaml")
 xgboost_predict_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Predict/component.yaml")
-upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
+upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/d5c9918850a6cc70004c4269dae066cfe2e664eb/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
 
 # Scikit-learn
 #train_linear_regression_model_using_scikit_learn_from_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/1f5cf6e06409b704064b2086c0a705e4e6b4fcde/community-content/pipeline_components/ML_frameworks/Scikit_learn/Train_linear_regression_model/from_CSV/component.yaml")

notebooks/official/pipelines/Train_tabular_models_with_many_frameworks_and_import_to_Vertex_AI_using_Pipelines/Train_tabular_regression_model_using_XGBoost_and_import_to_Vertex_AI/pipeline.py

@@ -8,7 +8,7 @@ fill_all_missing_values_using_Pandas_on_CSV_data_op = components.load_component_
 split_rows_into_subsets_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/dataset_manipulation/Split_rows_into_subsets/in_CSV/component.yaml")
 train_XGBoost_model_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Train/component.yaml")
 xgboost_predict_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Predict/component.yaml")
-upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
+upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/d5c9918850a6cc70004c4269dae066cfe2e664eb/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
 deploy_model_to_endpoint_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Deploy_to_endpoint/component.yaml")
 
 # %% Pipeline definition

notebooks/official/pipelines/Train_tabular_models_with_many_frameworks_and_import_to_Vertex_AI_using_Pipelines/Train_tabular_regression_model_using_all_frameworks_and_import_to_Vertex_AI/pipeline.py

@@ -22,7 +22,7 @@ upload_PyTorch_model_archive_to_Google_Cloud_Vertex_AI_op = components.load_comp
 # XGBoost
 train_XGBoost_model_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Train/component.yaml")
 xgboost_predict_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Predict/component.yaml")
-upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
+upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/d5c9918850a6cc70004c4269dae066cfe2e664eb/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
 
 # Scikit-learn
 train_linear_regression_model_using_scikit_learn_from_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/1f5cf6e06409b704064b2086c0a705e4e6b4fcde/community-content/pipeline_components/ML_frameworks/Scikit_learn/Train_linear_regression_model/from_CSV/component.yaml")

notebooks/official/pipelines/google_cloud_pipeline_components_dataproc_tabular.ipynb

@@ -53,6 +53,17 @@
         "</table>"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "962e636b5cee"
+      },
+      "source": [
+        "**_NOTE_**: This notebook has been tested in the following environment:\n",
+        "\n",
+        "* Python version = 3.9"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -128,61 +139,12 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "ze4-nDLfK4pw"
+        "id": "install_aip:mbsdk"
       },
       "source": [
-        "### Set up your local development environment\n",
+        "## Installation\n",
         "\n",
-        "**If you are using Colab or Vertex AI Workbench Notebooks**, your environment already meets\n",
-        "all the requirements to run this notebook. You can skip this step."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "gCuSR8GkAgzl"
-      },
-      "source": [
-        "**Otherwise**, make sure your environment meets this notebook's requirements.\n",
-        "You need the following:\n",
-        "\n",
-        "* The Google Cloud SDK\n",
-        "* Git\n",
-        "* Python 3\n",
-        "* virtualenv\n",
-        "* Jupyter notebook running in a virtual environment with Python 3\n",
-        "\n",
-        "The Google Cloud guide to [Setting up a Python development\n",
-        "environment](https://cloud.google.com/python/setup) and the [Jupyter\n",
-        "installation guide](https://jupyter.org/install) provide detailed instructions\n",
-        "for meeting these requirements. The following steps provide a condensed set of\n",
-        "instructions:\n",
-        "\n",
-        "1. [Install and initialize the Cloud SDK.](https://cloud.google.com/sdk/docs/)\n",
-        "\n",
-        "1. [Install Python 3.](https://cloud.google.com/python/setup#installing_python)\n",
-        "\n",
-        "1. [Install\n",
-        "   virtualenv](https://cloud.google.com/python/setup#installing_and_using_virtualenv)\n",
-        "   and create a virtual environment that uses Python 3. Activate the virtual environment.\n",
-        "\n",
-        "1. To install Jupyter, run `pip3 install jupyter` on the\n",
-        "command-line in a terminal shell.\n",
-        "\n",
-        "1. To launch Jupyter, run `jupyter notebook` on the command-line in a terminal shell.\n",
-        "\n",
-        "1. Open this notebook in the Jupyter Notebook Dashboard."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "i7EUnXsZhAGF"
-      },
-      "source": [
-        "## Install additional packages\n",
-        "\n",
-        "Install additional package dependencies not installed in your notebook environment, such as Pyspark, MLeap and others. Use the latest major GA version of each package."
+        "Install the packages required for executing this notebook."
       ]
     },
     {
@@ -195,59 +157,38 @@
       "source": [
         "import os\n",
         "\n",
-        "# The Vertex AI Workbench Notebook product has specific requirements\n",
-        "IS_WORKBENCH_NOTEBOOK = os.getenv(\"DL_ANACONDA_HOME\")\n",
-        "IS_USER_MANAGED_WORKBENCH_NOTEBOOK = os.path.exists(\n",
-        "    \"/opt/deeplearning/metadata/env_version\"\n",
-        ")\n",
+        "# (optional) update gcloud if needed\n",
+        "if os.getenv(\"IS_TESTING\"):\n",
+        "    ! gcloud components update --quiet\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 {USER_FLAG} --upgrade google-cloud-aiplatform==1.18.1 \\\n",
-        "                                     kfp==1.8.14 \\\n",
-        "                                     google-cloud-pipeline-components==1.0.24 --quiet --no-warn-conflicts"
+        "\n",
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform==1.30.1 \\\n",
+        "                                 kfp==1.8.14 \\\n",
+        "                                 google-cloud-pipeline-components==1.0.33 --no-warn-conflicts"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "hhq5zEbGg0XX"
+        "id": "58707a750154"
       },
       "source": [
-        "### Restart the kernel\n",
-        "\n",
-        "After you install 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": "EzrelQZ22IZj"
+        "id": "f200f10a1da3"
       },
       "outputs": [],
       "source": [
-        "# Automatically restart kernel after installs\n",
-        "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)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "lWEdiXsJg0XY"
-      },
-      "source": [
-        "## Before you begin"
+        "# app = IPython.Application.instance()\n",
+        "# app.kernel.do_shutdown(True)"
       ]
     },
     {
@@ -256,22 +197,19 @@
         "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",
+        "2. [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, Cloud Build, Cloud Storage, and Artifact Registry APIs.](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleleapis.com,%20cloudbuild.googleapis.com,%20storage-api.googleapis.com,%20artifactregistry.googleapis.com&_ga=2.11105652.1331351212.1648374605-1663626427.1647020959). \n",
+        "3. [Enable the Artifact Registry, Cloud Build, Container Registry, Dataproc and Vertex AI APIs](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,artifactregistry.googleapis.com,cloudbuild.googleapis.com,containerregistry.googleapis.com,dataproc.googleapis.com,aiplatform.googleapis.com).\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. 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."
+        "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
       ]
     },
     {
@@ -282,7 +220,10 @@
       "source": [
         "#### Set your project ID\n",
         "\n",
-        "**If you don't know your project ID**, you may be able to get your project ID using `gcloud`."
+        "**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)"
       ]
     },
     {
@@ -293,76 +234,32 @@
       },
       "outputs": [],
       "source": [
-        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "riG_qUokg0XZ"
-      },
-      "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)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "_B-bdB17cQ_3"
-      },
-      "outputs": [],
-      "source": [
-        "! gcloud config set project $PROJECT_ID --quiet"
+        "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": "6a5a0be04af6"
+        "id": "region"
       },
       "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)."
+        "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": "8e2a112253e5"
+        "id": "region"
       },
       "outputs": [],
       "source": [
-        "REGION = \"[your-region]\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "66cc6bbc7e31"
-      },
-      "outputs": [],
-      "source": [
-        "if REGION == \"[your-region]\":\n",
-        "    REGION = \"us-central1\""
+        "REGION = \"us-central1\"  # @param {type: \"string\"}"
       ]
     },
     {
@@ -399,86 +296,59 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "dr--iN2kAylZ"
+        "id": "sBCra4QMA2wR"
       },
       "source": [
         "### Authenticate your Google Cloud account\n",
         "\n",
-        "**If you are using Vertex AI Workbench Notebooks**, your environment is already\n",
-        "authenticated. Skip this step."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "sBCra4QMA2wR"
-      },
-      "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 the following role and select them\n",
-        "\n",
-        "      - Artifact Registry Administrator\n",
-        "      - Artifact Registry Repository Administrator\n",
-        "      - Cloud Build Editor\n",
-        "      - Compute Network Admin\n",
-        "      - Dataproc Administrator\n",
-        "      - Dataproc Worker\n",
-        "      - Service Account User\n",
-        "      - Storage Admin\n",
-        "      - Storage Object Admin\n",
-        "      - Vertex AI Administrator\n",
-        "\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."
+        "**2. Local JupyterLab instance, uncomment and run:**"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "PyQmSRbKA8r-"
+        "id": "254614fa0c46"
       },
       "outputs": [],
       "source": [
-        "# If you are running this notebook in Colab, run this cell and follow the\n",
-        "# instructions to authenticate your GCP account. This provides access to your\n",
-        "# Cloud Storage bucket and lets you submit training jobs and prediction\n",
-        "# requests.\n",
-        "\n",
-        "import os\n",
-        "import sys\n",
-        "\n",
-        "# If on Vertex AI Workbench, then don't execute this code\n",
-        "IS_COLAB = \"google.colab\" in sys.modules\n",
-        "if not os.path.exists(\"/opt/deeplearning/metadata/env_version\") and not os.getenv(\n",
-        "    \"DL_ANACONDA_HOME\"\n",
-        "):\n",
-        "    if \"google.colab\" in sys.modules:\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 ''"
+        "# ! gcloud auth login"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ef21552ccea8"
+      },
+      "source": [
+        "**3. Colab, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "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."
       ]
     },
     {
@@ -489,7 +359,7 @@
       "source": [
         "### Enable Google Cloud services\n",
         "\n",
-        "Enable the following services in your project:\n",
+        "Enable the following services in your project if not already done:\n",
         "\n",
         "* Artifact Registry\n",
         "* Cloud Build\n",
@@ -522,12 +392,7 @@
       "source": [
         "### Create a Cloud Storage bucket\n",
         "\n",
-        "**The following steps are required, regardless of your notebook environment.**\n",
-        "\n",
-        "When you submit a Spark job using the Cloud SDK, you need a bucket to read the python module associated and the data required to it. \n",
-        "\n",
-        "Set the name of your Cloud Storage bucket below. It must be unique across all\n",
-        "Cloud Storage buckets."
+        "Create a storage bucket to store intermediate artifacts such as datasets."
       ]
     },
     {
@@ -538,19 +403,7 @@
       },
       "outputs": [],
       "source": [
-        "BUCKET_URI = \"gs://[your-bucket-name]\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "cf221059d072"
-      },
-      "outputs": [],
-      "source": [
-        "if BUCKET_URI == \"\" or BUCKET_URI is None or BUCKET_URI == \"gs://[your-bucket-name]\":\n",
-        "    BUCKET_URI = \"gs://\" + PROJECT_ID + \"-aip-\" + UUID"
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
       ]
     },
     {
@@ -570,27 +423,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ucvCsknMCims"
-      },
-      "source": [
-        "Finally, validate access to your Cloud Storage bucket by examining its contents:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "vhOb7YnwClBb"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil ls -al $BUCKET_URI"
+        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
       ]
     },
     {
@@ -623,6 +456,11 @@
       },
       "outputs": [],
       "source": [
+        "import os\n",
+        "import sys\n",
+        "\n",
+        "IS_COLAB = \"google.colab\" in sys.modules\n",
+        "\n",
         "if (\n",
         "    SERVICE_ACCOUNT == \"\"\n",
         "    or SERVICE_ACCOUNT is None\n",
@@ -739,8 +577,10 @@
       },
       "outputs": [],
       "source": [
+        "# set repo name\n",
         "REPO_NAME = \"loan-eligibility-spark-demo\"\n",
         "\n",
+        "# create the repository\n",
         "!gcloud artifacts repositories create $REPO_NAME \\\n",
         "    --repository-format=docker \\\n",
         "    --location=$REGION \\\n",
@@ -759,7 +599,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 1,
       "metadata": {
         "id": "pRUOFELefqf1"
       },
@@ -784,7 +624,7 @@
       "outputs": [],
       "source": [
         "# Setup\n",
-        "DATAPROC_RUNTIME_VERSION = \"1.1.3\"\n",
+        "DATAPROC_RUNTIME_VERSION = \"1.1.20\"\n",
         "SRC = path(\"src\")\n",
         "BUILD_PATH = path(\"build\")\n",
         "DELIVERABLES = path(\"deliverables\")\n",
@@ -883,7 +723,12 @@
       },
       "outputs": [],
       "source": [
-        "vertex_ai.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
+        "vertex_ai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        "    staging_bucket=BUCKET_URI,\n",
+        "    experiment=EXPERIMENT_NAME,\n",
+        ")"
       ]
     },
     {
@@ -1956,10 +1801,9 @@
         "# The following packages are installed in the default image, it is strongly\n",
         "# recommended to include all of them.\n",
         "#\n",
-        "# Use mamba to install packages quickly.\n",
-        "RUN ${CONDA_HOME}/bin/conda install mamba -n base -c conda-forge \\\n",
-        "    && ${CONDA_HOME}/bin/mamba install \\\n",
-        "      conda \\\n",
+        "# Use mamba solver to install packages quickly.\n",
+        "RUN ${CONDA_HOME}/bin/conda install -n base conda-libmamba-solver\n",
+        "RUN ${CONDA_HOME}/bin/conda install \\\n",
         "      cython \\\n",
         "      fastavro \\\n",
         "      fastparquet \\\n",
@@ -1973,7 +1817,7 @@
         "      scikit-image \\\n",
         "      scikit-learn \\\n",
         "      scipy \\\n",
-        "      mleap\n",
+        "      mleap --solver=libmamba\n",
         "\n",
         "# (Required) Create the 'spark' group/user.\n",
         "# The GID and UID must be 1099. Home directory is required.\n",
@@ -2011,7 +1855,7 @@
       "source": [
         "#### Build the Dataproc Serverless custom runtime using Cloud Build\n",
         "\n",
-        "**Note:** this step may take approximately 5 to 10 minutes to complete."
+        "**Note:** this step may take approximately upto 20 minutes to complete."
       ]
     },
     {
@@ -2022,7 +1866,7 @@
       },
       "outputs": [],
       "source": [
-        "!gcloud builds submit --tag $RUNTIME_CONTAINER_IMAGE $BUILD_PATH --machine-type=N1_HIGHCPU_32 --timeout=900s --verbosity=info"
+        "!gcloud builds submit --tag $RUNTIME_CONTAINER_IMAGE $BUILD_PATH --machine-type=N1_HIGHCPU_32 --timeout=3600s --verbosity=info"
       ]
     },
     {
@@ -2117,7 +1961,6 @@
         "    metrics: Output[Metrics],\n",
         "    plots: Output[ClassificationMetrics],\n",
         ") -> NamedTuple(\"Outputs\", [(\"threshold_metric\", float)]):\n",
-        "\n",
         "    # Libraries --------------------------------------------------------------------------------------------------------------------------\n",
         "    import json\n",
         "\n",
@@ -2516,7 +2359,6 @@
         "        evaluate_model_op.outputs[\"threshold_metric\"] >= threshold,\n",
         "        name=AUPR_HYPERTUNE_CONDITION,\n",
         "    ):\n",
-        "\n",
         "        build_hpt_args_op = build_hpt_args(\n",
         "            train_path=train_path,\n",
         "            model_path=hpt_model_path,\n",
@@ -2668,7 +2510,11 @@
       },
       "outputs": [],
       "source": [
-        "experiment_df = vertex_ai.get_experiment_df(experiment=EXPERIMENT_NAME)\n",
+        "# get the experiment by name\n",
+        "experiment = vertex_ai.Experiment(experiment_name=EXPERIMENT_NAME)\n",
+        "\n",
+        "# export the data as a dataframe\n",
+        "experiment_df = experiment.get_data_frame()\n",
         "\n",
         "# Show successfully completed experiment runs, sorted by F1 score\n",
         "experiment_df.query('state == \"COMPLETE\"').sort_values(\n",
@@ -2773,7 +2619,15 @@
         "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
         "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 tutorial with the following code."
+        "Otherwise, you can delete the individual resources you created in this tutorial:\n",
+        "\n",
+        "- Vertex AI Pipeline\n",
+        "- Vertex AI Endpoint\n",
+        "- Vertex AI Model\n",
+        "- Vertex AI Experiment\n",
+        "- Artifact Repository\n",
+        "- Cloud Storage bucket\n",
+        "- Local src, build and cloned repo folders"
       ]
     },
     {
@@ -2796,7 +2650,10 @@
         "# Delete model\n",
         "model_list = vertex_ai.Model.list(filter=f'display_name=\"{MODEL_NAME}\"')\n",
         "for model in model_list:\n",
-        "    model.delete()"
+        "    model.delete()\n",
+        "\n",
+        "# Delete experiment\n",
+        "experiment.delete()"
       ]
     },
     {
@@ -2842,7 +2699,8 @@
       },
       "outputs": [],
       "source": [
-        "!rm -rf $SRC $BUILD_PATH"
+        "# remove the local src, build and repo folders\n",
+        "!rm -rf $SRC $BUILD_PATH cloud-builders-community vertex-ai-spark-ml-serving"
       ]
     }
   ],

notebooks/official/sdk/sdk2_bigframes_sklearn.ipynb

@@ -0,0 +1,725 @@
+{
+  "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": "JAPoU8Sm5E6e"
+      },
+      "source": [
+        "# Train a scikit-learn model with Vertex AI SDK 2.0 and Bigframes\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/sdk/sdk2_bigframes_sklearn.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/official/sdk/sdk2_bigframes_sklearn.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/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/sdk/sdk2_bigframes_sklearn.ipynb\">\n",
+        "       <img src=\"https://www.gstatic.com/cloud/images/navigation/vertex-ai.svg\" alt=\"Vertex AI logo\">Open in Vertex AI Workbench\n",
+        "    </a>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tvgnzT1CKxrO"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This tutorial demonstrates how to train a scikit-learn model using Vertex AI local-to-remote training with Vertex AI SDK 2.0 and BigQuery Bigframes as the data source."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "d975e698c9a4"
+      },
+      "source": [
+        "### Objective\n",
+        "\n",
+        "In this tutorial, you learn to use `Vertex AI SDK 2.0` with Bigframes as input data source.\n",
+        "\n",
+        "\n",
+        "This tutorial uses the following Google Cloud ML services:\n",
+        "\n",
+        "- `Vertex AI Training`\n",
+        "- `Vertex AI Remote Training`\n",
+        "\n",
+        "\n",
+        "The steps performed include:\n",
+        "\n",
+        "- Initialize a dataframe from a BigQuery table and split the dataset\n",
+        "- Perform transformations as a Vertex AI remote training.\n",
+        "- Train the model remotely and evaluate the model locally\n",
+        "\n",
+        "**Local-to-remote training**\n",
+        "\n",
+        "```\n",
+        "import vertexai\n",
+        "from my_module import MyModelClass\n",
+        "\n",
+        "vertexai.preview.init(remote=True, project=\"my-project\", location=\"my-location\", staging_bucket=\"gs://my-bucket\")\n",
+        "\n",
+        "# Wrap the model class with `vertex_ai.preview.remote`\n",
+        "MyModelClass = vertexai.preview.remote(MyModelClass)\n",
+        "\n",
+        "# Instantiate the class\n",
+        "model = MyModelClass(...)\n",
+        "\n",
+        "# Optional set remote config\n",
+        "model.fit.vertex.remote_config.display_name = \"MyModelClass-remote-training\"\n",
+        "model.fit.vertex.remote_config.staging_bucket = \"gs://my-bucket\"\n",
+        "\n",
+        "# This `fit` call will be executed remotely\n",
+        "model.fit(...)\n",
+        "```"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "08d289fa873f"
+      },
+      "source": [
+        "### Dataset\n",
+        "\n",
+        "This tutorial uses the <a href=\"https://scikit-learn.org/stable/auto_examples/datasets/plot_iris_dataset.html\">IRIS dataset</a>, which predicts the iris species."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aed92deeb4a0"
+      },
+      "source": [
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* BigQuery\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing),\n",
+        "[BigQuery pricing](https://cloud.google.com/bigquery/pricing),\n",
+        "and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), \n",
+        "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n",
+        "to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "i7EUnXsZhAGF"
+      },
+      "source": [
+        "## Installation\n",
+        "\n",
+        "Install the following packages required to execute this notebook. "
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "2b4ef9b72d43"
+      },
+      "outputs": [],
+      "source": [
+        "# Install the packages\n",
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform[preview]\n",
+        "! pip3 install --upgrade --quiet bigframes"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "58707a750154"
+      },
+      "source": [
+        "### Colab only: Uncomment the following cell to restart the kernel."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "f200f10a1da3"
+      },
+      "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)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "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",
+        "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 API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
+        "\n",
+        "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "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,
+      "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": "region"
+      },
+      "outputs": [],
+      "source": [
+        "REGION = \"us-central1\"  # @param {type: \"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "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."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "74ccc9e52986"
+      },
+      "source": [
+        "**1. Vertex AI Workbench**\n",
+        "* Do nothing as you are already authenticated."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "de775a3773ba"
+      },
+      "source": [
+        "**2. Local JupyterLab instance, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "254614fa0c46"
+      },
+      "outputs": [],
+      "source": [
+        "# ! gcloud auth login"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ef21552ccea8"
+      },
+      "source": [
+        "**3. Colab, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "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": "960505627ddf"
+      },
+      "source": [
+        "### Import libraries and define constants"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "PyQmSRbKA8r-"
+      },
+      "outputs": [],
+      "source": [
+        "import bigframes.pandas as bf\n",
+        "import vertexai\n",
+        "\n",
+        "bf.options.bigquery.location = \"us\"  # Dataset is in 'us' not 'us-central1'\n",
+        "bf.options.bigquery.project = PROJECT_ID\n",
+        "\n",
+        "from bigframes.ml.model_selection import \\\n",
+        "    train_test_split as bf_train_test_split\n",
+        "\n",
+        "REMOTE_JOB_NAME = \"sdk2-bigframes-sklearn\"\n",
+        "REMOTE_JOB_BUCKET = f\"{BUCKET_URI}/{REMOTE_JOB_NAME}\""
+      ]
+    },
+    {
+      "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 and corresponding bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "init_aip:mbsdk,all"
+      },
+      "outputs": [],
+      "source": [
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        "    staging_bucket=BUCKET_URI,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "105334524e96"
+      },
+      "source": [
+        "## Prepare the dataset\n",
+        "\n",
+        "Now load the Iris dataset and split the data into train and test sets."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "b44cdc4e03f1"
+      },
+      "outputs": [],
+      "source": [
+        "df = bf.read_gbq(\"bigquery-public-data.ml_datasets.iris\")\n",
+        "\n",
+        "species_categories = {\n",
+        "    \"versicolor\": 0,\n",
+        "    \"virginica\": 1,\n",
+        "    \"setosa\": 2,\n",
+        "}\n",
+        "df[\"species\"] = df[\"species\"].map(species_categories)\n",
+        "\n",
+        "# Assign an index column name\n",
+        "index_col = \"index\"\n",
+        "df.index.name = index_col"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9cb8616b1997"
+      },
+      "outputs": [],
+      "source": [
+        "feature_columns = df[[\"sepal_length\", \"sepal_width\", \"petal_length\", \"petal_width\"]]\n",
+        "label_columns = df[[\"species\"]]\n",
+        "train_X, test_X, train_y, test_y = bf_train_test_split(\n",
+        "    feature_columns, label_columns, test_size=0.2\n",
+        ")\n",
+        "\n",
+        "print(\"X_train size: \", train_X.size)\n",
+        "print(\"X_test size: \", test_X.size)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8306545fcc57"
+      },
+      "source": [
+        "## Feature transformation\n",
+        "\n",
+        "Next, you do feature transformations on the data using the Vertex AI remote training service.\n",
+        "\n",
+        "First, you re-initialize Vertex AI to enable remote training."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "55e701c31036"
+      },
+      "outputs": [],
+      "source": [
+        "# Switch to remote mode for training\n",
+        "vertexai.preview.init(remote=True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4a0e9d59b273"
+      },
+      "source": [
+        "### Execute remote job for fit_transform() on training data\n",
+        "\n",
+        "Next, indicate that the `StandardScalar` class is to be executed remotely. Then set up the data transform and call the `fit_transform()` method is executed remotely."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "90333089d362"
+      },
+      "outputs": [],
+      "source": [
+        "from sklearn.preprocessing import StandardScaler\n",
+        "\n",
+        "# Wrap classes to enable Vertex remote execution\n",
+        "StandardScaler = vertexai.preview.remote(StandardScaler)\n",
+        "\n",
+        "# Instantiate transformer\n",
+        "transformer = StandardScaler()\n",
+        "\n",
+        "# Set training config\n",
+        "transformer.fit_transform.vertex.remote_config.display_name = (\n",
+        "    f\"{REMOTE_JOB_NAME}-fit-transformer-bigframes\"\n",
+        ")\n",
+        "transformer.fit_transform.vertex.remote_config.staging_bucket = REMOTE_JOB_BUCKET\n",
+        "\n",
+        "# Execute transformer on Vertex (train_X is bigframes.dataframe.DataFrame, X_train is np.array)\n",
+        "X_train = transformer.fit_transform(train_X)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6bf95574c907"
+      },
+      "source": [
+        "### Remote transform on test data"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "da6eea22a89a"
+      },
+      "outputs": [],
+      "source": [
+        "# Transform test dataset before calculate test score\n",
+        "transformer.transform.vertex.remote_config.display_name = (\n",
+        "    REMOTE_JOB_NAME + \"-transformer\"\n",
+        ")\n",
+        "transformer.transform.vertex.remote_config.staging_bucket = REMOTE_JOB_BUCKET\n",
+        "\n",
+        "# Execute transformer on Vertex (test_X is bigframes.dataframe.DataFrame, X_test is np.array)\n",
+        "X_test = transformer.transform(test_X)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ddf906c886e4"
+      },
+      "source": [
+        "## Remote training\n",
+        "\n",
+        "First, train the scikit-learn model as a remote training job:\n",
+        "\n",
+        "- Set LogisticRegression for the remote training job.\n",
+        "- Invoke LogisticRegression locally which will launch the remote training job."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "c7b0116fa60c"
+      },
+      "outputs": [],
+      "source": [
+        "from sklearn.linear_model import LogisticRegression\n",
+        "\n",
+        "# Wrap classes to enable Vertex remote execution\n",
+        "LogisticRegression = vertexai.preview.remote(LogisticRegression)\n",
+        "\n",
+        "# Instantiate model, warm_start=True for uptraining\n",
+        "model = LogisticRegression(warm_start=True)\n",
+        "\n",
+        "# Set training config\n",
+        "model.fit.vertex.remote_config.display_name = REMOTE_JOB_NAME + \"-sklearn-model\"\n",
+        "model.fit.vertex.remote_config.staging_bucket = REMOTE_JOB_BUCKET\n",
+        "\n",
+        "# Train model on Vertex\n",
+        "model.fit(train_X, train_y)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ffe1d5903bcb"
+      },
+      "source": [
+        "## Remote prediction\n",
+        "\n",
+        "Obtain predictions from trained model."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "d00ce35920fa"
+      },
+      "outputs": [],
+      "source": [
+        "# Remote evaluation\n",
+        "vertexai.preview.init(remote=True)\n",
+        "\n",
+        "# Evaluate model's accuracy score\n",
+        "predictions = model.predict(test_X)\n",
+        "\n",
+        "print(f\"Remote predictions: {predictions}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "a8cd6cbd4403"
+      },
+      "source": [
+        "## Local evaluation\n",
+        "\n",
+        "Score model results locally."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "dc105dafdfb9"
+      },
+      "outputs": [],
+      "source": [
+        "# User must convert bigframes to pandas dataframe for local evaluation\n",
+        "train_X_pd = train_X.to_pandas().reset_index(drop=True)\n",
+        "train_y_pd = train_y.to_pandas().reset_index(drop=True)\n",
+        "\n",
+        "test_X_pd = test_X.to_pandas().reset_index(drop=True)\n",
+        "test_y_pd = test_y.to_pandas().reset_index(drop=True)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "25fec549de69"
+      },
+      "outputs": [],
+      "source": [
+        "# Switch to local mode for testing\n",
+        "vertexai.preview.init(remote=False)\n",
+        "\n",
+        "# Evaluate model's accuracy score\n",
+        "print(f\"Train accuracy: {model.score(train_X_pd, train_y_pd)}\")\n",
+        "\n",
+        "print(f\"Test accuracy: {model.score(test_X_pd, test_y_pd)}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "TpV-iwP9qw9c"
+      },
+      "source": [
+        "## Cleaning up\n",
+        "\n",
+        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
+        "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 tutorial:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "sx_vKniMq9ZX"
+      },
+      "outputs": [],
+      "source": [
+        "import os\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": {
+      "collapsed_sections": [],
+      "name": "sdk2_bigframes_sklearn.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}