code: def _get_inference_fn(model, tf_transform_output):
      """Defines the function used for inference."""
      model.tft_layer = tf_transform_output.transform_features_layer()
  
      @tf.function
      def inference_fn(cur_page, session_index):
          """Returns the output to be used in the serving signature."""
          return model(
              {
                  _CUR_PAGE_FEATURE_KEY: cur_page,
                  _SESSION_INDEX_FEATURE_KEY: session_index,
              }
          )
  
      return inference_fn
  
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 4 for calling a global function with example usage: _get_inference_fn(model, tf_transform_output) and returns: inference_fn, model