text: # /**
  #  * @author Sarath Chandra Asapu
  #  * @email sarath.chandra.asapu@accenture.com
  #  * @create date 2021-09-17 13:07:29
  #  * @modify date 2021-09-17 13:07:29
  #  * @desc [description]
  #  */
  
  from typing import List
  from absl import logging
  import tensorflow as tf
  from tensorflow import keras
  import tensorflow_data_validation as tfdv
  from tensorflow_transform.tf_metadata import schema_utils
  
  from tfx import v1 as tfx
  from tfx_bsl.public import tfxio
  from tensorflow_metadata.proto.v0 import schema_pb2
  
  _FEATURE_KEYS = [
      'age', 'job', 'marital', 'education', 'default', 'balance', 'housing',
         'loan', 'contact', 'day', 'month', 'duration', 'campaign', 'pdays',
         'previous', 'poutcome'
  ]
  _LABEL_KEY = 'target'
  
  _TRAIN_BATCH_SIZE = 20
  _EVAL_BATCH_SIZE = 10
  
  _FEATURE_SPEC = {
      **{
          feature: tf.io.FixedLenFeature(shape=[1], dtype=tf.int64)
             for feature in _FEATURE_KEYS
         },
      _LABEL_KEY: tf.io.FixedLenFeature(shape=[1], dtype=tf.int64)
  }
  
  
  def _input_fn(file_pattern: List[str],
                data_accessor: tfx.components.DataAccessor,
                schema: schema_pb2.Schema,
                batch_size: int = 200) -> tf.data.Dataset:
    return data_accessor.tf_dataset_factory(
        file_pattern,
        tfxio.TensorFlowDatasetOptions(
            batch_size=batch_size, label_key=_LABEL_KEY),
        schema=schema).repeat()
  
  
  def _build_keras_model() -> tf.keras.Model:
    inputs = [keras.layers.Input(shape=(1,), name=f) for f in _FEATURE_KEYS]
    d = keras.layers.concatenate(inputs)
    for _ in range(2):
      d = keras.layers.Dense(8, activation='relu')(d)
    outputs = keras.layers.Dense(2)(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):
    
    schema = schema_utils.schema_from_feature_spec(_FEATURE_SPEC)
  
    train_dataset = _input_fn(
        fn_args.train_files,
        fn_args.data_accessor,
        schema,
        batch_size=_TRAIN_BATCH_SIZE
    )
    eval_dataset = _input_fn(
        fn_args.eval_files,
        fn_args.data_accessor,
        schema,
        batch_size=_EVAL_BATCH_SIZE
    )
    print("train_dataset\n",train_dataset)
    print("type(train_dataset)\n",type(train_dataset))
    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)
    print("fn_args.serving_model_dir\n",fn_args.serving_model_dir)
    model.save(fn_args.serving_model_dir, save_format='tf')