def define_tfrecord_signature(
model,
tfrecord_type: str,
feature_config: FeatureConfig,
preprocessing_keys_to_fns: dict,
postprocessing_fn=None,
required_fields_only: bool = True,
pad_sequence: bool = False,
max_sequence_size: int = 0,
):
"""
Add signatures to the tf keras savedmodel
Returns:
Serving signature function that accepts a TFRecord string tensor and returns predictions
"""
# TFRecord Signature
# Define a parsing function for tfrecord protos
inputs = feature_config.get_all_features(key="node_name", include_label=False)
"""
NOTE:
Setting pad_sequence=False for tfrecord signature as it is used at inference time
and we do NOT want to score on padded records for performance reasons
Limitation: This limits the serving signature to only run inference on a single query
at a time given the current implementation. This is a tricky issue to fix because
there is no real way to generate a dense tensor of ranking scores from different queries,
as they might have varying number of records in each of them.
Workaround: To infer on multiple queries, run predict() on each of the queries separately.
"""
tfrecord_parse_fn = get_parse_fn(
feature_config=feature_config,
tfrecord_type=tfrecord_type,
preprocessing_keys_to_fns=preprocessing_keys_to_fns,
max_sequence_size=max_sequence_size,
required_fields_only=required_fields_only,
pad_sequence=pad_sequence,
)
dtype_map = dict()
for feature_info in feature_config.get_all_features(include_label=False):
feature_node_name = feature_info.get("node_name", feature_info["name"])
dtype_map[feature_node_name] = feature_config.get_dtype(feature_info)
# Define a serving signature for tfrecord
@tf.function(input_signature=[TensorSpec(shape=[None], dtype=tf.string)])
def _serve_tfrecord(protos):
input_size = tf.shape(protos)[0]
features_dict = {
feature: TensorArray(dtype=dtype_map[feature], size=input_size) for feature in inputs
}
# Define loop index
i = tf.constant(0)
# Define loop condition
def loop_condition(i, protos, features_dict):
return tf.less(i, input_size)
# Define loop body
def loop_body(i, protos, features_dict):
features, labels = tfrecord_parse_fn(protos[i])
for feature, feature_val in features.items():
features_dict[feature] = features_dict[feature].write(i, feature_val)
i += 1
return i, protos, features_dict
# Parse all SequenceExample protos to get features
_, _, features_dict = tf.while_loop(
cond=loop_condition, body=loop_body, loop_vars=[i, protos, features_dict],
)
# Convert TensorArray to tensor
features_dict = {k: v.stack() for k, v in features_dict.items()}
# Run the model to get predictions
predictions = model(inputs=features_dict)
# Define a post hook
if postprocessing_fn:
predictions = postprocessing_fn(predictions, features_dict)
return predictions
return _serve_tfrecord
def define_tfrecord_signature(
model,
tfrecord_type: str,
feature_config: FeatureConfig,
preprocessing_keys_to_fns: dict,
postprocessing_fn=None,
required_fields_only: bool = True,
pad_sequence: bool = False,
max_sequence_size: int = 0,
):
"""
Serving signature that wraps around the keras model trained as a RelevanceModel
with a pre-step to parse TFRecords and apply additional feature preprocessing
Parameters
----------
model : keras Model
Keras model object to be saved
tfrecord_type : {"example", "sequence_example"}
Type of the TFRecord protobuf that the saved model will be used on at serving time
feature_config : `FeatureConfig` object
FeatureConfig object that defines the input features into the model
and the corresponding feature preprocesing functions to be used
in the serving signature
preprocessing_keys_to_fns : dict
Dictionary mapping function names to tf.functions that should be saved in the preprocessing step of the tfrecord serving signature
postprocessing_fn: function
custom tensorflow compatible postprocessing function to be used at serving time.
Saved as part of the postprocessing layer of the tfrecord serving signature
required_fields_only: bool
boolean value defining if only required fields
need to be added to the tfrecord parsing function at serving time
pad_sequence: bool, optional
Value defining if sequences should be padded for SequenceExample proto inputs at serving time.
Set this to False if you want to not handle padded scores.
max_sequence_size : int, optional
Maximum sequence size for SequenceExample protobuf
The protobuf object will be padded or clipped to this value
Returns
-------
`tf.function`
Serving signature function that accepts a TFRecord string tensor and returns predictions
"""
# TFRecord Signature
# Define a parsing function for tfrecord protos
inputs = feature_config.get_all_features(key="node_name",
include_label=False)
"""
NOTE:
Setting pad_sequence=False for tfrecord signature as it is used at inference time
and we do NOT want to score on padded records for performance reasons
Limitation: This limits the serving signature to only run inference on a single query
at a time given the current implementation. This is a tricky issue to fix because
there is no real way to generate a dense tensor of ranking scores from different queries,
as they might have varying number of records in each of them.
Workaround: To infer on multiple queries, run predict() on each of the queries separately.
"""
tfrecord_parse_fn = get_parse_fn(
feature_config=feature_config,
tfrecord_type=tfrecord_type,
preprocessing_keys_to_fns=preprocessing_keys_to_fns,
max_sequence_size=max_sequence_size,
required_fields_only=required_fields_only,
pad_sequence=pad_sequence,
)
dtype_map = dict()
for feature_info in feature_config.get_all_features(include_label=False):
feature_node_name = feature_info.get("node_name", feature_info["name"])
dtype_map[feature_node_name] = feature_config.get_dtype(feature_info)
# Define a serving signature for tfrecord
@tf.function(input_signature=[TensorSpec(shape=[None], dtype=tf.string)])
def _serve_tfrecord(protos):
input_size = tf.shape(protos)[0]
features_dict = {
feature: TensorArray(dtype=dtype_map[feature], size=input_size)
for feature in inputs
}
# Define loop index
i = tf.constant(0)
# Define loop condition
def loop_condition(i, protos, features_dict):
return tf.less(i, input_size)
# Define loop body
def loop_body(i, protos, features_dict):
features, labels = tfrecord_parse_fn(protos[i])
for feature, feature_val in features.items():
features_dict[feature] = features_dict[feature].write(
i, feature_val)
i += 1
return i, protos, features_dict
# Parse all SequenceExample protos to get features
_, _, features_dict = tf.while_loop(
cond=loop_condition,
body=loop_body,
loop_vars=[i, protos, features_dict],
)
# Convert TensorArray to tensor
features_dict = {k: v.stack() for k, v in features_dict.items()}
# Run the model to get predictions
predictions = model(inputs=features_dict)
# Define a post hook
if postprocessing_fn:
predictions = postprocessing_fn(predictions, features_dict)
return predictions
return _serve_tfrecord