Update the server port in model_garden_keras_stable_diffusion.ipynb

Add new model Claude 3.5 Sonnet and update regions for other models (#3112 )
* add new model and update region for others * fix minor error --------- Co-authored-by: Huy Ngo <huyngo@google.com>
2026-09-26 22:51:56 +00:00 · 2024-06-21 10:49:11 -07:00 · 2024-06-20 18:22:34 +00:00 · 2024-06-20 12:39:50 +00:00 · 2024-06-18 21:47:37 +00:00 · 2024-06-18 12:36:46 +00:00
99 changed files with 5861 additions and 17409 deletions
@@ -3,8 +3,8 @@ ipython
 jupyter
 nbconvert
 black==24.4.2
-pyupgrade==3.15.2
+pyupgrade==3.16.0
 isort==5.13.2
-flake8==7.0.0
+flake8==7.1.0
 nbqa==1.8.5

@@ -1,23 +1,44 @@
-# 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 and learn
+
+You can explore, learn, and contribute to this repository to unleash the full potential of machine learning on Vertex AI! You can 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
+
+
+## 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
+
 ```

 ## Contributing
@@ -35,3 +56,6 @@ This is not an officially supported Google product. The code in this repository
 ## 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)
@@ -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}
+     
+
@@ -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
@@ -1,37 +1,51 @@
-from kfp.v2 import dsl
-
-@dsl.component(base_image='python:3.8',packages_to_install=['google-cloud-aiplatform==1.36.0'])
-def deploy_model(
-    model_id: str,
-    endpoint_id: str,
-    machine_type: str,
-    min_replica_count: int,
-    max_replica_count: int,
-):
-    import json
-    from google.cloud import aiplatform
-
-    model = aiplatform.Model(model_id)
-    endpoint = aiplatform.Endpoint(endpoint_id)
-
-    endpoint = model.deploy(
-        endpoint=endpoint,
-        machine_type=machine_type,
-        min_replica_count=min_replica_count,
-        max_replica_count=max_replica_count,
-    )
-
-
-@dsl.pipeline(name='deploy-model')
-def pipeline_deploy_model():
-    project = "projects/990000000009/locations/us-west1"
-    model_id = project + "/models/1100000000000000001"
-    endpoint_id = project + "/endpoints/2200000000000000002"
-    deploy_model(model_id, endpoint_id, "n1-standard-2", 1, 1)
-
-
-if __name__ == "__main__":
-    from kfp.v2 import compiler
-    compiler.Compiler().compile(
-        pipeline_func=pipeline_deploy_model,
-        package_path='deploy_model.json')
+from kfp.v2 import dsl
+
+@dsl.component(base_image='python:3.8',packages_to_install=['google-cloud-aiplatform==1.36.0'])
+def deploy_model(
+    model_id: str,
+    endpoint_id: str,
+    machine_type: str,
+    min_replica_count: int,
+    max_replica_count: int,
+):
+    import json
+    from google.cloud import aiplatform
+
+    model = aiplatform.Model(model_id)
+    endpoint = aiplatform.Endpoint(endpoint_id)
+
+    endpoint = model.deploy(
+        endpoint=endpoint,
+        machine_type=machine_type,
+        min_replica_count=min_replica_count,
+        max_replica_count=max_replica_count,
+    )
+
+
+@dsl.component(base_image='python:3.8',packages_to_install=['google-cloud-aiplatform==1.36.0'])
+def delete_endpoint(
+    endpoint_id: str,
+):
+    from google.cloud import aiplatform
+
+    endpoint = aiplatform.Endpoint(endpoint_id)
+    endpoint.undeploy_all()
+    endpoint.delete()
+
+
+@dsl.pipeline(name='deploy-model')
+def pipeline_deploy_model():
+    project = "projects/990000000009/locations/us-west1"
+    model_id = project + "/models/1100000000000000001"
+    endpoint_id = project + "/endpoints/2200000000000000002"
+    deploy_model(model_id, endpoint_id, "n1-standard-2", 1, 1)
+
+    # After serving predictions, recycling computing resources
+    delete_endpoint(endpoint_id)
+
+
+if __name__ == "__main__":
+    from kfp.v2 import compiler
+    compiler.Compiler().compile(
+        pipeline_func=pipeline_deploy_model,
+        package_path='deploy_model.json')
@@ -107,7 +107,6 @@
 /notebooks/community/model_garden/model_garden_movinet_clip_classification.ipynb @KCFindstr
 /notebooks/community/model_garden/model_garden_movinet_action_recognition.ipynb @KCFindstr
 /notebooks/community/model_garden/model_garden_pytorch_open_clip.ipynb @lydhr
-/notebooks/community/model_garden/model_garden_pytorch_llama2_peft.ipynb @genquan9
 /notebooks/community/model_garden/model_garden_pytorch_codellama.ipynb @xiangxu-google
 /notebooks/community/model_garden/model_garden_pytorch_nllb.ipynb @weigary
 /notebooks/community/model_garden/model_garden_pytorch_vicuna.ipynb @dstnluong-google
@@ -141,3 +140,4 @@
 /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_timesfm_deployment_on_vertex.ipynb @siriuz42
@@ -64,7 +64,7 @@
        "    - [intfloat/multilingual-e5-small](https://huggingface.co/intfloat/multilingual-e5-small): 118M params\n",
        "    - [intfloat/e5-base-v2](https://huggingface.co/intfloat/e5-base-v2): 109M params\n",
        "    - [intfloat/e5-small-v2](https://huggingface.co/intfloat/e5-small-v2): 33M params\n",
-        "- Run inference on the deployed Vertex AI Endpoint.\n",
+        "- Run inference on the deployed Vertex AI Endpoint\n",
        "\n",
        "\n",
        "### Costs\n",
@@ -39,7 +39,7 @@
        "  </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_gemma_fine_tuning_batch_deployment_on_rov.ipynb\">\n",
-        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\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",
@@ -99,6 +99,7 @@
        "\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",
        "import uuid\n",
        "from datetime import datetime\n",
@@ -346,7 +347,104 @@
        "        deploy_request_timeout=1800,\n",
        "        service_account=service_account,\n",
        "    )\n",
-        "    return model, endpoint"
+        "    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",
+        "        )"
      ]
    },
    {
@@ -484,6 +582,14 @@
        "\n",
        "    replica_count = 1\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=True,\n",
+        "    )\n",
+        "\n",
        "    # Setup training job.\n",
        "    job_name = get_job_name_with_datetime(\"gemma-lora-train\")\n",
        "\n",
@@ -560,6 +666,14 @@
        "    accelerator_type = \"NVIDIA_L4\"\n",
        "    accelerator_count = 1\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 a larger max_model_len will require more GPU memory.\n",
        "    max_model_len = 2048\n",
        "\n",
@@ -791,6 +905,14 @@
        "\n",
        "    replica_count = 1\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=True,\n",
+        "    )\n",
+        "\n",
        "    # Setup training job.\n",
        "    job_name = get_job_name_with_datetime(\"gemma-keras-lora-train\")\n",
        "\n",
@@ -933,9 +1055,23 @@
        "    if \"2b\" in KAGGLE_MODEL_ID:\n",
        "        # Sets ct5lp-hightpu-1t (1 TPU chip) to deploy Gemma 2B models.\n",
        "        machine_type = \"ct5lp-hightpu-1t\"\n",
+        "        accelerator_type = \"TPU_V5e\"\n",
+        "        # Note: 1 TPU-V5e chip has only 1 core.\n",
+        "        accelerator_count = 1\n",
        "    else:\n",
        "        # Sets ct5lp-hightpu-4t (4 TPU chips) to deploy Gemma 7B models.\n",
        "        machine_type = \"ct5lp-hightpu-4t\"\n",
+        "        accelerator_type = \"TPU_V5e\"\n",
+        "        # Note: 1 TPU-V5e chip has only 1 core.\n",
+        "        accelerator_count = 4\n",
+        "\n",
+        "    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",
@@ -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",
@@ -160,7 +160,7 @@
        "\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "# Define common functions.\n",
@@ -240,83 +240,96 @@
        "\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",
+        "    \"\"\"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",
-        "  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",
+        "    \"\"\"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",
-        "      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",
+        "        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",
+        "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",
+        "    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",
@@ -423,7 +436,10 @@
        "# The Gemma base model.\n",
        "MODEL_ID = \"google/gemma-1.1-2b-it\"  # @param[\"google/gemma-2b\", \"google/gemma-2b-it\", \"google/gemma-7b\", \"google/gemma-7b-it\", \"google/gemma-1.1-2b-it\", \"google/gemma-1.1-7b-it\"] {isTemplate:true}\n",
        "# The accelerator to use.\n",
+        "\n",
+        "# fmt: off\n",
        "ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"  # @param[\"NVIDIA_TESLA_V100\", \"NVIDIA_L4\", \"NVIDIA_TESLA_A100\"]\n",
+        "# fmt: on\n",
        "\n",
        "# Batch size for finetuning.\n",
        "per_device_train_batch_size = 1  # @param{type:\"integer\"}\n",
@@ -464,11 +480,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",
+        "check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=ACCELERATOR_TYPE,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=True,\n",
+        ")\n",
        "\n",
        "# Setup training job.\n",
        "job_name = get_job_name_with_datetime(\"gemma-lora-train\")\n",
@@ -542,11 +560,13 @@
        "accelerator_type = \"NVIDIA_L4\"\n",
        "accelerator_count = 1\n",
        "\n",
-        "check_quota(project_id=PROJECT_ID,\n",
-        "            region=REGION,\n",
-        "            accelerator_type=accelerator_type,\n",
-        "            accelerator_count=accelerator_count,\n",
-        "            is_for_training=False)\n",
+        "check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=False,\n",
+        ")\n",
        "\n",
        "# Note that larger token counts will require more GPU memory.\n",
        "max_input_length = 512\n",
@@ -149,7 +149,7 @@
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240429_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
@@ -122,7 +122,7 @@
        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240429_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
@@ -157,7 +157,7 @@
        "    auth.authenticate_user(project_id=PROJECT_ID)\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "# Define common functions.\n",
@@ -107,14 +107,14 @@
        "\n",
        "The workflow takes the following inputs:\n",
        "\n",
-        "* **preference_dataset**: Required, `str`. Cloud storage path to a human preference dataset used to train a reward model. Download the [example dataset](https://pantheon.corp.google.com/storage/browser/vertex-ai/generative-ai/rlhf/text_small/summarize_from_feedback_tfds/comparisons/train?pageState=(%22StorageObjectListTable%22:(%22f%22:%22%255B%255D%22))&e=13802955&jsmode=O&mods=-ai_platform_fake_service&project=vertex-ai&prefix=&forceOnObjectsSortingFiltering=false), and see its content.\n",
+        "* **preference_dataset**: Required, `str`. Cloud storage path to a human preference dataset used to train a reward model. Download the [example dataset](https://console.cloud.google.com/storage/browser/vertex-ai/generative-ai/rlhf/text_small/summarize_from_feedback_tfds/comparisons/train), and see its content.\n",
        "> The dataset format is `jsonl`. Each line in the file contains one example:\n",
        "> * An `input_text` field containing the prompt.\n",
        "> * `candidate_0`, `candidate_1` fields containing the candidate responses.\n",
        "> * A `choice` field that specifies which candidate is preferred. This field is an integer `0` or `1`. `0` means `candidate_0` is preferred, `1` means `candidate_1` is preferred.\n",
        "> * For example: `{\"input_text\": \"Create a description for Plantation Palms.\", \"candidate_0\": \"Enjoy some fun in the sun at Gulf Shores.\", \"candidate_1\": \"A Tranquil Oasis of Natural Beauty.\", \"choice\": 0}`\n",
        "\n",
-        "* **prompt_dataset**: Required, `str`. Cloud storage path to an unlabled prompt dataset used for reinforcement learning. Download the [example dataset](https://pantheon.corp.google.com/storage/browser/vertex-ai/generative-ai/rlhf/text_small/reddit_tfds/train?pageState=(%22StorageObjectListTable%22:(%22f%22:%22%255B%255D%22))&e=13802955&jsmode=O&mods=-ai_platform_fake_service&project=vertex-ai&prefix=&forceOnObjectsSortingFiltering=false), and see its content.\n",
+        "* **prompt_dataset**: Required, `str`. Cloud storage path to an unlabled prompt dataset used for reinforcement learning. Download the [example dataset](https://console.cloud.google.com/storage/browser/vertex-ai/generative-ai/rlhf/text_small/reddit_tfds/train), and see its content.\n",
        "> The dataset format is `jsonl`. Each line in the file contains one example:\n",
        "> * An `input_text` field containing the prompt.\n",
        "> * For example: `{\"input_text\": \"Create a description for Plantation Palms.\"}`\n",
@@ -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,22 @@
        "\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",
        "# Get the default cloud project id.\n",
        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
        "\n",
@@ -123,13 +131,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 +145,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 +187,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,20 +230,11 @@
        "# @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",
@@ -257,10 +258,7 @@
        "        \"--disable-log-stats\",\n",
        "    ]\n",
        "\n",
-        "    env_vars = {\n",
-        "        \"MODEL_ID\": model_id,\n",
-        "        \"DEPLOY_SOURCE\": \"notebook\",\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",
@@ -333,9 +331,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 +345,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 +380,13 @@
        "# @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. 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:20240607_0936_RC01\"\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",
@@ -393,55 +410,120 @@
        "\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",
+        "        accelerator_count = 4\n",
+        "        machine_type = \"g2-standard-48\"\n",
        "    elif accelerator_type == \"NVIDIA_TESLA_A100\":\n",
        "        accelerator_count = 1\n",
        "        machine_type = \"a2-highgpu-1g\"\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_TESLA_A100\":\n",
        "        accelerator_count = 4\n",
        "        machine_type = \"a2-highgpu-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_config_file_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 +531,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 +540,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 +574,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",
+        "model, endpoint = 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",
@@ -602,15 +672,15 @@
      "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",
+        "# @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 model.\n",
+        "# Delete models.\n",
        "model.delete()\n",
        "\n",
        "delete_bucket = False  # @param {type:\"boolean\"}\n",
@@ -304,7 +304,7 @@
        "id": "f0bd36bb-f988-479a-be40-18c54980ca1f"
      },
      "source": [
-        "You can manage your uploaded models in the [Model Registry](https://pantheon.corp.google.com/vertex-ai/models) and your endpoints in the [Endpoints](https://pantheon.corp.google.com/vertex-ai/endpoints)."
+        "You can manage your uploaded models in the [Model Registry](https://console.cloud.google.com/vertex-ai/models) and your endpoints in the [Endpoints](https://console.cloud.google.com/vertex-ai/endpoints)."
      ]
    },
    {
@@ -155,7 +155,7 @@
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
@@ -165,7 +165,7 @@
        "# The pre-built training docker images. They contain training scripts and models.\n",
        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20240318_0936_RC00\"\n",
        "# The pre-built serving docker images. They contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)\n",
        "\n",
@@ -168,7 +168,7 @@
        "# The pre-built training docker images. They contain training scripts and models.\n",
        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20240318_0936_RC00\"\n",
        "# The pre-built serving docker images. They contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)\n",
        "\n",
@@ -336,12 +336,12 @@
        "    # @markdown **[Optional]** If using own images, provide the GCS path to the images.\n",
        "    # @markdown Make sure you have permission to access the bucket.\n",
        "    source_image_gcs_dir = \"gs://\"  # @param {type:\"string\"}\n",
-        "    print(\n",
-        "        f\"Now downloading the images from the source gcs directory: {source_image_gcs_dir}\"\n",
-        "    )\n",
        "    assert source_image_gcs_dir.startswith(\n",
        "        \"gs://\"\n",
        "    ), \"source_image_gcs_dir must start with `gs://`.\"\n",
+        "    print(\n",
+        "        f\"Now downloading the images from the source gcs directory: {source_image_gcs_dir}\"\n",
+        "    )\n",
        "    download_gcs_dir_to_local(source_image_gcs_dir, local_dir)\n",
        "\n",
        "\n",
@@ -375,6 +375,8 @@
        "class_dir_fuse = class_dir.replace(\"gs://\", \"/gcs/\")\n",
        "output_dir_fuse = output_dir.replace(\"gs://\", \"/gcs/\")\n",
        "\n",
+        "\n",
+        "concept_prompt = \"A picture of a sks dog\"  # @param {type:\"string\"}\n",
        "model = job.run(\n",
        "    args=[\n",
        "        \"--task=text-to-image-dreambooth-lora-sdxl\",\n",
@@ -382,19 +384,19 @@
        "        f\"--instance_data_dir={instance_dir_fuse}\",\n",
        "        f\"--class_data_dir={class_dir_fuse}\",\n",
        "        f\"--output_dir={output_dir_fuse}\",\n",
+        "        f\"--instance_prompt={concept_prompt}\",\n",
        "        \"--pretrained_vae_model_name_or_path=madebyollin/sdxl-vae-fp16-fix\",\n",
        "        \"--mixed_precision=fp16\",\n",
-        "        \"--instance_prompt='a photo of sks dog'\",\n",
        "        \"--resolution=1024\",\n",
        "        \"--train_batch_size=1\",\n",
-        "        \"--gradient_accumulation_steps=1\",\n",
+        "        \"--gradient_accumulation_steps=4\",\n",
        "        \"--gradient_checkpointing\",\n",
-        "        \"--learning_rate=2e-6\",\n",
-        "        \"--lr_scheduler=constant\",\n",
+        "        \"--learning_rate=1e-4\",\n",
+        "        \"--lr_scheduler=linear\",\n",
        "        \"--lr_warmup_steps=0 \",\n",
        "        \"--use_8bit_adam\",\n",
-        "        f\"--max_train_steps={train_steps}\",\n",
-        "        f\"--checkpointing_steps={train_steps // 2}\",\n",
+        "        f\"--max_train_steps={int(train_steps)}\",\n",
+        "        f\"--checkpointing_steps={int(train_steps // 2)}\",\n",
        "        \"--seed=0\",\n",
        "    ],\n",
        "    replica_count=num_nodes,\n",
@@ -467,7 +469,8 @@
        "\n",
        "# @markdown Once deployment succeeds, you can generate images by sending text prompts to the endpoint.\n",
        "\n",
-        "# @markdown You can also batch send prompts by separating them with a comma.\n",
+        "# @markdown You can also batch send prompts by separating them with a comma. Be sure to include\n",
+        "# @markdown the `concept_prompt` of your instance in the prompt.\n",
        "# @markdown You may adjust the parameters below to achieve best image quality.\n",
        "\n",
        "comma_separated_prompt_list = \"A picture of a sks dog in a house, A picture of a sks dog catching a frisbee\"  # @param {type: \"string\"}\n",
@@ -157,7 +157,7 @@
        "    auth.authenticate_user(project_id=PROJECT_ID)\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "# Define common functions.\n",
@@ -154,8 +154,7 @@
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "TEXT_TO_IMAGE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
-        "IMAGE_TO_IMAGE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve:sd-optimized\"\n",
+        "DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
@@ -200,17 +199,10 @@
        "        \"TASK\": task,\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
-        "    serving_container_image_uri = None\n",
-        "    if task == \"text-to-image\":\n",
-        "        serving_container_image_uri = TEXT_TO_IMAGE_DOCKER_URI\n",
-        "    elif task == \"image-to-image\":\n",
-        "        serving_container_image_uri = IMAGE_TO_IMAGE_DOCKER_URI\n",
-        "    else:\n",
-        "        raise ValueError(f\"Unknown task: {task}\")\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
-        "        serving_container_image_uri=serving_container_image_uri,\n",
+        "        serving_container_image_uri=DOCKER_URI,\n",
        "        serving_container_ports=[7080],\n",
        "        serving_container_predict_route=\"/predictions/diffusers_serving\",\n",
        "        serving_container_health_route=\"/ping\",\n",
@@ -166,7 +166,7 @@
        "# The pre-built training docker images. They contain training scripts and models.\n",
        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20240318_0936_RC00\"\n",
        "# The pre-built serving docker images. They contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "# Define common functions.\n",
@@ -244,7 +244,7 @@
        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20240202_1425_RC00\"\n",
        "\n",
        "# The pre-built serving docker images. They contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\""
      ]
    },
    {
@@ -239,7 +239,7 @@
        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20231130_0936_RC00\"\n",
        "\n",
        "# The pre-built serving docker images. They contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\""
      ]
    },
    {
@@ -53,9 +53,9 @@
      "source": [
        "## Overview\n",
        "\n",
-        "This notebook demonstrates starting a playground based on [Gradio UI](https://www.gradio.app/), inspired by the famous [Stable Diffusion WebUI](https://github.com/AUTOMATIC1111/stable-diffusion-webui) project, which allows users to interact with the stable diffusion models more easily and intuitively. The playground now support `text-to-image`, `image-inpainting`, `controlnet-canny`, `instruct-pix2pix`, and `SD 4x upscaler` tasks.\n",
+        "This notebook demonstrates starting a playground based on [Gradio UI](https://www.gradio.app/), inspired by the popular [Stable Diffusion WebUI](https://github.com/AUTOMATIC1111/stable-diffusion-webui) project, which allows users to interact with the stable diffusion models more easily and intuitively. The playground now support `text-to-image`, `image-inpainting`, `controlnet-canny`, `instruct-pix2pix`, and `SD 4x upscaler` tasks.\n",
        "\n",
-        "This notebook also supports Dreambooth finetune the stable diffusion 1.5/2.1 models.\n",
+        "This notebook also supports Dreambooth finetune the stable diffusion 1.5/2.1/XL models.\n",
        "\n",
        "### Objective\n",
        "\n",
@@ -121,7 +121,7 @@
        "\n",
        "# Cloud Storage bucket for storing the experiment artifacts.\n",
        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
-        "# prefer using your own GCS bucket, please change the value yourself below.\n",
+        "# prefer using your own GCS bucket, change the value yourself below.\n",
        "now = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
        "BUCKET_URI = \"gs://\"  # @param {type: \"string\"}\n",
        "\n",
@@ -160,7 +160,7 @@
        "# The pre-built training docker images. They contain training scripts and models.\n",
        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20240318_0936_RC00\"\n",
        "# The pre-built serving docker images. They contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240429_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
      ]
@@ -177,7 +177,7 @@
        "# @title Start the playground\n",
        "\n",
        "# @markdown This is a playground for image generation similar to the popular [Stable Diffusion WebUI](https://github.com/AUTOMATIC1111/stable-diffusion-webui).\n",
-        "# @markdown After the cell runs, this playground is avaible in a separate browser tab if you click the public URL.\n",
+        "# @markdown After the cell runs, this playground is available in a separate browser tab if you click the public URL.\n",
        "# @markdown Something similar to [\"https://####.gradio.live\"](#) in the output of the cell.\n",
        "\n",
        "# @markdown For Image generations, five tasks `text-to-image`, `image-inpainting`, `controlnet-canny`, `instruct-pix2pix` and `SD 4x upscaler` are currently supported.\n",
@@ -194,7 +194,7 @@
        "# @markdown 1. Don't forget to undeploy the models after all the experiment to avoid continuous charges to the project.\n",
        "\n",
        "# @markdown **Note: we support the following models now:**\n",
-        "# @markdown Other models (with the same task) may work, but they are not tested please use with caution.\n",
+        "# @markdown Other models (with the same task) may work, but they are not tested so use with caution.\n",
        "# @markdown 1. Text-to-Image models\n",
        "# @markdown    > runwayml/stable-diffusion-v1-5 \\\n",
        "# @markdown    > stabilityai/stable-diffusion-2-1 \\\n",
@@ -424,7 +424,7 @@
        "\n",
        "    if model_name not in model_to_task_dict.keys():\n",
        "        print(model_name)\n",
-        "        raise gr.Error(\"Please select a valid model name for Endpoint creation.\")\n",
+        "        raise gr.Error(\"Select a valid model name for Endpoint creation.\")\n",
        "\n",
        "    return model_to_task_dict[model_name]\n",
        "\n",
@@ -432,7 +432,7 @@
        "def deploy_model(model_name: str) -> aiplatform.Endpoint:\n",
        "    \"\"\"Preprocess the model_name, before deploying the model to vertex.\"\"\"\n",
        "    if not model_name:\n",
-        "        raise gr.Error(\"Please select a valid model name for model list.\")\n",
+        "        raise gr.Error(\"Select a valid model name for model list.\")\n",
        "\n",
        "    model_id = model_name.split(\": \")[1]\n",
        "    task_name = get_task_name(model_id)\n",
@@ -440,7 +440,9 @@
        "    return deploy_model_vertex(model_id, task_name)\n",
        "\n",
        "\n",
-        "def deploy_model_vertex(model_id: str, task_name: str) -> aiplatform.Endpoint:\n",
+        "def deploy_model_vertex(\n",
+        "    model_id: str, task_name: str, lora_id: str = \"\"\n",
+        ") -> aiplatform.Endpoint:\n",
        "    \"\"\"\n",
        "    Creates a new Vertex prediction endpoint and deploys a model to it.\n",
        "    The `model_id` can either be a HF model resource id, or a GCS path where\n",
@@ -470,6 +472,11 @@
        "            **serving_env,\n",
        "            \"REFINER_MODEL_ID\": refiner_model_id,\n",
        "        }\n",
+        "    if lora_id:\n",
+        "        serving_env = {\n",
+        "            **serving_env,\n",
+        "            \"LORA_ID\": lora_id,\n",
+        "        }\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_id,\n",
@@ -589,7 +596,7 @@
        "    image_dimension=512,\n",
        ") -> list[Image.Image]:\n",
        "    if not endpoint_name:\n",
-        "        raise gr.Error(\"Please select (or deploy) a model first!\")\n",
+        "        raise gr.Error(\"Select (or deploy) a model first!\")\n",
        "\n",
        "    prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)\n",
        "    instances = [\n",
@@ -624,7 +631,7 @@
        "    image_dimension=512,\n",
        ") -> list[Image.Image]:\n",
        "    if not endpoint_name:\n",
-        "        raise gr.Error(\"Please select (or deploy) a model first!\")\n",
+        "        raise gr.Error(\"Select (or deploy) a model first!\")\n",
        "\n",
        "    prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)\n",
        "    default_dimension = 512\n",
@@ -677,7 +684,7 @@
        "    image_dimension=512,\n",
        ") -> list[Image.Image]:\n",
        "    if not endpoint_name:\n",
-        "        raise gr.Error(\"Please select (or deploy) a model first!\")\n",
+        "        raise gr.Error(\"Select (or deploy) a model first!\")\n",
        "\n",
        "    prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)\n",
        "    instances = [\n",
@@ -714,7 +721,7 @@
        ") -> list[Image.Image]:\n",
        "\n",
        "    if not endpoint_name:\n",
-        "        raise gr.Error(\"Please select (or deploy) a model first!\")\n",
+        "        raise gr.Error(\"Select (or deploy) a model first!\")\n",
        "\n",
        "    prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)\n",
        "    default_dimension = 256\n",
@@ -756,7 +763,7 @@
        "    canny_high_threshold=200,\n",
        ") -> list[Image.Image]:\n",
        "    if not endpoint_name:\n",
-        "        raise gr.Error(\"Please select (or deploy) a model first!\")\n",
+        "        raise gr.Error(\"Select (or deploy) a model first!\")\n",
        "\n",
        "    prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)\n",
        "    init_image = init_image.convert(\"RGB\").resize((image_dimension, image_dimension))\n",
@@ -1154,12 +1161,14 @@
        "\n",
        "\n",
        "def get_dreambooth_base_model_and_resolution(base_model_select: str) -> tuple[str, int]:\n",
-        "    if base_model_select == \"v2-1-768\":\n",
+        "    if base_model_select == \"sd-v2-1-768\":\n",
        "        return \"stabilityai/stable-diffusion-2-1\", 768\n",
-        "    elif base_model_select == \"v2-1-512\":\n",
+        "    elif base_model_select == \"sd-v2-1-512\":\n",
        "        return \"stabilityai/stable-diffusion-2-1\", 512\n",
-        "    elif base_model_select == \"v1-5-512\":\n",
+        "    elif base_model_select == \"sd-v1-5-512\":\n",
        "        return \"runwayml/stable-diffusion-v1-5\", 512\n",
+        "    elif base_model_select == \"sdxl-dreambooth-lora-1024\":\n",
+        "        return \"stabilityai/stable-diffusion-xl-base-1.0\", 1024\n",
        "    else:\n",
        "        raise gr.Error(f\"Invalid base model: {base_model_select}\")\n",
        "\n",
@@ -1169,21 +1178,22 @@
        "    def prepare_instance_images(\n",
        "        base_model_select: str,\n",
        "        concept_prompt: str,\n",
-        "        train_steps: int,\n",
        "        file_collection: list[gr.File],\n",
        "    ):\n",
        "        \"\"\"\n",
        "        Prepares the instance images for training from the provided dir of images.\n",
        "        \"\"\"\n",
-        "        if not concept_prompt:\n",
-        "            raise gr.Error(\"Please enter a concept prompt first!\")\n",
-        "\n",
-        "        resolution = 512\n",
-        "        if base_model_select == \"v2-1-768\":\n",
-        "            resolution = 768\n",
-        "\n",
        "        if not file_collection:\n",
-        "            raise gr.Error(\"Please provide a few valid instance images first!\")\n",
+        "            raise gr.Error(gr.Markdown(\"Provide a few valid instance images!\"))\n",
+        "\n",
+        "        if not concept_prompt:\n",
+        "            raise gr.Error(\"Provide a unique concept prompt!\")\n",
+        "\n",
+        "        resolution = 768\n",
+        "        if base_model_select == \"sd-v2-1-512\" or base_model_select == \"sd-v1-5-512\":\n",
+        "            resolution = 512\n",
+        "        elif base_model_select == \"sdxl-dreambooth-lora-1024\":\n",
+        "            resolution = 1024\n",
        "\n",
        "        local_tmp_dir = \"/tmp/instance_images\"\n",
        "        if os.path.exists(local_tmp_dir):\n",
@@ -1207,6 +1217,18 @@
        "        instant_image_dir = os.path.join(BUCKET_URI, f\"instance_images-{now}\")\n",
        "        upload_local_dir_to_gcs(local_tmp_dir, instant_image_dir)\n",
        "\n",
+        "    def deploy_finetuned_model(\n",
+        "        base_model_select: str, output_dir: str\n",
+        "    ) -> aiplatform.Endpoint:\n",
+        "        if base_model_select == \"sdxl-dreambooth-lora-1024\":\n",
+        "            return deploy_model_vertex(\n",
+        "                model_id=\"stabilityai/stable-diffusion-xl-base-1.0\",\n",
+        "                lora_id=output_dir,\n",
+        "                task_name=\"text-to-image-sdxl\",\n",
+        "            )\n",
+        "\n",
+        "        return deploy_model_vertex(model_id=output_dir, task_name=\"text-to-image\")\n",
+        "\n",
        "    def dreambooth_start_training(*inputs):\n",
        "        [\n",
        "            output_dir,\n",
@@ -1224,16 +1246,18 @@
        "            seed,\n",
        "        ] = inputs[:13]\n",
        "\n",
-        "        file_collection = list(chain.from_iterable(inputs[13:]))\n",
        "        if not concept_prompt:\n",
-        "            raise gr.Error(\"Please enter a concept prompt first!\")\n",
+        "            raise gr.Error(\"Enter a concept prompt first!\")\n",
        "\n",
-        "        if not file_collection:\n",
-        "            raise gr.Error(\"Please provide a local folder with (5-10) instance image!\")\n",
+        "        if not inputs or len(inputs) < 13:\n",
+        "            raise gr.Error(\"Provide a local folder with (5-10) instance image!\")\n",
        "\n",
        "        if not base_model_select:\n",
-        "            raise gr.Error(\"Please select a base model!\")\n",
+        "            raise gr.Error(\"Select a base model!\")\n",
        "\n",
+        "        gr.Info(f\"Dreambooth finetune with task {base_model_select} started!\")\n",
+        "\n",
+        "        file_collection = list(chain.from_iterable(inputs[13:]))\n",
        "        base_model_id, resolution = get_dreambooth_base_model_and_resolution(\n",
        "            base_model_select\n",
        "        )\n",
@@ -1252,7 +1276,6 @@
        "        prepare_instance_images(\n",
        "            base_model_select=base_model_select,\n",
        "            concept_prompt=concept_prompt,\n",
-        "            train_steps=train_steps,\n",
        "            file_collection=list(file_collection),\n",
        "        )\n",
        "\n",
@@ -1269,10 +1292,14 @@
        "            container_uri=TRAIN_DOCKER_URI,\n",
        "        )\n",
        "\n",
-        "        gr.Info(f\"Custom Training job {job_name} is being submitted. Please wait...\")\n",
+        "        gr.Info(f\"Custom Training job {job_name} is being submitted. Please wait.\")\n",
+        "\n",
+        "        task = \"text-to-image-dreambooth\"\n",
+        "        if base_model_select == \"sdxl-dreambooth-lora-1024\":\n",
+        "            task = \"text-to-image-dreambooth-lora-sdxl\"\n",
        "\n",
        "        args = [\n",
-        "            \"--task=text-to-image-dreambooth\",\n",
+        "            f\"--task={task}\",\n",
        "            f\"--pretrained_model_name_or_path={base_model_id}\",\n",
        "            f\"--instance_data_dir={image_dir_gcsfuse}\",\n",
        "            f\"--class_data_dir={image_dir_gcsfuse}\",\n",
@@ -1287,9 +1314,15 @@
        "            f\"--max_train_steps={int(train_steps)}\",\n",
        "            f\"--checkpointing_steps={int(checkpointing_steps)}\",\n",
        "            f\"--seed={int(seed)}\",\n",
+        "            \"--gradient_checkpointing\",\n",
        "            \"--mixed_precision=fp16\",\n",
        "            \"--use_8bit_adam\",\n",
        "        ]\n",
+        "\n",
+        "        if base_model_select == \"sdxl-dreambooth-lora-1024\":\n",
+        "            args.append(\n",
+        "                \"--pretrained_vae_model_name_or_path=madebyollin/sdxl-vae-fp16-fix\"\n",
+        "            )\n",
        "        if train_text_encoder:\n",
        "            args.append(\"--train_text_encoder\")\n",
        "\n",
@@ -1303,7 +1336,21 @@
        "        )\n",
        "\n",
        "        if deploy_after_tuning:\n",
-        "            deploy_model_vertex(model_id=output_dir, task_name=\"text-to-image\")\n",
+        "            deploy_finetuned_model(base_model_select, output_dir)\n",
+        "\n",
+        "    def update_default_lr(base_model_select: str):\n",
+        "        \"\"\"Updates the default learning rate based on the selected base model.\"\"\"\n",
+        "        lr = 2e-6\n",
+        "        if base_model_select == \"sdxl-dreambooth-lora-1024\":\n",
+        "            lr = 1e-4\n",
+        "\n",
+        "        return {\n",
+        "            learning_rate: gr.update(value=lr),\n",
+        "        }\n",
+        "\n",
+        "    def update_default_checkpoing_steps(train_steps: int):\n",
+        "        \"\"\"Updates the default checkpointing_steps based on the training steps.\"\"\"\n",
+        "        return {checkpointing_steps: gr.update(value=int(train_steps // 2))}\n",
        "\n",
        "    dreambooth_tip_text = r\"\"\"\n",
        "        1. Upload 5-10 images of the object/face on from different angles/perspectives.\n",
@@ -1321,8 +1368,13 @@
        "            with gr.Column():\n",
        "                base_model_select = gr.Dropdown(\n",
        "                    label=\"Select a base model\",\n",
-        "                    choices=[\"v2-1-768\", \"v2-1-512\", \"v1-5-512\"],\n",
-        "                    value=\"v2-1-768\",\n",
+        "                    choices=[\n",
+        "                        \"sd-v2-1-768\",\n",
+        "                        \"sd-v2-1-512\",\n",
+        "                        \"sd-v1-5-512\",\n",
+        "                        \"sdxl-dreambooth-lora-1024\",\n",
+        "                    ],\n",
+        "                    value=\"sd-v2-1-768\",\n",
        "                    interactive=True,\n",
        "                )\n",
        "                concept_prompt = gr.Textbox(\n",
@@ -1344,7 +1396,7 @@
        "                        label=\"Train text encoder\", value=False, interactive=True\n",
        "                    )\n",
        "                    checkpointing_steps = gr.Number(\n",
-        "                        label=\"Checkpointing steps\", value=100, interactive=True\n",
+        "                        label=\"Checkpointing steps\", value=400, interactive=True\n",
        "                    )\n",
        "                    train_batch_size = gr.Number(\n",
        "                        label=\"Training batch size\", value=1, interactive=True\n",
@@ -1393,6 +1445,22 @@
        "                    visible=True,\n",
        "                )\n",
        "\n",
+        "    base_model_select.change(\n",
+        "        update_default_lr,\n",
+        "        base_model_select,\n",
+        "        [\n",
+        "            learning_rate,\n",
+        "        ],\n",
+        "    )\n",
+        "\n",
+        "    train_steps.change(\n",
+        "        update_default_checkpoing_steps,\n",
+        "        train_steps,\n",
+        "        [\n",
+        "            checkpointing_steps,\n",
+        "        ],\n",
+        "    )\n",
+        "\n",
        "    train_button.click(\n",
        "        dreambooth_start_training,\n",
        "        inputs=[\n",
@@ -159,7 +159,7 @@
        "    auth.authenticate_user(project_id=PROJECT_ID)\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "# Define common functions.\n",
@@ -159,7 +159,7 @@
        "    auth.authenticate_user(project_id=PROJECT_ID)\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "# Expand an image to a square image.\n",
@@ -149,7 +149,7 @@
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
@@ -158,7 +158,7 @@
        "    auth.authenticate_user(project_id=PROJECT_ID)\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "def base64_to_image(image_str):\n",
@@ -158,7 +158,7 @@
        "\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240429_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "# Define common functions.\n",
@@ -158,7 +158,7 @@
        "    auth.authenticate_user(project_id=PROJECT_ID)\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "# Define common functions.\n",
@@ -158,7 +158,7 @@
        "\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "# Define common functions.\n",
@@ -247,7 +247,7 @@
      "source": [
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
        "\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240223_1230_RC00\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\""
      ]
    },
    {
@@ -158,7 +158,7 @@
        "    auth.authenticate_user(project_id=PROJECT_ID)\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "def deploy_model(model_id, task):\n",
@@ -156,7 +156,7 @@
        "    auth.authenticate_user(project_id=PROJECT_ID)\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
        "\n",
        "\n",
        "# Define common functions.\n",
@@ -671,7 +671,7 @@
        "id": "s2gnYFEJ2AML"
      },
      "source": [
-        "You can mange your uploaded models in the [Model Registry](https://pantheon.corp.google.com/vertex-ai/models) and your endpoints in the [Endpoints](https://pantheon.corp.google.com/vertex-ai/endpoints)."
+        "You can mange your uploaded models in the [Model Registry](https://console.cloud.google.com/vertex-ai/models) and your endpoints in the [Endpoints](https://console.cloud.google.com/vertex-ai/endpoints)."
      ]
    },
    {
@@ -0,0 +1,797 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "vbnQoeZABLr6"
+      },
+      "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": "dlUIoPa6BB5Z"
+      },
+      "source": [
+        "# Vertex AI Model Garden - TimesFM 1.0 (CPU/GPU 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_timesfm_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_timesfm_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>\n",
+        "\n",
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates deploying TimesFM 1.0 to a Vertex AI Endpoint and making online predictions for times series forecast.\n",
+        "\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Deploy TimesFM 1.0 to a Vertex AI Endpoint.\n",
+        "- Make predictions to the endpoint for times series forecast.\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": "o8KqOx2DC3Yc"
+      },
+      "source": [
+        "## Setup Google Cloud project"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "iNgqf6w1DCZF"
+      },
+      "outputs": [],
+      "source": [
+        "# @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",
+        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets)\n",
+        "# @markdown for storing experiment outputs. Set the BUCKET_URI for the\n",
+        "# @markdown experiment environment. The specified Cloud Storage bucket\n",
+        "# @markdown (`BUCKET_URI`) should be located in the same region as where the\n",
+        "# @markdown notebook was launched. Note that a multi-region bucket (eg. \"us\") is\n",
+        "# @markdown not considered a match for a single region covered by the\n",
+        "# @markdown multi-region range (eg. \"us-central1\"). If not set, a unique GCS\n",
+        "# @markdown bucket will be created instead.\n",
+        "\n",
+        "import json\n",
+        "import os\n",
+        "# Import the necessary packages\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",
+        "print(f\"Using this region: {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",
+        "# @markdown Cloud Storage bucket for storing the experiment artifacts.\n",
+        "# @markdown A unique GCS bucket will be created for the purpose of this\n",
+        "# @markdown notebook. If you prefer using your own GCS bucket, change the value\n",
+        "# @markdown 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",
+        "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",
+        "    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",
+        "            f\"Bucket region {bucket_region} is different from notebook region\"\n",
+        "            f\" {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, \"timesfm\")\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",
+        "# 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",
+        "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",
+        "\n",
+        "# @markdown ### **Choose a prebuilt checkpoint**\n",
+        "# @markdown Here we specify where to get the model checkpoint. TimesFM\n",
+        "# @markdown pretrained checkpoints are by default saved under\n",
+        "# @markdown `gs://vertex-model-garden-public-{region}/timesfm` and indexed by\n",
+        "# @markdown the checkpoint version.\n",
+        "\n",
+        "VERTEX_AI_MODEL_GARDEN_TIMESFM = \"gs://vertex-model-garden-public-us/timesfm\"  # @param {type:\"string\", isTemplate:true} [\"gs://vertex-model-garden-public-us/timesfm\", \"gs://vertex-model-garden-public-eu/timesfm\", \"gs://vertex-model-garden-public-asia/timesfm\"]\n",
+        "MODEL_VARIANT = \"timesfm-1.0-200m\"  # @param [\"timesfm-1.0-200m\"]\n",
+        "\n",
+        "print(\n",
+        "    \"Copying TimesFM model artifacts from\",\n",
+        "    f\"{VERTEX_AI_MODEL_GARDEN_TIMESFM}/{MODEL_VARIANT}\",\n",
+        "    \"to\",\n",
+        "    MODEL_BUCKET,\n",
+        ")\n",
+        "\n",
+        "! gsutil -m cp -r -R $VERTEX_AI_MODEL_GARDEN_TIMESFM/$MODEL_VARIANT $MODEL_BUCKET\n",
+        "\n",
+        "model_path_prefix = MODEL_BUCKET\n",
+        "\n",
+        "# The pre-built serving docker images.\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/jax-timesfm-serve:20240528_1310_RC00\"\n",
+        "\n",
+        "\n",
+        "def get_job_name_with_datetime(prefix: str) -> str:\n",
+        "    \"\"\"Gets the job name with date time when triggering training or deployment\n",
+        "\n",
+        "    jobs in Vertex AI.\n",
+        "    \"\"\"\n",
+        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
+        "\n",
+        "\n",
+        "def deploy_model(\n",
+        "    model_name: str,\n",
+        "    checkpoint_path: str,\n",
+        "    horizon: str,\n",
+        "    machine_type: str = \"g2-standard-4\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 1,\n",
+        "    deploy_source: str = \"notebook\",\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Create a Vertex AI Endpoint and deploy the specified model to the endpoint.\"\"\"\n",
+        "    model_name_with_time = get_job_name_with_datetime(model_name)\n",
+        "    endpoint = aiplatform.Endpoint.create(\n",
+        "        display_name=f\"{model_name_with_time}-endpoint\",\n",
+        "        credentials=aiplatform.initializer.global_config.credentials,\n",
+        "    )\n",
+        "\n",
+        "    if accelerator_type == \"ACCELERATOR_TYPE_UNSPECIFIED\":\n",
+        "        timesfm_backend = \"cpu\"\n",
+        "        accelerator_type = None\n",
+        "    elif accelerator_type.startswith(\"NVIDIA\"):\n",
+        "        timesfm_backend = \"gpu\"\n",
+        "    else:\n",
+        "        timesfm_backend = \"tpu\"\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name_with_time,\n",
+        "        artifact_uri=checkpoint_path,\n",
+        "        serving_container_image_uri=SERVE_DOCKER_URI,\n",
+        "        serving_container_ports=[8080],\n",
+        "        serving_container_predict_route=\"/predict\",\n",
+        "        serving_container_health_route=\"/health\",\n",
+        "        serving_container_environment_variables={\n",
+        "            \"DEPLOY_SOURCE\": deploy_source,\n",
+        "            \"TIMESFM_HORIZON\": str(horizon),\n",
+        "            \"TIMESFM_BACKEND\": timesfm_backend,\n",
+        "        },\n",
+        "        credentials=aiplatform.initializer.global_config.credentials,\n",
+        "    )\n",
+        "    print(\n",
+        "        f\"Deploying {model_name_with_time} on {machine_type} with\"\n",
+        "        f\" {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",
+        "        enable_access_logging=True,\n",
+        "        min_replica_count=1,\n",
+        "        sync=True,\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.\n",
+        "\n",
+        "    Returns -1 if can not figure out the quota.\n",
+        "    \"\"\"\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 (\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",
+        "\n",
+        "    Args:\n",
+        "      accelerator_type: The accelerator type.\n",
+        "      is_for_training: Whether the resource is used for training. Set false for\n",
+        "        serving use case.\n",
+        "\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",
+        "        \"ACCELERATOR_TYPE_UNSPECIFIED\": \"custom_model_training_cpus\",\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",
+        "        \"ACCELERATOR_TYPE_UNSPECIFIED\": \"custom_model_serving_cpus\",\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": "MLTYNURLi6LQ"
+      },
+      "source": [
+        "## Deploy TimesFM to a Vertex AI Endpoint"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "UwyXUWD9i9fe"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown This section uploads the prebuilt TimesFM model to Model Registry\n",
+        "# @markdown and deploys it to a Vertex AI Endpoint.\n",
+        "# @markdown It takes **approximately 20 minutes** to deploy.\n",
+        "\n",
+        "# @markdown ### **Step 1: Set the checkpoint path**\n",
+        "# @markdown Leave this blank to load the checkpoint we copied over earlier.\n",
+        "# @markdown If you've brought your own checkpoint, specify its path here.\n",
+        "# @markdown\n",
+        "# @markdown **Note**: Most of the time you should leave it blank (as is)\n",
+        "# @markdown when you've chosen to use a prebuilt checkpoint.\n",
+        "# @markdown\n",
+        "\n",
+        "custom_timesfm_model_uri = \"gs://\"  # @param {type: \"string\"}\n",
+        "\n",
+        "if custom_timesfm_model_uri == \"gs://\" or not custom_timesfm_model_uri:\n",
+        "    print(\"Deploying prebuilt TimesFM model. \")\n",
+        "    checkpoint_path = model_path_prefix\n",
+        "else:\n",
+        "    print(\"Deploying custom TimesFM model.\")\n",
+        "    checkpoint_path = custom_timesfm_model_uri\n",
+        "print(f\"Loading checkpoint from {checkpoint_path}.\")\n",
+        "\n",
+        "# @markdown ### **Step 2: Choose the accelerator**\n",
+        "# @markdown Select the accelerator type to use to deploy the model.\n",
+        "# @markdown\n",
+        "# @markdown **Note**: Most of the time you can go with CPU only. TimesFM is\n",
+        "# @markdown fast even with the CPU backend. You can only consider GPU if you\n",
+        "# @markdown need a dedicated endpoint to handle large queries per second.\n",
+        "# @markdown\n",
+        "# @markdown **Note**: After deployment, please take a look at the log to get\n",
+        "# @markdown the model / enpoint that you can use in another session.\n",
+        "# @markdown\n",
+        "\n",
+        "accelerator_type = \"NVIDIA_L4\"  # @param [\"CPU\", \"NVIDIA_L4\"]\n",
+        "if accelerator_type == \"NVIDIA_L4\":\n",
+        "    machine_type = \"g2-standard-4\"\n",
+        "    accelerator_count = 1\n",
+        "elif accelerator_type == \"CPU\":\n",
+        "    accelerator_type = \"ACCELERATOR_TYPE_UNSPECIFIED\"\n",
+        "    machine_type = \"n1-standard-8\"\n",
+        "    accelerator_count = 0\n",
+        "else:\n",
+        "    raise ValueError(\n",
+        "        f\"Recommended machine settings not found for: {accelerator_type}. To use\"\n",
+        "        \" another another accelerator, edit this code block to pass in an\"\n",
+        "        \" appropriate `machine_type`, `accelerator_type`, and\"\n",
+        "        \" `accelerator_count` to the deploy_model function by clicking `Show\"\n",
+        "        \" Code` and then modifying the code.\"\n",
+        "    )\n",
+        "\n",
+        "if accelerator_type != \"ACCELERATOR_TYPE_UNSPECIFIED\":\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",
+        "print(\"Quota is OK.\")\n",
+        "# @markdown If you want to use other accelerator types not listed above, please\n",
+        "# @markdown check other Vertex AI prediction supported accelerators and regions\n",
+        "# @markdown at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
+        "# @markdown You may need to manually set the `machine_type`, `accelerator_type`,\n",
+        "# @markdown and `accelerator_count` in the code by clicking `Show code` first.\n",
+        "\n",
+        "# @markdown ### **Step 3: Set the forecast horizon**\n",
+        "# @markdown We need to specify the forecast horizon TimesFM will be queried on\n",
+        "# @markdown to compile its computation. The endpoint will always predict this\n",
+        "# @markdown number of time points in the future, possibly after being rounded\n",
+        "# @markdown up to the closest multiplier of the model output patch length.\n",
+        "# @markdown Make sure to set it to the potential maximum for your usecase.\n",
+        "horizon = 256  # @param {type:\"number\"}\n",
+        "print(\"Creating endpoint.\")\n",
+        "model, endpoint = deploy_model(\n",
+        "    model_name=f\"timesfm-{MODEL_VARIANT}\",\n",
+        "    checkpoint_path=checkpoint_path,\n",
+        "    horizon=horizon,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "VzCdSZsilcEv"
+      },
+      "source": [
+        "## Query TimesFM\n",
+        "\n",
+        "An endpoint prediction request looks like\n",
+        "```python\n",
+        "endpoint.predict(instances=[{\"input\": [...], \"freq\": 0}, ...])\n",
+        "```\n",
+        "\n",
+        "Now we can query the endpoint to forecast on input time series. Let's first start with some sanity checks.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "dBIpE8eRarWL"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Create a helper function to visulize forecasts.\n",
+        "import matplotlib.pyplot as plt\n",
+        "import numpy as np\n",
+        "\n",
+        "\n",
+        "def visualize_forecast(\n",
+        "    context: list[float],\n",
+        "    horizon_mean: list[float],\n",
+        "    ground_truth: list[float] | None = None,\n",
+        "    horizon_lower: list[float] | None = None,\n",
+        "    horizon_upper: list[float] | None = None,\n",
+        "    ylabel: str | None = None,\n",
+        "    title: str | None = None,\n",
+        "):\n",
+        "    plt_range = list(range(len(context) + len(horizon_mean)))\n",
+        "    plt.figure(figsize=(5, 3))\n",
+        "    plt.plot(\n",
+        "        plt_range,\n",
+        "        context + [np.nan for _ in horizon_mean],\n",
+        "        color=\"tab:cyan\",\n",
+        "        label=\"context\",\n",
+        "    )\n",
+        "    plt.plot(\n",
+        "        plt_range,\n",
+        "        [np.nan for _ in context] + horizon_mean,\n",
+        "        color=\"tab:red\",\n",
+        "        label=\"forecast\",\n",
+        "    )\n",
+        "    if ground_truth:\n",
+        "        plt.plot(\n",
+        "            list(range(len(context) + len(ground_truth))),\n",
+        "            [np.nan for _ in context] + ground_truth,\n",
+        "            color=\"tab:purple\",\n",
+        "            label=\"ground truth\",\n",
+        "        )\n",
+        "    if horizon_upper and horizon_lower:\n",
+        "        plt.plot(\n",
+        "            plt_range,\n",
+        "            [np.nan for _ in context] + horizon_upper,\n",
+        "            color=\"tab:orange\",\n",
+        "            linestyle=\"--\",\n",
+        "            label=\"forecast, upper\",\n",
+        "        )\n",
+        "        plt.plot(\n",
+        "            plt_range,\n",
+        "            [np.nan for _ in context] + horizon_lower,\n",
+        "            color=\"tab:orange\",\n",
+        "            linestyle=\":\",\n",
+        "            label=\"forecast, lower\",\n",
+        "        )\n",
+        "        plt.fill_between(\n",
+        "            plt_range,\n",
+        "            [np.nan for _ in context] + horizon_upper,\n",
+        "            [np.nan for _ in context] + horizon_lower,\n",
+        "            color=\"tab:orange\",\n",
+        "            alpha=0.2,\n",
+        "        )\n",
+        "    if ylabel:\n",
+        "        plt.ylabel(ylabel)\n",
+        "    if title:\n",
+        "        plt.title(title)\n",
+        "    plt.xlabel(\"time\")\n",
+        "    plt.legend()\n",
+        "    plt.show()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2gf-ND9VwWb9"
+      },
+      "source": [
+        "### Sanity checks\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "CuUe-P0xgIuU"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown We first check TimesFM on some sinusoidals. Pay attention to how\n",
+        "# @markdown we are calling the endpoints.\n",
+        "# Prepare the context. Notice each of them has a different context length.\n",
+        "# Note: this is strictly how the query should be structed:\n",
+        "instances = [\n",
+        "    {\"input\": np.sin(np.linspace(0, 20, 100)).tolist(), \"freq\": 0},\n",
+        "    {\"input\": np.sin(np.linspace(0, 40, 500)).tolist(), \"freq\": 0},\n",
+        "    {\n",
+        "        \"input\": (\n",
+        "            np.sin(np.linspace(0, 50, 300)) + np.sin(np.linspace(1, 71, 300)) * 0.5\n",
+        "        ).tolist(),\n",
+        "        \"freq\": 0,\n",
+        "    },\n",
+        "]\n",
+        "\n",
+        "# Query the endpoint.\n",
+        "results = endpoint.predict(instances=instances)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "zJqLvZ0Qh-oe"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown Now we visualize the response. Make sure the model makes legit\n",
+        "# @markdown forecasts on those curves, and we move on to real world data.\n",
+        "\n",
+        "# There's bunch of important stuff in the results. Here we focus on results[0]:\n",
+        "# This is the TimesFM response.\n",
+        "print(results[0][0].keys())\n",
+        "visualize_forecast(\n",
+        "    instances[0][\"input\"], results[0][0][\"point_forecast\"], title=\"Sinusoidal 1\"\n",
+        ")\n",
+        "visualize_forecast(\n",
+        "    instances[1][\"input\"], results[0][1][\"point_forecast\"], title=\"Sinusoidal 2\"\n",
+        ")\n",
+        "visualize_forecast(\n",
+        "    instances[2][\"input\"], results[0][2][\"point_forecast\"], title=\"Sinusoidal 3\"\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "otoyBWG8vIfo"
+      },
+      "source": [
+        "### Point forecast\n",
+        "\n",
+        "Let's use a real world dataset from Kaggle on the [daily temperatures in Delhi, India](https://www.kaggle.com/datasets/sumanthvrao/daily-climate-time-series-data/data). Make sure you've set the Kaggle credentials following [these instructions](https://github.com/Kaggle/kaggle-api/blob/main/docs/README.md#api-credentials).\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NsezQiCfnCwH"
+      },
+      "outputs": [],
+      "source": [
+        "# Install the dependencies.\n",
+        "! pip install kaggle"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "_-zztflGnMpH"
+      },
+      "outputs": [],
+      "source": [
+        "# Download and prepare the dataset\n",
+        "! kaggle datasets download sumanthvrao/daily-climate-time-series-data\n",
+        "! unzip /content/daily-climate-time-series-data.zip"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "5QREpLhPnxIR"
+      },
+      "outputs": [],
+      "source": [
+        "import pandas as pd\n",
+        "\n",
+        "data = pd.read_csv(\"/content/DailyDelhiClimateTrain.csv\")\n",
+        "data"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "msYISUQtojZo"
+      },
+      "outputs": [],
+      "source": [
+        "# We manually prepare 3 forecast tasks:\n",
+        "# 1. Use day 0 - 199 to forecast day 200-299.\n",
+        "# 2. Use day 300 - 599 to forecast day 600-699.\n",
+        "# 3. Use day 700 - 1200 to forecast day 1200 - 1299.\n",
+        "temperature = data.meantemp.to_list()\n",
+        "inputs = [temperature[0:200], temperature[300:600], temperature[700:1200]]\n",
+        "ground_truths = [\n",
+        "    temperature[200:300],\n",
+        "    temperature[600:700],\n",
+        "    temperature[1200:1300],\n",
+        "]\n",
+        "response = endpoint.predict(\n",
+        "    instances=[{\"input\": each_input, \"freq\": 0} for each_input in inputs]\n",
+        ")\n",
+        "response[0][0].keys()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "gy55NO_OqIKD"
+      },
+      "source": [
+        "This `response` is structured that:\n",
+        "* `response[0][i]` is the forecast result of the ith input inside `instances`.\n",
+        "* `response[0][i]` has three keys:\n",
+        " - `point_forecast`: the mean point forecast\n",
+        " - `quantiles`: the schema of the quantile outputs\n",
+        " - `quantile_forecast`: for each time stamp in the horizon this will be a list whose elements are the corresponding quantiles as denoted in the `quantiles` schema.\n",
+        "   "
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "MWrOfE4A5vjo"
+      },
+      "outputs": [],
+      "source": [
+        "# Visualize the response\n",
+        "for task_i in range(3):\n",
+        "    visualize_forecast(\n",
+        "        inputs[task_i],\n",
+        "        response[0][task_i][\"point_forecast\"][:100],\n",
+        "        ground_truth=ground_truths[0],\n",
+        "        title=f\"Daily temperature in Delhi, India, Task {task_i+1}\",\n",
+        "        ylabel=\"Temperature (°C)\",\n",
+        "    )"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Tx8CvOjNvI0m"
+      },
+      "source": [
+        "### Anomaly detection\n",
+        "\n",
+        "As of checkpoint TimesFM-1.0-200m, TimesFM is capable of outputing quantile forecasts as well. These are uncalibrated forecasts and are experimental. But please feel free to play with them to see what you can do with them.\n",
+        "\n",
+        "Here we show how these outputs can potentially serve as anomaly detectors, when we define the anomaly as something beyond a certain range of TimesFM forecasts. In this example we are drawing bands defined by the 30th and the 70th percentiles on the same tasks we did in the last section. Anything outside of the bands could be an \"anomaly\"."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ab_NZ1hWiWwL"
+      },
+      "outputs": [],
+      "source": [
+        "# Visualize the response\n",
+        "for task_i in range(3):\n",
+        "    visualize_forecast(\n",
+        "        inputs[task_i],\n",
+        "        response[0][task_i][\"point_forecast\"][:100],\n",
+        "        ground_truth=ground_truths[0],\n",
+        "        horizon_lower=[x[3] for x in response[0][task_i][\"quantile_forecast\"]][:100],\n",
+        "        horizon_upper=[x[7] for x in response[0][task_i][\"quantile_forecast\"]][:100],\n",
+        "        title=f\"Daily temperature in Delhi, India, Task {task_i+1}\",\n",
+        "        ylabel=\"Temperature (°C)\",\n",
+        "    )"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0jcS6Tb0pg1x"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "5P_TAFcmlmdV"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Releasing endpoint and model\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",
+        "if model:\n",
+        "    model.delete()\n",
+        "\n",
+        "delete_bucket = False  # @param {type:\"boolean\"}\n",
+        "if delete_bucket:\n",
+        "    ! gsutil -m rm -r $BUCKET_URI"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "collapsed_sections": [
+        "2gf-ND9VwWb9",
+        "otoyBWG8vIfo",
+        "Tx8CvOjNvI0m",
+        "0jcS6Tb0pg1x"
+      ],
+      "name": "model_garden_timesfm_deployment_on_vertex.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
@@ -41,7 +41,7 @@
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fnotebook_template.ipynb\">\n",
-        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\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",
@@ -76,7 +76,7 @@
        "id": "61RBz8LLbxCR"
      },
      "source": [
-        "## Getting Started"
+        "## Get started"
      ]
    },
    {
@@ -85,7 +85,7 @@
        "id": "No17Cw5hgx12"
      },
      "source": [
-        "### Install Vertex AI SDK and other required packages\n"
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
@@ -135,7 +135,7 @@
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
-        "<b>⚠️ The kernel is going to restart. Please wait until it is finished before continuing to the next step. ⚠️</b>\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
@@ -173,7 +173,7 @@
        "id": "DF4l8DTdWgPY"
      },
      "source": [
-        "### Set Google Cloud project information and initialize Vertex AI SDK\n",
+        "### 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)."
      ]
@@ -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:
@@ -32,19 +32,24 @@
        "# Deploying Iris-detection model using FastAPI and Vertex AI custom container serving\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/SDK_Custom_Container_Prediction.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/SDK_Custom_Container_Prediction.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fcustom%2FSDK_Custom_Container_Prediction.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
-        "<a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/custom/SDK_Custom_Container_Prediction.ipynb\" target='_blank'>\n",
-        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
-        "      Open in Vertex AI Workbench\n",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/custom/SDK_Custom_Container_Prediction.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/SDK_Custom_Container_Prediction.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
@@ -58,7 +63,7 @@
      "source": [
        "## Overview\n",
        "\n",
-        "This tutorial walks you through building a custom container to serve a scikit-learn model on Vertex AI. You use the FastAPI Python web server framework to create a prediction and health endpoint. You also incorporate a pre-processor from training pipeline into your online serving application.\n",
+        "In this tutorial, you build a scikit-learn model and deploy it on Vertex AI using the custom container method. You use the FastAPI Python web server framework to create a prediction endpoint. You also incorporate a preprocessor from training pipeline into your online serving application.\n",
        "\n",
        "Learn more about [Custom training](https://cloud.google.com/vertex-ai/docs/training/custom-training) and [Vertex AI Prediction](https://cloud.google.com/vertex-ai/docs/predictions/get-predictions)."
      ]
@@ -74,10 +79,10 @@
        "In this notebook, you learn how to create, deploy and serve a custom classification model on Vertex AI. This notebook focuses more on deploying the model than on the design of the model itself. \n",
        "\n",
        "\n",
-        "This tutorial uses the following Google Cloud ML services and resources:\n",
+        "This tutorial uses the following Vertex AI services and resources:\n",
        "\n",
-        "- Vertex AI Models\n",
-        "- Vertex AI Endpoints\n",
+        "- Vertex AI models\n",
+        "- Vertex AI endpoints\n",
        "\n",
        "The steps performed include:\n",
        "\n",
@@ -96,10 +101,10 @@
      "source": [
        "### Dataset\n",
        "\n",
-        "This tutorial uses R.A. Fisher's Iris dataset, a small dataset that is a popular choice for trying out machine learning techniques. Each instance has four numerical features, which are different measurements of a flower, and a target label that\n",
-        "marks it as one of three types of iris: Iris setosa, Iris versicolour, or Iris virginica.\n",
+        "This tutorial uses R.A. Fisher's Iris dataset, a small and popular dataset for machine learning experiments. Each instance has four numerical features, which are different measurements of a flower, and a target label that\n",
+        "categorizes the flower into: **Iris setosa**, **Iris versicolour** and **Iris virginica**.\n",
        "\n",
-        "This tutorial uses [the copy of the Iris dataset included in the\n",
+        "This tutorial uses [a version of the Iris dataset available in the\n",
        "scikit-learn library](https://scikit-learn.org/stable/modules/generated/sklearn.datasets.load_iris.html#sklearn.datasets.load_iris)."
      ]
    },
@@ -128,12 +133,21 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "i7EUnXsZhAGF"
+        "id": "f0316df526f8"
      },
      "source": [
-        "## Install additional packages\n",
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9065e8d7f0fb"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n",
        "\n",
-        "Install additional package dependencies not installed in your notebook environment, such as NumPy, Scikit-learn, FastAPI, Uvicorn, and joblib. Use the latest major GA version of each package."
+        "Write the requirements needed for building container into a file."
      ]
    },
    {
@@ -151,6 +165,15 @@
        "google-cloud-storage>=1.26.0,<2.0.0dev"
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "479caccfb277"
+      },
+      "source": [
+        "Install the dependencies."
+      ]
+    },
    {
      "cell_type": "code",
      "execution_count": null,
@@ -175,7 +198,9 @@
        "id": "restart"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel"
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
@@ -186,11 +211,53 @@
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2d8e61720c2a"
+      },
+      "source": [
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "429a9f6d237d"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
+        "\n",
+        "Authenticate your environment on Google Colab.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "0c74b83ece6c"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
      ]
    },
    {
@@ -199,14 +266,8 @@
        "id": "yfEglUHQk9S3"
      },
      "source": [
-        "## Before you begin\n",
-        "\n",
-        "### Set your project ID\n",
-        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "* Run `gcloud config list`.\n",
-        "* Run `gcloud projects list`.\n",
-        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+        "### Set Google Cloud project information \n",
+        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
@@ -218,89 +279,7 @@
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# Set the project id\n",
-        "! gcloud config set project {PROJECT_ID}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "region"
-      },
-      "source": [
-        "#### Region\n",
-        "\n",
-        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "region"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "gcp_authenticate"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
-        "\n",
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated.\n",
-        "\n",
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ce6043da7b33"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0367eac06a10"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "21ad4dbb4a61"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c13224697bfb"
-      },
-      "source": [
-        "**4. Service account or other**\n",
-        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n"
      ]
    },
    {
@@ -331,7 +310,7 @@
        "id": "create_bucket"
      },
      "source": [
-        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+        "**If your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
@@ -342,7 +321,31 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3330b4f12a0d"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). "
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "e088ea8cd4a0"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -351,7 +354,7 @@
        "id": "d3938f6d37a1"
      },
      "source": [
-        "## Import libraries"
+        "### Import the required libraries"
      ]
    },
    {
@@ -363,30 +366,7 @@
      "outputs": [],
      "source": [
        "import os\n",
-        "\n",
-        "from google.cloud import aiplatform"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "750d53e37094"
-      },
-      "source": [
-        "### Initialize Vertex AI SDK for Python\n",
-        "\n",
-        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "1a3aa2d4a74f"
-      },
-      "outputs": [],
-      "source": [
-        "aiplatform.init(project=PROJECT_ID, location=REGION)"
+        "import sys"
      ]
    },
    {
@@ -397,7 +377,7 @@
      "source": [
        "### Configure resource names\n",
        "\n",
-        "Set a name for the following resources:\n",
+        "Set a name for the following parameters:\n",
        "\n",
        "`MODEL_ARTIFACT_DIR` - Folder directory path to your model artifacts within a Cloud Storage bucket, for example: \"my-models/fraud-detection/trial-4\"\n",
        "\n",
@@ -405,7 +385,7 @@
        "\n",
        "`IMAGE` - Name of the container image that is pushed to the repository.\n",
        "\n",
-        "`MODEL_DISPLAY_NAME` - Display name of Vertex AI Model resource."
+        "`MODEL_DISPLAY_NAME` - Display name of Vertex AI model resource."
      ]
    },
    {
@@ -442,9 +422,9 @@
      },
      "source": [
        "## Write your pre-processor\n",
-        "Standardize the training data so each numerical feature column has a mean of 0 and a standard deviation of 1 [can improve your model](https://developers.google.com/machine-learning/crash-course/representation/cleaning-data).\n",
+        "Standardize the training data so that each numerical feature column has a mean of 0 and a standard deviation of 1 [can improve your model](https://developers.google.com/machine-learning/crash-course/representation/cleaning-data).\n",
        "\n",
-        "Define a `app` folder and create `preprocess.py`, which contains a class to perform standardization."
+        "Define an `app` folder and create `preprocess.py`, which contains a class to perform standardization."
      ]
    },
    {
@@ -493,9 +473,9 @@
      },
      "source": [
        "## Train and store model with pre-processor\n",
-        "Next, use `preprocess.MySimpleScaler` to preprocess the iris data, then train a model using scikit-learn.\n",
+        "Use `preprocess.MySimpleScaler` to preprocess the iris data, and then train a model using scikit-learn.\n",
        "\n",
-        "At the end, export your trained model as a joblib (`.joblib`) file and export your `MySimpleScaler` instance as a pickle (`.pkl`) file."
+        "After training, export your trained model as a joblib (`.joblib`) file and export your `MySimpleScaler` instance as a pickle (`.pkl`) file."
      ]
    },
    {
@@ -633,7 +613,7 @@
      },
      "source": [
        "### Add pre-start script\n",
-        "FastAPI executes the following script before starting up the server. The `PORT` environment variable is set to equal to `AIP_HTTP_PORT` in order to run FastAPI on the same port expected by Vertex AI."
+        "FastAPI executes the following script before starting up the server. Set the environment variable `PORT` to `AIP_HTTP_PORT` for running the FastAPI server."
      ]
    },
    {
@@ -682,7 +662,7 @@
        "id": "51e149fdec1b"
      },
      "source": [
-        "## Build and push container to Artifact Registry"
+        "## Push the container image to Artifact Registry"
      ]
    },
    {
@@ -691,7 +671,7 @@
        "id": "240578ec9efe"
      },
      "source": [
-        "Write the `Dockerfile`, using `tiangolo/uvicorn-gunicorn-fastapi` as a base image. This automatically runs FastAPI for you using Gunicorn and Uvicorn. Visit [the FastAPI docs to read more about deploying FastAPI with Docker](https://fastapi.tiangolo.com/deployment/docker/) to learn more."
+        "Write the `Dockerfile`, using `tiangolo/uvicorn-gunicorn-fastapi` as a base image. This automatically runs FastAPI for you using Gunicorn and Uvicorn. Visit [the FastAPI docs about deploying with Docker](https://fastapi.tiangolo.com/deployment/docker/) to learn more."
      ]
    },
    {
@@ -712,17 +692,26 @@
        "RUN pip install -r requirements.txt"
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5cae10684d2c"
+      },
+      "source": [
+        "### Test the image locally (optional)"
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "04c988201499"
      },
      "source": [
-        "### Build the image locally (optional)\n",
+        "#### Build the image locally (optional)\n",
        "\n",
        "Build the image using docker to test it locally.\n",
        "\n",
-        "**Note:** Docker is only being used to test the container locally. For deployment to Artifact registry, Cloud-Build is used."
+        "**Note:** In this tutorial, docker is only being used to test the container locally. For deployment to Artifact Registry, Cloud-Build is used."
      ]
    },
    {
@@ -733,12 +722,11 @@
      },
      "outputs": [],
      "source": [
-        "import sys\n",
-        "\n",
        "IS_COLAB = \"google.colab\" in sys.modules\n",
+        "\n",
        "if not IS_COLAB and not os.getenv(\"IS_TESTING\"):\n",
        "    ! sudo docker build \\\n",
-        "            --tag=\"{REGION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE}\" \\\n",
+        "            --tag=\"{LOCATION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE}\" \\\n",
        "            ."
      ]
    },
@@ -748,7 +736,7 @@
        "id": "147a555f6c93"
      },
      "source": [
-        "### Run and test the container locally (optional)\n",
+        "#### Run and test the container locally (optional)\n",
        "\n",
        "Test running the container locally in detached mode and provide the environment variables that the container requires. These variables are provided to the container by Vertex AI once deployed. Test the `/health` and `/predict` routes and then stop the running image."
      ]
@@ -764,13 +752,13 @@
        "if not IS_COLAB and not os.getenv(\"IS_TESTING\"):\n",
        "    ! sudo docker stop local-iris\n",
        "    ! sudo docker rm local-iris\n",
-        "    ! sudo docker run -d -p 80:8080 \\\n",
+        "    ! docker run -d -p 80:8080 \\\n",
        "            --name=local-iris \\\n",
        "            -e AIP_HTTP_PORT=8080 \\\n",
        "            -e AIP_HEALTH_ROUTE=/health \\\n",
        "            -e AIP_PREDICT_ROUTE=/predict \\\n",
        "            -e AIP_STORAGE_URI={BUCKET_URI}/{MODEL_ARTIFACT_DIR} \\\n",
-        "            \"{REGION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE}\""
+        "            \"{LOCATION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE}\""
      ]
    },
    {
@@ -846,7 +834,7 @@
        "id": "93002a20a2a6"
      },
      "source": [
-        "### Build Custom Container\n"
+        "### Create a repository in Artifact Registry\n"
      ]
    },
    {
@@ -882,12 +870,12 @@
        "id": "23f3ba8346f4"
      },
      "source": [
-        "### Create a private Docker repository\n",
+        "#### Create a docker repository in Artifact Registry\n",
        "Your first step is to create your own Docker repository in Google Artifact Registry.\n",
        "\n",
-        "1 - Run the gcloud artifacts repositories create command to create a new Docker repository with your region with the description \"docker repository\".\n",
+        "1 - Run the `gcloud artifacts repositories create` command to create a new Docker repository with your region and the description as \"docker repository\".\n",
        "\n",
-        "2 - Run the gcloud artifacts repositories list command to verify that your repository was created."
+        "2 - Run the `gcloud artifacts repositories list` command to verify that your repository was created."
      ]
    },
    {
@@ -900,7 +888,7 @@
      "source": [
        "REPOSITORY = \"my-docker-repo-unique\"\n",
        "\n",
-        "! gcloud artifacts repositories create {REPOSITORY} --repository-format=docker --location={REGION} --description=\"Docker repository\"\n",
+        "! gcloud artifacts repositories create {REPOSITORY} --repository-format=docker --location={LOCATION} --description=\"Docker repository\"\n",
        "\n",
        "! gcloud artifacts repositories list"
      ]
@@ -911,9 +899,10 @@
        "id": "d2e0b8b700aa"
      },
      "source": [
-        "Push the image to the created artifact repository using Cloud-Build.\n",
+        "### Submit the image\n",
+        "Push the image to the created artifact repository using Cloud build.\n",
        "\n",
-        "**Note:** The following command automatically considers the Dockerfile from the directory it is being run from."
+        "**Note:** The following command automatically considers the Dockerfile from the directory it's being run from."
      ]
    },
    {
@@ -924,7 +913,7 @@
      },
      "outputs": [],
      "source": [
-        "!gcloud builds submit --region={REGION} --tag={REGION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE}"
+        "!gcloud builds submit --region={LOCATION} --tag={LOCATION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE}"
      ]
    },
    {
@@ -943,7 +932,7 @@
      },
      "source": [
        "### Create Vertex AI model using artifact uri\n",
-        "Use the Python SDK to upload and deploy your model from the artifact registry."
+        "Use the Python SDK to upload your model artifact in Vertex AI."
      ]
    },
    {
@@ -957,7 +946,7 @@
        "model = aiplatform.Model.upload(\n",
        "    display_name=MODEL_DISPLAY_NAME,\n",
        "    artifact_uri=f\"{BUCKET_URI}/{MODEL_ARTIFACT_DIR}\",\n",
-        "    serving_container_image_uri=f\"{REGION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE}\",\n",
+        "    serving_container_image_uri=f\"{LOCATION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE}\",\n",
        ")"
      ]
    },
@@ -967,9 +956,9 @@
        "id": "bd1b85afc7df"
      },
      "source": [
-        "### Deploy the model to Vertex AI Endpoints\n",
+        "### Deploy the model to Vertex AI Endpoint\n",
        "\n",
-        "Deploy the model to a Vertex AI Endpoint. After this step completes, the model is deployed and ready for online predictions."
+        "Once the deployment process gets done, the model is deployed and ready for serving online predictions."
      ]
    },
    {
@@ -991,7 +980,7 @@
      "source": [
        "## Request predictions\n",
        "\n",
-        "Send online requests to the model deployed to the endpoint and get predictions.\n",
+        "Send online requests to the endpoint and get predictions.\n",
        "\n",
        "### Using Python SDK\n",
        "\n",
@@ -1006,6 +995,7 @@
      },
      "outputs": [],
      "source": [
+        "# Send some sample data to the endpoint\n",
        "endpoint.predict(instances=[[6.7, 3.1, 4.7, 1.5], [4.6, 3.1, 1.5, 0.2]])"
      ]
    },
@@ -1028,6 +1018,7 @@
      },
      "outputs": [],
      "source": [
+        "# Fetch the endpoint name\n",
        "ENDPOINT_ID = endpoint.name"
      ]
    },
@@ -1039,11 +1030,12 @@
      },
      "outputs": [],
      "source": [
+        "# Send a prediction request using sample data \n",
        "! curl \\\n",
        "-H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
        "-H \"Content-Type: application/json\" \\\n",
        "-d @instances.json \\\n",
-        "https://{REGION}-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/{REGION}/endpoints/{ENDPOINT_ID}:predict"
+        "https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}:predict"
      ]
    },
    {
@@ -1066,7 +1058,7 @@
      "outputs": [],
      "source": [
        "!gcloud ai endpoints predict $ENDPOINT_ID \\\n",
-        "  --region=$REGION \\\n",
+        "  --region=$LOCATION \\\n",
        "  --json-request=instances.json"
      ]
    },
@@ -1112,16 +1104,39 @@
        "!gcloud artifacts docker images delete \\\n",
        "    --quiet \\\n",
        "    --delete-tags \\\n",
-        "    {REGION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE}\n",
+        "    {LOCATION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE}\n",
        "\n",
-        "# Delete the artifact registry\n",
-        "if delete_art_repo or os.getenv(\"IS_TESTING\"):\n",
-        "    ! gcloud artifacts repositories delete {REPOSITORY} --location=$REGION -q\n",
+        "# Delete the Artifact Repository\n",
+        "if delete_art_repo:\n",
+        "    ! gcloud artifacts repositories delete {REPOSITORY} --location=$LOCATION -q\n",
        "    \n",
        "# Delete the Cloud Storage bucket\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "31885b85cc2c"
+      },
+      "source": [
+        "Clean up the locally created files."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "234ea94b87c5"
+      },
+      "outputs": [],
+      "source": [
+        "! rm -rf app/\n",
+        "! rm requirements.txt\n",
+        "! rm instances.json\n",
+        "! rm Dockerfile"
+      ]
    }
  ],
  "metadata": {
@@ -4,6 +4,7 @@
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
+        "cellView": "form",
        "id": "copyright"
      },
      "outputs": [],
@@ -32,23 +33,26 @@
        "# Training a TensorFlow model on BigQuery data\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/custom-tabular-bq-managed-dataset.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/custom-tabular-bq-managed-dataset.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fcustom%2Fcustom-tabular-bq-managed-dataset.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/custom/custom-tabular-bq-managed-dataset.ipynb\">\n",
-        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
-        "      Open in Vertex AI Workbench\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
-        "  </td>  \n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/custom-tabular-bq-managed-dataset.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
        "</table>"
      ]
    },
@@ -128,12 +132,19 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "install_aip"
+        "id": "3b1ffd5ab768"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the latest version of Cloud Storage, Bigquery and Vertex AI SDKs for Python."
+        "## Get Started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dc848186ab0e"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages"
      ]
    },
    {
@@ -145,7 +156,7 @@
      "outputs": [],
      "source": [
        "# Install the packages\n",
-        "! pip3 install --upgrade google-cloud-aiplatform \\\n",
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
        "                        google-cloud-storage \\\n",
        "                        'google-cloud-bigquery[pandas]'"
      ]
@@ -153,48 +164,80 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "blGlVGFYW9Pt"
+        "id": "ff555b32bab8"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel."
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "0JrvuK6LUYnQ"
+        "id": "f09b4dff629a"
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "a47846030fef"
+        "id": "54c5ef8a8f43"
      },
      "source": [
-        "## Before you begin"
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "project_id"
+        "id": "f82e28c631cc"
      },
      "source": [
-        "#### Set your project ID\n",
+        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "* Run `gcloud config list`.\n",
-        "* Run `gcloud projects list`.\n",
-        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+        "Authenticate your environment on Google Colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "46604f70e831"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "107c51893a64"
+      },
+      "source": [
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
@@ -206,103 +249,7 @@
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# Set the project id\n",
-        "! gcloud config set project {PROJECT_ID}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "a54f9d7c1876"
-      },
-      "source": [
-        "#### Region\n",
-        "\n",
-        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "3aaadaaf9b30"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "5c0404984792"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "BaFKzJ_xXpvm"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "_fV-KyGAX4Xl"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**\n"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "7uXB1HAPX6L_"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "Ab_TRMQIYCCX"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "vx25htmYYExI"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "uZdA0-jBYGqt"
-      },
-      "source": [
-        "**4. Service account or other**\n",
-        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -344,7 +291,38 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
+        "! gsutil mb -l $LOCATION -p $PROJECT_ID $BUCKET_URI"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4f1319830b6e"
+      },
+      "source": [
+        "### Import libraries"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "poijnGfZCFYi"
+      },
+      "source": [
+        "Import the Vertex AI Python SDK and other required Python libraries."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "b9a0a5a74fa6"
+      },
+      "outputs": [],
+      "source": [
+        "import numpy as np\n",
+        "import pandas as pd\n",
+        "from google.cloud import aiplatform, bigquery"
      ]
    },
    {
@@ -366,10 +344,8 @@
      },
      "outputs": [],
      "source": [
-        "from google.cloud import aiplatform\n",
-        "\n",
-        "# Initialize the Vertex AI SDK\n",
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
+        "# Initialize the Vertex AI SDK for Python\n",
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -393,8 +369,6 @@
      },
      "outputs": [],
      "source": [
-        "from google.cloud import bigquery\n",
-        "\n",
        "# Set up BigQuery client\n",
        "bq_client = bigquery.Client(project=PROJECT_ID)"
      ]
@@ -405,7 +379,7 @@
        "id": "0a2c41bc91a6"
      },
      "source": [
-        "## Preprocess data and split data\n",
+        "### Preprocess data and split data\n",
        "First you should download and preprocess your data for training and testing.\n",
        "\n",
        "- Convert categorical features to numeric\n",
@@ -422,9 +396,6 @@
      },
      "outputs": [],
      "source": [
-        "import numpy as np\n",
-        "import pandas as pd\n",
-        "\n",
        "LABEL_COLUMN = \"species\"\n",
        "\n",
        "# Define the BigQuery source dataset\n",
@@ -456,7 +427,7 @@
        "id": "a39df4692a70"
      },
      "source": [
-        "## Create a Vertex AI Tabular Dataset from BigQuery dataset\n",
+        "### Create a Vertex AI Tabular dataset from BigQuery dataset\n",
        "\n",
        "Create a Vertex AI tabular dataset resource from your BigQuery training data.\n",
        "\n",
@@ -498,7 +469,7 @@
        "id": "train_custom_model"
      },
      "source": [
-        "## Train a model\n",
+        "### Train a model\n",
        "\n",
        "There are two ways you can train a model using a container image:\n",
        "\n",
@@ -811,7 +782,7 @@
        "id": "make_prediction"
      },
      "source": [
-        "## Make an online prediction request\n",
+        "### Make an online prediction request\n",
        "\n",
        "Send an online prediction request to your deployed model."
      ]
@@ -884,7 +855,7 @@
        "id": "undeploy_model"
      },
      "source": [
-        "## Undeploy models\n",
+        "### Undeploy models\n",
        "\n",
        "To undeploy all `Model` resources from the serving `Endpoint` resource, use the endpoint's `undeploy_all` method."
      ]
@@ -908,7 +879,8 @@
      "source": [
        "## Cleaning up\n",
        "\n",
-        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
+        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
+        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
        "\n",
        "Otherwise, you can delete the individual resources you created in this tutorial:\n",
        "\n",
@@ -926,8 +898,6 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
        "# Delete the training job\n",
        "job.delete()\n",
        "\n",
@@ -938,9 +908,9 @@
        "endpoint.delete()\n",
        "\n",
        "# Warning: Setting this to true deletes everything in your bucket\n",
-        "delete_bucket = False\n",
+        "delete_bucket = True\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil rm -r $BUCKET_URI"
      ]
    }
@@ -24,7 +24,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "l2mMvIUG9meX"
@@ -33,29 +32,30 @@
        "# Profile model training performance using Cloud Profiler\n",
        "\n",
        "<table align=\"left\">\n",
-        "\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/custom_training_tensorboard_profiler.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/custom_training_tensorboard_profiler.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fcustom%2Fcustom_training_tensorboard_profiler.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/custom/custom_training_tensorboard_profiler.ipynb\">\n",
-        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
-        "      Open in Vertex AI Workbench\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
-        "  </td>                                                                                               \n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/custom_training_tensorboard_profiler.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
        "</table>"
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
@@ -69,7 +69,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "dmfmQL6w84pS"
@@ -93,7 +92,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "zfXf0r-K81Y-"
@@ -105,7 +103,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "I3KFLvpq87rs"
@@ -126,22 +123,28 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
-        "id": "ze4-nDLfK4pw"
+        "id": "61RBz8LLbxCR"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook. "
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "No17Cw5hgx12"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "2b4ef9b72d43"
+        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
@@ -149,178 +152,169 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
-        "id": "aUw6ibN-n5Za"
+        "id": "R5Xep4W9lq-Z"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel."
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "FM12wbWhn7w0"
+        "id": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
-        "id": "LgFWLeJfoGQu"
+        "id": "SbmM4z7FOBpM"
      },
      "source": [
-        "## Before you begin\n",
-        "\n",
-        "### Set up your Google Cloud project\n",
-        "\n",
-        "**The following steps are required, regardless of your notebook environment.**\n",
-        "\n",
-        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
-        "\n",
-        "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
-        "\n",
-        "3. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). \n",
-        "\n",
-        "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
-        "id": "8ckyxpX_oSzD"
+        "id": "dmWOrTJ3gx13"
      },
      "source": [
-        "#### Set your project ID\n",
+        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "* Run `gcloud config list`.\n",
-        "* Run `gcloud projects list`.\n",
-        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "zY8DKBoVoVy3"
+        "id": "NyKGtVQjgx13"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "DF4l8DTdWgPY"
+      },
+      "source": [
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# Set the project id\n",
-        "! gcloud config set project {PROJECT_ID}"
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
-        "id": "mSQjVQmMosMl"
+        "id": "OKtKGmr9pfr6"
      },
      "source": [
-        "#### Region\n",
+        "#### Create a Cloud Storage bucket\n",
        "\n",
-        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
+        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "Se9FWWhLotvB"
+        "id": "In3aQanwYjFB"
      },
      "outputs": [],
      "source": [
-        "REGION = \"us-central1\"  # @param {type:\"string\"}"
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
-        "id": "IfJRIMBpo5Pg"
+        "id": "GOaOsIjxp0oB"
      },
      "source": [
-        "### Authenticate your Google Cloud account\n",
+        "**If your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Wn5QiIl2p16e"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil mb -l $LOCATION -p $PROJECT_ID $BUCKET_URI"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "OMrAJ8RGqBQu"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
-      ]
-    },
-    {
-      "attachments": {},
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "acFN0s3So9-Y"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "attachments": {},
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "dQ_mNwuapE5T"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**"
+        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "cR_MzpknpGgM"
+        "id": "AWRzBFExqERG"
      },
      "outputs": [],
      "source": [
-        "# ! gcloud auth login"
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
-      "attachments": {},
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "h-MuVI_ypJfw"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "BeaQlCwMpQUT"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "3ivZkPUjpaFz"
      },
      "source": [
-        "**4. Setup service account and permissions**\n",
+        "### Setup service account and permissions\n",
        "\n",
        "A service account will be used to create custom training jobs. If you do not want to use your project's Compute Engine service account, set SERVICE_ACCOUNT to another service account ID. You can create a service account by following the [instructions](https://cloud.google.com/iam/docs/creating-managing-service-accounts#creating)."
      ]
@@ -340,122 +334,56 @@
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "WWIxsCJFCg5Z"
+        "id": "autoset_service_account"
      },
      "outputs": [],
      "source": [
-        "# Grant Cloud Storage permission.\n",
-        "! gcloud projects add-iam-policy-binding $PROJECT_ID \\\n",
-        "        --member=\"serviceAccount:$SERVICE_ACCOUNT\" \\\n",
-        "        --role=\"roles/storage.admin\" \\\n",
-        "        --quiet"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "S_8_5jm-Gk6w"
-      },
-      "outputs": [],
-      "source": [
-        "# Grant AI Platform permission.\n",
-        "! gcloud projects add-iam-policy-binding $PROJECT_ID \\\n",
-        "        --member=\"serviceAccount:$SERVICE_ACCOUNT\" \\\n",
-        "        --role=\"roles/aiplatform.user\" \\\n",
-        "        --quiet"
-      ]
-    },
-    {
-      "attachments": {},
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "OKtKGmr9pfr6"
-      },
-      "source": [
-        "### Create a Cloud Storage bucket\n",
+        "import sys\n",
        "\n",
-        "Create a storage bucket to store intermediate artifacts such as datasets."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "In3aQanwYjFB"
-      },
-      "outputs": [],
-      "source": [
-        "BUCKET_URI = \"gs://your-bucket-name-unique\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "attachments": {},
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "GOaOsIjxp0oB"
-      },
-      "source": [
-        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "Wn5QiIl2p16e"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
-      ]
-    },
-    {
-      "attachments": {},
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ankcS-vtp7Wv"
-      },
-      "source": [
-        "### Import libraries"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "WffSImMvp-Po"
-      },
-      "outputs": [],
-      "source": [
-        "from google.cloud import aiplatform"
-      ]
-    },
-    {
-      "attachments": {},
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "OMrAJ8RGqBQu"
-      },
-      "source": [
-        "### Initialize Vertex AI SDK for Python\n",
+        "IS_COLAB = \"google.colab\" in sys.modules\n",
+        "if (\n",
+        "    SERVICE_ACCOUNT == \"\"\n",
+        "    or SERVICE_ACCOUNT is None\n",
+        "    or SERVICE_ACCOUNT == \"[your-service-account]\"\n",
+        "):\n",
+        "    # Get your service account from gcloud\n",
+        "    if not IS_COLAB:\n",
+        "        shell_output = !gcloud auth list 2>/dev/null\n",
+        "        SERVICE_ACCOUNT = shell_output[2].replace(\"*\", \"\").strip()\n",
        "\n",
-        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
+        "    else:  # IS_COLAB:\n",
+        "        shell_output = ! gcloud projects describe  $PROJECT_ID\n",
+        "        project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "        SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "\n",
+        "    print(\"Service Account:\", SERVICE_ACCOUNT)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "set_service_account:pipelines"
+      },
+      "source": [
+        "#### Set service account access for Vertex AI Pipelines\n",
+        "\n",
+        "Run the following commands to grant your service account access to read and write pipeline artifacts in the bucket that you created in the previous step. You only need to run this step once per service account."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "AWRzBFExqERG"
+        "id": "bad678f5afda"
      },
      "outputs": [],
      "source": [
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectCreator $BUCKET_URI\n",
+        "\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectViewer $BUCKET_URI"
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "MYzX9ZgjsyyY"
@@ -468,6 +396,27 @@
        "Learn more about [Enabling service](https://cloud.google.com/artifact-registry/docs/enable-service)."
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4e302d6c9a11"
+      },
+      "source": [
+        "**Set project in colab environment (Colab only)**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7b9f07d9d1f4"
+      },
+      "outputs": [],
+      "source": [
+        "if IS_COLAB:\n",
+        "    ! gcloud config set project $PROJECT_ID"
+      ]
+    },
    {
      "cell_type": "code",
      "execution_count": null,
@@ -476,11 +425,11 @@
      },
      "outputs": [],
      "source": [
+        "# Enable Artifact Registry API\n",
        "! gcloud services enable artifactregistry.googleapis.com --quiet"
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "-ayTbNdi62_t"
@@ -492,7 +441,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "9c3QrDTZdaxk"
@@ -513,7 +461,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "vJrWKK0mY7H7"
@@ -533,7 +480,7 @@
      "outputs": [],
      "source": [
        "tensorboard = aiplatform.Tensorboard.create(\n",
-        "    display_name=TENSORBOARD_NAME, project=PROJECT_ID, location=REGION\n",
+        "    display_name=TENSORBOARD_NAME, project=PROJECT_ID, location=LOCATION\n",
        ")\n",
        "\n",
        "TENSORBOARD_INSTANCE_NAME = tensorboard.resource_name\n",
@@ -542,7 +489,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "yoR29gW2S24w"
@@ -560,7 +506,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "OrpUIkAIs_uQ"
@@ -591,7 +536,7 @@
        "\n",
        "! gcloud artifacts repositories create {DOCKER_REPOSITORY} \\\n",
        "    --repository-format=docker \\\n",
-        "    --location={REGION} \\\n",
+        "    --location={LOCATION} \\\n",
        "    --description=\"Repository for TensorBoard Custom Training Job\" \\\n",
        "    --quiet\n",
        "\n",
@@ -599,7 +544,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "MPqbvhO6tK_e"
@@ -607,7 +551,7 @@
      "source": [
        "### Configure authentication to your private Docker repository\n",
        "\n",
-        "Before you push or pull container images, configure Docker to use the `gcloud` command-line tool to authenticate requests to `Artifact Registry` for your region."
+        "Before you push or pull container images, configure Docker to use the `gcloud` command-line tool to authenticate requests to Artifact Registry for your location."
      ]
    },
    {
@@ -623,11 +567,10 @@
        "IS_COLAB = \"google.colab\" in sys.modules\n",
        "\n",
        "if not IS_COLAB:\n",
-        "    ! gcloud auth configure-docker {REGION}-docker.pkg.dev --quiet"
+        "    ! gcloud auth configure-docker {LOCATION}-docker.pkg.dev --quiet"
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "pSXrzCWi-bTD"
@@ -654,7 +597,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "hyAwgsoQmaYI"
@@ -784,7 +726,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "LiOEvgkA_90f"
@@ -819,7 +760,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "ihYFahRAr6sj"
@@ -837,13 +777,12 @@
      "outputs": [],
      "source": [
        "IMAGE_NAME = \"tensorboard-custom-container\"\n",
-        "IMAGE_URI = f\"{REGION}-docker.pkg.dev/{PROJECT_ID}/{DOCKER_REPOSITORY}/{IMAGE_NAME}\"\n",
+        "IMAGE_URI = f\"{LOCATION}-docker.pkg.dev/{PROJECT_ID}/{DOCKER_REPOSITORY}/{IMAGE_NAME}\"\n",
        "\n",
-        "! gcloud builds submit --project {PROJECT_ID} --region={REGION} --tag {IMAGE_URI} --timeout=60m --quiet"
+        "! gcloud builds submit --project {PROJECT_ID} --region={LOCATION} --tag {IMAGE_URI} --timeout=60m --quiet"
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "k4e6OYmimqTR"
@@ -870,7 +809,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "51hKGTbU32Eg"
@@ -915,7 +853,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "JkEe2Nb_85UD"
@@ -929,7 +866,6 @@
      ]
    },
    {
-      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "TpV-iwP9qw9c"
@@ -958,9 +894,10 @@
      "source": [
        "delete_tensorboard = True\n",
        "delete_bucket = False\n",
+        "delete_generated_files_after_execution = False\n",
        "\n",
        "# Delete docker repository.\n",
-        "! gcloud artifacts repositories delete $DOCKER_REPOSITORY --project {PROJECT_ID} --location {REGION} --quiet\n",
+        "! gcloud artifacts repositories delete $DOCKER_REPOSITORY --project {PROJECT_ID} --location {LOCATION} --quiet\n",
        "\n",
        "job.delete()\n",
        "\n",
@@ -968,7 +905,10 @@
        "    tensorboard.delete()\n",
        "\n",
        "if delete_bucket and \"BUCKET_URI\" in globals():\n",
-        "    ! gsutil -m rm -r $BUCKET_URI"
+        "    ! gsutil -m rm -r $BUCKET_URI\n",
+        "\n",
+        "if delete_generated_files_after_execution:\n",
+        "    ! rm -rf $PYTHON_PACKAGE_APPLICATION_DIR Dockerfile Dockerfile"
      ]
    }
  ],
@@ -32,23 +32,29 @@
        "# Get started with Vertex AI Training for XGBoost\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/get_started_vertex_training_xgboost.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"><br> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/get_started_vertex_training_xgboost.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fcustom%2Fget_started_vertex_training_xgboost.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/custom/get_started_vertex_training_xgboost.ipynb\">\n",
-        "        <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "        <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br>\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/get_started_vertex_training_xgboost.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br>\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  \n",
        "</table>\n",
        "<br/><br/><br/>"
      ]
@@ -74,19 +80,19 @@
      "source": [
        "### Objective\n",
        "\n",
-        "In this tutorial, you learn how to use `Vertex AI Training` for training a XGBoost custom model.\n",
+        "In this tutorial, you learn how to use Vertex AI Training for training a XGBoost custom model.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
-        "- `Vertex AI Training`\n",
-        "- `Vertex AI Model` resource\n",
+        "- Vertex AI Training\n",
+        "- Vertex AI model resource\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Training using a Python package.\n",
        "- Report accuracy when hyperparameter tuning.\n",
        "- Save the model artifacts to Cloud Storage using GCSFuse.\n",
-        "- Create a `Vertex AI Model` resource."
+        "- Create a Vertex AI model resource."
      ]
    },
    {
@@ -127,9 +133,16 @@
        "id": "install_mlops"
      },
      "source": [
-        "## Installations\n",
-        "\n",
-        "Install the following packages to execute this notebook."
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "a2c2cb2109a0"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
@@ -140,47 +153,86 @@
      },
      "outputs": [],
      "source": [
-        "! pip3 install --upgrade google-cloud-aiplatform  -quiet"
+        "! pip3 install --upgrade google-cloud-aiplatform  --quiet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "e9255e3b156f"
+        "id": "ff555b32bab8"
      },
      "source": [
-        "### Colab Only: Uncomment the following cell to restart the kernel"
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "0c0b2427998a"
+        "id": "f09b4dff629a"
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "435b8e413535"
+        "id": "ee775571c2b5"
      },
      "source": [
-        "### Before you begin\n",
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "92e68cfc3a90"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
-        "#### Set your project ID\n",
+        "Authenticate your environment on Google Colab.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "46604f70e831"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "-  Run `gcloud config list`\n",
-        "-  Run `gcloud projects list`\n",
-        "-  See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "c066e7925949"
+      },
+      "source": [
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
@@ -192,90 +244,7 @@
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# set the project id\n",
-        "! gcloud config set project $PROJECT_ID"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "2e6b8b324ce1"
-      },
-      "source": [
-        "#### Region\n",
-        "\n",
-        "You can also change the `REGION` variable used by Vertex AI. \n",
-        "Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "region"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "6c43a8673066"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
-        "\n",
-        "**1. Vertex AI Workbench** \n",
-        "- Do nothing as you are already authenticated.\n",
-        "\n",
-        "**2. Local JupyterLab Instance,** uncomment and run."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "fbc9cd30cc4b"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "cd0da2c26879"
-      },
-      "source": [
-        "**3. Colab,** uncomment and run:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "a336a05c6149"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0461097edfa5"
-      },
-      "source": [
-        "**4. Service Account or other**\n",
-        "- See how to grant Cloud Storage permissions to your service account at [IAM Ch Examples](https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples)."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -306,7 +275,7 @@
        "id": "create_bucket"
      },
      "source": [
-        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+        "**If your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
@@ -317,27 +286,29 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION $BUCKET_URI"
+        "! gsutil mb -l $LOCATION $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "validate_bucket"
+        "id": "a56633b047ee"
      },
      "source": [
-        "Finally, validate access to your Cloud Storage bucket by examining its contents:"
+        "### Initialize Vertex AI SDK for Python"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "yWnghzKFox9S"
+        "id": "0e24370a7912"
      },
      "outputs": [],
      "source": [
-        "! gsutil ls -al $BUCKET_URI"
+        "import google.cloud.aiplatform as aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -348,43 +319,7 @@
      "source": [
        "### Set up variables\n",
        "\n",
-        "Next, set up some variables used throughout the tutorial.\n",
-        "### Import libraries and define constants"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "import_aip:mbsdk"
-      },
-      "outputs": [],
-      "source": [
-        "import os\n",
-        "\n",
-        "import google.cloud.aiplatform as aiplatform"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "init_aip:mbsdk"
-      },
-      "source": [
-        "### Initialize Vertex AI SDK for Python\n",
-        "\n",
-        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "JZg2sszQox9T"
-      },
-      "outputs": [],
-      "source": [
-        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
+        "Next, set up some variables used throughout the tutorial."
      ]
    },
    {
@@ -406,7 +341,7 @@
        "\n",
        "Learn more about [hardware accelerator support for your region](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators).\n",
        "\n",
-        "*Note*: TF releases before 2.3 for GPU support will fail to load the custom model in this tutorial. It is a known issue and fixed in TF 2.3. This is caused by static graph ops that are generated in the serving function. If you encounter this issue on your own custom models, use a container image for TF 2.3 with GPU support."
+        "*Note*: TF releases before 2.3 for GPU support is expected to fail to load the custom model in this tutorial. It's a known issue and fixed in TF 2.3. This is caused by static graph ops that are generated in the serving function. If you encounter this issue on your own custom models, use a container image for TF 2.3 with GPU support."
      ]
    },
    {
@@ -451,10 +386,10 @@
        "DEPLOY_VERSION = \"xgboost-cpu.1-1\"\n",
        "\n",
        "TRAIN_IMAGE = \"{}-docker.pkg.dev/vertex-ai/training/{}:latest\".format(\n",
-        "    REGION.split(\"-\")[0], TRAIN_VERSION\n",
+        "    LOCATION.split(\"-\")[0], TRAIN_VERSION\n",
        ")\n",
        "DEPLOY_IMAGE = \"{}-docker.pkg.dev/vertex-ai/prediction/{}:latest\".format(\n",
-        "    REGION.split(\"-\")[0], DEPLOY_VERSION\n",
+        "    LOCATION.split(\"-\")[0], DEPLOY_VERSION\n",
        ")"
      ]
    },
@@ -468,7 +403,7 @@
        "\n",
        "Next, set the machine type to use for training.\n",
        "\n",
-        "- Set the variable `TRAIN_COMPUTE` to configure  the compute resources for the VMs you will use for for training.\n",
+        "- Set the variable `TRAIN_COMPUTE` to configure  the compute resources for the VMs used for training.\n",
        " - `machine type`\n",
        "     - `n1-standard`: 3.75GB of memory per vCPU.\n",
        "     - `n1-highmem`: 6.5GB of memory per vCPU\n",
@@ -503,8 +438,8 @@
      "source": [
        "## Introduction to XGBoost training\n",
        "\n",
-        "Once you have trained a XGBoost model, you will want to save it at a Cloud Storage location, so it can subsequently be uploaded to a `Vertex AI Model` resource.\n",
-        "The XGBoost package does not have support to save the model to a Cloud Storage location. Instead, you will do the following steps to save to a Cloud Storage location.\n",
+        "Once you have trained a XGBoost model, save it at a Cloud Storage location, so it can subsequently be uploaded to a Vertex AI model resource.\n",
+        "The XGBoost package does not have support to save the model to a Cloud Storage location. Instead, do the following steps to save to a Cloud Storage location.\n",
        "\n",
        "1. Save the in-memory model to the local filesystem (e.g., model.bst).\n",
        "2. Use gsutil to copy the local copy to the specified Cloud Storage location.\n",
@@ -518,11 +453,18 @@
        "id": "examine_training_package:xgboost"
      },
      "source": [
-        "### Examine the training package\n",
-        "\n",
+        "### Examine the training package"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "06692103c9e5"
+      },
+      "source": [
        "#### Package layout\n",
        "\n",
-        "Before you start the training, you will look at how a Python package is assembled for a custom training job. When unarchived, the package contains the following directory/file layout.\n",
+        "Before you start the training, look at how a Python package is assembled for a custom training job. When unarchived, the package contains the following directory/file layout.\n",
        "\n",
        "- PKG-INFO\n",
        "- README.md\n",
@@ -534,11 +476,16 @@
        "\n",
        "The files `setup.cfg` and `setup.py` are the instructions for installing the package into the operating environment of the Docker image.\n",
        "\n",
-        "The file `trainer/task.py` is the Python script for executing the custom training job. *Note*, when we referred to it in the worker pool specification, we replace the directory slash with a dot (`trainer.task`) and dropped the file suffix (`.py`).\n",
-        "\n",
-        "#### Package Assembly\n",
-        "\n",
-        "In the following cells, you will assemble the training package."
+        "The file `trainer/task.py` is the Python script for executing the custom training job. *Note*, when `trainer/task.py` is referred to in the worker pool specification, the directory slash is replaced with a dot and the file suffix (`.py`) suffix is dropped: (`trainer.task`).\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "b4b1b6caa966"
+      },
+      "source": [
+        "#### Package Assembly"
      ]
    },
    {
@@ -581,7 +528,7 @@
        "Next, you create the `task.py` script for driving the training package. Some noteable steps include:\n",
        "\n",
        "- Command-line arguments:\n",
-        "    - `model-dir`: The location to save the trained model. When using Vertex AI custom training, the location will be specified in the environment variable: `AIP_MODEL_DIR`,\n",
+        "    - `model-dir`: The location to save the trained model. When using Vertex AI custom training, the location is specified in the environment variable: `AIP_MODEL_DIR`,\n",
        "    - `dataset_data_url`: The location of the training data to download.\n",
        "    - `dataset_labels_url`: The location of the training labels to download.\n",
        "    - `boost-rounds`: Tunable hyperparameter\n",
@@ -629,7 +576,7 @@
        "parser.add_argument(\"--dataset-labels-url\", dest=\"dataset_labels_url\",\n",
        "                    type=str, help=\"Download url for the training data labels.\")\n",
        "parser.add_argument(\"--boost-rounds\", dest=\"boost_rounds\",\n",
-        "                    default=20, type=int, help=\"Number of boosted rounds\")\n",
+        "                    default=20, type=int, help=\"Number of boosted rounds\") \n",
        "args = parser.parse_args()\n",
        "\n",
        "logging.getLogger().setLevel(logging.INFO)\n",
@@ -951,14 +898,14 @@
        "id": "cleanup:mbsdk"
      },
      "source": [
-        "# Cleaning up\n",
-        "\n",
+        "## Clean up\n",
        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
        "\n",
        "Otherwise, you can delete the individual resources you created in this tutorial:\n",
        "\n",
-        "- Custom Job (Custome Training job is remove in previous step)\n",
+        "- Custom Job (Custom Training job is removed in the previous step)\n",
+        "- Model\n",
        "- Cloud Storage Bucket"
      ]
    },
@@ -970,11 +917,11 @@
      },
      "outputs": [],
      "source": [
-        "delete_bucket = True\n",
+        "delete_bucket = False\n",
        "\n",
        "model.delete()\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil rm -r $BUCKET_URI"
      ]
    }
@@ -2,7 +2,7 @@
  "cells": [
    {
      "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 1,
      "metadata": {
        "id": "copyright"
      },
@@ -32,25 +32,27 @@
        "# Custom training and batch prediction\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/sdk-custom-image-classification-batch.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/sdk-custom-image-classification-batch.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fcustom%2Fsdk-custom-image-classification-batch.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
-        "  </td>\n",
-        " <td>\n",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/custom/sdk-custom-image-classification-batch.ipynb\">\n",
-        "        <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
-        "      Open in Vertex AI Workbench\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
-        "</table>\n",
-        "<br/><br/><br/>"
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/sdk-custom-image-classification-batch.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
      ]
    },
    {
@@ -78,7 +80,7 @@
        "In this tutorial, you learn to use Vertex AI Training to create a custom trained model and use Vertex AI Batch Prediction to do a batch prediction on the trained model.\n",
        "\n",
        "\n",
-        "create a custom-trained model from a Python script in a Docker container using the Vertex AI SDK for Python, and then do a prediction on the deployed model by sending data. Alternatively, you can create custom-trained models using `gcloud` command-line tool, or online using the Cloud Console.\n",
+        "Create a custom-trained model from a Python script in a Docker container using the Vertex AI SDK for Python, and then do a prediction on the deployed model by sending data. Alternatively, you can create custom-trained models using `gcloud` command-line tool, or online using the Cloud Console.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
@@ -128,23 +130,29 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "install_aip"
+        "id": "3b1ffd5ab768"
      },
      "source": [
-        "### Install additional packages\n",
-        "\n",
-        "Install the following packages for executing this notebook."
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aae9ca040eab"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "1fd00fa70a2a"
+        "id": "23e23ce735c9"
      },
      "outputs": [],
      "source": [
-        "# Install the packages\n",
        "! pip3 install --upgrade google-cloud-aiplatform \\\n",
        "                        google-cloud-storage \\\n",
        "                        pillow  \\\n",
@@ -154,156 +162,92 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "d98bc9fdd80d"
+        "id": "ff555b32bab8"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel."
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "0dbf29389c65"
+        "id": "f09b4dff629a"
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "013daf3de88e"
+        "id": "54c5ef8a8f43"
      },
      "source": [
-        "## Before you begin"
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "8bc8a29f9001"
+        "id": "92e68cfc3a90"
      },
      "source": [
-        "#### Set your project ID\n",
+        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "* Run `gcloud config list`.\n",
-        "* Run `gcloud projects list`.\n",
-        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "e61aaa036444"
+        "id": "46604f70e831"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "107c51893a64"
+      },
+      "source": [
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "294fe4e5a671"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# Set the project id\n",
-        "! gcloud config set project {PROJECT_ID}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c4a624c8099d"
-      },
-      "source": [
-        "#### Region\n",
-        "\n",
-        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "f83bd6013894"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "08bfd1eb44ef"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "af349043f23b"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ad1138a125ea"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ce6043da7b33"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0367eac06a10"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "21ad4dbb4a61"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c13224697bfb"
-      },
-      "source": [
-        "**4. Service account or other**\n",
-        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -312,7 +256,7 @@
        "id": "ddbea904fbe5"
      },
      "source": [
-        "### Create a Cloud Storage bucket\n",
+        "#### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
@@ -325,7 +269,7 @@
      },
      "outputs": [],
      "source": [
-        "BUCKET_URI = \"gs://your-bucket-name-unique\"  # @param {type:\"string\"}"
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -345,7 +289,31 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
+        "! gsutil mb -l $LOCATION -p $PROJECT_ID $BUCKET_URI"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "fb03963cdb69"
+      },
+      "source": [
+        "#### Initialize Vertex AI SDK for Python\n",
+        "\n",
+        "Initialize the Vertex AI SDK for Python for your project."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "a4f61991b160"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -365,31 +333,7 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
-        "from google.cloud import aiplatform"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "import_aip"
-      },
-      "source": [
-        "### Initialize Vertex AI \n",
-        "\n",
-        "Initialize the Vertex AI SDK for Python for your project."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "cNEiwLd0lugu"
-      },
-      "outputs": [],
-      "source": [
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
+        "import os"
      ]
    },
    {
@@ -970,9 +914,6 @@
      },
      "outputs": [],
      "source": [
-        "delete_training_job = True\n",
-        "delete_model = True\n",
-        "\n",
        "# Warning: Setting this to true will delete everything in your bucket\n",
        "delete_bucket = False\n",
        "\n",
@@ -982,7 +923,7 @@
        "# Delete the model\n",
        "model.delete()\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
@@ -32,25 +32,27 @@
        "# Custom training and online prediction\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/sdk-custom-image-classification-online.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/sdk-custom-image-classification-online.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fcustom%2Fsdk-custom-image-classification-online.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/custom/sdk-custom-image-classification-online.ipynb\">\n",
-        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
-        "      Open in Vertex AI Workbench\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
-        "</table>\n",
-        "<br/><br/><br/>"
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/sdk-custom-image-classification-online.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
      ]
    },
    {
@@ -128,23 +130,31 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "install_aip"
+        "id": "3b1ffd5ab768"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the packages required for executing this notebook."
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aae9ca040eab"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "1fd00fa70a2a"
+        "id": "0d3d0b3a3d88"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
+        "                                 google-cloud \\\n",
        "                                 google-cloud-storage \\\n",
        "                                 pillow \\\n",
        "                                 numpy"
@@ -153,142 +163,92 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "d98bc9fdd80d"
+        "id": "ff555b32bab8"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel."
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "0dbf29389c65"
+        "id": "f09b4dff629a"
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "eec0fc7a0963"
+        "id": "54c5ef8a8f43"
      },
      "source": [
-        "## Before you begin\n",
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "92e68cfc3a90"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
-        "#### Set your project ID\n",
-        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "* Run `gcloud config list`.\n",
-        "* Run `gcloud projects list`.\n",
-        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "e3ce64be5527"
+        "id": "46604f70e831"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "107c51893a64"
+      },
+      "source": [
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "294fe4e5a671"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# Set the project id\n",
-        "! gcloud config set project {PROJECT_ID}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "7f2e7f78a864"
-      },
-      "source": [
-        "#### Region\n",
-        "\n",
-        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "369a2258dd59"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "6a94297012d5"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
-        "\n",
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "28e3c4539627"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**\n"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "fbc9cd30cc4b"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "79efab26ad02"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "a336a05c6149"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "e3c28b6b796b"
-      },
-      "source": [
-        "**4. Service account or other**\n",
-        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -297,7 +257,7 @@
        "id": "70a42f1033a3"
      },
      "source": [
-        "### Create a Cloud Storage bucket\n",
+        "#### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
@@ -319,7 +279,7 @@
        "id": "0cb016da6de3"
      },
      "source": [
-        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+        "**If your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
@@ -330,7 +290,31 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
+        "! gsutil mb -l $LOCATION -p $PROJECT_ID $BUCKET_URI"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ae35ce9954e4"
+      },
+      "source": [
+        "#### Initialize Vertex AI SDK for Python\n",
+        "\n",
+        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "a4f61991b160"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -353,32 +337,9 @@
        "import os\n",
        "\n",
        "import numpy as np\n",
-        "from google.cloud import aiplatform\n",
        "from PIL import Image"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "import_aip"
-      },
-      "source": [
-        "### Initialize Vertex AI SDK for Python\n",
-        "\n",
-        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "cNEiwLd0lugu"
-      },
-      "outputs": [],
-      "source": [
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -882,7 +843,7 @@
        "# Delete the endpoint\n",
        "endpoint.delete()\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
@@ -33,20 +33,25 @@
        "\n",
        "<table align=\"left\">\n",
        "\n",
-        "  <td>\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=\"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/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>\n",
        "  </td>\n",
-        "  <td>\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\">\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>\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>\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/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\">\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",
@@ -61,7 +66,7 @@
      "source": [
        "## Overview\n",
        "\n",
-        "As a Data Scientist, you probably start running model experiments locally on your notebook. Depending on the framework you use, you would need to track parameters, training time series and evaluation metrics. In this way, you would be able to explain the modelling approach you would choose. \n",
+        "As a Data Scientist, you may start running model experiments locally on your notebook. Depending on the framework you use, you need to track parameters, training time series and evaluation metrics. In this way, you are able to explain the modelling approach you have choosen. \n",
        "\n",
        "Learn more about [Vertex AI Experiments](https://cloud.google.com/vertex-ai/docs/experiments/intro-vertex-ai-experiments)."
      ]
@@ -96,7 +101,7 @@
      "source": [
        "### Dataset\n",
        "\n",
-        "In this notebook, we will train a simple distributed neural network (DNN) model to predict automobile's miles per gallon (MPG) based on automobile information in the [auto-mpg dataset](https://www.kaggle.com/devanshbesain/exploration-and-analysis-auto-mpg).\n"
+        "In this notebook, you train a simple distributed neural network (DNN) model to predict an automobile's miles per gallon (MPG) based on automobile information in the [auto-mpg dataset](https://www.kaggle.com/devanshbesain/exploration-and-analysis-auto-mpg).\n"
      ]
    },
    {
@@ -39,7 +39,7 @@
        "  </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%2Fdelete_outdated_tensorboard_experiments.ipynb\">\n",
-        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\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",
@@ -29,28 +29,29 @@
        "id": "rbTKH4_6f1ux"
      },
      "source": [
-        "## Online feature serving and fetching of BigQuery data with Vertex AI Feature Store Optimized Serving\n",
-        "\n",
+        "# Online feature serving and fetching of BigQuery data with Vertex AI Feature Store Optimized Serving\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/online_feature_serving_and_fetching_bigquery_data_with_feature_store_optimized.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/feature_store/online_feature_serving_and_fetching_bigquery_data_with_feature_store_optimized.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_optimized.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" 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/online_feature_serving_and_fetching_bigquery_data_with_feature_store_optimized.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/feature_store/online_feature_serving_and_fetching_bigquery_data_with_feature_store_optimized.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>"
      ]
    },
@@ -77,9 +78,9 @@
        "\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'll 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",
+        "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 ML services and resources:\n",
+        "This tutorial uses the following Google Cloud Vertex AI services and resources:\n",
        "\n",
        "- `Vertex AI Feature Store`\n",
        "\n",
@@ -96,7 +97,7 @@
        "id": "XIMD3zCgDGop"
      },
      "source": [
-        "###Note\n",
+        "### Note\n",
        "This is a public Preview release. By using the feature, you acknowledge that you're aware of the open issues and that this preview is provided “as is” under the pre-GA terms of service.\n"
      ]
    },
@@ -139,13 +140,23 @@
        "to generate a cost estimate based on your projected usage."
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "7cc240e9c1a6"
+      },
+      "source": [
+        "## Getting Started"
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yTJiDCrYsOmT"
      },
      "source": [
-        "## Installation\n",
+        "### Install Vertex AI SDK and other required packages\n",
+        "\n",
        "\n",
        "Install the following packages required to execute this notebook."
      ]
@@ -170,7 +181,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."
      ]
    },
    {
@@ -181,32 +194,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": "mdqw6ADTtJRI"
+        "id": "972161a99004"
      },
      "source": [
-        "## Before you begin\n",
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Please wait until it is finished before continuing to the next step. ⚠️</b>\n",
+        "</div>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ff666ce4051c"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
-        "### Set up your Google Cloud project\n",
+        "Authenticate your environment on Google Colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "cc7251520a07"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
        "\n",
-        "**The following steps are required, regardless of your notebook environment.**\n",
+        "if \"google.colab\" in sys.modules:\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",
+        "    from google.colab import auth\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)."
+        "    auth.authenticate_user()"
      ]
    },
    {
@@ -215,12 +249,9 @@
        "id": "bFANidV0tPbo"
      },
      "source": [
-        "#### Set your project ID\n",
+        "### Set Google Cloud project information and initialize Vertex AI SDK\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 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)."
      ]
    },
    {
@@ -232,116 +263,13 @@
      "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",
+        "API_ENDPOINT = f\"{LOCATION}-aiplatform.googleapis.com\"\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-central1`\n",
-        "* `asia-southeast1`\n",
-        "* `europe-west1`\n",
-        "* `europe-west2`\n",
-        "* `europe-west3`\n",
-        "* `europe-west4`\n",
-        "* `us-east1`\n",
-        "* `us-east4`\n",
-        "* `us-west1`\n",
-        "* `us-west2`\n",
-        "* `us-west3`\n",
-        "* `europe-west8`\n",
-        "* `asia-northeast2`"
-      ]
-    },
-    {
-      "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 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)"
      ]
    },
    {
@@ -350,7 +278,7 @@
        "id": "OMCl0avIusKl"
      },
      "source": [
-        "### Import libraries"
+        "## Import libraries"
      ]
    },
    {
@@ -361,7 +289,7 @@
      },
      "outputs": [],
      "source": [
-        "from google.cloud import aiplatform, bigquery\n",
+        "from google.cloud import bigquery\n",
        "from google.cloud.aiplatform_v1beta1 import (\n",
        "    FeatureOnlineStoreAdminServiceClient, FeatureRegistryServiceClient)\n",
        "from google.cloud.aiplatform_v1beta1.types import feature as feature_pb2\n",
@@ -385,30 +313,6 @@
        "                                        FeatureViewBigQuerySource)"
      ]
    },
-    {
-      "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": {
@@ -463,7 +367,7 @@
        "id": "ATHTdLGIrV5H"
      },
      "source": [
-        "#### View the retrieved data"
+        "### View the retrieved data"
      ]
    },
    {
@@ -533,7 +437,7 @@
        "id": "ZugQMKdLtYm7"
      },
      "source": [
-        "#### Create a BigQuery view"
+        "### Create a BigQuery view"
      ]
    },
    {
@@ -570,12 +474,12 @@
        "    * Choose Optimized as the storage type.\n",
        "1. Define the data (`FeatureView`) to be served by the newly-created instance. This can map to either of the following:\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` that you create to host feature metadata.\n",
        "\n",
-        "We recommend NOT sending loads larger than 7500 QPS to one FeatureOnlineStore.\n",
-        "In general, we recommend creating multiple gRPC connections to one FeatureOnlineStore, and evenly distribute your loads across them. More connections and smaller per-connection QPS typically help with internal load balancing and scaling, reducing the chance of seeing higher tail-latencies. Specifically:\n",
+        "It is not recommended to send loads larger than 7500 QPS to one FeatureOnlineStore.\n",
+        "In general, you can create multiple gRPC connections to one FeatureOnlineStore, and evenly distribute your loads across them. More connections and smaller per-connection QPS typically help with internal load balancing and scaling, reducing the chance of seeing higher tail-latencies. Specifically:\n",
        "1. If your FetchFeatureValues response payload size is small (e.g. less than 1 kB), you may create one connection for up to every 2000 QPS.\n",
-        "2. If your FetchFeatureValues response payload size can be large (e.g. more than a few kB or 10s of kB), we recommend you to create one connection for up to every 250 QPS, and we also recommend avoiding sudden increases of loads.\n"
+        "2. If your FetchFeatureValues response payload size is large (e.g. more than a few kB or 10s of kB), you may create one connection for up to every 250 QPS. It is recommended to avoid sudden increases of loads.\n"
      ]
    },
    {
@@ -654,7 +558,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",
@@ -774,7 +678,7 @@
        "\n",
        "Use the following guidelines to understand the schema and constraints while creating the BigQuery source:\n",
        "\n",
-        "* The BigQuery table or view *must* have a column with `string` values to use as the (entity) IDs. You'll need to specify that this column is the ID column during `FeatureView` creation. Note that the size of each value in this column must be less than 4 KB.\n",
+        "* The BigQuery table or view *must* have a column with `string` values to use as the (entity) IDs. You need to specify that this column is the ID column during `FeatureView` creation. Note that the size of each value in this column must be less than 4 KB.\n",
        "* Feature Store ignores rows containing the duplicate values in the ID column and arbitrarily picks one to use for online serving. This might occur if the table or view has multiple feature values recorded at different timestamps. In this case, you need to sync the data in your table based on the latest timestamp before serving.\n",
        "* Feature Store considers each row as a complete record. For example, if a feature value is missing in one of the columns, then it's considered as a null value.\n",
        "* Provide values for each feature is a separate column. Supported data types are `bool`, `int`, `double`, `string`, timestamp, arrays of these data types, and bytes. Note that the timestamp data type is converted to `int64` during data sync.\n",
@@ -814,7 +718,7 @@
        "id": "l1CnCKzCSTXd"
      },
      "source": [
-        "#### Verify FeatureView instance creation\n",
+        "##### Verify FeatureView instance creation\n",
        "\n",
        "After the long-running operation (LRO) is complete, show the result."
      ]
@@ -850,7 +754,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 we would like to register.\n"
      ]
    },
    {
@@ -861,11 +765,11 @@
      "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",
-        "* The BigQuery table or view *must* have a column with `string` values to use as the (entity) IDs. You'll need to specify that this column is the ID column during `FeatureGroup` creation. Note that the size of each value in this column must be less than 4 KB.\n",
+        "* The BigQuery table or view *must* have a column with `string` values to use as the (entity) IDs. You need to specify that this column is the ID column during `FeatureGroup` creation. Note that the size of each value in this column must be less than 4 KB.\n",
        "* The BigQuery table or view *must* have a column named `feature_timestamp` with `timestamp` values to use as timestamp column.\n",
        "* Feature Registry sources are treated as sparse by default i.e. a point in time lookup (BQ.ML_FEATURES_AT_TIME()) to generate latest featureValues per entityId.\n",
        "* Provide values for each feature is a separate column. Supported data types are `bool`, `int`, `double`, `string`, timestamp, arrays of these data types, and bytes. Note that the timestamp data type is converted to `int64` during data sync.\n",
@@ -911,7 +815,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",
@@ -951,7 +855,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",
@@ -1003,7 +907,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",
@@ -1022,6 +926,7 @@
        "id": "wXk90azuJM0z"
      },
      "source": [
+        "##### Verify FeatureView instance creation\n",
        "Verify that the `FeatureView` instance is created by listing all the feature views within the online store."
      ]
    },
@@ -1035,7 +940,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",
        ")"
      ]
    },
@@ -1148,7 +1053,7 @@
      "source": [
        "#### Option 1: Optimized online serving with private endpoint\n",
        "\n",
-        "You'll need to connect to private endpoint over gRPC. Follow these instructions to set up [Private Service Connect](https://cloud.google.com/vpc/docs/private-service-connect)."
+        "You need to connect to private endpoint over gRPC. Follow these instructions to set up [Private Service Connect](https://cloud.google.com/vpc/docs/private-service-connect)."
      ]
    },
    {
@@ -1172,7 +1077,7 @@
      "source": [
        "# Get Optimized online store\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",
        ")"
      ]
    },
@@ -1182,7 +1087,7 @@
        "id": "kk7ec2RDN597"
      },
      "source": [
-        "You'll see output similar to the following:\n",
+        "You should see output similar to the following:\n",
        "\n",
        "```\n",
        "dedicated_serving_endpoint {\n",
@@ -1208,7 +1113,7 @@
        "\n",
        "where you created your `FeatureOnlineStore` instance.\n",
        "\n",
-        "* `service_attachment_string` represents the target service that you'll need to specify while [adding Private Service Connect to your network configuration]()."
+        "* `service_attachment_string` represents the target service that you need to specify while [adding Private Service Connect to your network configuration]()."
      ]
    },
    {
@@ -1226,7 +1131,7 @@
        "1. Enter your **endpoint name**.\n",
        "1. In the **Network** field, select **default**.\n",
        "1. In the **Subetwork** field, select **default**.\n",
-        "1. In the **IP address** list, click **Create IP address** to create an IP address that you'll use to connect the Feature Store API.\n",
+        "1. In the **IP address** list, click **Create IP address** to create an IP address that you use to connect the Feature Store API.\n",
        "1. Select **Enable global access**.\n",
        "1. Click **Add Endpoint**.\n",
        "\n",
@@ -1280,7 +1185,7 @@
        "#     feature_online_store_service as feature_online_store_service_pb2\n",
        "# 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",
        "#         id=\"16050\"))"
      ]
    },
@@ -1314,7 +1219,7 @@
      },
      "outputs": [],
      "source": [
-        "# It will take some time for the DNS to be fully ready\n",
+        "# It takes some time for the DNS to be fully ready\n",
        "time.sleep(300)\n",
        "\n",
        "# Depends on the FeatureView you created, the FEATURE_VIEW_ID here can be different\n",
@@ -1352,12 +1257,12 @@
        "# 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",
        "for fos in FeatureOnlineStore.list():\n",
@@ -4,6 +4,7 @@
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
+        "cellView": "form",
        "id": "8EcdxqUnftBL"
      },
      "outputs": [],
@@ -29,28 +30,31 @@
        "id": "rbTKH4_6f1ux"
      },
      "source": [
-        "## Online feature serving and vector retrieval of BigQuery data with Vertex AI Feature Store\n",
+        "# Online feature serving and vector retrieval of BigQuery data with Vertex AI Feature Store\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/feature_store/online_feature_serving_and_vector_retrieval_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",
+        "      <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_vector_retrieval_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%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",
+        "      <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/online_feature_serving_and_vector_retrieval_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",
+        "      <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/feature_store/online_feature_serving_and_vector_retrieval_bigquery_data_with_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>"
      ]
    },
@@ -128,15 +132,22 @@
        "to generate a cost estimate based on your projected usage."
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "7cc240e9c1a6"
+      },
+      "source": [
+        "## Get Started"
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yTJiDCrYsOmT"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook."
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
@@ -159,7 +170,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."
      ]
    },
    {
@@ -170,46 +183,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": "mdqw6ADTtJRI"
+        "id": "e669f8088ac3"
      },
      "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": "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": "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)."
      ]
    },
    {
@@ -221,125 +252,13 @@
      "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",
+        "from google.cloud import aiplatform\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": "G6iMMALZthFM"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "Ni5jx6RGtzG3"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION)\n",
        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "deb54cbe532e"
-      },
-      "outputs": [],
-      "source": [
-        "SERVICE_ACCOUNT = \"[your service account]\""
-      ]
-    },
-    {
-      "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(project_id=PROJECT_ID)"
-      ]
-    },
-    {
-      "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."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "90a7c39da629"
-      },
-      "outputs": [],
-      "source": [
-        "! gcloud auth list"
+        "API_ENDPOINT = f\"{LOCATION}-aiplatform.googleapis.com\""
      ]
    },
    {
@@ -359,37 +278,13 @@
      },
      "outputs": [],
      "source": [
-        "from google.cloud import aiplatform, bigquery\n",
+        "from google.cloud import bigquery\n",
        "from google.cloud.aiplatform_v1.types import NearestNeighborQuery\n",
        "from vertexai.resources.preview import (FeatureOnlineStore, FeatureView,\n",
        "                                        FeatureViewBigQuerySource)\n",
        "from vertexai.resources.preview.feature_store import utils"
      ]
    },
-    {
-      "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": {
@@ -401,7 +296,7 @@
        "### Requirements\n",
        "The data source has to be a BigQuery table or a BigQuery view, with the following requirements on columns:\n",
        "1. [*Required*] One entity id column, type: string\n",
-        "2. [*Required*] One embedding column, type: float array\n",
+        "2. [*Required*] One embedding column, type: double array\n",
        "3. [*Optional*] One or more filtering columns, type: string or string array\n",
        "4. [*Optional*] One crowding column, type: integer. Crowding ensures that results are diverse by returning at most k' < k neighbors with any single crowding attribute out of k total neighbors"
      ]
@@ -588,7 +483,7 @@
      "source": [
        "### Create Feature Online Store\n",
        "\n",
-        "Create an Optimized feature online store. After the long-running operation (LRO) is complete, the result will be shown.\n",
+        "Create an Optimized feature online store. After the long-running operation (LRO) is complete, the result is shown.\n",
        "\n",
        "> **Note:** This operation might take up to 10 minutes to complete."
      ]
@@ -704,7 +599,7 @@
        "* A data source (BigQuery table or view URI or FeatureGroup/features ) synced to the `FeatureOnlineStore` instance for serving.\n",
        "* The cron schedule to run the sync pipeline.\n",
        "\n",
-        "Within feature view creation, a sync job will be scheduled, either started immediately or following the cron schedule. In the sync job, data is exported, index is built and deployed to Feature Store backend."
+        "Within feature view creation, a sync job is scheduled, either started immediately or following the cron schedule. In the sync job, data is exported, index is built and deployed to Feature Store backend."
      ]
    },
    {
@@ -716,7 +611,7 @@
      "outputs": [],
      "source": [
        "FEATURE_VIEW_ID = \"feature_view_publications\"  # @param {type: \"string\"}\n",
-        "# A schedule will be created based on cron setting.\n",
+        "# A schedule is created based on cron setting.\n",
        "CRON_SCHEDULE = \"TZ=America/Los_Angeles 00 13 11 8 *\"  # @param {type: \"string\"}"
      ]
    },
@@ -789,7 +684,7 @@
      "source": [
        "### Verify feature view instance creation\n",
        "\n",
-        "Verify that the FeatureView instance is created by gettting the feature view."
+        "Verify that the FeatureView instance is created by getting the feature view."
      ]
    },
    {
@@ -978,7 +873,7 @@
      },
      "outputs": [],
      "source": [
-        "# It will take some time for the DNS to be fully ready\n",
+        "# It takes some time for the DNS to be fully ready\n",
        "time.sleep(300)\n",
        "\n",
        "my_fv.search(\n",
@@ -1071,7 +966,11 @@
        "my_fv.delete()\n",
        "\n",
        "# Delete Feature Online Store\n",
-        "my_fos.delete(force=True)"
+        "my_fos.delete(force=True)\n",
+        "\n",
+        "# Delete BigQuery dataset. Uncomment and run the command below if you want to delete the BigQuery set.\n",
+        "# Do this only if the dataset is created for this demo.\n",
+        "# bq_client.delete_dataset(BQ_DATASET_ID, delete_contents=True, not_found_ok=True)"
      ]
    }
  ],
@@ -347,18 +347,17 @@
        "import uuid\n",
        "\n",
        "from google.cloud import aiplatform, bigquery\n",
-        "from google.cloud.aiplatform_v1beta1 import (\n",
-        "    FeatureOnlineStoreAdminServiceClient, FeatureOnlineStoreServiceClient)\n",
-        "from google.cloud.aiplatform_v1beta1.types import NearestNeighborQuery\n",
-        "from google.cloud.aiplatform_v1beta1.types import \\\n",
+        "from google.cloud.aiplatform_v1 import (FeatureOnlineStoreAdminServiceClient,\n",
+        "                                        FeatureOnlineStoreServiceClient)\n",
+        "from google.cloud.aiplatform_v1.types import NearestNeighborQuery\n",
+        "from google.cloud.aiplatform_v1.types import \\\n",
        "    feature_online_store as feature_online_store_pb2\n",
-        "from google.cloud.aiplatform_v1beta1.types import \\\n",
+        "from google.cloud.aiplatform_v1.types import \\\n",
        "    feature_online_store_admin_service as \\\n",
        "    feature_online_store_admin_service_pb2\n",
-        "from google.cloud.aiplatform_v1beta1.types import \\\n",
+        "from google.cloud.aiplatform_v1.types import \\\n",
        "    feature_online_store_service as feature_online_store_service_pb2\n",
-        "from google.cloud.aiplatform_v1beta1.types import \\\n",
-        "    feature_view as feature_view_pb2"
+        "from google.cloud.aiplatform_v1.types import feature_view as feature_view_pb2"
      ]
    },
    {
@@ -707,14 +706,7 @@
      "outputs": [],
      "source": [
        "online_store_config = feature_online_store_pb2.FeatureOnlineStore(\n",
-        "    bigtable=feature_online_store_pb2.FeatureOnlineStore.Bigtable(\n",
-        "        auto_scaling=feature_online_store_pb2.FeatureOnlineStore.Bigtable.AutoScaling(\n",
-        "            min_node_count=1, max_node_count=3, cpu_utilization_target=50\n",
-        "        )\n",
-        "    ),\n",
-        "    embedding_management=feature_online_store_pb2.FeatureOnlineStore.EmbeddingManagement(\n",
-        "        enabled=True\n",
-        "    ),\n",
+        "    optimized=feature_online_store_pb2.FeatureOnlineStore.Optimized(),\n",
        ")\n",
        "\n",
        "create_store_lro = admin_client.create_feature_online_store(\n",
@@ -810,7 +802,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 to Cloud Bigtable, a index is built and deployed to GKE cluster."
+        "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."
      ]
    },
    {
@@ -835,7 +827,7 @@
      },
      "outputs": [],
      "source": [
-        "# Vector search configs\n",
+        "# Index building configs\n",
        "DIMENSIONS = 768  # @param {type: \"number\"}\n",
        "EMBEDDING_COLUMN = \"embedding\"  # @param {type: \"string\"}\n",
        "# Optional\n",
@@ -871,15 +863,15 @@
        "\n",
        "sync_config = feature_view_pb2.FeatureView.SyncConfig(cron=CRON_SCHEDULE)\n",
        "\n",
-        "vector_search_config = feature_view_pb2.FeatureView.VectorSearchConfig(\n",
+        "index_config = feature_view_pb2.FeatureView.IndexConfig(\n",
        "    embedding_column=EMBEDDING_COLUMN,\n",
        "    # filter_columns=FILTER_COLUMNS,\n",
        "    # crowding_column=CROWDING_COLUMN,\n",
        "    embedding_dimension=DIMENSIONS,\n",
-        "    tree_ah_config=feature_view_pb2.FeatureView.VectorSearchConfig.TreeAHConfig(),\n",
+        "    tree_ah_config=feature_view_pb2.FeatureView.IndexConfig.TreeAHConfig(),\n",
        ")\n",
        "\n",
-        "print(f\"vector_search_config: {vector_search_config}\")\n",
+        "print(f\"index_config: {index_config}\")\n",
        "\n",
        "create_view_lro = admin_client.create_feature_view(\n",
        "    feature_online_store_admin_service_pb2.CreateFeatureViewRequest(\n",
@@ -888,7 +880,7 @@
        "        feature_view=feature_view_pb2.FeatureView(\n",
        "            big_query_source=big_query_source,\n",
        "            sync_config=sync_config,\n",
-        "            vector_search_config=vector_search_config,\n",
+        "            index_config=index_config,\n",
        "        ),\n",
        "    )\n",
        ")"
@@ -1109,6 +1101,9 @@
      },
      "outputs": [],
      "source": [
+        "# It will take some time for the DNS to be fully ready\n",
+        "time.sleep(300)\n",
+        "\n",
        "data_client = FeatureOnlineStoreServiceClient(\n",
        "    client_options={\"api_endpoint\": PUBLIC_ENDPOINT}\n",
        ")"
@@ -1222,7 +1217,7 @@
        "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",
-        "        id=ENTITY_ID,\n",
+        "        data_key=feature_online_store_service_pb2.FeatureViewDataKey(key=ENTITY_ID),\n",
        "    )\n",
        ")"
      ]
@@ -26,22 +26,45 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "vBSXW09S05My"
+        "id": "103lT16d4tLE"
      },
      "source": [
-        "# Vertex AI Feature Store Feature View Service Agents Tutorial"
+        "# Vertex AI Feature Store Feature View Service Agents Tutorial\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_feature_view_service_agents.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%2Fvertex_ai_feature_store_feature_view_service_agents.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" 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/vertex_ai_feature_store_feature_view_service_agents.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/vertex_ai_feature_store_feature_view_service_agents.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": "103lT16d4tLE"
+        "id": "becda595b9b8"
      },
      "source": [
-        "<table align=\"left\">\n",
-        "<a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/feature_store/vertex_ai_feature_store_feature_view_service_agents.ipynb\\\"><img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\\\" alt=\"Colab logo\">Run in Colab\n",
-        "<a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/feature_store/vertex_ai_feature_store_feature_view_service_agents.ipynb\\\"><img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\\\" alt=\"GitHub logo\">View on GitHub\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_feature_view_service_agents.ipynb\\\"><img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">Open in Vertex AI Workbench\""
+        "## Overview\n",
+        "In this tutorial, you learn how to enable feature view Service Agents and grant each feature view access to the specific source data that is used.\n",
+        "\n",
+        "Learn more about [Vertex AI Feature Store](https://cloud.google.com/vertex-ai/docs/featurestore/latest/overview)."
      ]
    },
    {
@@ -50,12 +73,7 @@
        "id": "E5wnTSHgyuFr"
      },
      "source": [
-        "## Overview\n",
-        "In this tutorial, you learn how to enable feature view Service Agents and grant each feature view access to the specific source data that is used.\n",
-        "\n",
-        "Learn more about [Vertex AI Feature Store](https://cloud.google.com/vertex-ai/docs/featurestore/latest/overview).\n",
-        "\n",
-        "## Objective\n",
+        "### Objective\n",
        "In this tutorial, you learn how to use a dedicated service agent for a feature view in Vertex AI Feature Store. Using a dedicated service agent for a feature view, you grant access to extract feature data from BigQuery in an end-to-end feature serving workflow.\n",
        "\n",
        "This tutorial uses the following Google Cloud services and resources:\n",
@@ -69,7 +87,7 @@
        "## Note\n",
        "This is a Preview release. By using the feature, you acknowledge that you're aware of the open issues and that this preview is provided “as is” under the pre-GA terms of service.\n",
        "\n",
-        "## Costs\n",
+        "### Costs\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "* `Vertex AI`\n",
        "* `BigQuery`\n",
@@ -83,9 +101,16 @@
        "id": "-HGNMzgHFy34"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook."
+        "## Getting Started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-HGNMzgHFy34"
+      },
+      "source": [
+        "### Install Vertex AI SDK and other required packages"
      ]
    },
    {
@@ -108,7 +133,9 @@
        "id": "UGJ_-HL1_P7w"
      },
      "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."
      ]
    },
    {
@@ -119,32 +146,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": "TZP9eoNCAAi4"
+        "id": "e767418763cd"
      },
      "source": [
-        "## Before you begin\n",
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Please wait until it is finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "870814a62e87"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
-        "### Set up your Google Cloud project\n",
+        "Authenticate your environment on Google Colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "c97be6a73155"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
        "\n",
-        "**The following steps are required, regardless of your notebook environment.**\n",
+        "if \"google.colab\" in sys.modules:\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",
+        "    from google.colab import auth\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)."
+        "    auth.authenticate_user()"
      ]
    },
    {
@@ -153,12 +201,9 @@
        "id": "7zkyTnWjACyk"
      },
      "source": [
-        "#### Set your project ID\n",
+        "### Set Google Cloud project information and initialize Vertex AI SDK\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 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)."
      ]
    },
    {
@@ -170,116 +215,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": "Vbc1nfqhxvrS"
-      },
-      "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-central1`\n",
-        "* `asia-southeast1`\n",
-        "* `europe-west1`\n",
-        "* `europe-west2`\n",
-        "* `europe-west3`\n",
-        "* `europe-west4`\n",
-        "* `us-east1`\n",
-        "* `us-east4`\n",
-        "* `us-west1`\n",
-        "* `us-west2`\n",
-        "* `us-west3`\n",
-        "* `europe-west8`\n",
-        "* `asia-northeast2`"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "7Ywl-59Sxy4h"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "1Z_1vxrS4quc"
-      },
-      "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": "xYej1E-9Acgz"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "_ii5t0qTAhDu"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "JQ7TXbsUAgF7"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "M0VJ-rsK4tu8"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "gyUOLvGJAqYo"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "PsAbCyVSAs62"
-      },
-      "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)"
      ]
    },
    {
@@ -299,7 +240,6 @@
      },
      "outputs": [],
      "source": [
-        "from google.cloud import aiplatform\n",
        "from google.cloud.aiplatform_v1beta1 import (\n",
        "    FeatureOnlineStoreAdminServiceClient, FeatureOnlineStoreServiceClient,\n",
        "    FeatureRegistryServiceClient)\n",
@@ -328,9 +268,9 @@
        "id": "S4nZTv4sx_wS"
      },
      "source": [
-        "### Initialize Vertex AI SDK for Python\n",
+        "### Initialize AI Platform Feature Store Clients\n",
        "\n",
-        "Initialize the Vertex AI SDK for Python for your project."
+        "Initialize the AI platform Feature Store Clients for your project."
      ]
    },
    {
@@ -341,9 +281,7 @@
      },
      "outputs": [],
      "source": [
-        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
-        "\n",
-        "API_ENDPOINT = f\"{REGION}-aiplatform.googleapis.com\"\n",
+        "API_ENDPOINT = f\"{LOCATION}-aiplatform.googleapis.com\"\n",
        "\n",
        "admin_client = FeatureOnlineStoreAdminServiceClient(\n",
        "    client_options={\"api_endpoint\": API_ENDPOINT}\n",
@@ -373,10 +311,10 @@
      },
      "outputs": [],
      "source": [
-        "DATASET_ID = \"test_data\"+\"_\"+REGION.replace('-', '_')  # @param {type:\"string\"}\n",
+        "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={REGION}\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",
@@ -391,7 +329,7 @@
        "id": "85tGwbXKkKCo"
      },
      "source": [
-        "## Create a Feature Group\n",
+        "### Create a Feature Group\n",
        "First, create a feature group."
      ]
    },
@@ -429,7 +367,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",
@@ -442,7 +380,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",
@@ -471,7 +409,7 @@
      "source": [
        "# Verify feature group is created.\n",
        "registry_client.get_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",
        ")"
      ]
    },
@@ -494,7 +432,7 @@
      "source": [
        "# Use list to verify the features are created.\n",
        "registry_client.list_features(\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",
        ")"
      ]
    },
@@ -538,7 +476,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",
@@ -568,7 +506,7 @@
      "source": [
        "# Use list 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",
        ")"
      ]
    },
@@ -616,7 +554,7 @@
        "sync_config = feature_view_pb2.FeatureView.SyncConfig(cron=CRON_SCHEDULE)\n",
        "\n",
        "create_view_lro = admin_client.create_feature_view(\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=\"test_fv\",\n",
        "    feature_view=feature_view_pb2.FeatureView(\n",
        "        feature_registry_source=feature_registry_source,\n",
@@ -656,7 +594,7 @@
      "source": [
        "# Use list to verify the store is created.\n",
        "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",
        ")"
      ]
    },
@@ -687,7 +625,7 @@
        "\n",
        "# Get feature view\n",
        "feature_view = 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",
        ")\n",
        "SERVICE_ACCOUNT = feature_view.service_account_email"
      ]
@@ -747,7 +685,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",
        ")\n",
        "\n",
        "# Wait for sync completion\n",
@@ -827,7 +765,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=\"test\"),\n",
        "    )\n",
        ")"
@@ -839,16 +777,12 @@
        "id": "6J8bU-GmrV5i"
      },
      "source": [
-        "# Clean up"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "sXNEYNZiBcpj"
-      },
-      "source": [
-        "To clean up all the Google Cloud resources used in this project, delete the individual resources you created in this tutorial."
+        "## 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 cluster you created in this tutorial."
      ]
    },
    {
@@ -861,23 +795,23 @@
      "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 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",
        ")\n",
        "\n",
        "# Delete features\n",
        "for feature_id in FEATURE_IDS:\n",
        "    registry_client.delete_feature(\n",
-        "        name=f\"projects/{PROJECT_ID}/locations/{REGION}/featureGroups/{FEATURE_GROUP_ID}/features/{feature_id}\"\n",
+        "        name=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureGroups/{FEATURE_GROUP_ID}/features/{feature_id}\"\n",
        "    )\n",
        "\n",
        "# Delete feature group\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 test data\n",
@@ -4,6 +4,7 @@
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
+        "cellView": "form",
        "id": "9A9NkTRTfo2I"
      },
      "outputs": [],
@@ -31,24 +32,27 @@
      "source": [
        "# Getting Started with Claude 3 Models\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/generative_ai/anthropic_claude_3_intro.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",
-        "\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/generative_ai/anthropic_claude_3_intro.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%2Fblob%2Fmain%2Fnotebooks%2Fofficial%2Fgenerative_ai%2Fanthropic_claude_3_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>                                                                                               <td>\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/anthropic_claude_3_intro.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/generative_ai/anthropic_claude_3_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>"
      ]
    },
@@ -68,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",
@@ -80,7 +90,7 @@
        "\n",
        "This notebook shows how to use **Vertex AI API** and **Anthropic’s Vertex SDK for Python** to call the Claude on Vertex AI API with the Claude 3 Sonnet and Claude 3 Haiku model.\n",
        "\n",
-        "For more information, see the [Use Claude](https://cloud.devsite.corp.google.com/vertex-ai/generative-ai/docs/third-party-models/use-claude) documentation.\n"
+        "For more information, see the [Use Claude](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-claude) documentation.\n"
      ]
    },
    {
@@ -98,7 +108,95 @@
        "id": "nwYvaaW25jYS"
      },
      "source": [
-        "### Getting Started\n"
+        "## Get Started\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0660e339bf3f"
+      },
+      "source": [
+        "### Install required packages\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "754611260f53"
+      },
+      "outputs": [],
+      "source": [
+        "! pip3 install -U -q httpx"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "b9f4c57a43f6"
+      },
+      "source": [
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "3b9119a60525"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e767418763cd"
+      },
+      "source": [
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6a5bea26f60f"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
+        "\n",
+        "Authenticate your environment on Google Colab.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "c97be6a73155"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
      ]
    },
    {
@@ -118,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\"]"
      ]
    },
@@ -133,7 +233,7 @@
        "id": "bpuX3sKtexlK"
      },
      "source": [
-        "#### Select a region"
+        "#### Select a location"
      ]
    },
    {
@@ -149,20 +249,20 @@
        "\n",
        "dropdown = widgets.Dropdown(\n",
        "    options=available_regions,\n",
-        "    description=\"Select a region:\",\n",
+        "    description=\"Select a location:\",\n",
        "    font_weight=\"bold\",\n",
        "    style={\"description_width\": \"initial\"},\n",
        ")\n",
        "\n",
        "\n",
        "def dropdown_eventhandler(change):\n",
-        "    global REGION\n",
+        "    global LOCATION\n",
        "    if change[\"type\"] == \"change\" and change[\"name\"] == \"value\":\n",
-        "        REGION = change.new\n",
+        "        LOCATION = change.new\n",
        "        print(\"Selected:\", change.new)\n",
        "\n",
        "\n",
-        "REGION = dropdown.value\n",
+        "LOCATION = dropdown.value\n",
        "dropdown.observe(dropdown_eventhandler, names=\"value\")\n",
        "display(dropdown)"
      ]
@@ -173,7 +273,9 @@
        "id": "3q58icinBjoK"
      },
      "source": [
-        "#### Define Google Cloud project and model information"
+        "#### 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)."
      ]
    },
    {
@@ -185,41 +287,19 @@
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
-        "ENDPOINT = f\"https://{REGION}-aiplatform.googleapis.com\"\n",
+        "ENDPOINT = f\"https://{LOCATION}-aiplatform.googleapis.com\"\n",
        "\n",
        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
        "    raise ValueError(\"Please set your PROJECT_ID\")"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "hVi8v2mxBkeG"
-      },
-      "source": [
-        "#### Authenticate your notebook environment (Colab only)\n",
-        "If you are running this notebook on Google Colab, uncomment and run the following cell to authenticate your environment. This step is not required if you are using [Vertex AI Workbench](https://cloud.google.com/vertex-ai-workbench)."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "a4d82773077f"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user(project_id=PROJECT_ID)"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4NAstKRFBt4N"
      },
      "source": [
-        "#### Install and Import required module"
+        "#### Import required libraries"
      ]
    },
    {
@@ -230,8 +310,12 @@
      },
      "outputs": [],
      "source": [
-        "! pip3 install -U httpx\n",
-        "import json"
+        "import base64\n",
+        "import json\n",
+        "\n",
+        "import httpx\n",
+        "import requests\n",
+        "from IPython.display import Image"
      ]
    },
    {
@@ -249,7 +333,9 @@
        "id": "61107099357a"
      },
      "source": [
-        "#### Unary call"
+        "#### Unary call\n",
+        "\n",
+        "Sends a POST request to the specified API endpoint to get a response from the model for a banana bread recipe using the provided payload."
      ]
    },
    {
@@ -268,7 +354,7 @@
        "}\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/{REGION}/publishers/anthropic/models/{MODEL}:rawPredict -d '{request}'"
+        "!curl -X POST -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/publishers/anthropic/models/{MODEL}:rawPredict -d '{request}'"
      ]
    },
    {
@@ -277,7 +363,9 @@
        "id": "e6f52fae9379"
      },
      "source": [
-        "#### Streaming call"
+        "#### Streaming call\n",
+        "\n",
+        "Sends a POST request to the specified API endpoint to stream a response from the model for a banana bread recipe using the provided payload."
      ]
    },
    {
@@ -296,7 +384,7 @@
        "}\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/{REGION}/publishers/anthropic/models/{MODEL}:streamRawPredict -d '{request}'"
+        "!curl -X POST -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/publishers/anthropic/models/{MODEL}:streamRawPredict -d '{request}'"
      ]
    },
    {
@@ -316,7 +404,7 @@
      "source": [
        "#### Encode And Preview Image\n",
        "\n",
-        "We'll fetch sample images from Wikipedia using the httpx library, but you can use whatever image sources work for you."
+        "We fetch sample images from Wikipedia using the httpx library, but you can use whatever image sources work for you."
      ]
    },
    {
@@ -327,12 +415,6 @@
      },
      "outputs": [],
      "source": [
-        "import base64\n",
-        "\n",
-        "import httpx\n",
-        "import requests\n",
-        "from IPython.display import Image\n",
-        "\n",
        "image_url = \"https://upload.wikimedia.org/wikipedia/commons/thumb/a/a7/Camponotus_flavomarginatus_ant.jpg/300px-Camponotus_flavomarginatus_ant.jpg\"\n",
        "image_b64 = base64.b64encode(httpx.get(image_url).content).decode(\"utf-8\")\n",
        "\n",
@@ -348,7 +430,9 @@
        "id": "9771877508aa"
      },
      "source": [
-        "#### Unary call"
+        "#### Unary call\n",
+        "\n",
+        "Sends a POST request to the specified API endpoint to get a response from the model analyzing the content of an image, provided as base64-encoded data, along with the text prompt."
      ]
    },
    {
@@ -382,7 +466,7 @@
        "}\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/{REGION}/publishers/anthropic/models/{MODEL}:rawPredict -d '{request}'"
+        "!curl -X POST -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/publishers/anthropic/models/{MODEL}:rawPredict -d '{request}'"
      ]
    },
    {
@@ -391,7 +475,9 @@
        "id": "2ecd941bf6cc"
      },
      "source": [
-        "#### Streaming call"
+        "#### Streaming call\n",
+        "\n",
+        "Sends a POST request to the specified API endpoint to stream a response from the model analyzing the content of an image, provided as base64-encoded data, along with the text prompt."
      ]
    },
    {
@@ -425,7 +511,7 @@
        "}\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/{REGION}/publishers/anthropic/models/{MODEL}:streamRawPredict -d '{request}'"
+        "!curl -X POST -H \"Authorization: Bearer $(gcloud auth print-access-token)\" -H \"Content-Type: application/json\" {ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION}/publishers/anthropic/models/{MODEL}:streamRawPredict -d '{request}'"
      ]
    },
    {
@@ -443,7 +529,7 @@
        "id": "E0x3GO6M_O3_"
      },
      "source": [
-        "### Getting Started\n"
+        "## Get Started\n"
      ]
    },
    {
@@ -452,7 +538,7 @@
        "id": "_CJrqUvqAfR7"
      },
      "source": [
-        "#### Install Anthropic's Vertex SDK and required dependencies"
+        "### Install Anthropic's Vertex SDK for Python and other required packages"
      ]
    },
    {
@@ -463,8 +549,8 @@
      },
      "outputs": [],
      "source": [
-        "! pip3 install -U 'anthropic[vertex]'\n",
-        "! pip3 install -U httpx"
+        "! pip3 install -U -q 'anthropic[vertex]'\n",
+        "! pip3 install -U -q httpx"
      ]
    },
    {
@@ -473,9 +559,9 @@
        "id": "hUiAYUFbBCpR"
      },
      "source": [
-        "#### Restart current runtime\n",
+        "### Restart runtime (Colab only)\n",
        "\n",
-        "To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which will restart the current kernel."
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
@@ -496,6 +582,45 @@
        "    app.kernel.do_shutdown(True)"
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aa9169957a89"
+      },
+      "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": "bffc70d3e8be"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
+        "\n",
+        "Authenticate your environment on Google Colab.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "d1baa068f48a"
+      },
+      "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": {
@@ -513,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\"]"
      ]
    },
@@ -544,20 +671,20 @@
        "\n",
        "dropdown = widgets.Dropdown(\n",
        "    options=available_regions,\n",
-        "    description=\"Select a region:\",\n",
+        "    description=\"Select a location:\",\n",
        "    font_weight=\"bold\",\n",
        "    style={\"description_width\": \"initial\"},\n",
        ")\n",
        "\n",
        "\n",
        "def dropdown_eventhandler(change):\n",
-        "    global REGION\n",
+        "    global LOCATION\n",
        "    if change[\"type\"] == \"change\" and change[\"name\"] == \"value\":\n",
-        "        REGION = change.new\n",
+        "        LOCATION = change.new\n",
        "        print(\"Selected:\", change.new)\n",
        "\n",
        "\n",
-        "REGION = dropdown.value\n",
+        "LOCATION = dropdown.value\n",
        "dropdown.observe(dropdown_eventhandler, names=\"value\")\n",
        "display(dropdown)"
      ]
@@ -568,7 +695,9 @@
        "id": "shZgRl6qbZYP"
      },
      "source": [
-        "#### Define Google Cloud project and model information"
+        "#### 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)."
      ]
    },
    {
@@ -580,7 +709,7 @@
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
-        "ENDPOINT = f\"https://{REGION}-aiplatform.googleapis.com\"\n",
+        "ENDPOINT = f\"https://{LOCATION}-aiplatform.googleapis.com\"\n",
        "\n",
        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
        "    raise ValueError(\"Please set your PROJECT_ID\")"
@@ -589,25 +718,25 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "30K5Hs_hBNEw"
+        "id": "a6a543c2cd07"
      },
      "source": [
-        "#### Authenticate your notebook environment\n",
-        "If you are running this notebook on Google Colab, uncomment run the following cell to authenticate your environment.\n",
-        "\n",
-        "Note: To authenticate for the SDK, you must provide `project_id` parameter. See the example command below."
+        "#### Import required libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "tvD3KC7z_lh4"
+        "id": "c4a553b6a5c2"
      },
      "outputs": [],
      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user(project_id=PROJECT_ID)"
+        "import base64\n",
+        "\n",
+        "import httpx\n",
+        "import requests\n",
+        "from IPython.display import Image"
      ]
    },
    {
@@ -625,7 +754,9 @@
        "id": "2f7ee8ceb620"
      },
      "source": [
-        "#### Unary call"
+        "#### Unary call\n",
+        "\n",
+        "Initializes a client for Anthropic's Vertex AI, sends a request to generate the content, and prints the response in a formatted JSON"
      ]
    },
    {
@@ -638,7 +769,7 @@
      "source": [
        "from anthropic import AnthropicVertex\n",
        "\n",
-        "client = AnthropicVertex(region=REGION, project_id=PROJECT_ID)\n",
+        "client = AnthropicVertex(region=LOCATION, project_id=PROJECT_ID)\n",
        "message = client.messages.create(\n",
        "    max_tokens=1024,\n",
        "    messages=[\n",
@@ -658,7 +789,9 @@
        "id": "f8e56a2fb209"
      },
      "source": [
-        "#### Streaming call"
+        "#### Streaming call\n",
+        "\n",
+        "Initializes a client for Anthropic's Vertex AI, sends a streaming request to generate the content, and continuously prints the received text as it is streamed."
      ]
    },
    {
@@ -671,7 +804,7 @@
      "source": [
        "from anthropic import AnthropicVertex\n",
        "\n",
-        "client = AnthropicVertex(region=REGION, project_id=PROJECT_ID)\n",
+        "client = AnthropicVertex(region=LOCATION, project_id=PROJECT_ID)\n",
        "\n",
        "with client.messages.stream(\n",
        "    max_tokens=1024,\n",
@@ -704,7 +837,7 @@
      "source": [
        "#### Encode And Preview Image\n",
        "\n",
-        "We'll fetch sample images from Wikipedia using the httpx library, but you can use whatever image sources work for you."
+        "We fetch sample images from Wikipedia using the httpx library, but you can use whatever image sources work for you."
      ]
    },
    {
@@ -715,12 +848,6 @@
      },
      "outputs": [],
      "source": [
-        "import base64\n",
-        "\n",
-        "import httpx\n",
-        "import requests\n",
-        "from IPython.display import Image\n",
-        "\n",
        "image_url = \"https://upload.wikimedia.org/wikipedia/commons/thumb/a/a7/Camponotus_flavomarginatus_ant.jpg/300px-Camponotus_flavomarginatus_ant.jpg\"\n",
        "image_media_type = \"image/jpeg\"\n",
        "image_b64 = base64.b64encode(httpx.get(image_url).content).decode(\"utf-8\")\n",
@@ -737,7 +864,9 @@
        "id": "9fd5937f839b"
      },
      "source": [
-        "#### Unary call"
+        "#### Unary call\n",
+        "\n",
+        "Initializes a client for Anthropic's Vertex AI, sends a request to describe an image (provided as base64-encoded data) along with a text prompt, and prints the response in a formatted JSON."
      ]
    },
    {
@@ -750,7 +879,7 @@
      "source": [
        "from anthropic import AnthropicVertex\n",
        "\n",
-        "client = AnthropicVertex(region=REGION, project_id=PROJECT_ID)\n",
+        "client = AnthropicVertex(region=LOCATION, project_id=PROJECT_ID)\n",
        "\n",
        "message = client.messages.create(\n",
        "    max_tokens=1024,\n",
@@ -781,7 +910,9 @@
        "id": "1fb4855047e3"
      },
      "source": [
-        "#### Streaming call"
+        "#### Streaming call\n",
+        "\n",
+        "Initializes a client for Anthropic's Vertex AI, sends a streaming request to describe an image (provided as base64-encoded data) along with a text prompt, and continuously prints the received text as it is streamed."
      ]
    },
    {
@@ -794,7 +925,7 @@
      "source": [
        "from anthropic import AnthropicVertex\n",
        "\n",
-        "client = AnthropicVertex(region=REGION, project_id=PROJECT_ID)\n",
+        "client = AnthropicVertex(region=LOCATION, project_id=PROJECT_ID)\n",
        "\n",
        "with client.messages.stream(\n",
        "    max_tokens=1024,\n",
@@ -4,6 +4,7 @@
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
+        "cellView": "form",
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
@@ -33,21 +34,26 @@
        "# Vertex AI LLM Batch Inference with RLHF-tuned Models\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
+        "    <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.sandbox.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/generative_ai/batch_eval_llm.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/batch_eval_llm.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%2Fbatch_eval_llm.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" 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/generative_ai/batch_eval_llm.ipynb\">\n",
-        "       <img src=\"https://www.gstatic.com/cloud/images/navigation/vertex-ai.svg\" alt=\"Vertex AI logo\">Open in Vertex AI Workbench\n",
+        "      <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/batch_eval_llm.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>"
      ]
    },
@@ -74,7 +80,7 @@
      "source": [
        "### Objective\n",
        "\n",
-        "In this tutorial, you will use Vertex AI to get predictions from an RLHF-tuned large-language model.\n",
+        "In this tutorial, you use Vertex AI to get predictions from an RLHF-tuned large-language model.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
@@ -116,13 +122,13 @@
        "> Datasets for chat models must contain at least 1 message in a `messages` field.\n",
        "> * Each message must be valid JSON that contains `author` and `content` fields, where valid `author` values are `user` and `assistant` and `content` must be non-empty.\n",
        "> * Each row may contain multiple messages, but the first and last author must be the `user`.\n",
-        "> * An optional `context` field may be provided for each example in a chat dataset. If provided, the `context` will preprended to the message `content`.\n",
-        "> * The `instruction` serves as the default context. (Useful if most messages use the same system-level context.) Any context provided in the example will override the default value.\n",
+        "> * An optional `context` field may be provided for each example in a chat dataset. If provided, the `context` is preprended to the message `content`.\n",
+        "> * The `instruction` serves as the default context. (Useful if most messages use the same system-level context.) Any context provided in the example overrides the default value.\n",
        "> * For example: `{\"context\": \"I am a helpful assistant that can answer questions about Plantation Palms.\", \"messages\": [{\"author\": \"user\", \"content\": \"Hello\"}, {\"author\": \"assistant\", \"content\": \"Hello, how can I help you?\"}, {\"author\": \"user\", \"content\": \"Tell me about Plantation Palms.\"}]}`\n",
        "\n",
        "* **large_model_reference**: Required, `str`. Name of the base model. In this example, we use `t5-small`. The valid values are listed in the table above.\n",
        "\n",
-        "* **model_checkpoint**. Optional, `str`. The GCS path to the model's checkpoint. It should be a derivative of the model specified by **large_model_reference**. This can be output of an RLHF training pipeline. The tuned model is stored at the **Reinforcer** > **Output Parameters** > `output_model_path`. If unspecified, the base (i.e. untuned) model corresponding to **large_model_reference** will be used.\n",
+        "* **model_checkpoint**. Optional, `str`. The GCS path to the model's checkpoint. It should be a derivative of the model specified by **large_model_reference**. This can be output of an RLHF training pipeline. The tuned model is stored at the **Reinforcer** > **Output Parameters** > `output_model_path`. If unspecified, the base (i.e. untuned) model corresponding to **large_model_reference** is used.\n",
        "\n",
        "* **prompt_sequence_length**: Optional, `int`. Maximum tokenized sequence length for input text. Higher values increase memory overhead. This value should be at most `8192`. Default value is `512`.\n",
        "\n",
@@ -130,7 +136,7 @@
        "\n",
        "* **sampling_strategy**: Optional, `str`. This field specifies the sampling strategy. The valid options are `'greedy'` and `'temperature_sampling'`. Default is `'greedy'`.\n",
        "\n",
-        "* **instruction**: Optional, `str`. This field lets the model know what task it needs to perform. Base models have been trained over a large set of varied instructions. You can give a simple and intuitive description of the task and the model will follow it, e.g. `Classify this movie review as positive or negative` or `Translate this sentence to Danish`. See [here](https://ai.googleblog.com/2021/10/introducing-flan-more-generalizable.html) for more details on the instruction-tuned models. Do not specify this if your dataset already prepends the instruction to the inputs field.\n",
+        "* **instruction**: Optional, `str`. This field lets the model know what task it needs to perform. Base models have been trained over a large set of varied instructions. You can give a simple and intuitive description of the task and the model follows it, e.g. `Classify this movie review as positive or negative` or `Translate this sentence to Danish`. See [here](https://ai.googleblog.com/2021/10/introducing-flan-more-generalizable.html) for more details on the instruction-tuned models. Do not specify this if your dataset already prepends the instruction to the inputs field.\n",
        "\n",
        "* **accelerator_type**: Optional, `str`. One of `'TPU'` or `'GPU'`. If `'TPU'` is specified, tuning components run in `europe-west4` on 64 v3 TPUs. Otherwise tuning components run in `us-central1` on 8 Nvidia A100 80GB. Default is `'GPU'`."
      ]
@@ -155,15 +161,22 @@
        "to generate a cost estimate based on your projected usage."
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3b1ffd5ab768"
+      },
+      "source": [
+        "## Getting Started"
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "install_aip:mbsdk"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook."
+        "### Install Vertex AI SDK and other required packages\n"
      ]
    },
    {
@@ -182,25 +195,69 @@
    {
      "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": "54c5ef8a8f43"
+      },
+      "source": [
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Please wait until it is finished before continuing to the next step. ⚠️</b>\n",
+        "</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()"
      ]
    },
    {
@@ -209,14 +266,9 @@
        "id": "before_you_begin:nogpu"
      },
      "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)"
+        "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)."
      ]
    },
    {
@@ -228,91 +280,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": "2dw8q9fdQEH5"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"europe-west4\""
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "gcp_authenticate"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may need to authenticate manually. Follow the relevant instructions below.\n",
-        "\n",
-        "#### Vertex AI Workbench\n",
-        "Do nothing as you are already authenticated.\n",
-        "\n",
-        "#### Local JupyterLab instance\n",
-        "\n",
-        "**1. Uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ce6043da7b33"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0367eac06a10"
-      },
-      "source": [
-        "**2. 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": [
-        "**3. 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 = \"europe-west4\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -354,7 +322,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
@@ -438,10 +406,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 and initialize Vertex AI SDK"
      ]
    },
    {
@@ -452,8 +417,6 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
        "import google.cloud.aiplatform as aiplatform\n",
        "from google_cloud_pipeline_components.preview.llm import infer_pipeline\n",
        "from kfp import compiler"
@@ -465,8 +428,6 @@
        "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."
      ]
    },
@@ -478,7 +439,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)"
      ]
    },
    {
@@ -487,7 +448,7 @@
        "id": "SaMeaYNLH5L9"
      },
      "source": [
-        "## Compile the Inference Pipeline"
+        "## Compile the inference pipeline"
      ]
    },
    {
@@ -496,7 +457,7 @@
        "id": "6CIl3CcJH8hR"
      },
      "source": [
-        "Compile the pipeline into a YAML file that will be submitted to Vertex AI."
+        "Compile the pipeline into a YAML file that is submitted to Vertex AI."
      ]
    },
    {
@@ -520,12 +481,12 @@
        "id": "R_Gy4jW6wFAT"
      },
      "source": [
-        "## Construct the Pipeline Job and Run on Vertex AI\n",
+        "## Construct the pipeline job and run on Vertex AI\n",
        "\n",
-        "Define a pipeline job with the following code, which will:\n",
+        "Define a pipeline job with the following code, which:\n",
        "\n",
-        "- load the pipeline template that was compiled in the previous step, and\n",
-        "- set the pipeline parameters. \n"
+        "- loads the pipeline template that was compiled in the previous step, and\n",
+        "- sets the pipeline parameters. \n"
      ]
    },
    {
@@ -594,7 +555,7 @@
      "source": [
        "## View the inference result\n",
        "\n",
-        "Go to \"Vertex AI Pipelines\" in the Google Cloud UI to view the pipeline job. When the pipeline completes, click on **Bulk Inferrer** > **Output Parameters** > `output_prediction_gcs_path`. If you encounter \"❗ Requested entity was not found\", this is a known issue. Hit the backarrow at the top of the page, you will see your inference results there."
+        "Go to \"Vertex AI Pipelines\" in the Google Cloud UI to view the pipeline job. When the pipeline completes, click on **Bulk Inferrer** > **Output Parameters** > `output_prediction_gcs_path`. If you encounter \"❗ Requested entity was not found\", this is a known issue. Hit the backarrow at the top of the page to see your inference results."
      ]
    },
    {
@@ -632,7 +593,7 @@
        "\n",
        "job.delete()\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil rm -rf {BUCKET_URI}\n",
        "\n",
        "! rm infer_pipeline.yaml"
@@ -31,40 +31,30 @@
      "source": [
        "# Semantic Search using Embeddings\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/generative_ai/text_embedding_api_semantic_search_with_scann.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/text_embedding_api_semantic_search_with_scann.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%2Ftext_embedding_api_semantic_search_with_scann.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/text_embedding_api_semantic_search_with_scann.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/text_embedding_api_semantic_search_with_scann.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": {
@@ -149,169 +139,128 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "BF1j6f9HApxa"
+        "id": "7cc240e9c1a6"
      },
      "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)."
+        "## Get Started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "WReHDGG5g0XY"
+        "id": "a6ab0747c5e6"
      },
      "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)"
+        "### Install Vertex AI SDK for Python and other required packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "oM1iC_MfAts1"
+        "id": "bb3fc5b6c2c3"
+      },
+      "outputs": [],
+      "source": [
+        "!pip3 install google-cloud-aiplatform \"shapely<2.0.0\" --quiet\n",
+        "!pip install scann --quiet"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "16220914acc5"
+      },
+      "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": "157953ab28f0"
+      },
+      "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": "e669f8088ac3"
+      },
+      "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": "ff666ce4051c"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
+        "\n",
+        "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\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": "0e1a068fcb88"
      },
      "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": "RfLtjhGNmKvA"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "sBCra4QMA2wR"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
+        "import vertexai\n",
        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "74ccc9e52986"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "de775a3773ba"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "254614fa0c46"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ef21552ccea8"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "603adbbf0532"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "f6b2ccc891ed"
-      },
-      "source": [
-        "**4. Service account or other**\n",
-        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "FyyMdUeAJIVv"
-      },
-      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook.\n",
-        "\n",
-        "Remember to restart the runtime after installation."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "snBUuUamoJPz"
-      },
-      "outputs": [],
-      "source": [
-        "!pip3 install google-cloud-aiplatform>=1.25 \"shapely<2.0.0\" --quiet\n",
-        "!pip install scann --quiet"
+        "# Initiate Vertex AI\n",
+        "vertexai.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
@@ -337,20 +286,16 @@
        "import numpy as np\n",
        "import pandas as pd\n",
        "import scann\n",
-        "import vertexai\n",
        "from vertexai.preview.language_models import TextEmbeddingModel"
      ]
    },
    {
-      "cell_type": "code",
-      "execution_count": null,
+      "cell_type": "markdown",
      "metadata": {
-        "id": "zgVQcE0ewO8W"
+        "id": "d1a62cd003f5"
      },
-      "outputs": [],
      "source": [
-        "# Initiate Vertex AI\n",
-        "vertexai.init(project=PROJECT_ID, location=REGION)"
+        "Initialize a text embedding model using a pretrained model from Google."
      ]
    },
    {
@@ -370,7 +315,7 @@
        "id": "f50f22f3-ec85-463e-b6fe-5c8e6b80b07b"
      },
      "source": [
-        "## Create embedding dataset\n",
+        "### Create embedding dataset\n",
        "\n",
        "The dataset demonstrates the use of the Text Embedding API with a vector database. It is not intended to be used for any other purpose, such as evaluating models. The dataset is small and does not represent a comprehensive sample of all possible text."
      ]
@@ -459,7 +404,7 @@
        "id": "ba4f49b6-65e1-49ea-988b-c4e195219deb"
      },
      "source": [
-        "## Create an index"
+        "### Create an index"
      ]
    },
    {
@@ -498,7 +443,7 @@
        "id": "9f1689eb-c27c-4566-afc7-7fbc55552aad"
      },
      "source": [
-        "## Query the index\n",
+        "### Query the index\n",
        "This is a good example of how to use the ScaNN library to perform approximate nearest neighbor search. The function takes a query string as input and returns the top 3 neighbors of the query. The function is efficient and can be used to search large datasets quickly."
      ]
    },
@@ -32,38 +32,29 @@
        "# Text Embedding New API\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/text_embedding_new_api.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/text_embedding_new_api.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%2Ftext_embedding_new_api.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/text_embedding_new_api.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/text_embedding_new_api.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": {
@@ -118,19 +109,26 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "i7EUnXsZhAGF"
+        "id": "8925ff9e165e"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook."
+        "## Get Started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6dac73aa3f75"
+      },
+      "source": [
+        "### Install Vertex AI SDK for python and other required packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "2b4ef9b72d43"
+        "id": "6c944350c059"
      },
      "outputs": [],
      "source": [
@@ -141,181 +139,117 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "58707a750154"
+        "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": "f200f10a1da3"
+        "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": "BF1j6f9HApxa"
+        "id": "7b49231643e4"
      },
      "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>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "WReHDGG5g0XY"
+        "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": "oM1iC_MfAts1"
+        "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": "e487ec618b5e"
+      },
+      "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": "48d81914dc54"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
        "\n",
-        "!gcloud config set project {PROJECT_ID}"
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "region"
+        "id": "09720c707f1c"
      },
      "source": [
-        "#### Set your region\n",
-        "\n",
-        "You can also set the `REGION` variable to use with Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
+        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "CWUrmELRY_8E"
+        "id": "a38ef56b684a"
      },
      "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(project_id=PROJECT_ID)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "init_aip:mbsdk,all"
-      },
-      "source": [
-        "### Initialize Vertex AI SDK for Python\n",
-        "\n",
-        "Initialize the Vertex AI SDK for Python for your project and region."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "d6gh52v1Y_8E"
-      },
-      "outputs": [],
-      "source": [
-        "import vertexai\n",
-        "\n",
-        "vertexai.init(project=PROJECT_ID, location=REGION)"
+        "from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel"
      ]
    },
    {
@@ -390,9 +324,6 @@
      },
      "outputs": [],
      "source": [
-        "from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel\n",
-        "\n",
-        "\n",
        "def embed_text(\n",
        "    model_name: str,\n",
        "    task_type: str,\n",
@@ -431,7 +362,7 @@
        "id": "TpV-iwP9qw9c"
      },
      "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."
@@ -53,17 +53,6 @@
        "</table>"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ItngP1vsKGno"
-      },
-      "source": [
-        "**_NOTE_**: This notebook has been tested in the following environment:\n",
-        "\n",
-        "* Python version = 3.9"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -295,7 +284,8 @@
      },
      "source": [
        "**2. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
+        "\n",
+        "Make sure that [the Compute Engine default service account](https://cloud.google.com/compute/docs/access/service-accounts#default_service_account) running a Vertex AI workbench instance has the permission iam.serviceAccounts.actAs (most likely through [roles/iam.serviceAccountUser](https://cloud.google.com/iam/docs/understanding-roles#iam.serviceAccountUser)) at [the IAM & Admin page of the Cloud Console](https://console.cloud.google.com/iam-admin). This permission allows a workbench instance to act as the service account when interacting with other Google Cloud services."
      ]
    },
    {
@@ -33,37 +33,29 @@
        "\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/vertex_sdk_llm_snippets.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/vertex_sdk_llm_snippets.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%2Fgenerative_ai%2Fvertex_sdk_llm_snippets.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" 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/generative_ai/vertex_sdk_llm_snippets.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/generative_ai/vertex_sdk_llm_snippets.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": "962e636b5cee"
-      },
-      "source": [
-        "**_NOTE_**: This notebook has been tested in the following environment:\n",
-        "\n",
-        "* Python version = 3.9"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -72,7 +64,7 @@
      "source": [
        "## Overview\n",
        "\n",
-        "This tutorial demonstrates how to use the Vertex AI SDK to run Large Language Models on Vertex AI via the PaLM API. You will find sample code to test, tune, and deploy generative AI language models. Get started by exploring examples of content summarization, sentiment snalysis, and chat, as well as text embedding and prompt tuning. \n",
+        "This tutorial demonstrates how to use the Vertex AI SDK to run Large Language Models on Vertex AI via the PaLM API. You find sample code to test, tune, and deploy generative AI language models. Get started by exploring examples of content summarization, sentiment snalysis, and chat, as well as text embedding and prompt tuning. \n",
        "\n",
        "Learn more about [PaLM API](https://ai.google/discover/palm2/)."
      ]
@@ -114,15 +106,22 @@
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing)"
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8925ff9e165e"
+      },
+      "source": [
+        "## Getting Started"
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e150bf471f1b"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook. "
+        "### Install Vertex AI SDK and other required packages\n"
      ]
    },
    {
@@ -134,7 +133,7 @@
      "outputs": [],
      "source": [
        "# Install the packages\n",
-        "! pip3 install --upgrade --quiet google-cloud-aiplatform"
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform \"shapely<2.0.0\""
      ]
    },
    {
@@ -143,7 +142,9 @@
        "id": "d98bc9fdd80d"
      },
      "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."
      ]
    },
    {
@@ -154,46 +155,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": "8547fbbb241c"
+        "id": "7b49231643e4"
      },
      "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. Please wait until it is finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "8bc8a29f9001"
+        "id": "7176ea64999b"
      },
      "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": "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": "e487ec618b5e"
+      },
+      "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)."
      ]
    },
    {
@@ -205,126 +224,11 @@
      "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": "c4a624c8099d"
-      },
-      "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": "f83bd6013894"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "08bfd1eb44ef"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "af349043f23b"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ad1138a125ea"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ce6043da7b33"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0367eac06a10"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 1,
-      "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": {
-        "id": "b2fd1f2fd670"
-      },
-      "source": [
-        "## Installation\n",
-        "Install the following packages required to execute this notebook.\n",
-        "\n",
-        "Remember to restart the runtime after installation."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "41df9761a7ea"
-      },
-      "outputs": [],
-      "source": [
-        "!pip3 install google-cloud-aiplatform>=1.25 \"shapely<2.0.0\" --quiet"
+        "vertexai.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
@@ -336,17 +240,6 @@
        "### Import libraries"
      ]
    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "4b9ee6efcf5c"
-      },
-      "outputs": [],
-      "source": [
-        "import pandas as pd"
-      ]
-    },
    {
      "cell_type": "code",
      "execution_count": null,
@@ -355,41 +248,19 @@
      },
      "outputs": [],
      "source": [
-        "import vertexai\n",
+        "import pandas as pd\n",
        "from vertexai.preview.language_models import (ChatModel, InputOutputTextPair,\n",
        "                                              TextEmbeddingModel,\n",
        "                                              TextGenerationModel)"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0ae1b641d63f"
-      },
-      "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": "4efe77faaf16"
-      },
-      "outputs": [],
-      "source": [
-        "vertexai.init(project=PROJECT_ID, location=REGION)"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mOabLACbseoE"
      },
      "source": [
-        "# Summarization examples: transcript summarization"
+        "### Summarization examples: transcript summarization"
      ]
    },
    {
@@ -447,7 +318,7 @@
        "id": "AdxJX2dNE7t3"
      },
      "source": [
-        "# Classification examples: classification headline"
+        "### Classification examples: classification headline"
      ]
    },
    {
@@ -499,7 +370,7 @@
        "id": "fKsuhKli70q-"
      },
      "source": [
-        "# Classification examples: sentiment analysis"
+        "### Classification examples: sentiment analysis"
      ]
    },
    {
@@ -563,7 +434,7 @@
        "id": "Dh0tLRggE5H1"
      },
      "source": [
-        "# Extraction examples: extractive question answering\n"
+        "### Extraction examples: extractive question answering\n"
      ]
    },
    {
@@ -625,7 +496,7 @@
        "id": "vV08j9H-Rbr7"
      },
      "source": [
-        "# Ideation examples: interview questions"
+        "### Ideation examples: interview questions"
      ]
    },
    {
@@ -653,7 +524,7 @@
        "id": "tpeXOD_bsCAX"
      },
      "source": [
-        "# Chat examples: science tutoring"
+        "### Chat examples: science tutoring"
      ]
    },
    {
@@ -702,7 +573,7 @@
        "id": "ag5Og6Y3xfvM"
      },
      "source": [
-        "# Text embedding"
+        "### Text embedding"
      ]
    },
    {
@@ -726,7 +597,7 @@
        "id": "CI-dpBTr6LGH"
      },
      "source": [
-        "# List tuned models"
+        "### List tuned models"
      ]
    },
    {
@@ -748,7 +619,7 @@
        "id": "gYvbJ9PrFF2g"
      },
      "source": [
-        "# Tune a model"
+        "### Tune a model"
      ]
    },
    {
@@ -32,24 +32,26 @@
        "# Vertex AI Pipelines: Evaluating BatchPrediction results from a custom tabular classification 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/model_evaluation/custom_tabular_classification_model_evaluation.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%2Fmodel_evaluation%2Fcustom_tabular_classification_model_evaluation.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/model_evaluation/custom_tabular_classification_model_evaluation.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/model_evaluation/custom_tabular_classification_model_evaluation.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/model_evaluation/custom_tabular_classification_model_evaluation.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>"
      ]
    },
@@ -63,7 +65,7 @@
        "\n",
        "This notebook demonstrates how to use the Vertex AI classification model evaluation component to evaluate a custom-trained tabular classification model saved in Vertex AI Model Registry. Model evaluation helps you determine your model performance based on the evaluation metrics and improve the model if necessary. \n",
        "\n",
-        "Learn more about [Vertex AI Model Evaluation](https://cloud.google.com/vertex-ai/docs/evaluation/introduction) and [Vertex AI Training](https://cloud.google.com/vertex-ai/docs/training/custom-training)."
+        "Learn more about [Vertex AI custom training](https://cloud.google.com/vertex-ai/docs/training/custom-training) and [Vertex AI model evaluation](https://cloud.google.com/vertex-ai/docs/evaluation/introduction)."
      ]
    },
    {
@@ -74,9 +76,9 @@
      "source": [
        "### Objective\n",
        "\n",
-        "In this tutorial, you train a scikit-learn RandomForest model, save it in Vertex AI Model Registry and learn how to evaluate it through a Vertex AI pipeline job using `google_cloud_pipeline_components`.\n",
+        "In this tutorial, you train a scikit-learn RandomForest model, save the model in Vertex AI Model Registry and learn how to evaluate the model through a Vertex AI pipeline job using Google Cloud Pipeline Components Python SDK.\n",
        "\n",
-        "This tutorial uses the following Google Cloud ML services and resources:\n",
+        "This tutorial uses the following Vertex AI services and resources:\n",
        "\n",
        "- Vertex AI Model Registry\n",
        "- Vertex AI Pipelines\n",
@@ -134,19 +136,26 @@
    {
      "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"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "2b4ef9b72d43"
+        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
@@ -166,168 +175,91 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "hhq5zEbGg0XX"
+        "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": "EzrelQZ22IZj"
+        "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",
-        "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, Compute Engine, Artifact Registry, Cloud Build and Dataflow APIs](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute.googleapis.com,artifactregistry.googleapis.com,cloudbuild.googleapis.com,dataflow.googleapis.com).\n",
-        "\n",
-        "1. 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",
+        "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": "sduDOFQVF6kv"
-      },
-      "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": "af349043f23b"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ad1138a125ea"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ce6043da7b33"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0367eac06a10"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "21ad4dbb4a61"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c13224697bfb"
-      },
-      "source": [
-        "**4. Service account or other**\n",
-        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -349,7 +281,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\"}"
      ]
    },
    {
@@ -358,7 +290,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."
      ]
    },
    {
@@ -369,7 +301,32 @@
      },
      "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": "init_aip:mbsdk,all"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and enable the [Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). \n",
+        "Initialize the Vertex AI SDK for Python with your project and the created bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ksAefQcCF6ky"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -378,9 +335,9 @@
        "id": "set_service_account"
      },
      "source": [
-        "#### Service Account\n",
+        "### Service Account\n",
        "\n",
-        "You use a service account to create Vertex AI Pipeline jobs. If you do not want to use your project's Compute Engine service account, set `SERVICE_ACCOUNT` to another service account ID."
+        "You use a service account to create Vertex AI Pipeline jobs. If you don't want to use your project's Compute Engine service account, set `SERVICE_ACCOUNT` to another service account ID."
      ]
    },
    {
@@ -469,7 +426,7 @@
        "import joblib\n",
        "import kfp\n",
        "import matplotlib.pyplot as plt\n",
-        "from google.cloud import aiplatform, aiplatform_v1, bigquery\n",
+        "from google.cloud import aiplatform_v1, bigquery\n",
        "from kfp.v2 import compiler\n",
        "from sklearn.ensemble import RandomForestClassifier\n",
        "from sklearn.feature_selection import SelectKBest\n",
@@ -484,9 +441,9 @@
        "id": "init_aip:mbsdk,all"
      },
      "source": [
-        "### Initialize Vertex AI and BigQuery SDK for Python\n",
+        "### Initialize BigQuery SDK for Python\n",
        "\n",
-        "Initialize the Vertex AI and BigQuery SDK for Python with your project and the created bucket."
+        "Initialize the BigQuery SDK for Python with your project and the created bucket."
      ]
    },
    {
@@ -497,8 +454,6 @@
      },
      "outputs": [],
      "source": [
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)\n",
-        "\n",
        "bq_client = bigquery.Client(\n",
        "    project=PROJECT_ID,\n",
        "    credentials=aiplatform.initializer.global_config.credentials,\n",
@@ -1098,7 +1053,7 @@
      "source": [
        "!gcloud artifacts repositories create {REPOSITORY} \\\n",
        "    --repository-format=docker \\\n",
-        "    --location=$REGION"
+        "    --location=$LOCATION"
      ]
    },
    {
@@ -1120,7 +1075,7 @@
      "outputs": [],
      "source": [
        "%cd $SRC_DIR/\n",
-        "!gcloud builds submit --region={REGION} --tag={REGION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE} --suppress-logs\n",
+        "!gcloud builds submit --region={LOCATION} --tag={LOCATION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE} --suppress-logs\n",
        "%cd .."
      ]
    },
@@ -1154,7 +1109,7 @@
        "aip_model = aiplatform.Model.upload(\n",
        "    display_name=MODEL_DISPLAY_NAME,\n",
        "    artifact_uri=f\"{BUCKET_URI}/{MODEL_ARTIFACT_DIR}\",\n",
-        "    serving_container_image_uri=f\"{REGION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE}\",\n",
+        "    serving_container_image_uri=f\"{LOCATION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE}\",\n",
        "    serving_container_predict_route=\"/predict\",\n",
        "    serving_container_health_route=\"/health\",\n",
        ")"
@@ -1163,7 +1118,7 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "77f4838ece15"
+        "id": "7f463f9183ef"
      },
      "source": [
        "## Create and run the evaluation pipeline\n",
@@ -1174,9 +1129,37 @@
        "1. Removes the target field from the sampled test data.\n",
        "1. Runs the batch prediction job.\n",
        "1. Evaluates results from the batch prediction job using the ground-truth/target information.\n",
-        "1. Imports the generated evaluation metrics to the Vertex AI model.\n",
-        "\n",
+        "1. Imports the generated evaluation metrics to the Vertex AI model."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9f72ff491712"
+      },
+      "source": [
+        "### Enable the Dataflow API\n",
        "\n",
+        "The model evaluation components from the Google Cloud Pipeline Components create dataflow jobs internally to perform the underlying tasks. It's essential to enable the Dataflow API before running the evaluation pipeline."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "795fdfef5736"
+      },
+      "outputs": [],
+      "source": [
+        "!gcloud services enable dataflow.googleapis.com"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "77f4838ece15"
+      },
+      "source": [
        "### Define the pipeline\n",
        "To define the Vertex AI Pipeline for evaluating your model, you use the `google-cloud-pipeline-components` Python package. Google Cloud Pipeline Components provides an SDK with a set of pipeline components for users to interact with Google Cloud services such as Vertex AI, Dataflow and BigQuery. \n",
        "\n",
@@ -1378,7 +1361,7 @@
        "Before running your pipeline, set the following parameters :\n",
        "\n",
        "- `project`: Project ID of the Google Cloud project.\n",
-        "- `location`: Region where the pipeline needs to be run. If not set, the pipeline defaults to the region that Vertex AI SDK is configured with.\n",
+        "- `location`: Location where the pipeline needs to be run. If not set, the pipeline defaults to the location that Vertex AI SDK is configured with.\n",
        "- `root_dir`: The Cloud Storage directory for keeping the staged files and artifacts. A random subdirectory is created under the directory to keep the job information for resuming the job in case of a failure.\n",
        "- `model_name`: Resource name of the trained custom tabular classification model.\n",
        "- `target_field_name`: Name of the column to be used as the ground truth for evaluation.\n",
@@ -1400,7 +1383,7 @@
      "source": [
        "parameters = {\n",
        "    \"project\": PROJECT_ID,\n",
-        "    \"location\": REGION,\n",
+        "    \"location\": LOCATION,\n",
        "    \"root_dir\": PIPELINE_ROOT,\n",
        "    \"model_name\": aip_model.resource_name,\n",
        "    \"target_field_name\": TARGET,\n",
@@ -1554,8 +1537,6 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
        "# Delete model resource\n",
        "aip_model.delete()\n",
        "\n",
@@ -1570,8 +1551,12 @@
        "\n",
        "delete_bucket = False\n",
        "# Delete Cloud Storage objects\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
-        "    ! gsutil -m rm -r $BUCKET_URI"
+        "if delete_bucket:\n",
+        "    ! gsutil -m rm -r $BUCKET_URI\n",
+        "\n",
+        "# Delete the locally generated files and folders\n",
+        "! rm $PIPELINE_PACKAGE_PATH\n",
+        "! rm -rf $SRC_DIR"
      ]
    }
  ],
@@ -33,24 +33,26 @@
        "\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/model_evaluation/model_based_llm_evaluation/autosxs_check_alignment_against_human_preference_data.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/model_evaluation/model_based_llm_evaluation/autosxs_check_alignment_against_human_preference_data.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%2Fmodel_evaluation%2Fmodel_based_llm_evaluation%2Fautosxs_check_alignment_against_human_preference_data.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/model_evaluation/model_based_llm_evaluation/autosxs_check_alignment_against_human_preference_data.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/model_evaluation/model_based_llm_evaluation/autosxs_check_alignment_against_human_preference_data.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>"
      ]
    },
@@ -114,15 +116,22 @@
        "to generate a cost estimate based on your projected usage."
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6a571eed9ece"
+      },
+      "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"
      ]
    },
    {
@@ -133,7 +142,7 @@
      },
      "outputs": [],
      "source": [
-        "! pip3 install --upgrade \\\n",
+        "! pip3 install --upgrade --quiet \\\n",
        "    google-cloud-aiplatform \\\n",
        "    google-cloud-pipeline-components \\\n",
        "    gcsfs"
@@ -142,59 +151,80 @@
    {
      "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": [
-        "# 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": "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": "ef3990d0482a"
+      },
+      "source": [
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
@@ -206,110 +236,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 may change the `REGION` variable, which is used for operations\n",
-        "throughout the rest of this notebook.  Below are regions supported for AutoSxS.\n",
-        "\n",
-        "- Americas: `us-central1`\n",
-        "- Europe: `europe-west4`\n",
-        "- Asia Pacific: `asia-southeast1`\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": "veaqVTqxIxGy"
-      },
-      "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\"}"
      ]
    },
    {
@@ -320,7 +247,7 @@
      "source": [
        "### UUID\n",
        "\n",
-        "We define a UUID generation function to avoid resource name collisions on resources created within the notebook."
+        "Define a UUID generation function to avoid resource name collisions on resources created within the notebook."
      ]
    },
    {
@@ -389,7 +316,7 @@
        "):\n",
        "    BUCKET_URI = \"gs://\" + PROJECT_ID + \"-aip-\" + UUID\n",
        "\n",
-        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
+        "! gsutil mb -l $LOCATION -p $PROJECT_ID $BUCKET_URI"
      ]
    },
    {
@@ -445,7 +372,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)"
      ]
    },
    {
@@ -454,7 +381,7 @@
        "id": "ljpNHjC9NYtG"
      },
      "source": [
-        "## Tutorial\n",
+        "### Tutorial\n",
        "It is unlikely that the autorater will perform at the same level as human raters in all customer use cases, especially in cases where human raters are expected to have specialized knowledge.\n",
        "\n",
        "The tutorial below shows how AutoSxS helps to determine if you can trust the autorater once you have the ground-truth human-preference data.\n"
@@ -759,7 +686,7 @@
      "outputs": [],
      "source": [
        "# To use an existing pipeline, override job using the line below.\n",
-        "# job = aiplatform.PipelineJob.get('projects/[PROJECT_NUMBER]/locations/[REGION]/pipelineJobs/[PIPELINE_RUN_NAME]')\n",
+        "# job = aiplatform.PipelineJob.get('projects/[PROJECT_NUMBER]/locations/[LOCATION]/pipelineJobs/[PIPELINE_RUN_NAME]')\n",
        "\n",
        "for details in job.task_details:\n",
        "    if details.task_name == \"online-evaluation-pairwise\":\n",
@@ -832,8 +759,7 @@
      "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",
+        "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",
@@ -854,7 +780,7 @@
        "\n",
        "# Delete Cloud Storage objects that were created\n",
        "delete_bucket = False\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
@@ -31,26 +31,27 @@
      "source": [
        "# AutoSxS: Evaluate a LLM in Vertex AI Model Registry against a third-party model\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/model_evaluation/model_based_llm_evaluation/autosxs_llm_evaluation_for_summarization_task.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/model_evaluation/model_based_llm_evaluation/autosxs_llm_evaluation_for_summarization_task.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%2Fmodel_evaluation%2Fmodel_based_llm_evaluation%2Fautosxs_llm_evaluation_for_summarization_task.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/model_evaluation/model_based_llm_evaluation/autosxs_llm_evaluation_for_summarization_task.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/model_evaluation/model_based_llm_evaluation/autosxs_llm_evaluation_for_summarization_task.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>"
      ]
    },
@@ -155,12 +156,19 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "i7EUnXsZhAGF"
+        "id": "3b1ffd5ab768"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook."
+        "## Get Started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "install_aip"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages"
      ]
    },
    {
@@ -171,7 +179,7 @@
      },
      "outputs": [],
      "source": [
-        "! pip3 install --upgrade \\\n",
+        "! pip3 install --upgrade --quiet \\\n",
        "    google-cloud-aiplatform \\\n",
        "    google-cloud-pipeline-components \\\n",
        "    gcsfs \\\n",
@@ -182,59 +190,80 @@
    {
      "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": "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": "6a870411c189"
+      },
+      "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)."
      ]
    },
    {
@@ -246,110 +275,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 may change the `REGION` variable, which is used for operations\n",
-        "throughout the rest of this notebook.  Below are regions supported for AutoSxS.\n",
-        "\n",
-        "- Americas: `us-central1`\n",
-        "- Europe: `europe-west4`\n",
-        "- Asia Pacific: `asia-southeast1`\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": "veaqVTqxIxGy"
-      },
-      "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\"}"
      ]
    },
    {
@@ -360,7 +286,7 @@
      "source": [
        "### UUID\n",
        "\n",
-        "We define a UUID generation function to avoid resource name collisions on resources created within the notebook."
+        "Define a UUID generation function to avoid resource name collisions on resources created within the notebook."
      ]
    },
    {
@@ -429,7 +355,7 @@
        "):\n",
        "    BUCKET_URI = \"gs://\" + PROJECT_ID + \"-aip-\" + UUID\n",
        "\n",
-        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
+        "! gsutil mb -l $LOCATION -p $PROJECT_ID $BUCKET_URI"
      ]
    },
    {
@@ -452,7 +378,7 @@
    },
    {
      "cell_type": "code",
-      "execution_count": 1,
+      "execution_count": null,
      "metadata": {
        "id": "PyQmSRbKA8r-"
      },
@@ -485,7 +411,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)"
      ]
    },
    {
@@ -494,7 +420,7 @@
        "id": "ljpNHjC9NYtG"
      },
      "source": [
-        "## Tutorial"
+        "### Tutorial"
      ]
    },
    {
@@ -519,7 +445,7 @@
    },
    {
      "cell_type": "code",
-      "execution_count": 2,
+      "execution_count": null,
      "metadata": {
        "id": "_hvNRQRTN3gb"
      },
@@ -586,7 +512,7 @@
    },
    {
      "cell_type": "code",
-      "execution_count": 3,
+      "execution_count": null,
      "metadata": {
        "id": "vykmkhp-ODKg"
      },
@@ -684,7 +610,7 @@
    },
    {
      "cell_type": "code",
-      "execution_count": 4,
+      "execution_count": null,
      "metadata": {
        "id": "tRdA3ovUOV6j"
      },
@@ -737,14 +663,14 @@
    },
    {
      "cell_type": "code",
-      "execution_count": 5,
+      "execution_count": null,
      "metadata": {
        "id": "V_9yMfhrOZDk"
      },
      "outputs": [],
      "source": [
        "# To use an existing pipeline, override job using the line below.\n",
-        "# job = aiplatform.PipelineJob.get('projects/[PROJECT_NUMBER]/locations/[REGION]/pipelineJobs/[PIPELINE_RUN_NAME]')\n",
+        "# job = aiplatform.PipelineJob.get('projects/[PROJECT_NUMBER]/locations/[LOCATION]/pipelineJobs/[PIPELINE_RUN_NAME]')\n",
        "\n",
        "for details in job.task_details:\n",
        "    if details.task_name == \"online-evaluation-pairwise\":\n",
@@ -767,7 +693,7 @@
    },
    {
      "cell_type": "code",
-      "execution_count": 6,
+      "execution_count": null,
      "metadata": {
        "id": "xtP1RabVNrlT"
      },
@@ -794,7 +720,7 @@
    },
    {
      "cell_type": "code",
-      "execution_count": 7,
+      "execution_count": null,
      "metadata": {
        "id": "G7meI2Eq4muT"
      },
@@ -815,8 +741,8 @@
      "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",
+        "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",
        "\n",
        "Otherwise, you can delete the individual resources you created in this tutorial:\n",
        "\n",
@@ -837,7 +763,7 @@
        "\n",
        "# Delete Cloud Storage objects that were created\n",
        "delete_bucket = False\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
@@ -63,7 +63,7 @@
      "source": [
        "## Overview\n",
        "\n",
-        "This tutorial demonstrates how to use with the Vertex AI SDK for Python to set up Vertex AI Model Monitoring V2 for batch predictions."
+        "This tutorial demonstrates how to use the Vertex AI SDK for Python to set up Vertex AI Model Monitoring V2 for your model."
      ]
    },
    {
@@ -74,25 +74,25 @@
      "source": [
        "### Objective\n",
        "\n",
-        "The steps performed include the following:\n",
+        "In this tutorial, you'll complete the following steps:\n",
        "\n",
        "- Upload a custom model to Vertex AI Model Registry.\n",
-        "- Create a Model Monitor.\n",
-        "- Create Vertex AI BatchPredictionJob.\n",
-        "- Run an on-demand model monitoring job to analyze data drift between the BatchPredictionJob results and the training dataset.\n",
-        "- Create another Vertex AI BatchPredictionJob.\n",
-        "- Run an on-demand model monitoring job to analyze data drift between the BatchPredictionJob results and the previous BatchPredictionJob.\n",
-        "- Run an on-demand model monitoring job to analyze the feature attribution drift between the BatchPredictionJob results and a GCS baseline dataset.\n",
+        "- Create a model monitor.\n",
+        "- Create Vertex AI batch prediction job.\n",
+        "- Run an on-demand model monitoring job to analyze data drift between the batch prediction job results and the training dataset.\n",
+        "- Create another Vertex AI batch prediction job.\n",
+        "- Run an on-demand model monitoring job to analyze data drift between the batch prediction job results and the previous batch prediction job.\n",
+        "- Run an on-demand model monitoring job to analyze the feature attribution drift between the batch prediction job results and a baseline dataset in Google Cloud Storage.\n",
        "\n",
        "\n",
        "### Costs\n",
        "\n",
-        "Vertex AI Model Monitoring V2 is free for public preview, but you will still be billed for the following components of Google Cloud:\n",
+        "Vertex AI Model Monitoring v2 is free during the public preview period, but you will still be billed for the following Google Cloud services:\n",
        "\n",
-        "* BigQuery\n",
-        "* Cloud Storage\n",
-        "* Vertex AI Batch Prediction Job\n",
-        "* Vertex AI Batch Explanation Job (if you run the feature attribution drift example)."
+        "* [BigQuery](https://cloud.google.com/bigquery/pricing)\n",
+        "* [Cloud Storage](https://cloud.google.com/storage/pricing)\n",
+        "* [Vertex AI Online Prediction](https://cloud.google.com/vertex-ai/pricing#prediction-prices)\n",
+        "* [Vertex AI Batch Explanation Job](https://cloud.google.com/vertex-ai/pricing#prediction-prices) (if you run the feature attribution drift example)."
      ]
    },
    {
@@ -136,7 +136,7 @@
        "id": "02nrfqgSOsw2"
      },
      "source": [
-        "Check the version of google-cloud-aiplatform, it must >= 1.51.0"
+        "Check that the version of google-cloud-aiplatform is 1.51.0 or later."
      ]
    },
    {
@@ -260,7 +260,7 @@
        "id": "LMRdCWwqPBpR"
      },
      "source": [
-        "## Start Model Monitoring Tutorial"
+        "## Start Model Monitoring tutorial"
      ]
    },
    {
@@ -269,9 +269,9 @@
        "id": "l_BdNCfk7Hjz"
      },
      "source": [
-        "### Step 1: Create a Could Storage Bucket\n",
+        "### Step 1: Create a Cloud Storage bucket\n",
        "\n",
-        "Create a storage bucket to store intermediate artifacts such as datasets."
+        "Create a Cloud Storage bucket to store intermediate artifacts such as datasets."
      ]
    },
    {
@@ -312,7 +312,7 @@
        "id": "5ZLxDjB6lYXx"
      },
      "source": [
-        "### Step 2: Prepare a Model in Vertex AI Model Registry"
+        "### Step 2: Prepare a model in Vertex AI Model Registry"
      ]
    },
    {
@@ -321,8 +321,8 @@
        "id": "1_OchKcp60FE"
      },
      "source": [
-        "You can register a model in Vertex AI Model Registry with its artifacts, enabling you to perform online serving or batch prediction. Alternatively, you can register a placeholder model that includes only the model's name.\n",
-        "In this notebook, we will register a model with artifacts since we need to execute a batch prediction job."
+        "You can register a model in Vertex AI Model Registry with its artifacts, enabling you to perform online serving or batch prediction. Alternatively, you can register a referenced/placeholder model that includes only the model's name.\n",
+        "In this notebook, you register a model with artifacts because you'll run a batch prediction job."
      ]
    },
    {
@@ -355,7 +355,7 @@
        "id": "RuwAoHTilv1j"
      },
      "source": [
-        "### Step 3: Create a Model Monitor"
+        "### Step 3: Create a model monitor"
      ]
    },
    {
@@ -382,9 +382,9 @@
        "id": "jyg0cMAdfrAa"
      },
      "source": [
-        "Monitoring schema is a required configuration for model monitor. It consists of the information on the names of input features, prediction outputs and, if available, ground truth, along with their respective data type.\n",
+        "The monitoring schema is required for model monitors. It includes the names of input features, prediction outputs and, if available, ground truths, along with their respective data type.\n",
        "\n",
-        "**Note: For AutoML tables (Regression/Classification), schema is optional, it will be automatically fetched when available (If we are not able to get the information, you will need to provide your own)**"
+        "**Note: For AutoML tables (regression and classification), defining the schema is optional. The schema is automatically fetched when available. If Vertex AI cannot get the schema information, you must provide it.**"
      ]
    },
    {
@@ -457,7 +457,7 @@
        "id": "SvFZRodti3qN"
      },
      "source": [
-        "The training dataset can serve as the baseline dataset to calculate some monitoring metrics. You can register the training dataset in the model monitor."
+        "The training dataset can serve as the baseline dataset to calculate monitoring metrics. You can register the training dataset in the model monitor."
      ]
    },
    {
@@ -472,7 +472,9 @@
        "\n",
        "# Copy files to your projects gs bucket to avoid permission issues.\n",
        "# Ignore any error(s) for bucket already exists.\n",
-        "PUBLIC_TRAINING_DATASET = \"gs://cloud-samples-data/vertex-ai/model-monitoring/churn/churn_training.csv\"\n",
+        "PUBLIC_TRAINING_DATASET = (\n",
+        "    \"gs://cloud-samples-data/vertex-ai/model-monitoring/churn/churn_training.csv\"\n",
+        ")\n",
        "TRAINING_URI = f\"{BUCKET_URI}/model-monitoring/churn/churn_training.csv\"\n",
        "\n",
        "! gsutil copy $PUBLIC_TRAINING_DATASET $TRAINING_URI\n",
@@ -488,7 +490,7 @@
        "id": "5ZnBZUEMjtDD"
      },
      "source": [
-        "#### Create a model monitor"
+        "#### Create a model monitor resource"
      ]
    },
    {
@@ -497,7 +499,7 @@
        "id": "OPBdgqz6jv9_"
      },
      "source": [
-        "Model monitor is a top level resource to manage your metrics and model monitoring jobs."
+        "A model monitor is a top-level resource to manage your metrics and model monitoring jobs."
      ]
    },
    {
@@ -540,25 +542,25 @@
      "source": [
        "#### Define the monitoring objective configs\n",
        "\n",
-        "For tabular models, we support the following objectives:\n",
+        "For tabular models, Model Monitoring supports the following objectives:\n",
        "\n",
-        "*   **Input Feature Drift Detection**\n",
+        "*   **Input feature drift detection**\n",
        "\n",
-        "    The specification for the input feature drift detection. We offer drift analysis for both categorical and numeric feature types, with the following supported metrics:\n",
+        "    Model Monitoring offers drift analysis for both categorical and numeric feature types, with the following supported metrics:\n",
        "\n",
        "    *    Categorical Feature: `Jensen Shannon Divergence`, `L Infinity`\n",
        "    *    Numeric Feature: `Jensen Shannon Divergence`\n",
        "\n",
-        "    You can choose to analyze only the features of interest in the `features` fields of `ml_monitoring.spec.DataDriftSpec`. If this is not specified, all input features in the model schema will be analyzed. Additionally, you have the option to set default thresholds for categorical or numeric features, or you can specify thresholds for individual features. If the detected drift surpasses these thresholds, an alert will be sent to you via email or another notification channel.\n",
+        "    You can choose to analyze only the features of interest by specifying them in the `features` fields of the `ml_monitoring.spec.DataDriftSpec` specification. If not specified, all input features in the model schema are analyzed. Additionally, you have the option to set default thresholds for categorical or numeric features, or you can specify thresholds for individual features. If the detected drift surpasses a threshold, an alert is sent through email or another notification channel.\n",
        "\n",
-        "*  **Prediction Output Drift Detection**\n",
+        "*  **Prediction output drift detection**\n",
        "\n",
-        "    Similar to input feature drift detection, prediction output drift detection aims to identify data drift in the prediction outputs.\n",
+        "    Similar to input feature drift detection, prediction output drift detection identifies data drift in the prediction outputs.\n",
        "\n",
-        "*   **Feature Attribution drift detection**\n",
+        "*   **Feature attribution drift detection**\n",
        "\n",
-        "    Vertex Explainable AI adds another facet to model monitoring, which we call feature attribution monitoring. Explainable AI enables you to understand the relative contribution of each feature to a resulting prediction. In essence, it assesses the magnitude of each feature's influence.\n",
-        "    You must configure the `Explanation` spec with the feature attribution objectives configuration.\n"
+        "    Model Monitoring leverages Vertex Explainable AI to monitor feature attributions. Explainable AI enables you to understand the relative contribution of each feature to a resulting prediction. In essence, it assesses the magnitude of each feature's influence.\n",
+        "    You must configure the `Explanation` specification with the feature attribution objectives configuration.\n"
      ]
    },
    {
@@ -567,7 +569,7 @@
        "id": "n49LgtHBfpTE"
      },
      "source": [
-        "Input Feature Drift Specification"
+        "Input feature drift specification"
      ]
    },
    {
@@ -600,7 +602,7 @@
        "id": "DsuNTFAeftgj"
      },
      "source": [
-        "Prediction Output Drift Specification"
+        "Prediction output drift specification"
      ]
    },
    {
@@ -625,7 +627,7 @@
        "id": "ak54Y713f1aj"
      },
      "source": [
-        "Feature Attribution Specification"
+        "Feature attribution specification"
      ]
    },
    {
@@ -657,13 +659,13 @@
        "id": "8N9YlXzOvyOs"
      },
      "source": [
-        "We support various methods of notification:\n",
+        "Model Monitoring supports the following notification methods:\n",
        "\n",
        "*   Email\n",
        "*   [Notification Channel](https://cloud.google.com/monitoring/support/notification-options)\n",
-        "*   [Cloud Logging](https://cloud.google.com/logging/docs?_gl=1*tdcri2*_up*MQ..&gclid=Cj0KCQjwir2xBhC_ARIsAMTXk84diOnqqpDckjOZUas26cUXUgEAgEGT9uFpz9tTvkfUjmVnRs7lQuwaAjiwEALw_wcB&gclsrc=aw.ds)  \n",
+        "*   [Cloud Logging](https://cloud.google.com/logging/docs)  \n",
        "\n",
-        "In this notebook, we will use email as an example."
+        "This notebook uses email as an example."
      ]
    },
    {
@@ -672,7 +674,7 @@
        "id": "7jzY1mpdw-CY"
      },
      "source": [
-        "We will export generated metrics to the Google Cloud Storage location you defined or to a default bucket created by Vertex AI if you don't specify one."
+        "Export generated metrics to the Google Cloud Storage location that you specified or, if you don't specify a location, Vertex AI creates a default bucket to use."
      ]
    },
    {
@@ -713,7 +715,7 @@
        "id": "saA2MVZ-Yhy_"
      },
      "source": [
-        "##### **Example 1: Feature drift detection, compares the batch prediction job with training dataset.**"
+        "##### **Example 1: Detect feature drift by comparing a batch prediction job with the training dataset.**"
      ]
    },
    {
@@ -791,7 +793,7 @@
        "id": "ihfK5KxgWEXz"
      },
      "source": [
-        "##### **Example 2: Feature drift & prediction output drift detection, compares the batch prediction job with previous batch prediction job result.**"
+        "##### **Example 2: Detect feature drift and prediction output drift by comparing a batch prediction job with a previous batch prediction job result.**"
      ]
    },
    {
@@ -800,9 +802,9 @@
        "id": "kBhhRjuIa2iv"
      },
      "source": [
-        "You can set up multiple objectives within a single model monitoring job. All metrics will be computed using the same baseline and target dataset.\n",
+        "You can set up multiple objectives within a single model monitoring job. All metrics are computed by using the same baseline and target dataset.\n",
        "\n",
-        "Let's create another batch prediction job and compares with the batch prediction job we created above."
+        "Create another batch prediction job and compare it with the batch prediction job created previously."
      ]
    },
    {
@@ -884,15 +886,15 @@
        "id": "kP7t8p8Ih2a-"
      },
      "source": [
-        "For feature attribution monitoring, the dataset will be sent to the Vertex AI Batch Explanation Job in the following way:\n",
+        "For feature attribution monitoring, the dataset is sent to the Vertex AI batch explanation job in the following way:\n",
        "\n",
-        "*   Google Cloud Storage -> Sent directly as input to Vertex AI Batch Explanation Job.\n",
-        "*   BigQuery Table -> Sent directly as input to Vertex AI Batch Explanation Job.\n",
+        "*   Google Cloud Storage -> Sent directly as input to Vertex AI batch explanation job.\n",
+        "*   BigQuery table -> Sent directly as input to Vertex AI batch explanation job.\n",
        "*   BigQuery Query -> Not supported.\n",
-        "*   Vertex AI Batch Prediction Job -> Input of Batch Prediction Job will be used as input for Vertex AI Batch Explanation Job.\n",
-        "*   Vertex AI Endpoint Logging -> `request_payload` will be used as input for Vertex AI Batch Explanation Job.\n",
+        "*   Vertex AI batch explanation job -> Input of batch prediction job is used as input for the Vertex AI batch explanation job.\n",
+        "*   Vertex AI endpoint logging -> Request logging is used as input for Vertex AI batch explanation job.\n",
        "\n",
-        "So please make sure these datasets meets the requirements for Vertex AI Batch Explanation Job."
+        "Check that these datasets meet the requirements for a Vertex AI batch explanation job."
      ]
    },
    {
@@ -902,7 +904,7 @@
      },
      "source": [
        "###### Generate model metadata for Vertex Explainable AI\n",
-        "As we will use Vertex AI Batch Explanation Job, you must speicfy the explanation spec. Run the following cell to extract metadata from the exported model, which is needed for generating the explanations for a prediction request.\n"
+        "You must specify the explanation specification to use a Vertex AI batch explanation job. Run the following cell to extract metadata from the exported model, which is needed for generating the explanations for a prediction request."
      ]
    },
    {
@@ -1143,7 +1145,7 @@
        "id": "LrSU6d0xFfzv"
      },
      "source": [
-        "#### Check email"
+        "#### Verify results through email"
      ]
    },
    {
@@ -1152,7 +1154,7 @@
        "id": "zw5KodgebDVE"
      },
      "source": [
-        "Once the model monitoring job begins running (it will start after the batch prediction jobs have finished), you will receive an email as follows:\n",
+        "After the model monitoring job begins running, which starts after the batch prediction jobs have finished, you receive an email like the following one:\n",
        "\n",
        "<img src=\"https://services.google.com/fh/files/misc/create_job_email.png\" />"
      ]
@@ -1163,7 +1165,7 @@
        "id": "37SPm_majLVj"
      },
      "source": [
-        "Once the monitoring job is complete, should any anomalies be detected, you will receive an email similar to\n",
+        "After the monitoring job is complete, if any anomalies are detected, you receive an email similar to the following one:\n",
        "\n",
        "<img src=\"https://services.google.com/fh/files/misc/job_anomalies_email.png\" />"
      ]
@@ -1174,7 +1176,7 @@
        "id": "LPC6ZoerifM6"
      },
      "source": [
-        "#### Check GCP Console"
+        "#### Check monitoring metrics: Google Cloud Console"
      ]
    },
    {
@@ -1183,7 +1185,7 @@
        "id": "4Pw6Z-bVbZaE"
      },
      "source": [
-        "Check the \"Monitor\" tab under \"Vertex AI\"\n",
+        "To view Model Monitoring metrics in the [Google Cloud Console](https://console.cloud.google.com/vertex-ai/model-monitoring/model-monitors), go to the **Monitoring** tab under **Vertex AI.**\n",
        "\n",
        "<img src=\"https://storage.googleapis.com/cmm-public-data/images/bp_details.gif\" />"
      ]
@@ -1194,7 +1196,9 @@
        "id": "MFZcBnwsdSfa"
      },
      "source": [
-        "#### After jobs are completed, check output GCS bucket"
+        "#### Check monitoring metrics: Cloud Storage bucket\n",
+        "\n",
+        "Run the following to view Model Monitoring metrics stored in the Cloud Storage bucket.  "
      ]
    },
    {
@@ -1245,7 +1249,9 @@
        "id": "hovSbsmBmEKZ"
      },
      "source": [
-        "### Step 6: Clean Up (after job finished)"
+        "### Step 6: Clean Up\n",
+        "\n",
+        "If you no longer need your model monitoring resources, run the following to delete them:"
      ]
    },
    {
@@ -63,7 +63,7 @@
      "source": [
        "## Overview\n",
        "\n",
-        "This tutorial demonstrates how to use with the Vertex AI SDK for Python to set up Vertex AI Model Monitoring V2 for online predictions."
+        "This tutorial demonstrates how to use the Vertex AI SDK for Python to set up Vertex AI Model Monitoring V2 for your model."
      ]
    },
    {
@@ -74,12 +74,12 @@
      "source": [
        "### Objective\n",
        "\n",
-        "The steps performed include the following:\n",
+        "In this tutorial, you'll complete the following steps:\n",
        "\n",
        "- Upload a custom model to Vertex AI Model Registry.\n",
        "- Deploy the model to Vertex AI Endpoint with Request-Response Logging enabled.\n",
        "- Generate some online prediction traffic.\n",
-        "- Create a Model Monitor.\n",
+        "- Create a model monitor.\n",
        "- Run an on-demand model monitoring job to analyze data drift between the Online prediction traffic and the training dataset.\n",
        "- Run an on-demand model monitoring job to analyze feature attribution drift between the Online prediction traffic and a GCS baseline dataset.\n",
        "- Create a schedule to continuously run model monitoring jobs to analyze data drift between the Online prediction traffic and the training dataset.\n",
@@ -87,12 +87,12 @@
        "\n",
        "### Costs\n",
        "\n",
-        "Vertex AI Model Monitoring V2 is free for public preview, but you will still be billed for the following components of Google Cloud:\n",
+        "Vertex AI Model Monitoring v2 is free during the public preview period, but you will still be billed for the following Google Cloud services:\n",
        "\n",
-        "* BigQuery\n",
-        "* Cloud Storage\n",
-        "* Vertex AI Online Prediction\n",
-        "* Vertex AI Batch Explanation Job (if you run the feature attribution drift example)."
+        "* [BigQuery](https://cloud.google.com/bigquery/pricing)\n",
+        "* [Cloud Storage](https://cloud.google.com/storage/pricing)\n",
+        "* [Vertex AI Online Prediction](https://cloud.google.com/vertex-ai/pricing#prediction-prices)\n",
+        "* [Vertex AI Batch Explanation Job](https://cloud.google.com/vertex-ai/pricing#prediction-prices) (if you run the feature attribution drift example)."
      ]
    },
    {
@@ -136,7 +136,7 @@
        "id": "s8ol0OhLT8SQ"
      },
      "source": [
-        "Check the version of google-cloud-aiplatform, it must >= 1.51.0"
+        "Check that the version of google-cloud-aiplatform is 1.51.0 or later."
      ]
    },
    {
@@ -237,16 +237,7 @@
      "metadata": {
        "id": "7-qYfaAsUKux"
      },
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "Updated property [core/project].\n",
-            "Updated property [ai/region].\n"
-          ]
-        }
-      ],
+      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
@@ -269,7 +260,7 @@
        "id": "QNiB5LfBaft-"
      },
      "source": [
-        "## Start Model Monitoring Tutorial"
+        "## Start Model Monitoring tutorial"
      ]
    },
    {
@@ -278,9 +269,9 @@
        "id": "l_BdNCfk7Hjz"
      },
      "source": [
-        "### Step 1: Create a Could Storage Bucket\n",
+        "### Step 1: Create a Cloud Storage bucket\n",
        "\n",
-        "Create a storage bucket to store intermediate artifacts such as datasets."
+        "Create a Cloud Storage bucket to store intermediate artifacts such as datasets."
      ]
    },
    {
@@ -321,7 +312,7 @@
        "id": "5ZLxDjB6lYXx"
      },
      "source": [
-        "### Step 2: Prepare a Model in Vertex AI Model Registry"
+        "### Step 2: Prepare a model in Vertex AI Model Registry"
      ]
    },
    {
@@ -330,8 +321,8 @@
        "id": "q26J5e-D7p8e"
      },
      "source": [
-        "You can register a model in Vertex AI Model Registry with its artifacts, enabling you to perform online serving or batch prediction. Alternatively, you can register a placeholder model that includes only the model's name.\n",
-        "In this notebook, we will register a model with artifacts since we need to execute a batch prediction job."
+        "You can register a model in Vertex AI Model Registry with its artifacts, enabling you to perform online serving or batch prediction. Alternatively, you can register a referenced/placeholder model that includes only the model's name.\n",
+        "In this notebook, you register a model with artifacts because you'll run a batch prediction job."
      ]
    },
    {
@@ -364,7 +355,7 @@
        "id": "5p3buU6pC5B6"
      },
      "source": [
-        "### Step 3: Deploy Model to Vertex AI Endpoint with Req-Resp Logging Enabled"
+        "### Step 3: Deploy Model to Vertex AI Endpoint with Request-Respponse Logging Enabled"
      ]
    },
    {
@@ -373,7 +364,7 @@
        "id": "8h-EOsirC8dY"
      },
      "source": [
-        "#### Create an endpoint with logging enabled and deploy the model to this endpoint"
+        "#### Create an endpoint with logging enabled and deploy the model to this endpoint."
      ]
    },
    {
@@ -384,7 +375,7 @@
      },
      "outputs": [],
      "source": [
-        "# Create an endpoint with logging enabled. Specify the logging sampling rate and BigQuery Destination.\n",
+        "# Create an endpoint with logging enabled. Specify the logging sampling rate and BigQuery destination.\n",
        "import pandas as pd\n",
        "\n",
        "TIMESTAMP = pd.Timestamp.utcnow().strftime(\"%Y%m%d%H%M%S\")\n",
@@ -569,9 +560,9 @@
        "id": "jyg0cMAdfrAa"
      },
      "source": [
-        "Monitoring schema is a required configuration for model monitor. It consists of the information on the names of input features, prediction outputs and, if available, ground truth, along with their respective data type.\n",
+        "The monitoring schema is a required configuration for a model monitor. The schema consists of input features names, prediction outputs, and ground truth (if available), along with their respective data types.\n",
        "\n",
-        "**Note: For AutoML tables (Regression/Classification), schema is optional, it will be automatically fetched when available (If we are not able to get the information, you will need to provide your own)**"
+        "**Note: The schema is optional only for AutoML tables (Regression/Classification), as it will be automatically fetched when available (if Model Monitoring is unable to retrieve the schema, you will need to provide your own).**"
      ]
    },
    {
@@ -659,7 +650,9 @@
        "\n",
        "# Copy files to your projects gs bucket to avoid permission issues.\n",
        "# Ignore any error(s) for bucket already exists.\n",
-        "PUBLIC_TRAINING_DATASET = \"gs://cloud-samples-data/vertex-ai/model-monitoring/churn/churn_training.csv\"\n",
+        "PUBLIC_TRAINING_DATASET = (\n",
+        "    \"gs://cloud-samples-data/vertex-ai/model-monitoring/churn/churn_training.csv\"\n",
+        ")\n",
        "TRAINING_URI = f\"{BUCKET_URI}/model-monitoring/churn/churn_training.csv\"\n",
        "\n",
        "! gsutil copy $PUBLIC_TRAINING_DATASET $TRAINING_URI\n",
@@ -684,7 +677,7 @@
        "id": "OPBdgqz6jv9_"
      },
      "source": [
-        "Model monitor is a top level resource to manage your metrics and model monitoring jobs."
+        "A model monitor is a top-level resource to manage your metrics and model monitoring jobs."
      ]
    },
    {
@@ -727,25 +720,25 @@
      "source": [
        "#### Define the monitoring objective configs\n",
        "\n",
-        "For tabular models, we support the following objectives:\n",
+        "For tabular models, Model Monitoring supports the following objectives:\n",
        "\n",
-        "*   **Input Feature Drift Detection**\n",
+        "*   **Input feature drift detection**\n",
        "\n",
-        "    The specification for the input feature drift detection. We offer drift analysis for both categorical and numeric feature types, with the following supported metrics:\n",
+        "    Model Monitoring offers drift analysis for both categorical and numeric feature types, with the following supported metrics:\n",
        "\n",
        "    *    Categorical Feature: `Jensen Shannon Divergence`, `L Infinity`\n",
        "    *    Numeric Feature: `Jensen Shannon Divergence`\n",
        "\n",
-        "    You can choose to analyze only the features of interest in the `features` fields of `ml_monitoring.spec.DataDriftSpec`. If this is not specified, all input features in the model schema will be analyzed. Additionally, you have the option to set default thresholds for categorical or numeric features, or you can specify thresholds for individual features. If the detected drift surpasses these thresholds, an alert will be sent to you via email or another notification channel.\n",
+        "    You can choose to analyze only the features of interest by specifying them in the `features` fields of the `ml_monitoring.spec.DataDriftSpec` specification. If not specified, all input features in the model schema are analyzed. Additionally, you have the option to set default thresholds for categorical or numeric features, or you can specify thresholds for individual features. If the detected drift surpasses a threshold, an alert is sent through email or another notification channel.\n",
        "\n",
-        "*  **Prediction Output Drift Detection**\n",
+        "*  **Prediction output drift detection**\n",
        "\n",
-        "    Similar to input feature drift detection, prediction output drift detection aims to identify data drift in the prediction outputs.\n",
+        "    Similar to input feature drift detection, prediction output drift detection identifies data drift in the prediction outputs.\n",
        "\n",
-        "*   **Feature Attribution drift detection**\n",
+        "*   **Feature attribution drift detection**\n",
        "\n",
-        "    Vertex Explainable AI adds another facet to model monitoring, which we call feature attribution monitoring. Explainable AI enables you to understand the relative contribution of each feature to a resulting prediction. In essence, it assesses the magnitude of each feature's influence.\n",
-        "    You must configure the `Explanation` spec with the feature attribution objectives configuration.\n"
+        "    Model Monitoring leverages Vertex Explainable AI to monitor feature attributions. Explainable AI enables you to understand the relative contribution of each feature to a resulting prediction. In essence, it assesses the magnitude of each feature's influence.\n",
+        "    You must configure the `Explanation` specification with the feature attribution objectives configuration.\n"
      ]
    },
    {
@@ -754,7 +747,7 @@
        "id": "XZgrufEwzuFC"
      },
      "source": [
-        "Input Feature Drift Specification"
+        "Input feature drift specification"
      ]
    },
    {
@@ -787,7 +780,7 @@
        "id": "i6rqFXXWzxuQ"
      },
      "source": [
-        "Prediction Output Drift Specification"
+        "Prediction output drift specification"
      ]
    },
    {
@@ -812,7 +805,7 @@
        "id": "WzdHn15az3pg"
      },
      "source": [
-        "Feature Attribution Specification"
+        "Feature attribution specification"
      ]
    },
    {
@@ -844,13 +837,13 @@
        "id": "8N9YlXzOvyOs"
      },
      "source": [
-        "We support various methods of notification:\n",
+        "Model Monitoring supports the following notification methods:\n",
        "\n",
        "*   Email\n",
        "*   [Notification Channel](https://cloud.google.com/monitoring/support/notification-options)\n",
-        "*   [Cloud Logging](https://cloud.google.com/logging/docs?_gl=1*tdcri2*_up*MQ..&gclid=Cj0KCQjwir2xBhC_ARIsAMTXk84diOnqqpDckjOZUas26cUXUgEAgEGT9uFpz9tTvkfUjmVnRs7lQuwaAjiwEALw_wcB&gclsrc=aw.ds)  \n",
+        "*   [Cloud Logging](https://cloud.google.com/logging/docs)  \n",
        "\n",
-        "In this notebook, we will use email as an example."
+        "This notebook uses email as an example."
      ]
    },
    {
@@ -859,7 +852,7 @@
        "id": "7jzY1mpdw-CY"
      },
      "source": [
-        "We will export generated metrics to the Google Cloud Storage location you defined or to a default bucket created by Vertex AI if you don't specify one."
+        "Export generated metrics to the Google Cloud Storage location that you specified or, if you don't specify a location, Vertex AI creates a default bucket to use."
      ]
    },
    {
@@ -898,7 +891,7 @@
        "id": "JaAEOio70bTZ"
      },
      "source": [
-        "##### **Example 1: Feature drift & Prediction output drift detection, compares with training dataset.**"
+        "##### **Example 1: Feature drift & Prediction output drift detection, compares current data with the training dataset.**"
      ]
    },
    {
@@ -935,7 +928,7 @@
        "id": "bpbpwths0lf5"
      },
      "source": [
-        "##### **Example 2: Feature attribution drift detection, compares with a GCS baseline dataset**"
+        "##### **Example 2: Feature attribution drift detection, compares current data with a Cloud Storage baseline dataset**"
      ]
    },
    {
@@ -944,15 +937,15 @@
        "id": "Yw2DR7t70zkj"
      },
      "source": [
-        "For feature attribution monitoring, the dataset will be sent to the Vertex AI Batch Explanation Job in the following way:\n",
+        "For feature attribution monitoring, the dataset is sent to the Vertex AI batch explanation job in the following way:\n",
        "\n",
-        "*   Google Cloud Storage -> Sent directly as input to Vertex AI Batch Explanation Job.\n",
-        "*   BigQuery Table -> Sent directly as input to Vertex AI Batch Explanation Job.\n",
+        "*   Google Cloud Storage -> Sent directly as input to Vertex AI batch explanation job.\n",
+        "*   BigQuery table -> Sent directly as input to Vertex AI batch explanation job.\n",
        "*   BigQuery Query -> Not supported.\n",
-        "*   Vertex AI Batch Prediction Job -> Input of Batch Prediction Job will be used as input for Vertex AI Batch Explanation Job.\n",
-        "*   Vertex AI Endpoint Logging -> `request_payload` will be used as input for Vertex AI Batch Explanation Job.\n",
+        "*   Vertex AI batch explanation job -> Input of batch prediction job is used as input for the Vertex AI batch explanation job.\n",
+        "*   Vertex AI endpoint logging -> Request logging is used as input for Vertex AI batch explanation job.\n",
        "\n",
-        "So please make sure these datasets meets the requirements for Vertex AI Batch Explanation Job."
+        "Check that these datasets meet the requirements for a Vertex AI batch explanation job."
      ]
    },
    {
@@ -962,7 +955,7 @@
      },
      "source": [
        "###### Generate model metadata for Vertex Explainable AI\n",
-        "As we will use Vertex AI Batch Explanation Job, you must speicfy the explanation spec. Run the following cell to extract metadata from the exported model, which is needed for generating the explanations for a prediction request."
+        "You must specify the explanation specification to use a Vertex AI batch explanation job. Run the following cell to extract metadata from the exported model, which is needed for generating the explanations for a prediction request."
      ]
    },
    {
@@ -1201,7 +1194,7 @@
        "id": "LrSU6d0xFfzv"
      },
      "source": [
-        "#### Check email"
+        "#### Verify results through email"
      ]
    },
    {
@@ -1221,7 +1214,7 @@
        "id": "37SPm_majLVj"
      },
      "source": [
-        "Once the monitoring job is complete, should any anomalies be detected, you will receive an email similar to\n",
+        "After the monitoring job is complete, if any anomalies are detected, you receive an email similar to the following one:\n",
        "\n",
        "<img src=\"https://services.google.com/fh/files/misc/job_anomalies_email.png\" />"
      ]
@@ -1232,7 +1225,7 @@
        "id": "LPC6ZoerifM6"
      },
      "source": [
-        "#### Check GCP Console"
+        "#### Check monitoring metrics: Google Cloud Console"
      ]
    },
    {
@@ -1241,7 +1234,7 @@
        "id": "4Pw6Z-bVbZaE"
      },
      "source": [
-        "Check the \"Monitor\" tab under \"Vertex AI\"\n",
+        "To view Model Monitoring metrics in the [Google Cloud Console](https://console.cloud.google.com/vertex-ai/model-monitoring/model-monitors), go to the **Monitoring** tab under **Vertex AI.**\n",
        "\n",
        "<img src=\"https://storage.googleapis.com/cmm-public-data/images/endpoint_jobs_details.gif\" />"
      ]
@@ -1252,7 +1245,9 @@
        "id": "MFZcBnwsdSfa"
      },
      "source": [
-        "#### Check Output GCS bucket"
+        "#### Check monitoring metrics: Cloud Storage bucket\n",
+        "\n",
+        "Run the following to view Model Monitoring metrics stored in the Cloud Storage bucket.  "
      ]
    },
    {
@@ -1298,9 +1293,10 @@
        "id": "yGp6GlV29QAE"
      },
      "source": [
-        "If you are interested at trying continous model monitoring, please following the example below to create a schedule. You could create multiple schedules for your monitor.\n",
+        "To set up continous model monitoring, follow the example below to create a schedule. You can create multiple schedules for your model monitor.\n",
        "\n",
-        "This example is set up to monitor drift in both the input and the prediction output. The schedule is configured to activate the model monitoring job every hour on the hour, such as at 00:00, 01:00, and so on. Each job will analyze data from the preceding one-hour window. For instance, if a job is scheduled for 6:00 am, it will analyze the data collected from 5:00 am to 6:00 am."
+        "\n",
+        "The following example monitors drift in both input features and prediction outputs. The schedule is configured to activate the model monitoring job every hour on the hour, such as at 00:00, 01:00, and so on. Each job analyzes data from the preceding one-hour window. For instance, if a job is scheduled for 6:00 a.m., it analyzes the data collected from 5:00 a.m. to 6:00 a.m."
      ]
    },
    {
@@ -1366,9 +1362,8 @@
      },
      "outputs": [],
      "source": [
-        "CRON = (\n",
-        "    \"0 * * * *\"  # @param {type:\"string\"} Every 1 hour at :00, for example 1:00, 2:00..\n",
-        ")"
+        "# Every 1 hour at :00, for example 1:00, 2:00..\n",
+        "CRON = \"0 * * * *\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -1432,7 +1427,9 @@
        "id": "YI9ZejD4S4jK"
      },
      "source": [
-        "#### Pause Schedule"
+        "#### Pause Schedule\n",
+        "\n",
+        "Run the following to pause the model monitoring schedule:"
      ]
    },
    {
@@ -1452,7 +1449,9 @@
        "id": "0O4BN64KTt95"
      },
      "source": [
-        "#### Resume Schedule"
+        "#### Resume Schedule\n",
+        "\n",
+        "Run the following to resume a paused model monitoring schedule:"
      ]
    },
    {
@@ -1472,7 +1471,9 @@
        "id": "ydPzrEVuT1O9"
      },
      "source": [
-        "#### Update Schedule"
+        "#### Update Schedule\n",
+        "\n",
+        "Run the following to update the model monitoring schedule:"
      ]
    },
    {
@@ -1495,7 +1496,9 @@
        "id": "pOH_UNH5NyJG"
      },
      "source": [
-        "You could check your schedules in Console.\n",
+        "#### Check monitoring schedules in Google Cloud Console\n",
+        "\n",
+        "To check your Model Monitoring schedule in the Google Cloud Console, go to the Monitor tab under Vertex AI.\n",
        "\n",
        "\n",
        "<img src=\"https://services.google.com/fh/files/misc/endpoint_schedules.gif\" />"
@@ -1507,7 +1510,9 @@
        "id": "hovSbsmBmEKZ"
      },
      "source": [
-        "### Step 8: Clean Up (after job finished)"
+        "### Step 8: Clean Up\n",
+        "\n",
+        "If you no longer need your model monitoring resources, run the following to delete them:"
      ]
    },
    {
@@ -84,10 +84,10 @@
        "\n",
        "### Costs\n",
        "\n",
-        "Vertex AI Model Monitoring V2 is free for public preview, but you will still be billed for the following components of Google Cloud:\n",
+        "Vertex AI Model Monitoring v2 is free during the public preview period, but you will still be billed for the following Google Cloud services:\n",
        "\n",
-        "* BigQuery\n",
-        "* Cloud Storage"
+        "* [BigQuery](https://cloud.google.com/bigquery/pricing)\n",
+        "* [Cloud Storage](https://cloud.google.com/storage/pricing)"
      ]
    },
    {
@@ -105,7 +105,7 @@
        "id": "f3848df1e5b0"
      },
      "source": [
-        "### Install Vertex AI SDK and other required packages\n"
+        "### Install Vertex AI SDK and other required packages"
      ]
    },
    {
@@ -131,7 +131,7 @@
        "id": "fbOQ_xi8xhCk"
      },
      "source": [
-        "Check the version of google-cloud-aiplatform, it must >= 1.51.0"
+        "Check that the version of google-cloud-aiplatform is 1.51.0 or later."
      ]
    },
    {
@@ -255,7 +255,7 @@
        "id": "MVnR14LEzbvd"
      },
      "source": [
-        "## Start Model Monitoring Tutorial"
+        "## Start Model Monitoring tutorial"
      ]
    },
    {
@@ -264,9 +264,9 @@
        "id": "AxqAU5Ua43GR"
      },
      "source": [
-        "### Step 1: Create a Could Storage Bucket\n",
+        "### Step 1: Create a Cloud Storage bucket\n",
        "\n",
-        "Create a storage bucket to store intermediate artifacts such as datasets."
+        "Create a Cloud Storage bucket to store intermediate artifacts such as datasets."
      ]
    },
    {
@@ -309,7 +309,7 @@
      "source": [
        "### Step 2: Create a Placeholder Model in Vertex AI\n",
        "\n",
-        "You can create a referenced/placeholder model for your model outside Vertex AI that using only a display name."
+        "You can create a referenced/placeholder model for your model outside Vertex AI using only a display name."
      ]
    },
    {
@@ -342,14 +342,21 @@
        "id": "yUO5gFnCtHaZ"
      },
      "source": [
-        "We recommend using BigQuery to store your production dataset. Please ensure that features are stored in separate columns, example BigQuery schema:\n",
+        "We recommend using BigQuery to store your production dataset. Please ensure that features are stored in separate columns. The following is an example BigQuery schema:\n",
        "\n",
        "<!-- <img src=\"https://services.google.com/fh/gumdrop/preview/misc/example_bq_schema_2.png\" width=\"400\" height=\"300\"/> -->\n",
        "<img src=\"https://services.google.com/fh/files/misc/example_bq_schema_2.png\" width=\"400\" height=\"300\"/>\n",
        "\n",
-        "Note: If you want to setup continous monitoring with time specficiation, for example a time window, a timestamp column is required.\n",
-        "\n",
-        "**For running this tutorial, let's create some fake serving data**"
+        "Note: If you want to setup continous monitoring with time specification, for example a time window, a timestamp column is required."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "DT3OrZZlIXu5"
+      },
+      "source": [
+        "Create some fake serving data for this tutorial:"
      ]
    },
    {
@@ -486,11 +493,11 @@
      "source": [
        "#### Define Model Monitoring Schema\n",
        "\n",
-        "Monitoring schema is a required configuration for model monitor. It consists of the information on the names of input features, prediction outputs and, if available, ground truth, along with their respective data type.\n",
+        "The monitoring schema is a required configuration for a model monitor. The schema consists of input features names, prediction outputs, and ground truth (if available), along with their respective data types.\n",
        "\n",
-        "**Note: Only for AutoML tables (Regression/Classification), schema is optional, it will be automatically fetched when available (If we are not able to get the information, you will need to provide your own)**\n",
+        "**Note: The schema is optional only for AutoML tables (Regression/Classification), as it will be automatically fetched when available (if Model Monitoring is unable to retrieve the schema, you will need to provide your own).**\n",
        "\n",
-        "##### You could manually define the model schema as follows:"
+        "##### You can manually define the model schema as follows:"
      ]
    },
    {
@@ -526,7 +533,7 @@
        "id": "zn7YMjRMz43K"
      },
      "source": [
-        "##### Alternatively, if you have a large number of features, you can use the `transform_schema_from_bigquery` method to retrieve the schema and modify it as needed."
+        "##### Alternatively, if you have a large number of features, you can use the `transform_schema_from_bigquery` method to retrieve the schema and modify it as needed:"
      ]
    },
    {
@@ -552,7 +559,7 @@
        "id": "_zn_efoq0d6T"
      },
      "source": [
-        "You could then modify it."
+        "Modify the schema:"
      ]
    },
    {
@@ -588,7 +595,7 @@
        "id": "L8KuANNbiiEV"
      },
      "source": [
-        "#### (Optional)Define training dataset"
+        "#### (Optional) Define the training dataset"
      ]
    },
    {
@@ -597,7 +604,7 @@
        "id": "SvFZRodti3qN"
      },
      "source": [
-        "The training dataset can serve as the baseline dataset to calculate some monitoring metrics. You can register the training dataset in the model monitor."
+        "The training dataset can serve as the baseline dataset to calculate monitoring metrics. You can register the training dataset in the model monitor."
      ]
    },
    {
@@ -612,7 +619,9 @@
        "\n",
        "# Copy files to your projects gs bucket to avoid permission issues.\n",
        "# Ignore any error(s) for bucket already exists.\n",
-        "PUBLIC_TRAINING_DATASET = \"gs://cloud-samples-data/vertex-ai/model-monitoring/penguins/penguins_training.csv\"\n",
+        "PUBLIC_TRAINING_DATASET = (\n",
+        "    \"gs://cloud-samples-data/vertex-ai/model-monitoring/penguins/penguins_training.csv\"\n",
+        ")\n",
        "TRAINING_URI = f\"{BUCKET_URI}/model-monitoring/penguins/penguins_training.csv\"\n",
        "\n",
        "! gsutil copy $PUBLIC_TRAINING_DATASET $TRAINING_URI\n",
@@ -637,7 +646,7 @@
        "id": "OPBdgqz6jv9_"
      },
      "source": [
-        "Model monitor is a top level resource to manage your metrics and model monitoring jobs."
+        "A model monitor is a top-level resource to manage your metrics and model monitoring jobs."
      ]
    },
    {
@@ -680,22 +689,22 @@
      "source": [
        "#### Define the monitoring objective configs\n",
        "\n",
-        "For tabular models, we support the following objectives:\n",
+        "For tabular models, Model Monitoring supports the following objectives:\n",
        "\n",
-        "*   **Input Feature Drift Detection**\n",
+        "*   **Input feature drift detection**\n",
        "\n",
-        "    The specification for the input feature drift detection. We offer drift analysis for both categorical and numeric feature types, with the following supported metrics:\n",
+        "    Model Monitoring offers drift analysis for both categorical and numeric feature types, with the following supported metrics:\n",
        "\n",
        "    *    Categorical Feature: `Jensen Shannon Divergence`, `L Infinity`\n",
        "    *    Numeric Feature: `Jensen Shannon Divergence`\n",
        "\n",
-        "    You can choose to analyze only the features of interest in the `features` fields of `ml_monitoring.spec.DataDriftSpec`. If this is not specified, all input features in the model schema will be analyzed. Additionally, you have the option to set default thresholds for categorical or numeric features, or you can specify thresholds for individual features. If the detected drift surpasses these thresholds, an alert will be sent to you via email or another notification channel.\n",
+        "    You can choose to analyze only the features of interest by specifying them in the `features` fields of the `ml_monitoring.spec.DataDriftSpec` specification. If not specified, all input features in the model schema are analyzed. Additionally, you have the option to set default thresholds for categorical or numeric features, or you can specify thresholds for individual features. If the detected drift surpasses a threshold, an alert is sent through email or another notification channel.\n",
        "\n",
-        "*  **Prediction Output Drift Detection**\n",
+        "*  **Prediction output drift detection**\n",
        "\n",
-        "    Similar to input feature drift detection, prediction output drift detection aims to identify data drift in the prediction outputs.\n",
+        "    Similar to input feature drift detection, prediction output drift detection identifies data drift in the prediction outputs.\n",
        "\n",
-        "*  **Feature attribution score drift detection**\n",
+        "*   **Feature attribution drift detection**\n",
        "\n",
        "    For model outside Vertex AI which don't have the model artifact or container information, the feature attribution score drift detection is not supported."
      ]
@@ -706,7 +715,7 @@
        "id": "vPDk-0FUdudb"
      },
      "source": [
-        "In the following example, we will set the `FEATURE_DRIFT_SPEC` and `PREDICTION_OUTPUT_DRIFT_SPEC` which will be assembled in the `tabular_objective_spec`."
+        "In the following example, we set the `FEATURE_DRIFT_SPEC` and `PREDICTION_OUTPUT_DRIFT_SPEC`, which is assembled in the `tabular_objective_spec`."
      ]
    },
    {
@@ -763,7 +772,7 @@
        "*   [Notification Channel](https://cloud.google.com/monitoring/support/notification-options)\n",
        "*   [Cloud Logging](https://cloud.google.com/logging/docs?_gl=1*tdcri2*_up*MQ..&gclid=Cj0KCQjwir2xBhC_ARIsAMTXk84diOnqqpDckjOZUas26cUXUgEAgEGT9uFpz9tTvkfUjmVnRs7lQuwaAjiwEALw_wcB&gclsrc=aw.ds)  \n",
        "\n",
-        "In this notebook, we will use email as an example."
+        "In this notebook, we use email as an example."
      ]
    },
    {
@@ -804,11 +813,11 @@
        "id": "vQm3hJHZ12nK"
      },
      "source": [
-        "Let's start a monitoring job for the feature drift detection(training vs serving).\n",
-        "In this example, training data is a csv file from Google Cloud Storage and the serving data is from BigQury. We support two options for connection:\n",
+        "Let's start a monitoring job for the feature drift detection (comparing training data and serving data).\n",
+        "In this example, training data is a CSV file from Google Cloud Storage and the serving data is from BigQuery. We support two options for connection:\n",
        "\n",
-        "* table_uri: It will consume all the features from the table.\n",
-        "* query: It's SQL query, you could select the features you are interested for analysis, be sure to include the timestamp column if you'd like to specify the data window or want the continous monitoring."
+        "* table_uri: Consumes all the features from the table.\n",
+        "* query: Using SQL query, you can select the features you are interested for analysis. Be sure to include the timestamp column if you'd like to specify the data window or set up continous monitoring."
      ]
    },
    {
@@ -832,7 +841,7 @@
        "from vertexai.resources.preview import ml_monitoring\n",
        "\n",
        "TIMESTAMP = pd.Timestamp.utcnow().strftime(\"%Y%m%d%H%M%S\")\n",
-        "JOB_DISPLAY_NAME = f\"churn_model_monitoring_job_{TIMESTAMP}\"\n",
+        "JOB_DISPLAY_NAME = f\"penguins_model_monitoring_job_{TIMESTAMP}\"\n",
        "model_monitoring_job_1 = my_model_monitor.run(\n",
        "    display_name=JOB_DISPLAY_NAME,\n",
        "    baseline_dataset=TRAINING_DATASET,\n",
@@ -865,7 +874,7 @@
        "id": "bAY1AEBZ6sT6"
      },
      "source": [
-        "Let's create another model monitoring job with SQL query. Be sure to include the timestamp column if you'd like to specify the time specification."
+        "Let's create another model monitoring job using SQL query. Be sure to include the timestamp column if you'd like to specify the time specification."
      ]
    },
    {
@@ -877,7 +886,7 @@
      "outputs": [],
      "source": [
        "TIMESTAMP = pd.Timestamp.utcnow().strftime(\"%Y%m%d%H%M%S\")\n",
-        "JOB_DISPLAY_NAME = f\"churn_model_monitoring_job_{TIMESTAMP}\"\n",
+        "JOB_DISPLAY_NAME = f\"penguins_model_monitoring_job_{TIMESTAMP}\"\n",
        "model_monitoring_job_2 = my_model_monitor.run(\n",
        "    display_name=JOB_DISPLAY_NAME,\n",
        "    baseline_dataset=TRAINING_DATASET,\n",
@@ -943,7 +952,7 @@
        "id": "37SPm_majLVj"
      },
      "source": [
-        "##### Once the monitoring job is complete, should any anomalies be detected, you will receive an email similar to\n",
+        "##### Once the monitoring job is complete, should any anomalies be detected, you will receive an email similar to the following:\n",
        "\n",
        "<img src=\"https://services.google.com/fh/files/misc/place_holder_job_anomalies.png\" />"
      ]
@@ -954,7 +963,7 @@
        "id": "LPC6ZoerifM6"
      },
      "source": [
-        "#### Check GCP Console"
+        "#### Check monitoring metrics: Google Cloud Console"
      ]
    },
    {
@@ -963,7 +972,7 @@
        "id": "4Pw6Z-bVbZaE"
      },
      "source": [
-        "##### Check the \"Monitor\" tab under \"Vertex AI\"\n",
+        "To view Model Monitoring metrics in the [Google Cloud Console](https://console.cloud.google.com/vertex-ai/model-monitoring/model-monitors), go to the **Monitoring** tab under **Vertex AI.**\n",
        "\n",
        "<img src=\"https://services.google.com/fh/files/misc/place_holder_job_details.gif\" />"
      ]
@@ -974,7 +983,9 @@
        "id": "MFZcBnwsdSfa"
      },
      "source": [
-        "#### Check Output GCS bucket"
+        "#### Check monitoring metrics: Cloud Storage bucket\n",
+        "\n",
+        "Run the following to view Model Monitoring metrics stored in the Cloud Storage bucket.  "
      ]
    },
    {
@@ -1041,10 +1052,10 @@
      "source": [
        "### Step 7: Schedule Continous Model Monitoring\n",
        "\n",
-        "If you are interested at trying continous model monitoring, please following the example below to create a schedule. You can create multiple schedules for your monitor.\n",
+        "To set up continous model monitoring, follow the example below to create a schedule. You can create multiple schedules for your model monitor.\n",
        "\n",
        "\n",
-        "The following example is to monitor drift in both input features and prediction outputs. The schedule is configured to activate the model monitoring job every hour on the hour, such as at 00:00, 01:00, and so on. Each job will analyze data from the preceding one-hour window. For instance, if a job is scheduled for 6:00 am, it will analyze the data collected from 5:00 am to 6:00 am."
+        "The following example monitors drift in both input features and prediction outputs. The schedule is configured to activate the model monitoring job every hour on the hour, such as at 00:00, 01:00, and so on. Each job analyzes data from the preceding one-hour window. For instance, if a job is scheduled for 6:00 a.m., it analyzes the data collected from 5:00 a.m. to 6:00 a.m."
      ]
    },
    {
@@ -1055,9 +1066,8 @@
      },
      "outputs": [],
      "source": [
-        "CRON = (\n",
-        "    \"0 * * * *\"  # @param {type:\"string\"} Every 1 hour at :00, for example 1:00, 2:00..\n",
-        ")\n",
+        "# Every 1 hour at :00, for example 1:00, 2:00..\n",
+        "CRON = \"0 * * * *\"  # @param {type:\"string\"}\n",
        "SCHEDULE_DISPLAY_NAME = \"penguins-continous-drift-detection\""
      ]
    },
@@ -1105,7 +1115,9 @@
        "id": "MEPsa1ZhyYMc"
      },
      "source": [
-        "#### Pause Schedule"
+        "#### Pause Schedule\n",
+        "\n",
+        "Run the following to pause the model monitoring schedule:"
      ]
    },
    {
@@ -1125,7 +1137,9 @@
        "id": "h0UF10HQyazl"
      },
      "source": [
-        "#### Resume Schedule"
+        "#### Resume Schedule\n",
+        "\n",
+        "Run the following to resume a paused model monitoring schedule:"
      ]
    },
    {
@@ -1145,7 +1159,9 @@
        "id": "W-NNTbacyezA"
      },
      "source": [
-        "#### Update Schedule"
+        "#### Update Schedule\n",
+        "\n",
+        "Run the following to update the model monitoring schedule:"
      ]
    },
    {
@@ -1168,7 +1184,9 @@
        "id": "BWB251IR8Jn2"
      },
      "source": [
-        "#### You could check your schedules in Console\n",
+        "#### Check monitoring schedules in Google Cloud Console\n",
+        "\n",
+        "To check your Model Monitoring schedule in the Google Cloud Console, go to the Monitor tab under Vertex AI.\n",
        "\n",
        "<img src=\"https://services.google.com/fh/files/misc/place_holder_schedules.gif\" />"
      ]
@@ -1179,7 +1197,9 @@
        "id": "hovSbsmBmEKZ"
      },
      "source": [
-        "### Step 8: Clean Up (after job finished)"
+        "### Step 8: Clean Up\n",
+        "\n",
+        "If you no longer need your model monitoring resources, run the following to delete them:"
      ]
    },
    {
@@ -31,25 +31,27 @@
      "source": [
        "# Get started with Vertex AI Model Registry\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/model_registry/get_started_with_model_registry.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img width=\"32px\" 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%2Fmodel_registry%2Fget_started_with_model_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/blob/main/notebooks/official/model_registry/get_started_with_model_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 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/model_registry/get_started_with_model_registry.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\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",
-        "  <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/model_registry/get_started_with_model_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",
-        "    </a>\n",
-        "  </td>         \n",
-        "</table>\n",
-        "<br/><br/><br/>"
+        "</table>"
      ]
    },
    {
@@ -60,7 +62,7 @@
      "source": [
        "## Overview\n",
        "\n",
-        "This tutorial demonstrates how to use Vertex AI SDK to create and register multiple versions of a model using `Vertex AI Model Registry`.\n",
+        "This tutorial demonstrates how to use Vertex AI SDK to create and register multiple versions of a model using Vertex AI Model Registry.\n",
        "\n",
        "Learn more about [Vertex AI Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction)."
      ]
@@ -73,19 +75,19 @@
      "source": [
        "### Objective\n",
        "\n",
-        "In this tutorial, you learn how to use `Vertex AI Model Registry` to create and register multiple versions of a model.\n",
+        "In this tutorial, you learn how to use Vertex AI Model Registry to create and register multiple versions of a model.\n",
        "\n",
-        "This tutorial uses the following Google Cloud ML services and resources:\n",
+        "This tutorial uses the following Vertex AI services and resources:\n",
        "\n",
-        "- `Vertex AI Model Registry`\n",
-        "- `Vertex AI Model` resources\n",
-        "- `Vertex AI Endpoint` resources\n",
+        "- Vertex AI Model Registry\n",
+        "- Vertex AI model resources\n",
+        "- Vertex AI endpoint resources\n",
        "\n",
        "The steps performed include:\n",
        "\n",
-        "- Create and register a first version of a model to `Vertex AI Model Registry`.\n",
-        "- Create and register a second version of a model to `Vertex AI Model Registry`.\n",
-        "- Updating the model version which is the default (blessed).\n",
+        "- Create and register a first version of a model to Vertex AI Model Registry.\n",
+        "- Create and register a second version of a model to Vertex AI Model Registry.\n",
+        "- Updating the model version which is the default.\n",
        "- Deleting a model version.\n",
        "- Retraining the next model version."
      ]
@@ -30,39 +30,30 @@
      },
      "source": [
        "# Custom model batch prediction with feature filtering \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/prediction/custom_batch_prediction_feature_filter.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/prediction/custom_batch_prediction_feature_filter.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%2Fprediction%2Fcustom_batch_prediction_feature_filter.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/prediction/custom_batch_prediction_feature_filter.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/prediction/custom_batch_prediction_feature_filter.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": "overview:custom"
-      },
-      "source": [
-        "**_NOTE_**: This notebook has been tested in the following environment:\n",
-        "\n",
-        "* Python version = 3.9"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -134,15 +125,22 @@
        "to generate a cost estimate based on your projected usage."
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3b1ffd5ab768"
+      },
+      "source": [
+        "## Get Started"
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "install_aip"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook. "
+        "### Install Vertex AI SDK for Python and other required packages"
      ]
    },
    {
@@ -158,174 +156,98 @@
        "                                 google-cloud-storage \\\n",
        "                                 google-cloud-bigquery \\\n",
        "                                 pyarrow \\\n",
-        "                                 db-dtypes\n"
+        "                                 db-dtypes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "restart"
+        "id": "ff555b32bab8"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel.\n"
+        "### 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": "bzPxhxS5lugp"
+        "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": "d2qpIurSjmpT"
+        "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 following APIs: Vertex AI API, Cloud Resource Manager API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,cloudresourcemanager.googleapis.com).\n",
-        "\n",
-        "4. If you are running this notebook locally, you 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": "project_id"
+        "id": "f82e28c631cc"
      },
      "source": [
-        "#### Set your project ID\n",
+        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "* Run `gcloud config list`.\n",
-        "* Run `gcloud projects list`.\n",
-        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+        "Authenticate your environment on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "wsePm9c4jmpT"
+        "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": "6a870411c189"
+      },
+      "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": "60af003fe0ed"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# Set the project id\n",
-        "! gcloud config set project {PROJECT_ID}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "a54f9d7c1876"
-      },
-      "source": [
-        "#### Region\n",
-        "\n",
-        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "3aaadaaf9b30"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "5c0404984792"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "x2n5SeAAjmpU"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "6FDh38swjmpU"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "Nt8cEM2GjmpU"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "XUSL_JcpjmpU"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "_2zemfGvjmpU"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "TCPJ38n7jmpU"
-      },
-      "source": [
-        "**4. Service account or other**\n",
-        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -373,7 +295,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
@@ -394,7 +316,6 @@
      "outputs": [],
      "source": [
        "import json\n",
-        "import os\n",
        "\n",
        "import numpy as np\n",
        "from google.cloud import aiplatform, bigquery"
@@ -403,24 +324,23 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "750d53e37094"
+        "id": "init_aip:mbsdk,all"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
-        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
+        "Initialize the Vertex SDK for Python for your project and corresponding bucket.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "c9d3ac73dfbc"
+        "id": "856023874e71"
      },
      "outputs": [],
      "source": [
-        "# Initialize the Vertex AI SDK\n",
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -483,7 +403,7 @@
        "id": "59f24e7d2269"
      },
      "source": [
-        "## Prepare the data\n",
+        "### Prepare the data\n",
        "\n",
        "To improve the convergence of the custom deep learning model, normalize the data. To prepare for this, calculate the mean and standard deviation for each numeric column.\n",
        "\n",
@@ -570,7 +490,7 @@
        "id": "5c7732822757"
      },
      "source": [
-        "## Create a Vertex AI tabular Dataset from BigQuery dataset\n",
+        "### Create a Vertex AI tabular dataset from BigQuery dataset\n",
        "\n",
        "Your first step in training the model is to create a Vertex AI tabular dataset resource."
      ]
@@ -596,7 +516,7 @@
        "id": "train_custom_model"
      },
      "source": [
-        "## Train a model\n",
+        "### Train a model\n",
        "\n",
        "There are two ways you can train a model using a container image:\n",
        "\n",
@@ -1038,7 +958,7 @@
        "id": "make_prediction"
      },
      "source": [
-        "## Send Batch Prediction job request with feature filtering (instanceConfig field)\n",
+        "### Send Batch Prediction job request with feature filtering (instanceConfig field)\n",
        "\n",
        "Now that the model is ready, you can send batch prediction request directly from the model resource without needing to deploy the model to an endpoint. \n",
        "\n",
@@ -1436,7 +1356,7 @@
        "id": "cleanup:custom"
      },
      "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 project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
        "\n",
@@ -1463,7 +1383,7 @@
        "# Delete the model\n",
        "model.delete()\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil rm -r $BUCKET_URI\n",
        "\n",
        "# Delete the created BigQuery dataset\n",
@@ -32,25 +32,27 @@
        "# Get started with TensorFlow serving functions with Vertex AI Raw 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/prediction/get_started_with_raw_predict.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/prediction/get_started_with_raw_predict.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%2Fprediction%2Fget_started_with_raw_predict.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/prediction/get_started_with_raw_predict.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>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/prediction/get_started_with_raw_predict.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"
+        "id": "3b1ffd5ab768"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages 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": "install_aip"
+        "id": "0ad8462a9eeb"
      },
      "outputs": [],
      "source": [
@@ -150,163 +159,92 @@
    {
      "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": "before_you_begin"
+        "id": "54c5ef8a8f43"
      },
      "source": [
-        "## Before you begin\n",
-        "\n",
-        "### GPU runtime\n",
-        "\n",
-        "*Make sure you're running this notebook in a GPU runtime if you have that option. In Colab, select* **Runtime > Change Runtime Type > GPU**\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 following APIs: Vertex AI APIs, Compute Engine APIs, and Cloud Storage.](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component,storage-component.googleapis.com)\n",
-        "\n",
-        "4. If you are running this notebook locally, you need to install the [Cloud SDK]((https://cloud.google.com/sdk)).\n",
-        "\n",
-        "5. Enter your project ID in the cell below. Then run the  cell to make sure the\n",
-        "Cloud SDK uses the right project for all the commands in this notebook.\n",
-        "\n",
-        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$`."
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "project_id"
+        "id": "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": "107c51893a64"
+      },
+      "source": [
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "294fe4e5a671"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# Set the project id\n",
-        "! gcloud config set project {PROJECT_ID}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "region"
-      },
-      "source": [
-        "#### Region\n",
-        "\n",
-        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "region"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "sBCra4QMA2wR"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\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": "254614fa0c46"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "gcp_authenticate"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "a336a05c6149"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "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\"}"
      ]
    },
    {
@@ -315,7 +253,7 @@
        "id": "zgPO1eR3CYjk"
      },
      "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."
      ]
@@ -337,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."
      ]
    },
    {
@@ -348,52 +286,29 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "setup_vars"
+        "id": "b7c2a8a91840"
      },
      "source": [
-        "### Set up variables\n",
-        "\n",
-        "Next, set up some variables used throughout the tutorial.\n",
-        "### Import libraries and define constants"
+        "#### Set Google Cloud project"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "import_aip:mbsdk"
+        "id": "a4f61991b160"
      },
      "outputs": [],
      "source": [
-        "import google.cloud.aiplatform as aip"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "init_aip:mbsdk"
-      },
-      "source": [
-        "### Initialize Vertex AI SDK for Python\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, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -408,7 +323,7 @@
        "\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 Tesla K80 GPUs allocated to each VM, you would specify:\n",
        "\n",
-        "    (aip.gapic.AcceleratorType.NVIDIA_TESLA_K80, 4)\n",
+        "    (aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_K80, 4)\n",
        "\n",
        "\n",
        "Otherwise specify `(None, None)` to use a container image to run on a CPU.\n",
@@ -449,7 +364,7 @@
      },
      "outputs": [],
      "source": [
-        "TF = \"2.5\".replace(\".\", \"-\")\n",
+        "TF = \"2-5\"\n",
        "\n",
        "if DEPLOY_GPU:\n",
        "    DEPLOY_VERSION = \"tf2-gpu.{}\".format(TF)\n",
@@ -458,7 +373,7 @@
        "\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)"
@@ -548,7 +463,7 @@
      },
      "outputs": [],
      "source": [
-        "model = aip.Model.upload(\n",
+        "model = aiplatform.Model.upload(\n",
        "    display_name=\"example\",\n",
        "    artifact_uri=MODEL_DIR,\n",
        "    serving_container_image_uri=DEPLOY_IMAGE,\n",
@@ -565,13 +480,13 @@
      "source": [
        "## Creating an `Endpoint` resource\n",
        "\n",
-        "You create an `Endpoint` resource 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",
+        "You create an `Endpoint` resource using the `Endpoint.create()` method. At a minimum, you specify the display name for the endpoint. Optionally, you can specify the project and location; otherwise the settings are inherited by the values you set when you initialized the Vertex AI SDK with the `init()` method.\n",
        "\n",
        "In this example, the following parameters are specified:\n",
        "\n",
        "- `display_name`: A human readable name for the `Endpoint` resource.\n",
        "- `project`: Your project ID.\n",
-        "- `location`: Your region.\n",
+        "- `location`: Your location.\n",
        "- `labels`: (optional) User defined metadata for the `Endpoint` in the form of key/value pairs.\n",
        "\n",
        "This method returns an `Endpoint` object.\n",
@@ -587,10 +502,10 @@
      },
      "outputs": [],
      "source": [
-        "endpoint = aip.Endpoint.create(\n",
+        "endpoint = aiplatform.Endpoint.create(\n",
        "    display_name=\"example\",\n",
        "    project=PROJECT_ID,\n",
-        "    location=REGION,\n",
+        "    location=LOCATION,\n",
        "    labels={\"your_key\": \"your_value\"},\n",
        ")\n",
        "\n",
@@ -717,10 +632,10 @@
      },
      "outputs": [],
      "source": [
-        "API_ENDPOINT = \"{}-aiplatform.googleapis.com\".format(REGION)\n",
+        "API_ENDPOINT = \"{}-aiplatform.googleapis.com\".format(LOCATION)\n",
        "client_options = {\"api_endpoint\": API_ENDPOINT}\n",
        "\n",
-        "pred_client = aip.gapic.PredictionServiceClient(client_options=client_options)\n",
+        "pred_client = aiplatform.gapic.PredictionServiceClient(client_options=client_options)\n",
        "\n",
        "response = pred_client.raw_predict(req)\n",
        "print(response)"
@@ -748,8 +663,6 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
        "delete_bucket = False\n",
        "delete_model = True\n",
        "delete_endpoint = True\n",
@@ -767,7 +680,7 @@
        "    except Exception as e:\n",
        "        print(e)\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil rm -rf {BUCKET_URI}"
      ]
    }
@@ -4,6 +4,7 @@
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
+        "cellView": "form",
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
@@ -32,25 +33,27 @@
        "# Get started with TensorFlow Serving with Vertex AI 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/prediction/get_started_with_tf_serving.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/prediction/get_started_with_tf_serving.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%2Fprediction%2Fget_started_with_tf_serving.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/prediction/get_started_with_tf_serving.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>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/prediction/get_started_with_tf_serving.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>"
      ]
    },
    {
@@ -63,7 +66,7 @@
        "\n",
        "This tutorial demonstrates how to serve predictions from a `Vertex AI Endpoint` with `TensorFlow Serving` serving binary.\n",
        "\n",
-        "Learn more about [Get predictions from a custom trained model](https://cloud.google.com/vertex-ai/docs/predictions/get-predictions)."
+        "Learn more about [getting predictions from a custom trained model](https://cloud.google.com/vertex-ai/docs/predictions/get-predictions)."
      ]
    },
    {
@@ -76,22 +79,22 @@
        "\n",
        "In this tutorial, you learn how to use `Vertex AI Prediction` on a `Vertex AI Endpoint` resource with `TensorFlow Serving` serving binary.\n",
        "\n",
-        "This tutorial uses the following Google Cloud ML services and resources:\n",
+        "This tutorial uses the following Vertex AI services and resources:\n",
        "\n",
        "- `Vertex AI Prediction`\n",
        "- `Vertex AI Batch Prediction`\n",
        "- `Vertex AI Models`\n",
        "- `Vertex AI Endpoints`\n",
        "\n",
-        "The steps performed include:\n",
+        "You perform the following steps:\n",
        "\n",
        "- Download a pretrained image classification model from TensorFlow Hub.\n",
        "- Create a serving function to receive compressed image data, and output decomopressed preprocessed data for the model input.\n",
-        "- Upload the TensorFlow Hub model and serving function as a `Vertex AI Model` resource.\n",
-        "- Creating an `Endpoint` resource.\n",
-        "- Deploying the `Model` resource to an `Endpoint` resource with `TensorFlow Serving` serving binary.\n",
-        "- Make an online prediction to the `Model` resource instance deployed to the `Endpoint` resource.\n",
-        "- Make a batch prediction to the `Model` resource instance."
+        "- Upload the TensorFlow Hub model and serving function as a Vertex AI model resource.\n",
+        "- Create a Vertex AI endpoint resource.\n",
+        "- Deploy the model resource to the endpoint resource with `TensorFlow Serving` serving binary.\n",
+        "- Make an online prediction with the deployed model.\n",
+        "- Make a batch prediction with the model."
      ]
    },
    {
@@ -130,50 +133,101 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "install_aip"
+        "id": "3b1ffd5ab768"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages to execute this notebook."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "install_aip"
-      },
-      "outputs": [],
-      "source": [
-        "! pip3 install --upgrade google-cloud-aiplatform -q\n",
-        "! pip3 install --upgrade google-cloud-pipeline-components  -q\n",
-        "! pip3 install --upgrade tensorflow -q\n",
-        "! pip3 install tensorflow-hub -q"
+        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "58707a750154"
+        "id": "install_aip"
      },
      "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": "f200f10a1da3"
+        "id": "install_aip"
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "! pip3 install --upgrade google-cloud-aiplatform \\\n",
+        "                        google-cloud-pipeline-components \\\n",
+        "                        tensorflow==2.15.1 \\\n",
+        "                        tensorflow-hub -q"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ff555b32bab8"
+      },
+      "source": [
+        "### Restart runtime (Colab only)\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "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": "54c5ef8a8f43"
+      },
+      "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()"
      ]
    },
    {
@@ -182,12 +236,9 @@
        "id": "WReHDGG5g0XY"
      },
      "source": [
-        "#### Set your project ID\n",
+        "### Set Google Cloud project information\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)."
      ]
    },
    {
@@ -199,92 +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": [
-        "#### Set the region\n",
-        "\n",
-        "**Optional**: Update the 'REGION' variable to specify the region that you want to use. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "nsN5NJKSu-GU"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "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."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -315,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."
      ]
    },
    {
@@ -326,7 +292,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
@@ -335,10 +301,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"
      ]
    },
    {
@@ -349,11 +312,11 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
+        "import json\n",
        "\n",
-        "import google.cloud.aiplatform as aip\n",
        "import tensorflow as tf\n",
-        "import tensorflow_hub as hub"
+        "import tensorflow_hub as hub\n",
+        "from google.cloud import aiplatform"
      ]
    },
    {
@@ -364,7 +327,8 @@
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
-        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
+        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket.\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). "
      ]
    },
    {
@@ -375,7 +339,7 @@
      },
      "outputs": [],
      "source": [
-        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
+        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -388,16 +352,16 @@
        "\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",
+        "    (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 about [hardware accelerator support for your region](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators).\n",
        "\n",
-        "*Note*: TF releases before 2.3 for GPU support will fail to load the custom model in this tutorial. It is a known issue and fixed in TF 2.3. This is caused by static graph ops that are generated in the serving function. If you encounter this issue on your own custom models, use a container image for TF 2.3 with GPU support."
+        "**Note**: TF releases before 2.3 for GPU support fail to load the custom model in this tutorial. It's a known issue and fixed in TF 2.3. This is caused by static graph ops that are generated in the serving function. If you encounter this issue on your own custom models, use a container image for TF 2.3 with GPU support."
      ]
    },
    {
@@ -421,14 +385,14 @@
        "\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",
        "     - `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 do not support GPUs."
      ]
    },
    {
@@ -456,7 +420,7 @@
        "\n",
        "You must 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)."
+        "Learn more about [enabling services](https://cloud.google.com/artifact-registry/docs/enable-service)."
      ]
    },
    {
@@ -467,11 +431,7 @@
      },
      "outputs": [],
      "source": [
-        "! gcloud services enable artifactregistry.googleapis.com\n",
-        "\n",
-        "if os.getenv(\"IS_TESTING\"):\n",
-        "    ! sudo apt-get update --yes && sudo apt-get --only-upgrade --yes install google-cloud-sdk-cloud-run-proxy google-cloud-sdk-harbourbridge google-cloud-sdk-cbt google-cloud-sdk-gke-gcloud-auth-plugin google-cloud-sdk-kpt google-cloud-sdk-local-extract google-cloud-sdk-minikube google-cloud-sdk-app-engine-java google-cloud-sdk-app-engine-go google-cloud-sdk-app-engine-python google-cloud-sdk-spanner-emulator google-cloud-sdk-bigtable-emulator google-cloud-sdk-nomos google-cloud-sdk-package-go-module google-cloud-sdk-firestore-emulator kubectl google-cloud-sdk-datastore-emulator google-cloud-sdk-app-engine-python-extras google-cloud-sdk-cloud-build-local google-cloud-sdk-kubectl-oidc google-cloud-sdk-anthos-auth google-cloud-sdk-app-engine-grpc google-cloud-sdk-pubsub-emulator google-cloud-sdk-datalab google-cloud-sdk-skaffold google-cloud-sdk google-cloud-sdk-terraform-tools google-cloud-sdk-config-connector\n",
-        "    ! gcloud components update --quiet"
+        "! gcloud services enable artifactregistry.googleapis.com"
      ]
    },
    {
@@ -499,7 +459,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"
      ]
@@ -523,7 +483,7 @@
      },
      "outputs": [],
      "source": [
-        "! gcloud auth configure-docker {REGION}-docker.pkg.dev --quiet"
+        "! gcloud auth configure-docker {LOCATION}-docker.pkg.dev --quiet"
      ]
    },
    {
@@ -558,7 +518,7 @@
        "# Executes in Vertex AI Workbench\n",
        "if DEPLOY_GPU:\n",
        "    DEPLOY_IMAGE = (\n",
-        "        f\"{REGION}-docker.pkg.dev/\"\n",
+        "        f\"{LOCATION}-docker.pkg.dev/\"\n",
        "        + PROJECT_ID\n",
        "        + f\"/{PRIVATE_REPO}\"\n",
        "        + \"/tf_serving:gpu\"\n",
@@ -566,7 +526,7 @@
        "    TF_IMAGE = \"tensorflow/serving:2.5.4-gpu\"\n",
        "else:\n",
        "    DEPLOY_IMAGE = (\n",
-        "        f\"{REGION}-docker.pkg.dev/\"\n",
+        "        f\"{LOCATION}-docker.pkg.dev/\"\n",
        "        + PROJECT_ID\n",
        "        + f\"/{PRIVATE_REPO}\"\n",
        "        + \"/tf_serving:cpu\"\n",
@@ -613,7 +573,7 @@
        "dockerd -b none --iptables=0 -l warn &\n",
        "for i in $(seq 5); do [ ! -S \"/var/run/docker.sock\" ] && sleep 2 || break; done\n",
        "docker pull $3\n",
-        "docker tag tensorflow/serving $2\n",
+        "docker tag $3 $2\n",
        "docker push $2\n",
        "kill $(jobs -p)"
      ]
@@ -626,7 +586,7 @@
      "source": [
        "## Get pretrained model from TensorFlow Hub\n",
        "\n",
-        "For demonstration purposes, this tutorial uses a pretrained model from TensorFlow Hub (TFHub), which is then uploaded to a `Vertex AI Model` resource. Once you have a `Vertex AI Model` resource, the model can be deployed to a `Vertex AI Endpoint` resource.\n",
+        "For demonstration purposes, this tutorial uses a pretrained model from TensorFlow Hub (TFHub), which is then uploaded to a Vertex AI model resource. Once you have a Vertex AI model resource created, the model can then be deployed to a Vertex AI endpoint resource.\n",
        "\n",
        "### Download the pretrained model\n",
        "\n",
@@ -660,7 +620,7 @@
        "\n",
        "At this point, the model is in memory. Next, you save the model artifacts to a Cloud Storage location.\n",
        "\n",
-        "*Note:* For TF Serving, the MODEL_DIR must end in a subfolder that is a number, e.g., 1."
+        "**Note:** For TF Serving, the MODEL_DIR must end in a subfolder that is a number, e.g., 1."
      ]
    },
    {
@@ -683,7 +643,7 @@
      "source": [
        "## Upload the model for serving\n",
        "\n",
-        "Next, you will upload your TF.Keras model from the custom job to Vertex `Model` service, which will create a Vertex `Model` resource for your custom 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 TF.Keras model from the custom job to Vertex AI model service, which creates a model resource for your custom 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",
@@ -700,7 +660,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 can't 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 isn't supported."
      ]
    },
    {
@@ -713,7 +673,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",
@@ -778,9 +738,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 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 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."
      ]
    },
    {
@@ -805,9 +765,9 @@
        "id": "e8ce91147c93"
      },
      "source": [
-        "### Upload the TensorFlow Hub model to a `Vertex AI Model` resource\n",
+        "### Upload the TensorFlow Hub model to a Vertex AI model resource\n",
        "\n",
-        "Finally, you upload the model artifacts from the TFHub model and serving function into a `Vertex AI Model` resource. Since you are using a non Google pre-built serving binary -- i.e., TensorFlow Serving, you need to specify the following additional serving configuration settings:\n",
+        "Finally, you upload the model artifacts from the TFHub model and serving function into a Vertex AI model resource. Since you're using a non Google pre-built serving binary -- i.e., TensorFlow Serving, you need to specify the following additional serving configuration settings:\n",
        "\n",
        "- `serving_container_command`: The serving binary (HTTP Server) to start up.\n",
        "- `serving_container_args`: The arguments to pass to the serving binary. For TensorFlow Serving, the required arguments are:\n",
@@ -815,15 +775,15 @@
        "    - `--model_base_name`: Where to store the model artifacts in the container. The Vertex service sets the variable $(AIP_STORAGE_URI) to where the service installed the model artifacts in the container.\n",
        "    - `--rest_api_port`: The port to which to send REST based prediction requests. Can either be 8080 or 8501 (default for TensorFlow Serving).\n",
        "    - `--port`: The port to which to send gRPC based prediction requests. Should be 8500 for TensorFlow Serving.\n",
-        "- `serving_container_health_route`: The URL for the service to periodically ping for a response to verify that the serving binary is running. For TensorFlow Serving, this will be /v1/models/\\<model_name\\>.\n",
-        "- `serving_container_predict_route`: The URL for the service to route REST-based prediction requests to. For TF Serving, this will be /v1/models/[model_name]:predict.\n",
+        "- `serving_container_health_route`: The URL for the service to periodically ping for a response to verify that the serving binary is running. For TensorFlow Serving, this is /v1/models/\\<model_name\\>.\n",
+        "- `serving_container_predict_route`: The URL for the service to route REST-based prediction requests to. For TF Serving, this is /v1/models/[model_name]:predict.\n",
        "- `serving_container_ports`: A list of ports for the HTTP server to listen for requests.\n",
        "\n",
-        "Uploading a model into a Vertex Model resource returns a long running operation, since it may take a few moments. \n",
+        "Uploading a model into a Vertex model resource returns a long running operation, since it may take a few moments. \n",
        "\n",
-        "*Note:* You drop the ending number subfolder (e.g., /1) from the model path to upload. The Vertex service will upload the parent folder above the subfolder with the model artifacts -- which is what TensorFlow Serving binary expects.\n",
+        "**Note:** You drop the ending number subfolder (e.g., /1) from the model path to upload. The Vertex service uploads the parent folder above the subfolder with the model artifacts -- which is what TensorFlow Serving binary expects.\n",
        "\n",
-        "*Note:* When you upload the model artifacts to a `Vertex AI Model` resource, you specify the corresponding deployment container image."
+        "**Note:** When you upload the model artifacts to a Vertex AI model resource, you specify the corresponding deployment container image."
      ]
    },
    {
@@ -836,7 +796,7 @@
      "source": [
        "MODEL_NAME = \"example_\"\n",
        "\n",
-        "model = aip.Model.upload(\n",
+        "model = aiplatform.Model.upload(\n",
        "    display_name=\"example_\",\n",
        "    artifact_uri=MODEL_DIR[:-2],\n",
        "    serving_container_image_uri=DEPLOY_IMAGE,\n",
@@ -862,18 +822,18 @@
        "id": "628de0914ba1"
      },
      "source": [
-        "## Creating an `Endpoint` resource\n",
+        "## Create a Vertex AI endpoint resource\n",
        "\n",
-        "You create an `Endpoint` resource 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",
+        "You create a Vertex AI endpoint resource 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",
        "\n",
        "In this example, the following parameters are specified:\n",
        "\n",
-        "- `display_name`: A human readable name for the `Endpoint` resource.\n",
+        "- `display_name`: A human readable name for the endpoint resource.\n",
        "- `project`: Your project ID.\n",
        "- `location`: Your region.\n",
-        "- `labels`: (optional) User defined metadata for the `Endpoint` in the form of key/value pairs.\n",
+        "- `labels`: (optional) User defined metadata for the endpoint in the form of key/value pairs.\n",
        "\n",
-        "This method returns an `Endpoint` object.\n",
+        "This method returns a Vertex AI endpoint object.\n",
        "\n",
        "Learn more about [Vertex AI Endpoints](https://cloud.google.com/vertex-ai/docs/predictions/deploy-model-api)."
      ]
@@ -886,10 +846,10 @@
      },
      "outputs": [],
      "source": [
-        "endpoint = aip.Endpoint.create(\n",
+        "endpoint = aiplatform.Endpoint.create(\n",
        "    display_name=\"example_\",\n",
        "    project=PROJECT_ID,\n",
-        "    location=REGION,\n",
+        "    location=LOCATION,\n",
        "    labels={\"your_key\": \"your_value\"},\n",
        ")\n",
        "\n",
@@ -902,13 +862,13 @@
        "id": "ca3fa3f6a894"
      },
      "source": [
-        "## Deploying `Model` resources to an `Endpoint` resource.\n",
+        "## Deploy model to the endpoint\n",
        "\n",
-        "You can deploy one of more `Vertex AI Model` resource instances to the same endpoint. Each `Vertex AI Model` resource that is deployed will have its own deployment container for the serving binary. \n",
+        "You can deploy one or more Vertex AI model resource instances to the same endpoint. Each model resource that is deployed has its own deployment container for the serving binary. \n",
        "\n",
-        "*Note:* For this example, you specified the deployment container for the TFHub model in the previous step of uploading the model artifacts to a `Vertex AI Model` resource.\n",
+        "**Note:** For this example, you specified the deployment container for the TFHub model in the previous step of uploading the model artifacts to a Vertex AI model resource.\n",
        "\n",
-        "In the next example, you deploy the `Vertex AI Model` resource to a `Vertex AI Endpoint` resource. The `Vertex AI Model` resource already has defined for it the deployment container image. To deploy, you specify the following additional configuration settings:\n",
+        "In the next example, you deploy the Vertex AI model resource to a Vertex AI endpoint resource. The model resource already has a container image defined for deployment. To deploy, you specify the following additional configuration settings:\n",
        "\n",
        "- The machine type.\n",
        "- The (if any) type and number of GPUs.\n",
@@ -916,7 +876,7 @@
        "\n",
        "In this example, you deploy the model with the minimal amount of specified parameters, as follows:\n",
        "\n",
-        "- `model`: The `Model` resource.\n",
+        "- `model`: The Vertex AI model resource.\n",
        "- `deployed_model_displayed_name`: The human readable name for the deployed model instance.\n",
        "- `machine_type`: The machine type for each VM instance.\n",
        "\n",
@@ -948,7 +908,7 @@
      "source": [
        "### Prepare test data for prediction\n",
        "\n",
-        "Next, you will load a compressed JPEG image into memory and then base64 encode it. For demonstration purposes, you use an image from the Flowers dataset."
+        "Next, you load a compressed JPEG image into memory and then base64 encode it. For demonstration purposes, you use an image from the Flowers dataset."
      ]
    },
    {
@@ -985,7 +945,7 @@
      "source": [
        "### 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 do online predictions by sending prediction requests to the Endpoint resource.\n",
        "\n",
        "#### Request\n",
        "\n",
@@ -1028,11 +988,11 @@
        "id": "11e16f54bc90"
      },
      "source": [
-        "## Introduction to Batch Prediction\n",
+        "## Introduction to batch prediction\n",
        "\n",
        "Batch prediction provides the ability to do offline batch processing of large amounts of prediction requests. Resources are only provisioned during the batch process and then deprovisioned when the batch request is completed. The results are stored in Cloud Storage, in contrast to online prediction where the results are returned as a HTTP response packet.\n",
        "\n",
-        "The input format for your batch job is dependent on the format supported by your model server. Foremost, the web server in your model server must support a JSONL format, which the web server will convert to a format support either directly by the model input intertace or a serving function interface. For batch prediction, this JSONL format is referred to as the `pivot` format.\n",
+        "The input format for your batch job is dependent on the format supported by your model server. Foremost, the web server in your model server must support a JSONL format, which the web server converts to a format support either directly by the model input interface or a serving function interface. For batch prediction, this JSONL format is referred to as the `pivot` format.\n",
        "\n",
        "### Input format for batch prediction jobs\n",
        "\n",
@@ -1063,7 +1023,7 @@
        "\n",
        "**CSV**\n",
        "\n",
-        "The csv header in the first line will always be ignored. String fields are required to be double quoted explicitly, otherwise the row is discarded and parsing error messages are outputted to error files. Non-quoted values are always transferred as floats.\n",
+        "The csv header in the first line is always ignored. String fields are required to be double quoted explicitly, otherwise the row is discarded and parsing error messages are outputted to error files. Non-quoted values are always transferred as floats.\n",
        "\n",
        "    col1,col2,col3\n",
        "    1,3,\"cat1\"\n",
@@ -1120,7 +1080,7 @@
        "id": "9ccd46a186da"
      },
      "source": [
-        "### Make the batch input file\n",
+        "## Make the batch input file\n",
        "\n",
        "Next, make a batch input file, which you store in your local Cloud Storage bucket. For custom models, you format the batch input file in JSONL format. Each JSON object entry in the JSONL file is specified in the same format as you specified for the online prediction request.\n",
        "\n",
@@ -1148,9 +1108,7 @@
      "outputs": [],
      "source": [
        "# For demonstration purposes, you write the same image (instance[0]) request twice to the JSONL file.\n",
-        "# You will receive back two predictions, one for each instance.\n",
-        "\n",
-        "import json\n",
+        "# You receive back two predictions, one for each instance.\n",
        "\n",
        "with open(\"test.jsonl\", \"w\") as f:\n",
        "    json.dump(instances[0], f)\n",
@@ -1166,7 +1124,7 @@
        "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",
        "\n",
@@ -1178,7 +1136,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."
+        "- `sync`: If set to True, the call gets blocked while waiting for the asynchronous batch job to complete."
      ]
    },
    {
@@ -1214,9 +1172,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."
      ]
    },
    {
@@ -1236,7 +1194,7 @@
        "id": "get_batch_prediction:mbsdk,custom,icn"
      },
      "source": [
-        "### Get the predictions\n",
+        "## Get the predictions\n",
        "\n",
        "Next, get the results from the completed batch prediction job.\n",
        "\n",
@@ -1254,8 +1212,6 @@
      },
      "outputs": [],
      "source": [
-        "import json\n",
-        "\n",
        "bp_iter_outputs = batch_predict_job.iter_outputs()\n",
        "\n",
        "prediction_results = list()\n",
@@ -1295,33 +1251,23 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
+        "# set delete_bucket to True to delete your bucket\n",
        "delete_bucket = False\n",
-        "delete_model = True\n",
-        "delete_endpoint = True\n",
-        "delete_batch_job = True\n",
        "\n",
-        "if delete_endpoint:\n",
-        "    try:\n",
-        "        endpoint.undeploy_all()\n",
-        "        endpoint.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
+        "# undeploy the model from the endpoint\n",
+        "endpoint.undeploy_all()\n",
        "\n",
-        "if delete_model:\n",
-        "    try:\n",
-        "        model.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
+        "# delete the endpoint\n",
+        "endpoint.delete()\n",
        "\n",
+        "# delete the model\n",
+        "model.delete()\n",
+        "\n",
+        "# delete the batch job\n",
+        "batch_predict_job.delete()\n",
+        "\n",
+        "# delete the bucket\n",
        "if delete_bucket:\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 -rf {BUCKET_URI}"
      ]
    }
@@ -33,25 +33,27 @@
        "\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/prediction/pytorch_image_classification_with_prebuilt_serving_containers.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/prediction/pytorch_image_classification_with_prebuilt_serving_containers.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%2Fprediction%2Fpytorch_image_classification_with_prebuilt_serving_containers.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/prediction/pytorch_image_classification_with_prebuilt_serving_containers.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>"
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/prediction/pytorch_image_classification_with_prebuilt_serving_containers.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"
      ]
    },
    {
@@ -140,19 +142,26 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "i7EUnXsZhAGF"
+        "id": "7cc240e9c1a6"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook."
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9d6f71d52f3d"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "2b4ef9b72d43"
+        "id": "39883d03df60"
      },
      "outputs": [],
      "source": [
@@ -167,155 +176,92 @@
    {
      "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": "e669f8088ac3"
      },
      "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**, 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": "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)."
+      ]
+    },
+    {
+      "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}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "region"
-      },
-      "source": [
-        "#### Region\n",
-        "\n",
-        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "tTy1gX11kCJY"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "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()"
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -324,7 +270,7 @@
        "id": "zgPO1eR3CYjk"
      },
      "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."
      ]
@@ -346,7 +292,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."
      ]
    },
    {
@@ -357,7 +303,31 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "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": "BnBAXs5XkCJZ"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -385,33 +355,10 @@
        "\n",
        "import tensorflow as tf\n",
        "import torch\n",
-        "from google.cloud import aiplatform\n",
        "from PIL import Image\n",
        "from torchvision import models"
      ]
    },
-    {
-      "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": "BnBAXs5XkCJZ"
-      },
-      "outputs": [],
-      "source": [
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -34,23 +34,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/prediction/pytorch_train_deploy_models_with_prebuilt_containers.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/prediction/pytorch_train_deploy_models_with_prebuilt_containers.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%2Fprediction%2Fpytorch_train_deploy_models_with_prebuilt_containers.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/prediction/pytorch_train_deploy_models_with_prebuilt_containers.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/prediction/pytorch_train_deploy_models_with_prebuilt_containers.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>"
      ]
    },
@@ -80,13 +85,13 @@
        "\n",
        "In this tutorial, you learn how to build, train and deploy a PyTorch image classification model using prebuilt containers for custom training and prediction.\n",
        "\n",
-        "This tutorial uses the following Google Cloud ML services and resources:\n",
+        "This tutorial uses the following Vertex AI services and resources:\n",
        "\n",
-        "- `Vertex AI Training` service\n",
-        "- `Vertex AI Prediction` service\n",
-        "- `Vertex AI Model Registry`\n",
-        "- `Vertex AI Model` resources\n",
-        "- `Vertex AI Endpoint` resources\n",
+        "- Vertex AI Training service\n",
+        "- Vertex AI Prediction service\n",
+        "- Vertex AI Model Registry\n",
+        "- Vertex AI model resources\n",
+        "- Vertex AI endpoint resources\n",
        "\n",
        "\n",
        "The steps performed include:\n",
@@ -135,9 +140,16 @@
        "id": "i7EUnXsZhAGF"
      },
      "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"
      ]
    },
    {
@@ -160,7 +172,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."
      ]
    },
    {
@@ -171,46 +185,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": "BF1j6f9HApxa"
+        "id": "ee775571c2b5"
      },
      "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": "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": "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)."
      ]
    },
    {
@@ -222,114 +254,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": "gkaCh23Pg2h8"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "sBCra4QMA2wR"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "74ccc9e52986"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "de775a3773ba"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "254614fa0c46"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ef21552ccea8"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "603adbbf0532"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "f6b2ccc891ed"
-      },
-      "source": [
-        "**4. Service account or other**\n",
-        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "zgPO1eR3CYjk"
-      },
-      "source": [
-        "### Create a Cloud Storage bucket\n",
-        "\n",
-        "Create a storage bucket to store intermediate artifacts such as datasets."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -343,13 +268,24 @@
        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
+    {
+      "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": "markdown",
      "metadata": {
        "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."
      ]
    },
    {
@@ -360,7 +296,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": "a4f61991b160"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -384,34 +342,11 @@
        "import io\n",
        "import os\n",
        "\n",
-        "from google.cloud import aiplatform\n",
        "from matplotlib import pyplot as plt\n",
        "from torch.utils.data import DataLoader\n",
        "from torchvision import datasets, transforms"
      ]
    },
-    {
-      "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": "4KyAr2MHg2h9"
-      },
-      "outputs": [],
-      "source": [
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -454,7 +389,7 @@
      },
      "outputs": [],
      "source": [
-        "APP_NAME = \"pytorch-101-trainer\"\n",
+        "APP_NAME = \"pytorch_101_trainer\"\n",
        "PYTHON_PACKAGE_APPLICATION_DIR = \"python_package\"\n",
        "\n",
        "source_package_file_name = (\n",
@@ -501,7 +436,7 @@
        "]\n",
        "\n",
        "setup(\n",
-        "    name='pytorch-101-trainer',\n",
+        "    name='pytorch_101_trainer',\n",
        "    version='0.1',\n",
        "    install_requires=REQUIRED_PACKAGES,\n",
        "    packages=find_packages(),\n",
@@ -684,7 +619,6 @@
      "outputs": [],
      "source": [
        "! cd {PYTHON_PACKAGE_APPLICATION_DIR} && python3 setup.py sdist --formats=gztar\n",
-        "\n",
        "! gsutil cp {source_package_file_name} {python_package_gcs_uri}\n",
        "\n",
        "print(f\"Python source distribution package location: {python_package_gcs_uri}\")"
@@ -748,7 +682,7 @@
        "ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
        "ACCELERATOR_COUNT = 1\n",
        "\n",
-        "EPOCHS = 2\n",
+        "EPOCHS = 1\n",
        "BATCH_SIZE = 32\n",
        "MODEL_DIR = f\"{BUCKET_URI}/{JOB_DISPLAY_NAME}\"\n",
        "\n",
@@ -988,7 +922,7 @@
      "source": [
        "### Upload model for deployment\n",
        "\n",
-        "Next, you upload the model artifacts to `Vertex AI Model Registry`, which will create a `Vertex AI Model` resource for your model. This tutorial uses the PyTorch v1.11 container, but for your own use case, you can choose from the list of [PyTorch prebuilt containers](https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers#pytorch)."
+        "Next, you upload the model artifacts to `Vertex AI Model Registry`, which creates a Vertex AI model resource for your model. This tutorial uses the PyTorch v1.11 container, but for your own use case, you can choose from the list of [PyTorch prebuilt containers](https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers#pytorch)."
      ]
    },
    {
@@ -1016,7 +950,7 @@
      "source": [
        "### Deploy model for prediction\n",
        "\n",
-        "Next, deploy your model for online prediction. You set the variable `DEPLOY_COMPUTE` to configure the machine type for the [compute resources](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute) you will use for prediction."
+        "Next, deploy your model for online prediction. Set the variable `DEPLOY_COMPUTE` to configure the machine type for the [compute resources](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute) which is used for prediction."
      ]
    },
    {
@@ -1152,7 +1086,7 @@
        "uploaded_model.delete()\n",
        "\n",
        "# Delete Cloud Storage objects that were created\n",
-        "delete_bucket = False\n",
+        "delete_bucket = True\n",
        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
@@ -39,7 +39,7 @@
        "  </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%2Fray_on_vertex_ai%2Fget_started_with_pytorch_rov.ipynb\">\n",
-        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\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",
@@ -41,13 +41,7 @@
        "  </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%2Fray_on_vertex_ai%2Fray_cluster_management.ipynb\">\n",
-        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" 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/ray_on_vertex_ai/ray_cluster_management.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"> <br>\n",
-        "      View on GitHub\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",
@@ -55,7 +49,13 @@
        "      <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/ray_on_vertex_ai/ray_cluster_management.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>"
      ]
    },
@@ -82,7 +82,7 @@
        "\n",
        "In this tutorial, you learn how to create a cluster, list existing clusters, get a cluster, update (manually scaling) a cluster, and delete a cluster.\n",
        "\n",
-        "This tutorial uses the following Google Cloud ML services and resources:\n",
+        "This tutorial uses the following Vertex AI services and resources:\n",
        "\n",
        "- [Ray on Vertex AI](https://cloud.google.com/vertex-ai/docs/open-source/ray-on-vertex-ai/overview)\n",
        "\n",
@@ -113,15 +113,22 @@
        "to generate a cost estimate based on your projected usage."
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8925ff9e165e"
+      },
+      "source": [
+        "## Getting Started"
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i7EUnXsZhAGF"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook. \n"
+        "### Install Vertex AI SDK and other required packages"
      ]
    },
    {
@@ -141,7 +148,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."
      ]
    },
    {
@@ -152,31 +161,53 @@
      },
      "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": "7b49231643e4"
      },
      "source": [
-        "## Before you begin\n",
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Please wait until it is finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "7176ea64999b"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
-        "### Set up your Google Cloud project\n",
+        "Authenticate your environment on Google Colab.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7de6ef0fac42"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
        "\n",
-        "**The following steps are required, regardless of your notebook environment.**\n",
+        "if \"google.colab\" in sys.modules:\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 credit towards your compute/storage costs.\n",
+        "    from google.colab import auth\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)."
+        "    auth.authenticate_user()"
      ]
    },
    {
@@ -185,12 +216,9 @@
        "id": "WReHDGG5g0XY"
      },
      "source": [
-        "#### Set your project ID\n",
+        "### Set Google Cloud project information and initialize Vertex AI SDK\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 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)."
      ]
    },
    {
@@ -202,147 +230,39 @@
      "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}"
+        "# Retrieve the project number\n",
+        "PROJECT_NUMBER = !gcloud projects list --filter=\"PROJECT_ID:'{PROJECT_ID}'\" --format='value(PROJECT_NUMBER)'\n",
+        "PROJECT_NUMBER = PROJECT_NUMBER[0]\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "region"
+        "id": "init_aip:mbsdk,all"
      },
      "source": [
-        "#### Region\n",
+        "### Set network information\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)."
+        "[Set up a VPC peering network](https://cloud.google.com/vertex-ai/docs/general/vpc-peering) and private services connection to access Vertex AI."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "region"
+        "id": "3927074343e3"
      },
      "outputs": [],
      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "sBCra4QMA2wR"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "74ccc9e52986"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "de775a3773ba"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "254614fa0c46"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ef21552ccea8"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "603adbbf0532"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "f6b2ccc891ed"
-      },
-      "source": [
-        "**4. Service account or other**\n",
-        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "zgPO1eR3CYjk"
-      },
-      "source": [
-        "### Create a Cloud Storage bucket\n",
-        "\n",
-        "Create a storage bucket to store intermediate artifacts such as datasets.\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}*"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "MzGDU7TWdts_"
-      },
-      "outputs": [],
-      "source": [
-        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "-EcIXiGsCePi"
-      },
-      "source": [
-        "**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}"
+        "VPC_NETWORK = \"[your-network-name]\"  # @param {type:\"string\"}\n",
+        "VPC_NETWORK_FULL = \"projects/{}/global/networks/{}\".format(PROJECT_NUMBER, VPC_NETWORK)\n",
+        "VPC_NETWORK_FULL"
      ]
    },
    {
@@ -362,49 +282,7 @@
      },
      "outputs": [],
      "source": [
-        "import vertex_ray\n",
-        "from google.cloud import aiplatform"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "init_aip:mbsdk,all"
-      },
-      "source": [
-        "### Initialize Vertex AI SDK for Python\n",
-        "\n",
-        "Initialize the Vertex AI SDK for Python for your project.\n",
-        "\n",
-        "[Set up a VPC peering network](https://cloud.google.com/vertex-ai/docs/general/vpc-peering) and private services connection to access Vertex AI."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "3927074343e3"
-      },
-      "outputs": [],
-      "source": [
-        "# Retrieve the project number\n",
-        "PROJECT_NUMBER = !gcloud projects list --filter=\"PROJECT_ID:'{PROJECT_ID}'\" --format='value(PROJECT_NUMBER)'\n",
-        "PROJECT_NUMBER = PROJECT_NUMBER[0]\n",
-        "\n",
-        "VPC_NETWORK = \"[your-network-name]\"\n",
-        "VPC_NETWORK_FULL = \"projects/{}/global/networks/{}\".format(PROJECT_NUMBER, VPC_NETWORK)\n",
-        "VPC_NETWORK_FULL"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "init_aip:mbsdk,all"
-      },
-      "outputs": [],
-      "source": [
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
+        "import vertex_ray"
      ]
    },
    {
@@ -435,7 +313,7 @@
        "    vertex_ray.Resources(\n",
        "        machine_type=\"n1-standard-8\",\n",
        "        node_count=2,  # Can be > 1\n",
-        "        accelerator_type=\"NVIDIA_TESLA_K80\",\n",
+        "        accelerator_type=\"NVIDIA_TESLA_T4\",\n",
        "        accelerator_count=1,\n",
        "    )\n",
        "]\n",
@@ -614,15 +492,8 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
        "# Delete the cluster\n",
-        "vertex_ray.delete_ray_cluster(cluster.cluster_resource_name)\n",
-        "\n",
-        "# Delete Cloud Storage objects that were created\n",
-        "delete_bucket = False\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
-        "    ! gsutil -m rm -r $BUCKET_URI"
+        "vertex_ray.delete_ray_cluster(cluster.cluster_resource_name)"
      ]
    }
  ],
@@ -32,25 +32,30 @@
        "# Vertex AI TensorBoard custom training with custom container\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\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/tensorboard_custom_training_with_custom_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=\"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%2Ftensorboard%2Ftensorboard_custom_training_with_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://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/tensorboard/tensorboard_custom_training_with_custom_container.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/tensorboard/tensorboard_custom_training_with_custom_container.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",
-        "</table>\n",
-        "<br/><br/><br/>"
+        "  </td>                                                                                               \n",
+        "</table>"
      ]
    },
    {
@@ -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/tensorboard/tensorboard_hyperparameter_tuning_with_hparams.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%2Ftensorboard%2Ftensorboard_hyperparameter_tuning_with_hparams.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/tensorboard/tensorboard_hyperparameter_tuning_with_hparams.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/tensorboard/tensorboard_hyperparameter_tuning_with_hparams.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>                                                                                               \n",
        "</table>"
      ]
    },
@@ -34,23 +34,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/tensorboard/tensorboard_profiler_custom_training_with_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=\"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%2Ftensorboard%2Ftensorboard_profiler_custom_training_with_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://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/tensorboard/tensorboard_profiler_custom_training_with_prebuilt_container.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/tensorboard/tensorboard_profiler_custom_training_with_prebuilt_container.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>                                                                                               \n",
        "</table>"
      ]
    },
@@ -32,25 +32,30 @@
        "# Vertex AI TensorBoard integration with Vertex AI Pipelines\n",
        "\n",
        "<table align=\"left\">\n",
-        "  <td>\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/tensorboard_vertex_ai_pipelines_integration.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%2Ftensorboard%2Ftensorboard_vertex_ai_pipelines_integration.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/tensorboard/tensorboard_vertex_ai_pipelines_integration.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/tensorboard/tensorboard_vertex_ai_pipelines_integration.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",
-        "</table>\n",
-        "<br/><br/><br/>"
+        "  </td>                                                                                               \n",
+        "</table>"
      ]
    },
    {
@@ -29,34 +29,29 @@
        "id": "9Nmi2KIicB7S"
      },
      "source": [
-        "# Distributed Vertex AI Hyperparameter Tuning"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "JAPoU8Sm5E6e"
-      },
-      "source": [
-        "<table align=\"left\">\n",
+        "# Distributed Vertex AI Hyperparameter Tuning\n",
        "\n",
-        "  <td>\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/training/distributed_hyperparameter_tuning.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/training/distributed_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%2Fofficial%2Ftraining%2Fdistributed_hyperparameter_tuning.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/training/distributed_hyperparameter_tuning.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/training/distributed_hyperparameter_tuning.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>"
      ]
    },
@@ -123,12 +118,19 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "install_mlops"
+        "id": "3b1ffd5ab768"
      },
      "source": [
-        "## Installations\n",
-        "\n",
-        "Install the following packages to execute this notebook."
+        "## Get Started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "install_aip"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages"
      ]
    },
    {
@@ -145,136 +147,92 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "e9255e3b156f"
+        "id": "ff555b32bab8"
      },
      "source": [
-        "### Colab Only: Uncomment the following cell to restart the kernel"
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "0c0b2427998a"
+        "id": "f09b4dff629a"
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import sys\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "435b8e413535"
+        "id": "54c5ef8a8f43"
      },
      "source": [
-        "### Before you begin\n",
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f82e28c631cc"
+      },
+      "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."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "be175254a715"
+        "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": "ef3990d0482a"
+      },
+      "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": "oM1iC_MfAts1"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# set the project id\n",
-        "! gcloud config set project $PROJECT_ID"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "2e6b8b324ce1"
-      },
-      "source": [
-        "#### Region\n",
-        "\n",
-        "You can also change the `REGION` variable used by Vertex AI. \n",
-        "Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "region"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "6c43a8673066"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
-        "\n",
-        "**1. Vertex AI Workbench** \n",
-        "- Do nothing as you are already authenticated.\n",
-        "\n",
-        "**2. Local JupyterLab Instance,** uncomment and run."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "fbc9cd30cc4b"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "cd0da2c26879"
-      },
-      "source": [
-        "**3. Colab,** uncomment and run:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "a336a05c6149"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "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\"}"
      ]
    },
    {
@@ -296,7 +254,7 @@
      },
      "outputs": [],
      "source": [
-        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
+        "BUCKET_URI = f\"gs://dist-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -316,7 +274,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} {BUCKET_URI}"
      ]
    },
    {
@@ -325,7 +283,7 @@
        "id": "XoEqT2Y4DJmf"
      },
      "source": [
-        "### Import libraries and define constants"
+        "### Import libraries"
      ]
    },
    {
@@ -336,8 +294,6 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
        "from google.cloud import aiplatform\n",
        "from google.cloud.aiplatform import hyperparameter_tuning as hpt"
      ]
@@ -365,7 +321,7 @@
      "source": [
        "%%writefile Dockerfile\n",
        "\n",
-        "FROM gcr.io/deeplearning-platform-release/tf2-gpu.2-5\n",
+        "FROM gcr.io/deeplearning-platform-release/tf2-gpu.2-6\n",
        "WORKDIR /\n",
        "\n",
        "# Installs hypertune library\n",
@@ -741,7 +697,9 @@
        "\n",
        "hp_job.delete()\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "!gcloud container images delete $IMAGE_URI --force-delete-tags --quiet\n",
+        "\n",
+        "if delete_bucket:\n",
        "    ! gsutil rm -r $BUCKET_URI"
      ]
    }
@@ -33,37 +33,29 @@
        "\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/training/hyperparameter_tuning_tensorflow.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/training/hyperparameter_tuning_tensorflow.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%2Ftraining%2Fhyperparameter_tuning_tensorflow.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/training/hyperparameter_tuning_tensorflow.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/training/hyperparameter_tuning_tensorflow.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": {
@@ -76,7 +68,7 @@
        "\n",
        "And it’s not just about tracking the results from all these trials. You also want a way to efficiently search the space of possible values so you don’t waste as much time trying out combinations that yield low accuracy scores.\n",
        "\n",
-        "Vertex AI Training includes a hyperparameter tuning service. A Vertex AI Hyperparameter tuning job will run multiple trials of your training code. On each trial, it will use different values for your chosen hyperparameters, set within limits you specify. By default, the service uses Bayesian optimization to search the space of possible hyperparameter values. This means that information from prior experiments is used to select the next set of values, making the search more efficient. \n",
+        "Vertex AI Training includes a hyperparameter tuning service. A Vertex AI hyperparameter tuning job runs multiple trials of your training code. On each trial, it uses different values for your chosen hyperparameters, set within the limits you specify. By default, the service uses Bayesian optimization to search the space of possible hyperparameter values. This means that information from prior experiments is used to select the next set of values, making the search more efficient. \n",
        "\n",
        "Learn more about [Vertex AI Hyperparameter Tuning](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview)."
      ]
@@ -89,7 +81,7 @@
      "source": [
        "### Objective\n",
        "\n",
-        "In this tutorial, you learn how to run a Vertex AI Hyperparameter Tuning job for a TensorFlow model. While this example uses TensorFlow, you can also use this service for other ML frameworks.\n",
+        "In this tutorial, you learn how to run a Vertex AI hyperparameter tuning job for a TensorFlow model. While this example uses TensorFlow, you can also use this service for other ML frameworks.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
@@ -137,15 +129,22 @@
        "to generate a cost estimate based on your projected usage."
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "7cc240e9c1a6"
+      },
+      "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"
      ]
    },
    {
@@ -165,7 +164,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."
      ]
    },
    {
@@ -176,11 +177,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": "972161a99004"
+      },
+      "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": "ff666ce4051c"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
+        "\n",
+        "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()"
      ]
    },
    {
@@ -189,66 +232,43 @@
        "id": "WReHDGG5g0XY"
      },
      "source": [
-        "#### Set your project ID\n",
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\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 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).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "oM1iC_MfAts1"
+        "id": "496fcb0092a7"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "\n",
        "# Set the project id\n",
-        "! gcloud config set project {PROJECT_ID}"
+        "! gcloud config set project {PROJECT_ID}\n",
+        "\n",
+        "LOCATION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "region"
+        "id": "cc5ef354b46e"
      },
      "source": [
-        "#### Set the region\n",
+        "### UUID\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)."
+        "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, create a uuid for each instance session, and append it onto the name of resources you created 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": "06571eb4063b"
-      },
-      "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": "697568e92bd6"
+        "id": "1105933b5528"
      },
      "outputs": [],
      "source": [
@@ -264,67 +284,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": {
@@ -364,13 +323,13 @@
      },
      "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": "XoEqT2Y4DJmf"
+        "id": "a2e29a9eb385"
      },
      "source": [
        "### Import libraries"
@@ -380,12 +339,10 @@
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "pRUOFELefqf1"
+        "id": "378ffacc0348"
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
        "import google.cloud.aiplatform as aiplatform\n",
        "from google.cloud.aiplatform import hyperparameter_tuning as hpt"
      ]
@@ -393,23 +350,21 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "init_aip:mbsdk,all"
+        "id": "6246cc5582b5"
      },
      "source": [
-        "### Initialize Vertex AI SDK for Python\n",
-        "\n",
-        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
+        "### Initialize VertexAI SDK for Python"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "init_aip:mbsdk,all"
+        "id": "oM1iC_MfAts1"
      },
      "outputs": [],
      "source": [
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -429,16 +384,16 @@
      "source": [
        "Before you can run a hyperparameter tuning job, you must create a source code file (training script) and a Dockerfile.\n",
        "\n",
-        "The source code trains a model in the ML framework of your choice. In this example, you'll use TensorFlow to train a classification model.\n",
+        "The source code trains a model in the ML framework of your choice. In this example, you use TensorFlow to train a classification model.\n",
        "\n",
-        "The Dockerfile will include all the commands needed to run the image. It'll install all the libraries required by your training script, and set up the entry point for the training code. \n",
+        "The Dockerfile includes all the commands needed to run the image. It installs all the libraries required by your training script, and set up the entry point for the training code. \n",
        "\n",
-        "You'll create a directory for storing your source code file and Dockerfile."
+        "You create a directory for storing your source code file and Dockerfile."
      ]
    },
    {
      "cell_type": "code",
-      "execution_count": 4,
+      "execution_count": null,
      "metadata": {
        "id": "b07b0aeb9309"
      },
@@ -589,7 +544,7 @@
        "\n",
        "`import hypertune`\n",
        "\n",
-        "The **function `get_args()`** defines a command-line argument for each hyperparameter that you want to tune. In this example, the hyperparameters that will be tuned are the learning rate, the momentum value in the optimizer, and the number of units in the last hidden layer of the model. Feel free to experiment with others. The values passed in those arguments are then used to set the corresponding hyperparameter in the code.\n",
+        "The **function `get_args()`** defines a command-line argument for each hyperparameter that you want to tune. In this example, the hyperparameters tuned are the learning rate, the momentum value in the optimizer, and the number of units in the last hidden layer of the model. Feel free to experiment with others. The values passed in those arguments are then used to set the corresponding hyperparameter in the code.\n",
        "\n",
        "```\n",
        "def get_args():\n",
@@ -616,7 +571,7 @@
        "```\n",
        "\n",
        "\n",
-        "At the end of the `main()` function, the `hypertune` library is used to **define the metric you want to optimize**. In TensorFlow, the keras `model.fit` method returns a `History` object. The `History.history` attribute is a record of training loss values and metrics values at successive epochs. If you pass validation data to `model.fit` the `History.history` attribute will include validation loss and metrics values as well.\n",
+        "At the end of the `main()` function, the `hypertune` library is used to **define the metric you want to optimize**. In TensorFlow, the keras `model.fit` method returns a `History` object. The `History.history` attribute is a record of training loss values and metrics values at successive epochs. If you pass validation data to `model.fit` the `History.history` attribute includes validation loss and metrics values as well.\n",
        "\n",
        "```\n",
        "  hp_metric = history.history['val_accuracy'][-1]\n",
@@ -665,7 +620,7 @@
      "source": [
        "### Write Dockerfile\n",
        "\n",
-        "After writing your training code, you create a Dockerfile. In the Dockerfile, you include all the commands needed to run the image. It'll install all the necessary libraries, including the CloudML Hypertune library, and set up the entry point for the training code."
+        "After writing your training code, you create a Dockerfile. In the Dockerfile, include all the commands needed to run the image. It installs all the necessary libraries, including the CloudML Hypertune library, and set up the entry point for the training code."
      ]
    },
    {
@@ -692,6 +647,30 @@
        "ENTRYPOINT [\"python\", \"-m\", \"trainer.task\"]"
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "a2ed53b6e259"
+      },
+      "source": [
+        "### Enable Artifact Registry API\n",
+        "\n",
+        "First, you must 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)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "13970c971b51"
+      },
+      "outputs": [],
+      "source": [
+        "! gcloud services enable artifactregistry.googleapis.com"
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -700,7 +679,7 @@
      "source": [
        "### Build the container\n",
        "\n",
-        "You'll store the Docker image in Artifact Registry. First, create a Docker repository in Artifact Registry"
+        "You store the Docker image in Artifact Registry. First, create a Docker repository in Artifact Registry"
      ]
    },
    {
@@ -711,18 +690,38 @@
      },
      "outputs": [],
      "source": [
-        "REPO_NAME='horses-app'\n",
+        "import os\n",
        "\n",
-        "! gcloud services enable artifactregistry.googleapis.com\n",
+        "REPO_NAME='horses-app'\n",
        "\n",
        "if os.getenv(\"IS_TESTING\"):\n",
        "    ! sudo apt-get update --yes && sudo apt-get --only-upgrade --yes install google-cloud-sdk-cloud-run-proxy google-cloud-sdk-harbourbridge google-cloud-sdk-cbt google-cloud-sdk-gke-gcloud-auth-plugin google-cloud-sdk-kpt google-cloud-sdk-local-extract google-cloud-sdk-minikube google-cloud-sdk-app-engine-java google-cloud-sdk-app-engine-go google-cloud-sdk-app-engine-python google-cloud-sdk-spanner-emulator google-cloud-sdk-bigtable-emulator google-cloud-sdk-nomos google-cloud-sdk-package-go-module google-cloud-sdk-firestore-emulator kubectl google-cloud-sdk-datastore-emulator google-cloud-sdk-app-engine-python-extras google-cloud-sdk-cloud-build-local google-cloud-sdk-kubectl-oidc google-cloud-sdk-anthos-auth google-cloud-sdk-app-engine-grpc google-cloud-sdk-pubsub-emulator google-cloud-sdk-datalab google-cloud-sdk-skaffold google-cloud-sdk google-cloud-sdk-terraform-tools google-cloud-sdk-config-connector\n",
        "    ! gcloud components update --quiet\n",
        "\n",
        "!gcloud artifacts repositories create $REPO_NAME --repository-format=docker \\\n",
-        "--location=$REGION --description=\"Docker repository\"\n",
+        "--location=$LOCATION --description=\"Docker repository\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "94f77f953470"
+      },
+      "source": [
+        "### Configure authentication to your private repo\n",
        "\n",
-        "! gcloud auth configure-docker {REGION}-docker.pkg.dev --quiet"
+        "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."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "3a47ce54c71d"
+      },
+      "outputs": [],
+      "source": [
+        "! gcloud auth configure-docker {LOCATION}-docker.pkg.dev --quiet"
      ]
    },
    {
@@ -743,7 +742,7 @@
      "outputs": [],
      "source": [
        "IMAGE_URI = (\n",
-        "    f\"{REGION}-docker.pkg.dev/{PROJECT_ID}/{REPO_NAME}/horse_human_hptune:latest\"\n",
+        "    f\"{LOCATION}-docker.pkg.dev/{PROJECT_ID}/{REPO_NAME}/horse_human_hptune:latest\"\n",
        ")"
      ]
    },
@@ -753,7 +752,7 @@
        "id": "efd695af40f0"
      },
      "source": [
-        "Then, build the container by running the following:"
+        "Then, build the container and push it to Artifact Registry by running the following:"
      ]
    },
    {
@@ -767,6 +766,17 @@
        "cd $APPLICATION_DIR"
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0ba5d83813d6"
+      },
+      "source": [
+        "Push the image to the created artifact repository using Cloud-Build.\n",
+        "\n",
+        "**Note**: The following command automatically uses the Dockerfile from the directory it's being run from."
+      ]
+    },
    {
      "cell_type": "code",
      "execution_count": null,
@@ -775,27 +785,7 @@
      },
      "outputs": [],
      "source": [
-        "! docker build ./ -t $IMAGE_URI"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ba55cbcf52a5"
-      },
-      "source": [
-        "Finally, push it to Artifact Registry"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "a6062ab48d49"
-      },
-      "outputs": [],
-      "source": [
-        "! docker push $IMAGE_URI"
+        "! gcloud builds submit --region={LOCATION} --tag={IMAGE_URI}"
      ]
    },
    {
@@ -831,7 +821,7 @@
        "        },\n",
        "        \"replica_count\": 1,\n",
        "        \"container_spec\": {\n",
-        "            \"image_uri\": f\"{REGION}-docker.pkg.dev/{PROJECT_ID}/{REPO_NAME}/horse_human_hptune:latest\"\n",
+        "            \"image_uri\": f\"{LOCATION}-docker.pkg.dev/{PROJECT_ID}/{REPO_NAME}/horse_human_hptune:latest\"\n",
        "        },\n",
        "    }\n",
        "]"
@@ -847,7 +837,7 @@
        "\n",
        "Next, define the `parameter_spec`, which is a dictionary specifying the parameters you want to optimize. The **dictionary key** is the string you assigned to the command line argument for each hyperparameter, and the **dictionary value** is the parameter specification.\n",
        "\n",
-        "For each hyperparameter, you need to define the `Type` as well as the bounds for the values that the tuning service will try. Hyperparameters can be of type `Double`, `Integer`, `Categorical`, or `Discrete`. If you select the type `Double` or `Integer`, you need to provide a minimum and maximum value. And if you select `Categorical` or `Discrete` you need to provide the values. For the `Double` and `Integer` types, you also need to provide the scaling value. Learn more about [Using an Appropriate Scale](https://www.youtube.com/watch?v=cSoK_6Rkbfg)."
+        "For each hyperparameter, you need to define the `Type` as well as the bounds for the values that the tuning service tries. Hyperparameters can be of type `Double`, `Integer`, `Categorical`, or `Discrete`. If you select the type `Double` or `Integer`, you need to provide a minimum and maximum value. And if you select `Categorical` or `Discrete` you need to provide the values. For the `Double` and `Integer` types, you also need to provide the scaling value. Learn more about [Using an Appropriate Scale](https://www.youtube.com/watch?v=cSoK_6Rkbfg)."
      ]
    },
    {
@@ -898,7 +888,7 @@
        "id": "d3acfccaa053"
      },
      "source": [
-        "Once the specs are defined, you create a `CustomJob`, which is the common spec that will be used to run your job on each of the hyperparameter tuning trials."
+        "Once the specs are defined, you create a `CustomJob`, which is the common spec that's used to run your job on each of the hyperparameter tuning trials."
      ]
    },
    {
@@ -954,7 +944,7 @@
        "There are a few arguments to note:\n",
        "\n",
        "\n",
-        "* **max_trial_count**: You need to put an upper bound on the number of trials the service will run. More trials generally leads to better results, but there will be a point of diminishing returns, after which additional trials have little or no effect on the metric you're trying to optimize. It is a best practice to start with a smaller number of trials and get a sense of how impactful your chosen hyperparameters are before scaling up.\n",
+        "* **max_trial_count**: You need to put an upper bound on the number of trials the service runs. More trials generally leads to better results, but there is a point of diminishing returns, after which additional trials have little or no effect on the metric you're trying to optimize. The best practice is to start with a smaller number of trials and get a sense of how impactful your chosen hyperparameters are before scaling up.\n",
        "\n",
        "* **parallel_trial_count**: If you use parallel trials, the service provisions multiple training processing clusters. Increasing the number of parallel trials reduces the amount of time the hyperparameter tuning job takes to run; however, it can reduce the effectiveness of the job overall. This is because the default tuning strategy uses results of previous trials to inform the assignment of values in subsequent trials.\n",
        "\n",
@@ -969,7 +959,7 @@
      "source": [
        "## Examine results\n",
        "\n",
-        "Click on the generated link in the output to see your run in the Cloud Console. When the job completes, you will see the results of the tuning trials.\n",
+        "Click on the generated link in the output to see your run in the Cloud Console. When the job completes, you see the results of the tuning trials.\n",
        "\n",
        "You can sort the results by the optimization metric and then set the hyperparamers in your training application code to the values from the trial with the highest accuracy."
      ]
@@ -1006,9 +996,10 @@
      "outputs": [],
      "source": [
        "# Delete artifact registry repo\n",
-        "! gcloud artifacts repositories delete $REPO_NAME --location $REGION --quiet\n",
+        "! gcloud artifacts repositories delete $REPO_NAME --location $LOCATION --quiet\n",
        "\n",
        "delete_custom_job = True\n",
+        "delete_application_directory = True\n",
        "delete_bucket = False\n",
        "\n",
        "# Delete hptune job\n",
@@ -1018,8 +1009,12 @@
        "    except Exception as e:\n",
        "        print(e)\n",
        "\n",
+        "# Delete application directory\n",
+        "if delete_application_directory:\n",
+        "    !cd .. && rm -rf $APPLICATION_DIR\n",
+        "\n",
        "# Delete bucket\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 Hyperparameter Tuning for XGBoost\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/training/hyperparameter_tuning_xgboost.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"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/training/hyperparameter_tuning_xgboost.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Ftraining%2Fhyperparameter_tuning_xgboost.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
+        "  </td>    \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/training/hyperparameter_tuning_xgboost.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/training/hyperparameter_tuning_xgboost.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,9 +63,9 @@
      "source": [
        "## Overview\n",
        "\n",
-        "This tutorial demonstrates how to Vertex AI Hyperparameter Tuning with XGBoost.\n",
+        "This tutorial demonstrates how to Vertex AI hyperparameter tuning with XGBoost.\n",
        "\n",
-        "Learn more about [Vertex AI Hyperparameter Tuning](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview)."
+        "Learn more about [Vertex AI hyperparameter tuning](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview)."
      ]
    },
    {
@@ -74,21 +76,19 @@
      "source": [
        "### Objective\n",
        "\n",
-        "In this tutorial, you learn how to use `Vertex AI Hyperparameter Tuning` for training a XGBoost custom model.\n",
+        "In this tutorial, you learn how to use **Vertex AI hyperparameter tuning** service for training an XGBoost model.\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 Hyperparameter Tuning`\n",
-        "- `Vertex AI Vizier`\n",
-        "- `Vertex AI Model` resource\n",
+        "- **Vertex AI training**\n",
+        "- **Vertex AI hyperparameter tuning** (uses **Vertex AI Vizier**)\n",
        "\n",
        "The steps performed include:\n",
        "\n",
-        "- Training using a Python package.\n",
-        "- Report accuracy when hyperparameter tuning.\n",
+        "- Train using a Python training application package.\n",
+        "- Report accuracy during hyperparameter tuning.\n",
        "- Save the model artifacts to Cloud Storage using GCSFuse.\n",
-        "- Create a `Vertex AI Model` resource."
+        "- List the best model."
      ]
    },
    {
@@ -123,15 +123,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_mlops"
      },
      "source": [
-        "## Installations\n",
-        "\n",
-        "Install the following packages to execute this notebook."
+        "### Install Vertex AI SDK for Python\n"
      ]
    },
    {
@@ -143,7 +150,7 @@
      "outputs": [],
      "source": [
        "# install packages\n",
-        "! pip3 install --upgrade --quiet google-cloud-aiplatform -q"
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform"
      ]
    },
    {
@@ -152,7 +159,9 @@
        "id": "e9255e3b156f"
      },
      "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."
      ]
    },
    {
@@ -163,11 +172,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)"
      ]
    },
    {
@@ -176,109 +188,60 @@
        "id": "435b8e413535"
      },
      "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": "be175254a715"
+        "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": "1ee0f8020b0c"
+      },
+      "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": "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": "2e6b8b324ce1"
-      },
-      "source": [
-        "#### Region\n",
-        "\n",
-        "You can also change the `REGION` variable used by Vertex AI. \n",
-        "Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "region"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "6c43a8673066"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
-        "\n",
-        "**1. Vertex AI Workbench** \n",
-        "- Do nothing as you are already authenticated.\n",
-        "\n",
-        "**2. Local JupyterLab Instance,** uncomment and run."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "fbc9cd30cc4b"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "cd0da2c26879"
-      },
-      "source": [
-        "**3. Colab,** uncomment and run:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "a336a05c6149"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "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\"}"
      ]
    },
    {
@@ -309,7 +272,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."
      ]
    },
    {
@@ -320,52 +283,31 @@
      },
      "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 aiplatform"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "init_aip:mbsdk"
+        "id": "306bf9ab7f7f"
      },
      "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 have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "JZg2sszQox9T"
+        "id": "5878e859ba4d"
      },
      "outputs": [],
      "source": [
-        "aiplatform.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)"
      ]
    },
    {
@@ -378,16 +320,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",
-        "    (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",
        "\n",
        "Learn more about [hardware accelerator support for your region](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators).\n",
        "\n",
-        "*Note*: TF releases before 2.3 for GPU support will fail to load the custom model in this tutorial. It is a known issue and fixed in TF 2.3. This is caused by static graph ops that are generated in the serving function. If you encounter this issue on your own custom models, use a container image for TF 2.3 with GPU support."
+        "**Note**: TF releases before version 2.3 for GPU support are known to fail while loading the custom model in this tutorial. The issue is fixed in TF versions 2.3 and above. This is caused by static graph ops that are generated in the serving function. If you encounter this issue on your own custom models, use a container image for TF 2.3 with GPU support."
      ]
    },
    {
@@ -398,8 +340,6 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
        "TRAIN_GPU, TRAIN_NGPU = (None, None)\n",
        "DEPLOY_GPU, DEPLOY_NGPU = (None, None)"
      ]
@@ -433,10 +373,10 @@
        "DEPLOY_VERSION = \"xgboost-cpu.1-1\"\n",
        "\n",
        "TRAIN_IMAGE = \"{}-docker.pkg.dev/vertex-ai/training/{}:latest\".format(\n",
-        "    REGION.split(\"-\")[0], TRAIN_VERSION\n",
+        "    LOCATION.split(\"-\")[0], TRAIN_VERSION\n",
        ")\n",
        "DEPLOY_IMAGE = \"{}-docker.pkg.dev/vertex-ai/prediction/{}:latest\".format(\n",
-        "    REGION.split(\"-\")[0], DEPLOY_VERSION\n",
+        "    LOCATION.split(\"-\")[0], DEPLOY_VERSION\n",
        ")"
      ]
    },
@@ -448,21 +388,7 @@
      "source": [
        "#### Set machine type\n",
        "\n",
-        "Next, set the machine type to use for training.\n",
-        "\n",
-        "- Set the variable `TRAIN_COMPUTE` to configure  the compute resources for the VMs you will use for for training.\n",
-        " - `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",
-        "\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*."
+        "Set the variable `TRAIN_COMPUTE` to configure the compute resources for VMs that you use for training. Learn more about the [machine types supported for training](https://cloud.google.com/vertex-ai/docs/training/configure-compute#machine-types)."
      ]
    },
    {
@@ -477,38 +403,19 @@
        "print(\"Train machine type\", TRAIN_COMPUTE)"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "xgboost_intro"
-      },
-      "source": [
-        "## Introduction to XGBoost training\n",
-        "\n",
-        "Once you have trained a XGBoost model, you will want to save it at a Cloud Storage location, so it can subsequently be uploaded to a `Vertex AI Model` resource.\n",
-        "The XGBoost package does not have support to save the model to a Cloud Storage location. Instead, you will do the following steps to save to a Cloud Storage location.\n",
-        "\n",
-        "1. Save the in-memory model to the local filesystem (e.g., model.bst).\n",
-        "2. Use gsutil to copy the local copy to the specified Cloud Storage location.\n",
-        "\n",
-        "## Vertex AI Hyperparameter Tuning service\n",
-        "\n",
-        "The following example demonstrates how to setup, execute and evaluate trials using the Vertex AI Hyperparameter Tuning service with `Vizier` search algorithm.\n",
-        "\n",
-        "Learn more about [Overview of hyperparameter tuning](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview)"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "examine_training_package:xgboost"
      },
      "source": [
-        "### Examine the training package\n",
+        "## Python training application package\n",
        "\n",
-        "#### Package layout\n",
+        "In this example, you use Vertex AI hyperparameter tuning service with a training job that executes a Python training application package.\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",
+        "Learn more about [hyperparameter tuning in Vertex AI](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview).\n",
+        "\n",
+        "Take a look at how a Python package can be structured for running a custom training job in Vertex AI. The package contains the following directory structure:\n",
        "\n",
        "- PKG-INFO\n",
        "- README.md\n",
@@ -518,13 +425,18 @@
        "  - \\_\\_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* provide instructions for installing the package into the operating environment of the Docker image.\n",
        "\n",
-        "The file `trainer/task.py` is the Python script for executing the custom training job. *Note*, when we referred to it in the worker pool specification, we replace the directory slash with a dot (`trainer.task`) and dropped the file suffix (`.py`).\n",
-        "\n",
-        "#### Package Assembly\n",
-        "\n",
-        "In the following cells, you will assemble the training package."
+        "The file *trainer/task.py* is the Python script that is executed when you run the custom training job. "
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "509c7082b6e4"
+      },
+      "source": [
+        "### Create a folder structure as Python package\n"
      ]
    },
    {
@@ -562,23 +474,23 @@
        "id": "taskpy_contents:iris,xgboost"
      },
      "source": [
-        "### Create the task script for the Python training package\n",
+        "### Create a training script\n",
        "\n",
-        "Next, you create the `task.py` script for driving the training package. Some noteable steps include:\n",
+        "Next, you create *task.py* script for your training job. Some noteable steps include:\n",
        "\n",
-        "- Command-line arguments:\n",
-        "    - `model-dir`: The location to save the trained model. When using Vertex AI custom training, the location will be specified in the environment variable: `AIP_MODEL_DIR`,\n",
+        "1. <u>Handling command-line arguments:</u>\n",
+        "    - `model-dir`: The location to save the trained model. If no value is passed, the location to save the model is obtained from the environment variable `AIP_MODEL_DIR` that defaults to the staging bucket location.\n",
        "    - `dataset_data_url`: The location of the training data to download.\n",
        "    - `dataset_labels_url`: The location of the training labels to download.\n",
        "    - `boost-rounds`: Tunable hyperparameter\n",
-        "- Data preprocessing (`get_data()`):\n",
+        "1. <u>Data preprocessing (`get_data()`):</u>\n",
        "    - Download the dataset and split into training and test.\n",
-        "- Training (`train_model()`):\n",
+        "1. <u>Training (`train_model()`):</u>\n",
        "    - Trains the model\n",
-        "- Evaluation (`evaluate_model()`):\n",
+        "1. <u>Evaluation (`evaluate_model()`):</u>\n",
        "    - Evaluates the model.\n",
        "    - If hyperparameter tuning, reports the metric for accuracy.\n",
-        "- Model artifact saving\n",
+        "1. <u>Saving model artifacts:</u>\n",
        "    - Saves the model artifacts and evaluation metrics where the Cloud Storage location specified by `model-dir`."
      ]
    },
@@ -619,9 +531,8 @@
        "\n",
        "def get_data():\n",
        "    logging.info(\"Downloading training data and labelsfrom: {}, {}\".format(args.dataset_data_url, args.dataset_labels_url))\n",
-        "    # gsutil outputs everything to stderr so we need to divert it to stdout.\n",
+        "    # gsutil outputs everything to stderr. Hence, the need to divert it to stdout.\n",
        "    subprocess.check_call(['gsutil', 'cp', args.dataset_data_url, 'data.csv'], stderr=sys.stdout)\n",
-        "    # gsutil outputs everything to stderr so we need to divert it to stdout.\n",
        "    subprocess.check_call(['gsutil', 'cp', args.dataset_labels_url, 'labels.csv'], stderr=sys.stdout)\n",
        "\n",
        "\n",
@@ -692,9 +603,9 @@
        "id": "tarball_training_script"
      },
      "source": [
-        "#### Store training script on your Cloud Storage bucket\n",
+        "### Store training script on 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."
+        "Compress the whole training folder as a tar ball and then store it in a Cloud Storage bucket."
      ]
    },
    {
@@ -717,12 +628,14 @@
        "id": "train_custom_job_machine_specification"
      },
      "source": [
-        "### Prepare your machine specification\n",
+        "## Define machine configuration\n",
        "\n",
-        "Now define the machine specification for your custom hyperparameter tuning job. This tells Vertex what type of machine instance to provision for the hyperparameter tuning.\n",
+        "Define the machine configuration for your custom hyperparameter tuning job. This tells Vertex AI what type of machine instance to provision when the job gets started. \n",
+        "\n",
+        "You can specify the following parameters: \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 will use a CPU.\n",
-        "  - `accelerator_count`: The number of accelerators."
+        "  - `accelerator_type`: The type, if any, of hardware accelerator. In this tutorial if you previously set the variable `TRAIN_GPU`, you're using a GPU. Otherwise you're using a CPU.\n",
+        "  - `accelerator_count`: Number of accelerators."
      ]
    },
    {
@@ -749,10 +662,11 @@
        "id": "train_custom_job_disk_specification"
      },
      "source": [
-        "### Prepare your disk specification\n",
+        "## Define disk configuration (optional)\n",
        "\n",
-        "(optional) Now define the disk specification for your custom hyperparameter tuning job. This tells Vertex what type and size of disk to provision in each machine instance for the hyperparameter tuning.\n",
+        "Optionally, define the disk configuration for your custom hyperparameter tuning job. This tells Vertex AI what type and size of disk to provision in each machine instance for the hyperparameter tuning.\n",
        "\n",
+        "You can specify the following parameters:\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."
      ]
@@ -766,7 +680,7 @@
      "outputs": [],
      "source": [
        "DISK_TYPE = \"pd-ssd\"  # [ pd-ssd, pd-standard]\n",
-        "DISK_SIZE = 200  # GB\n",
+        "DISK_SIZE = 100  # GB\n",
        "\n",
        "disk_spec = {\"boot_disk_type\": DISK_TYPE, \"boot_disk_size_gb\": DISK_SIZE}"
      ]
@@ -777,17 +691,15 @@
        "id": "prepare_custom_cmdargs:iris,xgboost"
      },
      "source": [
-        "### Prepare your command-line arguments\n",
+        "## Set worker pool specs\n",
        "\n",
-        "Now define the command-line arguments for your custom training container:\n",
+        "Specify the following worker pool specs 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",
+        "  - `--model-dir`: Specifies where to store the model artifacts in the Cloud Storage bucket.\n",
        "  - `--dataset-data-url`: The location of the training data to download.\n",
        "  - `--dataset-labels-url`: The location of the training labels to download.\n",
-        "  - `--boost-rounds`: Tunable hyperparameter."
+        "  - `--boost-rounds`: Sets the tunable hyperparameter `num_boost_round` while training the XGBoost model."
      ]
    },
    {
@@ -798,15 +710,18 @@
      },
      "outputs": [],
      "source": [
+        "# Set path to save model\n",
        "MODEL_DIR = \"{}/aiplatform-custom-job\".format(BUCKET_URI)\n",
+        "# Set the source path to the dataset\n",
        "DATASET_DIR = \"gs://cloud-samples-data/ai-platform/iris\"\n",
        "\n",
-        "\n",
+        "# Set the command-line arguments\n",
        "CMDARGS = [\n",
        "    \"--dataset-data-url=\" + DATASET_DIR + \"/iris_data.csv\",\n",
        "    \"--dataset-labels-url=\" + DATASET_DIR + \"/iris_target.csv\",\n",
        "]\n",
        "\n",
+        "# Set the worker pool specs\n",
        "worker_pool_spec = [\n",
        "    {\n",
        "        \"replica_count\": 1,\n",
@@ -828,9 +743,9 @@
        "id": "create_custom_job:mbsdk"
      },
      "source": [
-        "## Create a custom job\n",
+        "## Create a custom training job\n",
        "\n",
-        "Use the class `CustomJob` to create a custom job, such as for hyperparameter tuning, with the following parameters:\n",
+        "Use the `CustomJob` class to create a custom training job with the following parameters:\n",
        "\n",
        "- `display_name`: A human readable name for the custom job.\n",
        "- `worker_pool_specs`: The specification for the corresponding VM instances.\n",
@@ -860,13 +775,13 @@
      "source": [
        "## Create a hyperparameter tuning job\n",
        "\n",
-        "Use the class `HyperparameterTuningJob` to create a hyperparameter tuning job, with the following parameters:\n",
+        "Use the `HyperparameterTuningJob` class to create a hyperparameter tuning job with the following parameters:\n",
        "\n",
        "- `display_name`: A human readable name for the custom job.\n",
-        "- `custom_job`: The worker pool spec from this custom job applies to the CustomJobs created in all the trials.\n",
+        "- `custom_job`: The CustomJob object created for training. The worker pool spec from this custom job applies to the CustomJobs created in all the trials.\n",
        "- `metrics_spec`: The metrics to optimize. The dictionary key is the metric_id, which is reported by your training job, and the dictionary value is the optimization goal of the metric('minimize' or 'maximize').\n",
        "- `parameter_spec`: The parameters to optimize. The dictionary key is the metric_id, which is passed into your training job as a command line key word argument, and the dictionary value is the parameter specification of the metric.\n",
-        "- `search_algorithm`: The search algorithm to use: `grid`, `random` and `None`. If `None` is specified, the `Vizier` service (Bayesian) is used.\n",
+        "- `search_algorithm`: The search algorithm to use. Takes `grid`, `random` and `None` as values. Hyperparameter tuning for custom training uses [Vertex AI Vizier](https://cloud.google.com/vertex-ai/docs/vizier/overview) for training jobs. \n",
        "- `max_trial_count`: The maximum number of trials to perform."
      ]
    },
@@ -925,7 +840,7 @@
      "source": [
        "### Display the hyperparameter tuning job trial results\n",
        "\n",
-        "After the hyperparameter tuning job has completed, the property `trials` will return the results for each trial."
+        "Once the hyperparameter tuning job successfully finishes, you can access the results from each trial using the `trials` property."
      ]
    },
    {
@@ -945,9 +860,9 @@
        "id": "best_trial:mbsdk"
      },
      "source": [
-        "### Best trial\n",
+        "## Fetch the best trial\n",
        "\n",
-        "Now look at which trial was the best:"
+        "Identify the best trial and print the details."
      ]
    },
    {
@@ -958,7 +873,9 @@
      },
      "outputs": [],
      "source": [
+        "# Initialize a tuple to identify the best configuration\n",
        "best = (None, None, None, 0.0)\n",
+        "# Iterate through the trails and update the best configuration\n",
        "for trial in hpt_job.trials:\n",
        "    # Keep track of the best outcome\n",
        "    if float(trial.final_measurement.metrics[0].value) > best[3]:\n",
@@ -977,6 +894,7 @@
        "                float(trial.final_measurement.metrics[0].value),\n",
        "            )\n",
        "\n",
+        "# print details of the best configuration\n",
        "print(best)"
      ]
    },
@@ -986,9 +904,9 @@
        "id": "get_best_model"
      },
      "source": [
-        "### Get the Best Model\n",
+        "## List the best model\n",
        "\n",
-        "If you used the method of having the service tell the tuning script where to save the model artifacts (`DIRECT = False`), then the model artifacts for the best model are saved at:\n",
+        "The model artifacts for the best model are saved at: \n",
        "\n",
        "    MODEL_DIR/<best_trial_id>/model"
      ]
@@ -1001,48 +919,28 @@
      },
      "outputs": [],
      "source": [
+        "# Fetch the best model\n",
        "BEST_MODEL_DIR = MODEL_DIR + \"/\" + best[0] + \"/model\"\n",
        "\n",
        "! gsutil ls {BEST_MODEL_DIR}"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "delete_hpt_job"
-      },
-      "source": [
-        "### Delete the hyperparameter tuning job\n",
-        "\n",
-        "The method 'delete()' will delete the hyperparameter tuning job."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "delete_hpt_job"
-      },
-      "outputs": [],
-      "source": [
-        "hpt_job.delete()"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
        "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",
        "\n",
        "Otherwise, you can delete the individual resources you created in this tutorial:\n",
        "\n",
-        "- Custom Job (Custome Training job is remove in previous step)\n",
-        "- Cloud Storage Bucket"
+        "- HyperparameterTuning Job \n",
+        "- Cloud Storage bucket\n",
+        "- Locally generated files"
      ]
    },
    {
@@ -1053,10 +951,18 @@
      },
      "outputs": [],
      "source": [
-        "delete_bucket = False\n",
+        "# Delete the hyperparameter tuning job\n",
+        "hpt_job.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 to delete the bucket\n",
+        "\n",
+        "if delete_bucket:\n",
+        "    ! gsutil rm -r $BUCKET_URI\n",
+        "\n",
+        "# Delete the locally generated files\n",
+        "! rm -rf custom/\n",
+        "! rm custom.tar.gz"
      ]
    }
  ],
@@ -29,27 +29,29 @@
        "id": "c414a395a19b"
      },
      "source": [
-        "# PyTorch image classification multi-node NCCL distributed data parallel training on cpu using Vertex training with custom container\n",
+        "# PyTorch image classification using multi-node NCCL distributed data parallel training on CPU and Vertex AI\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/training/multi_node_ddp_nccl_vertex_training_with_custom_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",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/training/multi_node_ddp_nccl_vertex_training_with_custom_container.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%2Ftraining%2Fmulti_node_ddp_nccl_vertex_training_with_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>\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/training/multi_node_ddp_nccl_vertex_training_with_custom_container.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/training/multi_node_ddp_nccl_vertex_training_with_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",
        "</table>"
      ]
    },
@@ -77,18 +79,16 @@
        "\n",
        "In this tutorial, you learn how to create a distributed PyTorch training job using Vertex AI SDK for Python and custom containers. You set up a custom container, a Vertex AI TensorBoard instance and run a custom training job. \n",
        "\n",
-        "This tutorial uses the following Google Cloud ML services:\n",
+        "This tutorial uses the following Vertex AI services:\n",
        "\n",
-        "- `Vertex AI SDK`\n",
-        "- `Vertex AI TensorBoard`\n",
-        "- `CustomContainerTrainingJob`\n",
-        "- `Artifact Registry`\n",
+        "- Vertex AI TensorBoard\n",
+        "- Vertex AI Training\n",
        "\n",
        "The steps performed include:\n",
        "\n",
-        "- Building Custom Container using Artifact Registry and Docker\n",
-        "- Create a Vertex AI tensorboard instance to store your Vertex AI experiment\n",
-        "- Run a Vertex AI SDK CustomContainerTrainingJob"
+        "- Building a custom container using Artifact Registry and Docker.\n",
+        "- Creating a Vertex AI tensorboard instance to store your Vertex AI experiment.\n",
+        "- Run a Vertex AI training job using Vertex AI SDK for Python."
      ]
    },
    {
@@ -99,7 +99,7 @@
      "source": [
        "### Dataset\n",
        "\n",
-        "The dataset used for this tutorial is the <a href=\"http://yann.lecun.com/exdb/mnist/\">MNIST database</a>. The MNIST database of handwritten digits has a training set of 60,000 examples, and a test set of 10,000 examples. It is a subset of a larger set available from NIST. The digits have been size-normalized and centered in a fixed-size image.\n"
+        "The dataset used for this tutorial is the <a href=\"http://yann.lecun.com/exdb/mnist/\">MNIST database</a>. The MNIST database of handwritten digits has a training set of 60,000 examples, and a test set of 10,000 examples. This dataset is a subset of a larger dataset constructed from  binary images of handwritten digits. The digits have been size-normalized and centered in a fixed-size image.\n"
      ]
    },
    {
@@ -122,15 +122,22 @@
        "        to generate a cost estimate based on your projected usage.\n"
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f0316df526f8"
+      },
+      "source": [
+        "## Get started"
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "03d216c7f7b1"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook. "
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
@@ -141,7 +148,13 @@
      },
      "outputs": [],
      "source": [
-        "! pip3 install --quiet --upgrade google-cloud-aiplatform"
+        "# Install the Vertex AI SDK for Python\n",
+        "! pip3 install --quiet --upgrade google-cloud-aiplatform\n",
+        "\n",
+        "# Install other required libraries\n",
+        "! pip3 install --quiet torch \\\n",
+        "                        torchvision \\\n",
+        "                        tensorboard"
      ]
    },
    {
@@ -150,7 +163,9 @@
        "id": "35942e320683"
      },
      "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."
      ]
    },
    {
@@ -161,32 +176,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": "0e3cab0cc491"
+        "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 up your Google Cloud project\n",
+        "Authenticate your environment on Google Colab.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "46604f70e831"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
        "\n",
-        "**The following steps are required, regardless of your notebook environment.**\n",
+        "if \"google.colab\" in sys.modules:\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",
+        "    from google.colab import auth\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)."
+        "    auth.authenticate_user()"
      ]
    },
    {
@@ -195,12 +231,9 @@
        "id": "80f70e8aa911"
      },
      "source": [
-        "#### Set your project ID\n",
+        "### Set Google Cloud project information\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)."
      ]
    },
    {
@@ -212,123 +245,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": "88dd74c4c84e"
-      },
-      "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": "5c615e53149f"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "070f83c35863"
-      },
-      "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": "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."
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -359,7 +276,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."
      ]
    },
    {
@@ -370,7 +287,31 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "306bf9ab7f7f"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and enable the [Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "4dc6b3ba241c"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI, location=LOCATION)"
      ]
    },
    {
@@ -379,9 +320,9 @@
        "id": "05d881f62170"
      },
      "source": [
-        "#### Service Account \n",
+        "### Service Account \n",
        "\n",
-        "You use a service account to run Vetex AI CustomContainerTrainingJob. If you do not want to use your project's Compute Engine service account, set `SERVICE_ACCOUNT` to another service account ID."
+        "You use a service account to run Vetex AI CustomContainerTrainingJob. If you don't want to use your project's Compute Engine service account, set `SERVICE_ACCOUNT` to another service account ID."
      ]
    },
    {
@@ -405,7 +346,7 @@
      "source": [
        "import sys\n",
        "\n",
-        "IS_COLAB = 'google.colab' in sys.modules\n",
+        "IS_COLAB = \"google.colab\" in sys.modules\n",
        "if (\n",
        "    SERVICE_ACCOUNT == \"\"\n",
        "    or SERVICE_ACCOUNT is None\n",
@@ -427,25 +368,30 @@
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "d48860504181"
+        "id": "2d242773d707"
      },
      "source": [
-        "### Import libraries and define constants"
+        "### Enable Artifact Registry API\n",
+        "You must enable the [Artifact Registry API](https://cloud.google.com/artifact-registry/docs/reference/rest) service for your project.\n",
+        "\n",
+        "<a href=\"https://cloud.google.com/artifact-registry/docs/enable-service\">Learn more about Enabling service</a>."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "40b9227cb6a1"
+        "id": "00ceac726715"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
-        "from google.cloud import aiplatform\n",
+        "! gcloud services enable artifactregistry.googleapis.com\n",
        "\n",
-        "content_name = \"pt-img-cls-multi-node-ddp-cust-cont\""
+        "if os.getenv(\"IS_TESTING\"):\n",
+        "    ! sudo apt-get update --yes && sudo apt-get --only-upgrade --yes install google-cloud-sdk-cloud-run-proxy google-cloud-sdk-harbourbridge google-cloud-sdk-cbt google-cloud-sdk-gke-gcloud-auth-plugin google-cloud-sdk-kpt google-cloud-sdk-local-extract google-cloud-sdk-minikube google-cloud-sdk-app-engine-java google-cloud-sdk-app-engine-go google-cloud-sdk-app-engine-python google-cloud-sdk-spanner-emulator google-cloud-sdk-bigtable-emulator google-cloud-sdk-nomos google-cloud-sdk-package-go-module google-cloud-sdk-firestore-emulator kubectl google-cloud-sdk-datastore-emulator google-cloud-sdk-app-engine-python-extras google-cloud-sdk-cloud-build-local google-cloud-sdk-kubectl-oidc google-cloud-sdk-anthos-auth google-cloud-sdk-app-engine-grpc google-cloud-sdk-pubsub-emulator google-cloud-sdk-datalab google-cloud-sdk-skaffold google-cloud-sdk google-cloud-sdk-terraform-tools google-cloud-sdk-config-connector\n",
+        "    ! gcloud components update --quiet"
      ]
    },
    {
@@ -454,11 +400,9 @@
        "id": "12c03eca5908"
      },
      "source": [
-        "# Create Custom Training Python Package\n",
+        "## Create training application\n",
        "\n",
-        "Before you can perform local training, you must create source code file, requirements file, docker file.\n",
-        "\n",
-        "You will create a directory and write all of our files into that folder."
+        "Before you can perform local training, create a source code file, a requirements file, and a docker file for your training application."
      ]
    },
    {
@@ -469,17 +413,9 @@
      },
      "outputs": [],
      "source": [
-        "PYTHON_PACKAGE_APPLICATION_DIR = \"trainer\""
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "5bc3955a567b"
-      },
-      "outputs": [],
-      "source": [
+        "PYTHON_PACKAGE_APPLICATION_DIR = \"trainer\"\n",
+        "\n",
+        "# Create the app folder\n",
        "!mkdir -p $PYTHON_PACKAGE_APPLICATION_DIR"
      ]
    },
@@ -489,7 +425,9 @@
        "id": "436cc7d0197e"
      },
      "source": [
-        "### Write the Training Script"
+        "### Create training script\n",
+        "\n",
+        "Create the training file that runs the PyTorch distributed training."
      ]
    },
    {
@@ -850,7 +788,7 @@
        "id": "253359c5153a"
      },
      "source": [
-        "### Write requirements file"
+        "### Create requirements file"
      ]
    },
    {
@@ -863,7 +801,6 @@
      "source": [
        "%%writefile {PYTHON_PACKAGE_APPLICATION_DIR}/requirements.txt\n",
        "\n",
-        "\n",
        "torch\n",
        "torchvision\n",
        "tensorboard\n"
@@ -875,7 +812,7 @@
        "id": "34bea792269a"
      },
      "source": [
-        "### Write the docker file"
+        "### Create Dockerfile"
      ]
    },
    {
@@ -913,21 +850,9 @@
        "id": "57bf6f8b4361"
      },
      "source": [
-        "## Local Training\n"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "e5d8a3443da0"
-      },
-      "outputs": [],
-      "source": [
-        "! ls trainer\n",
-        "! cat trainer/requirements.txt\n",
-        "! pip install -r trainer/requirements.txt\n",
-        "! cat trainer/task.py"
+        "## Train locally\n",
+        "\n",
+        "Run your training script locally to test your application."
      ]
    },
    {
@@ -949,72 +874,17 @@
      },
      "outputs": [],
      "source": [
+        "# List the outputs\n",
        "! ls ./tmp"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "3614fcdfcd62"
-      },
-      "source": [
-        "Clean up temporary files"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "48d56ec621cc"
-      },
-      "outputs": [],
-      "source": [
-        "! rm -rf ./tmp"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8f3ea1210749"
      },
      "source": [
-        "## Vertex AI Training using a custom container"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "93002a20a2a6"
-      },
-      "source": [
-        "### Build Custom Container\n"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "2d242773d707"
-      },
-      "source": [
-        "#### Enable Artifact Registry API\n",
-        "You must enable the Artifact Registry API service for your project.\n",
-        "\n",
-        "<a href=\"https://cloud.google.com/artifact-registry/docs/enable-service\">Learn more about Enabling service</a>."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "9d72f89cabd5"
-      },
-      "outputs": [],
-      "source": [
-        "! gcloud services enable artifactregistry.googleapis.com\n",
-        "\n",
-        "if os.getenv(\"IS_TESTING\"):\n",
-        "    ! sudo apt-get update --yes && sudo apt-get --only-upgrade --yes install google-cloud-sdk-cloud-run-proxy google-cloud-sdk-harbourbridge google-cloud-sdk-cbt google-cloud-sdk-gke-gcloud-auth-plugin google-cloud-sdk-kpt google-cloud-sdk-local-extract google-cloud-sdk-minikube google-cloud-sdk-app-engine-java google-cloud-sdk-app-engine-go google-cloud-sdk-app-engine-python google-cloud-sdk-spanner-emulator google-cloud-sdk-bigtable-emulator google-cloud-sdk-nomos google-cloud-sdk-package-go-module google-cloud-sdk-firestore-emulator kubectl google-cloud-sdk-datastore-emulator google-cloud-sdk-app-engine-python-extras google-cloud-sdk-cloud-build-local google-cloud-sdk-kubectl-oidc google-cloud-sdk-anthos-auth google-cloud-sdk-app-engine-grpc google-cloud-sdk-pubsub-emulator google-cloud-sdk-datalab google-cloud-sdk-skaffold google-cloud-sdk google-cloud-sdk-terraform-tools google-cloud-sdk-config-connector\n",
-        "    ! gcloud components update --quiet"
+        "## Create a custom image for training"
      ]
    },
    {
@@ -1024,11 +894,12 @@
      },
      "source": [
        "### Create a private Docker repository\n",
-        "Your first step is to create your own Docker repository in Google Artifact Registry.\n",
        "\n",
-        "1 - Run the gcloud artifacts repositories create command to create a new Docker repository with your region with the description \"docker repository\".\n",
+        "Your first step is to create your own Docker repository in the Artifact Registry.\n",
        "\n",
-        "2 - Run the gcloud artifacts repositories list command to verify that your repository was created."
+        "1. Run the `gcloud artifacts repositories create` command to create a new Docker repository with your region with the description \"docker repository\".\n",
+        "\n",
+        "2. Run the `gcloud artifacts repositories list` command to verify that your repository was created."
      ]
    },
    {
@@ -1039,10 +910,13 @@
      },
      "outputs": [],
      "source": [
-        "PRIVATE_REPO = \"my-docker-repo\"\n",
+        "# Set the name for your private repo\n",
+        "PRIVATE_REPO = \"my-docker-repo-unique\"\n",
        "\n",
-        "! gcloud artifacts repositories create {PRIVATE_REPO} --repository-format=docker --location={REGION} --description=\"Docker repository\"\n",
+        "# Create the repository\n",
+        "! gcloud artifacts repositories create {PRIVATE_REPO} --repository-format=docker --location={LOCATION} --description=\"Docker repository\"\n",
        "\n",
+        "# List the repositories and verify\n",
        "! gcloud artifacts repositories list"
      ]
    },
@@ -1055,7 +929,7 @@
      "outputs": [],
      "source": [
        "DEPLOY_IMAGE = (\n",
-        "    f\"{REGION}-docker.pkg.dev/\" + PROJECT_ID + f\"/{PRIVATE_REPO}\" + \"/tf_serving\"\n",
+        "    f\"{LOCATION}-docker.pkg.dev/\" + PROJECT_ID + f\"/{PRIVATE_REPO}\" + \"/tf_serving\"\n",
        ")"
      ]
    },
@@ -1076,7 +950,11 @@
        "id": "535932b8c42d"
      },
      "source": [
-        "## Executes in Workbench\n"
+        "## Executes in a notebook environment\n",
+        "\n",
+        "If you're running this tutorial in a notebook environment, for example Vertex AI workbench, proceed with the below steps. \n",
+        "\n",
+        "If you're running on Colab, skip to the \"[Executes in Colab](#Executes-in-Colab)\" section."
      ]
    },
    {
@@ -1098,7 +976,7 @@
      "outputs": [],
      "source": [
        "if not IS_COLAB:\n",
-        "    ! gcloud auth configure-docker {REGION}-docker.pkg.dev --quiet"
+        "    ! gcloud auth configure-docker {LOCATION}-docker.pkg.dev --quiet"
      ]
    },
    {
@@ -1107,12 +985,13 @@
        "id": "a822a118c210"
      },
      "source": [
-        "### Container (Docker) image for serving\n",
-        "Set the TensorFlow Serving Docker container image for serving prediction.\n",
+        "### Build and push the image for serving\n",
+        "Set the TensorFlow Serving image for serving predictions.\n",
        "\n",
-        "1. Pull the corresponding CPU or GPU Docker image for TF Serving from Docker Hub.\n",
-        "2. Create a tag for registering the image with Artifact Registry\n",
-        "3. Register the image with Artifact Registry.\n",
+        "The following code:\n",
+        "1. Pulls the corresponding CPU or GPU Docker image for TF Serving from Docker Hub.\n",
+        "2. Creates a tag for registering the image with Artifact Registry\n",
+        "3. Registers the image with Artifact Registry.\n",
        "\n",
        "<a href=\"https://www.tensorflow.org/tfx/serving/docker\">Learn more about TensorFlow Serving</a>."
      ]
@@ -1139,7 +1018,7 @@
      "source": [
        "## Executes in Colab\n",
        "\n",
-        "Build and push a Docker image with Cloud Build"
+        "Build and push a Docker image with Cloud Build."
      ]
    },
    {
@@ -1151,31 +1030,7 @@
      "outputs": [],
      "source": [
        "if IS_COLAB:\n",
-        "    ! cd trainer && gcloud builds submit --timeout=1800s --region={REGION} --tag $DEPLOY_IMAGE"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "10c8cc6b3334"
-      },
-      "source": [
-        "### Initialize Vertex AI SDK"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "42e981cefe41"
-      },
-      "outputs": [],
-      "source": [
-        "aiplatform.init(\n",
-        "    project=PROJECT_ID,\n",
-        "    staging_bucket=BUCKET_URI,\n",
-        "    location=REGION,\n",
-        ")"
+        "    ! cd trainer && gcloud builds submit --timeout=1800s --region={LOCATION} --tag $DEPLOY_IMAGE"
      ]
    },
    {
@@ -1184,9 +1039,16 @@
        "id": "73c92c9298e9"
      },
      "source": [
-        "### Create a Vertex AI Tensorboard instance\n",
+        "## Create a Vertex AI Tensorboard instance\n",
        "\n",
-        "NOTE: <a href=\"https://cloud.google.com/vertex-ai/pricing#tensorboard\">Vertex AI TensorBoard </a> charges a monthly fee of $300 per unique active user. Active users are measured through the Vertex AI TensorBoard UI. You also pay for Google Cloud resources you use with Vertex AI TensorBoard, such as TensorBoard logs stored in Cloud Storage.</a>Please check above link for latest prices."
+        "Set a display name for your Tensorboard instance below and run the create command.\n",
+        "\n",
+        "**Note:**  To use an existing Vertex AI Tensorboard instance, use the following method.\n",
+        "\n",
+        "```\n",
+        "tensorboard_name = \"Your Tensorboard resource name or Tensorboard instance ID\"\n",
+        "tensorboard = aiplatform.Tensorboard(tensorboard_name=tensorboard_name)\n",
+        "```"
      ]
    },
    {
@@ -1197,7 +1059,8 @@
      },
      "outputs": [],
      "source": [
-        "content_name = content_name + \"-cpu\" + \"_\" + UUID"
+        "# Set display name for Tensorboard instance\n",
+        "TENSORBOARD_DISPLAY_NAME = \"pytorch-image-classify-multi-node-ddp\""
      ]
    },
    {
@@ -1208,32 +1071,21 @@
      },
      "outputs": [],
      "source": [
+        "# Create the Tensorboard instance\n",
        "tensorboard = aiplatform.Tensorboard.create(\n",
-        "    display_name=content_name,\n",
+        "    display_name=TENSORBOARD_DISPLAY_NAME,\n",
        ")"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "a1f0a4f54037"
-      },
-      "source": [
-        "#### Option: Use a previously created Vertex AI Tensorboard instance\n",
-        "\n",
-        "```\n",
-        "tensorboard_name = \"Your Tensorboard Resource Name or Tensorboard ID\"\n",
-        "tensorboard = aiplatform.Tensorboard(tensorboard_name=tensorboard_name)\n",
-        "```"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a4cac84e04ac"
      },
      "source": [
-        "### Run a Vertex AI SDK CustomContainerTrainingJob"
+        "## Create a Vertex AI custom container training job\n",
+        "\n",
+        "Learn more about [configuring compute resources for custom training](https://cloud.google.com/vertex-ai/docs/training/configure-compute)."
      ]
    },
    {
@@ -1244,15 +1096,20 @@
      },
      "outputs": [],
      "source": [
-        "display_name = content_name\n",
-        "gcs_output_uri_prefix = f\"{BUCKET_URI}/{display_name}\"\n",
-        "\n",
-        "replica_count = 1\n",
-        "machine_type = \"n1-standard-4\"\n",
-        "accelerator_count = 4\n",
-        "accelerator_type = \"NVIDIA_TESLA_K80\"\n",
-        "\n",
-        "args = [\n",
+        "# Set display name for your training job\n",
+        "JOB_DISPLAY_NAME = \"pytorch-image-classify-multi-node-ddp-training\"\n",
+        "# Set a Cloud Storage output path\n",
+        "GCS_OUTPUT_URI_PREFIX = f\"{BUCKET_URI}/{JOB_DISPLAY_NAME}\"\n",
+        "# Set the number of replicas in the workerpool\n",
+        "REPLICA_COUNT = 1\n",
+        "# Set the machine type for running the training job\n",
+        "MACHINE_TYPE = \"n1-standard-4\"\n",
+        "# Set the number of accelerators needed for training\n",
+        "ACCELERATOR_COUNT = 1\n",
+        "# Set the accelerator type\n",
+        "ACCELERATOR_TYPE = \"NVIDIA_TESLA_T4\"\n",
+        "# Set the args to be passed to the training application\n",
+        "ARGS = [\n",
        "    \"--backend\",\n",
        "    \"nccl\",\n",
        "    \"--no-cuda\",\n",
@@ -1271,12 +1128,22 @@
      },
      "outputs": [],
      "source": [
+        "# Create the custom container training job\n",
        "custom_container_training_job = aiplatform.CustomContainerTrainingJob(\n",
-        "    display_name=display_name,\n",
+        "    display_name=JOB_DISPLAY_NAME,\n",
        "    container_uri=DEPLOY_IMAGE,\n",
        ")"
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "24e0350a3c68"
+      },
+      "source": [
+        "## Run the training job"
+      ]
+    },
    {
      "cell_type": "code",
      "execution_count": null,
@@ -1286,12 +1153,12 @@
      "outputs": [],
      "source": [
        "custom_container_training_job.run(\n",
-        "    args=args,\n",
-        "    base_output_dir=gcs_output_uri_prefix,\n",
-        "    replica_count=replica_count,\n",
-        "    machine_type=machine_type,\n",
-        "    accelerator_count=accelerator_count,\n",
-        "    accelerator_type=accelerator_type,\n",
+        "    args=ARGS,\n",
+        "    base_output_dir=GCS_OUTPUT_URI_PREFIX,\n",
+        "    replica_count=REPLICA_COUNT,\n",
+        "    machine_type=MACHINE_TYPE,\n",
+        "    accelerator_count=ACCELERATOR_COUNT,\n",
+        "    accelerator_type=ACCELERATOR_TYPE,\n",
        "    tensorboard=tensorboard.resource_name,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        ")"
@@ -1306,7 +1173,7 @@
      "outputs": [],
      "source": [
        "print(f\"Custom Training Job Name: {custom_container_training_job.resource_name}\")\n",
-        "print(f\"GCS Output URI Prefix: {gcs_output_uri_prefix}\")"
+        "print(f\"GCS Output URI Prefix: {GCS_OUTPUT_URI_PREFIX}\")"
      ]
    },
    {
@@ -1315,7 +1182,7 @@
        "id": "78398f52807b"
      },
      "source": [
-        "### View training output artifact"
+        "## View training output artifacts"
      ]
    },
    {
@@ -1326,7 +1193,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil ls $gcs_output_uri_prefix"
+        "! gsutil ls $GCS_OUTPUT_URI_PREFIX"
      ]
    },
    {
@@ -1335,15 +1202,17 @@
        "id": "5e99a6a05b10"
      },
      "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",
        "\n",
        "Otherwise, you can delete the individual resources you created in this tutorial:\n",
        "\n",
+        "- Vertex AI Training job\n",
        "- Vertex AI Tensorboard\n",
-        "- Cloud Storage Bucket"
+        "- Cloud Storage bucket\n",
+        "- Artifact Registry repository"
      ]
    },
    {
@@ -1354,17 +1223,29 @@
      },
      "outputs": [],
      "source": [
-        "# Set this to true only if you'd like to delete your bucket\n",
-        "delete_bucket = False\n",
-        "delete_tensorboard = False\n",
+        "# Delete the training job\n",
+        "custom_container_training_job.delete()\n",
        "\n",
-        "! gsutil rm -rf $gcs_output_uri_prefix\n",
+        "# Delete the Tensorboard instance\n",
+        "delete_tensorboard = False  # Set True for deletion\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_tensorboard:\n",
+        "    tensorboard.delete()\n",
+        "\n",
+        "# Delete the Cloud Storage bucket\n",
+        "delete_bucket = False  # Set True for deletion\n",
+        "\n",
+        "if delete_bucket:\n",
        "    ! gsutil rm -r $BUCKET_URI\n",
        "\n",
-        "if delete_tensorboard or os.getenv(\"IS_TESTING\"):\n",
-        "    tensorboard.delete()"
+        "# Delete the repository in Artifact Registry\n",
+        "! gcloud artifacts repositories delete {PRIVATE_REPO} --location {LOCATION} --quiet\n",
+        "\n",
+        "# Remove the outputs from the local execution\n",
+        "! rm -rf ./tmp\n",
+        "\n",
+        "# Remove the training app folder\n",
+        "! rm -rf ./trainer"
      ]
    }
  ],
@@ -4,6 +4,7 @@
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
+        "cellView": "form",
        "id": "q7r5vp1ZZROn"
      },
      "outputs": [],
@@ -32,24 +33,26 @@
        "# Train PyTorch model on Vertex AI with data from Cloud Storage\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/training/pytorch_gcs_data_training.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/training/pytorch_gcs_data_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://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Ftraining%2Fpytorch_gcs_data_training.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/training/pytorch_gcs_data_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",
+        "      <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/training/pytorch_gcs_data_training.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>"
      ]
    },
@@ -76,15 +79,14 @@
        "\n",
        "In this tutorial, you learn how to create a training job using PyTorch and a dataset stored on Cloud Storage. You build a custom training script that uses GCSFuse to load data from a bucket. The custom training script creates a simple neural network and saves the model artifact to a bucket on Cloud Storage.\n",
        "\n",
-        "This tutorial uses the following Google Cloud ML services and resources:\n",
+        "This tutorial uses the following Vertex AI services:\n",
        "\n",
        "- Vertex AI Training\n",
-        "- Cloud Storage\n",
        "\n",
        "The steps performed include:\n",
        "\n",
-        "- Write a custom training script that creates your train & test datasets and trains the model.\n",
-        "- Run a Vertex AI SDK `CustomTrainingJob`"
+        "- Writing a custom training script that creates your train & test datasets and trains the model.\n",
+        "- Runing a `CustomTrainingJob` using Vertex AI SDK for Python."
      ]
    },
    {
@@ -133,15 +135,22 @@
        "Colab notebooks meet the requirements (after installation and authentication). You might need to switch to a GPU-enabled runtime."
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "7cc240e9c1a6"
+      },
+      "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"
      ]
    },
    {
@@ -152,22 +161,9 @@
      },
      "outputs": [],
      "source": [
-        "%%writefile requirements.txt\n",
-        "\n",
-        "google-cloud-aiplatform\n",
-        "google-cloud-storage\n",
-        "torch==1.12"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "WfE-8h14a1za"
-      },
-      "outputs": [],
-      "source": [
-        "%pip install -q -r requirements.txt"
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
+        "                                google-cloud-storage \\\n",
+        "                                torch==1.13"
      ]
    },
    {
@@ -176,7 +172,9 @@
        "id": "67e101d52a4c"
      },
      "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."
      ]
    },
    {
@@ -187,34 +185,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": "BF1j6f9HApxa"
+        "id": "e669f8088ac3"
      },
      "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 up your Google Cloud project\n",
+        "Authenticate your environment on Google Colab.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "cc7251520a07"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
        "\n",
-        "**The following steps are required, regardless of your notebook environment.**\n",
+        "if \"google.colab\" in sys.modules:\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",
+        "    from google.colab import auth\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. [Enable the Storage API](https://console.cloud.google.com/flows/enableapi?apiid=storage.googleapis.com).\n",
-        "\n",
-        "5. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
+        "    auth.authenticate_user()"
      ]
    },
    {
@@ -223,12 +240,9 @@
        "id": "WReHDGG5g0XY"
      },
      "source": [
-        "#### Set your project ID\n",
+        "### Set Google Cloud project information \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)."
      ]
    },
    {
@@ -239,104 +253,8 @@
      },
      "outputs": [],
      "source": [
-        "PROJECT_ID = \"[your-project-id]\"  # @param\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": "0gtuztzPg0x8"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param"
-      ]
-    },
-    {
-      "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."
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -368,7 +286,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."
      ]
    },
    {
@@ -379,7 +297,31 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "39dbc83d5530"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). "
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "86fa44b5ab17"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -411,28 +353,6 @@
        "from torch.utils.data import DataLoader, Dataset"
      ]
    },
-    {
-      "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": [
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -443,7 +363,7 @@
        "\n",
        "As mentioned before, this tutorial uses the classic MNIST handwritten digits dataset as an input. The dataset has been stored for your use on a publicly available Cloud Storage location. You can use these CSV files directly in your training script.\n",
        "\n",
-        "**Note**: You can use the [PyTorch datasets library](https://pytorch.org/vision/stable/generated/torchvision.datasets.MNIST.html#torchvision.datasets.MNIST) to download a this dataset. For learning purposes, this tutorial uses a copy of the dataset on Cloud Storage."
+        "**Note**: You can use the [PyTorch datasets library](https://pytorch.org/vision/stable/generated/torchvision.datasets.MNIST.html#torchvision.datasets.MNIST) to download this dataset. For learning purposes, this tutorial uses a copy of the dataset on Cloud Storage."
      ]
    },
    {
@@ -467,7 +387,7 @@
        "id": "_7zjKtBsLOrL"
      },
      "source": [
-        "## [OPTIONAL] Inspect the dataset from GCS\n",
+        "## [OPTIONAL] Inspect the dataset from Cloud Storage\n",
        "\n",
        "Before creating the training script, take a quick look at the data contained in the CSV files on Cloud Storage. You can use the PyTorch [`Dataset`](https://pytorch.org/docs/stable/data.html?highlight=torch+utils+data+dataset#torch.utils.data.Dataset) and [`DataLoader`](https://pytorch.org/docs/stable/data.html?highlight=torch+utils+data+dataset#torch.utils.data.DataLoader) classes to instantiate a dataset and then plot the data using [matplotlib](https://matplotlib.org/stable/index.html).\n",
        "\n",
@@ -587,7 +507,7 @@
        "    first_label = y[0]\n",
        "\n",
        "    print(len(first_image))\n",
-        "    print(first_label)  # This will be a Tensor object with a single scalar value, 5\n",
+        "    print(first_label)  # This is a Tensor object with a single scalar value, 5\n",
        "    break"
      ]
    },
@@ -653,9 +573,7 @@
        "device = (\n",
        "    \"cuda\"\n",
        "    if torch.cuda.is_available()\n",
-        "    else \"mps\"\n",
-        "    if torch.backends.mps.is_available()\n",
-        "    else \"cpu\"\n",
+        "    else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n",
        ")\n",
        "print(f\"Using {device} device\")\n",
        "\n",
@@ -765,14 +683,14 @@
      },
      "outputs": [],
      "source": [
-        "TRAIN_GPU, TEST_GPU = (aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_K80, 1)\n",
+        "TRAIN_GPU, TRAIN_NGPU = (aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_T4, 1)\n",
        "if TRAIN_GPU:\n",
-        "    TRAIN_VERSION = \"pytorch-gpu.1-9\"\n",
+        "    TRAIN_VERSION = \"pytorch-gpu.1-13\"\n",
        "else:\n",
-        "    TRAIN_VERSION = \"pytorch-xla.1-9\"\n",
+        "    TRAIN_VERSION = \"pytorch-xla.1-11\"\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",
        "MACHINE_TYPE = \"n1-standard\"\n",
@@ -792,7 +710,7 @@
        "\n",
        "Next, you write the training script, 'task.py', to file before creating the training job. Note that the script includes the dataset, dataloader, and neural net module that you inspected previously.\n",
        "\n",
-        "In the training script, the training scripts are loaded from Cloud Storage using [Storage FUSE](https://cloud.google.com/storage/docs/gcs-fuse). FUSE mounts Cloud Storage buckets as folders in the training container's file system. This allows the training script to load files storaged in the bucket as a dataset. FUSE also allows the training script to store the output of training--the model artifact--in a Cloud Storage bucket.\n",
+        "In the training script, the training scripts are loaded from Cloud Storage using [Storage FUSE](https://cloud.google.com/storage/docs/gcs-fuse). FUSE mounts Cloud Storage buckets as folders in the training container's file system. This allows the training script to load files stored in the bucket as a dataset. FUSE also allows the training script to store the output of training i.e., the model artifact, in a Cloud Storage bucket.\n",
        "\n",
        "To use a bucket mounted to the container using FUSE, you replace the `gs://` portion of the bucket's URI with the folder path `/gcs/`."
      ]
@@ -814,26 +732,20 @@
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "uIvoEYeB1sJe"
+        "id": "361c5b0c9949"
      },
      "outputs": [],
      "source": [
        "%%writefile trainer/task.py\n",
-        "import sys\n",
        "import os\n",
        "import argparse\n",
        "import logging\n",
-        "import hypertune\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
-        "import torch.nn.functional as F\n",
-        "import numpy as np\n",
        "import pandas as pd\n",
-        "from torchvision.io import read_image\n",
        "from torch.utils.data import Dataset\n",
-        "from torchvision.transforms import ToTensor\n",
        "from torch.utils.data import DataLoader\n",
        "\n",
        "parser = argparse.ArgumentParser(description='PyTorch CNN Training')\n",
@@ -1026,23 +938,28 @@
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "GDaFQpMEMeVb"
+        "id": "1edd9128828e"
      },
      "outputs": [],
      "source": [
+        "# Use timestamped path to save your model in Cloud Storage\n",
        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
+        "# Set a display name for the training job\n",
        "JOB_DISPLAY_NAME = \"pytorch-custom-job\"\n",
-        "TRAIN_IMAGE_URI = \"us-docker.pkg.dev/vertex-ai/training/pytorch-gpu.1-13:latest\"\n",
        "\n",
+        "# Create a custom training job in Vertex AI\n",
        "job = aiplatform.CustomTrainingJob(\n",
        "    display_name=JOB_DISPLAY_NAME,\n",
        "    script_path=\"trainer/task.py\",\n",
-        "    container_uri=TRAIN_IMAGE_URI,\n",
+        "    container_uri=TRAIN_IMAGE,\n",
        ")\n",
        "\n",
+        "# Run the job\n",
        "job.run(\n",
        "    replica_count=1,\n",
-        "    machine_type=\"n1-standard-4\",\n",
+        "    machine_type=TRAIN_COMPUTE,\n",
+        "    accelerator_type=TRAIN_GPU.name,\n",
+        "    accelerator_count=TRAIN_NGPU,\n",
        "    args=[\n",
        "        \"--train_uri\",\n",
        "        TRAIN_URI,\n",
@@ -1076,14 +993,12 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
        "# Delete training job created\n",
        "job.delete(sync=False)\n",
        "\n",
        "# Delete Cloud Storage objects that were created\n",
        "delete_bucket = False\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
@@ -1,5 +1,5 @@

-[Using Vertex AI Multimodal Embeddings and Vector Search](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/matching_engine/sdk_matching_engine_create_multimodal_embeddings.ipynb)
+[Using Vertex AI Multimodal Embeddings and Vector Search](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vector_search/sdk_vector_search_create_multimodal_embeddings.ipynb)

 ```
 Learn how to encode custom text embeddings, create an  Approximate Nearest Neighbor  index, and query against indexes.
@@ -16,7 +16,7 @@ The steps performed include:
 ```


-[Using Vertex AI Matching Engine for StackOverflow Questions](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/matching_engine/sdk_matching_engine_create_stack_overflow_embeddings.ipynb)
+[Using Vertex AI Vector Search for StackOverflow Questions](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vector_search/sdk_vector_search_create_stack_overflow_embeddings.ipynb)

 ```
 Learn how to encode custom text embeddings, create an Approximate Nearest Neighbor  index, and query against indexes.
@@ -30,10 +30,10 @@ The steps performed include:

 ```

-&nbsp;&nbsp;&nbsp;Learn more about [Vertex AI Matching Engine](https://cloud.google.com/vertex-ai/docs/matching-engine/overview).
+&nbsp;&nbsp;&nbsp;Learn more about [Vertex AI Vector Search](https://cloud.google.com/vertex-ai/docs/vector-search/overview).


-[Using Vertex AI Vector Search and Vertex AI Embeddings for Text for StackOverflow Questions](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/matching_engine/sdk_matching_engine_create_stack_overflow_embeddings_vertex.ipynb)
+[Using Vertex AI Vector Search and Vertex AI Embeddings for Text for StackOverflow Questions](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vector_search/sdk_vector_search_create_stack_overflow_embeddings_vertex.ipynb)

 ```
 Learn how to encode text embeddings, create an Approximate Nearest Neighbor  index, and query against indexes.
@@ -49,12 +49,12 @@ The steps performed include:

 ```

-&nbsp;&nbsp;&nbsp;Learn more about [Vertex AI Vector Search](https://cloud.google.com/vertex-ai/docs/matching-engine/overview).
+&nbsp;&nbsp;&nbsp;Learn more about [Vertex AI Vector Search](https://cloud.google.com/vertex-ai/docs/vector-search/overview).

 &nbsp;&nbsp;&nbsp;Learn more about [Vertex AI Embeddings for Text](https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-text-embeddings).


-[Using Vertex AI Matching Engine for Text-to-Image Embeddings](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/matching_engine/sdk_matching_engine_create_text_to_image_embeddings.ipynb)
+[Using Vertex AI Vector Search for Text-to-Image Embeddings](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vector_search/sdk_vector_search_create_text_to_image_embeddings.ipynb)

 ```
 Learn how to encode custom text embeddings, create an  Approximate Nearest Neighbor  index, and query against indexes.
@@ -68,10 +68,10 @@ The steps performed include:

 ```

-&nbsp;&nbsp;&nbsp;Learn more about [Vertex AI Matching Engine](https://cloud.google.com/vertex-ai/docs/matching-engine/overview).
+&nbsp;&nbsp;&nbsp;Learn more about [Vertex AI Vector Search](https://cloud.google.com/vertex-ai/docs/vector-search/overview).


-[Using Vertex AI Vector Search for Text-to-Image Embeddings](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/matching_engine/sdk_matching_engine_create_text_to_image_embeddings.ipynb)
+[Using Vertex AI Vector Search for Text-to-Image Embeddings](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vector_search/sdk_vector_search_create_text_to_image_embeddings.ipynb)

 ```
 Learn how to encode custom text embeddings, create an  Approximate Nearest Neighbor  index, and query against indexes.
@@ -85,10 +85,10 @@ The steps performed include:

 ```

-&nbsp;&nbsp;&nbsp;Learn more about [Vertex AI Vector Search](https://cloud.google.com/vertex-ai/docs/matching-engine/overview).
+&nbsp;&nbsp;&nbsp;Learn more about [Vertex AI Vector Search](https://cloud.google.com/vertex-ai/docs/vector-search/overview).


-[Create Vertex AI Matching Engine index](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/matching_engine/sdk_matching_engine_for_indexing.ipynb)
+[Create Vertex AI Vector Search index](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vector_search/sdk_vector_search_for_indexing.ipynb)

 ```
 Learn how to create Approximate Nearest Neighbor  Index, query against indexes, and validate the performance of the index.
@@ -103,5 +103,5 @@ The steps performed include:

 ```

-&nbsp;&nbsp;&nbsp;Learn more about [Vertex AI Matching Engine](https://cloud.google.com/vertex-ai/docs/matching-engine/overview).
+&nbsp;&nbsp;&nbsp;Learn more about [Vertex AI Vector Search](https://cloud.google.com/vertex-ai/docs/vector-search/overview).

@@ -30,25 +30,30 @@
      },
      "source": [
        "# Using Vertex AI Multimodal Embeddings and Vector Search\n",
-        "![ ](https://www.google-analytics.com/collect?v=2&tid=G-L6X3ECH596&cid=1&en=page_view&sid=1&dt=sdk_matching_engine_create_multimodal_embeddings.ipynb&dl=notebooks%2Fofficial%2Fmatching_engine%2Fsdk_matching_engine_create_multimodal_embeddings.ipynb)\n",
+        "![ ](https://www.google-analytics.com/collect?v=2&tid=G-L6X3ECH596&cid=1&en=page_view&sid=1&dt=sdk_vector_search_create_multimodal_embeddings.ipynb&dl=notebooks%2Fofficial%2Fvector_search%2Fsdk_vector_search_create_multimodal_embeddings.ipynb)\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/matching_engine/sdk_matching_engine_create_multimodal_embeddings.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/vector_search/sdk_vector_search_create_multimodal_embeddings.ipynb\">\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/matching_engine/sdk_matching_engine_create_multimodal_embeddings.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%2fvector_search%2fsdk_vector_search_create_multimodal_embeddings.ipynb.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/matching_engine/sdk_matching_engine_create_multimodal_embeddings.ipynb\">\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/vector_search/sdk_vector_search_create_multimodal_embeddings.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",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vector_search/sdk_vector_search_create_multimodal_embeddings.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>"
      ]
    },
@@ -109,9 +114,16 @@
        "id": "f0f1bea346db"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the latest version of Cloud Storage and the Vertex AI SDK for Python."
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0a51c572aa72"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages"
      ]
    },
    {
@@ -123,8 +135,9 @@
      "outputs": [],
      "source": [
        "# Install the packages\n",
-        "! pip3 install --upgrade google-cloud-aiplatform==1.35.0 \\\n",
-        "                         google-cloud-storage"
+        "! pip3 install --upgrade google-cloud-aiplatform \\\n",
+        "                         google-cloud-storage \\\n",
+        "                         google-cloud-vision"
      ]
    },
    {
@@ -133,7 +146,9 @@
        "id": "97473593f37f"
      },
      "source": [
-        "Install the latest version of google-cloud-vision for filtering for safe images"
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
@@ -144,47 +159,64 @@
      },
      "outputs": [],
      "source": [
-        "# Install the packages\n",
-        "! pip install google-cloud-vision"
+        "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": "5b08ba354c6e"
+        "id": "ffcde4d56c00"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel."
+        "<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": "7176ea64999b"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
+        "\n",
+        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "bea801acf6b5"
+        "id": "7de6ef0fac42"
      },
      "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",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "dd28c9e4f067"
+        "id": "ef3990d0482a"
      },
      "source": [
-        "## Before you begin\n",
-        "#### Set your project ID\n",
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\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 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)."
      ]
    },
    {
@@ -196,112 +228,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": "4f4512bf63b3"
-      },
-      "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": "474be5183c27"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "949271bfebe3"
-      },
-      "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": "b65b4ce80d9a"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing since you're already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "985cdbfe7372"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "fbc9cd30cc4b"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "79efab26ad02"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "a336a05c6149"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0c0a44fa330f"
-      },
-      "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": "d3uj8x73nDX_"
-      },
-      "source": [
-        "* Authentication: Rerun the `gcloud auth login` command in the Vertex AI Workbench notebook terminal when you are logged out and need the credential again."
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
@@ -343,7 +275,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
@@ -595,7 +527,6 @@
        "def process_image(image_path: str) -> Optional[bool]:\n",
        "    try:\n",
        "        annotation = detect_safe_search(image_path)\n",
-        "\n",
        "        if annotation:\n",
        "            return convert_annotation_to_safety(safe_search_annotation=annotation)\n",
        "        else:\n",
@@ -614,7 +545,6 @@
        "\n",
        "    for future in futures:\n",
        "        is_safe_values_cloud_vision.append(future.result())\n",
-        "\n",
        "# Set Nones to False\n",
        "is_safe_values_cloud_vision = [\n",
        "    is_safe or False for is_safe in is_safe_values_cloud_vision\n",
@@ -706,7 +636,7 @@
        "        api_regional_endpoint: str = \"us-central1-aiplatform.googleapis.com\",\n",
        "    ):\n",
        "        client_options = {\"api_endpoint\": api_regional_endpoint}\n",
-        "        # Initialize client that will be used to create and send requests.\n",
+        "        # Initialize client which is used to create and send requests.\n",
        "        # This client only needs to be created once, and can be reused for multiple requests.\n",
        "        self.client = aiplatform.gapic.PredictionServiceClient(\n",
        "            client_options=client_options\n",
@@ -1141,7 +1071,7 @@
        "id": "mglUPwHpJH98"
      },
      "source": [
-        "## Create MatchingEngineIndex\n"
+        "## Create Vector Search Index\n"
      ]
    },
    {
@@ -1164,7 +1094,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)"
      ]
    },
    {
@@ -1173,7 +1103,7 @@
        "id": "fa5eb5cdb2f0"
      },
      "source": [
-        "#### Create the index configuration\n",
+        "#### Create the configuration\n",
        "\n",
        "For information on configuration settings, see the [Manage Indexes documentation](https://cloud.google.com/vertex-ai/docs/vector-search/create-manage-index)"
      ]
@@ -1217,7 +1147,7 @@
        "id": "0f1a9fbecabb"
      },
      "source": [
-        "Using the resource name, you can retrieve an existing MatchingEngineIndex."
+        "Using the resource name, you can retrieve an existing Index resource."
      ]
    },
    {
@@ -1263,7 +1193,7 @@
        "id": "qV2xjAnDDObD"
      },
      "source": [
-        "## Create an MatchingEngineIndexEndpoint"
+        "## Create a Vector Search IndexEndpoint"
      ]
    },
    {
@@ -1438,7 +1368,7 @@
        "\n",
        "# Delete Cloud Storage objects that were created\n",
        "delete_bucket = False\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
@@ -1446,7 +1376,7 @@
  "metadata": {
    "colab": {
      "collapsed_sections": [],
-      "name": "sdk_matching_engine_create_multimodal_embeddings.ipynb",
+      "name": "sdk_vector_search_create_multimodal_embeddings.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
@@ -30,23 +30,26 @@
      },
      "source": [
        "# Using Vertex AI Vector Search for StackOverflow Questions\n",
-        "![ ](https://www.google-analytics.com/collect?v=2&tid=G-L6X3ECH596&cid=1&en=page_view&sid=1&dt=sdk_matching_engine_create_stack_overflow_embeddings.ipynb&dl=notebooks%2Fofficial%2Fmatching_engine%2Fsdk_matching_engine_create_stack_overflow_embeddings.ipynb)\n",
+        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/matching_engine/sdk_matching_engine_create_stack_overflow_embeddings.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/vector_search/sdk_vector_search_create_stack_overflow_embeddings.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/matching_engine/sdk_matching_engine_create_stack_overflow_embeddings.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%2Fvector_search%2Fsdk_vector_search_create_stack_overflow_embeddings.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/vector_search/sdk_vector_search_create_stack_overflow_embeddings.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/matching_engine/sdk_matching_engine_create_stack_overflow_embeddings.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/vector_search/sdk_vector_search_create_stack_overflow_embeddings.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 +65,7 @@
        "\n",
        "This example demonstrates how to encode custom text embeddings using the StackOverflow dataset and the sentence-T5 model. These are uploaded to the Vertex AI Vector Search service. It's a high scale, low latency solution, to find similar vectors (or more specifically \"embeddings\") for a large corpus. Moreover, it's a fully managed offering, further reducing operational overhead. The Vertex AI Vector Search service is built upon [Approximate Nearest Neighbor (ANN) technology](https://ai.googleblog.com/2020/07/announcing-scann-efficient-vector.html) developed by Google Research.\n",
        "\n",
-        "**Pre-requisite**: This notebook requires you to already have a VPC network set up. See the \"Prepare a VPC network\" section in [Create Vertex AI Vector Search index notebook](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/matching_engine/sdk_matching_engine_for_indexing.ipynb).\n",
+        "**Pre-requisite**: This notebook requires you to already have a VPC network set up. See the \"Prepare a VPC network\" section in [Create Vertex AI Vector Search index notebook](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vector_search/sdk_vector_search_for_indexing.ipynb).\n",
        "\n",
        "Learn more about [Vertex AI Vector Search](https://cloud.google.com/vertex-ai/docs/vector-search/overview)."
      ]
@@ -102,15 +105,22 @@
        "> Stack Overflow is the largest online community for programmers to learn, share their knowledge, and advance their careers. Updated on a quarterly basis, this BigQuery dataset includes an archive of Stack Overflow content, including posts, votes, tags, and badges. This dataset is updated to mirror the Stack Overflow content on the Internet Archive, and is also available through the Stack Exchange Data Explorer."
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "d0be1c1c229a"
+      },
+      "source": [
+        "## Get started"
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f0f1bea346db"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the latest version of Cloud Storage, BigQuery, and the Vertex AI SDK for Python."
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
@@ -121,54 +131,19 @@
      },
      "outputs": [],
      "source": [
-        "# Install the packages\n",
-        "! pip3 install --upgrade google-cloud-aiplatform==1.35.0 \\\n",
+        "# Install the google-cloud packages\n",
+        "! pip3 install --upgrade google-cloud-aiplatform \\\n",
        "                         google-cloud-storage \\\n",
-        "                         'google-cloud-bigquery[pandas]'"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "1ae34c2a9ce7"
-      },
-      "source": [
-        "Install the latest version of tensorflow and tensorflow_text to encode embeddings"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "f329eff0fe39"
-      },
-      "outputs": [],
-      "source": [
-        "# Install the packages\n",
+        "                         'google-cloud-bigquery[pandas]' -q\n",
+        "\n",
+        "# Install the latest version of tensorflow packages\n",
        "! pip3 install --upgrade tensorflow \\\n",
        "                         tensorflow_text \\\n",
-        "                         tensorflow-hub"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "54ac7ebac10b"
-      },
-      "source": [
-        "Install the latest version of Redis for low-latency data retrieval"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "e3dd53b3c06c"
-      },
-      "outputs": [],
-      "source": [
-        "# Install the redis package\n",
-        "! pip install --upgrade redis"
+        "                         tensorflow-hub -q\n",
+        "\n",
+        "# Install the redis and tqdm packages\n",
+        "! pip install --upgrade redis \\\n",
+        "                        tqdm -q"
      ]
    },
    {
@@ -177,7 +152,9 @@
        "id": "5b08ba354c6e"
      },
      "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."
      ]
    },
    {
@@ -188,11 +165,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": "ffcde4d56c00"
+      },
+      "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": "7176ea64999b"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
+        "\n",
+        "Authenticate your environment on Google Colab.\n"
+      ]
+    },
+    {
+      "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()"
      ]
    },
    {
@@ -201,13 +220,9 @@
        "id": "dd28c9e4f067"
      },
      "source": [
-        "## Before you begin\n",
-        "#### Set your project ID\n",
+        "### Set Google Cloud project information\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 how to [set up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
@@ -218,137 +233,8 @@
      },
      "outputs": [],
      "source": [
-        "PROJECT_ID = \"[your-project-id]\"\n",
-        "\n",
-        "# Set the project id\n",
-        "! gcloud config set project {PROJECT_ID}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "4f4512bf63b3"
-      },
-      "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": "474be5183c27"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "949271bfebe3"
-      },
-      "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": "b65b4ce80d9a"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing since you're already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "985cdbfe7372"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "fbc9cd30cc4b"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "79efab26ad02"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "a336a05c6149"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0c0a44fa330f"
-      },
-      "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": "d3uj8x73nDX_"
-      },
-      "source": [
-        "* Authentication: Rerun the `gcloud auth login` command in the Vertex AI Workbench notebook terminal when you are logged out and need the credential again."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "hhq5zEbGg0XX"
-      },
-      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "EzrelQZ22IZj"
-      },
-      "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)"
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
@@ -390,7 +276,63 @@
      },
      "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": "e45533d84985"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). "
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "b519c49740cb"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "44b8dbd3f108"
+      },
+      "source": [
+        "### Import the required libraries"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bf6f53cd87ee"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "import os\n",
+        "import tempfile\n",
+        "from typing import List\n",
+        "\n",
+        "import numpy as np\n",
+        "import redis\n",
+        "import tensorflow as tf\n",
+        "import tensorflow_hub as hub\n",
+        "# Registers the ops.\n",
+        "import tensorflow_text as text  # noqa: F401\n",
+        "from google.cloud import bigquery\n",
+        "from tqdm.auto import tqdm"
      ]
    },
    {
@@ -401,9 +343,18 @@
      "source": [
        "## Prepare the data\n",
        "\n",
-        "You'll use the [Stack Overflow dataset](https://console.cloud.google.com/marketplace/product/stack-exchange/stack-overflow) of question and answers hosted on BigQuery.\n",
+        "For this tutorial, use the [Stack Overflow dataset](https://console.cloud.google.com/marketplace/product/stack-exchange/stack-overflow) of question and answers hosted on BigQuery.\n",
        "\n",
-        "> This public dataset is hosted in Google BigQuery and is included in BigQuery's 1TB/mo of free tier processing. This means that each user receives 1TB of free BigQuery processing every month, which can be used to run queries on this public dataset."
+        "> This public dataset is hosted in Google BigQuery and is included in BigQuery's 1TB/mo of free tier processing. Each user receives 1TB of free BigQuery processing every month, which can be used to run queries on this public dataset."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "62f6feea28a2"
+      },
+      "source": [
+        "Fetch the dataset from the BigQuery source."
      ]
    },
    {
@@ -415,8 +366,6 @@
      "outputs": [],
      "source": [
        "%%time\n",
-        "from google.cloud import bigquery\n",
-        "\n",
        "client = bigquery.Client(project=PROJECT_ID)\n",
        "\n",
        "NUM_ROWS = 1000\n",
@@ -470,13 +419,13 @@
        "id": "1124422cc200"
      },
      "source": [
-        "#### Instantiate the text encoding model\n",
+        "### Instantiate the text encoding model\n",
        "\n",
        "Use the [sentence-t5 encoder](https://tfhub.dev/google/sentence-t5/st5-base/1) developed by Google for converting text to embeddings.\n",
        "\n",
        "> The sentence-T5 family of models encode text into high-dimensional vectors that can be used for text classification, semantic similarity, clustering and other natural language processing tasks.\n",
        ">\n",
-        "> Our model is built on top of T5 (i.e. the Text-To-Text Transfer Transformer). It is trained on a variety of data sources and initialized from pre-trained T5 models with different model sizes. The input is variable-length English text and the output is a 768-dimensional vector. The sentence-T5 base model employs a 12-layer transformer architecture as the T5 base model does."
+        "> The model is built on top of T5 (i.e., the Text-To-Text Transfer Transformer). It's trained on a variety of data sources and initialized from pre-trained T5 models with different model sizes. The input is variable-length English text and the output is a 768-dimensional vector. The sentence-T5 base model employs a 12-layer transformer architecture as does the T5 base model."
      ]
    },
    {
@@ -487,11 +436,6 @@
      },
      "outputs": [],
      "source": [
-        "import tensorflow as tf\n",
-        "import tensorflow_hub as hub\n",
-        "# Registers the ops.\n",
-        "import tensorflow_text as text  # noqa: F401\n",
-        "\n",
        "hub_url = \"https://tfhub.dev/google/sentence-t5/st5-base/1\"\n",
        "\n",
        "encoder = hub.KerasLayer(hub_url)"
@@ -503,7 +447,7 @@
        "id": "43088937e820"
      },
      "source": [
-        "#### Defining an encoding function\n",
+        "### Define an encoding function\n",
        "\n",
        "Define a function, to be used later, that takes sentences and converts them to embeddings."
      ]
@@ -516,12 +460,6 @@
      },
      "outputs": [],
      "source": [
-        "from typing import List\n",
-        "\n",
-        "import numpy as np\n",
-        "from tqdm.auto import tqdm\n",
-        "\n",
-        "\n",
        "def encode_text_to_embedding(\n",
        "    text_encoder: hub.KerasLayer, sentences: List[str], batch_size: int = 100\n",
        ") -> np.ndarray:\n",
@@ -595,6 +533,7 @@
      "source": [
        "question_index = 0\n",
        "\n",
+        "# Print the query question\n",
        "print(f\"Query question = {questions[question_index]}\")\n",
        "scores = np.dot(question_embeddings[question_index], question_embeddings.T)\n",
        "\n",
@@ -611,11 +550,11 @@
        "id": "aQIQSyF9GtSv"
      },
      "source": [
-        "#### Save the train split in JSONL format.\n",
+        "### Save the train split in JSONL format.\n",
        "\n",
        "The data must be formatted in JSONL format, which means each embedding dictionary is written as a JSON string on its own line.\n",
        "\n",
-        "See more information in the docs at [Input data format and structure](https://cloud.google.com/vertex-ai/docs/matching-engine/match-eng-setup#input-data-format)."
+        "See more information in the docs for [input data format and structure](https://cloud.google.com/vertex-ai/docs/vector-search/setup/format-structure#data-file-formats)."
      ]
    },
    {
@@ -626,8 +565,6 @@
      },
      "outputs": [],
      "source": [
-        "import tempfile\n",
-        "\n",
        "# Create temporary file to write embeddings to\n",
        "embeddings_file = tempfile.NamedTemporaryFile(suffix=\".json\", delete=False)\n",
        "\n",
@@ -642,10 +579,10 @@
      },
      "outputs": [],
      "source": [
-        "import json\n",
-        "\n",
+        "# Set batch size\n",
        "BATCH_SIZE = 100\n",
        "\n",
+        "# Create embeddings and write to a file\n",
        "with open(embeddings_file.name, \"a\") as f:\n",
        "    for i in tqdm(range(0, len(questions), BATCH_SIZE)):\n",
        "        id_chunk = ids[i : i + BATCH_SIZE]\n",
@@ -728,20 +665,7 @@
      "source": [
        "Create the ANN index configuration:\n",
        "\n",
-        "To learn more about configuring the index, see [Input data format and structure](https://cloud.google.com/vertex-ai/docs/vector-search/setup/setup).\n"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "Y4zooldkGoM4"
-      },
-      "outputs": [],
-      "source": [
-        "from google.cloud import aiplatform\n",
-        "\n",
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
+        "To learn more about configuring the index, see [Vector Search input data format and structure](https://cloud.google.com/vertex-ai/docs/vector-search/setup/setup).\n"
      ]
    },
    {
@@ -782,7 +706,7 @@
        "id": "0f1a9fbecabb"
      },
      "source": [
-        "Using the resource name, you can retrieve an existing MatchingEngineIndex."
+        "Using the resource name, you can retrieve an existing MatchingEngineIndex resource."
      ]
    },
    {
@@ -804,7 +728,7 @@
      "source": [
        "## Setup VPC peering network\n",
        "\n",
-        "To use a Vector Search index, you setup a VPC peering network between your project and the Vertex AI Vector Search service project. This eliminates additional hops in network traffic and allows using efficient gRPC protocol.\n",
+        "To use a Vector Search index, set up a VPC peering network between your project and the Vertex AI Vector Search service project. This eliminates additional hops in network traffic and allows using efficient gRPC protocol.\n",
        "\n",
        "Learn more about [VPC peering](https://cloud.google.com/vertex-ai/docs/general/vpc-peering).\n",
        "\n",
@@ -819,9 +743,9 @@
      "source": [
        "### Create VPC peering\n",
        "\n",
-        "For simplicity, we setup VPC peering to the ucaip-haystack-vpc-network network. You can create a different network for your project.\n",
+        "For simplicity, set up VPC peering to the `ucaip-haystack-vpc-network` network. You can create a different network for your project.\n",
        "\n",
-        "If you setup VPC peering with any other network, make sure that the network already exists and that your VM is running on that network."
+        "If you set up VPC peering with any other network, make sure that the network already exists and that your VM is running on that network."
      ]
    },
    {
@@ -833,10 +757,10 @@
      "outputs": [],
      "source": [
        "# This is for display only; you can name the range anything.\n",
+        "NETWORK = \"ucaip-haystack-vpc-network\"  # @param {type:\"string\"}\n",
        "PEERING_RANGE_NAME = \"vertex-ai-prediction-peering-range\"\n",
-        "NETWORK = \"ucaip-haystack-vpc-network\"\n",
        "\n",
-        "# NOTE: `prefix-length=16` means a CIDR block with mask /16 will be\n",
+        "# NOTE: `prefix-length=16` means a CIDR block with mask /16 is\n",
        "# reserved for use by Google services, such as Vertex AI.\n",
        "! gcloud compute addresses create $PEERING_RANGE_NAME \\\n",
        "  --global \\\n",
@@ -856,14 +780,14 @@
        "\n",
        "Next, create the connection for VPC peering.\n",
        "\n",
-        "*Note:* If you get a PERMISSION DENIED, you may not have the neccessary role 'Compute Network Admin' set for your default service account. In the Cloud Console, do the following steps.\n",
+        "**Note:** If you get a PERMISSION DENIED, you may not have the neccessary 'Compute Network Admin' role set for your default service account. In the Cloud Console, do the following:\n",
        "\n",
-        "1. Goto `IAM & Admin`.\n",
+        "1. Goto **IAM & Admin**.\n",
        "2. Find your service account.\n",
        "3. Click edit icon.\n",
-        "4. Select `Add Another Role`.\n",
-        "5. Enter 'Compute Network Admin'.\n",
-        "6. Select `Save`."
+        "4. Select **Add Another Role**.\n",
+        "5. Enter **Compute Network Admin**.\n",
+        "6. Select **Save**."
      ]
    },
    {
@@ -922,18 +846,10 @@
      "source": [
        "# Retrieve the project number\n",
        "PROJECT_NUMBER = !gcloud projects list --filter=\"PROJECT_ID:'{PROJECT_ID}'\" --format='value(PROJECT_NUMBER)'\n",
-        "PROJECT_NUMBER = PROJECT_NUMBER[0]"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "cd58eb809f71"
-      },
-      "outputs": [],
-      "source": [
-        "full_network_name = f\"projects/{PROJECT_NUMBER}/global/networks/{NETWORK}\""
+        "PROJECT_NUMBER = PROJECT_NUMBER[0]\n",
+        "\n",
+        "full_network_name = f\"projects/{PROJECT_NUMBER}/global/networks/{NETWORK}\"\n",
+        "print(full_network_name)"
      ]
    },
    {
@@ -986,6 +902,7 @@
      },
      "outputs": [],
      "source": [
+        "# Set a unique id for your deployed index\n",
        "DEPLOYED_INDEX_ID = \"deployed_index_id_unique\""
      ]
    },
@@ -997,6 +914,7 @@
      },
      "outputs": [],
      "source": [
+        "# Deploy your ANN index to the index endpoint\n",
        "my_index_endpoint = my_index_endpoint.deploy_index(\n",
        "    index=tree_ah_index, deployed_index_id=DEPLOYED_INDEX_ID\n",
        ")\n",
@@ -1014,7 +932,7 @@
        "\n",
        "After you've built your indexes, you can query against the deployed index to find nearest neighbors.\n",
        "\n",
-        "Note: For the DOT_PRODUCT_DISTANCE distance type, the \"distance\" property returned with each MatchNeighbor actually refers to the similarity."
+        "**Note:** For the **DOT_PRODUCT_DISTANCE** distance type, the \"distance\" property returned with each MatchNeighbor actually refers to the similarity."
      ]
    },
    {
@@ -1086,13 +1004,13 @@
      },
      "source": [
        "## Storing and retrieving titles from a Redis data store\n",
-        "When you productionize this code into a service, you'll need to convert the nearest nearest id's returned from Vertex AI Vector Search into data usable by downstream services.\n",
+        "When you productionize this code into a service, you need to convert the nearest ids returned from Vertex AI Vector Search into usable data for downstream services.\n",
        "\n",
-        "In this case, you'll need to convert the id's to titles.\n",
+        "In this case, you need to convert the ids to titles.\n",
        "\n",
        "You can use Google Cloud's Memorystore to deploy a managed Redis instance to save the id-title key-value pairs.\n",
        "\n",
-        "See more information on [Memorystore](https://cloud.google.com/memorystore/docs/redis/create-manage-instances?hl=en)"
+        "See more information on [Memorystore](https://cloud.google.com/memorystore/docs/redis/create-manage-instances?hl=en)."
      ]
    },
    {
@@ -1103,10 +1021,11 @@
      },
      "outputs": [],
      "source": [
+        "# Set a display name for your Redis instance\n",
        "REDIS_INSTANCE_NAME = \"stackoverflow-questions-unique\"\n",
        "\n",
        "# Create a Redis instance\n",
-        "! gcloud redis instances create '{REDIS_INSTANCE_NAME}' --size=5 --region={REGION} --network={VPC_NETWORK_FULL} --connect-mode=private-service-access"
+        "! gcloud redis instances create '{REDIS_INSTANCE_NAME}' --size=5 --region={LOCATION} --network={VPC_NETWORK_FULL} --connect-mode=private-service-access"
      ]
    },
    {
@@ -1117,12 +1036,10 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
        "# Get host and port info\n",
        "if not os.getenv(\"IS_TESTING\"):\n",
-        "    REDIS_HOST = ! gcloud redis instances list --filter=\"INSTANCE_NAME:'{REDIS_INSTANCE_NAME}'\" --region {REGION}  --format='value(HOST)'\n",
-        "    REDIS_PORT = ! gcloud redis instances list --filter=\"INSTANCE_NAME:'{REDIS_INSTANCE_NAME}'\" --region {REGION} --format='value(PORT)'\n",
+        "    REDIS_HOST = ! gcloud redis instances list --filter=\"INSTANCE_NAME:'{REDIS_INSTANCE_NAME}'\" --region {LOCATION}  --format='value(HOST)'\n",
+        "    REDIS_PORT = ! gcloud redis instances list --filter=\"INSTANCE_NAME:'{REDIS_INSTANCE_NAME}'\" --region {LOCATION} --format='value(PORT)'\n",
        "\n",
        "    if isinstance(REDIS_HOST, list):\n",
        "        REDIS_HOST = REDIS_HOST[0]\n",
@@ -1143,8 +1060,6 @@
      "outputs": [],
      "source": [
        "# Connect to the instance\n",
-        "import redis\n",
-        "\n",
        "if not os.getenv(\"IS_TESTING\"):\n",
        "    redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT)"
      ]
@@ -1170,8 +1085,8 @@
      },
      "outputs": [],
      "source": [
+        "# Verify that redis can retrieve the correct information\n",
        "if not os.getenv(\"IS_TESTING\"):\n",
-        "    # Verify that redis can retrieve the correct information\n",
        "    [\n",
        "        f\"Actual = {title}, Retrieved = {redis_client.get(str(id))}\"\n",
        "        for id, title in list(zip(df.id, df.title))[:10]\n",
@@ -1205,28 +1120,20 @@
        "# Delete indexes\n",
        "tree_ah_index.delete()\n",
        "\n",
+        "# Delete cloud storage bucket\n",
        "delete_bucket = False\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
-        "    ! gsutil rm -rf {BUCKET_URI}"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "d2fcf9468031"
-      },
-      "outputs": [],
-      "source": [
+        "if delete_bucket:\n",
+        "    ! gsutil rm -rf {BUCKET_URI}\n",
+        "\n",
        "# Delete redis instance\n",
-        "! gcloud redis instances delete '{REDIS_INSTANCE_NAME}' --region {REGION} --quiet"
+        "! gcloud redis instances delete '{REDIS_INSTANCE_NAME}' --region {LOCATION} --quiet"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
-      "name": "sdk_matching_engine_create_stack_overflow_embeddings.ipynb",
+      "name": "sdk_vector_search_create_stack_overflow_embeddings.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
@@ -30,21 +30,21 @@
      },
      "source": [
        "# Using Vertex AI Vector Search and Vertex AI embeddings for text for StackOverflow Questions \n",
-        "![ ](https://www.google-analytics.com/collect?v=2&tid=G-L6X3ECH596&cid=1&en=page_view&sid=1&dt=sdk_matching_engine_create_stack_overflow_embeddings_vertex.ipynb&dl=notebooks%2Fofficial%2Fmatching_engine%2Fsdk_matching_engine_create_stack_overflow_embeddings_vertex.ipynb)\n",
+        "![ ](https://www.google-analytics.com/collect?v=2&tid=G-L6X3ECH596&cid=1&en=page_view&sid=1&dt=sdk_vector_search_create_stack_overflow_embeddings_vertex.ipynb&dl=notebooks%2Fofficial%2Fvector_search%2Fsdk_vector_search_create_stack_overflow_embeddings_vertex.ipynb)\n",
        "<table align=\"left\">\n",
        "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/matching_engine/sdk_matching_engine_create_stack_overflow_embeddings_vertex.ipynb\">\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vector_search/sdk_vector_search_create_stack_overflow_embeddings_vertex.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/matching_engine/sdk_matching_engine_create_stack_overflow_embeddings_vertex.ipynb\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vector_search/sdk_vector_search_create_stack_overflow_embeddings_vertex.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/matching_engine/sdk_matching_engine_create_stack_overflow_embeddings_vertex\">\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/vector_search/sdk_vector_search_create_stack_overflow_embeddings_vertex\">\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",
@@ -649,7 +649,7 @@
        "\n",
        "The data must be formatted in JSONL format, which means each embedding dictionary is written as an individual JSON object on its own line.\n",
        "\n",
-        "For more information, see [Input data format and structure](https://cloud.google.com/vertex-ai/docs/matching-engine/match-eng-setup/format-structure#data-file-formats)."
+        "For more information, see [Input data format and structure](https://cloud.google.com/vertex-ai/docs/vector-search/setup/format-structure#data-file-formats)."
      ]
    },
    {
@@ -856,7 +856,7 @@
        "id": "0f1a9fbecabb"
      },
      "source": [
-        "Use the resource name to retrieve an existing MatchingEngineIndex."
+        "Use the resource name to retrieve an existing Index resource."
      ]
    },
    {
@@ -1073,7 +1073,7 @@
  "metadata": {
    "colab": {
      "collapsed_sections": [],
-      "name": "sdk_matching_engine_create_stack_overflow_embeddings_vertex.ipynb",
+      "name": "sdk_vector_search_create_stack_overflow_embeddings_vertex.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
@@ -4,6 +4,7 @@
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
+        "cellView": "form",
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
@@ -30,23 +31,27 @@
      },
      "source": [
        "# Create Vertex AI Vector Search index\n",
-        "![ ](https://www.google-analytics.com/collect?v=2&tid=G-L6X3ECH596&cid=1&en=page_view&sid=1&dt=sdk_matching_engine_for_indexing.ipynb&dl=notebooks%2Fofficial%2Fmatching_engine%2Fsdk_matching_engine_for_indexing.ipynb)\n",
+        "\n",
+        "\n",
        "<table align=\"left\">\n",
-        "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/matching_engine/sdk_matching_engine_for_indexing.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/vector_search/sdk_vector_search_for_indexing.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/matching_engine/sdk_matching_engine_for_indexing.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%2Fvector_search%2Fsdk_vector_search_for_indexing.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/vector_search/sdk_vector_search_for_indexing.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/matching_engine/sdk_matching_engine_for_indexing.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/vector_search/sdk_vector_search_for_indexing.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,7 +80,7 @@
        "\n",
        "In this notebook, you learn how to create Approximate Nearest Neighbor (ANN) Index, query against indexes, and validate the performance of the index. \n",
        "\n",
-        "This tutorial uses the following Google Cloud ML services:\n",
+        "This tutorial uses the following Vertex AI services:\n",
        "\n",
        "- Vertex AI Vector Search\n",
        "\n",
@@ -98,7 +103,16 @@
        "\n",
        "The dataset used for this tutorial is the [GloVe dataset](https://nlp.stanford.edu/projects/glove/).\n",
        "\n",
-        "\"GloVe is an unsupervised learning algorithm for obtaining vector representations for words. Training is performed on aggregated global word-word co-occurrence statistics from a corpus. The resulting representations showcase interesting linear substructures of the word vector space.\"\n"
+        "GloVe is an unsupervised learning algorithm for obtaining vector representations for words. Training is performed on aggregated global word-word co-occurrence statistics from a corpus. The resulting representations showcase interesting linear substructures of the word vector space.\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8925ff9e165e"
+      },
+      "source": [
+        "## Get started"
      ]
    },
    {
@@ -107,9 +121,7 @@
        "id": "f0f1bea346db"
      },
      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the latest versions of packages required to execute this notebook."
+        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
@@ -121,7 +133,7 @@
      "outputs": [],
      "source": [
        "# Install the packages\n",
-        "! pip3 install --upgrade google-cloud-aiplatform==1.34.0 \\\n",
+        "! pip3 install --upgrade google-cloud-aiplatform \\\n",
        "                         google-cloud-storage \\\n",
        "                         grpcio-tools \\\n",
        "                         h5py"
@@ -133,7 +145,9 @@
        "id": "5b08ba354c6e"
      },
      "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."
      ]
    },
    {
@@ -144,26 +158,63 @@
      },
      "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": "dd28c9e4f067"
+        "id": "7b49231643e4"
      },
      "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": "7176ea64999b"
+      },
+      "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": "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": "ee2dcc922f3d"
+      },
+      "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)."
      ]
    },
    {
@@ -175,103 +226,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": "4f4512bf63b3"
-      },
-      "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": "474be5183c27"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "949271bfebe3"
-      },
-      "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": "b65b4ce80d9a"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "985cdbfe7372"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "fbc9cd30cc4b"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "79efab26ad02"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "a336a05c6149"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0c0a44fa330f"
-      },
-      "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\"}"
      ]
    },
    {
@@ -346,36 +301,36 @@
        "id": "d5de53b31bf1"
      },
      "source": [
-        "## Make sure the following cells are run from inside the VPC network that you created in the previous step.\n",
-        "\n",
-        "* **WARNING:** The `MatchingIndexEndpoint.match` method (to create online queries against your deployed index) has to be executed in a Vertex AI Workbench notebook instance that is created with the following requirements:\n",
-        "  * **In the same region as where your ANN service is deployed** (for example, if you set `REGION = \"us-central1\"` as same as the tutorial, the notebook instance has to be in `us-central1`).\n",
-        "  * **Make sure you select the VPC network you created for ANN service** (instead of using the \"default\" one). That is, you have to create the VPC network below and then create a new notebook instance that uses that VPC.  \n",
-        "  * If you run it in the colab or a Vertex AI Workbench notebook instance in a different VPC network or region, \"Create Online Queries\" section will fail."
+        "**WARNING:** The `MatchingIndexEndpoint.match` method (to create online queries against your deployed index) has to be executed in a Vertex AI Workbench notebook instance that is created with the following requirements:\n",
+        "  * **In the same region where your ANN service is deployed** (for example, if you set `LOCATION = \"us-central1\"` as same as the tutorial, the notebook instance has to be in `us-central1`).\n",
+        "  \n",
+        "  * **Make sure you select the VPC network you created for ANN service** (instead of using the \"default\" one). That is, you have to create a new notebook instance that uses the VPC network you created earlier. Run the rest of the tutorial from that notebook instance.\n",
+        "  * If you run it in a Colab or a Vertex AI Workbench notebook in a different VPC network or region, \"Create Online Queries\" section is expected to fail."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
-        "id": "hhq5zEbGg0XX"
+        "id": "17ba55acc776"
      },
      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel."
+        "## Import the required libraries\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
-        "id": "EzrelQZ22IZj"
+        "id": "56182a4d44c4"
      },
      "outputs": [],
      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "import json\n",
        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
+        "import h5py\n",
+        "from google.cloud import aiplatform\n",
+        "from google.cloud.aiplatform.matching_engine.matching_engine_index_endpoint import \\\n",
+        "    Namespace"
      ]
    },
    {
@@ -384,7 +339,7 @@
        "id": "zgPO1eR3CYjk"
      },
      "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."
      ]
@@ -417,7 +372,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": "d5868f2942b8"
+      },
+      "source": [
+        "## 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)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "0a2809ebac11"
+      },
+      "outputs": [],
+      "source": [
+        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
@@ -428,8 +405,7 @@
      "source": [
        "## Prepare the data\n",
        "\n",
-        "The GloVe dataset consists of a set of pre-trained embeddings. The embeddings are split into a \"train\" split, and a \"test\" split.\n",
-        "You'll create a vector search index from the \"train\" split, and use the embedding vectors in the \"test\" split as query vectors to test the vector search index.\n",
+        "The GloVe dataset consists of a set of pre-trained embeddings. The embeddings are split into a \"train\" split, and a \"test\" split. You create a vector search index from the \"train\" split, and use the embedding vectors in the \"test\" split as query vectors to test the vector search index.\n",
        "\n",
        "**Note:** While the data split uses the term \"train\", these are pre-trained embeddings and therefore are ready to be indexed for search. The terms \"train\" and \"test\" split are used just to be consistent with machine learning terminology.\n",
        "\n",
@@ -464,8 +440,6 @@
      },
      "outputs": [],
      "source": [
-        "import h5py\n",
-        "\n",
        "# The number of nearest neighbors to be retrieved from database for each query.\n",
        "NUM_NEIGHBOURS = 10\n",
        "\n",
@@ -482,6 +456,7 @@
      },
      "outputs": [],
      "source": [
+        "# check the first record\n",
        "train[0]"
      ]
    },
@@ -491,12 +466,12 @@
        "id": "aQIQSyF9GtSv"
      },
      "source": [
-        "#### Save the train split in JSONL format.\n",
+        "### Save the train split in JSONL format.\n",
        "\n",
        "The data must be formatted in JSONL format, which means each embedding dictionary is written as a JSON string on its own line.\n",
        "\n",
        "Additionally, to demonstrate the filtering functionality, the `restricts` key is set such that each embedding has a different `class`, `even` or `odd`. These are used during the later matching step to filter for results.\n",
-        "For additional information of filtering, see [Filter vector matches](https://cloud.google.com/vertex-ai/docs/matching-engine/filtering)"
+        "For additional information of filtering, see [Filter vector matches](https://cloud.google.com/vertex-ai/docs/vector-search/filtering)."
      ]
    },
    {
@@ -507,8 +482,6 @@
      },
      "outputs": [],
      "source": [
-        "import json\n",
-        "\n",
        "with open(\"glove100.json\", \"w\") as f:\n",
        "    embeddings_formatted = [\n",
        "        json.dumps(\n",
@@ -546,7 +519,7 @@
      },
      "outputs": [],
      "source": [
-        "EMBEDDINGS_INITIAL_URI = f\"{BUCKET_URI}/matching_engine/initial/\"\n",
+        "EMBEDDINGS_INITIAL_URI = f\"{BUCKET_URI}/vector_search/initial/\"\n",
        "! gsutil cp glove100.json {EMBEDDINGS_INITIAL_URI}"
      ]
    },
@@ -559,15 +532,6 @@
        "## Create Indexes\n"
      ]
    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "qhIBCQ7dDSbW"
-      },
-      "source": [
-        "### Create ANN Index (for Production Usage)"
-      ]
-    },
    {
      "cell_type": "code",
      "execution_count": null,
@@ -576,11 +540,23 @@
      },
      "outputs": [],
      "source": [
+        "# set no.of dimensions for your embeddings\n",
        "DIMENSIONS = 100\n",
+        "# set the dispaly name for ann index\n",
        "DISPLAY_NAME = \"glove_100_1\"\n",
+        "# set the display name for brute force index\n",
        "DISPLAY_NAME_BRUTE_FORCE = DISPLAY_NAME + \"_brute_force\""
      ]
    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "qhIBCQ7dDSbW"
+      },
+      "source": [
+        "### Create ANN Index (for Production Usage)"
+      ]
+    },
    {
      "cell_type": "markdown",
      "metadata": {
@@ -589,22 +565,7 @@
      "source": [
        "Create the ANN index configuration:\n",
        "\n",
-        "To learn more about configuring the index, see [Input data format and structure](https://cloud.google.com/vertex-ai/docs/matching-engine/match-eng-setup/format-structure).\n"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "Y4zooldkGoM4"
-      },
-      "outputs": [],
-      "source": [
-        "import os\n",
-        "\n",
-        "from google.cloud import aiplatform\n",
-        "\n",
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
+        "To learn more about configuring the index, see [Input data format and structure](https://cloud.google.com/vertex-ai/docs/vector-search/setup/format-structure).\n"
      ]
    },
    {
@@ -646,7 +607,7 @@
        "id": "0f1a9fbecabb"
      },
      "source": [
-        "Using the resource name, you can retrieve an existing MatchingEngineIndex."
+        "Using the resource name, you can retrieve an existing Vector Search Index."
      ]
    },
    {
@@ -724,7 +685,7 @@
      "source": [
        "## Update Indexes\n",
        "\n",
-        "Create incremental data file.\n"
+        "Create an incremental data file.\n"
      ]
    },
    {
@@ -771,7 +732,7 @@
      },
      "outputs": [],
      "source": [
-        "EMBEDDINGS_UPDATE_URI = f\"{BUCKET_URI}/matching-engine/incremental/\""
+        "EMBEDDINGS_UPDATE_URI = f\"{BUCKET_URI}/vector-search/incremental/\""
      ]
    },
    {
@@ -791,7 +752,7 @@
        "id": "aiXtF_x0GoM6"
      },
      "source": [
-        "Create update index request\n"
+        "Create update index request.\n"
      ]
    },
    {
@@ -825,7 +786,7 @@
        "id": "qV2xjAnDDObD"
      },
      "source": [
-        "## Create an IndexEndpoint with VPC Network"
+        "## Create an IndexEndpoint in your VPC Network"
      ]
    },
    {
@@ -839,8 +800,7 @@
        "# Retrieve the project number\n",
        "PROJECT_NUMBER = !gcloud projects list --filter=\"PROJECT_ID:'{PROJECT_ID}'\" --format='value(PROJECT_NUMBER)'\n",
        "PROJECT_NUMBER = PROJECT_NUMBER[0]\n",
-        "\n",
-        "VPC_NETWORK = \"[your-network-name]\"\n",
+        "# Get the full network resource name\n",
        "VPC_NETWORK_FULL = \"projects/{}/global/networks/{}\".format(PROJECT_NUMBER, VPC_NETWORK)\n",
        "VPC_NETWORK_FULL"
      ]
@@ -853,6 +813,7 @@
      },
      "outputs": [],
      "source": [
+        "# Create your IndexEndpoint\n",
        "my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint.create(\n",
        "    display_name=\"index_endpoint_for_demo\",\n",
        "    description=\"index endpoint description\",\n",
@@ -898,6 +859,7 @@
      },
      "outputs": [],
      "source": [
+        "# Set an id for your ann index deployment\n",
        "DEPLOYED_INDEX_ID = \"tree_ah_glove_deployed_unique\""
      ]
    },
@@ -909,6 +871,7 @@
      },
      "outputs": [],
      "source": [
+        "# Deploy your ann index\n",
        "my_index_endpoint = my_index_endpoint.deploy_index(\n",
        "    index=tree_ah_index, deployed_index_id=DEPLOYED_INDEX_ID\n",
        ")\n",
@@ -933,6 +896,7 @@
      },
      "outputs": [],
      "source": [
+        "# Set an id for your brute force index deployment\n",
        "DEPLOYED_BRUTE_FORCE_INDEX_ID = \"glove_brute_force_deployed_unique\""
      ]
    },
@@ -944,6 +908,7 @@
      },
      "outputs": [],
      "source": [
+        "# Deploy your brute force index\n",
        "my_index_endpoint = my_index_endpoint.deploy_index(\n",
        "    index=brute_force_index, deployed_index_id=DEPLOYED_BRUTE_FORCE_INDEX_ID\n",
        ")\n",
@@ -957,7 +922,7 @@
        "id": "6LCGvBNvBd8D"
      },
      "source": [
-        "## Create Online Queries\n",
+        "## Create online queries\n",
        "\n",
        "After you've built your indexes, you may query against the deployed index through the online querying gRPC API (Match service) within the virtual machine instances from the same region (for example, 'us-central1' in this tutorial).\n",
        "\n",
@@ -972,14 +937,10 @@
      },
      "outputs": [],
      "source": [
-        "# Test query\n",
-        "from google.cloud.aiplatform.matching_engine.matching_engine_index_endpoint import \\\n",
-        "    Namespace\n",
-        "\n",
-        "# Test query\n",
+        "# Use match service with a test query\n",
        "response = my_index_endpoint.match(\n",
        "    deployed_index_id=DEPLOYED_INDEX_ID,\n",
-        "    queries=test[:1],\n",
+        "    queries=test[:1].tolist(),\n",
        "    num_neighbors=NUM_NEIGHBOURS,\n",
        "    filter=[Namespace(\"class\", [\"even\"])],\n",
        ")\n",
@@ -993,7 +954,7 @@
        "id": "xeUZO3bAGoM-"
      },
      "source": [
-        "### Compute Recall\n",
+        "### Compute recall\n",
        "\n",
        "Use the deployed brute force Index as the ground truth to calculate the recall of ANN Index. You can run multiple queries in a single match call."
      ]
@@ -1006,15 +967,15 @@
      },
      "outputs": [],
      "source": [
-        "# Retrieve nearest neighbors for both the tree-AH index and the brute-force index\n",
+        "# Retrieve nearest neighbors for both the tree-AH index and the brute force index\n",
        "tree_ah_response_test = my_index_endpoint.match(\n",
        "    deployed_index_id=DEPLOYED_INDEX_ID,\n",
-        "    queries=list(test),\n",
+        "    queries=test[:].tolist(),\n",
        "    num_neighbors=NUM_NEIGHBOURS,\n",
        ")\n",
        "brute_force_response_test = my_index_endpoint.match(\n",
        "    deployed_index_id=DEPLOYED_BRUTE_FORCE_INDEX_ID,\n",
-        "    queries=list(test),\n",
+        "    queries=test[:].tolist(),\n",
        "    num_neighbors=NUM_NEIGHBOURS,\n",
        ")"
      ]
@@ -1027,7 +988,7 @@
      },
      "outputs": [],
      "source": [
-        "# Calculate recall by determining how many neighbors were correctly retrieved as compared to the brute-force option.\n",
+        "# Calculate recall by determining how many neighbors were correctly retrieved as compared to the brute force option.\n",
        "recalled_neighbors = 0\n",
        "for tree_ah_neighbors, brute_force_neighbors in zip(\n",
        "    tree_ah_response_test, brute_force_response_test\n",
@@ -1067,8 +1028,6 @@
      },
      "outputs": [],
      "source": [
-        "import os\n",
-        "\n",
        "delete_bucket = False\n",
        "\n",
        "# Force undeployment of indexes and delete endpoint\n",
@@ -1078,7 +1037,7 @@
        "tree_ah_index.delete()\n",
        "brute_force_index.delete()\n",
        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "if delete_bucket:\n",
        "    ! gsutil rm -rf {BUCKET_URI}"
      ]
    }
@@ -1086,7 +1045,7 @@
  "metadata": {
    "colab": {
      "collapsed_sections": [],
-      "name": "sdk_matching_engine_for_indexing.ipynb",
+      "name": "sdk_vector_search_for_indexing.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
@@ -1,705 +0,0 @@
-{
-  "cells": [
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ur8xi4C7S06n"
-      },
-      "outputs": [],
-      "source": [
-        "# Copyright 2023 Google LLC\n",
-        "#\n",
-        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
-        "# you may not use this file except in compliance with the License.\n",
-        "# You may obtain a copy of the License at\n",
-        "#\n",
-        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
-        "#\n",
-        "# Unless required by applicable law or agreed to in writing, software\n",
-        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
-        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
-        "# See the License for the specific language governing permissions and\n",
-        "# limitations under the License."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "JAPoU8Sm5E6e"
-      },
-      "source": [
-        "# Train a pytorch model with Vertex AI SDK and Bigframes\n",
-        "\n",
-        "<table align=\"left\">\n",
-        "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vertex_ai_sdk/remote_training_bigframes_pytorch.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
-        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vertex_ai_sdk/remote_training_bigframes_pytorch.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/vertex_ai_sdk/remote_training_bigframes_pytorch.ipynb\">\n",
-        "       <img src=\"https://www.gstatic.com/cloud/images/navigation/vertex-ai.svg\" alt=\"Vertex AI logo\">Open in Vertex AI Workbench\n",
-        "    </a>\n",
-        "</table>"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "tvgnzT1CKxrO"
-      },
-      "source": [
-        "## Overview\n",
-        "\n",
-        "This tutorial demonstrates how to train a pytorch model using Vertex AI local-to-remote training with Vertex AI SDK and BigQuery Bigframes as the data source.\n",
-        "\n",
-        "Learn more about [bigframes](https://cloud.google.com/bigquery/docs/)."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "d975e698c9a4"
-      },
-      "source": [
-        "### Objective\n",
-        "\n",
-        "In this tutorial, you learn to use `Vertex AI SDK` with Bigframes as input data source.\n",
-        "\n",
-        "\n",
-        "This tutorial uses the following Google Cloud ML services:\n",
-        "\n",
-        "- `Vertex AI Training`\n",
-        "- `Vertex AI Remote Training`\n",
-        "\n",
-        "\n",
-        "The steps performed include:\n",
-        "\n",
-        "- Initialize a dataframe from a BigQuery table and split the dataset\n",
-        "- Perform transformations as a Vertex AI remote training.\n",
-        "- Train the model remotely and evaluate the model locally\n",
-        "\n",
-        "**Local-to-remote training**\n",
-        "\n",
-        "```\n",
-        "import vertexai\n",
-        "from my_module import MyModelClass\n",
-        "\n",
-        "vertexai.preview.init(remote=True, project=\"my-project\", location=\"my-location\", staging_bucket=\"gs://my-bucket\")\n",
-        "\n",
-        "# Wrap the model class with `vertex_ai.preview.remote`\n",
-        "MyModelClass = vertexai.preview.remote(MyModelClass)\n",
-        "\n",
-        "# Instantiate the class\n",
-        "model = MyModelClass(...)\n",
-        "\n",
-        "# Optional set remote config\n",
-        "model.fit.vertex.remote_config.display_name = \"MyModelClass-remote-training\"\n",
-        "model.fit.vertex.remote_config.staging_bucket = \"gs://my-bucket\"\n",
-        "\n",
-        "# This `fit` call will be executed remotely\n",
-        "model.fit(...)\n",
-        "```"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "08d289fa873f"
-      },
-      "source": [
-        "### Dataset\n",
-        "\n",
-        "This tutorial uses the <a href=\"https://scikit-learn.org/stable/auto_examples/datasets/plot_iris_dataset.html\">IRIS dataset</a>, which predicts the iris species."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "aed92deeb4a0"
-      },
-      "source": [
-        "### Costs\n",
-        "\n",
-        "This tutorial uses billable components of Google Cloud:\n",
-        "\n",
-        "* Vertex AI\n",
-        "* BigQuery\n",
-        "* Cloud Storage\n",
-        "\n",
-        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing),\n",
-        "[BigQuery pricing](https://cloud.google.com/bigquery/pricing),\n",
-        "and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), \n",
-        "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n",
-        "to generate a cost estimate based on your projected usage."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "i7EUnXsZhAGF"
-      },
-      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook. "
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "2b4ef9b72d43"
-      },
-      "outputs": [],
-      "source": [
-        "# Install the packages\n",
-        "! pip3 install --upgrade --quiet google-cloud-aiplatform[preview]\n",
-        "! pip3 install --upgrade --quiet bigframes\n",
-        "! pip3 install --upgrade --quiet torch"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "58707a750154"
-      },
-      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "f200f10a1da3"
-      },
-      "outputs": [],
-      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
-        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "BF1j6f9HApxa"
-      },
-      "source": [
-        "## Before you begin\n",
-        "\n",
-        "### Set up your Google Cloud project\n",
-        "\n",
-        "**The following steps are required, regardless of your notebook environment.**\n",
-        "\n",
-        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
-        "\n",
-        "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
-        "\n",
-        "3. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
-        "\n",
-        "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "WReHDGG5g0XY"
-      },
-      "source": [
-        "#### Set your project ID\n",
-        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "* Run `gcloud config list`.\n",
-        "* Run `gcloud projects list`.\n",
-        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "oM1iC_MfAts1"
-      },
-      "outputs": [],
-      "source": [
-        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# Set the project id\n",
-        "! gcloud config set project {PROJECT_ID}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "region"
-      },
-      "source": [
-        "#### Region\n",
-        "\n",
-        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "region"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "sBCra4QMA2wR"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "74ccc9e52986"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "de775a3773ba"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "254614fa0c46"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ef21552ccea8"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "603adbbf0532"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "f6b2ccc891ed"
-      },
-      "source": [
-        "**4. Service account or other**\n",
-        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "zgPO1eR3CYjk"
-      },
-      "source": [
-        "### Create a Cloud Storage bucket\n",
-        "\n",
-        "Create a storage bucket to store intermediate artifacts such as datasets."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "MzGDU7TWdts_"
-      },
-      "outputs": [],
-      "source": [
-        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "-EcIXiGsCePi"
-      },
-      "source": [
-        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "NIq7R4HZCfIc"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "960505627ddf"
-      },
-      "source": [
-        "### Import libraries and define constants"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "PyQmSRbKA8r-"
-      },
-      "outputs": [],
-      "source": [
-        "import bigframes.pandas as bf\n",
-        "import torch\n",
-        "import vertexai\n",
-        "from vertexai.preview import VertexModel\n",
-        "\n",
-        "bf.options.bigquery.location = \"us\"  # Dataset is in 'us' not 'us-central1'\n",
-        "bf.options.bigquery.project = PROJECT_ID\n",
-        "\n",
-        "from bigframes.ml.model_selection import \\\n",
-        "    train_test_split as bf_train_test_split"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "init_aip:mbsdk,all"
-      },
-      "source": [
-        "## Initialize Vertex AI SDK for Python\n",
-        "\n",
-        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "init_aip:mbsdk,all"
-      },
-      "outputs": [],
-      "source": [
-        "vertexai.init(\n",
-        "    project=PROJECT_ID,\n",
-        "    location=REGION,\n",
-        "    staging_bucket=BUCKET_URI,\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "105334524e96"
-      },
-      "source": [
-        "## Prepare the dataset\n",
-        "\n",
-        "Now load the Iris dataset and split the data into train and test sets."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "b44cdc4e03f1"
-      },
-      "outputs": [],
-      "source": [
-        "df = bf.read_gbq(\"bigquery-public-data.ml_datasets.iris\")\n",
-        "\n",
-        "species_categories = {\n",
-        "    \"versicolor\": 0,\n",
-        "    \"virginica\": 1,\n",
-        "    \"setosa\": 2,\n",
-        "}\n",
-        "df[\"species\"] = df[\"species\"].map(species_categories)\n",
-        "\n",
-        "# Assign an index column name\n",
-        "index_col = \"index\"\n",
-        "df.index.name = index_col"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "9cb8616b1997"
-      },
-      "outputs": [],
-      "source": [
-        "feature_columns = df[[\"sepal_length\", \"sepal_width\", \"petal_length\", \"petal_width\"]]\n",
-        "label_columns = df[[\"species\"]]\n",
-        "train_X, test_X, train_y, test_y = bf_train_test_split(\n",
-        "    feature_columns, label_columns, test_size=0.2\n",
-        ")\n",
-        "\n",
-        "print(\"X_train size: \", train_X.size)\n",
-        "print(\"X_test size: \", test_X.size)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "23fe7b734b08"
-      },
-      "outputs": [],
-      "source": [
-        "# Switch to remote mode for training\n",
-        "vertexai.preview.init(remote=True)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "5904a0f1bb03"
-      },
-      "source": [
-        "## PyTorch remote training with CPU (Custom PyTorch model)\n",
-        "\n",
-        "First, train a PyTorch model as a remote training job:\n",
-        "\n",
-        "- Reinitialize Vertex AI for remote training.\n",
-        "- Set TorchLogisticRegression for the remote training job.\n",
-        "- Invoke TorchLogisticRegression locally which will launch the remote training job."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "2a1b85195a17"
-      },
-      "outputs": [],
-      "source": [
-        "# define the custom model\n",
-        "class TorchLogisticRegression(VertexModel, torch.nn.Module):\n",
-        "    def __init__(self, input_size: int, output_size: int):\n",
-        "        torch.nn.Module.__init__(self)\n",
-        "        VertexModel.__init__(self)\n",
-        "        self.linear = torch.nn.Linear(input_size, output_size)\n",
-        "        self.softmax = torch.nn.Softmax(dim=1)\n",
-        "\n",
-        "    def forward(self, x):\n",
-        "        return self.softmax(self.linear(x))\n",
-        "\n",
-        "    @vertexai.preview.developer.mark.train()\n",
-        "    def train(self, X, y, num_epochs, lr):\n",
-        "        X = X.to(torch.float32)\n",
-        "        y = torch.flatten(y)  # necessary to get 1D tensor\n",
-        "        dataloader = torch.utils.data.DataLoader(\n",
-        "            torch.utils.data.TensorDataset(X, y),\n",
-        "            batch_size=10,\n",
-        "            shuffle=True,\n",
-        "            generator=torch.Generator(device=X.device),\n",
-        "        )\n",
-        "\n",
-        "        criterion = torch.nn.CrossEntropyLoss()\n",
-        "        optimizer = torch.optim.SGD(self.parameters(), lr=lr)\n",
-        "\n",
-        "        for t in range(num_epochs):\n",
-        "            for batch, (X, y) in enumerate(dataloader):\n",
-        "                optimizer.zero_grad()\n",
-        "                pred = self(X)\n",
-        "                loss = criterion(pred, y)\n",
-        "                loss.backward()\n",
-        "                optimizer.step()\n",
-        "\n",
-        "    @vertexai.preview.developer.mark.predict()\n",
-        "    def predict(self, X):\n",
-        "        X = torch.tensor(X).to(torch.float32)\n",
-        "        with torch.no_grad():\n",
-        "            pred = torch.argmax(self(X), dim=1)\n",
-        "        return pred"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "4e35593f520a"
-      },
-      "outputs": [],
-      "source": [
-        "# Switch to remote mode for training\n",
-        "vertexai.preview.init(remote=True)\n",
-        "\n",
-        "# Instantiate model\n",
-        "model = TorchLogisticRegression(4, 3)\n",
-        "\n",
-        "# Train model on Vertex\n",
-        "model.train(train_X, train_y, num_epochs=200, lr=0.05)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "edf4d0708f02"
-      },
-      "source": [
-        "## Remote prediction\n",
-        "\n",
-        "Obtain predictions from the trained model."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "42dfbff0ca15"
-      },
-      "outputs": [],
-      "source": [
-        "# Remote evaluation\n",
-        "vertexai.preview.init(remote=True)\n",
-        "\n",
-        "predictions = model.predict(test_X)\n",
-        "\n",
-        "print(f\"Remote predictions: {predictions}\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "4340ed8316cd"
-      },
-      "source": [
-        "## Local evaluation\n",
-        "\n",
-        "Evaluate model results locally."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "eb27a31cec6f"
-      },
-      "outputs": [],
-      "source": [
-        "# User must convert bigframes to torch tensor for local evaluation\n",
-        "train_X_tensor = torch.from_numpy(\n",
-        "    train_X.to_pandas().reset_index().drop(columns=[\"index\"]).values.astype(float)\n",
-        ")\n",
-        "train_y_tensor = torch.from_numpy(\n",
-        "    train_y.to_pandas().reset_index().drop(columns=[\"index\"]).values.astype(float)\n",
-        ")\n",
-        "\n",
-        "test_X_tensor = torch.from_numpy(\n",
-        "    test_X.to_pandas().reset_index().drop(columns=[\"index\"]).values.astype(float)\n",
-        ")\n",
-        "test_y_tensor = torch.from_numpy(\n",
-        "    test_y.to_pandas().reset_index().drop(columns=[\"index\"]).values.astype(float)\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "7db44ad81389"
-      },
-      "outputs": [],
-      "source": [
-        "from sklearn.metrics import accuracy_score\n",
-        "\n",
-        "# Switch to local mode for evaluation\n",
-        "vertexai.preview.init(remote=False)\n",
-        "\n",
-        "# Evaluate model's accuracy score\n",
-        "print(\n",
-        "    f\"Train accuracy: {accuracy_score(train_y_tensor, model.predict(train_X_tensor))}\"\n",
-        ")\n",
-        "\n",
-        "print(f\"Test accuracy: {accuracy_score(test_y_tensor, model.predict(test_X_tensor))}\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "TpV-iwP9qw9c"
-      },
-      "source": [
-        "## Cleaning up\n",
-        "\n",
-        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
-        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
-        "\n",
-        "Otherwise, you can delete the individual resources you created in this tutorial:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "sx_vKniMq9ZX"
-      },
-      "outputs": [],
-      "source": [
-        "import os\n",
-        "\n",
-        "# Delete Cloud Storage objects that were created\n",
-        "delete_bucket = False\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
-        "    ! gsutil -m rm -r $BUCKET_URI"
-      ]
-    }
-  ],
-  "metadata": {
-    "colab": {
-      "collapsed_sections": [],
-      "name": "remote_training_bigframes_pytorch.ipynb",
-      "toc_visible": true
-    },
-    "kernelspec": {
-      "display_name": "Python 3",
-      "name": "python3"
-    }
-  },
-  "nbformat": 4,
-  "nbformat_minor": 0
-}
@@ -1,707 +0,0 @@
-{
-  "cells": [
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ur8xi4C7S06n"
-      },
-      "outputs": [],
-      "source": [
-        "# Copyright 2023 Google LLC\n",
-        "#\n",
-        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
-        "# you may not use this file except in compliance with the License.\n",
-        "# You may obtain a copy of the License at\n",
-        "#\n",
-        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
-        "#\n",
-        "# Unless required by applicable law or agreed to in writing, software\n",
-        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
-        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
-        "# See the License for the specific language governing permissions and\n",
-        "# limitations under the License."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "JAPoU8Sm5E6e"
-      },
-      "source": [
-        "# Train a scikit-learn model with Vertex AI SDK and Bigframes\n",
-        "\n",
-        "<table align=\"left\">\n",
-        "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vertex_ai_sdk/remote_training_bigframes_sklearn.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
-        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vertex_ai_sdk/remote_training_bigframes_sklearn.ipynb\">\n",
-        "        <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
-        "    </a>\n",
-        "  </td>\n",
-        "    <td>\n",
-        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/vertex_ai_sdk/remote_training_bigframes_sklearn.ipynb\">\n",
-        "       <img src=\"https://www.gstatic.com/cloud/images/navigation/vertex-ai.svg\" alt=\"Vertex AI logo\">Open in Vertex AI Workbench\n",
-        "    </a>\n",
-        "</table>"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "tvgnzT1CKxrO"
-      },
-      "source": [
-        "## Overview\n",
-        "\n",
-        "This tutorial demonstrates how to train a scikit-learn model using Vertex AI local-to-remote training with Vertex AI SDK and BigQuery Bigframes as the data source.\n",
-        "\n",
-        "Learn more about [bigframes](https://cloud.google.com/bigquery/docs/)."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "d975e698c9a4"
-      },
-      "source": [
-        "### Objective\n",
-        "\n",
-        "In this tutorial, you learn to use `Vertex AI SDK` with Bigframes as input data source.\n",
-        "\n",
-        "\n",
-        "This tutorial uses the following Google Cloud ML services:\n",
-        "\n",
-        "- `Vertex AI Training`\n",
-        "- `Vertex AI Remote Training`\n",
-        "\n",
-        "\n",
-        "The steps performed include:\n",
-        "\n",
-        "- Initialize a dataframe from a BigQuery table and split the dataset\n",
-        "- Perform transformations as a Vertex AI remote training.\n",
-        "- Train the model remotely and evaluate the model locally\n",
-        "\n",
-        "**Local-to-remote training**\n",
-        "\n",
-        "```\n",
-        "import vertexai\n",
-        "from my_module import MyModelClass\n",
-        "\n",
-        "vertexai.preview.init(remote=True, project=\"my-project\", location=\"my-location\", staging_bucket=\"gs://my-bucket\")\n",
-        "\n",
-        "# Wrap the model class with `vertex_ai.preview.remote`\n",
-        "MyModelClass = vertexai.preview.remote(MyModelClass)\n",
-        "\n",
-        "# Instantiate the class\n",
-        "model = MyModelClass(...)\n",
-        "\n",
-        "# Optional set remote config\n",
-        "model.fit.vertex.remote_config.display_name = \"MyModelClass-remote-training\"\n",
-        "model.fit.vertex.remote_config.staging_bucket = \"gs://my-bucket\"\n",
-        "\n",
-        "# This `fit` call will be executed remotely\n",
-        "model.fit(...)\n",
-        "```"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "08d289fa873f"
-      },
-      "source": [
-        "### Dataset\n",
-        "\n",
-        "This tutorial uses the <a href=\"https://scikit-learn.org/stable/auto_examples/datasets/plot_iris_dataset.html\">IRIS dataset</a>, which predicts the iris species."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "aed92deeb4a0"
-      },
-      "source": [
-        "### Costs\n",
-        "\n",
-        "This tutorial uses billable components of Google Cloud:\n",
-        "\n",
-        "* Vertex AI\n",
-        "* BigQuery\n",
-        "* Cloud Storage\n",
-        "\n",
-        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing),\n",
-        "[BigQuery pricing](https://cloud.google.com/bigquery/pricing),\n",
-        "and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), \n",
-        "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n",
-        "to generate a cost estimate based on your projected usage."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "i7EUnXsZhAGF"
-      },
-      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook. "
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "2b4ef9b72d43"
-      },
-      "outputs": [],
-      "source": [
-        "# Install the packages\n",
-        "! pip3 install --upgrade --quiet google-cloud-aiplatform[preview]\n",
-        "! pip3 install --upgrade --quiet bigframes"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "58707a750154"
-      },
-      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "f200f10a1da3"
-      },
-      "outputs": [],
-      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
-        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "BF1j6f9HApxa"
-      },
-      "source": [
-        "## Before you begin\n",
-        "\n",
-        "### Set up your Google Cloud project\n",
-        "\n",
-        "**The following steps are required, regardless of your notebook environment.**\n",
-        "\n",
-        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
-        "\n",
-        "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
-        "\n",
-        "3. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
-        "\n",
-        "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "WReHDGG5g0XY"
-      },
-      "source": [
-        "#### Set your project ID\n",
-        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "* Run `gcloud config list`.\n",
-        "* Run `gcloud projects list`.\n",
-        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "oM1iC_MfAts1"
-      },
-      "outputs": [],
-      "source": [
-        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# Set the project id\n",
-        "! gcloud config set project {PROJECT_ID}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "region"
-      },
-      "source": [
-        "#### Region\n",
-        "\n",
-        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "region"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "sBCra4QMA2wR"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "74ccc9e52986"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "de775a3773ba"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "254614fa0c46"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ef21552ccea8"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "603adbbf0532"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "f6b2ccc891ed"
-      },
-      "source": [
-        "**4. Service account or other**\n",
-        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "zgPO1eR3CYjk"
-      },
-      "source": [
-        "### Create a Cloud Storage bucket\n",
-        "\n",
-        "Create a storage bucket to store intermediate artifacts such as datasets."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "MzGDU7TWdts_"
-      },
-      "outputs": [],
-      "source": [
-        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "-EcIXiGsCePi"
-      },
-      "source": [
-        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "NIq7R4HZCfIc"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "960505627ddf"
-      },
-      "source": [
-        "### Import libraries and define constants"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "PyQmSRbKA8r-"
-      },
-      "outputs": [],
-      "source": [
-        "import bigframes.pandas as bf\n",
-        "import vertexai\n",
-        "\n",
-        "bf.options.bigquery.location = \"us\"  # Dataset is in 'us' not 'us-central1'\n",
-        "bf.options.bigquery.project = PROJECT_ID\n",
-        "\n",
-        "from bigframes.ml.model_selection import \\\n",
-        "    train_test_split as bf_train_test_split"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "init_aip:mbsdk,all"
-      },
-      "source": [
-        "## Initialize Vertex AI SDK for Python\n",
-        "\n",
-        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "init_aip:mbsdk,all"
-      },
-      "outputs": [],
-      "source": [
-        "vertexai.init(\n",
-        "    project=PROJECT_ID,\n",
-        "    location=REGION,\n",
-        "    staging_bucket=BUCKET_URI,\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "105334524e96"
-      },
-      "source": [
-        "## Prepare the dataset\n",
-        "\n",
-        "Now load the Iris dataset and split the data into train and test sets."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "b44cdc4e03f1"
-      },
-      "outputs": [],
-      "source": [
-        "df = bf.read_gbq(\"bigquery-public-data.ml_datasets.iris\")\n",
-        "\n",
-        "species_categories = {\n",
-        "    \"versicolor\": 0,\n",
-        "    \"virginica\": 1,\n",
-        "    \"setosa\": 2,\n",
-        "}\n",
-        "df[\"species\"] = df[\"species\"].map(species_categories)\n",
-        "\n",
-        "# Assign an index column name\n",
-        "index_col = \"index\"\n",
-        "df.index.name = index_col"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "9cb8616b1997"
-      },
-      "outputs": [],
-      "source": [
-        "feature_columns = df[[\"sepal_length\", \"sepal_width\", \"petal_length\", \"petal_width\"]]\n",
-        "label_columns = df[[\"species\"]]\n",
-        "train_X, test_X, train_y, test_y = bf_train_test_split(\n",
-        "    feature_columns, label_columns, test_size=0.2\n",
-        ")\n",
-        "\n",
-        "print(\"X_train size: \", train_X.size)\n",
-        "print(\"X_test size: \", test_X.size)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "8306545fcc57"
-      },
-      "source": [
-        "## Feature transformation\n",
-        "\n",
-        "Next, you do feature transformations on the data using the Vertex AI remote training service.\n",
-        "\n",
-        "First, you re-initialize Vertex AI to enable remote training."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "55e701c31036"
-      },
-      "outputs": [],
-      "source": [
-        "# Switch to remote mode for training\n",
-        "vertexai.preview.init(remote=True)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "4a0e9d59b273"
-      },
-      "source": [
-        "### Execute remote job for fit_transform() on training data\n",
-        "\n",
-        "Next, indicate that the `StandardScalar` class is to be executed remotely. Then set up the data transform and call the `fit_transform()` method is executed remotely."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "90333089d362"
-      },
-      "outputs": [],
-      "source": [
-        "from sklearn.preprocessing import StandardScaler\n",
-        "\n",
-        "# Wrap classes to enable Vertex remote execution\n",
-        "StandardScaler = vertexai.preview.remote(StandardScaler)\n",
-        "\n",
-        "# Instantiate transformer\n",
-        "transformer = StandardScaler()\n",
-        "\n",
-        "# Execute transformer on Vertex (train_X is bigframes.dataframe.DataFrame, X_train is np.array)\n",
-        "X_train = transformer.fit_transform(train_X)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "6bf95574c907"
-      },
-      "source": [
-        "### Remote transform on test data"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "da6eea22a89a"
-      },
-      "outputs": [],
-      "source": [
-        "# Execute transformer on Vertex (test_X is bigframes.dataframe.DataFrame, X_test is np.array)\n",
-        "X_test = transformer.transform(test_X)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ddf906c886e4"
-      },
-      "source": [
-        "## Remote training\n",
-        "\n",
-        "First, train the scikit-learn model as a remote training job:\n",
-        "\n",
-        "- Set LogisticRegression for the remote training job.\n",
-        "- Invoke LogisticRegression locally which will launch the remote training job."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "c7b0116fa60c"
-      },
-      "outputs": [],
-      "source": [
-        "from sklearn.linear_model import LogisticRegression\n",
-        "\n",
-        "# Wrap classes to enable Vertex remote execution\n",
-        "LogisticRegression = vertexai.preview.remote(LogisticRegression)\n",
-        "\n",
-        "# Instantiate model, warm_start=True for uptraining\n",
-        "model = LogisticRegression(warm_start=True)\n",
-        "\n",
-        "# Train model on Vertex\n",
-        "model.fit(train_X, train_y)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ffe1d5903bcb"
-      },
-      "source": [
-        "## Remote prediction\n",
-        "\n",
-        "Obtain predictions from the trained model."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "d00ce35920fa"
-      },
-      "outputs": [],
-      "source": [
-        "# Remote evaluation\n",
-        "vertexai.preview.init(remote=True)\n",
-        "\n",
-        "predictions = model.predict(test_X)\n",
-        "\n",
-        "print(f\"Remote predictions: {predictions}\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "a8cd6cbd4403"
-      },
-      "source": [
-        "## Local evaluation\n",
-        "\n",
-        "Score model results locally."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "dc105dafdfb9"
-      },
-      "outputs": [],
-      "source": [
-        "# User must convert bigframes to pandas dataframe for local evaluation\n",
-        "train_X_pd = train_X.to_pandas().reset_index(drop=True)\n",
-        "train_y_pd = train_y.to_pandas().reset_index(drop=True)\n",
-        "\n",
-        "test_X_pd = test_X.to_pandas().reset_index(drop=True)\n",
-        "test_y_pd = test_y.to_pandas().reset_index(drop=True)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "25fec549de69"
-      },
-      "outputs": [],
-      "source": [
-        "# Switch to local mode for testing\n",
-        "vertexai.preview.init(remote=False)\n",
-        "\n",
-        "# Evaluate model's accuracy score\n",
-        "print(f\"Train accuracy: {model.score(train_X_pd, train_y_pd)}\")\n",
-        "\n",
-        "print(f\"Test accuracy: {model.score(test_X_pd, test_y_pd)}\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "TpV-iwP9qw9c"
-      },
-      "source": [
-        "## Cleaning up\n",
-        "\n",
-        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
-        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
-        "\n",
-        "Otherwise, you can delete the individual resources you created in this tutorial:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "sx_vKniMq9ZX"
-      },
-      "outputs": [],
-      "source": [
-        "import os\n",
-        "\n",
-        "# Delete Cloud Storage objects that were created\n",
-        "delete_bucket = False\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
-        "    ! gsutil -m rm -r $BUCKET_URI"
-      ]
-    }
-  ],
-  "metadata": {
-    "colab": {
-      "collapsed_sections": [],
-      "name": "remote_training_bigframes_sklearn.ipynb",
-      "toc_visible": true
-    },
-    "kernelspec": {
-      "display_name": "Python 3",
-      "name": "python3"
-    }
-  },
-  "nbformat": 4,
-  "nbformat_minor": 0
-}
@@ -1,643 +0,0 @@
-{
-  "cells": [
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ur8xi4C7S06n"
-      },
-      "outputs": [],
-      "source": [
-        "# Copyright 2023 Google LLC\n",
-        "#\n",
-        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
-        "# you may not use this file except in compliance with the License.\n",
-        "# You may obtain a copy of the License at\n",
-        "#\n",
-        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
-        "#\n",
-        "# Unless required by applicable law or agreed to in writing, software\n",
-        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
-        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
-        "# See the License for the specific language governing permissions and\n",
-        "# limitations under the License."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "JAPoU8Sm5E6e"
-      },
-      "source": [
-        "# Train a Tensorflow Keras model with Vertex AI SDK and Bigframes \n",
-        "\n",
-        "<table align=\"left\">\n",
-        "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vertex_ai_sdk/remote_training_bigframes_tensorflow.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
-        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vertex_ai_sdk/remote_training_bigframes_tensorflow.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/vertex_ai_sdk/remote_training_bigframes_tensorflow.ipynb\">\n",
-        "       <img src=\"https://www.gstatic.com/cloud/images/navigation/vertex-ai.svg\" alt=\"Vertex AI logo\">Open in Vertex AI Workbench\n",
-        "    </a>\n",
-        "</table>"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "tvgnzT1CKxrO"
-      },
-      "source": [
-        "## Overview\n",
-        "\n",
-        "This tutorial demonstrates how to train a tensorflow keras model using Vertex AI local-to-remote training with Vertex AI SDK and BigQuery Bigframes as the data source.\n",
-        "\n",
-        "Learn more about [bigframes](https://cloud.google.com/bigquery/docs/)."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "d975e698c9a4"
-      },
-      "source": [
-        "### Objective\n",
-        "\n",
-        "In this tutorial, you learn to use `Vertex AI SDK` with Bigframes as input data source.\n",
-        "\n",
-        "\n",
-        "This tutorial uses the following Google Cloud ML services:\n",
-        "\n",
-        "- `Vertex AI Training`\n",
-        "- `Vertex AI Remote Training`\n",
-        "\n",
-        "\n",
-        "The steps performed include:\n",
-        "\n",
-        "- Initialize a dataframe from a BigQuery table and split the dataset\n",
-        "- Perform transformations as a Vertex AI remote training.\n",
-        "- Train the model remotely and evaluate the model locally\n",
-        "\n",
-        "**Local-to-remote training**\n",
-        "\n",
-        "```\n",
-        "import vertexai\n",
-        "from my_module import MyModelClass\n",
-        "\n",
-        "vertexai.preview.init(remote=True, project=\"my-project\", location=\"my-location\", staging_bucket=\"gs://my-bucket\")\n",
-        "\n",
-        "# Wrap the model class with `vertex_ai.preview.remote`\n",
-        "MyModelClass = vertexai.preview.remote(MyModelClass)\n",
-        "\n",
-        "# Instantiate the class\n",
-        "model = MyModelClass(...)\n",
-        "\n",
-        "# Optional set remote config\n",
-        "model.fit.vertex.remote_config.display_name = \"MyModelClass-remote-training\"\n",
-        "model.fit.vertex.remote_config.staging_bucket = \"gs://my-bucket\"\n",
-        "\n",
-        "# This `fit` call will be executed remotely\n",
-        "model.fit(...)\n",
-        "```"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "08d289fa873f"
-      },
-      "source": [
-        "### Dataset\n",
-        "\n",
-        "This tutorial uses the <a href=\"https://scikit-learn.org/stable/auto_examples/datasets/plot_iris_dataset.html\">IRIS dataset</a>, which predicts the iris species."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "aed92deeb4a0"
-      },
-      "source": [
-        "### Costs\n",
-        "\n",
-        "This tutorial uses billable components of Google Cloud:\n",
-        "\n",
-        "* Vertex AI\n",
-        "* BigQuery\n",
-        "* Cloud Storage\n",
-        "\n",
-        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing),\n",
-        "[BigQuery pricing](https://cloud.google.com/bigquery/pricing),\n",
-        "and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), \n",
-        "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n",
-        "to generate a cost estimate based on your projected usage."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "i7EUnXsZhAGF"
-      },
-      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook. "
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "2b4ef9b72d43"
-      },
-      "outputs": [],
-      "source": [
-        "# Install the packages\n",
-        "! pip3 install --upgrade --quiet google-cloud-aiplatform[preview]\n",
-        "! pip3 install --upgrade --quiet bigframes\n",
-        "! pip3 install --upgrade --quiet tensorflow==2.12.0"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "58707a750154"
-      },
-      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "f200f10a1da3"
-      },
-      "outputs": [],
-      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
-        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "BF1j6f9HApxa"
-      },
-      "source": [
-        "## Before you begin\n",
-        "\n",
-        "### Set up your Google Cloud project\n",
-        "\n",
-        "**The following steps are required, regardless of your notebook environment.**\n",
-        "\n",
-        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
-        "\n",
-        "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
-        "\n",
-        "3. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
-        "\n",
-        "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "WReHDGG5g0XY"
-      },
-      "source": [
-        "#### Set your project ID\n",
-        "\n",
-        "**If you don't know your project ID**, try the following:\n",
-        "* Run `gcloud config list`.\n",
-        "* Run `gcloud projects list`.\n",
-        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "oM1iC_MfAts1"
-      },
-      "outputs": [],
-      "source": [
-        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# Set the project id\n",
-        "! gcloud config set project {PROJECT_ID}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "region"
-      },
-      "source": [
-        "#### Region\n",
-        "\n",
-        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "region"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "sBCra4QMA2wR"
-      },
-      "source": [
-        "### Authenticate your Google Cloud account\n",
-        "\n",
-        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "74ccc9e52986"
-      },
-      "source": [
-        "**1. Vertex AI Workbench**\n",
-        "* Do nothing as you are already authenticated."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "de775a3773ba"
-      },
-      "source": [
-        "**2. Local JupyterLab instance, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "254614fa0c46"
-      },
-      "outputs": [],
-      "source": [
-        "# ! gcloud auth login"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ef21552ccea8"
-      },
-      "source": [
-        "**3. Colab, uncomment and run:**"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "603adbbf0532"
-      },
-      "outputs": [],
-      "source": [
-        "# from google.colab import auth\n",
-        "# auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "f6b2ccc891ed"
-      },
-      "source": [
-        "**4. Service account or other**\n",
-        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "zgPO1eR3CYjk"
-      },
-      "source": [
-        "### Create a Cloud Storage bucket\n",
-        "\n",
-        "Create a storage bucket to store intermediate artifacts such as datasets."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "MzGDU7TWdts_"
-      },
-      "outputs": [],
-      "source": [
-        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "-EcIXiGsCePi"
-      },
-      "source": [
-        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "NIq7R4HZCfIc"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "960505627ddf"
-      },
-      "source": [
-        "### Import libraries and define constants"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "PyQmSRbKA8r-"
-      },
-      "outputs": [],
-      "source": [
-        "import bigframes.pandas as bf\n",
-        "import tensorflow as tf\n",
-        "import vertexai\n",
-        "from tensorflow import keras\n",
-        "\n",
-        "bf.options.bigquery.location = \"us\"  # Dataset is in 'us' not 'us-central1'\n",
-        "bf.options.bigquery.project = PROJECT_ID\n",
-        "\n",
-        "from bigframes.ml.model_selection import \\\n",
-        "    train_test_split as bf_train_test_split"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "init_aip:mbsdk,all"
-      },
-      "source": [
-        "## Initialize Vertex AI SDK for Python\n",
-        "\n",
-        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "init_aip:mbsdk,all"
-      },
-      "outputs": [],
-      "source": [
-        "vertexai.init(\n",
-        "    project=PROJECT_ID,\n",
-        "    location=REGION,\n",
-        "    staging_bucket=BUCKET_URI,\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "105334524e96"
-      },
-      "source": [
-        "## Prepare the dataset\n",
-        "\n",
-        "Now load the Iris dataset and split the data into train and test sets."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "94576deccd8c"
-      },
-      "outputs": [],
-      "source": [
-        "df = bf.read_gbq(\"bigquery-public-data.ml_datasets.iris\")\n",
-        "\n",
-        "species_categories = {\n",
-        "    \"versicolor\": 0,\n",
-        "    \"virginica\": 1,\n",
-        "    \"setosa\": 2,\n",
-        "}\n",
-        "df[\"species\"] = df[\"species\"].map(species_categories)\n",
-        "\n",
-        "train, test = bf_train_test_split(df, test_size=0.2)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "cfcbce726efa"
-      },
-      "source": [
-        "## Remote training with GPU\n",
-        "\n",
-        "First, train a TensorFlow model as a remote training job:\n",
-        "\n",
-        "- Reinitialize Vertex AI for remote training.\n",
-        "- Instantiate the tensorflow keras model for the remote training job.\n",
-        "- Invoke the tensorflow keras model.fit() locally which will launch the remote training job."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "fd865b0c4e8b"
-      },
-      "outputs": [],
-      "source": [
-        "# Switch to remote mode for training\n",
-        "vertexai.preview.init(remote=True)\n",
-        "\n",
-        "Sequential = vertexai.preview.remote(keras.Sequential)\n",
-        "\n",
-        "# Instantiate model\n",
-        "model = Sequential([keras.layers.Dense(5, input_shape=(4,)), keras.layers.Softmax()])\n",
-        "\n",
-        "# Specify optimizer and loss function\n",
-        "model.compile(optimizer=\"adam\", loss=\"mean_squared_error\")\n",
-        "\n",
-        "# (Optional) Manually set GPU compute resources\n",
-        "model.fit.vertex.remote_config.enable_cuda = True\n",
-        "model.fit.vertex.remote_config.machine_type = \"n1-highmem-4\"\n",
-        "model.fit.vertex.remote_config.accelerator_type = \"NVIDIA_TESLA_K80\"\n",
-        "model.fit.vertex.remote_config.accelerator_count = 4\n",
-        "\n",
-        "# (Optional) Set batch_size, target_col\n",
-        "model.fit.vertex.remote_config.serializer_args[train] = {\n",
-        "    \"batch_size\": 32,\n",
-        "    \"target_col\": \"species\",\n",
-        "}\n",
-        "\n",
-        "# Train model on Vertex\n",
-        "model.fit(train, epochs=10)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "f1af94ac1477"
-      },
-      "source": [
-        "## Remote prediction\n",
-        "\n",
-        "Obtain predictions from the trained model."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "1d75879948b5"
-      },
-      "outputs": [],
-      "source": [
-        "# Remote prediction\n",
-        "vertexai.preview.init(remote=True)\n",
-        "\n",
-        "# Disable GPU for remote prediction\n",
-        "model.predict.vertex.remote_config.enable_cuda = False\n",
-        "\n",
-        "# (Optional) Set batch_size, target_col\n",
-        "model.predict.vertex.remote_config.serializer_args[train] = {\n",
-        "    \"batch_size\": 32,\n",
-        "    \"target_col\": \"species\",\n",
-        "}\n",
-        "\n",
-        "predictions = model.predict(train)\n",
-        "\n",
-        "print(f\"Remote predictions: {predictions}\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "798b77c95067"
-      },
-      "source": [
-        "## Local evaluation\n",
-        "\n",
-        "Evaluate model results locally."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "88e734e30791"
-      },
-      "outputs": [],
-      "source": [
-        "# User must convert bigframes to pandas dataframe for local evaluation\n",
-        "feature_columns = [\"sepal_length\", \"sepal_width\", \"petal_length\", \"petal_width\"]\n",
-        "label_columns = [\"species\"]\n",
-        "\n",
-        "train_X_np = train[feature_columns].to_pandas().values.astype(float)\n",
-        "train_y_np = train[label_columns].to_pandas().values.astype(float)\n",
-        "train_ds = tf.data.Dataset.from_tensor_slices((train_X_np, train_y_np))\n",
-        "\n",
-        "test_X_np = test[feature_columns].to_pandas().values.astype(float)\n",
-        "test_y_np = test[label_columns].to_pandas().values.astype(float)\n",
-        "test_ds = tf.data.Dataset.from_tensor_slices((test_X_np, test_y_np))"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "cb8637f783ad"
-      },
-      "outputs": [],
-      "source": [
-        "# Switch to local mode for evaluation\n",
-        "vertexai.preview.init(remote=False)\n",
-        "\n",
-        "# Evaluate model's mean square errors\n",
-        "print(f\"Train loss: {model.evaluate(train_ds.batch(32))}\")\n",
-        "\n",
-        "print(f\"Test loss: {model.evaluate(test_ds.batch(32))}\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "TpV-iwP9qw9c"
-      },
-      "source": [
-        "## Cleaning up\n",
-        "\n",
-        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
-        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
-        "\n",
-        "Otherwise, you can delete the individual resources you created in this tutorial:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "sx_vKniMq9ZX"
-      },
-      "outputs": [],
-      "source": [
-        "import os\n",
-        "\n",
-        "# Delete Cloud Storage objects that were created\n",
-        "delete_bucket = False\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
-        "    ! gsutil -m rm -r $BUCKET_URI"
-      ]
-    }
-  ],
-  "metadata": {
-    "colab": {
-      "collapsed_sections": [],
-      "name": "remote_training_bigframes_tensorflow.ipynb",
-      "toc_visible": true
-    },
-    "kernelspec": {
-      "display_name": "Python 3",
-      "name": "python3"
-    }
-  },
-  "nbformat": 4,
-  "nbformat_minor": 0
-}
@@ -1,672 +0,0 @@
-{
-  "cells": [
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "cellView": "form",
-        "id": "ur8xi4C7S06n"
-      },
-      "outputs": [],
-      "source": [
-        "# Copyright 2023 Google LLC\n",
-        "#\n",
-        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
-        "# you may not use this file except in compliance with the License.\n",
-        "# You may obtain a copy of the License at\n",
-        "#\n",
-        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
-        "#\n",
-        "# Unless required by applicable law or agreed to in writing, software\n",
-        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
-        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
-        "# See the License for the specific language governing permissions and\n",
-        "# limitations under the License."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "fsv4jGuU89rX"
-      },
-      "source": [
-        "# Vertex AI SDK 2.0 Remote Training for lightning model\n",
-        "\n",
-        "<table align=\"left\">\n",
-        "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vertex_ai_sdk/remote_training_lightning.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
-        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vertex_ai_sdk/remote_training_lightning.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/vertex_ai_sdk/remote_training_lightning.ipynb\">\n",
-        "       <img src=\"https://www.gstatic.com/cloud/images/navigation/vertex-ai.svg\" alt=\"Vertex AI logo\">Open in Vertex AI Workbench\n",
-        "    </a>\n",
-        "</table>"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "overview:automl"
-      },
-      "source": [
-        "## Overview\n",
-        "\n",
-        "This tutorial demonstrates how to use Vertex AI SDK 2.0 for remote model training of a local model training job for OSS ML frameworks."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "objective:pipelines,automl"
-      },
-      "source": [
-        "### Objective\n",
-        "\n",
-        "In this tutorial, you learn to use `Vertex AI SDK 2.0` to remotely training models of various ML frameworks as a local (on-prem) training job.\n",
-        "\n",
-        "This tutorial uses the following Google Cloud ML services:\n",
-        "\n",
-        "- `Vertex AI Training`\n",
-        "- `Vertex AI Remote Training`\n",
-        "\n",
-        "The steps performed include:\n",
-        "\n",
-        "- Download and split the dataset\n",
-        "- Perform transformations as a Vertex AI remote training.\n",
-        "- For scikit-learn, PyTorch, TensorFlow, PyTorch Lightning\n",
-        "    - Train the model remotely.\n",
-        "    - Uptrain the pretrained model remotely.\n",
-        "    - Evaluate both the pretrained and uptrained model.\n",
-        "\n",
-        "**Local-to-remote training**\n",
-        "\n",
-        "```\n",
-        "import vertexai\n",
-        "from my_module import MyModelClass\n",
-        "\n",
-        "vertexai.init(project=\"my-project\", location=\"my-location\", staging_bucket=\"gs://my-bucket\")\n",
-        "\n",
-        "# Switch to remote mode\n",
-        "vertexai.preview.init(remote=True)\n",
-        "\n",
-        "# Wrap the model class with `vertexai.preview.remote`\n",
-        "MyModelClass = vertexai.preview.remote(MyModelClass)\n",
-        "\n",
-        "# Instantiate the class\n",
-        "model = MyModelClass(...)\n",
-        "\n",
-        "# Optional set training config\n",
-        "model.fit.vertex.remote_config.display_name = \"MyModelClass-remote-training\"\n",
-        "model.fit.vertex.remote_config.staging_bucket = \"gs://my-bucket\"\n",
-        "\n",
-        "# This `fit` call will be executed remotely\n",
-        "model.fit(...)\n",
-        "```\n",
-        "\n",
-        "*Remote training supported OSS ML frameworks*\n",
-        "1.  scikit-learn\n",
-        "2.  TensorFlow\n",
-        "3.  PyTorch\n",
-        "4.  Pytorch Lightning\n",
-        "5.  Custom model\n",
-        "\n",
-        "\n",
-        "---\n",
-        "\n",
-        "**Uptraining**\n",
-        "```\n",
-        "...\n",
-        "model = MyModelClass(...)\n",
-        "model.fit(...)\n",
-        "\n",
-        "# Save the trained model to Model Registry\n",
-        "registered_model = vertexai.preview.register(model)\n",
-        "\n",
-        "# The model can be loaded to a new (or current) local runtime\n",
-        "loaded_model = vertexai.preview.from_pretrained(\"registered-model-resource-id\")\n",
-        "\n",
-        "# Loaded model can cuntinue perform local-to-remote training\n",
-        "loaded_model.fit(...)\n",
-        "\n",
-        "```\n",
-        "\n",
-        "*Remote training supported OSS ML frameworks*\n",
-        "1.  scikit-learn\n",
-        "2.  TensorFlow\n",
-        "3.  Custom model\n",
-        "4.  PyTorch\n",
-        "\n",
-        "\n",
-        "\n",
-        "---\n",
-        "\n",
-        "**GPU Training**\n",
-        "```\n",
-        "...\n",
-        "model = MyModelClass(...)\n",
-        "\n",
-        "# Set enable_cuda to True to enable GPU training.\n",
-        "model.fit.vertex.remote_config.enable_cuda = True\n",
-        "\n",
-        "# (Optional) Training image and compute resources will be automatically\n",
-        "# handled by Vertex, but you can also config by yourself.\n",
-        "model.fit.vertex.remote_config.container_uri = \"your-cuda-image\"\n",
-        "model.fit.vertex.remote_config.machine_type = \"a2-highgpu-8g\"\n",
-        "model.fit.vertex.remote_config.accelerator_type = \"NVIDIA_TESLA_A100\"\n",
-        "model.fit.vertex.remote_config.accelerator_count = 8\n",
-        "\n",
-        "# Model will be trained remotely using GPU\n",
-        "model.fit(...)\n",
-        "```\n",
-        "\n",
-        "*GPU remote training supported OSS ML frameworks*\n",
-        "1.  TensorFlow\n",
-        "2.  PyTorch"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "aef4f59195ad"
-      },
-      "source": [
-        "### Dataset\n",
-        "\n",
-        "This tutorial uses the <a href=\"https://scikit-learn.org/stable/auto_examples/datasets/plot_iris_dataset.html\">IRIS dataset</a>, which predicts the iris species."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "costs"
-      },
-      "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\n",
-        "pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage\n",
-        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
-        "Calculator](https://cloud.google.com/products/calculator/)\n",
-        "to generate a cost estimate based on your projected usage."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "install_aip:mbsdk"
-      },
-      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "Fpp5OjGI6BaR"
-      },
-      "outputs": [],
-      "source": [
-        "! pip3 install --upgrade --quiet google-cloud-aiplatform[preview]\n",
-        "! pip3 install --upgrade --quiet scikit-learn\n",
-        "! pip3 install --upgrade --quiet lightning"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "restart"
-      },
-      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "D-ZBOjErv5mM"
-      },
-      "outputs": [],
-      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
-        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "before_you_begin:nogpu"
-      },
-      "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": "2dw8q9fdQEH5"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\""
-      ]
-    },
-    {
-      "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": {
-        "id": "bucket:mbsdk"
-      },
-      "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": "bucket"
-      },
-      "outputs": [],
-      "source": [
-        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "autoset_bucket"
-      },
-      "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": "91c46850b49b"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "setup_vars"
-      },
-      "source": [
-        "### Import libraries and define constants"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "LF2bwT6q-of1"
-      },
-      "outputs": [],
-      "source": [
-        "import vertexai\n",
-        "from sklearn.datasets import load_iris\n",
-        "from sklearn.model_selection import train_test_split\n",
-        "from sklearn.preprocessing import StandardScaler"
-      ]
-    },
-    {
-      "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": "-p0YAOY64cA5"
-      },
-      "outputs": [],
-      "source": [
-        "REMOTE_JOB_NAME = \"remote-scalar\"\n",
-        "REMOTE_JOB_BUCKET = f\"{BUCKET_URI}/{REMOTE_JOB_NAME}\"\n",
-        "\n",
-        "vertexai.init(\n",
-        "    project=PROJECT_ID,\n",
-        "    location=REGION,\n",
-        "    staging_bucket=REMOTE_JOB_BUCKET,\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "Ud6sRD1WWFIr"
-      },
-      "source": [
-        "## Prepare the dataset\n",
-        "\n",
-        "Now load the Iris dataset and split the data into train, retrain and test sets."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "N_tbI76DWRWt"
-      },
-      "outputs": [],
-      "source": [
-        "dataset = load_iris()\n",
-        "\n",
-        "X, X_retrain, y, y_retrain = train_test_split(\n",
-        "    dataset.data, dataset.target, test_size=0.60, random_state=42\n",
-        ")\n",
-        "X_train, X_test, y_train, y_test = train_test_split(\n",
-        "    X, y, test_size=0.20, random_state=42\n",
-        ")\n",
-        "\n",
-        "\n",
-        "transformer = StandardScaler()\n",
-        "X_train = transformer.fit_transform(X_train)\n",
-        "X_test = transformer.transform(X_test)\n",
-        "X_retrain = transformer.transform(X_retrain)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "Z6amYysl-h_x"
-      },
-      "source": [
-        "## PyTorch Lightning"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "kr6WLoCIwXbJ"
-      },
-      "source": [
-        "### Remote training\n",
-        "\n",
-        "First, train a PyTorch Ligtning model as a remote training job:\n",
-        "\n",
-        "- Reinitialize Vertex AI for remote training.\n",
-        "- Set Trainer for the remote training job.\n",
-        "- Invoke Trainer locally which will launch the remote training job."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "H2kWuzcZLdVE"
-      },
-      "outputs": [],
-      "source": [
-        "# Switch to local mode for testing\n",
-        "vertexai.preview.init(remote=True)\n",
-        "\n",
-        "import lightning.pytorch as pl\n",
-        "import torch\n",
-        "\n",
-        "# Wrap classes to enable Vertex remote execution\n",
-        "pl.Trainer = vertexai.preview.remote(pl.Trainer)\n",
-        "\n",
-        "\n",
-        "# prepare data loaders\n",
-        "train_loader = torch.utils.data.DataLoader(\n",
-        "    torch.utils.data.TensorDataset(\n",
-        "        torch.tensor(X_train).to(torch.float32),\n",
-        "        torch.tensor(y_train),\n",
-        "    ),\n",
-        "    batch_size=10,\n",
-        "    shuffle=True,\n",
-        ")\n",
-        "\n",
-        "val_loader = torch.utils.data.DataLoader(\n",
-        "    torch.utils.data.TensorDataset(\n",
-        "        torch.tensor(X_test).to(torch.float32),\n",
-        "        torch.tensor(y_test),\n",
-        "    ),\n",
-        "    batch_size=10,\n",
-        ")\n",
-        "\n",
-        "\n",
-        "# define the model\n",
-        "class LitLogisticRegression(pl.LightningModule):\n",
-        "    def __init__(self, input_size: int, output_size: int):\n",
-        "        super().__init__()\n",
-        "        self.linear = torch.nn.Linear(input_size, output_size)\n",
-        "        self.softmax = torch.nn.Softmax(dim=1)\n",
-        "\n",
-        "    def forward(self, x):\n",
-        "        return self.softmax(self.linear(x))\n",
-        "\n",
-        "    def training_step(self, batch, batch_idx):\n",
-        "        x, y = batch\n",
-        "        y_hat = self(x)\n",
-        "        loss = torch.nn.functional.cross_entropy(y_hat, y)\n",
-        "        return loss\n",
-        "\n",
-        "    def validation_step(self, batch, batch_idx):\n",
-        "        x, y = batch\n",
-        "        y_hat = self(x)\n",
-        "        loss = torch.nn.functional.cross_entropy(y_hat, y)\n",
-        "        return loss\n",
-        "\n",
-        "    def configure_optimizers(self):\n",
-        "        return torch.optim.Adam(self.parameters(), lr=0.05)\n",
-        "\n",
-        "    def predict(self, X):\n",
-        "        X = torch.tensor(X).to(torch.float32)\n",
-        "        with torch.no_grad():\n",
-        "            pred = torch.argmax(self(X), dim=1)\n",
-        "        return pred\n",
-        "\n",
-        "\n",
-        "model = LitLogisticRegression(4, 3)\n",
-        "\n",
-        "# set up the trainer and training config\n",
-        "trainer = pl.Trainer(max_epochs=100, accelerator=\"cpu\")\n",
-        "\n",
-        "# Train model on Vertex\n",
-        "trainer.fit(model, train_dataloaders=train_loader, val_dataloaders=val_loader)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "b9GglFxvoeuw"
-      },
-      "source": [
-        "### Local evaluation\n",
-        "\n",
-        "Next, evaluate the pretrained version of the model, and compare the results."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "tG9sGXYb-nuU"
-      },
-      "outputs": [],
-      "source": [
-        "from sklearn.metrics import accuracy_score\n",
-        "\n",
-        "# Switch to local mode for testing\n",
-        "vertexai.preview.init(remote=False)\n",
-        "\n",
-        "# Evaluate model's accuracy score\n",
-        "print(f\"Train accuracy: {accuracy_score(y_train, model.predict(X_train))}\")\n",
-        "print(f\"Test accuracy: {accuracy_score(y_test, model.predict(X_test))}\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ArvOcC4B6BaY"
-      },
-      "source": [
-        "## Cleaning up\n",
-        "\n",
-        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
-        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
-        "\n",
-        "Otherwise, you can delete the individual resources you created in this tutorial."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "yYnCN1gV6BaZ"
-      },
-      "outputs": [],
-      "source": [
-        "import os\n",
-        "\n",
-        "delete_bucket = False\n",
-        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
-        "    ! gsutil rm -rf {BUCKET_URI}"
-      ]
-    }
-  ],
-  "metadata": {
-    "accelerator": "GPU",
-    "colab": {
-      "collapsed_sections": [
-        "-aNMacb-Ghxs",
-        "6fC-g1B7wGgZ"
-      ],
-      "name": "remote_training_lightning.ipynb",
-      "toc_visible": true
-    },
-    "kernelspec": {
-      "display_name": "Python 3",
-      "name": "python3"
-    }
-  },
-  "nbformat": 4,
-  "nbformat_minor": 0
-}
@@ -1,728 +0,0 @@
-{
-  "cells": [
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "cellView": "form",
-        "id": "ur8xi4C7S06n"
-      },
-      "outputs": [],
-      "source": [
-        "# Copyright 2023 Google LLC\n",
-        "#\n",
-        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
-        "# you may not use this file except in compliance with the License.\n",
-        "# You may obtain a copy of the License at\n",
-        "#\n",
-        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
-        "#\n",
-        "# Unless required by applicable law or agreed to in writing, software\n",
-        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
-        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
-        "# See the License for the specific language governing permissions and\n",
-        "# limitations under the License."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "fsv4jGuU89rX"
-      },
-      "source": [
-        "# Vertex AI SDK 2.0 Remote Training for pytorch model\n",
-        "\n",
-        "<table align=\"left\">\n",
-        "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vertex_ai_sdk/remote_training_pytorch.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
-        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vertex_ai_sdk/remote_training_pytorch.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/vertex_ai_sdk/remote_training_pytorch.ipynb\">\n",
-        "       <img src=\"https://www.gstatic.com/cloud/images/navigation/vertex-ai.svg\" alt=\"Vertex AI logo\">Open in Vertex AI Workbench\n",
-        "    </a>\n",
-        "</table>"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "overview:automl"
-      },
-      "source": [
-        "## Overview\n",
-        "\n",
-        "This tutorial demonstrates how to use Vertex AI SDK 2.0 for remote model training of a local model training job for OSS ML frameworks."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "objective:pipelines,automl"
-      },
-      "source": [
-        "### Objective\n",
-        "\n",
-        "In this tutorial, you learn to use `Vertex AI SDK 2.0` to remotely training models of various ML frameworks as a local (on-prem) training job.\n",
-        "\n",
-        "This tutorial uses the following Google Cloud ML services:\n",
-        "\n",
-        "- `Vertex AI Training`\n",
-        "- `Vertex AI Remote Training`\n",
-        "\n",
-        "The steps performed include:\n",
-        "\n",
-        "- Download and split the dataset\n",
-        "- Perform transformations as a Vertex AI remote training.\n",
-        "- For scikit-learn, PyTorch, TensorFlow, PyTorch Lightning\n",
-        "    - Train the model remotely.\n",
-        "    - Uptrain the pretrained model remotely.\n",
-        "    - Evaluate both the pretrained and uptrained model.\n",
-        "\n",
-        "**Local-to-remote training**\n",
-        "\n",
-        "```\n",
-        "import vertexai\n",
-        "from my_module import MyModelClass\n",
-        "\n",
-        "vertexai.init(project=\"my-project\", location=\"my-location\", staging_bucket=\"gs://my-bucket\")\n",
-        "\n",
-        "# Switch to remote mode\n",
-        "vertexai.preview.init(remote=True)\n",
-        "\n",
-        "# Wrap the model class with `vertexai.preview.remote`\n",
-        "MyModelClass = vertexai.preview.remote(MyModelClass)\n",
-        "\n",
-        "# Instantiate the class\n",
-        "model = MyModelClass(...)\n",
-        "\n",
-        "# Optional set training config\n",
-        "model.fit.vertex.remote_config.display_name = \"MyModelClass-remote-training\"\n",
-        "model.fit.vertex.remote_config.staging_bucket = \"gs://my-bucket\"\n",
-        "\n",
-        "# This `fit` call will be executed remotely\n",
-        "model.fit(...)\n",
-        "```\n",
-        "\n",
-        "*Remote training supported OSS ML frameworks*\n",
-        "1.  scikit-learn\n",
-        "2.  TensorFlow\n",
-        "3.  PyTorch\n",
-        "4.  Pytorch Lightning\n",
-        "5.  Custom model\n",
-        "\n",
-        "\n",
-        "---\n",
-        "\n",
-        "**Uptraining**\n",
-        "```\n",
-        "...\n",
-        "model = MyModelClass(...)\n",
-        "model.fit(...)\n",
-        "\n",
-        "# Save the trained model to Model Registry\n",
-        "registered_model = vertexai.preview.register(model)\n",
-        "\n",
-        "# The model can be loaded to a new (or current) local runtime\n",
-        "loaded_model = vertexai.preview.from_pretrained(\"registered-model-resource-id\")\n",
-        "\n",
-        "# Loaded model can cuntinue perform local-to-remote training\n",
-        "loaded_model.fit(...)\n",
-        "\n",
-        "```\n",
-        "\n",
-        "*Remote training supported OSS ML frameworks*\n",
-        "1.  scikit-learn\n",
-        "2.  TensorFlow\n",
-        "3.  Custom model\n",
-        "4.  PyTorch\n",
-        "\n",
-        "\n",
-        "\n",
-        "---\n",
-        "\n",
-        "**GPU Training**\n",
-        "```\n",
-        "...\n",
-        "model = MyModelClass(...)\n",
-        "\n",
-        "# Set enable_cuda to True to enable GPU training.\n",
-        "model.fit.vertex.remote_config.enable_cuda = True\n",
-        "\n",
-        "# (Optional) Training image and compute resources will be automatically\n",
-        "# handled by Vertex, but you can also config by yourself.\n",
-        "model.fit.vertex.remote_config.container_uri = \"your-cuda-image\"\n",
-        "model.fit.vertex.remote_config.machine_type = \"a2-highgpu-8g\"\n",
-        "model.fit.vertex.remote_config.accelerator_type = \"NVIDIA_TESLA_A100\"\n",
-        "model.fit.vertex.remote_config.accelerator_count = 8\n",
-        "\n",
-        "# Model will be trained remotely using GPU\n",
-        "model.fit(...)\n",
-        "```\n",
-        "\n",
-        "*GPU remote training supported OSS ML frameworks*\n",
-        "1.  TensorFlow\n",
-        "2.  PyTorch"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "aef4f59195ad"
-      },
-      "source": [
-        "### Dataset\n",
-        "\n",
-        "This tutorial uses the <a href=\"https://scikit-learn.org/stable/auto_examples/datasets/plot_iris_dataset.html\">IRIS dataset</a>, which predicts the iris species."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "costs"
-      },
-      "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\n",
-        "pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage\n",
-        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
-        "Calculator](https://cloud.google.com/products/calculator/)\n",
-        "to generate a cost estimate based on your projected usage."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "install_aip:mbsdk"
-      },
-      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "Fpp5OjGI6BaR"
-      },
-      "outputs": [],
-      "source": [
-        "! pip3 install --upgrade --quiet google-cloud-aiplatform[preview]\n",
-        "! pip3 install --upgrade --quiet scikit-learn\n",
-        "! pip3 install --upgrade --quiet torch"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "restart"
-      },
-      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "D-ZBOjErv5mM"
-      },
-      "outputs": [],
-      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
-        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "before_you_begin:nogpu"
-      },
-      "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": "2dw8q9fdQEH5"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\""
-      ]
-    },
-    {
-      "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": {
-        "id": "bucket:mbsdk"
-      },
-      "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": "bucket"
-      },
-      "outputs": [],
-      "source": [
-        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "autoset_bucket"
-      },
-      "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": "91c46850b49b"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "setup_vars"
-      },
-      "source": [
-        "### Import libraries and define constants"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "LF2bwT6q-of1"
-      },
-      "outputs": [],
-      "source": [
-        "import vertexai\n",
-        "from sklearn.datasets import load_iris\n",
-        "from sklearn.model_selection import train_test_split\n",
-        "from sklearn.preprocessing import StandardScaler"
-      ]
-    },
-    {
-      "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": "-p0YAOY64cA5"
-      },
-      "outputs": [],
-      "source": [
-        "REMOTE_JOB_NAME = \"remote-scalar\"\n",
-        "REMOTE_JOB_BUCKET = f\"{BUCKET_URI}/{REMOTE_JOB_NAME}\"\n",
-        "\n",
-        "vertexai.init(\n",
-        "    project=PROJECT_ID,\n",
-        "    location=REGION,\n",
-        "    staging_bucket=REMOTE_JOB_BUCKET,\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "Ud6sRD1WWFIr"
-      },
-      "source": [
-        "## Prepare the dataset\n",
-        "\n",
-        "Now load the Iris dataset and split the data into train, retrain and test sets."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "N_tbI76DWRWt"
-      },
-      "outputs": [],
-      "source": [
-        "dataset = load_iris()\n",
-        "\n",
-        "X, X_retrain, y, y_retrain = train_test_split(\n",
-        "    dataset.data, dataset.target, test_size=0.60, random_state=42\n",
-        ")\n",
-        "X_train, X_test, y_train, y_test = train_test_split(\n",
-        "    X, y, test_size=0.20, random_state=42\n",
-        ")\n",
-        "\n",
-        "\n",
-        "transformer = StandardScaler()\n",
-        "X_train = transformer.fit_transform(X_train)\n",
-        "X_test = transformer.transform(X_test)\n",
-        "X_retrain = transformer.transform(X_retrain)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "kdg47y-wjoxT"
-      },
-      "source": [
-        "## PyTorch"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "eTSejqbdv2RE"
-      },
-      "source": [
-        "### Remote training with CPU\n",
-        "\n",
-        "First, train a PyTorch model as a remote training job:\n",
-        "\n",
-        "- Reinitialize Vertex AI for remote training.\n",
-        "- Set TorchLogisticRegression for the remote training job.\n",
-        "- Invoke TorchLogisticRegression locally which will launch the remote training job."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "wtk_KqDX-5WI"
-      },
-      "outputs": [],
-      "source": [
-        "# Switch to remote mode for training\n",
-        "vertexai.preview.init(remote=True)\n",
-        "\n",
-        "import torch\n",
-        "from vertexai.preview import VertexModel\n",
-        "\n",
-        "\n",
-        "# define the custom model\n",
-        "class TorchLogisticRegression(VertexModel, torch.nn.Module):\n",
-        "    def __init__(self, input_size: int, output_size: int):\n",
-        "        torch.nn.Module.__init__(self)\n",
-        "        VertexModel.__init__(self)\n",
-        "        self.linear = torch.nn.Linear(input_size, output_size)\n",
-        "        self.softmax = torch.nn.Softmax(dim=1)\n",
-        "\n",
-        "    def forward(self, x):\n",
-        "        return self.softmax(self.linear(x))\n",
-        "\n",
-        "    @vertexai.preview.developer.mark.train()\n",
-        "    def train(self, X, y, num_epochs, lr):\n",
-        "        X, y = torch.tensor(X).to(torch.float32), torch.tensor(y)\n",
-        "        dataloader = torch.utils.data.DataLoader(\n",
-        "            torch.utils.data.TensorDataset(X, y),\n",
-        "            batch_size=10,\n",
-        "            shuffle=True,\n",
-        "            generator=torch.Generator(device=X.device),\n",
-        "        )\n",
-        "\n",
-        "        criterion = torch.nn.CrossEntropyLoss()\n",
-        "        optimizer = torch.optim.SGD(self.parameters(), lr=lr)\n",
-        "\n",
-        "        for t in range(num_epochs):\n",
-        "            for batch, (X, y) in enumerate(dataloader):\n",
-        "                optimizer.zero_grad()\n",
-        "                pred = self(X)\n",
-        "                loss = criterion(pred, y)\n",
-        "                loss.backward()\n",
-        "                optimizer.step()\n",
-        "\n",
-        "    @vertexai.preview.developer.mark.predict()\n",
-        "    def predict(self, X):\n",
-        "        X = torch.tensor(X).to(torch.float32)\n",
-        "        with torch.no_grad():\n",
-        "            pred = torch.argmax(self(X), dim=1)\n",
-        "        return pred\n",
-        "\n",
-        "\n",
-        "# Instantiate model\n",
-        "model = TorchLogisticRegression(4, 3)\n",
-        "\n",
-        "# Train model on Vertex\n",
-        "model.train(X_train, y_train, num_epochs=100, lr=0.05)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "qQ53E9NK-lYB"
-      },
-      "source": [
-        "### Uptrain the pretrained model with GPU\n",
-        "\n",
-        "Next, get the registered model from the Vertex AI Model Registry. Then request the pretrained version of the model."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "1dV4pWtbRdkq"
-      },
-      "outputs": [],
-      "source": [
-        "registered_model = vertexai.preview.register(model)\n",
-        "\n",
-        "pulled_model = vertexai.preview.from_pretrained(\n",
-        "    model_name=registered_model.resource_name\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "rYuE40b36BaW"
-      },
-      "source": [
-        "Now train the model remotely via Vertex AI Training."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "pbjDOXzk9AkG"
-      },
-      "outputs": [],
-      "source": [
-        "pulled_model.train(X_retrain, y_retrain, num_epochs=100, lr=0.05)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "l35WfW-dkERU"
-      },
-      "source": [
-        "### Local evaluation\n",
-        "\n",
-        "Next, evaluate the pretrained and uptrained versions of the model, and compare the results."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "GFo8usi_kB-F"
-      },
-      "outputs": [],
-      "source": [
-        "from sklearn.metrics import accuracy_score\n",
-        "\n",
-        "# Switch to local mode for testing\n",
-        "vertexai.preview.init(remote=False)\n",
-        "\n",
-        "# Evaluate model's accuracy score\n",
-        "print(f\"Train accuracy: {accuracy_score(y_train, model.predict(X_train))}\")\n",
-        "print(f\"Test accuracy: {accuracy_score(y_test, model.predict(X_test))}\")\n",
-        "\n",
-        "# Evaluate uptrained model's accuracy score\n",
-        "print(f\"Train accuracy: {accuracy_score(y_train, pulled_model.predict(X_train))}\")\n",
-        "print(f\"Test accuracy: {accuracy_score(y_test, pulled_model.predict(X_test))}\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "zZ6DLHxo6BaX"
-      },
-      "source": [
-        "#### Delete the registered model\n",
-        "\n",
-        "You can delete the registered model in the Vertex AI Model Registry with the delete() method."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "esqEbYAA6BaX"
-      },
-      "outputs": [],
-      "source": [
-        "registered_model.delete()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ArvOcC4B6BaY"
-      },
-      "source": [
-        "## Cleaning up\n",
-        "\n",
-        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
-        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
-        "\n",
-        "Otherwise, you can delete the individual resources you created in this tutorial."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "yYnCN1gV6BaZ"
-      },
-      "outputs": [],
-      "source": [
-        "import os\n",
-        "\n",
-        "delete_bucket = False\n",
-        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
-        "    ! gsutil rm -rf {BUCKET_URI}"
-      ]
-    }
-  ],
-  "metadata": {
-    "accelerator": "GPU",
-    "colab": {
-      "collapsed_sections": [
-        "-aNMacb-Ghxs",
-        "6fC-g1B7wGgZ"
-      ],
-      "name": "remote_training_pytorch.ipynb",
-      "toc_visible": true
-    },
-    "kernelspec": {
-      "display_name": "Python 3",
-      "name": "python3"
-    }
-  },
-  "nbformat": 4,
-  "nbformat_minor": 0
-}
@@ -1,787 +0,0 @@
-{
-  "cells": [
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "cellView": "form",
-        "id": "ur8xi4C7S06n"
-      },
-      "outputs": [],
-      "source": [
-        "# Copyright 2023 Google LLC\n",
-        "#\n",
-        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
-        "# you may not use this file except in compliance with the License.\n",
-        "# You may obtain a copy of the License at\n",
-        "#\n",
-        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
-        "#\n",
-        "# Unless required by applicable law or agreed to in writing, software\n",
-        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
-        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
-        "# See the License for the specific language governing permissions and\n",
-        "# limitations under the License."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "fsv4jGuU89rX"
-      },
-      "source": [
-        "# Vertex AI SDK 2.0 Remote Training for sklearn model\n",
-        "\n",
-        "<table align=\"left\">\n",
-        "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vertex_ai_sdk/remote_training_sklearn.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
-        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vertex_ai_sdk/remote_training_sklearn.ipynb\">\n",
-        "        <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
-        "    </a>\n",
-        "  </td>\n",
-        "    <td>\n",
-        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/vertex_ai_sdk/remote_training_sklearn.ipynb\">\n",
-        "       <img src=\"https://www.gstatic.com/cloud/images/navigation/vertex-ai.svg\" alt=\"Vertex AI logo\">Open in Vertex AI Workbench\n",
-        "    </a>\n",
-        "</table>"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "overview:automl"
-      },
-      "source": [
-        "## Overview\n",
-        "\n",
-        "This tutorial demonstrates how to use Vertex AI SDK 2.0 for remote model training of a local model training job for OSS ML frameworks."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "objective:pipelines,automl"
-      },
-      "source": [
-        "### Objective\n",
-        "\n",
-        "In this tutorial, you learn to use `Vertex AI SDK 2.0` to remotely training models of various ML frameworks as a local (on-prem) training job.\n",
-        "\n",
-        "This tutorial uses the following Google Cloud ML services:\n",
-        "\n",
-        "- `Vertex AI Training`\n",
-        "- `Vertex AI Remote Training`\n",
-        "\n",
-        "The steps performed include:\n",
-        "\n",
-        "- Download and split the dataset\n",
-        "- Perform transformations as a Vertex AI remote training.\n",
-        "- For scikit-learn, PyTorch, TensorFlow, PyTorch Lightning\n",
-        "    - Train the model remotely.\n",
-        "    - Uptrain the pretrained model remotely.\n",
-        "    - Evaluate both the pretrained and uptrained model.\n",
-        "\n",
-        "**Local-to-remote training**\n",
-        "\n",
-        "```\n",
-        "import vertexai\n",
-        "from my_module import MyModelClass\n",
-        "\n",
-        "vertexai.init(project=\"my-project\", location=\"my-location\", staging_bucket=\"gs://my-bucket\")\n",
-        "\n",
-        "# Switch to remote mode\n",
-        "vertexai.preview.init(remote=True)\n",
-        "\n",
-        "# Wrap the model class with `vertexai.preview.remote`\n",
-        "MyModelClass = vertexai.preview.remote(MyModelClass)\n",
-        "\n",
-        "# Instantiate the class\n",
-        "model = MyModelClass(...)\n",
-        "\n",
-        "# Optional set training config\n",
-        "model.fit.vertex.remote_config.display_name = \"MyModelClass-remote-training\"\n",
-        "model.fit.vertex.remote_config.staging_bucket = \"gs://my-bucket\"\n",
-        "\n",
-        "# This `fit` call will be executed remotely\n",
-        "model.fit(...)\n",
-        "```\n",
-        "\n",
-        "*Remote training supported OSS ML frameworks*\n",
-        "1.  scikit-learn\n",
-        "2.  TensorFlow\n",
-        "3.  PyTorch\n",
-        "4.  Pytorch Lightning\n",
-        "5.  Custom model\n",
-        "\n",
-        "\n",
-        "---\n",
-        "\n",
-        "**Uptraining**\n",
-        "```\n",
-        "...\n",
-        "model = MyModelClass(...)\n",
-        "model.fit(...)\n",
-        "\n",
-        "# Save the trained model to Model Registry\n",
-        "registered_model = vertexai.preview.register(model)\n",
-        "\n",
-        "# The model can be loaded to a new (or current) local runtime\n",
-        "loaded_model = vertexai.preview.from_pretrained(\"registered-model-resource-id\")\n",
-        "\n",
-        "# Loaded model can cuntinue perform local-to-remote training\n",
-        "loaded_model.fit(...)\n",
-        "\n",
-        "```\n",
-        "\n",
-        "*Remote training supported OSS ML frameworks*\n",
-        "1.  scikit-learn\n",
-        "2.  TensorFlow\n",
-        "3.  Custom model\n",
-        "4.  PyTorch\n",
-        "\n",
-        "\n",
-        "\n",
-        "---\n",
-        "\n",
-        "**GPU Training**\n",
-        "```\n",
-        "...\n",
-        "model = MyModelClass(...)\n",
-        "\n",
-        "# Set enable_cuda to True to enable GPU training.\n",
-        "model.fit.vertex.remote_config.enable_cuda = True\n",
-        "\n",
-        "# (Optional) Training image and compute resources will be automatically\n",
-        "# handled by Vertex, but you can also config by yourself.\n",
-        "model.fit.vertex.remote_config.container_uri = \"your-cuda-image\"\n",
-        "model.fit.vertex.remote_config.machine_type = \"a2-highgpu-8g\"\n",
-        "model.fit.vertex.remote_config.accelerator_type = \"NVIDIA_TESLA_A100\"\n",
-        "model.fit.vertex.remote_config.accelerator_count = 8\n",
-        "\n",
-        "# Model will be trained remotely using GPU\n",
-        "model.fit(...)\n",
-        "```\n",
-        "\n",
-        "*GPU remote training supported OSS ML frameworks*\n",
-        "1.  TensorFlow\n",
-        "2.  PyTorch"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "aef4f59195ad"
-      },
-      "source": [
-        "### Dataset\n",
-        "\n",
-        "This tutorial uses the <a href=\"https://scikit-learn.org/stable/auto_examples/datasets/plot_iris_dataset.html\">IRIS dataset</a>, which predicts the iris species."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "costs"
-      },
-      "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\n",
-        "pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage\n",
-        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
-        "Calculator](https://cloud.google.com/products/calculator/)\n",
-        "to generate a cost estimate based on your projected usage."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "install_aip:mbsdk"
-      },
-      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "Fpp5OjGI6BaR"
-      },
-      "outputs": [],
-      "source": [
-        "! pip3 install --upgrade --quiet google-cloud-aiplatform[preview]\n",
-        "! pip3 install --upgrade --quiet scikit-learn"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "restart"
-      },
-      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "D-ZBOjErv5mM"
-      },
-      "outputs": [],
-      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
-        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "before_you_begin:nogpu"
-      },
-      "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": "2dw8q9fdQEH5"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\""
-      ]
-    },
-    {
-      "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": {
-        "id": "bucket:mbsdk"
-      },
-      "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": "bucket"
-      },
-      "outputs": [],
-      "source": [
-        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "autoset_bucket"
-      },
-      "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": "91c46850b49b"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "setup_vars"
-      },
-      "source": [
-        "### Import libraries and define constants"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "LF2bwT6q-of1"
-      },
-      "outputs": [],
-      "source": [
-        "import vertexai\n",
-        "from sklearn.datasets import load_iris\n",
-        "from sklearn.linear_model import LogisticRegression\n",
-        "from sklearn.model_selection import train_test_split\n",
-        "from sklearn.preprocessing import StandardScaler"
-      ]
-    },
-    {
-      "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": "-p0YAOY64cA5"
-      },
-      "outputs": [],
-      "source": [
-        "REMOTE_JOB_NAME = \"remote-scalar\"\n",
-        "REMOTE_JOB_BUCKET = f\"{BUCKET_URI}/{REMOTE_JOB_NAME}\"\n",
-        "\n",
-        "vertexai.init(\n",
-        "    project=PROJECT_ID,\n",
-        "    location=REGION,\n",
-        "    staging_bucket=REMOTE_JOB_BUCKET,\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "Ud6sRD1WWFIr"
-      },
-      "source": [
-        "## Prepare the dataset\n",
-        "\n",
-        "Now load the Iris dataset and split the data into train, retrain and test sets."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "N_tbI76DWRWt"
-      },
-      "outputs": [],
-      "source": [
-        "dataset = load_iris()\n",
-        "\n",
-        "X, X_retrain, y, y_retrain = train_test_split(\n",
-        "    dataset.data, dataset.target, test_size=0.60, random_state=42\n",
-        ")\n",
-        "X_train, X_test, y_train, y_test = train_test_split(\n",
-        "    X, y, test_size=0.20, random_state=42\n",
-        ")\n",
-        "\n",
-        "print(\"Data size: \", len(dataset.target))\n",
-        "print(\"X_train size: \", len(X_train))\n",
-        "print(\"X_retrain size: \", len(X_retrain))\n",
-        "print(\"X_test size: \", len(X_test))"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "XzQfxEwuY22S"
-      },
-      "source": [
-        "## Feature transformation\n",
-        "\n",
-        "Next, you do feature transformations on the data using the Vertex AI remote training service.\n",
-        "\n",
-        "First, you re-initialize Vertex AI to enable remote training."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "mSR1x9C66BaU"
-      },
-      "outputs": [],
-      "source": [
-        "# Switch to remote mode for training\n",
-        "vertexai.preview.init(remote=True)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "3Kf6XbLqw1DY"
-      },
-      "source": [
-        "### Execute remote job for fit_transform() on training data\n",
-        "\n",
-        "Next, indicate that the `StandardScalar` class is to be executed remotely. Then set up the data transform and call the `fit_transform()` method is executed remotely."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "vGIuXQMBY8LO"
-      },
-      "outputs": [],
-      "source": [
-        "# Wrap classes to enable Vertex remote execution\n",
-        "StandardScaler = vertexai.preview.remote(StandardScaler)\n",
-        "\n",
-        "\n",
-        "# Instantiate transformer\n",
-        "transformer = StandardScaler()\n",
-        "\n",
-        "\n",
-        "# Execute transformer on Vertex\n",
-        "X_train = transformer.fit_transform(X_train)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "N15AC1Niw-Rf"
-      },
-      "source": [
-        "### Remote transform on test data"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "v3tBFU3IaZ-l"
-      },
-      "outputs": [],
-      "source": [
-        "# Transform test dataset before calculate test score\n",
-        "X_test = transformer.transform(X_test)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "wRmP4x0bxCdD"
-      },
-      "source": [
-        "### Local transform on retrain data"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "C4wzmSxJA5py"
-      },
-      "outputs": [],
-      "source": [
-        "# Switch to local transformation\n",
-        "vertexai.preview.init(remote=False)\n",
-        "\n",
-        "X_retrain = transformer.transform(X_retrain)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "-aNMacb-Ghxs"
-      },
-      "source": [
-        "## scikit-learn"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "G1uGe_lIviwf"
-      },
-      "source": [
-        "### Remote training\n",
-        "\n",
-        "First, train the scikit-learn model as a remote training job:\n",
-        "\n",
-        "- Reinitialize Vertex AI for remote training.\n",
-        "- Set LogisticRegression for the remote training job.\n",
-        "- Invoke LogisticRegression locally which will launch the remote training job."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "CNnzgrUKVQDP"
-      },
-      "outputs": [],
-      "source": [
-        "# Switch to remote mode for training\n",
-        "vertexai.preview.init(remote=True)\n",
-        "\n",
-        "# Wrap classes to enable Vertex remote execution\n",
-        "LogisticRegression = vertexai.preview.remote(LogisticRegression)\n",
-        "\n",
-        "# Instantiate model, warm_start=True for uptraining\n",
-        "model = LogisticRegression(warm_start=True)\n",
-        "\n",
-        "# (Optional) Customize the display name of the remote job\n",
-        "model.fit.vertex.remote_config.display_name = REMOTE_JOB_NAME + \"-sklearn-model\"\n",
-        "\n",
-        "# Train model on Vertex\n",
-        "model = model.fit(X_train, y_train)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "6mPjqC2Q_wsH"
-      },
-      "source": [
-        "### Uptrain the pretrained model\n",
-        "\n",
-        "Next, get the registered model from the Vertex AI Model Registry. Then request the pretrained version of the model."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "KRHxOjkiHYpd"
-      },
-      "outputs": [],
-      "source": [
-        "registered_model = vertexai.preview.register(model)\n",
-        "\n",
-        "pulled_model = vertexai.preview.from_pretrained(\n",
-        "    model_name=registered_model.resource_name\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "cjUjIyDC6BaV"
-      },
-      "source": [
-        "Now train the model remotely via Vertex AI Training."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ubIl_mVayYk6"
-      },
-      "outputs": [],
-      "source": [
-        "pulled_model.fit(X_retrain, y_retrain)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "BXZ7A-ViN2AL"
-      },
-      "source": [
-        "### Local evaluation\n",
-        "\n",
-        "Next, evaluate the pretrained and uptrained versions of the model, and compare the results."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ywlaPlzVzjWv"
-      },
-      "outputs": [],
-      "source": [
-        "# Switch to local mode for testing\n",
-        "vertexai.preview.init(remote=False)\n",
-        "\n",
-        "# Evaluate model's accuracy score\n",
-        "print(f\"Train accuracy: {model.score(X_train, y_train)}\")\n",
-        "print(f\"Test accuracy: {model.score(X_test, y_test)}\")\n",
-        "\n",
-        "# Evaluate uptrained model's accuracy score\n",
-        "print(f\"Train accuracy: {pulled_model.score(X_train, y_train)}\")\n",
-        "print(f\"Test accuracy: {pulled_model.score(X_test, y_test)}\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "NVi2j-cd6BaW"
-      },
-      "source": [
-        "#### Delete the registered model\n",
-        "\n",
-        "You can delete the registered model in the Vertex AI Model Registry with the delete() method."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "NH00-l-06BaW"
-      },
-      "outputs": [],
-      "source": [
-        "registered_model.delete()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ArvOcC4B6BaY"
-      },
-      "source": [
-        "## Cleaning up\n",
-        "\n",
-        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
-        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
-        "\n",
-        "Otherwise, you can delete the individual resources you created in this tutorial."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "yYnCN1gV6BaZ"
-      },
-      "outputs": [],
-      "source": [
-        "import os\n",
-        "\n",
-        "delete_bucket = False\n",
-        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
-        "    ! gsutil rm -rf {BUCKET_URI}"
-      ]
-    }
-  ],
-  "metadata": {
-    "accelerator": "GPU",
-    "colab": {
-      "collapsed_sections": [
-        "-aNMacb-Ghxs",
-        "6fC-g1B7wGgZ"
-      ],
-      "name": "remote_training_sklearn.ipynb",
-      "toc_visible": true
-    },
-    "kernelspec": {
-      "display_name": "Python 3",
-      "name": "python3"
-    }
-  },
-  "nbformat": 4,
-  "nbformat_minor": 0
-}
@@ -1,749 +0,0 @@
-{
-  "cells": [
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "cellView": "form",
-        "id": "ur8xi4C7S06n"
-      },
-      "outputs": [],
-      "source": [
-        "# Copyright 2023 Google LLC\n",
-        "#\n",
-        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
-        "# you may not use this file except in compliance with the License.\n",
-        "# You may obtain a copy of the License at\n",
-        "#\n",
-        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
-        "#\n",
-        "# Unless required by applicable law or agreed to in writing, software\n",
-        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
-        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
-        "# See the License for the specific language governing permissions and\n",
-        "# limitations under the License."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "fsv4jGuU89rX"
-      },
-      "source": [
-        "# Vertex AI SDK 2.0 Remote Training for keras model\n",
-        "\n",
-        "<table align=\"left\">\n",
-        "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vertex_ai_sdk/remote_training_tensorflow_with_autologging.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
-        "    </a>\n",
-        "  </td>\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vertex_ai_sdk/remote_training_tensorflow_with_autologging.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/vertex_ai_sdk/remote_training_tensorflow_with_autologging.ipynb\">\n",
-        "       <img src=\"https://www.gstatic.com/cloud/images/navigation/vertex-ai.svg\" alt=\"Vertex AI logo\">Open in Vertex AI Workbench\n",
-        "    </a>\n",
-        "</table>"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "overview:automl"
-      },
-      "source": [
-        "## Overview\n",
-        "\n",
-        "This tutorial demonstrates how to use Vertex AI SDK 2.0 for remote model training of a local model training job for OSS ML frameworks."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "objective:pipelines,automl"
-      },
-      "source": [
-        "### Objective\n",
-        "\n",
-        "In this tutorial, you learn to use `Vertex AI SDK 2.0` to remotely training models of various ML frameworks as a local (on-prem) training job.\n",
-        "\n",
-        "This tutorial uses the following Google Cloud ML services:\n",
-        "\n",
-        "- `Vertex AI Training`\n",
-        "- `Vertex AI Remote Training`\n",
-        "\n",
-        "The steps performed include:\n",
-        "\n",
-        "- Download and split the dataset\n",
-        "- Perform transformations as a Vertex AI remote training.\n",
-        "- For scikit-learn, PyTorch, TensorFlow, PyTorch Lightning\n",
-        "    - Train the model remotely.\n",
-        "    - Uptrain the pretrained model remotely.\n",
-        "    - Evaluate both the pretrained and uptrained model.\n",
-        "\n",
-        "**Local-to-remote training**\n",
-        "\n",
-        "```\n",
-        "import vertexai\n",
-        "from my_module import MyModelClass\n",
-        "\n",
-        "vertexai.init(project=\"my-project\", location=\"my-location\", staging_bucket=\"gs://my-bucket\")\n",
-        "\n",
-        "# Switch to remote mode\n",
-        "vertexai.preview.init(remote=True)\n",
-        "\n",
-        "# Wrap the model class with `vertexai.preview.remote`\n",
-        "MyModelClass = vertexai.preview.remote(MyModelClass)\n",
-        "\n",
-        "# Instantiate the class\n",
-        "model = MyModelClass(...)\n",
-        "\n",
-        "# Optional set training config\n",
-        "model.fit.vertex.remote_config.display_name = \"MyModelClass-remote-training\"\n",
-        "model.fit.vertex.remote_config.staging_bucket = \"gs://my-bucket\"\n",
-        "\n",
-        "# This `fit` call will be executed remotely\n",
-        "model.fit(...)\n",
-        "```\n",
-        "\n",
-        "*Remote training supported OSS ML frameworks*\n",
-        "1.  scikit-learn\n",
-        "2.  TensorFlow\n",
-        "3.  PyTorch\n",
-        "4.  Pytorch Lightning\n",
-        "5.  Custom model\n",
-        "\n",
-        "\n",
-        "---\n",
-        "\n",
-        "**Uptraining**\n",
-        "```\n",
-        "...\n",
-        "model = MyModelClass(...)\n",
-        "model.fit(...)\n",
-        "\n",
-        "# Save the trained model to Model Registry\n",
-        "registered_model = vertexai.preview.register(model)\n",
-        "\n",
-        "# The model can be loaded to a new (or current) local runtime\n",
-        "loaded_model = vertexai.preview.from_pretrained(\"registered-model-resource-id\")\n",
-        "\n",
-        "# Loaded model can cuntinue perform local-to-remote training\n",
-        "loaded_model.fit(...)\n",
-        "\n",
-        "```\n",
-        "\n",
-        "*Remote training supported OSS ML frameworks*\n",
-        "1.  scikit-learn\n",
-        "2.  TensorFlow\n",
-        "3.  Custom model\n",
-        "4.  PyTorch\n",
-        "\n",
-        "\n",
-        "\n",
-        "---\n",
-        "\n",
-        "**GPU Training**\n",
-        "```\n",
-        "...\n",
-        "model = MyModelClass(...)\n",
-        "\n",
-        "# Set enable_cuda to True to enable GPU training.\n",
-        "model.fit.vertex.remote_config.enable_cuda = True\n",
-        "\n",
-        "# (Optional) Training image and compute resources will be automatically\n",
-        "# handled by Vertex, but you can also config by yourself.\n",
-        "model.fit.vertex.remote_config.container_uri = \"your-cuda-image\"\n",
-        "model.fit.vertex.remote_config.machine_type = \"a2-highgpu-8g\"\n",
-        "model.fit.vertex.remote_config.accelerator_type = \"NVIDIA_TESLA_A100\"\n",
-        "model.fit.vertex.remote_config.accelerator_count = 8\n",
-        "\n",
-        "# Model will be trained remotely using GPU\n",
-        "model.fit(...)\n",
-        "```\n",
-        "\n",
-        "*GPU remote training supported OSS ML frameworks*\n",
-        "1.  TensorFlow\n",
-        "2.  PyTorch"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "aef4f59195ad"
-      },
-      "source": [
-        "### Dataset\n",
-        "\n",
-        "This tutorial uses the <a href=\"https://scikit-learn.org/stable/auto_examples/datasets/plot_iris_dataset.html\">IRIS dataset</a>, which predicts the iris species."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "costs"
-      },
-      "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\n",
-        "pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage\n",
-        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
-        "Calculator](https://cloud.google.com/products/calculator/)\n",
-        "to generate a cost estimate based on your projected usage."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "install_aip:mbsdk"
-      },
-      "source": [
-        "## Installation\n",
-        "\n",
-        "Install the following packages required to execute this notebook."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "Fpp5OjGI6BaR"
-      },
-      "outputs": [],
-      "source": [
-        "! pip3 install --upgrade --quiet google-cloud-aiplatform[preview,autologging]\n",
-        "! pip3 install --upgrade --quiet scikit-learn\n",
-        "! pip3 install --upgrade --quiet tensorflow==2.12"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "restart"
-      },
-      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "D-ZBOjErv5mM"
-      },
-      "outputs": [],
-      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
-        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "before_you_begin:nogpu"
-      },
-      "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": "2dw8q9fdQEH5"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\""
-      ]
-    },
-    {
-      "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": {
-        "id": "bucket:mbsdk"
-      },
-      "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": "bucket"
-      },
-      "outputs": [],
-      "source": [
-        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "autoset_bucket"
-      },
-      "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": "91c46850b49b"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "setup_vars"
-      },
-      "source": [
-        "### Import libraries and define constants"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "LF2bwT6q-of1"
-      },
-      "outputs": [],
-      "source": [
-        "import vertexai\n",
-        "from sklearn.datasets import load_iris\n",
-        "from sklearn.model_selection import train_test_split\n",
-        "from sklearn.preprocessing import StandardScaler"
-      ]
-    },
-    {
-      "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": "-p0YAOY64cA5"
-      },
-      "outputs": [],
-      "source": [
-        "REMOTE_JOB_NAME = \"remote-scalar\"\n",
-        "REMOTE_JOB_BUCKET = f\"{BUCKET_URI}/{REMOTE_JOB_NAME}\"\n",
-        "\n",
-        "vertexai.init(\n",
-        "    project=PROJECT_ID,\n",
-        "    location=REGION,\n",
-        "    staging_bucket=REMOTE_JOB_BUCKET,\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "Ud6sRD1WWFIr"
-      },
-      "source": [
-        "## Prepare the dataset\n",
-        "\n",
-        "Now load the Iris dataset turn it into tf dataset."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "N_tbI76DWRWt"
-      },
-      "outputs": [],
-      "source": [
-        "import tensorflow as tf\n",
-        "\n",
-        "dataset = load_iris()\n",
-        "\n",
-        "X, X_retrain, y, y_retrain = train_test_split(\n",
-        "    dataset.data, dataset.target, test_size=0.60, random_state=42\n",
-        ")\n",
-        "X_train, X_test, y_train, y_test = train_test_split(\n",
-        "    X, y, test_size=0.20, random_state=42\n",
-        ")\n",
-        "\n",
-        "\n",
-        "transformer = StandardScaler()\n",
-        "X_train = transformer.fit_transform(X_train)\n",
-        "X_test = transformer.transform(X_test)\n",
-        "X_retrain = transformer.transform(X_retrain)\n",
-        "\n",
-        "\n",
-        "tf_train_dataset = tf.data.Dataset.from_tensor_slices((X_train, y_train))\n",
-        "tf_train_dataset = tf_train_dataset.shuffle(buffer_size=64).batch(32)\n",
-        "\n",
-        "tf_retrain_dataset = tf.data.Dataset.from_tensor_slices((X_retrain, y_retrain))\n",
-        "tf_retrain_dataset = tf_retrain_dataset.shuffle(buffer_size=64).batch(32)\n",
-        "\n",
-        "tf_test_dataset = tf.data.Dataset.from_tensor_slices((X_test, y_test))\n",
-        "tf_test_dataset = tf_test_dataset.shuffle(buffer_size=64).batch(32)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "6fC-g1B7wGgZ"
-      },
-      "source": [
-        "## TensorFlow"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ZRbU6gMUzaxq"
-      },
-      "source": [
-        "### Remote training with GPU\n",
-        "\n",
-        "First, train a TensorFlow model as a remote training job:\n",
-        "\n",
-        "- Reinitialize Vertex AI for remote training.\n",
-        "- Set Sequential for the remote training job.\n",
-        "- Invoke Sequential locally which will launch the remote training job."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "hgV3626pGoH2"
-      },
-      "outputs": [],
-      "source": [
-        "# Switch to remote mode for training\n",
-        "vertexai.preview.init(remote=True)\n",
-        "\n",
-        "from tensorflow import keras\n",
-        "\n",
-        "# Wrap classes to enable Vertex remote execution\n",
-        "keras.Sequential = vertexai.preview.remote(keras.Sequential)\n",
-        "\n",
-        "# Instantiate model\n",
-        "model = keras.Sequential(\n",
-        "    [keras.layers.Dense(5, input_shape=(4,)), keras.layers.Softmax()]\n",
-        ")\n",
-        "\n",
-        "# Specify optimizer and loss function\n",
-        "model.compile(optimizer=\"adam\", loss=\"mean_squared_error\")\n",
-        "\n",
-        "# Enable GPU training in remote_config\n",
-        "model.fit.vertex.remote_config.enable_cuda = True\n",
-        "\n",
-        "# Train model on Vertex\n",
-        "model.fit(tf_train_dataset, epochs=10)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "7JAwOWXX-uTV"
-      },
-      "source": [
-        "### Uptrain the pretrained model with autologging feature\n",
-        "\n",
-        "Next, get the registered model from the Vertex AI Model Registry. Then request the pretrained version of the model."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "oD2Bw3QzFdVV"
-      },
-      "outputs": [],
-      "source": [
-        "registered_model = vertexai.preview.register(model)\n",
-        "\n",
-        "pulled_model = vertexai.preview.from_pretrained(\n",
-        "    model_name=registered_model.resource_name\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "XZaN8KVu6BaX"
-      },
-      "source": [
-        "Now train the model remotely via Vertex AI Training."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "oBsY7MRBjeM8"
-      },
-      "outputs": [],
-      "source": [
-        "# Config experiment and turn on autologging\n",
-        "vertexai.init(\n",
-        "    project=PROJECT_ID,\n",
-        "    location=REGION,\n",
-        "    staging_bucket=REMOTE_JOB_BUCKET,\n",
-        "    experiment=\"test-remote-training-autologging\",\n",
-        ")\n",
-        "vertexai.preview.init(remote=True, autolog=True)\n",
-        "\n",
-        "# service account is required since autolog is True\n",
-        "pulled_model.fit.vertex.remote_config.service_account = \"GCE\"\n",
-        "\n",
-        "# Turn off GPU training\n",
-        "pulled_model.fit.vertex.remote_config.enable_cuda = False\n",
-        "\n",
-        "# Train model on Vertex\n",
-        "pulled_model.fit(tf_retrain_dataset, epochs=10)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "yO9fryXI6BaX"
-      },
-      "source": [
-        "### Get experiments results\n",
-        "\n",
-        "Finally, get the Vertex AI Experiments results from the remote training job."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "1ANPcw0Q-Fx7"
-      },
-      "outputs": [],
-      "source": [
-        "# View logged metrics & params\n",
-        "vertexai.preview.get_experiment_df()\n",
-        "\n",
-        "# Turn off the autologging\n",
-        "vertexai.preview.init(autolog=False)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "gHrWH_r0e6Kw"
-      },
-      "source": [
-        "### Local evaluation\n",
-        "\n",
-        "Next, evaluate the pretrained and uptrained versions of the model, and compare the results."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "gMyrRUe9a2ci"
-      },
-      "outputs": [],
-      "source": [
-        "# Switch to local mode for testing\n",
-        "vertexai.preview.init(remote=False)\n",
-        "\n",
-        "# Evaluate model's mean square errors\n",
-        "print(f\"Train loss: {model.evaluate(tf_train_dataset)}\")\n",
-        "print(f\"Test loss: {model.evaluate(tf_test_dataset)}\")\n",
-        "\n",
-        "# Evaluate uptrained model's mean square errors\n",
-        "print(f\"Train loss: {pulled_model.evaluate(tf_retrain_dataset)}\")\n",
-        "print(f\"Test loss: {pulled_model.evaluate(tf_test_dataset)}\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "u3RU6l0X6BaY"
-      },
-      "source": [
-        "#### Delete the registered model\n",
-        "\n",
-        "You can delete the registered model in the Vertex AI Model Registry with the delete() method."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "nlyFEi8M6BaY"
-      },
-      "outputs": [],
-      "source": [
-        "registered_model.delete()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ArvOcC4B6BaY"
-      },
-      "source": [
-        "## Cleaning up\n",
-        "\n",
-        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
-        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
-        "\n",
-        "Otherwise, you can delete the individual resources you created in this tutorial."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "yYnCN1gV6BaZ"
-      },
-      "outputs": [],
-      "source": [
-        "import os\n",
-        "\n",
-        "delete_bucket = False\n",
-        "\n",
-        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
-        "    ! gsutil rm -rf {BUCKET_URI}"
-      ]
-    }
-  ],
-  "metadata": {
-    "accelerator": "GPU",
-    "colab": {
-      "collapsed_sections": [
-        "-aNMacb-Ghxs",
-        "6fC-g1B7wGgZ"
-      ],
-      "name": "remote_training_tensorflow_with_autologging.ipynb",
-      "toc_visible": true
-    },
-    "kernelspec": {
-      "display_name": "Python 3",
-      "name": "python3"
-    }
-  },
-  "nbformat": 4,
-  "nbformat_minor": 0
-}