examples/notebooks/beam-ml/bigtable_enrichment_transform.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "fFjof1NgAJwu"
      },
      "outputs": [],
      "source": [
        "# @title ###### Licensed to the Apache Software Foundation (ASF), Version 2.0 (the \"License\")\n",
        "\n",
        "# Licensed to the Apache Software Foundation (ASF) under one\n",
        "# or more contributor license agreements. See the NOTICE file\n",
        "# distributed with this work for additional information\n",
        "# regarding copyright ownership. The ASF licenses this file\n",
        "# to you under the Apache License, Version 2.0 (the\n",
        "# \"License\"); you may not use this file except in compliance\n",
        "# with the License. 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,\n",
        "# software distributed under the License is distributed on an\n",
        "# \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY\n",
        "# KIND, either express or implied. See the License for the\n",
        "# specific language governing permissions and limitations\n",
        "# under the License"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A8xNRyZMW1yK"
      },
      "source": [
        "# Use Apache Beam and Bigtable to enrich data\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/apache/beam/blob/master/examples/notebooks/beam-ml/bigtable_enrichment_transform.ipynb\"><img src=\"https://raw.githubusercontent.com/google/or-tools/main/tools/colab_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/apache/beam/blob/master/examples/notebooks/beam-ml/bigtable_enrichment_transform.ipynb\"><img src=\"https://raw.githubusercontent.com/google/or-tools/main/tools/github_32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "</table>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HrCtxslBGK8Z"
      },
      "source": [
        "This notebook shows how to enrich data by using the Apache Beam [enrichment transform](https://beam.apache.org/documentation/transforms/python/elementwise/enrichment/) with [Bigtable](https://cloud.google.com/bigtable/docs/overview). The enrichment transform is an Apache Beam turnkey transform that lets you enrich data by using a key-value lookup. This transform has the following features:\n",
        "\n",
        "- The transform has a built-in Apache Beam handler that interacts with Bigtable to get data to use in the enrichment.\n",
        "- The enrichment transform uses client-side throttling to manage rate limiting the requests. The requests are exponentially backed off with a default retry strategy. You can configure rate limiting to suit your use case."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ltn5zrBiGS9C"
      },
      "source": [
        "This notebook demonstrates the following ecommerce use case:\n",
        "\n",
        "A stream of online transaction from [Pub/Sub](https://cloud.google.com/pubsub/docs/guides) contains the following fields: `sale_id`, `product_id`, `customer_id`, `quantity`, and `price`. Additional customer demographic data is stored in a separate Bigtable cluster. The demographic data is used to enrich the event stream from Pub/Sub. Then, the enriched data is used to predict the next product to recommended to a customer."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gVCtGOKTHMm4"
      },
      "source": [
        "## Before you begin\n",
        "Set up your environment and download dependencies."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YDHPlMjZRuY0"
      },
      "source": [
        "### Install Apache Beam\n",
        "To use the enrichment transform with the built-in Bigtable handler, install the Apache Beam SDK version 2.54.0 or later."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jBakpNZnAhqk"
      },
      "outputs": [],
      "source": [
        "!pip install torch\n",
        "!pip install apache_beam[interactive,gcp]==2.54.0 --quiet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SiJii48A2Rnb"
      },
      "outputs": [],
      "source": [
        "import datetime\n",
        "import json\n",
        "import math\n",
        "\n",
        "from typing import Any\n",
        "from typing import Dict\n",
        "\n",
        "import torch\n",
        "from google.cloud import pubsub_v1\n",
        "from google.cloud.bigtable import Client\n",
        "from google.cloud.bigtable import column_family\n",
        "\n",
        "import apache_beam as beam\n",
        "import apache_beam.runners.interactive.interactive_beam as ib\n",
        "from apache_beam.ml.inference.base import RunInference\n",
        "from apache_beam.ml.inference.pytorch_inference import PytorchModelHandlerTensor\n",
        "from apache_beam.options import pipeline_options\n",
        "from apache_beam.runners.interactive.interactive_runner import InteractiveRunner\n",
        "from apache_beam.transforms.enrichment import Enrichment\n",
        "from apache_beam.transforms.enrichment_handlers.bigtable import BigTableEnrichmentHandler"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X80jy3FqHjK4"
      },
      "source": [
        "### Authenticate with Google Cloud\n",
        "This notebook reads data from Pub/Sub and Bigtable. To use your Google Cloud account, authenticate this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Kz9sccyGBqz3"
      },
      "outputs": [],
      "source": [
        "from google.colab import auth\n",
        "auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nAmGgUMt48o9"
      },
      "source": [
        "Replace `<PROJECT_ID>`, `<INSTANCE_ID>`, and `<TABLE_ID>` with the appropriate values for your setup. These fields are used with Bigtable."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wEXucyi2liij"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"<PROJECT_ID>\"\n",
        "INSTANCE_ID = \"<INSTANCE_ID>\"\n",
        "TABLE_ID = \"<TABLE_ID>\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RpqZFfFfA_Dt"
      },
      "source": [
        "### Train the model\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8cUpV7mkB_xE"
      },
      "source": [
        "Create sample data by using the format `[product_id, quantity, price, customer_id, customer_location, recommend_product_id]`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TpxDHGObBEsj"
      },
      "outputs": [],
      "source": [
        "data = [\n",
        "    [3, 5, 127, 9, 'China', 7], [1, 6, 167, 5, 'Peru', 4], [5, 4, 91, 2, 'USA', 8], [7, 2, 52, 1, 'India', 4], [1, 8, 118, 3, 'UK', 8], [4, 6, 132, 8, 'Mexico', 2],\n",
        "    [6, 3, 154, 6, 'Brazil', 3], [4, 7, 163, 1, 'India', 7], [5, 2, 80, 4, 'Egypt', 9], [9, 4, 107, 7, 'Bangladesh', 1], [2, 9, 192, 8, 'Mexico', 4], [4, 5, 116, 5, 'Peru', 8],\n",
        "    [8, 1, 195, 1, 'India', 7], [8, 6, 153, 5, 'Peru', 1], [5, 3, 120, 6, 'Brazil', 2], [2, 7, 187, 7, 'Bangladesh', 4], [1, 8, 103, 6, 'Brazil', 8], [2, 9, 181, 1, 'India', 8],\n",
        "    [6, 5, 166, 3, 'UK', 5], [3, 4, 115, 8, 'Mexico', 1], [4, 7, 170, 4, 'Egypt', 2], [9, 3, 141, 7, 'Bangladesh', 3], [9, 3, 157, 1, 'India', 2], [7, 6, 128, 9, 'China', 1],\n",
        "    [1, 8, 102, 3, 'UK', 4], [5, 2, 107, 4, 'Egypt', 6], [6, 5, 164, 8, 'Mexico', 9], [4, 7, 188, 5, 'Peru', 1], [8, 1, 184, 1, 'India', 2], [8, 6, 198, 2, 'USA', 5],\n",
        "    [5, 3, 105, 6, 'Brazil', 7], [2, 7, 162, 7, 'Bangladesh', 7], [1, 8, 133, 9, 'China', 3], [2, 9, 173, 1, 'India', 7], [6, 5, 183, 5, 'Peru', 8], [3, 4, 191, 3, 'UK', 6],\n",
        "    [4, 7, 123, 2, 'USA', 5], [9, 3, 159, 8, 'Mexico', 2], [9, 3, 146, 4, 'Egypt', 8], [7, 6, 194, 1, 'India', 8], [3, 5, 112, 6, 'Brazil', 1], [4, 6, 101, 7, 'Bangladesh', 2],\n",
        "    [8, 1, 192, 4, 'Egypt', 4], [7, 2, 196, 5, 'Peru', 6], [9, 4, 124, 9, 'China', 7], [3, 4, 129, 5, 'Peru', 6], [6, 3, 151, 8, 'Mexico', 9], [5, 7, 114, 7, 'Bangladesh', 4],\n",
        "    [4, 7, 175, 6, 'Brazil', 5], [1, 8, 121, 1, 'India', 2], [4, 6, 187, 2, 'USA', 5], [6, 5, 144, 9, 'China', 9], [9, 4, 103, 5, 'Peru', 3], [5, 3, 84, 3, 'UK', 1],\n",
        "    [3, 5, 193, 2, 'USA', 4], [4, 7, 135, 1, 'India', 1], [7, 6, 148, 8, 'Mexico', 8], [1, 6, 160, 5, 'Peru', 7], [8, 6, 155, 6, 'Brazil', 9], [5, 7, 183, 7, 'Bangladesh', 2],\n",
        "    [2, 9, 125, 4, 'Egypt', 4], [6, 3, 111, 9, 'China', 9], [5, 2, 132, 3, 'UK', 3], [4, 5, 104, 7, 'Bangladesh', 7], [2, 7, 177, 8, 'Mexico', 7]]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bQt1cB4-CSBd"
      },
      "outputs": [],
      "source": [
        "countries_to_id = {'India': 1, 'USA': 2, 'UK': 3, 'Egypt': 4, 'Peru': 5,\n",
        "                   'Brazil': 6, 'Bangladesh': 7, 'Mexico': 8, 'China': 9}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y0Duet4nCdN1"
      },
      "source": [
        "Preprocess the data:\n",
        "\n",
        "1.   Convert the lists to tensors.\n",
        "2.   Separate the features from the expected prediction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7TT1O7sBCaZN"
      },
      "outputs": [],
      "source": [
        "X = [torch.tensor(item[:4]+[countries_to_id[item[4]]], dtype=torch.float) for item in data]\n",
        "Y = [torch.tensor(item[-1], dtype=torch.float) for item in data]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q6wB_ZsXDjjd"
      },
      "source": [
        "Define a simple model that has five input features and predicts a single value."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nphNfhUnESES"
      },
      "outputs": [],
      "source": [
        "def build_model(n_inputs, n_outputs):\n",
        "  \"\"\"build_model builds and returns a model that takes\n",
        "  `n_inputs` features and predicts `n_outputs` value\"\"\"\n",
        "  return torch.nn.Sequential(\n",
        "      torch.nn.Linear(n_inputs, 8),\n",
        "      torch.nn.ReLU(),\n",
        "      torch.nn.Linear(8, 16),\n",
        "      torch.nn.ReLU(),\n",
        "      torch.nn.Linear(16, n_outputs))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_sBSzDllEmCz"
      },
      "source": [
        "Train the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CaYrplaPDayp"
      },
      "outputs": [],
      "source": [
        "model = build_model(n_inputs=5, n_outputs=1)\n",
        "\n",
        "loss_fn = torch.nn.MSELoss()\n",
        "optimizer = torch.optim.Adam(model.parameters())\n",
        "\n",
        "for epoch in range(1000):\n",
        "  print(f'Epoch {epoch}: ---')\n",
        "  optimizer.zero_grad()\n",
        "  for i in range(len(X)):\n",
        "    pred = model(X[i])\n",
        "    loss = loss_fn(pred, Y[i])\n",
        "    loss.backward()\n",
        "  optimizer.step()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_rJYv8fFFPYb"
      },
      "source": [
        "Save the model to the `STATE_DICT_PATH` variable."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W4t260o9FURP"
      },
      "outputs": [],
      "source": [
        "STATE_DICT_PATH = './model.pth'\n",
        "\n",
        "torch.save(model.state_dict(), STATE_DICT_PATH)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ouMQZ4sC4zuO"
      },
      "source": [
        "### Set up the Bigtable table\n",
        "\n",
        "Create a sample Bigtable table for this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "E7Y4ipuL5kFD"
      },
      "outputs": [],
      "source": [
        "# Connect to the Bigtable instance. If you don't have admin access, then drop `admin=True`.\n",
        "client = Client(project=PROJECT_ID, admin=True)\n",
        "instance = client.instance(INSTANCE_ID)\n",
        "\n",
        "# Create a column family.\n",
        "column_family_id = 'demograph'\n",
        "max_versions_rule = column_family.MaxVersionsGCRule(2)\n",
        "column_families = {column_family_id: max_versions_rule}\n",
        "\n",
        "# Create a table.\n",
        "table = instance.table(TABLE_ID)\n",
        "\n",
        "# You need admin access to use `.exists()`. If you don't have the admin access, then\n",
        "# comment out the if-else block.\n",
        "if not table.exists():\n",
        "  table.create(column_families=column_families)\n",
        "else:\n",
        "  print(\"Table %s already exists in %s:%s\" % (TABLE_ID, PROJECT_ID, INSTANCE_ID))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eQLkSg3p7WAm"
      },
      "source": [
        "Add rows to the table for the enrichment example."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LI6oYkZ97Vtu",
        "outputId": "c72b28b5-8692-40f5-f8da-85622437d8f7"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Inserted row for key: 1\n",
            "Inserted row for key: 2\n",
            "Inserted row for key: 3\n"
          ]
        }
      ],
      "source": [
        "# Define column names for the table.\n",
        "customer_id = 'customer_id'\n",
        "customer_name = 'customer_name'\n",
        "customer_location = 'customer_location'\n",
        "\n",
        "# The following data is sample data to insert into Bigtable.\n",
        "customers = [\n",
        "  {\n",
        "    'customer_id': 1, 'customer_name': 'Sam', 'customer_location': 'India'\n",
        "  },\n",
        "  {\n",
        "    'customer_id': 2, 'customer_name': 'John', 'customer_location': 'USA'\n",
        "  },\n",
        "  {\n",
        "    'customer_id': 3, 'customer_name': 'Travis', 'customer_location': 'UK'\n",
        "  },\n",
        "]\n",
        "\n",
        "for customer in customers:\n",
        "  row_key = str(customer[customer_id]).encode()\n",
        "  row = table.direct_row(row_key)\n",
        "  row.set_cell(\n",
        "    column_family_id,\n",
        "    customer_id.encode(),\n",
        "    str(customer[customer_id]),\n",
        "    timestamp=datetime.datetime.utcnow())\n",
        "  row.set_cell(\n",
        "    column_family_id,\n",
        "    customer_name.encode(),\n",
        "    customer[customer_name],\n",
        "    timestamp=datetime.datetime.utcnow())\n",
        "  row.set_cell(\n",
        "    column_family_id,\n",
        "    customer_location.encode(),\n",
        "    customer[customer_location],\n",
        "    timestamp=datetime.datetime.utcnow())\n",
        "  row.commit()\n",
        "  print('Inserted row for key: %s' % customer[customer_id])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pHODouJDwc60"
      },
      "source": [
        "### Publish messages to Pub/Sub\n",
        "\n",
        "Use the Pub/Sub Python client to publish messages.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QKCuwDioxw-f"
      },
      "outputs": [],
      "source": [
        "# Replace <TOPIC_NAME> with the name of your Pub/Sub topic.\n",
        "TOPIC = \"<TOPIC_NAME>\"\n",
        "\n",
        "# Replace <SUBSCRIPTION_PATH> with the subscription for your topic.\n",
        "SUBSCRIPTION = \"<SUBSCRIPTION_PATH>\"\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MaCJwaPexPKZ"
      },
      "outputs": [],
      "source": [
        "messages = [\n",
        "    {'sale_id': i, 'customer_id': i, 'product_id': i, 'quantity': i, 'price': i*100}\n",
        "    for i in range(1,4)\n",
        "  ]\n",
        "\n",
        "publisher = pubsub_v1.PublisherClient()\n",
        "topic_name = publisher.topic_path(PROJECT_ID, TOPIC)\n",
        "\n",
        "for message in messages:\n",
        "  data = json.dumps(message).encode('utf-8')\n",
        "  publish_future = publisher.publish(topic_name, data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zPSFEMm02omi"
      },
      "source": [
        "## Use the Bigtable enrichment handler\n",
        "\n",
        "The [`BigTableEnrichmentHandler`](https://beam.apache.org/releases/pydoc/current/apache_beam.transforms.enrichment_handlers.bigtable.html#apache_beam.transforms.enrichment_handlers.bigtable.BigTableEnrichmentHandler) is a built-in handler included in the Apache Beam SDK versions 2.54.0 and later."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K41xhvmA5yQk"
      },
      "source": [
        "Configure the `BigTableEnrichmentHandler` handler with the following required parameters:\n",
        "\n",
        "* `project_id`: the Google Cloud project ID for the Bigtable instance\n",
        "* `instance_id`: the instance name of the Bigtable cluster\n",
        "* `table_id`: the table ID of table containing relevant data\n",
        "* `row_key`: The field name from the input row that contains the row key to use when querying Bigtable.\n",
        "\n",
        "Optionally, you can use parameters to further configure the `BigTableEnrichmentHandler` handler. For more information about the available parameters, see the [enrichment handler module documentation](https://beam.apache.org/releases/pydoc/current/apache_beam.transforms.enrichment_handlers.bigtable.html#module-apache_beam.transforms.enrichment_handlers.bigtable)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yFMcaf8i7TbI"
      },
      "source": [
        "**Note:** When exceptions occur, by default, the logging severity is set to warning ([`ExceptionLevel.WARN`](https://beam.apache.org/releases/pydoc/current/apache_beam.transforms.enrichment_handlers.utils.html#apache_beam.transforms.enrichment_handlers.utils.ExceptionLevel.WARN)).  To configure the severity to raise exceptions, set `exception_level` to [`ExceptionLevel.RAISE`](https://beam.apache.org/releases/pydoc/current/apache_beam.transforms.enrichment_handlers.utils.html#apache_beam.transforms.enrichment_handlers.utils.ExceptionLevel.RAISE). To ignore exceptions, set `exception_level` to [`ExceptionLevel.QUIET`](https://beam.apache.org/releases/pydoc/current/apache_beam.transforms.enrichment_handlers.utils.html#apache_beam.transforms.enrichment_handlers.utils.ExceptionLevel.QUIET).\n",
        "\n",
        "The following example demonstrates how to set the exception level in the `BigTableEnrichmentHandler` handler:\n",
        "\n",
        "```\n",
        "bigtable_handler = BigTableEnrichmentHandler(project_id=PROJECT_ID,\n",
        "                                             instance_id=INSTANCE_ID,\n",
        "                                             table_id=TABLE_ID,\n",
        "                                             row_key=row_key,\n",
        "                                             exception_level=ExceptionLevel.RAISE)\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The `row_key` parameter represents the field in input schema (`beam.Row`) that contains the row key for a row in the table.\n",
        "\n",
        "Starting with Apache Beam version 2.54.0, you can perform either of the following tasks when a table uses composite row keys:\n",
        "* Modify the input schema to contain the row key in the format required by Bigtable.\n",
        "* Use a custom enrichment handler. For more information, see the [example handler with composite row key support](https://www.toptal.com/developers/paste-gd/BYFGUL08#)."
      ],
      "metadata": {
        "id": "UEpjy_IsW4P4"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3dB26jhI45gd"
      },
      "outputs": [],
      "source": [
        "row_key = 'customer_id'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cr1j_DHK4gA4"
      },
      "outputs": [],
      "source": [
        "bigtable_handler = BigTableEnrichmentHandler(project_id=PROJECT_ID,\n",
        "                                             instance_id=INSTANCE_ID,\n",
        "                                             table_id=TABLE_ID,\n",
        "                                             row_key=row_key)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-Lvo8O2V-0Ey"
      },
      "source": [
        "## Use the enrichment transform\n",
        "\n",
        "To use the [enrichment transform](https://beam.apache.org/releases/pydoc/current/apache_beam.transforms.enrichment.html#apache_beam.transforms.enrichment.Enrichment), the enrichment handler parameter is the only required parameter."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xJTCfSmiV1kv"
      },
      "source": [
        "The following example demonstrates the code needed to add this transform to your pipeline.\n",
        "\n",
        "\n",
        "```\n",
        "with beam.Pipeline() as p:\n",
        "  output = (p\n",
        "            ...\n",
        "            | \"Enrich with BigTable\" >> Enrichment(bigtable_handler)\n",
        "            | \"RunInference\" >> RunInference(model_handler)\n",
        "            ...\n",
        "            )\n",
        "```\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F-xjiP_pHWZr"
      },
      "source": [
        "By default, the enrichment transform performs a [`cross_join`](https://beam.apache.org/releases/pydoc/current/apache_beam.transforms.enrichment.html#apache_beam.transforms.enrichment.cross_join). This join returns the enriched row with the following fields: `sale_id`, `customer_id`, `product_id`, `quantity`, `price`, and `customer_location`.\n",
        "\n",
        "To make a prediction when running the ecommerce example, however, the trained model needs the following fields: `product_id`, `quantity`, `price`, `customer_id`, and `customer_location`.\n",
        "\n",
        "Therefore, to get the required fields for the ecommerce example, design a custom join function that takes two dictionaries as input and returns an enriched row that include these fields."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8LnCnEPNIPtg"
      },
      "outputs": [],
      "source": [
        "def custom_join(left: Dict[str, Any], right: Dict[str, Any]):\n",
        "  enriched = {}\n",
        "  enriched['product_id'] = left['product_id']\n",
        "  enriched['quantity'] = left['quantity']\n",
        "  enriched['price'] = left['price']\n",
        "  enriched['customer_id'] = left['customer_id']\n",
        "  enriched['customer_location'] = right['demograph']['customer_location']\n",
        "  return beam.Row(**enriched)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "To provide a `lambda` function for using a custom join with the enrichment transform, see the following example.\n",
        "\n",
        "```\n",
        "with beam.Pipeline() as p:\n",
        "  output = (p\n",
        "            ...\n",
        "            | \"Enrich with BigTable\" >> Enrichment(bigtable_handler, join_fn=custom_join)\n",
        "            | \"RunInference\" >> RunInference(model_handler)\n",
        "            ...\n",
        "            )\n",
        "```"
      ],
      "metadata": {
        "id": "fe3bIclV1jZ5"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Because the enrichment transform makes API calls to the remote service, use the `timeout` parameter to specify a timeout duration of 10 seconds:\n",
        "\n",
        "```\n",
        "with beam.Pipeline() as p:\n",
        "  output = (p\n",
        "            ...\n",
        "            | \"Enrich with BigTable\" >> Enrichment(bigtable_handler, join_fn=custom_join, timeout=10)\n",
        "            | \"RunInference\" >> RunInference(model_handler)\n",
        "            ...\n",
        "            )\n",
        "```"
      ],
      "metadata": {
        "id": "uilxdknE3ihO"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CX9Cqybu6scV"
      },
      "source": [
        "## Use the `PyTorchModelHandlerTensor` interface to run inference\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zy5Jl7_gLklX"
      },
      "source": [
        "Because the enrichment transform outputs data in the format `beam.Row`, to make it compatible with the [`PyTorchModelHandlerTensor`](https://beam.apache.org/releases/pydoc/current/apache_beam.ml.inference.pytorch_inference.html#apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerTensor) interface,  convert it to `torch.tensor`. Additionally, the enriched field `customer_location` is a `string` type, but the model requires a `float` type. Convert the `customer_location` field to a `float` type."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KBKoB06nL4LF"
      },
      "outputs": [],
      "source": [
        "def convert_row_to_tensor(element: beam.Row):\n",
        "  row_dict = element._asdict()\n",
        "  row_dict['customer_location'] = countries_to_id[row_dict['customer_location']]\n",
        "  return torch.tensor(list(row_dict.values()), dtype=torch.float)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-tGHyB_vL3rJ"
      },
      "source": [
        "Initialize the model handler with the preprocessing function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VqUUEwcU-r2e"
      },
      "outputs": [],
      "source": [
        "model_handler = PytorchModelHandlerTensor(state_dict_path=STATE_DICT_PATH,\n",
        "                                          model_class=build_model,\n",
        "                                          model_params={'n_inputs':5, 'n_outputs':1}\n",
        "                                          ).with_preprocess_fn(convert_row_to_tensor)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vNHI4gVgNec2"
      },
      "source": [
        "Define a `DoFn` to format the output."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rkN-_Yf4Nlwy"
      },
      "outputs": [],
      "source": [
        "class PostProcessor(beam.DoFn):\n",
        "  def process(self, element, *args, **kwargs):\n",
        "    print('Customer %d who bought product %d is recommended to buy product %d' % (element.example[3], element.example[0], math.ceil(element.inference[0])))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0a1zerXycQ0z"
      },
      "source": [
        "## Run the pipeline\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WrwY0_gV_IDK"
      },
      "source": [
        "Configure the pipeline to run in streaming mode."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t0425sYBsYtB"
      },
      "outputs": [],
      "source": [
        "options = pipeline_options.PipelineOptions()\n",
        "options.view_as(pipeline_options.StandardOptions).streaming = True # Streaming mode is set True"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DBNijQDY_dRe"
      },
      "source": [
        "Pub/Sub sends the data in bytes. Convert the data to `beam.Row` objects by using a `DoFn`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sRw9iL8pKP5O"
      },
      "outputs": [],
      "source": [
        "class DecodeBytes(beam.DoFn):\n",
        "  \"\"\"\n",
        "  The DecodeBytes `DoFn` converts the data read from Pub/Sub to `beam.Row`.\n",
        "  First, decode the encoded string. Convert the output to\n",
        "  a `dict` with `json.loads()`, which is used to create a `beam.Row`.\n",
        "  \"\"\"\n",
        "  def process(self, element, *args, **kwargs):\n",
        "    element_dict = json.loads(element.decode('utf-8'))\n",
        "    yield beam.Row(**element_dict)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xofUJym-_GuB"
      },
      "source": [
        "Use the following code to run the pipeline.\n",
        "\n",
        "**Note:** Because this pipeline is a streaming pipeline, you need to manually stop the cell. If you don't stop the cell, the pipeline continues to run."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "St07XoibcQSb",
        "outputId": "34e0a603-fb77-455c-e40b-d15b672edeb2"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Customer 1 who bought product 1 is recommended to buy product 3\n",
            "Customer 2 who bought product 2 is recommended to buy product 5\n",
            "Customer 3 who bought product 3 is recommended to buy product 7\n"
          ]
        }
      ],
      "source": [
        "with beam.Pipeline(options=options) as p:\n",
        "  _ = (p\n",
        "       | \"Read from Pub/Sub\" >> beam.io.ReadFromPubSub(subscription=SUBSCRIPTION)\n",
        "       | \"ConvertToRow\" >> beam.ParDo(DecodeBytes())\n",
        "       | \"Enrichment\" >> Enrichment(bigtable_handler, join_fn=custom_join, timeout=10)\n",
        "       | \"RunInference\" >> RunInference(model_handler)\n",
        "       | \"Format Output\" >> beam.ParDo(PostProcessor())\n",
        "       )\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true,
      "collapsed_sections": [
        "RpqZFfFfA_Dt"
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}