Renaming arguments to match features in serving function Keras DNN / TensorFlow

[Tushar Bisht]

We developed a Keras DNN based ML model on TF / TFX 

This model was trained on features, some of them listed below

maxlimit

minlimit

departday

maxload

Our serving function looks something like below:

def _get_serve_tf_examples_fn(model, tf_transform_output):

    """Returns a function that parses a serialized tf.Example and applies TFT."""

    model.tft_layer = tf_transform_output.transform_features_layer()

    @tf.function

    def serve_tf_examples_fn(serialized_tf_examples):

        """Returns the output to be used in the serving signature."""

        feature_spec = tf_transform_output.raw_feature_spec()

        feature_spec.pop(_LABEL_KEY)

        parsed_features = tf.io.parse_example(serialized_tf_examples, feature_spec)

        transformed_features = model.tft_layer(parsed_features)

        return model(transformed_features)

    return serve_tf_examples_fn

the model is current hosted using GCP AI serving

When a request is sent on the endpoint with arguments same as feature names, we have correct predictions  returned (something that is expected)

But some of our consumers, due to design restrains at their end, have prefixes on the arguments, like below

journey:maxlimit

journey:minlimit

departday

segment:maxload

This is leading  the model to ignore such unrecognized features in the request and come up with wrong prediciton

I was trying to rename the features in the serving function like below:

def _get_serve_tf_examples_fn(model, tf_transform_output):

    """Returns a function that parses a serialized tf.Example and applies TFT."""

    model.tft_layer = tf_transform_output.transform_features_layer()

    @tf.function

    def serve_tf_examples_fn(serialized_tf_examples):

        """Returns the output to be used in the serving signature."""

        feature_spec = tf_transform_output.raw_feature_spec()

        feature_spec.pop(_LABEL_KEY)

        

        old_max_limit = 'journey:maxlimit'

        old_min_limit = 'journey:minlimit'

        old_seg_max_load = 'segment:maxload'

        

        

        new_max_limit = 'maxlimit'

        new_min_limit = 'minlimit'

        

        new_seg_max_load = 'segment:maxload'

        

        feature_spec[new_max_limit] = feature_spec.pop(old_max_limit)

        feature_spec[new_min_limit] = feature_spec.pop(old_min_limit)

        feature_spec[new_seg_max_load] = feature_spec.pop(old_seg_max_load)

            

        parsed_features = tf.io.parse_example(serialized_tf_examples, feature_spec)

        transformed_features = model.tft_layer(parsed_features)

        return model(transformed_features)

    return serve_tf_examples_fn

 This gives me an error:

 KeyError: in user code:

 

    trainer.py:60 serve_tf_examples_fn  *

        feature_spec[new_max_limit] = feature_spec.pop(old_max_limit)

 

    KeyError: 'journey:maxlimit'

  And the model does not train

  Can some help be provided on how to rename arguments to matcht the  feature names in above problem statement