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.
      """
      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,
      )
  
      mirrored_strategy = tf.distribute.MirroredStrategy()
      with mirrored_strategy.scope():
          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,
          verbose=2,
      )
  
      signatures = {
          "serving_default": _get_inference_fn(
              model, tf_transform_output
          ).get_concrete_function(
              tf.TensorSpec(
                  shape=[None],
                  dtype=tf.int64,
                  name=_CUR_PAGE_FEATURE_KEY,
              ),
              tf.TensorSpec(
                  shape=[None],
                  dtype=tf.int64,
                  name=_SESSION_INDEX_FEATURE_KEY,
              ),
          ),
      }
  
      # Create the saved_model in a temporary directory.
      temp_saving_model_dir = os.path.join(
          fn_args.serving_model_dir, "temp"
      )
      model.save(
          temp_saving_model_dir, save_format="tf", signatures=signatures
      )
  
      # Convert the saved_model to a tfjs model and store it in the final directory.
      tfrw = rewriter_factory.create_rewriter(
          rewriter_factory.TFJS_REWRITER, name="tfjs_rewriter"
      )
      converters.rewrite_saved_model(
          temp_saving_model_dir,
          fn_args.serving_model_dir,
          tfrw,
          rewriter.ModelType.TFJS_MODEL,
      )
  
      # Copy the vocabulary computed by transform to the final directory.
      # The vocabulary is not included in the original savedmodel because vocab
      # lookups are currently not supported in TFJS and are expected to be done
      # independently by client code.
      fileio.copy(
          tf_transform_output.vocabulary_file_by_name(_VOCAB_FILENAME),
          os.path.join(fn_args.serving_model_dir, _VOCAB_FILENAME),
      )
  
      fileio.rmtree(temp_saving_model_dir)
  
file_path: tfx/tfx/examples/tfjs_next_page_prediction/tfjs_next_page_prediction_util.py
imports: import os
  from typing import List
  import absl
  import tensorflow
  from tensorflow import keras
  import tensorflow_transform
  from tfx.components.trainer.rewriting import converters
  from tfx.components.trainer.rewriting import rewriter
  from tfx.components.trainer.rewriting import rewriter_factory
  from tfx.dsl.io import fileio
  from tfx import v1
  from tfx_bsl.public import tfxio
label_desc: Use image node_id 5 for calling a global function with example usage: run_fn(fn_args) without return types