Update run_linter.sh

Fix spelling error `notebooked` → `notebooks`
Update vLLM version in Llama 3.1 and Guard deployment notebooks. (#3335 )
2026-09-27 15:42:05 +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
177 changed files with 31950 additions and 22695 deletions
@@ -5,6 +5,6 @@ nbconvert
 black==24.4.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')
@@ -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
@@ -115,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
@@ -123,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
@@ -134,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",
@@ -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,51 +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",
-        "import json\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",
@@ -174,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",
@@ -204,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",
@@ -221,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",
@@ -258,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",
@@ -273,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",
@@ -282,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",
@@ -323,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",
@@ -346,105 +390,10 @@
        "        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 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\"  # noqa: F841\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",
-        "        )"
+        "    return model, endpoint"
      ]
    },
    {
@@ -582,7 +531,7 @@
        "\n",
        "    replica_count = 1\n",
        "\n",
-        "    check_quota(\n",
+        "    common_util.check_quota(\n",
        "        project_id=PROJECT_ID,\n",
        "        region=REGION,\n",
        "        accelerator_type=ACCELERATOR_TYPE,\n",
@@ -591,7 +540,7 @@
        "    )\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",
@@ -600,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",
@@ -666,7 +615,7 @@
        "    accelerator_type = \"NVIDIA_L4\"\n",
        "    accelerator_count = 1\n",
        "\n",
-        "    check_quota(\n",
+        "    common_util.check_quota(\n",
        "        project_id=PROJECT_ID,\n",
        "        region=REGION,\n",
        "        accelerator_type=accelerator_type,\n",
@@ -678,7 +627,7 @@
        "    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",
@@ -700,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",
@@ -905,7 +856,7 @@
        "\n",
        "    replica_count = 1\n",
        "\n",
-        "    check_quota(\n",
+        "    common_util.check_quota(\n",
        "        project_id=PROJECT_ID,\n",
        "        region=REGION,\n",
        "        accelerator_type=accelerator_type,\n",
@@ -914,7 +865,7 @@
        "    )\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",
@@ -1004,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",
@@ -1065,7 +1016,7 @@
        "        # Note: 1 TPU-V5e chip has only 1 core.\n",
        "        accelerator_count = 4\n",
        "\n",
-        "    check_quota(\n",
+        "    common_util.check_quota(\n",
        "        project_id=PROJECT_ID,\n",
        "        region=REGION,\n",
        "        accelerator_type=accelerator_type,\n",
@@ -1084,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",
@@ -1099,6 +1050,7 @@
    {
      "cell_type": "code",
      "execution_count": null,
+      "language": "python",
      "metadata": {
        "cellView": "form",
        "id": "lgJPw3e9yL5a"
@@ -1183,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",
@@ -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,104 +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\"  # noqa: F841\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",
-        "        )"
+        "    return model, endpoint"
      ]
    },
    {
@@ -480,7 +399,7 @@
        "\n",
        "replica_count = 1\n",
        "\n",
-        "check_quota(\n",
+        "common_util.check_quota(\n",
        "    project_id=PROJECT_ID,\n",
        "    region=REGION,\n",
        "    accelerator_type=ACCELERATOR_TYPE,\n",
@@ -489,7 +408,7 @@
        ")\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",
@@ -498,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",
@@ -560,7 +479,7 @@
        "accelerator_type = \"NVIDIA_L4\"\n",
        "accelerator_count = 1\n",
        "\n",
-        "check_quota(\n",
+        "common_util.check_quota(\n",
        "    project_id=PROJECT_ID,\n",
        "    region=REGION,\n",
        "    accelerator_type=accelerator_type,\n",
@@ -573,8 +492,8 @@
        "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",
@@ -583,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)"
+        ")"
      ]
    },
    {
@@ -641,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)"
@@ -686,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"
      ]
    }
  ],
@@ -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",
@@ -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
+}
@@ -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"
      ]
    },
    {
@@ -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"
      ]
@@ -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,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"
      ]
    }
@@ -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"
      ]
    }
  ],
@@ -32,21 +32,24 @@
        "# Vertex SDK: Custom training tabular regression 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_tabular_regression_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=\"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%2Fexplainable_ai%2Fsdk_custom_tabular_regression_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 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_tabular_regression_batch_explain.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/explainable_ai/sdk_custom_tabular_regression_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",
-        "    </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/explainable_ai/sdk_custom_tabular_regression_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://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
@@ -61,8 +64,7 @@
      "source": [
        "## Overview\n",
        "\n",
-        "\n",
-        "This tutorial demonstrates how to use the Vertex AI SDK to train and deploy a custom tabular regression model for batch prediction with explanation.\n",
+        "This tutorial demonstrates how to use the Vertex AI SDK for Python to train and deploy a custom tabular regression model for batch prediction with explanations.\n",
        "\n",
        "Learn more about [Vertex Explainable AI](https://cloud.google.com/vertex-ai/docs/explainable-ai/overview) and [Vertex AI Batch Prediction](https://cloud.google.com/vertex-ai/docs/tabular-data/classification-regression/get-batch-predictions)."
      ]
@@ -70,45 +72,45 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "objective:custom,training,batch_prediction,xai"
+        "id": "82a273206a67"
      },
      "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 Batch Prediction` to make a batch 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 the Vertex AI batch prediction service to make a batch prediction request 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 Batch Prediction`\n",
-        "- `Vertex Explainable AI`\n",
-        "- `Vertex AI Model` resource\n",
+        "- Vertex AI Training\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 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",
+        "- Set explanation parameters for the model.\n",
+        "- Upload the trained model artifacts as a model resource.\n",
        "- Make a batch prediction with explanations."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "dataset:custom,boston,lrg"
+        "id": "543b0e9b3b33"
      },
      "source": [
        "### Dataset\n",
        "\n",
-        "The dataset used for this tutorial is the [Boston Housing Prices dataset](https://www.cs.toronto.edu/~delve/data/boston/bostonDetail.html). The version of the dataset you will use in this tutorial is built into TensorFlow. The trained model predicts the median price of a house in units of 1K USD."
+        "The dataset used for this tutorial is the [Boston Housing Prices dataset](https://www.cs.toronto.edu/~delve/data/boston/bostonDetail.html). The version of the dataset you use in this tutorial is the one that's available from TensorFlow SDK. The trained model predicts the median price of a house in units of 1K USD."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "costs"
+        "id": "181d4dfbf917"
      },
      "source": [
        "### Costs\n",
@@ -128,12 +130,19 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "install_aip:mbsdk"
+        "id": "f0316df526f8"
      },
      "source": [
-        "## Installation\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"
      ]
    },
    {
@@ -144,49 +153,88 @@
      },
      "outputs": [],
      "source": [
-        "! pip3 install --upgrade --quiet google-cloud-aiplatform==1.36.0 \\\n",
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
        "                                 google-cloud-storage \\\n",
-        "                                 tensorflow==2.11"
+        "                                 tensorflow==2.15.1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "restart"
+        "id": "install_aip:mbsdk"
      },
      "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": "f82e28c631cc"
+      },
+      "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": "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)."
      ]
    },
    {
@@ -198,89 +246,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\"}"
      ]
    },
    {
@@ -311,7 +277,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."
      ]
    },
    {
@@ -322,7 +288,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
@@ -331,10 +297,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"
      ]
    },
    {
@@ -347,7 +310,8 @@
      "source": [
        "import os\n",
        "\n",
-        "import google.cloud.aiplatform as aip"
+        "import tensorflow as tf\n",
+        "from google.cloud import aiplatform"
      ]
    },
    {
@@ -358,7 +322,7 @@
      "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). "
      ]
    },
    {
@@ -369,7 +333,7 @@
      },
      "outputs": [],
      "source": [
-        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -382,16 +346,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 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 regions](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."
      ]
    },
    {
@@ -404,7 +368,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",
@@ -412,7 +376,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",
@@ -465,10 +429,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",
@@ -485,19 +449,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 used 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."
      ]
    },
    {
@@ -527,17 +491,6 @@
        "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 Boston Housing."
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -548,7 +501,7 @@
        "\n",
        "#### 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 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",
@@ -560,11 +513,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",
+        "**Note**: When referring to the worker pool specification, the directory slash is replaced with a dot (`trainer.task`) and file suffix (`.py`) is dropped.\n",
        "\n",
        "#### Package Assembly\n",
        "\n",
-        "In the following cells, you will assemble the training package."
+        "In the following cells, you assemble the training package."
      ]
    },
    {
@@ -604,10 +559,12 @@
      "source": [
        "#### Task.py contents\n",
        "\n",
-        "In the next cell, you write the contents of the training script task.py. I 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 Boston Housing dataset from TF.Keras builtin datasets\n",
+        "- Loads Boston Housing dataset from TF.Keras builtin datasets.\n",
        "- Builds a simple deep neural network model using TF.Keras model API.\n",
        "- Compiles the model (`compile()`).\n",
        "- Sets a training distribution strategy according to the argument `args.distribute`.\n",
@@ -736,7 +693,7 @@
      "source": [
        "#### 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."
+        "Next, package the training folder into a compressed tar ball, and then store it in your Cloud Storage bucket."
      ]
    },
    {
@@ -782,7 +739,7 @@
      },
      "outputs": [],
      "source": [
-        "job = aip.CustomTrainingJob(\n",
+        "job = aiplatform.CustomTrainingJob(\n",
        "    display_name=\"boston\",\n",
        "    script_path=\"custom/trainer/task.py\",\n",
        "    container_uri=TRAIN_IMAGE,\n",
@@ -798,16 +755,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."
      ]
    },
    {
@@ -823,7 +780,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",
@@ -845,15 +802,15 @@
      "source": [
        "#### 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 the completion of the job."
      ]
    },
    {
@@ -894,9 +851,9 @@
      "source": [
        "## Load the saved model\n",
        "\n",
-        "Your model is stored in a TensorFlow SavedModel format in a Cloud Storage bucket. Now load it from the Cloud Storage bucket, and then you can do some things, like evaluate the model, and do a prediction.\n",
+        "Your model is stored in a TensorFlow SavedModel format in a Cloud Storage bucket. Now, load it from the Cloud Storage bucket, and then you can run model evaluation and prepare it for batch prediction requests.\n",
        "\n",
-        "To load, you use the TF.Keras `model.load_model()` method passing it the Cloud Storage path where the model is saved -- specified by `MODEL_DIR`."
+        "To load, use the TF.Keras `model.load_model()` method passing it the Cloud Storage path where the model is saved -- specified by `MODEL_DIR`."
      ]
    },
    {
@@ -907,8 +864,6 @@
      },
      "outputs": [],
      "source": [
-        "import tensorflow as tf\n",
-        "\n",
        "local_model = tf.keras.models.load_model(MODEL_DIR)"
      ]
    },
@@ -920,18 +875,15 @@
      "source": [
        "## Evaluate the model\n",
        "\n",
-        "Now let's 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 Boston Housing 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 feature data, and the corresponding labels (median value of owner-occupied home).\n",
+        "Load the Boston Housing 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 again a tuple of two elements: the feature data, and the corresponding labels (median value of owner-occupied home).\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 data in each column by dividing each value by the maximum value of that column. This replaces each single value with a 32-bit floating point number between 0 and 1."
+        "Before you can use the data for evaluation, you need to preprocess it. Normalize (rescale) the data in each column by dividing each value by the maximum value of that column. This replaces each single value with a 32-bit floating point number between 0 and 1."
      ]
    },
    {
@@ -976,7 +928,7 @@
      "source": [
        "### Perform the model evaluation\n",
        "\n",
-        "Now evaluate how well the model in the custom job did."
+        "Use the model's `evaluate()` method to perform the evaluation."
      ]
    },
    {
@@ -1000,7 +952,7 @@
        "\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. Therefore, you need to know the input layer name of 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."
      ]
@@ -1034,38 +986,38 @@
        "id": "explanation_spec"
      },
      "source": [
-        "### Explanation Specification\n",
+        "### 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 when doing a prediction, you must enable the explanation capability 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 be very long running)\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",
+        "#### 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",
        "Assigns credit for the outcome to each feature, and considers different permutations of the features. This method provides a sampling approximation of exact Shapley values.\n",
        "\n",
-        "Use Cases:\n",
+        "Use cases:\n",
        "  - Classification and regression on tabular data.\n",
        "\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",
+        "For any non-trival number of features, this is very compute expensive. You can reduce the number of paths over the features to M * `path_count`.\n",
        "\n",
        "#### Integrated Gradients\n",
        "\n",
        "A gradients-based method to efficiently compute feature attributions with the same axiomatic properties as the Shapley value.\n",
        "\n",
-        "Use Cases:\n",
+        "Use cases:\n",
        "  - Classification and regression on tabular data.\n",
        "  - Classification on image data.\n",
        "\n",
@@ -1077,7 +1029,7 @@
        "\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",
        "\n",
-        "Use Cases:\n",
+        "Use cases:\n",
        "\n",
        "  - Classification on image data.\n",
        "\n",
@@ -1085,7 +1037,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 the explainabilty algorithm you use on your custom model."
      ]
    },
    {
@@ -1105,7 +1057,7 @@
        "elif XAI == \"xrai\":\n",
        "    PARAMETERS = {\"xrai_attribution\": {\"step_count\": 50}}\n",
        "\n",
-        "parameters = aip.explain.ExplanationParameters(PARAMETERS)"
+        "parameters = aiplatform.explain.ExplanationParameters(PARAMETERS)"
      ]
    },
    {
@@ -1114,15 +1066,15 @@
        "id": "explanation_metadata:tabular"
      },
      "source": [
-        "#### Explanation Metadata\n",
+        "#### 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's what you want to explain.\n",
        "\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'd want to pick which of the two outputs to explain.\n",
        "\n",
        "    y, z = f(x)\n",
        "\n",
@@ -1136,10 +1088,10 @@
        "<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",
+        "- `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'd want to pick `a` (or some subset of) for the features, and not `b` since it doesn't contribute to the prediction.\n",
        "\n",
        "    y = f(a,b)\n",
        "\n",
@@ -1153,7 +1105,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",
@@ -1166,7 +1118,7 @@
        "\n",
        "<blockquote>\n",
        " - \"encoding\": \"BAG_OF_FEATURES\" : Indicates that the inputs are set of tabular features.<br/>\n",
-        " - \"index_feature_mapping\": [ feature-names ] : A list of human readable names for each feature. For this example, we use the feature names specified in the dataset.<br/>\n",
+        " - \"index_feature_mapping\": [ feature-names ] : A list of human readable names for each feature. For this example, you use the feature names specified in the dataset.<br/>\n",
        " - \"modality\": \"numeric\": Indicates the field values are numeric.\n",
        "</blockquote>"
      ]
@@ -1202,10 +1154,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={\"features\": input_metadata}, outputs={\"medv\": output_metadata}\n",
        ")"
      ]
@@ -1218,14 +1170,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."
      ]
@@ -1238,7 +1190,7 @@
      },
      "outputs": [],
      "source": [
-        "model = aip.Model.upload(\n",
+        "model = aiplatform.Model.upload(\n",
        "    display_name=\"boston\",\n",
        "    artifact_uri=MODEL_DIR,\n",
        "    serving_container_image_uri=DEPLOY_IMAGE,\n",
@@ -1258,7 +1210,7 @@
      "source": [
        "## Send a batch prediction request\n",
        "\n",
-        "Send a batch prediction to your deployed model."
+        "Create test items as your input to the batch prediction request. Then, send a batch prediction request to your deployed model. "
      ]
    },
    {
@@ -1280,7 +1232,7 @@
      "source": [
        "### Make the batch input file\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",
+        "Now make a batch input file which you 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",
        "\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",
@@ -1303,7 +1255,7 @@
        "    f.write(\"crim, zn, indus, chas, nox, rm, age, dis, rad, tax, ptratio, b, lstat\\n\")\n",
        "    f.write(str(x_test[0].tolist()).replace(\"[\", \"\").replace(\"]\", \"\"))\n",
        "    f.write(\"\\n\")\n",
-        "    f.write(str(x_test[1].tolist()))\n",
+        "    f.write(str(x_test[1].tolist()).replace(\"[\", \"\").replace(\"]\", \"\"))\n",
        "    f.write(\"\\n\")\n",
        "\n",
        "gcs_input_uri = BUCKET_URI + \"/test.csv\"\n",
@@ -1319,15 +1271,15 @@
      "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",
+        "Now that your model resource is trained, 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",
-        "- `gcs_destination_prefix`: The Cloud Storage location for storing the batch prediction resuls.\n",
+        "- `gcs_destination_prefix`: The Cloud Storage location for storing the batch prediction results.\n",
        "- `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 blocks while waiting for the asynchronous batch job to complete."
      ]
    },
    {
@@ -1363,9 +1315,9 @@
        "id": "batch_request_wait:mbsdk,testing"
      },
      "source": [
-        "### Wait for completion of batch prediction job\n",
+        "### Wait for the 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. Alternatively, you can set the `sync` parameter to `True` in the `batch_predict()` method to wait until the batch prediction job is completed."
      ]
    },
    {
@@ -1389,10 +1341,7 @@
        "\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",
-        "\n",
-        "- CSV header + predicted_label\n",
-        "- CSV row + explanation, per prediction request"
+        "The results are written to the Cloud Storage output bucket you specified in the batch prediction request. You call the `iter_outputs()` method to get a list of each Cloud Storage file generated with the results."
      ]
    },
    {
@@ -1440,17 +1389,18 @@
      },
      "outputs": [],
      "source": [
-        "delete_bucket = False\n",
-        "\n",
+        "# Delete the model\n",
        "model.delete()\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",
+        "# Delete the batch prediction job\n",
+        "batch_predict_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\n",
        "\n",
+        "# Remove the locally generated files\n",
        "! rm -rf batch.csv custom.tar.gz custom"
      ]
    }
@@ -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: Custom training tabular regression model for online 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_tabular_regression_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_tabular_regression_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_tabular_regression_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_tabular_regression_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_custom_tabular_regression_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>"
      ]
    },
    {
@@ -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 tabular regression model with explanations. Then, you learn how to use the Vertex AI 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 Prediction`\n",
-        "- `Vertex Explainable AI`\n",
-        "- `Vertex AI Model` resource\n",
-        "- `Vertex AI Endpoint` resource\n",
+        "- Vertex AI Training\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 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 [Boston Housing Prices dataset](https://www.cs.toronto.edu/~delve/data/boston/bostonDetail.html). The version of the dataset you will use in this tutorial is built into TensorFlow. The trained model predicts the median price of a house in units of 1K USD."
+        "The dataset used for this tutorial is the [Boston Housing Prices dataset](https://www.cs.toronto.edu/~delve/data/boston/bostonDetail.html). The version of the dataset you use in this tutorial is the one that's available from TensorFlow SDK. The trained model predicts the median price of a house in units of 1K USD."
      ]
    },
    {
@@ -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,7 +158,7 @@
      "source": [
        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
        "                                 google-cloud-storage \\\n",
-        "                                 tensorflow==2.11"
+        "                                 tensorflow==2.15.1"
      ]
    },
    {
@@ -158,7 +167,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."
      ]
    },
    {
@@ -169,63 +180,25 @@
      },
      "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",
-        "\n",
-        "### Set your project ID\n",
-        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "* Run `gcloud config list`.\n",
-        "* Run `gcloud projects list`.\n",
-        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "set_project_id"
-      },
-      "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\"}"
+        "<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"
      ]
    },
    {
@@ -234,14 +207,9 @@
        "id": "gcp_authenticate"
      },
      "source": [
-        "### Authenticate your Google Cloud account\n",
+        "### Authenticate your notebook environment (Colab only)\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:**"
+        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
@@ -252,38 +220,36 @@
      },
      "outputs": [],
      "source": [
-        "# ! gcloud auth login"
+        "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": "0367eac06a10"
+        "id": "1115531b8bec"
      },
      "source": [
-        "**3. Colab, uncomment and run:**"
+        "### 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)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "21ad4dbb4a61"
+        "id": "294fe4e5a671"
      },
      "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."
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -314,7 +280,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."
      ]
    },
    {
@@ -325,7 +291,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
@@ -334,10 +300,11 @@
        "id": "setup_vars"
      },
      "source": [
-        "### Set up variables\n",
+        "### initialize Vertex AI SDK for Python\n",
        "\n",
-        "Next, set up some variables used throughout the tutorial.\n",
-        "### Import libraries and define constants"
+        "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 and corresponding bucket."
      ]
    },
    {
@@ -350,29 +317,11 @@
      "source": [
        "import os\n",
        "\n",
-        "import google.cloud.aiplatform as aip"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "init_aip:mbsdk"
-      },
-      "source": [
-        "### Initialize Vertex AI SDK for Python\n",
+        "import numpy as np\n",
+        "import tensorflow as tf\n",
+        "from google.cloud import aiplatform\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": [
-        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
+        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -385,16 +334,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'd 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",
        "\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."
      ]
    },
    {
@@ -407,7 +356,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",
@@ -415,7 +364,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",
@@ -472,10 +421,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",
@@ -492,19 +441,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 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 is'nt 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."
      ]
    },
    {
@@ -540,9 +489,9 @@
        "id": "tutorial_start:custom"
      },
      "source": [
-        "# Tutorial\n",
+        "## Custom model creation and training\n",
        "\n",
-        "Now you are ready to start creating your own custom model and training for Boston Housing."
+        "Now you're ready to start creating your own custom model and training for Boston Housing price prediction."
      ]
    },
    {
@@ -555,7 +504,7 @@
        "\n",
        "#### 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",
@@ -567,11 +516,11 @@
        "\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",
-        "In the following cells, you will assemble the training package."
+        "In the following cells, you assemble the training package."
      ]
    },
    {
@@ -611,10 +560,12 @@
      "source": [
        "#### Task.py contents\n",
        "\n",
-        "In the next cell, you write the contents of the training script task.py. I won't go into detail, it's just there for you to browse. In summary:\n",
+        "In the next cell, write the contents of the training script *task.py*. \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 Boston Housing dataset from TF.Keras builtin datasets\n",
+        "To summarize, the script performs the following steps:\n",
+        "\n",
+        "- Gets 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 Boston Housing dataset from TF.Keras built-in datasets.\n",
        "- Builds a simple deep neural network model using TF.Keras model API.\n",
        "- Compiles the model (`compile()`).\n",
        "- Sets a training distribution strategy according to the argument `args.distribute`.\n",
@@ -688,7 +639,7 @@
        "  # Scaling Boston Housing data features\n",
        "  def scale(feature):\n",
        "    max = np.max(feature)\n",
-        "    feature = (feature / max).astype(np.float)\n",
+        "    feature = (feature / max).astype(float)\n",
        "    return feature, max\n",
        "\n",
        "  (x_train, y_train), (x_test, y_test) = tf.keras.datasets.boston_housing.load_data(\n",
@@ -741,9 +692,9 @@
        "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."
+        "Next, package the training folder into a compressed tar ball, and then store it in your Cloud Storage bucket."
      ]
    },
    {
@@ -769,7 +720,9 @@
        "### 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",
+        "1) Create a custom training job.\n",
+        "2) Specify your training parameters and run the job.\n",
        "\n",
        "#### Create custom training job\n",
        "\n",
@@ -789,12 +742,12 @@
      },
      "outputs": [],
      "source": [
-        "job = aip.CustomTrainingJob(\n",
+        "job = aiplatform.CustomTrainingJob(\n",
        "    display_name=\"boston\",\n",
        "    script_path=\"custom/trainer/task.py\",\n",
        "    container_uri=TRAIN_IMAGE,\n",
        "    requirements=[\"gcsfs==0.7.1\", \"tensorflow-datasets==4.4\"],\n",
-        "    location=REGION,\n",
+        "    location=LOCATION,\n",
        ")\n",
        "\n",
        "print(job)"
@@ -806,16 +759,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."
      ]
    },
    {
@@ -831,7 +784,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",
@@ -853,15 +806,15 @@
      "source": [
        "#### 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 write the model artifacts.\n",
+        "- `sync`: Set **True** to wait until the completion of the job."
      ]
    },
    {
@@ -902,7 +855,7 @@
      "source": [
        "## Load the saved model\n",
        "\n",
-        "Your model is stored in a TensorFlow SavedModel format in a Cloud Storage bucket. Now load it from the Cloud Storage bucket, and then you can do some things, like evaluate the model, and do a prediction.\n",
+        "Your model is stored in a TensorFlow SavedModel format in a Cloud Storage bucket. Now, load it from the Cloud Storage bucket, and then you can run model evaluation and prepare it for prediction requests.\n",
        "\n",
        "To load, you use the TF.Keras `model.load_model()` method passing it the Cloud Storage path where the model is saved -- specified by `MODEL_DIR`."
      ]
@@ -915,8 +868,6 @@
      },
      "outputs": [],
      "source": [
-        "import tensorflow as tf\n",
-        "\n",
        "local_model = tf.keras.models.load_model(MODEL_DIR)"
      ]
    },
@@ -928,18 +879,13 @@
      "source": [
        "## Evaluate the model\n",
        "\n",
-        "Now let's find out how good the model is.\n",
-        "\n",
        "### Load evaluation data\n",
        "\n",
-        "You will load the Boston Housing 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 feature data, and the corresponding labels (median value of owner-occupied home).\n",
+        "Load the Boston Housing 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 is the test data. Each element is also a tuple of two elements: the feature data, and the corresponding labels (median value of owner-occupied home).\n",
        "\n",
-        "You don't need the training data, and hence why we loaded it as `(_, _)`.\n",
+        "You don't need the training data, and hence load it 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 data in each column by dividing each value by the maximum value of that column. This replaces each single value with a 32-bit floating point number between 0 and 1."
+        "Before you can use the data for evaluation, you need to preprocess it. Normalize (rescale) the data in each column by dividing each value by the maximum value of that column. This replaces each single value with a 32-bit floating point number between 0 and 1."
      ]
    },
    {
@@ -950,7 +896,6 @@
      },
      "outputs": [],
      "source": [
-        "import numpy as np\n",
        "from tensorflow.keras.datasets import boston_housing\n",
        "\n",
        "(_, _), (x_test, y_test) = boston_housing.load_data(\n",
@@ -960,16 +905,16 @@
        "\n",
        "def scale(feature):\n",
        "    max = np.max(feature)\n",
-        "    feature = (feature / max).astype(np.float32)\n",
+        "    feature = (feature / max).astype(float)\n",
        "    return feature\n",
        "\n",
        "\n",
-        "# Let's save one data item that has not been scaled\n",
+        "# save one data item that has not been scaled\n",
        "x_test_notscaled = x_test[0:1].copy()\n",
        "\n",
        "for _ in range(13):\n",
        "    x_test[_] = scale(x_test[_])\n",
-        "x_test = x_test.astype(np.float32)\n",
+        "x_test = x_test.astype(float)\n",
        "\n",
        "print(x_test.shape, x_test.dtype, y_test.shape)\n",
        "print(\"scaled\", x_test[0])\n",
@@ -984,7 +929,7 @@
      "source": [
        "### Perform the model evaluation\n",
        "\n",
-        "Now evaluate how well the model in the custom job did."
+        "Use the model's `evaluate()` method to perform the evaluation."
      ]
    },
    {
@@ -1008,7 +953,7 @@
        "\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, therefore, you need to know the input layer name of 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."
      ]
@@ -1044,17 +989,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 when doing a prediction, you must enable the explanation capability and set corresponding settings when you upload your custom model to an 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 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 about the settings for the explainability algorithm.\n",
        "\n",
        "#### Shapley\n",
        "\n",
@@ -1065,7 +1010,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 that are 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",
@@ -1093,7 +1038,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 the explainabilty algorithm that you use on your custom model."
      ]
    },
    {
@@ -1113,7 +1058,7 @@
        "elif XAI == \"xrai\":\n",
        "    PARAMETERS = {\"xrai_attribution\": {\"step_count\": 50}}\n",
        "\n",
-        "parameters = aip.explain.ExplanationParameters(PARAMETERS)"
+        "parameters = aiplatform.explain.ExplanationParameters(PARAMETERS)"
      ]
    },
    {
@@ -1124,13 +1069,13 @@
      "source": [
        "#### Explanation Metadata\n",
        "\n",
-        "Let's first dive deeper into the explanation metadata, which consists of:\n",
+        "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 formulae, where the output is `y` and that's what needs to be explained.\n",
        "\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'd want to pick which of the two outputs to explain.\n",
        "\n",
        "    y, z = f(x)\n",
        "\n",
@@ -1144,10 +1089,10 @@
        "<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",
+        "- `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 doesn't contribute to the prediction.\n",
        "\n",
        "    y = f(a,b)\n",
        "\n",
@@ -1161,7 +1106,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",
@@ -1174,7 +1119,7 @@
        "\n",
        "<blockquote>\n",
        " - \"encoding\": \"BAG_OF_FEATURES\" : Indicates that the inputs are set of tabular features.<br/>\n",
-        " - \"index_feature_mapping\": [ feature-names ] : A list of human readable names for each feature. For this example, we use the feature names specified in the dataset.<br/>\n",
+        " - \"index_feature_mapping\": [ feature-names ] : A list of human readable names for each feature. For this example, you must use the feature names specified in the dataset.<br/>\n",
        " - \"modality\": \"numeric\": Indicates the field values are numeric.\n",
        "</blockquote>"
      ]
@@ -1210,10 +1155,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={\"features\": input_metadata}, outputs={\"medv\": output_metadata}\n",
        ")"
      ]
@@ -1226,14 +1171,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."
      ]
@@ -1246,7 +1191,7 @@
      },
      "outputs": [],
      "source": [
-        "model = aip.Model.upload(\n",
+        "model = aiplatform.Model.upload(\n",
        "    display_name=\"boston\",\n",
        "    artifact_uri=MODEL_DIR,\n",
        "    serving_container_image_uri=DEPLOY_IMAGE,\n",
@@ -1266,12 +1211,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, 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",
@@ -1324,7 +1269,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 out of the test (holdout) portion of the dataset as a test item."
      ]
    },
    {
@@ -1346,9 +1291,9 @@
        "id": "explain_request:mbsdk,custom,lrg"
      },
      "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 your endpoint, you can get online explanations by sending prediction requests to your endpoint.\n",
        "\n",
        "#### Request\n",
        "\n",
@@ -1360,12 +1305,12 @@
        "\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 which generated the predictions.\n",
+        "- `explanations`: The feature attributions returned by the Explainable AI."
      ]
    },
    {
@@ -1390,7 +1335,7 @@
      "source": [
        "### Understanding the explanations response\n",
        "\n",
-        "First, you will look what your model predicted and compare it to the actual value."
+        "First, look at what your model predicted and compare it to the actual value."
      ]
    },
    {
@@ -1413,7 +1358,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."
+        "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."
      ]
    },
    {
@@ -1457,9 +1402,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 better make sense of the feature attributions you receive, you should compare them with your model's baseline. In most cases, the sum of your attribution values and 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 the 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()` function.\n",
        "\n",
        "#### Get explanations"
      ]
@@ -1499,9 +1444,6 @@
      },
      "outputs": [],
      "source": [
-        "import numpy as np\n",
-        "\n",
-        "\n",
        "def sanity_check_explanations(\n",
        "    explanation, prediction, mean_tgt_value=None, variance_tgt_value=None\n",
        "):\n",
@@ -1536,28 +1478,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": {
@@ -1580,12 +1500,24 @@
      },
      "outputs": [],
      "source": [
-        "delete_bucket = False\n",
+        "# Undeploy the model from the endpoint\n",
+        "endpoint.undeploy_all()\n",
        "\n",
+        "# Delete the model endpoint\n",
        "endpoint.delete()\n",
+        "\n",
+        "# Delete the model resource\n",
        "model.delete()\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "# Delete the custom training job\n",
+        "job.delete()\n",
+        "\n",
+        "# Delete the locally generated files\n",
+        "! rm -rf custom custom.tar.gz\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"
      ]
    }
