model_card_toolkit/documentation/examples/MLMD_Model_Card_Toolkit_Demo.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "#@title 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": "23R0Z9RojXYW"
      },
      "source": [
        "# MLMD Model Card Toolkit Demo"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfBg1C5NB3X0"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/responsible_ai/model_card_toolkit/examples/MLMD_Model_Card_Toolkit_Demo\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/model-card-toolkit/blob/master/model_card_toolkit/documentation/examples/MLMD_Model_Card_Toolkit_Demo.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/model-card-toolkit/blob/master/model_card_toolkit/documentation/examples/MLMD_Model_Card_Toolkit_Demo.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/model-card-toolkit/model_card_toolkit/documentation/examples/MLMD_Model_Card_Toolkit_Demo.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sfSQ-kX-MLEr"
      },
      "source": [
        "## Background\n",
        "\n",
        "This notebook demonstrates how to generate a model card using the Model Card Toolkit with MLMD and TFX pipeline in a Jupyter/Colab environment. You can learn more about model cards at https://modelcards.withgoogle.com/about. \n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2GivNBNYjb3b"
      },
      "source": [
        "## Setup\n",
        "We first need to a) install and import the necessary packages, and b) download the data."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Fmgi8ZvQkScg"
      },
      "source": [
        "### Upgrade to Pip 21 (or later) and Install Model Card Toolkit"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install --upgrade pip==21.3\n",
        "!pip install model-card-toolkit"
      ],
      "metadata": {
        "id": "NYtxxdriz5VO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EwT0nov5QO1M"
      },
      "source": [
        "####*Did you restart the runtime?*\n",
        "\n",
        "If you are using Google Colab, the runtime must be restarted after installing new packages."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N-ePgV0Lj68Q"
      },
      "source": [
        "### Import packages\n",
        "\n",
        "We import necessary packages, including standard TFX component classes and check the library versions.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YIqpWK9efviJ"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import pprint\n",
        "import tempfile\n",
        "import urllib\n",
        "\n",
        "import absl\n",
        "import random\n",
        "import tensorflow.compat.v2 as tf\n",
        "import tensorflow_model_analysis as tfma\n",
        "tf.get_logger().propagate = False\n",
        "pp = pprint.PrettyPrinter()\n",
        "\n",
        "import tfx\n",
        "from tfx.components import CsvExampleGen\n",
        "from tfx.components import Evaluator\n",
        "from tfx.components import Pusher\n",
        "from tfx.components import SchemaGen\n",
        "from tfx.components import StatisticsGen\n",
        "from tfx.components import Trainer\n",
        "from tfx.components import Transform\n",
        "from tfx.components.trainer.executor import GenericExecutor\n",
        "from tfx.dsl.components.base import executor_spec\n",
        "from tfx.dsl.experimental import latest_blessed_model_resolver\n",
        "from tfx.orchestration import metadata\n",
        "from tfx.orchestration import pipeline\n",
        "from tfx.orchestration.experimental.interactive.interactive_context import InteractiveContext\n",
        "from tfx.proto import pusher_pb2\n",
        "from tfx.proto import trainer_pb2\n",
        "from tfx.types import Channel\n",
        "from tfx.types.standard_artifacts import Model\n",
        "from tfx.types.standard_artifacts import ModelBlessing\n",
        "\n",
        "import ml_metadata as mlmd"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Ensure TensorFlow 2 is running and executing eagerly."
      ],
      "metadata": {
        "id": "dDyZyTT02ysB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "tf.enable_v2_behavior()\n",
        "tf.executing_eagerly()"
      ],
      "metadata": {
        "id": "XuUdIuFIwvZR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eZ4K18_DN2D8"
      },
      "outputs": [],
      "source": [
        "print('TensorFlow version: {}'.format(tf.__version__))\n",
        "print('TFX version: {}'.format(tfx.version.__version__))\n",
        "print('MLMD version: {}'.format(mlmd.__version__))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ufJKQ6OvkJlY"
      },
      "source": [
        "### Set up pipeline paths"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ad5JLpKbf6sN"
      },
      "outputs": [],
      "source": [
        "# This is the root directory for your TFX pip package installation.\n",
        "_tfx_root = tfx.__path__\n",
        "\n",
        "# Set up logging.\n",
        "absl.logging.set_verbosity(absl.logging.INFO)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n2cMMAbSkGfX"
      },
      "source": [
        "### Download example data\n",
        "We download the example dataset for use in our TFX pipeline.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BywX6OUEhAqn"
      },
      "outputs": [],
      "source": [
        "DATA_PATH = 'https://archive.ics.uci.edu/ml/machine-learning-databases/adult/' \\\n",
        "   'adult.data'\n",
        "_data_root = tempfile.mkdtemp(prefix='tfx-data')\n",
        "_data_filepath = os.path.join(_data_root, \"data.csv\")\n",
        "urllib.request.urlretrieve(DATA_PATH, _data_filepath)\n",
        "\n",
        "columns = [\n",
        "  \"Age\", \"Workclass\", \"fnlwgt\", \"Education\", \"Education-Num\", \"Marital-Status\",\n",
        "  \"Occupation\", \"Relationship\", \"Race\", \"Sex\", \"Capital-Gain\", \"Capital-Loss\",\n",
        "  \"Hours-per-week\", \"Country\", \"Over-50K\"]\n",
        "\n",
        "with open(_data_filepath, 'r') as f:\n",
        "  content = f.read()\n",
        "  content = content.replace(\", <=50K\", ', 0').replace(\", >50K\", ', 1')\n",
        "\n",
        "with open(_data_filepath, 'w') as f:\n",
        "  f.write(','.join(columns) + '\\n' + content)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "blZC1sIQOWfH"
      },
      "source": [
        "Take a quick look at the CSV file."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c5YPeLPFOXaD"
      },
      "outputs": [],
      "source": [
        "!head {_data_filepath}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8ONIE_hdkPS4"
      },
      "source": [
        "### Create the InteractiveContext\n",
        "Last, we create an InteractiveContext, which will allow us to run TFX components interactively in this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0Rh6K5sUf9dd"
      },
      "outputs": [],
      "source": [
        "# Here, we create an InteractiveContext using default parameters. This will\n",
        "# use a temporary directory with an ephemeral ML Metadata database instance.\n",
        "# To use your own pipeline root or database, the optional properties\n",
        "# `pipeline_root` and `metadata_connection_config` may be passed to\n",
        "# InteractiveContext. Calls to InteractiveContext are no-ops outside of the\n",
        "# notebook.\n",
        "context = InteractiveContext(pipeline_name=\"Census Income Classification Pipeline\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HdQWxfsVkzdJ"
      },
      "source": [
        "## Run TFX components interactively\n",
        "In the cells that follow, we create TFX components one-by-one, run each of them, and visualize their output artifacts. In this notebook, we won’t provide detailed explanations of each TFX component, but you can see what each does at [TFX Colab workshop](https://github.com/tensorflow/workshops/blob/master/tfx_labs/Lab_1_Pipeline_in_Colab.ipynb)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L9fwt9gQk3BR"
      },
      "source": [
        "### ExampleGen\n",
        "\n",
        "Create the `ExampleGen` component to split data into training and evaluation sets, convert the data into `tf.Example` format, and copy data into the `_tfx_root` directory for other components to access. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PyXjuMt8f-9u"
      },
      "outputs": [],
      "source": [
        "example_gen = CsvExampleGen(input_base=_data_root)\n",
        "context.run(example_gen)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "880KkTAkPeUg"
      },
      "outputs": [],
      "source": [
        "artifact = example_gen.outputs['examples'].get()[0]\n",
        "print(artifact.split_names, artifact.uri)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J6vcbW_wPqvl"
      },
      "source": [
        "Let’s take a look at the first three training examples:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H4XIXjiCPwzQ"
      },
      "outputs": [],
      "source": [
        "# Get the URI of the output artifact representing the training examples, which is a directory\n",
        "train_uri = os.path.join(example_gen.outputs['examples'].get()[0].uri, 'Split-train')\n",
        "\n",
        "# Get the list of files in this directory (all compressed TFRecord files)\n",
        "tfrecord_filenames = [os.path.join(train_uri, name)\n",
        "                      for name in os.listdir(train_uri)]\n",
        "\n",
        "# Create a `TFRecordDataset` to read these files\n",
        "dataset = tf.data.TFRecordDataset(tfrecord_filenames, compression_type=\"GZIP\")\n",
        "\n",
        "# Iterate over the first 3 records and decode them.\n",
        "for tfrecord in dataset.take(3):\n",
        "  serialized_example = tfrecord.numpy()\n",
        "  example = tf.train.Example()\n",
        "  example.ParseFromString(serialized_example)\n",
        "  pp.pprint(example)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "csM6BFhtk5Aa"
      },
      "source": [
        "### StatisticsGen\n",
        "\n",
        "`StatisticsGen` takes as input the dataset we just ingested using `ExampleGen` and allows you to perform some analysis of your dataset using TensorFlow Data Validation (TFDV)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MAscCCYWgA-9"
      },
      "outputs": [],
      "source": [
        "statistics_gen = StatisticsGen(\n",
        "    examples=example_gen.outputs['examples'])\n",
        "context.run(statistics_gen)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0BDfOjGy048O"
      },
      "source": [
        "After `StatisticsGen` finishes running, we can visualize the outputted statistics. Try playing with the different plots!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tLjXy7K6Tp_G"
      },
      "outputs": [],
      "source": [
        "context.show(statistics_gen.outputs['statistics'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HLKLTO9Nk60p"
      },
      "source": [
        "### SchemaGen\n",
        "\n",
        "`SchemaGen` will take as input the statistics that we generated with `StatisticsGen`, looking at the training split by default."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ygQvZ6hsiQ_J"
      },
      "outputs": [],
      "source": [
        "schema_gen = SchemaGen(\n",
        "    statistics=statistics_gen.outputs['statistics'],\n",
        "    infer_feature_shape=False)\n",
        "context.run(schema_gen)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ec9vqDXpXeMb"
      },
      "outputs": [],
      "source": [
        "context.show(schema_gen.outputs['schema'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kZWWdbA-m7zp"
      },
      "source": [
        "To learn more about schemas, see [the SchemaGen documentation](https://www.tensorflow.org/tfx/guide/schemagen)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JPViEz5RlA36"
      },
      "source": [
        "### Transform\n",
        "\n",
        "`Transform` will take as input the data from `ExampleGen`, the schema from `SchemaGen`, as well as a module that contains user-defined Transform code.\n",
        "\n",
        "Let's see an example of user-defined Transform code below (for an introduction to the TensorFlow Transform APIs, [see the tutorial](https://www.tensorflow.org/tfx/tutorials/transform/simple)).\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PuNSiUKb4YJf"
      },
      "outputs": [],
      "source": [
        "_census_income_constants_module_file = 'census_income_constants.py'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HPjhXuIF4YJh",
        "jupyter": {
          "source_hidden": true
        }
      },
      "outputs": [],
      "source": [
        "%%writefile {_census_income_constants_module_file}\n",
        "\n",
        "# Categorical features are assumed to each have a maximum value in the dataset.\n",
        "MAX_CATEGORICAL_FEATURE_VALUES = [20]\n",
        "\n",
        "CATEGORICAL_FEATURE_KEYS = [\"Education-Num\"]\n",
        "\n",
        "\n",
        "DENSE_FLOAT_FEATURE_KEYS = [\"Capital-Gain\", \"Hours-per-week\", \"Capital-Loss\"]\n",
        "\n",
        "# Number of buckets used by tf.transform for encoding each feature.\n",
        "FEATURE_BUCKET_COUNT = 10\n",
        "\n",
        "BUCKET_FEATURE_KEYS = [\"Age\"]\n",
        "\n",
        "# Number of vocabulary terms used for encoding VOCAB_FEATURES by tf.transform\n",
        "VOCAB_SIZE = 200\n",
        "\n",
        "# Count of out-of-vocab buckets in which unrecognized VOCAB_FEATURES are hashed.\n",
        "OOV_SIZE = 10\n",
        "\n",
        "VOCAB_FEATURE_KEYS = [\"Workclass\", \"Education\", \"Marital-Status\", \"Occupation\", \n",
        "                      \"Relationship\", \"Race\", \"Sex\", \"Country\"]\n",
        "\n",
        "# Keys\n",
        "LABEL_KEY = \"Over-50K\"\n",
        "\n",
        "def transformed_name(key):\n",
        "  return key + '_xf'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4AJ9hBs94YJm"
      },
      "outputs": [],
      "source": [
        "_census_income_transform_module_file = 'census_income_transform.py'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MYmxxx9A4YJn",
        "jupyter": {
          "source_hidden": true
        }
      },
      "outputs": [],
      "source": [
        "%%writefile {_census_income_transform_module_file}\n",
        "\n",
        "import tensorflow as tf\n",
        "import tensorflow_transform as tft\n",
        "\n",
        "import census_income_constants\n",
        "\n",
        "_DENSE_FLOAT_FEATURE_KEYS = census_income_constants.DENSE_FLOAT_FEATURE_KEYS\n",
        "_VOCAB_FEATURE_KEYS = census_income_constants.VOCAB_FEATURE_KEYS\n",
        "_VOCAB_SIZE = census_income_constants.VOCAB_SIZE\n",
        "_OOV_SIZE = census_income_constants.OOV_SIZE\n",
        "_FEATURE_BUCKET_COUNT = census_income_constants.FEATURE_BUCKET_COUNT\n",
        "_BUCKET_FEATURE_KEYS = census_income_constants.BUCKET_FEATURE_KEYS\n",
        "_CATEGORICAL_FEATURE_KEYS = census_income_constants.CATEGORICAL_FEATURE_KEYS\n",
        "_LABEL_KEY = census_income_constants.LABEL_KEY\n",
        "_transformed_name = census_income_constants.transformed_name\n",
        "\n",
        "\n",
        "def preprocessing_fn(inputs):\n",
        "  \"\"\"tf.transform's callback function for preprocessing inputs.\n",
        "  Args:\n",
        "    inputs: map from feature keys to raw not-yet-transformed features.\n",
        "  Returns:\n",
        "    Map from string feature key to transformed feature operations.\n",
        "  \"\"\"\n",
        "  outputs = {}\n",
        "  for key in _DENSE_FLOAT_FEATURE_KEYS:\n",
        "    # Preserve this feature as a dense float, setting nan's to the mean.\n",
        "    outputs[_transformed_name(key)] = tft.scale_to_z_score(\n",
        "        _fill_in_missing(inputs[key]))\n",
        "\n",
        "  for key in _VOCAB_FEATURE_KEYS:\n",
        "    # Build a vocabulary for this feature.\n",
        "    outputs[_transformed_name(key)] = tft.compute_and_apply_vocabulary(\n",
        "        _fill_in_missing(inputs[key]),\n",
        "        top_k=_VOCAB_SIZE,\n",
        "        num_oov_buckets=_OOV_SIZE)\n",
        "\n",
        "  for key in _BUCKET_FEATURE_KEYS:\n",
        "    outputs[_transformed_name(key)] = tft.bucketize(\n",
        "        _fill_in_missing(inputs[key]), _FEATURE_BUCKET_COUNT)\n",
        "\n",
        "  for key in _CATEGORICAL_FEATURE_KEYS:\n",
        "    outputs[_transformed_name(key)] = _fill_in_missing(inputs[key])\n",
        "\n",
        "  label = _fill_in_missing(inputs[_LABEL_KEY])\n",
        "  outputs[_transformed_name(_LABEL_KEY)] = label\n",
        "  \n",
        "  return outputs\n",
        "\n",
        "\n",
        "def _fill_in_missing(x):\n",
        "  \"\"\"Replace missing values in a SparseTensor.\n",
        "  Fills in missing values of `x` with '' or 0, and converts to a dense tensor.\n",
        "  Args:\n",
        "    x: A `SparseTensor` of rank 2.  Its dense shape should have size at most 1\n",
        "      in the second dimension.\n",
        "  Returns:\n",
        "    A rank 1 tensor where missing values of `x` have been filled in.\n",
        "  \"\"\"\n",
        "  default_value = '' if x.dtype == tf.string else 0\n",
        "  return tf.squeeze(\n",
        "      tf.sparse.to_dense(\n",
        "          tf.SparseTensor(x.indices, x.values, [x.dense_shape[0], 1]),\n",
        "          default_value),\n",
        "      axis=1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jHfhth_GiZI9"
      },
      "outputs": [],
      "source": [
        "transform = Transform(\n",
        "    examples=example_gen.outputs['examples'],\n",
        "    schema=schema_gen.outputs['schema'],\n",
        "    module_file=os.path.abspath(_census_income_transform_module_file))\n",
        "context.run(transform)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SClrAaEGR1O5"
      },
      "outputs": [],
      "source": [
        "transform.outputs['transform_graph']"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OBJFtnl6lCg9"
      },
      "source": [
        "### Trainer\n",
        "Let's see an example of user-defined model code below (for an introduction to the TensorFlow Keras APIs, [see the tutorial](https://www.tensorflow.org/guide/keras)):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N1376oq04YJt"
      },
      "outputs": [],
      "source": [
        "_census_income_trainer_module_file = 'census_income_trainer.py'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nf9UuNng4YJu",
        "jupyter": {
          "source_hidden": true
        }
      },
      "outputs": [],
      "source": [
        "%%writefile {_census_income_trainer_module_file}\n",
        "\n",
        "from typing import List, Text\n",
        "\n",
        "import os\n",
        "import absl\n",
        "import datetime\n",
        "import tensorflow as tf\n",
        "import tensorflow_transform as tft\n",
        "\n",
        "from tfx.components.trainer.executor import TrainerFnArgs\n",
        "\n",
        "import census_income_constants\n",
        "\n",
        "_DENSE_FLOAT_FEATURE_KEYS = census_income_constants.DENSE_FLOAT_FEATURE_KEYS\n",
        "_VOCAB_FEATURE_KEYS = census_income_constants.VOCAB_FEATURE_KEYS\n",
        "_VOCAB_SIZE = census_income_constants.VOCAB_SIZE\n",
        "_OOV_SIZE = census_income_constants.OOV_SIZE\n",
        "_FEATURE_BUCKET_COUNT = census_income_constants.FEATURE_BUCKET_COUNT\n",
        "_BUCKET_FEATURE_KEYS = census_income_constants.BUCKET_FEATURE_KEYS\n",
        "_CATEGORICAL_FEATURE_KEYS = census_income_constants.CATEGORICAL_FEATURE_KEYS\n",
        "_MAX_CATEGORICAL_FEATURE_VALUES = census_income_constants.MAX_CATEGORICAL_FEATURE_VALUES\n",
        "_LABEL_KEY = census_income_constants.LABEL_KEY\n",
        "_transformed_name = census_income_constants.transformed_name\n",
        "\n",
        "\n",
        "def _transformed_names(keys):\n",
        "  return [_transformed_name(key) for key in keys]\n",
        "\n",
        "\n",
        "def _gzip_reader_fn(filenames):\n",
        "  \"\"\"Small utility returning a record reader that can read gzip'ed files.\"\"\"\n",
        "  return tf.data.TFRecordDataset(\n",
        "      filenames,\n",
        "      compression_type='GZIP')\n",
        "\n",
        "\n",
        "def _get_serve_tf_examples_fn(model, tf_transform_output):\n",
        "  \"\"\"Returns a function that parses a serialized tf.Example and applies TFT.\"\"\"\n",
        "\n",
        "  model.tft_layer = tf_transform_output.transform_features_layer()\n",
        "\n",
        "  @tf.function\n",
        "  def serve_tf_examples_fn(serialized_tf_examples):\n",
        "    \"\"\"Returns the output to be used in the serving signature.\"\"\"\n",
        "    feature_spec = tf_transform_output.raw_feature_spec()\n",
        "    feature_spec.pop(_LABEL_KEY)\n",
        "    parsed_features = tf.io.parse_example(serialized_tf_examples, feature_spec)\n",
        "\n",
        "    transformed_features = model.tft_layer(parsed_features)\n",
        "    if _transformed_name(_LABEL_KEY) in transformed_features:\n",
        "      transformed_features.pop(_transformed_name(_LABEL_KEY))\n",
        "\n",
        "    return model(transformed_features)\n",
        "\n",
        "  return serve_tf_examples_fn\n",
        "\n",
        "\n",
        "def _input_fn(file_pattern: List[Text],\n",
        "              tf_transform_output: tft.TFTransformOutput,\n",
        "              batch_size: int = 200) -> tf.data.Dataset:\n",
        "  \"\"\"Generates features and label for tuning/training.\n",
        "\n",
        "  Args:\n",
        "    file_pattern: List of paths or patterns of input tfrecord files.\n",
        "    tf_transform_output: A TFTransformOutput.\n",
        "    batch_size: representing the number of consecutive elements of returned\n",
        "      dataset to combine in a single batch\n",
        "\n",
        "  Returns:\n",
        "    A dataset that contains (features, indices) tuple where features is a\n",
        "      dictionary of Tensors, and indices is a single Tensor of label indices.\n",
        "  \"\"\"\n",
        "  transformed_feature_spec = (\n",
        "      tf_transform_output.transformed_feature_spec().copy())\n",
        "\n",
        "  dataset = tf.data.experimental.make_batched_features_dataset(\n",
        "      file_pattern=file_pattern,\n",
        "      batch_size=batch_size,\n",
        "      features=transformed_feature_spec,\n",
        "      reader=_gzip_reader_fn,\n",
        "      label_key=_transformed_name(_LABEL_KEY))\n",
        "\n",
        "  return dataset\n",
        "\n",
        "\n",
        "def _build_keras_model(hidden_units: List[int] = None) -> tf.keras.Model:\n",
        "  \"\"\"Creates a DNN Keras model.\n",
        "\n",
        "  Args:\n",
        "    hidden_units: [int], the layer sizes of the DNN (input layer first).\n",
        "\n",
        "  Returns:\n",
        "    A keras Model.\n",
        "  \"\"\"\n",
        "  real_valued_columns = [\n",
        "      tf.feature_column.numeric_column(key, shape=())\n",
        "      for key in _transformed_names(_DENSE_FLOAT_FEATURE_KEYS)\n",
        "  ]\n",
        "  categorical_columns = [\n",
        "      tf.feature_column.categorical_column_with_identity(\n",
        "          key, num_buckets=_VOCAB_SIZE + _OOV_SIZE, default_value=0)\n",
        "      for key in _transformed_names(_VOCAB_FEATURE_KEYS)\n",
        "  ]\n",
        "  categorical_columns += [\n",
        "      tf.feature_column.categorical_column_with_identity(\n",
        "          key, num_buckets=_FEATURE_BUCKET_COUNT, default_value=0)\n",
        "      for key in _transformed_names(_BUCKET_FEATURE_KEYS)\n",
        "  ]\n",
        "  categorical_columns += [\n",
        "      tf.feature_column.categorical_column_with_identity(  # pylint: disable=g-complex-comprehension\n",
        "          key,\n",
        "          num_buckets=num_buckets,\n",
        "          default_value=0) for key, num_buckets in zip(\n",
        "              _transformed_names(_CATEGORICAL_FEATURE_KEYS),\n",
        "              _MAX_CATEGORICAL_FEATURE_VALUES)\n",
        "  ]\n",
        "  indicator_column = [\n",
        "      tf.feature_column.indicator_column(categorical_column)\n",
        "      for categorical_column in categorical_columns\n",
        "  ]\n",
        "\n",
        "  model = _wide_and_deep_classifier(\n",
        "      # TODO(b/139668410) replace with premade wide_and_deep keras model\n",
        "      wide_columns=indicator_column,\n",
        "      deep_columns=real_valued_columns,\n",
        "      dnn_hidden_units=hidden_units or [100, 70, 50, 25])\n",
        "  return model\n",
        "\n",
        "\n",
        "def _wide_and_deep_classifier(wide_columns, deep_columns, dnn_hidden_units):\n",
        "  \"\"\"Build a simple keras wide and deep model.\n",
        "\n",
        "  Args:\n",
        "    wide_columns: Feature columns wrapped in indicator_column for wide (linear)\n",
        "      part of the model.\n",
        "    deep_columns: Feature columns for deep part of the model.\n",
        "    dnn_hidden_units: [int], the layer sizes of the hidden DNN.\n",
        "\n",
        "  Returns:\n",
        "    A Wide and Deep Keras model\n",
        "  \"\"\"\n",
        "  # Following values are hard coded for simplicity in this example,\n",
        "  # However prefarably they should be passsed in as hparams.\n",
        "\n",
        "  # Keras needs the feature definitions at compile time.\n",
        "  # TODO(b/139081439): Automate generation of input layers from FeatureColumn.\n",
        "  input_layers = {\n",
        "      colname: tf.keras.layers.Input(name=colname, shape=(), dtype=tf.float32)\n",
        "      for colname in _transformed_names(_DENSE_FLOAT_FEATURE_KEYS)\n",
        "  }\n",
        "  input_layers.update({\n",
        "      colname: tf.keras.layers.Input(name=colname, shape=(), dtype='int32')\n",
        "      for colname in _transformed_names(_VOCAB_FEATURE_KEYS)\n",
        "  })\n",
        "  input_layers.update({\n",
        "      colname: tf.keras.layers.Input(name=colname, shape=(), dtype='int32')\n",
        "      for colname in _transformed_names(_BUCKET_FEATURE_KEYS)\n",
        "  })\n",
        "  input_layers.update({\n",
        "      colname: tf.keras.layers.Input(name=colname, shape=(), dtype='int32')\n",
        "      for colname in _transformed_names(_CATEGORICAL_FEATURE_KEYS)\n",
        "  })\n",
        "\n",
        "  # TODO(b/161816639): SparseFeatures for feature columns + Keras.\n",
        "  deep = tf.keras.layers.DenseFeatures(deep_columns)(input_layers)\n",
        "  for numnodes in dnn_hidden_units:\n",
        "    deep = tf.keras.layers.Dense(numnodes)(deep)\n",
        "  wide = tf.keras.layers.DenseFeatures(wide_columns)(input_layers)\n",
        "\n",
        "  output = tf.keras.layers.Dense(\n",
        "      1, activation='sigmoid')(\n",
        "          tf.keras.layers.concatenate([deep, wide]))\n",
        "\n",
        "  model = tf.keras.Model(input_layers, output)\n",
        "  model.compile(\n",
        "      loss='binary_crossentropy',\n",
        "      optimizer=tf.keras.optimizers.Adam(lr=0.001),\n",
        "      metrics=[tf.keras.metrics.BinaryAccuracy()])\n",
        "  model.summary(print_fn=absl.logging.info)\n",
        "  return model\n",
        "\n",
        "\n",
        "# TFX Trainer will call this function.\n",
        "def run_fn(fn_args: TrainerFnArgs):\n",
        "  \"\"\"Train the model based on given args.\n",
        "\n",
        "  Args:\n",
        "    fn_args: Holds args used to train the model as name/value pairs.\n",
        "  \"\"\"\n",
        "  # Number of nodes in the first layer of the DNN\n",
        "  first_dnn_layer_size = 100\n",
        "  num_dnn_layers = 4\n",
        "  dnn_decay_factor = 0.7\n",
        "\n",
        "  tf_transform_output = tft.TFTransformOutput(fn_args.transform_output)\n",
        "\n",
        "  train_dataset = _input_fn(fn_args.train_files, tf_transform_output, 40)\n",
        "  eval_dataset = _input_fn(fn_args.eval_files, tf_transform_output, 40)\n",
        "\n",
        "  model = _build_keras_model(\n",
        "      # Construct layers sizes with exponetial decay\n",
        "      hidden_units=[\n",
        "          max(2, int(first_dnn_layer_size * dnn_decay_factor**i))\n",
        "          for i in range(num_dnn_layers)\n",
        "      ])\n",
        "\n",
        "  # This log path might change in the future.\n",
        "  log_dir = os.path.join(os.path.dirname(fn_args.serving_model_dir), 'logs')\n",
        "  tensorboard_callback = tf.keras.callbacks.TensorBoard(\n",
        "     log_dir=log_dir, update_freq='batch')\n",
        "  model.fit(\n",
        "      train_dataset,\n",
        "      steps_per_epoch=fn_args.train_steps,\n",
        "      validation_data=eval_dataset,\n",
        "      validation_steps=fn_args.eval_steps,\n",
        "      callbacks=[tensorboard_callback])\n",
        "\n",
        "  signatures = {\n",
        "      'serving_default':\n",
        "          _get_serve_tf_examples_fn(model,\n",
        "                                    tf_transform_output).get_concrete_function(\n",
        "                                        tf.TensorSpec(\n",
        "                                            shape=[None],\n",
        "                                            dtype=tf.string,\n",
        "                                            name='examples')),\n",
        "  }\n",
        "  model.save(fn_args.serving_model_dir, save_format='tf', signatures=signatures)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "429-vvCWibO0"
      },
      "outputs": [],
      "source": [
        "trainer = Trainer(\n",
        "    module_file=os.path.abspath(_census_income_trainer_module_file),\n",
        "    custom_executor_spec=executor_spec.ExecutorClassSpec(GenericExecutor),\n",
        "    examples=transform.outputs['transformed_examples'],\n",
        "    transform_graph=transform.outputs['transform_graph'],\n",
        "    schema=schema_gen.outputs['schema'],\n",
        "    train_args=trainer_pb2.TrainArgs(num_steps=100),\n",
        "    eval_args=trainer_pb2.EvalArgs(num_steps=50))\n",
        "context.run(trainer)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "trainer.outputs"
      ],
      "metadata": {
        "id": "cSb8fhbQDmyJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FmPftrv0lEQy"
      },
      "source": [
        "### Evaluator\n",
        "The `Evaluator` component computes model performance metrics over the evaluation set. It uses the [TensorFlow Model Analysis](https://www.tensorflow.org/tfx/model_analysis/get_started) library. \n",
        "\n",
        "`Evaluator` will take as input the data from `ExampleGen`, the trained model from `Trainer`, and slicing configuration. The slicing configuration allows you to slice your metrics on feature values. See an example of this configuration below:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fVhfzzh9PDEx"
      },
      "outputs": [],
      "source": [
        "from google.protobuf.wrappers_pb2 import BoolValue\n",
        "\n",
        "eval_config = tfma.EvalConfig(\n",
        "    model_specs=[\n",
        "        # This assumes a serving model with signature 'serving_default'. If\n",
        "        # using estimator based EvalSavedModel, add signature_name: 'eval' and \n",
        "        # remove the label_key.\n",
        "        tfma.ModelSpec(label_key=\"Over-50K\")\n",
        "    ],\n",
        "    metrics_specs=[\n",
        "        tfma.MetricsSpec(\n",
        "            # The metrics added here are in addition to those saved with the\n",
        "            # model (assuming either a keras model or EvalSavedModel is used).\n",
        "            # Any metrics added into the saved model (for example using\n",
        "            # model.compile(..., metrics=[...]), etc) will be computed\n",
        "            # automatically.\n",
        "            # To add validation thresholds for metrics saved with the model,\n",
        "            # add them keyed by metric name to the thresholds map.\n",
        "            metrics=[\n",
        "                tfma.MetricConfig(class_name='ExampleCount'),\n",
        "                tfma.MetricConfig(class_name='BinaryAccuracy'),\n",
        "                tfma.MetricConfig(class_name='FairnessIndicators',\n",
        "                                  config='{ \"thresholds\": [0.5] }'),\n",
        "            ]\n",
        "        )\n",
        "    ],\n",
        "    slicing_specs=[\n",
        "        # An empty slice spec means the overall slice, i.e. the whole dataset.\n",
        "        tfma.SlicingSpec(),\n",
        "        # Data can be sliced along a feature column. In this case, data is\n",
        "        # sliced by feature column Race and Sex.\n",
        "        tfma.SlicingSpec(feature_keys=['Race']),\n",
        "        tfma.SlicingSpec(feature_keys=['Sex']),\n",
        "        tfma.SlicingSpec(feature_keys=['Race', 'Sex']),\n",
        "    ],\n",
        "    options = tfma.Options(compute_confidence_intervals=BoolValue(value=True))\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Warning: the Evaluator Component may take 5-10 minutes to run due to errors regarding \"inconsistent references\". "
      ],
      "metadata": {
        "id": "QfbptwWQ4k0z"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Zjcx8g6mihSt"
      },
      "outputs": [],
      "source": [
        "# Use TFMA to compute a evaluation statistics over features of a model and\n",
        "# validate them against a baseline.\n",
        "\n",
        "# TODO(b/226656838) Fix the inconsistent references warnings.\n",
        "evaluator = Evaluator(\n",
        "    examples=example_gen.outputs['examples'],\n",
        "    model=trainer.outputs['model'],\n",
        "    eval_config=eval_config)\n",
        "context.run(evaluator)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "k4GghePOTJxL"
      },
      "outputs": [],
      "source": [
        "evaluator.outputs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y5TMskWe9LL0"
      },
      "source": [
        "Using the `evaluation` output we can show the default visualization of global metrics on the entire evaluation set."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U729j5X5QQUQ"
      },
      "outputs": [],
      "source": [
        "context.show(evaluator.outputs['evaluation'])"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "###Model Card Generator\n",
        "The `Model Card` component is a [TFX Component](https://www.tensorflow.org/tfx/guide/understanding_tfx_pipelines#component) that generates model cards-- short documentation that provides key information about a machine learning model-- from the StatisticGen outputs, the Evaluator outputs, and a prepared json annotation. Optionally, a pushed model or a template can be provided as well. \n",
        "\n",
        "The model cards assets are saved to a ModelCard artifact that can be fetched from the `outputs['model_card]'` property."
      ],
      "metadata": {
        "id": "cdTLb_-8dLcu"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Prepare Annotation Json for Model Card\n",
        "\n",
        "It is also important to document model information that might be important to downstream users, such as its limitations, intended use cases, trade offs, and ethical considerations. Thus, we will prepare this information in json format to be used in the model card generating step."
      ],
      "metadata": {
        "id": "l8LY5bMUJsuM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import json\n",
        "import model_card_toolkit as mctlib\n",
        "\n",
        "model_card_json = {'model_details': {'name': 'Census Income Classifier'}, \n",
        "                   'model_details': {'overview': \n",
        "                      'This is a wide and deep Keras model which aims to classify whether or not '\n",
        "                      'an individual has an income of over $50,000 based on various demographic '\n",
        "                      'features. The model is trained on the UCI Census Income Dataset. This is '\n",
        "                      'not a production model, and this dataset has traditionally only been used '\n",
        "                      'for research purposes. In this Model Card, you can review quantitative '\n",
        "                      'components of the model’s performance and data, as well as information '\n",
        "                      'about the model’s intended uses, limitations, and ethical considerations.'},\n",
        "                   'model_details': {'owners': [{\"name\": \"Model Cards Team\", \"contact\": \"model-cards@google.com\"}]},\n",
        "                   'considerations': {'use_cases':[{\"description\":'This dataset that this model was trained on was originally created to '\n",
        "                      'support the machine learning community in conducting empirical analysis '\n",
        "                      'of ML algorithms. The Adult Data Set can be used in fairness-related '\n",
        "                      'studies that compare inequalities across sex and race, based on '\n",
        "                      'people’s annual incomes.'}]},\n",
        "                   'considerations': {'limitations': [{'description':\n",
        "                      'This is a class-imbalanced dataset across a variety of sensitive classes.'\n",
        "                      ' The ratio of male-to-female examples is about 2:1 and there are far more'\n",
        "                      ' examples with the “white” attribute than every other race combined. '\n",
        "                      'Furthermore, the ratio of $50,000 or less earners to $50,000 or more '\n",
        "                      'earners is just over 3:1. Due to the imbalance across income levels, we '\n",
        "                      'can see that our true negative rate seems quite high, while our true '\n",
        "                      'positive rate seems quite low. This is true to an even greater degree '\n",
        "                      'when we only look at the “female” sub-group, because there are even '\n",
        "                      'fewer female examples in the $50,000+ earner group, causing our model to '\n",
        "                      'overfit these examples. To avoid this, we can try various remediation '\n",
        "                      'strategies in future iterations (e.g. undersampling, hyperparameter '\n",
        "                      'tuning, etc), but we may not be able to fix all of the fairness issues.'}]}, \n",
        "                    'considerations': {'ethical_considerations': [\n",
        "                                        {'name': 'We risk expressing the viewpoint that the attributes in this dataset '\n",
        "                                                  'are the only ones that are predictive of someone’s income, even '\n",
        "                                                  'though we know this is not the case.', \n",
        "                                         'mitigation_strategy': 'As mentioned, some interventions may need to be '\n",
        "                                            'performed to address the class imbalances in the dataset.'}]}\n",
        "                  }"
      ],
      "metadata": {
        "id": "SIUiTor4Johj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SOYofSZKOMZx"
      },
      "source": [
        "#### Generate the Model Card.\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from model_card_toolkit.tfx.component import ModelCardGenerator\n",
        "\n",
        "mct_gen = ModelCardGenerator(statistics=statistics_gen.outputs['statistics'],\n",
        "                             evaluation=evaluator.outputs['evaluation'],\n",
        "                             json=json.dumps(model_card_json))\n",
        "context.run(mct_gen)\n"
      ],
      "metadata": {
        "id": "bspjHq6u5aFf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "mct_gen.outputs['model_card']"
      ],
      "metadata": {
        "id": "euskJ-0qM6nV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "##Display Model Card\n",
        "\n",
        "Lastly, we isolate the uri from the model card generator artifact and use it to display the model card."
      ],
      "metadata": {
        "id": "0kj1szHNdqrX"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Sd68Ih928vr9"
      },
      "outputs": [],
      "source": [
        "from IPython import display\n",
        "\n",
        "mct_artifact = mct_gen.outputs['model_card'].get()[0]\n",
        "mct_uri = mct_artifact.uri\n",
        "\n",
        "print(os.listdir(mct_uri))\n",
        "\n",
        "mct_path = os.path.join(mct_uri, 'model_cards', 'model_card.html')\n",
        "with open(mct_path) as f:\n",
        "  mct_content = f.read()\n",
        "\n",
        "\n",
        "display.display(display.HTML(mct_content))"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "MLMD Model Card Toolkit Demo.ipynb",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}