try:
  import colab
  !pip install --upgrade pip
except:
  pass

"""
Explanation: Using TFX Pipeline and TensorFlow Transform with Feature Engineering
Learning objectives

Prepare example data.
Create a pipeline.
Run the pipeline.

Introduction
In this notebook, you will create and run a TFX pipeline
to ingest raw input data and preprocess it appropriately for ML training.
This notebook is based on the TFX pipeline built in
Data validation using TFX Pipeline and TensorFlow Data Validation Tutorial.
If you have not read that one yet, you should read it before proceeding with
this notebook.
You can increase the predictive quality of your data and/or reduce
dimensionality with feature engineering. One of the benefits of using TFX is
that you will write your transformation code once, and the resulting transforms
will be consistent between training and serving in
order to avoid training/serving skew.
You will add a Transform component to the pipeline. The Transform component is
implemented using the
tf.transform library.
Please see
Understanding TFX Pipelines
to learn more about various concepts in TFX.
Each learning objective will correspond to a #TODO in the student lab notebook -- try to complete that notebook first before reviewing this solution notebook.
Set Up
You first need to install the TFX Python package and download
the dataset which you will use for our model.
Upgrade Pip
To avoid upgrading Pip in a system when running locally,
check to make sure that you are running in Colab.
Local systems can of course be upgraded separately.
End of explanation
"""


!pip install -U tfx

"""
Explanation: Install TFX
End of explanation
"""


import tensorflow as tf
print('TensorFlow version: {}'.format(tf.__version__))
from tfx import v1 as tfx
print('TFX version: {}'.format(tfx.__version__))

"""
Explanation: Restart the kernel
Please ignore any incompatibility warnings and errors. Restart the kernel to use updated packages. (On the Notebook menu, select Kernel > Restart Kernel > Restart).
Check the TensorFlow and TFX versions.
End of explanation
"""


import os

PIPELINE_NAME = "penguin-transform"

# Output directory to store artifacts generated from the pipeline.
PIPELINE_ROOT = os.path.join('pipelines', PIPELINE_NAME)
# Path to a SQLite DB file to use as an MLMD storage.
METADATA_PATH = os.path.join('metadata', PIPELINE_NAME, 'metadata.db')
# Output directory where created models from the pipeline will be exported.
SERVING_MODEL_DIR = os.path.join('serving_model', PIPELINE_NAME)

from absl import logging
logging.set_verbosity(logging.INFO)  # Set default logging level.

"""
Explanation: Set up variables
There are some variables used to define a pipeline. You can customize these
variables as you want. By default all output from the pipeline will be
generated under the current directory.
End of explanation
"""


import urllib.request
import tempfile

DATA_ROOT = tempfile.mkdtemp(prefix='tfx-data')  # TODO 1: Create a temporary directory.
_data_path = 'https://storage.googleapis.com/download.tensorflow.org/data/palmer_penguins/penguins_size.csv'
_data_filepath = os.path.join(DATA_ROOT, "data.csv")
urllib.request.urlretrieve(_data_path, _data_filepath)

"""
Explanation: Prepare example data
You will download the example dataset for use in our TFX pipeline. The dataset
you are using is
Palmer Penguins dataset.
However, unlike previous notebooks which used an already preprocessed dataset,
you will use the raw Palmer Penguins dataset.
Because the TFX ExampleGen component reads inputs from a directory, you need
to create a directory and copy the dataset to it.
End of explanation
"""


!head {_data_filepath}

"""
Explanation: Take a quick look at what the raw data looks like.
End of explanation
"""


!sed -i '/\bNA\b/d' {_data_filepath}
!head {_data_filepath}

"""
Explanation: There are some entries with missing values which are represented as NA.
You will just delete those entries in this notebook.
End of explanation
"""


import shutil

SCHEMA_PATH = 'schema'

_schema_uri = 'https://raw.githubusercontent.com/tensorflow/tfx/master/tfx/examples/penguin/schema/raw/schema.pbtxt'
_schema_filename = 'schema.pbtxt'
_schema_filepath = os.path.join(SCHEMA_PATH, _schema_filename)

os.makedirs(SCHEMA_PATH, exist_ok=True)
urllib.request.urlretrieve(_schema_uri, _schema_filepath)

"""
Explanation: You should be able to see seven features which describe penguins. You will use
the same set of features as the previous notebooks - 'culmen_length_mm',
'culmen_depth_mm', 'flipper_length_mm', 'body_mass_g' - and will predict
the 'species' of a penguin.
The only difference will be that the input data is not preprocessed. Note
that you will not use other features like 'island' or 'sex' in this notebook.
Prepare a schema file
As described in
Data validation using TFX Pipeline and TensorFlow Data Validation Tutorial,
you need a schema file for the dataset. Because the dataset is different, you need to generate it again. In this notebook, you will skip those steps and just use a prepared schema file.
End of explanation
"""