@@ -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: Custom training tabular regression model for online prediction with explainabilty using get_metadata\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_tabular_regression_online_explain_get_metadata.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_tabular_regression_online_explain_get_metadata.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_tabular_regression_online_explain_get_metadata.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_tabular_regression_online_explain_get_metadata.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_tabular_regression_online_explain_get_metadata.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,28 +77,28 @@
      "source": [
        "### Objective\n",
        "\n",
-        "In this tutorial, you learn how to create a custom model from a Python script in a Google prebuilt Docker container using the Vertex AI SDK, and then do a prediction with explanations on the deployed model by sending data. You can alternatively create custom models using `gcloud` command-line tool or online using Cloud Console.\n",
+        "In this tutorial, you learn how to create a custom model from a Python script in a Google prebuilt Docker container using the Vertex AI SDK. Then, you learn how to use the Vertex AI 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",
        "\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",
        "\n",
        "The steps performed include:\n",
        "\n",
-        "- Create a Vertex custom job for training a model.\n",
+        "- Create a Vertex AI custom job for training a TensorFLow model.\n",
        "- Train a TensorFlow model.\n",
        "- Retrieve and load the model artifacts.\n",
-        "- View the model evaluation.\n",
+        "- View the model evaluation for the trained model.\n",
        "- Set explanation parameters.\n",
-        "- Upload the model as a Vertex `Model` resource.\n",
-        "- Deploy the `Model` resource to a serving `Endpoint` resource.\n",
+        "- Upload the model as a Vertex AI model 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."
      ]
    },
    {
@@ -107,7 +109,7 @@
      "source": [
        "### Dataset\n",
        "\n",
-        "The dataset used for this tutorial is the [Boston Housing Prices dataset](https://www.cs.toronto.edu/~delve/data/boston/bostonDetail.html). The version of the dataset you will use in this tutorial is built into TensorFlow. The trained model predicts the median price of a house in units of 1K USD."
+        "The dataset used for this tutorial is the [Boston Housing Prices dataset](https://www.cs.toronto.edu/~delve/data/boston/bostonDetail.html). The version of the dataset you use in this tutorial is the one that's available from TensorFlow SDK. The trained model predicts the median price of a house in units of 1K USD."
      ]
    },
    {
@@ -133,12 +135,19 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "install_aip:mbsdk"
+        "id": "61RBz8LLbxCR"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the latest version of Vertex SDK for Python."
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "No17Cw5hgx12"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
@@ -149,147 +158,100 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
-        "                                 google-cloud-storage\n",
-        "\n",
-        "if os.getenv(\"IS_TESTING\"):\n",
-        "    ! pip3 install --upgrade --quiet tensorflow "
+        "                                 google-cloud-storage \\\n",
+        "                                 tensorflow==2.15.1"
      ]
    },
    {
      "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\"}"
      ]
    },
    {
@@ -331,30 +293,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}"
      ]
    },
    {
@@ -363,9 +302,9 @@
        "id": "init_aip:mbsdk"
      },
      "source": [
-        "## Initialize Vertex SDK for Python\n",
+        "## Initialize Vertex SDK for Python and other libraries\n",
        "\n",
-        "Initialize the Vertex SDK for Python for your project and corresponding bucket."
+        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
      ]
    },
    {
@@ -376,7 +315,15 @@
      },
      "outputs": [],
      "source": [
-        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
+        "import os\n",
+        "\n",
+        "import numpy as np\n",
+        "import tensorflow as tf\n",
+        "from google.cloud import aiplatform\n",
+        "from google.cloud.aiplatform.explain.metadata.tf.v2 import \\\n",
+        "    saved_model_metadata_builder\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -389,16 +336,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.AcceleratorType.NVIDIA_TESLA_K80, 4)\n",
+        "    (aiplatform.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",
        "\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 will 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."
      ]
    },
    {
@@ -409,11 +356,9 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
        "if os.getenv(\"IS_TESTING_TRAIN_GPU\"):\n",
        "    TRAIN_GPU, TRAIN_NGPU = (\n",
-        "        aip.gapic.AcceleratorType.NVIDIA_TESLA_K80,\n",
+        "        aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_T4,\n",
        "        int(os.getenv(\"IS_TESTING_TRAIN_GPU\")),\n",
        "    )\n",
        "else:\n",
@@ -421,7 +366,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",
@@ -494,19 +439,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 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 is'nt 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,9 +487,9 @@
        "id": "tutorial_start:custom"
      },
      "source": [
-        "# Tutorial\n",
+        "## Custom model creation and training\n",
        "\n",
-        "Now you are ready to start creating your own custom model and training for Boston Housing."
+        "Now you're ready to start creating your own custom model and training for Boston Housing price prediction."
      ]
    },
    {
@@ -557,7 +502,7 @@
        "\n",
        "#### 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",
@@ -569,11 +514,11 @@
        "\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 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",
-        "In the following cells, you will assemble the training package."
+        "In the following cells, you assemble the training package."
      ]
    },
    {
@@ -613,16 +558,18 @@
      "source": [
        "#### Task.py contents\n",
        "\n",
-        "In the next cell, you write the contents of the training script task.py. I won't go into detail, it's just there for you to browse. In summary:\n",
+        "In the next cell, write the contents of the training script *task.py*.\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 Boston Housing dataset from TF.Keras builtin datasets\n",
+        "To summarize, the script performs the following steps:\n",
+        "\n",
+        "- Retrieves the directory to save the model artifacts from the `--model_dir` command line argument. If it's not specified, then uses the directory set in the environment variable `AIP_MODEL_DIR`.\n",
+        "- Loads Boston Housing dataset from TF.Keras builtin datasets.\n",
        "- Builds a simple deep neural network model using TF.Keras model API.\n",
        "- Compiles the model (`compile()`).\n",
-        "- Sets a training distribution strategy according to the argument `args.distribute`.\n",
-        "- Trains the model (`fit()`) with epochs specified by `args.epochs`.\n",
+        "- Sets a training distribution strategy according to the `args.distribute` argument.\n",
+        "- Trains the model (`fit()`) with epochs specified by the `args.epochs` argument.\n",
        "- Saves the trained model (`save(args.model_dir)`) to the specified model directory.\n",
-        "- Saves the maximum value for each feature `f.write(str(params))` to the specified parameters file."
+        "- Saves the maximum value for each feature (`f.write(str(params))`) to the specified parameters file."
      ]
    },
    {
@@ -743,9 +690,9 @@
        "id": "tarball_training_script"
      },
      "source": [
-        "#### Store training script on your Cloud Storage bucket\n",
+        "### Store training script in 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."
+        "Next, package the training folder into a compressed tar ball, and then store it in your Cloud Storage bucket."
      ]
    },
    {
@@ -771,11 +718,13 @@
        "### 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",
+        "1) Create a custom training job\n",
+        "2) Specify your training parameters and run the job.\n",
        "\n",
        "#### Create custom training job\n",
        "\n",
-        "A custom training job is created with the `CustomTrainingJob` class, with the following parameters:\n",
+        "A custom training job is created using `CustomTrainingJob` class, with the following parameters:\n",
        "\n",
        "- `display_name`: The human readable name for the custom training job.\n",
        "- `container_uri`: The training container image.\n",
@@ -791,7 +740,7 @@
      },
      "outputs": [],
      "source": [
-        "job = aip.CustomTrainingJob(\n",
+        "job = aiplatform.CustomTrainingJob(\n",
        "    display_name=\"boston\",\n",
        "    script_path=\"custom/trainer/task.py\",\n",
        "    container_uri=TRAIN_IMAGE,\n",
@@ -807,16 +756,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-1**(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."
      ]
    },
    {
@@ -832,7 +781,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,15 +803,15 @@
      "source": [
        "#### 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 write the model artifacts.\n",
+        "- `sync`: Set **True** to wait until the completion of the job."
      ]
    },
    {
@@ -903,7 +852,7 @@
      "source": [
        "## Load the saved model\n",
        "\n",
-        "Your model is stored in a TensorFlow SavedModel format in a Cloud Storage bucket. Now load it from the Cloud Storage bucket, and then you can do some things, like evaluate the model, and do a prediction.\n",
+        "Your model is stored in a TensorFlow SavedModel format in a Cloud Storage bucket. Now, load it from the Cloud Storage bucket, and then you can run model evaluation and prepare it for prediction requests.\n",
        "\n",
        "To load, you use the TF.Keras `model.load_model()` method passing it the Cloud Storage path where the model is saved -- specified by `MODEL_DIR`."
      ]
@@ -916,8 +865,6 @@
      },
      "outputs": [],
      "source": [
-        "import tensorflow as tf\n",
-        "\n",
        "local_model = tf.keras.models.load_model(MODEL_DIR)"
      ]
    },
@@ -929,18 +876,14 @@
      "source": [
        "## Evaluate the model\n",
        "\n",
-        "Now let's find out how good the model is.\n",
        "\n",
        "### Load evaluation data\n",
        "\n",
-        "You will load the Boston Housing 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 feature data, and the corresponding labels (median value of owner-occupied home).\n",
+        "Load the Boston Housing 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 is the test data. Each element is also a tuple of two elements: the feature data, and the corresponding labels (median value of owner-occupied home).\n",
        "\n",
-        "You don't need the training data, and hence why we loaded it as `(_, _)`.\n",
+        "You don't need the training data, and hence load it 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 data in each column by dividing each value by the maximum value of that column. This replaces each single value with a 32-bit floating point number between 0 and 1."
+        "Before you can use the data for evaluation, you need to preprocess it. Normalize (rescale) the data in each column by dividing each value by the maximum value of that column. This replaces each single value with a 32-bit floating point number between 0 and 1."
      ]
    },
    {
@@ -951,7 +894,6 @@
      },
      "outputs": [],
      "source": [
-        "import numpy as np\n",
        "from tensorflow.keras.datasets import boston_housing\n",
        "\n",
        "(_, _), (x_test, y_test) = boston_housing.load_data(\n",
@@ -985,7 +927,7 @@
      "source": [
        "### Perform the model evaluation\n",
        "\n",
-        "Now evaluate how well the model in the custom job did."
+        "Use the model's `evaluate()` method to perform the evaluation"
      ]
    },
    {
@@ -1009,7 +951,7 @@
        "\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, therefore, you need to know the input layer name of 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."
      ]
@@ -1045,17 +987,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 when doing a prediction, you must enable the explanation capability and set corresponding settings when you upload your custom model to an 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 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 about the settings for the explainability algorithm.\n",
        "\n",
        "#### Shapley\n",
        "\n",
@@ -1066,7 +1008,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 that are 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",
@@ -1094,7 +1036,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 the explainabilty algorithm that you use on your custom model."
      ]
    },
    {
@@ -1114,7 +1056,7 @@
        "elif XAI == \"xrai\":\n",
        "    PARAMETERS = {\"xrai_attribution\": {\"step_count\": 50}}\n",
        "\n",
-        "parameters = aip.explain.ExplanationParameters(PARAMETERS)"
+        "parameters = aiplatform.explain.ExplanationParameters(PARAMETERS)"
      ]
    },
    {
@@ -1125,9 +1067,9 @@
      "source": [
        "#### Explanation Metadata\n",
        "\n",
-        "Let's first dive deeper into the explanation metadata, which consists of:\n",
+        "The explanation metadata consists of:\n",
        "\n",
-        "- `outputs`: A scalar value in the output to attribute -- what to explain.\n",
+        "- `outputs`: A scalar value in the output to attribute -- what you want to explain.\n",
        "- `inputs`: The features for attribution -- how they contributed to the output.\n",
        "\n",
        "You can either customize your metadata -- what to explain and what to attribute, or automatically generate the metadata using the method `get_metadata_protobuf()`. This method will construct metadata for explaining all outputs and attributing all inputs."
@@ -1141,9 +1083,6 @@
      },
      "outputs": [],
      "source": [
-        "from google.cloud.aiplatform.explain.metadata.tf.v2 import \\\n",
-        "    saved_model_metadata_builder\n",
-        "\n",
        "builder = saved_model_metadata_builder.SavedModelMetadataBuilder(MODEL_DIR)\n",
        "metadata = builder.get_metadata_protobuf()\n",
        "print(metadata)"
@@ -1157,14 +1096,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."
      ]
@@ -1177,7 +1116,7 @@
      },
      "outputs": [],
      "source": [
-        "model = aip.Model.upload(\n",
+        "model = aiplatform.Model.upload(\n",
        "    display_name=\"boston\",\n",
        "    artifact_uri=MODEL_DIR,\n",
        "    serving_container_image_uri=DEPLOY_IMAGE,\n",
@@ -1197,12 +1136,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, 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",
@@ -1255,7 +1194,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 out of the test (holdout) portion of the dataset as a test item."
      ]
    },
    {
@@ -1277,9 +1216,9 @@
        "id": "explain_request:mbsdk,custom,lrg"
      },
      "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 resource is deployed to an endpoint resource, you can get online explanations by sending prediction requests to the endpoint resource.\n",
        "\n",
        "#### Request\n",
        "\n",
@@ -1295,8 +1234,8 @@
        "\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 which generated the predictions.\n",
+        "- `explanations`: The feature attributions returned by the Explainable AI."
      ]
    },
    {
@@ -1321,7 +1260,7 @@
      "source": [
        "### Understanding the explanations response\n",
        "\n",
-        "First, you will look what your model predicted and compare it to the actual value."
+        "First, look at what your model predicted and compare it to the actual value."
      ]
    },
    {
@@ -1344,7 +1283,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."
+        "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."
      ]
    },
    {
@@ -1390,9 +1329,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 better make sense of the feature attributions you receive, you should compare them with your model's baseline. In most cases, the sum of your attribution values and 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 the 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()` function.\n",
        "\n",
        "#### Get explanations"
      ]
@@ -1432,9 +1371,6 @@
      },
      "outputs": [],
      "source": [
-        "import numpy as np\n",
-        "\n",
-        "\n",
        "def sanity_check_explanations(\n",
        "    explanation, prediction, mean_tgt_value=None, variance_tgt_value=None\n",
        "):\n",
@@ -1469,28 +1405,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": {
@@ -1513,12 +1427,23 @@
      },
      "outputs": [],
      "source": [
-        "delete_bucket = False\n",
+        "# Undeploy the model from the endpoint\n",
+        "endpoint.undeploy_all()\n",
        "\n",
+        "# Delete the model endpoint\n",
        "endpoint.delete()\n",
+        "\n",
+        "# Delete the model resource\n",
        "model.delete()\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "# Delete the custom training job\n",
+        "job.delete()\n",
+        "\n",
+        "# Delete the locally generated files\n",
+        "! rm -rf custom custom.tar.gz\n",
+        "\n",
+        "delete_bucket = False  # Set True for deletion\n",
+        "if delete_bucket:\n",
        "    ! gsutil rm -r $BUCKET_URI"
      ]
    }
@@ -31,43 +31,30 @@
      "source": [
        "# Explaining image classification with Vertex Explainable AI\n",
        "\n",
-        "---\n",
-        "\n",
-        "\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/explainable_ai/xai_image_classification_feature_attributions.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/xai_image_classification_feature_attributions.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%2fxai_image_classification_feature_attributions.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/xai_image_classification_feature_attributions.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/explainable_ai/xai_image_classification_feature_attributions.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": "24743cf4a1e1"
-      },
-      "source": [
-        "**_NOTE_**: This notebook has been tested in the following environment:\n",
-        "\n",
-        "* Python version = 3.9\n"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -78,7 +65,7 @@
        "\n",
        "Vertex Explainable AI offers feature-based and example-based explanations to provide better understanding of model decision making. For feature-based explanations, Vertex Explainable AI integrates feature attributions into Vertex AI. Feature attributions indicate how much each feature in your model contributed to the predictions for each given instance. For an image classification model, when you request explanations, you get the predicted class along with an overlay for the image, showing which areas in the image contributed most strongly to the resulting prediction.\n",
        "\n",
-        "To use Vertex Explainable AI on a pre-trained or custom-trained model, you must configure certain options when you create the `Model` resource that you plan to request explanations from, when you deploy the model, or when you submit a batch explanation job. This tutorial demonstrates how to configure these options, and get and visualize explanations for online and batch predictions.\n",
+        "To use Vertex Explainable AI on a pre-trained or custom-trained model, you must configure certain options when you create the model resource that you plan to request explanations from, when you deploy the model, or when you submit a batch explanation job. This tutorial demonstrates how to configure these options, and get and visualize explanations for online and batch predictions.\n",
        "\n",
        "Learn more about [Vertex Explainable AI](https://cloud.google.com/vertex-ai/docs/explainable-ai/overview) and [Vertex AI Prediction](https://cloud.google.com/vertex-ai/docs/predictions/get-predictions)."
      ]
@@ -140,12 +127,19 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "i7EUnXsZhAGF"
+        "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"
      ]
    },
    {
@@ -158,175 +152,100 @@
      "source": [
        "# Install the packages\n",
        "! pip3 install --upgrade -q google-cloud-aiplatform \\\n",
-        "                            tensorflow \\\n",
-        "                            tensorflow-hub"
+        "                            tensorflow==2.15.1 \\\n",
+        "                            tensorflow-hub \\\n",
+        "                            matplotlib"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "58707a750154"
+        "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": "f200f10a1da3"
+        "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": "BF1j6f9HApxa"
+        "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"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "WReHDGG5g0XY"
+        "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": "oM1iC_MfAts1"
+        "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",
+        "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": "LxQmrc_AARaD"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "sBCra4QMA2wR"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "74ccc9e52986"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "de775a3773ba"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "254614fa0c46"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ef21552ccea8"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "603adbbf0532"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "f6b2ccc891ed"
-      },
-      "source": [
-        "**4. Service account or other**\n",
-        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -348,7 +267,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\"}"
      ]
    },
    {
@@ -357,7 +276,7 @@
        "id": "-EcIXiGsCePi"
      },
      "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."
      ]
    },
    {
@@ -368,7 +287,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
+        "! gsutil mb -l $LOCATION -p $PROJECT_ID $BUCKET_URI"
      ]
    },
    {
@@ -402,7 +321,7 @@
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
-        "Initialize the Vertex AI SDK for Python for your project."
+        "To get started using Vertex AI, you must [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com)."
      ]
    },
    {
@@ -413,7 +332,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)"
      ]
    },
    {
@@ -456,7 +375,7 @@
      "source": [
        "## Upload model for deployment\n",
        "\n",
-        "Next, you upload the model to `Vertex AI` Model Registry, which will create a `Vertex AI Model` resource for your model. Prior to uploading, you need to define a serving function to convert data to the format your model expects."
+        "Next, you upload the model to Vertex AI Model Registry, which will create a Vertex AI Model resource for your model. Prior to uploading, you need to define a serving function to convert data to the format your model expects."
      ]
    },
    {
@@ -467,9 +386,9 @@
      "source": [
        "### Define a serving function for image data\n",
        "\n",
-        "You define a serving function to convert image data to the format your model expects. When 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",
+        "You define a serving function to convert image data to the format your model expects. When 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",
        "\n",
-        "To enable `Vertex Explainable AI` in your custom models, you need to set two additional signatures from the serving function:\n",
+        "To enable Vertex Explainable AI in your custom models, you need to set two additional signatures from the serving function:\n",
        "\n",
        "- `xai_preprocess`: The preprocessing function in the serving function.\n",
        "- `xai_model`: The concrete function for calling the model."
@@ -558,9 +477,9 @@
        "id": "wnefcZK1ba4b"
      },
      "source": [
-        "### Configure explanation settings\n",
+        "## Configure explanation settings\n",
        "\n",
-        "To use `Vertex Explainable AI` with a custom-trained model, you must configure explanation settings when uploading the model. These settings include:\n",
+        "To use Vertex Explainable AI with a custom-trained model, you must configure explanation settings when uploading the model. These settings include:\n",
        "\n",
        "- `parameters`: The feature attribution method. Available methods include `shapley`, `ig`, `xrai`.\n",
        "- `metadata`: The model's input and output for explanation. **This field is optional for TensorFlow 2 models. If omitted, Vertex AI automatically infers the inputs and outputs from the model**. You don't need to configure this field in this tutorial.\n",
@@ -594,9 +513,9 @@
        "id": "aLslkKlneNyX"
      },
      "source": [
-        "### Upload the model to a `Vertex AI Model` resource\n",
+        "## Upload the model to a Vertex AI Model resource\n",
        "\n",
-        "Next, upload your model to a `Vertex AI Model` resource with the explanation configuration. Vertex AI provides [Docker container images](https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers) that you run as pre-built containers for serving predictions and explanations from trained model artifacts."
+        "Next, upload your model to a Vertex AI Model resource with the explanation configuration. Vertex AI provides [Docker container images](https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers) that you run as pre-built containers for serving predictions and explanations from trained model artifacts."
      ]
    },
    {
@@ -607,7 +526,7 @@
      },
      "outputs": [],
      "source": [
-        "MODEL_DISPLAY_NAME = \"inception_v3_model_unique\"\n",
+        "MODEL_DISPLAY_NAME = \"inception-v3-model-unique\"\n",
        "DEPLOY_IMAGE = \"us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-11:latest\"\n",
        "\n",
        "model = aiplatform.Model.upload(\n",
@@ -637,7 +556,7 @@
      },
      "outputs": [],
      "source": [
-        "DEPLOY_DISPLAY_NAME = \"inception_v3_deploy_unique\"\n",
+        "DEPLOY_DISPLAY_NAME = \"inception-v3-deploy-unique\"\n",
        "DEPLOY_COMPUTE = \"n1-standard-4\"\n",
        "\n",
        "endpoint = model.deploy(\n",
@@ -869,7 +788,7 @@
      },
      "outputs": [],
      "source": [
-        "JOB_DISPLAY_NAME = \"inception_v3_job_unique\"\n",
+        "JOB_DISPLAY_NAME = \"inception-v3-job-unique\"\n",
        "\n",
        "batch_predict_job = model.batch_predict(\n",
        "    job_display_name=JOB_DISPLAY_NAME,\n",
@@ -964,15 +883,18 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
+        "# Deploy the model from endpoint\n",
        "endpoint.undeploy_all()\n",
+        "# Delete the endpoint\n",
        "endpoint.delete()\n",
+        "# Delete the model\n",
        "model.delete()\n",
+        "# Delete the batch prediction job\n",
        "batch_predict_job.delete()\n",
        "\n",
-        "delete_bucket = False\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "# Delete Cloud Storage bucket\n",
+        "delete_bucket = False  # Set True to delete your bucket\n",
+        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
@@ -0,0 +1,598 @@
+{
+  "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": [
+        "## Fetch historical feature values\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/feature_store/offline_feature_serving_from_bigquery_with_feature_registry.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%2Ffeature_store%2Foffline_feature_serving_from_bigquery_with_feature_registry.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/feature_store/offline_feature_serving_from_bigquery_with_feature_registry.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/feature_store/offline_feature_serving_from_bigquery_with_feature_registry.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",
+        "In this tutorial, you will learn how to use the Vertex AI SDK for Python to retrieve historical values from the feature data source in BigQuery.\n",
+        "\n",
+        "This tutorial uses the following Google Cloud ML services and resources:\n",
+        "\n",
+        "* Vertex AI Feature Store\n",
+        "* BigQuery\n",
+        "\n",
+        "The steps performed include the following:\n",
+        "\n",
+        "* Setup BigQuery data\n",
+        "* Setup Feature Registry\n",
+        "* Fetch historical feature values from feature data source in BigQuery\n",
+        "* Clean up"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "19cf444ebb99"
+      },
+      "source": [
+        "### Objective"
+      ]
+    },
+    {
+      "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 --quiet google-cloud-aiplatform bigframes"
+      ]
+    },
+    {
+      "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 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",
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
+        "\n",
+        "\n",
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(project=PROJECT_ID, location=LOCATION)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "33067053f38b"
+      },
+      "source": [
+        "### Imports and IDs\n",
+        "\n",
+        "Import the packages required to use the`fetch_historical_feature_values()`\n",
+        "function in the Vertex AI SDK for Python."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "c8abe818393b"
+      },
+      "outputs": [],
+      "source": [
+        "import bigframes\n",
+        "import bigframes.pandas\n",
+        "import pandas as pd\n",
+        "from google.cloud import bigquery\n",
+        "from vertexai.resources.preview.feature_store import (Feature, FeatureGroup,\n",
+        "                                                      offline_store)\n",
+        "from vertexai.resources.preview.feature_store import utils as fs_utils"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4d0295f5d524"
+      },
+      "source": [
+        "The following variables set BigQuery and Feature Group resources that will be\n",
+        "used or created. If you'd like to use your own data source (CSV), please adjust\n",
+        "`DATA_SOURCE`."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ac036ecfbc32"
+      },
+      "outputs": [],
+      "source": [
+        "BQ_DATASET_ID = \"fhfv_dataset_unique\"  # @param {type:\"string\"}\n",
+        "BQ_TABLE_ID = \"fhfv_table_unique\"  # @param {type:\"string\"}\n",
+        "BQ_TABLE_URI = f\"{PROJECT_ID}.{BQ_DATASET_ID}.{BQ_TABLE_ID}\"\n",
+        "\n",
+        "FEATURE_GROUP_ID = \"fhfv_fg_unique\"  # @param {type:\"string\"}\n",
+        "\n",
+        "DATA_SOURCE = \"gs://cloud-samples-data-us-central1/vertex-ai/feature-store/datasets/movie_prediction.csv\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cd580a0679ce"
+      },
+      "source": [
+        "## Create BigQuery table containing feature data"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "1e2c688b844b"
+      },
+      "source": [
+        "First we'll use BigQuery DataFrames to load in our CSV data source. Then we'll\n",
+        "rename the `timestamp` column to `feature_timestamp` to support usage as a\n",
+        "BigQuery source in Feature Registry."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "1ff4481243a8"
+      },
+      "outputs": [],
+      "source": [
+        "session = bigframes.connect(\n",
+        "    bigframes.BigQueryOptions(\n",
+        "        project=PROJECT_ID,\n",
+        "        location=LOCATION,\n",
+        "    )\n",
+        ")\n",
+        "df = session.read_csv(DATA_SOURCE)\n",
+        "df[\"timestamp\"] = pd.to_datetime(df[\"timestamp\"], utc=True)\n",
+        "df = df.rename(columns={\"timestamp\": \"feature_timestamp\"})"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "de5008f71567"
+      },
+      "source": [
+        "Let's preview the data we'll write to the table."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "38b448c47657"
+      },
+      "outputs": [],
+      "source": [
+        "df.head()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "967ec2a0193f"
+      },
+      "source": [
+        "And finally we'll write the DataFrame to the target BigQuery table."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "4c11b88ab55d"
+      },
+      "outputs": [],
+      "source": [
+        "df.to_gbq(BQ_TABLE_URI, if_exists=\"replace\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "818a36b9da86"
+      },
+      "source": [
+        "## Create feature registry resources"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2fd96d0d8628"
+      },
+      "source": [
+        "Create a feature group backed by the BigQuery table created above."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "1f915ddd4669"
+      },
+      "outputs": [],
+      "source": [
+        "fg: FeatureGroup = FeatureGroup.create(\n",
+        "    f\"{FEATURE_GROUP_ID}\",\n",
+        "    fs_utils.FeatureGroupBigQuerySource(\n",
+        "        uri=f\"bq://{BQ_TABLE_URI}\", entity_id_columns=[\"users\"]\n",
+        "    ),\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ebb9179572f7"
+      },
+      "source": [
+        "Create the `movies` feature which corresponds to the `movies` column in the\n",
+        "recently created BigQuery table."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "0dba1c02883c"
+      },
+      "outputs": [],
+      "source": [
+        "movies_feature: Feature = fg.create_feature(\"movies\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "c71964219962"
+      },
+      "source": [
+        "## Fetch historical feature values"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8eb9b5257006"
+      },
+      "source": [
+        "### Fetch historical feature values for an entity"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "31edf36830b7"
+      },
+      "source": [
+        "The following will fetch historical feature values for the same entity (`alice`)\n",
+        "at two different timestamps. We expect the values of the `movies` feature at\n",
+        "each of those timestamps."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "8fb9774db748"
+      },
+      "outputs": [],
+      "source": [
+        "entity_df = pd.DataFrame(\n",
+        "    data={\n",
+        "        \"users\": [\"alice\", \"alice\"],\n",
+        "        \"timestamp\": [\n",
+        "            pd.Timestamp(\"2021-09-14T09:36\"),\n",
+        "            pd.Timestamp(\"2023-12-12T13:13\"),\n",
+        "        ],\n",
+        "    },\n",
+        ")\n",
+        "\n",
+        "offline_store.fetch_historical_feature_values(\n",
+        "    entity_df=entity_df,\n",
+        "    features=[movies_feature],\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f9ecae3005df"
+      },
+      "source": [
+        "### Fetch with multiple entities"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f463b68f2f4f"
+      },
+      "source": [
+        "The following will fetch historical feature values for two different entities\n",
+        "at different timestamps. We expect the values of the `movies` feature for each\n",
+        "entity at it's corresponding timestamp."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "80ab288afd22"
+      },
+      "outputs": [],
+      "source": [
+        "entity_df = pd.DataFrame(\n",
+        "    data={\n",
+        "        \"users\": [\"alice\", \"bob\"],\n",
+        "        \"timestamp\": [\n",
+        "            pd.Timestamp(\"2021-09-14T09:36\"),\n",
+        "            pd.Timestamp(\"2023-12-12T13:13\"),\n",
+        "        ],\n",
+        "    },\n",
+        ")\n",
+        "\n",
+        "offline_store.fetch_historical_feature_values(\n",
+        "    entity_df=entity_df,\n",
+        "    features=[movies_feature],\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2a4e033321ad"
+      },
+      "source": [
+        "## Cleaning up"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ffd3dc65a25f"
+      },
+      "source": [
+        "### Delete feature and feature group"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7517048d8510"
+      },
+      "outputs": [],
+      "source": [
+        "movies_feature.delete()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "8488287340ca"
+      },
+      "outputs": [],
+      "source": [
+        "fg.delete()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "753e85f60d06"
+      },
+      "source": [
+        "### Delete BigQuery dataset and table"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "99e984fe9c53"
+      },
+      "outputs": [],
+      "source": [
+        "client = bigquery.Client()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "6cc5fdf51c9e"
+      },
+      "outputs": [],
+      "source": [
+        "client.delete_table(f\"{BQ_TABLE_URI}\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "4bac93a9ffcb"
+      },
+      "outputs": [],
+      "source": [
+        "client.delete_dataset(f\"{PROJECT_ID}.{BQ_DATASET_ID}\")"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "offline_feature_serving_from_bigquery_with_feature_registry.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
@@ -29,29 +29,31 @@
        "id": "rbTKH4_6f1ux"
      },
      "source": [
-        "## Online feature serving and fetching of BigQuery data with Vertex AI Feature Store\n",
+        "# Online feature serving and fetching of BigQuery data with Vertex AI Feature Store\n",
        "\n",
        "\n",
        "<table align=\"left\">\n",
-        "\n",
-        "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/feature_store/online_feature_serving_and_fetching_bigquery_data_with_feature_store.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/feature_store/online_feature_serving_and_fetching_bigquery_data_with_feature_store_bigtable.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://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/feature_store/online_feature_serving_and_fetching_bigquery_data_with_feature_store.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%2Ffeature_store%2Fonline_feature_serving_and_fetching_bigquery_data_with_feature_store_bigtable.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/feature_store/online_feature_serving_and_fetching_bigquery_data_with_feature_store_bigtable.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>\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/feature_store/online_feature_serving_and_fetching_bigquery_data_with_feature_store.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://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/feature_store/online_feature_serving_and_fetching_bigquery_data_with_feature_store_bigtable.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>"
+        "  </td>\n",
+        "</table>\n"
      ]
    },
    {
@@ -62,7 +64,7 @@
      "source": [
        "## Overview\n",
        "\n",
-        "This tutorial demonstrates how to use `Vertex AI Feature Store` for online serving and fetching of feature values in `BigQuery`.\n",
+        "This tutorial demonstrates how to use Vertex AI Feature Store for online serving and fetching of feature values in BigQuery.\n",
        "\n",
        "Learn more about [Vertex AI Feature Store](https://cloud.google.com/vertex-ai/docs/featurestore/overview)."
      ]
@@ -75,16 +77,16 @@
      "source": [
        "### Objective\n",
        "\n",
-        "In this tutorial, you learn how to create and use an online feature store instance to host and serve data in `BigQuery` with `Vertex AI Feature Store` in an end to end workflow of feature values serving and fetching user journey.\n",
+        "In this tutorial, you learn how to create and use an online feature store instance to host and serve data in BigQuery with Vertex AI Feature Store in an end to end workflow of feature values serving and fetching user journey.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
-        "- `Vertex AI Feature Store`\n",
+        "- Vertex AI Feature Store\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Provision an online feature store instance to host and serve data.\n",
-        "- Register a `BigQuery` view with the online feature store instance and set up the sync job.\n",
+        "- Register a BigQuery view with the online feature store instance and set up the sync job.\n",
        "- Use the online server to fetch feature values for online prediction."
      ]
    },
@@ -96,7 +98,7 @@
      "source": [
        "### Dataset\n",
        "\n",
-        "This tutorial uses the [theLook_eCommerce](https://pantheon.corp.google.com/marketplace/product/bigquery-public-data/thelook-ecommerce) dataset from the `BigQuery` public datasets.\n",
+        "This tutorial uses the [theLook_eCommerce](https://console.cloud.google.com/marketplace/product/bigquery-public-data/thelook-ecommerce) dataset from the BigQuery public datasets.\n",
        "\n",
        "The Look eCommerce has developed a sales prediction model for their merchandise. The model uses the following data:\n",
        "\n",
@@ -118,8 +120,8 @@
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
-        "* `Vertex AI`\n",
-        "* `BigQuery`\n",
+        "* Vertex AI\n",
+        "* BigQuery\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and\n",
        "[BigQuery pricing](https://cloud.google.com/bigquery/pricing)\n",
@@ -130,23 +132,29 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "yTJiDCrYsOmT"
+        "id": "3b1ffd5ab768"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute 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": "GAqbE5Z2sTVM"
+        "id": "feec187f8410"
      },
      "outputs": [],
      "source": [
-        "# Install the packages\n",
        "! pip3 install --upgrade --quiet google-cloud-aiplatform\\\n",
        "                                 google-cloud-bigquery\\\n",
        "                                 db-dtypes"
@@ -155,180 +163,98 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "np60_uuCs7X5"
+        "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": "_u0aEgaSs-3v"
+        "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": "mdqw6ADTtJRI"
+        "id": "54c5ef8a8f43"
      },
      "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": "bFANidV0tPbo"
+        "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": "t5cutPRQtQ7m"
+        "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": "36840c73a5e4"
      },
      "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": "lfY8yWnbtZ0K"
-      },
-      "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). Note that the new API is currently only available in the following regions:\n",
-        "* `us-centra1`\n",
-        "* `us-east1`\n",
-        "* `us-east4`\n",
-        "* `us-west1`\n",
-        "* `us-west2`\n",
-        "* `us-west3`\n",
-        "* `europe-west2`\n",
-        "* `europe-west3`\n",
-        "* `europe-west4`\n",
-        "* `europe-west8`\n",
-        "* `asia-southeast1`\n",
-        "* `asia-northeast2`"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "G6iMMALZthFM"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-east1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "Ni5jx6RGtzG3"
-      },
-      "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": "UzsMphY2t4-v"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "R_OnJm_Yt8bw"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "x4ybIfzhuAOc"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "gmnRqX6BuBnx"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "uN9JoC1buE9P"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "YNAvMVJjuH5b"
-      },
-      "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)\n",
+        "API_ENDPOINT = f\"{LOCATION}-aiplatform.googleapis.com\""
      ]
    },
    {
@@ -348,7 +274,7 @@
      },
      "outputs": [],
      "source": [
-        "from google.cloud import aiplatform, bigquery\n",
+        "from google.cloud import bigquery\n",
        "from google.cloud.aiplatform_v1 import (FeatureOnlineStoreAdminServiceClient,\n",
        "                                        FeatureOnlineStoreServiceClient,\n",
        "                                        FeatureRegistryServiceClient)\n",
@@ -369,30 +295,6 @@
        "from google.cloud.aiplatform_v1.types import io as io_pb2"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "6ooJNSOvu6Q5"
-      },
-      "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": "mLTm3pquu9ar"
-      },
-      "outputs": [],
-      "source": [
-        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
-        "\n",
-        "API_ENDPOINT = f\"{REGION}-aiplatform.googleapis.com\""
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -499,7 +401,7 @@
      },
      "outputs": [],
      "source": [
-        "# First, create a dataset if it does not already exist. The source data for this demo is located in the US region, so the dataset must also be located in the US region.\n",
+        "# First, create a dataset if it doesn't already exist. The source data for this demo is located in the US region, so the dataset must also be located in the US region.\n",
        "\n",
        "BQ_DATASET_ID = \"featurestore_demo\"  # @param {type:\"string\"}\n",
        "dataset = bigquery.Dataset(f\"{PROJECT_ID}.{BQ_DATASET_ID}\")\n",
@@ -553,7 +455,7 @@
        "    * Create a `FeatureOnlineStore` instance with autoscaling.\n",
        "1. Define the data (`FeatureView`) to be served by the newly-created instance. This can either map to\n",
        "    * The BigQuery view that you just created for serving data.\n",
-        "    * The `FeatureGroup` and `Feature` we will create to host feature metadata.\n",
+        "    * The `FeatureGroup` and `Feature` you'll create to host feature metadata.\n",
        "\n",
        "Bigtable serving latency is affected by the (Bigtable) load. However, when Bigtable is not overloaded, benchmarks show that the expected server-side latency is around 30 ms, measured at around 100 qps. The client-side latency is expected to be more than 5 ms higher than the server-side latency."
      ]
@@ -615,7 +517,7 @@
        "\n",
        "create_store_lro = admin_client.create_feature_online_store(\n",
        "    feature_online_store_admin_service_pb2.CreateFeatureOnlineStoreRequest(\n",
-        "        parent=f\"projects/{PROJECT_ID}/locations/{REGION}\",\n",
+        "        parent=f\"projects/{PROJECT_ID}/locations/{LOCATION}\",\n",
        "        feature_online_store_id=FEATURE_ONLINE_STORE_ID,\n",
        "        feature_online_store=online_store_config,\n",
        "    )\n",
@@ -666,7 +568,7 @@
      "source": [
        "# Use list to verify the store is created.\n",
        "admin_client.list_feature_online_stores(\n",
-        "    parent=f\"projects/{PROJECT_ID}/locations/{REGION}\"\n",
+        "    parent=f\"projects/{PROJECT_ID}/locations/{LOCATION}\"\n",
        ")"
      ]
    },
