serving_signatures.py

def get_serve_raw_fn(model, tf_transform_output):
    """
        Returns a tf.function that preforms preprocessing and inference.
        For usage as the seving_raw signature.
    """
    model.tft_layer = tf_transform_output.transform_features_layer()

    @tf.function
    def serve_raw_fn(features):
        # you cancan also drop label key or non-needed keys here
        transformed_features = model.tft_layer(features)
        return model(transformed_features)

    return serve_raw_fn

serving_raw_entry = get_serve_raw_fn(
    model, tf_transform_output)

serving_raw_signature_tensorspecs = {x: tf.TensorSpec(
    shape=[None, 1], name=x) for x in INPUT_FEATURE_KEYS} # also specify dtype