code: def run_fn(trainer_fn_args):
      """TFX trainer entry point."""
  
      tf_transform_output = tft.TFTransformOutput(
          trainer_fn_args.transform_output
      )
      hparams = dict(
          batch_size=32,
          embedding_dimension=20,
          learning_rate=0.05,
          dropout_rate=0.8,
          hidden_layer_dims=[64, 32, 16],
          loss="approx_ndcg_loss",
          use_batch_norm=True,
          batch_norm_moment=0.99,
      )
  
      train_dataset = _input_fn(
          trainer_fn_args.train_files,
          trainer_fn_args.data_accessor,
          hparams["batch_size"],
      )
      eval_dataset = _input_fn(
          trainer_fn_args.eval_files,
          trainer_fn_args.data_accessor,
          hparams["batch_size"],
      )
  
      model = _create_ranking_model(tf_transform_output, hparams)
      model.summary()
      log_dir = trainer_fn_args.model_run_dir
      # Write logs to path
      tensorboard_callback = tf.keras.callbacks.TensorBoard(
          log_dir=log_dir, update_freq="epoch"
      )
      model.fit(
          train_dataset,
          steps_per_epoch=trainer_fn_args.train_steps,
          validation_data=eval_dataset,
          validation_steps=trainer_fn_args.eval_steps,
          callbacks=[tensorboard_callback],
      )
  
      # TODO(zhuo): Add support for Regress signature.
      @tf.function(
          input_signature=[tf.TensorSpec([None], tf.string)],
          autograph=False,
      )
      def predict_serving_fn(serialized_elwc_records):
          decode_fn = trainer_fn_args.data_accessor.data_view_decode_fn
          decoded = decode_fn(serialized_elwc_records)
          decoded.pop(features.LABEL)
          return {tf.saved_model.PREDICT_OUTPUTS: model(decoded)}
  
      model.save(
          trainer_fn_args.serving_model_dir,
          save_format="tf",
          signatures={
              "serving_default": predict_serving_fn.get_concrete_function(),
          },
      )
  
file_path: tfx/tfx/examples/ranking/ranking_utils.py
imports: import tensorflow
  import tensorflow_ranking
  import tensorflow_transform
  from tfx.examples.ranking import features
  from tfx.examples.ranking import struct2tensor_parsing_utils
  from tfx_bsl.public import tfxio
label_desc: Use image node_id 3 for calling a global function with example usage: run_fn(trainer_fn_args) and returns: dict