_module_file = 'penguin_utils.py'

%%writefile {_module_file}


from typing import List, Text
from absl import logging
import tensorflow as tf
from tensorflow import keras
from tensorflow_metadata.proto.v0 import schema_pb2
import tensorflow_transform as tft
from tensorflow_transform.tf_metadata import schema_utils

from tfx import v1 as tfx
from tfx_bsl.public import tfxio

# Specify features that you will use.
_FEATURE_KEYS = [
    'culmen_length_mm', 'culmen_depth_mm', 'flipper_length_mm', 'body_mass_g'
]
_LABEL_KEY = 'species'

_TRAIN_BATCH_SIZE = 20
_EVAL_BATCH_SIZE = 10


# NEW: TFX Transform will call this function.
def preprocessing_fn(inputs):
  """tf.transform's callback function for preprocessing inputs.

  Args:
    inputs: map from feature keys to raw not-yet-transformed features.

  Returns:
    Map from string feature key to transformed feature.
  """
  outputs = {}

  # Uses features defined in _FEATURE_KEYS only.
  for key in _FEATURE_KEYS:
    # tft.scale_to_z_score computes the mean and variance of the given feature
    # and scales the output based on the result.
    outputs[key] = tft.scale_to_z_score(inputs[key])

  # For the label column you provide the mapping from string to index.
  # You could instead use `tft.compute_and_apply_vocabulary()` in order to
  # compute the vocabulary dynamically and perform a lookup.
  # Since in this example there are only 3 possible values, you use a hard-coded
  # table for simplicity.
  table_keys = ['Adelie', 'Chinstrap', 'Gentoo']
  initializer = tf.lookup.KeyValueTensorInitializer(
      keys=table_keys,
      values=tf.cast(tf.range(len(table_keys)), tf.int64),
      key_dtype=tf.string,
      value_dtype=tf.int64)
  table = tf.lookup.StaticHashTable(initializer, default_value=-1)
  outputs[_LABEL_KEY] = table.lookup(inputs[_LABEL_KEY])

  return outputs


# NEW: This function will apply the same transform operation to training data
#      and serving requests.
def _apply_preprocessing(raw_features, tft_layer):
  transformed_features = tft_layer(raw_features)
  if _LABEL_KEY in raw_features:
    transformed_label = transformed_features.pop(_LABEL_KEY)
    return transformed_features, transformed_label
  else:
    return transformed_features, None


# NEW: This function will create a handler function which gets a serialized
#      tf.example, preprocess and run an inference with it.
def _get_serve_tf_examples_fn(model, tf_transform_output):
  # You must save the tft_layer to the model to ensure its assets are kept and
  # tracked.
  model.tft_layer = tf_transform_output.transform_features_layer()

  @tf.function(input_signature=[
      tf.TensorSpec(shape=[None], dtype=tf.string, name='examples')
  ])
  def serve_tf_examples_fn(serialized_tf_examples):
    # Expected input is a string which is serialized tf.Example format.
    feature_spec = tf_transform_output.raw_feature_spec()
    # Because input schema includes unnecessary fields like 'species' and
    # 'island', you filter feature_spec to include required keys only.
    required_feature_spec = {
        k: v for k, v in feature_spec.items() if k in _FEATURE_KEYS
    }
    parsed_features = tf.io.parse_example(serialized_tf_examples,
                                          required_feature_spec)

    # Preprocess parsed input with transform operation defined in
    # preprocessing_fn().
    transformed_features, _ = _apply_preprocessing(parsed_features,
                                                   model.tft_layer)
    # Run inference with ML model.
    return model(transformed_features)

  return serve_tf_examples_fn


def _input_fn(file_pattern: List[Text],
              data_accessor: tfx.components.DataAccessor,
              tf_transform_output: tft.TFTransformOutput,
              batch_size: int = 200) -> tf.data.Dataset:
  """Generates features and label for tuning/training.

  Args:
    file_pattern: List of paths or patterns of input tfrecord files.
    data_accessor: DataAccessor for converting input to RecordBatch.
    tf_transform_output: A TFTransformOutput.
    batch_size: representing the number of consecutive elements of returned
      dataset to combine in a single batch

  Returns:
    A dataset that contains (features, indices) tuple where features is a
      dictionary of Tensors, and indices is a single Tensor of label indices.
  """
  dataset = data_accessor.tf_dataset_factory(
      file_pattern,
      tfxio.TensorFlowDatasetOptions(batch_size=batch_size),
      schema=tf_transform_output.raw_metadata.schema)

  transform_layer = tf_transform_output.transform_features_layer()
  def apply_transform(raw_features):
    return _apply_preprocessing(raw_features, transform_layer)

  return dataset.map(apply_transform).repeat()


