Beam_Summit_Demo.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Write your own model handler for RunInference\n"
      ],
      "metadata": {
        "id": "Y7EyCkrYGl2-"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Install Dependencies"
      ],
      "metadata": {
        "id": "dWyrdQj7Gx24"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eWeqbZ6EGh96"
      },
      "outputs": [],
      "source": [
        "!pip install apache_beam\n",
        "!pip install tensorflow"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import sys\n",
        "from typing import Any\n",
        "from typing import Dict\n",
        "from typing import Iterable\n",
        "from typing import Optional\n",
        "from typing import Sequence\n",
        "\n",
        "import numpy\n",
        "import tensorflow as tf\n",
        "from tensorflow import keras\n",
        "import tensorflow_hub as hub\n",
        "\n",
        "import apache_beam as beam\n",
        "from apache_beam.ml.inference.base import ModelHandler\n",
        "from apache_beam.ml.inference.base import PredictionResult\n",
        "from apache_beam.ml.inference.base import RunInference\n",
        "from apache_beam.ml.inference import utils\n"
      ],
      "metadata": {
        "id": "n4xk_NY_H2Lz"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Define the Model Handler Class\n",
        "\n",
        "Typically the notion to set the name of a model handler class is\n",
        "\n",
        "`<Framework>ModelHandler<InputType>`\n",
        "\n",
        "Eg: For TensorFlow framework and tensor input, it would be\n",
        "\n",
        "`TFModelHandlerTensor`\n"
      ],
      "metadata": {
        "id": "gB5YQJGEG-Nn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def default_tensor_inference_fn(\n",
        "    model: tf.Module,\n",
        "    batch: Sequence[tf.Tensor],\n",
        "    inference_args: Dict[str, Any],\n",
        "    model_id: Optional[str] = None) -> Iterable[PredictionResult]:\n",
        "    vectorized_batch = tf.stack(batch, axis=0)\n",
        "    predictions = model(vectorized_batch, **inference_args)\n",
        "    return utils._convert_to_result(batch, predictions, model_id)\n",
        "\n",
        "class TFModelHandlerTensor(ModelHandler[tf.Tensor, PredictionResult, tf.Module]):\n",
        "    def __init__(\n",
        "        self,\n",
        "        model_uri: str,\n",
        "        *,\n",
        "        load_model_args: Optional[Dict[str, Any]] = None,\n",
        "        inference_fn = default_tensor_inference_fn,\n",
        "        **kwargs):\n",
        "        self._model_uri = model_uri\n",
        "        self._inference_fn = inference_fn\n",
        "        self._load_model_args = {} if not load_model_args else load_model_args\n",
        "\n",
        "    def load_model(self) -> tf.Module:\n",
        "        model = tf.keras.models.load_model(hub.resolve(self._model_uri),  **self._load_model_args)\n",
        "        return model\n",
        "\n",
        "    def run_inference(\n",
        "        self,\n",
        "        batch: Sequence[numpy.ndarray],\n",
        "        model: tf.Module,\n",
        "        inference_args: Optional[Dict[str, Any]] = None\n",
        "    ) -> Iterable[PredictionResult]:\n",
        "        inference_args = {} if not inference_args else inference_args\n",
        "        return self._inference_fn(model, batch, inference_args, self._model_uri)\n",
        "\n",
        "    def update_model_path(self, model_path: Optional[str] = None):\n",
        "        self._model_uri = model_path if model_path else self._model_uri\n",
        "\n",
        "    def get_num_bytes(self, batch: Sequence[numpy.ndarray]) -> int:\n",
        "        return sum(sys.getsizeof(element) for element in batch)\n"
      ],
      "metadata": {
        "id": "htDqs6NEG3ax"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Create a simple model for testing"
      ],
      "metadata": {
        "id": "GWAQ5MAVJQO7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Create training data that represents the 5 times multiplication table for the numbers 0 to 99.\n",
        "# x is the data and y is the labels.\n",
        "x = tf.range(0, 100)   # Examples\n",
        "y = x * 5              # Labels\n",
        "\n",
        "# Use create_model to build a simple linear regression model.\n",
        "# Note that the model has a shape of (1) for its input layer and expects a single int64 value.\n",
        "def create_model():\n",
        "  input_layer = keras.layers.Input(shape=(1), dtype=tf.float32, name='x')\n",
        "  output_layer= keras.layers.Dense(1)(input_layer)\n",
        "  model = keras.Model(input_layer, output_layer)\n",
        "  model.compile(optimizer=tf.optimizers.Adam(), loss='mean_absolute_error')\n",
        "  return model\n",
        "\n",
        "model = create_model()\n",
        "model.fit(x, y, epochs=2000, verbose=0)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "g49zycIsJQAd",
        "outputId": "7826405e-7d3a-4851-cf4e-1030b85493d3"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<keras.callbacks.History at 0x7f57ad2bf1c0>"
            ]
          },
          "metadata": {},
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Save the model and use the path in the model handler."
      ],
      "metadata": {
        "id": "CNjPJX3GJpu5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "saved_model_path = \"./saved_models/\"\n",
        "model.save(saved_model_path)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HVF56UAoJweh",
        "outputId": "ff3c62f6-0e1e-41e0-df42-a208a6ca114a"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:absl:Found untraced functions such as _update_step_xla while saving (showing 1 of 1). These functions will not be directly callable after loading.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Run the pipeline"
      ],
      "metadata": {
        "id": "fF3BqVEAI9eK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class FormatOutput(beam.DoFn):\n",
        "  def process(self, element, *args, **kwargs):\n",
        "     yield \"example is {example} prediction is {prediction}\".format(example=element.example, prediction=element.inference)\n",
        "\n",
        "test_examples = [20, 40, 60, 90]\n",
        "value_to_predict = tf.constant(test_examples, dtype=tf.float32)"
      ],
      "metadata": {
        "id": "MLCCsevnHeHN"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model_handler = TFModelHandlerTensor(model_uri=saved_model_path)\n",
        "\n",
        "with beam.Pipeline() as p:\n",
        "    _ = (p\n",
        "         | beam.Create(value_to_predict)\n",
        "         | RunInference(model_handler)\n",
        "         | beam.ParDo(FormatOutput())\n",
        "         | beam.Map(print)\n",
        "         )\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5LAfIETZAqRQ",
        "outputId": "ea5e82b6-bc0f-4f89-ed71-0b786065c1ec"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "example is 20.0 prediction is [102.58883]\n",
            "example is 40.0 prediction is [201.53615]\n",
            "example is 60.0 prediction is [300.48346]\n",
            "example is 90.0 prediction is [448.90442]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "KJsPWdWFIWtM"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}