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():
          bert_layer = hub.KerasLayer(_BERT_LINK, trainable=True)
          model = build_and_compile_bert_classifier(
              bert_layer, _MAX_LEN, 2
          )
  
      model.fit(
          train_dataset,
          epochs=_EPOCHS,
          steps_per_epoch=fn_args.train_steps,
          validation_data=eval_dataset,
          validation_steps=fn_args.eval_steps,
      )
  
      signatures = {
          "serving_default": _get_serve_tf_examples_fn(
              model, tf_transform_output
          ).get_concrete_function(
              tf.TensorSpec(
                  shape=[None], dtype=tf.string, name="examples"
              )
          ),
      }
      model.save(
          fn_args.serving_model_dir,
          save_format="tf",
          signatures=signatures,
      )
  
file_path: tfx/tfx/examples/bert/cola/bert_cola_utils.py
imports: from typing import List
  import tensorflow
  import tensorflow_data_validation
  import tensorflow_hub
  import tensorflow_transform
  from tfx import v1
  from tfx.components.transform import stats_options_util
  from tfx.examples.bert.utils.bert_models import build_and_compile_bert_classifier
  from tfx.examples.bert.utils.bert_tokenizer_utils import BertPreprocessor
  from tfx_bsl.public import tfxio
  from google.protobuf import text_format
label_desc: Use image node_id 6 for calling a global function with example usage: run_fn(fn_args) without return types