import tensorflow as tf
from tfx import v1 as tfxio

def _input_fn(file_pattern, data_accessor, schema, batch_size=200):
    raw_data = data_accessor.tf_dataset_factory(
        file_pattern, tfxio.TensorFlowDatasetOptions(batch_size=batch_size), schema)
    transformed_data = raw_data.map(_parse_fn)
    return transformed_data

def _build_keras_model():
    model = tf.keras.Sequential([
        tf.keras.layers.InputLayer(input_shape=(4,)),
        tf.keras.layers.Dense(10, activation='relu'),
        tf.keras.layers.Dense(10, activation='relu'),
        tf.keras.layers.Dense(3, activation='softmax')
    ])
    model.compile(optimizer=tf.keras.optimizers.Adam(lr=0.001),
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])
    return model

def run_fn(fn_args):
    schema = tfx_bsl.public.tfxio_utils.get_tfx_schema_from_tensorflow_metadata_schema(
        fn_args.schema)
    train_dataset = _input_fn(fn_args.train_files, fn_args.data_accessor, schema, batch_size=200)
    eval_dataset = _input_fn(fn_args.eval_files, fn_args.data_accessor, schema, batch_size=200)

    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)
    model.save(fn_args.serving_model_dir, save_format='tf')

        __
        __