code: def run_fn(fn_args: FnArgs):
      """Train the model based on given args.
  
      Args:
        fn_args: Holds args used to train the model as name/value pairs.
  
      Raises:
        ValueError: if invalid inputs.
      """
      tf_transform_output = tft.TFTransformOutput(
          fn_args.transform_output
      )
  
      train_dataset = _input_fn(
          fn_args.train_files,
          fn_args.data_accessor,
          tf_transform_output,
          is_train=True,
          batch_size=_TRAIN_BATCH_SIZE,
      )
      eval_dataset = _input_fn(
          fn_args.eval_files,
          fn_args.data_accessor,
          tf_transform_output,
          is_train=False,
          batch_size=_EVAL_BATCH_SIZE,
      )
  
      model, base_model = _build_keras_model()
  
      absl.logging.info(
          "Tensorboard logging to {}".format(fn_args.model_run_dir)
      )
      # Write logs to path
      tensorboard_callback = tf.keras.callbacks.TensorBoard(
          log_dir=fn_args.model_run_dir, update_freq="epoch"
      )
  
      # Our training regime has two phases: we first freeze the backbone and train
      # the newly added classifier only, then unfreeze part of the backbone and
      # fine-tune with classifier jointly.
      steps_per_epoch = int(_TRAIN_DATA_SIZE / _TRAIN_BATCH_SIZE)
      total_epochs = int(fn_args.train_steps / steps_per_epoch)
      if _CLASSIFIER_EPOCHS > total_epochs:
          raise ValueError(
              "Classifier epochs is greater than the total epochs"
          )
  
      absl.logging.info("Start training the top classifier")
      model.fit(
          train_dataset,
          epochs=_CLASSIFIER_EPOCHS,
          steps_per_epoch=steps_per_epoch,
          validation_data=eval_dataset,
          validation_steps=fn_args.eval_steps,
          callbacks=[tensorboard_callback],
      )
  
      absl.logging.info("Start fine-tuning the model")
      # Unfreeze the top MobileNet layers and do joint fine-tuning
      _freeze_model_by_percentage(base_model, 0.9)
  
      # We need to recompile the model because layer properties have changed
      model.compile(
          loss="sparse_categorical_crossentropy",
          optimizer=tf.keras.optimizers.RMSprop(
              lr=_FINETUNE_LEARNING_RATE
          ),
          metrics=["sparse_categorical_accuracy"],
      )
      model.summary(print_fn=absl.logging.info)
  
      model.fit(
          train_dataset,
          initial_epoch=_CLASSIFIER_EPOCHS,
          epochs=total_epochs,
          steps_per_epoch=steps_per_epoch,
          validation_data=eval_dataset,
          validation_steps=fn_args.eval_steps,
          callbacks=[tensorboard_callback],
      )
  
      # Prepare the TFLite model used for serving in MLKit
      signatures = {
          "serving_default": _get_serve_image_fn(
              model
          ).get_concrete_function(
              tf.TensorSpec(
                  shape=[None, 224, 224, 3],
                  dtype=tf.float32,
                  name=_transformed_name(_IMAGE_KEY),
              )
          )
      }
  
      temp_saving_model_dir = os.path.join(
          fn_args.serving_model_dir, "temp"
      )
      model.save(
          temp_saving_model_dir, save_format="tf", signatures=signatures
      )
  
      tfrw = rewriter_factory.create_rewriter(
          rewriter_factory.TFLITE_REWRITER, name="tflite_rewriter"
      )
      converters.rewrite_saved_model(
          temp_saving_model_dir,
          fn_args.serving_model_dir,
          tfrw,
          rewriter.ModelType.TFLITE_MODEL,
      )
  
      # Add necessary TFLite metadata to the model in order to use it within MLKit
      # TODO(dzats@): Handle label map file path more properly, currently
      # hard-coded.
      tflite_model_path = os.path.join(
          fn_args.serving_model_dir, _TFLITE_MODEL_NAME
      )
      # TODO(dzats@): Extend the TFLite rewriter to be able to add TFLite metadata
      # @ to the model.
      _write_metadata(
          model_path=tflite_model_path,
          label_map_path=fn_args.custom_config["labels_path"],
          mean=[127.5],
          std=[127.5],
      )
  
      fileio.rmtree(temp_saving_model_dir)
  
imports: import os
  from typing import List
  import absl
  import flatbuffers
  import tensorflow
  import tensorflow_transform
  from tfx.components.trainer.fn_args_utils import DataAccessor
  from tfx.components.trainer.fn_args_utils import FnArgs
  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_bsl.tfxio import dataset_options
  from tflite_support import metadata_schema_py_generated
  from tflite_support import metadata
label_desc: Use image node_id 10 for calling a global function with example usage: run_fn(fn_args) without return types