Update run_linter.sh

Fix spelling error `notebooked` → `notebooks`
Update vLLM version in Llama 3.1 and Guard deployment notebooks. (#3335 )
2026-09-27 07:31:58 +00:00 · 2024-07-29 17:33:58 +02:00 · 2024-07-27 00:25:53 +00:00 · 2024-07-26 17:45:40 +00:00 · 2024-07-26 16:54:57 +00:00 · 2024-07-26 16:51:19 +00:00
220 changed files with 34523 additions and 28780 deletions
@@ -3,8 +3,8 @@ ipython
 jupyter
 nbconvert
 black==24.4.2
-pyupgrade==3.15.2
+pyupgrade==3.16.0
 isort==5.13.2
-flake8==7.0.0
+flake8==7.1.0
 nbqa==1.8.5

@@ -58,7 +58,7 @@ done
 # Only check notebooks in test folders modified in this pull request.
 # Note: Use process substitution to persist the data in the array
 if [ ${#notebooks[@]} -eq 0 ]; then
-    echo "Checking for changed notebooked using git"
+    echo "Checking for changed notebooks using git"
    while read -r file || [ -n "$line" ]; do
        notebooks+=("$file")
    done < <(git diff --name-only main... | grep '\.ipynb$')
@@ -1,37 +1,176 @@
-# Google Cloud Vertex AI Samples
+# ![Google Cloud](https://avatars.githubusercontent.com/u/2810941?s=60&v=4) Google Cloud Vertex AI Samples

-[![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](LICENSE)
-
-Welcome to the Google Cloud [Vertex AI](https://cloud.google.com/vertex-ai/docs/) sample repository.
+This repository contains notebooks, code samples, sample apps, and other resources that demonstrate how to use, develop and manage machine learning and generative AI workflows using Google Cloud Vertex AI.

 ## Overview

-The repository contains [notebooks](https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/master/notebooks) and [community content](https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/master/community-content) that demonstrate how to develop and manage ML workflows using Google Cloud Vertex AI.
+[Vertex AI](https://cloud.google.com/vertex-ai) is a fully-managed, unified AI development platform for building and using generative AI. This repository is designed to help you get started with Vertex AI. Whether you're new to Vertex AI or an experienced ML practitioner, you'll find valuable resources here.
+
+For more Vertex AI Generative AI notebook samples, please visit the Vertex AI [Generative AI](https://github.com/GoogleCloudPlatform/generative-ai) GitHub repository.
+
+## Explore, learn and contribute
+
+You can explore, learn, and contribute to this repository to unleash the full potential of machine learning on Vertex AI! 
+
+### Explore and learn
+
+Explore this repository, follow the links in the header section of each of the notebooks to -
+
+![Colab](https://cloud.google.com/ml-engine/images/colab-logo-32px.png)  Open and run the notebook in [Colab](https://colab.google/)\
+![Colab Enterprise](https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png)  Open and run the notebook in [Colab Enterprise](https://cloud.google.com/colab/docs/introduction)\
+![Workbench](https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32)  Open and run the notebook in [Vertex AI Workbench](https://cloud.google.com/vertex-ai/docs/workbench/introduction)\
+![Github](https://cloud.google.com/ml-engine/images/github-logo-32px.png)  View the notebook on Github
+
+### Contribute
+
+See the [Contributing Guide](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/CONTRIBUTING.md).
+
+## Get started
+
+To get started using Vertex AI, you must have a Google Cloud project. 
+
+- If you don't have a Google Cloud project, you can learn and build on GCP for free using [Free Trail](https://cloud.google.com/free).
+- Once you have a Google Cloud project, you can learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment).
+

 ## Repository structure

 ```bash
-├── community-content - Sample code and tutorials contributed by the community
 ├── notebooks
-│   ├── community - Notebooks contributed by the community
 │   ├── official - Notebooks demonstrating use of each Vertex AI service
 │   │   ├── automl
 │   │   ├── custom
 │   │   ├── ...
+│   ├── community - Notebooks contributed by the community
+│   │   ├── model_garden
+│   │   ├── ...
+├── community-content - Sample code and tutorials contributed by the community
+
 ```
+## Examples

-## Contributing
+<!-- markdownlint-disable MD033 -->
+<table>

-Contributions welcome! See the [Contributing Guide](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/CONTRIBUTING.md).
+  <tr>
+    <th style="text-align: center;">Category</th>
+    <th style="text-align: center;">Product</th>
+    <th style="text-align: center;">Description</th>
+  </tr>
+  <tr>
+    <td>Model</td>
+    <td>
+      <a href="notebooks/community/model_garden"><code>Model Garden/</code></a>
+    </td>
+    <td>
+      Curated collection of first-party, open-source, and third-party models available on Vertex AI including Gemini, Gemma, Llama 3, Claude 3 and many more.
+    </td>
+  </tr>
+  <tr>
+    <td>Data</td>
+    <td>
+      <a href="notebooks/official/feature_store"><code>Feature Store/</code></a>
+    </td>
+    <td>
+      Set up and manage online serving using Vertex AI Feature Store.
+    </td>
+  </tr>
+  <tr>
+    <td></td>
+    <td>
+      <a href="notebooks/official/datasets"><code>datasets/</code></a>
+    </td>
+    <td>
+      Use BigQuery and Data Labeling service with Vertex AI.
+    </td>
+  </tr>
+  <tr>
+    <td>Model development</td>
+    <td>
+      <a href="notebooks/official/automl"><code>automl/</code></a>
+    </td>
+    <td>
+      Train and make predictions on AutoML models
+    </td>
+  </tr>
+  <tr>
+    <td></td>
+    <td>
+      <a href="notebooks/official/custom"><code>custom/</code></a>
+    </td>
+    <td>
+      Create, deploy and serve custom  models on Vertex AI
+    </td>
+  </tr>
+  <tr>
+    <td></td>
+    <td>
+      <a href="notebooks/official/ray_on_vertex_ai"><code>ray_on_vertex_ai/</code></a>
+    </td>
+    <td>
+      Use Colab Enterprise and Vertex AI SDK for Python to connect to the Ray Cluster.
+    </td>
+  </tr>
+  <tr>
+    <td>Deploy and use</td>
+    <td>
+      <a href="notebooks/official/prediction"><code>prediction/</code></a>
+    </td>
+    <td>
+      Build, train and deploy models using prebuilt containers for custom training and prediction.
+    </td>
+  </tr>
+  <tr>
+    <td></td>
+    <td>
+      <a href="notebooks/official/model_registry"><code>model_registry/</code></a>
+    </td>
+    <td>
+      Use Model Registry to create and register a model.
+    </td>
+  </tr>
+  <tr>
+    <td></td>
+    <td>
+      <a href="notebooks/official/explainable_ai"><code>Explainable AI/</code></a>
+    </td>
+    <td>
+      Use Vertex Explainable AI's feature-based and example-based explanations to explain how or why a model produced a specific prediction.
+    </td>
+  </tr>
+  <tr>
+    <td></td>
+    <td>
+      <a href="notebooks/official/ml_metadata"><code>ml_metadata/</code></a>
+    </td>
+    <td>
+      Record the metadata and artifacts and query that metadata to help analyze, debug, and audit the performance of your ML system.
+    </td>
+  </tr>
+  <tr>
+    <td>Tools</td>
+    <td>
+      <a href="notebooks/official/pipelines"><code>Pipelines/</code></a>
+    </td>
+    <td>
+      Use `Vertex AI Pipelines` and `Google Cloud Pipeline Components` to build, tune, or deploy a custom model.
+    </td>
+  </tr>
+</table>
+<!-- markdownlint-enable MD033 -->

-## Getting help

-Please use the [issues page](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues) to provide feedback or submit a bug report.
+## Get help
+
+Please use the [Issues page](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues) to provide feedback or submit a bug report.

 ## Disclaimer

 This is not an officially supported Google product. The code in this repository is for demonstrative purposes only.

-## Feedback

-Please feel free to fill out our [survey](https://bit.ly/vertex-ai-samples-survey) to give us feedback on the repo and its content.
+## References
+- [Vertex AI Jupyter Notebook tutorials](https://cloud.google.com/vertex-ai/docs/tutorials/jupyter-notebooks)
+- Vertex AI [Generative AI](https://github.com/GoogleCloudPlatform/generative-ai) GitHub repository
+- [Vertex AI documentaton](https://cloud.google.com/vertex-ai/docs)
+  
@@ -1,3 +1,3 @@
-torch==1.13.1
+torch==2.2.0
 torchvision==0.9.1
 tensorboard==2.5.0
@@ -0,0 +1,15 @@
+# Vertex AI custom prediction routines samples
+
+## Overview
+Vertex Custom Prediction Routines(CPR) simplify the process of building custom containers 
+and make local model testing easy. Here are the sameple codes for different libraries.
+
+
+### Objectives
+The objective is to provide various samples for Vertex Custom Prediction Routine(CPR).
+
+
+### Supporting libraries
+* torch
+* sklearn
+* xgboost
@@ -0,0 +1,73 @@
+import ast
+import json
+import os
+import pickle
+import torch
+
+from google.cloud.aiplatform.utils import prediction_utils
+from google.cloud.aiplatform.prediction.predictor import Predictor
+from transformers import AutoModelForQuestionAnswering
+from typing import Dict, List
+
+class TorchTransformersPredictor(Predictor):
+
+    def __init__(self):
+        return
+
+    def load(self, artifacts_uri: str) -> None:
+        prediction_utils.download_model_artifacts(artifacts_uri)
+        
+        if os.path.isfile("setup_config.json"):
+            with open("setup_config.json") as setup_config_file:
+                self.setup_config = json.load(setup_config_file)
+
+        if os.path.exists("model.pt"):
+            self.model = AutoModelForQuestionAnswering.from_pretrained("model.pt")
+            self.model.eval()
+        else:
+            raise ValueError("One of the following model files must be provided: model.pt.")
+
+    def preprocess(self, prediction_input: dict) -> torch.Tensor:
+        max_length = self.setup_config["max_length"]
+        instances = prediction_input["instances"]
+        question_context = ast.literal_eval(instances)
+        question = question_context["question"]
+        context = question_context["context"]
+        inputs = self.tokenizer.encode_plus(
+            question,
+            context,
+            max_length=int(max_length),
+            pad_to_max_length=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        input_ids = inputs["input_ids"]
+        attention_mask = inputs["attention_mask"]
+        return torch.Tensor(input_ids, attention_mask)
+    
+    @torch.inference_mode()
+    def predict(self, instances: torch.Tensor) -> List[str]:
+        input_ids, attention_mask = instances
+        outputs = self._model(input_ids, attention_mask)
+        answer_start_scores = outputs.start_logits
+        answer_end_scores = outputs.end_logits
+        
+        num_rows, num_cols = answer_start_scores.shape
+        inferences = []
+        for i in range(num_rows):
+            answer_start_scores_one_seq = answer_start_scores[i].unsqueeze(0)
+            answer_start = torch.argmax(answer_start_scores_one_seq)
+            answer_end_scores_one_seq = answer_end_scores[i].unsqueeze(0)
+            answer_end = torch.argmax(answer_end_scores_one_seq) + 1
+            prediction = self.tokenizer.convert_tokens_to_string(
+                self.tokenizer.convert_ids_to_tokens(
+                    input_ids[i].tolist()[answer_start:answer_end]
+                )
+            )
+            inferences.append(prediction)
+        return inferences
+
+    def postprocess(self, prediction_results: List[str]) -> Dict:
+        return {"predictions": prediction_results}
+     
+
@@ -1,227 +0,0 @@
-# Benchmark report on fine tuning the OpenLLaMA 7B model on Google Cloud Vertex Model Garden
-
-Gary Wei, Software Engineer, Google Cloud
-Dustin Luong, Software Engineer, Google Cloud
-Changyu Zhu, Software Engineer, Google Cloud
-Genquan Duan, Software Engineer, Google Cloud
-
-## Introduction
-
-Fine-tuning of LLMs can be non-trivial to find an optimal configuration of
-machine types, training parameters, and other hyperparameters that achieves a
-good balance between cost efficiency and model performance. To facilitate users
-in conducting tuning experiments, this report benchmarks OpenLLaMA 7B
-fine-tuning on Google Cloud Vertex Model Garden, demonstrating both efficiency
-and effectiveness. The observations are general and can be applied to other LLM
-models.
-
-We benchmarked fine tuning algorithms [LoRA](https://arxiv.org/abs/2106.09685)
-and [QLoRA](https://arxiv.org/abs/2305.14314) supported by
-[huggingface PEFT libraries](https://github.com/huggingface/peft). LoRA, short
-for Low-Rank Adaptation of Large Language Models, is an improved fine tuning
-method where instead of fine tuning all the weights that constitute the weight
-matrix of the pre-trained large language model, two smaller matrices that
-approximate this larger matrix are fine-tuned. QLoRA is an even more
-memory-efficient version of LoRA, where the pretrained model is loaded to GPU
-memory as quantized 4-bit weights, while preserving similar effectiveness to
-LoRA. We also provide simple scripts and parameter settings to reproduce the
-results reported in this report.
-
-In general, there are many factors that affect the performance of fine-tuning
-experiments, such as hardware settings, parameters, cost, and accuracy. It is
-impractical to obtain benchmarks for all possible combinations of these factors.
-Instead, we focus on tuning a subset of related parameters and evaluating their
-impact on a set of chosen metrics. The evaluation metrics are GPU memory usage,
-percentage of parameters tuned, tuning speed, cost, and accuracy. The tuning
-parameters are batch size, lora rank, maximum sequence length, and maximum
-training steps.
-
-## Key takeaways
-
-   **Use QLoRA to minimize the peak GPU requirements**: The QLoRA can
-    significantly reduce the peak GPU memory usage by ~75% compared to LoRA. For
-    OpenLLaMA7b, the peak memory is ~28G for LoRA and ~7G for QLoRA.
-   **Use LoRA to maximize the tuning speed and minimize the tuning cost**: LoRA
-    is ~66% faster than QLoRA in fine tuning speed. LoRA/QLoRA tuning cost is
-    low generally, while LoRA is even ~40% cheaper than QLoRA with the same
-    parameters. Suggest to use QLoRA for limited GPU memories, and LoRA for
-    limited training budgets. For OpenLLaMA7b, the tuning speed for LoRA/QLoRA
-    ~5 samples / 3 samples per second, and the tuning cost for LoRA/QLoRA in 500
-    steps is ~$1/$1.7 on `a2-highgpu-1g` with 1 A100 40G GPU. The tuning cost
-    for QLoRA in 500 steps is $6.75 on n1-standard-8 with 1 V100 GPU, while LoRA
-    could not run because of OOM.
-   **Use QLoRA to tune models with large sequence lengths**. For OpenLLaMA7b,
-    the max sequence length for QLoRA can be 2048 when consuming 16.3G GPU,
-    while the max sequence length for LoRA is 512 when consuming 28.2G GPU, and
-    encounter OOM when max sequence length is 1024.
-   **Both LoRA and QLoRA give similar accuracy improvement after fine tuning.**
-    For OpenLLaMA7b, both LoRA/QLoRA can improve the average accuracy by ~4%
-    evaluating on 3 typical tasks (ARC challenge, HellaSwag and TruthfulQA),
-    after training 1875 steps on dataset
-    [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco).
-   **Use a big batch size if GPU memory is not a constraint**. For OpenLLaMA7b
-    with other default parameters, we suggest using a batch size as 24 for
-    QLoRA, but 2 for LoRA when tuning with 1 A100 40G. We also suggest using a
-    batch size as 8 for QLoRA when tuning with 1 V100. Tuning with LoRA and
-    batch size as 1 got OOM and we don't recommend tuning LoRA with 1 V100.
-
-## Benchmark Details
-
-### Experiment Setup
-
-The benchmark dataset is
-[timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco).
-The training dataset is directly downloaded from hugging face to the VM, before
-every experiment.
-
-The default tuning parameters during benchmark are:
-
-   Host VM: a2-highgpu-1g
-   Accelerator type: 1 A100 40G
-   batch size: 2
-   lora_rank: 16
-   max_seq_length: 512
-   precision_mode: float16
-   max_train_steps: 500
-
-For simplicity, we set the precision mode to `float16` when tuning LoRA models,
-and set the precision to `4bit` for QLoRA.
-
-Sample script to start fine tuning dockers in a VM on GCP.
-
-```shell
-IMAGE_TAG=us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:latest
-docker run  --runtime=nvidia  -e NVIDIA_VISIBLE_DEVICES=0 \
--rm --name "test_gpu"  -it --pull=always ${IMAGE_TAG} \
--task=instruct-lora \
--pretrained_model_id=openlm-research/open_llama_7b \
--dataset_name="timdettmers/openassistant-guanaco"  \
--instruct_column_in_dataset="text" \
--precision_mode="float16" \
--output_dir=<OUTPUT DIR>  \
--lora_rank=2 \
--max_sequence_length=512 \
--learning_rate=2e-4 \
--max_steps=50
-```
-
-### GPU Memory
-
-In this benchmark, we investigated the impact of batch size, lora rank, and
-maximum sequence length on GPU memory, and then made recommendations on the
-maximum batch size for different GPUs.
-
-#### Peak GPU memory by batch size (GB)
-
-<img src="images/openllama_7b_fine_tune_benchmark_report/openllama-7b-peak-gpu-vs-batch-size.png" width="600">
-
-   The QLoRA can significantly reduce the peak GPU memory usage by ~75%
-    compared to LoRA. The peak GPU memory is ~28G for LoRA and ~7G for QLoRA
-    when batch size is 2.
-   QLoRA can support much larger batch sizes than LoRA
-    -   We can use a batch size as 32 for QLoRA, but only 2 for LoRA on 1 A100
-        40G.
-    -   We can use a batch size of 8 for QLoRA on 1 V100 GPU. LoRA will fail
-        with OOM even with a batch size of 1.
-
-#### Peak GPU memory by LoRA rank (GB)
-
-<img src="images/openllama_7b_fine_tune_benchmark_report/openllama-7b-peak-gpu-vs-lora-rank.png" width="600">
-
-   Peak GPU memories are quite similar for different LoRA ranks for both
-    LoRA/QLoRA.
-   The peak GPU memory increasing percentages are very small generally when
-    LoRA rank increases.
-    -   The peak GPU memory increases from 28G with LoRA rank 4 to 29.09G with
-        LoRA rank 64, and the increasing percentage is only ~3.9%.
-
-#### Peak GPU memory by max sequence length for LoRA/QLoRA (GB)
-
-<img src="images/openllama_7b_fine_tune_benchmark_report/openllama-7b-peak-gpu-vs-max-seq-length.png" width="600">
-
-   The peak GPU increases quickly when max sequence length increases for both
-    LoRA/QLoRA, and the increasing rate of LoRA is much faster than QLoRA.
-    -   For LoRA tuning, the GPU memory increased from 20.5G (max sequence
-        length=256) to 28.2G (max sequence length=512), an increase of ~37%.
-    -   For QLoRA tuning, the GPU memory increased from 6.94G (max sequence
-        length=256) to 7.57G (max sequence length=512), an increase of ~9%.
-   The max sequence length for QLoRA can be 2048 when consuming 16.3G GPU,
-    while the max sequence length for LoRA is 512 when consuming 28.2G GPU, and
-    encounter OOM when max sequence length is 1024.
-
-### Fine Tuning Parameters
-
-This section shows the number/percentage of trainable parameters, and the sizes
-of the fine tuned models. LoRA and QLoRA differ only in how they represent the
-precision of their parameters. The total number of parameters and the number of
-trainable parameters are the same for both methods.
-
-| LoRA Rank | Finetuned parameters | Total parameters | Trainable Parameter Percentage | Fine tuned model size (MB) |
-| --------- | -------------------- | ---------------- | ------------------------------ | -------------------------- |
-| 8         | 2.00E+07             | 6.76E+09         | 0.3%                           | 76.4                       |
-| 16        | 4.00E+07             | 6.78E+09         | 0.6%                           | 152.65                     |
-| 32        | 8.00E+07             | 6.82E+09         | 1.2%                           | 305.15                     |
-| 64        | 1.60E+08             | 6.90E+09         | 2.3%                           | 610.15                     |
-
-
-LoRA/QLoRA tunes quite a small fraction (only 0.3% with LoRA rank=8) of all
-parameters, and the tuned models are very small (only 76.4MB with LoRA rank=8).
-
-### Fine Tuning Speed And Costs
-
-The fine-tuning speed and cost are affected by various factors, such as the
-GPUs, LoRA ranks, and max sequence lengths.
-
-   LoRA is ~66% faster than QLoRA in fine tuning speed. The tuning speed for
-    LoRA/QLoRA ~5 samples / 3 samples per second on 1 A100 40G GPU
-   Higher LoRA ranks, slower tuning speed for both LoRA/QLoRA.
-    -   LoRA tuning speed reduces from ~5 samples per second with LoRA rank as 8
-        to ~4 samples per second with LoRA rank as 64, slowed down by 20%.
-    -   QLoRA tuning speed reduces from ~3 samples per second with LoRA rank as
-        8 to ~2.5 samples per second with LoRA rank as 64, slowed down by 17%.
-
-<img src="images/openllama_7b_fine_tune_benchmark_report/openllama-7b-tune-speed-vs-lora-rank.png" width="600">
-
-   Longer sequence lengths, slower tuning speed.
-    -   LoRA tuning speed reduces from ~5.56 samples per second with max
-        sequence length as 256 to ~4.84 samples per second with max sequence
-        length as 512 slowed down by 13%.
-    -   LoRA tuning speed reduces from ~2.95 samples per second with max
-        sequence length as 256 to ~2.88 samples per second with max sequence
-        length as 512 slowed down by ~2.4%.
-
-<img src="images/openllama_7b_fine_tune_benchmark_report/openllama-7b-tune-speed-lora-qlora.png" width="600">
-
-   LoRA/QLoRA tuning cost is low generally, while LoRA is even ~40% cheaper
-    than QLoRA with the same parameters.
-    -   The LoRA/QLoRA fine tuning cost for 500 steps is ~$1/$1.7 on 1 A100 40G.
-    -   The tuning cost for QLoRA in 500 steps is $6.75 on n1-standard-8 with 1
-        V100 GPU, while LoRA could not run because of OOM.
-
-<img src="images/openllama_7b_fine_tune_benchmark_report/openllama-7b-tune-cost-lora-qlora.png" width="600">
-
-### Accuracy
-
-We fine tuned Open Llama 7B model with
-[timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco),
-and report accuracy similar to the
-[HuggingFace leaderboard](https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard)
-using
-[Eleuther AI Language Model Evaluation Harness](https://github.com/EleutherAI/lm-evaluation-harness).
-[HuggingFace leaderboard](https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard)
-mainly compares models on ARC, HellaSwag, MMLU, and TruthfulQA. The authors did
-not publish OpenLLaMA 7B on MMLU
-([link](https://huggingface.co/openlm-research/open_llama_7b)). Therefore, we
-only benchmark accuracies on ARC, HellaSwag, and TruthfulQA.
-
-|                                                              | Mean | ARC  | HellaSwag | TruthfulQA | Tuning Parameters                                            |
-| ------------------------------------------------------------ | ---- | ---- | --------- | ---------- | ------------------------------------------------------------ |
-| OpenLLaMA7B ([Original Report](https://huggingface.co/openlm-research/open_llama_7b)) | 0.49 | 0.41 | 0.73      | 0.34       | n/a                                                          |
-| OpenLLaMA7B ([Re-run with lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness)) | 0.51 | 0.47 | 0.72      | 0.35       | n/a                                                          |
-| OpenLLaMA7B+LoRA                                             | 0.56 | 0.48 | 0.74      | 0.45       | LoRA Rank=16; Max Sequence Length=512;Learning Rate=1e-4; Train steps=1875 |
-| OpenLLaMA7B+QLoRA                                            | 0.53 | 0.45 | 0.73      | 0.42       | LoRA Rank=16; Max Sequence Length=512; Learning Rate=1e-4; Train steps=1875 |
-
-   The base OpenLLaMA7B model gets better performance (2%) when using the
-    [Eleuther AI Language Model Evaluation Harness](https://github.com/EleutherAI/lm-evaluation-harness).
-   LoRA/QLoRA can improve the performance by ~2-4% when trained for 1875 steps
-    with learning rate 1e-4.
@@ -16,7 +16,7 @@ Pillow==10.3.0
 rawpy==0.18.1
 scipy==1.11.3
 scikit-image==0.22.0
-scikit-learn==1.3.2
+scikit-learn==1.5.0
 tensorboard==2.15.0
 tensorboardX==2.6.2.2
 tqdm==4.66.3
@@ -231,6 +231,22 @@ def download_image(url: str) -> str:
  return Image.open(io.BytesIO(response.content))


+def resize_image(image: Any, new_width: int = 1000) -> Any:
+    """Resizes an image to a certain width.
+
+    Args:
+      image: The image which has to be resized.
+      new_width: New width of the image.
+
+    Returns:
+      New resized image.
+    """
+    width, height = image.size
+    new_height = int(height * new_width / width)
+    new_img = image.resize((new_width, new_height))
+    return new_img
+
+
 def load_img(path: str) -> Any:
  """Reads image from path and return PIL.Image instance.

@@ -372,6 +388,7 @@ def get_resource_id(accelerator_type: str, is_for_training: bool) -> str:
      "NVIDIA_TESLA_V100": "custom_model_training_nvidia_v100_gpus",
      "NVIDIA_L4": "custom_model_training_nvidia_l4_gpus",
      "NVIDIA_TESLA_A100": "custom_model_training_nvidia_a100_gpus",
+      "NVIDIA_A100_80GB": "custom_model_training_nvidia_a100_80gb_gpus",
      "NVIDIA_TESLA_T4": "custom_model_training_nvidia_t4_gpus",
      "TPU_V5e": "custom_model_training_tpu_v5e",
      "TPU_V3": "custom_model_training_tpu_v3",
@@ -380,6 +397,7 @@ def get_resource_id(accelerator_type: str, is_for_training: bool) -> str:
      "NVIDIA_TESLA_V100": "custom_model_serving_nvidia_v100_gpus",
      "NVIDIA_L4": "custom_model_serving_nvidia_l4_gpus",
      "NVIDIA_TESLA_A100": "custom_model_serving_nvidia_a100_gpus",
+      "NVIDIA_A100_80GB": "custom_model_serving_nvidia_a100_80gb_gpus",
      "NVIDIA_TESLA_T4": "custom_model_serving_nvidia_t4_gpus",
      "TPU_V5e": "custom_model_serving_tpu_v5e",
  }
@@ -0,0 +1,32 @@
+import numpy as np
+from kfp.v2 import dsl
+
+@dsl.component(base_image='python:3.8',packages_to_install=['google-cloud-aiplatform==1.36.0'])
+def async_predict(
+    endpoint_id: str,
+    instances: dict,
+) -> np.ndarray:
+    import numpy as np
+    from google.cloud import aiplatform
+
+    endpoint = aiplatform.Endpoint(endpoint_id)
+    response = await endpoint.predict_async(instances)
+    predictions = np.asarray(response.predictions)
+    print(predictions.tolist())
+    return predictions
+
+@dsl.pipeline(name='async-prediction')
+def pipeline_prediction():
+    project = "projects/990000000009/locations/us-west1"
+    endpoint_id = project + "/endpoints/2200000000000000002"
+    instances = [{
+        "key1": "value1",
+        "key2": 2
+    }]
+    async_predict(endpoint_id, instances)
+
+if __name__ == "__main__":
+    from kfp.v2 import compiler
+    compiler.Compiler().compile(
+        pipeline_func=pipeline_prediction,
+        package_path='async_prediction.json')
@@ -1,37 +1,51 @@
-from kfp.v2 import dsl
-
-@dsl.component(base_image='python:3.8',packages_to_install=['google-cloud-aiplatform==1.36.0'])
-def deploy_model(
-    model_id: str,
-    endpoint_id: str,
-    machine_type: str,
-    min_replica_count: int,
-    max_replica_count: int,
-):
-    import json
-    from google.cloud import aiplatform
-
-    model = aiplatform.Model(model_id)
-    endpoint = aiplatform.Endpoint(endpoint_id)
-
-    endpoint = model.deploy(
-        endpoint=endpoint,
-        machine_type=machine_type,
-        min_replica_count=min_replica_count,
-        max_replica_count=max_replica_count,
-    )
-
-
-@dsl.pipeline(name='deploy-model')
-def pipeline_deploy_model():
-    project = "projects/990000000009/locations/us-west1"
-    model_id = project + "/models/1100000000000000001"
-    endpoint_id = project + "/endpoints/2200000000000000002"
-    deploy_model(model_id, endpoint_id, "n1-standard-2", 1, 1)
-
-
-if __name__ == "__main__":
-    from kfp.v2 import compiler
-    compiler.Compiler().compile(
-        pipeline_func=pipeline_deploy_model,
-        package_path='deploy_model.json')
+from kfp.v2 import dsl
+
+@dsl.component(base_image='python:3.8',packages_to_install=['google-cloud-aiplatform==1.36.0'])
+def deploy_model(
+    model_id: str,
+    endpoint_id: str,
+    machine_type: str,
+    min_replica_count: int,
+    max_replica_count: int,
+):
+    import json
+    from google.cloud import aiplatform
+
+    model = aiplatform.Model(model_id)
+    endpoint = aiplatform.Endpoint(endpoint_id)
+
+    endpoint = model.deploy(
+        endpoint=endpoint,
+        machine_type=machine_type,
+        min_replica_count=min_replica_count,
+        max_replica_count=max_replica_count,
+    )
+
+
+@dsl.component(base_image='python:3.8',packages_to_install=['google-cloud-aiplatform==1.36.0'])
+def delete_endpoint(
+    endpoint_id: str,
+):
+    from google.cloud import aiplatform
+
+    endpoint = aiplatform.Endpoint(endpoint_id)
+    endpoint.undeploy_all()
+    endpoint.delete()
+
+
+@dsl.pipeline(name='deploy-model')
+def pipeline_deploy_model():
+    project = "projects/990000000009/locations/us-west1"
+    model_id = project + "/models/1100000000000000001"
+    endpoint_id = project + "/endpoints/2200000000000000002"
+    deploy_model(model_id, endpoint_id, "n1-standard-2", 1, 1)
+
+    # After serving predictions, recycling computing resources
+    delete_endpoint(endpoint_id)
+
+
+if __name__ == "__main__":
+    from kfp.v2 import compiler
+    compiler.Compiler().compile(
+        pipeline_func=pipeline_deploy_model,
+        package_path='deploy_model.json')
@@ -52,6 +52,7 @@
 /notebooks/community/model_garden/model_garden_jax_paligemma_deployment.ipynb @minwoo33park
 /notebooks/community/model_garden/model_garden_jax_paligemma_finetuning.ipynb @minwoo33park
 /notebooks/community/model_garden/model_garden_jax_stable_diffusion_xl.ipynb @weigary
+/notebooks/community/model_garden/model_garden_mammut.ipynb @ivywang9331
 /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
@@ -107,7 +108,6 @@
 /notebooks/community/model_garden/model_garden_movinet_clip_classification.ipynb @KCFindstr
 /notebooks/community/model_garden/model_garden_movinet_action_recognition.ipynb @KCFindstr
 /notebooks/community/model_garden/model_garden_pytorch_open_clip.ipynb @lydhr
-/notebooks/community/model_garden/model_garden_pytorch_llama2_peft.ipynb @genquan9
 /notebooks/community/model_garden/model_garden_pytorch_codellama.ipynb @xiangxu-google
 /notebooks/community/model_garden/model_garden_pytorch_nllb.ipynb @weigary
 /notebooks/community/model_garden/model_garden_pytorch_vicuna.ipynb @dstnluong-google
@@ -116,6 +116,7 @@
 /notebooks/community/model_garden/model_garden_pytorch_biomedclip.ipynb @KCFindstr
 /notebooks/community/model_garden/model_garden_pytorch_imagebind.ipynb @kathyyu-google
 /notebooks/community/persistent_resource/00_persistent_resource_getting_started_cli.ipynb @jbrache
+/notebooks/community/persistent_resource/00_persistent_resource_getting_started_sdk.ipynb @jbrache
 /notebooks/community/model_garden/model_garden_pytorch_llama2_deployment.ipynb @genquan9
 /notebooks/community/model_garden/model_garden_pytorch_llama2_peft_finetuning.ipynb @genquan9
 /notebooks/community/model_garden/model_garden_pytorch_llama2_quantization.ipynb @dstnluong-google
@@ -124,6 +125,8 @@
 /notebooks/community/model_garden/model_garden_pytorch_llama2_rlhf_tuning.ipynb @genquan9
 /notebooks/community/model_garden/model_garden_pytorch_llama3_deployment.ipynb @kathyyu-google
 /notebooks/community/model_garden/model_garden_pytorch_llama3_finetuning.ipynb @kathyyu-google
+/notebooks/community/model_garden/model_garden_pytorch_llama3_1_deployment.ipynb @xiangxu-google
+/notebooks/community/model_garden/model_garden_pytorch_llama3_1_finetuning.ipynb @wrzhao-work
 /notebooks/community/model_garden/model_garden_pytorch_wizard_coder.ipynb @KCFindstr
 /notebooks/community/model_registry/get_started_with_vertex_ai_deployer.ipynb angelmontero@ @inardini
 /notebooks/community/model_garden/model_garden_pytorch_wizard_lm.ipynb @KCFindstr
@@ -135,10 +138,17 @@
 /notebooks/community/model_garden/model_garden_pytorch_mixtral_deployment.ipynb @KCFindstr
 /notebooks/community/model_garden/model_garden_gemma_deployment_on_gke.ipynb @vilobhmm
 /notebooks/community/model_garden/model_garden_gemma_deployment_on_vertex.ipynb @kathyyu-google
+/notebooks/community/model_garden/model_garden_gemma2_deployment_on_vertex.ipynb @kathyyu-google
 /notebooks/community/model_garden/model_garden_gemma_finetuning_on_vertex.ipynb @KCFindstr
 /notebooks/community/model_garden/model_garden_pytorch_gemma_peft_finetuning_hf.ipynb @KCFindstr
 /notebooks/community/model_garden/model_garden_pytorch_stable_diffusion_deployment_1_5.ipynb @weigary
 /notebooks/community/model_garden/model_garden_pytorch_stable_diffusion_gradio.ipynb @weigary
 /notebooks/community/model_garden/model_garden_pytorch_sd_2_1_finetuning_dreambooth.ipynb @weigary
 /notebooks/community/model_garden/model_garden_pytorch_sd_xl_finetuning_dreambooth_lora.ipynb @weigary
+/notebooks/community/model_garden/model_garden_pytorch_sd_2_1_local_finetuning_dreambooth.ipynb @weigary
 /notebooks/community/model_garden/model_garden_timesfm_deployment_on_vertex.ipynb @siriuz42
+/notebooks/community/model_garden/model_garden_llama_guard_deployment.ipynb @kathyyu-google
+/notebooks/community/model_garden/model_garden_rag.ipynb @kathyyu-google
+/notebooks/community/model_garden/synthetic_data_generation_using_llama3_1.ipynb @xiangxu-google
+/notebooks/community/model_garden/model_garden_autosxs_evaluation_llama3_1.ipynb @inardini
+/notebooks/community/model_garden/model_garden_openai_api_llama3_1.ipynb @inardini
@@ -31,9 +31,9 @@
   "source": [
    "# Exploratory Data Analysis with R and BigQuery\n",
    "\n",
-    "**Authors**:  [Alok Pattani](https://github.com/alokpattani), [Khalid Salama](https://github.com/ksalama)\n",
+    "**Author**:  [Alok Pattani](https://github.com/alokpattani)\n",
    "\n",
-    "**Last Updated**:  February 2024\n",
+    "**Last Updated**:  July 2024\n",
    "\n",
    "## Overview\n",
    "\n",
@@ -79,7 +79,9 @@
  {
   "cell_type": "code",
   "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
   "outputs": [],
   "source": [
    "version"
@@ -135,7 +137,9 @@
  {
   "cell_type": "code",
   "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
   "outputs": [],
   "source": [
    "bq_auth(use_oob = TRUE)"
@@ -151,7 +155,9 @@
  {
   "cell_type": "code",
   "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
   "outputs": [],
   "source": [
    "# Set the project ID\n",
@@ -168,13 +174,27 @@
  {
   "cell_type": "code",
   "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
   "outputs": [],
   "source": [
    "# Set your Cloud Storage bucket name\n",
    "BUCKET_NAME <- \"[YOUR-BUCKET-NAME]\""
   ]
  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Set default height/width for plots generated\n",
+    "options(repr.plot.height = 9, repr.plot.width = 16)"
+   ]
+  },
  {
   "cell_type": "markdown",
   "metadata": {},
@@ -192,31 +212,64 @@
  {
   "cell_type": "code",
   "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
   "outputs": [],
   "source": [
    "sql_query_template <- \"\n",
    "    SELECT\n",
-    "      ROUND(weight_pounds, 2) AS weight_pounds,\n",
-    "      is_male,\n",
-    "      mother_age,\n",
-    "      plurality,\n",
-    "      gestation_weeks,\n",
-    "      cigarette_use,\n",
-    "      alcohol_use,\n",
-    "      CAST(ABS(FARM_FINGERPRINT(CONCAT(\n",
-    "        CAST(YEAR AS STRING), CAST(month AS STRING), \n",
-    "        CAST(weight_pounds AS STRING)))\n",
-    "        ) AS STRING) AS key\n",
+    "      TIMESTAMP_DIFF(dropoff_datetime, pickup_datetime, MINUTE) AS trip_time_minutes,  \n",
+    "\n",
+    "      passenger_count,\n",
+    "\n",
+    "      ROUND(trip_distance, 1) AS trip_distance_miles,\n",
+    "\n",
+    "      rate_code,\n",
+    "      /* Mapping from rate code to type from description column in BQ table schema */\n",
+    "      (CASE \n",
+    "        WHEN rate_code = '1.0'\n",
+    "          THEN 'Standard rate'\n",
+    "        WHEN rate_code = '2.0'\n",
+    "          THEN 'JFK'\n",
+    "        WHEN rate_code = '3.0'\n",
+    "          THEN 'Newark'\n",
+    "        WHEN rate_code = '4.0'\n",
+    "          THEN 'Nassau or Westchester'\n",
+    "        WHEN rate_code = '5.0'\n",
+    "          THEN 'Negotiated fare'\n",
+    "        WHEN rate_code = '6.0'\n",
+    "          THEN 'Group ride'\n",
+    "        /* Several NULL AND some '99.0' values go here */\n",
+    "        ELSE 'Unknown'\n",
+    "        END)\n",
+    "        AS rate_type,\n",
+    "\n",
+    "      fare_amount,\n",
+    "\n",
+    "      CAST(ABS(FARM_FINGERPRINT(\n",
+    "        CONCAT(\n",
+    "          CAST(trip_distance AS STRING), \n",
+    "          CAST(fare_amount AS STRING)\n",
+    "          )\n",
+    "        ))\n",
+    "        AS STRING)\n",
+    "        AS key\n",
+    "\n",
    "    FROM\n",
-    "        publicdata.samples.natality\n",
-    "    WHERE \n",
-    "      year > 2000\n",
-    "      AND weight_pounds > 0\n",
-    "      AND mother_age > 0\n",
-    "      AND plurality > 0\n",
-    "      AND gestation_weeks > 0\n",
-    "      AND month > 0\n",
+    "      `bigquery-public-data.new_york_taxi_trips.tlc_yellow_trips_2022`\n",
+    "\n",
+    "    /* Filter out some outlier or hard to understand values */\n",
+    "    WHERE\n",
+    "      (TIMESTAMP_DIFF(dropoff_datetime, pickup_datetime, MINUTE)\n",
+    "        BETWEEN 0.01 AND 120)\n",
+    "      AND\n",
+    "      (passenger_count BETWEEN 1 AND 10)\n",
+    "      AND\n",
+    "      (trip_distance BETWEEN 0.01 AND 100)\n",
+    "      AND\n",
+    "      (fare_amount BETWEEN 0.01 AND 250)\n",
+    "\n",
    "    LIMIT %s\n",
    "\""
   ]
@@ -232,14 +285,16 @@
  {
   "cell_type": "code",
   "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
   "outputs": [],
   "source": [
    "sample_size <- 10000\n",
    "\n",
    "sql_query <- sprintf(sql_query_template, sample_size)\n",
    "\n",
-    "natality_data <- bq_table_download(\n",
+    "taxi_trip_data <- bq_table_download(\n",
    "    bq_project_query(\n",
    "        PROJECT_ID, \n",
    "        query = sql_query\n",
@@ -257,31 +312,37 @@
  {
   "cell_type": "code",
   "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
   "outputs": [],
   "source": [
    "# View the query result\n",
-    "head(natality_data)"
+    "head(taxi_trip_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
   "outputs": [],
   "source": [
    "# Show # of rows and data types of each column\n",
-    "str(natality_data)"
+    "str(taxi_trip_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
   "outputs": [],
   "source": [
    "# View the results summary\n",
-    "summary(natality_data)"
+    "summary(taxi_trip_data)"
   ]
  },
  {
@@ -294,27 +355,31 @@
  {
   "cell_type": "code",
   "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
   "outputs": [],
   "source": [
-    "# Display the distribution of baby weights using a histogram\n",
+    "# Display the distribution of fare amounts using a histogram\n",
    "ggplot(\n",
-    "    data = natality_data, \n",
-    "    aes(x = weight_pounds)\n",
+    "    data = taxi_trip_data, \n",
+    "    aes(x = fare_amount)\n",
    "    ) + \n",
-    "geom_histogram(bins = 200)"
+    "geom_histogram(bins = 100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
   "outputs": [],
   "source": [
-    "# Display the relationship between gestation weeks and baby weights \n",
+    "# Display the relationship between trip distance and fare amount\n",
    "ggplot(\n",
-    "    data = natality_data, \n",
-    "    aes(x = gestation_weeks, y = weight_pounds)\n",
+    "    data = taxi_trip_data, \n",
+    "    aes(x = trip_distance_miles, y = fare_amount)\n",
    "    ) + \n",
    "geom_point() + \n",
    "geom_smooth(method = \"lm\")"
@@ -325,23 +390,41 @@
   "metadata": {},
   "source": [
    "### Performing the processing in BigQuery\n",
-    "Create a function that finds the number of records and the average weight for each value of the chosen column."
+    "Create a function that finds the number of trips and the average fare amount for each value of the chosen column."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
   "outputs": [],
   "source": [
-    "get_distinct_values <- function(column_name) {\n",
+    "get_distinct_value_aggregates <- function(column) {\n",
    "    query <- paste0(\n",
-    "        'SELECT ', column_name, ', \n",
-    "            COUNT(1) AS num_babies,\n",
-    "            AVG(weight_pounds) AS avg_wt\n",
-    "        FROM publicdata.samples.natality\n",
-    "        WHERE year > 2000\n",
-    "        GROUP BY ', column_name)\n",
+    "        'SELECT ', \n",
+    "          column, \n",
+    "          ', \n",
+    "          COUNT(1) AS num_trips,\n",
+    "          AVG(fare_amount) AS avg_fare_amount\n",
+    "        \n",
+    "        FROM\n",
+    "          `bigquery-public-data.new_york_taxi_trips.tlc_yellow_trips_2022`\n",
+    "        \n",
+    "        WHERE\n",
+    "          (TIMESTAMP_DIFF(dropoff_datetime, pickup_datetime, MINUTE) \n",
+    "            BETWEEN 0.01 AND 120)\n",
+    "          AND\n",
+    "          (passenger_count BETWEEN 1 AND 10)\n",
+    "          AND\n",
+    "          (trip_distance BETWEEN 0.01 AND 100)\n",
+    "          AND\n",
+    "          (fare_amount BETWEEN 0.01 AND 250)\n",
+    "        \n",
+    "        GROUP BY 1\n",
+    "        '\n",
+    "        )\n",
    "    \n",
    "    bq_table_download(\n",
    "        bq_project_query(\n",
@@ -362,20 +445,23 @@
  {
   "cell_type": "code",
   "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
   "outputs": [],
   "source": [
-    "df <- get_distinct_values('mother_age')\n",
+    "df <- get_distinct_value_aggregates(\n",
+    "  'TIMESTAMP_DIFF(dropoff_datetime, pickup_datetime, MINUTE) AS trip_time_minutes')\n",
    "\n",
    "ggplot(\n",
    "    data = df, \n",
-    "    aes(x = mother_age, y = num_babies)\n",
+    "    aes(x = trip_time_minutes, y = num_trips)\n",
    "    ) + \n",
    "geom_line()\n",
    "\n",
    "ggplot(\n",
-    "    data = df, \n",
-    "    aes(x = mother_age, y = avg_wt)\n",
+    "    data = df,\n",
+    "    aes(x = trip_time_minutes, y = avg_fare_amount)\n",
    "    ) + \n",
    "geom_line()"
   ]
@@ -383,64 +469,88 @@
  {
   "cell_type": "code",
   "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
   "outputs": [],
   "source": [
-    "df <- get_distinct_values('is_male')\n",
+    "df <- get_distinct_value_aggregates('passenger_count')\n",
    "\n",
    "ggplot(\n",
    "    data = df, \n",
-    "    aes(x = is_male, y = num_babies)\n",
+    "    aes(x = passenger_count, y = num_trips)\n",
+    "    ) + \n",
+    "geom_col() +\n",
+    "scale_x_continuous(breaks = 1:10)\n",
+    "\n",
+    "ggplot(\n",
+    "    data = df, \n",
+    "    aes(x = passenger_count, y = avg_fare_amount)\n",
+    "    ) + \n",
+    "geom_col() +\n",
+    "scale_x_continuous(breaks = 1:10)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "df <- get_distinct_value_aggregates('ROUND(trip_distance, 0) AS trip_distance_miles')\n",
+    "\n",
+    "ggplot(\n",
+    "    data = df, \n",
+    "    aes(x = trip_distance_miles, y = num_trips)\n",
+    "    ) + \n",
+    "geom_line()\n",
+    "\n",
+    "ggplot(\n",
+    "    data = df,\n",
+    "    aes(x = trip_distance_miles, y = avg_fare_amount)\n",
+    "    ) + \n",
+    "geom_line()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "df <- get_distinct_value_aggregates(\"\n",
+    "  (CASE \n",
+    "    WHEN rate_code = '1.0'\n",
+    "      THEN 'Standard rate'\n",
+    "    WHEN rate_code = '2.0'\n",
+    "      THEN 'JFK'\n",
+    "    WHEN rate_code = '3.0'\n",
+    "      THEN 'Newark'\n",
+    "    WHEN rate_code = '4.0'\n",
+    "      THEN 'Nassau or Westchester'\n",
+    "    WHEN rate_code = '5.0'\n",
+    "      THEN 'Negotiated fare'\n",
+    "    WHEN rate_code = '6.0'\n",
+    "      THEN 'Group ride'\n",
+    "    /* Several NULL AND some '99.0' values go here */\n",
+    "    ELSE 'Unknown'\n",
+    "    END)\n",
+    "    AS rate_type\n",
+    "  \")\n",
+    "\n",
+    "ggplot(\n",
+    "    data = df,\n",
+    "    aes(x = rate_type, y = num_trips)\n",
    "    ) + \n",
    "geom_col()\n",
    "\n",
    "ggplot(\n",
-    "    data = df, \n",
-    "    aes(x = is_male, y = avg_wt)\n",
-    "    ) + \n",
-    "geom_col()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "df <- get_distinct_values('plurality')\n",
-    "\n",
-    "ggplot(\n",
-    "    data = df, \n",
-    "    aes(x = plurality, y = num_babies)\n",
-    "    ) + \n",
-    "geom_col() + \n",
-    "scale_y_log10()\n",
-    "\n",
-    "ggplot(\n",
    "    data = df,\n",
-    "    aes(x = plurality, y = avg_wt)\n",
-    "    ) + \n",
-    "geom_col()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "df <- get_distinct_values('gestation_weeks')\n",
-    "\n",
-    "ggplot(\n",
-    "    data = df,\n",
-    "    aes(x = gestation_weeks, y = num_babies)\n",
-    "    ) + \n",
-    "geom_col() + \n",
-    "scale_y_log10()\n",
-    "\n",
-    "ggplot(\n",
-    "    data = df,\n",
-    "    aes(x = gestation_weeks, y = avg_wt)\n",
+    "    aes(x = rate_type, y = avg_fare_amount)\n",
    "    ) + \n",
    "geom_col()"
   ]
@@ -455,7 +565,9 @@
  {
   "cell_type": "code",
   "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
   "outputs": [],
   "source": [
    "# Prepare training and evaluation data from BigQuery\n",
@@ -489,7 +601,9 @@
  {
   "cell_type": "code",
   "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
   "outputs": [],
   "source": [
    "print(paste0(\"Training instances count: \", nrow(train_data)))\n",
@@ -500,23 +614,27 @@
  {
   "cell_type": "code",
   "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
   "outputs": [],
   "source": [
-    "# Write data frames to local CSV files, without headers or row names\n",
+    "# Write data frames to local CSV files, with headers\n",
    "dir.create(file.path('data'), showWarnings = FALSE)\n",
    "\n",
    "write.table(train_data, \"data/train_data.csv\", \n",
-    "   row.names = FALSE, col.names = FALSE, sep = \",\")\n",
+    "   row.names = FALSE, col.names = TRUE, sep = \",\")\n",
    "\n",
    "write.table(eval_data, \"data/eval_data.csv\", \n",
-    "   row.names = FALSE, col.names = FALSE, sep = \",\")"
+    "   row.names = FALSE, col.names = TRUE, sep = \",\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
   "outputs": [],
   "source": [
    "# Upload CSV data to Cloud Storage by passing gsutil commands to system\n",
@@ -541,9 +659,9 @@
 "metadata": {
  "environment": {
   "kernel": "conda-env-r-r",
-   "name": "workbench-notebooks.m115",
+   "name": "workbench-notebooks.m123",
   "type": "gcloud",
-   "uri": "gcr.io/deeplearning-platform-release/workbench-notebooks:m115"
+   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m123"
  },
  "kernelspec": {
   "display_name": "R (Local)",
@@ -0,0 +1,741 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ur8xi4C7S06n"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2024 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "JAPoU8Sm5E6e"
+      },
+      "source": [
+        "# Vertex AI Model Garden - Evaluate Llama 3.1 models using Vertex AI AutoSxS\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_autosxs_evaluation_llama3_1.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_autosxs_evaluation_llama3_1.ipynb\"\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_autosxs_evaluation_llama3_1.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_autosxs_evaluation_llama3_1.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tvgnzT1CKxrO"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demostrates how to use the Vertex AI automatic side-by-side (AutoSxS) tool to evaluate Llama 3.1 models for a question-answering task.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Choose the Llama 3.1 models you want to compare.\n",
+        "\n",
+        "- Create an evaluation dataset with question-answer data.\n",
+        "\n",
+        "- Create and run a Vertex AI AutoSxS pipeline that generates judgments and a set of AutoSxS metrics using the generated judgments.\n",
+        "\n",
+        "- Print the judgments and AutoSxS metrics.\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "61RBz8LLbxCR"
+      },
+      "source": [
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "No17Cw5hgx12"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tFy3H3aPgx12"
+      },
+      "outputs": [],
+      "source": [
+        "! pip3 install --upgrade --user --quiet google-cloud-aiplatform google-cloud-pipeline-components\n",
+        "! pip3 install --upgrade --user --quiet openai gcsfs"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "R5Xep4W9lq-Z"
+      },
+      "source": [
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "XRvKdaPDTznN"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "SbmM4z7FOBpM"
+      },
+      "source": [
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dmWOrTJ3gx13"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
+        "\n",
+        "Authenticate your environment on Google Colab.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NyKGtVQjgx13"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "DF4l8DTdWgPY"
+      },
+      "source": [
+        "### Set Google Cloud project information\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Nqwi-5ufWp_B"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"<your-project-id>\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Set the region of the instance\n",
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zgPO1eR3CYjk"
+      },
+      "source": [
+        "### Create a Cloud Storage bucket\n",
+        "\n",
+        "Create a storage bucket to store tutorial artifacts."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "MzGDU7TWdts_"
+      },
+      "outputs": [],
+      "source": [
+        "BUCKET_NAME = \"<your-bucket-name>\"  # @param {type:\"string\"}\n",
+        "\n",
+        "BUCKET_URI = f\"gs://{BUCKET_NAME}\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-EcIXiGsCePi"
+      },
+      "source": [
+        "**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 {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0Wn8ZkcV86KR"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "B8DawN9D9NLU"
+      },
+      "outputs": [],
+      "source": [
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "jVYoyDl165EE"
+      },
+      "source": [
+        "### Import libraries\n",
+        "\n",
+        "Import libraries to use in this tutorial."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "c1tEW-U968h8"
+      },
+      "outputs": [],
+      "source": [
+        "import random\n",
+        "import string\n",
+        "\n",
+        "import openai\n",
+        "import pandas as pd\n",
+        "from google.auth import default, transport\n",
+        "from google.cloud import aiplatform\n",
+        "from google_cloud_pipeline_components.v1 import model_evaluation\n",
+        "from kfp import compiler"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ZXnx1_CtEV5L"
+      },
+      "source": [
+        "### Set variables\n",
+        "\n",
+        "Before starting, you must decide how to access Llama 3.1 models. You can access Llama 3.1 models in just a few clicks using Model-as-a-Service (MaaS) without any setup or infrastructure hassles. You can also access Llama models for self-service in Vertex AI Model Garden, allowing you to choose your preferred infrastructure.\n",
+        "\n",
+        "This tutorial assumes that you deploy a self-managed instance of the Llama 3.1 model and compare it with Llama 3 405b using Model-as-a-Service (MaaS). Notice, only `us-central1` is supported region for Llama 3.1 models using Model-as-a-Service (MaaS).\n",
+        "\n",
+        "[Check out Llama 3 model card](https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama3?_ga=2.31261500.2048242469.1721714335-1107467625.1721655511) to learn how to deploy a Llama 3.1 models on Vertex AI."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "XASp0SPNEX10"
+      },
+      "outputs": [],
+      "source": [
+        "SELF_DEPLOYED_ENDPOINT_REGION = \"<your-endpoint-region>\"  # @param {type:\"string\"}\n",
+        "SELF_DEPLOYED_ENDPOINT_ID = \"<your-endpoint-id>\"  # @param {type:\"string\"}\n",
+        "MODEL_LOCATION = \"us-central1\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "_sdmrDed2aHd"
+      },
+      "source": [
+        "### Helpers"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "pACHdEUf2bfq"
+      },
+      "outputs": [],
+      "source": [
+        "def self_model_generate(\n",
+        "    question,\n",
+        "    context,\n",
+        "    endpoint_id=SELF_DEPLOYED_ENDPOINT_ID,\n",
+        "    endpoint_location=SELF_DEPLOYED_ENDPOINT_REGION,\n",
+        "    **model_kwargs,\n",
+        "):\n",
+        "    \"\"\"Generate a response from a self-managed Llama 3.1 model.\"\"\"\n",
+        "\n",
+        "    aiplatform.init(project=PROJECT_ID, location=endpoint_location)\n",
+        "\n",
+        "    prompt = \"\"\"You are an AI assistant. Your goal is to answer questions using the pieces of context. \"\"\"\n",
+        "    prompt += f\"\"\"Question: {question}.\"\"\"\n",
+        "    prompt += f\"\"\"Context: {context}.\"\"\"\n",
+        "    prompt += \"\"\"Answer:\"\"\"\n",
+        "\n",
+        "    instance = {\"prompt\": prompt}\n",
+        "    instance.update(model_kwargs)\n",
+        "    instances = [instance]\n",
+        "\n",
+        "    endpoint = aiplatform.Endpoint(endpoint_id)\n",
+        "    response = endpoint.predict(instances=instances)\n",
+        "    return response.predictions[0][len(prompt) + 1 :]\n",
+        "\n",
+        "\n",
+        "def maas_generate(\n",
+        "    question,\n",
+        "    context,\n",
+        "    model=\"meta/llama3-405b-instruct-maas\",\n",
+        "    model_location=MODEL_LOCATION,\n",
+        "    **model_kwargs,\n",
+        "):\n",
+        "    \"\"\"Generate a response from a MaaS Llama 3.1 model.\"\"\"\n",
+        "\n",
+        "    creds, _ = default()\n",
+        "    auth_req = transport.requests.Request()\n",
+        "    creds.refresh(auth_req)\n",
+        "    if model_kwargs is None:\n",
+        "        model_kwargs = {}\n",
+        "\n",
+        "    client = openai.OpenAI(\n",
+        "        base_url=f\"https://{model_location}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{model_location}/endpoints/openapi/chat/completions?\",\n",
+        "        api_key=creds.token,\n",
+        "    )\n",
+        "\n",
+        "    response = client.chat.completions.create(\n",
+        "        model=model,\n",
+        "        messages=[\n",
+        "            {\n",
+        "                \"role\": \"system\",\n",
+        "                \"content\": \"\"\"You are an AI assistant. Your goal is to answer questions using the pieces of context. If you don't know the answer, say that you don't know.\"\"\",\n",
+        "            },\n",
+        "            {\"role\": \"user\", \"content\": question},\n",
+        "            {\"role\": \"assistant\", \"content\": context},\n",
+        "        ],\n",
+        "        **model_kwargs,\n",
+        "    )\n",
+        "\n",
+        "    return response.choices[0].message.content\n",
+        "\n",
+        "\n",
+        "def generate_uuid(length: int = 8) -> str:\n",
+        "    \"\"\"Generate a uuid of a specified length (default=8).\"\"\"\n",
+        "    return \"\".join(random.choices(string.ascii_lowercase + string.digits, k=length))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "eha2l9nkNxZs"
+      },
+      "source": [
+        "### Generate evaluation dataset for AutoSxS\n",
+        "\n",
+        "Below you create your evaluation dataset, you specify a set of prompts to evaluate on.\n",
+        "\n",
+        "In this notebook, you:\n",
+        "\n",
+        "- Use 10 examples from the original dataset to create an evaluation dataset for AutoSxS.\n",
+        "  - Data in the `contexts` column will be treated as model context.\n",
+        "  - Data in the `questions` column will be treated as model instruction.\n",
+        "  - Data in the `response_a` column will be treated as responses for model A.\n",
+        "  - Data in the `response_b` will be treated as responses for model B.\n",
+        "\n",
+        "- Store the data in a JSON file in Google sCloud Storage.\n",
+        "\n",
+        "#### **Note: For the best results we recommend using at least 100 examples. There are diminishing returns when using more than 400 examples.**"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kuVd8Y7GHbp8"
+      },
+      "source": [
+        "#### Provide context and question"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "j_OuH0yh_PMe"
+      },
+      "outputs": [],
+      "source": [
+        "contexts = [\n",
+        "    \"Beginning in the late 1910s and early 1920s, Whitehead gradually turned his attention from mathematics to philosophy of science, and finally to metaphysics. He developed a comprehensive metaphysical system which radically departed from most of western philosophy. Whitehead argued that reality consists of processes rather than material objects, and that processes are best defined by their relations with other processes, thus rejecting the theory that reality is fundamentally constructed by bits of matter that exist independently of one another. Today Whitehead's philosophical works – particularly Process and Reality – are regarded as the foundational texts of process philosophy.\",\n",
+        "    \"The gills have an adnate attachment to the cap, are narrow to moderately broad, closely spaced, and eventually separate from the stem. Young gills are cinnamon-brown in color, with lighter edges, but darken in maturity because they become covered with the dark spores. The stem is 6 to 8 cm (2+3⁄8 to 3+1⁄8 in) long by 1.5 to 2 mm (1⁄16 to 3⁄32 in) thick, and roughly equal in width throughout except for a slightly enlarged base. The lower region of the stem is brownish in color and has silky 'hairs' pressed against the stem; the upper region is grayish and pruinose (lightly dusted with powdery white granules). The flesh turns slightly bluish or greenish where it has been injured. The application of a drop of dilute potassium hydroxide solution on the cap or flesh will cause a color change to pale to dark yellowish to reddish brown; a drop on the stem produces a less intense or no color change.\",\n",
+        "    \"Go to Device Support. Choose your device. Scroll to Getting started and select Hardware & phone details. Choose Insert or remove SIM card and follow the steps. Review the Account Summary page for details. Image 13 Activate online Go to att.com/activateprepaid ((att.com/activarprepaid for Spanish)) and follow the prompts. Activate over the phone Call us at 877.426.0525 for automated instructions. You will need to know your SIM/eSIM ICCID & IMEI number for activation. Note: Look for your SIM (( ICCID )) number on your box or SIM card Now youre ready to activate your phone 1. Start with your new device powered off. 2. To activate a new line of service or a replacement device, please go to the AT&T Activation site or call 866.895.1099. You download the eSIM to your device over Wi-Fi®. The eSIM connects your device to our wireless network. How do I activate my phone with an eSIM? Turn your phone on, connect to Wi-Fi, and follow the prompts. Swap active SIM cards AT&T Wireless SM SIM Card Turn your device off. Remove the old SIM card. Insert the new one. Turn on your device.\",\n",
+        "    \"According to chief astronaut Deke Slayton's autobiography, he chose Bassett for Gemini 9 because he was 'strong enough to carry' both himself and See. Slayton had also assigned Bassett as command module pilot for the second backup Apollo crew, alongside Frank Borman and William Anders.\",\n",
+        "    \"Adaptation of the endosymbiont to the host's lifestyle leads to many changes in the endosymbiont–the foremost being drastic reduction in its genome size. This is due to many genes being lost during the process of metabolism, and DNA repair and recombination. While important genes participating in the DNA to RNA transcription, protein translation and DNA/RNA replication are retained. That is, a decrease in genome size is due to loss of protein coding genes and not due to lessening of inter-genic regions or open reading frame (ORF) size. Thus, species that are naturally evolving and contain reduced sizes of genes can be accounted for an increased number of noticeable differences between them, thereby leading to changes in their evolutionary rates. As the endosymbiotic bacteria related with these insects are passed on to the offspring strictly via vertical genetic transmission, intracellular bacteria goes through many hurdles during the process, resulting in the decrease in effective population sizes when compared to the free living bacteria. This incapability of the endosymbiotic bacteria to reinstate its wild type phenotype via a recombination process is called as Muller's ratchet phenomenon. Muller's ratchet phenomenon together with less effective population sizes has led to an accretion of deleterious mutations in the non-essential genes of the intracellular bacteria. This could have been due to lack of selection mechanisms prevailing in the rich environment of the host.\",\n",
+        "    \"The National Archives Building in downtown Washington holds record collections such as all existing federal census records, ships' passenger lists, military unit records from the American Revolution to the Philippine–American War, records of the Confederate government, the Freedmen's Bureau records, and pension and land records.\",\n",
+        "    \"Standard 35mm photographic film used for cinema projection has a much higher image resolution than HDTV systems, and is exposed and projected at a rate of 24 frames per second (frame/s). To be shown on standard television, in PAL-system countries, cinema film is scanned at the TV rate of 25 frame/s, causing a speedup of 4.1 percent, which is generally considered acceptable. In NTSC-system countries, the TV scan rate of 30 frame/s would cause a perceptible speedup if the same were attempted, and the necessary correction is performed by a technique called 3:2 Pulldown: Over each successive pair of film frames, one is held for three video fields (1/20 of a second) and the next is held for two video fields (1/30 of a second), giving a total time for the two frames of 1/12 of a second and thus achieving the correct average film frame rate.\",\n",
+        "    \"Maria Deraismes was initiated into Freemasonry in 1882, then resigned to allow her lodge to rejoin their Grand Lodge. Having failed to achieve acceptance from any masonic governing body, she and Georges Martin started a mixed masonic lodge that actually worked masonic ritual. Annie Besant spread the phenomenon to the English speaking world. Disagreements over ritual led to the formation of exclusively female bodies of Freemasons in England, which spread to other countries. Meanwhile, the French had re-invented Adoption as an all-female lodge in 1901, only to cast it aside again in 1935. The lodges, however, continued to meet, which gave rise, in 1959, to a body of women practising continental Freemasonry.\",\n",
+        "    \"Excavation of the foundations began in November 1906, with an average of 275 workers during the day shift and 100 workers during the night shift. The excavation was required to be completed in 120 days. To remove the spoils from the foundation, three temporary wooden platforms were constructed to street level. Hoisting engines were installed to place the beams for the foundation, while the piers were sunk into the ground under their own weight. Because of the lack of space in the area, the contractors' offices were housed beneath the temporary platforms. During the process of excavation, the Gilsey Building's foundations were underpinned or shored up, because that building had relatively shallow foundations descending only 18 feet (5.5 m) below Broadway.\",\n",
+        "    \"Dopamine consumed in food cannot act on the brain, because it cannot cross the blood–brain barrier. However, there are also a variety of plants that contain L-DOPA, the metabolic precursor of dopamine. The highest concentrations are found in the leaves and bean pods of plants of the genus Mucuna, especially in Mucuna pruriens (velvet beans), which have been used as a source for L-DOPA as a drug. Another plant containing substantial amounts of L-DOPA is Vicia faba, the plant that produces fava beans (also known as 'broad beans'). The level of L-DOPA in the beans, however, is much lower than in the pod shells and other parts of the plant. The seeds of Cassia and Bauhinia trees also contain substantial amounts of L-DOPA.\",\n",
+        "]\n",
+        "\n",
+        "questions = [\n",
+        "    \"What was the predominant theory of reality that Whitehead opposed?\",\n",
+        "    \"Why do the gills on the Psilocybe pelliculosa mushroom darken as they mature?\",\n",
+        "    \"user: How do I provision my AT&T SIM card?\",\n",
+        "    \"Why did chief astronaut Deke Slayton choose Charles Bassett for Gemini 9, according to Slayton's autobiography?\",\n",
+        "    \"What is the main alteration in an endosymbiont when it adapts to a host?\",\n",
+        "    \"What's the earliest war The National Archives Building has military unit records for\",\n",
+        "    \"To be shown on SDTV in PAL-system countries, at what rate is cinema film scanned?\",\n",
+        "    \"What year was the all-female masonic lodge cast aside?\",\n",
+        "    \"Why did the Gilsey Building have underpinned and shored up foundations?\",\n",
+        "    \"Why can dopamine consumed in food not act on the brain?\",\n",
+        "]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Oiwr677h_cSk"
+      },
+      "outputs": [],
+      "source": [
+        "examples = pd.DataFrame(\n",
+        "    {\n",
+        "        \"questions\": questions,\n",
+        "        \"context\": contexts,\n",
+        "    }\n",
+        ")\n",
+        "examples.head()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "2je-Rs8e_65p"
+      },
+      "outputs": [],
+      "source": [
+        "examples[\"response_a\"] = examples.apply(\n",
+        "    lambda x: self_model_generate(\n",
+        "        x[\"questions\"], x[\"context\"], max_tokens=2500, temperature=0.5\n",
+        "    ),\n",
+        "    axis=1,\n",
+        ")\n",
+        "examples.head()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "VGIpgBk9Br_G"
+      },
+      "outputs": [],
+      "source": [
+        "examples[\"response_b\"] = examples.apply(\n",
+        "    lambda x: maas_generate(\n",
+        "        x[\"questions\"], x[\"context\"], max_tokens=2500, temperature=0.5\n",
+        "    ),\n",
+        "    axis=1,\n",
+        ")\n",
+        "examples.head()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "NY1Jsj4aOCe1"
+      },
+      "source": [
+        "#### Upload your dataset to Cloud Storage\n",
+        "\n",
+        "Finally, we upload our evaluation dataset to Cloud Storage to be used as input for AutoSxS."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "vykmkhp-ODKg"
+      },
+      "outputs": [],
+      "source": [
+        "examples.to_json(f\"{BUCKET_URI}/evaluation_dataset.json\", orient=\"records\", lines=True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Qgdk_qNIOFik"
+      },
+      "source": [
+        "### Create and run AutoSxS job\n",
+        "\n",
+        "In order to run AutoSxS, we need to define a `autosxs_pipeline` job with the following parameters.\n",
+        "\n",
+        "More details of the AutoSxS pipeline configuration can be found [here](https://google-cloud-pipeline-components.readthedocs.io/en/google-cloud-pipeline-components-2.9.0/api/preview/model_evaluation.html#preview.model_evaluation.autosxs_pipeline)."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "veq26QZ7OMoC"
+      },
+      "source": [
+        "First, compile the AutoSxS pipeline locally."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "C2NGZzOMOJPV"
+      },
+      "outputs": [],
+      "source": [
+        "template_uri = \"pipeline.yaml\"\n",
+        "compiler.Compiler().compile(\n",
+        "    pipeline_func=model_evaluation.autosxs_pipeline,\n",
+        "    package_path=template_uri,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "I0aMBhoqOTXF"
+      },
+      "source": [
+        "The following code starts a Vertex Pipeline job, viewable from the Vertex UI. This pipeline job will take ~15 mins. This pipeline is made for batch prediction at a much larger scale than this example, so the time won't scale up linearly if there were thousands of Q&A pairs."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tRdA3ovUOV6j"
+      },
+      "outputs": [],
+      "source": [
+        "display_name = f\"autosxs-qa-{generate_uuid()}\"\n",
+        "context_column = \"context\"\n",
+        "question_column = \"questions\"\n",
+        "response_column_a = \"response_a\"\n",
+        "response_column_b = \"response_b\"\n",
+        "\n",
+        "parameters = {\n",
+        "    \"evaluation_dataset\": BUCKET_URI + \"/evaluation_dataset.json\",\n",
+        "    \"id_columns\": [question_column],\n",
+        "    \"autorater_prompt_parameters\": {\n",
+        "        \"inference_context\": {\"column\": context_column},\n",
+        "        \"inference_instruction\": {\"column\": question_column},\n",
+        "    },\n",
+        "    \"task\": \"question_answering\",\n",
+        "    \"response_column_a\": response_column_a,\n",
+        "    \"response_column_b\": response_column_b,\n",
+        "}\n",
+        "\n",
+        "job = aiplatform.PipelineJob(\n",
+        "    job_id=display_name,\n",
+        "    display_name=display_name,\n",
+        "    pipeline_root=BUCKET_URI + \"/pipeline\",\n",
+        "    template_path=template_uri,\n",
+        "    parameter_values=parameters,\n",
+        "    enable_caching=False,\n",
+        "    project=PROJECT_ID,\n",
+        "    location=LOCATION,\n",
+        ")\n",
+        "job.run()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "EinPbr3XOYPQ"
+      },
+      "source": [
+        "### Get the judgments and AutoSxS metrics\n",
+        "Next, you can review judgments from the completed AutoSxS job."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "V_9yMfhrOZDk"
+      },
+      "outputs": [],
+      "source": [
+        "for details in job.task_details:\n",
+        "    if details.task_name == \"online-evaluation-pairwise\":\n",
+        "        break\n",
+        "\n",
+        "judgments_uri = details.outputs[\"judgments\"].artifacts[0].uri\n",
+        "judgments_df = pd.read_json(judgments_uri, lines=True)\n",
+        "judgments_df.head()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "BlKXu5Ze4tD3"
+      },
+      "source": [
+        "You can also review AutoSxS metrics computed from the judgments.\n",
+        "\n",
+        "You can find more details of AutoSxS metrics [here](https://cloud.google.com/vertex-ai/generative-ai/docs/models/side-by-side-eval#aggregate-metrics)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "G7meI2Eq4muT"
+      },
+      "outputs": [],
+      "source": [
+        "for details in job.task_details:\n",
+        "    if details.task_name == \"model-evaluation-text-generation-pairwise\":\n",
+        "        break\n",
+        "pd.DataFrame([details.outputs[\"autosxs_metrics\"].artifacts[0].metadata])"
+      ]
+    },
+    {
+      "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 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:\n",
+        "\n",
+        "Set `delete_bucket` to **True** to delete the Cloud Storage bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "sx_vKniMq9ZX"
+      },
+      "outputs": [],
+      "source": [
+        "delete_pipeline_job = False  # @param {type:\"boolean\"}\n",
+        "delete_bucket = False  # @param {type:\"boolean\"}\n",
+        "\n",
+        "if delete_pipeline_job:\n",
+        "    job.delete()\n",
+        "\n",
+        "if delete_bucket:\n",
+        "    ! gsutil rm -r gs://{BUCKET_NAME}"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_autosxs_evaluation_llama3_1.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
@@ -53,10 +53,13 @@
      "source": [
        "## Overview\n",
        "\n",
-        "This notebook demonstrates deploying CodeGemma models on GPU using [vLLM](https://github.com/vllm-project/vllm), the state-of-the-art open source LLM serving solution on GPU.\n",
+        "This notebook demonstrates deploying CodeGemma models\n",
+        " * on TPU using **Hex-LLM**, a **H**igh-**E**fficiency **L**arge **L**anguage **M**odel serving solution built with **XLA** that is being developed by Google Cloud, and\n",
+        " * on GPU using [vLLM](https://github.com/vllm-project/vllm), the state-of-the-art open source LLM serving solution on GPU.\n",
        "\n",
        "### Objective\n",
        "\n",
+        "- Deploy CodeGemma with Hex-LLM on TPU\n",
        "- Deploy CodeGemma with [vLLM](https://github.com/vllm-project/vllm) on GPU\n",
        "\n",
        "### Costs\n",
@@ -78,6 +81,20 @@
        "## Run the notebook"
      ]
    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "494697c28ee7"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Request for TPU quota\n",
+        "\n",
+        "# @markdown By default, the quota for TPU deployment `Custom model serving TPU v5e cores per region` is 4. TPU quota is only available in `us-west1`. You can request for higher TPU quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota)."
+      ]
+    },
    {
      "cell_type": "code",
      "execution_count": null,
@@ -155,6 +172,7 @@
        "# @markdown Alternatively, you can also load the original CodeGemma models for serving from Vertex AI after accepting the agreement.\n",
        "# @markdown **Please only select and fill one of the two following sections.**\n",
        "LOAD_MODEL_FROM = \"Google Cloud\"  # @param [\"Hugging Face\", \"Google Cloud\"] {isTemplate:true}\n",
+        "\n",
        "# @markdown #### Access CodeGemma models on Vertex AI\n",
        "# @markdown Accept the model agreement to access the models:\n",
        "# @markdown 1. Open the [CodeGemma model card](https://console.cloud.google.com/vertex-ai/publishers/google/model-garden/364) from [Vertex AI Model Garden](https://cloud.google.com/model-garden).\n",
@@ -170,7 +188,7 @@
        "\n",
        "# @markdown *--- Or ---*\n",
        "\n",
-        "# @markdown ##Enable the Vertex AI API## Access CodeGemma models on HuggingFace\n",
+        "# @markdown #### Access CodeGemma models on HuggingFace\n",
        "# @markdown You must provide a Hugging Face User Access Token (read) to access the CodeGemma models. You can follow the [Hugging Face documentation](https://huggingface.co/docs/hub/en/security-tokens) to create a **read** access token and put it in the `HF_TOKEN` field below.\n",
        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate:true}\n",
        "if LOAD_MODEL_FROM == \"Hugging Face\":\n",
@@ -193,7 +211,9 @@
        "    model_path_prefix = BUCKET_URI.strip(\"/\") + \"/codegemma\"\n",
        "else:\n",
        "    model_path_prefix = \"google/\"\n",
-        "# The pre-built serving docker images with vLLM\n",
+        "\n",
+        "# The pre-built serving docker images with Hex-LLM and vLLM\n",
+        "HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:deploy\"\n",
        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240220_0936_RC01\"\n",
        "\n",
        "\n",
@@ -202,6 +222,72 @@
        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
        "\n",
        "\n",
+        "def deploy_model_hexllm(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    service_account: str,\n",
+        "    machine_type: str = \"ct5lp-hightpu-1t\",\n",
+        "    tensor_parallel_size: int = 1,\n",
+        "    hbm_utilization_factor: float = 0.6,\n",
+        "    max_running_seqs: int = 256,\n",
+        "    endpoint_id: str = \"\",\n",
+        "    min_replica_count: int = 1,\n",
+        "    max_replica_count: int = 1,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys models with Hex-LLM on TPU in Vertex AI.\"\"\"\n",
+        "    if endpoint_id:\n",
+        "        aip_endpoint_name = (\n",
+        "            f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_id}\"\n",
+        "        )\n",
+        "        endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "    else:\n",
+        "        endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
+        "\n",
+        "    hexllm_args = [\n",
+        "        \"--host=0.0.0.0\",\n",
+        "        \"--port=7080\",\n",
+        "        \"--log_level=INFO\",\n",
+        "        \"--enable_jit\",\n",
+        "        f\"--model={model_id}\",\n",
+        "        \"--load_format=auto\",\n",
+        "        f\"--tensor_parallel_size={tensor_parallel_size}\",\n",
+        "        f\"--hbm_utilization_factor={hbm_utilization_factor}\",\n",
+        "        f\"--max_running_seqs={max_running_seqs}\",\n",
+        "    ]\n",
+        "    hexllm_envs = {\n",
+        "        \"PJRT_DEVICE\": \"TPU\",\n",
+        "        \"RAY_DEDUP_LOGS\": \"0\",\n",
+        "        \"RAY_USAGE_STATS_ENABLED\": \"0\",\n",
+        "        \"MODEL_ID\": model_id,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "    if HF_TOKEN:\n",
+        "        hexllm_envs.update({\"HF_TOKEN\": HF_TOKEN})\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=HEXLLM_DOCKER_URI,\n",
+        "        serving_container_command=[\"python\", \"-m\", \"hex_llm.server.api_server\"],\n",
+        "        serving_container_args=hexllm_args,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_predict_route=\"/generate\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        serving_container_environment_variables=hexllm_envs,\n",
+        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
+        "        serving_container_deployment_timeout=7200,\n",
+        "    )\n",
+        "\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=service_account,\n",
+        "        min_replica_count=min_replica_count,\n",
+        "        max_replica_count=max_replica_count,\n",
+        "    )\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
        "def deploy_model_vllm(\n",
        "    model_name: str,\n",
        "    model_id: str,\n",
@@ -259,6 +345,127 @@
        "    return model, endpoint"
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8neJc8CnDDpu"
+      },
+      "source": [
+        "## Deploy CodeGemma models with Hex-LLM on TPU\n",
+        "\n",
+        "**Hex-LLM** is a **H**igh-**E**fficiency **L**arge **L**anguage **M**odel (LLM) TPU serving solution built with **XLA**, which is being developed by Google Cloud.\n",
+        "\n",
+        "Refer to the \"Request for TPU quota\" section for TPU quota."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "E8OiHHNNE_wj"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "\n",
+        "# @markdown Set the model to deploy.\n",
+        "\n",
+        "MODEL_ID = \"codegemma-7b-it\"  # @param [\"codegemma-2b\", \"codegemma-7b\", \"codegemma-7b-it\"]\n",
+        "model_id = os.path.join(model_path_prefix, MODEL_ID)\n",
+        "\n",
+        "# @markdown Find Vertex AI prediction TPUv5e machine types in\n",
+        "# @markdown https://cloud.google.com/vertex-ai/docs/predictions/use-tpu#deploy_a_model.\n",
+        "if \"2b\" in model_id:\n",
+        "    # Sets ct5lp-hightpu-1t (1 TPU chip) to deploy CodeGemma 2B models.\n",
+        "    machine_type = \"ct5lp-hightpu-1t\"\n",
+        "    accelerator_type = \"TPU_V5e\"\n",
+        "    # Note: 1 TPU V5 chip has only one core.\n",
+        "    accelerator_count = 1\n",
+        "else:\n",
+        "    # Sets ct5lp-hightpu-4t (4 TPU chips) to deploy CodeGemma 7B models.\n",
+        "    machine_type = \"ct5lp-hightpu-4t\"\n",
+        "    accelerator_type = \"TPU_V5e\"\n",
+        "    # Note: 1 TPU V5 chip has only one core.\n",
+        "    accelerator_count = 4\n",
+        "\n",
+        "# Server parameters.\n",
+        "tensor_parallel_size = accelerator_count\n",
+        "hbm_utilization_factor = 0.6  # Fraction of HBM memory allocated for KV cache after model loading. A larger value improves throughput but gives higher risk of TPU out-of-memory errors with long prompts.\n",
+        "max_running_seqs = 256  # Maximum number of running sequences in a continuous batch.\n",
+        "\n",
+        "# Endpoint configurations.\n",
+        "min_replica_count = 1\n",
+        "max_replica_count = 1\n",
+        "\n",
+        "model_hexllm, endpoint_hexllm = deploy_model_hexllm(\n",
+        "    model_name=get_job_name_with_datetime(prefix=MODEL_ID),\n",
+        "    model_id=model_id,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    machine_type=machine_type,\n",
+        "    tensor_parallel_size=tensor_parallel_size,\n",
+        "    hbm_utilization_factor=hbm_utilization_factor,\n",
+        "    max_running_seqs=max_running_seqs,\n",
+        "    min_replica_count=min_replica_count,\n",
+        "    max_replica_count=max_replica_count,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "nkUaMxIus6Pv"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. The first few requests may have high latency. This is because the server needs to warm up with the initial requests. The following requests should not have the same delay.\n",
+        "\n",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown > What is a car?\n",
+        "# @markdown > A car is a four-wheeled vehicle designed for the transportation of passengers and their belongings.\n",
+        "# @markdown ```\n",
+        "\n",
+        "# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for more details.\n",
+        "\n",
+        "# Loads an existing endpoint instance using the endpoint name:\n",
+        "# - Using `endpoint_name = endpoint_hexllm.name` allows us to get the endpoint\n",
+        "#   name of the endpoint `endpoint_hexllm` created in the cell above.\n",
+        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
+        "#   an existing endpoint with the ID 1234567890123456789.\n",
+        "# You may uncomment the code below to load an existing endpoint:\n",
+        "# endpoint_name = endpoint_hexllm.name\n",
+        "# # endpoint_name = \"\"  # @param {type:\"string\"}\n",
+        "# aip_endpoint_name = (\n",
+        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
+        "# )\n",
+        "# endpoint_hexllm = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "\n",
+        "prompt = \"What is a car?\"  # @param {type: \"string\"}\n",
+        "max_tokens = 50  # @param {type: \"integer\"}\n",
+        "temperature = 1.0  # @param {type: \"number\"}\n",
+        "top_p = 1.0  # @param {type: \"number\"}\n",
+        "top_k = 1  # @param {type: \"integer\"}\n",
+        "instances = [\n",
+        "    {\n",
+        "        \"prompt\": prompt,\n",
+        "        \"max_tokens\": max_tokens,\n",
+        "        \"temperature\": temperature,\n",
+        "        \"top_p\": top_p,\n",
+        "        \"top_k\": top_k,\n",
+        "    },\n",
+        "]\n",
+        "response = endpoint_hexllm.predict(instances=instances)\n",
+        "\n",
+        "prediction = response.predictions[0]\n",
+        "print(prediction)"
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -279,7 +486,7 @@
      "execution_count": null,
      "metadata": {
        "cellView": "form",
-        "id": "E8OiHHNNE_wj"
+        "id": "tQIEisUajS6t"
      },
      "outputs": [],
      "source": [
@@ -313,7 +520,7 @@
        "max_model_len = 2048\n",
        "gpu_memory_utilization = 0.9\n",
        "\n",
-        "model, endpoint = deploy_model_vllm(\n",
+        "model_vllm, endpoint_vllm = deploy_model_vllm(\n",
        "    model_name=get_job_name_with_datetime(prefix=\"codegemma-serve-vllm\"),\n",
        "    model_id=model_id,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
@@ -348,20 +555,22 @@
        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. Sampling parameters supported by vLLM can be found [here](https://github.com/vllm-project/vllm/blob/2e8e49fce3775e7704d413b2f02da6d7c99525c9/vllm/sampling_params.py#L23-L64).\n",
        "\n",
        "# Loads an existing endpoint instance using the endpoint name:\n",
-        "# - Using `endpoint_name = endpoint.name` allows us to get the endpoint name of\n",
-        "#   the endpoint `endpoint` created in the cell above.\n",
+        "# - Using `endpoint_name = endpoint_vllm.name` allows us to get the endpoint name of\n",
+        "#   the endpoint `endpoint_vllm` created in the cell above.\n",
        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
        "#   an existing endpoint with the ID 1234567890123456789.\n",
        "# You may uncomment the code below to load an existing endpoint.\n",
        "\n",
-        "# endpoint_name = endpoint.name\n",
+        "# endpoint_name = endpoint_vllm.name\n",
        "# # endpoint_name = \"\"  # @param {type:\"string\"}\n",
        "# aip_endpoint_name = (\n",
        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
        "# )\n",
-        "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "# endpoint_vllm = aiplatform.Endpoint(aip_endpoint_name)\n",
        "\n",
-        "prompt = \"Write a function to list n Fibonacci numbers in Python.\"  # @param {type: \"string\"}\n",
+        "prompt = (\n",
+        "    \"Write a function to list n Fibonacci numbers in Python.\"  # @param {type: \"string\"}\n",
+        ")\n",
        "max_tokens = 500  # @param {type:\"integer\"}\n",
        "temperature = 1.0  # @param {type:\"number\"}\n",
        "top_p = 1.0  # @param {type:\"number\"}\n",
@@ -374,10 +583,10 @@
        "        \"temperature\": temperature,\n",
        "        \"top_p\": top_p,\n",
        "        \"top_k\": top_k,\n",
-        "        \"raw_response\": False,\n",
+        "        \"raw_response\": True,\n",
        "    },\n",
        "]\n",
-        "response = endpoint.predict(instances=instances)\n",
+        "response = endpoint_vllm.predict(instances=instances)\n",
        "\n",
        "# \"<|file_separator|>\" is the end of the file token.\n",
        "for prediction in response.predictions:\n",
@@ -397,10 +606,13 @@
        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
        "# @markdown  and avoid unnecessary continouous charges that may incur.\n",
        "\n",
-        "endpoint.delete(force=True)\n",
+        "# Undeploy models and delete endpoints.\n",
+        "endpoint_hexllm.delete(force=True)\n",
+        "endpoint_vllm.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
-        "model.delete()\n",
+        "model_hexllm.delete()\n",
+        "model_vllm.delete()\n",
        "\n",
        "# Delete Cloud Storage objects.\n",
        "delete_bucket = False  # @param {type:\"boolean\"}\n",
@@ -64,7 +64,7 @@
        "    - [intfloat/multilingual-e5-small](https://huggingface.co/intfloat/multilingual-e5-small): 118M params\n",
        "    - [intfloat/e5-base-v2](https://huggingface.co/intfloat/e5-base-v2): 109M params\n",
        "    - [intfloat/e5-small-v2](https://huggingface.co/intfloat/e5-small-v2): 33M params\n",
-        "- Run inference on the deployed Vertex AI Endpoint.\n",
+        "- Run inference on the deployed Vertex AI Endpoint\n",
        "\n",
        "\n",
        "### Costs\n",
@@ -0,0 +1,725 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7d9bbf86da5e"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2024 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "99c1c3fc2ca5"
+      },
+      "source": [
+        "# Vertex AI Model Garden - Gemma 2 (Deployment)\n",
+        "\n",
+        "<table><tbody><tr>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_gemma2_deployment_on_vertex.ipynb\">\n",
+        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_gemma2_deployment_on_vertex.ipynb\">\n",
+        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3de7470326a2"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates deploying Gemma 2 models\n",
+        " * on TPU using **Hex-LLM**, a **H**igh-**E**fficiency **L**arge **L**anguage **M**odel serving solution built with **XLA** that is being developed by Google Cloud, and\n",
+        " * on GPU using **TGI** ([text-generation-inference](https://github.com/huggingface/text-generation-inference)), the state-of-the-art open source LLM serving solution on GPU.\n",
+        "\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Deploy Gemma 2 with Hex-LLM on TPU\n",
+        "- Deploy Gemma with [TGI](https://github.com/huggingface/text-generation-inference) on GPU\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "264c07757582"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "C_wC61dhpWXj"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Request for TPU quota\n",
+        "\n",
+        "# @markdown By default, the quota for TPU deployment `Custom model serving TPU v5e cores per region` is 4. TPU quota is only available in `us-west1`. You can request for higher TPU quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "2707b02ef5df"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "# @markdown **[Optional]** Set the GCS BUCKET_URI to store the experiment artifacts, if you want to use your own bucket. **If not set, a unique GCS bucket will be created automatically on your behalf**.\n",
+        "\n",
+        "import json\n",
+        "import os\n",
+        "from datetime import datetime\n",
+        "from typing import Tuple\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "\n",
+        "# Cloud Storage bucket for storing the experiment artifacts.\n",
+        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
+        "# prefer using your own GCS bucket, please change the value yourself below.\n",
+        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
+        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
+        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
+        "    # Create a unique GCS bucket for this notebook if not specified\n",
+        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\n",
+        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
+        "else:\n",
+        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
+        "    bucket_region = shell_output[0].strip().lower()\n",
+        "    if bucket_region != REGION:\n",
+        "        raise ValueError(\n",
+        "            \"Bucket region %s is different from notebook region %s\"\n",
+        "            % (bucket_region, REGION)\n",
+        "        )\n",
+        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"gemma2\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
+        "\n",
+        "# Gets the default SERVICE_ACCOUNT.\n",
+        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
+        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
+        "\n",
+        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
+        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "\n",
+        "# Enable Vertex AI and Cloud Compute APIs.\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
+        "\n",
+        "# @markdown ## Access Gemma 2 Models\n",
+        "\n",
+        "# @markdown You must provide a Hugging Face User Access Token (read) to access the Gemma 2 models. You can follow the [Hugging Face documentation](https://huggingface.co/docs/hub/en/security-tokens) to create a **read** access token and put it in the `HF_TOKEN` field below.\n",
+        "\n",
+        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate:true}\n",
+        "assert (\n",
+        "    HF_TOKEN\n",
+        "), \"Please provide a read HF_TOKEN to load models from Hugging Face, or select a different model source.\"\n",
+        "\n",
+        "model_path_prefix = \"google/\"\n",
+        "\n",
+        "# The pre-built serving docker images.\n",
+        "HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:gemma2\"\n",
+        "TGI_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu121.2-1.ubuntu2204.py310\"\n",
+        "\n",
+        "SERVICE_ENDPOINT = \"aiplatform.googleapis.com\"\n",
+        "\n",
+        "\n",
+        "def get_job_name_with_datetime(prefix: str) -> str:\n",
+        "    \"\"\"Gets the job name with date time when triggering deployment jobs.\"\"\"\n",
+        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
+        "\n",
+        "\n",
+        "def deploy_model_hexllm(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    service_account: str,\n",
+        "    machine_type: str = \"ct5lp-hightpu-1t\",\n",
+        "    tensor_parallel_size: int = 1,\n",
+        "    hbm_utilization_factor: float = 0.6,\n",
+        "    max_running_seqs: int = 256,\n",
+        "    endpoint_id: str = \"\",\n",
+        "    min_replica_count: int = 1,\n",
+        "    max_replica_count: int = 1,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys models with Hex-LLM on TPU in Vertex AI.\"\"\"\n",
+        "    if endpoint_id:\n",
+        "        aip_endpoint_name = (\n",
+        "            f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_id}\"\n",
+        "        )\n",
+        "        endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "    else:\n",
+        "        endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
+        "\n",
+        "    hexllm_args = [\n",
+        "        \"--host=0.0.0.0\",\n",
+        "        \"--port=7080\",\n",
+        "        \"--log_level=INFO\",\n",
+        "        \"--enable_jit\",\n",
+        "        f\"--model={model_id}\",\n",
+        "        \"--load_format=auto\",\n",
+        "        f\"--tensor_parallel_size={tensor_parallel_size}\",\n",
+        "        f\"--hbm_utilization_factor={hbm_utilization_factor}\",\n",
+        "        f\"--max_running_seqs={max_running_seqs}\",\n",
+        "    ]\n",
+        "    hexllm_envs = {\n",
+        "        \"PJRT_DEVICE\": \"TPU\",\n",
+        "        \"RAY_DEDUP_LOGS\": \"0\",\n",
+        "        \"RAY_USAGE_STATS_ENABLED\": \"0\",\n",
+        "        \"MODEL_ID\": model_id,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "    if HF_TOKEN:\n",
+        "        hexllm_envs.update({\"HF_TOKEN\": HF_TOKEN})\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=HEXLLM_DOCKER_URI,\n",
+        "        serving_container_command=[\"python\", \"-m\", \"hex_llm.server.api_server\"],\n",
+        "        serving_container_args=hexllm_args,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_predict_route=\"/generate\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        serving_container_environment_variables=hexllm_envs,\n",
+        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
+        "        serving_container_deployment_timeout=7200,\n",
+        "    )\n",
+        "\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=service_account,\n",
+        "        min_replica_count=min_replica_count,\n",
+        "        max_replica_count=max_replica_count,\n",
+        "    )\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "def deploy_model_tgi(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    service_account: str,\n",
+        "    machine_type: str = \"g2-standard-24\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 2,\n",
+        "    max_input_length: int = 1562,\n",
+        "    max_total_tokens: int = 2048,\n",
+        "    max_batch_prefill_tokens: int = 2048,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys models with TGI on GPU in Vertex AI.\"\"\"\n",
+        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
+        "\n",
+        "    env_vars = {\n",
+        "        \"AIP_HTTP_PORT\": 7080,\n",
+        "        \"MODEL_ID\": model_id,\n",
+        "        \"NUM_SHARD\": f\"{accelerator_count}\",\n",
+        "        \"MAX_INPUT_LENGTH\": f\"{max_input_length}\",\n",
+        "        \"MAX_TOTAL_TOKENS\": f\"{max_total_tokens}\",\n",
+        "        \"MAX_BATCH_PREFILL_TOKENS\": f\"{max_batch_prefill_tokens}\",\n",
+        "        \"CUDA_MEMORY_FRACTION\": 0.93,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "\n",
+        "    if HF_TOKEN:\n",
+        "        env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=TGI_DOCKER_URI,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_environment_variables=env_vars,\n",
+        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
+        "    )\n",
+        "\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=service_account,\n",
+        "    )\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "def get_quota(project_id: str, region: str, resource_id: str) -> int:\n",
+        "    \"\"\"Returns the quota for a resource in a region. Returns -1 if can not figure out the quota.\"\"\"\n",
+        "    quota_list_output = !gcloud alpha services quota list --service=$SERVICE_ENDPOINT  --consumer=projects/$project_id --filter=\"$SERVICE_ENDPOINT/$resource_id\" --format=json\n",
+        "    # Use '.s' on the command output because it is an SList type.\n",
+        "    quota_data = json.loads(quota_list_output.s)\n",
+        "    if len(quota_data) == 0 or \"consumerQuotaLimits\" not in quota_data[0]:\n",
+        "        return -1\n",
+        "    if (\n",
+        "        len(quota_data[0][\"consumerQuotaLimits\"]) == 0\n",
+        "        or \"quotaBuckets\" not in quota_data[0][\"consumerQuotaLimits\"][0]\n",
+        "    ):\n",
+        "        return -1\n",
+        "    all_regions_data = quota_data[0][\"consumerQuotaLimits\"][0][\"quotaBuckets\"]\n",
+        "    for region_data in all_regions_data:\n",
+        "        if (\n",
+        "            region_data.get(\"dimensions\")\n",
+        "            and region_data[\"dimensions\"][\"region\"] == region\n",
+        "        ):\n",
+        "            if \"effectiveLimit\" in region_data:\n",
+        "                return int(region_data[\"effectiveLimit\"])\n",
+        "            else:\n",
+        "                return 0\n",
+        "    return -1\n",
+        "\n",
+        "\n",
+        "def get_resource_id(accelerator_type: str, is_for_training: bool) -> str:\n",
+        "    \"\"\"Returns the resource id for a given accelerator type and the use case.\n",
+        "    Args:\n",
+        "      accelerator_type: The accelerator type.\n",
+        "      is_for_training: Whether the resource is used for training. Set false\n",
+        "      for serving use case.\n",
+        "    Returns:\n",
+        "      The resource id.\n",
+        "    \"\"\"\n",
+        "    training_accelerator_map = {\n",
+        "        \"NVIDIA_TESLA_V100\": \"custom_model_training_nvidia_v100_gpus\",\n",
+        "        \"NVIDIA_L4\": \"custom_model_training_nvidia_l4_gpus\",\n",
+        "        \"NVIDIA_TESLA_A100\": \"custom_model_training_nvidia_a100_gpus\",\n",
+        "        \"NVIDIA_TESLA_T4\": \"custom_model_training_nvidia_t4_gpus\",\n",
+        "        \"TPU_V5e\": \"custom_model_training_tpu_v5e\",\n",
+        "        \"TPU_V3\": \"custom_model_training_tpu_v3\",\n",
+        "    }\n",
+        "    serving_accelerator_map = {\n",
+        "        \"NVIDIA_TESLA_V100\": \"custom_model_serving_nvidia_v100_gpus\",\n",
+        "        \"NVIDIA_L4\": \"custom_model_serving_nvidia_l4_gpus\",\n",
+        "        \"NVIDIA_TESLA_A100\": \"custom_model_serving_nvidia_a100_gpus\",\n",
+        "        \"NVIDIA_TESLA_T4\": \"custom_model_serving_nvidia_t4_gpus\",\n",
+        "        \"TPU_V5e\": \"custom_model_serving_tpu_v5e\",\n",
+        "    }\n",
+        "    if is_for_training:\n",
+        "        if accelerator_type in training_accelerator_map:\n",
+        "            return training_accelerator_map[accelerator_type]\n",
+        "        else:\n",
+        "            raise ValueError(\n",
+        "                f\"Could not find accelerator type: {accelerator_type} for training.\"\n",
+        "            )\n",
+        "    else:\n",
+        "        if accelerator_type in serving_accelerator_map:\n",
+        "            return serving_accelerator_map[accelerator_type]\n",
+        "        else:\n",
+        "            raise ValueError(\n",
+        "                f\"Could not find accelerator type: {accelerator_type} for serving.\"\n",
+        "            )\n",
+        "\n",
+        "\n",
+        "def check_quota(\n",
+        "    project_id: str,\n",
+        "    region: str,\n",
+        "    accelerator_type: str,\n",
+        "    accelerator_count: int,\n",
+        "    is_for_training: bool,\n",
+        "):\n",
+        "    \"\"\"Checks if the project and the region has the required quota.\"\"\"\n",
+        "    resource_id = get_resource_id(accelerator_type, is_for_training)\n",
+        "    quota = get_quota(project_id, region, resource_id)\n",
+        "    quota_request_instruction = (\n",
+        "        \"Either use \"\n",
+        "        \"a different region or request additional quota. Follow \"\n",
+        "        \"instructions here \"\n",
+        "        \"https://cloud.google.com/docs/quotas/view-manage#requesting_higher_quota\"\n",
+        "        \" to check quota in a region or request additional quota for \"\n",
+        "        \"your project.\"\n",
+        "    )\n",
+        "    if quota == -1:\n",
+        "        raise ValueError(\n",
+        "            f\"\"\"Quota not found for: {resource_id} in {region}.\n",
+        "            {quota_request_instruction}\"\"\"\n",
+        "        )\n",
+        "    if quota < accelerator_count:\n",
+        "        raise ValueError(\n",
+        "            f\"\"\"Quota not enough for {resource_id} in {region}:\n",
+        "            {quota} < {accelerator_count}.\n",
+        "            {quota_request_instruction}\"\"\"\n",
+        "        )"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8neJc8CnDDpu"
+      },
+      "source": [
+        "## Deploy Gemma 2 models with Hex-LLM on TPU\n",
+        "\n",
+        "**Hex-LLM** is a **H**igh-**E**fficiency **L**arge **L**anguage **M**odel (LLM) TPU serving solution built with **XLA**, which is being developed by Google Cloud.\n",
+        "\n",
+        "Refer to the \"Request for TPU quota\" section for TPU quota."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "E8OiHHNNE_wj"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "# @markdown Set the model ID. Model weights can be loaded from HuggingFace or from a GCS bucket.\n",
+        "\n",
+        "# @markdown Select one of the four model variations.\n",
+        "MODEL_ID = \"gemma-2-9b\"  # @param [\"gemma-2-9b\", \"gemma-2-9b-it\", \"gemma-2-27b\", \"gemma-2-27b-it\"] {allow-input: true, isTemplate: true}\n",
+        "model_id = os.path.join(model_path_prefix, MODEL_ID)\n",
+        "\n",
+        "# @markdown Find Vertex AI prediction TPUv5e machine types in\n",
+        "# @markdown https://cloud.google.com/vertex-ai/docs/predictions/use-tpu#deploy_a_model.\n",
+        "if \"9b\" in model_id:\n",
+        "    # Sets ct5lp-hightpu-4t (4 TPU chips) to deploy Gemma 2 9B models.\n",
+        "    machine_type = \"ct5lp-hightpu-4t\"\n",
+        "    accelerator_type = \"TPU_V5e\"\n",
+        "    # Note: 1 TPU V5 chip has only one core.\n",
+        "    accelerator_count = 4\n",
+        "else:\n",
+        "    # Sets ct5lp-hightpu-8t (8 TPU chips) to deploy Gemma 2 27B models.\n",
+        "    machine_type = \"ct5lp-hightpu-8t\"\n",
+        "    accelerator_type = \"TPU_V5e\"\n",
+        "    # Note: 1 TPU V5 chip has only one core.\n",
+        "    accelerator_count = 8\n",
+        "\n",
+        "check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=False,\n",
+        ")\n",
+        "\n",
+        "# Server parameters.\n",
+        "tensor_parallel_size = accelerator_count\n",
+        "hbm_utilization_factor = 0.6  # Fraction of HBM memory allocated for KV cache after model loading. A larger value improves throughput but gives higher risk of TPU out-of-memory errors with long prompts.\n",
+        "max_running_seqs = 256  # Maximum number of running sequences in a continuous batch.\n",
+        "\n",
+        "# Endpoint configurations.\n",
+        "min_replica_count = 1\n",
+        "max_replica_count = 1\n",
+        "\n",
+        "model_hexllm, endpoint_hexllm = deploy_model_hexllm(\n",
+        "    model_name=get_job_name_with_datetime(prefix=MODEL_ID),\n",
+        "    model_id=model_id,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    machine_type=machine_type,\n",
+        "    tensor_parallel_size=tensor_parallel_size,\n",
+        "    hbm_utilization_factor=hbm_utilization_factor,\n",
+        "    max_running_seqs=max_running_seqs,\n",
+        "    min_replica_count=min_replica_count,\n",
+        "    max_replica_count=max_replica_count,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "nkUaMxIus6Pv"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. The first few requests may have high latency. This is because the server needs to warm up with the initial requests. The following requests should not have the same delay.\n",
+        "\n",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown > What is a car?\n",
+        "# @markdown > A car is a four-wheeled vehicle designed for the transportation of passengers and their belongings.\n",
+        "# @markdown ```\n",
+        "\n",
+        "# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for more details.\n",
+        "\n",
+        "# Loads an existing endpoint instance using the endpoint name:\n",
+        "# - Using `endpoint_name = endpoint_hexllm.name` allows us to get the endpoint\n",
+        "#   name of the endpoint `endpoint_hexllm` created in the cell above.\n",
+        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
+        "#   an existing endpoint with the ID 1234567890123456789.\n",
+        "# You may uncomment the code below to load an existing endpoint:\n",
+        "# endpoint_name = endpoint_without_peft.name\n",
+        "# # endpoint_name = \"\"  # @param {type:\"string\"}\n",
+        "# aip_endpoint_name = (\n",
+        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
+        "# )\n",
+        "# endpoint_hexllm = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "\n",
+        "prompt = \"What is a car?\"  # @param {type: \"string\"}\n",
+        "max_tokens = 50  # @param {type: \"integer\"}\n",
+        "temperature = 1.0  # @param {type: \"number\"}\n",
+        "top_p = 1.0  # @param {type: \"number\"}\n",
+        "top_k = 1  # @param {type: \"integer\"}\n",
+        "instances = [\n",
+        "    {\n",
+        "        \"prompt\": prompt,\n",
+        "        \"max_tokens\": max_tokens,\n",
+        "        \"temperature\": temperature,\n",
+        "        \"top_p\": top_p,\n",
+        "        \"top_k\": top_k,\n",
+        "    },\n",
+        "]\n",
+        "response = endpoint_hexllm.predict(instances=instances)\n",
+        "\n",
+        "prediction = response.predictions[0]\n",
+        "print(prediction)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "GFfqpQm8BNwZ"
+      },
+      "source": [
+        "## Deploy Gemma models with TGI on GPU\n",
+        "\n",
+        "[TGI](https://github.com/huggingface/text-generation-inference) stands for Text Generation Inference. It's a powerful library designed specifically for running large language models on GPUs efficiently. TGI utilizes techniques like \"paged attention\" and \"continuous batching\" to improve the speed and throughput of LLMs.\n",
+        "\n",
+        "Currently, only L4 GPUs are demonstrated in this notebook. Functionality on other GPU types will be confirmed and added in the future.\n",
+        "\n",
+        "Gemma2 9B models require at least 2 L4 GPUs for deployment. Gemma2 27B models require at least 4 L4 GPUs for deployment."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "TBNJYZMlBNwZ"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "MODEL_ID = \"gemma-2-9b\"  # @param [\"gemma-2-9b\", \"gemma-2-9b-it\", \"gemma-2-27b\", \"gemma-2-27b-it\"] {allow-input: true, isTemplate: true}\n",
+        "model_id = os.path.join(model_path_prefix, MODEL_ID)\n",
+        "\n",
+        "# @markdown Finds Vertex AI prediction supported accelerators and regions in\n",
+        "# @markdown https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
+        "\n",
+        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\"] {isTemplate: true}\n",
+        "\n",
+        "if \"9b\" in MODEL_ID:\n",
+        "    if accelerator_type == \"NVIDIA_L4\":\n",
+        "        # Sets 2 L4 (24G) to deploy Gemma 9B models.\n",
+        "        machine_type = \"g2-standard-24\"\n",
+        "        accelerator_count = 2\n",
+        "    else:\n",
+        "        raise ValueError(\n",
+        "            \"Recommended machine settings not found for accelerator type: %s\"\n",
+        "            % accelerator_type\n",
+        "        )\n",
+        "elif \"27b\" in MODEL_ID:\n",
+        "    if accelerator_type == \"NVIDIA_L4\":\n",
+        "        # Sets 4 L4 (24G) to deploy Gemma 27B models.\n",
+        "        machine_type = \"g2-standard-48\"\n",
+        "        accelerator_count = 4\n",
+        "    else:\n",
+        "        raise ValueError(\n",
+        "            \"Recommended machine settings not found for accelerator type: %s\"\n",
+        "            % accelerator_type\n",
+        "        )\n",
+        "else:\n",
+        "    raise ValueError(\"Recommended machine settings not found for model: %s\" % MODEL_ID)\n",
+        "\n",
+        "check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=False,\n",
+        ")\n",
+        "\n",
+        "# Note that larger token counts will require more GPU memory. For example, if you'd\n",
+        "# like to increase the `max_total_tokens` and `max_batch_prefill_tokens` to 8192,\n",
+        "# you may need 4 L4s for the 9b model, and 8 L4s for the 27b model.\n",
+        "max_input_length = 1562\n",
+        "max_total_tokens = 2048\n",
+        "max_batch_prefill_tokens = 2048\n",
+        "\n",
+        "model_tgi, endpoint_tgi = deploy_model_tgi(\n",
+        "    model_name=get_job_name_with_datetime(prefix=MODEL_ID),\n",
+        "    model_id=model_id,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    max_input_length=max_input_length,\n",
+        "    max_total_tokens=max_total_tokens,\n",
+        "    max_batch_prefill_tokens=max_batch_prefill_tokens,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "ZCwACkjuBNwZ"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts.\n",
+        "\n",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown > What is a car?\n",
+        "# @markdown > A car is a four-wheeled vehicle designed for the transportation of passengers and their belongings.\n",
+        "# @markdown ```\n",
+        "\n",
+        "# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for more details.\n",
+        "\n",
+        "\n",
+        "# @markdown Please click \"Show Code\" to see more details.\n",
+        "\n",
+        "# Loads an existing endpoint instance using the endpoint name:\n",
+        "# - Using `endpoint_name = endpoint_tgi.name` allows us to get the\n",
+        "#   endpoint name of the endpoint `endpoint_tgi` created in the cell\n",
+        "#   above.\n",
+        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
+        "#   an existing endpoint with the ID 1234567890123456789.\n",
+        "# You may uncomment the code below to load an existing endpoint.\n",
+        "\n",
+        "# endpoint_name = \"\"  # @param {type:\"string\"}\n",
+        "# aip_endpoint_name = (\n",
+        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
+        "# )\n",
+        "# endpoint_tgi = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "\n",
+        "prompt = \"What is a car?\"  # @param {type: \"string\"}\n",
+        "max_new_tokens = 128  # @param {type:\"integer\"}\n",
+        "temperature = 1.0  # @param {type:\"number\"}\n",
+        "top_p = 0.9  # @param {type:\"number\"}\n",
+        "top_k = 1  # @param {type:\"integer\"}\n",
+        "\n",
+        "# Overides max_new_tokens and top_k parameters during inferences.\n",
+        "# If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`,\n",
+        "# you can reduce the max length, such as set max_new_tokens as 20.\n",
+        "instances = [\n",
+        "    {\n",
+        "        \"inputs\": f\"### Human: {prompt}### Assistant: \",\n",
+        "        \"parameters\": {\n",
+        "            \"max_new_tokens\": max_new_tokens,\n",
+        "            \"temperature\": temperature,\n",
+        "            \"top_p\": top_p,\n",
+        "            \"top_k\": top_k,\n",
+        "        },\n",
+        "    },\n",
+        "]\n",
+        "\n",
+        "response = endpoint_tgi.predict(instances=instances)\n",
+        "\n",
+        "for prediction in response.predictions:\n",
+        "    print(prediction)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "af21a3cff1e0"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "911406c1561e"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Delete the models and endpoints\n",
+        "\n",
+        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
+        "# @markdown  and avoid unnecessary continouous charges that may incur.\n",
+        "# Undeploy models and delete endpoints.\n",
+        "endpoint_hexllm.delete(force=True)\n",
+        "endpoint_tgi.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "model_hexllm.delete()\n",
+        "model_tgi.delete()\n",
+        "\n",
+        "# Delete Cloud Storage objects.\n",
+        "delete_bucket = False  # @param {type:\"boolean\", isTemplate: true}\n",
+        "if delete_bucket:\n",
+        "    ! gsutil -m rm -r $BUCKET_URI"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_gemma2_deployment_on_vertex.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
@@ -31,19 +31,18 @@
      "source": [
        "# Gemma deployment to GKE using TGI on GPU\n",
        "\n",
-        "\u003ctable align=\"left\"\u003e\u003ctbody\u003e\u003ctr\u003e\n",
-        "  \u003ctd\u003e\n",
-        "    \u003ca href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_gemma_deployment_on_gke.ipynb\"\u003e\n",
-        "      \u003cimg alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"\u003e\u003cbr\u003e Run in Colab Enterprise\n",
-        "    \u003c/a\u003e\n",
-        "  \u003c/td\u003e\n",
-        "  \u003ctd\u003e\n",
-        "    \u003ca href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_gemma_deployment_on_gke.ipynb\"\u003e\n",
-        "      \u003cimg src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"\u003e\u003cbr\u003e\n",
-        "      View on GitHub\n",
-        "    \u003c/a\u003e\n",
-        "  \u003c/td\u003e\n",
-        "\u003c/tr\u003e\u003c/tbody\u003e\u003c/table\u003e"
+        "<table><tbody><tr>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_gemma_deployment_on_gke.ipynb\">\n",
+        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_gemma_deployment_on_gke.ipynb\">\n",
+        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
      ]
    },
    {
@@ -157,7 +156,7 @@
        "# Create Kubernetes secret for Hugging Face credentials\n",
        "! kubectl create secret generic hf-secret \\\n",
        "    --from-literal=hf_api_token={HF_TOKEN} \\\n",
-        "    --dry-run=client -o yaml \u003e hf-secret.yaml\n",
+        "    --dry-run=client -o yaml > hf-secret.yaml\n",
        "\n",
        "! kubectl apply -f hf-secret.yaml"
      ]
@@ -340,13 +339,22 @@
        "command = f\"\"\"kubectl exec -t $( kubectl get pod -l app=gemma-server -o jsonpath=\"{{.items[0].metadata.name}}\" ) -c inference-server -- curl -X POST http://localhost:8000/generate \\\n",
        "   -H \"Content-Type: application/json\" \\\n",
        "   -d '{json.dumps(request)}' \\\n",
-        "   2\u003e /dev/null\"\"\"\n",
+        "   2> /dev/null\"\"\"\n",
        "\n",
        "output = !{command}\n",
        "print(\"Output:\")\n",
        "print(json.loads(output[0])[\"generated_text\"])"
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "wbRmgoOZF6es"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
    {
      "cell_type": "code",
      "execution_count": null,
@@ -356,8 +364,6 @@
      },
      "outputs": [],
      "source": [
-        "# @title Clean up resources\n",
-        "\n",
        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
        "# @markdown  and avoid unnecessary continouous charges that may incur.\n",
        "\n",
@@ -377,7 +383,6 @@
  "metadata": {
    "colab": {
      "name": "model_garden_gemma_deployment_on_gke.ipynb",
-      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
@@ -31,18 +31,18 @@
      "source": [
        "# Vertex AI Model Garden - Gemma (Deployment)\n",
        "\n",
-        "\u003ctable\u003e\u003ctbody\u003e\u003ctr\u003e\n",
-        "  \u003ctd style=\"text-align: center\"\u003e\n",
-        "    \u003ca href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_gemma_deployment_on_vertex.ipynb\"\u003e\n",
-        "      \u003cimg alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"\u003e\u003cbr\u003e Run in Colab Enterprise\n",
-        "    \u003c/a\u003e\n",
-        "  \u003c/td\u003e\n",
-        "  \u003ctd style=\"text-align: center\"\u003e\n",
-        "    \u003ca href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_gemma_deployment_on_vertex.ipynb\"\u003e\n",
-        "      \u003cimg alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"\u003e\u003cbr\u003e View on GitHub\n",
-        "    \u003c/a\u003e\n",
-        "  \u003c/td\u003e\n",
-        "\u003c/tr\u003e\u003c/tbody\u003e\u003c/table\u003e"
+        "<table><tbody><tr>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_gemma_deployment_on_vertex.ipynb\">\n",
+        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_gemma_deployment_on_vertex.ipynb\">\n",
+        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
      ]
    },
    {
@@ -57,8 +57,6 @@
        " * on TPU using **Hex-LLM**, a **H**igh-**E**fficiency **L**arge **L**anguage **M**odel serving solution built with **XLA** that is being developed by Google Cloud, and\n",
        " * on GPU using [vLLM](https://github.com/vllm-project/vllm), the state-of-the-art open source LLM serving solution on GPU.\n",
        "\n",
-        "This notebook also showcases how to use the [Text moderation API](https://cloud.google.com/natural-language/docs/moderating-text) to analyze model predictions against a predefined list of safety attributes.\n",
-        "\n",
        "\n",
        "### Objective\n",
        "\n",
@@ -84,17 +82,6 @@
        "## Before you begin"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "C_wC61dhpWXj"
-      },
-      "source": [
-        "### Request for TPU quota\n",
-        "\n",
-        "By default, the quota for TPU deployment `Custom model serving TPU v5e cores per region` is 0. If you would like to use Hex-LLM TPU deployment, please request TPU quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota) if you haven't done so already. Please note that vLLM GPU deployment does not need this step."
-      ]
-    },
    {
      "cell_type": "code",
      "execution_count": null,
@@ -108,28 +95,44 @@
        "\n",
        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
-        "# @markdown **[Optional]** Set the GCS BUCKET_URI to store the experiment artifacts, if you want to use your own bucket. **If not set, a unique GCS bucket will be created automatically on your behalf**.\n",
+        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
        "\n",
+        "# @markdown 3. By default, the quota for TPU deployment `Custom model serving TPU v5e cores per region` is 4. TPU quota is only available in `us-west1`. You can request for higher TPU quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
+        "\n",
+        "# Import the necessary packages\n",
+        "\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
+        "import importlib\n",
        "import os\n",
-        "import json\n",
        "from datetime import datetime\n",
        "from typing import Tuple\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "\n",
+        "common_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
+        ")\n",
+        "\n",
+        "models, endpoints = {}, {}\n",
+        "\n",
        "# Get the default cloud project id.\n",
        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
        "\n",
        "# Get the default region for launching jobs.\n",
        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
        "\n",
+        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
+        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
+        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
+        "\n",
        "# Cloud Storage bucket for storing the experiment artifacts.\n",
        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
-        "# prefer using your own GCS bucket, please change the value yourself below.\n",
+        "# prefer using your own GCS bucket, change the value yourself below.\n",
        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
+        "\n",
        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
-        "    # Create a unique GCS bucket for this notebook if not specified\n",
        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\n",
        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
        "else:\n",
@@ -146,6 +149,10 @@
        "\n",
        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"gemma\")\n",
+        "\n",
+        "\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
        "\n",
        "# Gets the default SERVICE_ACCOUNT.\n",
@@ -154,29 +161,28 @@
        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
        "\n",
+        "\n",
        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
-        "# Enable Vertex AI and Cloud Compute APIs.\n",
        "! gcloud config set project $PROJECT_ID\n",
-        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
        "\n",
        "# @markdown ## Access Gemma Models\n",
-        "# @markdown For GPU based serving, choose between accessing Gemma models on [Hugging Face](https://huggingface.co/)\n",
+        "# @markdown Choose between accessing Gemma models on [Hugging Face](https://huggingface.co/)\n",
        "# @markdown or Vertex AI as described below.\n",
        "\n",
        "# @markdown If you already obtained access to Gemma models on [Hugging Face](https://huggingface.co/), you can load models from there.\n",
        "# @markdown Alternatively, you can also load the original Gemma models for serving from Vertex AI after accepting the agreement.\n",
        "\n",
-        "# @markdown For TPU based serving with Hex-LLM, choose the Kaggle option.\n",
-        "\n",
-        "# @markdown **Please only select and fill one of the three following sections.**\n",
-        "LOAD_MODEL_FROM = \"Hugging Face\"  # @param [\"Hugging Face\", \"Google Cloud\", \"Kaggle\"] {isTemplate:true}\n",
+        "# @markdown **Please only select and fill one of the two following sections.**\n",
+        "LOAD_MODEL_FROM = (\n",
+        "    \"Hugging Face\"  # @param [\"Hugging Face\", \"Google Cloud\"] {isTemplate:true}\n",
+        ")\n",
        "\n",
        "# @markdown ---\n",
        "\n",
-        "# @markdown ### Access Gemma models on Hugging Face for GPU based serving\n",
+        "# @markdown ### Access Gemma models on Hugging Face\n",
        "# @markdown You must provide a Hugging Face User Access Token (read) to access the Gemma models. You can follow the [Hugging Face documentation](https://huggingface.co/docs/hub/en/security-tokens) to create a **read** access token and put it in the `HF_TOKEN` field below.\n",
        "\n",
        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate:true}\n",
@@ -186,7 +192,7 @@
        "    ), \"Please provide a read HF_TOKEN to load models from Hugging Face, or select a different model source.\"\n",
        "\n",
        "# @markdown *--- Or ---*\n",
-        "# @markdown ### Access Gemma models on Vertex AI for GPU based serving\n",
+        "# @markdown ### Access Gemma models on Vertex AI\n",
        "# @markdown Accept the model agreement to access the models:\n",
        "# @markdown 1. Open the [Gemma model card](https://console.cloud.google.com/vertex-ai/publishers/google/model-garden/335) from [Vertex AI Model Garden](https://cloud.google.com/model-garden).\n",
        "# @markdown 1. Review and accept the agreement in the pop-up window on the model card page. If you have previously accepted the model agreement, there will not be a pop-up window on the model card page and this step is not needed.\n",
@@ -214,72 +220,74 @@
        "else:\n",
        "    model_path_prefix = \"google/\"\n",
        "\n",
-        "\n",
-        "# @markdown *--- Or ---*\n",
-        "# @markdown ### Access Gemma models from Kaggle for TPU based serving\n",
-        "# @markdown Kaggle credentials are required for Hex-LLM deployment with TPUs.\n",
-        "# @markdown Generate the Kaggle username and key by following [these instructions](https://github.com/Kaggle/kaggle-api?tab=readme-ov-file#api-credentials).\n",
-        "# @markdown You will need to review and accept the model license.\n",
-        "KAGGLE_USERNAME = \"\"  # @param {type:\"string\", isTemplate:true}\n",
-        "KAGGLE_KEY = \"\"  # @param {type:\"string\", isTemplate:true}\n",
-        "if LOAD_MODEL_FROM == \"Kaggle\":\n",
-        "    assert (\n",
-        "        KAGGLE_USERNAME and KAGGLE_KEY\n",
-        "    ), \"Please provide Kaggle credentials to load models from Kaggle, or select a different model source.\"\n",
        "# @markdown ---\n",
        "\n",
        "# The pre-built serving docker images.\n",
+        "HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:deploy\"\n",
        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240508_0916_RC02\"\n",
-        "HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:20240426_0936_RC01\"\n",
-        "\n",
-        "\n",
-        "def get_job_name_with_datetime(prefix: str) -\u003e str:\n",
-        "    \"\"\"Gets the job name with date time when triggering deployment jobs.\"\"\"\n",
-        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
        "\n",
        "\n",
        "def deploy_model_hexllm(\n",
        "    model_name: str,\n",
        "    model_id: str,\n",
        "    service_account: str,\n",
+        "    base_model_id: str = None,\n",
+        "    tensor_parallel_size: int = 1,\n",
        "    machine_type: str = \"ct5lp-hightpu-1t\",\n",
-        "    max_num_batched_tokens: int = 11264,\n",
-        "    tokens_pad_multiple: int = 1024,\n",
-        "    seqs_pad_multiple: int = 32,\n",
-        ") -\u003e Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    hbm_utilization_factor: float = 0.6,\n",
+        "    max_running_seqs: int = 256,\n",
+        "    endpoint_id: str = \"\",\n",
+        "    min_replica_count: int = 1,\n",
+        "    max_replica_count: int = 1,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
        "    \"\"\"Deploys models with Hex-LLM on TPU in Vertex AI.\"\"\"\n",
-        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
+        "    if endpoint_id:\n",
+        "        aip_endpoint_name = (\n",
+        "            f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_id}\"\n",
+        "        )\n",
+        "        endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "    else:\n",
+        "        endpoint = aiplatform.Endpoint.create(\n",
+        "            display_name=f\"{model_name}-endpoint\",\n",
+        "            location=TPU_DEPLOYMENT_REGION,\n",
+        "        )\n",
+        "\n",
+        "    if not base_model_id:\n",
+        "        base_model_id = model_id\n",
+        "\n",
+        "    if not tensor_parallel_size:\n",
+        "        tensor_parallel_size = int(machine_type[-2])\n",
        "\n",
-        "    num_tpu_chips = int(machine_type[-2])\n",
        "    hexllm_args = [\n",
        "        \"--host=0.0.0.0\",\n",
        "        \"--port=7080\",\n",
        "        \"--log_level=INFO\",\n",
        "        f\"--model={model_id}\",\n",
-        "        f\"--tensor_parallel_size={num_tpu_chips}\",\n",
-        "        \"--num_nodes=1\",\n",
-        "        \"--use_ray\",\n",
-        "        \"--batch_mode=continuous\",\n",
-        "        f\"--max_num_batched_tokens={max_num_batched_tokens}\",\n",
-        "        f\"--tokens_pad_multiple={tokens_pad_multiple}\",\n",
-        "        f\"--seqs_pad_multiple={seqs_pad_multiple}\",\n",
+        "        f\"--tensor_parallel_size={tensor_parallel_size}\",\n",
+        "        \"--enable_jit\",\n",
+        "        \"--load_format=auto\",\n",
+        "        f\"--hbm_utilization_factor={hbm_utilization_factor}\",\n",
+        "        f\"--max_running_seqs={max_running_seqs}\",\n",
        "    ]\n",
        "\n",
        "    env_vars = {\n",
-        "        \"MODEL_ID\": model_id,\n",
+        "        \"MODEL_ID\": base_model_id,\n",
        "        \"PJRT_DEVICE\": \"TPU\",\n",
        "        \"RAY_DEDUP_LOGS\": \"0\",\n",
        "        \"RAY_USAGE_STATS_ENABLED\": \"0\",\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
-        "    if KAGGLE_USERNAME and KAGGLE_KEY:\n",
-        "        env_vars[\"KAGGLE_USERNAME\"] = KAGGLE_USERNAME\n",
-        "        env_vars[\"KAGGLE_KEY\"] = KAGGLE_KEY\n",
+        "\n",
+        "    try:\n",
+        "        if HF_TOKEN:\n",
+        "            env_vars.update({\"HF_TOKEN\": HF_TOKEN})\n",
+        "    except:\n",
+        "        pass\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
        "        serving_container_image_uri=HEXLLM_DOCKER_URI,\n",
-        "        serving_container_command=[\"python\", \"-m\", \"hex_llm.entrypoints.api_server\"],\n",
+        "        serving_container_command=[\"python\", \"-m\", \"hex_llm.server.api_server\"],\n",
        "        serving_container_args=hexllm_args,\n",
        "        serving_container_ports=[7080],\n",
        "        serving_container_predict_route=\"/generate\",\n",
@@ -287,6 +295,7 @@
        "        serving_container_environment_variables=env_vars,\n",
        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
        "        serving_container_deployment_timeout=7200,\n",
+        "        location=TPU_DEPLOYMENT_REGION,\n",
        "    )\n",
        "\n",
        "    model.deploy(\n",
@@ -294,6 +303,8 @@
        "        machine_type=machine_type,\n",
        "        deploy_request_timeout=1800,\n",
        "        service_account=service_account,\n",
+        "        min_replica_count=min_replica_count,\n",
+        "        max_replica_count=max_replica_count,\n",
        "    )\n",
        "    return model, endpoint\n",
        "\n",
@@ -302,33 +313,42 @@
        "    model_name: str,\n",
        "    model_id: str,\n",
        "    service_account: str,\n",
-        "    machine_type: str = \"g2-standard-12\",\n",
+        "    base_model_id: str = None,\n",
+        "    machine_type: str = \"g2-standard-8\",\n",
        "    accelerator_type: str = \"NVIDIA_L4\",\n",
        "    accelerator_count: int = 1,\n",
-        "    max_model_len: int = 8192,\n",
-        "    dtype: str = \"bfloat16\",\n",
-        ") -\u003e Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
-        "    \"\"\"Deploys models with vLLM on GPU in Vertex AI.\"\"\"\n",
+        "    gpu_memory_utilization: float = 0.9,\n",
+        "    max_model_len: int = 4096,\n",
+        "    dtype: str = \"auto\",\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",
+        "    if not base_model_id:\n",
+        "        base_model_id = model_id\n",
+        "\n",
        "    vllm_args = [\n",
        "        \"--host=0.0.0.0\",\n",
        "        \"--port=7080\",\n",
        "        f\"--model={model_id}\",\n",
        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
        "        \"--swap-space=16\",\n",
-        "        \"--gpu-memory-utilization=0.9\",\n",
+        "        f\"--gpu-memory-utilization={gpu_memory_utilization}\",\n",
        "        f\"--max-model-len={max_model_len}\",\n",
        "        f\"--dtype={dtype}\",\n",
        "        \"--disable-log-stats\",\n",
        "    ]\n",
        "\n",
        "    env_vars = {\n",
-        "        \"MODEL_ID\": model_id,\n",
+        "        \"MODEL_ID\": base_model_id,\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
-        "    if HF_TOKEN:\n",
-        "        env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
+        "\n",
+        "    try:\n",
+        "        if HF_TOKEN:\n",
+        "            env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
+        "    except:\n",
+        "        pass\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
@@ -342,7 +362,9 @@
        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
        "        serving_container_deployment_timeout=7200,\n",
        "    )\n",
-        "\n",
+        "    print(\n",
+        "        f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
+        "    )\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        machine_type=machine_type,\n",
@@ -351,91 +373,9 @@
        "        deploy_request_timeout=1800,\n",
        "        service_account=service_account,\n",
        "    )\n",
-        "    return model, endpoint\n",
+        "    print(\"endpoint_name:\", endpoint.name)\n",
        "\n",
-        "\n",
-        "def get_quota(project_id: str, region: str, resource_id: str) -> int:\n",
-        "  \"\"\"Returns the quota for a resource in a region. Returns -1 if can not figure out the quota.\"\"\"\n",
-        "  service_endpoint = \"aiplatform.googleapis.com\"\n",
-        "  quota_list_output = !gcloud alpha services quota list --service=$service_endpoint  --consumer=projects/$project_id --filter=\"$service_endpoint/$resource_id\" --format=json\n",
-        "  # Use '.s' on the command output because it is an SList type.\n",
-        "  quota_data = json.loads(quota_list_output.s)\n",
-        "  if len(quota_data) == 0 or \"consumerQuotaLimits\" not in quota_data[0]:\n",
-        "    return -1\n",
-        "  if len(quota_data[0][\"consumerQuotaLimits\"]) == 0 or \"quotaBuckets\" not in quota_data[0][\"consumerQuotaLimits\"][0]:\n",
-        "    return -1\n",
-        "  all_regions_data = quota_data[0][\"consumerQuotaLimits\"][0][\"quotaBuckets\"]\n",
-        "  for region_data in all_regions_data:\n",
-        "    if region_data.get('dimensions') and region_data['dimensions']['region'] == region:\n",
-        "      if 'effectiveLimit' in region_data:\n",
-        "        return int(region_data['effectiveLimit'])\n",
-        "      else:\n",
-        "        return 0\n",
-        "  return -1\n",
-        "\n",
-        "\n",
-        "def get_resource_id(accelerator_type: str, is_for_training: bool) -> str:\n",
-        "  \"\"\"Returns the resource id for a given accelerator type and the use case.\n",
-        "  Args:\n",
-        "    accelerator_type: The accelerator type.\n",
-        "    is_for_training: Whether the resource is used for training. Set false\n",
-        "    for serving use case.\n",
-        "  Returns:\n",
-        "    The resource id.\n",
-        "  \"\"\"\n",
-        "  training_accelerator_map = {\n",
-        "      \"NVIDIA_TESLA_V100\": \"custom_model_training_nvidia_v100_gpus\",\n",
-        "      \"NVIDIA_L4\": \"custom_model_training_nvidia_l4_gpus\",\n",
-        "      \"NVIDIA_TESLA_A100\": \"custom_model_training_nvidia_a100_gpus\",\n",
-        "      \"NVIDIA_TESLA_T4\": \"custom_model_training_nvidia_t4_gpus\",\n",
-        "      \"TPU_V5e\": \"custom_model_training_tpu_v5e\",\n",
-        "      \"TPU_V3\": \"custom_model_training_tpu_v3\",\n",
-        "  }\n",
-        "  serving_accelerator_map = {\n",
-        "      \"NVIDIA_TESLA_V100\": \"custom_model_serving_nvidia_v100_gpus\",\n",
-        "      \"NVIDIA_L4\": \"custom_model_serving_nvidia_l4_gpus\",\n",
-        "      \"NVIDIA_TESLA_A100\": \"custom_model_serving_nvidia_a100_gpus\",\n",
-        "      \"NVIDIA_TESLA_T4\": \"custom_model_serving_nvidia_t4_gpus\",\n",
-        "      \"TPU_V5e\": \"custom_model_serving_tpu_v5e\",\n",
-        "  }\n",
-        "  if is_for_training:\n",
-        "    if accelerator_type in training_accelerator_map:\n",
-        "      return training_accelerator_map[accelerator_type]\n",
-        "    else:\n",
-        "      raise ValueError(\n",
-        "          f\"Could not find accelerator type: {accelerator_type} for training.\"\n",
-        "      )\n",
-        "  else:\n",
-        "    if accelerator_type in serving_accelerator_map:\n",
-        "      return serving_accelerator_map[accelerator_type]\n",
-        "    else:\n",
-        "      raise ValueError(\n",
-        "          f\"Could not find accelerator type: {accelerator_type} for serving.\"\n",
-        "      )\n",
-        "\n",
-        "\n",
-        "def check_quota(project_id:str, region: str, accelerator_type: str,\n",
-        "                accelerator_count: int, is_for_training: bool):\n",
-        "  \"\"\"Checks if the project and the region has the required quota.\"\"\"\n",
-        "  resource_id = get_resource_id(accelerator_type, is_for_training)\n",
-        "  quota = get_quota(project_id, region, resource_id)\n",
-        "  quota_request_instruction = (\"Either use \"\n",
-        "            \"a different region or request additional quota. Follow \"\n",
-        "            \"instructions here \"\n",
-        "            \"https://cloud.google.com/docs/quotas/view-manage#requesting_higher_quota\"\n",
-        "            \" to check quota in a region or request additional quota for \"\n",
-        "            \"your project.\")\n",
-        "  if quota == -1:\n",
-        "    raise ValueError(\n",
-        "            f\"\"\"Quota not found for: {resource_id} in {region}.\n",
-        "            {quota_request_instruction}\"\"\"\n",
-        "        )\n",
-        "  if quota < accelerator_count:\n",
-        "    raise ValueError(\n",
-        "            f\"\"\"Quota not enough for {resource_id} in {region}:\n",
-        "            {quota} < {accelerator_count}.\n",
-        "            {quota_request_instruction}\"\"\"\n",
-        "        )"
+        "    return model, endpoint"
      ]
    },
    {
@@ -448,7 +388,7 @@
        "\n",
        "**Hex-LLM** is a **H**igh-**E**fficiency **L**arge **L**anguage **M**odel (LLM) TPU serving solution built with **XLA**, which is being developed by Google Cloud.\n",
        "\n",
-        "To request TPU quota, please follow the instructions at [Request a higher quota](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota)."
+        "Refer to the \"Request for TPU quota\" section for TPU quota."
      ]
    },
    {
@@ -461,34 +401,16 @@
      "outputs": [],
      "source": [
        "# @title Deploy\n",
-        "# @markdown Set the model ID. Model weights can be loaded from Kaggle [google/gemma](https://www.kaggle.com/models/google/gemma/frameworks/pyTorch) or from a GCS bucket.\n",
-        "\n",
-        "# @markdown Alternatively, you can specify a GCS folder that contains the original or\n",
-        "# @markdown customized PyTorch model checkpoint and tokenizer. In this case, the GCS\n",
-        "# @markdown folder is expected to contain \"gemma-2b\", \"gemma-7b\", \"gemma-1.1-2b\" or\n",
-        "# @markdown  \"gemma-1.1-7b\" in either the GCS folder path or the name of the model checkpoint\n",
-        "# @markdown file is expected to have the suffix \".ckpt\" and the tokenizer file is\n",
-        "# @markdown file. The model checkpointis expected to have the name \"tokenizer.model\" (same as\n",
-        "# @markdown the Kaggle files). An example structure for the GCS folder\n",
-        "# @markdown gs://my-deployment-bucket/pytorch-files is:\n",
-        "# @markdown - gs://my-deployment-bucket/pytorch-files/gemma-2b.ckpt\n",
-        "# @markdown - gs://my-deployment-bucket/pytorch-files/tokenizer.model\n",
-        "\n",
-        "# @markdown Select one of the four model variations or enter a GCS folder containing the original or customized PyTorch model checkpoint and tokenizer.\n",
-        "MODEL_ID = \"google/gemma-1.1-2b-it\"  # @param [\"google/gemma-2b\", \"google/gemma-2b-it\", \"google/gemma-7b\", \"google/gemma-7b-it\", \"google/gemma-1.1-2b-it\", \"google/gemma-1.1-7b-it\"] {allow-input: true, isTemplate: true}\n",
+        "# @markdown Set the model ID. Model weights can be loaded from HuggingFace or from a GCS bucket.\n",
        "\n",
+        "# @markdown Select one of the six model variations.\n",
+        "MODEL_ID = \"gemma-1.1-2b-it\"  # @param [\"gemma-2b\", \"gemma-2b-it\", \"gemma-7b\", \"gemma-7b-it\", \"gemma-1.1-2b-it\", \"gemma-1.1-7b-it\"] {allow-input: true, isTemplate: true}\n",
+        "TPU_DEPLOYMENT_REGION = \"us-west1\"  # @param [\"us-west1\"] {isTemplate:true}\n",
+        "model_id = os.path.join(model_path_prefix, MODEL_ID)\n",
        "\n",
        "# @markdown Find Vertex AI prediction TPUv5e machine types in\n",
        "# @markdown https://cloud.google.com/vertex-ai/docs/predictions/use-tpu#deploy_a_model.\n",
-        "\n",
-        "import re\n",
-        "\n",
-        "if MODEL_ID.startswith(\"gs://\"):\n",
-        "    MODEL_BUCKET_URI = re.search(\"gs://(.*?)/\", MODEL_ID).group()\n",
-        "    ! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $MODEL_BUCKET_URI\n",
-        "    MODEL_ID = MODEL_BUCKET_URI\n",
-        "\n",
-        "if \"2b\" in MODEL_ID:\n",
+        "if \"2b\" in model_id:\n",
        "    # Sets ct5lp-hightpu-1t (1 TPU chip) to deploy Gemma 2B models.\n",
        "    machine_type = \"ct5lp-hightpu-1t\"\n",
        "    accelerator_type = \"TPU_V5e\"\n",
@@ -501,29 +423,31 @@
        "    # Note: 1 TPU V5 chip has only one core.\n",
        "    accelerator_count = 4\n",
        "\n",
-        "check_quota(project_id=PROJECT_ID,\n",
-        "            region=REGION,\n",
-        "            accelerator_type=accelerator_type,\n",
-        "            accelerator_count=accelerator_count,\n",
-        "            is_for_training=False)\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=False,\n",
+        ")\n",
        "\n",
-        "# Note that a larger max_num_batched_tokens will require more TPU memory.\n",
-        "max_num_batched_tokens = 11264\n",
-        "# Multiple of tokens for padding alignment. A higher value can reduce\n",
-        "# re-compilation but can also increase the waste in computation.\n",
-        "tokens_pad_multiple = 1024\n",
-        "# Multiple of sequences for padding alignment. A higher value can reduce\n",
-        "# re-compilation but can also increase the waste in computation.\n",
-        "seqs_pad_multiple = 32\n",
+        "# Server parameters.\n",
+        "hbm_utilization_factor = 0.6  # A larger value improves throughput but gives higher risk of TPU out-of-memory errors with long prompts.\n",
+        "max_running_seqs = 256\n",
        "\n",
-        "model_hexllm, endpoint_hexllm = deploy_model_hexllm(\n",
-        "    model_name=get_job_name_with_datetime(prefix=\"gemma-serve-hexllm\"),\n",
-        "    model_id=MODEL_ID,\n",
+        "# Endpoint configurations.\n",
+        "min_replica_count = 1\n",
+        "max_replica_count = 1\n",
+        "\n",
+        "models[\"hexllm_tpu\"], endpoints[\"hexllm_tpu\"] = deploy_model_hexllm(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=MODEL_ID),\n",
+        "    model_id=model_id,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        "    machine_type=machine_type,\n",
-        "    max_num_batched_tokens=max_num_batched_tokens,\n",
-        "    tokens_pad_multiple=tokens_pad_multiple,\n",
-        "    seqs_pad_multiple=seqs_pad_multiple,\n",
+        "    hbm_utilization_factor=hbm_utilization_factor,\n",
+        "    max_running_seqs=max_running_seqs,\n",
+        "    min_replica_count=min_replica_count,\n",
+        "    max_replica_count=max_replica_count,\n",
        ")"
      ]
    },
@@ -538,20 +462,22 @@
      "source": [
        "# @title Predict\n",
        "\n",
-        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. Note that the first few prompts will take longer to execute.\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts based on your `template`. Note that the first few prompts will take longer to execute.\n",
+        "\n",
+        "# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for more details.\n",
        "\n",
        "# @markdown Example:\n",
        "\n",
        "# @markdown ```\n",
-        "# @markdown \u003e What is a car?\n",
-        "# @markdown \u003e A car is a four-wheeled vehicle designed for the transportation of passengers and their belongings.\n",
+        "# @markdown > What is a car?\n",
+        "# @markdown > A car is a four-wheeled vehicle designed for the transportation of passengers and their belongings.\n",
        "# @markdown ```\n",
        "\n",
        "# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for more details.\n",
        "\n",
        "# Loads an existing endpoint instance using the endpoint name:\n",
-        "# - Using `endpoint_name = endpoint_hexllm.name` allows us to get the endpoint\n",
-        "#   name of the endpoint `endpoint_hexllm` created in the cell above.\n",
+        "# - Using `endpoint_name = endpoint.name` allows us to get the endpoint\n",
+        "#   name of the endpoint `endpoint` created in the cell above.\n",
        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
        "#   an existing endpoint with the ID 1234567890123456789.\n",
        "# You may uncomment the code below to load an existing endpoint:\n",
@@ -560,13 +486,13 @@
        "# aip_endpoint_name = (\n",
        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
        "# )\n",
-        "# endpoint_hexllm = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
        "\n",
        "prompt = \"What is a car?\"  # @param {type: \"string\"}\n",
        "max_tokens = 50  # @param {type: \"integer\"}\n",
        "temperature = 1.0  # @param {type: \"number\"}\n",
        "top_p = 1.0  # @param {type: \"number\"}\n",
-        "top_k = 10  # @param {type: \"integer\"}\n",
+        "top_k = 1  # @param {type: \"integer\"}\n",
        "instances = [\n",
        "    {\n",
        "        \"prompt\": prompt,\n",
@@ -576,10 +502,10 @@
        "        \"top_k\": top_k,\n",
        "    },\n",
        "]\n",
-        "response = endpoint_hexllm.predict(instances=instances)\n",
+        "response = endpoints[\"hexllm_tpu\"].predict(instances=instances)\n",
        "\n",
-        "prediction = response.predictions[0]\n",
-        "print(prediction)"
+        "for prediction in response.predictions:\n",
+        "    print(prediction)"
      ]
    },
    {
@@ -595,25 +521,25 @@
        "The instruction tuned Gemma models were trained with a specific formatter that annotates instruction tuning examples with extra information, both during training and inference. The annotations (1) indicate roles in a conversation, and (2) delineate tunes in a conversation. Below we show a sample code snippet for formatting the model prompt using the user and model chat templates for a multi-turn conversation. The relevant tokens are:\n",
        "- `user`: user turn\n",
        "- `model`: model turn\n",
-        "- `\u003cstart_of_turn\u003e`: beginning of dialogue turn\n",
-        "- `\u003cend_of_turn\u003e`: end of dialogue turn\n",
+        "- `<start_of_turn>`: beginning of dialogue turn\n",
+        "- `<end_of_turn>`: end of dialogue turn\n",
        "\n",
        "An example set of dialogues is:\n",
        "```\n",
-        "\u003cstart_of_turn\u003euser\n",
-        "knock knock\u003cend_of_turn\u003e\n",
-        "\u003cstart_of_turn\u003emodel\n",
-        "who is there\u003cend_of_turn\u003e\n",
-        "\u003cstart_of_turn\u003euser\n",
-        "LaMDA\u003cend_of_turn\u003e\n",
-        "\u003cstart_of_turn\u003emodel\n",
-        "LaMDA who?\u003cend_of_turn\u003e\n",
+        "<start_of_turn>user\n",
+        "knock knock<end_of_turn>\n",
+        "<start_of_turn>model\n",
+        "who is there<end_of_turn>\n",
+        "<start_of_turn>user\n",
+        "LaMDA<end_of_turn>\n",
+        "<start_of_turn>model\n",
+        "LaMDA who?<end_of_turn>\n",
        "```\n",
-        "where `\u003cend_of_turn\u003e\\n` is the turn separator and `\u003cstart_of_turn\u003emodel\\n` is the prompt prefix. This means if we would like to prompt the model with a question like, `What is Cramer's Rule?`, we should use:\n",
+        "where `<end_of_turn>\\n` is the turn separator and `<start_of_turn>model\\n` is the prompt prefix. This means if we would like to prompt the model with a question like, `What is Cramer's Rule?`, we should use:\n",
        "```\n",
-        "\u003cstart_of_turn\u003euser\n",
-        "What is Cramer's Rule?\u003cend_of_turn\u003e\n",
-        "\u003cstart_of_turn\u003emodel\n",
+        "<start_of_turn>user\n",
+        "What is Cramer's Rule?<end_of_turn>\n",
+        "<start_of_turn>model\n",
        "```"
      ]
    },
@@ -626,15 +552,15 @@
      "outputs": [],
      "source": [
        "# Chat templates.\n",
-        "USER_CHAT_TEMPLATE = \"\u003cstart_of_turn\u003euser\\n{prompt}\u003cend_of_turn\u003e\\n\"\n",
-        "MODEL_CHAT_TEMPLATE = \"\u003cstart_of_turn\u003emodel\\n{prompt}\u003cend_of_turn\u003e\\n\"\n",
+        "USER_CHAT_TEMPLATE = \"<start_of_turn>user\\n{prompt}<end_of_turn>\\n\"\n",
+        "MODEL_CHAT_TEMPLATE = \"<start_of_turn>model\\n{prompt}<end_of_turn>\\n\"\n",
        "\n",
        "# Sample formatted prompt.\n",
        "prompt = (\n",
        "    USER_CHAT_TEMPLATE.format(prompt=\"What is a good place for travel in the US?\")\n",
        "    + MODEL_CHAT_TEMPLATE.format(prompt=\"California.\")\n",
        "    + USER_CHAT_TEMPLATE.format(prompt=\"What can I do in California?\")\n",
-        "    + \"\u003cstart_of_turn\u003emodel\\n\"\n",
+        "    + \"<start_of_turn>model\\n\"\n",
        ")\n",
        "print(\"Chat prompt:\\n\", prompt)\n",
        "\n",
@@ -644,10 +570,10 @@
        "        \"max_tokens\": 50,\n",
        "        \"temperature\": 1.0,\n",
        "        \"top_p\": 1.0,\n",
-        "        \"top_k\": 10,\n",
+        "        \"top_k\": 1,\n",
        "    },\n",
        "]\n",
-        "response = endpoint_hexllm.predict(instances=instances)\n",
+        "response = endpoints[\"hexllm_tpu\"].predict(instances=instances)\n",
        "\n",
        "prediction = response.predictions[0]\n",
        "print(prediction)"
@@ -734,17 +660,19 @@
        "        % accelerator_type\n",
        "    )\n",
        "\n",
-        "check_quota(project_id=PROJECT_ID,\n",
-        "            region=REGION,\n",
-        "            accelerator_type=accelerator_type,\n",
-        "            accelerator_count=accelerator_count,\n",
-        "            is_for_training=False)\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=False,\n",
+        ")\n",
        "\n",
        "# Note that a larger max_model_len will require more GPU memory.\n",
        "max_model_len = 2048\n",
        "\n",
-        "model_vllm, endpoint_vllm = deploy_model_vllm(\n",
-        "    model_name=get_job_name_with_datetime(prefix=\"gemma-serve-vllm\"),\n",
+        "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=\"gemma-serve-vllm\"),\n",
        "    model_id=model_id,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        "    machine_type=machine_type,\n",
@@ -775,26 +703,37 @@
      "source": [
        "# @title Predict\n",
        "\n",
-        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. Note that the first few prompts will take longer to execute.\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. Sampling parameters supported by vLLM can be found [here](https://docs.vllm.ai/en/latest/dev/sampling_params.html).\n",
+        "\n",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown Human: What is a car?\n",
+        "# @markdown Assistant:  A car, or a motor car, is a road-connected human-transportation system used to move people or goods from one place to another. The term also encompasses a wide range of vehicles, including motorboats, trains, and aircrafts. Cars typically have four wheels, a cabin for passengers, and an engine or motor. They have been around since the early 19th century and are now one of the most popular forms of transportation, used for daily commuting, shopping, and other purposes.\n",
+        "# @markdown ```\n",
+        "# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for more details.\n",
        "\n",
        "# Loads an existing endpoint instance using the endpoint name:\n",
-        "# - Using `endpoint_name = endpoint_vllm.name` allows us to get the endpoint\n",
-        "#   name of the endpoint `endpoint_vllm` created in the cell above.\n",
+        "# - Using `endpoint_name = endpoint.name` allows us to get the\n",
+        "#   endpoint name of the endpoint `endpoint` created in the cell\n",
+        "#   above.\n",
        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
        "#   an existing endpoint with the ID 1234567890123456789.\n",
-        "# You may uncomment the code below to load an existing endpoint:\n",
-        "# endpoint_name = endpoint_without_peft.name\n",
-        "# # endpoint_name = \"\"  # @param {type:\"string\"}\n",
+        "# You may uncomment the code below to load an existing endpoint.\n",
+        "\n",
+        "# endpoint_name = \"\"  # @param {type:\"string\"}\n",
        "# aip_endpoint_name = (\n",
        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
        "# )\n",
-        "# endpoint_vllm = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
        "\n",
        "prompt = \"What is a car?\"  # @param {type: \"string\"}\n",
-        "max_tokens = 50  # @param {type: \"integer\"}\n",
-        "temperature = 1.0  # @param {type: \"number\"}\n",
-        "top_p = 1.0  # @param {type: \"number\"}\n",
-        "top_k = 10  # @param {type: \"integer\"}\n",
+        "max_tokens = 50  # @param {type:\"integer\"}\n",
+        "temperature = 1.0  # @param {type:\"number\"}\n",
+        "top_p = 1.0  # @param {type:\"number\"}\n",
+        "top_k = 1  # @param {type:\"integer\"}\n",
+        "raw_response = False  # @param {type:\"boolean\"}\n",
+        "\n",
        "instances = [\n",
        "    {\n",
        "        \"prompt\": prompt,\n",
@@ -802,12 +741,13 @@
        "        \"temperature\": temperature,\n",
        "        \"top_p\": top_p,\n",
        "        \"top_k\": top_k,\n",
+        "        \"raw_response\": raw_response,\n",
        "    },\n",
        "]\n",
-        "response = endpoint_vllm.predict(instances=instances)\n",
+        "response = endpoints[\"vllm_gpu\"].predict(instances=instances)\n",
        "\n",
-        "prediction = response.predictions[0]\n",
-        "print(prediction)"
+        "for prediction in response.predictions:\n",
+        "    print(prediction)"
      ]
    },
    {
@@ -839,20 +779,18 @@
      },
      "outputs": [],
      "source": [
-        "# @title  Delete the models and endpoints\n",
-        "\n",
        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
        "# @markdown  and avoid unnecessary continouous charges that may incur.\n",
-        "# Undeploy models and delete endpoints.\n",
-        "endpoint_hexllm.delete(force=True)\n",
-        "endpoint_vllm.delete(force=True)\n",
+        "\n",
+        "# Undeploy model and delete endpoint.\n",
+        "for endpoint in endpoints.values():\n",
+        "    endpoint.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
-        "model_hexllm.delete()\n",
-        "model_vllm.delete()\n",
+        "for model in models.values():\n",
+        "    model.delete()\n",
        "\n",
-        "# Delete Cloud Storage objects.\n",
-        "delete_bucket = False  # @param {type:\"boolean\", isTemplate: true}\n",
+        "delete_bucket = False  # @param {type:\"boolean\"}\n",
        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
@@ -3,6 +3,7 @@
    {
      "cell_type": "code",
      "execution_count": null,
+      "language": "python",
      "metadata": {
        "id": "B8S-yo8qTIcO"
      },
@@ -91,6 +92,7 @@
    {
      "cell_type": "code",
      "execution_count": null,
+      "language": "python",
      "metadata": {
        "cellView": "form",
        "id": "81CC3tL1T_TL"
@@ -98,18 +100,26 @@
      "outputs": [],
      "source": [
        "# @title Setup Google Cloud project\n",
-        "\n",
        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
        "\n",
        "# Import the necessary packages\n",
+        "\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
+        "import importlib\n",
        "import os\n",
-        "import json\n",
        "from datetime import datetime\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "\n",
+        "common_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
+        ")\n",
+        "\n",
+        "models, endpoints = {}, {}\n",
+        "\n",
        "# Get the default cloud project id.\n",
        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
        "\n",
@@ -122,16 +132,16 @@
        "\n",
        "# Cloud Storage bucket for storing the experiment artifacts.\n",
        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
-        "# prefer using your own GCS bucket, please change the value yourself below.\n",
+        "# prefer using your own GCS bucket, change the value yourself below.\n",
        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
-        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
-        "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "\n",
        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
-        "    # Create a unique GCS bucket for this notebook, if not specified by the user\n",
        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\n",
        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
        "else:\n",
+        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
        "    bucket_region = shell_output[0].strip().lower()\n",
        "    if bucket_region != REGION:\n",
@@ -139,127 +149,37 @@
        "            \"Bucket region %s is different from notebook region %s\"\n",
        "            % (bucket_region, REGION)\n",
        "        )\n",
-        "\n",
        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
        "\n",
-        "# Gets the default BUCKET_URI and SERVICE_ACCOUNT if they were not specified by the user.\n",
-        "SERVICE_ACCOUNT = None\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"gemma\")\n",
+        "\n",
+        "\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
+        "\n",
+        "# Gets the default SERVICE_ACCOUNT.\n",
        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
        "\n",
-        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
-        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
-        "# Initialize Vertex AI API.\n",
-        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"staging\")\n",
-        "MODEL_BUCKET = os.path.join(STAGING_BUCKET, \"model\")\n",
-        "print(\"Initializing Vertex AI API.\")\n",
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
+        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
+        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
        "! gcloud config set project $PROJECT_ID\n",
        "\n",
        "# The evaluation docker image.\n",
-        "EVAL_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-lm-evaluation-harness:20240320_0655_RC00\"\n",
-        "\n",
-        "# Define common functions\n",
-        "\n",
-        "\n",
-        "def get_job_name_with_datetime(prefix: str) -> str:\n",
-        "    \"\"\"Gets the job name with date time when triggering training or deployment\n",
-        "    jobs in Vertex AI.\n",
-        "    \"\"\"\n",
-        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
-        "\n",
-        "\n",
-        "def get_quota(project_id: str, region: str, resource_id: str) -> int:\n",
-        "  \"\"\"Returns the quota for a resource in a region. Returns -1 if can not figure out the quota.\"\"\"\n",
-        "  service_endpoint = \"aiplatform.googleapis.com\"\n",
-        "  quota_list_output = !gcloud alpha services quota list --service=$service_endpoint  --consumer=projects/$project_id --filter=\"$service_endpoint/$resource_id\" --format=json\n",
-        "  # Use '.s' on the command output because it is an SList type.\n",
-        "  quota_data = json.loads(quota_list_output.s)\n",
-        "  if len(quota_data) == 0 or \"consumerQuotaLimits\" not in quota_data[0]:\n",
-        "    return -1\n",
-        "  if len(quota_data[0][\"consumerQuotaLimits\"]) == 0 or \"quotaBuckets\" not in quota_data[0][\"consumerQuotaLimits\"][0]:\n",
-        "    return -1\n",
-        "  all_regions_data = quota_data[0][\"consumerQuotaLimits\"][0][\"quotaBuckets\"]\n",
-        "  for region_data in all_regions_data:\n",
-        "    if region_data.get('dimensions') and region_data['dimensions']['region'] == region:\n",
-        "      if 'effectiveLimit' in region_data:\n",
-        "        return int(region_data['effectiveLimit'])\n",
-        "      else:\n",
-        "        return 0\n",
-        "  return -1\n",
-        "\n",
-        "\n",
-        "def get_resource_id(accelerator_type: str, is_for_training: bool) -> str:\n",
-        "  \"\"\"Returns the resource id for a given accelerator type and the use case.\n",
-        "  Args:\n",
-        "    accelerator_type: The accelerator type.\n",
-        "    is_for_training: Whether the resource is used for training. Set false\n",
-        "    for serving use case.\n",
-        "  Returns:\n",
-        "    The resource id.\n",
-        "  \"\"\"\n",
-        "  training_accelerator_map = {\n",
-        "      \"NVIDIA_TESLA_V100\": \"custom_model_training_nvidia_v100_gpus\",\n",
-        "      \"NVIDIA_L4\": \"custom_model_training_nvidia_l4_gpus\",\n",
-        "      \"NVIDIA_TESLA_A100\": \"custom_model_training_nvidia_a100_gpus\",\n",
-        "      \"NVIDIA_TESLA_T4\": \"custom_model_training_nvidia_t4_gpus\",\n",
-        "      \"TPU_V5e\": \"custom_model_training_tpu_v5e\",\n",
-        "      \"TPU_V3\": \"custom_model_training_tpu_v3\",\n",
-        "  }\n",
-        "  serving_accelerator_map = {\n",
-        "      \"NVIDIA_TESLA_V100\": \"custom_model_serving_nvidia_v100_gpus\",\n",
-        "      \"NVIDIA_L4\": \"custom_model_serving_nvidia_l4_gpus\",\n",
-        "      \"NVIDIA_TESLA_A100\": \"custom_model_serving_nvidia_a100_gpus\",\n",
-        "      \"NVIDIA_TESLA_T4\": \"custom_model_serving_nvidia_t4_gpus\",\n",
-        "      \"TPU_V5e\": \"custom_model_serving_tpu_v5e\",\n",
-        "  }\n",
-        "  if is_for_training:\n",
-        "    if accelerator_type in training_accelerator_map:\n",
-        "      return training_accelerator_map[accelerator_type]\n",
-        "    else:\n",
-        "      raise ValueError(\n",
-        "          f\"Could not find accelerator type: {accelerator_type} for training.\"\n",
-        "      )\n",
-        "  else:\n",
-        "    if accelerator_type in serving_accelerator_map:\n",
-        "      return serving_accelerator_map[accelerator_type]\n",
-        "    else:\n",
-        "      raise ValueError(\n",
-        "          f\"Could not find accelerator type: {accelerator_type} for serving.\"\n",
-        "      )\n",
-        "\n",
-        "\n",
-        "def check_quota(project_id:str, region: str, accelerator_type: str,\n",
-        "                accelerator_count: int, is_for_training: bool):\n",
-        "  \"\"\"Checks if the project and the region has the required quota.\"\"\"\n",
-        "  resource_id = get_resource_id(accelerator_type, is_for_training)\n",
-        "  quota = get_quota(project_id, region, resource_id)\n",
-        "  quota_request_instruction = (\"Either use \"\n",
-        "            \"a different region or request additional quota. Follow \"\n",
-        "            \"instructions here \"\n",
-        "            \"https://cloud.google.com/docs/quotas/view-manage#requesting_higher_quota\"\n",
-        "            \" to check quota in a region or request additional quota for \"\n",
-        "            \"your project.\")\n",
-        "  if quota == -1:\n",
-        "    raise ValueError(\n",
-        "            f\"\"\"Quota not found for: {resource_id} in {region}.\n",
-        "            {quota_request_instruction}\"\"\"\n",
-        "        )\n",
-        "  if quota < accelerator_count:\n",
-        "    raise ValueError(\n",
-        "            f\"\"\"Quota not enough for {resource_id} in {region}:\n",
-        "            {quota} < {accelerator_count}.\n",
-        "            {quota_request_instruction}\"\"\"\n",
-        "        )"
+        "EVAL_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-lm-evaluation-harness:20240320_0655_RC00\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
+      "language": "python",
      "metadata": {
        "cellView": "form",
        "id": "pNHMbjr0UjrK"
@@ -285,12 +205,14 @@
        "# Setup evaluation job.\n",
        "# @markdown Set the base model id.\n",
        "base_model_id = \"google/gemma-1.1-2b-it\"  # @param[\"google/gemma-2b\", \"google/gemma-2b-it\", \"google/gemma-7b\", \"google/gemma-7b-it\", \"google/gemma-1.1-2b-it\", \"google/gemma-1.1-7b-it\"] {isTemplate:true}\n",
-        "job_name = get_job_name_with_datetime(prefix=\"gemma-eval\")\n",
+        "job_name = common_util.get_job_name_with_datetime(prefix=\"gemma-eval\")\n",
        "eval_output_dir = os.path.join(MODEL_BUCKET, job_name)\n",
        "eval_output_dir_gcsfuse = eval_output_dir.replace(\"gs://\", \"/gcs/\")\n",
        "\n",
        "# @markdown Set the accelerator type.\n",
-        "accelerator_type = \"NVIDIA_L4\"  # @param[\"NVIDIA_TESLA_V100\", \"NVIDIA_L4\", \"NVIDIA_TESLA_A100\"]\n",
+        "accelerator_type = (\n",
+        "    \"NVIDIA_L4\"  # @param[\"NVIDIA_TESLA_V100\", \"NVIDIA_L4\", \"NVIDIA_TESLA_A100\"]\n",
+        ")\n",
        "\n",
        "# @markdown To evaluate a PEFT-finetuned model, enter the PEFT output directory to the LoRA adapter below.\n",
        "# @markdown Otherwise, leave it empty.\n",
@@ -313,11 +235,13 @@
        "\n",
        "replica_count = 1\n",
        "\n",
-        "check_quota(project_id=PROJECT_ID,\n",
-        "            region=REGION,\n",
-        "            accelerator_type=accelerator_type,\n",
-        "            accelerator_count=accelerator_count,\n",
-        "            is_for_training=True)\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=True,\n",
+        ")\n",
        "\n",
        "# Prepare evaluation command that runs the evaluation harness.\n",
        "# Set `trust_remote_code = True` because evaluating the model requires\n",
@@ -409,6 +333,17 @@
        "print(f\"Evaluation result:\\n{result_formatted}\")"
      ]
    },
+    {
+      "cell_type": "markdown",
+      "execution_count": null,
+      "metadata": {
+        "id": "unjukbcjEBOd"
+      },
+      "outputs": [],
+      "source": [
+        "## Clean up resources"
+      ]
+    },
    {
      "cell_type": "code",
      "execution_count": null,
@@ -418,7 +353,6 @@
      },
      "outputs": [],
      "source": [
-        "# @title Clean up resources\n",
        "# Delete evaluation job.\n",
        "\n",
        "delete_bucket = False  # @param {type:\"boolean\"}\n",
@@ -3,6 +3,7 @@
    {
      "cell_type": "code",
      "execution_count": null,
+      "language": "python",
      "metadata": {
        "id": "7d9bbf86da5e"
      },
@@ -25,6 +26,7 @@
    },
    {
      "cell_type": "markdown",
+      "language": "markdown",
      "metadata": {
        "id": "99c1c3fc2ca5"
      },
@@ -87,6 +89,7 @@
    {
      "cell_type": "code",
      "execution_count": null,
+      "language": "python",
      "metadata": {
        "cellView": "form",
        "id": "855d6b96f291"
@@ -94,50 +97,76 @@
      "outputs": [],
      "source": [
        "# @title Setup Google Cloud project\n",
-        "\n",
        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
-        "# @markdown **[Optional]** Set the GCS BUCKET_URI to store the experiment artifacts, if you want to use your own bucket. **If not set, a unique GCS bucket will be created automatically on your behalf**.\n",
+        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
        "\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
+        "import importlib\n",
        "import os\n",
-        "import uuid\n",
        "from datetime import datetime\n",
        "from typing import Tuple\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "\n",
+        "common_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
+        ")\n",
+        "\n",
+        "models, endpoints = {}, {}\n",
+        "\n",
        "# Get the default cloud project id.\n",
        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
        "\n",
        "# Get the default region for launching jobs.\n",
        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
        "\n",
+        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
+        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
+        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
+        "\n",
        "# Cloud Storage bucket for storing the experiment artifacts.\n",
        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
-        "# prefer using your own GCS bucket, please change the value yourself below.\n",
+        "# prefer using your own GCS bucket, change the value yourself below.\n",
        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
-        "if not BUCKET_URI.strip() or BUCKET_URI == \"gs://\":\n",
-        "    # Create a unique GCS bucket for this notebook if not specified\n",
-        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
+        "\n",
+        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
+        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\n",
        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
-        "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "else:\n",
+        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
+        "    bucket_region = shell_output[0].strip().lower()\n",
+        "    if bucket_region != REGION:\n",
+        "        raise ValueError(\n",
+        "            \"Bucket region %s is different from notebook region %s\"\n",
+        "            % (bucket_region, REGION)\n",
+        "        )\n",
+        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"gemma\")\n",
+        "\n",
+        "\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
        "\n",
        "# Gets the default SERVICE_ACCOUNT.\n",
        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
-        "\n",
        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
-        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
+        "\n",
        "\n",
        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
-        "# Enable Vertex AI and Cloud Compute APIs.\n",
        "! gcloud config set project $PROJECT_ID\n",
-        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
        "\n",
        "# @markdown ## Access Gemma Models\n",
        "# @markdown For GPU based finetuning and serving, choose between accessing Gemma models on [Hugging Face](https://huggingface.co/)\n",
@@ -173,10 +202,6 @@
        "\n",
        "VERTEX_AI_MODEL_GARDEN_GEMMA = \"\"  # @param {type:\"string\", isTemplate:true}\n",
        "\n",
-        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
-        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"gemma\")\n",
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
-        "\n",
        "if LOAD_MODEL_FROM == \"Google Cloud\":\n",
        "    assert (\n",
        "        VERTEX_AI_MODEL_GARDEN_GEMMA\n",
@@ -203,7 +228,7 @@
        "# and Hex-LLM serving.\n",
        "KERAS_TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/jax-keras-train-tpu:20240422_0939_RC00\"\n",
        "KERAS_MODEL_CONVERSION_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/jax-keras-model-conversion:20240422_0949_RC00\"\n",
-        "HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:20240220_0936_RC01\"\n",
+        "HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:deploy\"\n",
        "conversion_job = None\n",
        "\n",
        "# @markdown *--- Or ---*\n",
@@ -220,34 +245,31 @@
        "# @markdown ---\n",
        "\n",
        "\n",
-        "def get_job_name_with_datetime(prefix: str) -> str:\n",
-        "    \"\"\"Gets the job name with date time when triggering training or deployment\n",
-        "    jobs in Vertex AI.\n",
-        "    \"\"\"\n",
-        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
-        "\n",
-        "\n",
        "def deploy_model_vllm(\n",
        "    model_name: str,\n",
-        "    base_model_id: str,\n",
        "    model_id: str,\n",
        "    service_account: str,\n",
-        "    machine_type: str = \"g2-standard-12\",\n",
+        "    base_model_id: str = None,\n",
+        "    machine_type: str = \"g2-standard-8\",\n",
        "    accelerator_type: str = \"NVIDIA_L4\",\n",
        "    accelerator_count: int = 1,\n",
-        "    max_model_len: int = 8192,\n",
-        "    dtype: str = \"bfloat16\",\n",
+        "    gpu_memory_utilization: float = 0.9,\n",
+        "    max_model_len: int = 4096,\n",
+        "    dtype: str = \"auto\",\n",
        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
-        "    \"\"\"Deploys models with vLLM on GPU in Vertex AI.\"\"\"\n",
+        "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
        "\n",
+        "    if not base_model_id:\n",
+        "        base_model_id = model_id\n",
+        "\n",
        "    vllm_args = [\n",
        "        \"--host=0.0.0.0\",\n",
        "        \"--port=7080\",\n",
        "        f\"--model={model_id}\",\n",
        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
        "        \"--swap-space=16\",\n",
-        "        \"--gpu-memory-utilization=0.95\",\n",
+        "        f\"--gpu-memory-utilization={gpu_memory_utilization}\",\n",
        "        f\"--max-model-len={max_model_len}\",\n",
        "        f\"--dtype={dtype}\",\n",
        "        \"--disable-log-stats\",\n",
@@ -257,8 +279,12 @@
        "        \"MODEL_ID\": base_model_id,\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
-        "    if HF_TOKEN:\n",
-        "        env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
+        "\n",
+        "    try:\n",
+        "        if HF_TOKEN:\n",
+        "            env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
+        "    except:\n",
+        "        pass\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
@@ -272,7 +298,9 @@
        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
        "        serving_container_deployment_timeout=7200,\n",
        "    )\n",
-        "\n",
+        "    print(\n",
+        "        f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
+        "    )\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        machine_type=machine_type,\n",
@@ -281,38 +309,52 @@
        "        deploy_request_timeout=1800,\n",
        "        service_account=service_account,\n",
        "    )\n",
+        "    print(\"endpoint_name:\", endpoint.name)\n",
+        "\n",
        "    return model, endpoint\n",
        "\n",
        "\n",
        "def deploy_model_hexllm(\n",
        "    model_name: str,\n",
-        "    base_model_id: str,\n",
        "    model_id: str,\n",
        "    service_account: str,\n",
+        "    base_model_id: str = None,\n",
+        "    tensor_parallel_size: int = 1,\n",
        "    machine_type: str = \"ct5lp-hightpu-1t\",\n",
-        "    max_num_batched_tokens: int = 11264,\n",
-        "    tokens_pad_multiple: int = 1024,\n",
-        "    seqs_pad_multiple: int = 32,\n",
+        "    hbm_utilization_factor: float = 0.6,\n",
+        "    max_running_seqs: int = 256,\n",
+        "    endpoint_id: str = \"\",\n",
+        "    min_replica_count: int = 1,\n",
+        "    max_replica_count: int = 1,\n",
        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
        "    \"\"\"Deploys models with Hex-LLM on TPU in Vertex AI.\"\"\"\n",
-        "    endpoint = aiplatform.Endpoint.create(\n",
-        "        display_name=f\"{model_name}-endpoint\",\n",
-        "        location=TPU_DEPLOYMENT_REGION,\n",
-        "    )\n",
+        "    if endpoint_id:\n",
+        "        aip_endpoint_name = (\n",
+        "            f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_id}\"\n",
+        "        )\n",
+        "        endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "    else:\n",
+        "        endpoint = aiplatform.Endpoint.create(\n",
+        "            display_name=f\"{model_name}-endpoint\",\n",
+        "            location=TPU_DEPLOYMENT_REGION,\n",
+        "        )\n",
+        "\n",
+        "    if not base_model_id:\n",
+        "        base_model_id = model_id\n",
+        "\n",
+        "    if not tensor_parallel_size:\n",
+        "        tensor_parallel_size = int(machine_type[-2])\n",
        "\n",
-        "    num_tpu_chips = int(machine_type[-2])\n",
        "    hexllm_args = [\n",
        "        \"--host=0.0.0.0\",\n",
        "        \"--port=7080\",\n",
        "        \"--log_level=INFO\",\n",
        "        f\"--model={model_id}\",\n",
-        "        f\"--tensor_parallel_size={num_tpu_chips}\",\n",
-        "        \"--num_nodes=1\",\n",
-        "        \"--use_ray\",\n",
-        "        \"--batch_mode=continuous\",\n",
-        "        f\"--max_num_batched_tokens={max_num_batched_tokens}\",\n",
-        "        f\"--tokens_pad_multiple={tokens_pad_multiple}\",\n",
-        "        f\"--seqs_pad_multiple={seqs_pad_multiple}\",\n",
+        "        f\"--tensor_parallel_size={tensor_parallel_size}\",\n",
+        "        \"--enable_jit\",\n",
+        "        \"--load_format=auto\",\n",
+        "        f\"--hbm_utilization_factor={hbm_utilization_factor}\",\n",
+        "        f\"--max_running_seqs={max_running_seqs}\",\n",
        "    ]\n",
        "\n",
        "    env_vars = {\n",
@@ -322,14 +364,17 @@
        "        \"RAY_USAGE_STATS_ENABLED\": \"0\",\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
-        "    if KAGGLE_USERNAME and KAGGLE_KEY:\n",
-        "        env_vars[\"KAGGLE_USERNAME\"] = KAGGLE_USERNAME\n",
-        "        env_vars[\"KAGGLE_KEY\"] = KAGGLE_KEY\n",
+        "\n",
+        "    try:\n",
+        "        if HF_TOKEN:\n",
+        "            env_vars.update({\"HF_TOKEN\": HF_TOKEN})\n",
+        "    except:\n",
+        "        pass\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
        "        serving_container_image_uri=HEXLLM_DOCKER_URI,\n",
-        "        serving_container_command=[\"python\", \"-m\", \"hex_llm.entrypoints.api_server\"],\n",
+        "        serving_container_command=[\"python\", \"-m\", \"hex_llm.server.api_server\"],\n",
        "        serving_container_args=hexllm_args,\n",
        "        serving_container_ports=[7080],\n",
        "        serving_container_predict_route=\"/generate\",\n",
@@ -345,6 +390,8 @@
        "        machine_type=machine_type,\n",
        "        deploy_request_timeout=1800,\n",
        "        service_account=service_account,\n",
+        "        min_replica_count=min_replica_count,\n",
+        "        max_replica_count=max_replica_count,\n",
        "    )\n",
        "    return model, endpoint"
      ]
@@ -484,8 +531,16 @@
        "\n",
        "    replica_count = 1\n",
        "\n",
+        "    common_util.check_quota(\n",
+        "        project_id=PROJECT_ID,\n",
+        "        region=REGION,\n",
+        "        accelerator_type=ACCELERATOR_TYPE,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        is_for_training=True,\n",
+        "    )\n",
+        "\n",
        "    # Setup training job.\n",
-        "    job_name = get_job_name_with_datetime(\"gemma-lora-train\")\n",
+        "    job_name = common_util.get_job_name_with_datetime(\"gemma-lora-train\")\n",
        "\n",
        "    # Pass training arguments and launch job.\n",
        "    train_job = aiplatform.CustomContainerTrainingJob(\n",
@@ -494,12 +549,12 @@
        "    )\n",
        "\n",
        "    # Create a GCS folder to store the LORA adapter.\n",
-        "    lora_adapter_dir = get_job_name_with_datetime(\"gemma-lora-adapter\")\n",
+        "    lora_adapter_dir = common_util.get_job_name_with_datetime(\"gemma-lora-adapter\")\n",
        "    lora_output_dir = os.path.join(STAGING_BUCKET, lora_adapter_dir)\n",
        "\n",
        "    # Create a GCS folder to store the merged model with the base model and the\n",
        "    # finetuned LORA adapter.\n",
-        "    merged_model_dir = get_job_name_with_datetime(\"gemma-merged-model\")\n",
+        "    merged_model_dir = common_util.get_job_name_with_datetime(\"gemma-merged-model\")\n",
        "    merged_model_output_dir = os.path.join(STAGING_BUCKET, merged_model_dir)\n",
        "\n",
        "    train_job.run(\n",
@@ -560,11 +615,19 @@
        "    accelerator_type = \"NVIDIA_L4\"\n",
        "    accelerator_count = 1\n",
        "\n",
+        "    common_util.check_quota(\n",
+        "        project_id=PROJECT_ID,\n",
+        "        region=REGION,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        is_for_training=False,\n",
+        "    )\n",
+        "\n",
        "    # Note that a larger max_model_len will require more GPU memory.\n",
        "    max_model_len = 2048\n",
        "\n",
        "    model, endpoint = deploy_model_vllm(\n",
-        "        model_name=get_job_name_with_datetime(prefix=\"gemma-vllm-serve\"),\n",
+        "        model_name=common_util.get_job_name_with_datetime(prefix=\"gemma-vllm-serve\"),\n",
        "        base_model_id=f\"google/{MODEL_ID}\",\n",
        "        model_id=merged_model_output_dir,\n",
        "        service_account=SERVICE_ACCOUNT,\n",
@@ -586,6 +649,8 @@
      "outputs": [],
      "source": [
        "# @title Predict\n",
+        "\n",
+        "\n",
        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts.\n",
        "\n",
        "# @markdown Here we use an example from the [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) to show the finetuning outcome:\n",
@@ -791,8 +856,16 @@
        "\n",
        "    replica_count = 1\n",
        "\n",
+        "    common_util.check_quota(\n",
+        "        project_id=PROJECT_ID,\n",
+        "        region=REGION,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        is_for_training=True,\n",
+        "    )\n",
+        "\n",
        "    # Setup training job.\n",
-        "    job_name = get_job_name_with_datetime(\"gemma-keras-lora-train\")\n",
+        "    job_name = common_util.get_job_name_with_datetime(\"gemma-keras-lora-train\")\n",
        "\n",
        "    # Pass training arguments and launch job.\n",
        "    train_job = aiplatform.CustomContainerTrainingJob(\n",
@@ -882,7 +955,7 @@
        "    replica_count = 1\n",
        "\n",
        "    # Setup training job.\n",
-        "    job_name = get_job_name_with_datetime(\"gemma-keras-model-conversion\")\n",
+        "    job_name = common_util.get_job_name_with_datetime(\"gemma-keras-model-conversion\")\n",
        "\n",
        "    # Pass training arguments and launch job.\n",
        "    conversion_job = aiplatform.CustomContainerTrainingJob(\n",
@@ -933,9 +1006,23 @@
        "    if \"2b\" in KAGGLE_MODEL_ID:\n",
        "        # Sets ct5lp-hightpu-1t (1 TPU chip) to deploy Gemma 2B models.\n",
        "        machine_type = \"ct5lp-hightpu-1t\"\n",
+        "        accelerator_type = \"TPU_V5e\"\n",
+        "        # Note: 1 TPU-V5e chip has only 1 core.\n",
+        "        accelerator_count = 1\n",
        "    else:\n",
        "        # Sets ct5lp-hightpu-4t (4 TPU chips) to deploy Gemma 7B models.\n",
        "        machine_type = \"ct5lp-hightpu-4t\"\n",
+        "        accelerator_type = \"TPU_V5e\"\n",
+        "        # Note: 1 TPU-V5e chip has only 1 core.\n",
+        "        accelerator_count = 4\n",
+        "\n",
+        "    common_util.check_quota(\n",
+        "        project_id=PROJECT_ID,\n",
+        "        region=REGION,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        is_for_training=False,\n",
+        "    )\n",
        "\n",
        "    # Note that a larger max_num_batched_tokens will require more TPU memory.\n",
        "    max_num_batched_tokens = 11264\n",
@@ -948,7 +1035,7 @@
        "\n",
        "    print(\"Using model from: \", output_folder)\n",
        "    model, endpoint = deploy_model_hexllm(\n",
-        "        model_name=get_job_name_with_datetime(prefix=\"gemma-serve-hexllm\"),\n",
+        "        model_name=common_util.get_job_name_with_datetime(prefix=\"gemma-serve-hexllm\"),\n",
        "        base_model_id=f\"google/{KAGGLE_MODEL_ID}\",\n",
        "        model_id=output_folder,\n",
        "        service_account=SERVICE_ACCOUNT,\n",
@@ -963,6 +1050,7 @@
    {
      "cell_type": "code",
      "execution_count": null,
+      "language": "python",
      "metadata": {
        "cellView": "form",
        "id": "lgJPw3e9yL5a"
@@ -1047,14 +1135,19 @@
        "if conversion_job:\n",
        "    conversion_job.delete()\n",
        "\n",
+        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
+        "# @markdown  and avoid unnecessary continouous charges that may incur.\n",
+        "\n",
        "# Undeploy model and delete endpoint.\n",
-        "endpoint.delete(force=True)\n",
+        "for endpoint in endpoints.values():\n",
+        "    endpoint.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
-        "model.delete()\n",
+        "for model in models.values():\n",
+        "    model.delete()\n",
        "\n",
-        "# Delete Cloud Storage bucket that was created.\n",
-        "if BUCKET_URI == f\"gs://{PROJECT_ID}-tmp-{now}\":\n",
+        "delete_bucket = False  # @param {type:\"boolean\"}\n",
+        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
@@ -356,7 +356,7 @@
        "    model = aiplatform.Model.upload(\n",
        "        display_name=deploy_model_name,\n",
        "        serving_container_image_uri=PREDICTION_CONTAINER_URI,\n",
-        "        serving_container_ports=[8501],\n",
+        "        serving_container_ports=[8080],\n",
        "        serving_container_predict_route=\"/predict\",\n",
        "        serving_container_health_route=\"/ping\",\n",
        "        serving_container_environment_variables=serving_env,\n",
@@ -0,0 +1,503 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7d9bbf86da5e"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2024 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "99c1c3fc2ca5"
+      },
+      "source": [
+        "# Vertex AI Model Garden - Llama Guard\n",
+        "\n",
+        "<table><tbody><tr>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_llama_guard_deployment.ipynb\">\n",
+        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_llama_guard_deployment.ipynb\">\n",
+        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3de7470326a2"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates downloading and deploying prebuilt [Llama Guard models](https://huggingface.co/meta-llama) with [vLLM](https://github.com/vllm-project/vllm) on GPU, and demonstrates using the Llama Guard model to safeguard LLM inputs and outputs with the Vertex Llama 3.1 API service.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Download and deploy prebuilt Llama Guard models with [vLLM](https://github.com/vllm-project/vllm) on GPU\n",
+        "- Use the Llama Guard models to safeguard LLM inputs and outputs with the Vertex Llama 3.1 API service\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "264c07757582"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YXFGIp1l-qtT"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
+        "\n",
+        "# Import the necessary packages\n",
+        "\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
+        "import importlib\n",
+        "import os\n",
+        "import re\n",
+        "from datetime import datetime\n",
+        "from typing import Tuple\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "common_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
+        ")\n",
+        "\n",
+        "models, endpoints = {}, {}\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "\n",
+        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
+        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
+        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
+        "\n",
+        "# Cloud Storage bucket for storing the experiment artifacts.\n",
+        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
+        "# prefer using your own GCS bucket, change the value yourself below.\n",
+        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
+        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
+        "\n",
+        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
+        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\n",
+        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
+        "else:\n",
+        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
+        "    bucket_region = shell_output[0].strip().lower()\n",
+        "    if bucket_region != REGION:\n",
+        "        raise ValueError(\n",
+        "            \"Bucket region %s is different from notebook region %s\"\n",
+        "            % (bucket_region, REGION)\n",
+        "        )\n",
+        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "MODEL_BUCKET = BUCKET_URI\n",
+        "\n",
+        "\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
+        "\n",
+        "# Gets the default SERVICE_ACCOUNT.\n",
+        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
+        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
+        "\n",
+        "\n",
+        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
+        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "\n",
+        "# @markdown # Access Llama Guard models on Vertex AI\n",
+        "# @markdown The original models from Meta are converted into the Hugging Face format for serving in Vertex AI.\n",
+        "# @markdown Accept the model agreement to access the models:\n",
+        "# @markdown 1. Open the [Llama Guard model card](https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama-guard) from [Vertex AI Model Garden](https://cloud.google.com/model-garden).\n",
+        "# @markdown 2. Review and accept the agreement in the pop-up window on the model card page. If you have previously accepted the model agreement, there will not be a pop-up window on the model card page and this step is not needed.\n",
+        "# @markdown 3. After accepting the agreement, a `gs://` URI containing Llama Guard pretrained and finetuned models will be shared.\n",
+        "# @markdown 4. Paste the URI in the `VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD` field below.\n",
+        "# @markdown 5. The Llama Guard models will be copied into `BUCKET_URI`.\n",
+        "\n",
+        "\n",
+        "VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD = \"\"  # @param {type:\"string\", isTemplate:true}\n",
+        "assert (\n",
+        "    VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD\n",
+        "), \"Please click the agreement in Vertex AI Model Garden at https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama-guard, and get the GCS path of Llama Guard model artifacts.\"\n",
+        "parsed_gcs_url = re.search(\"gs://.*?(?=[ ]|$)\", VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD)\n",
+        "if parsed_gcs_url:\n",
+        "    VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD = parsed_gcs_url.group()\n",
+        "assert VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD.startswith(\n",
+        "    \"gs://\"\n",
+        "), \"VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD is expected to be a GCS URI and must start with `gs://`.\"\n",
+        "print(\n",
+        "    \"Copying Llama Guard model artifacts from\",\n",
+        "    VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD,\n",
+        "    \"to \",\n",
+        "    MODEL_BUCKET,\n",
+        ")\n",
+        "\n",
+        "! gsutil -m cp -R $VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD $MODEL_BUCKET\n",
+        "\n",
+        "# The pre-built serving docker images.\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240726_1329_RC00\"\n",
+        "\n",
+        "\n",
+        "def deploy_model_vllm(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    service_account: str,\n",
+        "    base_model_id: str = None,\n",
+        "    machine_type: str = \"g2-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 1,\n",
+        "    gpu_memory_utilization: float = 0.9,\n",
+        "    max_model_len: int = 4096,\n",
+        "    dtype: str = \"auto\",\n",
+        ") -> 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",
+        "    if not base_model_id:\n",
+        "        base_model_id = model_id\n",
+        "\n",
+        "    vllm_args = [\n",
+        "        \"python\",\n",
+        "        \"-m\",\n",
+        "        \"vllm.entrypoints.api_server\",\n",
+        "        \"--host=0.0.0.0\",\n",
+        "        \"--port=7080\",\n",
+        "        f\"--model={model_id}\",\n",
+        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
+        "        \"--swap-space=16\",\n",
+        "        f\"--gpu-memory-utilization={gpu_memory_utilization}\",\n",
+        "        f\"--max-model-len={max_model_len}\",\n",
+        "        f\"--dtype={dtype}\",\n",
+        "        \"--disable-log-stats\",\n",
+        "        \"--enforce-eager\",\n",
+        "        \"--disable-custom-all-reduce\",\n",
+        "        \"--enable-chunked-prefill\",\n",
+        "        \"--max-num-seqs=12\",\n",
+        "    ]\n",
+        "\n",
+        "    env_vars = {\n",
+        "        \"MODEL_ID\": base_model_id,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "\n",
+        "    try:\n",
+        "        if HF_TOKEN:\n",
+        "            env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
+        "    except:\n",
+        "        pass\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
+        "        serving_container_args=vllm_args,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_predict_route=\"/generate\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        serving_container_environment_variables=env_vars,\n",
+        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
+        "        serving_container_deployment_timeout=7200,\n",
+        "    )\n",
+        "    print(\n",
+        "        f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
+        "    )\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=service_account,\n",
+        "    )\n",
+        "    print(\"endpoint_name:\", endpoint.name)\n",
+        "\n",
+        "    return model, endpoint"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "z-XybZjtgF9M"
+      },
+      "source": [
+        "## Deploy Llama Guard with vLLM on GPU"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "E8OiHHNNE_wj"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "\n",
+        "# @markdown This section uploads prebuilt Llama Guard models to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model.\n",
+        "\n",
+        "# @markdown NVIDIA_L4 GPUs are used for demonstration. The serving efficiency of L4 GPUs is inferior to that of A100 GPUs, but L4 GPUs are nevertheless good serving solutions if you do not have A100 quota.\n",
+        "\n",
+        "# @markdown Set the model to deploy.\n",
+        "\n",
+        "MODEL_ID = \"Llama-Guard-3-8B\"  # @param [\"Llama-Guard-3-8B\"] {allow-input: true, isTemplate: true}\n",
+        "model_id = os.path.join(MODEL_BUCKET, MODEL_ID)\n",
+        "\n",
+        "# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
+        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_A100\"]\n",
+        "\n",
+        "if accelerator_type == \"NVIDIA_L4\":\n",
+        "    machine_type = \"g2-standard-12\"\n",
+        "    accelerator_count = 1\n",
+        "elif accelerator_type == \"NVIDIA_TESLA_A100\":\n",
+        "    machine_type = \"a2-highgpu-1g\"\n",
+        "    accelerator_count = 1\n",
+        "else:\n",
+        "    raise ValueError(\n",
+        "        f\"Recommended GPU setting not found for: {accelerator_type} and {MODEL_ID}.\"\n",
+        "    )\n",
+        "\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=False,\n",
+        ")\n",
+        "\n",
+        "gpu_memory_utilization = 0.9\n",
+        "max_model_len = 32768  # Maximum context length.\n",
+        "\n",
+        "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=MODEL_ID),\n",
+        "    model_id=model_id,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    gpu_memory_utilization=gpu_memory_utilization,\n",
+        "    max_model_len=max_model_len,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "192a021iB_DE"
+      },
+      "source": [
+        "## Use the Llama Guard models to safeguard LLM inputs and outputs with the Vertex Llama 3.1 API service\n",
+        "\n",
+        "We use [meta-llama/Llama-Guard-3-8B](https://huggingface.co/meta-llama/Llama-Guard-3-8B) to safeguard input and output conversations with the [Llama 3.1 405B Instruct model API service on Vertex](https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama3-405b-instruct-maas).\n",
+        "\n",
+        "Llama Guard 3 builds on the capabilities introduced with Llama Guard 2, adding three new categories, Defamation, Elections and Code Interpreter Abuse. Additionally this model is multilingual and a new prompt format is introduced, making Llama Guard 3’s prompt format consistent with Llama 3+ Instruct models.\n",
+        "\n",
+        "This section references [LlamaGuard.ipynb](https://colab.research.google.com/drive/16s0tlCSEDtczjPzdIK3jq0Le5LlnSYGf?usp=sharing) from [https://huggingface.co/meta-llama/LlamaGuard-7b](https://huggingface.co/meta-llama/LlamaGuard-7b)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "fHC7INgjB_DF"
+      },
+      "outputs": [],
+      "source": [
+        "!pip install --upgrade --quiet openai"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ajjcGNzhB_DF"
+      },
+      "outputs": [],
+      "source": [
+        "import google.auth\n",
+        "import openai\n",
+        "\n",
+        "# @markdown Set up the Llama 3.1 405B Instruct model API service.\n",
+        "\n",
+        "# Programmatically get an access token\n",
+        "creds, _ = google.auth.default(\n",
+        "    scopes=[\"https://www.googleapis.com/auth/cloud-platform\"]\n",
+        ")\n",
+        "auth_req = google.auth.transport.requests.Request()\n",
+        "creds.refresh(auth_req)\n",
+        "# Note: the credential lives for 1 hour by default (https://cloud.google.com/docs/authentication/token-types#at-lifetime); after expiration, it must be refreshed.\n",
+        "\n",
+        "client = openai.OpenAI(\n",
+        "    base_url=f\"https://us-central1-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{REGION}/endpoints/openapi\",\n",
+        "    api_key=creds.token,\n",
+        ")\n",
+        "LLAMA3_405B_INSTRUCT = \"meta/llama3-405b-instruct-maas\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NvSfBcUUB_DF"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown Define input message in conversation and get output message from model.\n",
+        "\n",
+        "message_role = \"user\"  # @param {type: \"string\"}\n",
+        "message_content = \"What is a car?\"  # @param {type: \"string\"}\n",
+        "\n",
+        "messages = [\n",
+        "    {\n",
+        "        \"role\": message_role,\n",
+        "        \"content\": message_content,\n",
+        "    }\n",
+        "]\n",
+        "print(\"Conversation [turn 1]:\", messages)\n",
+        "\n",
+        "response = client.chat.completions.create(\n",
+        "    model=LLAMA3_405B_INSTRUCT,\n",
+        "    messages=messages,\n",
+        ")\n",
+        "print(\"Response:\", response)\n",
+        "\n",
+        "messages.append(\n",
+        "    {\n",
+        "        \"role\": response.choices[0].message.role,\n",
+        "        \"content\": response.choices[0].message.content,\n",
+        "    }\n",
+        ")\n",
+        "print(\"Conversation [turn 2]:\", messages)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Y7-ym3GlB_DG"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown Use Llama Guard to classify the conversation: safe versus unsafe.\n",
+        "# @markdown Classification is performed on the last turn of the conversation.\n",
+        "# @markdown If the content is safe, the model will return `safe`. If the content is unsafe, the model will return `unsafe` and additionally the list of offending categories as a comma-separated list in a new line.\n",
+        "# @markdown Set `\"@requestFormat\": \"chatCompletions\"` to use the OpenAI chat completions format.\n",
+        "\n",
+        "instances = [\n",
+        "    {\n",
+        "        \"messages\": messages,\n",
+        "        \"@requestFormat\": \"chatCompletions\",\n",
+        "    },\n",
+        "]\n",
+        "response = endpoints[\"vllm_gpu\"].predict(instances=instances)\n",
+        "\n",
+        "prediction = response.predictions[0]\n",
+        "print(prediction)\n",
+        "print(\"Llama Guard prediction:\", prediction[\"choices\"][0][\"message\"][\"content\"])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "956x4r7rsrza"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "911406c1561e"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Delete the models and endpoints\n",
+        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
+        "# @markdown  and avoid unnecessary continouous charges that may incur.\n",
+        "\n",
+        "# Undeploy model and delete endpoint.\n",
+        "for endpoint in endpoints.values():\n",
+        "    endpoint.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "for model in models.values():\n",
+        "    model.delete()\n",
+        "\n",
+        "delete_bucket = False  # @param {type:\"boolean\"}\n",
+        "if delete_bucket:\n",
+        "    ! gsutil -m rm -r $BUCKET_URI"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_llama_guard_deployment.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
@@ -0,0 +1,749 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ur8xi4C7S06n"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2024 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "99c1c3fc2ca5"
+      },
+      "source": [
+        "# Vertex AI Model Garden - MaMMUT\n",
+        "\n",
+        "<table><tbody><tr>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_mammut.ipynb\">\n",
+        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_mammut.ipynb\">\n",
+        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tvgnzT1CKxrO"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates deploying MaMMUT to a Vertex AI Endpoint and making online predictions."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "d975e698c9a4"
+      },
+      "source": [
+        "### Objective\n",
+        "\n",
+        "- Deploy MaMMUT to a Vertex AI Endpoint.\n",
+        "- Make predictions to the endpoint including:\n",
+        "  - Answering questions about a given image.\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aed92deeb4a0"
+      },
+      "source": [
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "BF1j6f9HApxa"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "DU0WWEDqWJLy"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "# @markdown ### Prerequisites\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
+        "\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "! pip install -q gradio==4.21.0\n",
+        "\n",
+        "import importlib\n",
+        "import os\n",
+        "from datetime import datetime\n",
+        "from typing import Tuple\n",
+        "\n",
+        "import gradio as gr\n",
+        "import numpy as np\n",
+        "from google.cloud import aiplatform\n",
+        "from PIL import Image\n",
+        "\n",
+        "common_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
+        ")\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "\n",
+        "# Cloud Storage bucket for storing the experiment artifacts.\n",
+        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
+        "# prefer using your own GCS bucket, change the value yourself below.\n",
+        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
+        "BUCKET_URI = \"gs://\"  # @param {type: \"string\"}\n",
+        "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "\n",
+        "# Create a unique GCS bucket for this notebook, if not specified by the user.\n",
+        "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
+        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\n",
+        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
+        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "else:\n",
+        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
+        "    bucket_region = shell_output[0].strip().lower()\n",
+        "    if bucket_region != REGION:\n",
+        "        raise ValueError(\n",
+        "            \"Bucket region %s is different from notebook region %s\"\n",
+        "            % (bucket_region, REGION)\n",
+        "        )\n",
+        "\n",
+        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "! gcloud services enable language.googleapis.com\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "\n",
+        "# Set up default SERVICE_ACCOUNT\n",
+        "SERVICE_ACCOUNT = None\n",
+        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
+        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "\n",
+        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
+        "\n",
+        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
+        "\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-gpu.nightly:latest\"\n",
+        ")\n",
+        "\n",
+        "models, endpoints = {}, {}\n",
+        "\n",
+        "\n",
+        "def resize_image(image: Image.Image, new_width: int = 512) -> Image.Image:\n",
+        "    width, height = image.size\n",
+        "    new_height = int(height * new_width / width)\n",
+        "    new_image = image.resize((new_width, new_height))\n",
+        "    return new_image\n",
+        "\n",
+        "\n",
+        "def load_image(image_url):\n",
+        "    if image_url.startswith(\"gs://\"):\n",
+        "        local_image_path = \"./images/test_image.jpg\"\n",
+        "        common_util.download_gcs_file_to_local(image_url, local_image_path)\n",
+        "        image = common_util.load_img(local_image_path)\n",
+        "    else:\n",
+        "        image = common_util.download_image(image_url)\n",
+        "    return image\n",
+        "\n",
+        "\n",
+        "def deploy_mammut(\n",
+        "    task: str, machine_type: str, accelerator_type: str, accelerator_count: int\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploy the model to a Vertex endpoint for prediction.\"\"\"\n",
+        "    serving_env = {\n",
+        "        \"MODEL_ID\": \"mammut\",\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "\n",
+        "    if task == \"vqa\":\n",
+        "        model_dir = \"gs://vertex-model-garden-public-us/mammut/vqa\"\n",
+        "    else:\n",
+        "        model_dir = \"gs://vertex-model-garden-public-us/mammut/retrieval\"\n",
+        "\n",
+        "    upload_job_name = common_util.get_job_name_with_datetime(\n",
+        "        prefix=\"mammut-\" + task + \"-upload\"\n",
+        "    )\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=upload_job_name,\n",
+        "        artifact_uri=model_dir,\n",
+        "        serving_container_image_uri=OPTIMIZED_TF_RUNTIME_IMAGE_URI,\n",
+        "        serving_container_args=[],\n",
+        "        location=REGION,\n",
+        "        serving_container_environment_variables=serving_env,\n",
+        "    )\n",
+        "\n",
+        "    print(\"The uploaded model name is: \", upload_job_name)\n",
+        "\n",
+        "    deploy_model_name = common_util.get_job_name_with_datetime(\n",
+        "        prefix=\"mammut-\" + task + \"-deploy\"\n",
+        "    )\n",
+        "\n",
+        "    common_util.check_quota(\n",
+        "        project_id=PROJECT_ID,\n",
+        "        region=REGION,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        is_for_training=False,\n",
+        "    )\n",
+        "\n",
+        "    endpoint = model.deploy(\n",
+        "        deployed_model_display_name=deploy_model_name,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        min_replica_count=1,\n",
+        "        max_replica_count=1,\n",
+        "    )\n",
+        "\n",
+        "    print(\"The deployed job name is: \", deploy_model_name)\n",
+        "\n",
+        "    endpoint_id = endpoint.name\n",
+        "    print(\"endpoint id is: \", endpoint_id)\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "def predict(\n",
+        "    endpoint: aiplatform.Endpoint,\n",
+        "    image: Image.Image,\n",
+        "    prompt: str,\n",
+        "    new_width: int = 1000,\n",
+        "):\n",
+        "    \"\"\"Generates predictions based on the input image and text using an Endpoint.\"\"\"\n",
+        "    # Resize and convert image to base64 string.\n",
+        "    resized_image = resize_image(image, new_width)\n",
+        "    instances = [\n",
+        "        {\n",
+        "            \"image_bytes\": {\"b64\": common_util.image_to_base64(resized_image)},\n",
+        "            \"text\": prompt,\n",
+        "        }\n",
+        "    ]\n",
+        "\n",
+        "    response = endpoint.predict(instances=instances)\n",
+        "    return response.predictions[0]"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "iILhhP3TfO8B"
+      },
+      "source": [
+        "## Run online prediction\n",
+        "\n",
+        "Run online prediction with the TF SavedModel."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Lgq-FTe7wak_"
+      },
+      "source": [
+        "### Visual Question Answering"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "74yqis5ufO8B"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "# @markdown Upload TF SavedModel and deploy it to an endpoint for prediction. This step takes around 15 minutes to finish.\n",
+        "\n",
+        "# @markdown Select the accelerator type to use to deploy the model:\n",
+        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_V100\"]\n",
+        "# @markdown If you want to use other accelerator types not listed above, then check other Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute. You may need to manually set the `machine_type`, `accelerator_type`, and `accelerator_count` in the code by clicking `Show code` first.\n",
+        "\n",
+        "accelerator_count = 1\n",
+        "if accelerator_type == \"NVIDIA_L4\":\n",
+        "    machine_type = \"g2-standard-4\"\n",
+        "elif accelerator_type == \"NVIDIA_TESLA_V100\":\n",
+        "    machine_type = \"n1-standard-4\"\n",
+        "else:\n",
+        "    raise ValueError(\n",
+        "        f\"Recommended machine settings not found for: {accelerator_type}. To use another another accelerator, edit this code block to pass in an appropriate `machine_type`, `accelerator_type`, and `accelerator_count` to the deploy_model function by clicking `Show Code` and then modifying the code.\"\n",
+        "    )\n",
+        "\n",
+        "models[\"vqa\"], endpoints[\"vqa\"] = deploy_mammut(\n",
+        "    \"vqa\", machine_type, accelerator_type, accelerator_count\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "qxj4Xv_DhHXj"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict\n",
+        "# @markdown Use the deployed MaMMUT model to answer questions about a given image.\n",
+        "\n",
+        "# @markdown **Note: The first prediction can take up to 2 minutes due to one time JIT compilation of the model. This may cause a timeout error below. If you get a timeout error, then wait for 2 minutes and run the prediction again. You will not get the timeout error after that.**\n",
+        "\n",
+        "# @markdown This section uses images from [pexels.com](https://www.pexels.com/) for demoing purposes. All the images have the following license: https://www.pexels.com/license/.\n",
+        "\n",
+        "# @markdown Images will be resized to a width of 1000 pixels by default since requests made to a Vertex Endpoint are limited to 1.500MB.\n",
+        "\n",
+        "# @markdown ![](https://images.pexels.com/photos/4012966/pexels-photo-4012966.jpeg?w=1260&h=750)\n",
+        "\n",
+        "# @markdown This can be either a Cloud Storage path (gs://\\<image-path\\>) or a public url (http://\\<image-path\\>)\n",
+        "image_url = \"https://images.pexels.com/photos/4012966/pexels-photo-4012966.jpeg\"  # @param {type:\"string\"}\n",
+        "\n",
+        "image = load_image(image_url)\n",
+        "display(image)\n",
+        "\n",
+        "# @markdown You may leave question prompts empty and they will be ignored.\n",
+        "question_prompt_1 = \"Is there a person in the image?\"  # @param {type: \"string\"}\n",
+        "question_prompt_2 = \"What is the person doing in the image?\"  # @param {type: \"string\"}\n",
+        "question_prompt_3 = \"What's the color of the cup?\"  # @param {type: \"string\"}\n",
+        "question_prompt_4 = \"How many laptops are in the image?\"  # @param {type: \"string\"}\n",
+        "\n",
+        "questions_list = [\n",
+        "    question_prompt_1,\n",
+        "    question_prompt_2,\n",
+        "    question_prompt_3,\n",
+        "    question_prompt_4,\n",
+        "]\n",
+        "questions_list = [question for question in questions_list if question]\n",
+        "\n",
+        "for question in questions_list:\n",
+        "    answer = predict(endpoints[\"vqa\"], image, question)\n",
+        "    print(f\"Question: {question}\")\n",
+        "    print(f\"Answer: {answer}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "VNLXg2BxZli7"
+      },
+      "source": [
+        "#### Creating a webpage playground with Gradio"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "wRbUhytaY3Nt"
+      },
+      "outputs": [],
+      "source": [
+        "# @title How to use\n",
+        "\n",
+        "# @markdown **Prerequisites**\n",
+        "# @markdown -  Before you can upload an image to make a prediction, you need to select a Vertex prediction endpoint serving MaMMUT\n",
+        "# @markdown from the endpoint dropdown list that has been deployed in the current project and region.\n",
+        "# @markdown -  If no models have been deployed, you can create a new Vertex prediction\n",
+        "# @markdown endpoint by clicking \"Deploy to Vertex\" in the playground or running the `Deploy` cell above.\n",
+        "# @markdown   * New model deployment takes approximately 15 minutes. You can check the progress at [Vertex Online Prediction](https://console.cloud.google.com/vertex-ai/online-prediction/endpoints).\n",
+        "\n",
+        "# @markdown **How to use**\n",
+        "\n",
+        "# @markdown Just run this cell and a link to the playground formatted as `https://####.gradio.live` will be outputted.\n",
+        "# @markdown This link will take you to the playground in a separate browser tab.\n",
+        "\n",
+        "\n",
+        "def list_mammut_endpoints() -> list[str]:\n",
+        "    \"\"\"Returns all valid prediction endpoints for in the project and region.\"\"\"\n",
+        "    # Gets all the valid endpoints in the project and region.\n",
+        "    endpoints = aiplatform.Endpoint.list(order_by=\"create_time desc\")\n",
+        "    # Filters out the endpoints which do not have a deployed model, and the endpoint is for image generation\n",
+        "    endpoints = list(\n",
+        "        filter(\n",
+        "            lambda endpoint: endpoint.traffic_split\n",
+        "            and \"mammut-vqa\" in endpoint.display_name.lower(),\n",
+        "            endpoints,\n",
+        "        )\n",
+        "    )\n",
+        "\n",
+        "    endpoint_names = list(\n",
+        "        map(\n",
+        "            lambda endpoint: f\"{endpoint.name} - {endpoint.display_name[:40]}\",\n",
+        "            endpoints,\n",
+        "        )\n",
+        "    )\n",
+        "\n",
+        "    if not endpoint_names:\n",
+        "        gr.Warning(\"No prediction endpoints were found. Create an Endpoint first.\")\n",
+        "\n",
+        "    return endpoint_names\n",
+        "\n",
+        "\n",
+        "def deploy_model_handler() -> None:\n",
+        "    gr.Info(\"Starting model deployment.\")\n",
+        "    model, endpoint = deploy_mammut(\"vqa\", \"g2-standard-4\", \"NVIDIA_L4\", 1)\n",
+        "    gr.Info(f\"Deploying model ID: {model.name}, endpoint ID: {endpoint.name}\")\n",
+        "\n",
+        "\n",
+        "def get_endpoint(endpoint_name: str) -> aiplatform.Endpoint:\n",
+        "    \"\"\"Returns a Vertex endpoint for the given endpoint_name.\"\"\"\n",
+        "    endpoint_id = endpoint_name.split(\" - \")[0]\n",
+        "    endpoint = aiplatform.Endpoint(\n",
+        "        f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_id}\"\n",
+        "    )\n",
+        "    return endpoint\n",
+        "\n",
+        "\n",
+        "def predict_handler(\n",
+        "    endpoint_name: str,\n",
+        "    image: Image.Image,\n",
+        "    prompt: str,\n",
+        ") -> str:\n",
+        "    if not endpoint_name:\n",
+        "        raise gr.Error(\"Select (or deploy) a model first!\")\n",
+        "    if not image:\n",
+        "        raise gr.Error(\"You must upload an image!\")\n",
+        "    endpoint = get_endpoint(endpoint_name)\n",
+        "    return predict(endpoint, image, prompt)\n",
+        "\n",
+        "\n",
+        "tip_text = r\"\"\"\n",
+        "<b> Tips: </b>\n",
+        "1. Select a Vertex prediction endpoint with a deployed MaMMUT model or click `Deploy to Vertex` to deploy MaMMUT to Vertex.\n",
+        "2. New model deployment takes approximately 15 minutes. You can check the progress by examining the output section of the notebook cell that runs this playground. Your endpoint will show up at [Vertex Online Prediction](https://console.cloud.google.com/vertex-ai/online-prediction/endpoints) once the deployment is done.\n",
+        "3. After the model deployment is complete, click `Refresh Endpoints list` to view the new endpoint in the dropdown list.\n",
+        "4. Note: The first prediction can take up to 2 minutes due to one time JIT compilation of the model. This may cause a timeout error below. If you get a timeout error, then wait for 2 minutes and run the prediction again. You will not get the timeout error after that.\n",
+        "\"\"\"\n",
+        "\n",
+        "css = \"\"\"\n",
+        ".gradio-container {\n",
+        "  width: 85% !important\n",
+        "}\n",
+        "\"\"\"\n",
+        "with gr.Blocks(\n",
+        "    css=css, theme=gr.themes.Default(primary_hue=\"orange\", secondary_hue=\"blue\")\n",
+        ") as demo:\n",
+        "    gr.Markdown(\"# Model Garden Playground for MaMMUT\")\n",
+        "    with gr.Row(equal_height=True):\n",
+        "        with gr.Column(scale=3):\n",
+        "            gr.Markdown(tip_text)\n",
+        "        with gr.Column(scale=2):\n",
+        "            with gr.Row():\n",
+        "                endpoint_name = gr.Dropdown(\n",
+        "                    scale=7,\n",
+        "                    label=\"Select a model previously deployed on Vertex (Click inside the input box below)\",\n",
+        "                    choices=list_mammut_endpoints(),\n",
+        "                    value=None,\n",
+        "                )\n",
+        "                refresh_button = gr.Button(\n",
+        "                    \"Refresh Endpoints list\",\n",
+        "                    scale=1,\n",
+        "                    variant=\"primary\",\n",
+        "                    min_width=10,\n",
+        "                )\n",
+        "            with gr.Row():\n",
+        "                deploy_model_button = gr.Button(\n",
+        "                    \"Deploy a new model\",\n",
+        "                    scale=1,\n",
+        "                    variant=\"primary\",\n",
+        "                    min_width=10,\n",
+        "                )\n",
+        "    with gr.Row(equal_height=True):\n",
+        "        with gr.Column(scale=1):\n",
+        "            image_input = gr.Image(\n",
+        "                show_label=True,\n",
+        "                type=\"pil\",\n",
+        "                label=\"Upload\",\n",
+        "                visible=True,\n",
+        "                height=400,\n",
+        "            )\n",
+        "            with gr.Group():\n",
+        "                text_input_box = gr.Textbox(label=\"Question\", lines=1)\n",
+        "                submit_button = gr.Button(\"Answer\", variant=\"primary\")\n",
+        "        with gr.Column(scale=1):\n",
+        "            image_output = gr.Image(label=\"Image response:\", visible=False)\n",
+        "            text_output = gr.Textbox(label=\"Text response:\")\n",
+        "\n",
+        "    refresh_button.click(\n",
+        "        fn=lambda: gr.update(choices=list_mammut_endpoints()),\n",
+        "        outputs=[endpoint_name],\n",
+        "    )\n",
+        "    deploy_model_button.click(\n",
+        "        deploy_model_handler,\n",
+        "        outputs=[],\n",
+        "    )\n",
+        "    submit_button.click(\n",
+        "        fn=predict_handler,\n",
+        "        inputs=[\n",
+        "            endpoint_name,\n",
+        "            image_input,\n",
+        "            text_input_box,\n",
+        "        ],\n",
+        "        outputs=[text_output],\n",
+        "    )\n",
+        "show_debug_logs = True  # @param {type: \"boolean\"}\n",
+        "demo.queue()\n",
+        "demo.launch(\n",
+        "    share=True, inline=False, inbrowser=True, debug=show_debug_logs, show_error=True\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "97jEsBSfwm-3"
+      },
+      "source": [
+        "### Retrieval and Multimodal Embeddings"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "oJemp5enwl2g"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "# @markdown Upload TF SavedModel and deploy it to an endpoint for prediction. This step takes around 15 minutes to finish.\n",
+        "\n",
+        "# @markdown Select the accelerator type to use to deploy the model:\n",
+        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_V100\"]\n",
+        "# @markdown If you want to use other accelerator types not listed above, then check other Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute. You may need to manually set the `machine_type`, `accelerator_type`, and `accelerator_count` in the code by clicking `Show code` first.\n",
+        "\n",
+        "accelerator_count = 1\n",
+        "if accelerator_type == \"NVIDIA_L4\":\n",
+        "    machine_type = \"g2-standard-4\"\n",
+        "elif accelerator_type == \"NVIDIA_TESLA_V100\":\n",
+        "    machine_type = \"n1-standard-4\"\n",
+        "else:\n",
+        "    raise ValueError(\n",
+        "        f\"Recommended machine settings not found for: {accelerator_type}. To use another another accelerator, edit this code block to pass in an appropriate `machine_type`, `accelerator_type`, and `accelerator_count` to the deploy_model function by clicking `Show Code` and then modifying the code.\"\n",
+        "    )\n",
+        "\n",
+        "models[\"retrieval\"], endpoints[\"retrieval\"] = deploy_mammut(\n",
+        "    \"retrieval\", machine_type, accelerator_type, accelerator_count\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "fO0-YrOZi34k"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Image-Text Retrieval\n",
+        "# @markdown Given an image, use the deployed MaMMUT model to find the best matching text out of multiple options based on similarity scores of their embeddings. This example uses only 5 text options but you can modify the example to retrieve over as many text examples as needed.\n",
+        "\n",
+        "# @markdown **Note: The first prediction can take up to 2 minutes due to one time JIT compilation of the model. This may cause a timeout error below. If you get a timeout error, then wait for 2 minutes and run the prediction again. You will not get the timeout error after that.**\n",
+        "\n",
+        "# @markdown This section uses images from [pexels.com](https://www.pexels.com/) for demoing purposes. All the images have the following license: https://www.pexels.com/license/.\n",
+        "\n",
+        "# @markdown Images will be resized to a width of 1000 pixels by default since requests made to a Vertex Endpoint are limited to 1.500MB.\n",
+        "\n",
+        "# @markdown ![](https://images.pexels.com/photos/20427316/pexels-photo-20427316/free-photo-of-a-moped-parked-in-front-of-a-blue-door.jpeg?auto=compress&cs=tinysrgb&w=630&h=375&dpr=2)\n",
+        "\n",
+        "# @markdown This can be either a Cloud Storage path (gs://\\<image-path\\>) or a public url (http://\\<image-path\\>)\n",
+        "image_url = \"https://images.pexels.com/photos/20427316/pexels-photo-20427316/free-photo-of-a-moped-parked-in-front-of-a-blue-door.jpeg?auto=compress&cs=tinysrgb&w=630&h=375&dpr=2\"  # @param {type:\"string\"}\n",
+        "\n",
+        "image = load_image(image_url)\n",
+        "display(image)\n",
+        "\n",
+        "text_1 = \"A tennis player about to serve.\"  # @param {type: \"string\"}\n",
+        "text_2 = \"Green broccolis and fruit in a bowl on a table.\"  # @param {type: \"string\"}\n",
+        "text_3 = \"A moped parked in front of a blue door.\"  # @param {type: \"string\"}\n",
+        "text_4 = \"A baguette with some ham in it.\"  # @param {type: \"string\"}\n",
+        "text_5 = \"Three zebras in a dry land with some bush.\"  # @param {type: \"string\"}\n",
+        "\n",
+        "text_list = [text_1, text_2, text_3, text_4, text_5]\n",
+        "text_list = [text for text in text_list if text]\n",
+        "\n",
+        "image_embeddings = []\n",
+        "text_embeddings = []\n",
+        "for text in text_list:\n",
+        "    prediction = predict(endpoints[\"retrieval\"], image, text)\n",
+        "    image_embeddings.append(np.array(prediction[\"normalized_image_embedding\"]))\n",
+        "    text_embeddings.append(np.array(prediction[\"normalized_text_embedding\"]))\n",
+        "\n",
+        "# predictions = predict(endpoint, image, text_list)\n",
+        "# image_embeddings = [np.array(prediction[\"normalized_image_embedding\"]) for prediction in predictions]\n",
+        "\n",
+        "# text_embeddings = [np.array(prediction[\"normalized_text_embedding\"]) for prediction in predictions]\n",
+        "\n",
+        "image_embeddings = np.vstack(image_embeddings)\n",
+        "text_embeddings = np.vstack(text_embeddings)\n",
+        "similarity = np.matmul(image_embeddings, text_embeddings.T)\n",
+        "argmax_indices = np.argmax(similarity, axis=-1)\n",
+        "argmax = argmax_indices[0]\n",
+        "print(f\"The text that's most similar to the image is: {text_list[argmax]}\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "o2modPjMFYrm"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Text-Image Retrieval\n",
+        "# @markdown Given a text description, use the deployed MaMMUT model to find the best matching image out of multiple options based on similarity scores of their embeddings. This example uses only 5 image options but you can modify the example to retrieve over as many image examples as needed.\n",
+        "\n",
+        "# @markdown This section uses images from [pexels.com](https://www.pexels.com/) for demoing purposes. All the images have the following license: https://www.pexels.com/license/.\n",
+        "\n",
+        "# @markdown Images will be resized to a width of 1000 pixels by default since requests made to a Vertex Endpoint are limited to 1.500MB.\n",
+        "\n",
+        "text = \"A view of the city with many red roofs.\"  # @param {type: \"string\"}\n",
+        "\n",
+        "# @markdown Image URLs can be either a Cloud Storage path (gs://\\<image-path\\>) or a public url (http://\\<image-path\\>)\n",
+        "\n",
+        "image_url_1 = \"https://images.pexels.com/photos/4012966/pexels-photo-4012966.jpeg?w=1260&h=750\"  # @param {type:\"string\"}\n",
+        "# @markdown ![](https://images.pexels.com/photos/4012966/pexels-photo-4012966.jpeg?w=1260&h=750)\n",
+        "\n",
+        "image_url_2 = \"https://images.pexels.com/photos/24427993/pexels-photo-24427993/free-photo-of-a-group-of-cherries-arranged-in-a-row-on-a-white-wall.jpeg?auto=compress&cs=tinysrgb&w=1260&h=750&dpr=1\"  # @param {type:\"string\"}\n",
+        "# @markdown ![](https://images.pexels.com/photos/24427993/pexels-photo-24427993/free-photo-of-a-group-of-cherries-arranged-in-a-row-on-a-white-wall.jpeg?auto=compress&cs=tinysrgb&w=1260&h=750&dpr=1)\n",
+        "\n",
+        "image_url_3 = \"https://images.pexels.com/photos/20427316/pexels-photo-20427316/free-photo-of-a-moped-parked-in-front-of-a-blue-door.jpeg?auto=compress&cs=tinysrgb&w=630&h=375&dpr=2\"  # @param {type:\"string\"}\n",
+        "# @markdown ![](https://images.pexels.com/photos/20427316/pexels-photo-20427316/free-photo-of-a-moped-parked-in-front-of-a-blue-door.jpeg?auto=compress&cs=tinysrgb&w=630&h=375&dpr=2)\n",
+        "\n",
+        "image_url_4 = \"https://images.pexels.com/photos/18592009/pexels-photo-18592009/free-photo-of-a-view-of-the-city-with-many-red-roofs.jpeg?auto=compress&cs=tinysrgb&w=1260&h=750&dpr=1\"  # @param {type:\"string\"}\n",
+        "# @markdown ![](https://images.pexels.com/photos/18592009/pexels-photo-18592009/free-photo-of-a-view-of-the-city-with-many-red-roofs.jpeg?auto=compress&cs=tinysrgb&w=1260&h=750&dpr=1)\n",
+        "\n",
+        "image_url_5 = \"https://images.pexels.com/photos/1006293/pexels-photo-1006293.jpeg?auto=compress&cs=tinysrgb&w=1260&h=750&dpr=2\"  # @param {type:\"string\"}\n",
+        "# @markdown ![](https://images.pexels.com/photos/1006293/pexels-photo-1006293.jpeg?auto=compress&cs=tinysrgb&w=1260&h=750&dpr=2)\n",
+        "\n",
+        "image_url_list = [image_url_1, image_url_2, image_url_3, image_url_4, image_url_5]\n",
+        "image_url_list = [image_url for image_url in image_url_list if image_url]\n",
+        "\n",
+        "images = [load_image(image_url) for image_url in image_url_list]\n",
+        "\n",
+        "text_embeddings = []\n",
+        "image_embeddings = []\n",
+        "for image in images:\n",
+        "    prediction = predict(retrieval_endpoint, image, text)\n",
+        "    image_embeddings.append(np.array(prediction[\"normalized_image_embedding\"]))\n",
+        "    text_embeddings.append(np.array(prediction[\"normalized_text_embedding\"]))\n",
+        "\n",
+        "image_embeddings = np.vstack(image_embeddings)\n",
+        "text_embeddings = np.vstack(text_embeddings)\n",
+        "similarity = np.matmul(image_embeddings, text_embeddings.T)\n",
+        "argmax_indices = np.argmax(similarity, axis=0)\n",
+        "argmax = argmax_indices[0]\n",
+        "print(f\"The image that's most similar to the text is: {image_url_list[argmax]}\")\n",
+        "display(images[argmax])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3aD4PW3d1bG5"
+      },
+      "source": [
+        "## Clean up resources\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "sx_vKniMq9ZX"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Run\n",
+        "\n",
+        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
+        "# @markdown  and avoid unnecessary continuous charges that may incur.\n",
+        "\n",
+        "# Delete endpoint resource.\n",
+        "for endpoint in endpoints.values():\n",
+        "    endpoint.delete(force=True)\n",
+        "\n",
+        "# Delete model resource.\n",
+        "for model in models.values():\n",
+        "    model.delete()\n",
+        "\n",
+        "# Delete Cloud Storage objects that were created.\n",
+        "delete_bucket = False  # @param {type:\"boolean\"}\n",
+        "if delete_bucket:\n",
+        "    ! gsutil -m rm -r $BUCKET_URI"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_mammut.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
@@ -265,7 +265,7 @@
        "# Prediction constants.\n",
        "# You can adjust accelerator types and machine types to get faster predictions.\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_PORT = 8080\n",
        "PREDICTION_ACCELERATOR_COUNT = 1\n",
        "PREDICTION_ACCELERATOR_TYPE = \"NVIDIA_TESLA_T4\"\n",
        "PREDICTION_MACHINE_TYPE = \"n1-standard-4\"\n",
@@ -250,7 +250,7 @@
        "# Prediction constants.\n",
        "# You can adjust accelerator types and machine types to get faster predictions.\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_PORT = 8080\n",
        "PREDICTION_ACCELERATOR_COUNT = 1\n",
        "PREDICTION_ACCELERATOR_TYPE = \"NVIDIA_TESLA_T4\"\n",
        "PREDICTION_MACHINE_TYPE = \"n1-standard-4\"\n",
@@ -160,7 +160,7 @@
        "\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "# Define common functions.\n",
@@ -3,6 +3,7 @@
    {
      "cell_type": "code",
      "execution_count": null,
+      "language": "python",
      "metadata": {
        "id": "7d9bbf86da5e"
      },
@@ -94,55 +95,77 @@
      "outputs": [],
      "source": [
        "# @title Setup Google Cloud project\n",
-        "\n",
        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
-        "# @markdown **[Optional]** Set the GCS BUCKET_URI to store the experiment artifacts, if you want to use your own bucket. **If not set, a unique GCS bucket will be created automatically on your behalf**.\n",
+        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
        "\n",
-        "import json\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
+        "import importlib\n",
        "import os\n",
        "import sys\n",
-        "import uuid\n",
        "from datetime import datetime\n",
        "from typing import Tuple\n",
        "\n",
        "from google.cloud import aiplatform, language\n",
        "\n",
+        "common_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
+        ")\n",
+        "\n",
+        "models, endpoints = {}, {}\n",
+        "\n",
        "# Get the default cloud project id.\n",
        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
        "\n",
        "# Get the default region for launching jobs.\n",
        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
        "\n",
+        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
+        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
+        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
+        "\n",
        "# Cloud Storage bucket for storing the experiment artifacts.\n",
        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
-        "# prefer using your own GCS bucket, please change the value yourself below.\n",
+        "# prefer using your own GCS bucket, change the value yourself below.\n",
        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
-        "if not BUCKET_URI.strip() or BUCKET_URI == \"gs://\":\n",
-        "    # Create a unique GCS bucket for this notebook if not specified\n",
-        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
+        "\n",
+        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
+        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\n",
        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
-        "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "else:\n",
+        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
+        "    bucket_region = shell_output[0].strip().lower()\n",
+        "    if bucket_region != REGION:\n",
+        "        raise ValueError(\n",
+        "            \"Bucket region %s is different from notebook region %s\"\n",
+        "            % (bucket_region, REGION)\n",
+        "        )\n",
+        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"gemma\")\n",
+        "\n",
+        "\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
        "\n",
        "# Gets the default SERVICE_ACCOUNT.\n",
        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
-        "\n",
        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
-        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
+        "\n",
        "\n",
        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
-        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
-        "\n",
-        "# Enable Vertex AI, Cloud Compute, and Cloud Language APIs.\n",
        "! gcloud config set project $PROJECT_ID\n",
-        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com language.googleapis.com\n",
        "\n",
        "# @markdown ## Access Gemma Models\n",
        "\n",
@@ -161,13 +184,6 @@
        "    auth.authenticate_user(project_id=PROJECT_ID)\n",
        "\n",
        "\n",
-        "def get_job_name_with_datetime(prefix: str) -> str:\n",
-        "    \"\"\"Gets the job name with date time when triggering training or deployment\n",
-        "    jobs in Vertex AI.\n",
-        "    \"\"\"\n",
-        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
-        "\n",
-        "\n",
        "def moderate_text(text: str) -> language.ModerateTextResponse:\n",
        "    \"\"\"Calls Vertex AI APIs to analyze text moderations.\"\"\"\n",
        "    client = language.LanguageServiceClient()\n",
@@ -236,91 +252,7 @@
        "        deploy_request_timeout=1800,\n",
        "        service_account=service_account,\n",
        "    )\n",
-        "    return model, endpoint\n",
-        "\n",
-        "\n",
-        "def get_quota(project_id: str, region: str, resource_id: str) -> int:\n",
-        "  \"\"\"Returns the quota for a resource in a region. Returns -1 if can not figure out the quota.\"\"\"\n",
-        "  service_endpoint = \"aiplatform.googleapis.com\"\n",
-        "  quota_list_output = !gcloud alpha services quota list --service=$service_endpoint  --consumer=projects/$project_id --filter=\"$service_endpoint/$resource_id\" --format=json\n",
-        "  # Use '.s' on the command output because it is an SList type.\n",
-        "  quota_data = json.loads(quota_list_output.s)\n",
-        "  if len(quota_data) == 0 or \"consumerQuotaLimits\" not in quota_data[0]:\n",
-        "    return -1\n",
-        "  if len(quota_data[0][\"consumerQuotaLimits\"]) == 0 or \"quotaBuckets\" not in quota_data[0][\"consumerQuotaLimits\"][0]:\n",
-        "    return -1\n",
-        "  all_regions_data = quota_data[0][\"consumerQuotaLimits\"][0][\"quotaBuckets\"]\n",
-        "  for region_data in all_regions_data:\n",
-        "    if region_data.get('dimensions') and region_data['dimensions']['region'] == region:\n",
-        "      if 'effectiveLimit' in region_data:\n",
-        "        return int(region_data['effectiveLimit'])\n",
-        "      else:\n",
-        "        return 0\n",
-        "  return -1\n",
-        "\n",
-        "\n",
-        "def get_resource_id(accelerator_type: str, is_for_training: bool) -> str:\n",
-        "  \"\"\"Returns the resource id for a given accelerator type and the use case.\n",
-        "  Args:\n",
-        "    accelerator_type: The accelerator type.\n",
-        "    is_for_training: Whether the resource is used for training. Set false\n",
-        "    for serving use case.\n",
-        "  Returns:\n",
-        "    The resource id.\n",
-        "  \"\"\"\n",
-        "  training_accelerator_map = {\n",
-        "      \"NVIDIA_TESLA_V100\": \"custom_model_training_nvidia_v100_gpus\",\n",
-        "      \"NVIDIA_L4\": \"custom_model_training_nvidia_l4_gpus\",\n",
-        "      \"NVIDIA_TESLA_A100\": \"custom_model_training_nvidia_a100_gpus\",\n",
-        "      \"NVIDIA_TESLA_T4\": \"custom_model_training_nvidia_t4_gpus\",\n",
-        "      \"TPU_V5e\": \"custom_model_training_tpu_v5e\",\n",
-        "      \"TPU_V3\": \"custom_model_training_tpu_v3\",\n",
-        "  }\n",
-        "  serving_accelerator_map = {\n",
-        "      \"NVIDIA_TESLA_V100\": \"custom_model_serving_nvidia_v100_gpus\",\n",
-        "      \"NVIDIA_L4\": \"custom_model_serving_nvidia_l4_gpus\",\n",
-        "      \"NVIDIA_TESLA_A100\": \"custom_model_serving_nvidia_a100_gpus\",\n",
-        "      \"NVIDIA_TESLA_T4\": \"custom_model_serving_nvidia_t4_gpus\",\n",
-        "      \"TPU_V5e\": \"custom_model_serving_tpu_v5e\",\n",
-        "  }\n",
-        "  if is_for_training:\n",
-        "    if accelerator_type in training_accelerator_map:\n",
-        "      return training_accelerator_map[accelerator_type]\n",
-        "    else:\n",
-        "      raise ValueError(\n",
-        "          f\"Could not find accelerator type: {accelerator_type} for training.\"\n",
-        "      )\n",
-        "  else:\n",
-        "    if accelerator_type in serving_accelerator_map:\n",
-        "      return serving_accelerator_map[accelerator_type]\n",
-        "    else:\n",
-        "      raise ValueError(\n",
-        "          f\"Could not find accelerator type: {accelerator_type} for serving.\"\n",
-        "      )\n",
-        "\n",
-        "\n",
-        "def check_quota(project_id:str, region: str, accelerator_type: str,\n",
-        "                accelerator_count: int, is_for_training: bool):\n",
-        "  \"\"\"Checks if the project and the region has the required quota.\"\"\"\n",
-        "  resource_id = get_resource_id(accelerator_type, is_for_training)\n",
-        "  quota = get_quota(project_id, region, resource_id)\n",
-        "  quota_request_instruction = (\"Either use \"\n",
-        "            \"a different region or request additional quota. Follow \"\n",
-        "            \"instructions here \"\n",
-        "            \"https://cloud.google.com/docs/quotas/view-manage#requesting_higher_quota\"\n",
-        "            \" to check quota in a region or request additional quota for \"\n",
-        "            \"your project.\")\n",
-        "  if quota == -1:\n",
-        "    raise ValueError(\n",
-        "            f\"\"\"Quota not found for: {resource_id} in {region}.\n",
-        "            {quota_request_instruction}\"\"\"\n",
-        "        )\n",
-        "  if quota < accelerator_count:\n",
-        "    raise ValueError(\n",
-        "            f\"\"\"Quota not enough for {resource_id} in {region}:\n",
-        "            {quota} < {accelerator_count}.\n",
-        "            {quota_request_instruction}\"\"\"\n",
-        "        )"
+        "    return model, endpoint"
      ]
    },
    {
@@ -423,7 +355,10 @@
        "# The Gemma base model.\n",
        "MODEL_ID = \"google/gemma-1.1-2b-it\"  # @param[\"google/gemma-2b\", \"google/gemma-2b-it\", \"google/gemma-7b\", \"google/gemma-7b-it\", \"google/gemma-1.1-2b-it\", \"google/gemma-1.1-7b-it\"] {isTemplate:true}\n",
        "# The accelerator to use.\n",
+        "\n",
+        "# fmt: off\n",
        "ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"  # @param[\"NVIDIA_TESLA_V100\", \"NVIDIA_L4\", \"NVIDIA_TESLA_A100\"]\n",
+        "# fmt: on\n",
        "\n",
        "# Batch size for finetuning.\n",
        "per_device_train_batch_size = 1  # @param{type:\"integer\"}\n",
@@ -464,14 +399,16 @@
        "\n",
        "replica_count = 1\n",
        "\n",
-        "check_quota(project_id=PROJECT_ID,\n",
-        "            region=REGION,\n",
-        "            accelerator_type=accelerator_type,\n",
-        "            accelerator_count=accelerator_count,\n",
-        "            is_for_training=True)\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=ACCELERATOR_TYPE,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=True,\n",
+        ")\n",
        "\n",
        "# Setup training job.\n",
-        "job_name = get_job_name_with_datetime(\"gemma-lora-train\")\n",
+        "job_name = common_util.get_job_name_with_datetime(\"gemma-lora-train\")\n",
        "\n",
        "# Pass training arguments and launch job.\n",
        "train_job = aiplatform.CustomContainerTrainingJob(\n",
@@ -480,12 +417,12 @@
        ")\n",
        "\n",
        "# Create a GCS folder to store the LORA adapter.\n",
-        "lora_adapter_dir = get_job_name_with_datetime(\"gemma-lora-adapter\")\n",
+        "lora_adapter_dir = common_util.get_job_name_with_datetime(\"gemma-lora-adapter\")\n",
        "lora_output_dir = os.path.join(STAGING_BUCKET, lora_adapter_dir)\n",
        "\n",
        "# Create a GCS folder to store the merged model with the base model and the\n",
        "# finetuned LORA adapter.\n",
-        "merged_model_dir = get_job_name_with_datetime(\"gemma-merged-model\")\n",
+        "merged_model_dir = common_util.get_job_name_with_datetime(\"gemma-merged-model\")\n",
        "merged_model_output_dir = os.path.join(STAGING_BUCKET, merged_model_dir)\n",
        "\n",
        "train_job.run(\n",
@@ -542,19 +479,21 @@
        "accelerator_type = \"NVIDIA_L4\"\n",
        "accelerator_count = 1\n",
        "\n",
-        "check_quota(project_id=PROJECT_ID,\n",
-        "            region=REGION,\n",
-        "            accelerator_type=accelerator_type,\n",
-        "            accelerator_count=accelerator_count,\n",
-        "            is_for_training=False)\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=False,\n",
+        ")\n",
        "\n",
        "# Note that larger token counts will require more GPU memory.\n",
        "max_input_length = 512\n",
        "max_total_tokens = 1024\n",
        "max_batch_prefill_tokens = 2048\n",
        "\n",
-        "model, endpoint = deploy_model_tgi(\n",
-        "    model_name=get_job_name_with_datetime(prefix=\"gemma-tgi-serve\"),\n",
+        "models[\"tgi\"], endpoints[\"tgi\"] = deploy_model_tgi(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=\"gemma-tgi-serve\"),\n",
        "    model_id=merged_model_output_dir,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        "    machine_type=machine_type,\n",
@@ -563,8 +502,7 @@
        "    max_input_length=max_input_length,\n",
        "    max_total_tokens=max_total_tokens,\n",
        "    max_batch_prefill_tokens=max_batch_prefill_tokens,\n",
-        ")\n",
-        "print(\"endpoint_name:\", endpoint.name)"
+        ")"
      ]
    },
    {
@@ -621,7 +559,7 @@
        "        },\n",
        "    },\n",
        "]\n",
-        "response = endpoint.predict(instances=instances)\n",
+        "response = endpoints[\"tgi\"].predict(instances=instances)\n",
        "\n",
        "for prediction in response.predictions:\n",
        "    print(prediction)"
@@ -666,15 +604,20 @@
        "# Delete the train job.\n",
        "train_job.delete()\n",
        "\n",
+        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
+        "# @markdown  and avoid unnecessary continouous charges that may incur.\n",
+        "\n",
        "# Undeploy model and delete endpoint.\n",
-        "endpoint.delete(force=True)\n",
+        "for endpoint in endpoints.values():\n",
+        "    endpoint.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
-        "model.delete()\n",
+        "for model in models.values():\n",
+        "    model.delete()\n",
        "\n",
-        "# Delete Cloud Storage bucket that was created.\n",
-        "if BUCKET_URI == f\"gs://{PROJECT_ID}-tmp-{now}\":\n",
-        "    ! gsutil -m rm -r $STAGING_BUCKET"
+        "delete_bucket = False  # @param {type:\"boolean\"}\n",
+        "if delete_bucket:\n",
+        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
  ],
@@ -149,7 +149,7 @@
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240429_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
@@ -122,7 +122,7 @@
        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240429_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
@@ -157,7 +157,7 @@
        "    auth.authenticate_user(project_id=PROJECT_ID)\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "# Define common functions.\n",
@@ -107,14 +107,14 @@
        "\n",
        "The workflow takes the following inputs:\n",
        "\n",
-        "* **preference_dataset**: Required, `str`. Cloud storage path to a human preference dataset used to train a reward model. Download the [example dataset](https://pantheon.corp.google.com/storage/browser/vertex-ai/generative-ai/rlhf/text_small/summarize_from_feedback_tfds/comparisons/train?pageState=(%22StorageObjectListTable%22:(%22f%22:%22%255B%255D%22))&e=13802955&jsmode=O&mods=-ai_platform_fake_service&project=vertex-ai&prefix=&forceOnObjectsSortingFiltering=false), and see its content.\n",
+        "* **preference_dataset**: Required, `str`. Cloud storage path to a human preference dataset used to train a reward model. Download the [example dataset](https://console.cloud.google.com/storage/browser/vertex-ai/generative-ai/rlhf/text_small/summarize_from_feedback_tfds/comparisons/train), and see its content.\n",
        "> The dataset format is `jsonl`. Each line in the file contains one example:\n",
        "> * An `input_text` field containing the prompt.\n",
        "> * `candidate_0`, `candidate_1` fields containing the candidate responses.\n",
        "> * A `choice` field that specifies which candidate is preferred. This field is an integer `0` or `1`. `0` means `candidate_0` is preferred, `1` means `candidate_1` is preferred.\n",
        "> * For example: `{\"input_text\": \"Create a description for Plantation Palms.\", \"candidate_0\": \"Enjoy some fun in the sun at Gulf Shores.\", \"candidate_1\": \"A Tranquil Oasis of Natural Beauty.\", \"choice\": 0}`\n",
        "\n",
-        "* **prompt_dataset**: Required, `str`. Cloud storage path to an unlabled prompt dataset used for reinforcement learning. Download the [example dataset](https://pantheon.corp.google.com/storage/browser/vertex-ai/generative-ai/rlhf/text_small/reddit_tfds/train?pageState=(%22StorageObjectListTable%22:(%22f%22:%22%255B%255D%22))&e=13802955&jsmode=O&mods=-ai_platform_fake_service&project=vertex-ai&prefix=&forceOnObjectsSortingFiltering=false), and see its content.\n",
+        "* **prompt_dataset**: Required, `str`. Cloud storage path to an unlabled prompt dataset used for reinforcement learning. Download the [example dataset](https://console.cloud.google.com/storage/browser/vertex-ai/generative-ai/rlhf/text_small/reddit_tfds/train), and see its content.\n",
        "> The dataset format is `jsonl`. Each line in the file contains one example:\n",
        "> * An `input_text` field containing the prompt.\n",
        "> * For example: `{\"input_text\": \"Create a description for Plantation Palms.\"}`\n",
@@ -0,0 +1,692 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7d9bbf86da5e"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2024 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "99c1c3fc2ca5"
+      },
+      "source": [
+        "# Vertex AI Model Garden - Llama 3.1 (Deployment)\n",
+        "\n",
+        "<table><tbody><tr>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_pytorch_llama3_1_deployment.ipynb\">\n",
+        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_llama3_1_deployment.ipynb\">\n",
+        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3de7470326a2"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates downloading, deploying, and serving prebuilt Llama 3.1 models with [Hex-LLM](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-hex-llm) or [vLLM](https://github.com/vllm-project/vllm).\n",
+        "\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Deploy Llama 3.1 8B with Hex-LLM on TPU.\n",
+        "- Deploy Llama 3.1 70B and 405B with vLLM on GPU, optionally with dynamic LoRA adapters.\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "264c07757582"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "6fe2644d854f"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Accept the model agreement to access the models\n",
+        "\n",
+        "# @markdown 1. Open the [Llama 3.1 model card](https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama3_1) from [Vertex AI Model Garden](https://cloud.google.com/model-garden).\n",
+        "\n",
+        "# @markdown 2. Review and accept the agreement in the pop-up window on the model card page. If you have previously accepted the model agreement, there will not be a pop-up window on the model card page and this step is not needed."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "ax7zWynUDcjk"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Request for quota\n",
+        "\n",
+        "# @markdown By default, the quota for TPU deployment `Custom model serving TPU v5e cores per region` is 4. TPU quota is only available in `us-west1`. You can request for higher TPU quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
+        "\n",
+        "# @markdown By default, the quota for H100 deployment `Custom model serving per region` is 0. You need to request for H100 quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YXFGIp1l-qtT"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "# Import the necessary packages\n",
+        "\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
+        "import importlib\n",
+        "import os\n",
+        "import re\n",
+        "from datetime import datetime\n",
+        "from typing import Tuple\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "common_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
+        ")\n",
+        "\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "\n",
+        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
+        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
+        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
+        "\n",
+        "# Cloud Storage bucket for storing the experiment artifacts.\n",
+        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
+        "# prefer using your own GCS bucket, change the value yourself below.\n",
+        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
+        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
+        "\n",
+        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
+        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\n",
+        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
+        "else:\n",
+        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
+        "    bucket_region = shell_output[0].strip().lower()\n",
+        "    if bucket_region != REGION:\n",
+        "        raise ValueError(\n",
+        "            \"Bucket region %s is different from notebook region %s\"\n",
+        "            % (bucket_region, REGION)\n",
+        "        )\n",
+        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"llama_3_1\")\n",
+        "\n",
+        "\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
+        "\n",
+        "# Gets the default SERVICE_ACCOUNT.\n",
+        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
+        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
+        "\n",
+        "\n",
+        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
+        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "\n",
+        "# @markdown # Access Llama 3.1 models on Vertex AI for serving\n",
+        "# @markdown The original models from Meta are converted into the Hugging Face format for serving in Vertex AI.\n",
+        "# @markdown Accept the model agreement to access the models:\n",
+        "# @markdown 1. Open the [Llama 3.1 model card](https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama3_1) from [Vertex AI Model Garden](https://cloud.google.com/model-garden).\n",
+        "# @markdown 2. Review and accept the agreement in the pop-up window on the model card page. If you have previously accepted the model agreement, there will not be a pop-up window on the model card page and this step is not needed.\n",
+        "# @markdown 3. After accepting the agreement of Llama 3.1, a `gs://` URI containing Llama 3.1 pretrained and finetuned models will be shared.\n",
+        "# @markdown 4. Paste the URI in the `VERTEX_AI_MODEL_GARDEN_LLAMA_3_1` field below.\n",
+        "# @markdown 5. The Llama 3.1 models will be copied into `BUCKET_URI`.\n",
+        "\n",
+        "\n",
+        "VERTEX_AI_MODEL_GARDEN_LLAMA_3_1 = \"\"  # @param {type:\"string\", isTemplate:true}\n",
+        "assert (\n",
+        "    VERTEX_AI_MODEL_GARDEN_LLAMA_3_1\n",
+        "), \"Please click the agreement of Llama 3.1 in Vertex AI Model Garden, and get the GCS path of Llama 3.1 model artifacts.\"\n",
+        "parsed_gcs_url = re.search(\"gs://.*?(?=[ ]|$)\", VERTEX_AI_MODEL_GARDEN_LLAMA_3_1)\n",
+        "if parsed_gcs_url:\n",
+        "    VERTEX_AI_MODEL_GARDEN_LLAMA_3_1 = parsed_gcs_url.group()\n",
+        "assert VERTEX_AI_MODEL_GARDEN_LLAMA_3_1.startswith(\n",
+        "    \"gs://\"\n",
+        "), \"VERTEX_AI_MODEL_GARDEN_LLAMA_3_1 is expected to be a GCS URI and must start with `gs://`.\"\n",
+        "print(\n",
+        "    \"Copying LLaMA3 model artifacts from\",\n",
+        "    VERTEX_AI_MODEL_GARDEN_LLAMA_3_1,\n",
+        "    \"to \",\n",
+        "    MODEL_BUCKET,\n",
+        ")\n",
+        "\n",
+        "! gsutil -m cp -R $VERTEX_AI_MODEL_GARDEN_LLAMA_3_1/* $MODEL_BUCKET\n",
+        "\n",
+        "# The pre-built serving docker images.\n",
+        "HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:llama3.1\"\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240726_1329_RC00\"\n",
+        "\n",
+        "SERVICE_ENDPOINT = \"aiplatform.googleapis.com\"\n",
+        "\n",
+        "\n",
+        "def deploy_model_hexllm(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    service_account: str,\n",
+        "    machine_type: str = \"ct5lp-hightpu-4t\",\n",
+        "    tensor_parallel_size: int = 4,\n",
+        "    hbm_utilization_factor: float = 0.8,\n",
+        "    max_running_seqs: int = 256,\n",
+        "    max_model_len: int = 8192,\n",
+        "    endpoint_id: str = \"\",\n",
+        "    min_replica_count: int = 1,\n",
+        "    max_replica_count: int = 1,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys models with Hex-LLM on TPU in Vertex AI.\"\"\"\n",
+        "    if endpoint_id:\n",
+        "        aip_endpoint_name = (\n",
+        "            f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_id}\"\n",
+        "        )\n",
+        "        endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "    else:\n",
+        "        endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
+        "\n",
+        "    hexllm_args = [\n",
+        "        \"--host=0.0.0.0\",\n",
+        "        \"--port=7080\",\n",
+        "        \"--log_level=INFO\",\n",
+        "        \"--enable_jit\",\n",
+        "        f\"--model={model_id}\",\n",
+        "        \"--load_format=auto\",\n",
+        "        f\"--tensor_parallel_size={tensor_parallel_size}\",\n",
+        "        f\"--hbm_utilization_factor={hbm_utilization_factor}\",\n",
+        "        f\"--max_running_seqs={max_running_seqs}\",\n",
+        "        f\"--max_model_len={max_model_len}\",\n",
+        "        \"--max-num-seqs=12\",\n",
+        "    ]\n",
+        "    hexllm_envs = {\n",
+        "        \"PJRT_DEVICE\": \"TPU\",\n",
+        "        \"RAY_DEDUP_LOGS\": \"0\",\n",
+        "        \"RAY_USAGE_STATS_ENABLED\": \"0\",\n",
+        "        \"MODEL_ID\": model_id,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=HEXLLM_DOCKER_URI,\n",
+        "        serving_container_command=[\"python\", \"-m\", \"hex_llm.server.api_server\"],\n",
+        "        serving_container_args=hexllm_args,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_predict_route=\"/generate\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        serving_container_environment_variables=hexllm_envs,\n",
+        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
+        "        serving_container_deployment_timeout=7200,\n",
+        "    )\n",
+        "\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=service_account,\n",
+        "        min_replica_count=min_replica_count,\n",
+        "        max_replica_count=max_replica_count,\n",
+        "    )\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "def deploy_model_vllm(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    service_account: str,\n",
+        "    machine_type: str = \"g2-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 1,\n",
+        "    gpu_memory_utilization: float = 0.9,\n",
+        "    max_model_len: int = 8192,\n",
+        "    max_loras: int = 1,\n",
+        "    max_cpu_loras: int = 16,\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",
+        "    vllm_args = [\n",
+        "        \"python\",\n",
+        "        \"-m\",\n",
+        "        \"vllm.entrypoints.api_server\",\n",
+        "        \"--host=0.0.0.0\",\n",
+        "        \"--port=7080\",\n",
+        "        f\"--model={model_id}\",\n",
+        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
+        "        \"--swap-space=16\",\n",
+        "        f\"--gpu-memory-utilization={gpu_memory_utilization}\",\n",
+        "        f\"--max-model-len={max_model_len}\",\n",
+        "        \"--enable-lora\",\n",
+        "        \"--disable-custom-all-reduce\",\n",
+        "        f\"--max-loras={max_loras}\",\n",
+        "        f\"--max-cpu-loras={max_cpu_loras}\",\n",
+        "        \"--disable-log-stats\",\n",
+        "    ]\n",
+        "\n",
+        "    env_vars = {\"MODEL_ID\": model_id, \"DEPLOY_SOURCE\": \"notebook\"}\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
+        "        serving_container_args=vllm_args,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_predict_route=\"/generate\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        serving_container_environment_variables=env_vars,\n",
+        "    )\n",
+        "    print(\n",
+        "        f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
+        "    )\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=service_account,\n",
+        "    )\n",
+        "    print(\"endpoint_name:\", endpoint.name)\n",
+        "\n",
+        "    return model, endpoint"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "mAt6NcA5Dcjl"
+      },
+      "source": [
+        "## Deploy prebuilt Llama 3.1 8B with Hex-LLM\n",
+        "\n",
+        "**Hex-LLM** is a **H**igh-**E**fficiency **L**arge **L**anguage **M**odel (LLM) TPU serving solution built with **XLA**, which is being developed by Google Cloud.\n",
+        "\n",
+        "Refer to the \"Request for TPU quota\" section for TPU quota."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9-5obzXZDcjl"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "\n",
+        "# @markdown This section uploads prebuilt Llama 3.1 models to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model.\n",
+        "\n",
+        "# @markdown Select one of the four model variations. More model variants will be supported by Hex-LLM in the future.\n",
+        "MODEL_ID = \"Meta-Llama-3.1-8B\"  # @param [\"Meta-Llama-3.1-8B\", \"Meta-Llama-3.1-8B-Instruct\"] {allow-input: true, isTemplate: true}\n",
+        "model_id = os.path.join(MODEL_BUCKET, MODEL_ID)\n",
+        "\n",
+        "# @markdown Find Vertex AI prediction TPUv5e machine types in\n",
+        "# @markdown https://cloud.google.com/vertex-ai/docs/predictions/use-tpu#deploy_a_model.\n",
+        "\n",
+        "# Sets ct5lp-hightpu-4t (4 TPU chips) to deploy Llama 3.1 8B models.\n",
+        "machine_type = \"ct5lp-hightpu-4t\"\n",
+        "accelerator_type = \"TPU_V5e\"\n",
+        "# Note: 1 TPU V5 chip has only one core.\n",
+        "accelerator_count = 4\n",
+        "\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=False,\n",
+        ")\n",
+        "\n",
+        "# Server parameters.\n",
+        "tensor_parallel_size = accelerator_count\n",
+        "hbm_utilization_factor = 0.8  # Fraction of HBM memory allocated for KV cache after model loading. A larger value improves throughput but gives higher risk of TPU out-of-memory errors with long prompts.\n",
+        "max_running_seqs = 256  # Maximum number of running sequences in a continuous batch.\n",
+        "max_model_len = 8192\n",
+        "\n",
+        "# Endpoint configurations.\n",
+        "min_replica_count = 1\n",
+        "max_replica_count = 1\n",
+        "\n",
+        "model_hexllm, endpoint_hexllm = deploy_model_hexllm(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=\"llama_3_1-hexllm-serve\"),\n",
+        "    model_id=model_id,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    machine_type=machine_type,\n",
+        "    tensor_parallel_size=tensor_parallel_size,\n",
+        "    hbm_utilization_factor=hbm_utilization_factor,\n",
+        "    max_running_seqs=max_running_seqs,\n",
+        "    max_model_len=max_model_len,\n",
+        "    min_replica_count=min_replica_count,\n",
+        "    max_replica_count=max_replica_count,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "cEq8oadxDcjl"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. The first few requests may have high latency. This is because the server needs to warm up with the initial requests. The following requests should not have the same delay.\n",
+        "\n",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown > What is a car?\n",
+        "# @markdown > A car is a four-wheeled vehicle designed for the transportation of passengers and their belongings.\n",
+        "# @markdown ```\n",
+        "\n",
+        "# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for more details.\n",
+        "\n",
+        "# Loads an existing endpoint instance using the endpoint name:\n",
+        "# - Using `endpoint_name = endpoint_hexllm.name` allows us to get the endpoint\n",
+        "#   name of the endpoint `endpoint_hexllm` created in the cell above.\n",
+        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
+        "#   an existing endpoint with the ID 1234567890123456789.\n",
+        "# You may uncomment the code below to load an existing endpoint:\n",
+        "# endpoint_name = endpoint_without_peft.name\n",
+        "# # endpoint_name = \"\"  # @param {type:\"string\"}\n",
+        "# aip_endpoint_name = (\n",
+        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
+        "# )\n",
+        "# endpoint_hexllm = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "\n",
+        "prompt = \"What is a car?\"  # @param {type: \"string\"}\n",
+        "max_tokens = 50  # @param {type: \"integer\"}\n",
+        "temperature = 1.0  # @param {type: \"number\"}\n",
+        "top_p = 1.0  # @param {type: \"number\"}\n",
+        "top_k = 1  # @param {type: \"integer\"}\n",
+        "instances = [\n",
+        "    {\n",
+        "        \"prompt\": prompt,\n",
+        "        \"max_tokens\": max_tokens,\n",
+        "        \"temperature\": temperature,\n",
+        "        \"top_p\": top_p,\n",
+        "        \"top_k\": top_k,\n",
+        "    },\n",
+        "]\n",
+        "response = endpoint_hexllm.predict(instances=instances)\n",
+        "\n",
+        "prediction = response.predictions[0]\n",
+        "print(prediction)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "z-XybZjtgF9M"
+      },
+      "source": [
+        "## Deploy prebuilt Llama 3.1 70B and 405B with vLLM"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "E8OiHHNNE_wj"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "\n",
+        "# @markdown This section uploads prebuilt Llama 3.1 models to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model.\n",
+        "\n",
+        "# @markdown NVIDIA_L4 GPUs are used for demonstration. The serving efficiency of L4 GPUs is inferior to that of H100 GPUs, but L4 GPUs are nevertheless good serving solutions if you do not have H100 quota.\n",
+        "\n",
+        "# @markdown H100 is hard to get for now. It's recommended to use the deployment button in the model card. You can still try to deploy H100 endpoint through the notebook, but there is a chance that resource is not available.\n",
+        "\n",
+        "# @markdown Set the model to deploy.\n",
+        "\n",
+        "base_model_name = \"Meta-Llama-3.1-70B\"  # @param [\"Meta-Llama-3.1-70B\", \"Meta-Llama-3.1-70B-Instruct\", \"Meta-Llama-3.1-405B-FP8\", \"Meta-Llama-3.1-405B-Instruct-FP8\"] {isTemplate:true}\n",
+        "model_id = os.path.join(MODEL_BUCKET, base_model_name)\n",
+        "\n",
+        "# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
+        "\n",
+        "if \"70\" in base_model_name:\n",
+        "    accelerator_type = \"NVIDIA_L4\"\n",
+        "    machine_type = \"g2-standard-8\"\n",
+        "    accelerator_count = 8\n",
+        "elif \"405\" in base_model_name:\n",
+        "    accelerator_type = \"NVIDIA_H100_80GB\"\n",
+        "    machine_type = \"a3-highgpu-8g\"\n",
+        "    accelerator_count = 8\n",
+        "else:\n",
+        "    raise ValueError(\n",
+        "        f\"Recommended GPU setting not found for: {accelerator_type} and {base_model_name}.\"\n",
+        "    )\n",
+        "\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=False,\n",
+        ")\n",
+        "\n",
+        "gpu_memory_utilization = 0.9\n",
+        "max_model_len = 32768  # Maximum context length.\n",
+        "\n",
+        "model, endpoint = deploy_model_vllm(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=\"llama_3_1-vllm-serve\"),\n",
+        "    model_id=model_id,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    gpu_memory_utilization=gpu_memory_utilization,\n",
+        "    max_model_len=max_model_len,\n",
+        ")\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "rDHsCOqvFYBi"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. Sampling parameters supported by vLLM can be found [here](https://docs.vllm.ai/en/latest/dev/sampling_params.html).\n",
+        "\n",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown Q: What is a car?\n",
+        "# @markdown A:  A car, or a motor car, is a road-connected human-transportation system used to move people or goods from one place to another. The term also encompasses a wide range of vehicles, including motorboats, trains, and aircrafts. Cars typically have four wheels, a cabin for passengers, and an engine or motor. They have been around since the early 19th century and are now one of the most popular forms of transportation, used for daily commuting, shopping, and other purposes.\n",
+        "# @markdown ```\n",
+        "\n",
+        "# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for more details.\n",
+        "\n",
+        "# @markdown NOTE: For the raw predict (non chat-completion API), a template like \"user:<input> assistant:\" is needed in the prompt to get a meaningful response for an instruct tuned model.\n",
+        "\n",
+        "prompt = \"user:What is a car? assistant:\"  # @param {type: \"string\"}\n",
+        "max_tokens = 50  # @param {type:\"integer\"}\n",
+        "temperature = 1.0  # @param {type:\"number\"}\n",
+        "top_p = 1.0  # @param {type:\"number\"}\n",
+        "top_k = 1  # @param {type:\"integer\"}\n",
+        "raw_response = False  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# @markdown Optionally, you can apply LoRA weights to prediction. Set `lora_weight` to be either a GCS URI or a HuggingFace repo containing the LoRA weight.\n",
+        "lora_weight = \"\"  # @param {type:\"string\", isTemplate: true}\n",
+        "\n",
+        "# Overides parameters for inferences.\n",
+        "# If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`,\n",
+        "# you can reduce the maximum number of output tokens, such as set max_tokens as 20.\n",
+        "instances = [\n",
+        "    {\n",
+        "        \"prompt\": prompt,\n",
+        "        \"max_tokens\": max_tokens,\n",
+        "        \"temperature\": temperature,\n",
+        "        \"top_p\": top_p,\n",
+        "        \"top_k\": top_k,\n",
+        "        \"raw_response\": raw_response,\n",
+        "        \"dynamic-lora\": lora_weight,\n",
+        "    },\n",
+        "]\n",
+        "response = endpoint.predict(instances=instances)\n",
+        "\n",
+        "for prediction in response.predictions:\n",
+        "    print(prediction)\n",
+        "\n",
+        "# @markdown You can also use the `@requestFormat` parameter to send the OpenAI chat completions request.\n",
+        "\n",
+        "message_role = \"user\"  # @param {type: \"string\"}\n",
+        "message_content = \"What is a car?\"  # @param {type: \"string\"}\n",
+        "\n",
+        "messages = [\n",
+        "    {\n",
+        "        \"role\": message_role,\n",
+        "        \"content\": message_content,\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "instances = [\n",
+        "    {\n",
+        "        \"messages\": messages,\n",
+        "        \"@requestFormat\": \"chatCompletions\",\n",
+        "    },\n",
+        "]\n",
+        "response = endpoint.predict(instances=instances)\n",
+        "\n",
+        "for prediction in response.predictions:\n",
+        "    print(prediction)\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "vfK-kZmwV7Bb"
+      },
+      "source": [
+        "## Use Llama guard model\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Empf8V2GWJJI"
+      },
+      "source": [
+        "You can use the Llama Guard model together with the Llama 3.1 405B Instruct API. See the [Llama Guard model](https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama-guard) for details."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "JETd33jIDcjm"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "911406c1561e"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Delete the models and endpoints\n",
+        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
+        "# @markdown  and avoid unnecessary continouous charges that may incur.\n",
+        "\n",
+        "# Undeploy model and delete endpoint.\n",
+        "endpoint.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "model.delete()\n",
+        "\n",
+        "delete_bucket = False  # @param {type:\"boolean\"}\n",
+        "if delete_bucket:\n",
+        "    ! gsutil -m rm -r $BUCKET_URI"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_pytorch_llama3_1_deployment.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
@@ -0,0 +1,718 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "7d9bbf86da5e"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2024 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "99c1c3fc2ca5"
+      },
+      "source": [
+        "# Vertex AI Model Garden - Llama 3.1 Finetuning\n",
+        "\n",
+        "<table><tbody><tr>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_pytorch_llama3_1_finetuning.ipynb\">\n",
+        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_llama3_1_finetuning.ipynb\">\n",
+        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3de7470326a2"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates finetuning and deploying Llama 3.1 models with Vertex AI. All of the examples in this notebook use parameter efficient finetuning methods [PEFT (LoRA)](https://github.com/huggingface/peft) to reduce training and storage costs. LoRA (Low-Rank Adaptation) is one approach of Parameter Efficient FineTuning (PEFT), where pretrained model weights are frozen and rank decomposition matrices representing the change in model weights are trained during finetuning. Read more about LoRA in the following publication: [Hu, E.J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L. and Chen, W., 2021. Lora: Low-rank adaptation of large language models. *arXiv preprint arXiv:2106.09685*](https://arxiv.org/abs/2106.09685).\n",
+        "\n",
+        "After finetuning, we can deploy models on Vertex with GPU.\n",
+        "\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Finetune Llama 3.1 models with Vertex AI Custom Training Jobs.\n",
+        "- Deploy finetuned Llama 3.1 models on Vertex AI Prediction.\n",
+        "- Send prediction requests to your finetuned Llama 3.1 models.\n",
+        "\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "264c07757582"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "855d6b96f291"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
+        "\n",
+        "# @markdown 3. [Make sure that you have GPU quota for Vertex Training (finetuing) and Vertex Prediction (serving)](https://cloud.google.com/docs/quotas/view-manage). The quota name for Vertex Training is \"Custom model training your-gpu-type per region\" and the quota name for Vertex Prediction is \"Custom model serving your-gpu-type per region\" such as `Custom model training Nvidia L4 GPUs per region` and `Custom model serving Nvidia L4 GPUs per region` for L4 GPUs. [Submit a quota increase request](https://cloud.google.com/docs/quotas/view-manage#requesting_higher_quota) if additional quota is needed. At minimum, running this notebook requires 4 L4s for finetuning and 1 L4 for serving. More GPUs may be needed for larger models and different finetuning configurations. To secure GPUs for larger models, ask your customer engineer to get you allowlisted for a Shared Reservation or a Dynamic Workload Scheduler.\n",
+        "\n",
+        "# Import the necessary packages\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
+        "import importlib\n",
+        "import os\n",
+        "import uuid\n",
+        "from datetime import datetime\n",
+        "from typing import Tuple\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "common_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
+        ")\n",
+        "\n",
+        "models, endpoints = {}, {}\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "\n",
+        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
+        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
+        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
+        "\n",
+        "# Cloud Storage bucket for storing the experiment artifacts.\n",
+        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
+        "# prefer using your own GCS bucket, change the value yourself below.\n",
+        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
+        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
+        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "\n",
+        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
+        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
+        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
+        "else:\n",
+        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
+        "    bucket_region = shell_output[0].strip().lower()\n",
+        "    if bucket_region != REGION:\n",
+        "        raise ValueError(\n",
+        "            \"Bucket region %s is different from notebook region %s\"\n",
+        "            % (bucket_region, REGION)\n",
+        "        )\n",
+        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"llama3_1\")\n",
+        "\n",
+        "\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
+        "\n",
+        "# Gets the default SERVICE_ACCOUNT.\n",
+        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
+        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
+        "\n",
+        "\n",
+        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "36c21f10355f"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Access Llama 3.1 models\n",
+        "\n",
+        "# @markdown For GPU based finetuning and serving, choose between accessing Llama 3.1 models on [Hugging Face](https://huggingface.co/)\n",
+        "# @markdown or Vertex AI as described below.\n",
+        "\n",
+        "# @markdown If you already obtained access to Llama 3.1 models on [Hugging Face](https://huggingface.co/), you can load models from there.\n",
+        "# @markdown Alternatively, you can also load the original Llama 3.1 models for finetuning and serving from Vertex AI after accepting the agreement.\n",
+        "\n",
+        "# @markdown **Only select and fill one of the following sections.**\n",
+        "# fmt: off\n",
+        "LOAD_MODEL_FROM = \"Hugging Face\"  # @param [\"Hugging Face\", \"Google Cloud\"] {isTemplate:true}\n",
+        "# fmt: on\n",
+        "\n",
+        "# @markdown ---\n",
+        "\n",
+        "# @markdown ### Access Llama 3.1 models on Hugging Face for GPU based finetuning and serving\n",
+        "# @markdown You must provide a Hugging Face User Access Token (read) to access the Llama 3.1 models. You can follow the [Hugging Face documentation](https://huggingface.co/docs/hub/en/security-tokens) to create a **read** access token and put it in the `HF_TOKEN` field below.\n",
+        "\n",
+        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate:true}\n",
+        "if LOAD_MODEL_FROM == \"Hugging Face\":\n",
+        "    assert (\n",
+        "        HF_TOKEN\n",
+        "    ), \"Provide a read HF_TOKEN to load models from Hugging Face, or select a different model source.\"\n",
+        "\n",
+        "# @markdown *--- Or ---*\n",
+        "# @markdown ### Access Llama 3.1 models on Vertex AI for GPU based serving\n",
+        "# @markdown The original models from Meta are converted into the Hugging Face format for serving in Vertex AI.\n",
+        "# @markdown Accept the model agreement to access the models:\n",
+        "# @markdown 1. Open the [Llama 3.1 model card](https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama3_1) from [Vertex AI Model Garden](https://cloud.google.com/model-garden).\n",
+        "# @markdown 2. Review and accept the agreement in the pop-up window on the model card page. If you have previously accepted the model agreement, there will not be a pop-up window on the model card page and this step is not needed.\n",
+        "# @markdown 3. After accepting the agreement of Llama 3.1, a `gs://` URI containing Llama 3.1 pretrained and finetuned models will be shared.\n",
+        "# @markdown 4. Paste the URI in the `VERTEX_AI_MODEL_GARDEN_LLAMA3_1` field below.\n",
+        "\n",
+        "VERTEX_AI_MODEL_GARDEN_LLAMA3_1 = \"\"  # @param {type:\"string\", isTemplate:true}\n",
+        "MODEL_BUCKET = VERTEX_AI_MODEL_GARDEN_LLAMA3_1\n",
+        "\n",
+        "# @markdown ---\n",
+        "\n",
+        "\n",
+        "# The pre-built serving docker image.\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240721_0916_RC00\"\n",
+        "\n",
+        "\n",
+        "def deploy_model_vllm(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    service_account: str,\n",
+        "    machine_type: str = \"g2-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 1,\n",
+        "    gpu_memory_utilization: float = 0.9,\n",
+        "    max_model_len: int = 8192,\n",
+        "    max_loras: int = 1,\n",
+        "    max_cpu_loras: int = 16,\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",
+        "    vllm_args = [\n",
+        "        \"python\",\n",
+        "        \"-m\",\n",
+        "        \"vllm.entrypoints.api_server\",\n",
+        "        \"--host=0.0.0.0\",\n",
+        "        \"--port=7080\",\n",
+        "        f\"--model={model_id}\",\n",
+        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
+        "        \"--swap-space=16\",\n",
+        "        f\"--gpu-memory-utilization={gpu_memory_utilization}\",\n",
+        "        f\"--max-model-len={max_model_len}\",\n",
+        "        f\"--max-loras={max_loras}\",\n",
+        "        f\"--max-cpu-loras={max_cpu_loras}\",\n",
+        "        \"--disable-log-stats\",\n",
+        "    ]\n",
+        "\n",
+        "    env_vars = {\n",
+        "        \"MODEL_ID\": base_model_id,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "\n",
+        "    try:\n",
+        "        if HF_TOKEN:\n",
+        "            env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
+        "    except:\n",
+        "        pass\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
+        "        serving_container_args=vllm_args,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_predict_route=\"/generate\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        serving_container_environment_variables=env_vars,\n",
+        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
+        "        serving_container_deployment_timeout=7200,\n",
+        "    )\n",
+        "    print(\n",
+        "        f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
+        "    )\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=service_account,\n",
+        "    )\n",
+        "    print(\"endpoint_name:\", endpoint.name)\n",
+        "\n",
+        "    return model, endpoint"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cb56d402e84a"
+      },
+      "source": [
+        "## Finetune with HuggingFace PEFT and deploy with vLLM on GPUs"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "KwAW99YZHTdy"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Set dataset\n",
+        "\n",
+        "# @markdown Use the Vertex AI SDK to create and run the custom training jobs.\n",
+        "\n",
+        "# @markdown This notebook uses [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) dataset as an example.\n",
+        "# @markdown You can set `dataset_name` to any existing [Hugging Face dataset](https://huggingface.co/datasets) name, and set `instruct_column_in_dataset` to the name of the dataset column containing training data. The [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) has only one column `text`, and therefore we set `instruct_column_in_dataset` to `text` in this notebook.\n",
+        "\n",
+        "# @markdown ### (Optional) Prepare a custom JSONL dataset for finetuning\n",
+        "\n",
+        "# @markdown You can prepare a JSONL file where each line is a valid JSON string as your custom training dataset. For example, here is one line from the [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) dataset:\n",
+        "# @markdown ```\n",
+        "# @markdown {\"text\": \"### Human: Hola### Assistant: \\u00a1Hola! \\u00bfEn qu\\u00e9 puedo ayudarte hoy?\"}\n",
+        "# @markdown ```\n",
+        "\n",
+        "# @markdown The JSON object has a key `text`, which should match `instruct_column_in_dataset`; The value should be one training data point, i.e. a string. After you prepared your JSONL file, you can either upload it to [Hugging Face datasets](https://huggingface.co/datasets) or [Google Cloud Storage](https://cloud.google.com/storage).\n",
+        "\n",
+        "# @markdown - To upload a JSONL dataset to [Hugging Face datasets](https://huggingface.co/datasets), follow the instructions on [Uploading Datasets](https://huggingface.co/docs/hub/en/datasets-adding). Then, set `dataset_name` to the name of your newly created dataset on Hugging Face.\n",
+        "\n",
+        "# @markdown - To upload a JSONL dataset to [Google Cloud Storage](https://cloud.google.com/storage), follow the instructions on [Upload objects from a filesystem](https://cloud.google.com/storage/docs/uploading-objects). Then, set `dataset_name` to the `gs://` URI to your JSONL file. For example: `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`.\n",
+        "\n",
+        "# @markdown Optionally update the `instruct_column_in_dataset` field below if your JSON objects use a key other than the default `text`.\n",
+        "\n",
+        "# @markdown ### (Optional) Format your data with custom JSON template\n",
+        "\n",
+        "# @markdown Sometimes, your dataset might have multiple text columns and you want to construct the training data with a template. You can prepare a JSON template in the following format:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown {\n",
+        "# @markdown   \"description\": \"Template used by Llama 3.1, accepting text-bison format.\",\n",
+        "# @markdown   \"source\": \"https://cloud.google.com/vertex-ai/generative-ai/docs/models/tune-text-models-supervised#dataset-format\",\n",
+        "# @markdown   \"prompt_input\": \"<|start_header_id|>user<|end_header_id|>\\n\\n{input_text}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\\n\\n{output_text}<|eot_id|>\",\n",
+        "# @markdown   \"instruction_separator\": \"<|start_header_id|>user<|end_header_id|>\\n\\n\",\n",
+        "# @markdown   \"response_separator\": \"<|start_header_id|>assistant<|end_header_id|>\\n\\n\"\n",
+        "# @markdown }\n",
+        "# @markdown ```\n",
+        "\n",
+        "# @markdown As an example, the template above can be used to format the following training data (this line comes from `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`):\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown {\"input_text\":\"TRANSCRIPT: \\nREASON FOR EVALUATION:,\\n\\n LABEL:\",\"output_text\":\"Chiropractic\"}\n",
+        "# @markdown ```\n",
+        "\n",
+        "# @markdown This example template simply concatenates `input_text` with `output_text` with some special tokens in between.\n",
+        "# @markdown\n",
+        "# @markdown To try such custom dataset, you can make the following changes:\n",
+        "# @markdown 1. Set `template` to `llama3-text-bison`\n",
+        "# @markdown 1. Set `train_dataset_name` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`\n",
+        "# @markdown 1. Set `train_split_name` to `train`\n",
+        "# @markdown 1. Set `eval_dataset_name` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_eval_sample.jsonl`\n",
+        "# @markdown 1. Set `eval_split_name` to `train` (**NOT** `test`)\n",
+        "# @markdown 1. Set `instruct_column_in_dataset` as `input_text`.\n",
+        "\n",
+        "# Template name or gs:// URI to a custom template.\n",
+        "template = \"openassistant-guanaco\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Hugging Face dataset name or gs:// URI to a custom JSONL dataset.\n",
+        "train_dataset_name = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
+        "train_split_name = \"train\"  # @param {type:\"string\"}\n",
+        "eval_dataset_name = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
+        "eval_split_name = \"test\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Name of the dataset column containing training text input.\n",
+        "instruct_column_in_dataset = \"text\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "ivVGS9dHXPOz"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Finetune\n",
+        "# @markdown Use the Vertex AI SDK to create and run the custom training jobs.\n",
+        "\n",
+        "# @markdown **Note**:\n",
+        "# @markdown 1. We recommend setting `finetuning_precision_mode` to `4bit` because it enables using fewer hardware resources for finetuning.\n",
+        "# @markdown 1. We recommend using NVIDIA_L4 for 8B models and NVIDIA_A100_80GB for 70B models.\n",
+        "# @markdown 1. If `max_steps>0`, it will precedence over `epochs`. One can set a small `max_steps` value to quickly check the pipeline.\n",
+        "# @markdown 1. With the default setting, training takes between 1.5 ~ 2 hours.\n",
+        "\n",
+        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20240724_0936_RC00\"\n",
+        "\n",
+        "\n",
+        "# The Llama 3.1 base model.\n",
+        "MODEL_ID = \"meta-llama/Meta-Llama-3.1-8B-Instruct\"  # @param [\"meta-llama/Meta-Llama-3.1-8B\", \"meta-llama/Meta-Llama-3.1-8B-Instruct\", \"meta-llama/Meta-Llama-3.1-70B\", \"meta-llama/Meta-Llama-3.1-70B-Instruct\"] {isTemplate:true}\n",
+        "if LOAD_MODEL_FROM == \"Google Cloud\":\n",
+        "    base_model_id = os.path.join(MODEL_BUCKET, MODEL_ID.split(\"/\")[-1])\n",
+        "else:\n",
+        "    base_model_id = MODEL_ID\n",
+        "\n",
+        "# The accelerator to use.\n",
+        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_A100_80GB\"]\n",
+        "\n",
+        "# Batch size for finetuning.\n",
+        "per_device_train_batch_size = 1  # @param{type:\"integer\"}\n",
+        "gradient_accumulation_steps = 8  # @param{type:\"integer\"}\n",
+        "# Maximum sequence length.\n",
+        "max_seq_length = 4096  # @param{type:\"integer\"}\n",
+        "# Setting a positive `max_steps` here will override `num_epochs`\n",
+        "max_steps = -1  # @param{type:\"integer\"}\n",
+        "num_epochs = 1.0  # @param{type:\"number\"}\n",
+        "# Precision mode for finetuning.\n",
+        "finetuning_precision_mode = \"4bit\"  # @param [\"4bit\", \"8bit\", \"float16\"]\n",
+        "# Learning rate.\n",
+        "learning_rate = 5e-5  # @param{type:\"number\"}\n",
+        "lr_scheduler_type = \"cosine\"  # @param{type:\"string\"}\n",
+        "# LoRA parameters.\n",
+        "lora_rank = 16  # @param{type:\"integer\"}\n",
+        "lora_alpha = 32  # @param{type:\"integer\"}\n",
+        "lora_dropout = 0.05  # @param{type:\"number\"}\n",
+        "enable_gradient_checkpointing = True\n",
+        "attn_implementation = \"flash_attention_2\"\n",
+        "optimizer = \"paged_adamw_32bit\"\n",
+        "warmup_ratio = \"0.01\"\n",
+        "report_to = \"tensorboard\"\n",
+        "save_steps = 10\n",
+        "logging_steps = save_steps\n",
+        "\n",
+        "# Worker pool spec.\n",
+        "machine_type = None\n",
+        "if \"8b\" in MODEL_ID.lower():\n",
+        "    if accelerator_type == \"NVIDIA_L4\":\n",
+        "        accelerator_count = 4\n",
+        "        machine_type = \"g2-standard-48\"\n",
+        "    else:\n",
+        "        raise ValueError(\n",
+        "            f\"Recommended machine settings not found for: {accelerator_type}. To use another accelerator, edit this code block to pass in an appropriate `machine_type`, `accelerator_type`, and `accelerator_count` to the deploy_model_vllm function by clicking `Show Code` and then modifying the code.\"\n",
+        "        )\n",
+        "elif \"70b\" in MODEL_ID.lower():\n",
+        "    if accelerator_type == \"NVIDIA_A100_80GB\":\n",
+        "        accelerator_count = 4\n",
+        "        machine_type = \"a2-ultragpu-4g\"\n",
+        "    else:\n",
+        "        raise ValueError(\n",
+        "            f\"Recommended machine settings not found for: {accelerator_type}. To use another accelerator, edit this code block to pass in an appropriate `machine_type`, `accelerator_type`, and `accelerator_count` to the deploy_model_vllm function by clicking `Show Code` and then modifying the code.\"\n",
+        "        )\n",
+        "else:\n",
+        "    raise ValueError(f\"Unsupported model ID or GCS path: {MODEL_ID}.\")\n",
+        "\n",
+        "replica_count = 1\n",
+        "\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=True,\n",
+        ")\n",
+        "\n",
+        "job_name = common_util.get_job_name_with_datetime(\"llama3_1-lora-train\").replace(\n",
+        "    \"_\", \"-\"\n",
+        ")\n",
+        "\n",
+        "base_output_dir = os.path.join(STAGING_BUCKET, job_name)\n",
+        "# Create a GCS folder to store the LORA adapter.\n",
+        "lora_output_dir = os.path.join(base_output_dir, \"adapter\")\n",
+        "# Create a GCS folder to store the merged model with the base model and the\n",
+        "# finetuned LORA adapter.\n",
+        "merged_model_output_dir = os.path.join(base_output_dir, \"merged-model\")\n",
+        "\n",
+        "eval_args = [\n",
+        "    f\"--eval_dataset_path={eval_dataset_name}\",\n",
+        "    f\"--eval_column={instruct_column_in_dataset}\",\n",
+        "    f\"--eval_template={template}\",\n",
+        "    f\"--eval_split={eval_split_name}\",\n",
+        "    f\"--eval_steps={save_steps}\",\n",
+        "    \"--eval_tasks=builtin_eval\",\n",
+        "    \"--eval_metric_name=loss\",\n",
+        "]\n",
+        "\n",
+        "train_job_args = [\n",
+        "    \"--config_file=vertex_vision_model_garden_peft/deepspeed_zero2_4gpu.yaml\",\n",
+        "    \"--task=instruct-lora\",\n",
+        "    \"--completion_only=True\",\n",
+        "    f\"--pretrained_model_id={base_model_id}\",\n",
+        "    f\"--dataset_name={train_dataset_name}\",\n",
+        "    f\"--train_split_name={train_split_name}\",\n",
+        "    f\"--instruct_column_in_dataset={instruct_column_in_dataset}\",\n",
+        "    f\"--output_dir={lora_output_dir}\",\n",
+        "    f\"--merge_base_and_lora_output_dir={merged_model_output_dir}\",\n",
+        "    f\"--per_device_train_batch_size={per_device_train_batch_size}\",\n",
+        "    f\"--gradient_accumulation_steps={gradient_accumulation_steps}\",\n",
+        "    f\"--lora_rank={lora_rank}\",\n",
+        "    f\"--lora_alpha={lora_alpha}\",\n",
+        "    f\"--lora_dropout={lora_dropout}\",\n",
+        "    f\"--max_steps={max_steps}\",\n",
+        "    f\"--max_seq_length={max_seq_length}\",\n",
+        "    f\"--learning_rate={learning_rate}\",\n",
+        "    f\"--lr_scheduler_type={lr_scheduler_type}\",\n",
+        "    f\"--precision_mode={finetuning_precision_mode}\",\n",
+        "    f\"--enable_gradient_checkpointing={enable_gradient_checkpointing}\",\n",
+        "    f\"--num_epochs={num_epochs}\",\n",
+        "    f\"--attn_implementation={attn_implementation}\",\n",
+        "    f\"--optimizer={optimizer}\",\n",
+        "    f\"--warmup_ratio={warmup_ratio}\",\n",
+        "    f\"--report_to={report_to}\",\n",
+        "    f\"--logging_output_dir={base_output_dir}\",\n",
+        "    f\"--save_steps={save_steps}\",\n",
+        "    f\"--logging_steps={logging_steps}\",\n",
+        "    f\"--template={template}\",\n",
+        "    f\"--huggingface_access_token={HF_TOKEN}\",\n",
+        "] + eval_args\n",
+        "\n",
+        "# Create TensorBoard\n",
+        "tensorboard = aiplatform.Tensorboard.create(job_name)\n",
+        "exp = aiplatform.TensorboardExperiment.create(\n",
+        "    tensorboard_experiment_id=job_name, tensorboard_name=tensorboard.name\n",
+        ")\n",
+        "\n",
+        "# Pass training arguments and launch job.\n",
+        "train_job = aiplatform.CustomContainerTrainingJob(\n",
+        "    display_name=job_name,\n",
+        "    container_uri=TRAIN_DOCKER_URI,\n",
+        ")\n",
+        "\n",
+        "train_job.run(\n",
+        "    args=train_job_args,\n",
+        "    environment_variables={\"WANDB_DISABLED\": True},\n",
+        "    replica_count=replica_count,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    boot_disk_size_gb=500,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    tensorboard=tensorboard.resource_name,\n",
+        "    base_output_dir=base_output_dir,\n",
+        ")\n",
+        "\n",
+        "print(\"LoRA adapter was saved in: \", lora_output_dir)\n",
+        "print(\"Trained and merged models were saved in: \", merged_model_output_dir)\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "qmHW6m8xG_4U"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "# @markdown This section uploads the model to Model Registry and deploys it on the Endpoint. It takes 15 minutes to 1 hour to finish.\n",
+        "\n",
+        "print(\"Deploying models in: \", merged_model_output_dir)\n",
+        "\n",
+        "# Find Vertex AI prediction supported accelerators and regions in [here](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute).\n",
+        "if \"8b\" in MODEL_ID.lower():\n",
+        "    machine_type = \"g2-standard-12\"\n",
+        "    accelerator_type = \"NVIDIA_L4\"\n",
+        "    accelerator_count = 1\n",
+        "else:\n",
+        "    machine_type = \"g2-standard-96\"\n",
+        "    accelerator_type = \"NVIDIA_L4\"\n",
+        "    accelerator_count = 8\n",
+        "\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=False,\n",
+        ")\n",
+        "\n",
+        "gpu_memory_utilization = 0.85\n",
+        "max_model_len = 8192  # Maximum context length.\n",
+        "\n",
+        "# Ensure max_model_len does not exceed the limit\n",
+        "if max_model_len > 8192:\n",
+        "    raise ValueError(\"max_model_len cannot exceed 8192\")\n",
+        "\n",
+        "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=\"llama3_1-vllm-serve\"),\n",
+        "    model_id=merged_model_output_dir,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    gpu_memory_utilization=gpu_memory_utilization,\n",
+        "    max_model_len=max_model_len,\n",
+        ")\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "2UYUNn60G_4U"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. Sampling parameters supported by vLLM can be found [here](https://docs.vllm.ai/en/latest/dev/sampling_params.html).\n",
+        "\n",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown Human: What is a car?\n",
+        "# @markdown Assistant:  A car, or a motor car, is a road-connected human-transportation system used to move people or goods from one place to another. The term also encompasses a wide range of vehicles, including motorboats, trains, and aircrafts. Cars typically have four wheels, a cabin for passengers, and an engine or motor. They have been around since the early 19th century and are now one of the most popular forms of transportation, used for daily commuting, shopping, and other purposes.\n",
+        "# @markdown ```\n",
+        "# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for more details.\n",
+        "\n",
+        "# Loads an existing endpoint instance using the endpoint name:\n",
+        "# - Using `endpoint_name = endpoint.name` allows us to get the\n",
+        "#   endpoint name of the endpoint `endpoint` created in the cell\n",
+        "#   above.\n",
+        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
+        "#   an existing endpoint with the ID 1234567890123456789.\n",
+        "# You may uncomment the code below to load an existing endpoint.\n",
+        "\n",
+        "# endpoint_name = \"\"  # @param {type:\"string\"}\n",
+        "# aip_endpoint_name = (\n",
+        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
+        "# )\n",
+        "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "\n",
+        "prompt = \"What is a car?\"  # @param {type: \"string\"}\n",
+        "max_tokens = 50  # @param {type:\"integer\"}\n",
+        "temperature = 1.0  # @param {type:\"number\"}\n",
+        "top_p = 1.0  # @param {type:\"number\"}\n",
+        "top_k = 1  # @param {type:\"integer\"}\n",
+        "raw_response = False  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# Overrides parameters for inferences.\n",
+        "# If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`,\n",
+        "# you can reduce the maximum number of output tokens, such as set max_tokens as 20.\n",
+        "instances = [\n",
+        "    {\n",
+        "        \"prompt\": prompt,\n",
+        "        \"max_tokens\": max_tokens,\n",
+        "        \"temperature\": temperature,\n",
+        "        \"top_p\": top_p,\n",
+        "        \"top_k\": top_k,\n",
+        "        \"raw_response\": raw_response,\n",
+        "    },\n",
+        "]\n",
+        "response = endpoints[\"vllm_gpu\"].predict(instances=instances)\n",
+        "\n",
+        "for prediction in response.predictions:\n",
+        "    print(prediction)\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "af21a3cff1e0"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "911406c1561e"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Delete the model and endpoint\n",
+        "\n",
+        "train_job.delete()\n",
+        "\n",
+        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
+        "# @markdown  and avoid unnecessary continouous charges that may incur.\n",
+        "\n",
+        "# Undeploy model and delete endpoint.\n",
+        "for endpoint in endpoints.values():\n",
+        "    endpoint.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "for model in models.values():\n",
+        "    model.delete()\n",
+        "\n",
+        "delete_bucket = False  # @param {type:\"boolean\"}\n",
+        "if delete_bucket:\n",
+        "    ! gsutil -m rm -r $BUCKET_NAME"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_pytorch_llama3_1_finetuning.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
@@ -92,16 +92,23 @@
        "# @title Setup Google Cloud project\n",
        "# Import the necessary packages\n",
        "\n",
-        "import json\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
+        "import importlib\n",
        "import os\n",
+        "import re\n",
        "from datetime import datetime\n",
        "from typing import Tuple\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "\n",
+        "common_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
+        ")\n",
+        "\n",
        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
-        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the specified region (`REGION`). Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
+        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
        "\n",
        "# Get the default cloud project id.\n",
        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
@@ -115,7 +122,7 @@
        "\n",
        "# Cloud Storage bucket for storing the experiment artifacts.\n",
        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
-        "# prefer using your own GCS bucket, please change the value yourself below.\n",
+        "# prefer using your own GCS bucket, change the value yourself below.\n",
        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
        "\n",
@@ -169,6 +176,12 @@
        "assert (\n",
        "    VERTEX_AI_MODEL_GARDEN_LLAMA3\n",
        "), \"Please click the agreement of LLaMA3 in Vertex AI Model Garden, and get the GCS path of LLaMA3 model artifacts.\"\n",
+        "parsed_gcs_url = re.search(\"gs://.*?(?=[ ]|$)\", VERTEX_AI_MODEL_GARDEN_LLAMA3)\n",
+        "if parsed_gcs_url:\n",
+        "    VERTEX_AI_MODEL_GARDEN_LLAMA3 = parsed_gcs_url.group()\n",
+        "assert VERTEX_AI_MODEL_GARDEN_LLAMA3.startswith(\n",
+        "    \"gs://\"\n",
+        "), \"VERTEX_AI_MODEL_GARDEN_LLAMA3 is expected to be a GCS URI and must start with `gs://`.\"\n",
        "print(\n",
        "    \"Copying LLaMA3 model artifacts from\",\n",
        "    VERTEX_AI_MODEL_GARDEN_LLAMA3,\n",
@@ -182,13 +195,6 @@
        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240508_0916_RC02\"\n",
        "\n",
        "\n",
-        "def get_job_name_with_datetime(prefix: str) -> str:\n",
-        "    \"\"\"Gets the job name with date time when triggering deployment jobs in\n",
-        "    Vertex AI.\n",
-        "    \"\"\"\n",
-        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
-        "\n",
-        "\n",
        "def deploy_model_vllm(\n",
        "    model_name: str,\n",
        "    model_id: str,\n",
@@ -237,101 +243,7 @@
        "    )\n",
        "    print(\"endpoint_name:\", endpoint.name)\n",
        "\n",
-        "    return model, endpoint\n",
-        "\n",
-        "\n",
-        "SERVICE_ENDPOINT = \"aiplatform.googleapis.com\"\n",
-        "\n",
-        "\n",
-        "def get_quota(project_id: str, region: str, resource_id: str) -> int:\n",
-        "    \"\"\"Returns the quota for a resource in a region. Returns -1 if can not figure out the quota.\"\"\"\n",
-        "    quota_list_output = !gcloud alpha services quota list --service=$SERVICE_ENDPOINT  --consumer=projects/$project_id --filter=\"$SERVICE_ENDPOINT/$resource_id\" --format=json\n",
-        "    # Use '.s' on the command output because it is an SList type.\n",
-        "    quota_data = json.loads(quota_list_output.s)\n",
-        "    if len(quota_data) == 0 or \"consumerQuotaLimits\" not in quota_data[0]:\n",
-        "        return -1\n",
-        "    if (\n",
-        "        len(quota_data[0][\"consumerQuotaLimits\"]) == 0\n",
-        "        or \"quotaBuckets\" not in quota_data[0][\"consumerQuotaLimits\"][0]\n",
-        "    ):\n",
-        "        return -1\n",
-        "    all_regions_data = quota_data[0][\"consumerQuotaLimits\"][0][\"quotaBuckets\"]\n",
-        "    for region_data in all_regions_data:\n",
-        "        if (\n",
-        "            region_data.get(\"dimensions\")\n",
-        "            and region_data[\"dimensions\"][\"region\"] == region\n",
-        "        ):\n",
-        "            if \"effectiveLimit\" in region_data:\n",
-        "                return int(region_data[\"effectiveLimit\"])\n",
-        "            else:\n",
-        "                return 0\n",
-        "    return -1\n",
-        "\n",
-        "\n",
-        "def get_resource_id(accelerator_type: str, is_for_training: bool) -> str:\n",
-        "    \"\"\"Returns the resource id for a given accelerator type and the use case.\n",
-        "    Args:\n",
-        "      accelerator_type: The accelerator type.\n",
-        "      is_for_training: Whether the resource is used for training. Set false\n",
-        "      for serving use case.\n",
-        "    Returns:\n",
-        "      The resource id.\n",
-        "    \"\"\"\n",
-        "    training_accelerator_map = {\n",
-        "        \"NVIDIA_TESLA_V100\": \"custom_model_training_nvidia_v100_gpus\",\n",
-        "        \"NVIDIA_L4\": \"custom_model_training_nvidia_l4_gpus\",\n",
-        "        \"NVIDIA_TESLA_A100\": \"custom_model_training_nvidia_a100_gpus\",\n",
-        "    }\n",
-        "    serving_accelerator_map = {\n",
-        "        \"NVIDIA_TESLA_V100\": \"custom_model_serving_nvidia_v100_gpus\",\n",
-        "        \"NVIDIA_L4\": \"custom_model_serving_nvidia_l4_gpus\",\n",
-        "        \"NVIDIA_TESLA_A100\": \"custom_model_serving_nvidia_a100_gpus\",\n",
-        "    }\n",
-        "    if is_for_training:\n",
-        "        if accelerator_type in training_accelerator_map:\n",
-        "            return training_accelerator_map[accelerator_type]\n",
-        "        else:\n",
-        "            raise ValueError(\n",
-        "                f\"Could not find accelerator type: {accelerator_type} for training.\"\n",
-        "            )\n",
-        "    else:\n",
-        "        if accelerator_type in serving_accelerator_map:\n",
-        "            return serving_accelerator_map[accelerator_type]\n",
-        "        else:\n",
-        "            raise ValueError(\n",
-        "                f\"Could not find accelerator type: {accelerator_type} for serving.\"\n",
-        "            )\n",
-        "\n",
-        "\n",
-        "def check_quota(\n",
-        "    project_id: str,\n",
-        "    region: str,\n",
-        "    accelerator_type: str,\n",
-        "    accelerator_count: int,\n",
-        "    is_for_training: bool,\n",
-        "):\n",
-        "    \"\"\"Checks if the project and the region has the required quota.\"\"\"\n",
-        "    resource_id = get_resource_id(accelerator_type, is_for_training)\n",
-        "    quota = get_quota(project_id, region, resource_id)\n",
-        "    quota_request_instruction = (\n",
-        "        \"Either use \"\n",
-        "        \"a different region or request additional quota. Follow \"\n",
-        "        \"instructions here \"\n",
-        "        \"https://cloud.google.com/docs/quotas/view-manage#requesting_higher_quota\"\n",
-        "        \" to check quota in a region or request additional quota for \"\n",
-        "        \"your project.\"\n",
-        "    )\n",
-        "    if quota == -1:\n",
-        "        raise ValueError(\n",
-        "            f\"\"\"Quota not found for: {resource_id} in {region}.\n",
-        "            {quota_request_instruction}\"\"\"\n",
-        "        )\n",
-        "    if quota < accelerator_count:\n",
-        "        raise ValueError(\n",
-        "            f\"\"\"Quota not enough for {resource_id} in {region}:\n",
-        "            {quota} < {accelerator_count}.\n",
-        "            {quota_request_instruction}\"\"\"\n",
-        "        )"
+        "    return model, endpoint\n"
      ]
    },
    {
@@ -356,9 +268,10 @@
        "\n",
        "# @markdown This section uploads prebuilt LLaMA3 models to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model.\n",
        "\n",
-        "\n",
        "# @markdown NVIDIA_L4 GPUs are used for demonstration. The serving efficiency of L4 GPUs is inferior to that of A100 GPUs, but L4 GPUs are nevertheless good serving solutions if you do not have A100 quota.\n",
        "\n",
+        "# @markdown Llama 3 uses a context length of 8,192 tokens, double the context length of Llama 2. Please see this [Meta blog post](https://ai.meta.com/blog/meta-llama-3/) for more details.\n",
+        "\n",
        "# @markdown Set the model to deploy.\n",
        "\n",
        "base_model_name = \"llama3-8b-chat-hf\"  # @param [\"llama3-8b-hf\", \"llama3-8b-chat-hf\", \"llama3-70b-hf\", \"llama3-70b-chat-hf\"] {isTemplate:true}\n",
@@ -401,7 +314,7 @@
        "    raise ValueError(\n",
        "        f\"Recommended GPU setting not found for: {accelerator_type} and {base_model_name}.\"\n",
        "    )\n",
-        "check_quota(\n",
+        "common_util.check_quota(\n",
        "    project_id=PROJECT_ID,\n",
        "    region=REGION,\n",
        "    accelerator_type=accelerator_type,\n",
@@ -412,8 +325,12 @@
        "gpu_memory_utilization = 0.85\n",
        "max_model_len = 8192  # Maximum context length.\n",
        "\n",
+        "# Ensure max_model_len does not exceed the limit\n",
+        "if max_model_len > 8192:\n",
+        "    raise ValueError(\"max_model_len cannot exceed 8192\")\n",
+        "\n",
        "model, endpoint = deploy_model_vllm(\n",
-        "    model_name=get_job_name_with_datetime(prefix=\"llama3-serve\"),\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=\"llama3-serve\"),\n",
        "    model_id=model_id,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        "    machine_type=machine_type,\n",
@@ -436,7 +353,7 @@
      "source": [
        "# @title Predict\n",
        "\n",
-        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. Sampling parameters supported by vLLM can be found [here](https://github.com/vllm-project/vllm/blob/2e8e49fce3775e7704d413b2f02da6d7c99525c9/vllm/sampling_params.py#L23-L64).\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. Sampling parameters supported by vLLM can be found [here](https://docs.vllm.ai/en/latest/dev/sampling_params.html).\n",
        "\n",
        "# @markdown Example:\n",
        "\n",
@@ -4,6 +4,7 @@
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
+        "cellView": "form",
        "id": "7d9bbf86da5e"
      },
      "outputs": [],
@@ -99,15 +100,24 @@
        "\n",
        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
        "\n",
-        "# @markdown 3. [Make sure that you have GPU quota for Vertex Training (finetuing) and Vertex Prediction (serving)](https://cloud.google.com/docs/quotas/view-manage). The quota name for Vertex Training is \"Custom model training your-gpu-type per region\" and the quota name for Vertex Prediction is \"Custom model serving your-gpu-type per region\" such as `Custom model training Nvidia L4 GPUs per region` and `Custom model serving Nvidia L4 GPUs per region` for L4 GPUs. [Submit a quota increase request](https://cloud.google.com/docs/quotas/view-manage#requesting_higher_quota) if additional quota is needed. At minimum, running this notebook requires 1 L4 for finetuning and 1 L4 for serving. More GPUs may be needed for larger models and different finetuning configurations.\n",
+        "# @markdown 3. [Make sure that you have GPU quota for Vertex Training (finetuing) and Vertex Prediction (serving)](https://cloud.google.com/docs/quotas/view-manage). The quota name for Vertex Training is \"Custom model training your-gpu-type per region\" and the quota name for Vertex Prediction is \"Custom model serving your-gpu-type per region\" such as `Custom model training Nvidia L4 GPUs per region` and `Custom model serving Nvidia L4 GPUs per region` for L4 GPUs. [Submit a quota increase request](https://cloud.google.com/docs/quotas/view-manage#requesting_higher_quota) if additional quota is needed. At minimum, running this notebook requires 4 L4s for finetuning and 1 L4 for serving. More GPUs may be needed for larger models and different finetuning configurations.\n",
        "\n",
        "# Import the necessary packages\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
+        "import importlib\n",
        "import os\n",
        "from datetime import datetime\n",
        "from typing import Tuple\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "\n",
+        "common_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
+        ")\n",
+        "\n",
+        "models, endpoints = {}, {}\n",
+        "\n",
        "# Get the default cloud project id.\n",
        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
        "\n",
@@ -123,13 +133,13 @@
        "# prefer using your own GCS bucket, change the value yourself below.\n",
        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
-        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
-        "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "\n",
        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
-        "    # Create a unique GCS bucket for this notebook, if not specified by the user\n",
        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\n",
        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
        "else:\n",
+        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
        "    bucket_region = shell_output[0].strip().lower()\n",
        "    if bucket_region != REGION:\n",
@@ -137,28 +147,28 @@
        "            \"Bucket region %s is different from notebook region %s\"\n",
        "            % (bucket_region, REGION)\n",
        "        )\n",
-        "\n",
        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
        "\n",
-        "# Gets the default BUCKET_URI and SERVICE_ACCOUNT if they were not specified by the user.\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"llama3\")\n",
+        "\n",
+        "\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
+        "\n",
+        "# Gets the default SERVICE_ACCOUNT.\n",
        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
        "\n",
+        "\n",
        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
-        "# Initialize Vertex AI API.\n",
-        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"staging\")\n",
-        "MODEL_BUCKET = os.path.join(STAGING_BUCKET, \"model\")\n",
-        "print(\"Initializing Vertex AI API.\")\n",
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
-        "\n",
-        "# Enable Vertex AI and Cloud Compute APIs.\n",
-        "! gcloud config set project $PROJECT_ID\n",
-        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com"
+        "! gcloud config set project $PROJECT_ID"
      ]
    },
    {
@@ -179,7 +189,9 @@
        "# @markdown Alternatively, you can also load the original Llama 3 models for finetuning and serving from Vertex AI after accepting the agreement.\n",
        "\n",
        "# @markdown **Only select and fill one of the following sections.**\n",
+        "# fmt: off\n",
        "LOAD_MODEL_FROM = \"Hugging Face\"  # @param [\"Hugging Face\", \"Google Cloud\"] {isTemplate:true}\n",
+        "# fmt: on\n",
        "\n",
        "# @markdown ---\n",
        "\n",
@@ -220,32 +232,28 @@
        "# @markdown ---\n",
        "\n",
        "\n",
-        "# The pre-built training and serving docker images.\n",
-        "# The training image uses transformers 4.38.2 and tokenizers 0.15.2.\n",
-        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20240415_0936_RC00\"\n",
-        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240418_0936_RC01\"\n",
+        "# The pre-built serving docker image.\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240508_0916_RC02\"\n",
        "\n",
        "\n",
-        "def get_job_name_with_datetime(prefix: str) -> str:\n",
-        "    \"\"\"Gets the job name with date time when triggering training or deployment\n",
-        "    jobs in Vertex AI.\n",
-        "    \"\"\"\n",
-        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
-        "\n",
-        "\n",
-        "def deploy_model(\n",
+        "def deploy_model_vllm(\n",
        "    model_name: str,\n",
        "    model_id: str,\n",
        "    service_account: str,\n",
+        "    base_model_id: str = None,\n",
        "    machine_type: str = \"g2-standard-8\",\n",
        "    accelerator_type: str = \"NVIDIA_L4\",\n",
        "    accelerator_count: int = 1,\n",
        "    gpu_memory_utilization: float = 0.9,\n",
        "    max_model_len: int = 4096,\n",
+        "    dtype: str = \"auto\",\n",
        ") -> 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",
+        "    if not base_model_id:\n",
+        "        base_model_id = model_id\n",
+        "\n",
        "    vllm_args = [\n",
        "        \"--host=0.0.0.0\",\n",
        "        \"--port=7080\",\n",
@@ -254,13 +262,18 @@
        "        \"--swap-space=16\",\n",
        "        f\"--gpu-memory-utilization={gpu_memory_utilization}\",\n",
        "        f\"--max-model-len={max_model_len}\",\n",
+        "        f\"--dtype={dtype}\",\n",
        "        \"--disable-log-stats\",\n",
        "    ]\n",
        "\n",
        "    env_vars = {\n",
-        "        \"MODEL_ID\": model_id,\n",
+        "        \"MODEL_ID\": base_model_id,\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
+        "\n",
+        "    if HF_TOKEN:\n",
+        "        env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
+        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
        "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
@@ -270,6 +283,8 @@
        "        serving_container_predict_route=\"/generate\",\n",
        "        serving_container_health_route=\"/ping\",\n",
        "        serving_container_environment_variables=env_vars,\n",
+        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
+        "        serving_container_deployment_timeout=7200,\n",
        "    )\n",
        "    print(\n",
        "        f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
@@ -333,9 +348,11 @@
        "\n",
        "# @markdown ```\n",
        "# @markdown {\n",
-        "# @markdown   \"description\": \"A short template for vertex sample dataset.\",\n",
-        "# @markdown   \"prompt_input\": \"{input_text}{output_text}\",\n",
-        "# @markdown   \"prompt_no_input\": \"{input_text}{output_text}\"\n",
+        "# @markdown   \"description\": \"Template used by Llama 3, accepting text-bison format.\",\n",
+        "# @markdown   \"source\": \"https://cloud.google.com/vertex-ai/generative-ai/docs/models/tune-text-models-supervised#dataset-format\",\n",
+        "# @markdown   \"prompt_input\": \"<|start_header_id|>user<|end_header_id|>\\n\\n{input_text}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\\n\\n{output_text}<|eot_id|>\",\n",
+        "# @markdown   \"instruction_separator\": \"<|start_header_id|>user<|end_header_id|>\\n\\n\",\n",
+        "# @markdown   \"response_separator\": \"<|start_header_id|>assistant<|end_header_id|>\\n\\n\"\n",
        "# @markdown }\n",
        "# @markdown ```\n",
        "\n",
@@ -345,16 +362,27 @@
        "# @markdown {\"input_text\":\"TRANSCRIPT: \\nREASON FOR EVALUATION:,\\n\\n LABEL:\",\"output_text\":\"Chiropractic\"}\n",
        "# @markdown ```\n",
        "\n",
-        "# @markdown This example template simply concatenates `input_text` with `output_text`. You can set `template` to `vertex_sample` to try out this built-in template with the dataset `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`, or build more complicated JSON templates such as [the alpaca example](https://github.com/tloen/alpaca-lora/blob/main/templates/alpaca.json). To use your own JSON template, [upload it to Google Cloud Storage](https://cloud.google.com/storage/docs/uploading-objects) and put the `gs://` URI in the `template` field below. Leave `instruct_column_in_dataset` as `text`.\n",
+        "# @markdown This example template simply concatenates `input_text` with `output_text` with some special tokens in between.\n",
+        "# @markdown\n",
+        "# @markdown To try such custom dataset, you can make the following changes:\n",
+        "# @markdown 1. Set `template` to `llama3-text-bison`\n",
+        "# @markdown 1. Set `train_dataset_name` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`\n",
+        "# @markdown 1. Set `train_split_name` to `train`\n",
+        "# @markdown 1. Set `eval_dataset_name` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_eval_sample.jsonl`\n",
+        "# @markdown 1. Set `eval_split_name` to `train` (**NOT** `test`)\n",
+        "# @markdown 1. Set `instruct_column_in_dataset` as `input_text`.\n",
+        "\n",
+        "# Template name or gs:// URI to a custom template.\n",
+        "template = \"openassistant-guanaco\"  # @param {type:\"string\"}\n",
        "\n",
        "# Hugging Face dataset name or gs:// URI to a custom JSONL dataset.\n",
-        "dataset_name = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
+        "train_dataset_name = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
+        "train_split_name = \"train\"  # @param {type:\"string\"}\n",
+        "eval_dataset_name = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
+        "eval_split_name = \"test\"  # @param {type:\"string\"}\n",
        "\n",
        "# Name of the dataset column containing training text input.\n",
-        "instruct_column_in_dataset = \"text\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# Optional. Template name or gs:// URI to a custom template.\n",
-        "template = \"\"  # @param {type:\"string\"}"
+        "instruct_column_in_dataset = \"text\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -369,7 +397,14 @@
        "# @title Finetune\n",
        "# @markdown Use the Vertex AI SDK to create and run the custom training jobs.\n",
        "\n",
-        "# @markdown **Note**: We recommend setting `finetuning_precision_mode` to `4bit` because it enables using fewer hardware resources for finetuning.\n",
+        "# @markdown **Note**:\n",
+        "# @markdown 1. We recommend setting `finetuning_precision_mode` to `4bit` because it enables using fewer hardware resources for finetuning.\n",
+        "# @markdown 1. We recommend using NVIDIA_L4 for 8B models and NVIDIA_A100_80GB for 70B models.\n",
+        "# @markdown 1. If `max_steps>0`, it will precedence over `epochs`. One can set a small `max_steps` value to quickly check the pipeline.\n",
+        "# @markdown 1. With the default setting, training takes between 1.5 ~ 2 hours.\n",
+        "\n",
+        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20240625_0902_RC00\"\n",
+        "\n",
        "\n",
        "# The Llama 3 base model.\n",
        "MODEL_ID = \"meta-llama/Meta-Llama-3-8B-Instruct\"  # @param [\"meta-llama/Meta-Llama-3-8B\", \"meta-llama/Meta-Llama-3-8B-Instruct\", \"meta-llama/Meta-Llama-3-70B\", \"meta-llama/Meta-Llama-3-70B-Instruct\"] {isTemplate:true}\n",
@@ -389,59 +424,121 @@
        "    base_model_id = MODEL_ID\n",
        "\n",
        "# The accelerator to use.\n",
-        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_A100\"]\n",
+        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_A100_80GB\"]\n",
        "\n",
        "# Batch size for finetuning.\n",
        "per_device_train_batch_size = 1  # @param{type:\"integer\"}\n",
-        "# Runs 10 training steps as a minimal example.\n",
-        "max_steps = 10  # @param {type:\"integer\"}\n",
+        "gradient_accumulation_steps = 8  # @param{type:\"integer\"}\n",
+        "# Maximum sequence length.\n",
+        "max_seq_length = 4096  # @param{type:\"integer\"}\n",
+        "# Setting a positive `max_steps` here will override `num_epochs`\n",
+        "max_steps = -1  # @param{type:\"integer\"}\n",
+        "num_epochs = 1.0  # @param{type:\"number\"}\n",
        "# Precision mode for finetuning.\n",
        "finetuning_precision_mode = \"4bit\"  # @param [\"4bit\", \"8bit\", \"float16\"]\n",
        "# Learning rate.\n",
-        "learning_rate = 2e-4  # @param{type:\"number\"}\n",
+        "learning_rate = 5e-5  # @param{type:\"number\"}\n",
+        "lr_scheduler_type = \"cosine\"  # @param{type:\"string\"}\n",
        "# LoRA parameters.\n",
        "lora_rank = 16  # @param{type:\"integer\"}\n",
-        "lora_alpha = 64  # @param{type:\"integer\"}\n",
-        "lora_dropout = 0.1  # @param{type:\"number\"}\n",
-        "# Maximum sequence length.\n",
-        "max_seq_length = 8192\n",
+        "lora_alpha = 32  # @param{type:\"integer\"}\n",
+        "lora_dropout = 0.05  # @param{type:\"number\"}\n",
+        "enable_gradient_checkpointing = True\n",
+        "attn_implementation = \"flash_attention_2\"\n",
+        "optimizer = \"paged_adamw_32bit\"\n",
+        "warmup_ratio = \"0.01\"\n",
+        "report_to = \"tensorboard\"\n",
+        "save_steps = 10\n",
+        "logging_steps = save_steps\n",
        "\n",
        "# Worker pool spec.\n",
-        "\n",
        "machine_type = None\n",
        "if \"8b\" in MODEL_ID.lower():\n",
        "    if accelerator_type == \"NVIDIA_L4\":\n",
-        "        if finetuning_precision_mode == \"4bit\" or finetuning_precision_mode == \"8bit\":\n",
-        "            accelerator_count = 1\n",
-        "            machine_type = \"g2-standard-12\"\n",
-        "        else:\n",
-        "            accelerator_count = 2\n",
-        "            machine_type = \"g2-standard-24\"\n",
-        "    elif accelerator_type == \"NVIDIA_TESLA_A100\":\n",
-        "        accelerator_count = 1\n",
-        "        machine_type = \"a2-highgpu-1g\"\n",
+        "        accelerator_count = 4\n",
+        "        machine_type = \"g2-standard-48\"\n",
        "    else:\n",
        "        raise ValueError(\n",
-        "            f\"Recommended machine settings not found for: {accelerator_type}. To use another accelerator, edit this code block to pass in an appropriate `machine_type`, `accelerator_type`, and `accelerator_count` to the deploy_model function by clicking `Show Code` and then modifying the code.\"\n",
+        "            f\"Recommended machine settings not found for: {accelerator_type}. To use another accelerator, edit this code block to pass in an appropriate `machine_type`, `accelerator_type`, and `accelerator_count` to the deploy_model_vllm function by clicking `Show Code` and then modifying the code.\"\n",
        "        )\n",
        "elif \"70b\" in MODEL_ID.lower():\n",
-        "    if accelerator_type == \"NVIDIA_L4\":\n",
-        "        accelerator_count = 8\n",
-        "        machine_type = \"g2-standard-96\"\n",
-        "    elif accelerator_type == \"NVIDIA_TESLA_A100\":\n",
+        "    if accelerator_type == \"NVIDIA_A100_80GB\":\n",
        "        accelerator_count = 4\n",
-        "        machine_type = \"a2-highgpu-4g\"\n",
+        "        machine_type = \"a2-ultragpu-4g\"\n",
        "    else:\n",
        "        raise ValueError(\n",
-        "            f\"Recommended machine settings not found for: {accelerator_type}. To use another accelerator, edit this code block to pass in an appropriate `machine_type`, `accelerator_type`, and `accelerator_count` to the deploy_model function by clicking `Show Code` and then modifying the code.\"\n",
+        "            f\"Recommended machine settings not found for: {accelerator_type}. To use another accelerator, edit this code block to pass in an appropriate `machine_type`, `accelerator_type`, and `accelerator_count` to the deploy_model_vllm function by clicking `Show Code` and then modifying the code.\"\n",
        "        )\n",
        "else:\n",
        "    raise ValueError(f\"Unsupported model ID or GCS path: {MODEL_ID}.\")\n",
        "\n",
        "replica_count = 1\n",
        "\n",
-        "# Setup training job.\n",
-        "job_name = get_job_name_with_datetime(\"llama3-lora-train\")\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=True,\n",
+        ")\n",
+        "\n",
+        "job_name = common_util.get_job_name_with_datetime(\"llama3-lora-train\").replace(\"_\", \"-\")\n",
+        "\n",
+        "base_output_dir = os.path.join(STAGING_BUCKET, job_name)\n",
+        "# Create a GCS folder to store the LORA adapter.\n",
+        "lora_output_dir = os.path.join(base_output_dir, \"adapter\")\n",
+        "# Create a GCS folder to store the merged model with the base model and the\n",
+        "# finetuned LORA adapter.\n",
+        "merged_model_output_dir = os.path.join(base_output_dir, \"merged-model\")\n",
+        "\n",
+        "eval_args = [\n",
+        "    f\"--eval_dataset_path={eval_dataset_name}\",\n",
+        "    f\"--eval_column={instruct_column_in_dataset}\",\n",
+        "    f\"--eval_template={template}\",\n",
+        "    f\"--eval_split={eval_split_name}\",\n",
+        "    f\"--eval_steps={save_steps}\",\n",
+        "    \"--eval_tasks=builtin_eval\",\n",
+        "    \"--eval_metric_name=loss\",\n",
+        "]\n",
+        "\n",
+        "train_job_args = [\n",
+        "    \"--config_file=vertex_vision_model_garden_peft/deepspeed_zero2_4gpu.yaml\",\n",
+        "    \"--task=instruct-lora\",\n",
+        "    \"--completion_only=True\",\n",
+        "    f\"--pretrained_model_id={base_model_id}\",\n",
+        "    f\"--dataset_name={train_dataset_name}\",\n",
+        "    f\"--train_split_name={train_split_name}\",\n",
+        "    f\"--instruct_column_in_dataset={instruct_column_in_dataset}\",\n",
+        "    f\"--output_dir={lora_output_dir}\",\n",
+        "    f\"--merge_base_and_lora_output_dir={merged_model_output_dir}\",\n",
+        "    f\"--per_device_train_batch_size={per_device_train_batch_size}\",\n",
+        "    f\"--gradient_accumulation_steps={gradient_accumulation_steps}\",\n",
+        "    f\"--lora_rank={lora_rank}\",\n",
+        "    f\"--lora_alpha={lora_alpha}\",\n",
+        "    f\"--lora_dropout={lora_dropout}\",\n",
+        "    f\"--max_steps={max_steps}\",\n",
+        "    f\"--max_seq_length={max_seq_length}\",\n",
+        "    f\"--learning_rate={learning_rate}\",\n",
+        "    f\"--lr_scheduler_type={lr_scheduler_type}\",\n",
+        "    f\"--precision_mode={finetuning_precision_mode}\",\n",
+        "    f\"--enable_gradient_checkpointing={enable_gradient_checkpointing}\",\n",
+        "    f\"--num_epochs={num_epochs}\",\n",
+        "    f\"--attn_implementation={attn_implementation}\",\n",
+        "    f\"--optimizer={optimizer}\",\n",
+        "    f\"--warmup_ratio={warmup_ratio}\",\n",
+        "    f\"--report_to={report_to}\",\n",
+        "    f\"--logging_output_dir={base_output_dir}\",\n",
+        "    f\"--save_steps={save_steps}\",\n",
+        "    f\"--logging_steps={logging_steps}\",\n",
+        "    f\"--template={template}\",\n",
+        "    f\"--huggingface_access_token={HF_TOKEN}\",\n",
+        "] + eval_args\n",
+        "\n",
+        "# Create TensorBoard\n",
+        "tensorboard = aiplatform.Tensorboard.create(job_name)\n",
+        "exp = aiplatform.TensorboardExperiment.create(\n",
+        "    tensorboard_experiment_id=job_name, tensorboard_name=tensorboard.name\n",
+        ")\n",
        "\n",
        "# Pass training arguments and launch job.\n",
        "train_job = aiplatform.CustomContainerTrainingJob(\n",
@@ -449,34 +546,8 @@
        "    container_uri=TRAIN_DOCKER_URI,\n",
        ")\n",
        "\n",
-        "# Create a GCS folder to store the LORA adapter.\n",
-        "lora_adapter_dir = get_job_name_with_datetime(\"llama3-lora-adapter\")\n",
-        "lora_output_dir = os.path.join(STAGING_BUCKET, lora_adapter_dir)\n",
-        "\n",
-        "# Create a GCS folder to store the merged model with the base model and the\n",
-        "# finetuned LORA adapter.\n",
-        "merged_model_dir = get_job_name_with_datetime(\"llama3-merged-model\")\n",
-        "merged_model_output_dir = os.path.join(STAGING_BUCKET, merged_model_dir)\n",
-        "\n",
        "train_job.run(\n",
-        "    args=[\n",
-        "        \"--task=instruct-lora\",\n",
-        "        f\"--pretrained_model_id={base_model_id}\",\n",
-        "        f\"--dataset_name={dataset_name}\",\n",
-        "        f\"--instruct_column_in_dataset={instruct_column_in_dataset}\",\n",
-        "        f\"--output_dir={lora_output_dir}\",\n",
-        "        f\"--merge_base_and_lora_output_dir={merged_model_output_dir}\",\n",
-        "        f\"--per_device_train_batch_size={per_device_train_batch_size}\",\n",
-        "        f\"--lora_rank={lora_rank}\",\n",
-        "        f\"--lora_alpha={lora_alpha}\",\n",
-        "        f\"--lora_dropout={lora_dropout}\",\n",
-        "        f\"--max_steps={max_steps}\",\n",
-        "        f\"--max_seq_length={max_seq_length}\",\n",
-        "        f\"--learning_rate={learning_rate}\",\n",
-        "        f\"--precision_mode={finetuning_precision_mode}\",\n",
-        "        f\"--template={template}\",\n",
-        "        f\"--huggingface_access_token={HF_TOKEN}\",\n",
-        "    ],\n",
+        "    args=train_job_args,\n",
        "    environment_variables={\"WANDB_DISABLED\": True},\n",
        "    replica_count=replica_count,\n",
        "    machine_type=machine_type,\n",
@@ -484,6 +555,8 @@
        "    accelerator_count=accelerator_count,\n",
        "    boot_disk_size_gb=500,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
+        "    tensorboard=tensorboard.resource_name,\n",
+        "    base_output_dir=base_output_dir,\n",
        ")\n",
        "\n",
        "print(\"LoRA adapter was saved in: \", lora_output_dir)\n",
@@ -516,11 +589,23 @@
        "    accelerator_type = \"NVIDIA_L4\"\n",
        "    accelerator_count = 8\n",
        "\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=False,\n",
+        ")\n",
+        "\n",
        "gpu_memory_utilization = 0.85\n",
        "max_model_len = 8192  # Maximum context length.\n",
        "\n",
-        "model, endpoint = deploy_model(\n",
-        "    model_name=get_job_name_with_datetime(prefix=\"llama3-vllm-serve\"),\n",
+        "# Ensure max_model_len does not exceed the limit\n",
+        "if max_model_len > 8192:\n",
+        "    raise ValueError(\"max_model_len cannot exceed 8192\")\n",
+        "\n",
+        "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=\"llama3-vllm-serve\"),\n",
        "    model_id=merged_model_output_dir,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        "    machine_type=machine_type,\n",
@@ -554,6 +639,20 @@
        "# @markdown ```\n",
        "# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for more details.\n",
        "\n",
+        "# Loads an existing endpoint instance using the endpoint name:\n",
+        "# - Using `endpoint_name = endpoint.name` allows us to get the\n",
+        "#   endpoint name of the endpoint `endpoint` created in the cell\n",
+        "#   above.\n",
+        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
+        "#   an existing endpoint with the ID 1234567890123456789.\n",
+        "# You may uncomment the code below to load an existing endpoint.\n",
+        "\n",
+        "# endpoint_name = \"\"  # @param {type:\"string\"}\n",
+        "# aip_endpoint_name = (\n",
+        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
+        "# )\n",
+        "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "\n",
        "prompt = \"What is a car?\"  # @param {type: \"string\"}\n",
        "max_tokens = 50  # @param {type:\"integer\"}\n",
        "temperature = 1.0  # @param {type:\"number\"}\n",
@@ -561,9 +660,6 @@
        "top_k = 1  # @param {type:\"integer\"}\n",
        "raw_response = False  # @param {type:\"boolean\"}\n",
        "\n",
-        "# Overides parameters for inferences.\n",
-        "# If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`,\n",
-        "# you can reduce the maximum number of output tokens, such as set max_tokens as 20.\n",
        "instances = [\n",
        "    {\n",
        "        \"prompt\": prompt,\n",
@@ -574,7 +670,7 @@
        "        \"raw_response\": raw_response,\n",
        "    },\n",
        "]\n",
-        "response = endpoint.predict(instances=instances)\n",
+        "response = endpoints[\"vllm_gpu\"].predict(instances=instances)\n",
        "\n",
        "for prediction in response.predictions:\n",
        "    print(prediction)\n",
@@ -602,16 +698,18 @@
      "source": [
        "# @title Delete the model and endpoint\n",
        "\n",
-        "# @markdown Delete the model and endpoint to recycle the resources and avoid\n",
-        "# @markdown unnecessary continuous charges that may incur.\n",
-        "\n",
        "train_job.delete()\n",
        "\n",
-        "# Undeploy model and delete endpoint.\n",
-        "endpoint.delete(force=True)\n",
+        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
+        "# @markdown  and avoid unnecessary continouous charges that may incur.\n",
        "\n",
-        "# Delete model.\n",
-        "model.delete()\n",
+        "# Undeploy model and delete endpoint.\n",
+        "for endpoint in endpoints.values():\n",
+        "    endpoint.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "for model in models.values():\n",
+        "    model.delete()\n",
        "\n",
        "delete_bucket = False  # @param {type:\"boolean\"}\n",
        "if delete_bucket:\n",
@@ -1,874 +0,0 @@
-{
-  "cells": [
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "7d9bbf86da5e"
-      },
-      "outputs": [],
-      "source": [
-        "# Copyright 2023 Google LLC\n",
-        "#\n",
-        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
-        "# you may not use this file except in compliance with the License.\n",
-        "# You may obtain a copy of the License at\n",
-        "#\n",
-        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
-        "#\n",
-        "# Unless required by applicable law or agreed to in writing, software\n",
-        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
-        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
-        "# See the License for the specific language governing permissions and\n",
-        "# limitations under the License."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "99c1c3fc2ca5"
-      },
-      "source": [
-        "# Vertex AI Model Garden - Mistral and Mixtral 8x7B Models\n",
-        "\n",
-        "<table align=\"left\">\n",
-        "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_mistral.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
-        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_mistral.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
-        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_pytorch_mistral.ipynb\">\n",
-        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
-        "Open in Vertex AI Workbench\n",
-        "    </a> (A Python-3 CPU notebook is recommended)\n",
-        "  </td>\n",
-        "</table>"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "3de7470326a2"
-      },
-      "source": [
-        "## Overview\n",
-        "\n",
-        "This notebook demonstrates deploying prebuilt [Mistral](https://mistral.ai/) and Mixtral 8x7B models in Vertex AI.\n",
-        "\n",
-        "### Objective\n",
-        "\n",
-        "- Deploy prebuilt [Mistral models](https://huggingface.co/mistralai) with [vLLM](https://github.com/vllm-project/vllm) containers\n",
-        "    - [mistralai/Mistral-7B-v0.1](https://huggingface.co/mistralai/Mistral-7B-v0.1): pretrained generative text model with 7 billion parameters\n",
-        "    - [mistralai/Mistral-7B-Instruct-v0.1](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1): Instruction fine-tuned version of the Mistral-7B-v0.1 generative text model\n",
-        "    - [mistralai/Mistral-7B-Instruct-v0.2](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2): Improved instruction fine-tuned version of Mistral-7B-Instruct-v0.1 supporting 32k context length\n",
-        "- Deploy prebuit [Mixtral 8x7B model](https://huggingface.co/mistralai) with [vLLM](https://github.com/vllm-project/vllm) containers\n",
-        "    - [mistralai/Mixtral-8x7B-v0.1](https://huggingface.co/mistralai/Mixtral-8x7B-v0.1): pretrained Mixture of Experts (MoE) model with 8 branches\n",
-        "    - [mistralai/Mixtral-8x7B-Instruct-v0.1](https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1): Instruction fine-tuned version of the Mixture of Experts (MoE) model with 8 branches\n",
-        "\n",
-        "### Costs\n",
-        "\n",
-        "This tutorial uses billable components of Google Cloud:\n",
-        "\n",
-        "* Vertex AI\n",
-        "* Cloud Storage\n",
-        "\n",
-        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), [Cloud NL API pricing](https://cloud.google.com/natural-language/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "264c07757582"
-      },
-      "source": [
-        "## Before you begin"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ioensNKM8ned"
-      },
-      "source": [
-        "### Colab only\n",
-        "Run the following commands for Colab and skip this section if you are using Workbench."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "2707b02ef5df"
-      },
-      "outputs": [],
-      "source": [
-        "import sys\n",
-        "\n",
-        "if \"google.colab\" in sys.modules:\n",
-        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
-        "    from google.colab import auth as google_auth\n",
-        "\n",
-        "    google_auth.authenticate_user()\n",
-        "    # Install gdown for downloading example training images.\n",
-        "    ! pip3 install gdown\n",
-        "\n",
-        "    # Restart the notebook kernel after installs.\n",
-        "    import IPython\n",
-        "\n",
-        "    app = IPython.Application.instance()\n",
-        "    app.kernel.do_shutdown(True)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "46d25fe73955"
-      },
-      "source": [
-        "### Install dependencies"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "c75c2c1fa6e0"
-      },
-      "outputs": [],
-      "source": [
-        "! pip3 install transformers==4.36.0\n",
-        "! pip3 install accelerate==0.23.0"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "bb7adab99e41"
-      },
-      "source": [
-        "### Setup Google Cloud project\n",
-        "\n",
-        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
-        "\n",
-        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
-        "\n",
-        "1. [Enable the Vertex AI API, Compute Engine API and Cloud Natural Language API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component,language.googleapis.com).\n",
-        "\n",
-        "1. [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs.\n",
-        "\n",
-        "1. [Create a service account](https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console) with `Vertex AI User` and `Storage Object Admin` roles for deploying fine tuned model to Vertex AI endpoint."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "6c460088b873"
-      },
-      "source": [
-        "### Define environment variables\n",
-        "\n",
-        "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\")."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "855d6b96f291"
-      },
-      "outputs": [],
-      "source": [
-        "# Cloud project id.\n",
-        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# The region you want to launch jobs in.\n",
-        "# Select region based on the accelerators and regions supported by Vertex AI Prediction\n",
-        "# https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
-        "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 = \"gs://\"  # @param {type:\"string\"}\n",
-        "\n",
-        "! gcloud config set project $PROJECT_ID\n",
-        "\n",
-        "import os\n",
-        "\n",
-        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
-        "\n",
-        "# The service account looks like:\n",
-        "# '@.iam.gserviceaccount.com'\n",
-        "# Please go to https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console\n",
-        "# and create service account with `Vertex AI User` and `Storage Object Admin` roles.\n",
-        "# The service account for deploying fine tuned model.\n",
-        "SERVICE_ACCOUNT = \"\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "e828eb320337"
-      },
-      "source": [
-        "### Initialize Vertex AI API"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "12cd25839741"
-      },
-      "outputs": [],
-      "source": [
-        "from google.cloud import aiplatform\n",
-        "\n",
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "2cc825514deb"
-      },
-      "source": [
-        "### Define constants"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "b42bd4fa2b2d"
-      },
-      "outputs": [],
-      "source": [
-        "# The pre-built serving docker images with vLLM\n",
-        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240313_0916_RC00\""
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0c250872074f"
-      },
-      "source": [
-        "### Define common functions"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "354da31189dc"
-      },
-      "outputs": [],
-      "source": [
-        "import os\n",
-        "from datetime import datetime\n",
-        "from typing import Tuple\n",
-        "\n",
-        "from google.cloud import aiplatform\n",
-        "\n",
-        "\n",
-        "def get_job_name_with_datetime(prefix: str) -> str:\n",
-        "    \"\"\"Gets the job name with date time when triggering training or deployment\n",
-        "    jobs in Vertex AI.\n",
-        "    \"\"\"\n",
-        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
-        "\n",
-        "\n",
-        "def deploy_model_vllm(\n",
-        "    model_name: str,\n",
-        "    model_id: str,\n",
-        "    service_account: str,\n",
-        "    machine_type: str = \"g2-standard-8\",\n",
-        "    accelerator_type: str = \"NVIDIA_L4\",\n",
-        "    accelerator_count: int = 1,\n",
-        "    max_model_len: int = 4096,\n",
-        "    gpu_memory_utilization: float = 0.9,\n",
-        "    use_openai_server: bool = False,\n",
-        "    use_chat_completions_if_openai_server: bool = False,\n",
-        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
-        "    \"\"\"Deploys Mistral models with vLLM on Vertex AI.\n",
-        "\n",
-        "    Args:\n",
-        "        model_name: Display name of the model.\n",
-        "        model_id: Model ID or path to model weights.\n",
-        "        service_account: Service account for model uploading and deployment.\n",
-        "        machine_type: Deployment machine type.\n",
-        "        accelerator_type: Deployment accelerator type.\n",
-        "        accelerator_count: Number of accelerators to use.\n",
-        "        max_model_len: Maximum model length.\n",
-        "        gpu_memory_utilization: Fraction of GPU memory to be used for the model\n",
-        "            executor.\n",
-        "        use_openai_server: Whether to use the OpenAI-format vLLM model server.\n",
-        "        use_chat_completions_if_openai_server: If the OpenAI model server is\n",
-        "            used, whether to use the chat completion API as opposed to the text\n",
-        "            completion API. The vLLM text completion API mimics the OpenAI text\n",
-        "            completion API:\n",
-        "            https://platform.openai.com/docs/api-reference/completions/create.\n",
-        "            It has two required parameters: the model ID to direct requests to\n",
-        "            and the prompt. The response includes a \"choices\" field that\n",
-        "            contains the generated text and a \"usage\" field that contains token\n",
-        "            counts. The vLLM chat completion API mimics the OpenAI chat\n",
-        "            completion API:\n",
-        "            https://platform.openai.com/docs/api-reference/chat/create. It has\n",
-        "            two required parameters: the model ID to direct requests to and\n",
-        "            \"messages\" which is a sequence of system/user/assistant/tool\n",
-        "            messages that can represent a multi-turn chat conversation. The\n",
-        "            response includes a \"choices\" field that contains the generated\n",
-        "            message from a role and a \"usage\" field that contains token counts.\n",
-        "\n",
-        "    Returns:\n",
-        "        Model instance and endpoint instance.\n",
-        "    \"\"\"\n",
-        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
-        "\n",
-        "    dtype = \"bfloat16\"\n",
-        "    if accelerator_type in [\"NVIDIA_TESLA_T4\", \"NVIDIA_TESLA_V100\"]:\n",
-        "        dtype = \"float16\"\n",
-        "\n",
-        "    vllm_args = [\n",
-        "        \"--host=0.0.0.0\",\n",
-        "        \"--port=7080\",\n",
-        "        f\"--model={model_id}\",\n",
-        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
-        "        \"--swap-space=16\",\n",
-        "        f\"--dtype={dtype}\",\n",
-        "        f\"--gpu-memory-utilization={gpu_memory_utilization}\",\n",
-        "        f\"--max-model-len={max_model_len}\",\n",
-        "        \"--disable-log-stats\",\n",
-        "    ]\n",
-        "    serving_env = {\n",
-        "        \"MODEL_ID\": model_id,\n",
-        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
-        "    }\n",
-        "    if use_openai_server:\n",
-        "        if use_chat_completions_if_openai_server:\n",
-        "            serving_container_predict_route = \"/v1/chat/completions\"\n",
-        "        else:\n",
-        "            serving_container_predict_route = \"/v1/completions\"\n",
-        "    else:\n",
-        "        serving_container_predict_route = \"/generate\"\n",
-        "    model = aiplatform.Model.upload(\n",
-        "        display_name=model_name,\n",
-        "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
-        "        serving_container_command=[\n",
-        "            \"python\",\n",
-        "            \"-m\",\n",
-        "            (\n",
-        "                \"vllm.entrypoints.api_server\"\n",
-        "                if not use_openai_server\n",
-        "                else \"vllm.entrypoints.openai.api_server\"\n",
-        "            ),\n",
-        "        ],\n",
-        "        serving_container_args=vllm_args,\n",
-        "        serving_container_ports=[7080],\n",
-        "        serving_container_predict_route=serving_container_predict_route,\n",
-        "        serving_container_health_route=\"/health\",\n",
-        "        serving_container_environment_variables=serving_env,\n",
-        "    )\n",
-        "\n",
-        "    model.deploy(\n",
-        "        endpoint=endpoint,\n",
-        "        machine_type=machine_type,\n",
-        "        accelerator_type=accelerator_type,\n",
-        "        accelerator_count=accelerator_count,\n",
-        "        deploy_request_timeout=1800,\n",
-        "        service_account=service_account,\n",
-        "    )\n",
-        "    return model, endpoint"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "e057b5edcf81"
-      },
-      "source": [
-        "## Run inferences locally with prebuilt Mistral and Mixtral models\n",
-        "\n",
-        "You will need at least 24GB of memory to run inference with Mistral-7B. You can run locally or on Vertex AI Prediction endpoints with any of the following specs:\n",
-        "- g2-standard-8 with 1 L4 GPU\n",
-        "- n1-standard-16 with 2 V100 GPUs\n",
-        "- n1-standard-16 with 2 T4 GPUs\n",
-        "- a2-highgpu-1g with 1 A100 GPU\n",
-        "\n",
-        "You will need at least 96GB of memory to run inference with Mixtral 8x7B. You can run locally or on Vertex AI Prediction endpoints with any of the following specs:\n",
-        "- g2-standard-96 with 8 L4 GPUs\n",
-        "- n1-standard-32 with 8 V100 GPUs\n",
-        "- n1-standard-32 with 8 T4 GPUs\n",
-        "- a2-highgpu-4g with 4 A100 GPUs"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "31f6cc84efdd"
-      },
-      "outputs": [],
-      "source": [
-        "%%time\n",
-        "import torch\n",
-        "import transformers\n",
-        "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
-        "\n",
-        "device = \"cuda\"  # the device to load the model onto\n",
-        "model_name = \"mistralai/Mistral-7B-v0.1\"  # @param [\"mistralai/Mistral-7B-v0.1\", \"mistralai/Mistral-7B-Instruct-v0.1\", \"mistralai/Mistral-7B-Instruct-v0.2\", \"mistralai/Mixtral-8x7B-v0.1\", \"mistralai/Mixtral-8x7B-Instruct-v0.1\"]\n",
-        "model = AutoModelForCausalLM.from_pretrained(\n",
-        "    model_name, device_map=\"auto\", return_dict=True, torch_dtype=torch.float16\n",
-        ")\n",
-        "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
-        "\n",
-        "pipeline = transformers.pipeline(\"text-generation\", model=model, tokenizer=tokenizer)\n",
-        "\n",
-        "prompt = \"My favourite condiment is\"\n",
-        "\n",
-        "sequences = pipeline(\n",
-        "    prompt,\n",
-        "    max_length=200,\n",
-        "    do_sample=True,\n",
-        "    top_k=10,\n",
-        "    num_return_sequences=1,\n",
-        "    eos_token_id=tokenizer.eos_token_id,\n",
-        ")\n",
-        "\n",
-        "for seq in sequences:\n",
-        "    print(f\"Result: {seq['generated_text']}\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "YKZ4CBJ2kYaW"
-      },
-      "source": [
-        "## Deploy Prebuilt Mistral model with vLLM\n",
-        "\n",
-        "This section deploys the prebuilt Mistral model with [vLLM](https://github.com/vllm-project/vllm) on a Vertex endpoint. The model deployment step will take ~15 minutes to complete.\n",
-        "\n",
-        "vLLM is a highly optimized LLM serving framework which can significantly increase serving throughput. The higher QPS you have, the more benefits you get using vLLM."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "25b5b3a44cf8"
-      },
-      "source": [
-        "Set the prebuilt model id."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "10547af949fc"
-      },
-      "outputs": [],
-      "source": [
-        "prebuilt_model_id = \"mistralai/Mistral-7B-v0.1\"  # @param [\"mistralai/Mistral-7B-v0.1\", \"mistralai/Mistral-7B-Instruct-v0.1\", \"mistralai/Mistral-7B-Instruct-v0.2\"]"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "03d504bcd60b"
-      },
-      "outputs": [],
-      "source": [
-        "# Find Vertex AI prediction supported accelerators and regions in\n",
-        "# https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
-        "\n",
-        "# Sets 1 L4 to deploy Mistral 7B.\n",
-        "machine_type = \"g2-standard-8\"\n",
-        "accelerator_type = \"NVIDIA_L4\"\n",
-        "accelerator_count = 1\n",
-        "\n",
-        "# Sets 2 V100s to deploy Mistral 7B.\n",
-        "# machine_type = \"n1-standard-16\"\n",
-        "# accelerator_type = \"NVIDIA_TESLA_V100\"\n",
-        "# accelerator_count = 2\n",
-        "\n",
-        "# Sets 2 T4s to deploy Mistral 7B.\n",
-        "# machine_type = \"n1-standard-16\"\n",
-        "# accelerator_type = \"NVIDIA_TESLA_T4\"\n",
-        "# accelerator_count = 2\n",
-        "\n",
-        "# Sets 1 A100 (40G) to deploy Mistral 7B.\n",
-        "# machine_type = \"a2-highgpu-1g\"\n",
-        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
-        "# accelerator_count = 1\n",
-        "\n",
-        "# Larger setting of `max-model-len` can lead to higher requirements on\n",
-        "# `gpu-memory-utilization` and GPU configuration. Larger setting of\n",
-        "# `gpu-memory-utilization` increases the risk of running out of GPU memory with\n",
-        "# long prompts.\n",
-        "max_model_len = 4096\n",
-        "gpu_memory_utilization = 0.9\n",
-        "\n",
-        "model, endpoint = deploy_model_vllm(\n",
-        "    model_name=get_job_name_with_datetime(prefix=\"mistral-serve-vllm\"),\n",
-        "    model_id=prebuilt_model_id,\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
-        "    machine_type=machine_type,\n",
-        "    accelerator_type=accelerator_type,\n",
-        "    accelerator_count=accelerator_count,\n",
-        "    max_model_len=max_model_len,\n",
-        "    gpu_memory_utilization=gpu_memory_utilization,\n",
-        "    use_openai_server=False,\n",
-        "    use_chat_completions_if_openai_server=False,\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "RRR11SWykYaX"
-      },
-      "source": [
-        "NOTE: If you see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint, the model server is likely still initializing. Please retry later.\n",
-        "\n",
-        "NOTE: If you receive `InternalServerError: 500 System error` during the deployment, most likely the operation failed due to unavailability of resources. Either retry or use a different accelerator type.\n",
-        "\n",
-        "Once deployment succeeds, you can send requests to the endpoint with text prompts."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "3a7948c56e3d"
-      },
-      "source": [
-        "### Run sample prompt"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "3f5a1e1de60d"
-      },
-      "outputs": [],
-      "source": [
-        "# Loads an existing endpoint instance using the endpoint name:\n",
-        "# - Using `endpoint_name = endpoint.name` allows us to get the endpoint name of\n",
-        "#   the endpoint `endpoint` created in the cell above.\n",
-        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
-        "#   an existing endpoint with the ID 1234567890123456789.\n",
-        "# You may uncomment the code below to load an existing endpoint.\n",
-        "\n",
-        "# endpoint_name = endpoint_without_peft.name\n",
-        "# # endpoint_name = \"\"  # @param {type:\"string\"}\n",
-        "# aip_endpoint_name = (\n",
-        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
-        "# )\n",
-        "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
-        "\n",
-        "instances = [\n",
-        "    {\n",
-        "        \"prompt\": \"My favourite condiment is\",\n",
-        "        \"n\": 1,\n",
-        "        \"max_tokens\": 200,\n",
-        "        \"temperature\": 1.0,\n",
-        "        \"top_p\": 1.0,\n",
-        "        \"top_k\": 10,\n",
-        "    },\n",
-        "]\n",
-        "response = endpoint.predict(instances=instances)\n",
-        "\n",
-        "for prediction in response.predictions:\n",
-        "    print(prediction)\n",
-        "\n",
-        "# Reference the following code for using the OpenAI vLLM server.\n",
-        "# import json\n",
-        "# response = endpoint.raw_predict(\n",
-        "#     body=json.dumps({\n",
-        "#         \"model\": prebuilt_model_id,\n",
-        "#         \"prompt\": \"My favourite condiment is\",\n",
-        "#         \"n\": 1,\n",
-        "#         \"max_tokens\": 200,\n",
-        "#         \"temperature\": 1.0,\n",
-        "#         \"top_p\": 1.0,\n",
-        "#         \"top_k\": 10,\n",
-        "#     }),\n",
-        "#     headers={\"Content-Type\": \"application/json\"},\n",
-        "# )\n",
-        "# print(response.json())"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "wOh9irbqJ-MM"
-      },
-      "source": [
-        "## Deploy Prebuilt Mixtral 8x7B model with vLLM\n",
-        "\n",
-        "This section deploys the prebuilt Mixtral 8x7B model with [vLLM](https://github.com/vllm-project/vllm) on a Vertex endpoint. The model deployment step will take ~40 minutes to complete.\n",
-        "\n",
-        "vLLM is a highly optimized LLM serving framework which can significantly increase serving throughput. The higher QPS you have, the more benefits you get using vLLM."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "h2uCSnoaJ-MM"
-      },
-      "source": [
-        "Set the prebuilt model id."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "-X42gkGYJ-MM"
-      },
-      "outputs": [],
-      "source": [
-        "prebuilt_model_id = \"mistralai/Mixtral-8x7B-v0.1\"  # @param [\"mistralai/Mixtral-8x7B-v0.1\", \"mistralai/Mixtral-8x7B-Instruct-v0.1\"]"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "M-YiJXT3J-MM"
-      },
-      "outputs": [],
-      "source": [
-        "# Find Vertex AI prediction supported accelerators and regions in\n",
-        "# https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
-        "\n",
-        "# Sets 8 L4s to deploy Mixtral 8x7B.\n",
-        "machine_type = \"g2-standard-96\"\n",
-        "accelerator_type = \"NVIDIA_L4\"\n",
-        "accelerator_count = 8\n",
-        "\n",
-        "# Sets 4 A100s (40G) to deploy Mixtral 8x7B.\n",
-        "# machine_type = \"a2-highgpu-4g\"\n",
-        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
-        "# accelerator_count = 4\n",
-        "\n",
-        "# Larger setting of `max-model-len` can lead to higher requirements on\n",
-        "# `gpu-memory-utilization` and GPU configuration. Larger setting of\n",
-        "# `gpu-memory-utilization` increases the risk of running out of GPU memory with\n",
-        "# long prompts.\n",
-        "max_model_len = 4096\n",
-        "gpu_memory_utilization = 0.85\n",
-        "\n",
-        "model, endpoint = deploy_model_vllm(\n",
-        "    model_name=get_job_name_with_datetime(prefix=\"mixtral-serve-vllm\"),\n",
-        "    model_id=prebuilt_model_id,\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
-        "    machine_type=machine_type,\n",
-        "    accelerator_type=accelerator_type,\n",
-        "    accelerator_count=accelerator_count,\n",
-        "    max_model_len=max_model_len,\n",
-        "    gpu_memory_utilization=gpu_memory_utilization,\n",
-        "    use_openai_server=False,\n",
-        "    use_chat_completions_if_openai_server=False,\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "agDw0_7JJ-MM"
-      },
-      "source": [
-        "NOTE: If you see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint, the model server is likely still initializing. Please retry later.\n",
-        "\n",
-        "NOTE: If you receive `InternalServerError: 500 System error` during the deployment, most likely the operation failed due to unavailability of resources. Either retry or use a different accelerator type.\n",
-        "\n",
-        "Once deployment succeeds, you can send requests to the endpoint with text prompts."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "zLyWJK5aJ-MM"
-      },
-      "source": [
-        "### Run sample prompt"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "NYN1Z49SJ-MM"
-      },
-      "outputs": [],
-      "source": [
-        "# Loads an existing endpoint instance using the endpoint name:\n",
-        "# - Using `endpoint_name = endpoint.name` allows us to get the endpoint name of\n",
-        "#   the endpoint `endpoint` created in the cell above.\n",
-        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
-        "#   an existing endpoint with the ID 1234567890123456789.\n",
-        "# You may uncomment the code below to load an existing endpoint.\n",
-        "\n",
-        "# endpoint_name = endpoint_without_peft.name\n",
-        "# # endpoint_name = \"\"  # @param {type:\"string\"}\n",
-        "# aip_endpoint_name = (\n",
-        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
-        "# )\n",
-        "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
-        "\n",
-        "instances = [\n",
-        "    {\n",
-        "        \"prompt\": \"What is a car?\",\n",
-        "        \"max_tokens\": 50,\n",
-        "        \"temperature\": 1.0,\n",
-        "        \"top_p\": 1.0,\n",
-        "        \"top_k\": 10,\n",
-        "    },\n",
-        "]\n",
-        "response = endpoint.predict(instances=instances)\n",
-        "\n",
-        "for prediction in response.predictions:\n",
-        "    print(prediction)\n",
-        "\n",
-        "# Reference the following code for using the OpenAI vLLM server.\n",
-        "# import json\n",
-        "# response = endpoint.raw_predict(\n",
-        "#     body=json.dumps({\n",
-        "#         \"model\": prebuilt_model_id,\n",
-        "#         \"prompt\": \"My favourite condiment is\",\n",
-        "#         \"n\": 1,\n",
-        "#         \"max_tokens\": 200,\n",
-        "#         \"temperature\": 1.0,\n",
-        "#         \"top_p\": 1.0,\n",
-        "#         \"top_k\": 10,\n",
-        "#     }),\n",
-        "#     headers={\"Content-Type\": \"application/json\"},\n",
-        "# )\n",
-        "# print(response.json())"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "NRaBADRM4JEn"
-      },
-      "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": "markdown",
-      "metadata": {
-        "id": "0a9da70d4abc"
-      },
-      "source": [
-        "### Undeploy models and Delete endpoints"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "e53749ae6b2c"
-      },
-      "outputs": [],
-      "source": [
-        "# Set this flag to delete endpoint including undeploying models\n",
-        "delete_endpoint = False"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "65078f3ec44e"
-      },
-      "outputs": [],
-      "source": [
-        "def list_endpoints():\n",
-        "    return [\n",
-        "        (r.name, r.display_name)\n",
-        "        for r in aiplatform.Endpoint.list()\n",
-        "        if r.display_name.startswith(\"mistral-serve-vllm\")\n",
-        "        or r.display_name.startswith(\"mixtral-serve-vllm\")\n",
-        "    ]"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "cf56ac4cc73b"
-      },
-      "outputs": [],
-      "source": [
-        "# Delete the endpoint using the Vertex AI fully qualified identifier for the endpoint\n",
-        "try:\n",
-        "    if delete_endpoint:\n",
-        "        endpoints = list_endpoints()\n",
-        "        for endpoint_id, endpoint_name in endpoints:\n",
-        "            endpoint = aiplatform.Endpoint(endpoint_id)\n",
-        "            print(\n",
-        "                f\"Undeploying all deployed models and deleting endpoint {endpoint_id} [{endpoint_name}]\"\n",
-        "            )\n",
-        "            endpoint.delete(force=True)\n",
-        "except Exception as e:\n",
-        "    print(e)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "d25a89a34b5e"
-      },
-      "source": [
-        "### Delete Cloud Storage bucket"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "PAr4UWWx4JEo"
-      },
-      "outputs": [],
-      "source": [
-        "import os\n",
-        "\n",
-        "delete_bucket = False\n",
-        "\n",
-        "job.delete()\n",
-        "\n",
-        "if delete_bucket or os.getenv(\"ID_TESTING\"):\n",
-        "    ! gsutil rm -rf {BUCKET_URI}"
-      ]
-    }
-  ],
-  "metadata": {
-    "colab": {
-      "name": "model_garden_pytorch_mistral.ipynb",
-      "toc_visible": true
-    },
-    "kernelspec": {
-      "display_name": "Python 3",
-      "name": "python3"
-    }
-  },
-  "nbformat": 4,
-  "nbformat_minor": 0
-}
@@ -59,8 +59,10 @@
        "\n",
        "- Deploy prebuilt [Mistral models](https://huggingface.co/mistralai) with [vLLM](https://github.com/vllm-project/vllm) containers\n",
        "    - [mistralai/Mistral-7B-v0.1](https://huggingface.co/mistralai/Mistral-7B-v0.1): pretrained generative text model with 7 billion parameters\n",
-        "    - [mistralai/Mistral-7B-Instruct-v0.1](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1): Instruction fine-tuned version of the Mistral-7B-v0.1 generative text model\n",
-        "    - [mistralai/Mistral-7B-Instruct-v0.2](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2): Improved instruction fine-tuned version of Mistral-7B-Instruct-v0.1 supporting 32k context length\n",
+        "    - [mistralai/Mistral-7B-Instruct-v0.1](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1): instruction fine-tuned version of the Mistral-7B-v0.1 generative text model\n",
+        "    - [mistralai/Mistral-7B-Instruct-v0.2](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2): improved instruction fine-tuned version of Mistral-7B-Instruct-v0.1 supporting 32k context length\n",
+        "    - [mistralai/Mistral-7B-v0.3](https://huggingface.co/mistralai/Mistral-7B-v0.3): Mistral-7B-v0.2 with extended vocabulary of 32768 and supports function calling\n",
+        "    - [mistralai/Mistral-7B-Instruct-v0.3](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.3): instruction fine-tuned version of the Mistral-7B-v0.3 generative text model\n",
        "\n",
        "### Costs\n",
        "\n",
@@ -156,7 +158,7 @@
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)\n",
        "\n",
        "# The pre-built serving docker images with vLLM\n",
-        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240313_0916_RC00\"\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240620_1616_RC00\"\n",
        "\n",
        "\n",
        "def get_job_name_with_datetime(prefix: str) -> str:\n",
@@ -218,6 +220,13 @@
        "        dtype = \"float16\"\n",
        "\n",
        "    vllm_args = [\n",
+        "        \"python\",\n",
+        "        \"-m\",\n",
+        "        (\n",
+        "            \"vllm.entrypoints.openai.api_server\"\n",
+        "            if use_openai_server\n",
+        "            else \"vllm.entrypoints.api_server\"\n",
+        "        ),\n",
        "        \"--host=0.0.0.0\",\n",
        "        \"--port=7080\",\n",
        "        f\"--model=gs://vertex-model-garden-public-us/{model_id}\",\n",
@@ -242,15 +251,6 @@
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
        "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
-        "        serving_container_command=[\n",
-        "            \"python\",\n",
-        "            \"-m\",\n",
-        "            (\n",
-        "                \"vllm.entrypoints.openai.api_server\"\n",
-        "                if use_openai_server\n",
-        "                else \"vllm.entrypoints.api_server\"\n",
-        "            ),\n",
-        "        ],\n",
        "        serving_container_args=vllm_args,\n",
        "        serving_container_ports=[7080],\n",
        "        serving_container_predict_route=serving_container_predict_route,\n",
@@ -282,9 +282,9 @@
        "\n",
        "# @markdown This section deploys the prebuilt Mistral model with [vLLM](https://github.com/vllm-project/vllm) on a Vertex endpoint. It takes 15 minutes to 1 hour to finish depending on the model and the accelerator.\n",
        "\n",
-        "# @markdown Set the model to deploy.\n",
+        "# @markdown Set the model to deploy and the accelerator to use.\n",
        "\n",
-        "prebuilt_model_id = \"mistralai/Mistral-7B-v0.1\"  # @param [\"mistralai/Mistral-7B-v0.1\", \"mistralai/Mistral-7B-Instruct-v0.1\", \"mistralai/Mistral-7B-Instruct-v0.2\"]\n",
+        "prebuilt_model_id = \"mistralai/Mistral-7B-v0.3\"  # @param [\"mistralai/Mistral-7B-v0.1\", \"mistralai/Mistral-7B-Instruct-v0.1\", \"mistralai/Mistral-7B-Instruct-v0.2\", \"mistralai/Mistral-7B-v0.3\", \"mistralai/Mistral-7B-Instruct-v0.3\"]\n",
        "# Find Vertex AI prediction supported accelerators and regions in\n",
        "# https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
        "\n",
@@ -304,7 +304,7 @@
        "id": "f0bd36bb-f988-479a-be40-18c54980ca1f"
      },
      "source": [
-        "You can manage your uploaded models in the [Model Registry](https://pantheon.corp.google.com/vertex-ai/models) and your endpoints in the [Endpoints](https://pantheon.corp.google.com/vertex-ai/endpoints)."
+        "You can manage your uploaded models in the [Model Registry](https://console.cloud.google.com/vertex-ai/models) and your endpoints in the [Endpoints](https://console.cloud.google.com/vertex-ai/endpoints)."
      ]
    },
    {
@@ -155,7 +155,7 @@
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
@@ -165,7 +165,7 @@
        "# The pre-built training docker images. They contain training scripts and models.\n",
        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20240318_0936_RC00\"\n",
        "# The pre-built serving docker images. They contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)\n",
        "\n",
@@ -0,0 +1,293 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7d9bbf86da5e"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2024 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "99c1c3fc2ca5"
+      },
+      "source": [
+        "# Vertex AI Model Garden - Stable Diffusion V2.1 (Local Dreambooth Finetune)\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_pytorch_sd_2_1_local_finetuning_dreambooth.ipynb\">\n",
+        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_sd_2_1_local_finetuning_dreambooth.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",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3de7470326a2"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates how to finetune [stabilityai/stable-diffusion-2-1](https://huggingface.co/stabilityai/stable-diffusion-2-1) with [Dreambooth](https://huggingface.co/docs/diffusers/training/dreambooth) locally in a Colab notebook and to test it with a local `text-2-image` prediction pipeline.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Finetune the stabilityai/stable-diffusion-2-1 model with [Dreambooth](https://huggingface.co/docs/diffusers/training/dreambooth) locally in a notebook.\n",
+        "- Run predictions for text-to-image in a local pipeline.\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Colab Enterprise\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloab Enterprise pricing](https://cloud.google.com/colab/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6546207faba5"
+      },
+      "source": [
+        "## Dreambooth Finetune in a notebook"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "969e7ccbc033"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Check if the Colab VM has GPU\n",
+        "\n",
+        "# @markdown **Important:** This notebook requires a GPU runtime to function correctly.\n",
+        "# @markdown The default Colab runtime does not have a GPU and will not work. Please\n",
+        "# @markdown create a GPU runtime by following the instructions at:\n",
+        "# @markdown   1. [Create a runtime template](https://cloud.google.com/vertex-ai/docs/colab/create-runtime-template#create)\n",
+        "# @markdown   1. [Create a runtime](https://cloud.google.com/vertex-ai/docs/colab/create-runtime#create) \\\n",
+        "# @markdown\n",
+        "# @markdown Once you have created a GPU runtime, you can use this notebook to run Dreambooth training locally within Colab.\n",
+        "\n",
+        "import subprocess\n",
+        "\n",
+        "if subprocess.run(\"nvidia-smi\").returncode:\n",
+        "    raise RuntimeError(\n",
+        "        \"Cannot communicate with GPU. Make sure you are using a GPU Colab runtime. \"\n",
+        "        \"Go to the Runtimes menu and select/create a runtime with GPUs.\"\n",
+        "    )"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "2707b02ef5df"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Prepare the virtual environment\n",
+        "\n",
+        "! pip install --upgrade pip\n",
+        "# Git clone the Huggingface diffusers code repo\n",
+        "! git clone --depth 1 --branch v0.25.1 https://github.com/huggingface/diffusers.git\n",
+        "\n",
+        "print(\"Installing diffusers from the source\")\n",
+        "! pip install -e /content/diffusers\n",
+        "print(\"Installing the requirements for dreambooth finetune\")\n",
+        "! pip install -r /content/diffusers/examples/dreambooth/requirements.txt\n",
+        "\n",
+        "! pip install bitsandbytes==0.43.1"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "d133319b6a53"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Prepare the example dataset\n",
+        "\n",
+        "# @markdown For this example, we'll download some images from Huggingface. If you\n",
+        "# @markdown have already had a dataset you wish to use, please choose the option\n",
+        "# @markdown `Use your own` and upload from your local computer.\n",
+        "#\n",
+        "import glob\n",
+        "import shutil\n",
+        "\n",
+        "from google.colab import files\n",
+        "from huggingface_hub import snapshot_download\n",
+        "from PIL import Image\n",
+        "\n",
+        "local_dir = \"/content/dreambooth/dog/\"\n",
+        "! rm -rf $local_dir\n",
+        "! mkdir -p $local_dir\n",
+        "\n",
+        "dataset_source = \"Huggingface\"  # @param [\"Huggingface\", \"Use your own\"]\n",
+        "if dataset_source == \"Huggingface\":\n",
+        "    snapshot_download(\n",
+        "        \"diffusers/dog-example\",\n",
+        "        local_dir=local_dir,\n",
+        "        repo_type=\"dataset\",\n",
+        "        ignore_patterns=\".gitattributes\",\n",
+        "    )\n",
+        "else:\n",
+        "    uploaded = files.upload()\n",
+        "    for name, data in uploaded.items():\n",
+        "        shutil.copy2(name, local_dir)\n",
+        "\n",
+        "\n",
+        "def image_grid(imgs, rows, cols, resize=256):\n",
+        "    if resize is not None:\n",
+        "        imgs = [img.resize((resize, resize)) for img in imgs]\n",
+        "    w, h = imgs[0].size\n",
+        "    grid = Image.new(\"RGB\", size=(cols * w, rows * h))\n",
+        "\n",
+        "    for i, img in enumerate(imgs):\n",
+        "        grid.paste(img, box=(i % cols * w, i // cols * h))\n",
+        "    return grid\n",
+        "\n",
+        "\n",
+        "# change path to display images from your local dir\n",
+        "img_paths = \"/content/dreambooth/dog/*.jpeg\"\n",
+        "imgs = [Image.open(path) for path in glob.glob(img_paths)]\n",
+        "\n",
+        "num_imgs_to_preview = 5\n",
+        "image_grid(imgs[:num_imgs_to_preview], 1, num_imgs_to_preview)\n",
+        "\n",
+        "! rm -rf $local_dir/.huggingface\n",
+        "! ls -alt $local_dir"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "8cbdbbc97728"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Train\n",
+        "\n",
+        "import locale\n",
+        "\n",
+        "locale.getpreferredencoding = lambda: \"UTF-8\"\n",
+        "! accelerate config default\n",
+        "\n",
+        "model_id = \"stabilityai/stable-diffusion-2-1\" # @param {type:\"string\"}\n",
+        "instance_prompt = \"a photo of sks dog\" # @param {type:\"string\"}\n",
+        "learning_rate = 2e-6 # @param {type:\"number\"}\n",
+        "resolution = 768 # @param {type:\"number\"}\n",
+        "train_steps = 200 # @param {type:\"number\"}\n",
+        "output_dir = \"/content/dreambooth/output_dir\"\n",
+        "\n",
+        "local_dir = \"/content/dreambooth/dog/\"\n",
+        "\n",
+        "!accelerate launch \\\n",
+        "  /content/diffusers/examples/dreambooth/train_dreambooth.py \\\n",
+        "  --pretrained_model_name_or_path=\"$model_id\" \\\n",
+        "  --instance_data_dir=\"$local_dir\" \\\n",
+        "  --class_data_dir=\"$local_dir\" \\\n",
+        "  --output_dir=\"$output_dir\" \\\n",
+        "  --instance_prompt=\"$instance_prompt\" \\\n",
+        "  --resolution=\"$resolution\" \\\n",
+        "  --learning_rate=\"$learning_rate\" \\\n",
+        "  --max_train_steps=\"$train_steps\" \\\n",
+        "  --mixed_precision=\"fp16\" \\\n",
+        "  --train_batch_size=1 \\\n",
+        "  --gradient_accumulation_steps=1 \\\n",
+        "  --gradient_checkpointing \\\n",
+        "  --lr_scheduler=\"constant\" \\\n",
+        "  --lr_warmup_steps=0 \\\n",
+        "  --use_8bit_adam \\\n",
+        "  --checkpointing_steps=100 \\\n",
+        "  --seed=42"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "356db37ed6ce"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Load the finetuned model checkpoint to a local diffusion pipeline.\n",
+        "\n",
+        "# @markdown `text-to-image` lets you send text prompts to the pipeline to generate images.\n",
+        "import torch\n",
+        "from diffusers import AutoPipelineForText2Image\n",
+        "\n",
+        "pipe = AutoPipelineForText2Image.from_pretrained(\n",
+        "    pretrained_model_or_path=output_dir,\n",
+        "    torch_dtype=torch.float16,\n",
+        "    use_safetensors=True,\n",
+        ").to(\"cuda\")\n",
+        "\n",
+        "prompt = \"a photo of sks dog in a bucket\"  # @param {type: \"string\"}\n",
+        "height = 768  # @param {type:\"number\"}\n",
+        "width = 768  # @param {type:\"number\"}\n",
+        "num_inference_steps = 25  # @param {type:\"number\"}\n",
+        "guidance_scale = 7.5  # @param {type:\"number\"}\n",
+        "\n",
+        "images = pipe(\n",
+        "    prompt=prompt,\n",
+        "    height=height,\n",
+        "    width=width,\n",
+        "    num_inference_steps=num_inference_steps,\n",
+        "    guidance_scale=guidance_scale,\n",
+        ").images\n",
+        "display(images[0])"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "name": "model_garden_pytorch_sd_2_1_local_finetuning_dreambooth.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
@@ -168,7 +168,7 @@
        "# The pre-built training docker images. They contain training scripts and models.\n",
        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20240318_0936_RC00\"\n",
        "# The pre-built serving docker images. They contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)\n",
        "\n",
@@ -336,12 +336,12 @@
        "    # @markdown **[Optional]** If using own images, provide the GCS path to the images.\n",
        "    # @markdown Make sure you have permission to access the bucket.\n",
        "    source_image_gcs_dir = \"gs://\"  # @param {type:\"string\"}\n",
-        "    print(\n",
-        "        f\"Now downloading the images from the source gcs directory: {source_image_gcs_dir}\"\n",
-        "    )\n",
        "    assert source_image_gcs_dir.startswith(\n",
        "        \"gs://\"\n",
        "    ), \"source_image_gcs_dir must start with `gs://`.\"\n",
+        "    print(\n",
+        "        f\"Now downloading the images from the source gcs directory: {source_image_gcs_dir}\"\n",
+        "    )\n",
        "    download_gcs_dir_to_local(source_image_gcs_dir, local_dir)\n",
        "\n",
        "\n",
@@ -375,6 +375,8 @@
        "class_dir_fuse = class_dir.replace(\"gs://\", \"/gcs/\")\n",
        "output_dir_fuse = output_dir.replace(\"gs://\", \"/gcs/\")\n",
        "\n",
+        "\n",
+        "concept_prompt = \"A picture of a sks dog\"  # @param {type:\"string\"}\n",
        "model = job.run(\n",
        "    args=[\n",
        "        \"--task=text-to-image-dreambooth-lora-sdxl\",\n",
@@ -382,19 +384,19 @@
        "        f\"--instance_data_dir={instance_dir_fuse}\",\n",
        "        f\"--class_data_dir={class_dir_fuse}\",\n",
        "        f\"--output_dir={output_dir_fuse}\",\n",
+        "        f\"--instance_prompt={concept_prompt}\",\n",
        "        \"--pretrained_vae_model_name_or_path=madebyollin/sdxl-vae-fp16-fix\",\n",
        "        \"--mixed_precision=fp16\",\n",
-        "        \"--instance_prompt='a photo of sks dog'\",\n",
        "        \"--resolution=1024\",\n",
        "        \"--train_batch_size=1\",\n",
-        "        \"--gradient_accumulation_steps=1\",\n",
+        "        \"--gradient_accumulation_steps=4\",\n",
        "        \"--gradient_checkpointing\",\n",
-        "        \"--learning_rate=2e-6\",\n",
-        "        \"--lr_scheduler=constant\",\n",
+        "        \"--learning_rate=1e-4\",\n",
+        "        \"--lr_scheduler=linear\",\n",
        "        \"--lr_warmup_steps=0 \",\n",
        "        \"--use_8bit_adam\",\n",
-        "        f\"--max_train_steps={train_steps}\",\n",
-        "        f\"--checkpointing_steps={train_steps // 2}\",\n",
+        "        f\"--max_train_steps={int(train_steps)}\",\n",
+        "        f\"--checkpointing_steps={int(train_steps // 2)}\",\n",
        "        \"--seed=0\",\n",
        "    ],\n",
        "    replica_count=num_nodes,\n",
@@ -467,7 +469,8 @@
        "\n",
        "# @markdown Once deployment succeeds, you can generate images by sending text prompts to the endpoint.\n",
        "\n",
-        "# @markdown You can also batch send prompts by separating them with a comma.\n",
+        "# @markdown You can also batch send prompts by separating them with a comma. Be sure to include\n",
+        "# @markdown the `concept_prompt` of your instance in the prompt.\n",
        "# @markdown You may adjust the parameters below to achieve best image quality.\n",
        "\n",
        "comma_separated_prompt_list = \"A picture of a sks dog in a house, A picture of a sks dog catching a frisbee\"  # @param {type: \"string\"}\n",
@@ -157,7 +157,7 @@
        "    auth.authenticate_user(project_id=PROJECT_ID)\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "# Define common functions.\n",
@@ -154,8 +154,7 @@
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "TEXT_TO_IMAGE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
-        "IMAGE_TO_IMAGE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve:sd-optimized\"\n",
+        "DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
@@ -200,17 +199,10 @@
        "        \"TASK\": task,\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
-        "    serving_container_image_uri = None\n",
-        "    if task == \"text-to-image\":\n",
-        "        serving_container_image_uri = TEXT_TO_IMAGE_DOCKER_URI\n",
-        "    elif task == \"image-to-image\":\n",
-        "        serving_container_image_uri = IMAGE_TO_IMAGE_DOCKER_URI\n",
-        "    else:\n",
-        "        raise ValueError(f\"Unknown task: {task}\")\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
-        "        serving_container_image_uri=serving_container_image_uri,\n",
+        "        serving_container_image_uri=DOCKER_URI,\n",
        "        serving_container_ports=[7080],\n",
        "        serving_container_predict_route=\"/predictions/diffusers_serving\",\n",
        "        serving_container_health_route=\"/ping\",\n",
@@ -166,7 +166,7 @@
        "# The pre-built training docker images. They contain training scripts and models.\n",
        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20240318_0936_RC00\"\n",
        "# The pre-built serving docker images. They contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "# Define common functions.\n",
@@ -244,7 +244,7 @@
        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20240202_1425_RC00\"\n",
        "\n",
        "# The pre-built serving docker images. They contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\""
      ]
    },
    {
@@ -239,7 +239,7 @@
        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20231130_0936_RC00\"\n",
        "\n",
        "# The pre-built serving docker images. They contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\""
      ]
    },
    {
@@ -53,9 +53,9 @@
      "source": [
        "## Overview\n",
        "\n",
-        "This notebook demonstrates starting a playground based on [Gradio UI](https://www.gradio.app/), inspired by the famous [Stable Diffusion WebUI](https://github.com/AUTOMATIC1111/stable-diffusion-webui) project, which allows users to interact with the stable diffusion models more easily and intuitively. The playground now support `text-to-image`, `image-inpainting`, `controlnet-canny`, `instruct-pix2pix`, and `SD 4x upscaler` tasks.\n",
+        "This notebook demonstrates starting a playground based on [Gradio UI](https://www.gradio.app/), inspired by the popular [Stable Diffusion WebUI](https://github.com/AUTOMATIC1111/stable-diffusion-webui) project, which allows users to interact with the stable diffusion models more easily and intuitively. The playground now support `text-to-image`, `image-inpainting`, `controlnet-canny`, `instruct-pix2pix`, and `SD 4x upscaler` tasks.\n",
        "\n",
-        "This notebook also supports Dreambooth finetune the stable diffusion 1.5/2.1 models.\n",
+        "This notebook also supports Dreambooth finetune the stable diffusion 1.5/2.1/XL models.\n",
        "\n",
        "### Objective\n",
        "\n",
@@ -121,7 +121,7 @@
        "\n",
        "# Cloud Storage bucket for storing the experiment artifacts.\n",
        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
-        "# prefer using your own GCS bucket, please change the value yourself below.\n",
+        "# prefer using your own GCS bucket, change the value yourself below.\n",
        "now = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
        "BUCKET_URI = \"gs://\"  # @param {type: \"string\"}\n",
        "\n",
@@ -160,7 +160,7 @@
        "# The pre-built training docker images. They contain training scripts and models.\n",
        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20240318_0936_RC00\"\n",
        "# The pre-built serving docker images. They contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240429_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
      ]
@@ -177,7 +177,7 @@
        "# @title Start the playground\n",
        "\n",
        "# @markdown This is a playground for image generation similar to the popular [Stable Diffusion WebUI](https://github.com/AUTOMATIC1111/stable-diffusion-webui).\n",
-        "# @markdown After the cell runs, this playground is avaible in a separate browser tab if you click the public URL.\n",
+        "# @markdown After the cell runs, this playground is available in a separate browser tab if you click the public URL.\n",
        "# @markdown Something similar to [\"https://####.gradio.live\"](#) in the output of the cell.\n",
        "\n",
        "# @markdown For Image generations, five tasks `text-to-image`, `image-inpainting`, `controlnet-canny`, `instruct-pix2pix` and `SD 4x upscaler` are currently supported.\n",
@@ -194,7 +194,7 @@
        "# @markdown 1. Don't forget to undeploy the models after all the experiment to avoid continuous charges to the project.\n",
        "\n",
        "# @markdown **Note: we support the following models now:**\n",
-        "# @markdown Other models (with the same task) may work, but they are not tested please use with caution.\n",
+        "# @markdown Other models (with the same task) may work, but they are not tested so use with caution.\n",
        "# @markdown 1. Text-to-Image models\n",
        "# @markdown    > runwayml/stable-diffusion-v1-5 \\\n",
        "# @markdown    > stabilityai/stable-diffusion-2-1 \\\n",
@@ -424,7 +424,7 @@
        "\n",
        "    if model_name not in model_to_task_dict.keys():\n",
        "        print(model_name)\n",
-        "        raise gr.Error(\"Please select a valid model name for Endpoint creation.\")\n",
+        "        raise gr.Error(\"Select a valid model name for Endpoint creation.\")\n",
        "\n",
        "    return model_to_task_dict[model_name]\n",
        "\n",
@@ -432,7 +432,7 @@
        "def deploy_model(model_name: str) -> aiplatform.Endpoint:\n",
        "    \"\"\"Preprocess the model_name, before deploying the model to vertex.\"\"\"\n",
        "    if not model_name:\n",
-        "        raise gr.Error(\"Please select a valid model name for model list.\")\n",
+        "        raise gr.Error(\"Select a valid model name for model list.\")\n",
        "\n",
        "    model_id = model_name.split(\": \")[1]\n",
        "    task_name = get_task_name(model_id)\n",
@@ -440,7 +440,9 @@
        "    return deploy_model_vertex(model_id, task_name)\n",
        "\n",
        "\n",
-        "def deploy_model_vertex(model_id: str, task_name: str) -> aiplatform.Endpoint:\n",
+        "def deploy_model_vertex(\n",
+        "    model_id: str, task_name: str, lora_id: str = \"\"\n",
+        ") -> aiplatform.Endpoint:\n",
        "    \"\"\"\n",
        "    Creates a new Vertex prediction endpoint and deploys a model to it.\n",
        "    The `model_id` can either be a HF model resource id, or a GCS path where\n",
@@ -470,6 +472,11 @@
        "            **serving_env,\n",
        "            \"REFINER_MODEL_ID\": refiner_model_id,\n",
        "        }\n",
+        "    if lora_id:\n",
+        "        serving_env = {\n",
+        "            **serving_env,\n",
+        "            \"LORA_ID\": lora_id,\n",
+        "        }\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_id,\n",
@@ -589,7 +596,7 @@
        "    image_dimension=512,\n",
        ") -> list[Image.Image]:\n",
        "    if not endpoint_name:\n",
-        "        raise gr.Error(\"Please select (or deploy) a model first!\")\n",
+        "        raise gr.Error(\"Select (or deploy) a model first!\")\n",
        "\n",
        "    prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)\n",
        "    instances = [\n",
@@ -624,7 +631,7 @@
        "    image_dimension=512,\n",
        ") -> list[Image.Image]:\n",
        "    if not endpoint_name:\n",
-        "        raise gr.Error(\"Please select (or deploy) a model first!\")\n",
+        "        raise gr.Error(\"Select (or deploy) a model first!\")\n",
        "\n",
        "    prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)\n",
        "    default_dimension = 512\n",
@@ -677,7 +684,7 @@
        "    image_dimension=512,\n",
        ") -> list[Image.Image]:\n",
        "    if not endpoint_name:\n",
-        "        raise gr.Error(\"Please select (or deploy) a model first!\")\n",
+        "        raise gr.Error(\"Select (or deploy) a model first!\")\n",
        "\n",
        "    prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)\n",
        "    instances = [\n",
@@ -714,7 +721,7 @@
        ") -> list[Image.Image]:\n",
        "\n",
        "    if not endpoint_name:\n",
-        "        raise gr.Error(\"Please select (or deploy) a model first!\")\n",
+        "        raise gr.Error(\"Select (or deploy) a model first!\")\n",
        "\n",
        "    prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)\n",
        "    default_dimension = 256\n",
@@ -756,7 +763,7 @@
        "    canny_high_threshold=200,\n",
        ") -> list[Image.Image]:\n",
        "    if not endpoint_name:\n",
-        "        raise gr.Error(\"Please select (or deploy) a model first!\")\n",
+        "        raise gr.Error(\"Select (or deploy) a model first!\")\n",
        "\n",
        "    prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)\n",
        "    init_image = init_image.convert(\"RGB\").resize((image_dimension, image_dimension))\n",
@@ -1154,12 +1161,14 @@
        "\n",
        "\n",
        "def get_dreambooth_base_model_and_resolution(base_model_select: str) -> tuple[str, int]:\n",
-        "    if base_model_select == \"v2-1-768\":\n",
+        "    if base_model_select == \"sd-v2-1-768\":\n",
        "        return \"stabilityai/stable-diffusion-2-1\", 768\n",
-        "    elif base_model_select == \"v2-1-512\":\n",
+        "    elif base_model_select == \"sd-v2-1-512\":\n",
        "        return \"stabilityai/stable-diffusion-2-1\", 512\n",
-        "    elif base_model_select == \"v1-5-512\":\n",
+        "    elif base_model_select == \"sd-v1-5-512\":\n",
        "        return \"runwayml/stable-diffusion-v1-5\", 512\n",
+        "    elif base_model_select == \"sdxl-dreambooth-lora-1024\":\n",
+        "        return \"stabilityai/stable-diffusion-xl-base-1.0\", 1024\n",
        "    else:\n",
        "        raise gr.Error(f\"Invalid base model: {base_model_select}\")\n",
        "\n",
@@ -1169,21 +1178,22 @@
        "    def prepare_instance_images(\n",
        "        base_model_select: str,\n",
        "        concept_prompt: str,\n",
-        "        train_steps: int,\n",
        "        file_collection: list[gr.File],\n",
        "    ):\n",
        "        \"\"\"\n",
        "        Prepares the instance images for training from the provided dir of images.\n",
        "        \"\"\"\n",
-        "        if not concept_prompt:\n",
-        "            raise gr.Error(\"Please enter a concept prompt first!\")\n",
-        "\n",
-        "        resolution = 512\n",
-        "        if base_model_select == \"v2-1-768\":\n",
-        "            resolution = 768\n",
-        "\n",
        "        if not file_collection:\n",
-        "            raise gr.Error(\"Please provide a few valid instance images first!\")\n",
+        "            raise gr.Error(gr.Markdown(\"Provide a few valid instance images!\"))\n",
+        "\n",
+        "        if not concept_prompt:\n",
+        "            raise gr.Error(\"Provide a unique concept prompt!\")\n",
+        "\n",
+        "        resolution = 768\n",
+        "        if base_model_select == \"sd-v2-1-512\" or base_model_select == \"sd-v1-5-512\":\n",
+        "            resolution = 512\n",
+        "        elif base_model_select == \"sdxl-dreambooth-lora-1024\":\n",
+        "            resolution = 1024\n",
        "\n",
        "        local_tmp_dir = \"/tmp/instance_images\"\n",
        "        if os.path.exists(local_tmp_dir):\n",
@@ -1207,6 +1217,18 @@
        "        instant_image_dir = os.path.join(BUCKET_URI, f\"instance_images-{now}\")\n",
        "        upload_local_dir_to_gcs(local_tmp_dir, instant_image_dir)\n",
        "\n",
+        "    def deploy_finetuned_model(\n",
+        "        base_model_select: str, output_dir: str\n",
+        "    ) -> aiplatform.Endpoint:\n",
+        "        if base_model_select == \"sdxl-dreambooth-lora-1024\":\n",
+        "            return deploy_model_vertex(\n",
+        "                model_id=\"stabilityai/stable-diffusion-xl-base-1.0\",\n",
+        "                lora_id=output_dir,\n",
+        "                task_name=\"text-to-image-sdxl\",\n",
+        "            )\n",
+        "\n",
+        "        return deploy_model_vertex(model_id=output_dir, task_name=\"text-to-image\")\n",
+        "\n",
        "    def dreambooth_start_training(*inputs):\n",
        "        [\n",
        "            output_dir,\n",
@@ -1224,16 +1246,18 @@
        "            seed,\n",
        "        ] = inputs[:13]\n",
        "\n",
-        "        file_collection = list(chain.from_iterable(inputs[13:]))\n",
        "        if not concept_prompt:\n",
-        "            raise gr.Error(\"Please enter a concept prompt first!\")\n",
+        "            raise gr.Error(\"Enter a concept prompt first!\")\n",
        "\n",
-        "        if not file_collection:\n",
-        "            raise gr.Error(\"Please provide a local folder with (5-10) instance image!\")\n",
+        "        if not inputs or len(inputs) < 13:\n",
+        "            raise gr.Error(\"Provide a local folder with (5-10) instance image!\")\n",
        "\n",
        "        if not base_model_select:\n",
-        "            raise gr.Error(\"Please select a base model!\")\n",
+        "            raise gr.Error(\"Select a base model!\")\n",
        "\n",
+        "        gr.Info(f\"Dreambooth finetune with task {base_model_select} started!\")\n",
+        "\n",
+        "        file_collection = list(chain.from_iterable(inputs[13:]))\n",
        "        base_model_id, resolution = get_dreambooth_base_model_and_resolution(\n",
        "            base_model_select\n",
        "        )\n",
@@ -1252,7 +1276,6 @@
        "        prepare_instance_images(\n",
        "            base_model_select=base_model_select,\n",
        "            concept_prompt=concept_prompt,\n",
-        "            train_steps=train_steps,\n",
        "            file_collection=list(file_collection),\n",
        "        )\n",
        "\n",
@@ -1269,10 +1292,14 @@
        "            container_uri=TRAIN_DOCKER_URI,\n",
        "        )\n",
        "\n",
-        "        gr.Info(f\"Custom Training job {job_name} is being submitted. Please wait...\")\n",
+        "        gr.Info(f\"Custom Training job {job_name} is being submitted. Please wait.\")\n",
+        "\n",
+        "        task = \"text-to-image-dreambooth\"\n",
+        "        if base_model_select == \"sdxl-dreambooth-lora-1024\":\n",
+        "            task = \"text-to-image-dreambooth-lora-sdxl\"\n",
        "\n",
        "        args = [\n",
-        "            \"--task=text-to-image-dreambooth\",\n",
+        "            f\"--task={task}\",\n",
        "            f\"--pretrained_model_name_or_path={base_model_id}\",\n",
        "            f\"--instance_data_dir={image_dir_gcsfuse}\",\n",
        "            f\"--class_data_dir={image_dir_gcsfuse}\",\n",
@@ -1287,9 +1314,15 @@
        "            f\"--max_train_steps={int(train_steps)}\",\n",
        "            f\"--checkpointing_steps={int(checkpointing_steps)}\",\n",
        "            f\"--seed={int(seed)}\",\n",
+        "            \"--gradient_checkpointing\",\n",
        "            \"--mixed_precision=fp16\",\n",
        "            \"--use_8bit_adam\",\n",
        "        ]\n",
+        "\n",
+        "        if base_model_select == \"sdxl-dreambooth-lora-1024\":\n",
+        "            args.append(\n",
+        "                \"--pretrained_vae_model_name_or_path=madebyollin/sdxl-vae-fp16-fix\"\n",
+        "            )\n",
        "        if train_text_encoder:\n",
        "            args.append(\"--train_text_encoder\")\n",
        "\n",
@@ -1303,7 +1336,21 @@
        "        )\n",
        "\n",
        "        if deploy_after_tuning:\n",
-        "            deploy_model_vertex(model_id=output_dir, task_name=\"text-to-image\")\n",
+        "            deploy_finetuned_model(base_model_select, output_dir)\n",
+        "\n",
+        "    def update_default_lr(base_model_select: str):\n",
+        "        \"\"\"Updates the default learning rate based on the selected base model.\"\"\"\n",
+        "        lr = 2e-6\n",
+        "        if base_model_select == \"sdxl-dreambooth-lora-1024\":\n",
+        "            lr = 1e-4\n",
+        "\n",
+        "        return {\n",
+        "            learning_rate: gr.update(value=lr),\n",
+        "        }\n",
+        "\n",
+        "    def update_default_checkpoing_steps(train_steps: int):\n",
+        "        \"\"\"Updates the default checkpointing_steps based on the training steps.\"\"\"\n",
+        "        return {checkpointing_steps: gr.update(value=int(train_steps // 2))}\n",
        "\n",
        "    dreambooth_tip_text = r\"\"\"\n",
        "        1. Upload 5-10 images of the object/face on from different angles/perspectives.\n",
@@ -1321,8 +1368,13 @@
        "            with gr.Column():\n",
        "                base_model_select = gr.Dropdown(\n",
        "                    label=\"Select a base model\",\n",
-        "                    choices=[\"v2-1-768\", \"v2-1-512\", \"v1-5-512\"],\n",
-        "                    value=\"v2-1-768\",\n",
+        "                    choices=[\n",
+        "                        \"sd-v2-1-768\",\n",
+        "                        \"sd-v2-1-512\",\n",
+        "                        \"sd-v1-5-512\",\n",
+        "                        \"sdxl-dreambooth-lora-1024\",\n",
+        "                    ],\n",
+        "                    value=\"sd-v2-1-768\",\n",
        "                    interactive=True,\n",
        "                )\n",
        "                concept_prompt = gr.Textbox(\n",
@@ -1344,7 +1396,7 @@
        "                        label=\"Train text encoder\", value=False, interactive=True\n",
        "                    )\n",
        "                    checkpointing_steps = gr.Number(\n",
-        "                        label=\"Checkpointing steps\", value=100, interactive=True\n",
+        "                        label=\"Checkpointing steps\", value=400, interactive=True\n",
        "                    )\n",
        "                    train_batch_size = gr.Number(\n",
        "                        label=\"Training batch size\", value=1, interactive=True\n",
@@ -1393,6 +1445,22 @@
        "                    visible=True,\n",
        "                )\n",
        "\n",
+        "    base_model_select.change(\n",
+        "        update_default_lr,\n",
+        "        base_model_select,\n",
+        "        [\n",
+        "            learning_rate,\n",
+        "        ],\n",
+        "    )\n",
+        "\n",
+        "    train_steps.change(\n",
+        "        update_default_checkpoing_steps,\n",
+        "        train_steps,\n",
+        "        [\n",
+        "            checkpointing_steps,\n",
+        "        ],\n",
+        "    )\n",
+        "\n",
        "    train_button.click(\n",
        "        dreambooth_start_training,\n",
        "        inputs=[\n",
@@ -159,7 +159,7 @@
        "    auth.authenticate_user(project_id=PROJECT_ID)\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "# Define common functions.\n",
@@ -159,7 +159,7 @@
        "    auth.authenticate_user(project_id=PROJECT_ID)\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "# Expand an image to a square image.\n",
@@ -149,7 +149,7 @@
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
@@ -158,7 +158,7 @@
        "    auth.authenticate_user(project_id=PROJECT_ID)\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "def base64_to_image(image_str):\n",
@@ -158,7 +158,7 @@
        "\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240429_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "# Define common functions.\n",
@@ -158,7 +158,7 @@
        "    auth.authenticate_user(project_id=PROJECT_ID)\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "# Define common functions.\n",
@@ -158,7 +158,7 @@
        "\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "# Define common functions.\n",
@@ -247,7 +247,7 @@
      "source": [
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
        "\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240223_1230_RC00\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\""
      ]
    },
    {
@@ -158,7 +158,7 @@
        "    auth.authenticate_user(project_id=PROJECT_ID)\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "def deploy_model(model_id, task):\n",
@@ -156,7 +156,7 @@
        "    auth.authenticate_user(project_id=PROJECT_ID)\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "# Define common functions.\n",
@@ -671,7 +671,7 @@
        "id": "s2gnYFEJ2AML"
      },
      "source": [
-        "You can mange your uploaded models in the [Model Registry](https://pantheon.corp.google.com/vertex-ai/models) and your endpoints in the [Endpoints](https://pantheon.corp.google.com/vertex-ai/endpoints)."
+        "You can mange your uploaded models in the [Model Registry](https://console.cloud.google.com/vertex-ai/models) and your endpoints in the [Endpoints](https://console.cloud.google.com/vertex-ai/endpoints)."
      ]
    },
    {
@@ -0,0 +1,522 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Rdr9qXnG1HaN"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2024 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "E_X6E0Jl1NbD"
+      },
+      "source": [
+        "## Model Garden RAG API\n",
+        "\n",
+        "Last updated: 7/24/2024\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_openai_api_llama3_1.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_openai_api_llama3_1.ipynb\"\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_openai_api_llama3_1.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_openai_api_llama3_1.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "P4tsSWku3wCh"
+      },
+      "source": [
+        "## Onboarding\n",
+        "If you have any questions, please reach out to *Vertex RAG API * team vertex-rag-eng@google.com, for the onboarding process.\n",
+        "\n",
+        "## 0. Set up the Environment and Test Project"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "u9mTxNC41S_1"
+      },
+      "outputs": [],
+      "source": [
+        "!pip3 install --force-reinstall google-cloud-aiplatform \"numpy<2.0.0\" --user"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ILNZ8_hw1WaC"
+      },
+      "outputs": [],
+      "source": [
+        "from google.colab import auth\n",
+        "\n",
+        "auth.authenticate_user()\n",
+        "\n",
+        "# Install gcloud\n",
+        "!pip install google-cloud"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "WuzltFSF1ZWw"
+      },
+      "source": [
+        "**Remember to restart after pip install.**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "yj8SdgZM1cOP"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "o7x73NiprHhJ"
+      },
+      "source": [
+        "## Initialization\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ldoxVA24qnAF"
+      },
+      "outputs": [],
+      "source": [
+        "import vertexai\n",
+        "from vertexai.preview import rag\n",
+        "from vertexai.preview.generative_models import GenerativeModel, Tool"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tB-fhAybq0T2"
+      },
+      "outputs": [],
+      "source": [
+        "# Set Project\n",
+        "PROJECT_ID = \"<your-project-id>\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "O2Rr1Ymlq3Uq"
+      },
+      "outputs": [],
+      "source": [
+        "vertexai.init(project=PROJECT_ID, location=\"us-central1\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cmxhW-LArK2L"
+      },
+      "source": [
+        "## Create a RAG corpus\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "5E1tVMx3rAXF"
+      },
+      "outputs": [],
+      "source": [
+        "# Currently supports Google first-party embedding models\n",
+        "embedding_model_config = rag.EmbeddingModelConfig(\n",
+        "    publisher_model=\"publishers/google/models/text-embedding-004\"\n",
+        ")\n",
+        "\n",
+        "# Name your corpus\n",
+        "DISPLAY_NAME = \"<your-corpus-display-name>\"  # @param {type:\"string\"}\n",
+        "\n",
+        "rag_corpus = rag.create_corpus(\n",
+        "    display_name=DISPLAY_NAME, embedding_model_config=embedding_model_config\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "16i1ZInQrFnL"
+      },
+      "outputs": [],
+      "source": [
+        "# Check the corpus just created\n",
+        "rag.list_corpora()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "sSLWYGF8rfMf"
+      },
+      "source": [
+        "## Upload a file to the corpus"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "4G5uyvbdraMY"
+      },
+      "outputs": [],
+      "source": [
+        "%%writefile test.txt\n",
+        "\n",
+        "Here's a demo for Llama3 RAG"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Z2vnvVO9rtDF"
+      },
+      "outputs": [],
+      "source": [
+        "rag_file = rag.upload_file(\n",
+        "    corpus_name=rag_corpus.name,\n",
+        "    path=\"test.txt\",\n",
+        "    display_name=\"test.txt\",\n",
+        "    description=\"my test\",\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "oOhO9-G2r1Gc"
+      },
+      "source": [
+        "## Import files from Google Cloud Storage\n",
+        "Remember to grant \"Viewer\" access to the \"Vertex RAG Data Service Agent\" (with the format of service-{project_number}@gcp-sa-vertex-rag.iam.gserviceaccount.com) for your Google Cloud Storage bucket"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "y64Hdd_9r5H9"
+      },
+      "outputs": [],
+      "source": [
+        "GS_BUCKET = \"gs://<your-gs-bucket-name>\"  # @param {type:\"string\"}\n",
+        "\n",
+        "response = await rag.import_files_async(  # noqa: F704\n",
+        "    corpus_name=rag_corpus.name,\n",
+        "    paths=[GS_BUCKET],\n",
+        "    chunk_size=512,\n",
+        "    chunk_overlap=50,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YiTAFiEasHLX"
+      },
+      "outputs": [],
+      "source": [
+        "# Check the files just imported. It may take a few seconds to process the imported files.\n",
+        "list(rag.list_files(corpus_name=rag_corpus.name))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bBsPDR-jsML5"
+      },
+      "source": [
+        "## Import files from Google Drive\n",
+        "Eligible paths can be https://drive.google.com/drive/folders/{folder_id} or https://drive.google.com/file/d/{file_id}.\n",
+        "\n",
+        "Remember to grant \"Viewer\" access to the \"Vertex RAG Data Service Agent\" (with the format of `service-{project_number}@gcp-sa-vertex-rag.iam.gserviceaccount.com`) for your Drive folder/files."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "u16-LvjT2Thi"
+      },
+      "outputs": [],
+      "source": [
+        "FILE_ID = \"<your-file-id>\"  # @param {type:\"string\"}\n",
+        "FILE_PATH = f\"https://drive.google.com/file/d/{FILE_ID}\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "iY4_6tshsPSA"
+      },
+      "outputs": [],
+      "source": [
+        "rag.import_files(\n",
+        "    corpus_name=rag_corpus.name,\n",
+        "    paths=[FILE_PATH],\n",
+        "    chunk_size=1024,\n",
+        "    chunk_overlap=100,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Dl8gPm9l4DQ3"
+      },
+      "outputs": [],
+      "source": [
+        "# Check the files just imported. It may take a few seconds to process the imported files.\n",
+        "list(rag.list_files(corpus_name=rag_corpus.name))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "I1-joFPys-FS"
+      },
+      "source": [
+        "## Using Rag Retrieval Tool and Generate Content API for non-self-deployed Llama3 MaaS endpoint\n",
+        "\n",
+        "When retrieval query similarity distance < vector_distance_threshold, generate content will cite the retrieved context.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "LvTPfijxtAQO"
+      },
+      "outputs": [],
+      "source": [
+        "rag_resource = rag.RagResource(\n",
+        "    rag_corpus=rag_corpus.name,\n",
+        "    # Need to manually get the ids from rag.list_files.\n",
+        "    # rag_file_ids=[],\n",
+        ")\n",
+        "\n",
+        "rag_retrieval_tool = Tool.from_retrieval(\n",
+        "    retrieval=rag.Retrieval(\n",
+        "        source=rag.VertexRagStore(\n",
+        "            rag_resources=[rag_resource],  # Currently only 1 corpus is allowed.\n",
+        "            similarity_top_k=5,\n",
+        "            vector_distance_threshold=0.4,\n",
+        "        ),\n",
+        "    )\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "T3IffRu42nRp"
+      },
+      "outputs": [],
+      "source": [
+        "ENDPOINT = f\"projects/{PROJECT_ID}/locations/us-central1/publishers/meta/models/llama3-405b-instruct-maas\"  # @param {type:\"string\"}\n",
+        "\n",
+        "rag_model = GenerativeModel(ENDPOINT, tools=[rag_retrieval_tool])"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "SFEEv2u0tVNz"
+      },
+      "outputs": [],
+      "source": [
+        "QUERY = \"What is RAG and why it is helpful?\"  # @param {type:\"string\"}\n",
+        "\n",
+        "response = rag_model.generate_content(QUERY)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "0I2EniAZtiUt"
+      },
+      "outputs": [],
+      "source": [
+        "response"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cje1WHtOtt2q"
+      },
+      "source": [
+        "## Using Rag Retrieval Tool with other generation API for non-self-deployed Llama3 MaaS endpoint\n",
+        "\n",
+        "The retrieved contexts can be passed to any SDK or model generation API to generate final results.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "dK7YmoIGtyki"
+      },
+      "outputs": [],
+      "source": [
+        "QUERY = \"What is RAG and why it is helpful?\"  # @param {type:\"string\"}\n",
+        "\n",
+        "rag_resource = rag.RagResource(\n",
+        "    rag_corpus=rag_corpus.name,\n",
+        "    # Need to manually get the ids from rag.list_files.\n",
+        "    # rag_file_ids=[],\n",
+        ")\n",
+        "\n",
+        "response = rag.retrieval_query(\n",
+        "    rag_resources=[rag_resource],  # Currently only 1 corpus is allowed.\n",
+        "    text=QUERY,\n",
+        "    similarity_top_k=5,\n",
+        "    vector_distance_threshold=0.4,\n",
+        ")\n",
+        "\n",
+        "# The retrieved context can be passed to any SDK or model generation API to generate final results.\n",
+        "retrieved_context = \" \".join(\n",
+        "    [context.text for context in response.contexts.contexts]\n",
+        ").replace(\"\\n\", \"\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "--GwKlAO29bZ"
+      },
+      "outputs": [],
+      "source": [
+        "retrieved_context"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "y9wTk2vW8tq0"
+      },
+      "source": [
+        "## Cleaning up\n",
+        "\n",
+        "Clean up resources created in this notebook."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "MRP8GZkw80aT"
+      },
+      "outputs": [],
+      "source": [
+        "delete_rag_corpus = True  # @param {type:\"boolean\"}\n",
+        "delete_bucket = False  # @param {type:\"boolean\"}\n",
+        "\n",
+        "if delete_rag_corpus:\n",
+        "    rag_corpus_list = rag.list_corpora()\n",
+        "    for rag_corpus in rag_corpus_list:\n",
+        "        rag.delete_corpus(name=rag_corpus.name)\n",
+        "\n",
+        "if delete_bucket:\n",
+        "    ! gsutil rm -r gs://{BUCKET_NAME}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kTHflSAhupiN"
+      },
+      "source": [
+        "## API reference\n",
+        "\n",
+        "For more details on RAG corpus/file management and detailed support please visit https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/rag-api\n"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_rag.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
@@ -0,0 +1,464 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NoEDALsivNDl"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2024 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "jCo1J-gNwJM2"
+      },
+      "source": [
+        "# Vertex AI Model Garden - Synthetic Data Generation using Llama 3.1\n",
+        "\n",
+        "<table><tbody><tr>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fsynthetic_data_generation_using_llama3_1.ipynb\">\n",
+        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/synthetic_data_generation_using_llama3_1.ipynb\">\n",
+        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3FlmJqqHyYwm"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates generating synthetic data using the [Llama 3.1 405B service API](https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama3-405b-instruct-maas).\n",
+        "\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "Leverage the Llama 3.1 405B service API to gnerate synthetic data. The framework is based on [Snowfakery](https://snowfakery.readthedocs.io/en/latest/) which is itself based on [Faker](https://faker.readthedocs.io/en/master/). It requires the expected outputs to be codified in a YAML file per Snowfakery specs, detailing all the required fields and their respective data generation strategies.\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "RlelO-pw5xaT"
+      },
+      "source": [
+        "## Steps"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "AjA_UYD25_1w"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "\n",
+        "!pip install --upgrade --user -q openai snowfakery==3.6.2 wikipedia-api==0.6.0\n",
+        "\n",
+        "import os\n",
+        "from datetime import datetime\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "# Define project information\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "\n",
+        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
+        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
+        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
+        "\n",
+        "# Cloud Storage bucket for storing the experiment artifacts.\n",
+        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
+        "# prefer using your own GCS bucket, change the value yourself below.\n",
+        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
+        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
+        "\n",
+        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
+        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\n",
+        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
+        "else:\n",
+        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
+        "    bucket_region = shell_output[0].strip().lower()\n",
+        "    if bucket_region != REGION:\n",
+        "        raise ValueError(\n",
+        "            \"Bucket region %s is different from notebook region %s\"\n",
+        "            % (bucket_region, REGION)\n",
+        "        )\n",
+        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"llama_3_1\")\n",
+        "\n",
+        "\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
+        "\n",
+        "from google.colab import auth\n",
+        "\n",
+        "auth.authenticate_user(project_id=PROJECT_ID)\n",
+        "\n",
+        "\n",
+        "import google.auth\n",
+        "\n",
+        "# Programmatically get an access token\n",
+        "creds, _ = google.auth.default(\n",
+        "    scopes=[\"https://www.googleapis.com/auth/cloud-platform\"]\n",
+        ")\n",
+        "auth_req = google.auth.transport.requests.Request()\n",
+        "creds.refresh(auth_req)\n",
+        "# Note: the credential lives for 1 hour by default (https://cloud.google.com/docs/authentication/token-types#at-lifetime); after expiration, it must be refreshed."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Yl8gDtmA75hD"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Creating Plugins and Prompts\n",
+        "\n",
+        "# @markdown The following cells create the 2 custom plugins we need for this use case along with the needed prompts.\n",
+        "\n",
+        "import logging\n",
+        "import sys\n",
+        "import types\n",
+        "from io import StringIO\n",
+        "\n",
+        "import jinja2\n",
+        "import openai\n",
+        "import wikipediaapi\n",
+        "from snowfakery import generate_data\n",
+        "from snowfakery.plugins import SnowfakeryPlugin\n",
+        "\n",
+        "MODEL_ID = \"meta/llama3-405b-instruct-maas\"\n",
+        "ENDPOINT = \"aiplatform.googleapis.com\"\n",
+        "\n",
+        "# Pass the Vertex endpoint and authentication to the OpenAI SDK\n",
+        "client = openai.OpenAI(\n",
+        "    base_url=f\"https://us-central1-{ENDPOINT}/v1beta1/projects/{PROJECT_ID}/locations/{REGION}/endpoints/openapi\",\n",
+        "    api_key=creds.token,\n",
+        ")\n",
+        "\n",
+        "\n",
+        "class SyntheticDataGeneration:\n",
+        "    \"\"\"\n",
+        "    Implements all the extra functionality needed for this use-case\n",
+        "    \"\"\"\n",
+        "\n",
+        "    # The first plugin allows us to interact with the Llama 3.1 405B service API.\n",
+        "    class Plugins(types.ModuleType):\n",
+        "        \"\"\"\n",
+        "        Provides the plugins needed to extend Snowfakery\n",
+        "        \"\"\"\n",
+        "\n",
+        "        class Llama3(SnowfakeryPlugin):\n",
+        "            \"\"\"\n",
+        "            Plugin for interacting with Llama3 API service.\n",
+        "            \"\"\"\n",
+        "\n",
+        "            class Functions:\n",
+        "                \"\"\"\n",
+        "                Functions to implement field / object level data generation\n",
+        "                \"\"\"\n",
+        "\n",
+        "                def fill_prompt(self, prompt_name: str, **kwargs) -> str:\n",
+        "                    \"\"\"\n",
+        "                    Returns a formatted prompt\n",
+        "                    \"\"\"\n",
+        "                    return (\n",
+        "                        jinja2.Environment(\n",
+        "                            loader=jinja2.FileSystemLoader(searchpath=\"./\")\n",
+        "                        )\n",
+        "                        .get_template(prompt_name)\n",
+        "                        .render(**kwargs)\n",
+        "                    )\n",
+        "\n",
+        "                def generate(\n",
+        "                    self,\n",
+        "                    prompt_name: str,\n",
+        "                    model=\"Llama3\",\n",
+        "                    temperature=0.9,\n",
+        "                    top_p=1,\n",
+        "                    **kwargs,\n",
+        "                ) -> str | None:\n",
+        "                    \"\"\"\n",
+        "                    A wrapper around Llama3 plugin\n",
+        "                    \"\"\"\n",
+        "                    prompt = self.fill_prompt(prompt_name, **kwargs)\n",
+        "                    try:\n",
+        "                        response = client.chat.completions.create(\n",
+        "                            model=MODEL_ID,\n",
+        "                            messages=[{\"role\": \"user\", \"content\": prompt}],\n",
+        "                            temperature=temperature,\n",
+        "                            top_p=top_p,\n",
+        "                        )\n",
+        "                        return response.choices[0].message.content\n",
+        "                    except Exception as e:\n",
+        "                        logging.trace(\n",
+        "                            (\n",
+        "                                \"Unable to generate text using %s.\\n\"\n",
+        "                                \"Prepared Prompt: \\n%s\\n\\nError: %s\"\n",
+        "                            ),\n",
+        "                            prompt_name,\n",
+        "                            prompt,\n",
+        "                            e,\n",
+        "                        )\n",
+        "                        return None\n",
+        "\n",
+        "        # The second plugin gives us the ability to interact with Wikipedia and fetch the contents for a given page.\n",
+        "        class Wikipedia(SnowfakeryPlugin):\n",
+        "            \"\"\"\n",
+        "            Plugin for interacting with Wikipedia.\n",
+        "            \"\"\"\n",
+        "\n",
+        "            class Functions:\n",
+        "                \"\"\"\n",
+        "                Implements a single function to fetch a Wikipedia page\n",
+        "                \"\"\"\n",
+        "\n",
+        "                def get_page(self, title: str):\n",
+        "                    \"\"\"\n",
+        "                    Returns the title, URL and sections of the given wikipedia page\n",
+        "                    \"\"\"\n",
+        "                    logging.info(\"Parsing Wikipedia Page %s\", title)\n",
+        "                    page = wikipediaapi.Wikipedia(\n",
+        "                        \"Snowfakery (example@google.com)\", \"en\"\n",
+        "                    ).page(title)\n",
+        "                    results = {\"sections\": {}, \"title\": page.title, \"url\": page.fullurl}\n",
+        "                    sections = [(s.title, s) for s in page.sections]\n",
+        "                    while sections:\n",
+        "                        sec_title, sec_obj = sections.pop()\n",
+        "                        if sec_title in [\n",
+        "                            \"External links\",\n",
+        "                            \"References\",\n",
+        "                            \"See also\",\n",
+        "                            \"Further reading\",\n",
+        "                        ]:\n",
+        "                            continue\n",
+        "                        if sec_obj.text:\n",
+        "                            results[\"sections\"][sec_title] = sec_obj.text\n",
+        "                        for sub_sec in sec_obj.sections:\n",
+        "                            sections.append((f\"{sec_title} - {sub_sec.title}\", sub_sec))\n",
+        "                    logging.info(\"Parsing Wikipedia Page %s Complete\", title)\n",
+        "                    return results"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "WSryene29Dan"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Making plugins discoverable\n",
+        "\n",
+        "# @markdown We add the created class to sys.modules to ensure Snowfakery can find them and import them as modules as needed.\n",
+        "\n",
+        "sys.modules[\"SyntheticDataGeneration.Plugins\"] = SyntheticDataGeneration.Plugins(\n",
+        "    name=\"SyntheticDataGeneration.Plugins\"\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "cJaPNlBi9bJr"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Creating Prompt Templates\n",
+        "\n",
+        "%%writefile blog_generator.jinja\n",
+        "You are an expert content creator who writes detailed, factual blogs.\n",
+        "You have been asked to write a blog about {{idea_title}}.\n",
+        "To get you started, you have also been given the following context about the topic:\n",
+        "\n",
+        "{{idea_body}}\n",
+        "\n",
+        "Ensure the blog that you write is interesting,detailed and factual.\n",
+        "Take a deep breath and start writing:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "BbOYKhD5-Bht"
+      },
+      "outputs": [],
+      "source": [
+        "%%writefile comment_generator.jinja\n",
+        "You are {{first_name}} {{last_name}}. You are {{age}} years old. You are interested in {{interests}}. You work at {{organization}} as a {{profession}}.\n",
+        "You came across the following article:\n",
+        "\n",
+        "{{blog_title}}\n",
+        "\n",
+        "{{blog_body}}\n",
+        "\n",
+        "Present your thoughts and feelings about the article in a short comment.\n",
+        "\n",
+        "Comment:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "sPe-UoD_-Ilo"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Creating the Recipe\n",
+        "\n",
+        "# @markdown In order to generate synthetic data, the schema of the synthetic data must be defined first. This is done by creating a recipe in a YAML format as demonstrated below, more details on writing recipes can be found [here](https://snowfakery.readthedocs.io/en/latest/#central-concepts).\n",
+        "\n",
+        "recipe = \"\"\"\n",
+        "- plugin: SyntheticDataGeneration.Plugins.Wikipedia\n",
+        "- plugin: SyntheticDataGeneration.Plugins.Llama3\n",
+        "- option: wiki_title\n",
+        "- var: __seed\n",
+        "  value:\n",
+        "    - Wikipedia.get_page :\n",
+        "      title : ${{wiki_title}}\n",
+        "\n",
+        "- object : users\n",
+        "  count : ${{random_number(min=100, max=500)}}\n",
+        "  fields :\n",
+        "    first_name : ${{fake.FirstName}}\n",
+        "    last_name : ${{fake.FirstName}}\n",
+        "    age:\n",
+        "      random_number:\n",
+        "        min: 18\n",
+        "        max: 95\n",
+        "    email : ${{fake.Email}}\n",
+        "    phone : ${{fake.PhoneNumber}}\n",
+        "    interests : ${{fake.Bs}}\n",
+        "    postal_code : ${{fake.Postalcode}}\n",
+        "    organization : ${{fake.Company}}\n",
+        "    profession : ${{fake.Job}}\n",
+        "\n",
+        "- object : seeds\n",
+        "  fields :\n",
+        "    title : ${{__seed['title']}}\n",
+        "    url : ${{__seed['url']}}\n",
+        "    section_count : ${{__seed['sections'] | length}}\n",
+        "\n",
+        "  friends:\n",
+        "    - object : blog_ideas\n",
+        "      count : ${{seeds.section_count}}\n",
+        "      fields :\n",
+        "        seed_id : ${{seeds.id}}\n",
+        "        section : ${{(__seed.sections.keys() | list)[child_index]}}\n",
+        "        body : ${{__seed.sections[section]}}\n",
+        "\n",
+        "      friends:\n",
+        "        - object : blog_posts\n",
+        "          fields :\n",
+        "            blog_idea_id : ${{blog_ideas.id}}\n",
+        "            title : ${{seeds.title}} - ${{blog_ideas.section}}\n",
+        "            body :\n",
+        "              - Llama3.generate:\n",
+        "                prompt_name : blog_generator.jinja\n",
+        "                idea_title : ${{title}}\n",
+        "                idea_body : ${{blog_ideas.body}}\n",
+        "            author : Llama3\n",
+        "\n",
+        "          friends:\n",
+        "            - object : blog_post_comments\n",
+        "              fields :\n",
+        "                blog_post_id : ${{blog_posts.id}}\n",
+        "                author_id :\n",
+        "                  random_reference : users\n",
+        "                author_email : ${{author_id.email}}\n",
+        "                comment :\n",
+        "                  - Llama3.generate:\n",
+        "                    prompt_name : comment_generator.jinja\n",
+        "                    first_name : ${{author_id.first_name}}\n",
+        "                    last_name : ${{author_id.last_name}}\n",
+        "                    age : ${{author_id.age}}\n",
+        "                    interests : ${{author_id.interests}}\n",
+        "                    organization : ${{author_id.organization}}\n",
+        "                    profession : ${{author_id.profession}}\n",
+        "                    blog_title : ${{blog_posts.title}}\n",
+        "                    blog_body : ${{blog_posts.body | truncate(1000)}}\n",
+        "\"\"\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "JTEQU-u--ei7"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Generating Data\n",
+        "\n",
+        "generate_data(\n",
+        "    StringIO(recipe),\n",
+        "    output_format=\"csv\",\n",
+        "    output_folder=\"outputs\",\n",
+        "    user_options={\"wiki_title\": \"Python_(programming_language)\"},\n",
+        ")\n",
+        "\n",
+        "# @markdown Results The synthetic data has been generated and stored as CSV files in the `outputs` folder."
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "synthetic_data_generation_using_llama3_1.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
@@ -96,7 +96,7 @@
      },
      "outputs": [],
      "source": [
-        "! pip3 install --upgrade --user --quiet google-cloud-aiplatform"
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform"
      ]
    },
    {
@@ -135,7 +135,7 @@
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
-        "<b>⚠️ The kernel is going to restart. Please wait until it is finished before continuing to the next step. ⚠️</b>\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
@@ -111,6 +111,7 @@ class ErrorCode(Enum):
    ERROR_LINK_GIT_BAD = 7,
    ERROR_LINK_COLAB_BAD = 8,
    ERROR_LINK_WORKBENCH_BAD = 9,
+    ERROR_LINK_COLAB_ENTERPRISE_BAD = 102,

    # Overview cells
    #   Overview cell required
@@ -320,7 +321,8 @@ def parse_notebook(path: str,
                  objective.steps, 
                  links.git_link, 
                  links.colab_link, 
-                  links.workbench_link
+                  links.colab_enterprise_link, 
+                  links.workbench_link,
        )
        
    if args.fix:
@@ -535,6 +537,7 @@ class LinksRule(NotebookRule):

        self.git_link = None
        self.colab_link = None
+        self.colab_enterprise_link = None
        self.workbench_link = None
        source = ''
        ret = True
@@ -565,6 +568,17 @@ class LinksRule(NotebookRule):
                        cell['source'][ix] = fix_link
                    else:
                        ret = notebook.report_error(ErrorCode.ERROR_LINK_COLAB_BAD, f"bad Colab link: {self.colab_link}")
+            
+            if '<a href="https://console.cloud.google.com/vertex-ai/colab/' in line:
+                self.colab_enterprise_link = line.strip()[9:-2].replace('" target="_blank', '').replace('" target=\'_blank', '')
+                modified_notebook_path = notebook.path.replace("/", "%2F")
+                derived_link = os.path.join('https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2F', modified_notebook_path)
+                if self.workbench_link != derived_link:
+                    if notebook.report_fix(FixCode.FIX_BAD_LINK, f"fixed Colab Enterprise link: {derived_link}"):
+                        fix_link = f"<a href=\"{derived_link}\" target='_blank'>\n"
+                        cell['source'][ix] = fix_link
+                    else:
+                        ret = notebook.report_error(ErrorCode.ERROR_LINK_COLAB_ENTERPRISE_BAD, f"bad Colab Enterprise link: {self.colab_enterprise_link}")

            if '<a href="https://console.cloud.google.com/vertex-ai/workbench/' in line:
                self.workbench_link = line.strip()[9:-2].replace('" target="_blank', '').replace('" target=\'_blank', '')
@@ -581,7 +595,7 @@ class LinksRule(NotebookRule):
        if 'View on GitHub' not in source or not self.git_link:
            ret = notebook.report_error(ErrorCode.ERROR_LINK_GIT_MISSING, 'Missing link for GitHub')
        if 'Run in Colab' not in source or not self.colab_link:
-            ret = notebook.report_error(ErrorCode.ERROR_LINK_COLAB_MISSING, 'Missing link for Colab')    
+            ret = notebook.report_error(ErrorCode.ERROR_LINK_COLAB_MISSING, 'Missing link for Colab')  
        if 'Open in Vertex AI Workbench' not in source or not self.workbench_link:
            ret = notebook.report_error(ErrorCode.ERROR_LINK_WORKBENCH_MISSING, 'Missing link for Workbench')
        
@@ -1120,6 +1134,7 @@ def add_index(path: str,
              steps: str, 
              git_link: str, 
              colab_link: str, 
+              colab_enterprise_link: str,
              workbench_link: str
             ):
    """
@@ -1133,6 +1148,7 @@ def add_index(path: str,
        steps: The steps specified by the notebook
        git_link: The link to the notebook in the git repo
        colab_link: Link to launch notebook in Colab
+        colab_enterpise_link: Link to launch notebook in Colab Enterprise
        workbench_link: Link to launch notebook in Workbench
        linkbacks: The linkbacks per tag
    """
@@ -1188,6 +1204,8 @@ def add_index(path: str,
        print('        <td>')
        if colab_link:
            print(f'            <a href="{colab_link}" target="_blank" class="external" track-type="notebookTutorial" track-name="colabLink">Colab</a><br/>\n')
+        if colab_enterprise_link:
+            print(f'            <a href="{colab_enterprise_link}" target="_blank" class="external" track-type="notebookTutorial" track-name="colabEnterpriseLink">Colab Enterprise</a><br/>\n')
        if git_link:
            print(f'            <a href="{git_link}" target="_blank" class="external" track-type="notebookTutorial" track-name="gitHubLink">GitHub</a><br/>\n')
        if workbench_link:
@@ -1278,7 +1296,7 @@ def replace_cl(text : str ) -> str:
        'Vertex Vizier': '{{vertex_vizier_name}}',
        'Vertex AI Vizier': '{{vertex_vizier_name}}',
        'Vizier': '{{vertex_vizier_name}}',
-        'Vertex Explainable AI': '{{vertex_xai_name}}',
+        'Vertex Explainable AI': '{{xai_name_short}}',
        'Explainable AI': '{{vertex_xai_name}}',
        'NAS': '{{vertex_nas_name_short}}',
        'Vertex AI Neural Architectural Search': '{{vertex_nas_name}}',
@@ -1288,6 +1306,10 @@ def replace_cl(text : str ) -> str:
        #'Vertex SDK': '{{vertex_sdk_name}}',
        #'Vertex AI SDK': '{{vertex_sdk_name}}',
        'Vertex AI': '{{vertex_ai_name}}',
+        'Vertex AI batch prediction': '{{vertex_ai_name}} {{batch_prediction_name}}',
+        'Vertex AI SDK for Python': '{{vertex_sdk_python}}',
+        'Ray on Vertex AI': '{{ray_vertex_ai_name}}',
+        'Google Cloud console': '{{console_name}}',
        
        'Cloud Storage': '{{storage_name}}',
        'GCS': '{{storage_name}}',
@@ -2,7 +2,7 @@
 # the repo. Unless a later match takes precedence,
 # @global-owner1 and @global-owner2 will be requested for
 # review when someone opens a pull request.
-* @GoogleCloudPlatform/vertex-ai-samples-contributors  @GoogleCloudPlatform/caiis-tw
+* @GoogleCloudPlatform/vertex-ai-samples-contributors @GoogleCloudPlatform/caiis-tw

 # matching_engine folder
 /matching_engine @shenzhimo2 @ivanmkc
@@ -27,6 +27,7 @@
 /training/xgboost_data_parallel_training_on_cpu_using_dask.ipynb @halio-g
 /feature_store/gapic-feature-store.ipynb @diemtvu
 /feature_store/vertex_ai_feature_store_feature_view_service_agents.ipynb @ethan-gordon
+/feature_store/vertex_ai_feature_store_iam_policy.ipynb @ethan-gordon
 /managed_notebooks @GoogleCloudPlatform/notebooks-team
 /pipelines/google_cloud_pipeline_components_bqml_text.ipynb @inardini
 /pipelines/google_cloud_pipelines_dataproc_tabular @inardini
@@ -53,4 +54,7 @@
 /model_evaluation/model_based_llm_evaluation/autosxs_check_alignment_against_human_preference_data.ipynb @sipingh
 /ray_on_vertex_ai/get_started_with_pytorch_rov.ipynb @inardini
 /training/tpuv5e_gemma_peft_finetuning_and_serving.ipynb @brianchunkang
-/training/tpuv5e_llama2_pytorch_finetuning_and_serving.ipynb @brianchunkang @chiefkarlin
+/training/tpuv5e_llama2_pytorch_finetuning_and_serving.ipynb @brianchunkang @chiefkarlin
+/prediction/get_started_with_psc_private_endpoint.ipynb @tianjiaoliu
+/ray_on_vertex_ai/spark_on_ray_on_vertex_ai.ipynb @ravi-dalal
+/generative_ai/mistralai_intro.ipynb @sujituk
@@ -23,6 +23,16 @@
        "# limitations under the License."
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0bdfa24a80ed"
+      },
+      "source": [
+        "Starting on September 15, 2024, you can only customize classification, entity extraction, and sentiment analysis models by moving to Vertex AI Gemini prompts and tuning. Training or updating models for Vertex AI AutoML for Text classification, entity extraction, and sentiment analysis objectives will no longer be available. You can continue using existing Vertex AI AutoML Text objectives until June 15, 2025. For more information about how Gemini offers enhanced user experience through improved prompting capabilities, see \n",
+        "[Introduction to tuning](https://cloud.google.com/vertex-ai/generative-ai/docs/models/tune-gemini-overview)"
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -32,25 +32,27 @@
        "# AutoML training image classification model for online prediction\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/automl/automl_image_classification_online_prediction.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/automl/automl_image_classification_online_prediction.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fautoml%2Fautoml_image_classification_online_prediction.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/automl//automl_image_classification_online_prediction.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",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
-        "</table>\n",
-        "<br/><br/><br/>"
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/automl/automl_image_classification_online_prediction.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
      ]
    },
    {
@@ -95,7 +97,7 @@
      "source": [
        "### Dataset\n",
        "\n",
-        "The dataset used for this tutorial is the [Flowers dataset](https://www.tensorflow.org/datasets/catalog/tf_flowers) from [TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/overview). The version of the dataset you will use in this tutorial is stored in a public Cloud Storage bucket. The trained model predicts the type of flower an image is from a class of five flowers: daisy, dandelion, rose, sunflower, or tulip."
+        "The dataset used for this tutorial is the [Flowers dataset](https://www.tensorflow.org/datasets/catalog/tf_flowers) from [TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/overview). The version of the dataset you use in this tutorial is stored in a public Cloud Storage bucket. The trained model predicts the type of flower an image is from a class of five flowers: daisy, dandelion, rose, sunflower, or tulip."
      ]
    },
    {
@@ -118,15 +120,22 @@
        "to generate a cost estimate based on your projected usage."
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f0316df526f8"
+      },
+      "source": [
+        "## Get started"
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "install_aip:mbsdk"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the latest version of Vertex AI SDK for Python."
+        "### Install Vertex AI SDK for Python and other required packages"
      ]
    },
    {
@@ -138,54 +147,86 @@
      "outputs": [],
      "source": [
        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
-        "                                 tensorflow"
+        "                                 tensorflow==2.15.1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "restart"
+        "id": "ff555b32bab8"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel"
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "D-ZBOjErv5mM"
+        "id": "f09b4dff629a"
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "before_you_begin:nogpu"
+        "id": "4a2b7b59bbf7"
      },
      "source": [
-        "## Before you begin"
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "before_you_begin:nogpu"
+        "id": "f82e28c631cc"
      },
      "source": [
-        "### Set your project ID\n",
+        "### Authenticate your notebook environment (Colab only)\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)"
+        "Authenticate your environment on Google Colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "46604f70e831"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "191d1345e064"
+      },
+      "source": [
+        "### Set Google Cloud project information\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project. Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
@@ -197,103 +238,7 @@
      "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": "gcp_authenticate"
-      },
-      "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": "FvQeFm3Gv5mR"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ad1138a125ea"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ce6043da7b33"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0367eac06a10"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "21ad4dbb4a61"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c13224697bfb"
-      },
-      "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."
+        "LOCATION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
@@ -335,7 +280,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
+        "! gsutil mb -l $LOCATION -p $PROJECT_ID $BUCKET_URI"
      ]
    },
    {
@@ -358,8 +303,6 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
        "import google.cloud.aiplatform as aiplatform"
      ]
    },
@@ -393,7 +336,7 @@
      "source": [
        "# Tutorial\n",
        "\n",
-        "Now you are ready to start creating your own AutoML image classification model."
+        "Now you're ready to start creating your own AutoML image classification model."
      ]
    },
    {
@@ -415,9 +358,7 @@
      },
      "outputs": [],
      "source": [
-        "IMPORT_FILE = (\n",
-        "    \"gs://cloud-samples-data/vision/automl_classification/flowers/all_data_v2.csv\"\n",
-        ")"
+        "IMPORT_FILE = \"gs://cloud-samples-data/ai-platform/flowers/flowers.csv\""
      ]
    },
    {
@@ -557,7 +498,7 @@
        "\n",
        "The `run` method when completed returns the `Model` resource.\n",
        "\n",
-        "The execution of the training pipeline will take upto 20 minutes."
+        "The execution of the training pipeline takes upto 20 minutes."
      ]
    },
    {
@@ -587,7 +528,7 @@
      "source": [
        "## Review model evaluation scores\n",
        "\n",
-        "After your model training has finished, you can review the evaluation scores for it using the `list_model_evaluations()` method. This method will return an iterator for each evaluation slice."
+        "After your model training has finished, you can review the evaluation scores for it using the `list_model_evaluations()` method. This method returns an iterator for each evaluation slice."
      ]
    },
    {
@@ -645,7 +586,7 @@
      "source": [
        "### Get test item\n",
        "\n",
-        "You will use an arbitrary example out of the dataset as a test item. Don't be concerned that the example was likely used in training the model -- we just want to demonstrate how to make a prediction."
+        "You use an arbitrary example out of the dataset as a test item. Don't be concerned that the example was likely used in training the model -- the point is to demonstrate how to make a prediction."
      ]
    },
    {
@@ -726,7 +667,7 @@
      "source": [
        "## Undeploy the model\n",
        "\n",
-        "When you are done doing predictions, you undeploy the model from the `Endpoint` resouce. This deprovisions all compute resources and ends billing for the deployed model."
+        "When you're done doing predictions, you undeploy the model from the `Endpoint` resouce. This deprovisions all compute resources and ends billing for the deployed model."
      ]
    },
    {
@@ -762,16 +703,11 @@
      },
      "outputs": [],
      "source": [
-        "delete_bucket = False\n",
-        "\n",
        "# Delete the dataset using the Vertex dataset object\n",
        "dataset.delete()\n",
        "\n",
-        "try:\n",
-        "    endpoint.undeploy_all()\n",
-        "    endpoint.delete()\n",
-        "except Exception as e:\n",
-        "    print(e)\n",
+        "# Delete the endpoint using the Vertex endpoint object\n",
+        "endpoint.delete()\n",
        "\n",
        "# Delete the model using the Vertex model object\n",
        "model.delete()\n",
@@ -779,8 +715,10 @@
        "# Delete the AutoML trainig job\n",
        "dag.delete()\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
-        "    ! gsutil rm -r $BUCKET_URI"
+        "# Delete Cloud Storage objects that were created\n",
+        "delete_bucket = False  # Set True for deletion\n",
+        "if delete_bucket:\n",
+        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
  ],
@@ -23,6 +23,16 @@
        "# limitations under the License."
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "148c71404373"
+      },
+      "source": [
+        "Starting on September 15, 2024, you can only customize classification, entity extraction, and sentiment analysis models by moving to Vertex AI Gemini prompts and tuning. Training or updating models for Vertex AI AutoML for Text classification, entity extraction, and sentiment analysis objectives will no longer be available. You can continue using existing Vertex AI AutoML Text objectives until June 15, 2025. For more information about how Gemini offers enhanced user experience through improved prompting capabilities, see \n",
+        "[Introduction to tuning](https://cloud.google.com/vertex-ai/generative-ai/docs/models/tune-gemini-overview)"
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -23,6 +23,16 @@
        "# limitations under the License."
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0bdfa24a80ed"
+      },
+      "source": [
+        "Starting on September 15, 2024, you can only customize classification, entity extraction, and sentiment analysis models by moving to Vertex AI Gemini prompts and tuning. Training or updating models for Vertex AI AutoML for Text classification, entity extraction, and sentiment analysis objectives will no longer be available. You can continue using existing Vertex AI AutoML Text objectives until June 15, 2025. For more information about how Gemini offers enhanced user experience through improved prompting capabilities, see \n",
+        "[Introduction to tuning](https://cloud.google.com/vertex-ai/generative-ai/docs/models/tune-gemini-overview)"
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -32,25 +32,27 @@
        "# Vertex AI SDK: AutoML tabular forecasting model for batch prediction\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/automl/sdk_automl_tabular_forecasting_batch.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/automl/sdk_automl_tabular_forecasting_batch.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fautoml%2Fsdk_automl_tabular_forecasting_batch.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/automl/sdk_automl_tabular_forecasting_batch.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",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
-        "</table>\n",
-        "<br/><br/><br/>"
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/automl/sdk_automl_tabular_forecasting_batch.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>\n"
      ]
    },
    {
@@ -62,9 +64,9 @@
        "## Overview\n",
        "\n",
        "\n",
-        "This tutorial demonstrates how to use the Vertex AI SDK to create tabular forecasting models and do batch prediction using a Google Cloud [AutoML](https://cloud.google.com/vertex-ai/docs/start/automl-users) model.\n",
+        "This tutorial demonstrates how to use the Vertex AI SDK to create tabular forecasting models and generate batch prediction using a Google Cloud [AutoML](https://cloud.google.com/vertex-ai/docs/start/automl-users) model.\n",
        "\n",
-        "Learn more about [Forecasting for tabular data](https://cloud.google.com/vertex-ai/docs/tabular-data/forecasting/overview)."
+        "Learn more about [Forecasting with AutoML](https://cloud.google.com/vertex-ai/docs/tabular-data/forecasting/overview)."
      ]
    },
    {
@@ -75,19 +77,19 @@
      "source": [
        "### Objective\n",
        "\n",
-        "In this tutorial, you learn how to create an `AutoML` tabular forecasting model from a Python script, and then do a batch prediction using the Vertex AI SDK. You can alternatively create and deploy models using the `gcloud` command-line tool or online using the Cloud Console.\n",
+        "In this tutorial, you learn how to create an AutoML tabular forecasting model from a Python script, and then generate batch prediction using the Vertex AI SDK. You can alternatively create and deploy models using the `gcloud` command-line tool or online using the Cloud Console.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
-        "- `AutoML Training`\n",
-        "- `Vertex AI Batch Prediction`\n",
-        "- `Vertex AI Model` resource\n",
+        "- AutoML Training\n",
+        "- Vertex AI batch prediction\n",
+        "- Vertex AI model resource\n",
        "\n",
        "The steps performed include:\n",
        "\n",
-        "- Create a `Vertex AI Dataset` resource.\n",
-        "- Train an `AutoML` tabular forecasting `Model` resource.\n",
-        "- Obtain the evaluation metrics for the `Model` resource.\n",
+        "- Create a Vertex AI dataset resource.\n",
+        "- Train an AutoML tabular forecasting model resource.\n",
+        "- Obtain the evaluation metrics for the model resource.\n",
        "- Make a batch prediction."
      ]
    },
@@ -99,7 +101,7 @@
      "source": [
        "### Dataset\n",
        "\n",
-        "The dataset used for this tutorial is a time series dataset containing samples drawn from the Iowa Liquor Retail Sales dataset. Data is made available by the Iowa Department of Commerce. It is provided under the Creative Commons Zero v1.0 Universal license. For more details, see: https://console.cloud.google.com/marketplace/product/iowa-department-of-commerce/iowa-liquor-sales. This dataset does not require any feature engineering. The version of the dataset you use in this tutorial is stored in BigQuery."
+        "The dataset used for this tutorial is a time series dataset containing samples drawn from the Iowa Liquor Retail Sales dataset. Data is made available by the Iowa Department of Commerce. It's provided under the Creative Commons Zero v1.0 Universal license. For more details, see: https://console.cloud.google.com/marketplace/product/iowa-department-of-commerce/iowa-liquor-sales. This dataset doesn't require any feature engineering. The version of the dataset used in this tutorial is stored in BigQuery."
      ]
    },
    {
@@ -125,12 +127,19 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "install_aip:mbsdk"
+        "id": "61RBz8LLbxCR"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook. "
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "No17Cw5hgx12"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
@@ -141,49 +150,86 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
        "! pip3 install --upgrade --quiet google-cloud-aiplatform"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "restart"
+        "id": "R5Xep4W9lq-Z"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel"
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "D-ZBOjErv5mM"
+        "id": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "yfEglUHQk9S3"
+        "id": "SbmM4z7FOBpM"
      },
      "source": [
-        "## Before you begin\n",
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dmWOrTJ3gx13"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
-        "### Set your project ID\n",
+        "Authenticate your environment on Google Colab.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NyKGtVQjgx13"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\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)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "DF4l8DTdWgPY"
+      },
+      "source": [
+        "### Set Google Cloud project information\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project. Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
@@ -195,91 +241,12 @@
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "LOCATION = \"us-central1\"  # @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": "gcp_authenticate"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
-        "\n",
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated.\n",
-        "\n",
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ce6043da7b33"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0367eac06a10"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "21ad4dbb4a61"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c13224697bfb"
-      },
-      "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": {
@@ -308,7 +275,7 @@
        "id": "create_bucket"
      },
      "source": [
-        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+        "**If your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
@@ -319,7 +286,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
+        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
@@ -341,8 +308,7 @@
      "source": [
        "import urllib\n",
        "\n",
-        "import google.cloud.aiplatform as aiplatform\n",
-        "from google.cloud import bigquery"
+        "from google.cloud import aiplatform, bigquery"
      ]
    },
    {
@@ -373,9 +339,9 @@
        "id": "tutorial_start:automl"
      },
      "source": [
-        "# Tutorial\n",
+        "## Tutorial\n",
        "\n",
-        "Now you are ready to start creating your own AutoML tabular forecasting model."
+        "Now you're ready to begin creating your own AutoML tabular forecasting model."
      ]
    },
    {
@@ -410,11 +376,11 @@
      "source": [
        "### Create the Dataset\n",
        "\n",
-        "Next, create the `Dataset` resource using the `create` method for the `TimeSeriesDataset` class, which takes the following parameters:\n",
+        "Next, create the dataset resource by using the `create` method of the `TimeSeriesDataset` class, which takes the following parameters:\n",
        "\n",
-        "- `display_name`: The human readable name for the `Dataset` resource.\n",
-        "- `gcs_source`: A list of one or more dataset index files to import the data items into the `Dataset` resource.\n",
-        "- `bq_source`: Alternatively, import data items from a BigQuery table into the `Dataset` resource.\n",
+        "- `display_name`: The human readable name for the dataset resource.\n",
+        "- `gcs_source`: A list of one or more dataset index files to import the data items into the dataset resource.\n",
+        "- `bq_source`: Alternatively, import data items from a BigQuery table into the dataset resource.\n",
        "\n",
        "This operation may take several minutes."
      ]
@@ -464,13 +430,15 @@
      "source": [
        "### Create and run training job\n",
        "\n",
-        "To train an AutoML model, you perform two steps: 1) create a training job, and 2) run the job.\n",
+        "To train an AutoML model, you perform two steps: \n",
+        "1) Create a training job.\n",
+        "2) Specify your training parameters and run the job.\n",
        "\n",
        "#### Create training job\n",
        "\n",
-        "An AutoML training job is created with the `AutoMLForecastingTrainingJob` class, with the following parameters:\n",
+        "An AutoML training job is created using the `AutoMLForecastingTrainingJob` class, with the following parameters:\n",
        "\n",
-        "- `display_name`: The human readable name for the `TrainingJob` resource.\n",
+        "- `display_name`: The human readable name for the training job resource.\n",
        "- `column_transformations`: (Optional): Transformations to apply to the input columns\n",
        "- `optimization_objective`: The optimization objective to minimize or maximize.\n",
        "    - `minimize-rmse`\n",
@@ -505,20 +473,27 @@
      "source": [
        "#### Run the training pipeline\n",
        "\n",
-        "Next, you start the training job by invoking the method `run`, with the following parameters:\n",
+        "Next, start the training job by invoking the `run` method, with the following parameters:\n",
        "\n",
-        "- `dataset`: The `Dataset` resource to train the model.\n",
-        "- `model_display_name`: The human readable name for the trained model.\n",
-        "- `training_fraction_split`: The percentage of the dataset to use for training.\n",
-        "- `test_fraction_split`: The percentage of the dataset to use for test (holdout data).\n",
-        "- `target_column`: The name of the column to train as the label.\n",
-        "- `budget_milli_node_hours`: (optional) Maximum training time specified in unit of millihours (1000 = hour).\n",
+        "- `dataset`: The dataset resource to train the model.\n",
+        "- `target_column`: The column in the dataset that contains the values the model is trying to forecast.\n",
        "- `time_column`: Time-series column for the forecast model.\n",
        "- `time_series_identifier_column`: ID column for the time-series column.\n",
+        "- `available_at_forecast_columns`: List of columns that are available at the time of forecasting.\n",
+        "- `unavailable_at_forecast_columns`: List of columns that aren't available at the time of forecasting.\n",
+        "- `time_series_attribute_columns`: Columns that contain attributes or metadata related to the time series data, such as \"city,\" \"zip_code,\" and \"county\" in this example. These attributes can help the model understand the context of the time series.\n",
+        "- `forecast_horizon`: It determines how far into the future you want to predict, representing the number of time steps ahead for which the model generates predictions.\n",
+        "- `context_window`: The number of historical time steps the model uses as context for making predictions. A context window of 30 means the model uses data from the past 30 time steps to forecast future values.\n",
+        "- `data_granularity_unit`: The unit of time used for granularity in the data, such as \"day\" or \"hour.\" This specifies the time interval between data points.\n",
+        "- `data_granularity_count`: The count of the granularity unit. For example, a data_granularity_count of 1 with a `data_granularity_unit` of \"day\" means each data point represents one day.\n",
+        "- `weight_column`: This parameter lets you assign different weights to different data points in your training set.\n",
+        "- `budget_milli_node_hours`: Maximum training time specified in unit of millihours (1000 = hour).\n",
+        "- `model_display_name`: The human readable name for the trained model.\n",
+        "- `predefined_split_column_name`: The name of a column used to specify predefined splits for training and evaluation. If not used, it’s set to `None`.\n",
        "\n",
-        "The `run` method when completed returns the `Model` resource.\n",
+        "The `run` method when completed returns the model resource.\n",
        "\n",
-        "The execution of the training pipeline will take up to one hour."
+        "The execution of the training pipeline may take up to one hour."
      ]
    },
    {
@@ -556,7 +531,7 @@
      "source": [
        "## Review model evaluation scores\n",
        "\n",
-        "After your model training has finished, you can review the evaluation scores for "
+        "Once your model training is complete, you can examine the evaluation scores to assess the model performance"
      ]
    },
    {
@@ -592,15 +567,15 @@
      "source": [
        "### Make the batch prediction request\n",
        "\n",
-        "Now that your Model resource is trained, you can make a batch prediction by invoking the batch_predict() method using a BigQuery source and destination, with the following parameters:\n",
+        "Now that your Model resource is trained, you can make a batch prediction by invoking the `batch_predict()` method using a BigQuery source and destination, with the following parameters:\n",
        "\n",
        "- `job_display_name`: The human readable name for the batch prediction job.\n",
        "- `bigquery_source`: BigQuery URI to a table, up to 2000 characters long. For example: `bq://projectId.bqDatasetId.bqTableId`\n",
-        "- `bigquery_destination_prefix`: The BigQuery dataset or table for storing the batch prediction resuls.\n",
+        "- `bigquery_destination_prefix`: The BigQuery dataset or table for storing the batch prediction results.\n",
        "- `instances_format`: The format for the input instances. Since a BigQuery source is used here, this should be set to `bigquery`.\n",
        "- `predictions_format`: The format for the output predictions, `bigquery` is used here to output to a BigQuery table.\n",
        "- `generate_explanations`: Set to `True` to generate explanations.\n",
-        "- `sync`: If set to True, the call will block while waiting for the asynchronous batch job to complete."
+        "- `sync`: Set **True** to wait until the completion of the job."
      ]
    },
    {
@@ -618,11 +593,11 @@
        "batch_predict_bq_output_uri_prefix = \"bq://{}.{}\".format(\n",
        "    PROJECT_ID, batch_predict_bq_output_dataset_name\n",
        ")\n",
-        "# Must be the same region as batch_predict_bq_input_uri\n",
+        "# Must be the same location as batch_predict_bq_input_uri\n",
        "client = bigquery.Client(project=PROJECT_ID)\n",
        "bq_dataset_id = bigquery.Dataset(batch_predict_bq_output_dataset_path)\n",
-        "dataset_region = \"US\"  # @param {type : \"string\"}\n",
-        "bq_dataset_id.location = dataset_region\n",
+        "dataset_location = \"US\"  # @param {type : \"string\"}\n",
+        "bq_dataset_id.location = dataset_location\n",
        "# delete any existing dataset\n",
        "try:\n",
        "    client.delete_dataset(bq_dataset_id, delete_contents=True)\n",
@@ -631,7 +606,7 @@
        "bq_dataset = client.create_dataset(bq_dataset_id)\n",
        "print(\n",
        "    \"Created bigquery dataset {} in {}\".format(\n",
-        "        batch_predict_bq_output_dataset_path, dataset_region\n",
+        "        batch_predict_bq_output_dataset_path, dataset_location\n",
        "    )\n",
        ")"
      ]
@@ -642,9 +617,9 @@
        "id": "99b7a9287ba6"
      },
      "source": [
-        "For AutoML models, manual scaling can be adjusted by setting both min and max nodes i.e., `starting_replica_count` and `max_replica_count` as the same value(in this example, set to 1). The node count can be increased or decreased as required by load.\n",
+        "For AutoML models, manual scaling can be adjusted by setting both min and max nodes i.e., `starting_replica_count` and `max_replica_count` as the same value(in this example, set to 1). The node count can be increased or decreased as required by the load\n",
        " \n",
-        "`batch_predict` can export predictions either to BigQuery or GCS. This example exports to BigQuery."
+        "The `batch_predict` method can export predictions either to BigQuery or GCS. In this example, the predictions are exported to BigQuery"
      ]
    },
    {
@@ -680,7 +655,7 @@
      "source": [
        "### Wait for completion of batch prediction job\n",
        "\n",
-        "Next, wait for the batch job to complete. Alternatively, you can set the parameter `sync` to `True` in the `batch_predict()` method to block until the batch prediction job is completed."
+        "Next, wait for the batch job to complete. Alternatively, you can set the `sync` parameter to `True` in the `batch_predict()` method to wait until the batch prediction job is completed."
      ]
    },
    {
@@ -702,7 +677,7 @@
      "source": [
        "### Get the predictions and explanations\n",
        "\n",
-        "Next, get the results from the completed batch prediction job and print them out. Each result row will include the prediction and explanation."
+        "Next, get the results from the completed batch prediction job and print them out. Each result row includes the prediction and explanation."
      ]
    },
    {
@@ -864,14 +839,14 @@
        "\n",
        "# Delete the dataset\n",
        "try:\n",
-        "    client.delete_dataset(bq_dataset_id)\n",
+        "    client.delete_dataset(bq_dataset_id, delete_contents=True, not_found_ok=True)\n",
        "except Exception as e:\n",
        "    print(e)\n",
        "\n",
        "# Set this to true only if you'd like to delete your bucket\n",
-        "delete_bucket = False\n",
+        "delete_bucket = False  # set True for deletion\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil rm -r $BUCKET_URI"
      ]
    }
@@ -32,24 +32,27 @@
        "# Vertex AI SDK for Python: AutoML training tabular regression model for batch prediction using BigQuery\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/automl/sdk_automl_tabular_regression_batch_bq.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fautoml%2Fsdk_automl_tabular_regression_batch_bq.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/automl/sdk_automl_tabular_regression_batch_bq.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/automl/sdk_automl_tabular_regression_batch_bq.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
-        "  <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/automl/sdk_automl_tabular_regression_batch_bq.ipynb\" target='_blank'>\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>"
+        "</table>\n"
      ]
    },
    {
@@ -61,7 +64,7 @@
        "## Overview\n",
        "\n",
        "\n",
-        "This tutorial demonstrates how to use the Vertex AI SDK for Python to create tabular regression models and do batch prediction using a Google Cloud [AutoML](https://cloud.google.com/vertex-ai/docs/start/automl-users) model.\n",
+        "This tutorial demonstrates how to use the Vertex AI SDK for Python to create tabular regression models and generate batch prediction using a Google Cloud [AutoML](https://cloud.google.com/vertex-ai/docs/start/automl-users) model.\n",
        "\n",
        "Learn more about [Regression for tabular data](https://cloud.google.com/vertex-ai/docs/tabular-data/classification-regression/overview)."
      ]
@@ -81,17 +84,14 @@
        "- Vertex AI Datasets (Tabular)\n",
        "- Vertex AI Training (AutoML Tabular Training)\n",
        "- Vertex AI Model Registry\n",
-        "- Vertex AI Endpoint\n",
-        "- Vertex AI Batch predictions\n",
+        "- Vertex AI batch prediction\n",
        "\n",
        "The steps performed include:\n",
        "\n",
-        "- Create a Vertex AI `Dataset` resource.\n",
-        "- Train the model.\n",
-        "- View the model evaluation.\n",
-        "- Deploy the `Model` resource to a serving `Endpoint` resource.\n",
-        "- Make a prediction.\n",
-        "- Undeploy the `Model`."
+        "- Create a Vertex AI dataset resource.\n",
+        "- Train an AutoML tabular regression model resource.\n",
+        "- Obtain the evaluation metrics for the model resource.\n",
+        "- Make a batch prediction."
      ]
    },
    {
@@ -102,7 +102,7 @@
      "source": [
        "### Dataset\n",
        "\n",
-        "The dataset used for this tutorial is the [GSOD dataset](https://console.cloud.google.com/marketplace/product/noaa-public/gsod) from [BigQuery public datasets](https://cloud.google.com/bigquery/public-data). In this version of the dataset, you use the year, month, and day fields to predict the mean daily temperature (mean_temp)."
+        "The dataset used for this tutorial is the [GSOD dataset](https://console.cloud.google.com/marketplace/product/noaa-public/gsod) from [BigQuery public datasets](https://cloud.google.com/bigquery/public-data). In this version of the dataset, you use the year, month, and day fields to predict the mean daily temperature (`mean_temp`)."
      ]
    },
    {
@@ -116,7 +116,6 @@
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
-        "* Cloud Storage\n",
        "* BigQuery / BigQuery ML\n",
        "\n",
        "Learn about [Vertex AI\n",
@@ -127,12 +126,19 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "install_aip:mbsdk"
+        "id": "61RBz8LLbxCR"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook."
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "No17Cw5hgx12"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
@@ -153,40 +159,80 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "restart"
+        "id": "R5Xep4W9lq-Z"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel."
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "restart"
+        "id": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "before_you_begin:nogpu"
+        "id": "SbmM4z7FOBpM"
      },
      "source": [
-        "## Before you begin\n",
-        "#### Set your project ID\n",
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dmWOrTJ3gx13"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\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)"
+        "Authenticate your environment on Google Colab.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NyKGtVQjgx13"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "DF4l8DTdWgPY"
+      },
+      "source": [
+        "### Set Google Cloud project information\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project. Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
@@ -198,90 +244,12 @@
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "LOCATION = \"us-central1\"  # @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 = \"[your-region]\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "gcp_authenticate"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
-        "\n",
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated.\n",
-        "**2. Local JupyterLab instance, uncomment and run:**\n"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "457c78b08293"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "d3e571ce6c56"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "984a0526fb68"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "2c549a59cca4"
-      },
-      "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": {
@@ -302,8 +270,7 @@
      },
      "outputs": [],
      "source": [
-        "import google.cloud.aiplatform as aiplatform\n",
-        "from google.cloud import bigquery"
+        "from google.cloud import aiplatform, bigquery"
      ]
    },
    {
@@ -325,7 +292,7 @@
      },
      "outputs": [],
      "source": [
-        "aiplatform.init(project=PROJECT_ID, location=REGION)"
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
@@ -334,9 +301,9 @@
        "id": "tutorial_start:automl"
      },
      "source": [
-        "# Tutorial\n",
+        "## Tutorial\n",
        "\n",
-        "Now you are ready to start creating your own AutoML tabular regression model."
+        "Now you're ready to start creating your own AutoML tabular regression model."
      ]
    },
    {
@@ -345,7 +312,7 @@
        "id": "import_file:u_dataset,bq"
      },
      "source": [
-        "#### Location of BigQuery training data.\n",
+        "### Location of BigQuery training data.\n",
        "\n",
        "Set the `IMPORT_File` variable to the location of the data table in BigQuery."
      ]
@@ -380,7 +347,7 @@
      },
      "outputs": [],
      "source": [
-        "# Create client in default region\n",
+        "# Create client in default location\n",
        "bq_client = bigquery.Client(\n",
        "    project=PROJECT_ID,\n",
        "    credentials=aiplatform.initializer.global_config.credentials,\n",
@@ -395,13 +362,13 @@
      },
      "outputs": [],
      "source": [
-        "# Create training dataset in default region\n",
+        "# Create training dataset in default location\n",
        "TRAINING_INPUT_DATASET_ID = \"gsod_training_unique\"\n",
        "bq_dataset = bigquery.Dataset(f\"{PROJECT_ID}.{TRAINING_INPUT_DATASET_ID}\")\n",
        "bq_dataset = bq_client.create_dataset(bq_dataset)\n",
        "print(f\"Created dataset {bq_client.project}.{bq_dataset.dataset_id}\")\n",
        "\n",
-        "# Create test dataset in default region\n",
+        "# Create test dataset in default location\n",
        "PREDICTION_INPUT_DATASET_ID = \"gsod_prediction_unique\"\n",
        "bq_dataset = bigquery.Dataset(f\"{PROJECT_ID}.{PREDICTION_INPUT_DATASET_ID}\")\n",
        "bq_dataset = bq_client.create_dataset(bq_dataset)\n",
@@ -463,11 +430,11 @@
      "source": [
        "### Create the Dataset\n",
        "\n",
-        "Use `TabularDataset.create()` to create a `TabularDataset` resource, which takes the following parameters:\n",
+        "Use `TabularDataset.create()` method to create a tabular dataset resource, which takes the following parameters:\n",
        "\n",
-        "- `display_name`: The human readable name for the `Dataset` resource.\n",
-        "- `gcs_source`: A list of one or more dataset index files to import the data items into the `Dataset` resource.\n",
-        "- `bq_source`: Alternatively, import data items from a BigQuery table into the `Dataset` resource.\n",
+        "- `display_name`: The human readable name for the dataset resource.\n",
+        "- `gcs_source`: A list of one or more dataset index files to import the data items into the dataset resource.\n",
+        "- `bq_source`: Alternatively, import data items from a BigQuery table into the dataset resource.\n",
        "\n",
        "This operation may take several minutes."
      ]
@@ -521,11 +488,11 @@
        "\n",
        "Create an AutoML training pipeline using the `AutoMLTabularTrainingJob` class, with the following parameters:\n",
        "\n",
-        "- `display_name`: The human readable name for the `TrainingJob` resource.\n",
+        "- `display_name`: The human readable name for the training job resource.\n",
        "- `optimization_prediction_type`: The type task to train the model for.\n",
        "  - `classification`: A tabular classification model.\n",
        "  - `regression`: A tabular regression model.\n",
-        "- `column_transformations`: (Optional): Transformations to apply to the input columns\n",
+        "- `column_transformations`: (Optional): Transformations to apply to the input columns.\n",
        "- `optimization_objective`: The optimization objective (minimize or maximize).\n",
        "  - binary classification:\n",
        "    - `minimize-log-loss`\n",
@@ -571,18 +538,18 @@
        "\n",
        "Run the training job by invoking the `run` method with the following parameters:\n",
        "\n",
-        "- `dataset`: The `Dataset` resource to train the model.\n",
+        "- `dataset`: The dataset resource to train the model.\n",
        "- `model_display_name`: The human readable name for the trained model.\n",
        "- `training_fraction_split`: The percentage of the dataset to use for training.\n",
        "- `test_fraction_split`: The percentage of the dataset to use for test (holdout data).\n",
        "- `validation_fraction_split`: The percentage of the dataset to use for validation.\n",
        "- `target_column`: The name of the column to train as the label.\n",
        "- `budget_milli_node_hours`: (optional) Maximum training time specified in unit of millihours (1000 = hour).\n",
-        "- `disable_early_stopping`: If `True`, training maybe completed before using the entire budget if the service believes it cannot further improve on the model objective measurements.\n",
+        "- `disable_early_stopping`: By default, the model training stops early if the model performance doesn't improve. Setting `disable_early_stopping` = `True` overrides this behavior, allowing the model to train for the entire specified duration.\n",
        "\n",
-        "The `run` method when completed returns the `Model` resource.\n",
+        "The `run` method, upon completion, returns the model resource.\n",
        "\n",
-        "The execution of the training pipeline will take upto 3 hours."
+        "The execution of the training pipeline may take upto 3 hours."
      ]
    },
    {
@@ -612,7 +579,7 @@
      },
      "source": [
        "## Review model evaluation scores\n",
-        "After your model has finished training, you can review its evaluation scores."
+        "After model training is complete, you can review its evaluation scores."
      ]
    },
    {
@@ -660,7 +627,7 @@
      },
      "outputs": [],
      "source": [
-        "# Create results dataset in default region\n",
+        "# Create results dataset in default location\n",
        "RESULTS_DATASET_ID = \"gsod_results_unique\"\n",
        "bq_dataset = bigquery.Dataset(f\"{PROJECT_ID}.{RESULTS_DATASET_ID}\")\n",
        "bq_dataset = bq_client.create_dataset(bq_dataset)\n",
@@ -675,7 +642,7 @@
      "source": [
        "### Make the batch prediction request\n",
        "\n",
-        "You can make a batch prediction by invoking the batch_predict() method, with the following parameters:\n",
+        "You can make a batch prediction by invoking the `batch_predict()` method, with the following parameters:\n",
        "\n",
        "- `job_display_name`: The human readable name for the batch prediction job.\n",
        "- `gcs_source`: A list of one or more batch request input files.\n",
@@ -685,7 +652,7 @@
        "- `machine_type`: The type of machine to use for training.\n",
        "- `accelerator_type`: The hardware accelerator type.\n",
        "- `accelerator_count`: The number of accelerators to attach to a worker replica.\n",
-        "- `sync`: If set to True, the call will block while waiting for the asynchronous batch job to complete.\n",
+        "- `sync`: Set `True` to wait until the completion of the job.\n",
        "\n",
        "Batch prediction job takes roughly 1 hour to finish."
      ]
@@ -698,7 +665,7 @@
      },
      "outputs": [],
      "source": [
-        "# Note: The bigquery_source and bigquery_destination_prefix must be in the same region\n",
+        "# Note: The bigquery_source and bigquery_destination_prefix must be in the same location\n",
        "PREDICTION_RESULTS_DATASET_ID = f\"{PROJECT_ID}.{RESULTS_DATASET_ID}\"\n",
        "\n",
        "batch_predict_job = model.batch_predict(\n",
@@ -753,8 +720,7 @@
        "\n",
        "- Model\n",
        "- AutoML Training Job\n",
-        "- Batch Job\n",
-        "- Cloud Storage Bucket"
+        "- Batch Job"
      ]
    },
    {
@@ -23,6 +23,16 @@
        "# limitations under the License."
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "64a4f198313d"
+      },
+      "source": [
+        "Starting on September 15, 2024, you can only customize classification, entity extraction, and sentiment analysis models by moving to Vertex AI Gemini prompts and tuning. Training or updating models for Vertex AI AutoML for Text classification, entity extraction, and sentiment analysis objectives will no longer be available. You can continue using existing Vertex AI AutoML Text objectives until June 15, 2025. For more information about how Gemini offers enhanced user experience through improved prompting capabilities, see                                                            \n",
+        "[Introduction to tuning](https://cloud.google.com/vertex-ai/generative-ai/docs/models/tune-gemini-overview)"
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -23,6 +23,16 @@
        "# limitations under the License."
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "148c71404373"
+      },
+      "source": [
+        "Starting on September 15, 2024, you can only customize classification, entity extraction, and sentiment analysis models by moving to Vertex AI Gemini prompts and tuning. Training or updating models for Vertex AI AutoML for Text classification, entity extraction, and sentiment analysis objectives will no longer be available. You can continue using existing Vertex AI AutoML Text objectives until June 15, 2025. For more information about how Gemini offers enhanced user experience through improved prompting capabilities, see \n",
+        "[Introduction to tuning](https://cloud.google.com/vertex-ai/generative-ai/docs/models/tune-gemini-overview)"
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -33,23 +33,28 @@
        "\n",
        "<table align=\"left\">\n",
        "\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/automl/sdk_automl_video_action_recognition_batch.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/automl/sdk_automl_video_action_recognition_batch.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fautoml%2Fsdk_automl_video_action_recognition_batch.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\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/automl/sdk_automl_video_action_recognition_batch.ipynb\" target='_blank'>\n",
-        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/automl/sdk_automl_video_action_recognition_batch.ipynb\" target='_blank'>\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br>\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
-        "  </td>                                                                                               \n",
+        "  </td> \n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/automl/sdk_automl_video_action_recognition_batch.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br>\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>                 \n",
        "</table>\n",
        "<br/><br/><br/>"
      ]
@@ -78,25 +83,25 @@
        "\n",
        "In this tutorial, you learn how to create an AutoML video action recognition model from a Python script, and then do a batch prediction using the Vertex AI SDK. You can alternatively create and deploy models using the `gcloud` command-line tool or online using the Cloud Console.\n",
        "\n",
-        "This tutorial uses the following Google Cloud ML services and resources:\n",
+        "This tutorial uses the following Google Cloud Vertex AI services and resources:\n",
        "\n",
-        "- Vertex AI Dataset\n",
-        "- Vertex AI Model\n",
-        "- Vertex AI Batch Prediction\n",
+        "- Vertex AI dataset\n",
+        "- Vertex AI model\n",
+        "- Vertex AI batch prediction\n",
        "\n",
        "\n",
        "The steps performed include:\n",
        "\n",
-        "- Create a `Vertex AI Dataset` resource.\n",
+        "- Create a Vertex AI dataset resource.\n",
        "- Train the model.\n",
        "- View the model evaluation.\n",
        "- Make a batch prediction.\n",
        "\n",
        "There is one key difference between using batch prediction and using online prediction:\n",
        "\n",
-        "* Prediction Service: Does an on-demand prediction for the entire set of instances (i.e., one or more data items) and returns the results in real-time.\n",
+        "* Prediction service: Does an on-demand prediction for the entire set of instances (i.e., one or more data items) and returns the results in real-time.\n",
        "\n",
-        "* Batch Prediction Service: Does a queued (batch) prediction for the entire set of instances in the background and stores the results in a Cloud Storage bucket when ready."
+        "* Batch prediction service: Does a queued (batch) prediction for the entire set of instances in the background and stores the results in a Cloud Storage bucket when ready."
      ]
    },
    {
@@ -133,160 +138,122 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "install_aip:mbsdk"
+        "id": "f0316df526f8"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook. \n"
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "a2c2cb2109a0"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "install_aip:mbsdk"
+        "id": "dae013b807c7"
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
-        "                                google-cloud-storage"
+        "                                 google-cloud-storage"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "restart"
+        "id": "ff555b32bab8"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel"
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "D-ZBOjErv5mM"
+        "id": "f09b4dff629a"
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "yfEglUHQk9S3"
+        "id": "ee775571c2b5"
      },
      "source": [
-        "## Before you begin\n",
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "92e68cfc3a90"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
-        "### Set your project ID\n",
-        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "* Run `gcloud config list`.\n",
-        "* Run `gcloud projects list`.\n",
-        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "set_project_id"
+        "id": "46604f70e831"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4f872cd812d0"
+      },
+      "source": [
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "294fe4e5a671"
      },
      "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": "gcp_authenticate"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
-        "\n",
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated.\n",
-        "\n",
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ce6043da7b33"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0367eac06a10"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "21ad4dbb4a61"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c13224697bfb"
-      },
-      "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."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -317,7 +284,7 @@
        "id": "autoset_bucket"
      },
      "source": [
-        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+        "**If your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
@@ -328,7 +295,29 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
+        "! gsutil mb -l $LOCATION -p $PROJECT_ID $BUCKET_URI"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "a56633b047ee"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "673ef17c3cff"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -349,34 +338,10 @@
      "outputs": [],
      "source": [
        "import json\n",
-        "import os\n",
        "\n",
-        "import google.cloud.aiplatform as aiplatform\n",
        "from google.cloud import storage"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "init_aip:mbsdk"
-      },
-      "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"
-      },
-      "outputs": [],
-      "source": [
-        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -442,12 +407,12 @@
        "id": "create_dataset:video,var"
      },
      "source": [
-        "### Create the Dataset\n",
+        "### Create the dataset\n",
        "\n",
-        "Next, create the `Dataset` resource using the `create` method for the `VideoDataset` class, which takes the following parameters:\n",
+        "Next, create the dataset resource using the `create` method for the `VideoDataset` class, which takes the following parameters:\n",
        "\n",
-        "- `display_name`: The human readable name for the `Dataset` resource.\n",
-        "- `gcs_source`: A list of one or more dataset index files to import the data items into the `Dataset` resource.\n",
+        "- `display_name`: The human readable name for the dataset resource.\n",
+        "- `gcs_source`: A list of one or more dataset index files to import the data items into the dataset resource.\n",
        "\n",
        "This operation may take several minutes."
      ]
@@ -485,7 +450,7 @@
        "\n",
        "An AutoML training pipeline is created with the `AutoMLVideoTrainingJob` class, with the following parameters:\n",
        "\n",
-        "- `display_name`: The human readable name for the `TrainingJob` resource.\n",
+        "- `display_name`: The human readable name for the TrainingJob resource.\n",
        "- `prediction_type`: The type task to train the model for.\n",
        "    - `classification`: A video classification model.\n",
        "    - `object_tracking`: A video object tracking model.\n",
@@ -518,12 +483,12 @@
        "\n",
        "Next, you run the training job by invoking the method `run`, with the following parameters:\n",
        "\n",
-        "- `dataset`: The `Dataset` resource to train the model.\n",
+        "- `dataset`: The dataset resource to train the model.\n",
        "- `model_display_name`: The human readable name for the trained model.\n",
        "- `training_fraction_split`: The percentage of the dataset to use for training.\n",
        "- `test_fraction_split`: The percentage of the dataset to use for test (holdout data).\n",
        "\n",
-        "The `run` method when completed returns the `Model` resource.\n",
+        "The `run` method when completed returns the model resource.\n",
        "\n",
        "The execution of the training pipeline can take upto 40 minutes."
      ]
@@ -802,7 +767,7 @@
        "\n",
        "# Delete Cloud Storage objects\n",
        "delete_bucket = False\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
@@ -32,20 +32,25 @@
        "# Vertex AI SDK: AutoML training video classification model for batch prediction\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/automl/sdk_automl_video_classification_batch.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fautoml%2Fsdk_automl_video_classification_batch.ipynb\"\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/automl/sdk_automl_video_classification_batch.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br>\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "<a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/automl/sdk_automl_video_classification_batch.ipynb\" target='_blank'>            \n",
-        "    <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"> \n",
+        "    <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> \n",
        "Open in Vertex AI Workbench \n",
        "    </a>\n",
        "  </td>\n",
@@ -132,160 +137,121 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "install_aip:mbsdk"
+        "id": "f0316df526f8"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the latest versions of Vertex AI and Cloud Storage SDK for Python."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "2abdd254e90f"
-      },
-      "outputs": [],
-      "source": [
-        "import os\n",
-        "\n",
-        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
-        "                                 google-cloud-storage"
+        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "restart"
+        "id": "a2c2cb2109a0"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel"
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "D-ZBOjErv5mM"
+        "id": "514a03ed1a82"
      },
      "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)"
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform google-cloud-storage"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "yfEglUHQk9S3"
+        "id": "ff555b32bab8"
      },
      "source": [
-        "## Before you begin\n",
+        "### Restart runtime (Colab only)\n",
        "\n",
-        "### Set your project ID\n",
-        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "* Run `gcloud config list`.\n",
-        "* Run `gcloud projects list`.\n",
-        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "set_project_id"
+        "id": "f09b4dff629a"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ee775571c2b5"
+      },
+      "source": [
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "92e68cfc3a90"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
+        "\n",
+        "Authenticate your environment on Google Colab.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "46604f70e831"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4f872cd812d0"
+      },
+      "source": [
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "294fe4e5a671"
      },
      "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": "gcp_authenticate"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
-        "\n",
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated.\n",
-        "\n",
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ce6043da7b33"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0367eac06a10"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "21ad4dbb4a61"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c13224697bfb"
-      },
-      "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."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -327,27 +293,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "setup_vars"
-      },
-      "source": [
-        "### Import libraries"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 13,
-      "metadata": {
-        "id": "import_aip:mbsdk"
-      },
-      "outputs": [],
-      "source": [
-        "from google.cloud import aiplatform"
+        "! gsutil mb -l $LOCATION -p $PROJECT_ID $BUCKET_URI"
      ]
    },
    {
@@ -369,6 +315,8 @@
      },
      "outputs": [],
      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
      ]
    },
@@ -842,7 +790,7 @@
        "batch_predict_job.delete()\n",
        "\n",
        "# Delete the Cloud storage bucket\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil rm -r $BUCKET_URI"
      ]
    }
@@ -32,23 +32,26 @@
        "# Get started with BigQuery ML Training\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/bigquery_ml/get_started_with_bqml_training.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/pipelines/lightweight_functions_component_io_kfp.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "    <td>\n",
-        "        <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/bigquery_ml/get_started_with_bqml_training.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://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/bigquery_ml/get_started_with_bqml_training.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",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fnotebook_template.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/pipelines/lightweight_functions_component_io_kfp.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "<a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/pipelines/lightweight_functions_component_io_kfp.ipynb\" target='_blank'>\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
+        "     </a>\n",
+        "   </td>\n",
        "</table>\n",
        "<br/><br/><br/>"
      ]
@@ -117,7 +120,10 @@
        "- Cloud Storage\n",
        "- BigQuery\n",
        "\n",
-        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing) and [BigQuery pricing](https://cloud.google.com/bigquery/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing),\n",
+        "[Cloud Storage pricing](https://cloud.google.com/storage/pricing) and\n",
+        "[BigQuery pricing](https://cloud.google.com/bigquery/pricing) and use the \n",
+        "[Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
@@ -126,9 +132,8 @@
        "id": "install_mlops"
      },
      "source": [
-        "## Installations\n",
-        "\n",
-        "Install the following packages for executing this notebook."
+        "## Get started\n",
+        "Install Vertex AI SDK for Python and other required packages"
      ]
    },
    {
@@ -153,7 +158,8 @@
        "id": "restart"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel"
+        "### Restart runtime (Colab only)\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
@@ -164,11 +170,52 @@
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4de1bd77992b"
+      },
+      "source": [
+        "<div class=\"alert alert-block alert-warning\">,\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>,\n",
+        "</div>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "befa6ca14bc0"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
+        "Authenticate your environment on Google Colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7de6ef0fac42"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
      ]
    },
    {
@@ -177,108 +224,20 @@
        "id": "yfEglUHQk9S3"
      },
      "source": [
-        "## Before you begin\n",
-        "\n",
-        "### Set your project ID\n",
-        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "* Run `gcloud config list`.\n",
-        "* Run `gcloud projects list`.\n",
-        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+        "### Set Google Cloud project information\n",
+        "Learn more about [setting up a project and a development environment.](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "set_project_id"
+        "id": "region"
      },
      "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": "gcp_authenticate"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
-        "\n",
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated.\n",
-        "\n",
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ce6043da7b33"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0367eac06a10"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "21ad4dbb4a61"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c13224697bfb"
-      },
-      "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."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -320,7 +279,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} {BUCKET_URI}"
      ]
    },
    {
@@ -404,7 +363,7 @@
      "source": [
        "### Initialize Vertex AI and BigQuery SDKs for Python\n",
        "\n",
-        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
+        "Initialize the Vertex AI SDK for Python and BigQuery SDK  with your project and the created bucket."
      ]
    },
    {
@@ -448,13 +407,13 @@
        "\n",
        "You can set hardware accelerators for prediction.\n",
        "\n",
-        "Set the variable `DEPLOY_GPU/DEPLOY_NGPU` to use a container image supporting a GPU and the number of GPUs allocated to the virtual machine (VM) instance. For example, to use a GPU container image with 4 Nvidia Telsa K80 GPUs allocated to each VM, you would specify:\n",
+        "Set the variable `DEPLOY_GPU/DEPLOY_NGPU` to use a container image supporting a GPU and the number of GPUs allocated to the virtual machine (VM) instance. For example, to use a GPU container image with 4 Nvidia Telsa T4 GPUs allocated to each VM, you would specify:\n",
        "\n",
-        "    (aiplatform.AcceleratorType.NVIDIA_TESLA_K80, 4)\n",
+        "    (aiplatform.AcceleratorType.NVIDIA_TESLA_T4, 4)\n",
        "\n",
        "Otherwise specify `(None, None)` to use a container image to run on a CPU.\n",
        "\n",
-        "Learn more [here](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators) hardware accelerator support for your region."
+        "Learn more [about hardware accelerator support for your region.](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators)"
      ]
    },
    {
@@ -469,11 +428,11 @@
        "\n",
        "if os.getenv(\"IS_TESTING_DEPLOY_GPU\"):\n",
        "    DEPLOY_GPU, DEPLOY_NGPU = (\n",
-        "        aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_K80,\n",
+        "        aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_T4,\n",
        "        int(os.getenv(\"IS_TESTING_DEPLOY_GPU\")),\n",
        "    )\n",
        "else:\n",
-        "    DEPLOY_GPU, DEPLOY_NGPU = (aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_K80, 1)"
+        "    DEPLOY_GPU, DEPLOY_NGPU = (aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_T4, 1)"
      ]
    },
    {
@@ -482,14 +441,14 @@
        "id": "container:prediction"
      },
      "source": [
-        "### Set pre-built containers\n",
+        "### Set prebuilt containers\n",
        "\n",
-        "Set the pre-built Docker container image for prediction.\n",
+        "Set the prebuilt Docker container image for prediction.\n",
        "\n",
-        "- Set the variable `TF` to the TensorFlow version of the container image. For example, `2-1` would be version 2.1, and `1-15` would be version 1.15. The following list shows some of the pre-built images available:\n",
+        "- Set the variable `TF` to the TensorFlow version of the container image. The following list shows some of the prebuilt images available:\n",
        "\n",
        "\n",
-        "For the latest list, see [Pre-built containers for prediction](https://cloud.google.com/ai-platform-unified/docs/predictions/pre-built-containers)."
+        "For the latest list, see [prebuilt containers for prediction](https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers)."
      ]
    },
    {
@@ -517,7 +476,7 @@
        "        DEPLOY_VERSION = \"tf-cpu.{}\".format(TF)\n",
        "\n",
        "DEPLOY_IMAGE = \"{}-docker.pkg.dev/vertex-ai/prediction/{}:latest\".format(\n",
-        "    REGION.split(\"-\")[0], DEPLOY_VERSION\n",
+        "    LOCATION.split(\"-\")[0], DEPLOY_VERSION\n",
        ")\n",
        "\n",
        "print(\"Deployment:\", DEPLOY_IMAGE, DEPLOY_GPU)"
@@ -540,7 +499,7 @@
        "     - `n1-highcpu`: 0.9 GB of memory per vCPU\n",
        " - `vCPUs`: number of \\[2, 4, 8, 16, 32, 64, 96 \\]\n",
        "\n",
-        "*Note: You may also use n2 and e2 machine types for training and deployment, but they do not support GPUs*"
+        "*Note: You may also use n2 and e2 machine types for training and deployment, but they don't support GPUs*"
      ]
    },
    {
@@ -569,7 +528,7 @@
      "source": [
        "## BigQuery ML introduction\n",
        "\n",
-        "BigQuery ML (BQML) provides the capability to train ML tabular models, such as classification and regression, in BigQuery using SQL syntax.\n",
+        "BigQuery ML (BQML) provides the capability to train ML tabular models, such as classification and regression in BigQuery using SQL syntax.\n",
        "\n",
        "Learn more about [BigQuery ML documentation](https://cloud.google.com/bigquery-ml/docs)."
      ]
@@ -623,7 +582,7 @@
        "Next, you create and train a BigQuery ML tabular classification model from the public dataset penguins and store the model in your project using the `CREATE MODEL` statement. The model configuration is specified in the `OPTIONS` statement as follows:\n",
        "\n",
        "- `model_type`: The type and archictecture of tabular model to train, e.g., DNN classification.\n",
-        "- `labels`: The column which are the labels.\n",
+        "- `labels`: The column labels.\n",
        "\n",
        "Learn more about [The CREATE MODEL statement](https://cloud.google.com/bigquery-ml/docs/reference/standard-sql/bigqueryml-syntax-create)."
      ]
@@ -710,7 +669,7 @@
      "source": [
        "### Export the model from BigQuery ML\n",
        "\n",
-        "The model you trained in BigQuery ML is a TensorFlow model. Next, you export the TensorFlow model artifacts in TF.SavedModel format."
+        "The model you trained in BigQuery ML is a TensorFlow model. Next, export the TensorFlow model artifacts in TF.SavedModel format."
      ]
    },
    {
@@ -736,7 +695,7 @@
      "source": [
        "## Upload the BigQuery ML model to a Vertex AI Model resource\n",
        "\n",
-        "Finally, now that you have the BigQuery ML model exported, you upload the model artifacts to Vertex AI Model resource, in the same way as if you were uploading a custom trained model.\n",
+        "Finally, now that you have the BigQuery ML model exported, you upload the model artifacts to Vertex AI model resource, in the same way as if you were uploading a custom trained model.\n",
        "\n",
        "Below is a partial list of mapping BigQuery ML model types to their corresponding exported model format:\n",
        "\n",
@@ -782,12 +741,12 @@
      "source": [
        "## Deploy the model\n",
        "\n",
-        "Next, deploy your model for online prediction. To deploy the model, you invoke the `deploy` method, with the following parameters:\n",
+        "Next, deploy your model for online prediction. To deploy the model, invoke the `deploy` method, with the following parameters:\n",
        "\n",
        "- `deployed_model_display_name`: A human readable name for the deployed model.\n",
        "- `traffic_split`: Percent of traffic at the endpoint that goes to this model, which is specified as a dictionary of one or more key/value pairs.\n",
        "If only one model, then specify as { \"0\": 100 }, where \"0\" refers to this model being uploaded and 100 means 100% of the traffic.\n",
-        "If there are existing models on the endpoint, for which the traffic needs to be split, then use model_id to specify as { \"0\": percent, model_id: percent, ... }, where model_id is the model id of an existing model to the deployed endpoint. The percents must add up to 100.\n",
+        "If there are existing models on the endpoint, for which the traffic needs to split, then use model_id to specify as { \"0\": percent, model_id: percent, ... }, where model_id is the model id of an existing model to the deployed endpoint. The percent must add up to 100.\n",
        "- `machine_type`: The type of machine to use for training.\n",
        "- `accelerator_type`: The hardware accelerator type.\n",
        "- `accelerator_count`: The number of accelerators to attach to a worker replica.\n",
@@ -841,7 +800,7 @@
      "source": [
        "#### Undeploy the model\n",
        "\n",
-        "When you are done doing predictions, you undeploy the model from the `Endpoint` resouce. This deprovisions all compute resources and ends billing for the deployed model."
+        "When you are done doing predictions, you undeploy the model from the endpoint resource. This deprovisions all compute resources and ends billing for the deployed model."
      ]
    },
    {
@@ -864,9 +823,9 @@
        "id": "model_delete:mbsdk"
      },
      "source": [
-        "#### Delete the `Vertex AI Model` resource\n",
+        "#### Delete the Vertex AI model resource\n",
        "\n",
-        "The method 'delete()' deletes the model."
+        "The method `delete()` deletes the model."
      ]
    },
    {
@@ -889,9 +848,9 @@
        "id": "7890ae6f6410"
      },
      "source": [
-        "### Delete the `BigQuery ML` model\n",
+        "### Delete the BigQuery ML model\n",
        "\n",
-        "Next, delete the `BigQuery ML` instance of the model."
+        "Next, delete the BigQuery ML instance of the model."
      ]
    },
    {
@@ -1040,7 +999,7 @@
      "source": [
        "### Train a BigQuery ML model with Explainability\n",
        "\n",
-        "Next, you train the same BigQuery ML model, but this time you enable Vertex AI Explainability on the model predictions by adding the option:\n",
+        "Next, you train the same BigQuery ML model, this time you enable Vertex Explainable AI on the model predictions by adding the option:\n",
        "\n",
        "- `ENABLE_GLOBAL_EXPLAIN`"
      ]
@@ -1110,23 +1069,6 @@
        "job = bqclient.query(MODEL_QUERY)"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "2b4498ca6fea"
-      },
-      "source": [
-        "## Model Registry\n",
-        "\n",
-        "Alternatively, you can implicitly upload your BigQuery ML model as a Vertex AI model resource with exporting and importing the model artifacts. In this method, you add additional options when training the model that tells BigQuery ML to automatically upload and register the trained model as a model resource.\n",
-        "\n",
-        "### Setting permissions to automatically register the model\n",
-        "\n",
-        "You need to set some additional IAM permissions for BigQuery ML to automatically upload and register the model after training. Depending on your service account, the setting of the permissions below may fail. In this case, we recommend executing the permissions in a Cloud Shell.\n",
-        "\n",
-        "Learn more about [Setting permissions for Model Registry](https://cloud.google.com/bigquery-ml/docs/managing-models-vertex)\n"
-      ]
-    },
    {
      "cell_type": "code",
      "execution_count": null,
@@ -1149,9 +1091,9 @@
        "\n",
        "Next, you train the model and automatically register the model to the Vertex AI Model Registry, by adding the following parameters as options:\n",
        "\n",
-        "- `model_registry`: Set to \"vertex_ai\" to indicate automatic registration to Vertex AI Model Registry.\n",
+        "- `model_registry`: Set to `vertex_ai` to indicate automatic registration to Vertex AI Model Registry.\n",
        "- `vertex_ai_model_id`: The human readable display name for the registered model.\n",
-        "- `vertex_ai_model_version_aliases`: Alternate names for the model."
+        "- `vertex_ai_model_version_aliases`: Alternate name for the model."
      ]
    },
    {
@@ -1201,9 +1143,9 @@
        "id": "5b4970272040"
      },
      "source": [
-        "### Find the model in the Vertex Model Registry\n",
+        "### Find the model in the Vertex AI Model Registry\n",
        "\n",
-        "Finally, you can use the Vertex AI Model `list()` method with a filter query to find the automatically registered model."
+        "Finally, you can use the Vertex AI model `list()` method with a filter query to find the automatically registered model."
      ]
    },
    {
@@ -1301,7 +1243,7 @@
        "! bq rm -r -f $PROJECT_ID:$BQ_DATASET_NAME\n",
        "\n",
        "delete_storage = False\n",
-        "if delete_storage or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_storage:\n",
        "    # Delete the created GCS bucket\n",
        "    ! gsutil rm -r $BUCKET_URI"
      ]
@@ -32,19 +32,24 @@
        "# Deploying Iris-detection model using FastAPI and Vertex AI custom container serving\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/SDK_Custom_Container_Prediction.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/SDK_Custom_Container_Prediction.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fcustom%2FSDK_Custom_Container_Prediction.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
-        "<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/custom/SDK_Custom_Container_Prediction.ipynb\" target='_blank'>\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",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/custom/SDK_Custom_Container_Prediction.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/SDK_Custom_Container_Prediction.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
@@ -58,7 +63,7 @@
      "source": [
        "## Overview\n",
        "\n",
-        "This tutorial walks you through building a custom container to serve a scikit-learn model on Vertex AI. You use the FastAPI Python web server framework to create a prediction and health endpoint. You also incorporate a pre-processor from training pipeline into your online serving application.\n",
+        "In this tutorial, you build a scikit-learn model and deploy it on Vertex AI using the custom container method. You use the FastAPI Python web server framework to create a prediction endpoint. You also incorporate a preprocessor from training pipeline into your online serving application.\n",
        "\n",
        "Learn more about [Custom training](https://cloud.google.com/vertex-ai/docs/training/custom-training) and [Vertex AI Prediction](https://cloud.google.com/vertex-ai/docs/predictions/get-predictions)."
      ]
@@ -74,10 +79,10 @@
        "In this notebook, you learn how to create, deploy and serve a custom classification model on Vertex AI. This notebook focuses more on deploying the model than on the design of the model itself. \n",
        "\n",
        "\n",
-        "This tutorial uses the following Google Cloud ML services and resources:\n",
+        "This tutorial uses the following Vertex AI services and resources:\n",
        "\n",
-        "- Vertex AI Models\n",
-        "- Vertex AI Endpoints\n",
+        "- Vertex AI models\n",
+        "- Vertex AI endpoints\n",
        "\n",
        "The steps performed include:\n",
        "\n",
@@ -96,10 +101,10 @@
      "source": [
        "### Dataset\n",
        "\n",
-        "This tutorial uses R.A. Fisher's Iris dataset, a small dataset that is a popular choice for trying out machine learning techniques. Each instance has four numerical features, which are different measurements of a flower, and a target label that\n",
-        "marks it as one of three types of iris: Iris setosa, Iris versicolour, or Iris virginica.\n",
+        "This tutorial uses R.A. Fisher's Iris dataset, a small and popular dataset for machine learning experiments. Each instance has four numerical features, which are different measurements of a flower, and a target label that\n",
+        "categorizes the flower into: **Iris setosa**, **Iris versicolour** and **Iris virginica**.\n",
        "\n",
-        "This tutorial uses [the copy of the Iris dataset included in the\n",
+        "This tutorial uses [a version of the Iris dataset available in the\n",
        "scikit-learn library](https://scikit-learn.org/stable/modules/generated/sklearn.datasets.load_iris.html#sklearn.datasets.load_iris)."
      ]
    },
@@ -128,12 +133,21 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "i7EUnXsZhAGF"
+        "id": "f0316df526f8"
      },
      "source": [
-        "## Install additional packages\n",
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9065e8d7f0fb"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n",
        "\n",
-        "Install additional package dependencies not installed in your notebook environment, such as NumPy, Scikit-learn, FastAPI, Uvicorn, and joblib. Use the latest major GA version of each package."
+        "Write the requirements needed for building container into a file."
      ]
    },
    {
@@ -151,6 +165,15 @@
        "google-cloud-storage>=1.26.0,<2.0.0dev"
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "479caccfb277"
+      },
+      "source": [
+        "Install the dependencies."
+      ]
+    },
    {
      "cell_type": "code",
      "execution_count": null,
@@ -175,7 +198,9 @@
        "id": "restart"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel"
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
@@ -186,11 +211,53 @@
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2d8e61720c2a"
+      },
+      "source": [
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "429a9f6d237d"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
+        "\n",
+        "Authenticate your environment on Google Colab.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "0c74b83ece6c"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
      ]
    },
    {
@@ -199,14 +266,8 @@
        "id": "yfEglUHQk9S3"
      },
      "source": [
-        "## Before you begin\n",
-        "\n",
-        "### Set your project ID\n",
-        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "* Run `gcloud config list`.\n",
-        "* Run `gcloud projects list`.\n",
-        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+        "### Set Google Cloud project information \n",
+        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
@@ -218,89 +279,7 @@
      "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": "gcp_authenticate"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
-        "\n",
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated.\n",
-        "\n",
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ce6043da7b33"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0367eac06a10"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "21ad4dbb4a61"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c13224697bfb"
-      },
-      "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."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n"
      ]
    },
    {
@@ -331,7 +310,7 @@
        "id": "create_bucket"
      },
      "source": [
-        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+        "**If your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
@@ -342,7 +321,31 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3330b4f12a0d"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). "
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "e088ea8cd4a0"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -351,7 +354,7 @@
        "id": "d3938f6d37a1"
      },
      "source": [
-        "## Import libraries"
+        "### Import the required libraries"
      ]
    },
    {
@@ -363,30 +366,7 @@
      "outputs": [],
      "source": [
        "import os\n",
-        "\n",
-        "from google.cloud import aiplatform"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "750d53e37094"
-      },
-      "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": "1a3aa2d4a74f"
-      },
-      "outputs": [],
-      "source": [
-        "aiplatform.init(project=PROJECT_ID, location=REGION)"
+        "import sys"
      ]
    },
    {
@@ -397,7 +377,7 @@
      "source": [
        "### Configure resource names\n",
        "\n",
-        "Set a name for the following resources:\n",
+        "Set a name for the following parameters:\n",
        "\n",
        "`MODEL_ARTIFACT_DIR` - Folder directory path to your model artifacts within a Cloud Storage bucket, for example: \"my-models/fraud-detection/trial-4\"\n",
        "\n",
@@ -405,7 +385,7 @@
        "\n",
        "`IMAGE` - Name of the container image that is pushed to the repository.\n",
        "\n",
-        "`MODEL_DISPLAY_NAME` - Display name of Vertex AI Model resource."
+        "`MODEL_DISPLAY_NAME` - Display name of Vertex AI model resource."
      ]
    },
    {
@@ -442,9 +422,9 @@
      },
      "source": [
        "## Write your pre-processor\n",
-        "Standardize the training data so each numerical feature column has a mean of 0 and a standard deviation of 1 [can improve your model](https://developers.google.com/machine-learning/crash-course/representation/cleaning-data).\n",
+        "Standardize the training data so that each numerical feature column has a mean of 0 and a standard deviation of 1 [can improve your model](https://developers.google.com/machine-learning/crash-course/representation/cleaning-data).\n",
        "\n",
-        "Define a `app` folder and create `preprocess.py`, which contains a class to perform standardization."
+        "Define an `app` folder and create `preprocess.py`, which contains a class to perform standardization."
      ]
    },
    {
@@ -493,9 +473,9 @@
      },
      "source": [
        "## Train and store model with pre-processor\n",
-        "Next, use `preprocess.MySimpleScaler` to preprocess the iris data, then train a model using scikit-learn.\n",
+        "Use `preprocess.MySimpleScaler` to preprocess the iris data, and then train a model using scikit-learn.\n",
        "\n",
-        "At the end, export your trained model as a joblib (`.joblib`) file and export your `MySimpleScaler` instance as a pickle (`.pkl`) file."
+        "After training, export your trained model as a joblib (`.joblib`) file and export your `MySimpleScaler` instance as a pickle (`.pkl`) file."
      ]
    },
    {
@@ -633,7 +613,7 @@
      },
      "source": [
        "### Add pre-start script\n",
-        "FastAPI executes the following script before starting up the server. The `PORT` environment variable is set to equal to `AIP_HTTP_PORT` in order to run FastAPI on the same port expected by Vertex AI."
+        "FastAPI executes the following script before starting up the server. Set the environment variable `PORT` to `AIP_HTTP_PORT` for running the FastAPI server."
      ]
    },
    {
@@ -682,7 +662,7 @@
        "id": "51e149fdec1b"
      },
      "source": [
-        "## Build and push container to Artifact Registry"
+        "## Push the container image to Artifact Registry"
      ]
    },
    {
@@ -691,7 +671,7 @@
        "id": "240578ec9efe"
      },
      "source": [
-        "Write the `Dockerfile`, using `tiangolo/uvicorn-gunicorn-fastapi` as a base image. This automatically runs FastAPI for you using Gunicorn and Uvicorn. Visit [the FastAPI docs to read more about deploying FastAPI with Docker](https://fastapi.tiangolo.com/deployment/docker/) to learn more."
+        "Write the `Dockerfile`, using `tiangolo/uvicorn-gunicorn-fastapi` as a base image. This automatically runs FastAPI for you using Gunicorn and Uvicorn. Visit [the FastAPI docs about deploying with Docker](https://fastapi.tiangolo.com/deployment/docker/) to learn more."
      ]
    },
    {
@@ -712,17 +692,26 @@
        "RUN pip install -r requirements.txt"
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5cae10684d2c"
+      },
+      "source": [
+        "### Test the image locally (optional)"
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "04c988201499"
      },
      "source": [
-        "### Build the image locally (optional)\n",
+        "#### Build the image locally (optional)\n",
        "\n",
        "Build the image using docker to test it locally.\n",
        "\n",
-        "**Note:** Docker is only being used to test the container locally. For deployment to Artifact registry, Cloud-Build is used."
+        "**Note:** In this tutorial, docker is only being used to test the container locally. For deployment to Artifact Registry, Cloud-Build is used."
      ]
    },
    {
@@ -733,12 +722,11 @@
      },
      "outputs": [],
      "source": [
-        "import sys\n",
-        "\n",
        "IS_COLAB = \"google.colab\" in sys.modules\n",
+        "\n",
        "if not IS_COLAB and not os.getenv(\"IS_TESTING\"):\n",
        "    ! sudo docker build \\\n",
-        "            --tag=\"{REGION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE}\" \\\n",
+        "            --tag=\"{LOCATION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE}\" \\\n",
        "            ."
      ]
    },
@@ -748,7 +736,7 @@
        "id": "147a555f6c93"
      },
      "source": [
-        "### Run and test the container locally (optional)\n",
+        "#### Run and test the container locally (optional)\n",
        "\n",
        "Test running the container locally in detached mode and provide the environment variables that the container requires. These variables are provided to the container by Vertex AI once deployed. Test the `/health` and `/predict` routes and then stop the running image."
      ]
@@ -764,13 +752,13 @@
        "if not IS_COLAB and not os.getenv(\"IS_TESTING\"):\n",
        "    ! sudo docker stop local-iris\n",
        "    ! sudo docker rm local-iris\n",
-        "    ! sudo docker run -d -p 80:8080 \\\n",
+        "    ! docker run -d -p 80:8080 \\\n",
        "            --name=local-iris \\\n",
        "            -e AIP_HTTP_PORT=8080 \\\n",
        "            -e AIP_HEALTH_ROUTE=/health \\\n",
        "            -e AIP_PREDICT_ROUTE=/predict \\\n",
        "            -e AIP_STORAGE_URI={BUCKET_URI}/{MODEL_ARTIFACT_DIR} \\\n",
-        "            \"{REGION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE}\""
+        "            \"{LOCATION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE}\""
      ]
    },
    {
@@ -846,7 +834,7 @@
        "id": "93002a20a2a6"
      },
      "source": [
-        "### Build Custom Container\n"
+        "### Create a repository in Artifact Registry\n"
      ]
    },
    {
@@ -882,12 +870,12 @@
        "id": "23f3ba8346f4"
      },
      "source": [
-        "### Create a private Docker repository\n",
+        "#### Create a docker repository in Artifact Registry\n",
        "Your first step is to create your own Docker repository in Google Artifact Registry.\n",
        "\n",
-        "1 - Run the gcloud artifacts repositories create command to create a new Docker repository with your region with the description \"docker repository\".\n",
+        "1 - Run the `gcloud artifacts repositories create` command to create a new Docker repository with your region and the description as \"docker repository\".\n",
        "\n",
-        "2 - Run the gcloud artifacts repositories list command to verify that your repository was created."
+        "2 - Run the `gcloud artifacts repositories list` command to verify that your repository was created."
      ]
    },
    {
@@ -900,7 +888,7 @@
      "source": [
        "REPOSITORY = \"my-docker-repo-unique\"\n",
        "\n",
-        "! gcloud artifacts repositories create {REPOSITORY} --repository-format=docker --location={REGION} --description=\"Docker repository\"\n",
+        "! gcloud artifacts repositories create {REPOSITORY} --repository-format=docker --location={LOCATION} --description=\"Docker repository\"\n",
        "\n",
        "! gcloud artifacts repositories list"
      ]
@@ -911,9 +899,10 @@
        "id": "d2e0b8b700aa"
      },
      "source": [
-        "Push the image to the created artifact repository using Cloud-Build.\n",
+        "### Submit the image\n",
+        "Push the image to the created artifact repository using Cloud build.\n",
        "\n",
-        "**Note:** The following command automatically considers the Dockerfile from the directory it is being run from."
+        "**Note:** The following command automatically considers the Dockerfile from the directory it's being run from."
      ]
    },
    {
@@ -924,7 +913,7 @@
      },
      "outputs": [],
      "source": [
-        "!gcloud builds submit --region={REGION} --tag={REGION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE}"
+        "!gcloud builds submit --region={LOCATION} --tag={LOCATION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE}"
      ]
    },
    {
@@ -943,7 +932,7 @@
      },
      "source": [
        "### Create Vertex AI model using artifact uri\n",
-        "Use the Python SDK to upload and deploy your model from the artifact registry."
+        "Use the Python SDK to upload your model artifact in Vertex AI."
      ]
    },
    {
@@ -957,7 +946,7 @@
        "model = aiplatform.Model.upload(\n",
        "    display_name=MODEL_DISPLAY_NAME,\n",
        "    artifact_uri=f\"{BUCKET_URI}/{MODEL_ARTIFACT_DIR}\",\n",
-        "    serving_container_image_uri=f\"{REGION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE}\",\n",
+        "    serving_container_image_uri=f\"{LOCATION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE}\",\n",
        ")"
      ]
    },
@@ -967,9 +956,9 @@
        "id": "bd1b85afc7df"
      },
      "source": [
-        "### Deploy the model to Vertex AI Endpoints\n",
+        "### Deploy the model to Vertex AI Endpoint\n",
        "\n",
-        "Deploy the model to a Vertex AI Endpoint. After this step completes, the model is deployed and ready for online predictions."
+        "Once the deployment process gets done, the model is deployed and ready for serving online predictions."
      ]
    },
    {
@@ -991,7 +980,7 @@
      "source": [
        "## Request predictions\n",
        "\n",
-        "Send online requests to the model deployed to the endpoint and get predictions.\n",
+        "Send online requests to the endpoint and get predictions.\n",
        "\n",
        "### Using Python SDK\n",
        "\n",
@@ -1006,6 +995,7 @@
      },
      "outputs": [],
      "source": [
+        "# Send some sample data to the endpoint\n",
        "endpoint.predict(instances=[[6.7, 3.1, 4.7, 1.5], [4.6, 3.1, 1.5, 0.2]])"
      ]
    },
@@ -1028,6 +1018,7 @@
      },
      "outputs": [],
      "source": [
+        "# Fetch the endpoint name\n",
        "ENDPOINT_ID = endpoint.name"
      ]
    },
@@ -1039,11 +1030,12 @@
      },
      "outputs": [],
      "source": [
+        "# Send a prediction request using sample data \n",
        "! curl \\\n",
        "-H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
        "-H \"Content-Type: application/json\" \\\n",
        "-d @instances.json \\\n",
-        "https://{REGION}-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/{REGION}/endpoints/{ENDPOINT_ID}:predict"
+        "https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}:predict"
      ]
    },
    {
@@ -1066,7 +1058,7 @@
      "outputs": [],
      "source": [
        "!gcloud ai endpoints predict $ENDPOINT_ID \\\n",
-        "  --region=$REGION \\\n",
+        "  --region=$LOCATION \\\n",
        "  --json-request=instances.json"
      ]
    },
@@ -1112,16 +1104,39 @@
        "!gcloud artifacts docker images delete \\\n",
        "    --quiet \\\n",
        "    --delete-tags \\\n",
-        "    {REGION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE}\n",
+        "    {LOCATION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE}\n",
        "\n",
-        "# Delete the artifact registry\n",
-        "if delete_art_repo or os.getenv(\"IS_TESTING\"):\n",
-        "    ! gcloud artifacts repositories delete {REPOSITORY} --location=$REGION -q\n",
+        "# Delete the Artifact Repository\n",
+        "if delete_art_repo:\n",
+        "    ! gcloud artifacts repositories delete {REPOSITORY} --location=$LOCATION -q\n",
        "    \n",
        "# Delete the Cloud Storage bucket\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "31885b85cc2c"
+      },
+      "source": [
+        "Clean up the locally created files."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "234ea94b87c5"
+      },
+      "outputs": [],
+      "source": [
+        "! rm -rf app/\n",
+        "! rm requirements.txt\n",
+        "! rm instances.json\n",
+        "! rm Dockerfile"
+      ]
    }
  ],
  "metadata": {
@@ -4,6 +4,7 @@
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
+        "cellView": "form",
        "id": "copyright"
      },
      "outputs": [],
@@ -32,23 +33,26 @@
        "# Training a TensorFlow model on BigQuery data\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/custom-tabular-bq-managed-dataset.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/custom-tabular-bq-managed-dataset.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fcustom%2Fcustom-tabular-bq-managed-dataset.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/custom/custom-tabular-bq-managed-dataset.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",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
-        "  </td>  \n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/custom-tabular-bq-managed-dataset.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
        "</table>"
      ]
    },
@@ -128,12 +132,19 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "install_aip"
+        "id": "3b1ffd5ab768"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the latest version of Cloud Storage, Bigquery and Vertex AI SDKs for Python."
+        "## Get Started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dc848186ab0e"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages"
      ]
    },
    {
@@ -145,7 +156,7 @@
      "outputs": [],
      "source": [
        "# Install the packages\n",
-        "! pip3 install --upgrade google-cloud-aiplatform \\\n",
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
        "                        google-cloud-storage \\\n",
        "                        'google-cloud-bigquery[pandas]'"
      ]
@@ -153,48 +164,80 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "blGlVGFYW9Pt"
+        "id": "ff555b32bab8"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel."
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "0JrvuK6LUYnQ"
+        "id": "f09b4dff629a"
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "a47846030fef"
+        "id": "54c5ef8a8f43"
      },
      "source": [
-        "## Before you begin"
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "project_id"
+        "id": "f82e28c631cc"
      },
      "source": [
-        "#### Set your project ID\n",
+        "### Authenticate your notebook environment (Colab only)\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)"
+        "Authenticate your environment on Google Colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "46604f70e831"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "107c51893a64"
+      },
+      "source": [
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
@@ -206,103 +249,7 @@
      "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": "a54f9d7c1876"
-      },
-      "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": "3aaadaaf9b30"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "5c0404984792"
-      },
-      "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": "BaFKzJ_xXpvm"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "_fV-KyGAX4Xl"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**\n"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "7uXB1HAPX6L_"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "Ab_TRMQIYCCX"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "vx25htmYYExI"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "uZdA0-jBYGqt"
-      },
-      "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."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -344,7 +291,38 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
+        "! gsutil mb -l $LOCATION -p $PROJECT_ID $BUCKET_URI"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4f1319830b6e"
+      },
+      "source": [
+        "### Import libraries"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "poijnGfZCFYi"
+      },
+      "source": [
+        "Import the Vertex AI Python SDK and other required Python libraries."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "b9a0a5a74fa6"
+      },
+      "outputs": [],
+      "source": [
+        "import numpy as np\n",
+        "import pandas as pd\n",
+        "from google.cloud import aiplatform, bigquery"
      ]
    },
    {
@@ -366,10 +344,8 @@
      },
      "outputs": [],
      "source": [
-        "from google.cloud import aiplatform\n",
-        "\n",
-        "# Initialize the Vertex AI SDK\n",
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
+        "# Initialize the Vertex AI SDK for Python\n",
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -393,8 +369,6 @@
      },
      "outputs": [],
      "source": [
-        "from google.cloud import bigquery\n",
-        "\n",
        "# Set up BigQuery client\n",
        "bq_client = bigquery.Client(project=PROJECT_ID)"
      ]
@@ -405,7 +379,7 @@
        "id": "0a2c41bc91a6"
      },
      "source": [
-        "## Preprocess data and split data\n",
+        "### Preprocess data and split data\n",
        "First you should download and preprocess your data for training and testing.\n",
        "\n",
        "- Convert categorical features to numeric\n",
@@ -422,9 +396,6 @@
      },
      "outputs": [],
      "source": [
-        "import numpy as np\n",
-        "import pandas as pd\n",
-        "\n",
        "LABEL_COLUMN = \"species\"\n",
        "\n",
        "# Define the BigQuery source dataset\n",
@@ -456,7 +427,7 @@
        "id": "a39df4692a70"
      },
      "source": [
-        "## Create a Vertex AI Tabular Dataset from BigQuery dataset\n",
+        "### Create a Vertex AI Tabular dataset from BigQuery dataset\n",
        "\n",
        "Create a Vertex AI tabular dataset resource from your BigQuery training data.\n",
        "\n",
@@ -498,7 +469,7 @@
        "id": "train_custom_model"
      },
      "source": [
-        "## Train a model\n",
+        "### Train a model\n",
        "\n",
        "There are two ways you can train a model using a container image:\n",
        "\n",
@@ -811,7 +782,7 @@
        "id": "make_prediction"
      },
      "source": [
-        "## Make an online prediction request\n",
+        "### Make an online prediction request\n",
        "\n",
        "Send an online prediction request to your deployed model."
      ]
@@ -884,7 +855,7 @@
        "id": "undeploy_model"
      },
      "source": [
-        "## Undeploy models\n",
+        "### Undeploy models\n",
        "\n",
        "To undeploy all `Model` resources from the serving `Endpoint` resource, use the endpoint's `undeploy_all` method."
      ]
@@ -908,7 +879,8 @@
      "source": [
        "## Cleaning up\n",
        "\n",
-        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
+        "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:\n",
        "\n",
@@ -926,8 +898,6 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
        "# Delete the training job\n",
        "job.delete()\n",
        "\n",
@@ -938,9 +908,9 @@
        "endpoint.delete()\n",
        "\n",
        "# Warning: Setting this to true deletes everything in your bucket\n",
-        "delete_bucket = False\n",
+        "delete_bucket = True\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil rm -r $BUCKET_URI"
      ]
    }
@@ -24,7 +24,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "l2mMvIUG9meX"
@@ -33,29 +32,30 @@
        "# Profile model training performance using Cloud Profiler\n",
        "\n",
        "<table align=\"left\">\n",
-        "\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/custom_training_tensorboard_profiler.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/custom_training_tensorboard_profiler.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fcustom%2Fcustom_training_tensorboard_profiler.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/custom/custom_training_tensorboard_profiler.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",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
-        "  </td>                                                                                               \n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/custom_training_tensorboard_profiler.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
        "</table>"
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
@@ -69,7 +69,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "dmfmQL6w84pS"
@@ -93,7 +92,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "zfXf0r-K81Y-"
@@ -105,7 +103,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "I3KFLvpq87rs"
@@ -126,22 +123,28 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
-        "id": "ze4-nDLfK4pw"
+        "id": "61RBz8LLbxCR"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook. "
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "No17Cw5hgx12"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "2b4ef9b72d43"
+        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
@@ -149,178 +152,169 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
-        "id": "aUw6ibN-n5Za"
+        "id": "R5Xep4W9lq-Z"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel."
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "FM12wbWhn7w0"
+        "id": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
-        "id": "LgFWLeJfoGQu"
+        "id": "SbmM4z7FOBpM"
      },
      "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)."
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
-        "id": "8ckyxpX_oSzD"
+        "id": "dmWOrTJ3gx13"
      },
      "source": [
-        "#### Set your project ID\n",
+        "### Authenticate your notebook environment (Colab only)\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)"
+        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "zY8DKBoVoVy3"
+        "id": "NyKGtVQjgx13"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "DF4l8DTdWgPY"
+      },
+      "source": [
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# Set the project id\n",
-        "! gcloud config set project {PROJECT_ID}"
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
-        "id": "mSQjVQmMosMl"
+        "id": "OKtKGmr9pfr6"
      },
      "source": [
-        "#### Region\n",
+        "#### Create a Cloud Storage bucket\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)."
+        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "Se9FWWhLotvB"
+        "id": "In3aQanwYjFB"
      },
      "outputs": [],
      "source": [
-        "REGION = \"us-central1\"  # @param {type:\"string\"}"
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
-        "id": "IfJRIMBpo5Pg"
+        "id": "GOaOsIjxp0oB"
      },
      "source": [
-        "### Authenticate your Google Cloud account\n",
+        "**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": "Wn5QiIl2p16e"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil mb -l $LOCATION -p $PROJECT_ID $BUCKET_URI"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "OMrAJ8RGqBQu"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
-      ]
-    },
-    {
-      "attachments": {},
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "acFN0s3So9-Y"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "attachments": {},
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "dQ_mNwuapE5T"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**"
+        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "cR_MzpknpGgM"
+        "id": "AWRzBFExqERG"
      },
      "outputs": [],
      "source": [
-        "# ! gcloud auth login"
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
-      "attachments": {},
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "h-MuVI_ypJfw"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "BeaQlCwMpQUT"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "3ivZkPUjpaFz"
      },
      "source": [
-        "**4. Setup service account and permissions**\n",
+        "### Setup service account and permissions\n",
        "\n",
        "A service account will be used to create custom training jobs. If you do not want to use your project's Compute Engine service account, set SERVICE_ACCOUNT to another service account ID. You can create a service account by following the [instructions](https://cloud.google.com/iam/docs/creating-managing-service-accounts#creating)."
      ]
@@ -340,122 +334,56 @@
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "WWIxsCJFCg5Z"
+        "id": "autoset_service_account"
      },
      "outputs": [],
      "source": [
-        "# Grant Cloud Storage permission.\n",
-        "! gcloud projects add-iam-policy-binding $PROJECT_ID \\\n",
-        "        --member=\"serviceAccount:$SERVICE_ACCOUNT\" \\\n",
-        "        --role=\"roles/storage.admin\" \\\n",
-        "        --quiet"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "S_8_5jm-Gk6w"
-      },
-      "outputs": [],
-      "source": [
-        "# Grant AI Platform permission.\n",
-        "! gcloud projects add-iam-policy-binding $PROJECT_ID \\\n",
-        "        --member=\"serviceAccount:$SERVICE_ACCOUNT\" \\\n",
-        "        --role=\"roles/aiplatform.user\" \\\n",
-        "        --quiet"
-      ]
-    },
-    {
-      "attachments": {},
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "OKtKGmr9pfr6"
-      },
-      "source": [
-        "### Create a Cloud Storage bucket\n",
+        "import sys\n",
        "\n",
-        "Create a storage bucket to store intermediate artifacts such as datasets."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "In3aQanwYjFB"
-      },
-      "outputs": [],
-      "source": [
-        "BUCKET_URI = \"gs://your-bucket-name-unique\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "attachments": {},
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "GOaOsIjxp0oB"
-      },
-      "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": "Wn5QiIl2p16e"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
-      ]
-    },
-    {
-      "attachments": {},
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ankcS-vtp7Wv"
-      },
-      "source": [
-        "### Import libraries"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "WffSImMvp-Po"
-      },
-      "outputs": [],
-      "source": [
-        "from google.cloud import aiplatform"
-      ]
-    },
-    {
-      "attachments": {},
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "OMrAJ8RGqBQu"
-      },
-      "source": [
-        "### Initialize Vertex AI SDK for Python\n",
+        "IS_COLAB = \"google.colab\" in sys.modules\n",
+        "if (\n",
+        "    SERVICE_ACCOUNT == \"\"\n",
+        "    or SERVICE_ACCOUNT is None\n",
+        "    or SERVICE_ACCOUNT == \"[your-service-account]\"\n",
+        "):\n",
+        "    # Get your service account from gcloud\n",
+        "    if not IS_COLAB:\n",
+        "        shell_output = !gcloud auth list 2>/dev/null\n",
+        "        SERVICE_ACCOUNT = shell_output[2].replace(\"*\", \"\").strip()\n",
        "\n",
-        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
+        "    else:  # IS_COLAB:\n",
+        "        shell_output = ! gcloud projects describe  $PROJECT_ID\n",
+        "        project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "        SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "\n",
+        "    print(\"Service Account:\", SERVICE_ACCOUNT)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "set_service_account:pipelines"
+      },
+      "source": [
+        "#### Set service account access for Vertex AI Pipelines\n",
+        "\n",
+        "Run the following commands to grant your service account access to read and write pipeline artifacts in the bucket that you created in the previous step. You only need to run this step once per service account."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "AWRzBFExqERG"
+        "id": "bad678f5afda"
      },
      "outputs": [],
      "source": [
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectCreator $BUCKET_URI\n",
+        "\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectViewer $BUCKET_URI"
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "MYzX9ZgjsyyY"
@@ -468,6 +396,27 @@
        "Learn more about [Enabling service](https://cloud.google.com/artifact-registry/docs/enable-service)."
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4e302d6c9a11"
+      },
+      "source": [
+        "**Set project in colab environment (Colab only)**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7b9f07d9d1f4"
+      },
+      "outputs": [],
+      "source": [
+        "if IS_COLAB:\n",
+        "    ! gcloud config set project $PROJECT_ID"
+      ]
+    },
    {
      "cell_type": "code",
      "execution_count": null,
@@ -476,11 +425,11 @@
      },
      "outputs": [],
      "source": [
+        "# Enable Artifact Registry API\n",
        "! gcloud services enable artifactregistry.googleapis.com --quiet"
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "-ayTbNdi62_t"
@@ -492,7 +441,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "9c3QrDTZdaxk"
@@ -513,7 +461,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "vJrWKK0mY7H7"
@@ -533,7 +480,7 @@
      "outputs": [],
      "source": [
        "tensorboard = aiplatform.Tensorboard.create(\n",
-        "    display_name=TENSORBOARD_NAME, project=PROJECT_ID, location=REGION\n",
+        "    display_name=TENSORBOARD_NAME, project=PROJECT_ID, location=LOCATION\n",
        ")\n",
        "\n",
        "TENSORBOARD_INSTANCE_NAME = tensorboard.resource_name\n",
@@ -542,7 +489,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "yoR29gW2S24w"
@@ -560,7 +506,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "OrpUIkAIs_uQ"
@@ -591,7 +536,7 @@
        "\n",
        "! gcloud artifacts repositories create {DOCKER_REPOSITORY} \\\n",
        "    --repository-format=docker \\\n",
-        "    --location={REGION} \\\n",
+        "    --location={LOCATION} \\\n",
        "    --description=\"Repository for TensorBoard Custom Training Job\" \\\n",
        "    --quiet\n",
        "\n",
@@ -599,7 +544,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "MPqbvhO6tK_e"
@@ -607,7 +551,7 @@
      "source": [
        "### Configure authentication to your private Docker repository\n",
        "\n",
-        "Before you push or pull container images, configure Docker to use the `gcloud` command-line tool to authenticate requests to `Artifact Registry` for your region."
+        "Before you push or pull container images, configure Docker to use the `gcloud` command-line tool to authenticate requests to Artifact Registry for your location."
      ]
    },
    {
@@ -623,11 +567,10 @@
        "IS_COLAB = \"google.colab\" in sys.modules\n",
        "\n",
        "if not IS_COLAB:\n",
-        "    ! gcloud auth configure-docker {REGION}-docker.pkg.dev --quiet"
+        "    ! gcloud auth configure-docker {LOCATION}-docker.pkg.dev --quiet"
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "pSXrzCWi-bTD"
@@ -654,7 +597,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "hyAwgsoQmaYI"
@@ -784,7 +726,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "LiOEvgkA_90f"
@@ -819,7 +760,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "ihYFahRAr6sj"
@@ -837,13 +777,12 @@
      "outputs": [],
      "source": [
        "IMAGE_NAME = \"tensorboard-custom-container\"\n",
-        "IMAGE_URI = f\"{REGION}-docker.pkg.dev/{PROJECT_ID}/{DOCKER_REPOSITORY}/{IMAGE_NAME}\"\n",
+        "IMAGE_URI = f\"{LOCATION}-docker.pkg.dev/{PROJECT_ID}/{DOCKER_REPOSITORY}/{IMAGE_NAME}\"\n",
        "\n",
-        "! gcloud builds submit --project {PROJECT_ID} --region={REGION} --tag {IMAGE_URI} --timeout=60m --quiet"
+        "! gcloud builds submit --project {PROJECT_ID} --region={LOCATION} --tag {IMAGE_URI} --timeout=60m --quiet"
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "k4e6OYmimqTR"
@@ -870,7 +809,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "51hKGTbU32Eg"
@@ -915,7 +853,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "JkEe2Nb_85UD"
@@ -929,7 +866,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "TpV-iwP9qw9c"
@@ -958,9 +894,10 @@
      "source": [
        "delete_tensorboard = True\n",
        "delete_bucket = False\n",
+        "delete_generated_files_after_execution = False\n",
        "\n",
        "# Delete docker repository.\n",
-        "! gcloud artifacts repositories delete $DOCKER_REPOSITORY --project {PROJECT_ID} --location {REGION} --quiet\n",
+        "! gcloud artifacts repositories delete $DOCKER_REPOSITORY --project {PROJECT_ID} --location {LOCATION} --quiet\n",
        "\n",
        "job.delete()\n",
        "\n",
@@ -968,7 +905,10 @@
        "    tensorboard.delete()\n",
        "\n",
        "if delete_bucket and \"BUCKET_URI\" in globals():\n",
-        "    ! gsutil -m rm -r $BUCKET_URI"
+        "    ! gsutil -m rm -r $BUCKET_URI\n",
+        "\n",
+        "if delete_generated_files_after_execution:\n",
+        "    ! rm -rf $PYTHON_PACKAGE_APPLICATION_DIR Dockerfile Dockerfile"
      ]
    }
  ],
@@ -32,23 +32,29 @@
        "# Get started with Vertex AI Training for XGBoost\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/get_started_vertex_training_xgboost.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"><br> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/get_started_vertex_training_xgboost.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fcustom%2Fget_started_vertex_training_xgboost.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/custom/get_started_vertex_training_xgboost.ipynb\">\n",
-        "        <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "        <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br>\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/get_started_vertex_training_xgboost.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br>\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  \n",
        "</table>\n",
        "<br/><br/><br/>"
      ]
@@ -74,19 +80,19 @@
      "source": [
        "### Objective\n",
        "\n",
-        "In this tutorial, you learn how to use `Vertex AI Training` for training a XGBoost custom model.\n",
+        "In this tutorial, you learn how to use Vertex AI Training for training a XGBoost custom model.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
-        "- `Vertex AI Training`\n",
-        "- `Vertex AI Model` resource\n",
+        "- Vertex AI Training\n",
+        "- Vertex AI model resource\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Training using a Python package.\n",
        "- Report accuracy when hyperparameter tuning.\n",
        "- Save the model artifacts to Cloud Storage using GCSFuse.\n",
-        "- Create a `Vertex AI Model` resource."
+        "- Create a Vertex AI model resource."
      ]
    },
    {
@@ -127,9 +133,16 @@
        "id": "install_mlops"
      },
      "source": [
-        "## Installations\n",
-        "\n",
-        "Install the following packages to execute this notebook."
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "a2c2cb2109a0"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
@@ -140,47 +153,86 @@
      },
      "outputs": [],
      "source": [
-        "! pip3 install --upgrade google-cloud-aiplatform  -quiet"
+        "! pip3 install --upgrade google-cloud-aiplatform  --quiet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "e9255e3b156f"
+        "id": "ff555b32bab8"
      },
      "source": [
-        "### Colab Only: Uncomment the following cell to restart the kernel"
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "0c0b2427998a"
+        "id": "f09b4dff629a"
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "435b8e413535"
+        "id": "ee775571c2b5"
      },
      "source": [
-        "### Before you begin\n",
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "92e68cfc3a90"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
-        "#### Set your project ID\n",
+        "Authenticate your environment on Google Colab.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "46604f70e831"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\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)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "c066e7925949"
+      },
+      "source": [
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
@@ -192,90 +244,7 @@
      "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": "2e6b8b324ce1"
-      },
-      "source": [
-        "#### Region\n",
-        "\n",
-        "You can also change the `REGION` variable used by Vertex AI. \n",
-        "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": "6c43a8673066"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
-        "\n",
-        "**1. Vertex AI Workbench** \n",
-        "- Do nothing as you are already authenticated.\n",
-        "\n",
-        "**2. Local JupyterLab Instance,** uncomment and run."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "fbc9cd30cc4b"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "cd0da2c26879"
-      },
-      "source": [
-        "**3. Colab,** uncomment and run:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "a336a05c6149"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0461097edfa5"
-      },
-      "source": [
-        "**4. Service Account or other**\n",
-        "- See how to grant Cloud Storage permissions to your service account at [IAM Ch Examples](https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples)."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -306,7 +275,7 @@
        "id": "create_bucket"
      },
      "source": [
-        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+        "**If your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
@@ -317,27 +286,29 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION $BUCKET_URI"
+        "! gsutil mb -l $LOCATION $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "validate_bucket"
+        "id": "a56633b047ee"
      },
      "source": [
-        "Finally, validate access to your Cloud Storage bucket by examining its contents:"
+        "### Initialize Vertex AI SDK for Python"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "yWnghzKFox9S"
+        "id": "0e24370a7912"
      },
      "outputs": [],
      "source": [
-        "! gsutil ls -al $BUCKET_URI"
+        "import google.cloud.aiplatform as aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -348,43 +319,7 @@
      "source": [
        "### Set up variables\n",
        "\n",
-        "Next, set up some variables used throughout the tutorial.\n",
-        "### Import libraries and define constants"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "import_aip:mbsdk"
-      },
-      "outputs": [],
-      "source": [
-        "import os\n",
-        "\n",
-        "import google.cloud.aiplatform as aiplatform"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "init_aip:mbsdk"
-      },
-      "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": "JZg2sszQox9T"
-      },
-      "outputs": [],
-      "source": [
-        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
+        "Next, set up some variables used throughout the tutorial."
      ]
    },
    {
@@ -406,7 +341,7 @@
        "\n",
        "Learn more about [hardware accelerator support for your region](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators).\n",
        "\n",
-        "*Note*: TF releases before 2.3 for GPU support will fail to load the custom model in this tutorial. It is a known issue and fixed in TF 2.3. This is caused by static graph ops that are generated in the serving function. If you encounter this issue on your own custom models, use a container image for TF 2.3 with GPU support."
+        "*Note*: TF releases before 2.3 for GPU support is expected to fail to load the custom model in this tutorial. It's a known issue and fixed in TF 2.3. This is caused by static graph ops that are generated in the serving function. If you encounter this issue on your own custom models, use a container image for TF 2.3 with GPU support."
      ]
    },
    {
@@ -451,10 +386,10 @@
        "DEPLOY_VERSION = \"xgboost-cpu.1-1\"\n",
        "\n",
        "TRAIN_IMAGE = \"{}-docker.pkg.dev/vertex-ai/training/{}:latest\".format(\n",
-        "    REGION.split(\"-\")[0], TRAIN_VERSION\n",
+        "    LOCATION.split(\"-\")[0], TRAIN_VERSION\n",
        ")\n",
        "DEPLOY_IMAGE = \"{}-docker.pkg.dev/vertex-ai/prediction/{}:latest\".format(\n",
-        "    REGION.split(\"-\")[0], DEPLOY_VERSION\n",
+        "    LOCATION.split(\"-\")[0], DEPLOY_VERSION\n",
        ")"
      ]
    },
@@ -468,7 +403,7 @@
        "\n",
        "Next, set the machine type to use for training.\n",
        "\n",
-        "- Set the variable `TRAIN_COMPUTE` to configure  the compute resources for the VMs you will use for for training.\n",
+        "- Set the variable `TRAIN_COMPUTE` to configure  the compute resources for the VMs used for training.\n",
        " - `machine type`\n",
        "     - `n1-standard`: 3.75GB of memory per vCPU.\n",
        "     - `n1-highmem`: 6.5GB of memory per vCPU\n",
@@ -503,8 +438,8 @@
      "source": [
        "## Introduction to XGBoost training\n",
        "\n",
-        "Once you have trained a XGBoost model, you will want to save it at a Cloud Storage location, so it can subsequently be uploaded to a `Vertex AI Model` resource.\n",
-        "The XGBoost package does not have support to save the model to a Cloud Storage location. Instead, you will do the following steps to save to a Cloud Storage location.\n",
+        "Once you have trained a XGBoost model, save it at a Cloud Storage location, so it can subsequently be uploaded to a Vertex AI model resource.\n",
+        "The XGBoost package does not have support to save the model to a Cloud Storage location. Instead, do the following steps to save to a Cloud Storage location.\n",
        "\n",
        "1. Save the in-memory model to the local filesystem (e.g., model.bst).\n",
        "2. Use gsutil to copy the local copy to the specified Cloud Storage location.\n",
@@ -518,11 +453,18 @@
        "id": "examine_training_package:xgboost"
      },
      "source": [
-        "### Examine the training package\n",
-        "\n",
+        "### Examine the training package"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "06692103c9e5"
+      },
+      "source": [
        "#### Package layout\n",
        "\n",
-        "Before you start the training, you will look at how a Python package is assembled for a custom training job. When unarchived, the package contains the following directory/file layout.\n",
+        "Before you start the training, look at how a Python package is assembled for a custom training job. When unarchived, the package contains the following directory/file layout.\n",
        "\n",
        "- PKG-INFO\n",
        "- README.md\n",
@@ -534,11 +476,16 @@
        "\n",
        "The files `setup.cfg` and `setup.py` are the instructions for installing the package into the operating environment of the Docker image.\n",
        "\n",
-        "The file `trainer/task.py` is the Python script for executing the custom training job. *Note*, when we referred to it in the worker pool specification, we replace the directory slash with a dot (`trainer.task`) and dropped the file suffix (`.py`).\n",
-        "\n",
-        "#### Package Assembly\n",
-        "\n",
-        "In the following cells, you will assemble the training package."
+        "The file `trainer/task.py` is the Python script for executing the custom training job. *Note*, when `trainer/task.py` is referred to in the worker pool specification, the directory slash is replaced with a dot and the file suffix (`.py`) suffix is dropped: (`trainer.task`).\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "b4b1b6caa966"
+      },
+      "source": [
+        "#### Package Assembly"
      ]
    },
    {
@@ -581,7 +528,7 @@
        "Next, you create the `task.py` script for driving the training package. Some noteable steps include:\n",
        "\n",
        "- Command-line arguments:\n",
-        "    - `model-dir`: The location to save the trained model. When using Vertex AI custom training, the location will be specified in the environment variable: `AIP_MODEL_DIR`,\n",
+        "    - `model-dir`: The location to save the trained model. When using Vertex AI custom training, the location is specified in the environment variable: `AIP_MODEL_DIR`,\n",
        "    - `dataset_data_url`: The location of the training data to download.\n",
        "    - `dataset_labels_url`: The location of the training labels to download.\n",
        "    - `boost-rounds`: Tunable hyperparameter\n",
@@ -629,7 +576,7 @@
        "parser.add_argument(\"--dataset-labels-url\", dest=\"dataset_labels_url\",\n",
        "                    type=str, help=\"Download url for the training data labels.\")\n",
        "parser.add_argument(\"--boost-rounds\", dest=\"boost_rounds\",\n",
-        "                    default=20, type=int, help=\"Number of boosted rounds\")\n",
+        "                    default=20, type=int, help=\"Number of boosted rounds\") \n",
        "args = parser.parse_args()\n",
        "\n",
        "logging.getLogger().setLevel(logging.INFO)\n",
@@ -951,14 +898,14 @@
        "id": "cleanup:mbsdk"
      },
      "source": [
-        "# Cleaning up\n",
-        "\n",
+        "## Clean up\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:\n",
        "\n",
-        "- Custom Job (Custome Training job is remove in previous step)\n",
+        "- Custom Job (Custom Training job is removed in the previous step)\n",
+        "- Model\n",
        "- Cloud Storage Bucket"
      ]
    },
@@ -970,11 +917,11 @@
      },
      "outputs": [],
      "source": [
-        "delete_bucket = True\n",
+        "delete_bucket = False\n",
        "\n",
        "model.delete()\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil rm -r $BUCKET_URI"
      ]
    }
@@ -32,21 +32,24 @@
        "# Get started with Endpoint and shared VM\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/get_started_with_vertex_endpoint_and_shared_vm.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",
+        "  <td style=\"text-align: center\">\n",
        "        <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/get_started_with_vertex_endpoint_and_shared_vm.ipynb\">\n",
-        "        <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "        <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "        </a>\n",
        "  </td>\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fgithub.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fblob%2Fmain%2Fnotebooks%2Fofficial%2Fcustom%2Fget_started_with_vertex_endpoint_and_shared_vm.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>   \n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/custom/get_started_with_vertex_endpoint_and_shared_vm.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",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/get_started_with_vertex_endpoint_and_shared_vm.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
@@ -133,9 +136,16 @@
        "id": "install_mlops"
      },
      "source": [
-        "## Installations\n",
-        "\n",
-        "Install the packages required for executing this notebook."
+        "## Get Started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "deda65b8cc6e"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages"
      ]
    },
    {
@@ -146,51 +156,88 @@
      },
      "outputs": [],
      "source": [
-        "# Install the packages\n",
-        "\n",
        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
-        "                                 tensorflow \\\n",
-        "                                 tensorflow-hub"
+        "                                 tensorflow==2.15.1 \\\n",
+        "                                 tensorflow-hub==0.16.1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "restart"
+        "id": "ff555b32bab8"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel"
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "D-ZBOjErv5mM"
+        "id": "f09b4dff629a"
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "yfEglUHQk9S3"
+        "id": "4a2b7b59bbf7"
      },
      "source": [
-        "## Before you begin\n",
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "92e68cfc3a90"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
-        "### Set your project ID\n",
+        "Authenticate your environment on Google Colab.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "46604f70e831"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\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)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4f872cd812d0"
+      },
+      "source": [
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
@@ -204,87 +251,9 @@
        "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",
+        "! gcloud config set project {PROJECT_ID}\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": "gcp_authenticate"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
-        "\n",
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated.\n",
-        "\n",
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ce6043da7b33"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0367eac06a10"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "21ad4dbb4a61"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c13224697bfb"
-      },
-      "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."
+        "LOCATION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
@@ -326,7 +295,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION $BUCKET_URI"
+        "! gsutil mb -l $LOCATION $BUCKET_URI"
      ]
    },
    {
@@ -335,10 +304,7 @@
        "id": "setup_vars"
      },
      "source": [
-        "### Set up variables\n",
-        "\n",
-        "Next, set up some variables used throughout the tutorial.\n",
-        "### Import libraries and define constants"
+        "### Import libraries"
      ]
    },
    {
@@ -363,9 +329,7 @@
        "id": "init_aip:mbsdk"
      },
      "source": [
-        "### Initialize Vertex AI SDK for Python\n",
-        "\n",
-        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
+        "### Initialize Vertex AI SDK for Python"
      ]
    },
    {
@@ -389,7 +353,8 @@
        "\n",
        "Setup up the following constants for Vertex AI:\n",
        "\n",
-        "- `API_ENDPOINT`: The Vertex AI API service endpoint for `Endpoint` services."
+        "- `API_ENDPOINT`: The Vertex AI API service endpoint for `Endpoint` services.\n",
+        "- `PARENT`: The base resource path for all Vertex AI resources within your project and location."
      ]
    },
    {
@@ -401,10 +366,10 @@
      "outputs": [],
      "source": [
        "# API service endpoint\n",
-        "API_ENDPOINT = \"{}-aiplatform.googleapis.com\".format(REGION)\n",
+        "API_ENDPOINT = \"{}-aiplatform.googleapis.com\".format(LOCATION)\n",
        "\n",
        "# Vertex location root path for your dataset, model and endpoint resources\n",
-        "PARENT = \"projects/\" + PROJECT_ID + \"/locations/\" + REGION"
+        "PARENT = \"projects/\" + PROJECT_ID + \"/locations/\" + LOCATION"
      ]
    },
    {
@@ -415,9 +380,9 @@
      "source": [
        "## Set up clients\n",
        "\n",
-        "The Vertex  works as a client/server model. On your side (the Python script) you will create a client that sends requests and receives responses from the Vertex AI server.\n",
+        "The Vertex AI  works as a client/server model. On your side (the Python script) you create a client that sends requests and receives responses from the Vertex AI server.\n",
        "\n",
-        "You will use different clients in this tutorial for different steps in the workflow. So set them all up upfront.\n",
+        "You use different clients in this tutorial for different steps in the workflow. This means you have to set them all up at the start.\n",
        "\n",
        "- Endpoint Service for creating endpoints, and deploying models to endpoints."
      ]
@@ -456,14 +421,14 @@
        "\n",
        "You can set hardware accelerators for training and prediction.\n",
        "\n",
-        "Set the variables `DEPLOY_GPU/DEPLOY_NGPU` to use a container image supporting a GPU and the number of GPUs allocated to the virtual machine (VM) instance. For example, to use a GPU container image with 4 Nvidia Telsa K80 GPUs allocated to each VM, you would specify:\n",
+        "Set the variables `DEPLOY_GPU/DEPLOY_NGPU` to use a container image supporting a GPU and the number of GPUs allocated to the virtual machine (VM) instance. For example, to use a GPU container image with 4 Nvidia Telsa T4 GPUs allocated to each VM, you would specify:\n",
        "\n",
-        "    (aip.gapic.AcceleratorType.NVIDIA_TESLA_K80, 4)\n",
+        "    (aip.gapic.AcceleratorType.NVIDIA_TESLA_T4, 4)\n",
        "\n",
        "\n",
        "Otherwise specify `(None, None)` to use a container image to run on a CPU.\n",
        "\n",
-        "Learn more about [hardware accelerator support for your region](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators)."
+        "Learn more about [hardware accelerator support for your location](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators)."
      ]
    },
    {
@@ -476,7 +441,7 @@
      "source": [
        "if os.getenv(\"IS_TESTING_DEPLOY_GPU\"):\n",
        "    DEPLOY_GPU, DEPLOY_NGPU = (\n",
-        "        aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_K80,\n",
+        "        aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_T4,\n",
        "        int(os.getenv(\"IS_TESTING_DEPLOY_GPU\")),\n",
        "    )\n",
        "else:\n",
@@ -505,7 +470,7 @@
      },
      "outputs": [],
      "source": [
-        "TF = \"2.5\".replace(\".\", \"-\")\n",
+        "TF = \"2.13\".replace(\".\", \"-\")\n",
        "\n",
        "if DEPLOY_GPU:\n",
        "    DEPLOY_VERSION = \"tf2-gpu.{}\".format(TF)\n",
@@ -513,7 +478,7 @@
        "    DEPLOY_VERSION = \"tf2-cpu.{}\".format(TF)\n",
        "\n",
        "DEPLOY_IMAGE = \"{}-docker.pkg.dev/vertex-ai/prediction/{}:latest\".format(\n",
-        "    REGION.split(\"-\")[0], DEPLOY_VERSION\n",
+        "    LOCATION.split(\"-\")[0], DEPLOY_VERSION\n",
        ")\n",
        "\n",
        "print(\"Deployment:\", DEPLOY_IMAGE, DEPLOY_GPU, DEPLOY_NGPU)"
@@ -529,7 +494,7 @@
        "\n",
        "Next, set the machine type to use for prediction.\n",
        "\n",
-        "- Set the variable `DEPLOY_COMPUTE` to configure  the compute resources for the VMs you will use for for prediction.\n",
+        "- Set the variable `DEPLOY_COMPUTE` to configure  the compute resources for the VMs you use for for prediction.\n",
        " - `machine type`\n",
        "     - `n1-standard`: 3.75GB of memory per vCPU.\n",
        "     - `n1-highmem`: 6.5GB of memory per vCPU\n",
@@ -614,7 +579,7 @@
      "source": [
        "## Upload the model for serving\n",
        "\n",
-        "Next, you will upload your TFHub image classification model to Vertex AI `Model` service, which will create a Vertex AI `Model` resource for your model. During upload, you need to define a serving function to convert data to the format your model expects. If you send encoded data to Vertex AI, your serving function ensures that the data is decoded on the model server before it is passed as input to your model.\n",
+        "Next, you upload your TFHub image classification model to Vertex AI `Model` service, which creates a Vertex AI `Model` resource for your model. During upload, you need to define a serving function to convert data to the format your model expects. If you send encoded data to Vertex AI, your serving function ensures that the data is decoded on the model server before it's passed as input to your model.\n",
        "\n",
        "### How does the serving function work\n",
        "\n",
@@ -631,7 +596,7 @@
        "\n",
        "Both the preprocessing and post-processing functions are converted to static graphs which are fused to the model. The output from the underlying model is passed to the post-processing function. The post-processing function passes the converted/packaged output back to the HTTP server. The HTTP server returns the output as the HTTP response content.\n",
        "\n",
-        "One consideration you need to consider when building serving functions for TF.Keras models is that they run as static graphs. That means, you cannot use TF graph operations that require a dynamic graph. If you do, you will get an error during the compile of the serving function which will indicate that you are using an EagerTensor which is not supported."
+        "One consideration you need to consider when building serving functions for TF.Keras models is that they run as static graphs. That means, you cannot use TF graph operations that require a dynamic graph. If you do, you get an error during the compile of the serving function which indicates that you're using an EagerTensor which is'nt supported."
      ]
    },
    {
@@ -644,7 +609,7 @@
        "\n",
        "#### Preprocessing\n",
        "\n",
-        "To pass images to the prediction service, you encode the compressed (e.g., JPEG) image bytes into base 64 -- which makes the content safe from modification while transmitting binary data over the network. Since this deployed model expects input data as raw (uncompressed) bytes, you need to ensure that the base 64 encoded data gets converted back to raw bytes, and then preprocessed to match the model input requirements, before it is passed as input to the deployed model.\n",
+        "To pass images to the prediction service, you encode the compressed (e.g., JPEG) image bytes into base 64 -- which makes the content safe from modification while transmitting binary data over the network. Since this deployed model expects input data as raw (uncompressed) bytes, you need to ensure that the base 64 encoded data gets converted back to raw bytes, and then preprocessed to match the model input requirements, before it's passed as input to the deployed model.\n",
        "\n",
        "To resolve this, you define a serving function (`serving_fn`) and attach it to the model as a preprocessing step. Add a `@tf.function` decorator so the serving function is fused to the underlying model (instead of upstream on a CPU).\n",
        "\n",
@@ -711,9 +676,9 @@
        "\n",
        "You can get the signatures of your model's input and output layers by reloading the model into memory, and querying it for the signatures corresponding to each layer.\n",
        "\n",
-        "For your purpose, you need the signature of the serving function. Why? Well, when we send our data for prediction as a HTTP request packet, the image data is base64 encoded, and our TF.Keras model takes numpy input. Your serving function will do the conversion from base64 to a numpy array.\n",
+        "For your purpose, you need the signature of the serving function. Why? Well, when you send your data for prediction as a HTTP request packet, the image data is base64 encoded, and your TF.Keras model takes numpy input. Your serving function does the conversion from base64 to a numpy array.\n",
        "\n",
-        "When making a prediction request, you need to route the request to the serving function instead of the model, so you need to know the input layer name of the serving function -- which you will use later when you make a prediction request."
+        "When making a prediction request, you need to route the request to the serving function instead of the model, so you need to know the input layer name of the serving function -- which you use later when you make a prediction request."
      ]
    },
    {
@@ -824,7 +789,7 @@
        "\n",
        "For your purpose, you need the signature of the serving function. \n",
        "\n",
-        "When making a prediction request, you need to route the request to the serving function instead of the model, so you need to know the input layer name of the serving function -- which you will use later when you make a prediction request."
+        "When making a prediction request, you need to route the request to the serving function instead of the model, so you need to know the input layer name of the serving function -- which you use later when you make a prediction request."
      ]
    },
    {
@@ -865,7 +830,7 @@
      "outputs": [],
      "source": [
        "model_use = aiplatform.Model.upload(\n",
-        "    display_name=\"icn\",\n",
+        "    display_name=\"use\",\n",
        "    artifact_uri=MODEL_USE_DIR,\n",
        "    serving_container_image_uri=DEPLOY_IMAGE,\n",
        ")\n",
@@ -889,7 +854,7 @@
        "- `min_replica_count`: Auto-scaling, the minimum number of compute nodes.\n",
        "- `max_replica_count`: Auto-scaling, the maximum number of compute nodes.\n",
        "\n",
-        "Learn more about [Deployment Resource Pools]()."
+        "Learn more about [Deployment Resource Pools](https://cloud.google.com/vertex-ai/docs/predictions/model-co-hosting)."
      ]
    },
    {
@@ -937,7 +902,7 @@
        "-X POST \\\n",
        "-H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
        "-H \"Content-Type: application/json\" \\\n",
-        "https://{REGION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{REGION}/deploymentResourcePools \\\n",
+        "https://{LOCATION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{LOCATION}/deploymentResourcePools \\\n",
        "-d '{CREATE_RP_REQUEST}'"
      ]
    },
@@ -965,7 +930,7 @@
        "! curl -X GET \\\n",
        "-H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
        "-H \"Content-Type: application/json\" \\\n",
-        "https://{REGION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{REGION}/deploymentResourcePools/{DEPLOYMENT_RESOURCE_POOL_ID}"
+        "https://{LOCATION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{LOCATION}/deploymentResourcePools/{DEPLOYMENT_RESOURCE_POOL_ID}"
      ]
    },
    {
@@ -992,7 +957,7 @@
        "! curl -X GET \\\n",
        "-H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
        "-H \"Content-Type: application/json\" \\\n",
-        "https://{REGION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{REGION}/deploymentResourcePools"
+        "https://{LOCATION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{LOCATION}/deploymentResourcePools"
      ]
    },
    {
@@ -1003,7 +968,7 @@
      "source": [
        "## Creating two `Endpoint` resource\n",
        "\n",
-        "Next, you create two `Endpoint` resources using the `Endpoint.create()` method. At a minimum, you specify the display name for the endpoint. Optionally, you can specify the project and location (region); otherwise the settings are inherited by the values you set when you initialized the Vertex AI SDK with the `init()` method.\n",
+        "Next, you create two `Endpoint` resources using the `Endpoint.create()` method. At a minimum, you specify the display name for the endpoint. Optionally, you can specify the project and location (location); otherwise the settings are inherited by the values you set when you initialized the Vertex AI SDK for Python with the `init()` method.\n",
        "\n",
        "In this example, the following parameters are specified:\n",
        "\n",
@@ -1039,7 +1004,7 @@
      "source": [
        "## Deploy Model in a Deployment Resource Pool\n",
        "\n",
-        "After you have created a Model and an Endpoint, you are ready to deploy using the DeployModel API. See an example of the CURL command below. Notice how you specified the `shared_resources` of DeployedModel with the deployment resource name of the resource pool that was created. \n",
+        "After you have created a Model and an Endpoint, you're ready to deploy using the DeployModel API. See an example of the CURL command below. Notice how you specified the `shared_resources` of DeployedModel with the deployment resource name of the resource pool that was created. \n",
        "\n",
        "Model deployments for the same deployment resource pool can be started concurrently.\n",
        "\n",
@@ -1056,9 +1021,9 @@
      },
      "outputs": [],
      "source": [
-        "SHARED_RESOURCE = \"projects/{project_id}/locations/{region}/deploymentResourcePools/{deployment_resource_pool_id}\".format(\n",
+        "SHARED_RESOURCE = \"projects/{project_id}/locations/{location}/deploymentResourcePools/{deployment_resource_pool_id}\".format(\n",
        "    project_id=PROJECT_ID,\n",
-        "    region=REGION,\n",
+        "    location=LOCATION,\n",
        "    deployment_resource_pool_id=DEPLOYMENT_RESOURCE_POOL_ID,\n",
        ")\n",
        "\n",
@@ -1066,6 +1031,7 @@
        "    \"deployedModel\": {\n",
        "        \"model\": model_icn.resource_name,\n",
        "        \"shared_resources\": SHARED_RESOURCE,\n",
+        "        \"enable_container_logging\": True,\n",
        "    },\n",
        "    \"trafficSplit\": {\"0\": 100},\n",
        "}\n",
@@ -1086,14 +1052,14 @@
        "output = ! curl -X POST \\\n",
        " -H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
        " -H \"Content-Type: application/json\" \\\n",
-        "https://{REGION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{REGION}/endpoints/{ENDPOINT_ID}:deployModel \\\n",
+        "https://{LOCATION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}:deployModel \\\n",
        "-d '{DEPLOY_MODEL_REQUEST}'\n",
        "\n",
        "for line in output:\n",
        "    if '\"name\"' in line:\n",
        "        operation_id = line.split(\":\")[-1].strip()[:-1]\n",
-        "        break\n",
-        "print(operation_id)"
+        "        print(operation_id)\n",
+        "        break"
      ]
    },
    {
@@ -1104,7 +1070,7 @@
      "source": [
        "### Wait for deployment to complete\n",
        "\n",
-        "Next, you will query the status of the operation, waiting for the operation state `done` to be set to `true`."
+        "Next, you query the status of the operation, waiting for the operation state `done` to be set to `true`."
      ]
    },
    {
@@ -1122,7 +1088,7 @@
        "    status = ! curl -X GET \\\n",
        " -H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
        " -H \"Content-Type: application/json\" \\\n",
-        "https://{REGION}-aiplatform.googleapis.com/v1beta1/{operation_id}\n",
+        "https://{LOCATION}-aiplatform.googleapis.com/v1beta1/{operation_id}\n",
        "    for line in status:\n",
        "        if '\"done\"' in line.strip():\n",
        "            done = line.strip()[0:-1]\n",
@@ -1153,6 +1119,7 @@
        "    \"deployedModel\": {\n",
        "        \"model\": model_use.resource_name,\n",
        "        \"shared_resources\": SHARED_RESOURCE,\n",
+        "        \"enable_container_logging\": True,\n",
        "    },\n",
        "    \"trafficSplit\": {\"0\": 100},\n",
        "}\n",
@@ -1173,14 +1140,14 @@
        "output = ! curl -X POST \\\n",
        " -H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
        " -H \"Content-Type: application/json\" \\\n",
-        "https://{REGION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{REGION}/endpoints/{ENDPOINT_ID}:deployModel \\\n",
+        "https://{LOCATION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}:deployModel \\\n",
        "-d '{DEPLOY_MODEL_REQUEST}'\n",
        "\n",
        "for line in output:\n",
        "    if '\"name\"' in line:\n",
        "        operation_id = line.split(\":\")[-1].strip()[:-1]\n",
-        "        break\n",
-        "print(operation_id)"
+        "        print(operation_id)\n",
+        "        break"
      ]
    },
    {
@@ -1191,7 +1158,7 @@
      "source": [
        "### Wait for deployment to complete\n",
        "\n",
-        "Next, you will query the status of the operation, waiting for the operation state `done` to be set to `true`."
+        "Next, you query the status of the operation, waiting for the operation state `done` to be set to `true`."
      ]
    },
    {
@@ -1207,7 +1174,7 @@
        "    status = ! curl -X GET \\\n",
        " -H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
        " -H \"Content-Type: application/json\" \\\n",
-        "https://{REGION}-aiplatform.googleapis.com/v1beta1/{operation_id}\n",
+        "https://{LOCATION}-aiplatform.googleapis.com/v1beta1/{operation_id}\n",
        "    for line in status:\n",
        "        if '\"done\"' in line.strip():\n",
        "            done = line.strip()[0:-1]\n",
@@ -1253,7 +1220,7 @@
      "source": [
        "### Make the prediction request for the image classification model\n",
        "\n",
-        "Finally, you make a prediction request. Since the model was trained on ImageNet, the prediction will return the probabilities for the corresponding 1000 classes."
+        "Finally, you make a prediction request. Since the model was trained on ImageNet, the prediction returns the probabilities for the corresponding 1000 classes."
      ]
    },
    {
@@ -1304,7 +1271,7 @@
      "source": [
        "### Make the prediction request for the text sentence encoder model\n",
        "\n",
-        "Finally, you make a prediction request. The prediction will return an embedding which is a 500 element vector."
+        "Finally, you make a prediction request. The prediction returns an embedding which is a 500 element vector."
      ]
    },
    {
@@ -1318,100 +1285,6 @@
        "endpoint_use.predict([instance])"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "undeploy_model:mbsdk"
-      },
-      "source": [
-        "#### Undeploy the models\n",
-        "\n",
-        "When you are done doing predictions, you undeploy the model from the `Endpoint` resouce. This deprovisions all compute resources and ends billing for the deployed model."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "undeploy_model:mbsdk"
-      },
-      "outputs": [],
-      "source": [
-        "endpoint_icn.undeploy_all()\n",
-        "endpoint_use.undeploy_all()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "model_delete:mbsdk"
-      },
-      "source": [
-        "#### Delete the `Model` resources\n",
-        "\n",
-        "The method 'delete()' will delete the model."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "model_delete:mbsdk"
-      },
-      "outputs": [],
-      "source": [
-        "model_icn.delete()\n",
-        "model_use.delete()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "endpoint_delete:mbsdk"
-      },
-      "source": [
-        "#### Delete the `Endpoint` resources\n",
-        "\n",
-        "The method 'delete()' will delete the endpoint."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "endpoint_delete:mbsdk"
-      },
-      "outputs": [],
-      "source": [
-        "endpoint_icn.delete()\n",
-        "endpoint_use.delete()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c3baa3a7496e"
-      },
-      "source": [
-        "#### Delete the `DeploymentResourcePool`\n",
-        "\n",
-        "The method 'delete()' will delete your deployment resource pool."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ac40cc1d594a"
-      },
-      "outputs": [],
-      "source": [
-        "! curl -X DELETE \\\n",
-        "-H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
-        "-H \"Content-Type: application/json\" \\\n",
-        "https://{REGION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{REGION}/deploymentResourcePools/{DEPLOYMENT_RESOURCE_POOL_ID}"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -1437,9 +1310,32 @@
        "# Set this to true only if you'd like to delete your bucket\n",
        "delete_bucket = False\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil rm -r $BUCKET_URI\n",
        "\n",
+        "# Undeploy the models\n",
+        "# When you're done doing predictions, you undeploy the model from the `Endpoint` resouce. \n",
+        "# This deprovisions all compute resources and ends billing for the deployed model.\n",
+        "endpoint_icn.undeploy_all()\n",
+        "endpoint_use.undeploy_all()\n",
+        "\n",
+        "# Delete the `Endpoint` resources\n",
+        "# The method 'delete()' deletes the model.\n",
+        "model_icn.delete()\n",
+        "model_use.delete()\n",
+        "\n",
+        "# Delete the `Endpoint` resources\n",
+        "# The method 'delete()' deletes the endpoint.\n",
+        "endpoint_icn.delete()\n",
+        "endpoint_use.delete()\n",
+        "\n",
+        "# Delete the `DeploymentResourcePool`\n",
+        "# The method 'delete()' deletes your deployment resource pool.\n",
+        "! curl -X DELETE \\\n",
+        "-H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
+        "-H \"Content-Type: application/json\" \\\n",
+        "https://{LOCATION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{LOCATION}/deploymentResourcePools/{DEPLOYMENT_RESOURCE_POOL_ID}\n",
+        "\n",
        "!rm -f test.jpg"
      ]
    }
@@ -2,7 +2,7 @@
  "cells": [
    {
      "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 1,
      "metadata": {
        "id": "copyright"
      },
@@ -32,25 +32,27 @@
        "# Custom training and batch prediction\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/sdk-custom-image-classification-batch.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/sdk-custom-image-classification-batch.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fcustom%2Fsdk-custom-image-classification-batch.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
-        "  </td>\n",
-        " <td>\n",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/custom/sdk-custom-image-classification-batch.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",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
-        "</table>\n",
-        "<br/><br/><br/>"
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/sdk-custom-image-classification-batch.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
      ]
    },
    {
@@ -78,7 +80,7 @@
        "In this tutorial, you learn to use Vertex AI Training to create a custom trained model and use Vertex AI Batch Prediction to do a batch prediction on the trained model.\n",
        "\n",
        "\n",
-        "create a custom-trained model from a Python script in a Docker container using the Vertex AI SDK for Python, and then do a prediction on the deployed model by sending data. Alternatively, you can create custom-trained models using `gcloud` command-line tool, or online using the Cloud Console.\n",
+        "Create a custom-trained model from a Python script in a Docker container using the Vertex AI SDK for Python, and then do a prediction on the deployed model by sending data. Alternatively, you can create custom-trained models using `gcloud` command-line tool, or online using the Cloud Console.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
@@ -128,23 +130,29 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "install_aip"
+        "id": "3b1ffd5ab768"
      },
      "source": [
-        "### Install additional packages\n",
-        "\n",
-        "Install the following packages for executing this notebook."
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aae9ca040eab"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "1fd00fa70a2a"
+        "id": "23e23ce735c9"
      },
      "outputs": [],
      "source": [
-        "# Install the packages\n",
        "! pip3 install --upgrade google-cloud-aiplatform \\\n",
        "                        google-cloud-storage \\\n",
        "                        pillow  \\\n",
@@ -154,156 +162,92 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "d98bc9fdd80d"
+        "id": "ff555b32bab8"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel."
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "0dbf29389c65"
+        "id": "f09b4dff629a"
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "013daf3de88e"
+        "id": "54c5ef8a8f43"
      },
      "source": [
-        "## Before you begin"
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "8bc8a29f9001"
+        "id": "92e68cfc3a90"
      },
      "source": [
-        "#### Set your project ID\n",
+        "### Authenticate your notebook environment (Colab only)\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)"
+        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "e61aaa036444"
+        "id": "46604f70e831"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "107c51893a64"
+      },
+      "source": [
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "294fe4e5a671"
      },
      "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": "c4a624c8099d"
-      },
-      "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": "f83bd6013894"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "08bfd1eb44ef"
-      },
-      "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": "af349043f23b"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ad1138a125ea"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ce6043da7b33"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0367eac06a10"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "21ad4dbb4a61"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c13224697bfb"
-      },
-      "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."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -312,7 +256,7 @@
        "id": "ddbea904fbe5"
      },
      "source": [
-        "### Create a Cloud Storage bucket\n",
+        "#### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
@@ -325,7 +269,7 @@
      },
      "outputs": [],
      "source": [
-        "BUCKET_URI = \"gs://your-bucket-name-unique\"  # @param {type:\"string\"}"
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -345,7 +289,31 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
+        "! gsutil mb -l $LOCATION -p $PROJECT_ID $BUCKET_URI"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "fb03963cdb69"
+      },
+      "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": "a4f61991b160"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -365,31 +333,7 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
-        "from google.cloud import aiplatform"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "import_aip"
-      },
-      "source": [
-        "### Initialize Vertex AI \n",
-        "\n",
-        "Initialize the Vertex AI SDK for Python for your project."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "cNEiwLd0lugu"
-      },
-      "outputs": [],
-      "source": [
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
+        "import os"
      ]
    },
    {
@@ -970,9 +914,6 @@
      },
      "outputs": [],
      "source": [
-        "delete_training_job = True\n",
-        "delete_model = True\n",
-        "\n",
        "# Warning: Setting this to true will delete everything in your bucket\n",
        "delete_bucket = False\n",
        "\n",
@@ -982,7 +923,7 @@
        "# Delete the model\n",
        "model.delete()\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
@@ -32,25 +32,27 @@
        "# Custom training and online prediction\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/sdk-custom-image-classification-online.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/sdk-custom-image-classification-online.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fcustom%2Fsdk-custom-image-classification-online.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/custom/sdk-custom-image-classification-online.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",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
-        "</table>\n",
-        "<br/><br/><br/>"
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/sdk-custom-image-classification-online.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
      ]
    },
    {
@@ -128,23 +130,31 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "install_aip"
+        "id": "3b1ffd5ab768"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the packages required for executing this notebook."
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aae9ca040eab"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "1fd00fa70a2a"
+        "id": "0d3d0b3a3d88"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
+        "                                 google-cloud \\\n",
        "                                 google-cloud-storage \\\n",
        "                                 pillow \\\n",
        "                                 numpy"
@@ -153,142 +163,92 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "d98bc9fdd80d"
+        "id": "ff555b32bab8"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel."
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "0dbf29389c65"
+        "id": "f09b4dff629a"
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "eec0fc7a0963"
+        "id": "54c5ef8a8f43"
      },
      "source": [
-        "## Before you begin\n",
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "92e68cfc3a90"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
-        "#### Set your project ID\n",
-        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "* Run `gcloud config list`.\n",
-        "* Run `gcloud projects list`.\n",
-        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "e3ce64be5527"
+        "id": "46604f70e831"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "107c51893a64"
+      },
+      "source": [
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "294fe4e5a671"
      },
      "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": "7f2e7f78a864"
-      },
-      "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": "369a2258dd59"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "6a94297012d5"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
-        "\n",
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "28e3c4539627"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**\n"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "fbc9cd30cc4b"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "79efab26ad02"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "a336a05c6149"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "e3c28b6b796b"
-      },
-      "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."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -297,7 +257,7 @@
        "id": "70a42f1033a3"
      },
      "source": [
-        "### Create a Cloud Storage bucket\n",
+        "#### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
@@ -319,7 +279,7 @@
        "id": "0cb016da6de3"
      },
      "source": [
-        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+        "**If your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
@@ -330,7 +290,31 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
+        "! gsutil mb -l $LOCATION -p $PROJECT_ID $BUCKET_URI"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ae35ce9954e4"
+      },
+      "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": "a4f61991b160"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -353,32 +337,9 @@
        "import os\n",
        "\n",
        "import numpy as np\n",
-        "from google.cloud import aiplatform\n",
        "from PIL import Image"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "import_aip"
-      },
-      "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": "cNEiwLd0lugu"
-      },
-      "outputs": [],
-      "source": [
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -882,7 +843,7 @@
        "# Delete the endpoint\n",
        "endpoint.delete()\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
@@ -32,22 +32,28 @@
        "# Get started with BigQuery datasets\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/datasets/get_started_bq_datasets.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 style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/datasets/get_started_bq_datasets.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fdatasets%2Fget_started_bq_datasets.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/datasets/get_started_bq_datasets.ipynb\">\n",
-        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br>\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/datasets/get_started_bq_datasets.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br>\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
        "</table>\n",
        "<br/><br/><br/>"
      ]
@@ -63,7 +69,7 @@
        "\n",
        "This tutorial demonstrates how to use Vertex AI in production. This tutorial covers data management: get started with BigQuery datasets.\n",
        "\n",
-        "Learn more about [BigQuery Datasets](https://cloud.google.com/bigquery/docs/datasets-intro) and [Vertex AI for BigQuery users](https://cloud.google.com/vertex-ai/docs/beginner/bqml)."
+        "Learn more about [BigQuery datasets](https://cloud.google.com/bigquery/docs/datasets-intro) and [Vertex AI for BigQuery users](https://cloud.google.com/vertex-ai/docs/beginner/bqml)."
      ]
    },
    {
@@ -74,22 +80,22 @@
      "source": [
        "### Objective\n",
        "\n",
-        "In this tutorial, you learn how to use `BigQuery` as a dataset for training with `Vertex AI`.\n",
+        "In this tutorial, you learn how to use BigQuery as a dataset for training with Vertex AI.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
-        "- `Vertex AI Datasets`\n",
-        "- `BigQuery Datasets`\n",
+        "- Vertex AI datasets\n",
+        "- BigQuery datasets\n",
        "\n",
        "The steps performed include:\n",
        "\n",
-        "- Create a Vertex AI `Dataset` resource from `BigQuery` table -- compatible for `AutoML` training.\n",
-        "- Extract a copy of the dataset from `BigQuery` to a CSV file in Cloud Storage -- compatible for `AutoML` or custom training.\n",
-        "- Select rows from a `BigQuery` dataset into a `pandas` dataframe -- compatible for custom training.\n",
-        "- Select rows from a `BigQuery` dataset into a `tf.data.Dataset` -- compatible for custom training `TensorFlow` models.\n",
-        "- Select rows from extracted CSV files into a `tf.data.Dataset` -- compatible for custom training `TensorFlow` models.\n",
-        "- Create a `BigQuery` dataset from CSV files.\n",
-        "- Extract data from `BigQuery` table into a `DMatrix` -- compatible for custom training `XGBoost` models."
+        "- Create a Vertex AI dataset resource from BigQuery table -- compatible for AutoML training.\n",
+        "- Extract a copy of the dataset from BigQuery to a CSV file in Cloud Storage -- compatible for AutoML or custom training.\n",
+        "- Select rows from a BigQuery dataset into a pandas dataframe -- compatible for custom training.\n",
+        "- Select rows from a BigQuery dataset into a `tf.data.Dataset` -- compatible for custom training TensorFlow models.\n",
+        "- Select rows from extracted CSV files into a `tf.data.Dataset` -- compatible for custom training TensorFlow models.\n",
+        "- Create a BigQuery dataset from CSV files.\n",
+        "- Extract data from BigQuery table into a DMatrix -- compatible for custom training XGBoost models."
      ]
    },
    {
@@ -103,7 +109,7 @@
        "When doing E2E MLOps on Google Cloud, following are the best practices when dealing with structured (tabular) data in BigQuery:\n",
        "\n",
        "- For AutoML training:\n",
-        "  - Create a managed dataset with Vertex AI `TabularDataset`.\n",
+        "  - Create a managed dataset with Vertex AI TabularDataset.\n",
        "  - Use the BigQuery table as the input to the dataset.\n",
        "  - Specify columns and columns transformations when running the AutoML training pipeline job.\n",
        "\n",
@@ -164,164 +170,134 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "install_mlops"
+        "id": "d1ea81ac77f0"
      },
      "source": [
-        "## Installations\n",
-        "\n",
-        "Install the following packages to execute this notebook."
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e5d353aa47ac"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "install_mlops"
+        "id": "9bbcbe73e685"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
        "                                 google-cloud-bigquery \\\n",
        "                                 tensorflow \\\n",
-        "                                 tensorflow-io==0.18 \\\n",
+        "                                 tensorflow-io \\\n",
        "                                 xgboost \\\n",
        "                                 numpy \\\n",
        "                                 pandas \\\n",
-        "                                 pyarrow"
+        "                                 pyarrow \\\n",
+        "                                 db-dtypes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "restart"
+        "id": "16220914acc5"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel"
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "D-ZBOjErv5mM"
+        "id": "157953ab28f0"
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "yfEglUHQk9S3"
+        "id": "c87a2a5d7e35"
      },
      "source": [
-        "## Before you begin\n",
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5dccb1c8feb6"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
-        "### Set your project ID\n",
-        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "* Run `gcloud config list`.\n",
-        "* Run `gcloud projects list`.\n",
-        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "set_project_id"
+        "id": "cc7251520a07"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "c2fc3d7b6bfa"
+      },
+      "source": [
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "45769dc0c97d"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "LOCATION = \"us-central1\"  # @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": "gcp_authenticate"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
+        "from google.cloud import aiplatform\n",
        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
-        "\n",
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated.\n",
-        "\n",
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ce6043da7b33"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0367eac06a10"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "21ad4dbb4a61"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c13224697bfb"
-      },
-      "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."
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
@@ -352,7 +328,7 @@
        "id": "create_bucket"
      },
      "source": [
-        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+        "**If your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
@@ -363,7 +339,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION $BUCKET_URI"
+        "! gsutil mb -l $LOCATION $BUCKET_URI"
      ]
    },
    {
@@ -383,34 +359,11 @@
      },
      "outputs": [],
      "source": [
-        "import google.cloud.aiplatform as aiplatform\n",
        "import pandas as pd\n",
        "import xgboost as xgb\n",
        "from google.cloud import bigquery"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "init_aip:mbsdk,region"
-      },
-      "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,region"
-      },
-      "outputs": [],
-      "source": [
-        "aiplatform.init(project=PROJECT_ID, location=REGION)"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -462,14 +415,14 @@
        "id": "create_dataset:tabular,bq,lrg"
      },
      "source": [
-        "### Create the Dataset\n",
+        "### Create the dataset\n",
        "\n",
        "#### BigQuery input data\n",
        "\n",
-        "Next, create the `Dataset` resource using the `create` method for the `TabularDataset` class, which takes the following parameters:\n",
+        "Next, create the dataset resource using the `create` method for the `TabularDataset` class, which takes the following parameters:\n",
        "\n",
-        "- `display_name`: The human readable name for the `Dataset` resource.\n",
-        "- `bq_source`: Import data items from a BigQuery table into the `Dataset` resource.\n",
+        "- `display_name`: The human readable name for the dataset resource.\n",
+        "- `bq_source`: Import data items from a BigQuery table into the dataset resource.\n",
        "- `labels`: User defined metadata. In this example, you store the location of the Cloud Storage bucket containing the user defined data.\n",
        "\n",
        "Learn more about [TabularDataset from BigQuery table](https://cloud.google.com/vertex-ai/docs/datasets/create-dataset-api#aiplatform_create_dataset_tabular_bigquery_sample-python)."
@@ -535,14 +488,14 @@
        "id": "create_dataset:tabular,lrg"
      },
      "source": [
-        "### Create the Dataset\n",
+        "### Create the dataset\n",
        "\n",
        "#### CSV input data\n",
        "\n",
-        "Next, create the `Dataset` resource using the `create` method for the `TabularDataset` class, which takes the following parameters:\n",
+        "Next, create the dataset resource using the `create` method for the `TabularDataset` class, which takes the following parameters:\n",
        "\n",
-        "- `display_name`: The human readable name for the `Dataset` resource.\n",
-        "- `gcs_source`: A list of one or more dataset index files to import the data items into the `Dataset` resource.\n",
+        "- `display_name`: The human readable name for the dataset resource.\n",
+        "- `gcs_source`: A list of one or more dataset index files to import the data items into the dataset resource.\n",
        "- `labels`: User defined metadata. In this example, you store the location of the Cloud Storage bucket containing the user defined data.\n",
        "\n",
        "Learn more about [TabularDataset from CSV files](https://cloud.google.com/vertex-ai/docs/datasets/create-dataset-api#aiplatform_create_dataset_tabular_gcs_sample-python)"
@@ -1045,9 +998,9 @@
        "\n",
        "Otherwise, you can delete the individual resources you created in this tutorial:\n",
        "\n",
-        "- Vertex AI Dataset resource\n",
+        "- Vertex AI dataset resource\n",
        "- Cloud Storage Bucket\n",
-        "- BigQuery Dataset\n",
+        "- BigQuery dataset\n",
        "\n",
        "Set `delete_storage` to _True_ to delete the storage resources used in this notebook."
      ]
@@ -33,22 +33,27 @@
        "\n",
        "<table align=\"left\">\n",
        "\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/experiments/build_model_experimentation_lineage_with_prebuild_code.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/experiments/build_model_experimentation_lineage_with_prebuild_code.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fexperiments%2Fbuild_model_experimentation_lineage_with_prebuild_code.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
+        "  </td> \n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/experiments/build_model_experimentation_lineage_with_prebuild_code.ipynb\">\n",
-        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br>\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/experiments/build_model_experimentation_lineage_with_prebuild_code.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br>\n",
+        "      View on GitHub\n",
+        "    </a>\n",
        "  </td>                                                                                               \n",
        "</table>\n"
      ]
@@ -135,12 +140,19 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "i7EUnXsZhAGF"
+        "id": "f0316df526f8"
      },
      "source": [
-        "### Install additional packages\n",
-        "\n",
-        "Install additional package dependencies not installed in your notebook environment."
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "a2c2cb2109a0"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
@@ -152,141 +164,98 @@
      "outputs": [],
      "source": [
        "! pip3 install --upgrade --quiet joblib fsspec gcsfs scikit-learn \n",
-        "! pip3 install --upgrade --quiet google-cloud-aiplatform==1.35 "
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "restart"
+        "id": "ff555b32bab8"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel"
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "D-ZBOjErv5mM"
+        "id": "f09b4dff629a"
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "yfEglUHQk9S3"
+        "id": "ee775571c2b5"
      },
      "source": [
-        "## Before you begin\n",
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "92e68cfc3a90"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
-        "### Set your project ID\n",
-        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "* Run `gcloud config list`.\n",
-        "* Run `gcloud projects list`.\n",
-        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "set_project_id"
+        "id": "46604f70e831"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4f872cd812d0"
+      },
+      "source": [
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "294fe4e5a671"
      },
      "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": "gcp_authenticate"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
-        "\n",
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated.\n",
-        "\n",
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ce6043da7b33"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0367eac06a10"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "21ad4dbb4a61"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c13224697bfb"
-      },
-      "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."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -317,7 +286,7 @@
        "id": "create_bucket"
      },
      "source": [
-        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+        "**If your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
@@ -328,7 +297,40 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} {BUCKET_URI}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "70bb458c7580"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "906e31206b45"
+      },
+      "outputs": [],
+      "source": [
+        "import uuid\n",
+        "\n",
+        "from google.cloud import aiplatform as vertex_ai\n",
+        "\n",
+        "# Experiments\n",
+        "TASK = \"classification\"\n",
+        "MODEL_TYPE = \"naivebayes\"\n",
+        "EXPERIMENT_NAME = f\"{TASK}-{MODEL_TYPE}-{uuid.uuid1()}\"\n",
+        "EXPERIMENT_RUN_NAME = \"run-1\"\n",
+        "\n",
+        "\n",
+        "vertex_ai.init(\n",
+        "    project=PROJECT_ID, experiment=EXPERIMENT_NAME, staging_bucket=BUCKET_URI\n",
+        ")"
      ]
    },
    {
@@ -404,13 +406,9 @@
        "import collections\n",
        "import tempfile\n",
        "import time\n",
-        "import uuid\n",
        "from json import dumps\n",
        "\n",
-        "collections.Iterable = collections.abc.Iterable\n",
-        "\n",
-        "# Vertex AI\n",
-        "from google.cloud import aiplatform as vertex_ai"
+        "collections.abc.Iterable = collections.abc.Iterable"
      ]
    },
    {
@@ -436,11 +434,6 @@
        "DATASET_NAME = \"news_corpora\"\n",
        "DATASET_URI = f\"{BUCKET_URI}/{DATA_PATH}/raw/newsCorpora.csv\"\n",
        "\n",
-        "# Experiments\n",
-        "TASK = \"classification\"\n",
-        "MODEL_TYPE = \"naivebayes\"\n",
-        "EXPERIMENT_NAME = f\"{TASK}-{MODEL_TYPE}-{uuid.uuid1()}\"\n",
-        "EXPERIMENT_RUN_NAME = \"run-1\"\n",
        "\n",
        "# Preprocessing\n",
        "PREPROCESSED_DATASET_NAME = f\"preprocessed_{DATASET_NAME}\"\n",
@@ -472,30 +465,6 @@
        "MODEL_NAME = f\"{EXPERIMENT_NAME}-model\""
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "inR70nh38PeK"
-      },
-      "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": "Nz0nasrh8T3c"
-      },
-      "outputs": [],
-      "source": [
-        "vertex_ai.init(\n",
-        "    project=PROJECT_ID, experiment=EXPERIMENT_NAME, staging_bucket=BUCKET_URI\n",
-        ")"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -32,29 +32,26 @@
        "# Vertex AI: Track parameters and metrics for locally trained models\n",
        "\n",
        "<table align=\"left\">\n",
-        "\n",
        "  <td style=\"text-align: center\">\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/tensorboard/comparing_local_trained_models.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> <br> Open in Colab\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/experiments/comparing_local_trained_models.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
-        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Ftensorboard%2Fcomparing_local_trained_models.ipynb\">\n",
-        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"> <br> Open in Colab Enterprise\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fexperiments%2Fcomparing_local_trained_models.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/experiments/comparing_local_trained_models.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/tensorboard/comparing_local_trained_models.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"> <br>\n",
-        "      View on GitHub\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/experiments/comparing_local_trained_models.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td style=\"text-align: center\">\n",
-        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/tensorboard/comparing_local_trained_models.ipynb\">\n",
-        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"> <br>\n",
-        "      Open in Vertex AI Workbench\n",
-        "    </a>\n",
-        "  </td>                                                                                               \n",
        "</table>"
      ]
    },
@@ -128,12 +125,19 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "i7EUnXsZhAGF"
+        "id": "61RBz8LLbxCR"
      },
      "source": [
-        "### Install additional packages\n",
-        "\n",
-        "Install additional package dependencies not installed in your notebook environment."
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "No17Cw5hgx12"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
@@ -144,165 +148,103 @@
      },
      "outputs": [],
      "source": [
-        "! pip3 install tensorflow==2.8 google-cloud-aiplatform -q --no-warn-conflicts\n",
-        "! pip3 install protobuf==3.20.3"
+        "! pip3 install --upgrade tensorflow==2.8 \\\n",
+        "                        protobuf==3.20.3 \\\n",
+        "                        google-cloud-aiplatform \\\n",
+        "                        matplotlib \\\n",
+        "                        pandas \\\n",
+        "                        'numpy<2' -q --no-warn-conflicts"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "restart"
+        "id": "R5Xep4W9lq-Z"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel"
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "D-ZBOjErv5mM"
+        "id": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "yfEglUHQk9S3"
+        "id": "SbmM4z7FOBpM"
      },
      "source": [
-        "## Before you begin\n",
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dmWOrTJ3gx13"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
-        "### Set your project ID\n",
-        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "* Run `gcloud config list`.\n",
-        "* Run `gcloud projects list`.\n",
-        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "set_project_id"
+        "id": "NyKGtVQjgx13"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "DF4l8DTdWgPY"
+      },
+      "source": [
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Nqwi-5ufWp_B"
      },
      "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": "717d548a9e90"
-      },
-      "source": [
-        "#### Get your project number\n",
-        "Now that the project ID is set, you get your corresponding project number."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "7647b76e5fd7"
-      },
-      "outputs": [],
-      "source": [
-        "shell_output = ! gcloud projects list --filter=\"PROJECT_ID:'{PROJECT_ID}'\" --format='value(PROJECT_NUMBER)'\n",
-        "PROJECT_NUMBER = shell_output[0]\n",
-        "print(\"Project Number:\", PROJECT_NUMBER)"
-      ]
-    },
-    {
-      "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": "gcp_authenticate"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
-        "\n",
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated.\n",
-        "\n",
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ce6043da7b33"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0367eac06a10"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "21ad4dbb4a61"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c13224697bfb"
-      },
-      "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."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -333,7 +275,7 @@
        "id": "create_bucket"
      },
      "source": [
-        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+        "**If your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
@@ -344,7 +286,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION $BUCKET_URI"
+        "! gsutil mb -l $LOCATION -p $PROJECT_ID $BUCKET_URI"
      ]
    },
    {
@@ -434,7 +376,7 @@
      },
      "outputs": [],
      "source": [
-        "vertex_ai.init(project=PROJECT_ID, location=REGION)"
+        "vertex_ai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -443,26 +385,13 @@
        "id": "bba904753a2c"
      },
      "source": [
-        "### Create TensorBoard instance using Vertex AI TensorBoard\n",
+        "### Create a Vertex AI Experiment with backing Vertex AI TensorBoard\n",
        "\n",
-        "You can upload your TensorBoard logs by first creating a Vertex AI TensorBoard instance.\n",
-        "\n",
-        "**Notice** that if you did not activate yet, Vertex AI TensorBoard charges a monthly fee of $300 per unique active user. \n",
+        "You can create a Vertex AI Experiment by using the init() method. This automatically gets or creates the default Vertex AI TensorBoard and associates it with your Experiment.\n",
        "\n",
        "Learn more about [TensorBoard overview](https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-overview)."
      ]
    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "d9800e386d0c"
-      },
-      "outputs": [],
-      "source": [
-        "vertex_ai_tb = vertex_ai.Tensorboard.create()"
-      ]
-    },
    {
      "cell_type": "code",
      "execution_count": null,
@@ -471,7 +400,7 @@
      },
      "outputs": [],
      "source": [
-        "vertex_ai.init(experiment=EXPERIMENT_NAME, experiment_tensorboard=vertex_ai_tb)"
+        "vertex_ai.init(experiment=EXPERIMENT_NAME)"
      ]
    },
    {
@@ -537,7 +466,7 @@
        "    dataset[\"Origin\"] = dataset[\"Origin\"].map(\n",
        "        lambda x: {1: \"USA\", 2: \"Europe\", 3: \"Japan\"}.get(x)\n",
        "    )\n",
-        "    dataset = pd.get_dummies(dataset, prefix=\"\", prefix_sep=\"\")\n",
+        "    dataset = pd.get_dummies(dataset, prefix=\"\", prefix_sep=\"\", dtype=float)\n",
        "    return dataset\n",
        "\n",
        "\n",
@@ -826,19 +755,21 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
        "# Delete experiment\n",
        "exp = vertex_ai.Experiment(EXPERIMENT_NAME)\n",
+        "backing_tensorboard = exp.get_backing_tensorboard_resource()\n",
        "exp.delete(delete_backing_tensorboard_runs=True)\n",
        "\n",
        "# Delete Tensorboard\n",
-        "vertex_ai_tb.delete()\n",
+        "delete_tensorboard = False  # Set True for deletion\n",
+        "\n",
+        "if delete_tensorboard:\n",
+        "    backing_tensorboard.delete()\n",
        "\n",
        "# Delete Cloud Storage objects that were created\n",
-        "delete_bucket = False\n",
+        "delete_bucket = False  # Set True for deletion\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil rm -rf {BUCKET_URI}"
      ]
    }
@@ -33,20 +33,25 @@
        "\n",
        "<table align=\"left\">\n",
        "\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/experiments/comparing_pipeline_runs.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"><br> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fgithub.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fblob%2Fmain%2Fnotebooks%2Fofficial%2Fexperiments%2Fcomparing_pipeline_runs.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\\\" alt=\"Google Cloud Colab Enterprise logo\\\"><br> Open in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/experiments/comparing_pipeline_runs.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br>\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/experiments/comparing_pipeline_runs.ipynb\">\n",
-        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br>\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>                                                                                               \n",
@@ -61,7 +66,7 @@
      "source": [
        "## Overview\n",
        "\n",
-        "Depending on the model life cycle of your data science team, you would like to experiment and track training pipeline runs and its associated parameters. You then want to compare the runs of these Pipelines to each other in order to figure out which has the best configuration for generating the model you plan to register in the Vertex AI Model Registry.\n",
+        "As a Data Scientist, depending on the model life cycle, you may like to experiment and track training pipeline runs and its associated parameters. You then want to compare the runs of these Pipelines to each other in order to figure out which has the best configuration for generating the model you plan to register in the Vertex AI Model Registry.\n",
        "\n",
        "Learn more about [Vertex AI Experiments](https://cloud.google.com/vertex-ai/docs/experiments/intro-vertex-ai-experiments) and [Vertex AI Pipelines](https://cloud.google.com/vertex-ai/docs/pipelines/introduction)."
      ]
@@ -74,7 +79,7 @@
      "source": [
        "### Objective\n",
        "\n",
-        "In this notebook, you learn how to use Vertex AI Experiments to log a pipeline job and then compare different pipeline jobs.\n",
+        "In this tutorial, you learn how to use Vertex AI Experiments to log a pipeline job and then compare different pipeline jobs.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
@@ -85,8 +90,8 @@
        "\n",
        "* Formalize a training component\n",
        "* Build a training pipeline\n",
-        "* Run several Pipeline jobs and log their results\n",
-        "* Compare different Pipeline jobs"
+        "* Run several pipeline jobs and log their results\n",
+        "* Compare different pipeline jobs"
      ]
    },
    {
@@ -106,7 +111,7 @@
        "id": "de76bb18c85b"
      },
      "source": [
-        "### Costs \n",
+        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
@@ -116,8 +121,16 @@
        "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."
+        "Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "xlpUSF8AIj63"
+      },
+      "source": [
+        "### Get started"
      ]
    },
    {
@@ -126,9 +139,7 @@
        "id": "i7EUnXsZhAGF"
      },
      "source": [
-        "### Install additional packages\n",
-        "\n",
-        "Install additional package dependencies not installed in your notebook environment."
+        "### Install Vertex AI SDK for Python and other required packages"
      ]
    },
    {
@@ -139,7 +150,7 @@
      },
      "outputs": [],
      "source": [
-        "! pip3 install --force-reinstall 'google-cloud-aiplatform>=1.15' -q --no-warn-conflicts\n",
+        "! pip3 install --upgrade google-cloud-aiplatform -q --no-warn-conflicts\n",
        "! pip3 install kfp -q --no-warn-conflicts"
      ]
    },
@@ -149,7 +160,9 @@
        "id": "restart"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel."
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
@@ -160,11 +173,14 @@
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
@@ -173,14 +189,48 @@
        "id": "yfEglUHQk9S3"
      },
      "source": [
-        "## Before you begin\n",
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "LnTLX7UGKql8"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
-        "### Set your project ID\n",
+        "Authenticate your environment on Google Colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "GhWCxscxKyeG"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\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)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "CdBaOmh9K3tC"
+      },
+      "source": [
+        "### Set Google Cloud project information\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
@@ -192,89 +242,7 @@
      "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": "gcp_authenticate"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
-        "\n",
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated.\n",
-        "\n",
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ce6043da7b33"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0367eac06a10"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "21ad4dbb4a61"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c13224697bfb"
-      },
-      "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."
+        "LOCATION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
@@ -305,18 +273,18 @@
        "id": "create_bucket"
      },
      "source": [
-        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+        "**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": "create_bucket"
+        "id": "SVwNOMp39Qtz"
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION $BUCKET_URI"
+        "! gsutil mb -l $LOCATION $BUCKET_URI"
      ]
    },
    {
@@ -379,7 +347,7 @@
      "source": [
        "#### Set service account access for Vertex AI Pipelines\n",
        "\n",
-        "Run the following commands to grant your service account access to read and write pipeline artifacts in the bucket that you created in the previous step -- you only need to run these once per service account."
+        "Run the following commands to grant your service account access to read and write pipeline artifacts in the bucket that you created in the previous step. You only need to run this step once per service account."
      ]
    },
    {
@@ -436,20 +404,19 @@
      "source": [
        "import logging\n",
        "# General\n",
-        "import os\n",
        "import time\n",
        "import uuid\n",
        "\n",
        "logger = logging.getLogger(\"logger\")\n",
        "logging.basicConfig(level=logging.INFO)\n",
        "\n",
-        "import kfp.v2.compiler as compiler\n",
+        "import kfp.compiler as compiler\n",
        "# Pipeline Experiments\n",
-        "import kfp.v2.dsl as dsl\n",
+        "import kfp.dsl as dsl\n",
        "# Vertex AI\n",
        "from google.cloud import aiplatform as vertex_ai\n",
        "from google.cloud.aiplatform_v1.types.pipeline_state import PipelineState\n",
-        "from kfp.v2.dsl import Metrics, Model, Output, component"
+        "from kfp.dsl import Metrics, Model, Output, component"
      ]
    },
    {
@@ -506,10 +473,10 @@
        "Set the pre-built Docker container image for training and prediction.\n",
        "\n",
        "\n",
-        "For the latest list, see [Pre-built containers for training](https://cloud.google.com/ai-platform-unified/docs/training/pre-built-containers).\n",
+        "For the latest list, see [Pre-built containers for training](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers).\n",
        "\n",
        "\n",
-        "For the latest list, see [Pre-built containers for prediction](https://cloud.google.com/ai-platform-unified/docs/predictions/pre-built-containers)."
+        "For the latest list, see [Pre-built containers for prediction](https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers)."
      ]
    },
    {
@@ -531,7 +498,7 @@
        "id": "t1NLYz1R-KWv"
      },
      "source": [
-        "## Formalize your training as pipeline component\n"
+        "## Formalize the training as pipeline component\n"
      ]
    },
    {
@@ -542,7 +509,7 @@
      "source": [
        "Before you start running your pipeline experiments, you have to formalize your training as pipeline component.\n",
        "\n",
-        "To do that, build the pipeline by using the `kfp.v2.dsl.component` decorator to convert your training task into a pipeline component. "
+        "To do that, build the pipeline by using the `kfp.dsl.component` decorator to convert your training task into a pipeline component. The below example specifies a base image for the component (python:3.8)."
      ]
    },
    {
@@ -554,12 +521,13 @@
      "outputs": [],
      "source": [
        "@component(\n",
+        "    base_image=\"python:3.8\",\n",
        "    packages_to_install=[\n",
        "        \"numpy==1.18.5\",\n",
        "        \"pandas==1.0.4\",\n",
        "        \"scikit-learn==0.23.1\",\n",
        "        \"xgboost==1.1.1\",\n",
-        "    ]\n",
+        "    ],\n",
        ")\n",
        "def custom_trainer(\n",
        "    train_uri: str,\n",
@@ -692,7 +660,7 @@
      "source": [
        "## Build a pipeline\n",
        "\n",
-        "Below code will perform creating pipelineJob in associated project."
+        "Next, create the pipelineJob in associated project."
      ]
    },
    {
@@ -729,7 +697,9 @@
        "id": "RkfZ7qVAVjBO"
      },
      "source": [
-        "### Compile your pipeline into a JSON file"
+        "### Compile the pipeline\n",
+        "\n",
+        "Next, compile the pipeline to a JSON file."
      ]
    },
    {
@@ -758,9 +728,9 @@
        "id": "ZNb6kZ2l5t-O"
      },
      "source": [
-        "### Submit Pipeline runs\n",
+        "### Submit pipeline runs\n",
        "\n",
-        "Now that you have the pipeline, define its training configuration depending on the defined parameters. Below you have an example and how you can submit several pipeline runs. "
+        "Now that you have the pipeline, define its training configuration depending on the defined parameters. In the following example, you can see how to submit several pipeline runs."
      ]
    },
    {
@@ -812,7 +782,7 @@
      "source": [
        "### Check Pipeline run states\n",
        "\n",
-        "Vertex AI SDK provides you `get_experiment_df` method to monitor the status of pipeline runs. You can use it either to return parameters and metrics of the Pipeline Runs in the Vertex AI Experiment or in combination with `get` method of `PipelineJob` to return the pipeline job in Vertex AI Pipeline.\n"
+        "Vertex AI SDK provides you `get_experiment_df` method to monitor the status of pipeline runs. You can use it either to return parameters and metrics of the pipeline runs in the Vertex AI Experiment or in combination with `get` method of `PipelineJob` to return the pipeline job in Vertex AI Pipeline.\n"
      ]
    },
    {
@@ -915,20 +885,18 @@
        "exp = vertex_ai.Experiment(EXPERIMENT_NAME)\n",
        "exp.delete()\n",
        "\n",
-        "# Delete bucket\n",
-        "delete_bucket = False\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "# Delete the Cloud Storage bucket\n",
+        "delete_bucket = False  # Set True for deletion\n",
+        "if delete_bucket:\n",
        "    ! gsutil rm -rf {BUCKET_URI}\n",
        "\n",
        "# Remove local files\n",
-        "\n",
        "!rm {PIPELINE_TEMPLATE_FILE}"
      ]
    }
  ],
  "metadata": {
    "colab": {
-      "collapsed_sections": [],
      "name": "comparing_pipeline_runs.ipynb",
      "toc_visible": true
    },
@@ -43,13 +43,13 @@
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/experiments/delete_outdated_tensorboard_experiments.ipynb\">\n",
-        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/vertex-ai-samples/main/notebooks/official/experiments/delete_outdated_tensorboard_experiments.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
-        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/vertex-ai-samples/main/notebooks/official/experiments/delete_outdated_tensorboard_experiments.ipynb\">\n",
-        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/experiments/delete_outdated_tensorboard_experiments.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
@@ -131,15 +131,22 @@
        "to generate a cost estimate based on your projected usage."
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "d1ea81ac77f0"
+      },
+      "source": [
+        "## Get started"
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i7EUnXsZhAGF"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook."
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
@@ -166,7 +173,9 @@
        "id": "58707a750154"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel."
+        "### Colab only: Uncomment the ### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab.following cell to restart the kernel."
      ]
    },
    {
@@ -177,174 +186,81 @@
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "BF1j6f9HApxa"
+        "id": "c87a2a5d7e35"
      },
      "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)."
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "WReHDGG5g0XY"
+        "id": "5dccb1c8feb6"
      },
      "source": [
-        "#### Set your project ID\n",
+        "### Authenticate your notebook environment (Colab only)\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)"
+        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "oM1iC_MfAts1"
+        "id": "cc7251520a07"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "c2fc3d7b6bfa"
+      },
+      "source": [
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "45769dc0c97d"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "LOCATION = \"us-central1\"  # @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": "QAup21nz6LHk"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "sBCra4QMA2wR"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
+        "from google.cloud import aiplatform\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 permissions to your service account at https://cloud.google.com/marketplace/docs/grant-service-account-access."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "960505627ddf"
-      },
-      "source": [
-        "### Import libraries"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "PyQmSRbKA8r-"
-      },
-      "outputs": [],
-      "source": [
-        "from google.cloud import aiplatform"
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
@@ -473,28 +389,6 @@
        "## Delete outdated Vertex AI TensorBoard Experiments"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "EGfFVHhWDjkl"
-      },
-      "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": "G1l-7Wft3jb6"
-      },
-      "outputs": [],
-      "source": [
-        "aiplatform.init(project=PROJECT_ID, location=REGION)"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -543,7 +437,9 @@
      },
      "outputs": [],
      "source": [
-        "tensorboard_instances = aiplatform.Tensorboard.list(project=PROJECT_ID, location=REGION)\n",
+        "tensorboard_instances = aiplatform.Tensorboard.list(\n",
+        "    project=PROJECT_ID, location=LOCATION\n",
+        ")\n",
        "print(tensorboard_instances)"
      ]
    },
@@ -569,7 +465,7 @@
        "    TENSORBOARD_INSTANCE_ID = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "    TENSORBOARD_INSTANCE = aiplatform.Tensorboard(\n",
-        "        project=PROJECT_ID, location=REGION, tensorboard_name=TENSORBOARD_INSTANCE_ID\n",
+        "        project=PROJECT_ID, location=LOCATION, tensorboard_name=TENSORBOARD_INSTANCE_ID\n",
        "    )\n",
        "    print(TENSORBOARD_INSTANCE)"
      ]
@@ -90,7 +90,7 @@
      "source": [
        "### Objective\n",
        "\n",
-        "In this tutorial, you learn how to autolog paramenters and metrics of an ML experiment running on Vertex AI Training by leveraging the integration with Vertex AI Experiments.\n",
+        "In this tutorial, you learn how to autolog parameters and metrics of an ML experiment running on Vertex AI Training by leveraging the integration with Vertex AI Experiments.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
@@ -41,25 +41,27 @@
        "# Get started with Vertex AI Experiments\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/experiments/get_started_with_vertex_experiments.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/experiments/get_started_with_vertex_experiments.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "    <td>\n",
-        "        <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/experiments/get_started_with_vertex_experiments.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",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fexperiments%2Fget_started_with_vertex_experiments.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/experiments/get_started_with_vertex_experiments.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",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
-        "</table>\n",
-        "<br/><br/><br/>"
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/experiments/get_started_with_vertex_experiments.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
      ]
    },
    {
@@ -163,12 +165,19 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "install_mlops"
+        "id": "61RBz8LLbxCR"
      },
      "source": [
-        "## Installations\n",
-        "\n",
-        "Install the following packages for executing this notebook."
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "No17Cw5hgx12"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
@@ -179,141 +188,98 @@
      },
      "outputs": [],
      "source": [
-        "! pip3 install --upgrade --quiet google-cloud-aiplatform\n"
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "restart"
+        "id": "R5Xep4W9lq-Z"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel"
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "D-ZBOjErv5mM"
+        "id": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "yfEglUHQk9S3"
+        "id": "SbmM4z7FOBpM"
      },
      "source": [
-        "## Before you begin\n",
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dmWOrTJ3gx13"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
-        "### Set your project ID\n",
-        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "* Run `gcloud config list`.\n",
-        "* Run `gcloud projects list`.\n",
-        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "set_project_id"
+        "id": "NyKGtVQjgx13"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "DF4l8DTdWgPY"
+      },
+      "source": [
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Nqwi-5ufWp_B"
      },
      "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": "gcp_authenticate"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
-        "\n",
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated.\n",
-        "\n",
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ce6043da7b33"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0367eac06a10"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "21ad4dbb4a61"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c13224697bfb"
-      },
-      "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."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -355,7 +321,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION $BUCKET_URI"
+        "! gsutil mb -l $LOCATION $BUCKET_URI"
      ]
    },
    {
@@ -454,7 +420,7 @@
      },
      "outputs": [],
      "source": [
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -528,7 +494,7 @@
        "\n",
        "\n",
        "TRAIN_IMAGE = \"{}-docker.pkg.dev/vertex-ai/training/{}:latest\".format(\n",
-        "    REGION.split(\"-\")[0], TRAIN_VERSION\n",
+        "    LOCATION.split(\"-\")[0], TRAIN_VERSION\n",
        ")\n",
        "\n",
        "print(\"Training:\", TRAIN_IMAGE, TRAIN_GPU, TRAIN_NGPU)"
@@ -1484,7 +1450,7 @@
        "\n",
        "delete_bucket = False\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil rm -rf {BUCKET_URI}"
      ]
    }
@@ -33,23 +33,28 @@
        "\n",
        "<table align=\"left\">\n",
        "\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/experiments/get_started_with_vertex_experiments_autologging.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/experiments/get_started_with_vertex_experiments_autologging.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fexperiments%2Fget_started_with_vertex_experiments_autologging.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/experiments/get_started_with_vertex_experiments_autologging.ipynb\">\n",
-        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br>\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/experiments/get_started_with_vertex_experiments_autologging.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br>\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
        "</table>"
      ]
    },
@@ -83,7 +88,7 @@
      "source": [
        "### Objective\n",
        "\n",
-        "In this tutorial, you learn how to use `Vertex AI Autologging`.\n",
+        "In this tutorial, you learn how to use Vertex AI Autologging.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
@@ -133,9 +138,16 @@
        "id": "i7EUnXsZhAGF"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook."
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e5d353aa47ac"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
@@ -149,7 +161,7 @@
        "# Install the packages\n",
        "USER = \"\"\n",
        "! pip3 install {USER} --upgrade google-cloud-aiplatform tensorflow\n",
-        "! pip3 install {USER} --upgrade pandas scikit-learn category_encoders torch torchdata torchmetrics mlflow\n",
+        "! pip3 install {USER} --upgrade pandas==2.0.* scikit-learn category_encoders torch torchdata torchmetrics mlflow\n",
        "! pip3 install {USER} --upgrade protobuf==3.20.3\n"
      ]
    },
@@ -159,7 +171,9 @@
        "id": "58707a750154"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel."
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
@@ -170,82 +184,140 @@
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "BF1j6f9HApxa"
+        "id": "c87a2a5d7e35"
      },
      "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)."
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "WReHDGG5g0XY"
+        "id": "5dccb1c8feb6"
      },
      "source": [
-        "#### Set your project ID\n",
+        "### Authenticate your notebook environment (Colab only)\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)"
+        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "oM1iC_MfAts1"
+        "id": "cc7251520a07"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "c2fc3d7b6bfa"
+      },
+      "source": [
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "f02130bff721"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# Set the project id\n",
-        "! gcloud config set project {PROJECT_ID}"
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "region"
+        "id": "zgPO1eR3CYjk"
      },
      "source": [
-        "#### Region\n",
+        "### Create a Cloud Storage bucket\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)."
+        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "jWz5v12n3jb3"
+        "id": "MzGDU7TWdts_"
      },
      "outputs": [],
      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-EcIXiGsCePi"
+      },
+      "source": [
+        "**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 $LOCATION -p $PROJECT_ID $BUCKET_URI"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3d5191a94246"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "d61513af29a4"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform as vertex_ai\n",
+        "\n",
+        "vertex_ai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)\n"
      ]
    },
    {
@@ -279,121 +351,6 @@
        "UUID = generate_uuid()"
      ]
    },
-    {
-      "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",
-        "\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 = \"gs://your-bucket-name-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": {
@@ -428,7 +385,7 @@
        "id": "9dAKzp-i1UJX"
      },
      "source": [
-        "### Download dataset\n",
+        "### Download dataset\n",
        "\n",
        "Download the car evaluation dataset from the public Cloud Storage bucket."
      ]
@@ -468,28 +425,6 @@
        "!head {processed_data_filepath} -n 5"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "960505627ddf"
-      },
-      "source": [
-        "### Import libraries\n",
-        "\n",
-        "Import the Vertex AI SDK to log experiments in Vertex AI Experiments."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "PyQmSRbKA8r-"
-      },
-      "outputs": [],
-      "source": [
-        "from google.cloud import aiplatform as vertex_ai"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -660,11 +595,14 @@
        "    print(\"Building model...\")\n",
        "    inputs_layer = get_input_layer(features_values)\n",
        "    features_layer = get_features_layer(inputs_layer, features_values)\n",
-        "    x = tf.keras.layers.Concatenate()(features_layer.values())\n",
+        "    print(\"FEATURE LAYER TYPE: \", type(features_layer.values()))\n",
+        "    print(\"FEATURE LAYER TYPE Mod: \", type(features_layer))\n",
+        "    print(\"FEATURE LAYER: \", features_layer)\n",
+        "    x = tf.keras.layers.Concatenate()(list(features_layer.values()))\n",
        "    x = tf.keras.layers.Dense(10, activation=\"relu\")(x)\n",
        "    x = tf.keras.layers.Dense(5, activation=\"relu\")(x)\n",
        "    output_layer = tf.keras.layers.Dense(1)(x)\n",
-        "    model = tf.keras.Model(inputs=inputs_layer.values(), outputs=output_layer)\n",
+        "    model = tf.keras.Model(inputs=list(inputs_layer.values()), outputs=output_layer)\n",
        "\n",
        "    # Compile model\n",
        "    model.compile(\n",
@@ -802,9 +740,9 @@
        "id": "init_aip:mbsdk,all"
      },
      "source": [
-        "### Initialize Vertex AI SDK for Python and set seed for reproducibility\n",
+        "### Set seed for reproducibility\n",
        "\n",
-        "Initialize the Vertex AI SDK for Python for your project and set seed to guarantee reproducibility."
+        "Set seed to guarantee reproducibility."
      ]
    },
    {
@@ -815,7 +753,6 @@
      },
      "outputs": [],
      "source": [
-        "vertex_ai.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)\n",
        "set_seed(8)"
      ]
    },
@@ -878,7 +815,7 @@
        "experiment_tensorboard = vertex_ai.Tensorboard.create()\n",
        "vertex_ai.init(\n",
        "    project=PROJECT_ID,\n",
-        "    location=REGION,\n",
+        "    location=LOCATION,\n",
        "    staging_bucket=BUCKET_URI,\n",
        "    experiment=autologged_experiment_name,\n",
        "    experiment_tensorboard=experiment_tensorboard,\n",
@@ -29,28 +29,30 @@
        "id": "title"
      },
      "source": [
-        "# Vertex AI SDK: AutoML training tabular binary classification model for batch explanation\n",
+        "# Vertex AI SDK: Batch explanation for AutoML tabular binary classification model\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/explainable_ai/sdk_automl_tabular_binary_classification_batch_explain.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/explainable_ai/sdk_automl_tabular_binary_classification_batch_explain.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fexplainable_ai%2Fsdk_automl_tabular_binary_classification_batch_explain.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/explainable_ai/sdk_automl_tabular_binary_classification_batch_explain.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",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
-        "</table>\n",
-        "<br/><br/><br/>"
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/explainable_ai/sdk_automl_tabular_binary_classification_batch_explain.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
      ]
    },
    {
@@ -61,8 +63,7 @@
      "source": [
        "## Overview\n",
        "\n",
-        "\n",
-        "This tutorial demonstrates how to use the Vertex AI SDK to create tabular binary classification models and do batch prediction with explanation using a Google Cloud [AutoML](https://cloud.google.com/vertex-ai/docs/start/automl-users) model.\n",
+        "This tutorial demonstrates how to use the Vertex AI SDK to create a tabular binary classification model and do batch prediction with explanations using Google Cloud [AutoML](https://cloud.google.com/vertex-ai/docs/start/automl-users).\n",
        "\n",
        "Learn more about [Classification for tabular data](https://cloud.google.com/vertex-ai/docs/tabular-data/classification-regression/overview). Learn more about [Vertex Explainable AI](https://cloud.google.com/vertex-ai/docs/explainable-ai/overview)."
      ]
@@ -77,26 +78,26 @@
        "\n",
        "In this tutorial, you learn to use `AutoML` to create a tabular binary classification model from a Python script, and then learn to use `Vertex AI Batch Prediction` to make predictions with explanations. You can alternatively create and deploy models using the `gcloud` command-line tool or online using the Cloud Console.\n",
        "\n",
-        "This tutorial uses the following Google Cloud ML services:\n",
+        "This tutorial uses the following Vertex AI services:\n",
        "\n",
-        "- `Vertex AI AutoML`\n",
-        "- `Vertex AI Batch Prediction`\n",
-        "- `Vertex Explainable AI`\n",
-        "- `Vertex AI Model` resource\n",
+        "- Vertex AI AutoML\n",
+        "- Vertex AI Batch Prediction\n",
+        "- Vertex Explainable AI\n",
+        "- Vertex AI Model resource\n",
        "\n",
        "\n",
        "The steps performed include:\n",
        "\n",
-        "- Create a `Vertex Dataset` resource.\n",
-        "- Train an `AutoML` tabular binary classification model.\n",
+        "- Create a Vertex AI managed dataset resource.\n",
+        "- Train an AutoML tabular binary classification model.\n",
        "- View the model evaluation metrics for the trained model.\n",
        "- Make a batch prediction request with explainability.\n",
        "\n",
        "There is one key difference between using batch prediction and using online prediction:\n",
        "\n",
-        "* Prediction Service: Does an on-demand prediction for the entire set of instances (i.e., one or more data items) and returns the results in real-time.\n",
+        "* **Prediction Service**: Does an on-demand prediction for the entire set of instances (i.e., one or more data items) and returns the results in real-time.\n",
        "\n",
-        "* Batch Prediction Service: Does a queued (batch) prediction for the entire set of instances in the background and stores the results in a Cloud Storage bucket when ready."
+        "* **Batch Prediction Service**: Does a queued (batch) prediction for the entire set of instances in the background and stores the results in a Cloud Storage bucket when ready."
      ]
    },
    {
@@ -107,7 +108,7 @@
      "source": [
        "### Dataset\n",
        "\n",
-        "The dataset used for this tutorial is the Bank Marketing. This dataset does not require any feature engineering. The version of the dataset you use in this tutorial is stored in a public Cloud Storage bucket."
+        "This tutorial uses a bank marketing dataset. This dataset doesn't require any feature engineering. The version of the dataset used in this tutorial is stored in a public Cloud Storage bucket."
      ]
    },
    {
@@ -130,15 +131,22 @@
        "to generate a cost estimate based on your projected usage."
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f0316df526f8"
+      },
+      "source": [
+        "## Get started"
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "install_aip:mbsdk"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the packages required for executing this notebook."
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
@@ -160,7 +168,9 @@
        "id": "restart"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel"
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
@@ -171,27 +181,64 @@
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "yfEglUHQk9S3"
+        "id": "ee775571c2b5"
      },
      "source": [
-        "## Before you begin\n",
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "92e68cfc3a90"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
-        "### Set your project ID\n",
+        "Authenticate your environment on Google Colab.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "46604f70e831"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\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)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "c708f022953d"
+      },
+      "source": [
+        "### Set Google Cloud project information \n",
+        "\n",
+        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
@@ -203,89 +250,7 @@
      "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": "gcp_authenticate"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
-        "\n",
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated.\n",
-        "\n",
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ce6043da7b33"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0367eac06a10"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "21ad4dbb4a61"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c13224697bfb"
-      },
-      "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."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -316,7 +281,7 @@
        "id": "create_bucket"
      },
      "source": [
-        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+        "**If your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
@@ -327,30 +292,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION $BUCKET_URI"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "setup_vars"
-      },
-      "source": [
-        "### Set up variables\n",
-        "\n",
-        "Next, set up some variables used throughout the tutorial.\n",
-        "### Import libraries and define constants"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "import_aip:mbsdk"
-      },
-      "outputs": [],
-      "source": [
-        "import google.cloud.aiplatform as aip"
+        "! gsutil mb -l $LOCATION -p $PROJECT_ID $BUCKET_URI"
      ]
    },
    {
@@ -361,7 +303,9 @@
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
-        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
+        "To get started using Vertex AI, you must [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
+        "\n",
+        "Initialize the Vertex AI SDK for Python using the location and staging bucket."
      ]
    },
    {
@@ -372,7 +316,9 @@
      },
      "outputs": [],
      "source": [
-        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -381,9 +327,9 @@
        "id": "tutorial_start:automl"
      },
      "source": [
-        "# Tutorial\n",
+        "## Take a quick peek at the dataset\n",
        "\n",
-        "Now you are ready to start creating your own AutoML tabular binary classification model."
+        "You use a version of the Bank Marketing dataset that's stored in a public Cloud Storage bucket, using a CSV index file."
      ]
    },
    {
@@ -392,9 +338,7 @@
        "id": "import_file:u_dataset,csv"
      },
      "source": [
-        "#### Location of Cloud Storage training data.\n",
-        "\n",
-        "Now set the variable `IMPORT_FILE` to the location of the CSV index file in Cloud Storage."
+        "Set the variable `IMPORT_FILE` to the Cloud Storage location of the CSV index file."
      ]
    },
    {
@@ -411,16 +355,12 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "quick_peek:tabular"
+        "id": "4a5007bb4010"
      },
      "source": [
-        "#### Quick peek at your data\n",
+        "Count the number of examples by counting the number of rows in the CSV index file  (`wc -l`) and then peek at the first few rows.\n",
        "\n",
-        "You use a version of the Bank Marketing dataset that is stored in a public Cloud Storage bucket, using a CSV index file.\n",
-        "\n",
-        "Start by doing a quick peek at the data. You count the number of examples by counting the number of rows in the CSV index file  (`wc -l`) and then peek at the first few rows.\n",
-        "\n",
-        "You also need for training to know the heading name of the label column, which is save as `label_column`. For this dataset, it is the last column in the CSV file."
+        "For training using AutoML, you need to know the title of the label column, which is later saved as `label_column`. In this dataset, it's saved as the last column."
      ]
    },
    {
@@ -452,11 +392,11 @@
      "source": [
        "### Create the Dataset\n",
        "\n",
-        "Next, create the `Dataset` resource using the `create` method for the `TabularDataset` class, which takes the following parameters:\n",
+        "Next, create the Vertex AI managed dataset resource using the `create` method for the `TabularDataset` class, which takes the following parameters:\n",
        "\n",
-        "- `display_name`: The human readable name for the `Dataset` resource.\n",
-        "- `gcs_source`: A list of one or more dataset index files to import the data items into the `Dataset` resource.\n",
-        "- `bq_source`: Alternatively, import data items from a BigQuery table into the `Dataset` resource.\n",
+        "- `display_name`: The human readable name for the resource.\n",
+        "- `gcs_source`: A list of one or more dataset index files to import the data items into the resource.\n",
+        "- `bq_source`: Alternatively, import data items from a BigQuery table into the resource.\n",
        "\n",
        "This operation may take several minutes."
      ]
@@ -469,7 +409,7 @@
      },
      "outputs": [],
      "source": [
-        "dataset = aip.TabularDataset.create(\n",
+        "dataset = aiplatform.TabularDataset.create(\n",
        "    display_name=\"Bank Marketing\", gcs_source=[IMPORT_FILE]\n",
        ")\n",
        "\n",
@@ -484,11 +424,13 @@
      "source": [
        "### Create and run training pipeline\n",
        "\n",
-        "To train an AutoML model, you perform two steps: 1) create a training pipeline, and 2) run the pipeline.\n",
+        "To train the AutoML model on your dataset, you perform the following steps: \n",
+        "1) Create a training pipeline.\n",
+        "2) Run the pipeline.\n",
        "\n",
        "#### Create training pipeline\n",
        "\n",
-        "An AutoML training pipeline is created with the `AutoMLTabularTrainingJob` class, with the following parameters:\n",
+        "An AutoML training pipeline is created using the `AutoMLTabularTrainingJob` class, with the following parameters:\n",
        "\n",
        "- `display_name`: The human readable name for the `TrainingJob` resource.\n",
        "- `optimization_prediction_type`: The type task to train the model for.\n",
@@ -496,20 +438,20 @@
        "  - `regression`: A tabular regression model.\n",
        "- `column_transformations`: (Optional): Transformations to apply to the input columns\n",
        "- `optimization_objective`: The optimization objective to minimize or maximize.\n",
-        "  - binary classification:\n",
+        "  - For binary classification:\n",
        "    - `minimize-log-loss`\n",
        "    - `maximize-au-roc`\n",
        "    - `maximize-au-prc`\n",
        "    - `maximize-precision-at-recall`\n",
        "    - `maximize-recall-at-precision`\n",
-        "  - multi-class classification:\n",
+        "  - For multi-class classification:\n",
        "    - `minimize-log-loss`\n",
-        "  - regression:\n",
+        "  - For regression:\n",
        "    - `minimize-rmse`\n",
        "    - `minimize-mae`\n",
        "    - `minimize-rmsle`\n",
        "\n",
-        "The instantiated object is the DAG (directed acyclic graph) for the training pipeline."
+        "The instantiated object is a DAG (directed acyclic graph) for the training pipeline."
      ]
    },
    {
@@ -520,7 +462,7 @@
      },
      "outputs": [],
      "source": [
-        "dag = aip.AutoMLTabularTrainingJob(\n",
+        "dag = aiplatform.AutoMLTabularTrainingJob(\n",
        "    display_name=\"bank\",\n",
        "    optimization_prediction_type=\"classification\",\n",
        "    optimization_objective=\"minimize-log-loss\",\n",
@@ -537,20 +479,20 @@
      "source": [
        "#### Run the training pipeline\n",
        "\n",
-        "Next, you run the DAG to start the training job by invoking the method `run`, with the following parameters:\n",
+        "Next, you run the DAG to start the training job by invoking the `run` method with the following parameters:\n",
        "\n",
-        "- `dataset`: The `Dataset` resource to train the model.\n",
+        "- `dataset`: The managed dataset resource to train the model.\n",
        "- `model_display_name`: The human readable name for the trained model.\n",
        "- `training_fraction_split`: The percentage of the dataset to use for training.\n",
        "- `test_fraction_split`: The percentage of the dataset to use for test (holdout data).\n",
        "- `validation_fraction_split`: The percentage of the dataset to use for validation.\n",
        "- `target_column`: The name of the column to train as the label.\n",
-        "- `budget_milli_node_hours`: (optional) Maximum training time specified in unit of millihours (1000 = hour).\n",
-        "- `disable_early_stopping`: If `True`, training maybe completed before using the entire budget if the service believes it cannot further improve on the model objective measurements.\n",
+        "- `budget_milli_node_hours`: (optional) Maximum training time specified in unit of millihours (1000 = 1 hour). The training cost of the model doesn't exceed this budget.\n",
+        "- `disable_early_stopping`: If set to `True`, training maybe completed before using the entire budget if the service believes it can't further improve on the model objective measurements.\n",
        "\n",
-        "The `run` method when completed returns the `Model` resource.\n",
+        "The `run` method when completed returns the Vertex AI model resource.\n",
        "\n",
-        "The execution of the training pipeline will take upto 20 minutes."
+        "The execution of the training pipeline takes upto 20 minutes."
      ]
    },
    {
@@ -567,7 +509,7 @@
        "    training_fraction_split=0.6,\n",
        "    validation_fraction_split=0.2,\n",
        "    test_fraction_split=0.2,\n",
-        "    budget_milli_node_hours=8000,\n",
+        "    budget_milli_node_hours=1000,\n",
        "    disable_early_stopping=False,\n",
        "    target_column=label_column,\n",
        ")"
@@ -581,7 +523,7 @@
      "source": [
        "## Review model evaluation scores\n",
        "\n",
-        "After your model training has finished, you can review the evaluation scores for it using the `list_model_evaluations()` method. This method will return an iterator for each evaluation slice."
+        "After your model training has finished, you can review the evaluation scores using the `list_model_evaluations()` method. This method returns an iterator for each evaluation slice."
      ]
    },
    {
@@ -606,18 +548,7 @@
      "source": [
        "## Send a batch prediction request\n",
        "\n",
-        "Send a batch prediction to your deployed model."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "make_test_items:automl,batch_prediction"
-      },
-      "source": [
-        "### Make test items\n",
-        "\n",
-        "You use synthetic data as a test data items. Don't be concerned that we are using synthetic data."
+        "Send a batch prediction to your registered model."
      ]
    },
    {
@@ -626,9 +557,11 @@
        "id": "make_batch_file:automl,tabular,alt"
      },
      "source": [
-        "### Make the batch input file\n",
+        "### Prepare test data\n",
        "\n",
-        "Now make a batch input file, which you will store in your local Cloud Storage bucket. Unlike image, video and text, the batch input file for tabular is only supported for CSV. For CSV file, you make:\n",
+        "You use synthetic data as a test data items. Don't be concerned that we are using synthetic data.\n",
+        "\n",
+        "Make a batch input file, which is later stored in your Cloud Storage bucket. Unlike image, video and text, the batch input file for tabular model only supports CSV format. To prepare the CSV file, you make sure that:\n",
        "\n",
        "- The first line is the heading with the feature (fields) heading names.\n",
        "- Each remaining line is a separate prediction request with the corresponding feature values.\n",
@@ -665,7 +598,7 @@
      "source": [
        "### Make the batch explanation request\n",
        "\n",
-        "Now that your Model resource is trained, you can make a batch prediction by invoking the batch_predict() method, with the following parameters:\n",
+        "As your model resource is created, you can make a batch prediction by invoking the `batch_predict()` method with the following parameters:\n",
        "\n",
        "- `job_display_name`: The human readable name for the batch prediction job.\n",
        "- `gcs_source`: A list of one or more batch request input files.\n",
@@ -673,7 +606,7 @@
        "- `instances_format`: The format for the input instances, either 'csv' or 'jsonl'. Defaults to 'jsonl'.\n",
        "- `predictions_format`: The format for the output predictions, either 'csv' or 'jsonl'. Defaults to 'jsonl'.\n",
        "- `generate_explanations`: Set to `True` to generate explanations.\n",
-        "- `sync`: If set to True, the call will block while waiting for the asynchronous batch job to complete."
+        "- `sync`: If set to True, the call waits until the batch job is completed."
      ]
    },
    {
@@ -705,7 +638,9 @@
      "source": [
        "### Wait for completion of batch prediction job\n",
        "\n",
-        "Next, wait for the batch job to complete. Alternatively, one can set the parameter `sync` to `True` in the `batch_predict()` method to block until the batch prediction job is completed."
+        "Next, wait for the batch job to complete. \n",
+        "\n",
+        "Alternatively, you can set `sync` to `True` in the `batch_predict()` method earlier to perform the wait operation."
      ]
    },
    {
@@ -729,10 +664,9 @@
        "\n",
        "Next, get the explanation results from the completed batch prediction job.\n",
        "\n",
-        "The results are written to the Cloud Storage output bucket you specified in the batch prediction request. You call the method iter_outputs() to get a list of each Cloud Storage file generated with the results. Each file contains one or more explanation requests in a CSV format:\n",
+        "The results are written to the output Cloud Storage location you specified in the batch prediction request. \n",
        "\n",
-        "- CSV header + predicted_label\n",
-        "- CSV row + explanation, per prediction request"
+        "Call the `iter_outputs()` method to get a list of each Cloud Storage file generated in the results. Each file contains one or more explanation responses in the specified prediction format i.e., JSONL format."
      ]
    },
    {
@@ -782,17 +716,23 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
+        "# Delete the training job\n",
+        "dag.delete()\n",
+        "# Delete the tabular dataset\n",
+        "dataset.delete()\n",
+        "# Delete the model resource\n",
+        "model.delete()\n",
+        "# Delete the batch prediction job\n",
+        "batch_predict_job.delete()\n",
        "\n",
        "# Set this to true only if you'd like to delete your bucket\n",
        "delete_bucket = False\n",
        "\n",
-        "dataset.delete()\n",
-        "model.delete()\n",
-        "batch_predict_job.delete()\n",
+        "if delete_bucket:\n",
+        "    ! gsutil rm -r $BUCKET_URI\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
-        "    ! gsutil rm -r $BUCKET_URI"
+        "# Delete the locally generated files\n",
+        "! rm batch.csv tmp.csv"
      ]
    }
  ],
@@ -8,7 +8,7 @@
      },
      "outputs": [],
      "source": [
-        "# Copyright 2021 Google LLC\n",
+        "# Copyright 2024 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
@@ -32,25 +32,27 @@
        "# Vertex SDK: AutoML training tabular classification model for online explanation\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/explainable_ai/sdk_automl_tabular_classification_online_explain.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/explainable_ai/sdk_automl_tabular_classification_online_explain.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fexplainable_ai%2Fsdk_automl_tabular_classification_online_explain.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/explainable_ai/sdk_automl_tabular_classification_online_explain.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",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
-        "</table>\n",
-        "<br/><br/><br/>"
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/explainable_ai/sdk_automl_tabular_classification_online_explain.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
      ]
    },
    {
@@ -62,7 +64,7 @@
        "## Overview\n",
        "\n",
        "\n",
-        "This tutorial demonstrates how to use the Vertex AI SDK to create tabular classification models and do online prediction with explanation using a Google Cloud [AutoML](https://cloud.google.com/vertex-ai/docs/start/automl-users) model.\n",
+        "This tutorial demonstrates how to use the Vertex AI SDK to create tabular classification models and get online predictions with explanation using a Google Cloud [AutoML](https://cloud.google.com/vertex-ai/docs/start/automl-users) model.\n",
        "\n",
        "Learn more about [Classification for tabular data](https://cloud.google.com/vertex-ai/docs/tabular-data/classification-regression/overview) and [Vertex Explainable AI](https://cloud.google.com/vertex-ai/docs/explainable-ai/overview)."
      ]
@@ -75,25 +77,25 @@
      "source": [
        "### Objective\n",
        "\n",
-        "In this tutorial, you learn how to use `AutoML` to create a tabular binary classification model from a Python script, and then learn to use `Vertex AI Online Prediction` to make online predictions with explanations. You can alternatively create and deploy models using the `gcloud` command-line tool or online using the Cloud Console.\n",
+        "In this tutorial, you learn how to use AutoML to create a tabular binary classification model from a Python script. Then, you learn to use the Vertex AI prediction service to make a online prediction request with explanations. You can alternatively create and deploy models using the `gcloud` command-line tool or online using the Cloud Console.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
-        "- `Vertex AI AutoML`\n",
-        "- `Vertex AI Prediction`\n",
-        "- `Vertex Explainable AI`\n",
-        "- `Vertex AI Model` resource\n",
-        "- `Vertex AI Endpoint` resource\n",
+        "- Vertex AI AutoML\n",
+        "- Vertex AI Prediction\n",
+        "- Vertex Explainable AI\n",
+        "- Vertex AI model resource\n",
+        "- Vertex AI endpoint resource\n",
        "\n",
        "The steps performed include:\n",
        "\n",
-        "- Create a `Vertex Dataset` resource.\n",
-        "- Train an `AutoML` tabular binary classification model.\n",
+        "- Create a Vertex AI dataset resource.\n",
+        "- Train an AutoML tabular binary classification model.\n",
        "- View the model evaluation metrics for the trained model.\n",
-        "- Create a serving `Endpoint` resource.\n",
-        "- Deploy the `Model` resource to a serving `Endpoint` resource.\n",
+        "- Create a serving endpoint resource.\n",
+        "- Deploy the Model resource to a serving endpoint resource.\n",
        "- Make an online prediction request with explainability.\n",
-        "- Undeploy the `Model` resource."
+        "- Undeploy the Model resource."
      ]
    },
    {
@@ -104,7 +106,7 @@
      "source": [
        "### Dataset\n",
        "\n",
-        "The dataset used for this tutorial is the [Iris dataset](https://www.tensorflow.org/datasets/catalog/iris) from [TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/overview). This dataset does not require any feature engineering. The version of the dataset you will use in this tutorial is stored in a public Cloud Storage bucket. The trained model predicts the type of Iris flower species from a class of three species: setosa, virginica, or versicolor."
+        "The dataset used for this tutorial is the [Iris dataset](https://www.tensorflow.org/datasets/catalog/iris) from [TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/overview). The dataset does'nt require any feature engineering. The version of the dataset you use in this tutorial is stored in a public Cloud Storage bucket. The trained model predicts the type of Iris flower species from a class of three species: setosa, virginica, or versicolor."
      ]
    },
    {
@@ -130,12 +132,19 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "install_aip:mbsdk"
+        "id": "61RBz8LLbxCR"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the packages required for executing this notebook."
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "No17Cw5hgx12"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
@@ -155,135 +164,92 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "restart"
+        "id": "R5Xep4W9lq-Z"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel"
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "D-ZBOjErv5mM"
+        "id": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "yfEglUHQk9S3"
+        "id": "SbmM4z7FOBpM"
      },
      "source": [
-        "## Before you begin\n",
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dmWOrTJ3gx13"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
-        "### Set your project ID\n",
-        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "* Run `gcloud config list`.\n",
-        "* Run `gcloud projects list`.\n",
-        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "set_project_id"
+        "id": "NyKGtVQjgx13"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "DF4l8DTdWgPY"
+      },
+      "source": [
+        "### Set Google Cloud project information\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project. Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Nqwi-5ufWp_B"
      },
      "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": "gcp_authenticate"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
-        "\n",
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated.\n",
-        "\n",
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ce6043da7b33"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0367eac06a10"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "21ad4dbb4a61"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c13224697bfb"
-      },
-      "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."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -325,36 +291,13 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION $BUCKET_URI"
+        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "setup_vars"
-      },
-      "source": [
-        "### Set up variables\n",
-        "\n",
-        "Next, set up some variables used throughout the tutorial.\n",
-        "### Import libraries and define constants"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "import_aip:mbsdk"
-      },
-      "outputs": [],
-      "source": [
-        "import google.cloud.aiplatform as aip"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "init_aip:mbsdk"
+        "id": "53c6c71ac857"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
@@ -366,11 +309,13 @@
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "init_aip:mbsdk"
+        "id": "673ef17c3cff"
      },
      "outputs": [],
      "source": [
-        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -381,7 +326,7 @@
      "source": [
        "# Tutorial\n",
        "\n",
-        "Now you are ready to start creating your own AutoML tabular classification model."
+        "Now you're ready to create your own AutoML tabular classification model."
      ]
    },
    {
@@ -390,9 +335,9 @@
        "id": "import_file:u_dataset,csv"
      },
      "source": [
-        "#### Location of Cloud Storage training data.\n",
+        "#### Location of Cloud Storage training data\n",
        "\n",
-        "Now set the variable `IMPORT_FILE` to the location of the CSV index file in Cloud Storage."
+        "Set `IMPORT_FILE` variable to the location of the CSV index file in Cloud Storage."
      ]
    },
    {
@@ -414,11 +359,9 @@
      "source": [
        "#### Quick peek at your data\n",
        "\n",
-        "You will use a version of the Iris dataset that is stored in a public Cloud Storage bucket, using a CSV index file.\n",
+        "You can count the number of examples by counting the number of rows in the CSV index file  (`wc -l`) and then take a peek at the first few rows.\n",
        "\n",
-        "Start by doing a quick peek at the data. You count the number of examples by counting the number of rows in the CSV index file  (`wc -l`) and then peek at the first few rows.\n",
-        "\n",
-        "You also need for training to know the heading name of the label column, which is save as `label_column`. For this dataset, it is the last column in the CSV file."
+        "The heading name of the label column, is `label_column`, which is the last column in the CSV file."
      ]
    },
    {
@@ -450,11 +393,11 @@
      "source": [
        "### Create the Dataset\n",
        "\n",
-        "Next, create the `Dataset` resource using the `create` method for the `TabularDataset` class, which takes the following parameters:\n",
+        "Next, create the dataset using the `create` method for the `TabularDataset` class, which takes the following parameters:\n",
        "\n",
-        "- `display_name`: The human readable name for the `Dataset` resource.\n",
-        "- `gcs_source`: A list of one or more dataset index files to import the data items into the `Dataset` resource.\n",
-        "- `bq_source`: Alternatively, import data items from a BigQuery table into the `Dataset` resource.\n",
+        "- `display_name`: The human readable name for the dataset.\n",
+        "- `gcs_source`: A list of one or more dataset index files to import the data items into the dataset.\n",
+        "- `bq_source`: Alternatively, you can import data items from a BigQuery table into the dataset.\n",
        "\n",
        "This operation may take several minutes."
      ]
@@ -467,7 +410,9 @@
      },
      "outputs": [],
      "source": [
-        "dataset = aip.TabularDataset.create(display_name=\"Iris\", gcs_source=[IMPORT_FILE])\n",
+        "dataset = aiplatform.TabularDataset.create(\n",
+        "    display_name=\"Iris\", gcs_source=[IMPORT_FILE]\n",
+        ")\n",
        "\n",
        "print(dataset.resource_name)"
      ]
@@ -480,14 +425,16 @@
      "source": [
        "### Create and run training pipeline\n",
        "\n",
-        "To train an AutoML model, you perform two steps: 1) create a training pipeline, and 2) run the pipeline.\n",
+        "To train an AutoML model, you need perform two steps:\n",
+        "1) Create a training pipeline.\n",
+        "2) Specify your training parameters and run the job.\n",
        "\n",
        "#### Create training pipeline\n",
        "\n",
-        "An AutoML training pipeline is created with the `AutoMLTabularTrainingJob` class, with the following parameters:\n",
+        "Create an AutoML training pipeline using the `AutoMLTabularTrainingJob` class, with the following parameters:\n",
        "\n",
-        "- `display_name`: The human readable name for the `TrainingJob` resource.\n",
-        "- `optimization_prediction_type`: The type task to train the model for.\n",
+        "- `display_name`: The human readable name for the training job.\n",
+        "- `optimization_prediction_type`: The type of prediction the model is to produce.\n",
        "  - `classification`: A tabuar classification model.\n",
        "  - `regression`: A tabular regression model.\n",
        "- `column_transformations`: (Optional): Transformations to apply to the input columns\n",
@@ -503,9 +450,7 @@
        "  - regression:\n",
        "    - `minimize-rmse`\n",
        "    - `minimize-mae`\n",
-        "    - `minimize-rmsle`\n",
-        "\n",
-        "The instantiated object is the DAG (directed acyclic graph) for the training pipeline."
+        "    - `minimize-rmsle`\n"
      ]
    },
    {
@@ -516,13 +461,13 @@
      },
      "outputs": [],
      "source": [
-        "dag = aip.AutoMLTabularTrainingJob(\n",
+        "job = aiplatform.AutoMLTabularTrainingJob(\n",
        "    display_name=\"iris\",\n",
        "    optimization_prediction_type=\"classification\",\n",
        "    optimization_objective=\"minimize-log-loss\",\n",
        ")\n",
        "\n",
-        "print(dag)"
+        "print(job)"
      ]
    },
    {
@@ -533,9 +478,9 @@
      "source": [
        "#### Run the training pipeline\n",
        "\n",
-        "Next, you run the DAG to start the training job by invoking the method `run`, with the following parameters:\n",
+        "Start the training job by invoking the `run` method, with the following parameters:\n",
        "\n",
-        "- `dataset`: The `Dataset` resource to train the model.\n",
+        "- `dataset`: The dataset to train the model.\n",
        "- `model_display_name`: The human readable name for the trained model.\n",
        "- `training_fraction_split`: The percentage of the dataset to use for training.\n",
        "- `test_fraction_split`: The percentage of the dataset to use for test (holdout data).\n",
@@ -544,9 +489,9 @@
        "- `budget_milli_node_hours`: (optional) Maximum training time specified in unit of millihours (1000 = hour).\n",
        "- `disable_early_stopping`: If `True`, training maybe completed before using the entire budget if the service believes it cannot further improve on the model objective measurements.\n",
        "\n",
-        "The `run` method when completed returns the `Model` resource.\n",
+        "The `run` method when completed returns the model resource.\n",
        "\n",
-        "The execution of the training pipeline will take upto 20 minutes."
+        "The execution of the training pipeline may take upto 20 minutes."
      ]
    },
    {
@@ -557,7 +502,7 @@
      },
      "outputs": [],
      "source": [
-        "model = dag.run(\n",
+        "model = job.run(\n",
        "    dataset=dataset,\n",
        "    model_display_name=\"iris\",\n",
        "    training_fraction_split=0.6,\n",
@@ -577,7 +522,7 @@
      "source": [
        "## Review model evaluation scores\n",
        "\n",
-        "After your model training has finished, you can review the evaluation scores for it using the `list_model_evaluations()` method. This method will return an iterator for each evaluation slice."
+        "After your model training is complete, you can review the evaluation scores using the `list_model_evaluations()` method. This method returns an iterator for each evaluation slice."
      ]
    },
    {
@@ -602,7 +547,7 @@
      "source": [
        "## Deploy the model\n",
        "\n",
-        "Next, deploy your model for online prediction. To deploy the model, you invoke the `deploy` method, with the following parameters:\n",
+        "Next, deploy your model for online prediction. To deploy the model, invoke the `deploy` method, with the following parameters:\n",
        "\n",
        "- `machine_type`: The type of compute machine."
      ]
@@ -624,9 +569,9 @@
        "id": "make_prediction:xai"
      },
      "source": [
-        "## Send a online prediction request with explainability\n",
+        "## Send an online prediction request\n",
        "\n",
-        "Send a online prediction with explainability to your deployed model. In this method, the predicted response will include an explanation on how the features contributed to the explanation."
+        "Send a online prediction with explainability to your deployed model. In this method, the predicted response includes an explanation on how the features contributed to the explanation."
      ]
    },
    {
@@ -635,9 +580,9 @@
        "id": "make_test_item:automl,online_prediction"
      },
      "source": [
-        "### Make test item\n",
+        "### Create test item\n",
        "\n",
-        "You will use synthetic data as a test data item. Don't be concerned that we are using synthetic data -- we just want to demonstrate how to make a prediction."
+        "For demostration purpose, use synthetic data to generate predictions."
      ]
    },
    {
@@ -662,9 +607,9 @@
        "id": "explain_request:mbsdk,lcn"
      },
      "source": [
-        "### Make the prediction with explanation\n",
+        "### Generate predictions with explanations\n",
        "\n",
-        "Now that your `Model` resource is deployed to an `Endpoint` resource, one can do online explanations by sending prediction requests to the `Endpoint` resource.\n",
+        "Now that your model is deployed to an endpoint, you can get online explanations by sending prediction requests to your endpoint.\n",
        "\n",
        "#### Request\n",
        "\n",
@@ -672,18 +617,18 @@
        "\n",
        "    [feature_list]\n",
        "\n",
-        "Since the explain() method can take multiple items (instances), send your single test item as a list of one test item.\n",
+        "Since the `explain()` method can take multiple items (instances), send your single test item as a list of one test item.\n",
        "\n",
        "#### Response\n",
        "\n",
-        "The response from the explain() call is a Python dictionary with the following entries:\n",
+        "The response from the `explain()` call is a Python dictionary with the following entries:\n",
        "\n",
        "- `ids`: The internal assigned unique identifiers for each prediction request.\n",
        "- `displayNames`: The class names for each class label.\n",
        "- `confidences`: For classification, the predicted confidence, between 0 and 1, per class label.\n",
        "- `values`: For regression, the predicted value.\n",
-        "- `deployed_model_id`: The Vertex AI identifier for the deployed `Model` resource which did the predictions.\n",
-        "- `explanations`: The feature attributions"
+        "- `deployed_model_id`: The Vertex AI identifier for the deployed model resource which generated the predictions.\n",
+        "- `explanations`: The feature attributions returned by Explainable AI."
      ]
    },
    {
@@ -708,7 +653,7 @@
      "source": [
        "### Understanding the explanations response\n",
        "\n",
-        "First, you will look what your model predicted and compare it to the actual value."
+        "Take a look at model predictions and compare it to the actual value."
      ]
    },
    {
@@ -737,7 +682,7 @@
      "source": [
        "### Examine feature attributions\n",
        "\n",
-        "Next you will look at the feature attributions for this particular example. Positive attribution values mean a particular feature pushed your model prediction up by that amount, and vice versa for negative attribution values."
+        "Next look at the feature attributions for this particular example. Positive attribution values mean a particular feature pushed your model prediction up by that amount, and vice versa for negative attribution values."
      ]
    },
    {
@@ -767,9 +712,9 @@
      "source": [
        "### Check your explanations and baselines\n",
        "\n",
-        "To better make sense of the feature attributions you're getting, you should compare them with your model's baseline. In most cases, the sum of your attribution values + the baseline should be very close to your model's predicted value for each input. Also note that for regression models, the `baseline_score` returned from AI Explanations will be the same for each example sent to your model. For classification models, each class will have its own baseline.\n",
+        "To make better sense of the feature attributions you're getting, you should compare them with your model's baseline. In most cases, the sum of your attribution values + the baseline should be very close to your model's predicted value for each input. Also note that for regression models, the `baseline_score` returned from explanations remain the same for each example sent to your model. For classification models, each class has its own baseline.\n",
        "\n",
-        "In this section you'll send 10 test examples to your model for prediction in order to compare the feature attributions with the baseline. Then you'll run each test example's attributions through a sanity check in the `sanity_check_explanations` method.\n",
+        "In this section, send 10 test examples to your model for prediction in order to compare the feature attributions with the baseline. Then, run each test example's attributions through a sanity check in the `sanity_check_explanations` method.\n",
        "\n",
        "#### Get explanations"
      ]
@@ -855,28 +800,6 @@
        "    i += 1"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "undeploy_model:mbsdk"
-      },
-      "source": [
-        "## Undeploy the model\n",
-        "\n",
-        "When you are done doing predictions, you undeploy the model from the `Endpoint` resouce. This deprovisions all compute resources and ends billing for the deployed model."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "undeploy_model:mbsdk"
-      },
-      "outputs": [],
-      "source": [
-        "endpoint.undeploy_all()"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -899,15 +822,24 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
-        "delete_bucket = False\n",
+        "# Undeploy the model from the endpoint\n",
+        "endpoint.undeploy_all()\n",
        "\n",
+        "# Delete the model\n",
        "model.delete()\n",
+        "\n",
+        "# Delete the endpoint\n",
        "endpoint.delete()\n",
+        "\n",
+        "# Delete the dataset\n",
        "dataset.delete()\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "# Delete the automl training job\n",
+        "job.delete()\n",
+        "\n",
+        "# Delete the Cloud Storage bucket\n",
+        "delete_bucket = False  # set True for deletion\n",
+        "if delete_bucket:\n",
        "    ! gsutil rm -r $BUCKET_URI"
      ]
    }
@@ -32,25 +32,27 @@
        "# Vertex AI SDK: Custom training image classification model for batch prediction with explainabilty\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/explainable_ai/sdk_custom_image_classification_batch_explain.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/explainable_ai/sdk_custom_image_classification_batch_explain.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fexplainable_ai%2Fsdk_custom_image_classification_batch_explain.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/explainable_ai/sdk_custom_image_classification_batch_explain.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",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
-        "</table>\n",
-        "<br/><br/><br/> "
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/explainable_ai/sdk_custom_image_classification_batch_explain.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
      ]
    },
    {
@@ -127,19 +129,26 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "install_aip:mbsdk"
+        "id": "61RBz8LLbxCR"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the packages required for executing this notebook."
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "No17Cw5hgx12"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "install_aip:mbsdk"
+        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
@@ -147,7 +156,7 @@
        "\n",
        "! pip3 install --upgrade --quiet google-cloud-aiplatform\n",
        "! pip3 install --upgrade --quiet google-cloud-storage\n",
-        "! pip3 install --upgrade --quiet tensorflow\n",
+        "! pip3 install --upgrade --quiet tensorflow==2.15.1\n",
        "! pip3 install --upgrade --quiet opencv-python-headless\n",
        "\n",
        "if os.getenv(\"IS_TESTING\"):\n",
@@ -158,135 +167,92 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "restart"
+        "id": "R5Xep4W9lq-Z"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel"
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "D-ZBOjErv5mM"
+        "id": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "yfEglUHQk9S3"
+        "id": "SbmM4z7FOBpM"
      },
      "source": [
-        "## Before you begin\n",
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dmWOrTJ3gx13"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
-        "### Set your project ID\n",
-        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "* Run `gcloud config list`.\n",
-        "* Run `gcloud projects list`.\n",
-        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "set_project_id"
+        "id": "NyKGtVQjgx13"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "DF4l8DTdWgPY"
+      },
+      "source": [
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Nqwi-5ufWp_B"
      },
      "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": "gcp_authenticate"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
-        "\n",
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated.\n",
-        "\n",
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ce6043da7b33"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0367eac06a10"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "21ad4dbb4a61"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c13224697bfb"
-      },
-      "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."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -295,7 +261,7 @@
        "id": "bucket:mbsdk"
      },
      "source": [
-        "### Create a Cloud Storage bucket\n",
+        "#### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
@@ -317,7 +283,7 @@
        "id": "create_bucket"
      },
      "source": [
-        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+        "**If your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
@@ -328,32 +294,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION $BUCKET_URI"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "setup_vars"
-      },
-      "source": [
-        "### Set up variables\n",
-        "\n",
-        "Next, set up some variables used throughout the tutorial.\n",
-        "### Import libraries and define constants"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "import_aip:mbsdk"
-      },
-      "outputs": [],
-      "source": [
-        "import os\n",
-        "\n",
-        "from google.cloud import aiplatform"
+        "! gsutil mb -l $LOCATION -p $PROJECT_ID $BUCKET_URI"
      ]
    },
    {
@@ -362,7 +303,7 @@
        "id": "init_aip:mbsdk"
      },
      "source": [
-        "### Initialize Vertex AI SDK for Python\n",
+        "#### Initialize Vertex AI SDK for Python\n",
        "\n",
        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
      ]
@@ -375,6 +316,8 @@
      },
      "outputs": [],
      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
      ]
    },
@@ -395,7 +338,7 @@
        "\n",
        "Otherwise specify `(None, None)` to use a container image to run on a CPU.\n",
        "\n",
-        "Learn more [here](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators) hardware accelerator support for your region\n",
+        "Learn more about [hardware accelerator support](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators) for your location\n",
        "\n",
        "*Note*: TF releases before 2.3 for GPU support fail to load the custom model in this tutorial. It is a known issue and is fixed in TF 2.3 -- which is caused by static graph ops that are generated in the serving function. If you encounter this issue on your own custom models, use a container image for TF 2.3 with GPU support."
      ]
@@ -476,10 +419,10 @@
        "\n",
        "\n",
        "TRAIN_IMAGE = \"{}-docker.pkg.dev/vertex-ai/training/{}:latest\".format(\n",
-        "    REGION.split(\"-\")[0], TRAIN_VERSION\n",
+        "    LOCATION.split(\"-\")[0], TRAIN_VERSION\n",
        ")\n",
        "DEPLOY_IMAGE = \"{}-docker.pkg.dev/vertex-ai/prediction/{}:latest\".format(\n",
-        "    REGION.split(\"-\")[0], DEPLOY_VERSION\n",
+        "    LOCATION.split(\"-\")[0], DEPLOY_VERSION\n",
        ")\n",
        "\n",
        "print(\"Training:\", TRAIN_IMAGE, TRAIN_GPU, TRAIN_NGPU)\n",
@@ -571,7 +514,7 @@
        "\n",
        "The files `setup.cfg` and `setup.py` are the instructions for installing the package into the operating environment of the Docker image.\n",
        "\n",
-        "The file `trainer/task.py` is the Python script for executing the custom training job. *Note*, when we referred to it in the worker pool specification, we replace the directory slash with a dot (`trainer.task`) and dropped the file suffix (`.py`).\n",
+        "The file `trainer/task.py` is the Python script for executing the custom training job. *Note*, when `trainer/task.py` is referred to in the worker pool specification, the directory slash is replaced with a dot and the file suffix (`.py`) is dropped (`trainer.task`).\n",
        "\n",
        "#### Package Assembly\n",
        "\n",
@@ -615,7 +558,7 @@
      "source": [
        "#### Task.py contents\n",
        "\n",
-        "In the next cell, you write the contents of the training script task.py. We won't go into detail, it's just there for you to browse. In summary:\n",
+        "In the next cell, provide the contents of the training script task.py file. You don't have to go into detail. It's just there for you to browse. In summary:\n",
        "\n",
        "- Get the directory where to save the model artifacts from the command line (`--model_dir`), and if not specified, then from the environment variable `AIP_MODEL_DIR`.\n",
        "- Loads CIFAR10 dataset from TF Datasets (tfds).\n",
@@ -769,7 +712,10 @@
        "### Create and run custom training job\n",
        "\n",
        "\n",
-        "To train a custom model, you perform two steps: 1) create a custom training job, and 2) run the job.\n",
+        "To train a custom model, you perform two steps: \n",
+        "\n",
+        "1. create a custom training job\n",
+        "2. run the job.\n",
        "\n",
        "#### Create custom training job\n",
        "\n",
@@ -810,7 +756,7 @@
        "Now define the command-line arguments for your custom training container:\n",
        "\n",
        "- `args`: The command-line arguments to pass to the executable that is set as the entry point into the container.\n",
-        "  - `--model-dir` : For our demonstrations, we use this command-line argument to specify where to store the model artifacts.\n",
+        "  - `--model-dir` : For these demonstrations, use this command-line argument to specify where to store the model artifacts.\n",
        "      - direct: You pass the Cloud Storage location as a command line argument to your training script (set variable `DIRECT = True`), or\n",
        "      - indirect: The service passes the Cloud Storage location as the environment variable `AIP_MODEL_DIR` to your training script (set variable `DIRECT = False`). In this case, you tell the service the model artifact location in the job specification.\n",
        "  - `\"--epochs=\" + EPOCHS`: The number of epochs for training.\n",
@@ -1150,7 +1096,7 @@
        "\n",
        "#### XRAI\n",
        "\n",
-        "Based on the integrated gradients method, XRAI assesses overlapping regions of the image to create a saliency map, which highlights relevant regions of the image rather than pixels.\n",
+        "Based on the integrated gradients method, XRAI assesses overlapping locations of the image to create a saliency map, which highlights relevant locations of the image rather than pixels.\n",
        "\n",
        "Use Cases:\n",
        "\n",
@@ -1197,7 +1143,7 @@
        "\n",
        "    y = f(x)\n",
        "\n",
-        "Consider the following formulae, where the outputs are `y` and `z`. Since we can only do attribution for one scalar value, we have to pick whether we want to explain the output `y` or `z`. Assume in this example the model is object detection and y and z are the bounding box and the object classification. You would want to pick which of the two outputs to explain.\n",
+        "Consider the following formulae, where the outputs are `y` and `z`. Since you can only do attribution for one scalar value, you have to pick whether you want to explain the output `y` or `z`. Assume in this example the model is object detection and y and z are the bounding box and the object classification. You would want to pick which of the two outputs to explain.\n",
        "\n",
        "    y, z = f(x)\n",
        "\n",
@@ -1214,7 +1160,7 @@
        " -  [layer]: The output layer to explain. In a single task model, like a tabular regressor, it is the last (topmost) layer in the model.\n",
        "</blockquote>\n",
        "\n",
-        "- `inputs`: The features for attribution -- how they contributed to the output. Consider the following formulae, where `a` and `b` are the features. We have to pick which features to explain how the contributed. Assume that this model is deployed for A/B testing, where `a` are the data_items for the prediction and `b` identifies whether the model instance is A or B. You would want to pick `a` (or some subset of) for the features, and not `b` since it does not contribute to the prediction.\n",
+        "- `inputs`: The features for attribution -- how they contributed to the output. Consider the following formulae, where `a` and `b` are the features. You have to pick which features to explain how the contributed. Assume that this model is deployed for A/B testing, where `a` are the data_items for the prediction and `b` identifies whether the model instance is A or B. You would want to pick `a` (or some subset of) for the features, and not `b` since it does not contribute to the prediction.\n",
        "\n",
        "    y = f(a,b)\n",
        "\n",
@@ -1567,14 +1513,26 @@
      },
      "outputs": [],
      "source": [
+        "delete_model = True\n",
+        "delete_batch_predict_job = True\n",
        "delete_bucket = False\n",
+        "delete_application_directory = False\n",
        "\n",
-        "model.delete()\n",
-        "batch_predict_job.delete()\n",
+        "try:\n",
+        "    model.delete()\n",
+        "except Exception as e:\n",
+        "    print(e)\n",
        "\n",
+        "try:\n",
+        "    batch_predict_job.delete()\n",
+        "except Exception as e:\n",
+        "    print(e)\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
-        "    ! gsutil rm -r $BUCKET_URI"
+        "if delete_bucket:\n",
+        "    ! gsutil rm -r $BUCKET_URI\n",
+        "\n",
+        "if delete_application_directory:\n",
+        "    ! rm -rf custom custom.tar.gz temp1.jpg temp2.jpg"
      ]
    }
  ],
@@ -32,25 +32,27 @@
        "# Vertex AI SDK: Custom training image classification model for online prediction with explainability\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/explainable_ai/sdk_custom_image_classification_online_explain.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/explainable_ai/sdk_custom_image_classification_online_explain.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fexplainable_ai%2Fsdk_custom_image_classification_online_explain.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/explainable_ai/sdk_custom_image_classification_online_explain.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",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
-        "</table>\n",
-        "<br/><br/><br/>"
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/notebook_template.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
      ]
    },
    {
@@ -75,26 +77,26 @@
      "source": [
        "### Objective\n",
        "\n",
-        "In this tutorial, you learn how to use `Vertex AI Training and Explainable AI` to create a custom image classification model with explanations, and then you learn to use `Vertex AI Prediction` to make an online prediction request with explanations. You can alternatively create custom models using `gcloud` command-line tool or online using Cloud Console.\n",
+        "In this tutorial, you learn how to use Vertex AI training and Explainable AI to create a custom image classification model with explanations. Then, you learn how to use Vertex AI online prediction service to make online prediction requests with explanations. Alternatively, you can create custom models using `gcloud` command-line tool or online using Cloud Console.\n",
        "\n",
-        "This tutorial uses the following Google Cloud ML services:\n",
+        "This tutorial uses the following Vertex AI services:\n",
        "\n",
-        "- `Vertex AI Training`\n",
-        "- `Vertex AI Online Prediction`\n",
-        "- `Vertex Explainable AI`\n",
-        "- `Vertex AI Model` resource\n",
-        "- `Vertex AI Endpoint` resource\n",
+        "- Vertex AI Training\n",
+        "- Vertex AI online prediction\n",
+        "- Vertex Explainable AI\n",
+        "- Vertex AI Model resource\n",
+        "- Vertex AI Endpoint resource\n",
        "\n",
        "The steps performed include:\n",
        "\n",
-        "- Create a `Vertex AI` custom job for training a TensorFlow model.\n",
+        "- Create a Vertex AI custom job for training a TensorFlow model.\n",
        "- View the model evaluation for the trained model.\n",
        "- Set explanation parameters for when the model is deployed.\n",
-        "- Upload the trained model artifacts and explanations as a `Model` resource.\n",
-        "- Create a serving `Endpoint` resource.\n",
-        "- Deploy the `Model` resource to a serving `Endpoint` resource.\n",
+        "- Upload the trained model artifacts and explanations as a model resource.\n",
+        "- Create a serving endpoint resource.\n",
+        "- Deploy the model resource to a serving endpoint resource.\n",
        "- Make a prediction with explanation.\n",
-        "- Undeploy the `Model` resource."
+        "- Undeploy the model resource."
      ]
    },
    {
@@ -105,7 +107,7 @@
      "source": [
        "### Dataset\n",
        "\n",
-        "The dataset used for this tutorial is the [CIFAR10 dataset](https://www.tensorflow.org/datasets/catalog/cifar10) from [TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/overview). The version of the dataset you will use is built into TensorFlow. The trained model predicts which type of class an image is from ten classes: airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck."
+        "The dataset used for this tutorial is the [CIFAR10 dataset](https://www.tensorflow.org/datasets/catalog/cifar10) from [TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/overview). The version of the dataset that you use is built into TensorFlow. The trained model predicts which type of class an image is from ten classes: airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck."
      ]
    },
    {
@@ -128,15 +130,22 @@
        "to generate a cost estimate based on your projected usage."
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f0316df526f8"
+      },
+      "source": [
+        "## Get started"
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "install_aip:mbsdk"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the packages required for executing this notebook."
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
@@ -149,10 +158,11 @@
      "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 \n",
-        "! pip3 install --upgrade --quiet opencv-python \n",
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform\n",
+        "! pip3 install --upgrade --quiet google-cloud-storage\n",
+        "! pip3 install --upgrade --quiet opencv-python\n",
+        "! pip3 install --upgrade --quiet matplotlib\n",
+        "! pip3 install --quiet tensorflow==2.15.1\n",
        "\n",
        "if os.getenv(\"IS_TESTING\"):\n",
        "    ! apt-get update && apt-get install -y python3-opencv-headless\n",
@@ -166,7 +176,9 @@
        "id": "restart"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel"
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
@@ -177,11 +189,53 @@
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ee775571c2b5"
+      },
+      "source": [
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "92e68cfc3a90"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
+        "\n",
+        "Authenticate your environment on Google Colab.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "46604f70e831"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
      ]
    },
    {
@@ -190,14 +244,9 @@
        "id": "yfEglUHQk9S3"
      },
      "source": [
-        "## Before you begin\n",
+        "### Set Google Cloud project information\n",
        "\n",
-        "### Set your project ID\n",
-        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "* Run `gcloud config list`.\n",
-        "* Run `gcloud projects list`.\n",
-        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
@@ -209,89 +258,7 @@
      "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": "gcp_authenticate"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
-        "\n",
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated.\n",
-        "\n",
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ce6043da7b33"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0367eac06a10"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "21ad4dbb4a61"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c13224697bfb"
-      },
-      "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."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -322,7 +289,7 @@
        "id": "create_bucket"
      },
      "source": [
-        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+        "**If your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
@@ -333,7 +300,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION $BUCKET_URI"
+        "! gsutil mb -l $LOCATION -p $PROJECT_ID $BUCKET_URI"
      ]
    },
    {
@@ -342,10 +309,7 @@
        "id": "setup_vars"
      },
      "source": [
-        "### Set up variables\n",
-        "\n",
-        "Next, set up some variables used throughout the tutorial.\n",
-        "### Import libraries and define constants"
+        "### Import the required libraries"
      ]
    },
    {
@@ -356,9 +320,18 @@
      },
      "outputs": [],
      "source": [
+        "import base64\n",
+        "import io\n",
        "import os\n",
+        "from io import BytesIO\n",
        "\n",
-        "import google.cloud.aiplatform as aip"
+        "import cv2\n",
+        "import matplotlib.image as mpimg\n",
+        "import matplotlib.pyplot as plt\n",
+        "import numpy as np\n",
+        "import tensorflow as tf\n",
+        "from google.cloud import aiplatform\n",
+        "from tensorflow.keras.datasets import cifar10"
      ]
    },
    {
@@ -369,7 +342,9 @@
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
-        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
+        "To get started using Vertex AI, you must [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
+        "\n",
+        "Initialize the Vertex AI SDK for Python for your project using the location and the Cloud Storage bucket."
      ]
    },
    {
@@ -380,7 +355,7 @@
      },
      "outputs": [],
      "source": [
-        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -393,16 +368,16 @@
        "\n",
        "You can set hardware accelerators for training and prediction.\n",
        "\n",
-        "Set the variables `TRAIN_GPU/TRAIN_NGPU` and `DEPLOY_GPU/DEPLOY_NGPU` to use a container image supporting a GPU and the number of GPUs allocated to the virtual machine (VM) instance. For example, to use a GPU container image with 4 Nvidia Telsa K80 GPUs allocated to each VM, you would specify:\n",
+        "Set the variables `TRAIN_GPU/TRAIN_NGPU` and `DEPLOY_GPU/DEPLOY_NGPU` to use a container image supporting a GPU and the number of GPUs allocated to the virtual machine (VM) instance. For example, to use a GPU container image with 4 Nvidia Telsa T4 GPUs allocated to each VM, you would specify:\n",
        "\n",
-        "    (aip.gapic.AcceleratorType.NVIDIA_TESLA_K80, 4)\n",
+        "    (aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_T4, 4)\n",
        "\n",
        "\n",
        "Otherwise specify `(None, None)` to use a container image to run on a CPU.\n",
        "\n",
-        "Learn more [here](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators) hardware accelerator support for your region\n",
+        "Learn more about [hardware accelerator support for your region](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators).\n",
        "\n",
-        "*Note*: TF releases before 2.3 for GPU support will fail to load the custom model in this tutorial. It is a known issue and fixed in TF 2.3 -- which is caused by static graph ops that are generated in the serving function. If you encounter this issue on your own custom models, use a container image for TF 2.3 with GPU support."
+        "**Note**: TF releases before 2.3 for GPU support fail to load the custom model in this tutorial. It's a known issue and fixed in TF 2.3 -- which is caused by static graph ops that are generated in the serving function. If you encounter this issue on your own custom models, use a container image for TF 2.3 with GPU support."
      ]
    },
    {
@@ -415,7 +390,7 @@
      "source": [
        "if os.getenv(\"IS_TESTING_TRAIN_GPU\"):\n",
        "    TRAIN_GPU, TRAIN_NGPU = (\n",
-        "        aip.gapic.AcceleratorType.NVIDIA_TESLA_K80,\n",
+        "        aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_T4,\n",
        "        int(os.getenv(\"IS_TESTING_TRAIN_GPU\")),\n",
        "    )\n",
        "else:\n",
@@ -423,7 +398,7 @@
        "\n",
        "if os.getenv(\"IS_TESTING_DEPLOY_GPU\"):\n",
        "    DEPLOY_GPU, DEPLOY_NGPU = (\n",
-        "        aip.gapic.AcceleratorType.NVIDIA_TESLA_K80,\n",
+        "        aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_T4,\n",
        "        int(os.getenv(\"IS_TESTING_DEPLOY_GPU\")),\n",
        "    )\n",
        "else:\n",
@@ -480,10 +455,10 @@
        "        DEPLOY_VERSION = \"tf-cpu.{}\".format(TF)\n",
        "\n",
        "TRAIN_IMAGE = \"{}-docker.pkg.dev/vertex-ai/training/{}:latest\".format(\n",
-        "    REGION.split(\"-\")[0], TRAIN_VERSION\n",
+        "    LOCATION.split(\"-\")[0], TRAIN_VERSION\n",
        ")\n",
        "DEPLOY_IMAGE = \"{}-docker.pkg.dev/vertex-ai/prediction/{}:latest\".format(\n",
-        "    REGION.split(\"-\")[0], DEPLOY_VERSION\n",
+        "    LOCATION.split(\"-\")[0], DEPLOY_VERSION\n",
        ")\n",
        "\n",
        "print(\"Training:\", TRAIN_IMAGE, TRAIN_GPU, TRAIN_NGPU)\n",
@@ -500,19 +475,19 @@
        "\n",
        "Next, set the machine type to use for training and prediction.\n",
        "\n",
-        "- Set the variables `TRAIN_COMPUTE` and `DEPLOY_COMPUTE` to configure  the compute resources for the VMs you will use for for training and prediction.\n",
+        "- Set the variables `TRAIN_COMPUTE` and `DEPLOY_COMPUTE` to configure  the compute resources for the VMs you use for for training and prediction.\n",
        " - `machine type`\n",
        "     - `n1-standard`: 3.75GB of memory per vCPU.\n",
        "     - `n1-highmem`: 6.5GB of memory per vCPU\n",
        "     - `n1-highcpu`: 0.9 GB of memory per vCPU\n",
        " - `vCPUs`: number of \\[2, 4, 8, 16, 32, 64, 96 \\]\n",
        "\n",
-        "*Note: The following is not supported for training:*\n",
+        "**Note**: The following isn't supported for training:\n",
        "\n",
        " - `standard`: 2 vCPUs\n",
        " - `highcpu`: 2, 4 and 8 vCPUs\n",
        "\n",
-        "*Note: You may also use n2 and e2 machine types for training and deployment, but they do not support GPUs*."
+        "**Note**: You may also use n2 and e2 machine types for training and deployment, but they don't support GPUs."
      ]
    },
    {
@@ -542,28 +517,21 @@
        "print(\"Deploy machine type\", DEPLOY_COMPUTE)"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "tutorial_start:custom"
-      },
-      "source": [
-        "# Tutorial\n",
-        "\n",
-        "Now you are ready to start creating your own custom model and training for CIFAR10."
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "examine_training_package"
      },
      "source": [
-        "### Examine the training package\n",
+        "## Examine the training package\n",
        "\n",
-        "#### Package layout\n",
+        "Now you're ready to start creating your own custom model and training for CIFAR10.\n",
        "\n",
-        "Before you start the training, you will look at how a Python package is assembled for a custom training job. When unarchived, the package contains the following directory/file layout.\n",
+        "Before you start the training, examine the training application package below.\n",
+        "\n",
+        "### Package layout\n",
+        "\n",
+        "Before you start training, look at how a Python package is assembled for a custom training job. When unarchived, the package contains the following directory/file layout.\n",
        "\n",
        "- PKG-INFO\n",
        "- README.md\n",
@@ -575,11 +543,13 @@
        "\n",
        "The files `setup.cfg` and `setup.py` are the instructions for installing the package into the operating environment of the Docker image.\n",
        "\n",
-        "The file `trainer/task.py` is the Python script for executing the custom training job. *Note*, when we referred to it in the worker pool specification, we replace the directory slash with a dot (`trainer.task`) and dropped the file suffix (`.py`).\n",
+        "The file `trainer/task.py` is the Python script for executing the custom training job. \n",
        "\n",
-        "#### Package Assembly\n",
+        "**Note**: When referring to the worker pool specification, the directory slash is replaced with a dot (`trainer.task`) and the file suffix (`.py`) is dropped.\n",
        "\n",
-        "In the following cells, you will assemble the training package."
+        "### Package Assembly\n",
+        "\n",
+        "In the following cells, you assemble the training package."
      ]
    },
    {
@@ -617,9 +587,11 @@
        "id": "taskpy_contents:cifar10"
      },
      "source": [
-        "#### Task.py contents\n",
+        "### Task.py contents\n",
        "\n",
-        "In the next cell, you write the contents of the training script task.py. We won't go into detail, it's just there for you to browse. In summary:\n",
+        "In the next cell, you write the contents of the training script *task.py*. \n",
+        "\n",
+        "In summary, *task.py* script does the following:\n",
        "\n",
        "- Get the directory where to save the model artifacts from the command line (`--model_dir`), and if not specified, then from the environment variable `AIP_MODEL_DIR`.\n",
        "- Loads CIFAR10 dataset from TF Datasets (tfds).\n",
@@ -745,7 +717,7 @@
        "id": "tarball_training_script"
      },
      "source": [
-        "#### Store training script on your Cloud Storage bucket\n",
+        "## Store training script on your Cloud Storage bucket\n",
        "\n",
        "Next, you package the training folder into a compressed tar ball, and then store it in your Cloud Storage bucket."
      ]
@@ -770,12 +742,12 @@
        "id": "create_custom_training_job:mbsdk,no_model"
      },
      "source": [
-        "### Create and run custom training job\n",
+        "## Create and run custom training job\n",
        "\n",
        "\n",
        "To train a custom model, you perform two steps: 1) create a custom training job, and 2) run the job.\n",
        "\n",
-        "#### Create custom training job\n",
+        "### Create custom training job\n",
        "\n",
        "A custom training job is created with the `CustomTrainingJob` class, with the following parameters:\n",
        "\n",
@@ -793,7 +765,7 @@
      },
      "outputs": [],
      "source": [
-        "job = aip.CustomTrainingJob(\n",
+        "job = aiplatform.CustomTrainingJob(\n",
        "    display_name=\"cifar10\",\n",
        "    script_path=\"custom/trainer/task.py\",\n",
        "    container_uri=TRAIN_IMAGE,\n",
@@ -809,16 +781,16 @@
        "id": "prepare_custom_cmdargs"
      },
      "source": [
-        "### Prepare your command-line arguments\n",
+        "### Prepare your training parameters\n",
        "\n",
        "Now define the command-line arguments for your custom training container:\n",
        "\n",
-        "- `args`: The command-line arguments to pass to the executable that is set as the entry point into the container.\n",
-        "  - `--model-dir` : For our demonstrations, we use this command-line argument to specify where to store the model artifacts.\n",
-        "      - direct: You pass the Cloud Storage location as a command line argument to your training script (set variable `DIRECT = True`), or\n",
-        "      - indirect: The service passes the Cloud Storage location as the environment variable `AIP_MODEL_DIR` to your training script (set variable `DIRECT = False`). In this case, you tell the service the model artifact location in the job specification.\n",
-        "  - `\"--epochs=\" + EPOCHS`: The number of epochs for training.\n",
-        "  - `\"--steps=\" + STEPS`: The number of steps per epoch."
+        "- `args`: The command-line arguments to pass to the executable that's set as the entry point into the container.\n",
+        "  - `--model-dir`: Command-line argument to specify where to store the model artifacts. You can use either of the following methods to specify the storage location for artifacts.\n",
+        "      - **method-1**(set `DIRECT` to `True`): You pass the Cloud Storage location as a command line argument to your training script.\n",
+        "      - **method-2**(set `DIRECT` to `False`): The service passes the Cloud Storage location as the environment variable `AIP_MODEL_DIR` to your training script. In this case, you tell the service the model artifact location in the job specification.\n",
+        "  - `--epochs`: The number of epochs for training.\n",
+        "  - `--steps`: The number of steps per epoch."
      ]
    },
    {
@@ -834,7 +806,7 @@
        "EPOCHS = 20\n",
        "STEPS = 100\n",
        "\n",
-        "DIRECT = True\n",
+        "DIRECT = True  # Set False to use AIP_MODEL_DIR\n",
        "if DIRECT:\n",
        "    CMDARGS = [\n",
        "        \"--model-dir=\" + MODEL_DIR,\n",
@@ -854,17 +826,17 @@
        "id": "run_custom_job:mbsdk,no_model"
      },
      "source": [
-        "#### Run the custom training job\n",
+        "### Run the custom training job\n",
        "\n",
-        "Next, you run the custom job to start the training job by invoking the method `run`, with the following parameters:\n",
+        "Next, you run the custom job to start the training job by invoking the `run()` method, with the following parameters:\n",
        "\n",
        "- `args`: The command-line arguments to pass to the training script.\n",
        "- `replica_count`: The number of compute instances for training (replica_count = 1 is single node training).\n",
        "- `machine_type`: The machine type for the compute instances.\n",
        "- `accelerator_type`: The hardware accelerator type.\n",
        "- `accelerator_count`: The number of accelerators to attach to a worker replica.\n",
-        "- `base_output_dir`: The Cloud Storage location to write the model artifacts to.\n",
-        "- `sync`: Whether to block until completion of the job."
+        "- `base_output_dir`: The Cloud Storage location to store the model artifacts.\n",
+        "- `sync`: Set **True** to wait until completion of the job."
      ]
    },
    {
@@ -918,8 +890,6 @@
      },
      "outputs": [],
      "source": [
-        "import tensorflow as tf\n",
-        "\n",
        "local_model = tf.keras.models.load_model(MODEL_DIR)"
      ]
    },
@@ -931,21 +901,19 @@
      "source": [
        "## Evaluate the model\n",
        "\n",
-        "Now find out how good the model is.\n",
+        "Now, find out how good the model performs. Prepare some test data for evaluation and run the `evaluate()` method.\n",
        "\n",
        "### Load evaluation data\n",
        "\n",
-        "You will load the CIFAR10 test (holdout) data from `tf.keras.datasets`, using the method `load_data()`. This returns the dataset as a tuple of two elements. The first element is the training data and the second is the test data. Each element is also a tuple of two elements: the image data, and the corresponding labels.\n",
+        "Load the CIFAR10 test (holdout) data from `tf.keras.datasets`, using the `load_data()` method. This returns the dataset as a tuple of two elements. The first element is the training data and the second one is the test data. Each element is also a tuple of two elements: the image data, and the corresponding labels.\n",
        "\n",
-        "You don't need the training data, and hence why we loaded it as `(_, _)`.\n",
+        "For this step, you don't need the training data. Hence, you skip it by loading as `(_, _)`.\n",
        "\n",
        "Before you can run the data through evaluation, you need to preprocess it:\n",
        "\n",
-        "`x_test`:\n",
-        "1. Normalize (rescale) the pixel data by dividing each pixel by 255. This replaces each single byte integer pixel with a 32-bit floating point number between 0 and 1.\n",
+        "`x_test`: Normalize (rescale) the pixel data by dividing each pixel by 255. This replaces each single byte integer pixel with a 32-bit floating point number between 0 and 1.\n",
        "\n",
-        "`y_test`:<br/>\n",
-        "2. The labels are currently scalar (sparse). If you look back at the `compile()` step in the `trainer/task.py` script, you will find that it was compiled for sparse labels. So we don't need to do anything more."
+        "`y_test`: The labels are currently scalar (sparse). If you look back at the `compile()` step in the `trainer/task.py` script, you find that it was compiled for sparse labels. So you don't need to do anything more."
      ]
    },
    {
@@ -956,9 +924,6 @@
      },
      "outputs": [],
      "source": [
-        "import numpy as np\n",
-        "from tensorflow.keras.datasets import cifar10\n",
-        "\n",
        "(_, _), (x_test, y_test) = cifar10.load_data()\n",
        "x_test = (x_test / 255.0).astype(np.float32)\n",
        "\n",
@@ -995,7 +960,7 @@
      "source": [
        "### Serving function for image data\n",
        "\n",
-        "To pass images to the prediction service, you encode the compressed (e.g., JPEG) image bytes into base 64 -- which makes the content safe from modification while transmitting binary data over the network. Since this deployed model expects input data as raw (uncompressed) bytes, you need to ensure that the base 64 encoded data gets converted back to raw bytes before it is passed as input to the deployed model.\n",
+        "To pass images to the prediction service, you encode the compressed (e.g., JPEG) image bytes into base 64 -- which makes the content safe from modification while transmitting binary data over the network. Since this deployed model expects input data as raw (uncompressed) bytes, you need to ensure that the base 64 encoded data gets converted back to raw bytes before it's passed as input to the deployed model.\n",
        "\n",
        "To resolve this, define a serving function (`serving_fn`) and attach it to the model as a preprocessing step. Add a `@tf.function` decorator so the serving function is fused to the underlying model (instead of upstream on a CPU).\n",
        "\n",
@@ -1079,9 +1044,9 @@
        "\n",
        "You can get the signatures of your model's input and output layers by reloading the model into memory, and querying it for the signatures corresponding to each layer.\n",
        "\n",
-        "When making a prediction request, you need to route the request to the serving function instead of the model, so you need to know the input layer name of the serving function -- which you will use later when you make a prediction request.\n",
+        "When making a prediction request, you need to route the request to the serving function instead of the model. So, you need to know the input layer's name for the serving function which you use later when you make a prediction request.\n",
        "\n",
-        "You also need to know the name of the serving function's input and output layer for constructing the explanation metadata -- which is discussed subsequently."
+        "You also need to know the name of the serving function's input and output layer for constructing the explanation metadata which is described subsequently."
      ]
    },
    {
@@ -1115,17 +1080,17 @@
      "source": [
        "### Explanation Specification\n",
        "\n",
-        "To get explanations when doing a prediction, you must enable the explanation capability and set corresponding settings when you upload your custom model to an Vertex `Model` resource. These settings are referred to as the explanation metadata, which consists of:\n",
+        "To get explanations along with the predictions, you must enable the explanations feature and set corresponding settings when you upload your custom model to a Vertex AI model resource. These settings are referred to as the explanation metadata, which consists of:\n",
        "\n",
        "- `parameters`: This is the specification for the explainability algorithm to use for explanations on your model. You can choose between:\n",
-        "  - Shapley - *Note*, not recommended for image data -- can be very long running\n",
+        "  - Shapley(Note: not recommended for image data since it can involve a long-running operation)\n",
        "  - XRAI\n",
        "  - Integrated Gradients\n",
        "- `metadata`: This is the specification for how the algoithm is applied on your custom model.\n",
        "\n",
        "#### Explanation Parameters\n",
        "\n",
-        "Let's first dive deeper into the settings for the explainability algorithm.\n",
+        "Take a look at a detailed overview of the settings for the explainability algorithm.\n",
        "\n",
        "#### Shapley\n",
        "\n",
@@ -1136,7 +1101,7 @@
        "\n",
        "Parameters:\n",
        "\n",
-        "- `path_count`: This is the number of paths over the features that will be processed by the algorithm. An exact approximation of the Shapley values requires M! paths, where M is the number of features. For the CIFAR10 dataset, this would be 784 (28*28).\n",
+        "- `path_count`: This is the number of paths over the features processed by the algorithm. An exact approximation of the Shapley values requires M! paths, where M is the number of features. For the CIFAR10 dataset, this would be 784 (28*28).\n",
        "\n",
        "For any non-trival number of features, this is too compute expensive. You can reduce the number of paths over the features to M * `path_count`.\n",
        "\n",
@@ -1164,7 +1129,7 @@
        "\n",
        "- `step_count`: This is the number of steps to approximate the remaining sum. The more steps, the more accurate the integral approximation. The general rule of thumb is 50 steps, but as you increase so does the compute time.\n",
        "\n",
-        "In the next code cell, set the variable `XAI` to which explainabilty algorithm you will use on your custom model."
+        "In the next code cell, set the variable `XAI` to which explainabilty algorithm you use on your custom model."
      ]
    },
    {
@@ -1184,7 +1149,7 @@
        "elif XAI == \"xrai\":\n",
        "    PARAMETERS = {\"xrai_attribution\": {\"step_count\": 50}}\n",
        "\n",
-        "parameters = aip.explain.ExplanationParameters(PARAMETERS)"
+        "parameters = aiplatform.explain.ExplanationParameters(PARAMETERS)"
      ]
    },
    {
@@ -1195,9 +1160,9 @@
      "source": [
        "#### Explanation Metadata\n",
        "\n",
-        "Let's first dive deeper into the explanation metadata, which consists of:\n",
+        "Now, let's dive deeper into the explanation metadata, which consists of:\n",
        "\n",
-        "- `outputs`: A scalar value in the output to attribute -- what to explain. For example, in a probability output \\[0.1, 0.2, 0.7\\] for classification, one wants an explanation for 0.7. Consider the following formulae, where the output is `y` and that is what we want to explain.\n",
+        "- `outputs`: A scalar value in the output to attribute -- what to explain. For example, in a probability output \\[0.1, 0.2, 0.7\\] for classification, one wants an explanation for 0.7. Consider the following formula, where the output is `y` and that is what we want to explain.\n",
        "\n",
        "    y = f(x)\n",
        "\n",
@@ -1215,7 +1180,7 @@
        "<blockquote>\n",
        " -  [your_display_name]: A human readable name you assign to the output to explain. A common example is \"probability\".<br/>\n",
        " -  \"output_tensor_name\": The key/value field to identify the output layer to explain. <br/>\n",
-        " -  [layer]: The output layer to explain. In a single task model, like a tabular regressor, it is the last (topmost) layer in the model.\n",
+        " -  [layer]: The output layer to explain. In a single task model, like a tabular regressor, it's the last (topmost) layer in the model.\n",
        "</blockquote>\n",
        "\n",
        "- `inputs`: The features for attribution -- how they contributed to the output. Consider the following formulae, where `a` and `b` are the features. We have to pick which features to explain how the contributed. Assume that this model is deployed for A/B testing, where `a` are the data_items for the prediction and `b` identifies whether the model instance is A or B. You would want to pick `a` (or some subset of) for the features, and not `b` since it does not contribute to the prediction.\n",
@@ -1232,7 +1197,7 @@
        "<blockquote>\n",
        " -  [your_display_name]: A human readable name you assign to the input to explain. A common example is \"features\".<br/>\n",
        " -  \"input_tensor_name\": The key/value field to identify the input layer for the feature attribution. <br/>\n",
-        " -  [layer]: The input layer for feature attribution. In a single input tensor model, it is the first (bottom-most) layer in the model.\n",
+        " -  [layer]: The input layer for feature attribution. In a single input tensor model, it's the first (bottom-most) layer in the model.\n",
        "</blockquote>\n",
        "\n",
        "Since the inputs to the model are tabular, you can specify the following two additional fields as reporting/visualization aids:\n",
@@ -1263,10 +1228,10 @@
        "\n",
        "OUTPUT_METADATA = {\"output_tensor_name\": serving_output}\n",
        "\n",
-        "input_metadata = aip.explain.ExplanationMetadata.InputMetadata(INPUT_METADATA)\n",
-        "output_metadata = aip.explain.ExplanationMetadata.OutputMetadata(OUTPUT_METADATA)\n",
+        "input_metadata = aiplatform.explain.ExplanationMetadata.InputMetadata(INPUT_METADATA)\n",
+        "output_metadata = aiplatform.explain.ExplanationMetadata.OutputMetadata(OUTPUT_METADATA)\n",
        "\n",
-        "metadata = aip.explain.ExplanationMetadata(\n",
+        "metadata = aiplatform.explain.ExplanationMetadata(\n",
        "    inputs={\"image\": input_metadata}, outputs={\"class\": output_metadata}\n",
        ")"
      ]
@@ -1279,14 +1244,14 @@
      "source": [
        "## Upload the model\n",
        "\n",
-        "Next, upload your model to a `Model` resource using `Model.upload()` method, with the following parameters:\n",
+        "Next, upload your model to Vertex AI Model Registry using `Model.upload()` method, with the following parameters:\n",
        "\n",
-        "- `display_name`: The human readable name for the `Model` resource.\n",
+        "- `display_name`: The human readable name for the model resource.\n",
        "- `artifact`: The Cloud Storage location of the trained model artifacts.\n",
        "- `serving_container_image_uri`: The serving container image.\n",
        "- `sync`: Whether to execute the upload asynchronously or synchronously.\n",
-        "- `explanation_parameters`: Parameters to configure explaining for `Model`'s predictions.\n",
-        "- `explanation_metadata`: Metadata describing the `Model`'s input and output for explanation.\n",
+        "- `explanation_parameters`: Parameters to configure explaining for model's predictions.\n",
+        "- `explanation_metadata`: Metadata describing the model's input and output for explanation.\n",
        "\n",
        "If the `upload()` method is run asynchronously, you can subsequently block until completion with the `wait()` method."
      ]
@@ -1299,7 +1264,7 @@
      },
      "outputs": [],
      "source": [
-        "model = aip.Model.upload(\n",
+        "model = aiplatform.Model.upload(\n",
        "    display_name=\"cifar10\",\n",
        "    artifact_uri=MODEL_DIR,\n",
        "    serving_container_image_uri=DEPLOY_IMAGE,\n",
@@ -1319,12 +1284,12 @@
      "source": [
        "## Deploy the model\n",
        "\n",
-        "Next, deploy your model for online prediction. To deploy the model, you invoke the `deploy` method, with the following parameters:\n",
+        "Next, deploy your model for online predictions. To deploy the model, you invoke the `deploy()` method, with the following parameters:\n",
        "\n",
        "- `deployed_model_display_name`: A human readable name for the deployed model.\n",
        "- `traffic_split`: Percent of traffic at the endpoint that goes to this model, which is specified as a dictionary of one or more key/value pairs.\n",
        "If only one model, then specify as { \"0\": 100 }, where \"0\" refers to this model being uploaded and 100 means 100% of the traffic.\n",
-        "If there are existing models on the endpoint, for which the traffic will be split, then use model_id to specify as { \"0\": percent, model_id: percent, ... }, where model_id is the model id of an existing model to the deployed endpoint. The percents must add up to 100.\n",
+        "If there are existing models on the endpoint, for which the traffic is split, then use model_id to specify as { \"0\": percent, model_id: percent, ... }, where model_id is the model id of an existing model to the deployed endpoint. The percents must add up to 100.\n",
        "- `machine_type`: The type of machine to use for training.\n",
        "- `accelerator_type`: The hardware accelerator type.\n",
        "- `accelerator_count`: The number of accelerators to attach to a worker replica.\n",
@@ -1377,7 +1342,7 @@
      "source": [
        "### Get test item\n",
        "\n",
-        "You will use an example out of the test (holdout) portion of the dataset as a test item."
+        "Use an example from the test (holdout) portion of the dataset as a test item."
      ]
    },
    {
@@ -1417,10 +1382,6 @@
      },
      "outputs": [],
      "source": [
-        "import base64\n",
-        "\n",
-        "import cv2\n",
-        "\n",
        "cv2.imwrite(\"tmp.jpg\", (test_image * 255).astype(np.uint8))\n",
        "\n",
        "bytes = tf.io.read_file(\"tmp.jpg\")\n",
@@ -1435,7 +1396,7 @@
      "source": [
        "### Make the prediction with explanation\n",
        "\n",
-        "Now that your `Model` resource is deployed to an `Endpoint` resource, one can do online explanations by sending prediction requests to the `Endpoint` resource.\n",
+        "Now that your model resource is deployed to an edpoint resource, you can get online explanations by sending prediction requests to your endpoint.\n",
        "\n",
        "#### Request\n",
        "\n",
@@ -1443,16 +1404,16 @@
        "\n",
        "    [{serving_input: {'b64': bytes}]\n",
        "\n",
-        "Since the explain() method can take multiple items (instances), send your single test item as a list of one test item.\n",
+        "Since the `explain()` method can take multiple items (instances), send your single test item as a list of one test item.\n",
        "\n",
        "#### Response\n",
        "\n",
-        "The response from the explain() call is a Python dictionary with the following entries:\n",
+        "The response from the `explain()` call is a Python dictionary with the following entries:\n",
        "\n",
        "- `ids`: The internal assigned unique identifiers for each prediction request.\n",
        "- `predictions`: The prediction per instance.\n",
-        "- `deployed_model_id`: The Vertex AI identifier for the deployed `Model` resource which did the predictions.\n",
-        "- `explanations`: The feature attributions"
+        "- `deployed_model_id`: The Vertex AI identifier for the deployed model resource.\n",
+        "- `explanations`: The feature attributions returned by the Explainable AI."
      ]
    },
    {
@@ -1475,9 +1436,9 @@
        "id": "understanding_explanations:cifar10"
      },
      "source": [
-        "### Understanding the explanations response\n",
+        "### Visualize the predictions\n",
        "\n",
-        "Preview the images and their predicted classes without the explanations. Why did the model predict these classes?"
+        "Preview the images and their predicted classes without the explanations."
      ]
    },
    {
@@ -1488,11 +1449,6 @@
      },
      "outputs": [],
      "source": [
-        "from io import BytesIO\n",
-        "\n",
-        "import matplotlib.image as mpimg\n",
-        "import matplotlib.pyplot as plt\n",
-        "\n",
        "CLASSES = [\n",
        "    \"airplane\",\n",
        "    \"automobile\",\n",
@@ -1548,8 +1504,6 @@
      },
      "outputs": [],
      "source": [
-        "import io\n",
-        "\n",
        "for explanation in response.explanations:\n",
        "    attributions = dict(explanation.attributions[0].feature_attributions)\n",
        "    label_index = explanation.attributions[0].output_index[0]\n",
@@ -1571,7 +1525,7 @@
      "source": [
        "## Undeploy the model\n",
        "\n",
-        "When you are done doing predictions, you undeploy the model from the `Endpoint` resouce. This deprovisions all compute resources and ends billing for the deployed model."
+        "When you are done with the model predictions, undeploy the model from the endpoint resouce. This deprovisions all compute resources and ends billing for the deployed model."
      ]
    },
    {
@@ -1591,7 +1545,7 @@
        "id": "cleanup:mbsdk"
      },
      "source": [
-        "# Cleaning up\n",
+        "## 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",
@@ -1607,13 +1561,24 @@
      },
      "outputs": [],
      "source": [
-        "delete_bucket = False\n",
+        "# Delete the training job\n",
+        "job.delete()\n",
        "\n",
+        "# Delete the endpoint\n",
        "endpoint.delete()\n",
+        "\n",
+        "# Delete the model\n",
        "model.delete()\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
-        "    ! gsutil rm -r $BUCKET_URI"
+        "# Delete the Cloud Storage bucket\n",
+        "delete_bucket = False  # Set True for deletion\n",
+        "if delete_bucket:\n",
+        "    ! gsutil rm -r $BUCKET_URI\n",
+        "\n",
+        "# Delete locally generated files\n",
+        "! rm -rf custom/\n",
+        "! rm custom.tar.gz\n",
+        "! rm tmp.jpg"
      ]
    }
  ],
--- a/Show More
+++ b/Show More