two_tower_model_introduction.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "two-tower-model-introduction.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "source": [
        "# Copyright 2021 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."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "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/tree/master/notebooks/official/matching_engine/two-tower-model-introduction.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/tree/master/notebooks/official/matching_engine/two-tower-model-introduction.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial demonstrates how to use the Two-Tower built-in algorithm on the Vertex AI platform.\n",
        "\n",
        "Two-tower models learn to represent two items of various types (such as user profiles, search queries, web documents, answer passages, or images) in the same vector space, so that similar or related items are close to each other. These two items are referred to as the query and candidate object, since when paired with a nearest neighbor search service such as Vertex Matching Engine, the two-tower model can retrieve candidate objects related to an input query object. These objects are encoded by a query and candidate encoder (the two \"towers\") respectively, which are trained on pairs of relevant items. This built-in algorithm exports trained query and candidate encoders as model artifacts, which can be deployed in Vertex Prediction for usage in a recommendation system.\n",
        "\n",
        "### Dataset\n",
        "\n",
        "This tutorial uses the `movielens_100k sample dataset` in the public bucket `gs://cloud-samples-data/vertex-ai/matching-engine/two-tower`, which was generated from the [MovieLens movie rating dataset](https://grouplens.org/datasets/movielens/100k/). For simplicity, the data for this tutorial only includes the user id feature for users, and the movie id and movie title features for movies. In this example, the user is the query object and the movie is the candidate object, and each training example in the dataset contains a user and a movie they rated (we only include positive ratings in the dataset). The two-tower model will embed the user and the movie in the same embedding space, so that given a user, the model will recommend movies it thinks the user will like.\n",
        "\n",
        "### Objective\n",
        "\n",
        "In this notebook, you will learn how to run the two-tower model.\n",
        "The tutorial covers the following steps:\n",
        "1. **Setup**: Importing the required libraries and setting your global variables.\n",
        "2. **Configure parameters**: Setting the appropriate parameter values for the training job.\n",
        "3. **Train on Vertex Training**: Submitting a training job.\n",
        "4. **Deploy on Vertex Prediction**: Importing and deploying the trained model to a callable endpoint.\n",
        "5. **Predict**: Calling the deployed endpoint using online or batch prediction.\n",
        "6. **Hyperparameter tuning**: Running a hyperparameter tuning job.\n",
        "7. **Cleaning up**: Deleting resources created by this tutorial.\n",
        "\n",
        "\n",
        "### Costs \n",
        "\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\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 Google Cloud 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"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2b4ef9b72d43"
      },
      "source": [
        "import os\n",
        "\n",
        "# The Google Cloud Notebook product has specific requirements\n",
        "IS_GOOGLE_CLOUD_NOTEBOOK = os.path.exists(\"/opt/deeplearning/metadata/env_version\")\n",
        "\n",
        "# Google Cloud Notebook requires dependencies to be installed with '--user'\n",
        "USER_FLAG = \"\"\n",
        "if IS_GOOGLE_CLOUD_NOTEBOOK:\n",
        "    USER_FLAG = \"--user\""
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wyy5Lbnzg5fi",
        "outputId": "1f9f25e7-fb97-45c6-988e-718b1990b7df",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "source": [
        "! pip3 install {USER_FLAG} --upgrade tensorflow\n",
        "! pip3 install {USER_FLAG} --upgrade google-cloud-aiplatform tensorboard-plugin-profile\n",
        "! gcloud components update --quiet"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: tensorflow in /usr/local/lib/python3.7/dist-packages (2.6.0)\n",
            "Requirement already satisfied: termcolor~=1.1.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (1.1.0)\n",
            "Requirement already satisfied: wheel~=0.35 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (0.37.0)\n",
            "Requirement already satisfied: numpy~=1.19.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (1.19.5)\n",
            "Requirement already satisfied: grpcio<2.0,>=1.37.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (1.40.0)\n",
            "Requirement already satisfied: h5py~=3.1.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (3.1.0)\n",
            "Requirement already satisfied: opt-einsum~=3.3.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (3.3.0)\n",
            "Requirement already satisfied: keras-preprocessing~=1.1.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (1.1.2)\n",
            "Requirement already satisfied: wrapt~=1.12.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (1.12.1)\n",
            "Requirement already satisfied: typing-extensions~=3.7.4 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (3.7.4.3)\n",
            "Requirement already satisfied: astunparse~=1.6.3 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (1.6.3)\n",
            "Requirement already satisfied: keras~=2.6 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (2.6.0)\n",
            "Requirement already satisfied: flatbuffers~=1.12.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (1.12)\n",
            "Requirement already satisfied: tensorboard~=2.6 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (2.6.0)\n",
            "Requirement already satisfied: absl-py~=0.10 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (0.12.0)\n",
            "Requirement already satisfied: google-pasta~=0.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (0.2.0)\n",
            "Requirement already satisfied: tensorflow-estimator~=2.6 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (2.6.0)\n",
            "Requirement already satisfied: clang~=5.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (5.0)\n",
            "Requirement already satisfied: six~=1.15.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (1.15.0)\n",
            "Requirement already satisfied: gast==0.4.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (0.4.0)\n",
            "Requirement already satisfied: protobuf>=3.9.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (3.17.3)\n",
            "Requirement already satisfied: cached-property in /usr/local/lib/python3.7/dist-packages (from h5py~=3.1.0->tensorflow) (1.5.2)\n",
            "Requirement already satisfied: requests<3,>=2.21.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.6->tensorflow) (2.23.0)\n",
            "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.6->tensorflow) (1.8.0)\n",
            "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.6->tensorflow) (0.4.6)\n",
            "Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.6->tensorflow) (0.6.1)\n",
            "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.6->tensorflow) (3.3.4)\n",
            "Requirement already satisfied: setuptools>=41.0.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.6->tensorflow) (57.4.0)\n",
            "Requirement already satisfied: werkzeug>=0.11.15 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.6->tensorflow) (1.0.1)\n",
            "Requirement already satisfied: google-auth<2,>=1.6.3 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.6->tensorflow) (1.35.0)\n",
            "Requirement already satisfied: cachetools<5.0,>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from google-auth<2,>=1.6.3->tensorboard~=2.6->tensorflow) (4.2.2)\n",
            "Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.7/dist-packages (from google-auth<2,>=1.6.3->tensorboard~=2.6->tensorflow) (4.7.2)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.7/dist-packages (from google-auth<2,>=1.6.3->tensorboard~=2.6->tensorflow) (0.2.8)\n",
            "Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.7/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard~=2.6->tensorflow) (1.3.0)\n",
            "Requirement already satisfied: importlib-metadata in /usr/local/lib/python3.7/dist-packages (from markdown>=2.6.8->tensorboard~=2.6->tensorflow) (4.8.1)\n",
            "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.7/dist-packages (from pyasn1-modules>=0.2.1->google-auth<2,>=1.6.3->tensorboard~=2.6->tensorflow) (0.4.8)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard~=2.6->tensorflow) (3.0.4)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard~=2.6->tensorflow) (2.10)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard~=2.6->tensorflow) (1.24.3)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard~=2.6->tensorflow) (2021.5.30)\n",
            "Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.7/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard~=2.6->tensorflow) (3.1.1)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata->markdown>=2.6.8->tensorboard~=2.6->tensorflow) (3.5.0)\n",
            "Requirement already satisfied: google-cloud-aiplatform in /usr/local/lib/python3.7/dist-packages (1.4.3)\n",
            "Requirement already satisfied: tensorboard-plugin-profile in /usr/local/lib/python3.7/dist-packages (2.5.0)\n",
            "Requirement already satisfied: proto-plus>=1.10.1 in /usr/local/lib/python3.7/dist-packages (from google-cloud-aiplatform) (1.19.0)\n",
            "Requirement already satisfied: google-api-core[grpc]<3.0.0dev,>=1.26.0 in /usr/local/lib/python3.7/dist-packages (from google-cloud-aiplatform) (1.26.3)\n",
            "Requirement already satisfied: google-cloud-storage<2.0.0dev,>=1.32.0 in /usr/local/lib/python3.7/dist-packages (from google-cloud-aiplatform) (1.41.1)\n",
            "Requirement already satisfied: google-cloud-bigquery<3.0.0dev,>=1.15.0 in /usr/local/lib/python3.7/dist-packages (from google-cloud-aiplatform) (2.18.0)\n",
            "Requirement already satisfied: packaging>=14.3 in /usr/local/lib/python3.7/dist-packages (from google-cloud-aiplatform) (21.0)\n",
            "Requirement already satisfied: protobuf>=3.12.0 in /usr/local/lib/python3.7/dist-packages (from google-api-core[grpc]<3.0.0dev,>=1.26.0->google-cloud-aiplatform) (3.17.3)\n",
            "Requirement already satisfied: six>=1.13.0 in /usr/local/lib/python3.7/dist-packages (from google-api-core[grpc]<3.0.0dev,>=1.26.0->google-cloud-aiplatform) (1.15.0)\n",
            "Requirement already satisfied: pytz in /usr/local/lib/python3.7/dist-packages (from google-api-core[grpc]<3.0.0dev,>=1.26.0->google-cloud-aiplatform) (2018.9)\n",
            "Requirement already satisfied: setuptools>=40.3.0 in /usr/local/lib/python3.7/dist-packages (from google-api-core[grpc]<3.0.0dev,>=1.26.0->google-cloud-aiplatform) (57.4.0)\n",
            "Requirement already satisfied: googleapis-common-protos<2.0dev,>=1.6.0 in /usr/local/lib/python3.7/dist-packages (from google-api-core[grpc]<3.0.0dev,>=1.26.0->google-cloud-aiplatform) (1.53.0)\n",
            "Requirement already satisfied: google-auth<2.0dev,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from google-api-core[grpc]<3.0.0dev,>=1.26.0->google-cloud-aiplatform) (1.35.0)\n",
            "Requirement already satisfied: requests<3.0.0dev,>=2.18.0 in /usr/local/lib/python3.7/dist-packages (from google-api-core[grpc]<3.0.0dev,>=1.26.0->google-cloud-aiplatform) (2.23.0)\n",
            "Requirement already satisfied: grpcio<2.0dev,>=1.29.0 in /usr/local/lib/python3.7/dist-packages (from google-api-core[grpc]<3.0.0dev,>=1.26.0->google-cloud-aiplatform) (1.40.0)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.7/dist-packages (from google-auth<2.0dev,>=1.21.1->google-api-core[grpc]<3.0.0dev,>=1.26.0->google-cloud-aiplatform) (0.2.8)\n",
            "Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.7/dist-packages (from google-auth<2.0dev,>=1.21.1->google-api-core[grpc]<3.0.0dev,>=1.26.0->google-cloud-aiplatform) (4.7.2)\n",
            "Requirement already satisfied: cachetools<5.0,>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from google-auth<2.0dev,>=1.21.1->google-api-core[grpc]<3.0.0dev,>=1.26.0->google-cloud-aiplatform) (4.2.2)\n",
            "Requirement already satisfied: google-cloud-core<2.0dev,>=1.4.1 in /usr/local/lib/python3.7/dist-packages (from google-cloud-bigquery<3.0.0dev,>=1.15.0->google-cloud-aiplatform) (1.7.2)\n",
            "Requirement already satisfied: google-resumable-media<2.0dev,>=0.6.0 in /usr/local/lib/python3.7/dist-packages (from google-cloud-bigquery<3.0.0dev,>=1.15.0->google-cloud-aiplatform) (1.3.3)\n",
            "Requirement already satisfied: google-crc32c<2.0dev,>=1.0 in /usr/local/lib/python3.7/dist-packages (from google-resumable-media<2.0dev,>=0.6.0->google-cloud-bigquery<3.0.0dev,>=1.15.0->google-cloud-aiplatform) (1.2.0)\n",
            "Requirement already satisfied: pyparsing>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging>=14.3->google-cloud-aiplatform) (2.4.7)\n",
            "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.7/dist-packages (from pyasn1-modules>=0.2.1->google-auth<2.0dev,>=1.21.1->google-api-core[grpc]<3.0.0dev,>=1.26.0->google-cloud-aiplatform) (0.4.8)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests<3.0.0dev,>=2.18.0->google-api-core[grpc]<3.0.0dev,>=1.26.0->google-cloud-aiplatform) (2021.5.30)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests<3.0.0dev,>=2.18.0->google-api-core[grpc]<3.0.0dev,>=1.26.0->google-cloud-aiplatform) (3.0.4)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests<3.0.0dev,>=2.18.0->google-api-core[grpc]<3.0.0dev,>=1.26.0->google-cloud-aiplatform) (2.10)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests<3.0.0dev,>=2.18.0->google-api-core[grpc]<3.0.0dev,>=1.26.0->google-cloud-aiplatform) (1.24.3)\n",
            "Requirement already satisfied: werkzeug>=0.11.15 in /usr/local/lib/python3.7/dist-packages (from tensorboard-plugin-profile) (1.0.1)\n",
            "Requirement already satisfied: gviz-api>=1.9.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard-plugin-profile) (1.9.0)\n",
            "\u001b[1;31mERROR:\u001b[0m (gcloud.components.update) The components management action could not be performed because the installation root of the Cloud SDK could not be located. If you previously used the Cloud SDK installer, you could re-install the SDK and retry again.\n"
          ]
        }
      ]
    },
    {
      "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",
      "metadata": {
        "id": "EzrelQZ22IZj"
      },
      "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)"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "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 the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=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 do not know your project ID**, you may be able to get your project ID using `gcloud`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oM1iC_MfAts1"
      },
      "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)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qJYoRfYng0XZ"
      },
      "source": [
        "Otherwise, set your project ID here."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "riG_qUokg0XZ"
      },
      "source": [
        "if PROJECT_ID == \"\" or PROJECT_ID is None:\n",
        "    PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "\n",
        "! gcloud config set project {PROJECT_ID}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "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",
      "metadata": {
        "id": "697568e92bd6"
      },
      "source": [
        "from datetime import datetime\n",
        "\n",
        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dr--iN2kAylZ"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "**If you are using Google Cloud 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",
      "metadata": {
        "id": "PyQmSRbKA8r-"
      },
      "source": [
        "import os\n",
        "import sys\n",
        "\n",
        "# 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",
        "# The Google Cloud Notebook product has specific requirements\n",
        "IS_GOOGLE_CLOUD_NOTEBOOK = os.path.exists(\"/opt/deeplearning/metadata/env_version\")\n",
        "\n",
        "# If on Google Cloud Notebooks, then don't execute this code\n",
        "if not IS_GOOGLE_CLOUD_NOTEBOOK:\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 ''"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "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",
        "Before you submit a training job for the two-tower model, you need to upload your training data and schema to Cloud Storage. Vertex AI trains the model using this input data. In this tutorial, the Two-Tower built-in algorithm also saves the trained model that results from your job in the same bucket. Using this model artifact, you can then create Vertex AI model and endpoint resources in order to serve online predictions.\n",
        "\n",
        "Set the name of your Cloud Storage bucket below. It must be unique across all\n",
        "Cloud Storage buckets.\n",
        "\n",
        "You may also change the `REGION` variable, which is used for operations\n",
        "throughout the rest of this notebook. Make sure to [choose a region where Vertex AI services are\n",
        "available](https://cloud.google.com/vertex-ai/docs/general/locations#available_regions). You may\n",
        "not use a Multi-Regional Storage bucket for training with Vertex AI."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MzGDU7TWdts_"
      },
      "source": [
        "BUCKET_NAME = \"gs://[your-bucket-name]\"  # @param {type:\"string\"}\n",
        "REGION = \"us-central1\"  # @param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cf221059d072"
      },
      "source": [
        "if BUCKET_NAME == \"\" or BUCKET_NAME is None or BUCKET_NAME == \"gs://[your-bucket-name]\":\n",
        "    BUCKET_NAME = \"gs://\" + PROJECT_ID + \"aip-\" + TIMESTAMP"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "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",
      "metadata": {
        "id": "NIq7R4HZCfIc"
      },
      "source": [
        "! gsutil mb -l $REGION $BUCKET_NAME"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ucvCsknMCims"
      },
      "source": [
        "Finally, validate access to your Cloud Storage bucket by examining its contents:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vhOb7YnwClBb"
      },
      "source": [
        "! gsutil ls -al $BUCKET_NAME"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XoEqT2Y4DJmf"
      },
      "source": [
        "### Import libraries and define constants"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pRUOFELefqf1"
      },
      "source": [
        "import json\n",
        "import os\n",
        "import re\n",
        "import time\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "\n",
        "%load_ext tensorboard"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ixC92jeHQMxk"
      },
      "source": [
        "## Configure parameters"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GgsNm8aim0Ym"
      },
      "source": [
        "The following table shows parameters that are common to all Vertex Training jobs created using the `gcloud ai custom-jobs create` command. See the [official documentation](https://cloud.google.com/sdk/gcloud/reference/ai/custom-jobs/create) for all the possible arguments.\n",
        "\n",
        "| Parameter | Data type | Description | Required |\n",
        "|--|--|--|--|\n",
        "| `display-name` | string | Name of the job. | Yes |\n",
        "| `worker-pool-spec` | string | Comma-separated list of arguments specifying a worker pool configuration (see below). | Yes |\n",
        "| `region` | string | Region to submit the job to. | No |\n",
        "\n",
        "The `worker-pool-spec` flag can be specified multiple times, one for each worker pool. The following table shows the arguments used to specify a worker pool.\n",
        "\n",
        "| Parameter | Data type | Description | Required |\n",
        "|--|--|--|--|\n",
        "| `machine-type` | string | Machine type for the pool. See the [official documentation](https://cloud.google.com/vertex-ai/docs/training/configure-compute) for supported machines. | Yes |\n",
        "| `replica-count` | int | The number of replicas of the machine in the pool. | No |\n",
        "| `container-image-uri` | string | Docker image to run on each worker. | No |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0MvQ22Sbm8lh"
      },
      "source": [
        "The following table shows the parameters for the two-tower model training job:\n",
        "\n",
        "| Parameter | Data type | Description | Required |\n",
        "|--|--|--|--|\n",
        "| `training_data_path` | string | Cloud Storage pattern where training data is stored. | Yes |\n",
        "| `input_schema_path` | string | Cloud Storage path where the JSON input schema is stored. | Yes |\n",
        "| `input_file_format` | string | The file format of input. Currently supports `jsonl` and `tfrecord`. | No - default is `jsonl`. |\n",
        "| `job_dir` | string | Cloud Storage directory where the model output files will be stored. | Yes |\n",
        "| `eval_data_path` | string | Cloud Storage pattern where eval data is stored. | No |\n",
        "| `candidate_data_path` | string | Cloud Storage pattern where candidate data is stored. Only used for top_k_categorical_accuracy metrics. If not set, it's generated from training/eval data. | No |\n",
        "| `train_batch_size` | int | Batch size for training. | No - Default is 100. |\n",
        "| `eval_batch_size` | int | Batch size for evaluation. | No - Default is 100. |\n",
        "| `eval_split` | float | Split fraction to use for the evaluation dataset, if `eval_data_path` is not provided. | No - Default is 0.2 |\n",
        "| `optimizer` | string | Training optimizer. Lowercase string name of any TF2.3 Keras optimizer is supported ('sgd', 'nadam', 'ftrl', etc.). See [TensorFlow documentation](https://www.tensorflow.org/api_docs/python/tf/keras/optimizers). | No - Default is 'adagrad'. |\n",
        "| `learning_rate` | float | Learning rate for training. | No - Default is the default learning rate of the specified optimizer. |\n",
        "| `momentum` | float | Momentum for optimizer, if specified. | No - Default is the default momentum value for the specified optimizer. |\n",
        "| `metrics` | string | Metrics used to evaluate the model. Can be either `auc`, `top_k_categorical_accuracy` or `precision_at_1`. | No - Default is `auc`. |\n",
        "| `num_epochs` | int | Number of epochs for training. | No - Default is 10. |\n",
        "| `num_hidden_layers` | int | Number of hidden layers. | No |\n",
        "| `num_nodes_hidden_layer{index}` | int | Num of nodes in hidden layer {index}. The range of index is 1 to 20. | No |\n",
        "| `output_dim` | int | The output embedding dimension for each encoder tower of the two-tower model. | No - Default is 64. |\n",
        "| `training_steps_per_epoch` | int | Number of steps per epoch to run the training for.  Only needed if you are using more than 1 machine or using a master machine with more than 1 gpu. | No - Default is None. |\n",
        "| `eval_steps_per_epoch` | int | Number of steps per epoch to run the evaluation for.  Only needed if you are using more than 1 machine or using a master machine with more than 1 gpu. | No - Default is None. |\n",
        "| `gpu_memory_alloc` | int | Amount of memory allocated per GPU (in MB). | No - Default is no limit. |"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2sEfn2ZVnI_s"
      },
      "source": [
        "DATASET_NAME = \"movielens_100k\"  # Change to your dataset name.\n",
        "\n",
        "# Change to your data and schema paths. These are paths to the movielens_100k\n",
        "# sample data.\n",
        "TRAINING_DATA_PATH = f\"gs://cloud-samples-data/vertex-ai/matching-engine/two-tower/{DATASET_NAME}/training_data/*\"\n",
        "INPUT_SCHEMA_PATH = f\"gs://cloud-samples-data/vertex-ai/matching-engine/two-tower/{DATASET_NAME}/input_schema.json\"\n",
        "\n",
        "# URI of the two-tower training Docker image.\n",
        "LEARNER_IMAGE_URI = \"us-docker.pkg.dev/vertex-ai-restricted/builtin-algorithm/two-tower\"\n",
        "\n",
        "# Change to your output location.\n",
        "OUTPUT_DIR = f\"{BUCKET_NAME}/experiment/output\"\n",
        "\n",
        "TRAIN_BATCH_SIZE = 100  # Batch size for training.\n",
        "NUM_EPOCHS = 3  # Number of epochs for training.\n",
        "\n",
        "print(f\"Dataset name: {DATASET_NAME}\")\n",
        "print(f\"Training data path: {TRAINING_DATA_PATH}\")\n",
        "print(f\"Input schema path: {INPUT_SCHEMA_PATH}\")\n",
        "print(f\"Output directory: {OUTPUT_DIR}\")\n",
        "print(f\"Train batch size: {TRAIN_BATCH_SIZE}\")\n",
        "print(f\"Number of epochs: {NUM_EPOCHS}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "upLZ8kcankwj"
      },
      "source": [
        "## Train on Vertex Training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6M_O_L55nwQ0"
      },
      "source": [
        "Submit the two-tower training job to Vertex Training. The following command uses a single CPU machine for training. When using single node training, `training_steps_per_epoch` and `eval_steps_per_epoch` do not need to be set."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1gXZRq80nl2S"
      },
      "source": [
        "learning_job_name = f\"two_tower_cpu_{DATASET_NAME}_{TIMESTAMP}\"\n",
        "\n",
        "CREATION_LOG = ! gcloud ai custom-jobs create \\\n",
        "  --display-name={learning_job_name} \\\n",
        "  --worker-pool-spec=machine-type=n1-standard-8,replica-count=1,container-image-uri={LEARNER_IMAGE_URI} \\\n",
        "  --region={REGION} \\\n",
        "  --args=--training_data_path={TRAINING_DATA_PATH} \\\n",
        "  --args=--input_schema_path={INPUT_SCHEMA_PATH} \\\n",
        "  --args=--job-dir={OUTPUT_DIR} \\\n",
        "  --args=--train_batch_size={TRAIN_BATCH_SIZE} \\\n",
        "  --args=--num_epochs={NUM_EPOCHS}\n",
        "\n",
        "print(CREATION_LOG)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RaIkcFT2n4_U"
      },
      "source": [
        "If you want to train using GPUs, you need to write configuration to a YAML file:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nAod1hbSn5yw"
      },
      "source": [
        "learning_job_name = f\"two_tower_gpu_{DATASET_NAME}_{TIMESTAMP}\"\n",
        "\n",
        "config = f\"\"\"workerPoolSpecs:\n",
        "  -\n",
        "    machineSpec:\n",
        "      machineType: n1-highmem-4\n",
        "      acceleratorType: NVIDIA_TESLA_K80\n",
        "      acceleratorCount: 1\n",
        "    replicaCount: 1\n",
        "    containerSpec:\n",
        "      imageUri: {LEARNER_IMAGE_URI}\n",
        "      args:\n",
        "      - --training_data_path={TRAINING_DATA_PATH}\n",
        "      - --input_schema_path={INPUT_SCHEMA_PATH}\n",
        "      - --job-dir={OUTPUT_DIR}\n",
        "      - --training_steps_per_epoch=1500\n",
        "      - --eval_steps_per_epoch=1500\n",
        "\"\"\"\n",
        "\n",
        "!echo $'{config}' > ./config.yaml\n",
        "\n",
        "CREATION_LOG = ! gcloud ai custom-jobs create \\\n",
        "  --display-name={learning_job_name} \\\n",
        "  --region={REGION} \\\n",
        "  --config=config.yaml\n",
        "\n",
        "print(CREATION_LOG)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "94tmU59YrKfe"
      },
      "source": [
        "If you want to use TFRecord input file format, you can try the following command:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8wZbRgUhrLD0"
      },
      "source": [
        "TRAINING_DATA_PATH = f\"gs://cloud-samples-data/vertex-ai/matching-engine/two-tower/{DATASET_NAME}/tfrecord/*\"\n",
        "\n",
        "learning_job_name = f\"two_tower_cpu_tfrecord_{DATASET_NAME}_{TIMESTAMP}\"\n",
        "\n",
        "CREATION_LOG = ! gcloud ai custom-jobs create \\\n",
        "  --display-name={learning_job_name} \\\n",
        "  --worker-pool-spec=machine-type=n1-standard-8,replica-count=1,container-image-uri={LEARNER_IMAGE_URI} \\\n",
        "  --region={REGION} \\\n",
        "  --args=--training_data_path={TRAINING_DATA_PATH} \\\n",
        "  --args=--input_schema_path={INPUT_SCHEMA_PATH} \\\n",
        "  --args=--job-dir={OUTPUT_DIR} \\\n",
        "  --args=--train_batch_size={TRAIN_BATCH_SIZE} \\\n",
        "  --args=--num_epochs={NUM_EPOCHS} \\\n",
        "  --args=--input_file_format=tfrecord\n",
        "\n",
        "print(CREATION_LOG)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yceUSlyWrWes"
      },
      "source": [
        "After the job is submitted successfully, you can view its details and logs:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XXDC7F_8rXWM"
      },
      "source": [
        "JOB_ID = re.search(r\"(?<=/customJobs/)\\d+\", CREATION_LOG[1]).group(0)\n",
        "print(JOB_ID)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NkbwWCEMoQcy"
      },
      "source": [
        "# View the job's configuration and state.\n",
        "STATE = \"state: JOB_STATE_PENDING\"\n",
        "\n",
        "while STATE not in [\"state: JOB_STATE_SUCCEEDED\", \"state: JOB_STATE_FAILED\"]:\n",
        "    DESCRIPTION = ! gcloud ai custom-jobs describe {JOB_ID} --region={REGION}\n",
        "    STATE = DESCRIPTION[-2]\n",
        "    print(STATE)\n",
        "    time.sleep(60)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wgs0qV_Nr-RN"
      },
      "source": [
        "When the training starts, you can view the logs in TensorBoard. Colab users can use the TensorBoard widget below:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8SweSrkhr_DP"
      },
      "source": [
        "TENSORBOARD_DIR = os.path.join(OUTPUT_DIR, \"tensorboard\")\n",
        "%tensorboard --logdir {TENSORBOARD_DIR}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RzCrdxgAsGll"
      },
      "source": [
        "For Google CLoud Notebooks users, the TensorBoard widget above won't work. We recommend you to launch TensorBoard through the Cloud Shell.\n",
        "\n",
        "1. In your Cloud Shell, launch Tensorboard on port 8080:\n",
        "\n",
        "    ```\n",
        "    export TENSORBOARD_DIR=gs://xxxxx/tensorboard\n",
        "    tensorboard --logdir=${TENSORBOARD_DIR} --port=8080\n",
        "    ```\n",
        "\n",
        "2. Click the \"Web Preview\" button at the top-right of the Cloud Shell window (looks like an eye in a rectangle). \n",
        "\n",
        "3. Select \"Preview on port 8080\". This should launch the TensorBoard webpage in a new tab in your browser.\n",
        "\n",
        "After the job finishes successfully, you can view the output directory:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mFPLfY4gsK1V"
      },
      "source": [
        "! gsutil ls {OUTPUT_DIR}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZhY0h8ijsPlP"
      },
      "source": [
        "## Deploy on Vertex Prediction"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6oquDjRgsS2V"
      },
      "source": [
        "### Import the model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gidmXBWysaeP"
      },
      "source": [
        "Our training job will export two TF SavedModels under `gs://<job_dir>/query_model` and `gs://<job_dir>/candidate_model`. These exported models can be used for online or batch prediction in Vertex Prediction. First, import the query (or candidate) model:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yEFd3Og_sbdm"
      },
      "source": [
        "# The following imports the query (user) encoder model.\n",
        "MODEL_TYPE = \"query\"\n",
        "# Use the following instead to import the candidate (movie) encoder model.\n",
        "# MODEL_TYPE = 'candidate'\n",
        "\n",
        "DISPLAY_NAME = f\"{DATASET_NAME}_{MODEL_TYPE}\"  # The display name of the model.\n",
        "MODEL_NAME = f\"{MODEL_TYPE}_model\"  # Used by the deployment container."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GCrhxf7GsdZS"
      },
      "source": [
        "aiplatform.init(\n",
        "    project=PROJECT_ID,\n",
        "    location=REGION,\n",
        "    staging_bucket=BUCKET_NAME,\n",
        ")\n",
        "\n",
        "model = aiplatform.Model.upload(\n",
        "    display_name=DISPLAY_NAME,\n",
        "    artifact_uri=OUTPUT_DIR,\n",
        "    serving_container_image_uri=\"us-central1-docker.pkg.dev/cloud-ml-algos/two-tower/deploy\",\n",
        "    serving_container_health_route=f\"/v1/models/{MODEL_NAME}\",\n",
        "    serving_container_predict_route=f\"/v1/models/{MODEL_NAME}:predict\",\n",
        "    serving_container_environment_variables={\n",
        "        \"MODEL_BASE_PATH\": \"$(AIP_STORAGE_URI)\",\n",
        "        \"MODEL_NAME\": MODEL_NAME,\n",
        "    },\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0x-pZJzUsh22"
      },
      "source": [
        "### Deploy the model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QJECrcTdskix"
      },
      "source": [
        "After importing the model, you must deploy it to an endpoint so that you can get online predictions. More information about this process can be found in the [official documentation](https://cloud.google.com/vertex-ai/docs/predictions/deploy-model-api)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jB3yT5xassCt"
      },
      "source": [
        "! gcloud ai models list --region={REGION} --filter={DISPLAY_NAME}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zkJh2rmysu2M"
      },
      "source": [
        "Create a model endpoint:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "er94Wp82sxYW"
      },
      "source": [
        "endpoint = aiplatform.Endpoint.create(display_name=DATASET_NAME)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PICvm8PhqtMw"
      },
      "source": [
        "Deploy model to the endpoint"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OUEF7Yces4uD"
      },
      "source": [
        "model.deploy(\n",
        "    endpoint=endpoint,\n",
        "    machine_type=\"n1-standard-4\",\n",
        "    traffic_split={\"0\": 100},\n",
        "    deployed_model_display_name=DISPLAY_NAME,\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y_LMW1rjtMM6"
      },
      "source": [
        "## Predict"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QAkHJlY7tOmu"
      },
      "source": [
        "Now that you have deployed the query/candidate encoder model on Vertex Prediction, you can call the model to calculate embeddings for live data. There are two methods of getting predictions, online and batch, which are shown below."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rwiJRfJRtQ1V"
      },
      "source": [
        "### Online prediction"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "THz5Gn5ftTsm"
      },
      "source": [
        "[Online prediction](https://cloud.google.com/vertex-ai/docs/predictions/online-predictions-custom-models) is used to synchronously query a model on a small batch of instances with minimal latency. The following function calls the deployed Vertex Prediction model endpoint using Vertex SDK for Python:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bFVnzRzltfDa"
      },
      "source": [
        "The input data you want predictions on should be provided as a stringified JSON in the `data` field. Note that you should also provide a unique `key` field (of type str) for each input instance so that you can associate each output embedding with its corresponding input."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E8Wt7wYgtg_f"
      },
      "source": [
        "# Input items for the query model:\n",
        "input_items = [\n",
        "    {\"data\": '{\"user_id\": [\"1\"]}', \"key\": \"key1\"},\n",
        "    {\"data\": '{\"user_id\": [\"2\"]}', \"key\": \"key2\"},\n",
        "]\n",
        "\n",
        "# Input items for the candidate model:\n",
        "# input_items = [{\n",
        "#     'data' : '{\"movie_id\": [\"1\"], \"movie_title\": [\"fake title\"]}',\n",
        "#     'key': 'key1'\n",
        "# }]\n",
        "\n",
        "encodings = endpoint.predict(input_items)\n",
        "print(f\"Number of encodings: {len(encodings.predictions)}\")\n",
        "print(encodings.predictions[0][\"encoding\"])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1k-_XJzlthfP"
      },
      "source": [
        "You can also do online prediction using the gcloud CLI, as shown below:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Tn5L9V0utkpA"
      },
      "source": [
        "request = json.dumps({\"instances\": input_items})\n",
        "with open(\"request.json\", \"w\") as writer:\n",
        "    writer.write(f\"{request}\\n\")\n",
        "\n",
        "ENDPOINT_ID = endpoint.resource_name\n",
        "\n",
        "! gcloud ai endpoints predict {ENDPOINT_ID} \\\n",
        "  --region={REGION} \\\n",
        "  --json-request=request.json"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ocLE_U6ftnA3"
      },
      "source": [
        "### Batch prediction"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U__JdxfPto-_"
      },
      "source": [
        "[Batch prediction](https://cloud.google.com/vertex-ai/docs/predictions/batch-predictions) is used to asynchronously make predictions on a batch of input data.  This is recommended if you have a large input size and do not need an immediate response, such as getting embeddings for candidate objects in order to create an index for a nearest neighbor search service such as [Vertex Matching Engine](https://cloud.google.com/vertex-ai/docs/matching-engine/overview).\n",
        "\n",
        "The input data needs to be on Cloud Storage and in JSONL format. You can use the sample query object file provided below. Like with online prediction, it's recommended to have the `key` field so that you can associate each output embedding with its corresponding input."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ar13RZ4VtquX"
      },
      "source": [
        "QUERY_SAMPLE_PATH = f\"gs://cloud-samples-data/vertex-ai/matching-engine/two-tower/{DATASET_NAME}/query_sample.jsonl\"\n",
        "\n",
        "! gsutil cat {QUERY_SAMPLE_PATH}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z4f5yEmatu8P"
      },
      "source": [
        "The following function calls the deployed Vertex Prediction model using the sample query object input file. Note that it uses the model resource directly and doesn't require a deployed endpoint. Once you start the job, you can track its status on the [Cloud Console](https://console.cloud.google.com/vertex-ai/batch-predictions)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EYEOOPS8txYY"
      },
      "source": [
        "model.batch_predict(\n",
        "    job_display_name=f\"batch_predict_{DISPLAY_NAME}\",\n",
        "    gcs_source=[QUERY_SAMPLE_PATH],\n",
        "    gcs_destination_prefix=OUTPUT_DIR,\n",
        "    machine_type=\"n1-standard-4\",\n",
        "    starting_replica_count=1,\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zImnlP2Yt6Sv"
      },
      "source": [
        "## Hyperparameter tuning"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I2nRxtLTt8xn"
      },
      "source": [
        "After successfully training your model, deploying it, and calling it to make predictions, you may want to optimize the hyperparameters used during training to improve your model's accuracy and performance. See the Vertex AI documentation for an [overview of hyperparameter tuning](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview) and [how to use it in your Vertex Training jobs](https://cloud.google.com/vertex-ai/docs/training/using-hyperparameter-tuning).\n",
        "\n",
        "For this example, the following command runs a Vertex AI hyperparameter tuning job with 8 trials that attempts to maximize the validation AUC metric. The hyperparameters it optimizes are the number of hidden layers, the size of the hidden layers, and the learning rate."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z_ea5wfjt_XD"
      },
      "source": [
        "PARALLEL_TRIAL_COUNT = 4\n",
        "MAX_TRIAL_COUNT = 8\n",
        "METRIC = \"val_auc\"\n",
        "hyper_tune_job_name = f\"hyper_tune_{DATASET_NAME}_{TIMESTAMP}\"\n",
        "\n",
        "config = json.dumps(\n",
        "    {\n",
        "        \"displayName\": hyper_tune_job_name,\n",
        "        \"studySpec\": {\n",
        "            \"metrics\": [{\"metricId\": METRIC, \"goal\": \"MAXIMIZE\"}],\n",
        "            \"parameters\": [\n",
        "                {\n",
        "                    \"parameterId\": \"num_hidden_layers\",\n",
        "                    \"scaleType\": \"UNIT_LINEAR_SCALE\",\n",
        "                    \"integerValueSpec\": {\"minValue\": 0, \"maxValue\": 2},\n",
        "                    \"conditionalParameterSpecs\": [\n",
        "                        {\n",
        "                            \"parameterSpec\": {\n",
        "                                \"parameterId\": \"num_nodes_hidden_layer1\",\n",
        "                                \"scaleType\": \"UNIT_LOG_SCALE\",\n",
        "                                \"integerValueSpec\": {\"minValue\": 1, \"maxValue\": 128},\n",
        "                            },\n",
        "                            \"parentIntValues\": {\"values\": [1, 2]},\n",
        "                        },\n",
        "                        {\n",
        "                            \"parameterSpec\": {\n",
        "                                \"parameterId\": \"num_nodes_hidden_layer2\",\n",
        "                                \"scaleType\": \"UNIT_LOG_SCALE\",\n",
        "                                \"integerValueSpec\": {\"minValue\": 1, \"maxValue\": 128},\n",
        "                            },\n",
        "                            \"parentIntValues\": {\"values\": [2]},\n",
        "                        },\n",
        "                    ],\n",
        "                },\n",
        "                {\n",
        "                    \"parameterId\": \"learning_rate\",\n",
        "                    \"scaleType\": \"UNIT_LOG_SCALE\",\n",
        "                    \"doubleValueSpec\": {\"minValue\": 0.0001, \"maxValue\": 1.0},\n",
        "                },\n",
        "            ],\n",
        "            \"algorithm\": \"ALGORITHM_UNSPECIFIED\",\n",
        "        },\n",
        "        \"maxTrialCount\": MAX_TRIAL_COUNT,\n",
        "        \"parallelTrialCount\": PARALLEL_TRIAL_COUNT,\n",
        "        \"maxFailedTrialCount\": 3,\n",
        "        \"trialJobSpec\": {\n",
        "            \"workerPoolSpecs\": [\n",
        "                {\n",
        "                    \"machineSpec\": {\n",
        "                        \"machineType\": \"n1-standard-4\",\n",
        "                    },\n",
        "                    \"replicaCount\": 1,\n",
        "                    \"containerSpec\": {\n",
        "                        \"imageUri\": LEARNER_IMAGE_URI,\n",
        "                        \"args\": [\n",
        "                            f\"--training_data_path={TRAINING_DATA_PATH}\",\n",
        "                            f\"--input_schema_path={INPUT_SCHEMA_PATH}\",\n",
        "                            f\"--job-dir={OUTPUT_DIR}\",\n",
        "                        ],\n",
        "                    },\n",
        "                }\n",
        "            ]\n",
        "        },\n",
        "    }\n",
        ")\n",
        "\n",
        "\n",
        "! curl -X POST -H \"Authorization: Bearer \"$(gcloud auth print-access-token) \\\n",
        " -H \"Content-Type: application/json; charset=utf-8\"  \\\n",
        " -d '{config}' https://us-central1-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/{REGION}/hyperparameterTuningJobs"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "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",
      "metadata": {
        "id": "sx_vKniMq9ZX"
      },
      "source": [
        "# Delete endpoint resource\n",
        "endpoint.delete(force=True)\n",
        "\n",
        "# Delete model resource\n",
        "model.delete()\n",
        "\n",
        "# Delete Cloud Storage objects that were created\n",
        "! gsutil -m rm -r $OUTPUT_DIR"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}