remove data science package samples since they are deprecated

* Update build_model_experimentation_lineage_with_prebuild_code.ipynb

added colab enterprise logo and link

* Update build_model_experimentation_lineage_with_prebuild_code.ipynb

hope I fixed the JSON issue

* Update notebook_template.ipynb

- changed "Getting Started" to "Get started" to be in compliance with style guide
- added "for Python" to "Vertex AI SDK" to be in compliance with product guidelines

* fix: reset changes

---------

Co-authored-by: Katie Nguyen <21978337+katiemn@users.noreply.github.com>

notebooks/community/CODEOWNERS

@@ -141,3 +141,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

notebooks/community/model_garden/model_garden_gemma_fine_tuning_batch_deployment_on_rov.ipynb

@@ -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",

notebooks/community/model_garden/model_garden_timesfm_deployment_on_vertex.ipynb

@@ -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
+}

notebooks/notebook_template.ipynb

@@ -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"
       ]
     },
     {
@@ -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)."
       ]

notebooks/official/experiments/comparing_local_trained_models.ipynb

@@ -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"
       ]
     },
     {

notebooks/official/experiments/delete_outdated_tensorboard_experiments.ipynb

@@ -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",

notebooks/official/feature_store/online_feature_serving_and_fetching_bigquery_data_with_feature_store_optimized.ipynb

@@ -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",

notebooks/official/feature_store/vertex_ai_feature_store_based_llm_grounding_tutorial.ipynb

@@ -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",
         ")"
       ]

notebooks/official/feature_store/vertex_ai_feature_store_feature_view_service_agents.ipynb

@@ -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",

notebooks/official/generative_ai/batch_eval_llm.ipynb

@@ -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"

notebooks/official/generative_ai/tuned_text-embeddings.ipynb

@@ -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."
       ]
     },
     {

notebooks/official/generative_ai/vertex_sdk_llm_snippets.ipynb

@@ -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"
       ]
     },
     {

notebooks/official/model_monitoring_v2/model_monitoring_for_custom_model_batch_prediction_job.ipynb

@@ -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:"
       ]
     },
     {

notebooks/official/model_monitoring_v2/model_monitoring_for_custom_model_online_prediction.ipynb

@@ -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:"
       ]
     },
     {

notebooks/official/model_monitoring_v2/model_monitoring_for_model_outside_vertex.ipynb

@@ -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:"
       ]
     },
     {

notebooks/official/model_registry/get_started_with_model_registry.ipynb

@@ -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."
       ]

notebooks/official/ray_on_vertex_ai/ray_cluster_management.ipynb

@@ -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)"
       ]
     }
   ],

notebooks/official/tensorboard/tensorboard_custom_training_with_custom_container.ipynb

@@ -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>"
       ]
     },
     {

notebooks/official/tensorboard/tensorboard_hyperparameter_tuning_with_hparams.ipynb

@@ -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>"
       ]
     },

notebooks/official/tensorboard/tensorboard_profiler_custom_training_with_prebuilt_container.ipynb

@@ -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>"
       ]
     },

notebooks/official/vertex_ai_sdk/remote_training_bigframes_pytorch.ipynb

@@ -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
-}

notebooks/official/vertex_ai_sdk/remote_training_bigframes_sklearn.ipynb

@@ -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
-}

notebooks/official/vertex_ai_sdk/remote_training_bigframes_tensorflow.ipynb

@@ -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
-}

notebooks/official/vertex_ai_sdk/remote_training_lightning.ipynb

@@ -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
-}

notebooks/official/vertex_ai_sdk/remote_training_pytorch.ipynb

@@ -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
-}

notebooks/official/vertex_ai_sdk/remote_training_sklearn.ipynb

@@ -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
-}

notebooks/official/vertex_ai_sdk/remote_training_tensorflow_with_autologging.ipynb

@@ -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
-}