code: 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.
      """
      if fn_args.transform_output is None:  # Transform is not used.
          tf_transform_output = None
          schema = tfx.utils.parse_pbtxt_file(
              fn_args.schema_file, schema_pb2.Schema()
          )
          feature_list = features.FEATURE_KEYS
          label_key = features.LABEL_KEY
      else:
          tf_transform_output = tft.TFTransformOutput(
              fn_args.transform_output
          )
          schema = tf_transform_output.transformed_metadata.schema
          feature_list = [
              features.transformed_name(f)
              for f in features.FEATURE_KEYS
          ]
          label_key = features.transformed_name(features.LABEL_KEY)
  
      mirrored_strategy = tf.distribute.MirroredStrategy()
      train_batch_size = (
          constants.TRAIN_BATCH_SIZE
          * mirrored_strategy.num_replicas_in_sync
      )
      eval_batch_size = (
          constants.EVAL_BATCH_SIZE
          * mirrored_strategy.num_replicas_in_sync
      )
  
      train_dataset = _input_fn(
          fn_args.train_files,
          fn_args.data_accessor,
          schema,
          label_key,
          batch_size=train_batch_size,
      )
      eval_dataset = _input_fn(
          fn_args.eval_files,
          fn_args.data_accessor,
          schema,
          label_key,
          batch_size=eval_batch_size,
      )
  
      with mirrored_strategy.scope():
          model = _build_keras_model(feature_list)
  
      # Write logs to path
      tensorboard_callback = tf.keras.callbacks.TensorBoard(
          log_dir=fn_args.model_run_dir, update_freq="epoch"
      )
  
      model.fit(
          train_dataset,
          steps_per_epoch=fn_args.train_steps,
          validation_data=eval_dataset,
          validation_steps=fn_args.eval_steps,
          callbacks=[tensorboard_callback],
      )
  
      signatures = {
          "serving_default": _get_tf_examples_serving_signature(
              model, schema, tf_transform_output
          ),
          "transform_features": _get_transform_features_signature(
              model, schema, tf_transform_output
          ),
      }
      model.save(
          fn_args.serving_model_dir,
          save_format="tf",
          signatures=signatures,
      )
  
file_path: tfx/tfx/experimental/templates/penguin/models/model.py
imports: from typing import List
  from absl import logging
  import tensorflow
  from tensorflow import keras
  import tensorflow_transform
  from tensorflow_transform.tf_metadata import schema_utils
  from tfx import v1
  from tfx.experimental.templates.penguin.models import constants
  from tfx.experimental.templates.penguin.models import features
  from tfx_bsl.public import tfxio
  from tensorflow_metadata.proto.v0 import schema_pb2
label_desc: Use image node_id 5 for calling a global function with example usage: run_fn(fn_args) without return types