@@ -701,7 +603,7 @@
      "source": [
        "##### Data source preparation guidelines for BigQuery data source\n",
        "\n",
-        "If the feature data for online serving is not already available in a BigQuery table or view, you can create a BigQuery dataset and table, and then load the data into it.\n",
+        "If the feature data for online serving isn't already available in a BigQuery table or view, you can create a BigQuery dataset and table, and then load the data into it.\n",
        "\n",
        "Note that if you choose to use BigQuery source, Feature Store only provides the option to snapshot and sync the entire BigQuery table or view for online serving. You need to construct this table or view to reflect the latest data to be served. Therefore, a timestamp column is not needed, since timestamps aren't used to differentiate the feature values.\n",
        "\n",
@@ -736,7 +638,7 @@
        "\n",
        "create_view_lro = admin_client.create_feature_view(\n",
        "    feature_online_store_admin_service_pb2.CreateFeatureViewRequest(\n",
-        "        parent=f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}\",\n",
+        "        parent=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}\",\n",
        "        feature_view_id=FEATURE_VIEW_ID,\n",
        "        feature_view=feature_view_pb2.FeatureView(\n",
        "            big_query_source=big_query_source,\n",
@@ -768,7 +670,7 @@
      "source": [
        "#### [Optional] Create FeatureGroup/Features\n",
        "\n",
-        "Create a FeatureGroup pointing to the created BigQuery view for the demo. We will then create features for each column we would like to register.\n"
+        "Create a FeatureGroup pointing to the created BigQuery view for the demo. You then create features for each column you'd like to register.\n"
      ]
    },
    {
@@ -779,7 +681,7 @@
      "source": [
        "##### Data source preparation guidelines for Feature Registry data source\n",
        "\n",
-        "Note that if you choose to use Feature Registry source, Feature Store only provides the option to support time-series sources for which Feature Store will generate latest featureValues.\n",
+        "Note that if you choose to use Feature Registry source, Feature Store only provides the option to support time-series sources for which Feature Store generates latest featureValues.\n",
        "\n",
        "Use the following guidelines to understand the schema and constraints while creating the BigQuery source:\n",
        "\n",
@@ -829,7 +731,7 @@
        "\n",
        "create_group_lro = registry_client.create_feature_group(\n",
        "    feature_registry_service_pb2.CreateFeatureGroupRequest(\n",
-        "        parent=f\"projects/{PROJECT_ID}/locations/{REGION}\",\n",
+        "        parent=f\"projects/{PROJECT_ID}/locations/{LOCATION}\",\n",
        "        feature_group_id=FEATURE_GROUP_ID,\n",
        "        feature_group=feature_group_config,\n",
        "    )\n",
@@ -869,7 +771,7 @@
        "    create_feature_lros.append(\n",
        "        registry_client.create_feature(\n",
        "            featurestore_service_pb2.CreateFeatureRequest(\n",
-        "                parent=f\"projects/{PROJECT_ID}/locations/{REGION}/featureGroups/{FEATURE_GROUP_ID}\",\n",
+        "                parent=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureGroups/{FEATURE_GROUP_ID}\",\n",
        "                feature_id=id,\n",
        "                feature=feature_pb2.Feature(),\n",
        "            )\n",
@@ -921,7 +823,7 @@
        "\n",
        "create_view_lro = admin_client.create_feature_view(\n",
        "    feature_online_store_admin_service_pb2.CreateFeatureViewRequest(\n",
-        "        parent=f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}\",\n",
+        "        parent=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}\",\n",
        "        feature_view_id=FEATURE_VIEW_ID,\n",
        "        feature_view=feature_view_pb2.FeatureView(\n",
        "            feature_registry_source=feature_registry_source,\n",
@@ -953,7 +855,7 @@
      "source": [
        "# Again, list all feature view under the FEATURE_ONLINE_STORE_ID to confirm\n",
        "admin_client.list_feature_views(\n",
-        "    parent=f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}\"\n",
+        "    parent=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}\"\n",
        ")"
      ]
    },
@@ -989,7 +891,7 @@
      "outputs": [],
      "source": [
        "sync_response = admin_client.sync_feature_view(\n",
-        "    feature_view=f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\"\n",
+        "    feature_view=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\"\n",
        ")"
      ]
    },
@@ -1045,7 +947,7 @@
      "outputs": [],
      "source": [
        "admin_client.list_feature_view_syncs(\n",
-        "    parent=f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\"\n",
+        "    parent=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\"\n",
        ")"
      ]
    },
@@ -1092,7 +994,7 @@
      "source": [
        "data_client.fetch_feature_values(\n",
        "    request=feature_online_store_service_pb2.FetchFeatureValuesRequest(\n",
-        "        feature_view=f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\",\n",
+        "        feature_view=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\",\n",
        "        data_key=feature_online_store_service_pb2.FeatureViewDataKey(key=\"28098\"),\n",
        "    )\n",
        ")"
@@ -1117,7 +1019,7 @@
      "source": [
        "data_client.fetch_feature_values(\n",
        "    request=feature_online_store_service_pb2.FetchFeatureValuesRequest(\n",
-        "        feature_view=f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\",\n",
+        "        feature_view=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\",\n",
        "        data_key=feature_online_store_service_pb2.FeatureViewDataKey(key=\"28098\"),\n",
        "        data_format=feature_online_store_service_pb2.FeatureViewDataFormat.PROTO_STRUCT,\n",
        "    )\n",
@@ -1149,22 +1051,22 @@
        "# Delete Features\n",
        "for id in FEATURE_IDS:\n",
        "    registry_client.delete_feature(\n",
-        "        name=f\"projects/{PROJECT_ID}/locations/{REGION}/featureGroups/{FEATURE_GROUP_ID}/features/{id}\"\n",
+        "        name=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureGroups/{FEATURE_GROUP_ID}/features/{id}\"\n",
        "    )\n",
        "\n",
        "# Delete Featuregroup\n",
        "registry_client.delete_feature_group(\n",
-        "    name=f\"projects/{PROJECT_ID}/locations/{REGION}/featureGroups/{FEATURE_GROUP_ID}\"\n",
+        "    name=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureGroups/{FEATURE_GROUP_ID}\"\n",
        ")\n",
        "\n",
        "# Delete FeatureViews\n",
        "admin_client.delete_feature_view(\n",
-        "    name=f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\"\n",
+        "    name=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\"\n",
        ")\n",
        "\n",
        "# Delete OnlineStore\n",
        "admin_client.delete_feature_online_store(\n",
-        "    name=f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}\",\n",
+        "    name=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}\",\n",
        "    force=True,\n",
        ")\n",
        "\n",
@@ -63,7 +63,7 @@
      "source": [
        "## Overview\n",
        "\n",
-        "This tutorial demonstrates how to use Optimized online serving in `Vertex AI Feature Store` for serving and fetching of feature values from `BigQuery`.\n",
+        "This tutorial demonstrates how to use Optimized online serving in Vertex AI Feature Store for serving and fetching of feature values from BigQuery.\n",
        "\n",
        "Learn more about [Vertex AI Feature Store](https://cloud.google.com/vertex-ai/docs/featurestore/latest/overview)."
      ]
@@ -76,18 +76,18 @@
      "source": [
        "### Objective\n",
        "\n",
-        "In this tutorial, you learn how to create and use an online feature store instance to host and serve data in `BigQuery` with `Vertex AI Feature Store` in an end-to-end workflow of serving and fetching feature values.\n",
+        "In this tutorial, you learn how to create and use an online feature store instance to host and serve data in BigQuery with Vertex AI Feature Store in an end-to-end workflow of serving and fetching feature values.\n",
        "\n",
        "You serve features using Optimized online serving, which lets you serve features at ultra-low latencies. It's ideal for online serving scenarios where super-low latency is critical. Optimized online serving provides an online serving infrastructure that's faster, more scalable, and more responsive to increased data loads than Cloud Bigtable online serving.\n",
        "\n",
        "This tutorial uses the following Google Cloud Vertex AI services and resources:\n",
        "\n",
-        "- `Vertex AI Feature Store`\n",
+        "- Vertex AI Feature Store\n",
        "\n",
        "The steps performed include the following:\n",
        "\n",
        "- Provision an online feature store instance to host and serve data using Optimized online serving with Public or Private endpoint.\n",
-        "- Register a `BigQuery` view with the online feature store instance and set up the sync job.\n",
+        "- Register a BigQuery view with the online feature store instance and set up the sync job.\n",
        "- Use the online server to fetch feature values for online prediction."
      ]
    },
@@ -109,7 +109,7 @@
      "source": [
        "### Dataset\n",
        "\n",
-        "This tutorial uses the [theLook_eCommerce](https://pantheon.corp.google.com/marketplace/product/bigquery-public-data/thelook-ecommerce) dataset from the `BigQuery` public datasets.\n",
+        "This tutorial uses the [theLook_eCommerce](https://console.cloud.google.com/marketplace/product/bigquery-public-data/thelook-ecommerce) dataset from the BigQuery public datasets.\n",
        "\n",
        "The Look eCommerce has developed a sales prediction model for their merchandise. The model uses the following data:\n",
        "\n",
@@ -131,8 +131,8 @@
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
-        "* `Vertex AI`\n",
-        "* `BigQuery`\n",
+        "* Vertex AI\n",
+        "* BigQuery\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and\n",
        "[BigQuery pricing](https://cloud.google.com/bigquery/pricing)\n",
@@ -41,7 +41,7 @@
        "    </a>\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%2Ffeature_store%2Fonline_feature_serving_and_vector_retrieval_bigquery_data_with_feature_store.ipynb\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Ffeature_store%2Fonline_feature_serving_and_vector_retrieval_bigquery_data_with_feature_store.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",
@@ -66,7 +66,7 @@
      "source": [
        "## Overview\n",
        "\n",
-        "This tutorial demonstrates how to use `Vertex AI Feature Store` for online serving and vector retrieval of feature values in `BigQuery`.\n",
+        "This tutorial demonstrates how to use Vertex AI Feature Store for online serving and vector retrieval of feature values in BigQuery.\n",
        "\n",
        "Learn more about [Vertex AI Feature Store](https://cloud.google.com/vertex-ai/docs/featurestore/overview)."
      ]
@@ -79,16 +79,16 @@
      "source": [
        "### Objective\n",
        "\n",
-        "In this tutorial, you learn how to create and use an online feature store instance to host and serve data in `BigQuery` with `Vertex AI Feature Store` in an end to end workflow of features serving and vector retrieval user journey.\n",
+        "In this tutorial, you learn how to create and use an online feature store instance to host and serve data in BigQuery with Vertex AI Feature Store in an end to end workflow of features serving and vector retrieval user journey.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
-        "- `Vertex AI Feature Store`\n",
+        "- Vertex AI Feature Store\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Provision an online feature store instance to host and serve data.\n",
-        "- Create an online feature store instance to serve a `BigQuery` table.\n",
+        "- Create an online feature store instance to serve a BigQuery table.\n",
        "- Use the online server to search nearest neighbors."
      ]
    },
@@ -110,7 +110,7 @@
      "source": [
        "### Dataset\n",
        "\n",
-        "This tutorial uses the [Google Patents Public Data](https://pantheon.corp.google.com/marketplace/product/google_patents_public_datasets/google-patents-public-data) dataset from the `BigQuery` public datasets.\n"
+        "This tutorial uses the [Google Patents Public Data](https://console.cloud.google.com/marketplace/product/google_patents_public_datasets/google-patents-public-data) dataset from the BigQuery public datasets.\n"
      ]
    },
    {
@@ -123,8 +123,8 @@
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
-        "* `Vertex AI`\n",
-        "* `BigQuery`\n",
+        "* Vertex AI\n",
+        "* BigQuery\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and\n",
        "[BigQuery pricing](https://cloud.google.com/bigquery/pricing)\n",
@@ -240,7 +240,7 @@
      "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)."
+        "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)."
      ]
    },
    {
@@ -339,7 +339,7 @@
        "id": "ZVec1FjhpwP6"
      },
      "source": [
-        "This data source has some filtering columns (e.g. country) and crowding columns (e.g. cited_by_filing_date). Below is the schema of publications data table we use in this guide:\n",
+        "This data source has some filtering columns (e.g. country) and crowding columns (for example, cited_by_filing_date). Below is the schema of publications data table we use in this guide:\n",
        "\n",
        "|Column name |  Type  |   Mode   |\n",
        "|------------|--------|----------|\n",
@@ -422,7 +422,7 @@
      },
      "outputs": [],
      "source": [
-        "# First, create a dataset to keep the feature store source data if it does not already exist.\n",
+        "# First, create a dataset to keep the feature store source data if it doesn't already exist.\n",
        "BQ_DATASET_ID = \"featurestore_demo_us\"  # @param {type:\"string\"}\n",
        "create_bq_dataset(BQ_DATASET_ID, \"US\")"
      ]
@@ -684,7 +684,7 @@
      "source": [
        "### Verify feature view instance creation\n",
        "\n",
-        "Verify that the FeatureView instance is created by getting the feature view."
+        "Verify that the `FeatureView` instance is created by getting the feature view."
      ]
    },
    {
@@ -706,7 +706,7 @@
        "id": "sqpu4nHAO4pW"
      },
      "source": [
-        "Verify that the FeatureView instance is created by listing all the feature views within the online store."
+        "Verify that the `FeatureView` instance is created by listing all the feature views within the online store."
      ]
    },
    {
@@ -29,27 +29,32 @@
        "id": "awoLZ5dc5bcG"
      },
      "source": [
-        "# Vertex AI Feature Store Based LLM Grounding Tutorial\n",
+        "# Vertex AI Feature Store Based LLM Grounding tutorial\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/feature_store/vertex_ai_feature_store_based_llm_grounding_tutorial.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/feature_store/vertex_ai_feature_store_based_llm_grounding_tutorial.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%2Ffeature_store%2Fvertex_ai_feature_store_based_llm_grounding_tutorial.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/feature_store/vertex_ai_feature_store_based_llm_grounding_tutorial.ipynb.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/feature_store/vertex_ai_feature_store_based_llm_grounding_tutorial.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>"
      ]
    },
@@ -74,16 +79,16 @@
      "source": [
        "### Objective\n",
        "\n",
-        "In this tutorial, you learn how to create and use an online feature store instance to host and serve data in `BigQuery` with `Vertex AI Feature Store` in an end to end workflow of features serving and vector retrieval user journey.\n",
+        "In this tutorial, you learn how to create and use an online feature store instance to host and serve data in BigQuery with Vertex AI Feature Store in an end to end workflow of features serving and vector retrieval user journey.\n",
        "\n",
-        "This tutorial uses the following Google Cloud ML services and resources:\n",
+        "This tutorial uses the following Google Vertex AI services and resources:\n",
        "\n",
-        "- `Vertex AI Feature Store`\n",
+        "- Vertex AI Feature Store\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Provision an online feature store instance to host and serve data.\n",
-        "- Create an online feature store instance to serve a `BigQuery` table.\n",
+        "- Create an online feature store instance to serve a BigQuery table.\n",
        "- Use the online server to search nearest neighbors."
      ]
    },
@@ -105,7 +110,7 @@
      "source": [
        "### Dataset\n",
        "\n",
-        "This tutorial uses the [Google Patents Public Data](https://console.cloud.google.com/marketplace/product/google_patents_public_datasets/google-patents-public-data) dataset from the `BigQuery` public datasets.\n"
+        "This tutorial uses the [Google Patents Public Data](https://console.cloud.google.com/marketplace/product/google_patents_public_datasets/google-patents-public-data) dataset from the BigQuery public datasets.\n"
      ]
    },
    {
@@ -118,9 +123,9 @@
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
-        "* `Vertex AI`\n",
-        "* `BigQuery`\n",
-        "* `Cloud Storage`\n",
+        "* Vertex AI\n",
+        "* BigQuery\n",
+        "* Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing),\n",
        "[BigQuery pricing](https://cloud.google.com/bigquery/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing),\n",
@@ -128,15 +133,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": "yTJiDCrYsOmT"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to run this notebook."
+        "### Install Vertex AI SDK for Python and other required packages"
      ]
    },
    {
@@ -161,7 +173,9 @@
        "id": "np60_uuCs7X5"
      },
      "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."
      ]
    },
    {
@@ -172,46 +186,62 @@
      },
      "outputs": [],
      "source": [
-        "# # Automatically restart the kernel after installation 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": "mdqw6ADTtJRI"
+        "id": "01e1cc61b578"
      },
      "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're running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "bFANidV0tPbo"
+        "id": "ff666ce4051c"
      },
      "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": "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": "60763ee24ce0"
+      },
+      "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)."
      ]
    },
    {
@@ -223,108 +253,14 @@
      "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": "lfY8yWnbtZ0K"
-      },
-      "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). Note that the new Feature Store capability showed in the colab is currently only available in the following regions:\n",
-        "* `us-central1`\n",
-        "* `us-east1`\n",
-        "* `us-west1`\n",
-        "* `europe-west4`\n",
-        "* `asia-southeast1`"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "G6iMMALZthFM"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "Ni5jx6RGtzG3"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
+        "from google.cloud import aiplatform\n",
        "\n",
-        "Depending on your Jupyter environment, you might have to manually authenticate. Follow these instructions:"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "UzsMphY2t4-v"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "R_OnJm_Yt8bw"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "x4ybIfzhuAOc"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "gmnRqX6BuBnx"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "uN9JoC1buE9P"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "YNAvMVJjuH5b"
-      },
-      "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)\n",
+        "\n",
+        "API_ENDPOINT = f\"{LOCATION}-aiplatform.googleapis.com\""
      ]
    },
    {
@@ -346,7 +282,7 @@
      "source": [
        "import uuid\n",
        "\n",
-        "from google.cloud import aiplatform, bigquery\n",
+        "from google.cloud import bigquery\n",
        "from google.cloud.aiplatform_v1 import (FeatureOnlineStoreAdminServiceClient,\n",
        "                                        FeatureOnlineStoreServiceClient)\n",
        "from google.cloud.aiplatform_v1.types import NearestNeighborQuery\n",
@@ -356,32 +292,7 @@
        "    feature_online_store_admin_service as \\\n",
        "    feature_online_store_admin_service_pb2\n",
        "from google.cloud.aiplatform_v1.types import \\\n",
-        "    feature_online_store_service as feature_online_store_service_pb2\n",
-        "from google.cloud.aiplatform_v1.types import feature_view as feature_view_pb2"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "6ooJNSOvu6Q5"
-      },
-      "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": "mLTm3pquu9ar"
-      },
-      "outputs": [],
-      "source": [
-        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
-        "\n",
-        "API_ENDPOINT = f\"{REGION}-aiplatform.googleapis.com\""
+        "    feature_online_store_service as feature_online_store_service_pb2"
      ]
    },
    {
@@ -425,7 +336,7 @@
        "id": "_OV9dADJb63o"
      },
      "source": [
-        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket.\n"
+        "**If your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket.\n"
      ]
    },
    {
@@ -436,7 +347,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} -p {PROJECT_ID} {GCS_BUCKET}"
+        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {GCS_BUCKET}"
      ]
    },
    {
@@ -504,7 +415,7 @@
        "\n",
        "BQ_DATASET_ID = \"fs_grounding\"  # @param {type:\"string\"}\n",
        "dataset = bigquery.Dataset(f\"{PROJECT_ID}.{BQ_DATASET_ID}\")\n",
-        "dataset.location = REGION\n",
+        "dataset.location = LOCATION\n",
        "dataset = bq_client.create_dataset(\n",
        "    dataset, exists_ok=True, timeout=30\n",
        ")  # Make an API request.\n",
@@ -538,7 +449,7 @@
        "\n",
        "PARAMS = {\n",
        "    \"project\": PROJECT_ID,\n",
-        "    \"location\": REGION,\n",
+        "    \"location\": LOCATION,\n",
        "    \"bigquery_bp_input_uri\": BIGQUERY_BP_INPUT_URI,\n",
        "    \"bigquery_bp_output_uri\": BIGQUERY_BP_OUTPUT_URI,\n",
        "    \"input_text_gcs_dir\": INPUT_TEXT_GCS_DIR,\n",
@@ -582,7 +493,7 @@
        "    parameters=PARAMS,\n",
        "    project=PROJECT_ID,\n",
        "    pipeline_root=f\"{GCS_BUCKET}/fs_based/pipeline_root\",\n",
-        "    location=REGION,\n",
+        "    location=LOCATION,\n",
        ")\n",
        "job.wait()"
      ]
@@ -611,7 +522,7 @@
        "\n",
        "    Args:\n",
        "      bigquery_bp_input_uri: The URI to a bigquery table as the input for the\n",
-        "        batch prediction component. The chunking component will populate data to\n",
+        "        batch prediction component. The chunking component populates data to\n",
        "        this uri first before batch prediction.\n",
        "      bigquery_bp_output_uri: The URI to a bigquery table as the output for the\n",
        "        batch prediction component.\n",
@@ -711,7 +622,7 @@
        "\n",
        "create_store_lro = admin_client.create_feature_online_store(\n",
        "    feature_online_store_admin_service_pb2.CreateFeatureOnlineStoreRequest(\n",
-        "        parent=f\"projects/{PROJECT_ID}/locations/{REGION}\",\n",
+        "        parent=f\"projects/{PROJECT_ID}/locations/{LOCATION}\",\n",
        "        feature_online_store_id=FEATURE_ONLINE_STORE_ID,\n",
        "        feature_online_store=online_store_config,\n",
        "    )\n",
@@ -762,7 +673,7 @@
      "source": [
        "# Use get to verify the store is created.\n",
        "admin_client.get_feature_online_store(\n",
-        "    name=f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}\"\n",
+        "    name=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}\"\n",
        ")"
      ]
    },
@@ -785,7 +696,7 @@
      "source": [
        "# Use list to verify the store is created.\n",
        "admin_client.list_feature_online_stores(\n",
-        "    parent=f\"projects/{PROJECT_ID}/locations/{REGION}\"\n",
+        "    parent=f\"projects/{PROJECT_ID}/locations/{LOCATION}\"\n",
        ")"
      ]
    },
@@ -802,7 +713,7 @@
        "* A data source (BigQuery table or view URI or `FeatureGroup/features`) synced to the `FeatureOnlineStore` instance for serving.\n",
        "* The [cron](https://en.wikipedia.org/wiki/Cron) schedule to run the sync pipeline.\n",
        "\n",
-        "During feature view creation, a sync job will be scheduled, and either started immediately or following the cron schedule. In the sync job, data is exported, a index is built and deployed to GKE cluster."
+        "During feature view creation, a sync job is scheduled, and either started immediately or following the cron schedule. In the sync job, data is exported, an index is built and deployed to GKE cluster."
      ]
    },
    {
@@ -814,7 +725,7 @@
      "outputs": [],
      "source": [
        "FEATURE_VIEW_ID = \"fs_grounding_test_new\"  # @param {type: \"string\"}\n",
-        "# A schedule will be created based on cron setting.\n",
+        "# A schedule is created based on cron setting.\n",
        "# If cron is unspecified, a sync job is started immediately.\n",
        "CRON_SCHEDULE = \"TZ=America/Los_Angeles 00 13 11 8 *\"  # @param {type: \"string\"}"
      ]
@@ -857,6 +768,8 @@
      },
      "outputs": [],
      "source": [
+        "from google.cloud.aiplatform_v1.types import feature_view as feature_view_pb2\n",
+        "\n",
        "big_query_source = feature_view_pb2.FeatureView.BigQuerySource(\n",
        "    uri=DATA_SOURCE, entity_id_columns=[\"vertex_generated_chunk_id\"]\n",
        ")\n",
@@ -875,7 +788,7 @@
        "\n",
        "create_view_lro = admin_client.create_feature_view(\n",
        "    feature_online_store_admin_service_pb2.CreateFeatureViewRequest(\n",
-        "        parent=f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}\",\n",
+        "        parent=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}\",\n",
        "        feature_view_id=FEATURE_VIEW_ID,\n",
        "        feature_view=feature_view_pb2.FeatureView(\n",
        "            big_query_source=big_query_source,\n",
@@ -926,7 +839,7 @@
      "outputs": [],
      "source": [
        "admin_client.get_feature_view(\n",
-        "    name=f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\"\n",
+        "    name=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\"\n",
        ")"
      ]
    },
@@ -936,7 +849,7 @@
        "id": "sqpu4nHAO4pW"
      },
      "source": [
-        "Verify that the FeatureView instance is created by listing all the feature views within the online store."
+        "Verify that the `FeatureView` instance is created by listing all the feature views within the online store."
      ]
    },
    {
@@ -948,7 +861,7 @@
      "outputs": [],
      "source": [
        "admin_client.list_feature_views(\n",
-        "    parent=f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}\"\n",
+        "    parent=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}\"\n",
        ")"
      ]
    },
@@ -962,7 +875,7 @@
      "source": [
        "# Optional: Delete feature views to avoid exceeding the deployed index nodes quota.\n",
        "# views = admin_client.list_feature_views(\n",
-        "#     parent=f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}\"\n",
+        "#     parent=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}\"\n",
        "# )\n",
        "# for view in views:\n",
        "#     admin_client.delete_feature_view(name=view.name)"
@@ -990,7 +903,7 @@
      "outputs": [],
      "source": [
        "sync_response = admin_client.sync_feature_view(\n",
-        "    feature_view=f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\"\n",
+        "    feature_view=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\"\n",
        ")"
      ]
    },
@@ -1048,7 +961,7 @@
      "outputs": [],
      "source": [
        "admin_client.list_feature_view_syncs(\n",
-        "    parent=f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\"\n",
+        "    parent=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\"\n",
        ")"
      ]
    },
@@ -1075,7 +988,7 @@
      "source": [
        "# Verify online store creation.\n",
        "featore_online_store_instance = admin_client.get_feature_online_store(\n",
-        "    name=f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}\"\n",
+        "    name=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}\"\n",
        ")\n",
        "PUBLIC_ENDPOINT = (\n",
        "    featore_online_store_instance.dedicated_serving_endpoint.public_endpoint_domain_name\n",
@@ -1147,7 +1060,7 @@
        "# A vertex_generated_chunk_id for testing\n",
        "data_client.search_nearest_entities(\n",
        "    request=feature_online_store_service_pb2.SearchNearestEntitiesRequest(\n",
-        "        feature_view=f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\",\n",
+        "        feature_view=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\",\n",
        "        query=NearestNeighborQuery(\n",
        "            entity_id=ENTITY_ID,\n",
        "            neighbor_count=5,\n",
@@ -1187,7 +1100,7 @@
      "source": [
        "data_client.search_nearest_entities(\n",
        "    request=feature_online_store_service_pb2.SearchNearestEntitiesRequest(\n",
-        "        feature_view=f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\",\n",
+        "        feature_view=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\",\n",
        "        query=NearestNeighborQuery(\n",
        "            embedding=NearestNeighborQuery.Embedding(value=EMBEDDINGS),\n",
        "            neighbor_count=10,\n",
@@ -1216,7 +1129,7 @@
      "source": [
        "data_client.fetch_feature_values(\n",
        "    request=feature_online_store_service_pb2.FetchFeatureValuesRequest(\n",
-        "        feature_view=f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\",\n",
+        "        feature_view=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\",\n",
        "        data_key=feature_online_store_service_pb2.FeatureViewDataKey(key=ENTITY_ID),\n",
        "    )\n",
        ")"
@@ -1246,20 +1159,18 @@
      "source": [
        "# Delete Feature View\n",
        "admin_client.delete_feature_view(\n",
-        "    name=f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\"\n",
+        "    name=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\"\n",
        ")\n",
        "\n",
        "# Delete Feature Online Store\n",
        "admin_client.delete_feature_online_store(\n",
-        "    name=f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}\",\n",
+        "    name=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}\",\n",
        "    force=True,\n",
        ")\n",
        "\n",
        "# Delete Cloud Storage objects that were created\n",
-        "import os\n",
-        "\n",
-        "delete_bucket = False\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "delete_bucket = True\n",
+        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $GCS_BUCKET"
      ]
    }
@@ -89,8 +89,8 @@
        "\n",
        "### Costs\n",
        "This tutorial uses billable components of Google Cloud:\n",
-        "* `Vertex AI`\n",
-        "* `BigQuery`\n",
+        "* Vertex AI\n",
+        "* BigQuery\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/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."
      ]
@@ -0,0 +1,708 @@
+{
+  "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": [
+        "# Configure IAM Policy in Vertex AI Feature Store\n",
+        "\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/feature_store/vertex_ai_feature_store_iam_policy.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%2Fblob%2Fmain%2Fnotebooks%2Fofficial%2Ffeature_store%2Fvertex_ai_feature_store_iam_policy.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/tree/main/notebooks/official/feature_store/vertex_ai_feature_store_iam_policy.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/tree/main/notebooks/official/feature_store/vertex_ai_feature_store_iam_policy.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 tutorial demonstrates how to configure an IAM Policy to control access to resources and data stored within Vertex AI Feature Store.\n",
+        "\n",
+        "Learn more about [Vertex AI Feature Store](https://cloud.google.com/vertex-ai/docs/featurestore/overview).\n",
+        "\n",
+        "The steps performed include:\n",
+        "- Create a BigQuery dataset and set up a `FeatureView` to run periodic sync jobs.\n",
+        "- Provision an online store instance and set up online serving.\n",
+        "- Configure an IAM Policy to manage access controls."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "61RBz8LLbxCR"
+      },
+      "source": [
+        "## Getting Started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "No17Cw5hgx12"
+      },
+      "source": [
+        "### Install Vertex AI SDK 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-bigquery"
+      ]
+    },
+    {
+      "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 and initialize Vertex AI SDK\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",
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
+        "\n",
+        "\n",
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(project=PROJECT_ID, location=LOCATION)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "EdvJRUWRNGHE"
+      },
+      "source": [
+        "## Configure IAM Policy in Vertex AI Feature Store"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "GKnDEUulaLLy"
+      },
+      "outputs": [],
+      "source": [
+        "LOCATION_PARENT = \"projects/\" + PROJECT_ID + \"/locations/\" + LOCATION"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6vbsjb91aJ1V"
+      },
+      "source": [
+        "Configure the necessary service clients:\n",
+        "- [FeaturestoreServiceClient](https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform_v1.services.featurestore_service.FeaturestoreServiceClient)\n",
+        "- [FeatureOnlineStoreAdminServiceClient](https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform_v1beta1.services.feature_online_store_admin_service.FeatureOnlineStoreAdminServiceClient)\n",
+        "- [FeatureRegistryServiceClient](https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform_v1beta1.services.feature_registry_service.FeatureRegistryServiceClient)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "zL8aOMAtaEUg"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud.aiplatform_v1beta1 import (\n",
+        "    FeatureOnlineStoreAdminServiceClient, FeatureOnlineStoreServiceClient,\n",
+        "    FeatureRegistryServiceClient, FeaturestoreServiceClient)\n",
+        "from google.cloud.aiplatform_v1beta1.types import \\\n",
+        "    feature_online_store_service as feature_online_store_service_pb2\n",
+        "from google.iam.v1 import iam_policy_pb2, policy_pb2\n",
+        "\n",
+        "ENDPOINT = LOCATION + \"-aiplatform.googleapis.com\"\n",
+        "featurestore_service_client = FeaturestoreServiceClient(\n",
+        "    client_options={\"api_endpoint\": ENDPOINT}\n",
+        ")\n",
+        "feature_online_store_admin_service_client = FeatureOnlineStoreAdminServiceClient(\n",
+        "    client_options={\"api_endpoint\": ENDPOINT}\n",
+        ")\n",
+        "feature_registry_service_client = FeatureRegistryServiceClient(\n",
+        "    client_options={\"api_endpoint\": ENDPOINT}\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "NkxyUncIaCca"
+      },
+      "source": [
+        "Generate sample data for this colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "XWAHq1fcZKsJ"
+      },
+      "outputs": [],
+      "source": [
+        "DATASET_ID = \"test_data\"+\"_\"+LOCATION.replace('-', '_')  # @param {type:\"string\"}\n",
+        "TABLE_ID = \"tableA\"  # @param {type:\"string\"}\n",
+        "\n",
+        "!bq mk --dataset_id={DATASET_ID} --location={LOCATION}\n",
+        "!bq query --nouse_legacy_sql \\\n",
+        "\"CREATE TABLE {DATASET_ID}.{TABLE_ID} AS (\" \\\n",
+        "\"SELECT * FROM UNNEST(ARRAY<STRUCT<entity_id STRING, feature_timestamp TIMESTAMP, feature1 INT64, feature2 INT64>>[\" \\\n",
+        "\"('test', TIMESTAMP('2024-02-26 08:00:00 UTC'), 10, 20),\" \\\n",
+        "\"('test', TIMESTAMP('2024-02-27 08:00:00 UTC'), 30, 40),\" \\\n",
+        "\"('test', TIMESTAMP('2024-02-28 08:00:00 UTC'), 50, 60)]))\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ZCS7_Zr4ZL6y"
+      },
+      "source": [
+        "### Create `FeatureOnlineStore` and `FeatureView` resources"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3KH56EmWZJnp"
+      },
+      "source": [
+        "Provision a `FeatureOnlineStore` instance."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "jrRqB3VpZF3C"
+      },
+      "outputs": [],
+      "source": [
+        "FEATURE_ONLINE_STORE_ID = \"test_feature_online_store\"\n",
+        "\n",
+        "feature_online_store_admin_service_client.create_feature_online_store(\n",
+        "    {\n",
+        "        \"parent\": LOCATION_PARENT,\n",
+        "        \"feature_online_store_id\": FEATURE_ONLINE_STORE_ID,\n",
+        "        \"feature_online_store\": {\n",
+        "            \"bigtable\": {\n",
+        "                \"auto_scaling\": {\n",
+        "                    \"min_node_count\": 1,\n",
+        "                    \"max_node_count\": 3,\n",
+        "                    \"cpu_utilization_target\": 50,\n",
+        "                }\n",
+        "            }\n",
+        "        },\n",
+        "    }\n",
+        ").result()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "QK6Q3tNHZFa8"
+      },
+      "source": [
+        "Configure periodic data synchronization from the offline store in BigQuery to the `FeatureOnlineStore` instance."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "kKl7vb8fZDyu"
+      },
+      "outputs": [],
+      "source": [
+        "FEATURE_VIEW_ID = \"test_feature_view\"\n",
+        "FEATURE_ONLINE_STORE_RESOURCE_ID = (\n",
+        "    LOCATION_PARENT + \"/featureOnlineStores/\" + FEATURE_ONLINE_STORE_ID\n",
+        ")\n",
+        "FEATURE_VIEW_RESOURCE_ID = (\n",
+        "    FEATURE_ONLINE_STORE_RESOURCE_ID + \"/featureViews/\" + FEATURE_VIEW_ID\n",
+        ")\n",
+        "create_feature_view_result = (\n",
+        "    feature_online_store_admin_service_client.create_feature_view(\n",
+        "        parent=FEATURE_ONLINE_STORE_RESOURCE_ID,\n",
+        "        feature_view_id=FEATURE_VIEW_ID,\n",
+        "        feature_view={\n",
+        "            \"big_query_source\": {\n",
+        "                \"uri\": f\"bq://{PROJECT_ID}.{DATASET_ID}.{TABLE_ID}\",\n",
+        "                \"entity_id_columns\": [\"entity_id\"],\n",
+        "            },\n",
+        "            \"sync_config\": {\"cron\": \"0 12 * * *\"},\n",
+        "        },\n",
+        "    ).result()\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "BU1_rOgVZpNj"
+      },
+      "source": [
+        "## Configure an IAM Policy"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "79enQT97ZBBK"
+      },
+      "source": [
+        "Configure IAM Policy for a feature online store."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ST1M5no-ZAgV"
+      },
+      "outputs": [],
+      "source": [
+        "# A specific user email or a group email can be used to add members to the IAM Policy.\n",
+        "EXAMPLE_RESOURCE_VIEWER_EMAIL = (\n",
+        "    \"google-cloud-eng-fte@google.com\"  # @param {type:\"string\"}\n",
+        ")\n",
+        "EXAMPLE_DATA_VIEWER_EMAIL = \"python-sample-owners@google.com\"  # @param {type:\"string\"}\n",
+        "# A Service Account can also be added to the IAM Policy.\n",
+        "EXAMPLE_ADMIN_SERVICE_ACCOUNT = \"samples@python-docs-samples-tests.iam.gserviceaccount.com\"  # @param {type:\"string\"}\n",
+        "feature_online_store_admin_service_client.set_iam_policy(\n",
+        "    request=iam_policy_pb2.SetIamPolicyRequest(\n",
+        "        resource=FEATURE_ONLINE_STORE_RESOURCE_ID,\n",
+        "        policy=policy_pb2.Policy(\n",
+        "            bindings=[\n",
+        "                policy_pb2.Binding(\n",
+        "                    role=\"roles/aiplatform.featurestoreResourceViewer\",\n",
+        "                    members=[f\"group:{EXAMPLE_RESOURCE_VIEWER_EMAIL}\"],\n",
+        "                ),\n",
+        "                policy_pb2.Binding(\n",
+        "                    role=\"roles/aiplatform.featurestoreDataViewer\",\n",
+        "                    members=[f\"group:{EXAMPLE_DATA_VIEWER_EMAIL}\"],\n",
+        "                ),\n",
+        "                policy_pb2.Binding(\n",
+        "                    role=\"roles/aiplatform.admin\",\n",
+        "                    members=[f\"serviceAccount:{EXAMPLE_ADMIN_SERVICE_ACCOUNT}\"],\n",
+        "                ),\n",
+        "            ]\n",
+        "        ),\n",
+        "    )\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3ct-qQeaY-HQ"
+      },
+      "source": [
+        "Confirm the updated IAM Policy of the feature online store."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "nPkYu3D_Y9q1"
+      },
+      "outputs": [],
+      "source": [
+        "feature_online_store_admin_service_client.get_iam_policy(\n",
+        "    request=iam_policy_pb2.GetIamPolicyRequest(\n",
+        "        resource=FEATURE_ONLINE_STORE_RESOURCE_ID\n",
+        "    )\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "gOldDbeVY6SK"
+      },
+      "source": [
+        "Configure IAM Policy for a feature view."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "IuMXMJQjY32o"
+      },
+      "outputs": [],
+      "source": [
+        "# A specific user email or a group email can be used to add members to the IAM Policy.\n",
+        "EXAMPLE_RESOURCE_VIEWER_EMAIL = (\n",
+        "    \"google-cloud-eng-fte@google.com\"  # @param {type:\"string\"}\n",
+        ")\n",
+        "EXAMPLE_DATA_VIEWER_EMAIL = \"python-sample-owners@google.com\"  # @param {type:\"string\"}\n",
+        "# A Service Account can also be added to the IAM Policy.\n",
+        "EXAMPLE_ADMIN_SERVICE_ACCOUNT = \"samples@python-docs-samples-tests.iam.gserviceaccount.com\"  # @param {type:\"string\"}\n",
+        "feature_online_store_admin_service_client.set_iam_policy(\n",
+        "    request=iam_policy_pb2.SetIamPolicyRequest(\n",
+        "        resource=FEATURE_VIEW_RESOURCE_ID,\n",
+        "        policy=policy_pb2.Policy(\n",
+        "            bindings=[\n",
+        "                policy_pb2.Binding(\n",
+        "                    role=\"roles/aiplatform.featurestoreResourceViewer\",\n",
+        "                    members=[f\"group:{EXAMPLE_RESOURCE_VIEWER_EMAIL}\"],\n",
+        "                ),\n",
+        "                policy_pb2.Binding(\n",
+        "                    role=\"roles/aiplatform.featurestoreDataViewer\",\n",
+        "                    members=[f\"group:{EXAMPLE_DATA_VIEWER_EMAIL}\"],\n",
+        "                ),\n",
+        "                policy_pb2.Binding(\n",
+        "                    role=\"roles/aiplatform.admin\",\n",
+        "                    members=[f\"serviceAccount:{EXAMPLE_ADMIN_SERVICE_ACCOUNT}\"],\n",
+        "                ),\n",
+        "            ]\n",
+        "        ),\n",
+        "    )\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "IdDAiz6TEWMI"
+      },
+      "source": [
+        "Wait some time for the newly-added IAM policy binding to take effect.\n",
+        "\n",
+        "* Note: The policy change typically takes 2 minutes, so this delay may be reduced. See IAM docs for more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "yGeqDWFgEX3s"
+      },
+      "outputs": [],
+      "source": [
+        "# Wait 10 minutes for the newly updated IAM policy binding to become effective.\n",
+        "from time import sleep\n",
+        "\n",
+        "sleep(600)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "RVhAFV0MY3Ko"
+      },
+      "source": [
+        "Confirm the updated IAM Policy."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "awfeS9z5Y0sy"
+      },
+      "outputs": [],
+      "source": [
+        "feature_online_store_admin_service_client.get_iam_policy(\n",
+        "    request=iam_policy_pb2.GetIamPolicyRequest(resource=FEATURE_VIEW_RESOURCE_ID)\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "t6A7fWzlY0Pc"
+      },
+      "source": [
+        "Run on-demand batch sync."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YLJUPFPmYsa3"
+      },
+      "outputs": [],
+      "source": [
+        "sync_response = feature_online_store_admin_service_client.sync_feature_view(\n",
+        "    feature_view=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\"\n",
+        ")\n",
+        "\n",
+        "# Wait for sync completion\n",
+        "while True:\n",
+        "    feature_view_sync = feature_online_store_admin_service_client.get_feature_view_sync(\n",
+        "        name=sync_response.feature_view_sync\n",
+        "    )\n",
+        "    if feature_view_sync.run_time.end_time.seconds > 0:\n",
+        "        status = \"Succeed\" if feature_view_sync.final_status.code == 0 else \"Failed\"\n",
+        "        print(f\"Sync {status} for {feature_view_sync.name}.\")\n",
+        "        break\n",
+        "    else:\n",
+        "        print(\"Sync ongoing, waiting for 30 seconds.\")\n",
+        "    sleep(30)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-k4qNUFcrYzC"
+      },
+      "source": [
+        "Confirm the status of batch sync."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "6mnM08ZLrT8f"
+      },
+      "outputs": [],
+      "source": [
+        "feature_online_store_admin_service_client.get_feature_view_sync(\n",
+        "    name=sync_response.feature_view_sync\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "z-PuxGYDYph2"
+      },
+      "source": [
+        "### Start online serving"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "uyZrlw9pMGze"
+      },
+      "source": [
+        "After the data sync is complete, use the `FetchFeatureValues` API to retrieve the data."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "PmOPvky1Ylbv"
+      },
+      "outputs": [],
+      "source": [
+        "data_client = FeatureOnlineStoreServiceClient(client_options={\"api_endpoint\": ENDPOINT})"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "27e6iqlTMNDv"
+      },
+      "source": [
+        "Read the synced data from feature online store."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "1m4OpNPsMMqP"
+      },
+      "outputs": [],
+      "source": [
+        "data_client.fetch_feature_values(\n",
+        "    request=feature_online_store_service_pb2.FetchFeatureValuesRequest(\n",
+        "        feature_view=FEATURE_VIEW_RESOURCE_ID,\n",
+        "        data_key=feature_online_store_service_pb2.FeatureViewDataKey(key=\"test\"),\n",
+        "    )\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2a4e033321ad"
+      },
+      "source": [
+        "## Cleaning up\n",
+        "\n",
+        "Cleaning up\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 cluster you created in this tutorial."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "aqLqumA9YeBh"
+      },
+      "outputs": [],
+      "source": [
+        "# Delete feature view\n",
+        "feature_online_store_admin_service_client.delete_feature_view(\n",
+        "    name=FEATURE_VIEW_RESOURCE_ID\n",
+        ")\n",
+        "\n",
+        "# Delete online store\n",
+        "feature_online_store_admin_service_client.delete_feature_online_store(\n",
+        "    name=FEATURE_ONLINE_STORE_RESOURCE_ID\n",
+        ")\n",
+        "\n",
+        "# Delete test data\n",
+        "!bq rm -f {DATASET_ID}.{TABLE_ID}"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "vertex_ai_feature_store_iam_policy.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
@@ -32,22 +32,24 @@
        "# Vertex AI Feature Store (Legacy): Streaming import SDK\n",
        "\n",
        "<table align=\"left\">\n",
-        "\n",
-        "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/feature_store/feature_store_streaming_ingestion_sdk.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/feature_store_legacy/feature_store_streaming_ingestion_sdk.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://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/feature_store/feature_store_streaming_ingestion_sdk.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%2Ffeature_store_legacy%2Ffeature_store_streaming_ingestion_sdk.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/feature_store_legacy/feature_store_streaming_ingestion_sdk.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>\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/feature_store/feature_store_streaming_ingestion_sdk.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://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/feature_store_legacy/feature_store_streaming_ingestion_sdk.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,15 +130,22 @@
        "to generate a cost estimate based on your projected usage.\n"
      ]
    },