def _build_keras_model() -> tf.keras.Model:
  """Creates a DNN Keras model for classifying penguin data.

  Returns:
    A Keras Model.
  """
  # The model below is built with Functional API, please refer to
  # https://www.tensorflow.org/guide/keras/overview for all API options.
  inputs = [
      keras.layers.Input(shape=(1,), name=key)
      for key in _FEATURE_KEYS
  ]
  d = keras.layers.concatenate(inputs)
  for _ in range(2):
    d = keras.layers.Dense(8, activation='relu')(d)
  outputs = keras.layers.Dense(3)(d)

  model = keras.Model(inputs=inputs, outputs=outputs)
  model.compile(
      optimizer=keras.optimizers.Adam(1e-2),
      loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
      metrics=[keras.metrics.SparseCategoricalAccuracy()])

  model.summary(print_fn=logging.info)
  return model


# TFX Trainer will call this function.
def run_fn(fn_args: tfx.components.FnArgs):
  """Train the model based on given args.

  Args:
    fn_args: Holds args used to train the model as name/value pairs.
  """
  tf_transform_output = tft.TFTransformOutput(fn_args.transform_output)

  train_dataset = _input_fn(
      fn_args.train_files,
      fn_args.data_accessor,
      tf_transform_output,
      batch_size=_TRAIN_BATCH_SIZE)
  eval_dataset = _input_fn(
      fn_args.eval_files,
      fn_args.data_accessor,
      tf_transform_output,
      batch_size=_EVAL_BATCH_SIZE)

  model = _build_keras_model()
  model.fit(
      train_dataset,
      steps_per_epoch=fn_args.train_steps,
      validation_data=eval_dataset,
      validation_steps=fn_args.eval_steps)

  # NEW: Save a computation graph including transform layer.
  signatures = {
      'serving_default': _get_serve_tf_examples_fn(model, tf_transform_output),
  }
  model.save(fn_args.serving_model_dir, save_format='tf', signatures=signatures)

"""
Explanation: This schema file was created with the same pipeline as in the previous notebook
without any manual changes.
Create a pipeline
TFX pipelines are defined using Python APIs. You will add Transform
component to the pipeline you created in the
Data Validation tutorial.
A Transform component requires input data from an ExampleGen component and
a schema from a SchemaGen component, and produces a "transform graph". The
output will be used in a Trainer component. Transform can optionally
produce "transformed data" in addition, which is the materialized data after
transformation.
However, you will transform data during training in this notebook without
materialization of the intermediate transformed data.
One thing to note is that you need to define a Python function,
preprocessing_fn to describe how input data should be transformed. This is
similar to a Trainer component which also requires user code for model
definition.
Write preprocessing and training code
You need to define two Python functions. One for Transform and one for Trainer.
preprocessing_fn
The Transform component will find a function named preprocessing_fn in the
given module file as you did for Trainer component. You can also specify a
specific function using the
preprocessing_fn parameter
of the Transform component.
In this example, you will do two kinds of transformation. For continuous numeric
features like culmen_length_mm and body_mass_g, you will normalize these
values using the
tft.scale_to_z_score
function. For the label feature, you need to convert string labels into numeric
index values. You will use
tf.lookup.StaticHashTable
for conversion.
To identify transformed fields easily, you append a _xf suffix to the
transformed feature names.
run_fn
The model itself is almost the same as in the previous notebooks, but this time
you will transform the input data using the transform graph from the Transform
component.
One more important difference compared to the previous notebook is that now you
export a model for serving which includes not only the computation graph of the
model, but also the transform graph for preprocessing, which is generated in
Transform component. You need to define a separate function which will be used
for serving incoming requests. You can see that the same function
_apply_preprocessing was used for both of the training data and the
serving request.
End of explanation
"""


