<commit_before>{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2022 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": [
        "<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",
        "    </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",
        "      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://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/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",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>                                                                                               \n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "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",
        "\n",
        "### Dataset\n",
        "\n",
        "The dataset is the Tensorflow Dataset's Large Yelp Review Dataset. The Yelp reviews dataset consists of reviews from Yelp. For more information, please refer to this [link](http://www.yelp.com/dataset).\n",
        "\n",
        "### Objective\n",
        "\n",
        "In this notebook, you will learn how to \n",
        "\n",
        "- log the model parameters\n",
        "- log the loss and metrics on every epoch to TensorBoard\n",
        "- log the evaluation metrics\n",
        "- compare two experiments\n",
        "\n",
        "in Vertex AI Experiment of a recurrent neural network (RNN) for sentiment analysis.\n",
        "\n",
        "### Costs \n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI\n",
        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
        "Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ze4-nDLfK4pw"
      },
      "source": [
        "### Set up your local development environment\n",
        "\n",
        "**If you are using Colab or Vertex AI Workbench Notebooks**, your environment already meets\n",
        "all the requirements to run this notebook. You can skip this step."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gCuSR8GkAgzl"
      },
      "source": [
        "**Otherwise**, make sure your environment meets this notebook's requirements.\n",
        "You need the following:\n",
        "\n",
        "* The Google Cloud SDK\n",
        "* Git\n",
        "* Python 3\n",
        "* virtualenv\n",
        "* Jupyter notebook running in a virtual environment with Python 3\n",
        "\n",
        "The Google Cloud guide to [Setting up a Python development\n",
        "environment](https://cloud.google.com/python/setup) and the [Jupyter\n",
        "installation guide](https://jupyter.org/install) provide detailed instructions\n",
        "for meeting these requirements. The following steps provide a condensed set of\n",
        "instructions:\n",
        "\n",
        "1. [Install and initialize the Cloud SDK.](https://cloud.google.com/sdk/docs/)\n",
        "\n",
        "1. [Install Python 3.](https://cloud.google.com/python/setup#installing_python)\n",
        "\n",
        "1. [Install\n",
        "   virtualenv](https://cloud.google.com/python/setup#installing_and_using_virtualenv)\n",
        "   and create a virtual environment that uses Python 3. Activate the virtual environment.\n",
        "\n",
        "1. To install Jupyter, run `pip3 install jupyter` on the\n",
        "command-line in a terminal shell.\n",
        "\n",
        "1. To launch Jupyter, run `jupyter notebook` on the command-line in a terminal shell.\n",
        "\n",
        "1. Open this notebook in the Jupyter Notebook Dashboard."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i7EUnXsZhAGF"
      },
      "source": [
        "### Install additional packages\n",
        "\n",
        "Install additional package dependencies not installed in your notebook environment, such as TensorFlow, TensorFlow Datasets and Vertex AI SDK. Use the latest major GA version of each package."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2b4ef9b72d43"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# The Vertex AI Workbench Notebook product has specific requirements\n",
        "IS_WORKBENCH_NOTEBOOK = os.getenv(\"DL_ANACONDA_HOME\")\n",
        "IS_USER_MANAGED_WORKBENCH_NOTEBOOK = os.path.exists(\n",
        "    \"/opt/deeplearning/metadata/env_version\"\n",
        ")\n",
        "\n",
        "# Vertex AI Notebook requires dependencies to be installed with '--user'\n",
        "USER_FLAG = \"\"\n",
        "if IS_WORKBENCH_NOTEBOOK:\n",
        "    USER_FLAG = \"--user\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wyy5Lbnzg5fi"
      },
      "outputs": [],
      "source": [
        "! pip3 install {USER_FLAG} --upgrade tensorflow==2.8.0 tensorflow_datasets==4.5.2 -q\n",
        "! pip3 install --user --force-reinstall 'google-cloud-aiplatform>=1.15' -q"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hhq5zEbGg0XX"
      },
      "source": [
        "### Restart the kernel\n",
        "\n",
        "After you install the additional packages, you need to restart the notebook kernel so it can find the packages."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EzrelQZ22IZj"
      },
      "outputs": [],
      "source": [
        "# Automatically restart kernel after installs\n",
        "import os\n",
        "\n",
        "if not os.getenv(\"IS_TESTING\"):\n",
        "    # Automatically restart kernel after installs\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lWEdiXsJg0XY"
      },
      "source": [
        "## Before you begin"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BF1j6f9HApxa"
      },
      "source": [
        "### Set up your Google Cloud project\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
        "\n",
        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "1. [Enable APIs](https://console.cloud.google.com/flows/enableapi?apiid=cloudresourcemanager.googleapis.com,aiplatform.googleapis.com).\n",
        "\n",
        "1. If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
        "\n",
        "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
        "Cloud SDK uses the right project for all the commands in this notebook.\n",
        "\n",
        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WReHDGG5g0XY"
      },
      "source": [
        "#### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, you may be able to get your project ID using `gcloud`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oM1iC_MfAts1"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "PROJECT_ID = \"\"\n",
        "\n",
        "# Get your Google Cloud project ID from gcloud\n",
        "if not os.getenv(\"IS_TESTING\"):\n",
        "    shell_output = !gcloud config list --format 'value(core.project)' 2>/dev/null\n",
        "    PROJECT_ID = shell_output[0]\n",
        "    print(\"Project ID: \", PROJECT_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qJYoRfYng0XZ"
      },
      "source": [
        "Otherwise, set your project ID here."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "riG_qUokg0XZ"
      },
      "outputs": [],
      "source": [
        "if PROJECT_ID == \"\" or PROJECT_ID is None:\n",
        "    PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WXaOgYQVYya8"
      },
      "outputs": [],
      "source": [
        "!gcloud config set project $PROJECT_ID"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "03d8d65b914d"
      },
      "source": [
        "#### Region\n",
        "\n",
        "You can also change the `REGION` variable, which is used for operations\n",
        "throughout the rest of this notebook.  Below are regions supported for Vertex AI. We recommend that you choose the region closest to you.\n",
        "\n",
        "- Americas: `us-central1`\n",
        "- Europe: `europe-west4`\n",
        "- Asia Pacific: `asia-east1`\n",
        "\n",
        "You may not use a multi-regional bucket for training with Vertex AI. Not all regions provide support for all Vertex AI services.\n",
        "\n",
        "Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3281bedf6d3c"
      },
      "outputs": [],
      "source": [
        "REGION = \"[your-region]\"  # @param {type: \"string\"}\n",
        "\n",
        "if REGION == \"[your-region]\":\n",
        "    REGION = \"us-central1\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "06571eb4063b"
      },
      "source": [
        "#### Timestamp\n",
        "\n",
        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a timestamp for each instance session, and append it onto the name of resources you create in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "697568e92bd6"
      },
      "outputs": [],
      "source": [
        "from datetime import datetime\n",
        "\n",
        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dr--iN2kAylZ"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "**If you are using Vertex AI Workbench Notebooks**, your environment is already\n",
        "authenticated. Skip this step."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sBCra4QMA2wR"
      },
      "source": [
        "**If you are using Colab**, run the cell below and follow the instructions\n",
        "when prompted to authenticate your account via oAuth.\n",
        "\n",
        "**Otherwise**, follow these steps:\n",
        "\n",
        "1. In the Cloud Console, go to the [**Create service account key**\n",
        "   page](https://console.cloud.google.com/apis/credentials/serviceaccountkey).\n",
        "\n",
        "2. Click **Create service account**.\n",
        "\n",
        "3. In the **Service account name** field, enter a name, and\n",
        "   click **Create**.\n",
        "\n",
        "4. In the **Grant this service account access to project** section, click the **Role** drop-down list. Type \"Vertex AI\"\n",
        "into the filter box, and select\n",
        "   **Vertex AI Administrator**. Type \"Storage Object Admin\" into the filter box, and select **Storage Object Admin**.\n",
        "\n",
        "5. Click *Create*. A JSON file that contains your key downloads to your\n",
        "local environment.\n",
        "\n",
        "6. Enter the path to your service account key as the\n",
        "`GOOGLE_APPLICATION_CREDENTIALS` variable in the cell below and run the cell."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PyQmSRbKA8r-"
      },
      "outputs": [],
      "source": [
        "# If you are running this notebook in Colab, run this cell and follow the\n",
        "# instructions to authenticate your GCP account. This provides access to your\n",
        "# Cloud Storage bucket and lets you submit training jobs and prediction\n",
        "# requests.\n",
        "\n",
        "import os\n",
        "import sys\n",
        "\n",
        "# If on Vertex AI Workbench, then don't execute this code\n",
        "\n",
        "IS_COLAB = False\n",
        "\n",
        "if not os.path.exists(\"/opt/deeplearning/metadata/env_version\") and not os.getenv(\n",
        "    \"DL_ANACONDA_HOME\"\n",
        "):\n",
        "    if \"google.colab\" in sys.modules:\n",
        "\n",
        "        IS_COLAB = True\n",
        "\n",
        "        from google.colab import auth as google_auth\n",
        "\n",
        "        google_auth.authenticate_user()\n",
        "\n",
        "    # If you are running this notebook locally, replace the string below with the\n",
        "    # path to your service account key and run this cell to authenticate your GCP\n",
        "    # account.\n",
        "\n",
        "    elif not os.getenv(\"IS_TESTING\"):\n",
        "\n",
        "        %env GOOGLE_APPLICATION_CREDENTIALS ''"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgPO1eR3CYjk"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "When you initialize the Vertex SDK for Python, you specify a Cloud Storage staging bucket. The staging bucket is where all the data associated with your dataset and model resources are retained across sessions.\n",
        "\n",
        "Set the name of your Cloud Storage bucket below. Bucket names must be globally unique across all Google Cloud projects, including those outside of your organization."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MzGDU7TWdts_"
      },
      "outputs": [],
      "source": [
        "BUCKET_NAME = \"[your-bucket-name]\"  # @param {type:\"string\"}\n",
        "BUCKET_URI = f\"gs://{BUCKET_NAME}\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cf221059d072"
      },
      "outputs": [],
      "source": [
        "if BUCKET_NAME == \"\" or BUCKET_NAME is None or BUCKET_NAME == \"[your-bucket-name]\":\n",
        "    BUCKET_NAME = PROJECT_ID + \"aip-\" + TIMESTAMP\n",
        "    BUCKET_URI = f\"gs://{BUCKET_NAME}\""
      ]
    },
    {
      "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": "ucvCsknMCims"
      },
      "source": [
        "Finally, validate access to your Cloud Storage bucket by examining its contents:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vhOb7YnwClBb"
      },
      "outputs": [],
      "source": [
        "! gsutil ls -al $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XoEqT2Y4DJmf"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pRUOFELefqf1"
      },
      "outputs": [],
      "source": [
        "# General\n",
        "import logging\n",
        "\n",
        "logger = logging.getLogger(\"logger\")\n",
        "logging.basicConfig(level=logging.INFO)\n",
        "\n",
        "# Vertex AI\n",
        "from google.cloud import aiplatform as vertex_ai"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xAoqKXTsR2Zd"
      },
      "source": [
        "### Define constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hl6BCLNnXrSD"
      },
      "outputs": [],
      "source": [
        "# Training\n",
        "DATA_DIR = \"data\"\n",
        "LOG_DIR = \"logs\"\n",
        "\n",
        "# Experiments\n",
        "TASK = \"classification\"\n",
        "MODEL_TYPE = \"rnn\"\n",
        "EXPERIMENT_NAME = f\"{PROJECT_ID}-{TASK}-{MODEL_TYPE}-{TIMESTAMP}\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "inR70nh38PeK"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nz0nasrh8T3c"
      },
      "outputs": [],
      "source": [
        "vertex_ai.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tUH9-Vitix5K"
      },
      "source": [
        "### Create TensorBoard instance using Vertex AI TensorBoard\n",
        "\n",
        "You can upload your TensorBoard logs by first creating a TensorBoard instance.\n",
        "\n",
        "Learn more about [TensorBoard overview](https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-overview)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NBpfzfA9iv3j"
      },
      "outputs": [],
      "source": [
        "vertex_ai_tb = vertex_ai.Tensorboard.create()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jotNH6scnV4S"
      },
      "outputs": [],
      "source": [
        "vertex_ai.init(experiment=EXPERIMENT_NAME, experiment_tensorboard=vertex_ai_tb)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tensorboard"
      },
      "source": [
        "## Training with Vertex AI Experiments\n",
        "\n",
        "Vertex AI enables users to track the steps (for example, preprocessing, training) of an experiment run, and track inputs (for example, algorithm, parameters, datasets) and outputs (for example, models, checkpoints, metrics) of those steps. \n",
        "\n",
        "Below you have some example of how track experiments to train recurrent neural network for sentiment analysis."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VVyNpW9Uqiyt"
      },
      "outputs": [],
      "source": [
        "!mkdir -p -m 777 $DATA_DIR $LOG_DIR"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cTZGTxAXYKAE"
      },
      "outputs": [],
      "source": [
        "import tensorflow_datasets as tfds\n",
        "\n",
        "tfds.disable_progress_bar()\n",
        "import tensorflow as tf\n",
        "\n",
        "\n",
        "# Helpers ----------------------------------------------------------------------\n",
        "def get_dataset(data_dir, buffer_size, batch_size):\n",
        "    \"\"\"\n",
        "    Returns a tf.data.Dataset object containing the training data\n",
        "    Returns:\n",
        "        tf.data.Dataset: A tf.data.Dataset object containing the training data\n",
        "        buffer_size: The buffer size for prefetch data\n",
        "        batch_size: The batch size of dataset\n",
        "    \"\"\"\n",
        "\n",
        "    dataset = tfds.load(\n",
        "        \"imdb_reviews\",\n",
        "        download=True,\n",
        "        data_dir=data_dir,\n",
        "        with_info=False,\n",
        "        as_supervised=True,\n",
        "    )\n",
        "\n",
        "    train_dataset, test_dataset = dataset[\"train\"], dataset[\"test\"]\n",
        "\n",
        "    train_dataset = (\n",
        "        train_dataset.shuffle(buffer_size).batch(batch_size).prefetch(tf.data.AUTOTUNE)\n",
        "    )\n",
        "    test_dataset = test_dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)\n",
        "    return train_dataset, test_dataset\n",
        "\n",
        "\n",
        "def get_encoder(train_dataset, vocab_size):\n",
        "    \"\"\"\n",
        "    Returns a TextVectorization object for the encoder\n",
        "    Args:\n",
        "        train_dataset: A tf.data.Dataset object containing the training data\n",
        "    Returns:\n",
        "        TextVectorization: A TextVectorization object for the encoder\n",
        "    \"\"\"\n",
        "\n",
        "    encoder = tf.keras.layers.TextVectorization(max_tokens=vocab_size)\n",
        "    encoder.adapt(\n",
        "        train_dataset.map(lambda text, label: text, num_parallel_calls=tf.data.AUTOTUNE)\n",
        "    )\n",
        "    return encoder\n",
        "\n",
        "\n",
        "def get_model(encoder, model_params, role):\n",
        "    \"\"\"\n",
        "    Returns a tf.keras.Model object for the model\n",
        "    Args:\n",
        "        encoder: A TextVectorization object for the encoder\n",
        "        model_params: A dictionary with model parameters\n",
        "        role: A variable to set the role of model\n",
        "    Returns:\n",
        "        tf.keras.Model: A tf.keras.Model object for the model\n",
        "    \"\"\"\n",
        "    model = tf.keras.Sequential()\n",
        "    model.add(encoder)\n",
        "    model.add(\n",
        "        tf.keras.layers.Embedding(\n",
        "            input_dim=len(encoder.get_vocabulary()), output_dim=64, mask_zero=True\n",
        "        )\n",
        "    )\n",
        "    if role == \"baseline\":\n",
        "        model.add(tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(64)))\n",
        "        model.add(tf.keras.layers.Dense(64, activation=\"relu\"))\n",
        "        model.add(tf.keras.layers.Dense(1))\n",
        "    else:\n",
        "        model.add(\n",
        "            tf.keras.layers.Bidirectional(\n",
        "                tf.keras.layers.LSTM(64, return_sequences=True)\n",
        "            )\n",
        "        )\n",
        "        model.add(tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(32)))\n",
        "        model.add(tf.keras.layers.Dropout(0.5))\n",
        "        model.add(tf.keras.layers.Dense(1))\n",
        "\n",
        "    model.compile(\n",
        "        loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),\n",
        "        optimizer=tf.keras.optimizers.Adam(\n",
        "            learning_rate=model_params[\"learning_rate\"],\n",
        "            beta_1=model_params[\"beta_1\"],\n",
        "            beta_2=model_params[\"beta_2\"],\n",
        "            epsilon=model_params[\"epsilon\"],\n",
        "        ),\n",
        "        metrics=[\"accuracy\"],\n",
        "    )\n",
        "    return model\n",
        "\n",
        "\n",
        "def train(train_dataset, test_dataset, model, epochs, steps, log_dir=LOG_DIR):\n",
        "    \"\"\"\n",
        "    Trains the model\n",
        "    Args:\n",
        "        train_dataset: A tf.data.Dataset object containing the training data\n",
        "        test_dataset: A tf.data.Dataset object containing the testing data\n",
        "        model: A tf.keras.Model object for the model\n",
        "        epochs: The number of epochs\n",
        "        steps: The number of validation steps\n",
        "        log_dir: The location of tf training logs\n",
        "    \"\"\"\n",
        "    tensorboard = tf.keras.callbacks.TensorBoard(log_dir=log_dir)\n",
        "    history = model.fit(\n",
        "        train_dataset,\n",
        "        validation_data=test_dataset,\n",
        "        epochs=epochs,\n",
        "        validation_steps=steps,\n",
        "        callbacks=[tensorboard],\n",
        "    )\n",
        "    return history"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Uz7MSBcwVHgc"
      },
      "source": [
        "### Define a baseline\n",
        "\n",
        "As part of the experimentation phase, you would like to define a baseline model to have a performance benchmark which it is represented by recurrent neural network (RNN) in this case. Then you try to improve the ML solution with more complex training configuration that produce better results."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TM1xGPRkjocQ"
      },
      "source": [
        "#### Run experiment and evaluate experiment runs using `with` statement\n",
        "\n",
        "This step would takes **10 min** approx. to finish.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zRp-gJWhPKf0"
      },
      "outputs": [],
      "source": [
        "# Experiment Settings ----------------------------------------------------------\n",
        "ID_1 = \"run-1\"\n",
        "BUFFER_SIZE = 10000\n",
        "BATCH_SIZE = 64\n",
        "VOCAB_SIZE = 1000\n",
        "EPOCHS = 2\n",
        "STEPS = 20\n",
        "ROLE = \"baseline\"\n",
        "LR = 1e-4\n",
        "B_1 = 0.9\n",
        "B_2 = 0.999\n",
        "EPS = 1e-07\n",
        "\n",
        "# Initialize the experiment\n",
        "logging.info(\"Initialize the experiment.\")\n",
        "with vertex_ai.start_run(ID_1) as run:\n",
        "\n",
        "    # Get the training and testing datasets\n",
        "    logging.info(\"Get the training and testing datasets.\")\n",
        "    data_params = {\"buffer_size\": BUFFER_SIZE, \"batch_size\": BATCH_SIZE}\n",
        "    train_dataset, test_dataset = get_dataset(\n",
        "        \"./data\",\n",
        "        buffer_size=data_params[\"buffer_size\"],\n",
        "        batch_size=data_params[\"batch_size\"],\n",
        "    )\n",
        "    run.log_params(data_params)\n",
        "\n",
        "    # Get the encoder\n",
        "    logging.info(\"Get the encoder.\")\n",
        "    encoder_params = {\"vocab_size\": VOCAB_SIZE}\n",
        "    encoder = get_encoder(\n",
        "        train_dataset=train_dataset, vocab_size=encoder_params[\"vocab_size\"]\n",
        "    )\n",
        "    run.log_params(encoder_params)\n",
        "\n",
        "    # Get the model\n",
        "    logging.info(\"Get the model.\")\n",
        "    run.log_params({\"role\": ROLE})\n",
        "    model_params = {\"learning_rate\": LR, \"beta_1\": B_1, \"beta_2\": B_2, \"epsilon\": EPS}\n",
        "    model = get_model(encoder=encoder, model_params=model_params, role=ROLE)\n",
        "    run.log_params(model_params)\n",
        "\n",
        "    # Train the model\n",
        "    logging.info(\"Train the model.\")\n",
        "    history = train(\n",
        "        model=model,\n",
        "        train_dataset=train_dataset,\n",
        "        test_dataset=test_dataset,\n",
        "        epochs=EPOCHS,\n",
        "        steps=STEPS,\n",
        "    )\n",
        "\n",
        "    run.log_params(history.params)\n",
        "    for idx in range(0, history.params[\"epochs\"]):\n",
        "        run.log_time_series_metrics(\n",
        "            {\n",
        "                \"train_loss\": history.history[\"loss\"][idx],\n",
        "                \"train_accuracy\": history.history[\"accuracy\"][idx],\n",
        "            }\n",
        "        )\n",
        "\n",
        "    # Evaluate model\n",
        "    logging.info(\"Evaluate model.\")\n",
        "    test_loss, test_accuracy = model.evaluate(test_dataset)\n",
        "    run.log_metrics({\"test_loss\": test_loss, \"test_accuracy\": test_accuracy})\n",
        "\n",
        "    # Exit the experiment\n",
        "    logging.info(\"Exit the run.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f8bd7ea02120"
      },
      "source": [
        "#### Evaluate the baseline model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xmaumLrSqo1e"
      },
      "outputs": [],
      "source": [
        "# Get experiment\n",
        "logging.info(\"Get experiment status.\")\n",
        "experiment_df = vertex_ai.get_experiment_df()\n",
        "experiment_df"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "evxi4ROLqpc_"
      },
      "outputs": [],
      "source": [
        "# Get time series metrics\n",
        "logging.info(\"Get time series metrics.\")\n",
        "ts_runs_df = run.get_time_series_data_frame()\n",
        "ts_runs_df"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eoGfiTobVSSB"
      },
      "source": [
        "### Make experimental improvement to the model\n",
        "\n",
        "As challenger model you add a LSTM layer to the baseline model. This step would takes **10 min** approx. to finish.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZjEww1ESZFTn"
      },
      "source": [
        "#### Run experiment and evaluate experiment runs with `experiment` instance\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kV8DIuUXVXkC"
      },
      "outputs": [],
      "source": [
        "# Experiment Settings ----------------------------------------------------------\n",
        "ID_2 = \"run-2\"\n",
        "ROLE = \"stacked\"\n",
        "\n",
        "# Initialize the experiment\n",
        "logger.info(\"Initialize the experiment.\")\n",
        "vertex_ai.start_run(ID_2)\n",
        "\n",
        "# Get the model\n",
        "logging.info(\"Get the model.\")\n",
        "run.log_params({\"role\": ROLE})\n",
        "model = get_model(encoder=encoder, model_params=model_params, role=ROLE)\n",
        "vertex_ai.log_params(model_params)\n",
        "\n",
        "# Train the model\n",
        "logging.info(\"Train the model.\")\n",
        "history = train(\n",
        "    model=model,\n",
        "    train_dataset=train_dataset,\n",
        "    test_dataset=test_dataset,\n",
        "    epochs=EPOCHS,\n",
        "    steps=STEPS,\n",
        ")\n",
        "vertex_ai.log_params(history.params)\n",
        "for idx in range(0, history.params[\"epochs\"]):\n",
        "    vertex_ai.log_time_series_metrics(\n",
        "        {\n",
        "            \"train_loss\": history.history[\"loss\"][idx],\n",
        "            \"train_accuracy\": history.history[\"accuracy\"][idx],\n",
        "        }\n",
        "    )\n",
        "\n",
        "# Evaluate model\n",
        "logging.info(\"Evaluate model.\")\n",
        "test_loss, test_accuracy = model.evaluate(test_dataset)\n",
        "vertex_ai.log_metrics({\"test_loss\": test_loss, \"test_accuracy\": test_accuracy})\n",
        "\n",
        "# Exit the experiment\n",
        "logging.info(\"Exit the experiment.\")\n",
        "vertex_ai.end_run()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ceaf7b22157d"
      },
      "source": [
        "#### Compare the baseline with the new model implementation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bUYoJXexrpmu"
      },
      "outputs": [],
      "source": [
        "# Get experiment\n",
        "logging.info(\"Get experiment status.\")\n",
        "experiment_df = vertex_ai.get_experiment_df()\n",
        "experiment_df"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fORBLdvLrqFU"
      },
      "outputs": [],
      "source": [
        "# Get time series metrics\n",
        "exp_run = vertex_ai.ExperimentRun(ID_2, experiment=EXPERIMENT_NAME)\n",
        "logging.info(\"Get time series metrics.\")\n",
        "ts_runs_df = exp_run.get_time_series_data_frame()\n",
        "print(ts_runs_df)"
      ]
    },
    {
      "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": "uGvJieN5BjPG"
      },
      "outputs": [],
      "source": [
        "# Delete experiment\n",
        "exp = vertex_ai.Experiment(EXPERIMENT_NAME)\n",
        "exp.delete(delete_backing_tensorboard_runs=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dde8937123d4"
      },
      "outputs": [],
      "source": [
        "# Delete Tensorboard\n",
        "vertex_ai_tb.delete()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sx_vKniMq9ZX"
      },
      "outputs": [],
      "source": [
        "# Delete Cloud Storage objects that were created\n",
        "delete_bucket = False\n",
        "\n",
        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
        "    ! gsutil rm -rf {BUCKET_URI}"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "comparing_local_trained_models.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}<commit_msg>Vertex AI Experiments - Compare Models - Notebook review (#852)<commit_after>{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2022 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": [
        "<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",
        "    </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",
        "      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://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/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",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>                                                                                               \n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "# Compare trained and evaluated model experiments using Vertex AI Experiments"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3a0651225470"
      },
      "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. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f3021b2963a6"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn how to use Vertex AI Experiments to compare and evaluate model experiments.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
        "- Vertex AI Workbench\n",
        "- Vertex AI Experiments\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- log the model parameters\n",
        "- log the loss and metrics on every epoch to TensorBoard\n",
        "- log the evaluation metrics\n",
        "- compare two experiments\n",
        "\n",
        "in Vertex AI Experiment of a recurrent neural network (RNN) for sentiment analysis."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "20a5168cf05e"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "The dataset is the Tensorflow Dataset's Large Yelp Review Dataset. The Yelp reviews dataset consists of reviews from Yelp. For more information, please refer to this [link](http://www.yelp.com/dataset).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "de76bb18c85b"
      },
      "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) and [Cloud Storage\n",
        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
        "Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ze4-nDLfK4pw"
      },
      "source": [
        "### Set up your local development environment\n",
        "\n",
        "**If you are using Colab or Vertex AI Workbench Notebooks**, your environment already meets\n",
        "all the requirements to run this notebook. You can skip this step.\n",
        "\n",
        "**Otherwise**, make sure your environment meets this notebook's requirements.\n",
        "You need the following:\n",
        "\n",
        "* The Google Cloud SDK\n",
        "* Git\n",
        "* Python 3\n",
        "* virtualenv\n",
        "* Jupyter notebook running in a virtual environment with Python 3\n",
        "\n",
        "The Google Cloud guide to [Setting up a Python development\n",
        "environment](https://cloud.google.com/python/setup) and the [Jupyter\n",
        "installation guide](https://jupyter.org/install) provide detailed instructions\n",
        "for meeting these requirements. The following steps provide a condensed set of\n",
        "instructions:\n",
        "\n",
        "1. [Install and initialize the Cloud SDK.](https://cloud.google.com/sdk/docs/)\n",
        "\n",
        "1. [Install Python 3.](https://cloud.google.com/python/setup#installing_python)\n",
        "\n",
        "1. [Install\n",
        "   virtualenv](https://cloud.google.com/python/setup#installing_and_using_virtualenv)\n",
        "   and create a virtual environment that uses Python 3. Activate the virtual environment.\n",
        "\n",
        "1. To install Jupyter, run `pip3 install jupyter` on the\n",
        "command-line in a terminal shell.\n",
        "\n",
        "1. To launch Jupyter, run `jupyter notebook` on the command-line in a terminal shell.\n",
        "\n",
        "1. Open this notebook in the Jupyter Notebook Dashboard."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i7EUnXsZhAGF"
      },
      "source": [
        "### Install additional packages\n",
        "\n",
        "Install additional package dependencies not installed in your notebook environment, such as TensorFlow, TensorFlow Datasets and Vertex AI SDK. Use the latest major GA version of each package."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2b4ef9b72d43"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# The Vertex AI Workbench Notebook product has specific requirements\n",
        "IS_WORKBENCH_NOTEBOOK = os.getenv(\"DL_ANACONDA_HOME\")\n",
        "IS_USER_MANAGED_WORKBENCH_NOTEBOOK = os.path.exists(\n",
        "    \"/opt/deeplearning/metadata/env_version\"\n",
        ")\n",
        "\n",
        "# Vertex AI Notebook requires dependencies to be installed with '--user'\n",
        "USER_FLAG = \"\"\n",
        "if IS_WORKBENCH_NOTEBOOK:\n",
        "    USER_FLAG = \"--user\"\n",
        "\n",
        "! pip3 install {USER_FLAG} --upgrade tensorflow==2.8.0 tensorflow_datasets==4.5.2 -q\n",
        "! pip3 install --user --force-reinstall 'google-cloud-aiplatform>=1.15' -q"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hhq5zEbGg0XX"
      },
      "source": [
        "### Restart the kernel\n",
        "\n",
        "After you install the additional packages, you need to restart the notebook kernel so it can find the packages."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EzrelQZ22IZj"
      },
      "outputs": [],
      "source": [
        "# Automatically restart kernel after installs\n",
        "import os\n",
        "\n",
        "if not os.getenv(\"IS_TESTING\"):\n",
        "    # Automatically restart kernel after installs\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lWEdiXsJg0XY"
      },
      "source": [
        "## Before you begin"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BF1j6f9HApxa"
      },
      "source": [
        "### Set up your Google Cloud project\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
        "\n",
        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "1. [Enable APIs](https://console.cloud.google.com/flows/enableapi?apiid=cloudresourcemanager.googleapis.com,aiplatform.googleapis.com).\n",
        "\n",
        "1. If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
        "\n",
        "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
        "Cloud SDK uses the right project for all the commands in this notebook.\n",
        "\n",
        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WReHDGG5g0XY"
      },
      "source": [
        "#### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, you may be able to get your project ID using `gcloud`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oM1iC_MfAts1"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "PROJECT_ID = \"\"\n",
        "\n",
        "# Get your Google Cloud project ID from gcloud\n",
        "if not os.getenv(\"IS_TESTING\"):\n",
        "    shell_output = !gcloud config list --format 'value(core.project)' 2>/dev/null\n",
        "    PROJECT_ID = shell_output[0]\n",
        "    print(\"Project ID: \", PROJECT_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qJYoRfYng0XZ"
      },
      "source": [
        "Otherwise, set your project ID here."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "riG_qUokg0XZ"
      },
      "outputs": [],
      "source": [
        "if PROJECT_ID == \"\" or PROJECT_ID is None:\n",
        "    PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WXaOgYQVYya8"
      },
      "outputs": [],
      "source": [
        "!gcloud config set project $PROJECT_ID"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "03d8d65b914d"
      },
      "source": [
        "#### Region\n",
        "\n",
        "You can also change the `REGION` variable, which is used for operations\n",
        "throughout the rest of this notebook.  Below are regions supported for Vertex AI. We recommend that you choose the region closest to you.\n",
        "\n",
        "- Americas: `us-central1`\n",
        "- Europe: `europe-west4`\n",
        "- Asia Pacific: `asia-east1`\n",
        "\n",
        "You may not use a multi-regional bucket for training with Vertex AI. Not all regions provide support for all Vertex AI services.\n",
        "\n",
        "Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3281bedf6d3c"
      },
      "outputs": [],
      "source": [
        "REGION = \"[your-region]\"  # @param {type: \"string\"}\n",
        "\n",
        "if REGION == \"[your-region]\":\n",
        "    REGION = \"us-central1\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "06571eb4063b"
      },
      "source": [
        "#### Timestamp\n",
        "\n",
        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a timestamp for each instance session, and append it onto the name of resources you create in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "697568e92bd6"
      },
      "outputs": [],
      "source": [
        "from datetime import datetime\n",
        "\n",
        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dr--iN2kAylZ"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "**If you are using Vertex AI Workbench Notebooks**, your environment is already\n",
        "authenticated. Skip this step.\n",
        "\n",
        "**If you are using Colab**, run the cell below and follow the instructions\n",
        "when prompted to authenticate your account via oAuth.\n",
        "\n",
        "**Otherwise**, follow these steps:\n",
        "\n",
        "1. In the Cloud Console, go to the [**Create service account key**\n",
        "   page](https://console.cloud.google.com/apis/credentials/serviceaccountkey).\n",
        "\n",
        "2. Click **Create service account**.\n",
        "\n",
        "3. In the **Service account name** field, enter a name, and\n",
        "   click **Create**.\n",
        "\n",
        "4. In the **Grant this service account access to project** section, click the **Role** drop-down list. Type and select\n",
        "the following role into the filter box:\n",
        "\n",
        "    *   Storage Admin\n",
        "    *   Storage Object Admin\n",
        "    *   Service Account User\n",
        "    *   Vertex AI Administrator\n",
        "\n",
        "\n",
        "5. Click *Create*. A JSON file that contains your key downloads to your\n",
        "local environment.\n",
        "\n",
        "6. Enter the path to your service account key as the\n",
        "`GOOGLE_APPLICATION_CREDENTIALS` variable in the cell below and run the cell."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PyQmSRbKA8r-"
      },
      "outputs": [],
      "source": [
        "# If you are running this notebook in Colab, run this cell and follow the\n",
        "# instructions to authenticate your GCP account. This provides access to your\n",
        "# Cloud Storage bucket and lets you submit training jobs and prediction\n",
        "# requests.\n",
        "\n",
        "import os\n",
        "import sys\n",
        "\n",
        "# If on Vertex AI Workbench, then don't execute this code\n",
        "IS_COLAB = \"google.colab\" in sys.modules\n",
        "if not os.path.exists(\"/opt/deeplearning/metadata/env_version\") and not os.getenv(\n",
        "    \"DL_ANACONDA_HOME\"\n",
        "):\n",
        "    if \"google.colab\" in sys.modules:\n",
        "        from google.colab import auth as google_auth\n",
        "\n",
        "        google_auth.authenticate_user()\n",
        "\n",
        "    # If you are running this notebook locally, replace the string below with the\n",
        "    # path to your service account key and run this cell to authenticate your GCP\n",
        "    # account.\n",
        "    elif not os.getenv(\"IS_TESTING\"):\n",
        "        %env GOOGLE_APPLICATION_CREDENTIALS ''"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgPO1eR3CYjk"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "When you initialize the Vertex SDK for Python, you specify a Cloud Storage staging bucket. The staging bucket is where all the data associated with your dataset and model resources are retained across sessions.\n",
        "\n",
        "Set the name of your Cloud Storage bucket below. Bucket names must be globally unique across all Google Cloud projects, including those outside of your organization."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MzGDU7TWdts_"
      },
      "outputs": [],
      "source": [
        "BUCKET_NAME = \"[your-bucket-name]\"  # @param {type:\"string\"}\n",
        "BUCKET_URI = f\"gs://{BUCKET_NAME}\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cf221059d072"
      },
      "outputs": [],
      "source": [
        "if BUCKET_NAME == \"\" or BUCKET_NAME is None or BUCKET_NAME == \"[your-bucket-name]\":\n",
        "    BUCKET_NAME = PROJECT_ID + \"aip-\" + TIMESTAMP\n",
        "    BUCKET_URI = f\"gs://{BUCKET_NAME}\""
      ]
    },
    {
      "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": "ucvCsknMCims"
      },
      "source": [
        "Finally, validate access to your Cloud Storage bucket by examining its contents:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vhOb7YnwClBb"
      },
      "outputs": [],
      "source": [
        "! gsutil ls -al $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XoEqT2Y4DJmf"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pRUOFELefqf1"
      },
      "outputs": [],
      "source": [
        "# General\n",
        "import logging\n",
        "\n",
        "logger = logging.getLogger(\"logger\")\n",
        "logging.basicConfig(level=logging.INFO)\n",
        "\n",
        "# Vertex AI\n",
        "from google.cloud import aiplatform as vertex_ai"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xAoqKXTsR2Zd"
      },
      "source": [
        "### Define constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hl6BCLNnXrSD"
      },
      "outputs": [],
      "source": [
        "# Training\n",
        "DATA_DIR = \"data\"\n",
        "LOG_DIR = \"logs\"\n",
        "\n",
        "# Experiments\n",
        "TASK = \"classification\"\n",
        "MODEL_TYPE = \"rnn\"\n",
        "EXPERIMENT_NAME = f\"{PROJECT_ID}-{TASK}-{MODEL_TYPE}-{TIMESTAMP}\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "inR70nh38PeK"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nz0nasrh8T3c"
      },
      "outputs": [],
      "source": [
        "vertex_ai.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tUH9-Vitix5K"
      },
      "source": [
        "### Create TensorBoard instance using Vertex AI TensorBoard\n",
        "\n",
        "You can upload your TensorBoard logs by first creating a TensorBoard instance.\n",
        "\n",
        "Learn more about [TensorBoard overview](https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-overview)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NBpfzfA9iv3j"
      },
      "outputs": [],
      "source": [
        "vertex_ai_tb = vertex_ai.Tensorboard.create()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jotNH6scnV4S"
      },
      "outputs": [],
      "source": [
        "vertex_ai.init(experiment=EXPERIMENT_NAME, experiment_tensorboard=vertex_ai_tb)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tensorboard"
      },
      "source": [
        "## Training with Vertex AI Experiments\n",
        "\n",
        "Vertex AI enables users to track the steps (for example, preprocessing, training) of an experiment run, and track inputs (for example, algorithm, parameters, datasets) and outputs (for example, models, checkpoints, metrics) of those steps. \n",
        "\n",
        "Below you have some example of how track experiments to train recurrent neural network for sentiment analysis. \n",
        "\n",
        "To simplify the code, here you have helper function to cover the following steps:\n",
        "\n",
        "- Collect training data\n",
        "- Create text encoder\n",
        "- Build a RNN as baseline model\n",
        "- Build a LSTM as challenger model\n",
        "- Train the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VVyNpW9Uqiyt"
      },
      "outputs": [],
      "source": [
        "!mkdir -p -m 777 $DATA_DIR $LOG_DIR"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cTZGTxAXYKAE"
      },
      "outputs": [],
      "source": [
        "import tensorflow_datasets as tfds\n",
        "\n",
        "tfds.disable_progress_bar()\n",
        "import tensorflow as tf\n",
        "\n",
        "\n",
        "# Helpers ----------------------------------------------------------------------\n",
        "def get_dataset(data_dir, buffer_size, batch_size):\n",
        "    \"\"\"\n",
        "    Returns a tf.data.Dataset object containing the training data\n",
        "    Returns:\n",
        "        tf.data.Dataset: A tf.data.Dataset object containing the training data\n",
        "        buffer_size: The buffer size for prefetch data\n",
        "        batch_size: The batch size of dataset\n",
        "    \"\"\"\n",
        "\n",
        "    dataset = tfds.load(\n",
        "        \"imdb_reviews\",\n",
        "        download=True,\n",
        "        data_dir=data_dir,\n",
        "        with_info=False,\n",
        "        as_supervised=True,\n",
        "    )\n",
        "\n",
        "    train_dataset, test_dataset = dataset[\"train\"], dataset[\"test\"]\n",
        "\n",
        "    train_dataset = (\n",
        "        train_dataset.shuffle(buffer_size).batch(batch_size).prefetch(tf.data.AUTOTUNE)\n",
        "    )\n",
        "    test_dataset = test_dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)\n",
        "    return train_dataset, test_dataset\n",
        "\n",
        "\n",
        "def get_encoder(train_dataset, vocab_size):\n",
        "    \"\"\"\n",
        "    Returns a TextVectorization object for the encoder\n",
        "    Args:\n",
        "        train_dataset: A tf.data.Dataset object containing the training data\n",
        "    Returns:\n",
        "        TextVectorization: A TextVectorization object for the encoder\n",
        "    \"\"\"\n",
        "\n",
        "    encoder = tf.keras.layers.TextVectorization(max_tokens=vocab_size)\n",
        "    encoder.adapt(\n",
        "        train_dataset.map(lambda text, label: text, num_parallel_calls=tf.data.AUTOTUNE)\n",
        "    )\n",
        "    return encoder\n",
        "\n",
        "\n",
        "def get_baseline_model(encoder, model_params):\n",
        "    \"\"\"\n",
        "    Returns a tf.keras.Model object for the model as baseline\n",
        "    Args:\n",
        "        encoder: A TextVectorization object for the encoder\n",
        "        model_params: A dictionary with model parameters\n",
        "    Returns:\n",
        "        tf.keras.Model: A tf.keras.Model object for the model\n",
        "    \"\"\"\n",
        "\n",
        "    model = tf.keras.Sequential()\n",
        "    model.add(encoder)\n",
        "    model.add(\n",
        "        tf.keras.layers.Embedding(\n",
        "            input_dim=len(encoder.get_vocabulary()), output_dim=64, mask_zero=True\n",
        "        )\n",
        "    )\n",
        "    model.add(tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(64)))\n",
        "    model.add(tf.keras.layers.Dense(64, activation=\"relu\"))\n",
        "    model.add(tf.keras.layers.Dense(1))\n",
        "    model.compile(\n",
        "        loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),\n",
        "        optimizer=tf.keras.optimizers.Adam(\n",
        "            learning_rate=model_params[\"learning_rate\"],\n",
        "            beta_1=model_params[\"beta_1\"],\n",
        "            beta_2=model_params[\"beta_2\"],\n",
        "            epsilon=model_params[\"epsilon\"],\n",
        "        ),\n",
        "        metrics=[\"accuracy\"],\n",
        "    )\n",
        "    return model\n",
        "\n",
        "\n",
        "def get_stacked_model(encoder, model_params):\n",
        "    \"\"\"\n",
        "    Returns a tf.keras.Model object for the model as challenger\n",
        "        Args:\n",
        "            encoder: A TextVectorization object for the encoder\n",
        "            model_params: A dictionary with model parameters\n",
        "        Returns:\n",
        "            tf.keras.Model: A tf.keras.Model object for the model\n",
        "    \"\"\"\n",
        "\n",
        "    model = tf.keras.Sequential()\n",
        "    model.add(encoder)\n",
        "    model.add(\n",
        "        tf.keras.layers.Embedding(\n",
        "            input_dim=len(encoder.get_vocabulary()), output_dim=64, mask_zero=True\n",
        "        )\n",
        "    )\n",
        "    model.add(\n",
        "        tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(64, return_sequences=True))\n",
        "    )\n",
        "    model.add(tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(32)))\n",
        "    model.add(tf.keras.layers.Dense(64, activation=\"relu\"))\n",
        "    model.add(tf.keras.layers.Dropout(0.5))\n",
        "    model.add(tf.keras.layers.Dense(1))\n",
        "    model.compile(\n",
        "        loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),\n",
        "        optimizer=tf.keras.optimizers.Adam(\n",
        "            learning_rate=model_params[\"learning_rate\"],\n",
        "            beta_1=model_params[\"beta_1\"],\n",
        "            beta_2=model_params[\"beta_2\"],\n",
        "            epsilon=model_params[\"epsilon\"],\n",
        "        ),\n",
        "        metrics=[\"accuracy\"],\n",
        "    )\n",
        "    return model\n",
        "\n",
        "\n",
        "def train(train_dataset, test_dataset, model, epochs, steps, log_dir=LOG_DIR):\n",
        "    \"\"\"\n",
        "    Trains the model\n",
        "    Args:\n",
        "        train_dataset: A tf.data.Dataset object containing the training data\n",
        "        test_dataset: A tf.data.Dataset object containing the testing data\n",
        "        model: A tf.keras.Model object for the model\n",
        "        epochs: The number of epochs\n",
        "        steps: The number of validation steps\n",
        "        log_dir: The location of tf training logs\n",
        "    \"\"\"\n",
        "    tensorboard = tf.keras.callbacks.TensorBoard(log_dir=log_dir)\n",
        "    history = model.fit(\n",
        "        train_dataset,\n",
        "        validation_data=test_dataset,\n",
        "        epochs=epochs,\n",
        "        validation_steps=steps,\n",
        "        callbacks=[tensorboard],\n",
        "    )\n",
        "    return history"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Uz7MSBcwVHgc"
      },
      "source": [
        "### Define a baseline\n",
        "\n",
        "As part of the experimentation phase, you would like to define a baseline model to have a performance benchmark which it is represented by recurrent neural network (RNN) in this case. Then you try to improve the ML solution with more complex training configuration that produce better results."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TM1xGPRkjocQ"
      },
      "source": [
        "#### Run experiment and evaluate experiment runs using `with` statement\n",
        "\n",
        "This step would takes **10 min** approx. to finish. And it covers the following steps:\n",
        "\n",
        "- Initialize an experiment run\n",
        "- Log the parameters associated to training data\n",
        "- Log the parameters of the encoder\n",
        "- Log the parameters of the model\n",
        "- Train the model\n",
        "- Log the metrics for each epochs\n",
        "- Log the overall training metrics"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zRp-gJWhPKf0"
      },
      "outputs": [],
      "source": [
        "# Experiment Settings ----------------------------------------------------------\n",
        "RUN_ID_1 = \"run-1\"\n",
        "BUFFER_SIZE = 10000\n",
        "BATCH_SIZE = 64\n",
        "VOCAB_SIZE = 1000\n",
        "EPOCHS = 2\n",
        "STEPS = 20\n",
        "ROLE = \"baseline\"\n",
        "LR = 1e-4\n",
        "B_1 = 0.9\n",
        "B_2 = 0.999\n",
        "EPS = 1e-07\n",
        "\n",
        "# Initialize the experiment\n",
        "logging.info(\"Initialize the experiment.\")\n",
        "with vertex_ai.start_run(RUN_ID_1) as run:\n",
        "\n",
        "    # Get the training and testing datasets\n",
        "    logging.info(\"Get the training and testing datasets.\")\n",
        "    data_params = {\"buffer_size\": BUFFER_SIZE, \"batch_size\": BATCH_SIZE}\n",
        "    train_dataset, test_dataset = get_dataset(\n",
        "        \"./data\",\n",
        "        buffer_size=data_params[\"buffer_size\"],\n",
        "        batch_size=data_params[\"batch_size\"],\n",
        "    )\n",
        "    run.log_params(data_params)\n",
        "\n",
        "    # Get the encoder\n",
        "    logging.info(\"Get the encoder.\")\n",
        "    encoder_params = {\"vocab_size\": VOCAB_SIZE}\n",
        "    encoder = get_encoder(\n",
        "        train_dataset=train_dataset, vocab_size=encoder_params[\"vocab_size\"]\n",
        "    )\n",
        "    run.log_params(encoder_params)\n",
        "\n",
        "    # Get the model\n",
        "    logging.info(\"Get the model.\")\n",
        "    run.log_params({\"role\": ROLE})\n",
        "    model_params = {\"learning_rate\": LR, \"beta_1\": B_1, \"beta_2\": B_2, \"epsilon\": EPS}\n",
        "    model = get_baseline_model(encoder=encoder, model_params=model_params)\n",
        "    run.log_params(model_params)\n",
        "\n",
        "    # Train the model\n",
        "    logging.info(\"Train the model.\")\n",
        "    history = train(\n",
        "        model=model,\n",
        "        train_dataset=train_dataset,\n",
        "        test_dataset=test_dataset,\n",
        "        epochs=EPOCHS,\n",
        "        steps=STEPS,\n",
        "    )\n",
        "\n",
        "    run.log_params(history.params)\n",
        "    for idx in range(0, history.params[\"epochs\"]):\n",
        "        run.log_time_series_metrics(\n",
        "            {\n",
        "                \"train_loss\": history.history[\"loss\"][idx],\n",
        "                \"train_accuracy\": history.history[\"accuracy\"][idx],\n",
        "            }\n",
        "        )\n",
        "\n",
        "    # Evaluate model\n",
        "    logging.info(\"Evaluate model.\")\n",
        "    test_loss, test_accuracy = model.evaluate(test_dataset)\n",
        "    run.log_metrics({\"test_loss\": test_loss, \"test_accuracy\": test_accuracy})\n",
        "\n",
        "    # Exit the experiment\n",
        "    logging.info(\"Exit the run.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f8bd7ea02120"
      },
      "source": [
        "#### Evaluate the baseline model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xmaumLrSqo1e"
      },
      "outputs": [],
      "source": [
        "# Get experiment\n",
        "logging.info(\"Get experiment status.\")\n",
        "experiment_df = vertex_ai.get_experiment_df()\n",
        "experiment_df.T"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "evxi4ROLqpc_"
      },
      "outputs": [],
      "source": [
        "# Get time series metrics\n",
        "logging.info(\"Get time series metrics.\")\n",
        "ts_runs_df = run.get_time_series_data_frame()\n",
        "ts_runs_df"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eoGfiTobVSSB"
      },
      "source": [
        "### Make experimental improvement to the model\n",
        "\n",
        "As challenger model you add a LSTM layer to the baseline model. This step would takes **10 min** approx. to finish.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZjEww1ESZFTn"
      },
      "source": [
        "#### Run experiment and evaluate experiment runs with `experiment` instance\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kV8DIuUXVXkC"
      },
      "outputs": [],
      "source": [
        "# Experiment Settings ----------------------------------------------------------\n",
        "RUN_ID_2 = \"run-2\"\n",
        "ROLE = \"stacked\"\n",
        "\n",
        "# Initialize the experiment\n",
        "logger.info(\"Initialize the experiment.\")\n",
        "vertex_ai.start_run(RUN_ID_2)\n",
        "\n",
        "# Get the model\n",
        "logging.info(\"Get the model.\")\n",
        "run.log_params({\"role\": ROLE})\n",
        "model = get_stacked_model(encoder=encoder, model_params=model_params)\n",
        "vertex_ai.log_params(model_params)\n",
        "\n",
        "# Train the model\n",
        "logging.info(\"Train the model.\")\n",
        "history = train(\n",
        "    model=model,\n",
        "    train_dataset=train_dataset,\n",
        "    test_dataset=test_dataset,\n",
        "    epochs=EPOCHS,\n",
        "    steps=STEPS,\n",
        ")\n",
        "vertex_ai.log_params(history.params)\n",
        "for idx in range(0, history.params[\"epochs\"]):\n",
        "    vertex_ai.log_time_series_metrics(\n",
        "        {\n",
        "            \"train_loss\": history.history[\"loss\"][idx],\n",
        "            \"train_accuracy\": history.history[\"accuracy\"][idx],\n",
        "        }\n",
        "    )\n",
        "\n",
        "# Evaluate model\n",
        "logging.info(\"Evaluate model.\")\n",
        "test_loss, test_accuracy = model.evaluate(test_dataset)\n",
        "vertex_ai.log_metrics({\"test_loss\": test_loss, \"test_accuracy\": test_accuracy})\n",
        "\n",
        "# Exit the experiment\n",
        "logging.info(\"Exit the experiment.\")\n",
        "vertex_ai.end_run()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ceaf7b22157d"
      },
      "source": [
        "#### Compare the baseline with the new model implementation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bUYoJXexrpmu"
      },
      "outputs": [],
      "source": [
        "# Get experiment\n",
        "logging.info(\"Get experiment status.\")\n",
        "experiment_df = vertex_ai.get_experiment_df()\n",
        "experiment_df.T"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fORBLdvLrqFU"
      },
      "outputs": [],
      "source": [
        "# Get time series metrics\n",
        "exp_run = vertex_ai.ExperimentRun(RUN_ID_2, experiment=EXPERIMENT_NAME)\n",
        "logging.info(\"Get time series metrics.\")\n",
        "ts_runs_df = exp_run.get_time_series_data_frame()\n",
        "ts_runs_df"
      ]
    },
    {
      "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": "uGvJieN5BjPG"
      },
      "outputs": [],
      "source": [
        "# Delete experiment\n",
        "exp = vertex_ai.Experiment(EXPERIMENT_NAME)\n",
        "exp.delete(delete_backing_tensorboard_runs=True)\n",
        "\n",
        "# Delete Tensorboard\n",
        "vertex_ai_tb.delete()\n",
        "\n",
        "# Delete Cloud Storage objects that were created\n",
        "delete_bucket = False\n",
        "\n",
        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
        "    ! gsutil rm -rf {BUCKET_URI}\n",
        "\n",
        "!rm -Rf $DATA_DIR $LOG_DIR"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "comparing_local_trained_models.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
{'fraction_non_alphanumeric': 0.16601779038287998, 'fraction_numerical': 0.007618924842078123, 'mean_word_length': 1.7818181818181817, 'pattern_counts': {'":': 643, '<': 37, '<?xml version=': 0, '>': 41, 'https://': 48, 'lorem ipsum': 0, 'www.': 4, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 14, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3872816', 'n_tokens_mistral': 25308, 'n_tokens_neox': 24684, 'n_words': 5547}