+    {
+      "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"
      ]
    },
    {
@@ -159,60 +168,80 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "58707a750154"
+        "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": "f200f10a1da3"
+        "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": "BF1j6f9HApxa"
+        "id": "b96b39fd4d7b"
      },
      "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>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "WReHDGG5g0XY"
+        "id": "ff666ce4051c"
      },
      "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": "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": "e8575d303471"
+      },
+      "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)."
      ]
    },
    {
@@ -224,103 +253,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": "kljmKgilI_de"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "sBCra4QMA2wR"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "74ccc9e52986"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "de775a3773ba"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "254614fa0c46"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ef21552ccea8"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "603adbbf0532"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "f6b2ccc891ed"
-      },
-      "source": [
-        "**4. Service account or other**\n",
-        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
+        "LOCATION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
@@ -331,7 +264,7 @@
      "source": [
        "### UUID\n",
        "\n",
-        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a uuid for each instance session, and append it onto the name of resources you create in this tutorial."
+        "If you're in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a uuid for each instance session, and append it onto the name of resources you create in this tutorial."
      ]
    },
    {
@@ -395,7 +328,7 @@
      },
      "outputs": [],
      "source": [
-        "aiplatform.init(project=PROJECT_ID, location=REGION)"
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
@@ -459,7 +392,7 @@
        "\n",
        "`{entity_id : {feature_id : feature_value}, ...},`\n",
        "\n",
-        "or a pandas `Dataframe`, where the `index` column holds the unique entity ID strings and each remaining column represents a feature.  In this notebook, since you use a pandas `DataFrame` for ingesting features we convert the index column data type to `string` to be used as `Entity ID`."
+        "or a pandas `Dataframe`, where the `index` column holds the unique entity ID strings and each remaining column represents a feature.  In this notebook, since you use a pandas `DataFrame` for ingesting features, convert the index column data type to `string` for your `Entity ID` usage."
      ]
    },
    {
@@ -484,7 +417,7 @@
      "source": [
        "# Remove null values\n",
        "NA_VALUES = [\"NA\", \".\"]\n",
-        "penguins_df = penguins_df.replace(to_replace=NA_VALUES, value=np.NaN).dropna()"
+        "penguins_df = penguins_df.replace(to_replace=NA_VALUES, value=np.nan).dropna()"
      ]
    },
    {
@@ -514,7 +447,7 @@
        "\n",
        "You create a featurestore using `aiplatform.Featurestore.create` with the following parameters:\n",
        "\n",
-        "* `featurestore_id (str)`: The ID to use for this featurestore, which will become the final component of the `featurestore` resource name. The value must be unique within the project and location.\n",
+        "* `featurestore_id (str)`: The ID to use for this featurestore, which becomes the final component of the `featurestore` resource name. The value must be unique within the project and location.\n",
        "* `online_store_fixed_node_count`: Configuration for online serving resources.\n",
        "* `project`: Project to create the `EntityType` in. If not set, project set in `aiplatform.init` is used.\n",
        "* `location`: Location to create the `EntityType` in. If not set, location set in `aiplatform.init` is used.\n",
@@ -535,7 +468,7 @@
        "    featurestore_id=FEATURESTORE_ID,\n",
        "    online_store_fixed_node_count=1,\n",
        "    project=PROJECT_ID,\n",
-        "    location=REGION,\n",
+        "    location=LOCATION,\n",
        "    sync=True,\n",
        ")"
      ]
@@ -561,7 +494,7 @@
        "fs = aiplatform.Featurestore(\n",
        "    featurestore_name=FEATURESTORE_ID,\n",
        "    project=PROJECT_ID,\n",
-        "    location=REGION,\n",
+        "    location=LOCATION,\n",
        ")\n",
        "print(fs.gca_resource)"
      ]
@@ -577,7 +510,7 @@
        "An entity type is a collection of semantically related features. You define your own entity types, based on the concepts that are relevant to your use case. For example, a movie service might have the entity types `movie` and `user`, which group related features that correspond to movies or users.\n",
        "\n",
        "Here, you create an entity type entity type named `penguin_entity_type` using `create_entity_type` with the following parameters:\n",
-        "* `entity_type_id (str)`: The ID to use for the `EntityType`, which will become the final component of the `EntityType` resource name. The value must be unique within a featurestore.\n",
+        "* `entity_type_id (str)`: The ID to use for the `EntityType`, which becomes the final component of the `EntityType` resource name. The value must be unique within a featurestore.\n",
        "* `description`: Description of the `EntityType`."
      ]
    },
@@ -670,7 +603,7 @@
        "id": "AKRXJCPijM8w"
      },
      "source": [
-        "You can create features either using `create_feature` or `batch_create_features`. Here, for convinience, you have added all feature configs in one variabel, so we use `batch_create_features`."
+        "You can create features either using `create_feature` or `batch_create_features`. Here, for convinience, you have added all feature configs in one variable, so you use `batch_create_features`."
      ]
    },
    {
@@ -32,23 +32,24 @@
        "# Using Vertex AI Feature Store (Legacy) with Pandas Dataframe\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/feature_store/sdk-feature-store-pandas.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> \n",
-        "        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",
-        "    \n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/feature_store/sdk-feature-store-pandas.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%2Ffeature_store%2Fsdk-feature-store-pandas.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/feature_store/sdk-feature-store-pandas.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/feature_store/sdk-feature-store-pandas.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>"
@@ -136,15 +137,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"
      ]
    },
    {
@@ -168,60 +176,80 @@
    {
      "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": "b96b39fd4d7b"
      },
      "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>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "21e3cac35e75"
+        "id": "ff666ce4051c"
      },
      "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": "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": "e8575d303471"
+      },
+      "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)."
      ]
    },
    {
@@ -233,89 +261,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\"}"
      ]
    },
    {
@@ -357,7 +303,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
@@ -393,7 +339,7 @@
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
-        "Initialize the Vertex AI SDK for Python for your project and region."
+        "Initialize the Vertex AI SDK for Python for your project and location."
      ]
    },
    {
@@ -404,7 +350,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)"
      ]
    },
    {
@@ -497,7 +443,7 @@
        "\n",
        "Add the defined features to the entity types `users` and `movies` using the following methods.\n",
        "\n",
-        "### Add features using *`create_feature`* method\n",
+        "### Add features using `create_feature` method\n",
        "\n",
        "You provide the following parameters for creating features:\n",
        "\n",
@@ -542,7 +488,7 @@
        "id": "ecb141839033"
      },
      "source": [
-        "### Add features using *`batch_create_features`* method\n",
+        "### Add features using `batch_create_features` method\n",
        "\n",
        "You can also add multiple features at a time using a config map in a dictionary format. For this, you use the `batch_create_features` method. \n",
        "\n",
@@ -1179,8 +1125,8 @@
        "! rm {USERS_AVRO_FN} {MOVIES_AVRO_FN}\n",
        "\n",
        "# Delete Cloud Storage objects that were created\n",
-        "delete_bucket = True\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "delete_bucket = False  # Set True for deletion\n",
+        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
@@ -32,21 +32,24 @@
        "# Online and Batch predictions using Vertex AI Feature Store (Legacy)\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/feature_store/sdk-feature-store.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/feature_store_legacy/sdk-feature-store.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://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/feature_store/sdk-feature-store.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%2Ffeature_store_legacy%2Fsdk-feature-store.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/feature_store/sdk-feature-store.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>    \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/feature_store_legacy/sdk-feature-store.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/feature_store_legacy/sdk-feature-store.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>"
@@ -122,15 +125,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": "s3Jje0B5zglA"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook."
+        "### Install Vertex AI SDK for Python and other required packages"
      ]
    },
    {
@@ -147,60 +157,80 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "GlWoVi7xz1TL"
+        "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": "CFS6OPNWz3KZ"
+        "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": "7RMhe6650CyB"
+        "id": "4a2b7b59bbf7"
      },
      "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). {TODO: Update the APIs needed for your tutorial. Edit the API names, and update the link to append the API IDs, separating each one with a comma. For example, container.googleapis.com,cloudbuild.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>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "T7C_dgnR0L_l"
+        "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": "91842ef41bbd"
+      },
+      "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)."
      ]
    },
    {
@@ -212,103 +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": "ybtwdOp40TVK"
-      },
-      "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": "oLUOopdB0UkU"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "G_ZkpZnv0a0b"
-      },
-      "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": "rfsExLao0b49"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ovUeYbbM0nmK"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "l_AmeEXr0pE1"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "fsl-OPfF0sUO"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "mOh0DLZP0vUI"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "qML_uytf0ymm"
-      },
-      "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\"}"
      ]
    },
    {
@@ -330,7 +264,7 @@
      },
      "outputs": [],
      "source": [
-        "BUCKET_URI = \"gs://your-bucket-name-unique\"  # @param {type:\"string\"}"
+        "BUCKET_URI = \"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -350,7 +284,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
+        "! gsutil mb -l $LOCATION -p $PROJECT_ID $BUCKET_URI"
      ]
    },
    {
@@ -395,7 +329,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)"
      ]
    },
    {
@@ -466,7 +400,7 @@
        "    featurestore_id=FEATURESTORE_ID,\n",
        "    online_store_fixed_node_count=ONLINE_STORE_FIXED_NODE_COUNT,\n",
        "    project=PROJECT_ID,\n",
-        "    location=REGION,\n",
+        "    location=LOCATION,\n",
        "    sync=True,\n",
        ")"
      ]
@@ -491,7 +425,7 @@
        "fs = Featurestore(\n",
        "    featurestore_name=FEATURESTORE_ID,\n",
        "    project=PROJECT_ID,\n",
-        "    location=REGION,\n",
+        "    location=LOCATION,\n",
        ")\n",
        "print(fs.gca_resource)"
      ]
@@ -1150,7 +1084,7 @@
        "client = bigquery.Client(project=PROJECT_ID)\n",
        "dataset_id = \"{}.{}\".format(client.project, DESTINATION_DATA_SET)\n",
        "dataset = bigquery.Dataset(dataset_id)\n",
-        "dataset.location = REGION\n",
+        "dataset.location = LOCATION\n",
        "dataset = client.create_dataset(dataset)\n",
        "print(\"Created dataset {}.{}\".format(client.project, dataset.dataset_id))"
      ]
@@ -1236,7 +1170,7 @@
        "    featurestore_online_service as featurestore_online_service_pb2\n",
        "from google.cloud.aiplatform_v1beta1.types import types as types_pb2\n",
        "\n",
-        "API_ENDPOINT = \"{}-aiplatform.googleapis.com\".format(REGION)\n",
+        "API_ENDPOINT = \"{}-aiplatform.googleapis.com\".format(LOCATION)\n",
        "# Create client connection\n",
        "admin_client = FeaturestoreServiceClient(client_options={\"api_endpoint\": API_ENDPOINT})\n",
        "data_client = FeaturestoreOnlineServingServiceClient(\n",
@@ -1255,7 +1189,7 @@
        "# Call `write_feature_values` to import data to `users` entity type.\n",
        "data_client.write_feature_values(\n",
        "    entity_type=admin_client.entity_type_path(\n",
-        "        PROJECT_ID, REGION, FEATURESTORE_ID, \"users\"\n",
+        "        PROJECT_ID, LOCATION, FEATURESTORE_ID, \"users\"\n",
        "    ),\n",
        "    payloads=[\n",
        "        featurestore_online_service_pb2.WriteFeatureValuesPayload(\n",
@@ -1312,18 +1246,6 @@
        "You can also keep the project, but delete the featurestore and the BigQuery dataset by running the following code:"
      ]
    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "NBTNfN8vxz4x"
-      },
-      "outputs": [],
-      "source": [
-        "# Delete Featurestore\n",
-        "fs.delete(force=True)"
-      ]
-    },
    {
      "cell_type": "code",
      "execution_count": null,
@@ -1332,13 +1254,21 @@
      },
      "outputs": [],
      "source": [
+        "# Delete Featurestore\n",
+        "fs.delete(force=True)\n",
+        "\n",
        "# Delete BigQuery dataset\n",
        "client = bigquery.Client(project=PROJECT_ID)\n",
        "client.delete_dataset(\n",
        "    DESTINATION_DATA_SET, delete_contents=True, not_found_ok=True\n",
        ")  # Make an API request.\n",
        "\n",
-        "print(\"Deleted dataset '{}'.\".format(DESTINATION_DATA_SET))"
+        "print(\"Deleted dataset '{}'.\".format(DESTINATION_DATA_SET))\n",
+        "\n",
+        "# 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"
      ]
    }
  ],
@@ -72,11 +72,17 @@
        "\n",
        "### Available Anthropic Claude models\n",
        "\n",
-        "#### Claude 3 Sonnet\n",
-        "Anthropic Claude 3 Sonnet provides a balance between intelligence and speed for enterprise workloads. It's a high-endurance model for scaled AI that's available at a competitive price.\n",
+        "#### Claude 3.5 Sonnet\n",
+        "Anthropic's most powerful AI model. Claude 3.5 Sonnet outperforms competitor models and Claude 3 Opus at higher speeds and lower cost.\n",
+        "\n",
+        "#### Claude 3 Opus\n",
+        "Claude 3 Opus is Anthropic's second-most intelligent AI model, with top-level performance on highly complex tasks.\n",
        "\n",
        "#### Claude 3 Haiku\n",
-        "Anthropic Claude 3 Haiku is the fastest, most compact model available from Anthropic. It is designed to answer simple queries and requests quickly. You can use it to build AI experiences that mimic human interactions.\n",
+        "Anthropic Claude 3 Haiku is Anthropic's fastest, most compact vision and text model for near-instant responses to simple queries, meant for seamless AI experiences mimicking human interactions.\n",
+        "\n",
+        "#### Claude 3 Sonnet\n",
+        "Anthropic Claude 3 Sonnet is engineered to be dependable for scaled AI deployments across a variety of use cases.\n",
        "\n",
        "All Claude 3 models can process images and return text outputs, and feature a 200K context window.\n",
        "\n",
@@ -210,12 +216,14 @@
      },
      "outputs": [],
      "source": [
-        "MODEL = \"claude-3-sonnet@20240229\"  # @param [\"claude-3-sonnet@20240229\", \"claude-3-haiku@20240307\", \"claude-3-opus@20240229\"]\n",
-        "if MODEL == \"claude-3-sonnet@20240229\":\n",
-        "    available_regions = [\"us-central1\", \"asia-southeast1\"]\n",
+        "MODEL = \"claude-3-5-sonnet@20240620\"  # @param [\"claude-3-5-sonnet@20240620\", \"claude-3-opus@20240229\", \"claude-3-haiku@20240307\", \"claude-3-sonnet@20240229\" ]\n",
+        "if MODEL == \"claude-3-5-sonnet@20240620\":\n",
+        "    available_regions = [\"us-east5\", \"europe-west1\"]\n",
+        "elif MODEL == \"claude-3-opus@20240229\":\n",
+        "    available_regions = [\"us-east5\"]\n",
        "elif MODEL == \"claude-3-haiku@20240307\":\n",
-        "    available_regions = [\"us-central1\", \"europe-west4\"]\n",
-        "else:\n",
+        "    available_regions = [\"us-east5\", \"europe-west1\"]\n",
+        "elif MODEL == \"claude-3-sonnet@20240229\":\n",
        "    available_regions = [\"us-east5\"]"
      ]
    },
@@ -630,12 +638,14 @@
      },
      "outputs": [],
      "source": [
-        "MODEL = \"claude-3-sonnet@20240229\"  # @param [\"claude-3-sonnet@20240229\", \"claude-3-haiku@20240307\", \"claude-3-opus@20240229\"]\n",
-        "if MODEL == \"claude-3-sonnet@20240229\":\n",
-        "    available_regions = [\"us-central1\", \"asia-southeast1\"]\n",
+        "MODEL = \"claude-3-5-sonnet@20240620\"  # @param [\"claude-3-5-sonnet@20240620\", \"claude-3-opus@20240229\", \"claude-3-haiku@20240307\", \"claude-3-sonnet@20240229\" ]\n",
+        "if MODEL == \"claude-3-5-sonnet@20240620\":\n",
+        "    available_regions = [\"us-east5\", \"europe-west1\"]\n",
+        "elif MODEL == \"claude-3-opus@20240229\":\n",
+        "    available_regions = [\"us-east5\"]\n",
        "elif MODEL == \"claude-3-haiku@20240307\":\n",
-        "    available_regions = [\"us-central1\", \"europe-west4\"]\n",
-        "else:\n",
+        "    available_regions = [\"us-east5\", \"europe-west1\"]\n",
+        "elif MODEL == \"claude-3-sonnet@20240229\":\n",
        "    available_regions = [\"us-east5\"]"
      ]
    },
@@ -29,41 +29,32 @@
        "id": "JAPoU8Sm5E6e"
      },
      "source": [
-        "# Vertex AI Distill a model\n",
+        "# Vertex AI: Distill a large language model\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/generative_ai/distillation.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/generative_ai/distillation.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%2Fgenerative_ai%2Fdistillation.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/generative_ai/distillation.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/generative_ai/distillation.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": "24743cf4a1e1"
-      },
-      "source": [
-        "**_NOTE_**: This notebook has been tested in the following environment:\n",
-        "\n",
-        "* Python version = 3.9"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -72,11 +63,15 @@
      "source": [
        "## Overview\n",
        "\n",
-        "This tutorial demonstrates how to use the Distilling Step by Step on the Vertex AI.\n",
+        "This tutorial demonstrates how to use the distilling Step by Step on the Vertex AI.\n",
        "\n",
-        "We developed the distilling step-by-step (DSS) method ([paper](https://arxiv.org/abs/2305.02301v1)) that can enrich customer’s data by eliciting the reasoning process (rationales) from a large language model (LLM). This new mechanism has shown to be able to (a) train smaller models that outperform LLMs, and (b) achieves so by leveraging less training data needed by fine-tuning or distillation. Our method extracts LLM rationales as additional supervision within a multi-task training framework.\n",
+        "The distilling step-by-step (DSS) method ([paper](https://arxiv.org/abs/2305.02301v1)) can enrich customer’s data by eliciting the reasoning process (rationales) from a large language model (LLM). This new mechanism has shown to be able to (a) train smaller models that outperform LLMs, and (b) achieves so by leveraging less training data needed by fine-tuning or distillation. This method extracts LLM rationales as additional supervision within a multi-task training framework.\n",
        "\n",
-        "Learn more about [distill-text-models](https://cloud.google.com/vertex-ai/docs/generative-ai/models/distill-text-models)."
+        "Learn more about [distill-text-models](https://cloud.google.com/vertex-ai/docs/generative-ai/models/distill-text-models).\n",
+        "\n",
+        "**_NOTE_**: This notebook is tested in the following environment:\n",
+        "\n",
+        "* Python version = 3.9"
      ]
    },
    {
@@ -87,23 +82,60 @@
      "source": [
        "### Objective\n",
        "\n",
-        "In this tutorial, you learn to use `Vertex AI LLM` to distill and deploy a large language model.\n",
+        "In this tutorial, you learn how to distill and deploy a large language model using Vertex AI LLM.\n",
        "\n",
+        "This tutorial uses the following Vertex AI services:\n",
        "\n",
-        "This tutorial uses the following Google Cloud ML services:\n",
-        "\n",
-        "- `Vertex AI LLM`\n",
-        "- `Vertex AI Model Garden`\n",
-        "- `Vertex AI Prediction`\n",
+        "- Vertex AI LLM\n",
+        "- Vertex AI Model Garden\n",
+        "- Vertex AI Online prediction\n",
        "\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Get the Vertex AI LLM model.\n",
-        "- Distill the model.\n",
-        "  - This will automatically create a Vertex AI endpoint and deploy the model to it.\n",
-        "- Make a prediction using `Vertex AI LLM`.\n",
-        "- Make a prediction using `Vertex AI Prediction`"
+        "- Distill the model(this automatically creates a Vertex AI endpoint and deploys the model to the endpoint). \n",
+        "- Make a prediction using Vertex AI LLM."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3d269b76353d"
+      },
+      "source": [
+        "### Dataset\n",
+        "\n",
+        "Distillation works on a labeled or an unlabeled dataset. If you have a high quality labeled dataset with hundreds of examples, then it's recommended that you use the labeled dataset. Otherwise, you can use an unlabeled prompt dataset. If you use an unlabeled dataset, then the teacher model generates the labels and the rationale for distillation. More than 1,000 examples are recommended if you use an unlabeled dataset.\n",
+        "\n",
+        "For this tutorial, you use a dataset stored in a public Cloud Storage bucket at the below paths. \n",
+        "- Train sample: `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`\n",
+        "- Validation sample: `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_eval_sample.jsonl`"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e7c0aa7c769f"
+      },
+      "source": [
+        "#### Input format requirement\n",
+        "\n",
+        "The labeled or unlabeled distillation dataset must be in JSON Lines (JSONL) format where each line contains a single tuning example. Before you distill your model, upload your dataset to a Cloud Storage bucket.\n",
+        "\n",
+        "Each dataset example contains an `input_text` field with the model prompt and an optional `output_text` field that contains an example response that the distilled model is expected to produce.\n",
+        "\n",
+        "The maximum token length for `input_text` is 7,168 and the maximum token length for `output_text` is 1,024. If either field exceeds the maximum token length, the excess tokens are truncated.\n",
+        "\n",
+        "The maximum number of examples that a dataset for a text generation model can contain is 10,000.\n",
+        "\n",
+        "\n",
+        "Example:\n",
+        "\n",
+        "```\n",
+        "{\"input_text\": \"question: How many people live in Beijing? context: With over 21 million residents, Beijing is the world's most populous national capital city and is China's second largest city after Shanghai. It is located in Northern China, and is governed as a municipality under the direct administration of the State Council with 16 urban, suburban, and rural districts.[14] Beijing is mostly surrounded by Hebei Province with the exception of neighboring Tianjin to the southeast; together, the three divisions form the Jingjinji megalopolis and the national capital region of China.\", \"output_text\": \"over 21 million people\"}\n",
+        "{\"input_text\": \"question: How many parishes are there in Louisiana? context: The U.S. state of Louisiana is divided into 64 parishes (French: paroisses) in the same manner that 48 other states of the United States are divided into counties, and Alaska is divided into boroughs.\", \"output_text\": \"64\"}\n",
+        "```"
      ]
    },
    {
@@ -129,189 +161,127 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "i7EUnXsZhAGF"
+        "id": "f0316df526f8"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook. "
+        "## 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": "2b4ef9b72d43"
+        "id": "89d404f6cc9d"
      },
      "outputs": [],
      "source": [
-        "! pip3 install --upgrade --quiet google-cloud-aiplatform \"shapely<2.0.0\""
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
+        "                                 \"shapely<2.0.0\" \\\n",
+        "                                 PyYAML"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "58707a750154"
+        "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": "f200f10a1da3"
+        "id": "f09b4dff629a"
      },
      "outputs": [],
      "source": [
-        "# 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": "54c5ef8a8f43"
      },
      "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": "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": "oM1iC_MfAts1"
+        "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": "5ff1a3cc4e1d"
      },
      "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",
+        "import vertexai\n",
        "\n",
-        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "region"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "sBCra4QMA2wR"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "74ccc9e52986"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "de775a3773ba"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "254614fa0c46"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ef21552ccea8"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "603adbbf0532"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "f6b2ccc891ed"
-      },
-      "source": [
-        "**4. Service account or other**\n",
-        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
+        "vertexai.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
@@ -322,9 +292,7 @@
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
-        "Create a storage bucket to store intermediate artifacts such as datasets.\n",
-        "\n",
-        "- *{Note to notebook author: For any user-provided strings that need to be unique (like bucket names or model ID's), append \"-unique\" to the end so proper testing can occur}*"
+        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
    },
    {
@@ -344,7 +312,7 @@
        "id": "-EcIXiGsCePi"
      },
      "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."
      ]
    },
    {
@@ -355,7 +323,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
@@ -364,9 +332,11 @@
        "id": "08d289fa873f"
      },
      "source": [
-        "### Dataset\n",
+        "#### Copy the dataset to your bucket\n",
        "\n",
-        "We've provided the below sample data for you to get started.\n"
+        "Before you start the distillation, copy the dataset from the source to your Cloud Storage bucket.\n",
+        "\n",
+        "**Note**: Alternatively, you can directly specify the source path for the data when you perform distillation. Copying the data to your Google Cloud project is only optional."
      ]
    },
    {
@@ -381,33 +351,6 @@
        "! gsutil cp gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl {BUCKET_URI}/peft_train_sample.jsonl"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0e6f90042af6"
-      },
-      "source": [
-        "#### Data Input format requirement\n",
-        "\n",
-        "Distillation works on a labeled or an unlabeled dataset. If you have a high quality labeled dataset with hundreds of examples, then we recommend that you use that. Otherwise, you can use an unlabeled prompt dataset. If you use an unlabeled dataset, then the teacher model generates the labels and the rationale for distillation. More than 1,000 examples are recommended if you use an unlabeled dataset.\n",
-        "\n",
-        "The labeled or unlabeled distillation dataset must be in JSON Lines (JSONL) format where each line contains a single tuning example. Before you distill your model, you upload your dataset to a Cloud Storage bucket.\n",
-        "\n",
-        "Each dataset example contains an `input_text` field with the model prompt and an optional `output_text` field that contains an example response that the distilled model is expected to produce.\n",
-        "\n",
-        "The maximum token length for `input_text` is 7,168 and the maximum token length for `output_text` is 1,024. If either field exceeds the maximum token length, the excess tokens are truncated.\n",
-        "\n",
-        "The maximum number of examples that a dataset for a text generation model can contain is 10,000.\n",
-        "\n",
-        "\n",
-        "Sample dataset:\n",
-        "\n",
-        "```\n",
-        "{\"input_text\": \"question: How many people live in Beijing? context: With over 21 million residents, Beijing is the world's most populous national capital city and is China's second largest city after Shanghai. It is located in Northern China, and is governed as a municipality under the direct administration of the State Council with 16 urban, suburban, and rural districts.[14] Beijing is mostly surrounded by Hebei Province with the exception of neighboring Tianjin to the southeast; together, the three divisions form the Jingjinji megalopolis and the national capital region of China.\", \"output_text\": \"over 21 million people\"}\n",
-        "{\"input_text\": \"question: How many parishes are there in Louisiana? context: The U.S. state of Louisiana is divided into 64 parishes (French: paroisses) in the same manner that 48 other states of the United States are divided into counties, and Alaska is divided into boroughs.\", \"output_text\": \"64\"}\n",
-        "```"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -425,44 +368,21 @@
      },
      "outputs": [],
      "source": [
-        "import vertexai\n",
        "from google.cloud import aiplatform\n",
        "from vertexai.preview.language_models import (TextGenerationModel,\n",
        "                                              TuningEvaluationSpec)"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "init_aip:mbsdk,all"
-      },
-      "source": [
-        "### Initialize Vertex AI SDK for Python\n",
-        "\n",
-        "Initialize the Vertex AI SDK for Python for your project."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "init_aip:mbsdk,all"
-      },
-      "outputs": [],
-      "source": [
-        "vertexai.init(project=PROJECT_ID, location=REGION)"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a0275c76a30f"
      },
      "source": [
-        "### Load pretrained model\n",
+        "## Load pretrained model\n",
        "\n",
        "Load the pretrained BISON model from Vertex AI LLM Model Garden.\n",
-        "See the models that supports distillation in [here](https://cloud.google.com/vertex-ai/docs/generative-ai/models/distill-text-models#supported_models)."
+        "See the [list of models that support distillation](https://cloud.google.com/vertex-ai/docs/generative-ai/models/distill-text-models#supported_models)."
      ]
    },
    {
@@ -485,20 +405,19 @@
        "id": "3f35db20ac38"
      },
      "source": [
-        "### Distill the model\n",
+        "## Distill the model\n",
        "\n",
        "Next, you distill the model using the `distill_from()` method, with the following parameters:\n",
        "\n",
-        "`teacher_model`: The teacher model that you would like to distill the knowledge from.\n",
-        "`dataset`: A pandas Dataframe or Cloud Storage location of the training data for tuning the model.<br>\n",
-        "`learning_rate_multiplier`: A multiplier to apply to the recommended learning rate. To use the recommended learning rate, use 1.0. <br>\n",
-        "`train_steps`: The number of steps to run for model tuning. The batch size varies by tuning location:<br>\n",
-        "- us-central1 has a batch size of 8.\n",
-        "- europe-west4 has a batch size of 24.<br>\n",
+        "- `teacher_model`: The teacher model that you would like to distill the knowledge from.\n",
+        "- `dataset`: A pandas Dataframe or Cloud Storage location of the training data for tuning the model.\n",
+        "- `learning_rate_multiplier`: A multiplier to apply to the recommended learning rate. To use the recommended learning rate, use 1.0.\n",
+        "- `train_steps`: The number of steps to run for model tuning. The default value is 300. The batch size varies by tuning location as below for 8k models such as `text-bison@002`:\n",
+        "    \n",
+        "    - us-central1 has a batch size of 8.\n",
+        "    - europe-west4 has a batch size of 24.\n",
        "\n",
-        "If there are 240 examples in a training dataset, in europe-west4, it takes 240 / 24 = 10 steps to process the entire dataset once. In us-central1, it takes 240 / 8 = 30 steps to process the entire dataset once. The default value is 300.<br>\n",
-        "\n",
-        "For more context, see this [doc](https://cloud.google.com/vertex-ai/docs/generative-ai/models/distill-text-models#create_a_text_model_distilling_job) for definition of the parameters. "
+        "For parameter definitions and further context, see [Create a text model distilling job](https://cloud.google.com/vertex-ai/docs/generative-ai/models/distill-text-models#create_a_text_model_distilling_job). "
      ]
    },
    {
@@ -517,6 +436,17 @@
        "eval_spec.evaluation_interval = 20"
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "495abb2c72c8"
+      },
+      "source": [
+        "Set a display name for your model resource and the endpoint resource using the `DISPLAY_NAME` parameter.\n",
+        "\n",
+        "**Note**: In the tuning pipeline, the model and endpoint share the same display name."
+      ]
+    },
    {
      "cell_type": "code",
      "execution_count": null,
@@ -525,15 +455,22 @@
      },
      "outputs": [],
      "source": [
-        "student_model.distill_from(\n",
+        "# Set the display name\n",
+        "DISPLAY_NAME = \"vertex-distillation-model-unique\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Create the tuning pipeline job\n",
+        "pipeline = student_model.distill_from(\n",
        "    teacher_model=teacher_model,\n",
        "    dataset=f\"{BUCKET_URI}/peft_train_sample.jsonl\",\n",
        "    train_steps=200,\n",
        "    learning_rate_multiplier=1,\n",
        "    accelerator_type=\"TPU\",\n",
-        "    model_display_name=\"test-vertex-distillation\",\n",
+        "    model_display_name=DISPLAY_NAME,\n",
        "    evaluation_spec=eval_spec,\n",
-        ")"
+        ")\n",
+        "\n",
+        "# Wait until the tuning pipeline job finishes\n",
+        "pipeline._job.wait()"
      ]
    },
    {
@@ -542,7 +479,7 @@
        "id": "d432a5238785"
      },
      "source": [
-        "### Make a prediction with Vertex AI LLM\n",
+        "## Make a prediction with Vertex AI LLM\n",
        "\n",
        "Now, make a prediction using the `predict()` method from the Vertex AI LLM interface."
      ]