def _create_pipeline(pipeline_name: str, pipeline_root: str, data_root: str,
                     schema_path: str, module_file: str, serving_model_dir: str,
                     metadata_path: str) -> tfx.dsl.Pipeline:
  """Implements the penguin pipeline with TFX."""
  # Brings data into the pipeline or otherwise joins/converts training data.
  example_gen = tfx.components.CsvExampleGen(input_base=data_root)

  # TODO 2: Computes statistics over data for visualization and example validation.
  statistics_gen = tfx.components.StatisticsGen(
      examples=example_gen.outputs['examples'])

  # Import the schema.
  schema_importer = tfx.dsl.Importer(
      source_uri=schema_path,
      artifact_type=tfx.types.standard_artifacts.Schema).with_id(
          'schema_importer')

  # Performs anomaly detection based on statistics and data schema.
  example_validator = tfx.components.ExampleValidator(
      statistics=statistics_gen.outputs['statistics'],
      schema=schema_importer.outputs['result'])

  # NEW: Transforms input data using preprocessing_fn in the 'module_file'.
  transform = tfx.components.Transform(
      examples=example_gen.outputs['examples'],
      schema=schema_importer.outputs['result'],
      materialize=False,
      module_file=module_file)

  # Uses user-provided Python function that trains a model.
  trainer = tfx.components.Trainer(
      module_file=module_file,
      examples=example_gen.outputs['examples'],

      # NEW: Pass transform_graph to the trainer.
      transform_graph=transform.outputs['transform_graph'],

      train_args=tfx.proto.TrainArgs(num_steps=100),
      eval_args=tfx.proto.EvalArgs(num_steps=5))

  # Pushes the model to a filesystem destination.
  pusher = tfx.components.Pusher(
      model=trainer.outputs['model'],
      push_destination=tfx.proto.PushDestination(
          filesystem=tfx.proto.PushDestination.Filesystem(
              base_directory=serving_model_dir)))

  components = [
      example_gen,
      statistics_gen,
      schema_importer,
      example_validator,

      transform,  # NEW: Transform component was added to the pipeline.

      trainer,
      pusher,
  ]

  return tfx.dsl.Pipeline(
      pipeline_name=pipeline_name,
      pipeline_root=pipeline_root,
      metadata_connection_config=tfx.orchestration.metadata
      .sqlite_metadata_connection_config(metadata_path),
      components=components)

"""
Explanation: Now you have completed all of the preparation steps to build a TFX pipeline.
Write a pipeline definition
You define a function to create a TFX pipeline. A Pipeline object
represents a TFX pipeline, which can be run using one of the pipeline
orchestration systems that TFX supports.
End of explanation
"""


tfx.orchestration.LocalDagRunner().run(
  _create_pipeline(
      pipeline_name=PIPELINE_NAME,
      pipeline_root=PIPELINE_ROOT,
      data_root=DATA_ROOT,
      schema_path=SCHEMA_PATH,
      module_file=_module_file,
      serving_model_dir=SERVING_MODEL_DIR,
      metadata_path=METADATA_PATH))

"""
Explanation: Run the pipeline
You will use LocalDagRunner as in the previous notebook.
End of explanation
"""


# List files in created model directory.
!find {SERVING_MODEL_DIR}

"""
Explanation: You should see "INFO:absl:Component Pusher is finished." if the pipeline
finished successfully.
The pusher component pushes the trained model to the SERVING_MODEL_DIR which
is the serving_model/penguin-transform directory if you did not change
the variables in the previous steps. You can see the result from the file
browser in the left-side panel in Colab, or using the following command:
End of explanation
"""


!saved_model_cli show --dir {SERVING_MODEL_DIR}/$(ls -1 {SERVING_MODEL_DIR} | sort -nr | head -1) --tag_set serve --signature_def serving_default

"""
Explanation: You can also check the signature of the generated model using the
saved_model_cli tool.
End of explanation
"""


# Find a model with the latest timestamp.
model_dirs = (item for item in os.scandir(SERVING_MODEL_DIR) if item.is_dir())
model_path = max(model_dirs, key=lambda i: int(i.name)).path

# TODO 3: Load a model saved via model.save()
loaded_model = tf.keras.models.load_model(model_path)
inference_fn = loaded_model.signatures['serving_default']

# Prepare an example and run inference.
features = {
  'culmen_length_mm': tf.train.Feature(float_list=tf.train.FloatList(value=[49.9])),
  'culmen_depth_mm': tf.train.Feature(float_list=tf.train.FloatList(value=[16.1])),
  'flipper_length_mm': tf.train.Feature(int64_list=tf.train.Int64List(value=[213])),
  'body_mass_g': tf.train.Feature(int64_list=tf.train.Int64List(value=[5400])),
}
example_proto = tf.train.Example(features=tf.train.Features(feature=features))
examples = example_proto.SerializeToString()

result = inference_fn(examples=tf.constant([examples]))
print(result['output_0'].numpy())

"""
Explanation: Because you defined serving_default with our own serve_tf_examples_fn
function, the signature shows that it takes a single string.
This string is a serialized string of tf.Examples and will be parsed with the
tf.io.parse_example()
function as you defined earlier (learn more about tf.Examples here).
You can load the exported model and try some inferences with a few examples.
End of explanation
"""


