notebooks/official/prediction/custom_batch_prediction_feature_filter.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "copyright"
      },
      "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": "title"
      },
      "source": [
        "# Custom model batch prediction with feature filtering \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/prediction/custom_batch_prediction_feature_filter.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%2Fprediction%2Fcustom_batch_prediction_feature_filter.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>    \n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://github.com/raw/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/prediction/custom_batch_prediction_feature_filter.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><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/prediction/custom_batch_prediction_feature_filter.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mEXYaDoTjmpR"
      },
      "source": [
        "## Overview\n",
        "\n",
        "\n",
        "This tutorial demonstrates how to use the Vertex AI SDK for Python to train a custom tabular classification model and perform batch prediction with feature filtering. This means that you can run batch prediction on a list of selected features or exclude a list of features from prediction.\n",
        "\n",
        "Learn more about [Vertex AI Batch Prediction](https://cloud.google.com/vertex-ai/docs/tabular-data/classification-regression/get-batch-predictions)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "objective:custom,training,online_prediction"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this notebook, you learn how to create a custom-trained model from a Python script in a Docker container using the Vertex AI SDK for Python, and then run a batch prediction job by including or excluding a list of features. \n",
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
        "- BigQuery\n",
        "- Cloud Storage\n",
        "- Vertex AI managed Datasets\n",
        "- Vertex AI Training\n",
        "- Vertex AI BatchPrediction\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Create a Vertex AI custom `TrainingPipeline` for training a model.\n",
        "- Train a TensorFlow model.\n",
        "- Send batch prediction job."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dataset:custom,cifar10,icn"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "The dataset used for this tutorial is the penguins dataset from [BigQuery public datasets](https://cloud.google.com/bigquery/public-data). This dataset has the following fields: `culmen_length_mm`, `culmen_depth_mm`, `flipper_length_mm`, `body_mass_g` from the dataset to predict the penguins species (`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",
        "* BigQuery\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), [BigQuery pricing](https://cloud.google.com/bigquery/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": "3b1ffd5ab768"
      },
      "source": [
        "## Get Started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "install_aip"
      },
      "source": [
        "### Install Vertex AI SDK for Python and other required packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1fd00fa70a2a"
      },
      "outputs": [],
      "source": [
        "# Install the packages\n",
        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
        "                                 google-cloud-storage \\\n",
        "                                 google-cloud-bigquery \\\n",
        "                                 pyarrow \\\n",
        "                                 db-dtypes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ff555b32bab8"
      },
      "source": [
        "### Restart runtime (Colab only)\n",
        "\n",
        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f09b4dff629a"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "54c5ef8a8f43"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f82e28c631cc"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "Authenticate your environment on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "46604f70e831"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6a870411c189"
      },
      "source": [
        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "60af003fe0ed"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bucket:custom"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets.\n",
        "\n",
        "When you submit a training job using the Cloud SDK, you upload a Python package\n",
        "containing your training code to a Cloud Storage bucket. Vertex AI runs\n",
        "the code from this package. In this tutorial, Vertex AI also saves the\n",
        "trained model that results from your job in the same bucket. Using this model artifact, you can then\n",
        "create Vertex AI Model resource and use for prediction."
      ]
    },
    {
      "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": "create_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": "Oz8J0vmSlugt"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "import_aip"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cNEiwLd0lugu"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "\n",
        "import numpy as np\n",
        "from google.cloud import aiplatform, bigquery"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "init_aip:mbsdk,all"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
        "Initialize the Vertex SDK for Python for your project and corresponding bucket.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "856023874e71"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7c163842eabd"
      },
      "source": [
        "### Initialize BigQuery Client\n",
        "\n",
        "Initialize the BigQuery Python client for your project."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fad2ba1ad7c3"
      },
      "outputs": [],
      "source": [
        "# Set up BigQuery client\n",
        "bqclient = bigquery.Client(project=PROJECT_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "container:training,prediction"
      },
      "source": [
        "### Set pre-built containers\n",
        "\n",
        "Vertex AI provides pre-built containers to run training and prediction.\n",
        "\n",
        "For the latest list, see [Pre-built containers for training](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers) and [Pre-built containers for prediction](https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1u1mr18jlugv"
      },
      "outputs": [],
      "source": [
        "TRAIN_VERSION = \"tf-cpu.2-8\"\n",
        "DEPLOY_VERSION = \"tf2-cpu.2-8\"\n",
        "\n",
        "TRAIN_IMAGE = \"us-docker.pkg.dev/vertex-ai/training/{}:latest\".format(TRAIN_VERSION)\n",
        "DEPLOY_IMAGE = \"us-docker.pkg.dev/vertex-ai/prediction/{}:latest\".format(DEPLOY_VERSION)\n",
        "\n",
        "print(\"Training:\", TRAIN_IMAGE)\n",
        "print(\"Deployment:\", DEPLOY_IMAGE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "59f24e7d2269"
      },
      "source": [
        "### Prepare the data\n",
        "\n",
        "To improve the convergence of the custom deep learning model, normalize the data. To prepare for this, calculate the mean and standard deviation for each numeric column.\n",
        "\n",
        "Pass these summary statistics to the training script to normalize the data before training. Later, during prediction, use these summary statistics again to normalize the testing data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8e52b7832cd3"
      },
      "outputs": [],
      "source": [
        "# Calculate mean and std across all rows\n",
        "\n",
        "# Define NA values\n",
        "NA_VALUES = [\"NA\", \".\"]\n",
        "\n",
        "\n",
        "# Download a table\n",
        "def download_table(bq_table_uri: str):\n",
        "    # Remove bq:// prefix if present\n",
        "    prefix = \"bq://\"\n",
        "    if bq_table_uri.startswith(prefix):\n",
        "        bq_table_uri = bq_table_uri[len(prefix) :]\n",
        "\n",
        "    table = bigquery.TableReference.from_string(bq_table_uri)\n",
        "    rows = bqclient.list_rows(\n",
        "        table,\n",
        "    )\n",
        "    return rows.to_dataframe()\n",
        "\n",
        "\n",
        "# Remove NA values\n",
        "def clean_dataframe(df):\n",
        "    return df.replace(to_replace=NA_VALUES, value=np.NaN).dropna()\n",
        "\n",
        "\n",
        "def calculate_mean_and_std(df):\n",
        "    # Calculate mean and std for each applicable column\n",
        "    mean_and_std = {}\n",
        "    dtypes = list(zip(df.dtypes.index, map(str, df.dtypes)))\n",
        "    # Normalize numeric columns.\n",
        "    for column, dtype in dtypes:\n",
        "        if dtype == \"float32\" or dtype == \"float64\":\n",
        "            mean_and_std[column] = {\n",
        "                \"mean\": df[column].mean(),\n",
        "                \"std\": df[column].std(),\n",
        "            }\n",
        "\n",
        "    return mean_and_std"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yi_10xgYjmpW"
      },
      "outputs": [],
      "source": [
        "# Define the BigQuery source dataset\n",
        "BQ_SOURCE = \"bq://bigquery-public-data.ml_datasets.penguins\"\n",
        "\n",
        "dataframe = download_table(BQ_SOURCE)\n",
        "dataframe = clean_dataframe(dataframe)\n",
        "mean_and_std = calculate_mean_and_std(dataframe)\n",
        "print(f\"The mean and stds for each column are: {str(mean_and_std)}\")\n",
        "\n",
        "# Write to a file\n",
        "MEAN_AND_STD_JSON_FILE = \"mean_and_std.json\"\n",
        "\n",
        "with open(MEAN_AND_STD_JSON_FILE, \"w\") as outfile:\n",
        "    json.dump(mean_and_std, outfile)\n",
        "\n",
        "# Save to the staging bucket\n",
        "! gsutil cp {MEAN_AND_STD_JSON_FILE} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5c7732822757"
      },
      "source": [
        "### Create a Vertex AI tabular dataset from BigQuery dataset\n",
        "\n",
        "Your first step in training the model is to create a Vertex AI tabular dataset resource."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1d26b159106f"
      },
      "outputs": [],
      "source": [
        "DATASET_DISPLAY_NAME = \"sample-penguins-unique\"\n",
        "\n",
        "dataset = aiplatform.TabularDataset.create(\n",
        "    display_name=DATASET_DISPLAY_NAME, bq_source=BQ_SOURCE\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "train_custom_model"
      },
      "source": [
        "### Train a model\n",
        "\n",
        "There are two ways you can train a model using a container image:\n",
        "\n",
        "- **Use a Vertex AI pre-built container**. If you use a pre-built training container, you must additionally specify a Python package to install into the container image. This Python package contains your training code.\n",
        "\n",
        "- **Use your own custom container image**. If you use your own container, the container image must contain your training code."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "train_custom_job_args"
      },
      "source": [
        "### Define the command args for the training script\n",
        "\n",
        "Prepare the command-line arguments to pass to your training script.\n",
        "* `args`: The command line arguments to pass to the corresponding Python module. In this example, they are:\n",
        "  * `--epochs`: The number of epochs for training.\n",
        "  * `--batch_size`: The number of batch size for training.\n",
        "  * `--distribute` : The training distribution strategy to use for single or distributed training.\n",
        "     * `\"single\"`: single device.\n",
        "     * `\"mirror\"`: all GPU devices on a single compute instance.\n",
        "     * `\"multi\"`: all GPU devices on all compute instances.\n",
        "  * `--mean_and_std_json_file`: The file on Cloud Storage with pre-calculated means and standard deviations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1npiDcUtlugw"
      },
      "outputs": [],
      "source": [
        "JOB_NAME = \"penquins-custom-job-unique\"\n",
        "EPOCHS = 20\n",
        "BATCH_SIZE = 10\n",
        "TRAIN_STRATEGY = \"single\"\n",
        "\n",
        "CMDARGS = [\n",
        "    \"--epochs=\" + str(EPOCHS),\n",
        "    \"--batch_size=\" + str(BATCH_SIZE),\n",
        "    \"--distribute=\" + TRAIN_STRATEGY,\n",
        "    \"--mean_and_std_json_file=\" + f\"{BUCKET_URI}/{MEAN_AND_STD_JSON_FILE}\",\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "taskpy_contents"
      },
      "source": [
        "### Training script\n",
        "\n",
        "In the next cell, write the contents of the training script, `task.py`. In summary, the script does the following:\n",
        "\n",
        "- Loads the data from the BigQuery table using the BigQuery Python client library.\n",
        "- Loads the pre-calculated mean and standard deviation from the Cloud Storage bucket.\n",
        "- Builds a model using TF.Keras model API.\n",
        "- Compiles the model by calling `compile()`.\n",
        "- Sets a training distribution strategy according to the argument `args.distribute`.\n",
        "- Trains the model by calling `fit()` with epochs and batch size according to the arguments `args.epochs` and `args.batch_size`\n",
        "- Gets the directory where to save the model artifacts from the environment variable `AIP_MODEL_DIR`. This variable is [set by the training service](https://cloud.google.com/vertex-ai/docs/training/code-requirements#environment-variables).\n",
        "- Saves the trained model to the model directory."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "72rUqXNFlugx"
      },
      "outputs": [],
      "source": [
        "%%writefile task.py\n",
        "\n",
        "import argparse\n",
        "import os\n",
        "from typing import Tuple, Optional\n",
        "\n",
        "import pandas as pd\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "\n",
        "from google.cloud import bigquery\n",
        "from google.cloud import storage\n",
        "\n",
        "# Read environmental variables\n",
        "training_data_uri = os.getenv(\"AIP_TRAINING_DATA_URI\")\n",
        "validation_data_uri = os.getenv(\"AIP_VALIDATION_DATA_URI\")\n",
        "test_data_uri = os.getenv(\"AIP_TEST_DATA_URI\")\n",
        "\n",
        "# Read args\n",
        "parser = argparse.ArgumentParser()\n",
        "parser.add_argument('--epochs', dest='epochs',\n",
        "                    default=10, type=int,\n",
        "                    help='Number of epochs.')\n",
        "parser.add_argument('--batch_size', dest='batch_size',\n",
        "                    default=10, type=int,\n",
        "                    help='Batch size.')\n",
        "parser.add_argument('--distribute', dest='distribute', type=str, default='single',\n",
        "                    help='Distributed training strategy.')\n",
        "parser.add_argument('--mean_and_std_json_file', dest='mean_and_std_json_file', type=str,\n",
        "                    help='GCS URI to the JSON file with pre-calculated column means and standard deviations.')\n",
        "args = parser.parse_args()\n",
        "\n",
        "# Set up BigQuery clients\n",
        "bqclient = bigquery.Client()\n",
        "\n",
        "\n",
        "def download_blob(\n",
        "  bucket_name: str, \n",
        "  source_blob_name: str, \n",
        "  destination_file_name: str\n",
        "  ) -> None:\n",
        "    \"\"\"Downloads a blob from the bucket to a local path.\n",
        "    Args:\n",
        "        - bucket_name: \"your-bucket-name\"\n",
        "        - source_blob_name: \"storage-object-name\"\n",
        "        - destination_file_name: \"local/path/to/file\"\n",
        "    \"\"\"\n",
        "\n",
        "    storage_client = storage.Client()\n",
        "    bucket = storage_client.bucket(bucket_name)\n",
        "\n",
        "    # Construct a client side representation of a blob.\n",
        "    # Note `Bucket.blob` differs from `Bucket.get_blob` as it doesn't retrieve\n",
        "    # any content from Cloud Storage. As we don't need additional data,\n",
        "    # using `Bucket.blob` is preferred here.\n",
        "    blob = bucket.blob(source_blob_name)\n",
        "    blob.download_to_filename(destination_file_name)\n",
        "\n",
        "    print(\n",
        "        \"Blob {} downloaded to {}.\".format(\n",
        "            source_blob_name, destination_file_name\n",
        "        )\n",
        "    )\n",
        "\n",
        "def extract_bucket_and_prefix_from_gcs_path(gcs_path: str) -> Tuple[str, Optional[str]]:\n",
        "    \"\"\"Given a complete GCS path, return the bucket name and prefix as a tuple.\n",
        "\n",
        "    Example Usage:\n",
        "\n",
        "        bucket, prefix = extract_bucket_and_prefix_from_gcs_path(\n",
        "            \"gs://example-bucket/path/to/folder\"\n",
        "        )\n",
        "\n",
        "        # bucket = \"example-bucket\"\n",
        "        # prefix = \"path/to/folder\"\n",
        "\n",
        "    Args:\n",
        "        gcs_path (str):\n",
        "            Required. A full path to a Cloud Storage folder or resource.\n",
        "            Can optionally include \"gs://\" prefix or end in a trailing slash \"/\".\n",
        "\n",
        "    Returns:\n",
        "        Tuple[str, Optional[str]]\n",
        "            A (bucket, prefix) pair from provided GCS path. If a prefix is not\n",
        "            present, None is returned in its place.\n",
        "    \"\"\"\n",
        "    if gcs_path.startswith(\"gs://\"):\n",
        "        gcs_path = gcs_path[5:]\n",
        "    if gcs_path.endswith(\"/\"):\n",
        "        gcs_path = gcs_path[:-1]\n",
        "\n",
        "    gcs_parts = gcs_path.split(\"/\", 1)\n",
        "    gcs_bucket = gcs_parts[0]\n",
        "    gcs_blob_prefix = None if len(gcs_parts) == 1 else gcs_parts[1]\n",
        "\n",
        "    return (gcs_bucket, gcs_blob_prefix)\n",
        "\n",
        "\n",
        "# Download means and std\n",
        "def download_mean_and_std(mean_and_std_json_file):\n",
        "    \"\"\"Download mean and std for each column\"\"\"\n",
        "    import json\n",
        "    \n",
        "    bucket, file_path = extract_bucket_and_prefix_from_gcs_path(mean_and_std_json_file)\n",
        "    download_blob(bucket_name=bucket, source_blob_name=file_path, destination_file_name=file_path)\n",
        "    \n",
        "    with open(file_path, 'r') as file:\n",
        "        return json.loads(file.read())\n",
        "\n",
        "        \n",
        "# # Download a table\n",
        "def download_table(bq_table_uri: str):\n",
        "    # Remove bq:// prefix if present\n",
        "    prefix = \"bq://\"\n",
        "    if bq_table_uri.startswith(prefix):\n",
        "        bq_table_uri = bq_table_uri[len(prefix):]\n",
        "\n",
        "    table = bigquery.TableReference.from_string(bq_table_uri)\n",
        "    rows = bqclient.list_rows(table)\n",
        "    \n",
        "    return rows.to_dataframe(create_bqstorage_client=False)\n",
        "\n",
        "\n",
        "def standardize(df, mean_and_std):\n",
        "    \"\"\"Scales numerical columns using their means and standard deviation to get\n",
        "    z-scores: the mean of each numerical column becomes 0, and the standard\n",
        "    deviation becomes 1. This can help the model converge during training.\n",
        "\n",
        "    Args:\n",
        "      df: Pandas df\n",
        "\n",
        "    Returns:\n",
        "      Input df with the numerical columns scaled to z-scores\n",
        "    \"\"\"\n",
        "    dtypes = list(zip(df.dtypes.index, map(str, df.dtypes)))\n",
        "    # Normalize numeric columns.\n",
        "    for column, dtype in dtypes:\n",
        "        if dtype == \"float32\":\n",
        "            df[column] -= mean_and_std[column][\"mean\"]\n",
        "            df[column] /= mean_and_std[column][\"std\"]\n",
        "    return df\n",
        "\n",
        "\n",
        "def preprocess(df):\n",
        "    \"\"\"Converts categorical features to numeric. Removes unused columns.\n",
        "\n",
        "    Args:\n",
        "      df: Pandas df with raw data\n",
        "\n",
        "    Returns:\n",
        "      df with preprocessed data\n",
        "    \"\"\"\n",
        "    df = df.drop(columns=UNUSED_COLUMNS)\n",
        "\n",
        "    # Drop rows with NaN's\n",
        "    df = df.dropna()\n",
        "\n",
        "    # Convert integer valued (numeric) columns to floating point\n",
        "    numeric_columns = df.select_dtypes([\"int32\", \"float32\", \"float64\"]).columns\n",
        "    df[numeric_columns] = df[numeric_columns].astype(\"float32\")\n",
        "\n",
        "    # Convert categorical columns to numeric\n",
        "    cat_columns = df.select_dtypes([\"object\"]).columns\n",
        "\n",
        "    df[cat_columns] = df[cat_columns].apply(\n",
        "        lambda x: x.astype(_CATEGORICAL_TYPES[x.name])\n",
        "    )\n",
        "    df[cat_columns] = df[cat_columns].apply(lambda x: x.cat.codes)\n",
        "    return df\n",
        "\n",
        "\n",
        "def convert_dataframe_to_dataset(\n",
        "    df_train,\n",
        "    df_validation,\n",
        "    mean_and_std\n",
        "):\n",
        "    df_train = preprocess(df_train)\n",
        "    df_validation = preprocess(df_validation)\n",
        "\n",
        "    df_train_x, df_train_y = df_train, df_train.pop(LABEL_COLUMN)\n",
        "    df_validation_x, df_validation_y = df_validation, df_validation.pop(LABEL_COLUMN)\n",
        "\n",
        "    # Join train_x and eval_x to normalize on overall means and standard\n",
        "    # deviations. Then separate them again.\n",
        "    all_x = pd.concat([df_train_x, df_validation_x], keys=[\"train\", \"eval\"])\n",
        "    all_x = standardize(all_x, mean_and_std)\n",
        "    df_train_x, df_validation_x = all_x.xs(\"train\"), all_x.xs(\"eval\")\n",
        "\n",
        "    y_train = np.asarray(df_train_y).astype(\"float32\")\n",
        "    y_validation = np.asarray(df_validation_y).astype(\"float32\")\n",
        "\n",
        "    # Convert to numpy representation\n",
        "    x_train = np.asarray(df_train_x)\n",
        "    x_test = np.asarray(df_validation_x)\n",
        "\n",
        "    # Convert to one-hot representation\n",
        "    y_train = tf.keras.utils.to_categorical(y_train, num_classes=len(SPECIES))\n",
        "    y_validation = tf.keras.utils.to_categorical(y_validation, num_classes=len(SPECIES))\n",
        "\n",
        "    dataset_train = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
        "    dataset_validation = tf.data.Dataset.from_tensor_slices((x_test, y_validation))\n",
        "    return (dataset_train, dataset_validation)\n",
        "\n",
        "\n",
        "# Remove NA values\n",
        "def clean_dataframe(df):\n",
        "    return df.replace(to_replace=NA_VALUES, value=np.NaN).dropna()\n",
        "\n",
        "\n",
        "def create_model(num_features):\n",
        "    # Create model\n",
        "    Dense = tf.keras.layers.Dense\n",
        "    model = tf.keras.Sequential(\n",
        "        [\n",
        "            Dense(\n",
        "                100,\n",
        "                activation=tf.nn.relu,\n",
        "                kernel_initializer=\"uniform\",\n",
        "                input_dim=num_features,\n",
        "            ),\n",
        "            Dense(75, activation=tf.nn.relu),\n",
        "            Dense(50, activation=tf.nn.relu),\n",
        "            Dense(25, activation=tf.nn.relu),\n",
        "            Dense(3, activation=tf.nn.softmax),\n",
        "        ]\n",
        "    )\n",
        "    \n",
        "    # Compile Keras model\n",
        "    optimizer = tf.keras.optimizers.RMSprop(lr=0.001)\n",
        "    model.compile(\n",
        "        loss=\"categorical_crossentropy\", metrics=[\"accuracy\"], optimizer=optimizer\n",
        "    )\n",
        "    \n",
        "    return model\n",
        "\n",
        "\n",
        "mean_and_std = download_mean_and_std(args.mean_and_std_json_file)\n",
        "\n",
        "# Single Machine, single compute device\n",
        "if args.distribute == 'single':\n",
        "    if tf.test.is_gpu_available():\n",
        "        strategy = tf.distribute.OneDeviceStrategy(device=\"/gpu:0\")\n",
        "    else:\n",
        "        strategy = tf.distribute.OneDeviceStrategy(device=\"/cpu:0\")\n",
        "# Single Machine, multiple compute device\n",
        "elif args.distribute == 'mirror':\n",
        "    strategy = tf.distribute.MirroredStrategy()\n",
        "# Multiple Machine, multiple compute device\n",
        "elif args.distribute == 'multi':\n",
        "    strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()\n",
        "\n",
        "# Set up training variables\n",
        "LABEL_COLUMN = \"species\"\n",
        "UNUSED_COLUMNS = []\n",
        "NA_VALUES = [\"NA\", \".\"]\n",
        "\n",
        "# Possible categorical values\n",
        "SPECIES = ['Adelie Penguin (Pygoscelis adeliae)',\n",
        "           'Chinstrap penguin (Pygoscelis antarctica)',\n",
        "           'Gentoo penguin (Pygoscelis papua)']\n",
        "ISLANDS = ['Dream', 'Biscoe', 'Torgersen']\n",
        "SEXES = ['FEMALE', 'MALE']\n",
        "\n",
        "df_train = download_table(training_data_uri)\n",
        "df_validation = download_table(validation_data_uri)\n",
        "df_test = download_table(test_data_uri)\n",
        "\n",
        "df_train = clean_dataframe(df_train)\n",
        "df_validation = clean_dataframe(df_validation)\n",
        "\n",
        "_CATEGORICAL_TYPES = {\n",
        "    \"island\": pd.api.types.CategoricalDtype(categories=ISLANDS),\n",
        "    \"species\": pd.api.types.CategoricalDtype(categories=SPECIES),\n",
        "    \"sex\": pd.api.types.CategoricalDtype(categories=SEXES),\n",
        "}\n",
        "\n",
        "# Create datasets\n",
        "dataset_train, dataset_validation = convert_dataframe_to_dataset(\n",
        "  df_train, \n",
        "  df_validation, \n",
        "  mean_and_std\n",
        ")\n",
        "\n",
        "# Shuffle train set\n",
        "dataset_train = dataset_train.shuffle(len(df_train))\n",
        "\n",
        "# Create the model\n",
        "with strategy.scope():\n",
        "    model = create_model(num_features=dataset_train._flat_shapes[0].dims[0].value)\n",
        "\n",
        "# Set up datasets\n",
        "NUM_WORKERS = strategy.num_replicas_in_sync\n",
        "# Here the batch size scales up by number of workers since\n",
        "# `tf.data.Dataset.batch` expects the global batch size.\n",
        "GLOBAL_BATCH_SIZE = args.batch_size * NUM_WORKERS\n",
        "dataset_train = dataset_train.batch(GLOBAL_BATCH_SIZE)\n",
        "dataset_validation = dataset_validation.batch(GLOBAL_BATCH_SIZE)\n",
        "\n",
        "# Train the model\n",
        "model.fit(dataset_train, epochs=args.epochs, validation_data=dataset_validation)\n",
        "\n",
        "tf.saved_model.save(model, os.getenv(\"AIP_MODEL_DIR\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "train_custom_job"
      },
      "source": [
        "### Train the model\n",
        "\n",
        "Define your custom `TrainingPipeline` on Vertex AI.\n",
        "\n",
        "Use the `CustomTrainingJob` class to define the `TrainingPipeline`. The class takes the following parameters:\n",
        "\n",
        "- `display_name`: The user-defined name of this training pipeline.\n",
        "- `script_path`: The local path to the training script.\n",
        "- `container_uri`: The URI of the training container image.\n",
        "- `requirements`: The list of Python package dependencies of the script.\n",
        "- `model_serving_container_image_uri`: The URI of a container that can serve predictions for your model — either a pre-built container or a custom container.\n",
        "\n",
        "Use the `run` function to start training. The function takes the following parameters:\n",
        "\n",
        "- `dataset`: Vertex AI Dataset to fit this training against.\n",
        "- `model_display_name`: The display name of the `Model` if the script produces a managed `Model`.\n",
        "- `bigquery_destination`: The BigQuery project location where the training data is to be written to.\n",
        "- `args`: The command line arguments to be passed to the Python script.\n",
        "\n",
        "The `run` function creates a training pipeline that trains and creates a `Model` object. After the training pipeline completes, the `run` function returns the `Model` object."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mxIxvDdglugx"
      },
      "outputs": [],
      "source": [
        "job = aiplatform.CustomTrainingJob(\n",
        "    display_name=JOB_NAME,\n",
        "    script_path=\"task.py\",\n",
        "    container_uri=TRAIN_IMAGE,\n",
        "    requirements=[\"google-cloud-bigquery>=2.20.0\", \"db-dtypes\", \"protobuf==3.20.3\"],\n",
        "    model_serving_container_image_uri=DEPLOY_IMAGE,\n",
        ")\n",
        "\n",
        "MODEL_DISPLAY_NAME = \"penguins-unique\"\n",
        "\n",
        "# Start the training\n",
        "model = job.run(\n",
        "    dataset=dataset,\n",
        "    model_display_name=MODEL_DISPLAY_NAME,\n",
        "    bigquery_destination=f\"bq://{PROJECT_ID}\",\n",
        "    args=CMDARGS,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "make_prediction"
      },
      "source": [
        "### Send Batch Prediction job request with feature filtering (instanceConfig field)\n",
        "\n",
        "Now that the model is ready, you can send batch prediction request directly from the model resource without needing to deploy the model to an endpoint. \n",
        "\n",
        "Sometimes, your input data does not match the data format that the predictor accepts. Feature filtering lets you either exclude certain fields (such as identifiers or metadata) that are in the input data from your prediction request, or include only a subset of fields from the input data in your prediction request, without having to do any custom pre/post-processing in the prediction container.\n",
        "You can filter and/or transform your batch input \n",
        "\n",
        "In this notebook you learn how to send batch prediction request by including or excluding a list of features by specifying `instanceConfig` in your `BatchPredictionJob` request (**v1beta1 only**).\n",
        "\n",
        "Learn more about [Prediction on Vertex AI](https://cloud.google.com/vertex-ai/docs/predictions/overview)<br>\n",
        "Learn more about [feature filtering](https://cloud.google.com/vertex-ai/docs/predictions/get-predictions#filter_and_transform_input_data_preview)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "get_test_item:test"
      },
      "source": [
        "### Prepare test data\n",
        "\n",
        "Prepare test data by normalizing it and converting categorical values to numeric values.\n",
        "You must normalize these values in the same way that your normalized training data.\n",
        "\n",
        "In this example, we add an extra column called `id` to the test dataset which was not used for training. We show how to exclude this feature at prediction. \n",
        "Here, you perform testing with the same dataset that you used for training. In practice, you generally want to use a separate test dataset to verify your results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e3a2449cfcf1"
      },
      "outputs": [],
      "source": [
        "import pandas as pd\n",
        "from google.cloud import bigquery\n",
        "\n",
        "UNUSED_COLUMNS = []\n",
        "LABEL_COLUMN = \"species\"\n",
        "\n",
        "# Possible categorical values\n",
        "SPECIES = [\n",
        "    \"Adelie Penguin (Pygoscelis adeliae)\",\n",
        "    \"Chinstrap penguin (Pygoscelis antarctica)\",\n",
        "    \"Gentoo penguin (Pygoscelis papua)\",\n",
        "]\n",
        "ISLANDS = [\"Dream\", \"Biscoe\", \"Torgersen\"]\n",
        "SEXES = [\"FEMALE\", \"MALE\"]\n",
        "\n",
        "_CATEGORICAL_TYPES = {\n",
        "    \"island\": pd.api.types.CategoricalDtype(categories=ISLANDS),\n",
        "    \"species\": pd.api.types.CategoricalDtype(categories=SPECIES),\n",
        "    \"sex\": pd.api.types.CategoricalDtype(categories=SEXES),\n",
        "}\n",
        "\n",
        "\n",
        "def standardize(df, mean_and_std):\n",
        "    \"\"\"Scales numerical columns using their means and standard deviation to get\n",
        "    z-scores: the mean of each numerical column becomes 0, and the standard\n",
        "    deviation becomes 1. This can help the model converge during training.\n",
        "\n",
        "    Args:\n",
        "      df: Pandas df\n",
        "\n",
        "    Returns:\n",
        "      Input df with the numerical columns scaled to z-scores\n",
        "    \"\"\"\n",
        "    dtypes = list(zip(df.dtypes.index, map(str, df.dtypes)))\n",
        "    # Normalize numeric columns.\n",
        "    for column, dtype in dtypes:\n",
        "        if dtype == \"float32\":\n",
        "            df[column] -= mean_and_std[column][\"mean\"]\n",
        "            df[column] /= mean_and_std[column][\"std\"]\n",
        "    return df\n",
        "\n",
        "\n",
        "def preprocess(df, mean_and_std):\n",
        "    \"\"\"Converts categorical features to numeric. Removes unused columns.\n",
        "\n",
        "    Args:\n",
        "      df: Pandas df with raw data\n",
        "\n",
        "    Returns:\n",
        "      df with preprocessed data\n",
        "    \"\"\"\n",
        "    df = df.drop(columns=UNUSED_COLUMNS)\n",
        "\n",
        "    # Drop rows with NaN's\n",
        "    df = df.dropna()\n",
        "\n",
        "    # Convert integer valued (numeric) columns to floating point\n",
        "    numeric_columns = df.select_dtypes([\"int32\", \"float32\", \"float64\"]).columns\n",
        "    df[numeric_columns] = df[numeric_columns].astype(\"float32\")\n",
        "\n",
        "    # Convert categorical columns to numeric\n",
        "    cat_columns = df.select_dtypes([\"object\"]).columns\n",
        "\n",
        "    df[cat_columns] = df[cat_columns].apply(\n",
        "        lambda x: x.astype(_CATEGORICAL_TYPES[x.name])\n",
        "    )\n",
        "    df[cat_columns] = df[cat_columns].apply(lambda x: x.cat.codes)\n",
        "    return df\n",
        "\n",
        "\n",
        "def convert_dataframe_to_list(df, mean_and_std):\n",
        "    df = preprocess(df, mean_and_std)\n",
        "\n",
        "    df_x, df_y = df, df.pop(LABEL_COLUMN)\n",
        "\n",
        "    # Normalize on overall means and standard deviations.\n",
        "    df = standardize(df, mean_and_std)\n",
        "\n",
        "    y = np.asarray(df_y).astype(\"float32\")\n",
        "\n",
        "    # Convert to numpy representation\n",
        "    x = np.asarray(df_x)\n",
        "\n",
        "    # Convert to one-hot representation\n",
        "    return x.tolist(), y.tolist(), df_x\n",
        "\n",
        "\n",
        "x_test, y_test, df_x = convert_dataframe_to_list(dataframe, mean_and_std)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VO2klROGjmpX"
      },
      "outputs": [],
      "source": [
        "# Add id column to the test dataframe\n",
        "ID_COLUMN_NAME = \"id\"\n",
        "df_x_with_id = df_x.copy()\n",
        "df_x_with_id[ID_COLUMN_NAME] = [i for i in range(0, df_x_with_id.shape[0])]\n",
        "\n",
        "# Print columns of the datafram\n",
        "print(f\"Test dataset columns: {df_x_with_id.columns}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PKyddetDjmpX"
      },
      "source": [
        "### Upload the test DataFrame to BigQuery "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PCXSaQ2OjmpX"
      },
      "outputs": [],
      "source": [
        "def save_dataframe_to_bigquery(\n",
        "    dataframe: pd.DataFrame, dataset_name: str, table_name: str\n",
        ") -> str:\n",
        "    \"\"\"This function loads a dataframe to a new bigquery table\n",
        "\n",
        "    Args:\n",
        "        dataframe (pd.Dataframe): dataframe to be loaded to bigquery\n",
        "        dataset_name (str): name of the BigQuery dataset for storing the data\n",
        "        table_name (str): name of the BigQuery table that is being created\n",
        "\n",
        "    Returns:\n",
        "        str: table id of the destination bigquery table\n",
        "    \"\"\"\n",
        "    client = bigquery.Client(PROJECT_ID)\n",
        "\n",
        "    bq_dataset = bigquery.Dataset(f\"{PROJECT_ID}.{dataset_name}\")\n",
        "    bq_dataset = client.create_dataset(bq_dataset, exists_ok=True)\n",
        "\n",
        "    job_config = bigquery.LoadJobConfig(\n",
        "        # Optionally, set the write disposition. BigQuery appends loaded rows\n",
        "        # to an existing table by default, but with WRITE_TRUNCATE write\n",
        "        # disposition it replaces the table with the loaded data.\n",
        "        write_disposition=\"WRITE_TRUNCATE\",\n",
        "    )\n",
        "\n",
        "    # Reference: https://cloud.google.com/bigquery/docs/samples/bigquery-load-table-dataframe\n",
        "    job = client.load_table_from_dataframe(\n",
        "        dataframe=dataframe,\n",
        "        destination=f\"{PROJECT_ID}.{dataset_name}.{table_name}\",\n",
        "        job_config=job_config,\n",
        "    )\n",
        "\n",
        "    job.result()\n",
        "\n",
        "    return str(job.destination)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cAluSCrsjmpX"
      },
      "outputs": [],
      "source": [
        "# Upload the Dataframe to a BigQuery table\n",
        "\n",
        "DATASET_NAME = \"test_dataset\"\n",
        "TABLE_NAME = \"test-data-unique\"\n",
        "\n",
        "TABLE_ID = save_dataframe_to_bigquery(\n",
        "    dataframe=df_x_with_id, dataset_name=DATASET_NAME, table_name=TABLE_NAME\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "send_prediction_request:image"
      },
      "source": [
        "### Send the BatchPredictionJob request using REST API\n",
        "\n",
        "Now that you have test data, you can use it to send a batch prediction request using REST API. To do that you need to create a `JSON` request with the following information:\n",
        "\n",
        "- `BATCH_JOB_NAME`: Display name for the batch prediction job.\n",
        "- `MODEL_URI`: The URI for the Model resource to use for making predictions.\n",
        "- `INPUT_FORMAT`: The format of your input data: bigquery, jsonl, csv, tf-record, tf-record-gzip, or file-list.\n",
        "- `INPUT_URI`: Cloud Storage URI of your input data. May contain wildcards.\n",
        "- `OUTPUT_URI`: Cloud Storage URI of a directory where you want Vertex AI to save output.\n",
        "- `MACHINE_TYPE`: The machine resources to be used for this batch prediction job.\n",
        "\n",
        "In this example, we create two versions of the same JSON request: one with `excludedFields` and the other with `includeFields` to show how to include or exclude certain features. Note that these two requests do the same job in this example!\n",
        "\n",
        "Learn more about [request a batch prediction](https://cloud.google.com/vertex-ai/docs/predictions/get-predictions#api_1)<br>\n",
        "Learn more about [instanceconfig](https://cloud.google.com/vertex-ai/docs/reference/rest/v1beta1/projects.locations.batchPredictionJobs#instanceconfig)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ofwR2WTcjmpY"
      },
      "outputs": [],
      "source": [
        "BATCH_JOB_NAME = \"penguins-test\"\n",
        "MODEL_URI = model.resource_name\n",
        "INPUT_FORMAT = \"bigquery\"\n",
        "INPUT_URI = f\"bq://{TABLE_ID}\"\n",
        "OUTPUT_FORMAT = \"bigquery\"\n",
        "OUTPUT_URI = f\"bq://{PROJECT_ID}\"\n",
        "MACHINE_TYPE = \"n1-standard-2\"\n",
        "EXCLUDED_FIELDS = [ID_COLUMN_NAME]\n",
        "\n",
        "# Create a list of columns to be included\n",
        "ALL_COLUMNS = list(df_x_with_id.columns)\n",
        "INCLUDED_FIELDS = ALL_COLUMNS.copy()\n",
        "INCLUDED_FIELDS.remove(ID_COLUMN_NAME)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YJdI9L5fjmpY"
      },
      "source": [
        "### Create JSON body requests"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hfVWSzwkjmpY"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "\n",
        "request_with_excluded_fields = {\n",
        "    \"displayName\": f\"{BATCH_JOB_NAME}-excluded_fields\",\n",
        "    \"model\": MODEL_URI,\n",
        "    \"inputConfig\": {\n",
        "        \"instancesFormat\": INPUT_FORMAT,\n",
        "        \"bigquerySource\": {\"inputUri\": INPUT_URI},\n",
        "    },\n",
        "    \"outputConfig\": {\n",
        "        \"predictionsFormat\": OUTPUT_FORMAT,\n",
        "        \"bigqueryDestination\": {\"outputUri\": OUTPUT_URI},\n",
        "    },\n",
        "    \"dedicatedResources\": {\n",
        "        \"machineSpec\": {\n",
        "            \"machineType\": MACHINE_TYPE,\n",
        "        }\n",
        "    },\n",
        "    \"instanceConfig\": {\"excludedFields\": EXCLUDED_FIELDS},\n",
        "}\n",
        "\n",
        "with open(\"request_with_excluded_fields.json\", \"w\") as outfile:\n",
        "    json.dump(request_with_excluded_fields, outfile)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e95Rk92fjmpZ"
      },
      "outputs": [],
      "source": [
        "request_with_included_fields = {\n",
        "    \"displayName\": f\"{BATCH_JOB_NAME}-included_fields\",\n",
        "    \"model\": MODEL_URI,\n",
        "    \"inputConfig\": {\n",
        "        \"instancesFormat\": INPUT_FORMAT,\n",
        "        \"bigquerySource\": {\"inputUri\": INPUT_URI},\n",
        "    },\n",
        "    \"outputConfig\": {\n",
        "        \"predictionsFormat\": OUTPUT_FORMAT,\n",
        "        \"bigqueryDestination\": {\"outputUri\": OUTPUT_URI},\n",
        "    },\n",
        "    \"dedicatedResources\": {\n",
        "        \"machineSpec\": {\n",
        "            \"machineType\": MACHINE_TYPE,\n",
        "        }\n",
        "    },\n",
        "    \"instanceConfig\": {\"includedFields\": INCLUDED_FIELDS},\n",
        "}\n",
        "\n",
        "with open(\"request_with_included_fields.json\", \"w\") as outfile:\n",
        "    json.dump(request_with_included_fields, outfile)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jPIYouoVjmpZ"
      },
      "source": [
        "### Send the requests\n",
        "\n",
        "To send the requests, specify the API version you want to use. In this case you use `v1beta1` to be able to use `instanceConfig`.\n",
        "\n",
        "#### Exclude fields\n",
        "\n",
        "Here, we send the request with `excludedFields`. After running the follwing cell you should receive a JSON response with your provided information. Then wait for the job to complete (you can check your job status on your Vertex AI Batch Predictions menu or use the Python SDK)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ykduvrX3jmpZ"
      },
      "outputs": [],
      "source": [
        "! curl \\\n",
        "  -X POST \\\n",
        "  -H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
        "  -H \"Content-Type: application/json\" \\\n",
        "  -d @request_with_excluded_fields.json \\\n",
        "  https://{REGION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{REGION}/batchPredictionJobs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QPK1jQRejmpZ"
      },
      "source": [
        "#### Include fields\n",
        "\n",
        "Here, we send the request with `includedFields`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1vXC7fNBjmpZ"
      },
      "outputs": [],
      "source": [
        "! curl \\\n",
        "  -X POST \\\n",
        "  -H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
        "  -H \"Content-Type: application/json\" \\\n",
        "  -d @request_with_included_fields.json \\\n",
        "  https://{REGION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{REGION}/batchPredictionJobs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cleanup:custom"
      },
      "source": [
        "### Cleaning up\n",
        "\n",
        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
        "\n",
        "Otherwise, you can delete the individual resources you created in this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NNmebHf7lug0"
      },
      "outputs": [],
      "source": [
        "# Warning: Setting this to true deletes everything in your bucket\n",
        "delete_bucket = True\n",
        "\n",
        "# Delete the training job\n",
        "job.delete()\n",
        "\n",
        "# Delete the dataset\n",
        "dataset.delete()\n",
        "\n",
        "# Delete the model\n",
        "model.delete()\n",
        "\n",
        "if delete_bucket:\n",
        "    ! gsutil rm -r $BUCKET_URI\n",
        "\n",
        "# Delete the created BigQuery dataset\n",
        "! bq rm -r -f $PROJECT_ID:$DATASET_NAME"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "custom_batch_prediction_feature_filter.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}