@@ -555,6 +492,7 @@
      },
      "outputs": [],
      "source": [
+        "# Define the prompt\n",
        "prompt = \"TRANSCRIPT: \\nPROCEDURE PERFORMED: , Umbilical hernia repair.,PROCEDURE:,  After informed consent was obtained, the patient was brought to the operative suite and placed supine on the operating table.  The patient was sedated, and an adequate local anesthetic was administered using 1% lidocaine without epinephrine.  The patient was prepped and draped in the usual sterile manner.,A standard curvilinear umbilical incision was made, and dissection was carried down to the hernia sac using a combination of Metzenbaum scissors and Bovie electrocautery.  The sac was cleared of overlying adherent tissue, and the fascial defect was delineated.  The fascia was cleared of any adherent tissue for a distance of 1.5 cm from the defect.  The sac was then placed into the abdominal cavity and the defect was closed primarily using simple interrupted 0 Vicryl sutures.  The umbilicus was then re-formed using 4-0 Vicryl to tack the umbilical skin to the fascia.,The wound was then irrigated using sterile saline, and hemostasis was obtained using Bovie electrocautery.  The skin was approximated with 4-0 Vicryl in a subcuticular fashion.  The skin was prepped with benzoin, and Steri-Strips were applied.  A dressing was then applied.  All surgical counts were reported as correct.,Having tolerated the procedure well, the patient was subsequently taken to the recovery room in good and stable condition.\\n\\n LABEL: \""
      ]
    },
@@ -566,6 +504,7 @@
      },
      "outputs": [],
      "source": [
+        "# Print the prompt\n",
        "print(student_model.predict(prompt))"
      ]
    },
@@ -591,27 +530,30 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
+        "# Fetch the endpoint resource using the display name and create time\n",
+        "endpoints = aiplatform.Endpoint.list(\n",
+        "    filter=f\"display_name={DISPLAY_NAME}\", order_by=\"create_time\"\n",
+        ")\n",
+        "if len(endpoints) > 0:\n",
+        "    # Undeploy the model from the endpoint\n",
+        "    endpoints[0].undeploy_all()\n",
+        "    # Delete the endpoint\n",
+        "    endpoints[0].delete()\n",
        "\n",
-        "delete_bucket = False\n",
+        "# Fetch the model resource using the display name and create time\n",
+        "models = aiplatform.Model.list(\n",
+        "    filter=f\"display_name={DISPLAY_NAME}\", order_by=\"create_time\"\n",
+        ")\n",
+        "if len(models) > 0:\n",
+        "    # Delete the model\n",
+        "    models[0].delete()\n",
        "\n",
-        "if not os.getenv(\"IS_TESTING\"):\n",
-        "    # Delete endpoint resource\n",
-        "    endpoint = aiplatform.Endpoint(student_model._endpoint.resource_name)\n",
-        "    endpoint.undeploy_all()\n",
-        "    endpoint.delete()"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "d473314cae55"
-      },
-      "outputs": [],
-      "source": [
-        "delete_bucket = False\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "# Delete the pipeline job\n",
+        "pipeline._job.delete()\n",
+        "\n",
+        "# Delete the Cloud Storage bucket\n",
+        "delete_bucket = True\n",
+        "if delete_bucket:\n",
        "    ! gsutil rm -rf {BUCKET_URI}"
      ]
    }
@@ -0,0 +1,352 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "9A9NkTRTfo2I"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2024 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "IPprg6Oz0QDs"
+      },
+      "source": [
+        "# Getting Started with Mistral AI Models\n",
+        "<table align=\"left\">\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/generative_ai/mistralai_intro.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fgenerative_ai%2Fmistralai_intro.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">                                                                             \n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/generative_ai/mistralai_intro.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/generative_ai/mistralai_intro.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  \n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8fK_rdvvx1iZ"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "### Mistral AI on Vertex AI\n",
+        "\n",
+        "Mistral AI models on Vertex AI offer fully managed and serverless models are offered as managed APIs. To use a Mistral AI model on Vertex AI, send a request directly to the Vertex AI API endpoint.\n",
+        "\n",
+        "You can stream your Mistral AI model responses to reduce the end-user latency perception. A streamed response uses server-sent events (SSE) to incrementally stream the response.\n",
+        "\n",
+        "Learn more about [Vertex AI](https://cloud.google.com/vertex-ai). \n",
+        "\n",
+        "### Available Mistral AI models\n",
+        "\n",
+        "*   ### Mistral Large (2407)\n",
+        "Complex tasks that require large reasoning capabilities or are highly specialized (synthetic text Generation, code generation, RAG, or agents).\n",
+        "\n",
+        "*   ### Mistral Nemo\n",
+        "Reasoning, world knowledge, and coding performance are state-of-the-art in its size category.\n",
+        "\n",
+        "*   ### Codestral\n",
+        "Coding specific tasks to enhance developers productivity with code completion and fill-in-the-middle capabilities.\n",
+        "\n",
+        "\n",
+        "## Objective\n",
+        "\n",
+        "This notebook shows how to use **Vertex AI API** to call the Mistral AI models on Vertex AI API with the Large, Nemo, and Codestral models.\n",
+        "\n",
+        "For more information, see the [Use Mistral's](https://docs.mistral.ai/) documentation and [Mistral's models](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/mistral) on Google Cloud.\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "HcJCV6Dw5usD"
+      },
+      "source": [
+        "## Vertex AI API"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "nwYvaaW25jYS"
+      },
+      "source": [
+        "## Get Started\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6a5bea26f60f"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
+        "\n",
+        "Authenticate your environment on Google Colab.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "c97be6a73155"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2fxZn4SAbxdl"
+      },
+      "source": [
+        "#### Select Mistral AI model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Y8X70FTSbx7U"
+      },
+      "outputs": [],
+      "source": [
+        "MODEL = \"mistral-large\"  # @param [\"mistral-large\", \"mistral-nemo\", \"codestral\"]\n",
+        "if MODEL == \"mistral-large\":\n",
+        "    available_regions = [\"europe-west4\", \"us-central1\"]\n",
+        "    available_versions = [\"latest\", \"2407\"]\n",
+        "elif MODEL == \"mistral-nemo\":\n",
+        "    available_regions = [\"europe-west4\", \"us-central1\"]\n",
+        "    available_versions = [\"latest\", \"2407\"]\n",
+        "elif MODEL == \"codestral\":\n",
+        "    available_regions = [\"europe-west4\", \"us-central1\"]\n",
+        "    available_versions = [\"latest\", \"2405\"]"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bpuX3sKtexlK"
+      },
+      "source": [
+        "#### Select a location"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "dHl8xW45ex_O"
+      },
+      "outputs": [],
+      "source": [
+        "import ipywidgets as widgets\n",
+        "from IPython.display import display\n",
+        "\n",
+        "dropdown_loc = widgets.Dropdown(\n",
+        "    options=available_regions,\n",
+        "    description=\"Select a location:\",\n",
+        "    font_weight=\"bold\",\n",
+        "    style={\"description_width\": \"initial\"},\n",
+        ")\n",
+        "\n",
+        "dropdown_ver = widgets.Dropdown(\n",
+        "    options=available_versions,\n",
+        "    description=\"Select a Model version (optional):\",\n",
+        "    font_weight=\"bold\",\n",
+        "    style={\"description_width\": \"initial\"},\n",
+        ")\n",
+        "\n",
+        "def dropdown_loc_eventhandler(change):\n",
+        "    global LOCATION\n",
+        "    if change[\"type\"] == \"change\" and change[\"name\"] == \"value\":\n",
+        "        LOCATION = change.new\n",
+        "        print(\"Selected:\", change.new)\n",
+        "\n",
+        "def dropdown_ver_eventhandler(change):\n",
+        "    global MODEL_VERSION\n",
+        "    if change[\"type\"] == \"change\" and change[\"name\"] == \"value\":\n",
+        "        MODEL_VERSION = change.new\n",
+        "        print(\"Selected:\", change.new)\n",
+        "\n",
+        "LOCATION = dropdown_loc.value\n",
+        "dropdown_loc.observe(dropdown_loc_eventhandler, names=\"value\")\n",
+        "display(dropdown_loc)\n",
+        "\n",
+        "MODEL_VERSION = dropdown_ver.value\n",
+        "dropdown_ver.observe(dropdown_ver_eventhandler, names=\"value\")\n",
+        "display(dropdown_ver)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3q58icinBjoK"
+      },
+      "source": [
+        "#### Set Google Cloud project and model 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": "hltNx33t6cSZ"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "ENDPOINT = f\"https://{LOCATION}-aiplatform.googleapis.com\"\n",
+        "SELECTED_MODEL_VERSION = \"\" if MODEL_VERSION == \"latest\" else f\"@{MODEL_VERSION}\"\n",
+        "\n",
+        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
+        "    raise ValueError(\"Please set your PROJECT_ID\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4NAstKRFBt4N"
+      },
+      "source": [
+        "#### Import required libraries"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "QZEFLE6a6bqy"
+      },
+      "outputs": [],
+      "source": [
+        "import json"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5ahw-uFjCAbo"
+      },
+      "source": [
+        "### Text generation"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "61107099357a"
+      },
+      "source": [
+        "#### Unary call\n",
+        "\n",
+        "Sends a POST request to the specified API endpoint to get a response from the model using the provided payload."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "4zFz260B50oi"
+      },
+      "outputs": [],
+      "source": [
+        "PAYLOAD = {\n",
+        "    \"model\": MODEL,\n",
+        "    \"messages\": [{\"role\": \"user\", \"content\": \"who is the best French painter?\"}],\n",
+        "    \"max_tokens\": 100,\n",
+        "    \"stream\": False,\n",
+        "}\n",
+        "\n",
+        "request = json.dumps(PAYLOAD)\n",
+        "!curl -X POST -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/publishers/mistralai/models/{MODEL}{SELECTED_MODEL_VERSION}:rawPredict -d '{request}'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e6f52fae9379"
+      },
+      "source": [
+        "#### Streaming call\n",
+        "\n",
+        "Sends a POST request to the specified API endpoint to stream a response from the model using the provided payload."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "c99761dcd7da"
+      },
+      "outputs": [],
+      "source": [
+        "PAYLOAD = {\n",
+        "    \"model\": MODEL,\n",
+        "    \"messages\": [{\"role\": \"user\", \"content\": \"who is the best French painter?\"}],\n",
+        "    \"max_tokens\": 100,\n",
+        "    \"stream\": True,\n",
+        "}\n",
+        "\n",
+        "request = json.dumps(PAYLOAD)\n",
+        "!curl -X POST -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/publishers/mistralai/models/{MODEL}{SELECTED_MODEL_VERSION}:streamRawPredict -d '{request}'"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "mistralai_intro.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
@@ -138,31 +138,40 @@
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "cellView": "form",
-        "id": "lcblD6_FJGXt"
+        "id": "ARIslNGX4IR3"
      },
      "outputs": [],
      "source": [
-        "# @title (Required) Install AI platform package { run: \"auto\" }\n",
-        "AI_PLATFORM_REQUIREMENT = \"google-cloud-aiplatform>=1.51.0\"  # @param {type:\"raw\"}\n",
+        "! pip3 install --upgrade --user --quiet google-cloud-aiplatform"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ngQYqEkV4aBS"
+      },
+      "source": [
+        "### Restart runtime (Colab only)\n",
        "\n",
-        "import builtins\n",
-        "import os\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "VXFRhykV3jyI"
+      },
+      "outputs": [],
+      "source": [
        "import sys\n",
        "\n",
-        "try:\n",
-        "    import vertexai\n",
+        "if \"google.colab\" in sys.modules:\n",
        "\n",
-        "    print(f\"Got {vertexai.__version__=}.\")\n",
-        "except:\n",
-        "    !pip3 install {AI_PLATFORM_REQUIREMENT}\n",
+        "    import IPython\n",
        "\n",
-        "    running_in_colab = \"google.colab\" in sys.modules and hasattr(\n",
-        "        builtins, \"get_ipython\"\n",
-        "    )\n",
-        "    if running_in_colab and not os.getenv(\"IS_TESTING\"):\n",
-        "        print(\"Please wait a minute. Colab session is rebooting.\")\n",
-        "        get_ipython().instance().kernel.do_shutdown(True)"
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
@@ -325,6 +334,7 @@
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
+        "cellView": "form",
        "id": "3VO4AFDbhnnx"
      },
      "outputs": [],
@@ -342,22 +352,8 @@
      "outputs": [],
      "source": [
        "# @title (Required) Resume an existing or start a fresh tuning session (depending on TUNING_JOB_ID)\n",
-        "from datetime import datetime, timezone\n",
-        "\n",
-        "\n",
-        "def now_in_zulu():\n",
-        "    iso_time = datetime.now(timezone.utc).isoformat(timespec=\"microseconds\")\n",
-        "    return iso_time.replace(\"+00:00\", \"Z\")\n",
-        "\n",
-        "\n",
        "BASE_MODEL = \"text-embedding-004\"  # @param [\"textembedding-gecko@003\", \"text-embedding-004\", \"textembedding-gecko-multilingual@001\", \"text-multilingual-embedding-002\"]\n",
        "TASK = \"DEFAULT\"  # @param [\"DEFAULT\", \"RETRIEVAL_QUERY\", \"RETRIEVAL_DOCUMENT\", \"SEMANTIC_SIMILARITY\", \"CLASSIFICATION\", \"CLUSTERING\", \"QUESTION_ANSWERING\", \"FACT_VERIFICATION\"]\n",
-        "if TASK in [\"QUESTION_ANSWERING\", \"FACT_VERIFICATION\"] and BASE_MODEL not in [\n",
-        "    \"text-embedding-004\",\n",
-        "    \"text-multilingual-embedding-002\",\n",
-        "]:\n",
-        "    raise ValueError(f\"TASK '{TASK}' is not valid for model '{BASE_MODEL}'.\")\n",
-        "\n",
        "CORPUS_DATA = \"gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/corpus.jsonl\"  # @param {type: \"string\"}\n",
        "QUERIES_DATA = \"gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/queries.jsonl\"  # @param {type: \"string\"}\n",
        "TRAINING_DATA = \"gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/train.tsv\"  # @param{type: \"string\"}\n",
@@ -365,6 +361,16 @@
        "TEST_DATA = \"gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/test.tsv\"  # @param{type: \"string\"}\n",
        "BATCH_SIZE = 128  # @param {type: \"number\"}\n",
        "TRAIN_STEPS = 1000  # @param {type: \"number\"}\n",
+        "OUTPUT_DIMENSIONALITY = 768  # @param {type: \"number\"}\n",
+        "LEARNING_RATE_MULTIPLIER = 1.0  # @param {type: \"number\"}\n",
+        "\n",
+        "# Synchronously validate some edge cases that will cause async validation to fail.\n",
+        "if BASE_MODEL not in [\"text-embedding-004\", \"text-multilingual-embedding-002\"]:\n",
+        "    if TASK in [\"QUESTION_ANSWERING\", \"FACT_VERIFICATION\"]:\n",
+        "        raise ValueError(f\"TASK '{TASK}' is not valid for model '{BASE_MODEL}'.\")\n",
+        "\n",
+        "    if OUTPUT_DIMENSIONALITY not in [-1, 768]:\n",
+        "        raise ValueError(f\"Model '{BASE_MODEL}' does not support the output_dimensionality parameter.\")\n",
        "\n",
        "base_model = TextEmbeddingModel.from_pretrained(BASE_MODEL)\n",
        "if \"TUNING_JOB_ID\" in locals() and TUNING_JOB_ID:\n",
@@ -374,7 +380,6 @@
        "        f\"Got an existing tuning job '{tuning_job.name}' (state: {tuning_job.state.name}).\"\n",
        "    )\n",
        "else:\n",
-        "    tuning_job_created_at = now_in_zulu()\n",
        "    tuning_result = base_model.tune_model(\n",
        "        task_type=TASK,\n",
        "        corpus_data=CORPUS_DATA,\n",
@@ -385,12 +390,10 @@
        "        batch_size=BATCH_SIZE,\n",
        "        train_steps=TRAIN_STEPS,\n",
        "        tuned_model_location=REGION,\n",
+        "        learning_rate_multiplier=LEARNING_RATE_MULTIPLIER,\n",
+        "        output_dimensionality=OUTPUT_DIMENSIONALITY,\n",
        "    )\n",
-        "    tuning_job_updated_at = now_in_zulu()\n",
-        "    filter = f'pipelineJobUserId:\"tune-text-embedding-model-*\" AND createTime>=\"{tuning_job_created_at}\" AND createTime<=\"{tuning_job_updated_at}\"'\n",
-        "    tuning_job = next(\n",
-        "        iter(pipeline_jobs.PipelineJob.list(filter=filter, order_by=\"create_time\"))\n",
-        "    )\n",
+        "    tuning_job = pipeline_jobs.PipelineJob.get(tuning_result.pipeline_job_name)\n",
        "    print(\n",
        "        f\"Got a fresh tuning job '{tuning_job.name}' (state: {tuning_job.state.name}).\"\n",
        "    )\n",
@@ -32,26 +32,27 @@
        "# Vertex AI Migration: AutoML Image Classification\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/migration/sdk-automl-image-classification-batch-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",
+        "  <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%2Fmigration%2Fsdk-automl-image-classification-batch-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 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/migration/sdk-automl-image-classification-batch-online.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/migration/sdk-automl-image-classification-batch-online.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/migration/sdk-automl-image-classification-batch-online.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>\n",
-        "<br/><br/><br/>"
+        "</table>"
      ]
    },
    {
@@ -80,10 +81,10 @@
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
-        "- `AutoML`\n",
-        "- `Vertex AI batch prediction`\n",
-        "- `Vertex AI Model` resource\n",
-        "- `Vertex AI Endpoint` resource\n",
+        "- AutoML\n",
+        "- Vertex AI batch prediction\n",
+        "- Vertex AI model resource\n",
+        "- Vertex AI endpoint resource\n",
        "\n",
        "The steps performed include:\n",
        "\n",
@@ -145,54 +146,86 @@
      "source": [
        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
        "                                 google-cloud-storage \\\n",
-        "                                 tensorflow"
+        "                                 tensorflow==2.15.1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "restart"
+        "id": "6e0c0cdffff3"
      },
      "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": "dcc98768955f"
      },
      "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": "294459ee3484"
      },
      "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": "23b421f88a9b"
      },
      "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": "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()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "df44641f2fdb"
+      },
+      "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)."
      ]
    },
    {
@@ -204,134 +237,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": "timestamp"
-      },
-      "source": [
-        "#### UUID\n",
-        "\n",
-        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a uuid for each instance session, and append it onto the name of resources you create in this tutorial."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "timestamp"
-      },
-      "outputs": [],
-      "source": [
-        "import random\n",
-        "import string\n",
-        "\n",
-        "\n",
-        "# Generate a uuid of a specifed length(default=8)\n",
-        "def generate_uuid(length: int = 8) -> str:\n",
-        "    return \"\".join(random.choices(string.ascii_lowercase + string.digits, k=length))\n",
-        "\n",
-        "\n",
-        "UUID = generate_uuid()"
-      ]
-    },
-    {
-      "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\"}"
      ]
    },
    {
@@ -373,7 +279,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} {BUCKET_URI}"
      ]
    },
    {
@@ -396,9 +302,7 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
-        "import google.cloud.aiplatform as aip"
+        "import google.cloud.aiplatform as aiplatform"
      ]
    },
    {
@@ -420,7 +324,7 @@
      },
      "outputs": [],
      "source": [
-        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
+        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -442,9 +346,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\""
      ]
    },
    {
@@ -523,10 +425,10 @@
      },
      "outputs": [],
      "source": [
-        "dataset = aip.ImageDataset.create(\n",
-        "    display_name=\"Flowers\" + \"_\" + UUID,\n",
+        "dataset = aiplatform.ImageDataset.create(\n",
+        "    display_name=\"Flowers\" + \"_\" + \"unique\",\n",
        "    gcs_source=[IMPORT_FILE],\n",
-        "    import_schema_uri=aip.schema.dataset.ioformat.image.single_label_classification,\n",
+        "    import_schema_uri=aiplatform.schema.dataset.ioformat.image.single_label_classification,\n",
        ")\n",
        "\n",
        "print(dataset.resource_name)"
@@ -608,8 +510,8 @@
      },
      "outputs": [],
      "source": [
-        "dag = aip.AutoMLImageTrainingJob(\n",
-        "    display_name=\"flowers_\" + UUID,\n",
+        "dag = aiplatform.AutoMLImageTrainingJob(\n",
+        "    display_name=\"flowers_\" + \"unique\",\n",
        "    prediction_type=\"classification\",\n",
        "    multi_label=False,\n",
        "    model_type=\"CLOUD\",\n",
@@ -663,10 +565,10 @@
      "source": [
        "model = dag.run(\n",
        "    dataset=dataset,\n",
-        "    model_display_name=\"flowers_\" + UUID,\n",
-        "    training_fraction_split=0.8,\n",
-        "    validation_fraction_split=0.1,\n",
-        "    test_fraction_split=0.1,\n",
+        "    model_display_name=\"flowers_\" + \"unique\",\n",
+        "    training_fraction_split=0.4,\n",
+        "    validation_fraction_split=0.3,\n",
+        "    test_fraction_split=0.3,\n",
        "    budget_milli_node_hours=8000,\n",
        "    disable_early_stopping=False,\n",
        ")"
@@ -736,11 +638,13 @@
      "outputs": [],
      "source": [
        "# Get model resource ID\n",
-        "models = aip.Model.list(filter=\"display_name=flowers_\" + UUID)\n",
+        "models = aiplatform.Model.list(filter=\"display_name=flowers_\" + \"unique\")\n",
        "\n",
        "# Get a reference to the Model Service client\n",
-        "client_options = {\"api_endpoint\": f\"{REGION}-aiplatform.googleapis.com\"}\n",
-        "model_service_client = aip.gapic.ModelServiceClient(client_options=client_options)\n",
+        "client_options = {\"api_endpoint\": f\"{LOCATION}-aiplatform.googleapis.com\"}\n",
+        "model_service_client = aiplatform.gapic.ModelServiceClient(\n",
+        "    client_options=client_options\n",
+        ")\n",
        "\n",
        "model_evaluations = model_service_client.list_model_evaluations(\n",
        "    parent=models[0].resource_name\n",
@@ -933,7 +837,7 @@
      "outputs": [],
      "source": [
        "batch_predict_job = model.batch_predict(\n",
-        "    job_display_name=\"flowers_\" + UUID,\n",
+        "    job_display_name=\"flowers_\" + \"unique\",\n",
        "    gcs_source=gcs_input_uri,\n",
        "    gcs_destination_prefix=BUCKET_URI,\n",
        "    sync=False,\n",
@@ -1279,11 +1183,12 @@
      },
      "outputs": [],
      "source": [
-        "delete_bucket = False\n",
-        "\n",
        "# Delete the dataset using the Vertex dataset object\n",
        "dataset.delete()\n",
        "\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",
        "\n",
@@ -1293,8 +1198,10 @@
        "# Delete the batch prediction job\n",
        "batch_predict_job.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 = True  # Set True for deletion\n",
+        "if delete_bucket:\n",
+        "    ! gsutil -m rm -r $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/migration/sdk-automl-object-tracking-batch-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=\"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%2Fmigration%2Fsdk-automl-object-tracking-batch-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 style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/migration/sdk-automl-object-tracking-batch-prediction.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/migration/sdk-automl-object-tracking-batch-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",
+        "      <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",
@@ -75,15 +80,15 @@
      "source": [
        "### Objective\n",
        "\n",
-        "In this tutorial, you learn to use `AutoML` to train a video model and use `Vertex AI Batch Prediction` to do batch predictions.\n",
+        "In this tutorial, you learn to use AutoML to train a video model and use Vertex AI batch prediction to do batch predictions.\n",
        "\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
-        "- `AutoML`\n",
-        "- `Vertex AI Batch Prediction`\n",
-        "- `Vertex AI Model` resource\n",
-        "- `Vertex AI Endpoint` resource\n",
+        "- AutoML\n",
+        "- Vertex AI batch prediction\n",
+        "- Vertex AI model resource\n",
+        "- Vertex AI endpoint resource\n",
        "\n",
        "The steps performed include:\n",
        "\n",
@@ -125,19 +130,26 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "install_aip:mbsdk"
+        "id": "f0316df526f8"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the packages required for executing this notebook."
+        "## 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": "5143c7e7acab"
      },
      "outputs": [],
      "source": [
@@ -149,84 +161,97 @@
    {
      "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": "ee775571c2b5"
      },
      "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": "before_you_begin:nogpu"
+        "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": "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": "b49704f4eeca"
      },
      "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\"}"
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
@@ -260,64 +285,6 @@
        "UUID = generate_uuid()"
      ]
    },
-    {
-      "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": {
@@ -346,7 +313,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."
      ]
    },
    {
@@ -357,54 +324,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",
-        "import google.cloud.aiplatform as aip"
-      ]
-    },
-    {
-      "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": [
-        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
+        "! gsutil mb -l {LOCATION} {BUCKET_URI}"
      ]
    },
    {
@@ -486,12 +406,12 @@
        "id": "create_dataset:video,vot"
      },
      "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."
      ]
@@ -504,10 +424,12 @@
      },
      "outputs": [],
      "source": [
-        "dataset = aip.VideoDataset.create(\n",
+        "import os\n",
+        "\n",
+        "dataset = aiplatform.VideoDataset.create(\n",
        "    display_name=\"Traffic\" + \"_\" + UUID,\n",
        "    gcs_source=[IMPORT_FILE],\n",
-        "    import_schema_uri=aip.schema.dataset.ioformat.video.object_tracking,\n",
+        "    import_schema_uri=aiplatform.schema.dataset.ioformat.video.object_tracking,\n",
        ")\n",
        "if os.getenv(\"IS_TESTING\"):\n",
        "    import time\n",
@@ -560,13 +482,15 @@
      "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 perform two steps: \n",
+        "1) create a training pipeline, and \n",
+        "2) run the pipeline.\n",
        "\n",
        "#### Create training pipeline\n",
        "\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",
@@ -583,7 +507,7 @@
      },
      "outputs": [],
      "source": [
-        "dag = aip.AutoMLVideoTrainingJob(\n",
+        "dag = aiplatform.AutoMLVideoTrainingJob(\n",
        "    display_name=\"traffic_\" + UUID,\n",
        "    prediction_type=\"object_tracking\",\n",
        ")\n",
@@ -612,12 +536,12 @@
        "\n",
        "Next, you run the DAG to start 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 take upto 20 minutes."
      ]
@@ -707,11 +631,13 @@
      "outputs": [],
      "source": [
        "# Get model resource ID\n",
-        "models = aip.Model.list(filter=\"display_name=traffic_\" + UUID)\n",
+        "models = aiplatform.Model.list(filter=\"display_name=traffic_\" + UUID)\n",
        "\n",
        "# Get a reference to the Model Service client\n",
-        "client_options = {\"api_endpoint\": f\"{REGION}-aiplatform.googleapis.com\"}\n",
-        "model_service_client = aip.gapic.ModelServiceClient(client_options=client_options)\n",
+        "client_options = {\"api_endpoint\": f\"{LOCATION}-aiplatform.googleapis.com\"}\n",
+        "model_service_client = aiplatform.gapic.ModelServiceClient(\n",
+        "    client_options=client_options\n",
+        ")\n",
        "\n",
        "model_evaluations = model_service_client.list_model_evaluations(\n",
        "    parent=models[0].resource_name\n",
@@ -828,7 +754,7 @@
        "Now make a batch input file, which you store in your local Cloud Storage bucket. The batch input file can be either CSV or JSONL. You use JSONL in this tutorial. For JSONL file, you make one dictionary entry per line for each video. The dictionary contains the key/value pairs:\n",
        "\n",
        "- `content`: The Cloud Storage path to the video.\n",
-        "- `mimeType`: The content type. In our example, it is a `avi` file.\n",
+        "- `mimeType`: The content type. In our example, it is a avi file.\n",
        "- `timeSegmentStart`: The start timestamp in the video to do prediction on. *Note*, the timestamp must be specified as a string and followed by s (second), m (minute) or h (hour).\n",
        "- `timeSegmentEnd`: The end timestamp in the video to do prediction on."
      ]
@@ -32,25 +32,28 @@
        "# Vertex AI migration: AutoML tabular binary classification\n",
        "\n",
        "<table align=\"left\">\n",
-        "\n",
-        "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/migration/sdk-automl-tabular-binary-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",
+        "  <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://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/migration/sdk-automl-tabular-binary-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%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>\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/migration/sdk-automl-tabular-binary-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",
+        "  <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",
-        "</table>"
+        "  </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/>\n"
      ]
    },
    {
@@ -75,23 +78,23 @@
      "source": [
        "### Objective\n",
        "\n",
-        "In this tutorial, you create an AutoML tabular binary classification model and deploy for online prediction from a Python script 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 create an AutoML tabular binary classification model and deploy for online prediction from a Python script using the Vertex AI SDK. You can alternatively create and deploy models using the `gcloud` command-line tool or using the online Cloud Console.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
-        "- Vertex AI managed Datasets\n",
+        "- Vertex AI managed dataset\n",
        "- Vertex AI Training\n",
-        "- Vertex AI Endpoints\n",
+        "- Vertex AI endpoints\n",
        "- Vertex AI prediction\n",
        "\n",
        "The steps performed include:\n",
        "\n",
-        "- Create a Vertex `Dataset` resource.\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",
+        "- Deploy the model resource to a serving endpoint resource.\n",
        "- Make a prediction.\n",
-        "- Undeploy the `Model`"
+        "- Undeploy the model"
      ]
    },
    {
@@ -102,7 +105,7 @@
      "source": [
        "### Dataset\n",
        "\n",
-        "The dataset used for this tutorial is the [Bank Marketing](https://pantheon.corp.google.com/storage/browser/_details/cloud-ml-tables-data/bank-marketing.csv) . 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."
+        "The dataset used for this tutorial is the [Bank Marketing](https://pantheon.corp.google.com/storage/browser/_details/cloud-ml-tables-data/bank-marketing.csv) . This dataset doesn't require any feature engineering. The version of the dataset you use in this tutorial is stored in a public Cloud Storage bucket."
      ]
    },
    {
@@ -118,11 +121,9 @@
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
-        "Learn about [Vertex AI\n",
-        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
-        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
-        "Calculator](https://cloud.google.com/products/calculator/)\n",
-        "to generate a cost estimate based on your projected usage."
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and \n",
+        "[Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the \n",
+        "[Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
@@ -131,9 +132,8 @@
        "id": "install_aip:mbsdk"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook. "
+        "### Get started\n",
+        "Install Vertex AI SDK for Python and other required packages"
      ]
    },
    {
@@ -144,7 +144,8 @@
      },
      "outputs": [],
      "source": [
-        "! pip3 install --upgrade --quiet google-cloud-aiplatform fsspec gcsfs"
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform fsspec gcsfs \\\n",
+        "                                 pandas"
      ]
    },
    {
@@ -153,7 +154,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,159 +166,74 @@
      },
      "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": "4de1bd77992b"
      },
      "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": "befa6ca14bc0"
      },
      "source": [
-        "### Set your project ID\n",
-        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "* Run `gcloud config list`.\n",
-        "* Run `gcloud projects list`.\n",
-        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+        "### Authenticate your notebook environment (Colab only)\n",
+        "Authenticate your environment on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "set_project_id"
+        "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()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bd0e79ceaea2"
+      },
+      "source": [
+        "### 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": "8940d70dfdef"
      },
      "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": "timestamp"
-      },
-      "source": [
-        "#### UUID\n",
-        "\n",
-        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a uuid for each instance session, and append it onto the name of resources you create in this tutorial."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "timestamp"
-      },
-      "outputs": [],
-      "source": [
-        "import random\n",
-        "import string\n",
-        "\n",
-        "\n",
-        "# Generate a uuid of a specifed length(default=8)\n",
-        "def generate_uuid(length: int = 8) -> str:\n",
-        "    return \"\".join(random.choices(string.ascii_lowercase + string.digits, k=length))\n",
-        "\n",
-        "\n",
-        "UUID = generate_uuid()"
-      ]
-    },
-    {
-      "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\"}"
      ]
    },
    {
@@ -347,7 +264,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."
      ]
    },
    {
@@ -358,7 +275,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} {BUCKET_URI}"
      ]
    },
    {
@@ -378,7 +295,7 @@
      },
      "outputs": [],
      "source": [
-        "import google.cloud.aiplatform as aip\n",
+        "import google.cloud.aiplatform as aiplatform\n",
        "import pandas as pd"
      ]
    },
