def export_serving_model(
    tf_transform_output: tft.TFTransformOutput,
    model: tf_keras.Model,
    output_dir: str,
):
  """Creates an input function reading from raw data.

  Args:
    tf_transform_output: Wrapper around output of tf.Transform.
    model: The keras model to export.
    output_dir: A path to export the model to.

  Returns:
    The serving input function.
  """
  # The layer has to be saved to the model for keras tracking purpases.
  model.tft_layer = tf_transform_output.transform_features_layer()

  raw_feature_spec = RAW_DATA_FEATURE_SPEC.copy()
  # Remove label since it is not available during serving.
  raw_feature_spec.pop(LABEL_KEY)

  @tf.function
  def serve_tf_examples_fn(serialized_tf_examples):
    """Serving tf.function model wrapper."""
    parsed_features = tf.io.parse_example(
        serialized_tf_examples, raw_feature_spec
    )
    transformed_features = model.tft_layer(parsed_features)
    outputs = model(transformed_features)
    classes_names = tf.constant([['0', '1']])
    classes = tf.tile(classes_names, [tf.shape(outputs)[0], 1])
    return {'classes': classes, 'scores': outputs}

  concrete_serving_fn = serve_tf_examples_fn.get_concrete_function(
      tf.TensorSpec(shape=[None], dtype=tf.string, name='inputs')
  )
  signatures = {'serving_default': concrete_serving_fn}

  # This is required in order to make this model servable with model_server.
  versioned_output_dir = os.path.join(output_dir, '1')
  model.save(versioned_output_dir, save_format='tf', signatures=signatures)
Creates an input function reading from raw data.

  Args:
    tf_transform_output: Wrapper around output of tf.Transform.
    model: The keras model to export.
    output_dir: A path to export the model to.

  Returns:
    The serving input function.
  
examples/sentiment_example_v2.py