Labs/MLMD_Labs/Lab2/C2_W3_Lab_2_IterativeSchema.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "23R0Z9RojXYW"
   },
   "source": [
    "# Ungraded Lab: Iterative Schema with TFX and ML Metadata\n",
    "\n",
    "\n",
    "In this notebook, you will get to review how to update an inferred schema and save the result to the metadata store used by TFX. As mentioned before, the TFX components get information from this database before running executions. Thus, if you will be curating a schema, you will need to save this as an artifact in the metadata store. You will get to see how that is done in the following exercise.\n",
    "\n",
    "Afterwards, you will also practice accessing the TFX metadata store and see how you can track the lineage of an artifact."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2GivNBNYjb3b"
   },
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "N-ePgV0Lj68Q"
   },
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YIqpWK9efviJ"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow_data_validation as tfdv\n",
    "\n",
    "from tfx import v1 as tfx\n",
    "\n",
    "from tfx.types import standard_artifacts\n",
    "\n",
    "from tfx.orchestration.experimental.interactive.interactive_context import InteractiveContext\n",
    "from google.protobuf.json_format import MessageToDict\n",
    "from tensorflow_metadata.proto.v0 import schema_pb2\n",
    "\n",
    "import os\n",
    "import pprint\n",
    "pp = pprint.PrettyPrinter()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define paths\n",
    "\n",
    "For familiarity, you will again be using the [Census Income dataset](https://archive.ics.uci.edu/ml/datasets/Adult) from the previous weeks' ungraded labs. You will use the same paths to your raw data and pipeline files as shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# location of the pipeline metadata store\n",
    "_pipeline_root = './pipeline/'\n",
    "\n",
    "# directory of the raw data files\n",
    "_data_root = './data/census_data'\n",
    "\n",
    "# path to the raw training data\n",
    "_data_filepath = os.path.join(_data_root, 'adult.data')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Pipeline\n",
    "\n",
    "Each TFX component you use accepts and generates artifacts which are instances of the different artifact types TFX has configured in the metadata store. The properties of these instances are shown neatly in a table in the outputs of `context.run()`. TFX does all of these for you so you only need to inspect the output of each component to know which property of the artifact you can pass on to the next component (e.g. the `outputs['examples']` of `ExampleGen` can be passed to `StatisticsGen`).\n",
    "\n",
    "Since you've already used this dataset before, we will just quickly go over `ExampleGen`, `StatisticsGen`, and `SchemaGen`. The new concepts will be discussed after the said components."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8ONIE_hdkPS4"
   },
   "source": [
    "### Create the Interactive Context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0Rh6K5sUf9dd"
   },
   "outputs": [],
   "source": [
    "# Initialize the InteractiveContext.\n",
    "# If you leave `_pipeline_root` blank, then the db will be created in a temporary directory.\n",
    "context = InteractiveContext(pipeline_root=_pipeline_root)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "L9fwt9gQk3BR"
   },
   "source": [
    "### ExampleGen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PyXjuMt8f-9u",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Instantiate ExampleGen with the input CSV dataset\n",
    "example_gen = tfx.components.CsvExampleGen(input_base=_data_root)\n",
    "\n",
    "# Execute the component\n",
    "context.run(example_gen)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "csM6BFhtk5Aa"
   },
   "source": [
    "### StatisticsGen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "MAscCCYWgA-9"
   },
   "outputs": [],
   "source": [
    "# Instantiate StatisticsGen with the ExampleGen ingested dataset\n",
    "statistics_gen = tfx.components.StatisticsGen(\n",
    "    examples=example_gen.outputs['examples'])\n",
    "\n",
    "# Execute the component\n",
    "context.run(statistics_gen)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HLKLTO9Nk60p"
   },
   "source": [
    "### SchemaGen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ygQvZ6hsiQ_J"
   },
   "outputs": [],
   "source": [
    "# Instantiate SchemaGen with the StatisticsGen ingested dataset\n",
    "schema_gen = tfx.components.SchemaGen(\n",
    "    statistics=statistics_gen.outputs['statistics'],\n",
    "    )\n",
    "\n",
    "# Run the component\n",
    "context.run(schema_gen)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ec9vqDXpXeMb"
   },
   "outputs": [],
   "source": [
    "# Visualize the schema\n",
    "context.show(schema_gen.outputs['schema'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Curating the Schema"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that you have the inferred schema, you can proceed to revising it to be more robust. For instance, you can restrict the age as you did in Week 1. First, you have to load the `Schema` protocol buffer from the metadata store. You can do this by getting the schema uri from the output of `SchemaGen` then use TFDV's `load_schema_text()` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the schema uri\n",
    "schema_uri = schema_gen.outputs['schema']._artifacts[0].uri\n",
    "\n",
    "# Get the schema pbtxt file from the SchemaGen output\n",
    "schema = tfdv.load_schema_text(os.path.join(schema_uri, 'schema.pbtxt'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With that, you can now make changes to the schema as before. For the purpose of this exercise, you will only modify the age domain but feel free to add more if you want."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Restrict the range of the `age` feature\n",
    "tfdv.set_domain(schema, 'age', schema_pb2.IntDomain(name='age', min=17, max=90))\n",
    "\n",
    "# Display the modified schema. Notice the `Domain` column of `age`.\n",
    "tfdv.display_schema(schema)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Schema Environments\n",
    "\n",
    "By default, your schema will watch for all the features declared above including the label. However, when the model is served for inference, it will get datasets that will not have the label because that is the feature that the model will be trying to predict. You need to configure the pipeline to not raise an alarm when this kind of dataset is received.\n",
    "\n",
    "You can do that with [schema environments](https://www.tensorflow.org/tfx/tutorials/data_validation/tfdv_basic#schema_environments). First, you will need to declare training and serving environments, then configure the serving schema to not watch for the presence of labels. See how it is implemented below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create schema environments for training and serving\n",
    "schema.default_environment.append('TRAINING')\n",
    "schema.default_environment.append('SERVING')\n",
    "\n",
    "# Omit label from the serving environment\n",
    "tfdv.get_feature(schema, 'label').not_in_environment.append('SERVING')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can now freeze the curated schema and save to a local directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Declare the path to the updated schema directory\n",
    "_updated_schema_dir = f'{_pipeline_root}/updated_schema'\n",
    "\n",
    "# Create the said directory\n",
    "!mkdir -p {_updated_schema_dir}\n",
    "\n",
    "# Declare the path to the schema file\n",
    "schema_file = os.path.join(_updated_schema_dir, 'schema.pbtxt')\n",
    "\n",
    "# Save the curated schema to the said file\n",
    "tfdv.write_schema_text(schema, schema_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ImportSchemaGen\n",
    "\n",
    "Now that the schema has been saved, you need to create an artifact in the metadata store that will point to it. TFX provides the [ImporterSchemaGen](https://www.tensorflow.org/tfx/api_docs/python/tfx/v1/components/ImportSchemaGen) component used to import a curated schema to ML Metadata. You simply need to specify the URL of the revised schema file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use ImportSchemaGen to put the curated schema to ML Metadata\n",
    "user_schema_importer = tfx.components.ImportSchemaGen(schema_file=schema_file)\n",
    "\n",
    "# Run the component\n",
    "context.run(user_schema_importer, enable_cache=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you pass in the component output to `context.show()`, then you should see the schema."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# See the result\n",
    "context.show(user_schema_importer.outputs['schema'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "V1qcUuO9k9f8"
   },
   "source": [
    "### ExampleValidator\n",
    "\n",
    "You can then use this new artifact as input to the other components of the pipeline. See how it is used as the `schema` argument in `ExampleValidator` below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XRlRUuGgiXks"
   },
   "outputs": [],
   "source": [
    "# Instantiate ExampleValidator with the StatisticsGen and SchemaGen ingested data\n",
    "example_validator = tfx.components.ExampleValidator(\n",
    "    statistics=statistics_gen.outputs['statistics'],\n",
    "    schema=user_schema_importer.outputs['schema'])\n",
    "\n",
    "# Run the component.\n",
    "context.run(example_validator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TDyAAozQcrk3"
   },
   "outputs": [],
   "source": [
    "# Visualize the results\n",
    "context.show(example_validator.outputs['anomalies'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Practice with ML Metadata\n",
    "\n",
    "At this point, you should now take some time exploring the contents of the metadata store saved by your component runs. This will let you practice tracking artifacts and how they are related to each other. This involves looking at artifacts, executions, and events. This skill will let you recover related artifacts even without seeing the code of the training run. All you need is access to the metadata store. \n",
    "\n",
    "See how the input artifact IDs to an instance of `ExampleAnomalies` are tracked in the following cells. If you have this notebook, then you will already know that it uses the output of StatisticsGen for this run and also the curated schema you imported. However, if you already have hundreds of training runs and parameter iterations, you may find it hard to track which is which. That's where the metadata store will be useful. Since it records information about a specific pipeline run, you will be able to track the inputs and outputs of a particular artifact.\n",
    "\n",
    "You will start by setting the connection config to the metadata store. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import mlmd and utilities\n",
    "import ml_metadata as mlmd\n",
    "from ml_metadata.proto import metadata_store_pb2\n",
    "\n",
    "# Get the connection config to connect to the context's metadata store\n",
    "connection_config = context.metadata_connection_config\n",
    "\n",
    "# Instantiate a MetadataStore instance with the connection config\n",
    "store = mlmd.MetadataStore(connection_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let's see what artifact types are available in the metadata store."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get artifact types\n",
    "artifact_types = store.get_artifact_types()\n",
    "\n",
    "# Print the results\n",
    "[artifact_type.name for artifact_type in artifact_types]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you get the artifacts of type `Schema`, you will see that there are two entries. One is the inferred and the other is the one you imported. At the end of this exercise, you can verify that the curated schema is the one used for the `ExampleValidator` run we will be investigating."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get artifact types\n",
    "schema_list = store.get_artifacts_by_type('Schema')\n",
    "\n",
    "[(f'schema uri: {schema.uri}', f'schema id:{schema.id}') for schema in schema_list]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's get the first instance of `ExampleAnomalies` to get the output of `ExampleValidator`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get 1st instance of ExampleAnomalies\n",
    "example_anomalies = store.get_artifacts_by_type('ExampleAnomalies')[0]\n",
    "\n",
    "# Print the artifact id\n",
    "print(f'Artifact id: {example_anomalies.id}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You will use the artifact ID to get events related to it. Let's just get the first instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get first event related to the ID\n",
    "anomalies_id_event = store.get_events_by_artifact_ids([example_anomalies.id])[0]\n",
    "\n",
    "# Print results\n",
    "print(anomalies_id_event)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, the event type will be an `OUTPUT` because this is the output of the `ExampleValidator` component. Since we want to get the inputs, we can track it through the execution id."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get execution ID\n",
    "anomalies_execution_id = anomalies_id_event.execution_id\n",
    "\n",
    "# Get events by the execution ID\n",
    "events_execution = store.get_events_by_execution_ids([anomalies_execution_id])\n",
    "\n",
    "# Print results\n",
    "print(events_execution)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see the artifacts which are marked as `INPUT` above representing the statistics and schema inputs. We can extract their IDs programmatically like this. You will see that you will get the artifact ID of the curated schema you printed out earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Filter INPUT type events\n",
    "inputs_to_exval = [event.artifact_id for event in events_execution \n",
    "                       if event.type == metadata_store_pb2.Event.INPUT]\n",
    "\n",
    "# Print results\n",
    "print(inputs_to_exval)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Congratulations!** You have now completed this notebook on iterative schemas and saw how it can be used in a TFX pipeline. You were also able to track an artifact's lineage by looking at the artifacts, events, and executions in the metadata store. These will come in handy in this week's assignment!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}