@@ -401,7 +318,7 @@
      },
      "outputs": [],
      "source": [
-        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
+        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -436,9 +353,9 @@
        "\n",
        "This tutorial uses 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",
+        "Start by doing a quick peek at the data. 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."
+        "You also need for training to know the heading name of the label column, which is save as `label_column`. For this dataset, it's the last column in the CSV file."
      ]
    },
    {
@@ -488,13 +405,13 @@
      "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 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 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."
+        "This operation may takes several minutes."
      ]
    },
    {
@@ -505,8 +422,8 @@
      },
      "outputs": [],
      "source": [
-        "dataset = aip.TabularDataset.create(\n",
-        "    display_name=\"Bank Marketing\" + \"_\" + UUID, gcs_source=[IMPORT_FILE]\n",
+        "dataset = aiplatform.TabularDataset.create(\n",
+        "    display_name=\"Bank Marketing\" + \"_\" + \"unique\", gcs_source=[IMPORT_FILE]\n",
        ")\n",
        "\n",
        "print(dataset.resource_name)"
@@ -560,24 +477,24 @@
        "\n",
        "An AutoML training pipeline is created with 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",
-        "  - `classification`: A tabuar classification model.\n",
-        "  - `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",
+        "- display_name: The human readable name for the `TrainingJob` resource.\n",
+        "- optimization_prediction_type: The type of the task to train the model for.\n",
+        "  - classification: A tabuar classification model.\n",
+        "  - 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",
-        "    - `minimize-log-loss`\n",
-        "    - `maximize-au-roc`\n",
-        "    - `maximize-au-prc`\n",
-        "    - `maximize-precision-at-recall`\n",
-        "    - `maximize-recall-at-precision`\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",
-        "    - `minimize-log-loss`\n",
+        "    - minimize-log-loss\n",
        "  - regression:\n",
-        "    - `minimize-rmse`\n",
-        "    - `minimize-mae`\n",
-        "    - `minimize-rmsle`\n",
+        "    - minimize-rmse\n",
+        "    - minimize-mae\n",
+        "    - minimize-rmsle\n",
        "\n",
        "The instantiated object is the DAG (directed acyclic graph) for the training pipeline."
      ]
@@ -590,8 +507,8 @@
      },
      "outputs": [],
      "source": [
-        "job = aip.AutoMLTabularTrainingJob(\n",
-        "    display_name=\"bank_\" + UUID,\n",
+        "job = aiplatform.AutoMLTabularTrainingJob(\n",
+        "    display_name=\"bank_\" + \"unique\",\n",
        "    optimization_prediction_type=\"classification\",\n",
        "    optimization_objective=\"minimize-log-loss\",\n",
        ")\n",
@@ -620,18 +537,18 @@
        "\n",
        "Next, you run the DAG to start the training job by invoking the method `run`, 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",
-        "- `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",
+        "- 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 can't further improve on the model objective measurements.\n",
        "\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."
      ]
    },
    {
@@ -644,7 +561,7 @@
      "source": [
        "model = job.run(\n",
        "    dataset=dataset,\n",
-        "    model_display_name=\"bank_\" + UUID,\n",
+        "    model_display_name=\"bank_\" + \"unique\",\n",
        "    training_fraction_split=0.6,\n",
        "    validation_fraction_split=0.2,\n",
        "    test_fraction_split=0.2,\n",
@@ -720,11 +637,13 @@
      "outputs": [],
      "source": [
        "# Get model resource ID\n",
-        "models = aip.Model.list(filter=\"display_name=bank_\" + UUID)\n",
+        "models = aiplatform.Model.list(filter=\"display_name=bank_\" + \"unique\")\n",
        "\n",
        "# Get a reference to the Model Service client\n",
-        "client_options = {\"api_endpoint\": f\"{REGION}-aiplatform.googleapis.com\"}\n",
-        "model_service_client = aip.gapic.ModelServiceClient(client_options=client_options)\n",
+        "client_options = {\"api_endpoint\": f\"{LOCATION}-aiplatform.googleapis.com\"}\n",
+        "model_service_client = aiplatform.gapic.ModelServiceClient(\n",
+        "    client_options=client_options\n",
+        ")\n",
        "\n",
        "model_evaluations = model_service_client.list_model_evaluations(\n",
        "    parent=models[0].resource_name\n",
@@ -799,7 +718,7 @@
      "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 -- we just want to demonstrate how to make a prediction."
+        "You use synthetic data as a test data items. Don't be concerned that you are using synthetic data -- you just want to demonstrate how to make a prediction."
      ]
    },
    {
@@ -847,7 +766,7 @@
      "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, with the following parameters:\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",
        "\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",
@@ -866,7 +785,7 @@
      "outputs": [],
      "source": [
        "batch_predict_job = model.batch_predict(\n",
-        "    job_display_name=\"bank_\" + UUID,\n",
+        "    job_display_name=\"bank_\" + \"unique\",\n",
        "    gcs_source=gcs_input_uri,\n",
        "    gcs_destination_prefix=BUCKET_URI,\n",
        "    instances_format=\"csv\",\n",
@@ -1093,7 +1012,7 @@
      "source": [
        "### Make test item\n",
        "\n",
-        "You 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."
+        "You use synthetic data as a test data item. Don't be concerned that you are using synthetic data -- you just want to demonstrate how to make a prediction."
      ]
    },
    {
@@ -1149,7 +1068,7 @@
        "- `ids`: The internal assigned unique identifiers for each prediction request.\n",
        "- `displayNames`: The class names for each class label.\n",
        "- `confidences`: The predicted confidence, between 0 and 1, per class label.\n",
-        "- `deployed_model_id`: The Vertex AI identifier for the deployed `Model` resource which did the predictions."
+        "- `deployed_model_id`: The Vertex AI identifier for the deployed model resource which did the predictions."
      ]
    },
    {
@@ -1185,7 +1104,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, undeploy the model from the `Endpoint` resouce. This deprovisions all compute resources and ends billing for the deployed model."
      ]
    },
    {
@@ -1228,8 +1147,6 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
        "# Delete the dataset using the Vertex dataset object\n",
        "dataset.delete()\n",
        "\n",
@@ -1245,8 +1162,12 @@
        "# Delete the batch prediction job using the Vertex batch prediction object\n",
        "batch_predict_job.delete()\n",
        "\n",
+        "# remove locally generated files\n",
+        "! rm -r batch.csv\n",
+        "! rm -r tmp.csv\n",
+        "\n",
        "delete_bucket = False\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
@@ -23,6 +23,16 @@
        "# limitations under the License."
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "90a6064798e9"
+      },
+      "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": "90a6064798e9"
+      },
+      "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": {
@@ -33,23 +33,26 @@
        "\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/migration/sdk-automl-video-classification-batch-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/migration/sdk-automl-video-classification-batch-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%2Fgithub.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fblob%2Fmain%2Fnotebooks%2Fofficial%2Fmigration%2Fsdk-automl-video-classification-batch-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 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/migration/sdk-automl-video-classification-batch-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",
+        "      <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/migration/sdk-automl-video-classification-batch-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>\n",
        "<br/><br/><br/>"
      ]
@@ -75,15 +78,15 @@
      "source": [
        "### Objective\n",
        "\n",
-        "In this tutorial, you learn to use `AutoML` to train a video model and use `Vertex AI Batch Prediction` to do batch predictions.\n",
+        "In this tutorial, you learn to use AutoML to train a video model and use Vertex AI batch prediction to do batch predictions.\n",
        "\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
-        "- `AutoML`\n",
-        "- `Vertex AI Batch Prediction`\n",
-        "- `Vertex AI Model` resource\n",
-        "- `Vertex AI Endpoint` resource\n",
+        "- AutoML\n",
+        "- Vertex AI batch prediction\n",
+        "- Vertex AI model resource\n",
+        "- Vertex AI endpoint resource\n",
        "\n",
        "The steps performed include:\n",
        "\n",
@@ -99,7 +102,7 @@
      "source": [
        "### Dataset\n",
        "\n",
-        "The dataset used for this tutorial is the [Human Motion dataset](https://TODO) from [MIT](http://cbcl.mit.edu/publications/ps/Kuehne_etal_iccv11.pdf). The version of the dataset you will use in this tutorial is stored in a public Cloud Storage bucket."
+        "The dataset used for this tutorial is the [Human Motion dataset](https://TODO) from [MIT](http://cbcl.mit.edu/publications/ps/Kuehne_etal_iccv11.pdf). The version of the dataset you use in this tutorial is stored in a public Cloud Storage bucket."
      ]
    },
    {
@@ -125,12 +128,19 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "install_aip:mbsdk"
+        "id": "f0316df526f8"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the packages required for executing this notebook."
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "a2c2cb2109a0"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
@@ -143,7 +153,9 @@
      "source": [
        "import os\n",
        "\n",
-        "! pip3 install --upgrade --quiet google-cloud-aiplatform \n",
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform\n",
+        "\n",
+        "! pip3 install --quiet tensorflow==2.15.1\n",
        "\n",
        "if os.getenv(\"IS_TESTING\"):\n",
        "    ! pip3 install --upgrade --quiet google-cloud-storage tensorflow"
@@ -152,48 +164,80 @@
    {
      "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": "ee775571c2b5"
      },
      "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": "before_you_begin:nogpu"
+        "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": "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)."
      ]
    },
    {
@@ -207,29 +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\"}"
+        "LOCATION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
@@ -240,7 +264,8 @@
      "source": [
        "#### UUID\n",
        "\n",
-        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a uuid for each instance session, and append it onto the name of resources you create in this tutorial."
+        "To avoid name collisions between users on created resources, create a uuid for each session instance. Append these uuids to the respective names of the resources \n",
+        "created in this tutorial."
      ]
    },
    {
@@ -263,64 +288,6 @@
        "UUID = generate_uuid()"
      ]
    },
-    {
-      "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": {
@@ -360,7 +327,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} {BUCKET_URI}"
      ]
    },
    {
@@ -386,7 +353,7 @@
        "import os\n",
        "import sys\n",
        "\n",
-        "import google.cloud.aiplatform as aip"
+        "from google.cloud import aiplatform"
      ]
    },
    {
@@ -408,7 +375,7 @@
      },
      "outputs": [],
      "source": [
-        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
+        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -521,10 +488,10 @@
      },
      "outputs": [],
      "source": [
-        "dataset = aip.VideoDataset.create(\n",
+        "dataset = aiplatform.VideoDataset.create(\n",
        "    display_name=\"MIT Human Motion\" + \"_\" + UUID,\n",
        "    gcs_source=[IMPORT_FILE],\n",
-        "    import_schema_uri=aip.schema.dataset.ioformat.video.classification,\n",
+        "    import_schema_uri=aiplatform.schema.dataset.ioformat.video.classification,\n",
        ")\n",
        "\n",
        "print(dataset.resource_name)"
@@ -596,7 +563,7 @@
      },
      "outputs": [],
      "source": [
-        "dag = aip.AutoMLVideoTrainingJob(\n",
+        "dag = aiplatform.AutoMLVideoTrainingJob(\n",
        "    display_name=\"hmdb_\" + UUID,\n",
        "    prediction_type=\"classification\",\n",
        ")\n",
@@ -744,11 +711,13 @@
      "outputs": [],
      "source": [
        "# Get model resource ID\n",
-        "models = aip.Model.list(filter=\"display_name=hmdb_\" + UUID)\n",
+        "models = aiplatform.Model.list(filter=\"display_name=hmdb_\" + UUID)\n",
        "\n",
        "# Get a reference to the Model Service client\n",
-        "client_options = {\"api_endpoint\": f\"{REGION}-aiplatform.googleapis.com\"}\n",
-        "model_service_client = aip.gapic.ModelServiceClient(client_options=client_options)\n",
+        "client_options = {\"api_endpoint\": f\"{LOCATION}-aiplatform.googleapis.com\"}\n",
+        "model_service_client = aiplatform.gapic.ModelServiceClient(\n",
+        "    client_options=client_options\n",
+        ")\n",
        "\n",
        "model_evaluations = model_service_client.list_model_evaluations(\n",
        "    parent=models[0].resource_name\n",
@@ -823,7 +792,7 @@
      "source": [
        "### Get test item(s)\n",
        "\n",
-        "Now do a batch prediction to your Vertex model. You will use arbitrary examples out of the dataset as a test items. Don't be concerned that the examples were likely used in training the model -- we just want to demonstrate how to make a prediction."
+        "Now do a batch prediction to your Vertex model. You use arbitrary examples out of the dataset as a test items. Don't be concerned that the examples were likely used in training the model -- the purpose here is to demonstrate how to make a prediction."
      ]
    },
    {
@@ -856,10 +825,10 @@
      "source": [
        "### Make a batch input file\n",
        "\n",
-        "Now make a batch input file, which you store in your local Cloud Storage bucket. The batch input file can be either CSV or JSONL. You will use JSONL in this tutorial. For JSONL file, you make one dictionary entry per line for each video. The dictionary contains the key/value pairs:\n",
+        "Now make a batch input file, which you store in your local Cloud Storage bucket. The batch input file can be either CSV or JSONL. You use JSONL in this tutorial. For JSONL file, you make one dictionary entry per line for each video. The dictionary contains the key/value pairs:\n",
        "\n",
        "- `content`: The Cloud Storage path to the video.\n",
-        "- `mimeType`: The content type. In our example, it is a `avi` file.\n",
+        "- `mimeType`: The content type. In our example, it's a `avi` file.\n",
        "- `timeSegmentStart`: The start timestamp in the video to do prediction on. *Note*, the timestamp must be specified as a string and followed by s (second), m (minute) or h (hour).\n",
        "- `timeSegmentEnd`: The end timestamp in the video to do prediction on."
      ]
@@ -910,7 +879,7 @@
        "- `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",
        "- `gcs_destination_prefix`: The Cloud Storage location for storing the batch prediction resuls.\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 blocks while waiting for the asynchronous batch job to complete."
      ]
    },
    {
@@ -1100,13 +1069,12 @@
        "model.delete()\n",
        "\n",
        "# Delete the AutoML or Pipeline trainig job\n",
-        "\n",
        "dag.delete()\n",
        "\n",
        "# Delete the batch prediction job using the Vertex batch prediction object\n",
        "batch_predict_job.delete()\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil rm -r $BUCKET_URI"
      ]
    }
@@ -32,24 +32,26 @@
        "# Vertex AI migration: Custom image classification with a custom training container\n",
        "\n",
        "<table align=\"left\">\n",
-        "\n",
-        "  <td>\n",
-        "<a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/migration/sdk-custom-image-classification-custom-container.ipynb\" target='_blank'>\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/migration/sdk-custom-image-classification-custom-container.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://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/migration/sdk-custom-image-classification-custom-container.ipynb\" target='_blank'>\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%2Fmigration%2Fsdk-custom-image-classification-custom-container.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/migration/sdk-custom-image-classification-custom-container.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>\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/migration/sdk-custom-image-classification-custom-container.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 style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/migration/sdk-custom-image-classification-custom-container.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>\n",
        "</table>"
      ]
    },
@@ -78,23 +80,23 @@
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
-        "- *Vertex AI Training*\n",
-        "- *Vertex AI Model Registry*\n",
-        "- *Vertex AI Batch Predictions*\n",
-        "- *Vertex AI Endpoints*\n",
+        "- Vertex AI Training\n",
+        "- Vertex AI Model Registry\n",
+        "- Vertex AI batch Predictions\n",
+        "- Vertex AI endpoints\n",
        "\n",
        "\n",
        "The steps performed include:\n",
        "\n",
-        "- *Package the training code into a python application.*\n",
-        "- *Containerize the training application using Cloud Build and Artifact Registry.*\n",
-        "- *Create a custom container training job in Vertex AI and run it.*\n",
-        "- *Evaluate the model generated from the training job.*\n",
-        "- *Create a model resource for the trained model in Vertex AI Model Registry.*\n",
-        "- *Run a Vertex AI batch prediction job.*\n",
-        "- *Deploy the model resource to a Vertex AI Endpoint.*\n",
-        "- *Run a online prediction job on the model resource.*\n",
-        "- *Clean up the resources created.*"
+        "- Package the training code into a python application.\n",
+        "- Containerize the training application using Cloud Build and Artifact Registry.\n",
+        "- Create a custom container training job in Vertex AI and run it.\n",
+        "- Evaluate the model generated from the training job.\n",
+        "- Create a model resource for the trained model in Vertex AI Model Registry.\n",
+        "- Run a Vertex AI batch prediction job.\n",
+        "- Deploy the model resource to a Vertex AI endpoint.\n",
+        "- Run a online prediction job on the model resource.\n",
+        "- Clean up the resources created."
      ]
    },
    {
@@ -130,15 +132,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": "iMHz63rPbq6P"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the packages required for executing this notebook."
+        "### Install Vertex AI SDK for Python and other required packages"
      ]
    },
    {
@@ -150,106 +159,124 @@
      "outputs": [],
      "source": [
        "! pip3 install --upgrade google-cloud-aiplatform \\\n",
-        "                         google-cloud-storage -q\n",
+        "                         google-cloud-storage \\\n",
+        "                         opencv-python-headless \\\n",
+        "                         tensorflow==2.15.1 -q\n",
        "\n",
        "import os\n",
        "if os.getenv(\"IS_TESTING\"):\n",
        "    ! apt-get update && apt-get install -y python3-opencv-headless\n",
        "    ! apt-get install -y libgl1-mesa-dev\n",
        "    ! pip3 install --upgrade opencv-python-headless -q\n",
-        "    ! pip3 install tensorflow==2.9 -q"
+        "    ! pip3 install tensorflow==2.15.1 -q"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "58707a750154"
+        "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": "f200f10a1da3"
+        "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": "WReHDGG5g0XY"
+        "id": "4a2b7b59bbf7"
      },
      "source": [
-        "#### 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>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f82e28c631cc"
+      },
+      "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."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "oM1iC_MfAts1"
+        "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": "91842ef41bbd"
+      },
+      "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": "294fe4e5a671"
      },
      "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": "0bd7a5e762cc"
      },
      "source": [
-        "#### Set the region\n",
-        "\n",
-        "**Optional**: Update the 'REGION' variable to specify the region that you want to use. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
+        "### UUID\n",
+        "If you're in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, create a uuid for each instance session, and append it onto the name of resources you create in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "nsN5NJKSu-GU"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "timestamp"
-      },
-      "source": [
-        "#### UUID\n",
-        "\n",
-        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a uuid for each instance session, and append it onto the name of resources you create in this tutorial.\n"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "timestamp"
+        "id": "3ee72715c0fd"
      },
      "outputs": [],
      "source": [
@@ -265,67 +292,6 @@
        "UUID = generate_uuid()"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "sBCra4QMA2wR"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "To authenticate your Google Cloud account, follow the instructions for your Jupyter environment:\n",
-        "\n",
-        "**1. Vertex AI Workbench**\n",
-        "<br>You are already authenticated.\n",
-        "\n",
-        "**2. Local JupyterLab instance**\n",
-        "<br>Uncomment and run the following code:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "254614fa0c46"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ef21552ccea8"
-      },
-      "source": [
-        "**3. Colab**\n",
-        "<br>Uncomment and run the following code:"
-      ]
-    },
-    {
-      "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": "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": {
@@ -365,7 +331,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
+        "! gsutil mb -l $LOCATION -p $PROJECT_ID $BUCKET_URI"
      ]
    },
    {
@@ -416,7 +382,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)"
      ]
    },
    {
@@ -434,9 +400,9 @@
        "\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",
-        "    (aiplatform.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",
@@ -519,7 +485,7 @@
        " - `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*."
      ]
    },
    {
@@ -530,6 +496,8 @@
      },
      "outputs": [],
      "source": [
+        "import os\n",
+        "\n",
        "MACHINE_TYPE = \"n1-standard\"\n",
        "\n",
        "VCPU = \"4\"\n",
@@ -753,7 +721,7 @@
        "For this step, your Dockerfile does:\n",
        "1. Install a pre-defined container image from TensorFlow repository for deep learning images.\n",
        "2. Copies in the Python training code, to be shown subsequently.\n",
-        "3. Sets the entry into the Python training script as `trainer/task.py`. Note that the `.py` is dropped in the ENTRYPOINT command, as it is implied."
+        "3. Sets the entry into the Python training script as `trainer/task.py`. Note that the `.py` is dropped in the ENTRYPOINT command, as it's implied."
      ]
    },
    {
@@ -831,7 +799,7 @@
      "source": [
        "REPOSITORY = \"my-docker-repo\"\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"
      ]
@@ -863,7 +831,7 @@
        "\n",
        "TAG = \"latest\"\n",
        "TRAIN_IMAGE = (\n",
-        "    f\"{REGION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{CONTAINER_NAME}:{TAG}\"\n",
+        "    f\"{LOCATION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{CONTAINER_NAME}:{TAG}\"\n",
        ")"
      ]
    },
@@ -898,7 +866,7 @@
      "outputs": [],
      "source": [
        "%cd custom\n",
-        "!gcloud builds submit --region={REGION} --tag=$TRAIN_IMAGE\n",
+        "!gcloud builds submit --region={LOCATION} --tag=$TRAIN_IMAGE\n",
        "%cd .."
      ]
    },
@@ -957,7 +925,7 @@
      "outputs": [],
      "source": [
        "job = aiplatform.CustomContainerTrainingJob(\n",
-        "    display_name=\"cifar10_\" + UUID, container_uri=TRAIN_IMAGE\n",
+        "    display_name=\"cifar10_\" + \"unique\", container_uri=TRAIN_IMAGE\n",
        ")\n",
        "\n",
        "print(job)"
@@ -1079,7 +1047,7 @@
        "\n",
        "For model evaluation, you 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",
        "\n",
-        "You don't need the training data, and hence why it is loaded as `(_, _)`.\n",
+        "You don't need the training data, and hence why it's loaded as `(_, _)`.\n",
        "\n",
        "Before you run the data through evaluation, you need to preprocess it:\n",
        "\n",
@@ -1132,7 +1100,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",
@@ -1250,7 +1218,7 @@
      "outputs": [],
      "source": [
        "model = aiplatform.Model.upload(\n",
-        "    display_name=\"cifar10_\" + UUID,\n",
+        "    display_name=\"cifar10_\" + \"unique\",\n",
        "    artifact_uri=MODEL_DIR,\n",
        "    serving_container_image_uri=DEPLOY_IMAGE,\n",
        "    sync=False,\n",
@@ -1269,7 +1237,7 @@
        "\n",
        "### Get test items\n",
        "\n",
-        "You will use examples out of the test (holdout) portion of the dataset as a test items."
+        "You use examples out of the test (holdout) portion of the dataset as a test items."
      ]
    },
    {
@@ -1294,7 +1262,7 @@
      },
      "source": [
        "### Prepare the request content\n",
-        "You are going to send the CIFAR10 images as compressed JPG image, instead of the raw uncompressed bytes:\n",
+        "You're going to send the CIFAR10 images as compressed JPG image, instead of the raw uncompressed bytes:\n",
        "\n",
        "- `cv2.imwrite`: Use openCV to write the uncompressed image to disk as a compressed JPEG image.\n",
        " - Denormalize the image data from \\[0,1) range back to [0,255).\n",
@@ -1418,7 +1386,7 @@
        "MAX_NODES = 1\n",
        "\n",
        "batch_predict_job = model.batch_predict(\n",
-        "    job_display_name=\"cifar10_\" + UUID,\n",
+        "    job_display_name=\"cifar10_\" + \"unique\",\n",
        "    gcs_source=gcs_input_uri,\n",
        "    gcs_destination_prefix=BUCKET_URI,\n",
        "    instances_format=\"jsonl\",\n",
@@ -1535,7 +1503,7 @@
      },
      "outputs": [],
      "source": [
-        "DEPLOYED_NAME = \"cifar10-\" + UUID\n",
+        "DEPLOYED_NAME = \"cifar10-\" + \"unique\"\n",
        "\n",
        "TRAFFIC_SPLIT = {\"0\": 100}\n",
        "\n",
@@ -1671,7 +1639,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."
      ]
    },
    {
@@ -1709,8 +1677,6 @@
      },
      "outputs": [],
      "source": [
-        "delete_bucket = False\n",
-        "\n",
        "# Delete the model using the Vertex model object\n",
        "model.delete()\n",
        "\n",
@@ -1724,10 +1690,19 @@
        "batch_predict_job.delete()\n",
        "\n",
        "# Delete artifact repository\n",
-        "! gcloud artifacts repositories delete $REPOSITORY --location=$REGION --quiet\n",
+        "! gcloud artifacts repositories delete $REPOSITORY --location=$LOCATION --quiet\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\n",
+        "\n",
+        "# remove the local users and movies avro files\n",
+        "! rm -rf custom\n",
+        "! rm -f custom.tar.gz\n",
+        "! rm tmp.jpg\n",
+        "! rm tmp1.jpg\n",
+        "! rm tmp2.jpg"
      ]
    }
  ],
@@ -32,26 +32,27 @@
        "# Vertex AI Migration: Custom image classification with a pre-built training container\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/migration/sdk-custom-image-classification-prebuilt-container.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%2Fmigration%2Fsdk-custom-image-classification-prebuilt-container.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/migration/sdk-custom-image-classification-prebuilt-container.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/migration/sdk-custom-image-classification-prebuilt-container.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/migration/sdk-custom-image-classification-prebuilt-container.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>\n",
-        "<br/><br/><br/>"
+        "</table>"
      ]
    },
    {
@@ -64,7 +65,7 @@
        "\n",
        "This tutorial demonstrates how to use the Vertex AI SDK for Python to train using a pre-built container and deploy a custom image classification model for online and batch prediction.\n",
        "\n",
-        "Learn more about [Migrate to Vertex AI](https://cloud.google.com/vertex-ai/docs/start/migrating-to-vertex-ai) and [Custom training](https://cloud.google.com/vertex-ai/docs/training/custom-training)."
+        "Learn more about [Migrate to Vertex AI](https://cloud.google.com/vertex-ai/docs/start/migrating-to-vertex-ai) and [Custom training overview](https://cloud.google.com/vertex-ai/docs/training/custom-training)."
      ]
    },
    {
@@ -81,22 +82,21 @@
        "\n",
        "- Vertex AI Training\n",
        "- Vertex AI Model Registry\n",
-        "- Vertex AI Predictions\n",
-        "- Vertex AI Batch Predictions\n",
-        "- Vertex AI Endpoints\n",
+        "- Vertex AI batch predictions\n",
+        "- Vertex AI endpoints\n",
        "\n",
        "\n",
        "The steps performed include:\n",
        "\n",
-        "- *Package the training code into a python application.*\n",
-        "- *Containerize the training application using Cloud Build and Artifact Registry.*\n",
-        "- *Create a custom container training job in Vertex AI and run it.*\n",
-        "- *Evaluate the model generated from the training job.*\n",
-        "- *Create a model resource for the trained model in Vertex AI Model Registry.*\n",
-        "- *Run a Vertex AI batch prediction job.*\n",
-        "- *Deploy the model resource to a Vertex AI Endpoint.*\n",
-        "- *Run a online prediction job on the model resource.*\n",
-        "- *Clean up the resources created.*"
+        "- Package the training code into a python application.\n",
+        "- Containerize the training application using Cloud Build and Artifact Registry.\n",
+        "- Create a custom container training job in Vertex AI and run it.\n",
+        "- Evaluate the model generated from the training job.\n",
+        "- Create a model resource for the trained model in Vertex AI Model Registry.\n",
+        "- Run a Vertex AI batch prediction job.\n",
+        "- Deploy the model resource to a Vertex AI endpoint.\n",
+        "- Run a online prediction job on the model resource.\n",
+        "- Clean up the resources created."
      ]
    },
    {
@@ -107,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 you use in this tutorial is the one that's available from TensorFlow SDK. The trained model classifies images into one of ten categories: *airplane*, *automobile*, *bird*, *cat*, *deer*, *dog*, *frog*, *horse*, *ship*, or *truck*."
      ]
    },
    {
@@ -133,12 +133,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"
      ]
    },
    {
@@ -154,54 +161,86 @@
        "! apt-get update && apt-get install -y python3-opencv-headless\n",
        "! apt-get install -y libgl1-mesa-dev\n",
        "! pip3 install --upgrade opencv-python-headless  \n",
-        "! pip3 install --upgrade tensorflow  "
+        "! pip3 install --upgrade tensorflow==2.15.1  "
      ]
    },
    {
      "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": "before_you_begin:nogpu"
+        "id": "SbmM4z7FOBpM"
      },
      "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": "before_you_begin:nogpu"
+        "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": "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)."
      ]
    },
    {
@@ -213,122 +252,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": "timestamp"
-      },
-      "source": [
-        "#### UUID\n",
-        "\n",
-        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a uuid for each instance session, and append it onto the name of resources you create in this tutorial."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "e87d5856317d"
-      },
-      "outputs": [],
-      "source": [
-        "import random\n",
-        "import string\n",
-        "\n",
-        "\n",
-        "# Generate a uuid of a specifed length(default=8)\n",
-        "def generate_uuid(length: int = 8) -> str:\n",
-        "    return \"\".join(random.choices(string.ascii_lowercase + string.digits, k=length))\n",
-        "\n",
-        "\n",
-        "UUID = generate_uuid()"
-      ]
-    },
-    {
-      "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": {
@@ -357,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."
      ]
    },
    {
@@ -368,7 +297,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
@@ -377,9 +306,6 @@
        "id": "setup_vars"
      },
      "source": [
-        "### Set up variables\n",
-        "\n",
-        "Next, set up some variables used throughout the tutorial.\n",
        "### Import libraries and define constants"
      ]
    },
@@ -391,7 +317,10 @@
      },
      "outputs": [],
      "source": [
-        "import google.cloud.aiplatform as aip"
+        "import numpy as np\n",
+        "import tensorflow as tf\n",
+        "from google.cloud import aiplatform\n",
+        "from tensorflow.keras.datasets import cifar10"
      ]
    },
    {
@@ -400,7 +329,7 @@
        "id": "init_aip:mbsdk"
      },
      "source": [
-        "## Initialize Vertex SDK for Python\n",
+        "### Initialize Vertex SDK for Python\n",
        "\n",
        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
      ]
@@ -413,7 +342,7 @@
      },
      "outputs": [],
      "source": [
-        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
+        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -426,16 +355,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'd 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 [here](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators) hardware accelerator support for your location\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."
      ]
    },
    {
@@ -504,19 +433,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 used 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."
      ]
    },
    {
@@ -544,11 +473,12 @@
        "id": "examine_training_package"
      },
      "source": [
-        "### Examine the training package\n",
+        "## Tutorial\n",
        "\n",
-        "#### 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",
+        "### Package layout\n",
+        "\n",
+        "Before you start the training, take a 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",
@@ -560,11 +490,11 @@
        "\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",
-        "In the following cells, you assemble the training package."
+        "In the following cells, assemble the training package."
      ]
    },
    {
@@ -602,11 +532,13 @@
        "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, write the contents of the training script *task.py*.\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",
+        "To summarize, the script performs the following steps:\n",
+        "\n",
+        "- Gets 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",
        "- Builds a model using TF.Keras model API.\n",
        "- Compiles the model (`compile()`).\n",
@@ -730,9 +662,9 @@
        "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."
+        "Next, package the training folder into a compressed tar ball, and then store it in your Cloud Storage bucket."
      ]
    },
    {
@@ -755,16 +687,13 @@
        "id": "train_a_model:migration"
      },
      "source": [
-        "## Train a model"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "custom_create:migration,new,mbsdk,prebuilt"
-      },
-      "source": [
-        "### [training.create-python-pre-built-container](https://cloud.google.com/vertex-ai/docs/training/create-python-pre-built-container)"
+        "### Create and run custom training job\n",
+        "\n",
+        "Learn more about how to [Create a Python training application for a prebuilt container](https://cloud.google.com/vertex-ai/docs/training/create-python-pre-built-container)\n",
+        "\n",
+        "To train a custom model, you perform two steps:\n",
+        "1) Create a custom training job.\n",
+        "2) Specify your training parameters and run the job."
      ]
    },
    {
@@ -773,14 +702,9 @@
        "id": "create_custom_training_job:mbsdk,no_model"
      },
      "source": [
-        "### 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",
        "\n",
-        "A custom training job is created with the `CustomTrainingJob` class, with the following parameters:\n",
+        "A custom training job is created using the `CustomTrainingJob` class, with the following parameters:\n",
        "\n",
        "- `display_name`: The human readable name for the custom training job.\n",
        "- `container_uri`: The training container image.\n",
@@ -796,8 +720,8 @@
      },
      "outputs": [],
      "source": [
-        "job = aip.CustomTrainingJob(\n",
-        "    display_name=\"cifar10_\" + UUID,\n",
+        "job = aiplatform.CustomTrainingJob(\n",
+        "    display_name=\"cifar10-unique\",\n",
        "    script_path=\"custom/trainer/task.py\",\n",
        "    container_uri=TRAIN_IMAGE,\n",
        "    requirements=[\"gcsfs==0.7.1\", \"tensorflow-datasets==4.4\"],\n",
@@ -825,7 +749,7 @@
      "source": [
        "#### 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, 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",
@@ -833,7 +757,7 @@
        "- `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."
+        "- `sync`: Set **True** to wait until the completion of the job."
      ]
    },
    {
@@ -844,7 +768,7 @@
      },
      "outputs": [],
      "source": [
-        "MODEL_DIR = \"{}/{}\".format(BUCKET_URI, UUID)\n",
+        "MODEL_DIR = \"{}/{}\".format(BUCKET_URI, \"unique\")\n",
        "\n",
        "EPOCHS = 20\n",
        "STEPS = 100\n",
@@ -890,18 +814,9 @@
        "id": "run_custom_job:mbsdk,no_model"
      },
      "source": [
-        "### Wait for completion of custom training job\n",
+        "#### Wait for completion of custom training job\n",
        "\n",
-        "Next, wait for the custom training job to complete. Alternatively, one can set the parameter `sync` to `True` in the `run()` methid to block until the custom training job is completed."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "evaluate_the_model:migration"
-      },
-      "source": [
-        "## Evaluate the model"
+        "Next, wait for the custom training job to complete. Alternatively, you can set the parameter `sync` to `True` in the `run()` method to block until the custom training job is completed."
      ]
    },
    {
@@ -910,11 +825,11 @@
        "id": "load_saved_model"
      },
      "source": [
-        "## Load the saved model\n",
+        "### Load the saved model\n",
        "\n",
-        "Your model is stored in a TensorFlow SavedModel format in a Cloud Storage bucket. Now load it from the Cloud Storage bucket, and then you can do some things, like evaluate the model, and do a prediction.\n",
+        "Your model is stored in a TensorFlow SavedModel format in a Cloud Storage bucket. Now, load the model from the Cloud Storage bucket and run model evaluation, preparing it for prediction requests.\n",
        "\n",
-        "To load, you use the TF.Keras `model.load_model()` method passing it the Cloud Storage path where the model is saved -- specified by `MODEL_DIR`."
+        "To load, use the TF.Keras `model.load_model()` method passing it the Cloud Storage path where the model is saved -- specified by `MODEL_DIR`."
      ]
    },
    {
@@ -925,8 +840,6 @@
      },
      "outputs": [],
      "source": [
-        "import tensorflow as tf\n",
-        "\n",
        "local_model = tf.keras.models.load_model(MODEL_DIR)"
      ]
    },
