notebooks/official/pipelines/google_cloud_pipeline_components_bqml_text.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2022 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JAPoU8Sm5E6e"
      },
      "source": [
        "# Training an acquisition-prediction model using Swivel, BigQuery ML and Vertex AI Pipelines\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/pipelines/google_cloud_pipeline_components_bqml_text.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.luolix.top%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fpipelines%2Fgoogle_cloud_pipeline_components_bqml_text.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>    \n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/pipelines/google_cloud_pipeline_components_bqml_text.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/pipelines/google_cloud_pipeline_components_bqml_text.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": "991ab00f3d75"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates the usage of `DataflowPythonJobOp` and BigQuery ML components through buidling a Text Categorization model and running it on Vertex AI Pipelines. \n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "1. Read raw text (HTML) documents stored in Google Cloud Storage.\n",
        "2. Extract title, content and topic of (HTML) documents using Dataflow and ingest into BigQuery.\n",
        "3. Apply the Swivel model to generate embeddings of your document’s content.\n",
        "4. Train a Logistic regression model to classify if an article is about corporate acquisitions (`acq` category). \n",
        "5. Evaluate the model.\n",
        "6. Apply the model to a dataset in order to generate predictions.\n",
        "\n",
        "Learn more about [Vertex AI Pipelines](https://cloud.google.com/vertex-ai/docs/pipelines/introduction) and [BigQuery ML components](https://cloud.google.com/vertex-ai/docs/pipelines/bigqueryml-component)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "acc98a3361cc"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this notebook, you learn how to build a simple BigQuery ML pipeline using Vertex AI pipelines in order to calculate text embeddings of content from articles and classify them\n",
        "into the *corporate acquisitions* category.\n",
        "\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
        "- Vertex AI Pipelines\n",
        "- BigQuery ML\n",
        "\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Creating a component for Dataflow job that ingests data to BigQuery.\n",
        "- Creating a component for preprocessing steps to run on the data in BigQuery.\n",
        "- Creating a component for training a logistic regression model using BigQuery ML.\n",
        "- Building and configuring a Kubeflow DSL pipeline with all the created components.\n",
        "- Compiling and running the pipeline in Vertex AI Pipelines."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9bc6d52899ba"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "The dataset used in this notebook is the [Reuters-21578 Text Categorization Collection Data Set](https://archive.ics.uci.edu/ml/datasets/reuters-21578+text+categorization+collection). This dataset is a collection of publicly available news articles appeared on the Reuters newswire in 1987. They were assembled and indexed with categories by personnel from Reuters Ltd. and Carnegie Group, Inc. in 1987."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "### Costs \n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "* BigQuery\n",
        "* Dataflow\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\n",
        "pricing](https://cloud.google.com/bigquery/pricing), [Dataflow\n",
        "pricing](https://cloud.google.com/dataflow/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": "f0316df526f8"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "install_aip:mbsdk"
      },
      "source": [
        "### Install Vertex AI SDK for Python and other required packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2b4ef9b72d43"
      },
      "outputs": [],
      "source": [
        "# Install dependencies\n",
        "! pip3 install --upgrade --quiet    google-cloud-aiplatform \\\n",
        "                                    google_cloud_pipeline_components \\\n",
        "                                    google-api-core \\\n",
        "                                    google-auth\n",
        "\n",
        "! pip3 install --upgrade --quiet    tensorflow==2.13.1 \\\n",
        "                                    tensorflow-hub==0.16.0 \\\n",
        "                                    kfp"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "613b1be08c68"
      },
      "source": [
        "### Restart runtime (Colab only)\n",
        "\n",
        "To use the newly installed packages, you must restart the runtime on Google "
      ]
    },
    {
      "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": "4a2b7b59bbf7"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
        "</div>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "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": "91842ef41bbd"
      },
      "source": [
        "### Set Google Cloud project information\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oM1iC_MfAts1"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "06571eb4063b"
      },
      "source": [
        "#### UUID\n",
        "\n",
        "If you're in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a uuid for each instance session, and append it onto the name of resources you create in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "697568e92bd6"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "import string\n",
        "\n",
        "\n",
        "# Generate a uuid of length 8\n",
        "def generate_uuid():\n",
        "    return \"\".join(random.choices(string.ascii_lowercase + string.digits, k=8))\n",
        "\n",
        "\n",
        "UUID = generate_uuid()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgPO1eR3CYjk"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MzGDU7TWdts_"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-EcIXiGsCePi"
      },
      "source": [
        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NIq7R4HZCfIc"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3708bd0b1855"
      },
      "source": [
        "#### Service Account\n",
        "\n",
        "You use a service account to create Vertex AI Pipeline jobs. If you don't want to use your project's Compute Engine service account, set `SERVICE_ACCOUNT` to another service account ID."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "199a32a35466"
      },
      "outputs": [],
      "source": [
        "SERVICE_ACCOUNT = \"[your-service-account]\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "abb872bb98c1"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import sys\n",
        "\n",
        "IS_COLAB = \"google.colab\" in sys.modules\n",
        "\n",
        "if (\n",
        "    SERVICE_ACCOUNT == \"\"\n",
        "    or SERVICE_ACCOUNT is None\n",
        "    or SERVICE_ACCOUNT == \"[your-service-account]\"\n",
        "):\n",
        "    # Get your service account from gcloud\n",
        "    if not IS_COLAB:\n",
        "        shell_output = !gcloud auth list 2>/dev/null\n",
        "        SERVICE_ACCOUNT = shell_output[2].replace(\"*\", \"\").strip()\n",
        "\n",
        "    else:  # IS_COLAB:\n",
        "        shell_output = ! gcloud projects describe  $PROJECT_ID\n",
        "        project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
        "        SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
        "\n",
        "    print(\"Service Account:\", SERVICE_ACCOUNT)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4a63f2d1cd52"
      },
      "source": [
        "#### Set service account access for Vertex AI Pipelines\n",
        "\n",
        "Run the following commands to grant your service account access to read and write pipeline artifacts in the bucket that you created in the previous step. You only need to run this step once per service account."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "749c598c5f5d"
      },
      "outputs": [],
      "source": [
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectCreator $BUCKET_URI\n",
        "\n",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectViewer $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FutXJbmAa1NC"
      },
      "source": [
        "### Set project template"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3Kms5lCRn7N-"
      },
      "outputs": [],
      "source": [
        "DATA_PATH = \"data\"\n",
        "KFP_COMPONENTS_PATH = \"components\"\n",
        "SRC = \"src\"\n",
        "BUILD = \"build\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hZRbBCLra0po"
      },
      "outputs": [],
      "source": [
        "!mkdir -m 777 -p {DATA_PATH} {KFP_COMPONENTS_PATH} {SRC} {BUILD}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Jpg5bPrNbpkl"
      },
      "source": [
        "### Prepare input data\n",
        "\n",
        "In the following cell, you:\n",
        "\n",
        "1) Get the dataset from UCI archive.\n",
        "2) Untar the dataset.\n",
        "3) Copy the dataset to the Cloud Storage location."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "un5pHVJtcNsQ"
      },
      "outputs": [],
      "source": [
        "!wget --no-parent https://archive.ics.uci.edu/ml/machine-learning-databases/reuters21578-mld/reuters21578.tar.gz --directory-prefix={DATA_PATH}/raw\n",
        "!mkdir -m 777 -p {DATA_PATH}/raw/temp {DATA_PATH}/raw\n",
        "!tar -zxvf {DATA_PATH}/raw/reuters21578.tar.gz -C {DATA_PATH}/raw/temp/\n",
        "!mv {DATA_PATH}/raw/temp/*.sgm {DATA_PATH}/raw && rm -rf {DATA_PATH}/raw/temp && rm -f {DATA_PATH}/raw/reuters21578.tar.gz"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ERaAoEO8bsl-"
      },
      "outputs": [],
      "source": [
        "!gsutil -m cp -R {DATA_PATH}/raw $BUCKET_URI/{DATA_PATH}/raw"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XoEqT2Y4DJmf"
      },
      "source": [
        "### Import libraries "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pRUOFELefqf1"
      },
      "outputs": [],
      "source": [
        "from pathlib import Path as path\n",
        "from urllib.parse import urlparse\n",
        "\n",
        "import tensorflow_hub as hub\n",
        "from google.cloud import aiplatform as vertex_ai\n",
        "from kfp import compiler, dsl\n",
        "from kfp.dsl import component\n",
        "\n",
        "os.environ[\"TFHUB_MODEL_LOAD_FORMAT\"] = \"UNCOMPRESSED\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vlQwSmhIam3l"
      },
      "source": [
        "### Define constants\n",
        "\n",
        "About the model you're going to use in preprocessing, you use the [Swivel](https://tfhub.dev/google/tf2-preview/gnews-swivel-20dim/1) embedding which was trained on English Google News 130GB corpus and has 20 dimensions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fAnAKF6KamY3"
      },
      "outputs": [],
      "source": [
        "JOB_NAME = f\"reuters-ingest-{UUID}\"\n",
        "SETUP_FILE_URI = urlparse(BUCKET_URI)._replace(path=\"setup.py\").geturl()\n",
        "RUNNER = \"DataflowRunner\"\n",
        "STAGING_LOCATION_URI = urlparse(BUCKET_URI)._replace(path=\"staging\").geturl()\n",
        "TMP_LOCATION_URI = urlparse(BUCKET_URI)._replace(path=\"temp\").geturl()\n",
        "INPUTS_URI = urlparse(BUCKET_URI)._replace(path=f\"{DATA_PATH}/raw/*.sgm\").geturl()\n",
        "BQ_DATASET = \"mlops_bqml_text_analyisis\"\n",
        "BQ_TABLE = \"reuters_ingested\"\n",
        "MODEL_NAME = \"swivel_text_embedding_model\"\n",
        "EMBEDDINGS_TABLE = f\"reuters_text_embeddings_{UUID}\"\n",
        "MODEL_PATH = (\n",
        "    f'{hub.resolve(\"https://tfhub.dev/google/tf2-preview/gnews-swivel-20dim/1\")}/*'\n",
        ")\n",
        "PREPROCESSED_TABLE = f\"reuters_text_preprocessed_{UUID}\"\n",
        "CLASSIFICATION_MODEL_NAME = \"logistic_reg\"\n",
        "PREDICT_TABLE = f\"reuters_text_predict_{UUID}\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NlyOjKrjCXsI"
      },
      "source": [
        "### Initialize Vertex AI SDK client"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IsT5FfP0CaXt"
      },
      "outputs": [],
      "source": [
        "vertex_ai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZrgOD30o7HcL"
      },
      "source": [
        "## Pipeline formalization\n",
        "\n",
        "In this step, you create various components for the pipeline and build the final pipeline."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LLs96NR97YAr"
      },
      "source": [
        "### Data ingestion component"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nR3r6ZD47oYb"
      },
      "source": [
        "#### Create Dataflow Python module\n",
        "\n",
        "The following module contains a Dataflow pipeline that\n",
        "\n",
        "1) Read the files from Cloud Storage.\n",
        "2) Extract the article and generate title, topics, and content from files.\n",
        "3) Load the structured data to BigQuery.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mO3ggLEtkD2R"
      },
      "outputs": [],
      "source": [
        "!touch {SRC}/__init__.py "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h5G6BFF874Ys"
      },
      "outputs": [],
      "source": [
        "%%writefile src/ingest_pipeline.py\n",
        "# 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.\n",
        "\n",
        "# General imports\n",
        "from __future__ import absolute_import\n",
        "import argparse\n",
        "import logging\n",
        "import os\n",
        "import string\n",
        "\n",
        "# Preprocessing imports\n",
        "import tensorflow as tf\n",
        "import bs4\n",
        "import nltk\n",
        "\n",
        "import apache_beam as beam\n",
        "from apache_beam.io.gcp.internal.clients import bigquery\n",
        "from apache_beam.options.pipeline_options import PipelineOptions\n",
        "from apache_beam.options.pipeline_options import SetupOptions\n",
        "\n",
        "\n",
        "# Helpers ---------------------------------------------------------------------\n",
        "\n",
        "def get_args():\n",
        "    \"\"\"\n",
        "    Get command line arguments.\n",
        "    Returns:\n",
        "      args: The parsed arguments.\n",
        "    \"\"\"\n",
        "    parser = argparse.ArgumentParser()\n",
        "    parser.add_argument('--inputs', dest='inputs', default='data/raw/reuters/*.sgm',\n",
        "                        help='A directory location of input data')\n",
        "    parser.add_argument('--bq-dataset', dest='bq_dataset', required=False,\n",
        "                        default='reuters_dataset', help='Dataset name used in BigQuery.')\n",
        "    parser.add_argument('--bq-table', dest='bq_table', required=False,\n",
        "                        default='reuters_ingested_table', help='Table name used in BigQuery.')\n",
        "    args, pipeline_args = parser.parse_known_args()\n",
        "    return args, pipeline_args\n",
        "\n",
        "def get_paths(data_pattern):\n",
        "    \"\"\"\n",
        "  A function to get all the paths of the files in the data directory.\n",
        "  Args:\n",
        "    data_pattern: A directory location of input data.\n",
        "  Returns:\n",
        "    A list of file paths.\n",
        "  \"\"\"\n",
        "    data_paths = tf.io.gfile.glob(data_pattern)\n",
        "    return data_paths\n",
        "\n",
        "\n",
        "def get_title(article):\n",
        "    \"\"\"\n",
        "    A function to get the title of an article.\n",
        "    Args:\n",
        "        article: A BeautifulSoup object of an article.\n",
        "    Returns:\n",
        "        A string of the title of the article.\n",
        "    \"\"\"\n",
        "    title = article.find('text').title\n",
        "    if title is not None:\n",
        "        title = ''.join(filter(lambda x: x in set(string.printable), title.text))\n",
        "        title = title.encode('ascii', 'ignore')\n",
        "    return title\n",
        "\n",
        "\n",
        "def get_content(article):\n",
        "    \"\"\"\n",
        "    A function to get the content of an article.\n",
        "    Args:\n",
        "        article: A BeautifulSoup object of an article.\n",
        "    Returns:\n",
        "        A string of the content of the article.\n",
        "    \"\"\"\n",
        "    content = article.find('text').body\n",
        "    if content is not None:\n",
        "        content = ''.join(filter(lambda x: x in set(string.printable), content.text))\n",
        "        content = ' '.join(content.split())\n",
        "        try:\n",
        "            content = '\\n'.join(nltk.sent_tokenize(content))\n",
        "        except LookupError:\n",
        "            nltk.download('punkt')\n",
        "            content = '\\n'.join(nltk.sent_tokenize(content))\n",
        "        content = content.encode('ascii', 'ignore')\n",
        "    return content\n",
        "\n",
        "\n",
        "def get_topics(article):\n",
        "    \"\"\"\n",
        "    A function to get the topics of an article.\n",
        "    Args:\n",
        "        article: A BeautifulSoup object of an article.\n",
        "    Returns:\n",
        "        A list of strings of the topics of the article.\n",
        "    \"\"\"\n",
        "    topics = []\n",
        "    for topic in article.topics.children:\n",
        "        topic = ''.join(filter(lambda x: x in set(string.printable), topic.text))\n",
        "        topics.append(topic.encode('ascii', 'ignore'))\n",
        "    return topics\n",
        "\n",
        "\n",
        "def get_articles(data_paths):\n",
        "    \"\"\"\n",
        "    Args:\n",
        "        data_paths: A list of file paths.\n",
        "    Returns:\n",
        "        A list of articles.\n",
        "    \"\"\"\n",
        "    data = tf.io.gfile.GFile(data_paths, 'rb').read()\n",
        "    soup = bs4.BeautifulSoup(data, \"html.parser\")\n",
        "    articles = []\n",
        "    for raw_article in soup.find_all('reuters'):\n",
        "        article = {\n",
        "            'title': get_title(raw_article),\n",
        "            'content': get_content(raw_article),\n",
        "            'topics': get_topics(raw_article)\n",
        "        }\n",
        "        if None not in article.values():\n",
        "            if [] not in article.values():\n",
        "                articles.append(article)\n",
        "    return articles\n",
        "\n",
        "\n",
        "def get_bigquery_schema():\n",
        "    \"\"\"\n",
        "    A function to get the BigQuery schema.\n",
        "    Returns:\n",
        "        A list of BigQuery schema.\n",
        "    \"\"\"\n",
        "\n",
        "    table_schema = bigquery.TableSchema()\n",
        "    columns = (('topics', 'string', 'repeated'),\n",
        "               ('title', 'string', 'nullable'),\n",
        "               ('content', 'string', 'nullable'))\n",
        "\n",
        "    for column in columns:\n",
        "        column_schema = bigquery.TableFieldSchema()\n",
        "        column_schema.name = column[0]\n",
        "        column_schema.type = column[1]\n",
        "        column_schema.mode = column[2]\n",
        "        table_schema.fields.append(column_schema)\n",
        "\n",
        "    return table_schema\n",
        "\n",
        "\n",
        "# Pipeline runner\n",
        "def run(args, pipeline_args=None):\n",
        "    \"\"\"\n",
        "    A function to run the pipeline.\n",
        "    Args:\n",
        "        args: The parsed arguments.\n",
        "    Returns:\n",
        "        None\n",
        "    \"\"\"\n",
        "\n",
        "    options = PipelineOptions(pipeline_args)\n",
        "    options.view_as(SetupOptions).save_main_session = True\n",
        "\n",
        "    pipeline = beam.Pipeline(options=options)\n",
        "    articles = (\n",
        "            pipeline\n",
        "            | 'Get Paths' >> beam.Create(get_paths(args.inputs))\n",
        "            | 'Get Articles' >> beam.Map(get_articles)\n",
        "            | 'Get Article' >> beam.FlatMap(lambda x: x)\n",
        "    )\n",
        "    if options.get_all_options()['runner'] == 'DirectRunner':\n",
        "        articles | 'Dry run' >> beam.io.WriteToText('data/processed/reuters', file_name_suffix=\".jsonl\")\n",
        "    else:\n",
        "        (articles\n",
        "         | 'Write to BigQuery' >> beam.io.WriteToBigQuery(\n",
        "                    project=options.get_all_options()['project'],\n",
        "                    dataset=args.bq_dataset,\n",
        "                    table=args.bq_table,\n",
        "                    schema=get_bigquery_schema(),\n",
        "                    create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,\n",
        "                    write_disposition=beam.io.BigQueryDisposition.WRITE_TRUNCATE)\n",
        "         )\n",
        "    job = pipeline.run()\n",
        "\n",
        "    if options.get_all_options()['runner'] == 'DirectRunner':\n",
        "        job.wait_until_finish()\n",
        "\n",
        "\n",
        "if __name__ == '__main__':\n",
        "    args, pipeline_args = get_args()\n",
        "    logging.getLogger().setLevel(logging.INFO)\n",
        "    run(args, pipeline_args)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nEsufwCX-h08"
      },
      "source": [
        "#### Create requirements\n",
        "\n",
        "Next, create the requirements.txt file with Python modules that are needed for Apache Beam pipeline. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "S13Ma74p7bey"
      },
      "outputs": [],
      "source": [
        "%%writefile requirements.txt\n",
        "apache-beam[gcp]==2.47.0\n",
        "bs4==0.0.2\n",
        "nltk==3.7\n",
        "tensorflow==2.13.1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3TlCcRKsqZut"
      },
      "source": [
        "#### Create Setup file\n",
        "\n",
        "And add the setup file with Python modules that are needed for executing the Dataflow workers. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "USKNGS0qqZu1"
      },
      "outputs": [],
      "source": [
        "%%writefile setup.py\n",
        "# !/usr/bin/python\n",
        "\n",
        "# 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",
        "#      http://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.\n",
        "\n",
        "import setuptools\n",
        "\n",
        "REQUIRED_PACKAGES = [\n",
        "    'bs4==0.0.2',\n",
        "    'nltk==3.7',\n",
        "    'tensorflow==2.13.1']\n",
        "\n",
        "setuptools.setup(\n",
        "    name='ingest',\n",
        "    version='0.0.1',\n",
        "    author='author',\n",
        "    author_email='author@google.com',\n",
        "    install_requires=REQUIRED_PACKAGES,\n",
        "    packages=setuptools.find_packages())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5Nc7ByK1AEe8"
      },
      "source": [
        "#### Copy the setup, module and requirements files to Cloud Storage\n",
        "\n",
        "Finally, copy the Python module, requirements and setup file to your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V0qslZVeAHqo"
      },
      "outputs": [],
      "source": [
        "# !gsutil cp -R {SRC}/preprocess_pipeline.py {BUCKET_URI}/preprocess_pipeline.py\n",
        "!gsutil cp -R {SRC} {BUCKET_URI}/{SRC}\n",
        "!gsutil cp requirements.txt {BUCKET_URI}/requirements.txt\n",
        "!gsutil cp setup.py {BUCKET_URI}/setup.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MI_wYYwdAZZs"
      },
      "source": [
        "### BigQuery ML components\n",
        "\n",
        "For the next steps in building the pipeline, you define a set of queries to:\n",
        "\n",
        "1) Create the BigQuery dataset schema.\n",
        "2) Preprocess your text data and generate the embeddings using Swivel model.\n",
        "2) Train the BigQuery ML Logistic Regression model.\n",
        "3) Evaluate the model.\n",
        "4) Run a batch prediction.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ytsqqHZnksHY"
      },
      "outputs": [],
      "source": [
        "!mkdir -m 777 -p {KFP_COMPONENTS_PATH}/bq_dataset_component\n",
        "!mkdir -m 777 -p {KFP_COMPONENTS_PATH}/bq_preprocess_component\n",
        "!mkdir -m 777 -p {KFP_COMPONENTS_PATH}/bq_model_component\n",
        "!mkdir -m 777 -p {KFP_COMPONENTS_PATH}/bq_prediction_component"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8tYtZMRiepKe"
      },
      "source": [
        "#### Create BigQuery Dataset query\n",
        "\n",
        "With this query, you create the Bigquery dataset schema that you're going to use to train your model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YP2K8Nc4HqrQ"
      },
      "outputs": [],
      "source": [
        "create_bq_dataset_query = f\"\"\"\n",
        "CREATE SCHEMA IF NOT EXISTS {BQ_DATASET}\n",
        "\"\"\"\n",
        "\n",
        "with open(\n",
        "    f\"{KFP_COMPONENTS_PATH}/bq_dataset_component/create_bq_dataset.sql\", \"w\"\n",
        ") as q:\n",
        "    q.write(create_bq_dataset_query)\n",
        "q.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T2_kuUSZUBDY"
      },
      "source": [
        "#### Create BigQuery Preprocess query\n",
        "\n",
        "The following query uses the TFHub Swivel model to generate embeddings for your text data and splits the dataset for training and serving purposes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XdyboQnLUBDg"
      },
      "outputs": [],
      "source": [
        "create_bq_preprocess_query = f\"\"\"\n",
        "-- create the embedding model\n",
        "CREATE OR REPLACE MODEL\n",
        "  `{PROJECT_ID}.{BQ_DATASET}.{MODEL_NAME}` OPTIONS(model_type='tensorflow',\n",
        "    model_path='{MODEL_PATH}');\n",
        "\n",
        "-- create the preprocessed table\n",
        "CREATE OR REPLACE TABLE `{PROJECT_ID}.{BQ_DATASET}.{PREPROCESSED_TABLE}`\n",
        "AS (\n",
        "  WITH\n",
        "    -- Apply the model for embedding generation\n",
        "    get_embeddings AS (\n",
        "      SELECT\n",
        "        title,\n",
        "        sentences,\n",
        "        output_0 as content_embeddings,\n",
        "        topics\n",
        "      FROM ML.PREDICT(MODEL `{PROJECT_ID}.{BQ_DATASET}.{MODEL_NAME}`,(\n",
        "        SELECT topics, title, content AS sentences\n",
        "        FROM `{PROJECT_ID}.{BQ_DATASET}.{BQ_TABLE}`\n",
        "      ))),\n",
        "    -- Get label\n",
        "    get_label AS (\n",
        "        SELECT\n",
        "            *,\n",
        "            STRUCT( CASE WHEN 'acq' in UNNEST(topics) THEN 1 ELSE 0 END as acq ) AS label,\n",
        "        FROM get_embeddings\n",
        "    ),\n",
        "    -- Train-serve splitting\n",
        "    get_split AS (\n",
        "        SELECT\n",
        "            *,\n",
        "            CASE WHEN ABS(MOD(FARM_FINGERPRINT(title), 10)) < 8 THEN 'TRAIN' ELSE 'PREDICT' END AS split\n",
        "        FROM get_label\n",
        "    )\n",
        "    -- create training table\n",
        "    SELECT\n",
        "        title,\n",
        "        sentences,\n",
        "        STRUCT( content_embeddings[OFFSET(0)] AS content_embed_0,\n",
        "                content_embeddings[OFFSET(1)] AS content_embed_1,\n",
        "                content_embeddings[OFFSET(2)] AS content_embed_2,\n",
        "                content_embeddings[OFFSET(3)] AS content_embed_3,\n",
        "                content_embeddings[OFFSET(4)] AS content_embed_4,\n",
        "                content_embeddings[OFFSET(5)] AS content_embed_5,\n",
        "                content_embeddings[OFFSET(6)] AS content_embed_6,\n",
        "                content_embeddings[OFFSET(7)] AS content_embed_7,\n",
        "                content_embeddings[OFFSET(8)] AS content_embed_8,\n",
        "                content_embeddings[OFFSET(9)] AS content_embed_9,\n",
        "                content_embeddings[OFFSET(10)] AS content_embed_10,\n",
        "                content_embeddings[OFFSET(11)] AS content_embed_11,\n",
        "                content_embeddings[OFFSET(12)] AS content_embed_12,\n",
        "                content_embeddings[OFFSET(13)] AS content_embed_13,\n",
        "                content_embeddings[OFFSET(14)] AS content_embed_14,\n",
        "                content_embeddings[OFFSET(15)] AS content_embed_15,\n",
        "                content_embeddings[OFFSET(16)] AS content_embed_16,\n",
        "                content_embeddings[OFFSET(17)] AS content_embed_17,\n",
        "                content_embeddings[OFFSET(18)] AS content_embed_18,\n",
        "                content_embeddings[OFFSET(19)] AS content_embed_19) AS feature,\n",
        "        label.acq as label,\n",
        "        split\n",
        "    FROM\n",
        "      get_split)\n",
        "\"\"\"\n",
        "\n",
        "with open(\n",
        "    f\"{KFP_COMPONENTS_PATH}/bq_preprocess_component/bq_preprocess_query.sql\", \"w\"\n",
        ") as q:\n",
        "    q.write(create_bq_preprocess_query)\n",
        "q.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yF9W5x4HgUQb"
      },
      "source": [
        "#### Create BigQuery Model query\n",
        "\n",
        "Below, you have a simple query to build a BigQuery ML Logistic Classifier model for topic's articles classification."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tjStSASggUQb"
      },
      "outputs": [],
      "source": [
        "create_bq_model_query = f\"\"\"\n",
        "CREATE OR REPLACE MODEL `{PROJECT_ID}.{BQ_DATASET}.{CLASSIFICATION_MODEL_NAME}`\n",
        "  OPTIONS (\n",
        "      model_type='logistic_reg',\n",
        "      input_label_cols=['label']) AS\n",
        "  SELECT\n",
        "      label,\n",
        "      feature.*\n",
        "  FROM\n",
        "     `{PROJECT_ID}.{BQ_DATASET}.{PREPROCESSED_TABLE}`\n",
        "  WHERE split = 'TRAIN';\n",
        "\"\"\"\n",
        "\n",
        "with open(f\"{KFP_COMPONENTS_PATH}/bq_model_component/create_bq_model.sql\", \"w\") as q:\n",
        "    q.write(create_bq_model_query)\n",
        "q.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RlLTcuUhdzFU"
      },
      "source": [
        "#### Create BigQuery Prediction query\n",
        "\n",
        "With the following query, you run a prediction job using the table with the preprocessing query."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rMwgXCyfdzFc"
      },
      "outputs": [],
      "source": [
        "create_bq_prediction_query = f\"\"\"SELECT title, sentences, feature.* FROM `{PROJECT_ID}.{BQ_DATASET}.{PREPROCESSED_TABLE}` WHERE split = 'PREDICT' \"\"\"\n",
        "\n",
        "with open(\n",
        "    f\"{KFP_COMPONENTS_PATH}/bq_prediction_component/create_bq_prediction_query.sql\", \"w\"\n",
        ") as q:\n",
        "    q.write(create_bq_prediction_query)\n",
        "q.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OxemUVCxAiSo"
      },
      "source": [
        "### Build the pipeline\n",
        "\n",
        "In this step, you build the pipeline using the individual components.\n",
        "\n",
        "Define the `JOB_NAME` and `JOB_CONFIG` below. `JOB_CONFIG` consists of the following parameters for the destination table:\n",
        "- `PROJECT_ID`: Id of the project.\n",
        "- `BQ_DATASET`: Id of the BigQuery dataset.\n",
        "- `PREDICT_TABLE`: Id of the BigQuery table where predictions are stored."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JzcBvqdpzjiz"
      },
      "outputs": [],
      "source": [
        "ID = random.randint(1, 10000)\n",
        "JOB_NAME = f\"reuters-preprocess-{UUID}-{ID}\"\n",
        "JOB_CONFIG = {\n",
        "    \"destinationTable\": {\n",
        "        \"projectId\": PROJECT_ID,\n",
        "        \"datasetId\": BQ_DATASET,\n",
        "        \"tableId\": PREDICT_TABLE,\n",
        "    }\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mdO8st_gLKBZ"
      },
      "source": [
        "#### Create a custom component for arguments\n",
        "\n",
        "Next, you create a component to pass arguments to the `DataflowPythonJobOp` component."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V0d0TYFzLUWK"
      },
      "outputs": [],
      "source": [
        "@component(base_image=\"python:3.8-slim\")\n",
        "def build_dataflow_args(\n",
        "    # destination_table: Input[Artifact],\n",
        "    bq_dataset: str,\n",
        "    bq_table: str,\n",
        "    job_name: str,\n",
        "    setup_file_uri: str,\n",
        "    runner: str,\n",
        "    inputs_uri: str,\n",
        ") -> list:\n",
        "    return [\n",
        "        \"--job_name\",\n",
        "        job_name,\n",
        "        \"--setup_file\",\n",
        "        setup_file_uri,\n",
        "        \"--runner\",\n",
        "        runner,\n",
        "        \"--inputs\",\n",
        "        inputs_uri,\n",
        "        \"--bq-dataset\",\n",
        "        bq_dataset,\n",
        "        \"--bq-table\",\n",
        "        bq_table,\n",
        "    ]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pcSL1FHk69KT"
      },
      "source": [
        "#### Create the pipeline\n",
        "\n",
        "Define the workflow of the pipeline and build the pipeline. The parameters passed to the pipeline include:\n",
        "\n",
        "- `create_bq_dataset_query`: SQL query to create the dataset in BigQuery.\n",
        "- `job_name`: Name of the Cloud Dataflow job to be configured in `PipelineOptions`.\n",
        "- `inputs_uri`: A directory location of input data.\n",
        "-  `bq_dataset`: Dataset name used in BigQuery.\n",
        "-  `bq_table`: Table name used in BigQuery for ingestion.\n",
        "-  `requirements_file_path`: The GCS path to the pip requirements file.\n",
        "-  `python_file_path`: The GCS path to the python file to run.\n",
        "-  `setup_file_uri`: Path to a Python setup file containing package dependencies.\n",
        "-  `temp_location`:  GCS path for Dataflow to stage temporary job files created during the execution of the pipeline.\n",
        "-  `runner`: Pipeline runner used to execute the workflow.\n",
        "-  `create_bq_preprocess_query`: SQL query to preprocess the data in BigQuery.\n",
        "-  `create_bq_model_query`: SQL query to create the BigQuery ML model.\n",
        "-  `create_bq_prediction_query`: SQL query for prediction.\n",
        "-  `job_config`: A json formatted string describing the job configuration. For more information, vist this [page]( https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#JobConfigurationQuery).\n",
        "-  `project`: Project ID.\n",
        "-  `region`: Selected region to run the Dataflow job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AlFXqsPIAk0l"
      },
      "outputs": [],
      "source": [
        "@dsl.pipeline(\n",
        "    name=\"mlops-bqml-text-generate-embeddings\",\n",
        "    description=\"A batch pipeline to generate embeddings\",\n",
        ")\n",
        "def pipeline(\n",
        "    create_bq_dataset_query: str,\n",
        "    job_name: str,\n",
        "    inputs_uri: str,\n",
        "    bq_dataset: str,\n",
        "    bq_table: str,\n",
        "    requirements_file_path: str,\n",
        "    python_file_path: str,\n",
        "    setup_file_uri: str,\n",
        "    temp_location: str,\n",
        "    runner: str,\n",
        "    create_bq_preprocess_query: str,\n",
        "    create_bq_model_query: str,\n",
        "    create_bq_prediction_query: str,\n",
        "    job_config: dict,\n",
        "    project: str = PROJECT_ID,\n",
        "    region: str = LOCATION,\n",
        "):\n",
        "\n",
        "    from google_cloud_pipeline_components.v1.bigquery import (\n",
        "        BigqueryCreateModelJobOp, BigqueryEvaluateModelJobOp,\n",
        "        BigqueryPredictModelJobOp, BigqueryQueryJobOp)\n",
        "    from google_cloud_pipeline_components.v1.dataflow import \\\n",
        "        DataflowPythonJobOp\n",
        "    from google_cloud_pipeline_components.v1.wait_gcp_resources import \\\n",
        "        WaitGcpResourcesOp\n",
        "\n",
        "    # create the dataset\n",
        "    bq_dataset_op = BigqueryQueryJobOp(\n",
        "        query=create_bq_dataset_query,\n",
        "        project=project,\n",
        "        location=\"US\",\n",
        "    )\n",
        "    # instantiate dataflow args\n",
        "    build_dataflow_args_op = build_dataflow_args(\n",
        "        job_name=job_name,\n",
        "        inputs_uri=inputs_uri,\n",
        "        # destination_table = bq_dataset_op.outputs['destination_table'],\n",
        "        bq_dataset=bq_dataset,\n",
        "        bq_table=bq_table,\n",
        "        setup_file_uri=setup_file_uri,\n",
        "        runner=runner,\n",
        "    ).after(bq_dataset_op)\n",
        "\n",
        "    # run dataflow job\n",
        "    dataflow_python_op = DataflowPythonJobOp(\n",
        "        requirements_file_path=requirements_file_path,\n",
        "        python_module_path=python_file_path,\n",
        "        args=build_dataflow_args_op.output,\n",
        "        project=project,\n",
        "        location=region,\n",
        "        temp_location=temp_location,\n",
        "    ).after(build_dataflow_args_op)\n",
        "\n",
        "    dataflow_wait_op = WaitGcpResourcesOp(\n",
        "        gcp_resources=dataflow_python_op.outputs[\"gcp_resources\"]\n",
        "    ).after(dataflow_python_op)\n",
        "\n",
        "    # run preprocessing job\n",
        "    bq_preprocess_op = BigqueryQueryJobOp(\n",
        "        query=create_bq_preprocess_query,\n",
        "        project=project,\n",
        "        location=\"US\",\n",
        "    ).after(dataflow_wait_op)\n",
        "\n",
        "    # create the logistic regression\n",
        "    bq_model_op = BigqueryCreateModelJobOp(\n",
        "        query=create_bq_model_query,\n",
        "        project=project,\n",
        "        location=\"US\",\n",
        "    ).after(bq_preprocess_op)\n",
        "\n",
        "    # evaluate the logistic regression\n",
        "    bq_evaluate_op = BigqueryEvaluateModelJobOp(\n",
        "        project=project, location=\"US\", model=bq_model_op.outputs[\"model\"]\n",
        "    ).after(bq_model_op)\n",
        "\n",
        "    # similuate prediction\n",
        "    BigqueryPredictModelJobOp(\n",
        "        model=bq_model_op.outputs[\"model\"],\n",
        "        query_statement=create_bq_prediction_query,\n",
        "        job_configuration_query=job_config,\n",
        "        project=project,\n",
        "        location=\"US\",\n",
        "    ).after(bq_evaluate_op)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nghLONQX7JNg"
      },
      "source": [
        "## Compile and Run the pipeline\n",
        "\n",
        "Pass the necessary constants and parameters to the pipeline and compile it to a yaml file."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8l6IR7OoADJV"
      },
      "outputs": [],
      "source": [
        "PIPELINE_ROOT = urlparse(BUCKET_URI)._replace(path=\"pipeline_root\").geturl()\n",
        "PIPELINE_PACKAGE = str(path(BUILD) / \"mlops_bqml_text_analyisis_pipeline.yaml\")\n",
        "REQUIREMENTS_URI = urlparse(BUCKET_URI)._replace(path=\"requirements.txt\").geturl()\n",
        "PYTHON_FILE_URI = urlparse(BUCKET_URI)._replace(path=\"src/ingest_pipeline.py\").geturl()\n",
        "MODEL_URI = urlparse(BUCKET_URI)._replace(path=\"swivel_text_embedding_model\").geturl()\n",
        "\n",
        "compiler.Compiler().compile(pipeline_func=pipeline, package_path=PIPELINE_PACKAGE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b5c1b7b8b290"
      },
      "source": [
        "Using the compiled yaml file, create Vertex AI Pipeline Job and run it by passing the `SERVICE_ACCOUNT` details configured earlier."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6rJ0NRuWwfnp"
      },
      "outputs": [],
      "source": [
        "pipeline = vertex_ai.PipelineJob(\n",
        "    display_name=f\"data_preprocess_{UUID}\",\n",
        "    template_path=PIPELINE_PACKAGE,\n",
        "    pipeline_root=PIPELINE_ROOT,\n",
        "    parameter_values={\n",
        "        \"create_bq_dataset_query\": create_bq_dataset_query,\n",
        "        \"bq_dataset\": BQ_DATASET,\n",
        "        \"job_name\": JOB_NAME,\n",
        "        \"inputs_uri\": INPUTS_URI,\n",
        "        \"bq_table\": BQ_TABLE,\n",
        "        \"requirements_file_path\": REQUIREMENTS_URI,\n",
        "        \"python_file_path\": PYTHON_FILE_URI,\n",
        "        \"setup_file_uri\": SETUP_FILE_URI,\n",
        "        \"temp_location\": PIPELINE_ROOT,\n",
        "        \"runner\": RUNNER,\n",
        "        \"create_bq_preprocess_query\": create_bq_preprocess_query,\n",
        "        \"create_bq_model_query\": create_bq_model_query,\n",
        "        \"create_bq_prediction_query\": create_bq_prediction_query,\n",
        "        \"job_config\": JOB_CONFIG,\n",
        "    },\n",
        "    enable_caching=False,\n",
        ")\n",
        "\n",
        "pipeline.run(service_account=SERVICE_ACCOUNT)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bcb582e65740"
      },
      "source": [
        "Once the pipeline job gets finished successfully, the trained model can be found created in the BigQuery dataset. Run the following cell to see the model listed in the output."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c0d194c006ae"
      },
      "outputs": [],
      "source": [
        "! bq ls $PROJECT_ID:$BQ_DATASET"
      ]
    },
    {
      "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 in the following cell. Set `delete_bucket` and `delete_dataset` to **True** to delete the Cloud Storage bucket and the Bigquery dataset used in this notebook respectively."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sx_vKniMq9ZX"
      },
      "outputs": [],
      "source": [
        "# delete the pipeline job\n",
        "pipeline.delete()\n",
        "\n",
        "# Delete the Cloud Storage bucket\n",
        "delete_bucket = False  # Set True for deletion\n",
        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI\n",
        "\n",
        "# delete dataset\n",
        "delete_dataset = False  # Set True for deletion\n",
        "if delete_dataset:\n",
        "    ! bq rm -r -f -d $PROJECT_ID:$BQ_DATASET\n",
        "\n",
        "# Remove the locally generated files\n",
        "!rm -r build\n",
        "!rm -r data\n",
        "!rm -r components\n",
        "!rm -r src\n",
        "!rm requirements.txt\n",
        "!rm setup.py"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "google_cloud_pipeline_components_bqml_text.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}