@@ -936,23 +849,21 @@
        "id": "evaluate_custom_model:image"
      },
      "source": [
-        "## Evaluate the model\n",
+        "### Evaluate the model\n",
        "\n",
        "Now find out how good the model is.\n",
        "\n",
-        "### Load evaluation data\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 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",
        "\n",
-        "You don't need the training data, and hence why we loaded it as `(_, _)`.\n",
+        "You don't need the training data, and hence load it 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 can find that it was compiled for sparse labels. So we don't need to do anything more."
      ]
    },
    {
@@ -963,9 +874,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",
@@ -978,9 +886,9 @@
        "id": "perform_evaluation_custom"
      },
      "source": [
-        "### Perform the model evaluation\n",
+        "#### Perform the model evaluation\n",
        "\n",
-        "Now evaluate how well the model in the custom job did."
+        "Use the model's `evaluate()` method to perform the evaluation."
      ]
    },
    {
@@ -1000,7 +908,9 @@
        "id": "import_model:migration,new"
      },
      "source": [
-        "### [general.import-model](https://cloud.google.com/vertex-ai/docs/general/import-model)"
+        "### Import models to Vertex\n",
+        "\n",
+        "Learn more on how to [Import models to Vertex](https://cloud.google.com/vertex-ai/docs/general/import-model)"
      ]
    },
    {
@@ -1009,11 +919,11 @@
        "id": "serving_function_image"
      },
      "source": [
-        "### Serving function for image data\n",
+        "#### 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, 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",
+        "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 used by the underlying model (instead of upstream on a CPU).\n",
        "\n",
        "When you send a prediction or explanation request, the content of the request is base 64 decoded into a Tensorflow string (`tf.string`), which is passed to the serving function (`serving_fn`). The serving function preprocesses the `tf.string` into raw (uncompressed) numpy bytes (`preprocess_fn`) to match the input requirements of the model:\n",
        "- `io.decode_jpeg`- Decompresses the JPG image which is returned as a Tensorflow tensor with three channels (RGB).\n",
@@ -1075,13 +985,13 @@
        "id": "serving_function_signature:image"
      },
      "source": [
-        "## Get the serving function signature\n",
+        "#### Get the serving function signature\n",
        "\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 this purpose, you need the signature of the serving function. When the data is sent for prediction as an HTTP request packet, the image data is base64 encoded, and the 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 is used later when you make a prediction request."
      ]
    },
    {
@@ -1106,14 +1016,14 @@
        "id": "upload_model:mbsdk"
      },
      "source": [
-        "## Upload the model\n",
+        "### 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 a model resource 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",
+        "- `sync`: Set **True** to wait until the completion of the job.\n",
        "\n",
        "If the `upload()` method is run asynchronously, you can subsequently block until completion with the `wait()` method."
      ]
@@ -1126,8 +1036,8 @@
      },
      "outputs": [],
      "source": [
-        "model = aip.Model.upload(\n",
-        "    display_name=\"cifar10_\" + UUID,\n",
+        "model = aiplatform.Model.upload(\n",
+        "    display_name=\"cifar10-unique\",\n",
        "    artifact_uri=MODEL_DIR,\n",
        "    serving_container_image_uri=DEPLOY_IMAGE,\n",
        "    sync=False,\n",
@@ -1157,16 +1067,9 @@
        "id": "make_batch_predictions:migration"
      },
      "source": [
-        "## Make batch predictions"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "batchpredictionjobs_create:migration,new,mbsdk"
-      },
-      "source": [
-        "### [predictions.batch-prediction](https://cloud.google.com/vertex-ai/docs/predictions/batch-predictions)"
+        "### Generate batch predictions\n",
+        "\n",
+        "Here is the [Overview of getting predictions on Vertex AI](https://cloud.google.com/vertex-ai/docs/predictions/overview)"
      ]
    },
    {
@@ -1175,9 +1078,9 @@
        "id": "get_test_items:test"
      },
      "source": [
-        "### Get test items\n",
+        "#### Get test items\n",
        "\n",
-        "You will use examples out of the test (holdout) portion of the dataset as a test items."
+        "Use examples out of the test (holdout) portion of the dataset as a test items."
      ]
    },
    {
@@ -1201,8 +1104,8 @@
        "id": "prepare_test_items:test,image"
      },
      "source": [
-        "### Prepare the request content\n",
-        "You are going to send the CIFAR10 images as compressed JPG image, instead of the raw uncompressed bytes:\n",
+        "#### Prepare the request content\n",
+        "You're going to send the CIFAR10 images as compressed JPG image, instead of the raw uncompressed bytes:\n",
        "\n",
        "- `cv2.imwrite`: Use openCV to write the uncompressed image to disk as a compressed JPEG image.\n",
        " - Denormalize the image data from \\[0,1) range back to [0,255).\n",
@@ -1229,7 +1132,7 @@
        "id": "copy_test_items:test"
      },
      "source": [
-        "### Copy test item(s)\n",
+        "#### Copy test item(s)\n",
        "\n",
        "For the batch prediction, copy the test items over to your Cloud Storage bucket."
      ]
@@ -1255,9 +1158,9 @@
        "id": "make_batch_file:custom,image"
      },
      "source": [
-        "### Make the batch input file\n",
+        "#### Create batch input file\n",
        "\n",
-        "Now make a batch input file, which you will store in your local Cloud Storage bucket. The batch input file can only be in JSONL format. For JSONL file, you make one dictionary entry per line for each data item (instance). The dictionary contains the key/value pairs:\n",
+        "Now create batch input file, and store it in your local Cloud Storage bucket. The batch input file can only be in JSONL format. For JSONL file, make one dictionary entry per line for each data item (instance). The dictionary contains the key/value pairs:\n",
        "\n",
        "- `input_name`: the name of the input layer of the underlying model.\n",
        "- `'b64'`: A key that indicates the content is base64 encoded.\n",
@@ -1267,7 +1170,7 @@
        "\n",
        "                        {serving_input: {'b64': content}}\n",
        "\n",
-        "To pass the image data to the prediction service you encode the bytes into base64 -- which makes the content safe from modification when transmitting binary data over the network.\n",
+        "To pass the image data to the prediction service encode the bytes into base64 -- which makes the content safe from modification when transmitting binary data over the network.\n",
        "\n",
        "- `tf.io.read_file`: Read the compressed JPG images into memory as raw bytes.\n",
        "- `base64.b64encode`: Encode the raw bytes into a base64 encoded string."
@@ -1302,19 +1205,19 @@
        "id": "batch_request:mbsdk,jsonl,custom"
      },
      "source": [
-        "### Make the batch prediction request\n",
+        "#### 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, with the following parameters:\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",
        "\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",
        "- `gcs_destination_prefix`: The Cloud Storage location for storing the batch prediction resuls.\n",
-        "- `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",
+        "- `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",
        "- `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."
+        "- `sync`: Set **True** to wait until the completion of the job."
      ]
    },
    {
@@ -1329,7 +1232,7 @@
        "MAX_NODES = 1\n",
        "\n",
        "batch_predict_job = model.batch_predict(\n",
-        "    job_display_name=\"cifar10_\" + UUID,\n",
+        "    job_display_name=\"cifar10-unique\",\n",
        "    gcs_source=gcs_input_uri,\n",
        "    gcs_destination_prefix=BUCKET_URI,\n",
        "    instances_format=\"jsonl\",\n",
@@ -1355,14 +1258,8 @@
        "*Example output:*\n",
        "\n",
        "    INFO:google.cloud.aiplatform.jobs:Creating BatchPredictionJob\n",
-        "    <google.cloud.aiplatform.jobs.BatchPredictionJob object at 0x7f806a6112d0> is waiting for upstream dependencies to complete.\n",
-        "    INFO:google.cloud.aiplatform.jobs:BatchPredictionJob created. Resource name: projects/759209241365/locations/us-central1/batchPredictionJobs/5110965452507447296\n",
-        "    INFO:google.cloud.aiplatform.jobs:To use this BatchPredictionJob in another session:\n",
-        "    INFO:google.cloud.aiplatform.jobs:bpj = aiplatform.BatchPredictionJob('projects/759209241365/locations/us-central1/batchPredictionJobs/5110965452507447296')\n",
-        "    INFO:google.cloud.aiplatform.jobs:View Batch Prediction Job:\n",
-        "    https://console.cloud.google.com/ai/platform/locations/us-central1/batch-predictions/5110965452507447296?project=759209241365\n",
-        "    INFO:google.cloud.aiplatform.jobs:BatchPredictionJob projects/759209241365/locations/us-central1/batchPredictionJobs/5110965452507447296 current state:\n",
-        "    JobState.JOB_STATE_RUNNING"
+        "    <google.cloud.aiplatform.jobs.BatchPredictionJob object at 0x7874ad735a80> is waiting for upstream dependencies to complete.\n",
+        "    "
      ]
    },
    {
@@ -1371,9 +1268,9 @@
        "id": "batch_request_wait:mbsdk"
      },
      "source": [
-        "### Wait for completion of batch prediction job\n",
+        "#### 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. Alternatively, you can set the parameter `sync` to `True` in the `batch_predict()` method to block until the batch prediction job is completed."
      ]
    },
    {
@@ -1395,11 +1292,6 @@
      "source": [
        "*Example Output:*\n",
        "\n",
-        "    INFO:google.cloud.aiplatform.jobs:BatchPredictionJob created. Resource name: projects/759209241365/locations/us-central1/batchPredictionJobs/181835033978339328\n",
-        "    INFO:google.cloud.aiplatform.jobs:To use this BatchPredictionJob in another session:\n",
-        "    INFO:google.cloud.aiplatform.jobs:bpj = aiplatform.BatchPredictionJob('projects/759209241365/locations/us-central1/batchPredictionJobs/181835033978339328')\n",
-        "    INFO:google.cloud.aiplatform.jobs:View Batch Prediction Job:\n",
-        "    https://console.cloud.google.com/ai/platform/locations/us-central1/batch-predictions/181835033978339328?project=759209241365\n",
        "    INFO:google.cloud.aiplatform.jobs:BatchPredictionJob projects/759209241365/locations/us-central1/batchPredictionJobs/181835033978339328 current state:\n",
        "    JobState.JOB_STATE_RUNNING\n",
        "    INFO:google.cloud.aiplatform.jobs:BatchPredictionJob projects/759209241365/locations/us-central1/batchPredictionJobs/181835033978339328 current state:\n",
@@ -1431,7 +1323,7 @@
        "\n",
        "Next, get the 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 prediction requests in a JSON format:\n",
+        "The results are written to the Cloud Storage output bucket specified in the batch prediction request. Call the method `iter_outputs()` to get a list of each Cloud Storage file generated with the results. Each file contains one or more prediction requests in a JSON format:\n",
        "\n",
        "- `instance`: The prediction request.\n",
        "- `prediction`: The prediction response."
@@ -1481,16 +1373,7 @@
        "id": "make_online_predictions:migration"
      },
      "source": [
-        "## Make online predictions"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "deploy_model:migration,new,mbsdk"
-      },
-      "source": [
-        "### [predictions.deploy-model-api](https://cloud.google.com/vertex-ai/docs/predictions/deploy-model-api)"
+        "### Generate online predictions \n"
      ]
    },
    {
@@ -1499,13 +1382,13 @@
        "id": "deploy_model:mbsdk,all"
      },
      "source": [
-        "## Deploy the model\n",
+        "#### 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 is only one model, then specify `traffic_split` 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",
        "- `machine_type`: The type of machine to use for training.\n",
        "- `accelerator_type`: The hardware accelerator type.\n",
@@ -1522,7 +1405,7 @@
      },
      "outputs": [],
      "source": [
-        "DEPLOYED_NAME = \"cifar10-\" + UUID\n",
+        "DEPLOYED_NAME = \"cifar10-unique\"\n",
        "\n",
        "TRAFFIC_SPLIT = {\"0\": 100}\n",
        "\n",
@@ -1575,7 +1458,7 @@
        "id": "endpoints_predict:migration,new,mbsdk"
      },
      "source": [
-        "### [predictions.online-prediction-automl](https://cloud.google.com/vertex-ai/docs/predictions/online-predictions-automl)"
+        "Learn more about how to [Train and use your own models](https://cloud.google.com/vertex-ai/docs/training-overview)"
      ]
    },
    {
@@ -1586,7 +1469,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."
      ]
    },
    {
@@ -1608,8 +1491,8 @@
        "id": "prepare_test_item:test,image"
      },
      "source": [
-        "### Prepare the request content\n",
-        "You are going to send the CIFAR10 image as compressed JPG image, instead of the raw uncompressed bytes:\n",
+        "#### Prepare the request content\n",
+        "You're going to send the CIFAR10 image as compressed JPG image, instead of the raw uncompressed bytes:\n",
        "\n",
        "- `cv2.imwrite`: Use openCV to write the uncompressed image to disk as a compressed JPEG image.\n",
        " - Denormalize the image data from \\[0,1) range back to [0,255).\n",
@@ -1642,13 +1525,13 @@
        "id": "predict_request:mbsdk,custom,icn"
      },
      "source": [
-        "### Make the prediction\n",
+        "### Get the prediction\n",
        "\n",
-        "Now that your `Model` resource is deployed to an `Endpoint` resource, you can do online predictions by sending prediction requests to the Endpoint resource.\n",
+        "Now that your model resource is deployed to an endpoint resource, you can get online predictions by sending prediction requests to the endpoint resource.\n",
        "\n",
        "#### Request\n",
        "\n",
-        "Since in this example your test item is in a Cloud Storage bucket, you open and read the contents of the image using `tf.io.gfile.Gfile()`. To pass the test data to the prediction service, you encode the bytes into base64 -- which makes the content safe from modification while transmitting binary data over the network.\n",
+        "Since in this example your test item is in a Cloud Storage bucket, open and read the contents of the image using `tf.io.gfile.Gfile()`. To pass the test data to the prediction service, you need to encode the bytes into base64 -- which makes the content safe from modification while transmitting binary data over the network.\n",
        "\n",
        "The format of each instance is:\n",
        "\n",
@@ -1662,7 +1545,7 @@
        "\n",
        "- `ids`: The internal assigned unique identifiers for each prediction request.\n",
        "- `predictions`: The predicted confidence, between 0 and 1, per class label.\n",
-        "- `deployed_model_id`: The Vertex AI identifier for the deployed `Model` resource which did the predictions."
+        "- `deployed_model_id`: The Vertex AI identifier for the deployed model resource which did the predictions."
      ]
    },
    {
@@ -1700,7 +1583,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, undeploy the model from the endpoint resouce. This deprovisions all compute resources and ends billing for the deployed model."
      ]
    },
    {
@@ -1720,7 +1603,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",
@@ -1736,10 +1619,6 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
-        "delete_bucket = False\n",
-        "\n",
        "# Delete the model using the Vertex model object\n",
        "model.delete()\n",
        "\n",
@@ -1752,7 +1631,12 @@
        "# Delete the batch prediction job using the Vertex batch prediction object\n",
        "batch_predict_job.delete()\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "# delete locally generated files\n",
+        "! rm -rf custom custom.tar.gz tmp1.jpg tmp2.jpg tmp.jpg\n",
+        "\n",
+        "# delete 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 Migration: Custom Scikit-Learn model with pre-built training container\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
-        "<a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/migration/sdk-custom-scikit-learn-prebuilt-container.ipynb\" target='_blank'>\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/migration/sdk-custom-scikit-learn-prebuilt-container.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://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/migration/sdk-custom-scikit-learn-prebuilt-container.ipynb\" target='_blank'>\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%2Fmigration%2Fsdk-custom-scikit-learn-prebuilt-container.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/migration/sdk-custom-scikit-learn-prebuilt-container.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>\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/migration/sdk-custom-scikit-learn-prebuilt-container.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 style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/migration/sdk-custom-scikit-learn-prebuilt-container.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",
-        "<br/><br/><br/>"
+        "  </td>\n",
+        "</table>"
      ]
    },
    {
@@ -64,7 +66,7 @@
        "\n",
        "This tutorial demonstrates how to use the Vertex AI SDK for Python to train and deploy a custom tabular classification scikit-learn model for batch prediction.\n",
        "\n",
-        "Learn more about [Migrate to Vertex AI](https://cloud.google.com/vertex-ai/docs/start/migrating-to-vertex-ai) and [Custom training](https://cloud.google.com/vertex-ai/docs/training/custom-training)."
+        "Learn more about [Migrate to Vertex AI](https://cloud.google.com/vertex-ai/docs/start/migrating-to-vertex-ai) and [Custom training overview](https://cloud.google.com/vertex-ai/docs/training/custom-training)."
      ]
    },
    {
@@ -75,25 +77,26 @@
      "source": [
        "### Objective\n",
        "\n",
-        "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",
+        "In this tutorial, you learn how to use Vertex AI Training to create a custom trained model. Then, you learn to use Vertex AI batch prediction to generate batch prediction on the trained model.\n",
        "\n",
        "\n",
        "You learn how to 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.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
-        "- `Vertex AI Training`\n",
-        "- `Vertex AI Batch Prediction`\n",
-        "- `Vertex AI Model` resource\n",
-        "- `Vertex AI Endpoint` resource\n",
+        "- Vertex AI Training\n",
+        "- Vertex AI batch prediction\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 scikit-learn model.\n",
-        "- Upload the trained model artifacts as a `Model` resource.\n",
-        "- Make a batch prediction.\n",
-        "- Deploy model to a endpoint\n",
-        "- Make a online prediction"
+        "- Create a Vertex AI custom job for training a scikit-learn model.\n",
+        "- Upload the trained model artifacts as a model resource.\n",
+        "- Generate batch predictions.\n",
+        "- Deploy the model resource to a serving endpoint resource.\n",
+        "- Generate online predictions.\n",
+        "- Undeploy the model resource. "
      ]
    },
    {
@@ -150,54 +153,86 @@
      "source": [
        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
        "                                 google-cloud-storage \\\n",
-        "                                 tensorflow"
+        "                                 tensorflow==2.15.1"
      ]
    },
    {
      "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": "before_you_begin:nogpu"
+        "id": "SbmM4z7FOBpM"
      },
      "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": "before_you_begin:nogpu"
+        "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": "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)."
      ]
    },
    {
@@ -209,33 +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 = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -244,7 +258,7 @@
      "source": [
        "#### UUID\n",
        "\n",
-        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a uuid for each instance session, and append it onto the name of resources you create in this tutorial."
+        "If you're in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a uuid for each instance session, and append it onto the name of resources you create in this tutorial."
      ]
    },
    {
@@ -267,64 +281,6 @@
        "UUID = generate_uuid()"
      ]
    },
-    {
-      "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": {
@@ -364,7 +320,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
@@ -387,8 +343,8 @@
        "import json\n",
        "import os\n",
        "\n",
-        "import google.cloud.aiplatform as aip\n",
-        "import tensorflow as tf"
+        "import tensorflow as tf\n",
+        "from google.cloud import aiplatform"
      ]
    },
    {
@@ -397,7 +353,7 @@
        "id": "init_aip:mbsdk"
      },
      "source": [
-        "## Initialize Vertex SDK for Python\n",
+        "### Initialize Vertex AI SDK for Python\n",
        "\n",
        "Initialize the Vertex SDK for Python for your project and corresponding bucket."
      ]
@@ -410,7 +366,7 @@
      },
      "outputs": [],
      "source": [
-        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
+        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -424,10 +380,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 custom training](https://cloud.google.com/ai-platform-unified/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 and explanation](https://cloud.google.com/ai-platform-unified/docs/predictions/pre-built-containers)."
      ]
    },
    {
@@ -458,19 +414,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 use for for training and prediction.\n",
+        "- Set the variables `TRAIN_COMPUTE` and `DEPLOY_COMPUTE` to configure the compute resources for the VMs used 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 can also use n2 and e2 machine types for training and deployment, but they don't support GPUs."
      ]
    },
    {
@@ -506,11 +462,13 @@
        "id": "examine_training_package"
      },
      "source": [
-        "### Examine the training package\n",
+        "### Tutorial\n",
+        "\n",
+        "Now you're ready to create your own custom model and training US census data.\n",
        "\n",
        "#### Package layout\n",
        "\n",
-        "Before you start the training, you 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",
@@ -522,11 +480,11 @@
        "\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",
-        "In the following cells, you assemble the training package."
+        "In the following cells, assemble the training package."
      ]
    },
    {
@@ -636,13 +594,13 @@
        "with open('./adult.data', 'r') as train_data:\n",
        "    raw_training_data = pd.read_csv(train_data, header=None, names=COLUMNS)\n",
        "\n",
-        "# Remove the column we are trying to predict ('income-level') from our features list\n",
+        "# Remove the column you're trying to predict ('income-level') from our features list\n",
        "# Convert the Dataframe to a lists of lists\n",
        "train_features = raw_training_data.drop('income-level', axis=1).values.tolist()\n",
        "# Create our training labels list, convert the Dataframe to a lists of lists\n",
        "train_labels = (raw_training_data['income-level'] == ' >50K').values.tolist()\n",
        "\n",
-        "# Since the census data set has categorical features, we need to convert\n",
+        "# Since the census data set has categorical features, you need to convert\n",
        "# them to numerical values. We'll use a list of pipelines to convert each\n",
        "# categorical column and then use FeatureUnion to combine them before calling\n",
        "# the RandomForestClassifier.\n",
@@ -652,7 +610,7 @@
        "# To do this, each categorical column will use a pipeline that extracts one feature column via\n",
        "# SelectKBest(k=1) and a LabelBinarizer() to convert the categorical value to a numerical one.\n",
        "# A scores array (created below) will select and extract the feature column. The scores array is\n",
-        "# created by iterating over the COLUMNS and checking if it is a CATEGORICAL_COLUMN.\n",
+        "# created by iterating over the COLUMNS and checking if it's a CATEGORICAL_COLUMN.\n",
        "for i, col in enumerate(COLUMNS[:-1]):\n",
        "    if col in CATEGORICAL_COLUMNS:\n",
        "        # Create a scores array to get the individual categorical column.\n",
@@ -664,7 +622,7 @@
        "        # Returns: [['State-gov']]\n",
        "        # Build the scores array.\n",
        "        scores = [0] * len(COLUMNS[:-1])\n",
-        "        # This column is the categorical column we want to extract.\n",
+        "        # This column is the categorical column you want to extract.\n",
        "        scores[i] = 1\n",
        "        skb = SelectKBest(k=1)\n",
        "        skb.scores_ = scores\n",
@@ -725,9 +683,9 @@
        "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."
+        "Next, package the training folder into a compressed tar ball, and then store it in your Cloud Storage bucket."
      ]
    },
    {
@@ -750,7 +708,7 @@
        "id": "train_a_model:migration"
      },
      "source": [
-        "## Train a model"
+        "### Create and run custom training job"
      ]
    },
    {
@@ -759,7 +717,7 @@
        "id": "custom_create:migration,new,mbsdk,prebuilt"
      },
      "source": [
-        "### [training.create-python-pre-built-container](https://cloud.google.com/vertex-ai/docs/training/create-python-pre-built-container)"
+        "Learn how to [Create a Python training application for a prebuilt container](https://cloud.google.com/vertex-ai/docs/training/create-python-pre-built-container)"
      ]
    },
    {
@@ -768,14 +726,13 @@
        "id": "create_custom_training_job:mbsdk,no_model"
      },
      "source": [
-        "### 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",
+        "1) Create a custom training job.\n",
+        "2) Specify your training parameters and run the job.\n",
        "\n",
        "#### Create custom training job\n",
        "\n",
-        "A custom training job is created with the `CustomTrainingJob` class, with the following parameters:\n",
+        "A custom training job is created using the `CustomTrainingJob` class, with the following parameters:\n",
        "\n",
        "- `display_name`: The human readable name for the custom training job.\n",
        "- `container_uri`: The training container image.\n",
@@ -791,7 +748,7 @@
      },
      "outputs": [],
      "source": [
-        "job = aip.CustomTrainingJob(\n",
+        "job = aiplatform.CustomTrainingJob(\n",
        "    display_name=\"census_\" + UUID,\n",
        "    script_path=\"custom/trainer/task.py\",\n",
        "    container_uri=TRAIN_IMAGE,\n",
@@ -820,12 +777,12 @@
      "source": [
        "#### 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, run the custom job to start the training job by invoking the `run()` method, with the following parameters:\n",
        "\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",
-        "- `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 write the model artifacts.\n",
+        "- `sync`: Set **True** to wait until the completion of the job."
      ]
    },
    {
@@ -850,28 +807,28 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "import_model:migration,new"
+        "id": "upload_model:mbsdk"
      },
      "source": [
-        "### [general.import-model](https://cloud.google.com/vertex-ai/docs/general/import-model)"
+        "### Upload the model\n",
+        "\n",
+        "Next, upload your model to a model resource using the `Model.upload()` method, with the following parameters:\n",
+        "\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`: Set **True** to wait until the completion of the job.\n",
+        "\n",
+        "If the `upload()` method is run asynchronously, you can subsequently block until completion with the `wait()` method."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "upload_model:mbsdk"
+        "id": "import_model:migration,new"
      },
      "source": [
-        "## Upload the model\n",
-        "\n",
-        "Next, upload your model to a `Model` resource using `Model.upload()` method, with the following parameters:\n",
-        "\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",
-        "\n",
-        "If the `upload()` method is run asynchronously, you can subsequently block until completion with the `wait()` method."
+        " Learn more about how to [Import models to Vertex AI](https://cloud.google.com/vertex-ai/docs/general/import-model)"
      ]
    },
    {
@@ -882,7 +839,7 @@
      },
      "outputs": [],
      "source": [
-        "model = aip.Model.upload(\n",
+        "model = aiplatform.Model.upload(\n",
        "    display_name=\"census_\" + UUID,\n",
        "    artifact_uri=MODEL_DIR,\n",
        "    serving_container_image_uri=DEPLOY_IMAGE,\n",
@@ -913,7 +870,7 @@
        "id": "make_batch_predictions:migration"
      },
      "source": [
-        "## Make batch predictions"
+        "### Generate batch predictions"
      ]
    },
    {
@@ -922,7 +879,7 @@
        "id": "batchpredictionjobs_create:migration,new,mbsdk"
      },
      "source": [
-        "### [predictions.batch-prediction](https://cloud.google.com/vertex-ai/docs/predictions/batch-predictions)"
+        "To learn more about batch predictions refer [Overview of getting predictions on Vertex AI](https://cloud.google.com/vertex-ai/docs/predictions/batch-predictions)"
      ]
    },
    {
@@ -931,9 +888,9 @@
        "id": "make_test_items:scilearn,tabular,census"
      },
      "source": [
-        "### Make test items\n",
+        "#### Create test items\n",
        "\n",
-        "You use synthetic data as test data items. Don't be concerned that we are using synthetic data -- we just want to demonstrate how to make a prediction."
+        "Use synthetic data as test data items. Don’t be concerned about using synthetic data – as it's just for demonstration of generating predictions. "
      ]
    },
    {
@@ -986,9 +943,9 @@
        "id": "make_batch_file:custom,tabular,list,jsonl"
      },
      "source": [
-        "### Make the batch input file\n",
+        "#### Create a batch input file\n",
        "\n",
-        "Now make a batch input file, which you store in your local Cloud Storage bucket.  Each instance in the prediction request is a list of the form:\n",
+        "Now create a batch input file, which is stored in your local Cloud Storage bucket. Each instance in the prediction request is a list of the form:\n",
        "\n",
        "                        [ [ content_1], [content_2] ]\n",
        "\n",
@@ -1017,9 +974,9 @@
        "id": "batch_request:mbsdk,jsonl,custom,cpu"
      },
      "source": [
-        "### Make the batch prediction request\n",
+        "#### 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, with the following parameters:\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",
        "\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",
@@ -1027,7 +984,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",
        "- `machine_type`: The type of machine to use for training.\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."
      ]
    },
    {
@@ -1082,9 +1039,9 @@
        "id": "batch_request_wait:mbsdk"
      },
      "source": [
-        "### Wait for completion of batch prediction job\n",
+        "#### 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. Alternatively, you can set the parameter `sync` to `True` in the `batch_predict()` method to block until the batch prediction job is completed."
      ]
    },
    {
@@ -1138,11 +1095,11 @@
        "id": "get_batch_prediction:mbsdk,custom,lbn"
      },
      "source": [
-        "### Get the predictions\n",
+        "#### Get the predictions\n",
        "\n",
        "Next, get the 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 prediction requests in a JSON format:\n",
+        "The results are written to the Cloud Storage output bucket you specified in the batch prediction request. Call the `iter_outputs()` to get a list of each Cloud Storage file generated with the results. Each file contains one or more prediction requests in a JSON format:\n",
        "\n",
        "- `instance`: The prediction request.\n",
        "- `prediction`: The prediction response."
@@ -1189,7 +1146,7 @@
        "id": "make_online_predictions:migration"
      },
      "source": [
-        "## Make online predictions"
+        "### Generate online predictions"
      ]
    },
    {
@@ -1198,7 +1155,7 @@
        "id": "deploy_model:migration,new,mbsdk"
      },
      "source": [
-        "### [predictions.deploy-model-api](https://cloud.google.com/vertex-ai/docs/predictions/deploy-model-api)"
+        "To learn more of online predictions refer, [Overview of getting predictions on Vertex AI](https://cloud.google.com/vertex-ai/docs/predictions/deploy-model-api)"
      ]
    },
    {
@@ -1207,9 +1164,9 @@
        "id": "deploy_model:mbsdk,cpu"
      },
      "source": [
-        "## Deploy the model\n",
+        "#### 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, 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",
@@ -1241,7 +1198,8 @@
        "    machine_type=DEPLOY_COMPUTE,\n",
        "    min_replica_count=MIN_NODES,\n",
        "    max_replica_count=MAX_NODES,\n",
-        ")"
+        ")\n",
+        "endpoint.wait()"
      ]
    },
    {
@@ -1268,9 +1226,9 @@
        "id": "make_test_item:scilearn,tabular,census"
      },
      "source": [
-        "### Make test item\n",
+        "#### Create a test item\n",
        "\n",
-        "You 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."
+        "Use synthetic data as a test data item. Don’t be concerned about using synthetic data – since it's just for demonstration purposes"
      ]
    },
    {
@@ -1305,9 +1263,9 @@
        "id": "predict_request:mbsdk,custom,lbn"
      },
      "source": [
-        "### Make the prediction\n",
+        "#### Make the prediction\n",
        "\n",
-        "Now that your `Model` resource is deployed to an `Endpoint` resource, you can do online predictions by sending prediction requests to the `Endpoint` resource.\n",
+        "Now that your model resource is deployed to an endpoint resource, you can make online predictions by sending prediction requests to the endpoint resource.\n",
        "\n",
        "#### Request\n",
        "\n",
@@ -1315,15 +1273,15 @@
        "\n",
        "    [feature_list]\n",
        "\n",
-        "Since the predict() method can take multiple items (instances), send your single test item as a list of one test item.\n",
+        "Since the `predict()` 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 predict() call is a Python dictionary with the following entries:\n",
+        "The response from the `predict()` call is a Python dictionary with the following entries:\n",
        "\n",
        "- `ids`: The internal assigned unique identifiers for each prediction request.\n",
        "- `predictions`: The predicted confidence, between 0 and 1, per class label.\n",
-        "- `deployed_model_id`: The Vertex AI identifier for the deployed `Model` resource which did the predictions."
+        "- `deployed_model_id`: The Vertex AI identifier for the deployed model resource which did the predictions."
      ]
    },
    {
@@ -1357,9 +1315,9 @@
        "id": "undeploy_model:mbsdk"
      },
      "source": [
-        "## Undeploy the model\n",
+        "### 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 generating predictions, simply undeploy the model from the endpoint resouce. This deprovisions all compute resources and ends billing for the deployed model."
      ]
    },
    {
@@ -1379,7 +1337,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",
@@ -1402,7 +1360,7 @@
      },
      "outputs": [],
      "source": [
-        "# delete endpoint\n",
+        "# Delete endpoint\n",
        "endpoint.delete()\n",
        "\n",
        "# Delete the model using the Vertex model object\n",
@@ -1414,8 +1372,12 @@
        "# Delete the batch prediction job using the Vertex batch prediction object\n",
        "batch_predict_job.delete()\n",
        "\n",
-        "delete_bucket = False\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "# Delete the locally generated files\n",
+        "! rm -rf custom custom.tar.gz\n",
+        "\n",
+        "# Delete the Cloud Storage bucket\n",
+        "delete_bucket = False  # set True for deletion\n",
+        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
@@ -32,25 +32,27 @@
        "# Vertex AI: Vertex AI Migration: Custom XGBoost model with pre-built training container\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
-        "<a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/migration/sdk-custom-xgboost-prebuilt-container.ipynb\" target='_blank'>\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/migration/sdk-custom-xgboost-prebuilt-container.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://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/migration/sdk-custom-xgboost-prebuilt-container.ipynb\" target='_blank'>\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%2Fmigration%2Fsdk-custom-xgboost-prebuilt-container.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/migration/sdk-custom-xgboost-prebuilt-container.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>\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/migration/sdk-custom-xgboost-prebuilt-container.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 style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/migration/sdk-custom-xgboost-prebuilt-container.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",
-        "<br/><br/><br/>"
+        "  </td>\n",
+        "</table>\n"
      ]
    },
    {
@@ -64,7 +66,7 @@
        "\n",
        "This tutorial demonstrates how to use the Vertex AI SDK for Python to train and deploy a custom tabular classification XGBoost model for batch prediction.\n",
        "\n",
-        "Learn more about [Migrate to Vertex AI](https://cloud.google.com/vertex-ai/docs/start/migrating-to-vertex-ai) and [Custom training](https://cloud.google.com/vertex-ai/docs/training/custom-training)."
+        "Learn more about [Migrate to Vertex AI](https://cloud.google.com/vertex-ai/docs/start/migrating-to-vertex-ai) and [Custom training overview](https://cloud.google.com/vertex-ai/docs/training/custom-training)."
      ]
    },
    {
@@ -75,25 +77,26 @@
      "source": [
        "### Objective\n",
        "\n",
-        "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",
+        "In this tutorial, you learn to use Vertex AI Training to create a custom trained model. Then, you learn to use Vertex AI batch prediction to generate batch prediction on the trained model.\n",
        "\n",
        "\n",
        "You learn how to 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.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
-        "- `Vertex AI Training`\n",
-        "- `Vertex AI Batch Prediction`\n",
-        "- `Vertex AI Model` resource\n",
-        "- `Vertex AI Endpoint` resource\n",
+        "- Vertex AI Training\n",
+        "- Vertex AI batch prediction\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 scikit-learn model.\n",
-        "- Upload the trained model artifacts as a `Model` resource.\n",
-        "- Make a batch prediction.\n",
-        "- Deploy model to a endpoint\n",
-        "- Make a online prediction"
+        "- Create a Vertex AI custom job for training a xgboost model.\n",
+        "- Upload the trained model artifacts as a model resource.\n",
+        "- Generate batch predictions.\n",
+        "- Deploy the model resource to a serving endpoint resource.\n",
+        "- Generate online predictions.\n",
+        "- Undeploy the model resource."
      ]
    },
    {
@@ -104,7 +107,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). 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. The trained model predicts the type of Iris flower species from a class of three species: *setosa*, *virginica*, or *versicolor*."
      ]
    },
    {
@@ -130,12 +133,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"
      ]
    },
    {
@@ -146,60 +156,88 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
-        "                                 google-cloud-storage\n",
-        "\n",
-        "if os.getenv(\"IS_TESTING\"):\n",
-        "    ! pip3 install --upgrade --quiet tensorflow "
+        "                                 google-cloud-storage \\\n",
+        "                                 tensorflow==2.15.1"
      ]
    },
    {
      "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": "before_you_begin:nogpu"
+        "id": "SbmM4z7FOBpM"
      },
      "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": "before_you_begin:nogpu"
+        "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": "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)."
      ]
    },
    {
@@ -211,122 +249,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": "timestamp"
-      },
-      "source": [
-        "#### UUID\n",
-        "\n",
-        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a uuid for each instance session, and append it onto the name of resources you create in this tutorial."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "e87d5856317d"
-      },
-      "outputs": [],
-      "source": [
-        "import random\n",
-        "import string\n",
-        "\n",
-        "\n",
-        "# Generate a uuid of a specifed length(default=8)\n",
-        "def generate_uuid(length: int = 8) -> str:\n",
-        "    return \"\".join(random.choices(string.ascii_lowercase + string.digits, k=length))\n",
-        "\n",
-        "\n",
-        "UUID = generate_uuid()"
-      ]
-    },
-    {
-      "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": {
@@ -366,7 +294,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
@@ -375,9 +303,6 @@
        "id": "setup_vars"
      },
      "source": [
-        "### Set up variables\n",
-        "\n",
-        "Next, set up some variables used throughout the tutorial.\n",
        "### Import libraries and define constants"
      ]
    },
@@ -389,7 +314,10 @@
      },
      "outputs": [],
      "source": [
-        "import google.cloud.aiplatform as aip"
+        "import os\n",
+        "\n",
+        "import tensorflow as tf\n",
+        "from google.cloud import aiplatform"
      ]
    },
    {
@@ -398,7 +326,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."
      ]
@@ -411,7 +339,7 @@
      },
      "outputs": [],
      "source": [
-        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
+        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -425,10 +353,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 custom training](https://cloud.google.com/ai-platform-unified/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 and explaination](https://cloud.google.com/ai-platform-unified/docs/predictions/pre-built-containers)."
      ]
    },
    {
@@ -456,19 +384,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 used 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 can also use n2 and e2 machine types for training and deployment, but they don't support GPUs."
      ]
    },
    {
@@ -504,11 +432,11 @@
        "id": "examine_training_package"
      },
      "source": [
-        "### Examine the training package\n",
+        "### Tutorial\n",
        "\n",
-        "#### Package layout\n",
+        "#### Examine training 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, take a 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",
@@ -518,13 +446,13 @@
        "  - \\_\\_init\\_\\_.py\n",
        "  - task.py\n",
        "\n",
-        "The files `setup.cfg` and `setup.py` are the instructions for installing the package into the operating environment of the Docker image.\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",
-        "In the following cells, you will assemble the training package."
+        "In the following cells, assemble the training package."
      ]
    },
    {
@@ -586,7 +514,7 @@
        "iris_target_filename = 'iris_target.csv'\n",
        "data_dir = 'gs://cloud-samples-data/ai-platform/iris'\n",
        "\n",
-        "# gsutil outputs everything to stderr so we need to divert it to stdout.\n",
+        "# gsutil outputs everything to stderr so you need to divert it to stdout.\n",
        "subprocess.check_call(['gsutil', 'cp', os.path.join(data_dir,\n",
        "                                                    iris_data_filename),\n",
        "                       iris_data_filename], stderr=sys.stdout)\n",
@@ -628,7 +556,7 @@
      "source": [
        "#### 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."
+        "Next, package the training folder into a compressed tar ball, and then store it in your Cloud Storage bucket."
      ]
    },
    {
@@ -645,24 +573,6 @@
        "! gsutil cp custom.tar.gz $BUCKET_URI/trainer_iris.tar.gz"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "train_a_model:migration"
-      },
-      "source": [
-        "## Train a model"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "custom_create:migration,new,mbsdk,prebuilt"
-      },
-      "source": [
-        "### [training.create-python-pre-built-container](https://cloud.google.com/vertex-ai/docs/training/create-python-pre-built-container)"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -671,12 +581,15 @@
      "source": [
        "### Create and run custom training job\n",
        "\n",
+        "Learn how to [Create a Python training application for a prebuilt container](https://cloud.google.com/vertex-ai/docs/training/create-python-pre-built-container).\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",
+        "1) Create a custom training job.\n",
+        "2) Specify your training parameters and run the job.\n",
        "\n",
        "#### Create custom training job\n",
        "\n",
-        "A custom training job is created with the `CustomTrainingJob` class, with the following parameters:\n",
+        "A custom training job is created using the `CustomTrainingJob` class, with the following parameters:\n",
        "\n",
        "- `display_name`: The human readable name for the custom training job.\n",
        "- `container_uri`: The training container image.\n",
@@ -692,8 +605,8 @@
      },
      "outputs": [],
      "source": [
-        "job = aip.CustomTrainingJob(\n",
-        "    display_name=\"iris_\" + UUID,\n",
+        "job = aiplatform.CustomTrainingJob(\n",
+        "    display_name=\"iris-unique\",\n",
        "    script_path=\"custom/trainer/task.py\",\n",
        "    container_uri=TRAIN_IMAGE,\n",
        "    requirements=[\"gcsfs==0.7.1\", \"tensorflow-datasets==4.4\"],\n",
@@ -721,12 +634,12 @@
      "source": [
        "#### 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, run the custom job to start the training job by invoking the `run()` method, with the following parameters:\n",
        "\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",
        "- `base_output_dir`: The Cloud Storage location to write the model artifacts to.\n",
-        "- `sync`: Whether to block until completion of the job."
+        "- `sync`: Set **True** to wait until the completion of the job."
      ]
    },
    {
@@ -737,7 +650,7 @@
      },
      "outputs": [],
      "source": [
-        "MODEL_DIR = \"{}/{}\".format(BUCKET_URI, UUID)\n",
+        "MODEL_DIR = \"{}/{}\".format(BUCKET_URI, \"unique\")\n",
        "\n",
        "\n",
        "job.run(\n",
@@ -748,31 +661,24 @@
        "model_path_to_deploy = MODEL_DIR"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "import_model:migration,new"
-      },
-      "source": [
-        "### [general.import-model](https://cloud.google.com/vertex-ai/docs/general/import-model)"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "upload_model:mbsdk"
      },
      "source": [
-        "## Upload the model\n",
+        "### 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 a model resource using the `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",
+        "- `sync`: Set **True** to wait until the completion of job\n",
        "\n",
-        "If the `upload()` method is run asynchronously, you can subsequently block until completion with the `wait()` method."
+        "If the `upload()` method is run asynchronously, you can subsequently block until completion with the `wait()` method.\n",
+        "\n",
+        "Learn more about how to [Import models to Vertex AI](https://cloud.google.com/vertex-ai/docs/general/import-model)"
      ]
    },
    {
@@ -783,8 +689,8 @@
      },
      "outputs": [],
      "source": [
-        "model = aip.Model.upload(\n",
-        "    display_name=\"iris_\" + UUID,\n",
+        "model = aiplatform.Model.upload(\n",
+        "    display_name=\"iris-unique\",\n",
        "    artifact_uri=MODEL_DIR,\n",
        "    serving_container_image_uri=DEPLOY_IMAGE,\n",
        "    sync=False,\n",
@@ -814,16 +720,9 @@
        "id": "make_batch_predictions:migration"
      },
      "source": [
-        "## Make batch predictions"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "batchpredictionjobs_create:migration,new,mbsdk"
-      },
-      "source": [
-        "### [predictions.batch-prediction](https://cloud.google.com/vertex-ai/docs/predictions/batch-predictions)"
+        "### Generate batch predictions\n",
+        "\n",
+        "To learn more about batch predictions refer [Overview of getting predictions on Vertex AI](https://cloud.google.com/vertex-ai/docs/predictions/overview)."
      ]
    },
    {
@@ -832,9 +731,9 @@
        "id": "make_test_items:xgboost,tabular,iris"
      },
      "source": [
-        "### Make test items\n",
+        "#### Create test items\n",
        "\n",
-        "You will use synthetic data as a test data items. Don't be concerned that we are using synthetic data -- we just want to demonstrate how to make a prediction."
+        "Use synthetic data as test data items. Don’t be concerned about using synthetic data – as it's just for demonstration of generating predictions."
      ]
    },
    {
@@ -856,7 +755,7 @@
      "source": [
        "### Make the batch input file\n",
        "\n",
-        "Now make a batch input file, which you will store in your local Cloud Storage bucket.  Each instance in the prediction request is a list of the form:\n",
+        "Now create a batch input file, which is stored in your Cloud Storage bucket. Each instance in the prediction request is a list of the form:\n",
        "\n",
        "                        [ [ content_1], [content_2] ]\n",
        "\n",
@@ -871,8 +770,6 @@
      },
      "outputs": [],
      "source": [
-        "import tensorflow as tf\n",
-        "\n",
        "gcs_input_uri = BUCKET_URI + \"/\" + \"test.jsonl\"\n",
        "with tf.io.gfile.GFile(gcs_input_uri, \"w\") as f:\n",
        "    for i in INSTANCES:\n",
@@ -887,17 +784,17 @@
        "id": "batch_request:mbsdk,jsonl,custom,cpu"
      },
      "source": [
-        "### Make the batch prediction request\n",
+        "#### 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, with the following parameters:\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",
        "\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",
        "- `gcs_destination_prefix`: The Cloud Storage location for storing the batch prediction resuls.\n",
-        "- `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",
+        "- `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",
        "- `machine_type`: The type of machine to use for training.\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."
      ]
    },
    {
@@ -912,7 +809,7 @@
        "MAX_NODES = 1\n",
        "\n",
        "batch_predict_job = model.batch_predict(\n",
-        "    job_display_name=\"iris_\" + UUID,\n",
+        "    job_display_name=\"iris-unique\",\n",
        "    gcs_source=gcs_input_uri,\n",
        "    gcs_destination_prefix=BUCKET_URI,\n",
        "    instances_format=\"jsonl\",\n",
@@ -952,9 +849,9 @@
        "id": "batch_request_wait:mbsdk"
      },
      "source": [
-        "### Wait for completion of batch prediction job\n",
+        "#### 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. Alternatively, you can set the parameter `sync` to `True` in the `batch_predict()` method to block until the batch prediction job is completed."
      ]
    },
    {
@@ -1008,11 +905,11 @@
        "id": "get_batch_prediction:mbsdk,custom,lcn"
      },
      "source": [
-        "### Get the predictions\n",
+        "#### Get the predictions\n",
        "\n",
        "Next, get the 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 prediction requests in a JSON format:\n",
+        "The results are written to the Cloud Storage output bucket specified in the batch prediction request. Call the `iter_outputs()` method to get a list of each Cloud Storage file generated with the results. Each file contains one or more prediction requests in a JSON format:\n",
        "\n",
        "- `instance`: The prediction request.\n",
        "- `prediction`: The prediction response."
@@ -1062,16 +959,9 @@
        "id": "make_online_predictions:migration"
      },
      "source": [
-        "## Make online predictions"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "deploy_model:migration,new,mbsdk"
-      },
-      "source": [
-        "### [predictions.deploy-model-api](https://cloud.google.com/vertex-ai/docs/predictions/deploy-model-api)"
+        "### Generate online predictions\n",
+        "\n",
+        "To learn more about online predictions refer, [Overview of getting predictions on Vertex AI](https://cloud.google.com/vertex-ai/docs/predictions/overview)."
      ]
    },
    {
@@ -1080,13 +970,13 @@
        "id": "deploy_model:mbsdk,cpu"
      },
      "source": [
-        "## Deploy the model\n",
+        "#### 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, 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 is only one model, then specify `traffic_split` 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",
        "- `machine_type`: The type of machine to use for training.\n",
        "- `starting_replica_count`: The number of compute instances to initially provision.\n",
@@ -1101,7 +991,7 @@
      },
      "outputs": [],
      "source": [
-        "DEPLOYED_NAME = \"iris-\" + UUID\n",
+        "DEPLOYED_NAME = \"iris-unique\"\n",
        "\n",
        "TRAFFIC_SPLIT = {\"0\": 100}\n",
        "\n",
@@ -1114,7 +1004,8 @@
        "    machine_type=DEPLOY_COMPUTE,\n",
        "    min_replica_count=MIN_NODES,\n",
        "    max_replica_count=MAX_NODES,\n",
-        ")"
+        ")\n",
+        "endpoint.wait()"
      ]
    },
    {
@@ -1135,24 +1026,15 @@
        "    INFO:google.cloud.aiplatform.models:Endpoint model deployed. Resource name: projects/759209241365/locations/us-central1/endpoints/4867177336350441472"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "endpoints_predict:migration,new,mbsdk"
-      },
-      "source": [
-        "### [predictions.online-prediction-automl](https://cloud.google.com/vertex-ai/docs/predictions/online-predictions-automl)"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "make_test_item:xgboost,tabular,iris"
      },
      "source": [
-        "### Make test item\n",
+        "#### Create a 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."
+        "Use synthetic data as a test data item. Don’t be concerned about using synthetic data – since it's just for demonstration purposes."
      ]
    },
    {
@@ -1172,9 +1054,9 @@
        "id": "predict_request:mbsdk,custom,lcn"
      },
      "source": [
-        "### Make the prediction\n",
+        "#### Make the prediction\n",
        "\n",
-        "Now that your `Model` resource is deployed to an `Endpoint` resource, you can do online predictions by sending prediction requests to the `Endpoint` resource.\n",
+        "Now that your model resource is deployed to an endpoint resource, you can make online predictions by sending prediction requests to the endpoint resource.\n",
        "\n",
        "#### Request\n",
        "\n",
@@ -1182,15 +1064,15 @@
        "\n",
        "    [feature_list]\n",
        "\n",
-        "Since the predict() method can take multiple items (instances), send your single test item as a list of one test item.\n",
+        "Since the `predict()` 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 predict() call is a Python dictionary with the following entries:\n",
+        "The response from the `predict()` call is a Python dictionary with the following entries:\n",
        "\n",
        "- `ids`: The internal assigned unique identifiers for each prediction request.\n",
        "- `predictions`: The predicted confidence, between 0 and 1, per class label.\n",
-        "- `deployed_model_id`: The Vertex AI identifier for the deployed `Model` resource which did the predictions."
+        "- `deployed_model_id`: The Vertex AI identifier for the deployed model resource which did the predictions.\n"
      ]
    },
    {
@@ -1215,7 +1097,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 generating predictions, undeploy the model from the endpoint resouce. This deprovisions all compute resources and ends billing for the deployed model."
      ]
    },
    {
@@ -1251,10 +1133,6 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
-        "delete_bucket = False\n",
-        "\n",
        "# Delete the model using the Vertex model object\n",
        "model.delete()\n",
        "\n",
@@ -1267,7 +1145,12 @@
        "# Delete the batch prediction job using the Vertex batch prediction object\n",
        "batch_predict_job.delete()\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "# Delete locally generated files\n",
+        "! rm -rf custom custom.tar.gz\n",
+        "\n",
+        "# Delete cloud storage bucket\n",
+        "delete_bucket = False  # set True for deletion\n",
+        "if delete_bucket:\n",
        "    ! gsutil rm -r $BUCKET_URI"
      ]
    }
@@ -30,26 +30,27 @@
      },
      "source": [
        "# Vertex AI Migration: Hyperparameter Tuning\n",
-        "\n",
        "<table align=\"left\">\n",
-        "\n",
-        "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/migration/sdk-hyperparameter-tuning.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\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://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/migration/sdk-hyperparameter-tuning.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%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>\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/migration/sdk-hyperparameter-tuning.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 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>\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/>"
      ]
@@ -82,12 +83,12 @@
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
-        "- `Vertex AI Training`\n",
-        "- `Vertex AI Hyperparameter Tuning`\n",
+        "- Vertex AI Training\n",
+        "- Vertex AI hyperparameter tuning\n",
        "\n",
        "The steps performed include:\n",
        "\n",
-        "- Create a `Vertex AI` hyperparameter tuning job for training a TensorFlow model."
+        "- Create a Vertex AI hyperparameter tuning job for training a TensorFlow model."
      ]
    },
    {
@@ -114,11 +115,9 @@
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
-        "Learn about [Vertex AI\n",
-        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
-        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
-        "Calculator](https://cloud.google.com/products/calculator/)\n",
-        "to generate a cost estimate based on your projected usage."
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and\n",
+        "[Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the \n",
+        "[Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
@@ -127,9 +126,8 @@
        "id": "install_aip:mbsdk"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the packages required for executing this notebook."
+        "## Get started\n",
+        "Install Vertex AI SDK for Python and other required packages"
      ]
    },
    {
@@ -150,7 +148,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."
      ]
    },
    {
@@ -161,159 +160,74 @@
      },
      "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": "4de1bd77992b"
      },
      "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": "56e219dbcb9a"
      },
      "source": [
-        "### Set your project ID\n",
-        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "* Run `gcloud config list`.\n",
-        "* Run `gcloud projects list`.\n",
-        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+        "### Authenticate your notebook environment (Colab only)\n",
+        "Authenticate your environment on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "set_project_id"
+        "id": "c97be6a73155"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "before_you_begin:nogpu"
+      },
+      "source": [
+        "## 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": "2d96e0c47bed"
      },
      "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": "timestamp"
-      },
-      "source": [
-        "#### UUID\n",
-        "\n",
-        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a uuid for each instance session, and append it onto the name of resources you create in this tutorial."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "timestamp"
-      },
-      "outputs": [],
-      "source": [
-        "import random\n",
-        "import string\n",
-        "\n",
-        "\n",
-        "# Generate a uuid of a specifed length(default=8)\n",
-        "def generate_uuid(length: int = 8) -> str:\n",
-        "    return \"\".join(random.choices(string.ascii_lowercase + string.digits, k=length))\n",
-        "\n",
-        "\n",
-        "UUID = generate_uuid()"
-      ]
-    },
-    {
-      "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\"}"
      ]
    },
    {
@@ -344,7 +258,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."
      ]
    },
    {
@@ -355,7 +269,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} {BUCKET_URI}"
      ]
    },
    {
@@ -378,7 +292,7 @@
      },
      "outputs": [],
      "source": [
-        "import google.cloud.aiplatform as aip"
+        "import google.cloud.aiplatform as aiplatform"
      ]
    },
    {
@@ -400,7 +314,7 @@
      },
      "outputs": [],
      "source": [
-        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
+        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -420,7 +334,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"
+        "Learn more [about hardware accelerator support for your region](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators)"
      ]
    },
    {
@@ -442,15 +356,15 @@
        "id": "container:training,prediction"
      },
      "source": [
-        "#### Set pre-built containers\n",
+        "#### Set prebuilt containers\n",
        "\n",
-        "Set the pre-built Docker container image for training and prediction.\n",
+        "Set the prebuilt 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 [prebuilt containers for training](https://cloud.google.com/ai-platform-unified/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 [prebuilt containers for prediction](https://cloud.google.com/ai-platform-unified/docs/predictions/pre-built-containers)."
      ]
    },
    {
@@ -489,7 +403,7 @@
        "\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 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 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",
@@ -501,7 +415,7 @@
        " - `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*."
      ]
    },
    {
@@ -543,7 +457,9 @@
        "\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",
+        "*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",
@@ -587,10 +503,10 @@
      "source": [
        "#### Task.py contents\n",
        "\n",
-        "In the next cell, you write the contents of the hyperparameter tuning script task.py. I 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 hyperparameter tuning script task.py. It's there for you to browse. In summary:\n",
        "\n",
        "- Parse the command line arguments for the hyperparameter settings for the current trial.\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",
+        " - Get the directory where the model artifacts from the command line (--model_dir) are to be saved. If unspecified, then from the environment variable AIP_MODEL_DIR is used.\n",
        "- Download and preprocess the Boston Housing dataset.\n",
        "- Build a DNN model.\n",
        "- The number of units per dense layer and learning rate hyperparameter values are used during the build and compile of the model.\n",
@@ -714,7 +630,7 @@
      "source": [
        "#### 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."
+        "Next, package the training folder into a compressed tar ball, and then store it in your Cloud Storage bucket."
      ]
    },
    {
@@ -757,9 +673,9 @@
      "source": [
        "### Prepare your machine specification\n",
        "\n",
-        "Now define the machine specification for your custom training job. This tells Vertex what type of machine instance to provision for the training.\n",
+        "Now define the machine specification for your custom training job. Specify the machine type to provision for training. This informs Vertex AI about the computational resources you need for your training job.\n",
        "  - `machine_type`: The type of GCP instance to provision -- e.g., n1-standard-8.\n",
-        "  - `accelerator_type`: The type, if any, of hardware accelerator. In this tutorial if you previously set the variable `TRAIN_GPU != None`, you are using a GPU; otherwise you use a CPU.\n",
+        "  - `accelerator_type`: The type, if any hardware accelerator. In this tutorial if you previously set the variable `TRAIN_GPU != None`, you are using a GPU; otherwise you use a CPU.\n",
        "  - `accelerator_count`: The number of accelerators."
      ]
    },
@@ -789,7 +705,7 @@
      "source": [
        "### Prepare your disk specification\n",
        "\n",
-        "(optional) Now define the disk specification for your custom training job. This tells Vertex what type and size of disk to provision in each machine instance for the training.\n",
+        "(optional) Now define the disk specification for your custom training job. Configure the disk options , this tells Vertex AI exactly what type and size of disk to allocate for each machine instance during training.\n",
        "\n",
        "  - `boot_disk_type`: Either SSD or Standard. SSD is faster, and Standard is less expensive. Defaults to SSD.\n",
        "  - `boot_disk_size_gb`: Size of disk in GB."
@@ -831,9 +747,9 @@
        "\n",
        "-`package_uris`: This is a list of the locations (URIs) of your python training packages to install on the provisioned instance. The locations need to be in a Cloud Storage bucket. These can be either individual python files or a zip (archive) of an entire package. In the later case, the job service  unzip (unarchive) the contents into the docker image.\n",
        "\n",
-        "-`python_module`: The Python module (script) to invoke for running the custom training job. In this example, you  be invoking `trainer.task.py` -- note that it was not neccessary to append the `.py` suffix.\n",
+        "-`python_module`: The Python module (script) to invoke for running the custom training job. In this example, you  be invoking `trainer.task.py` -- note that it wasn't neccessary to append the `.py` suffix.\n",
        "\n",
-        "-`args`: The command line arguments to pass to the corresponding Pythom module. In this example, you  be setting:\n",
+        "-`args`: The command line arguments to pass to the corresponding Python module. In this example, you are setting:\n",
        "  - `\"--model-dir=\" + MODEL_DIR` : The Cloud Storage location where to store the model artifacts. There are two ways to tell the training script where to save 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",
@@ -853,7 +769,7 @@
      },
      "outputs": [],
      "source": [
-        "JOB_NAME = \"custom_job_\" + UUID\n",
+        "JOB_NAME = \"custom_job_\" + \"unique\"\n",
        "MODEL_DIR = \"{}/{}\".format(BUCKET_URI, JOB_NAME)\n",
        "\n",
        "if not TRAIN_NGPU or TRAIN_NGPU < 2:\n",
@@ -925,7 +841,9 @@
      },
      "outputs": [],
      "source": [
-        "job = aip.CustomJob(display_name=\"boston_\" + UUID, worker_pool_specs=worker_pool_spec)\n",
+        "job = aiplatform.CustomJob(\n",
+        "    display_name=\"boston_\" + \"unique\", worker_pool_specs=worker_pool_spec\n",
+        ")\n",
        "\n",
        "# print(job)"
      ]
@@ -956,8 +874,8 @@
      "source": [
        "from google.cloud.aiplatform import hyperparameter_tuning as hpt\n",
        "\n",
-        "hpt_job = aip.HyperparameterTuningJob(\n",
-        "    display_name=\"boston_\" + UUID,\n",
+        "hpt_job = aiplatform.HyperparameterTuningJob(\n",
+        "    display_name=\"boston_\" + \"unique\",\n",
        "    custom_job=job,\n",
        "    metric_spec={\n",
        "        \"val_loss\": \"minimize\",\n",
@@ -1159,10 +1077,10 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
        "delete_bucket = False\n",
        "\n",
+        "# # remove locally generated  directory\n",
+        "! rm -r custom\n",
        "# Delete the training job\n",
        "try:\n",
        "    job.delete()\n",
@@ -1172,8 +1090,8 @@
        "# Delete the HPT job using the Vertex batch prediction object\n",
        "hpt_job.delete()\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
-        "    ! gsutil rm -r $BUCKET_URI\n"
+        "# if delete_bucket:\n",
+        "# #     ! gsutil rm -r $BUCKET_URI"
      ]
    }
  ],
@@ -32,23 +32,26 @@
        "# Get started with Google Artifact Registry\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/ml_metadata/get_started_with_google_artifact_registry.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\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://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/ml_metadata/get_started_with_google_artifact_registry.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%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>\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/ml_metadata/get_started_with_google_artifact_registry.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://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>\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/>"
      ]
@@ -75,11 +78,11 @@
      "source": [
        "### Objective\n",
        "\n",
-        "In this tutorial, you learn how to use `Google Artifact Registry`.\n",
+        "In this tutorial, you learn how to use Google Artifact Registry.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
-        "- `Google Artifact Registry`\n",
+        "- Google Artifact Registry\n",
        "\n",
        "The steps performed include:\n",
        "\n",
@@ -90,17 +93,6 @@
        "- Deleting a private Docker repository."
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "4ced09c1b4ce"
-      },
-      "source": [
-        "### Dataset\n",
-        "\n",
-        "No dataset is used in this tutorial. References to an example dataset are for demonstration purposes."
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -114,7 +106,9 @@
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
-        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and\n",
+        "[Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the\n",
+        "[Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
@@ -123,9 +117,8 @@
        "id": "install_mlops"
      },
      "source": [
-        "## Installations\n",
-        "\n",
-        "Install the packages required for executing the notebook."
+        "## Get started\n",
+        "Install Vertex AI SDK for Python and other required packages"
      ]
    },
    {
@@ -142,48 +135,75 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "hhq5zEbGg0XX"
+        "id": "8d726e21c0bb"
      },
      "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."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "EzrelQZ22IZj"
+        "id": "3b9119a60525"
      },
      "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": "lWEdiXsJg0XY"
+        "id": "96254aa096b5"
      },
      "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": "8bc8a29f9001"
+        "id": "1d7064423926"
      },
      "source": [
-        "#### Set your project ID\n",
+        "### Authenticate your notebook environment (Colab only)\n",
+        "Authenticate your environment on Google Colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "405401bbd1c7"
+      },
+      "outputs": [],
+      "source": [
+        "if \"google.colab\" in sys.modules:\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)"
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "d183adfc792a"
+      },
+      "source": [
+        "### Set Google Cloud project information\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,103 +215,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": "47bc07d4231b"
-      },
-      "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": "959545da671a"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "dr--iN2kAylZ"
-      },
-      "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": "sBCra4QMA2wR"
-      },
-      "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": "PyQmSRbKA8r-"
-      },
-      "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 +246,7 @@
        "id": "-EcIXiGsCePi"
      },
      "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,19 +257,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} -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"
+        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
@@ -356,7 +268,7 @@
      "source": [
        "## Introduction to Google Artifact Registry\n",
        "\n",
-        "The `Google Artifact Registry` is a service for storing and managing artifacts in private repositories, including container images, Helm charts, and language packages. It is the recommended container image registry for Google Cloud.\n",
+        "The Google Artifact Registry is a service for storing and managing artifacts in private repositories, including container images, Helm charts, and language packages. It's the recommended container image registry for Google Cloud.\n",
        "\n",
        "Learn more about [Quick start for Docker](https://cloud.google.com/artifact-registry/docs/docker/quickstart)"
      ]
@@ -369,7 +281,7 @@
      "source": [
        "### Enable Artifact Registry API\n",
        "\n",
-        "First, you must enable the Artifact Registry API service for your project.\n",
+        "First, enable the Artifact Registry API service for your project.\n",
        "\n",
        "Learn more about [Enabling service](https://cloud.google.com/artifact-registry/docs/enable-service)."
      ]
@@ -410,7 +322,7 @@
      "source": [
        "PRIVATE_REPO = \"my-docker-repo\"\n",
        "\n",
-        "! gcloud artifacts repositories create {PRIVATE_REPO} --repository-format=docker --location={REGION} --description=\"Docker repository\"\n",
+        "! gcloud artifacts repositories create {PRIVATE_REPO} --repository-format=docker --location={LOCATION} --description=\"Docker repository\"\n",
        "\n",
        "! gcloud artifacts repositories list"
      ]
@@ -423,7 +335,7 @@
      "source": [
        "### Configure authentication to your private repo\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 region."
      ]
    },
    {
@@ -434,7 +346,7 @@
      },
      "outputs": [],
      "source": [
-        "! gcloud auth configure-docker {REGION}-docker.pkg.dev --quiet"
+        "! gcloud auth configure-docker {LOCATION}-docker.pkg.dev --quiet"
      ]
    },
    {
@@ -471,7 +383,7 @@
        "\n",
        "- Tagging the Docker image with a repository name configures the docker push command to push the image to a specific location, e.g., us-central1-docker.pkg.dev.\n",
        "\n",
-        "- `:my-tag` is a tag you're adding to the Docker image. If a tag is not specified, it defaults to `:latest`."
+        "- :my-tag is a tag you're adding to the Docker image. It defaults to :latest."
      ]
    },
    {
@@ -506,7 +418,7 @@
      },
      "outputs": [],
      "source": [
-        "! docker push {REGION}-docker.pkg.dev/{PROJECT_ID}/{PRIVATE_REPO}/{CONTAINER_NAME}"
+        "! docker push {LOCATION}-docker.pkg.dev/{PROJECT_ID}/{PRIVATE_REPO}/{CONTAINER_NAME}"
      ]
    },
    {
@@ -528,7 +440,7 @@
      },
      "outputs": [],
      "source": [
-        "! docker pull {REGION}-docker.pkg.dev/{PROJECT_ID}/{PRIVATE_REPO}/{CONTAINER_NAME}"
+        "! docker pull {LOCATION}-docker.pkg.dev/{PROJECT_ID}/{PRIVATE_REPO}/{CONTAINER_NAME}"
      ]
    },
    {
@@ -539,7 +451,7 @@
      "source": [
        "### Deleting your private Docker repostory\n",
        "\n",
-        "Finally, once your private repository becomes obsolete, use the command `gcloud artifacts repositories delete` to delete it `Google Artifact Registry`."
+        "Finally, once your private repository becomes obsolete, use the `gcloud artifacts repositories delete` command to remove the repository from the Google Artifact Registry."
      ]
    },
    {
@@ -550,7 +462,7 @@
      },
      "outputs": [],
      "source": [
-        "! gcloud artifacts repositories delete {PRIVATE_REPO} --location={REGION} --quiet"
+        "! gcloud artifacts repositories delete {PRIVATE_REPO} --location={LOCATION} --quiet"
      ]
    }
  ],
@@ -32,25 +32,28 @@
        "# Vertex AI: Track parameters and metrics for custom training jobs\n",
        "\n",
        "<table align=\"left\">\n",
-        "\n",
-        "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/ml_metadata/sdk-metric-parameter-tracking-for-custom-jobs.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\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://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/ml_metadata/sdk-metric-parameter-tracking-for-custom-jobs.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%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>\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/ml_metadata/sdk-metric-parameter-tracking-for-custom-jobs.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://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",
-        "</table>"
+        "  </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/>"
      ]
    },
    {
@@ -63,7 +66,9 @@
        "\n",
        "This notebook demonstrates how to track metrics and parameters for Vertex AI custom training jobs, and how to perform detailed analysis using this data.\n",
        "\n",
-        "Learn more about [Vertex ML Metadata](https://cloud.google.com/vertex-ai/docs/ml-metadata), [Custom training](https://cloud.google.com/vertex-ai/docs/training/custom-training), and [Vertex AI Experiments](https://cloud.google.com/vertex-ai/docs/experiments/intro-vertex-ai-experiments)."
+        "Learn more about [Vertex ML Metadata](https://cloud.google.com/vertex-ai/docs/ml-metadata),\n",
+        "[Custom training](https://cloud.google.com/vertex-ai/docs/training/custom-training), and \n",
+        "[Vertex AI Experiments](https://cloud.google.com/vertex-ai/docs/experiments/intro-vertex-ai-experiments)."
      ]
    },
    {
@@ -113,11 +118,9 @@
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
-        "Learn about [Vertex AI\n",
-        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
-        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
-        "Calculator](https://cloud.google.com/products/calculator/)\n",
-        "to generate a cost estimate based on your projected usage."
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and \n",
+        "[Cloud Storage pricing](https://cloud.google.com/storage/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 +129,8 @@
        "id": "i7EUnXsZhAGF"
      },
      "source": [
-        "### Installation\n",
-        "\n",
-        "Install the packages required for executing this notebook."
+        "### Get Started\n",
+        "Install Vertex AI SDK for Python and other required packages"
      ]
    },
    {
@@ -141,54 +143,84 @@
      "source": [
        "! pip3 install --upgrade tensorflow \\\n",
        "                         google-cloud-aiplatform \\\n",
-        "                         scikit-learn -q"
+        "                         scikit-learn -q \\\n",
+        "                         pandas"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "hhq5zEbGg0XX"
+        "id": "5eec42e37bcf"
      },
      "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."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "EzrelQZ22IZj"
+        "id": "dcc98768955f"
      },
      "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": "lWEdiXsJg0XY"
+        "id": "4de1bd77992b"
      },
      "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": "8bc8a29f9001"
+        "id": "56e219dbcb9a"
      },
      "source": [
-        "#### Set your project ID\n",
+        "### Authenticate your notebook environment (Colab only)\n",
+        "Authenticate your environment on Google Colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "c97be6a73155"
+      },
+      "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": "442da99b7efa"
+      },
+      "source": [
+        "### 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)"
      ]
    },
    {
@@ -200,103 +232,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": "47bc07d4231b"
-      },
-      "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": "959545da671a"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "dr--iN2kAylZ"
-      },
-      "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": "sBCra4QMA2wR"
-      },
-      "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": "PyQmSRbKA8r-"
-      },
-      "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\""
      ]
    },
    {
@@ -327,7 +263,7 @@
        "id": "-EcIXiGsCePi"
      },
      "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."
      ]
    },
    {
@@ -338,7 +274,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
@@ -367,8 +303,6 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
        "import pandas as pd\n",
        "from google.cloud import aiplatform\n",
        "from sklearn.metrics import mean_absolute_error, mean_squared_error\n",
@@ -381,7 +315,7 @@
        "id": "O8XJZB3gR8eL"
      },
      "source": [
-        "## Initialize Vertex AI and set an _experiment_\n"
+        "## Initialize Vertex AI and set an experiment\n"
      ]
    },
    {
@@ -423,7 +357,7 @@
      "source": [
        "aiplatform.init(\n",
        "    project=PROJECT_ID,\n",
-        "    location=REGION,\n",
+        "    location=LOCATION,\n",
        "    staging_bucket=BUCKET_URI,\n",
        "    experiment=EXPERIMENT_NAME,\n",
        ")"
@@ -435,7 +369,7 @@
        "id": "6PlilQPFeS_h"
      },
      "source": [
-        "## Tracking parameters and metrics in Vertex AI custom training jobs"
+        "### Tracking parameters and metrics in Vertex AI custom training jobs"
      ]
    },
    {
@@ -444,7 +378,7 @@
        "id": "f8fd397cc4f6"
      },
      "source": [
-        "### Download the Dataset to Cloud Storage"
+        "# Download the Dataset to Cloud Storage"
      ]
    },
    {
@@ -580,7 +514,7 @@
        "id": "k_QorXXztzPH"
      },
      "source": [
-        "Start a new experiment run to track training parameters and start the training job. Note that this operation will take around 10 mins."
+        "Start a new experiment run to track training parameters and start the training job. Note that this operation takes around 10 minutes."
      ]
    },
    {
@@ -620,7 +554,7 @@
        "id": "O-uCOL3Naap4"
      },
      "source": [
-        "Next, deploy your Vertex AI Model resource to a Vertex AI Endpoint resource. This operation will take 10-20 mins."
+        "Next, deploy your Vertex AI Model resource to a Vertex AI endpoint resource. This operation takes 10-20 minutes."
      ]
    },
    {
@@ -824,7 +758,7 @@
      },
      "outputs": [],
      "source": [
-        "# Warning: Setting this to true will delete everything in your bucket\n",
+        "# Warning: Setting this to true deletes everything in your bucket\n",
        "delete_bucket = False\n",
        "\n",
        "# Delete dataset\n",
@@ -832,7 +766,7 @@
        "\n",
        "# Delete experiment\n",
        "experiment = aiplatform.Experiment(\n",
-        "    experiment_name=EXPERIMENT_NAME, project=PROJECT_ID, location=REGION\n",
+        "    experiment_name=EXPERIMENT_NAME, project=PROJECT_ID, location=LOCATION\n",
        ")\n",
        "experiment.delete()\n",
        "\n",
@@ -849,7 +783,7 @@
        "model.delete()\n",
        "\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
--- a/Show More
+++ b/Show More