def test_create_keras_inputs_sparse_features(self):
context_feature_columns = {
'query':
tf.feature_column.categorical_column_with_vocabulary_list(
'query',
vocabulary_list=[
'ranking', 'regression', 'classification', 'ordinal'
])
}
example_feature_columns = {
'title':
tf.feature_column.categorical_column_with_vocabulary_list(
'title',
vocabulary_list=[
'ranking', 'regression', 'classification', 'ordinal'
])
}
keras_inputs = feature.create_keras_inputs(
context_feature_columns=context_feature_columns,
example_feature_columns=example_feature_columns,
size_feature_name='example_list_size')
self.assertCountEqual(
keras_inputs.keys(),
list(context_feature_columns.keys()) +
list(example_feature_columns.keys()) + ['example_list_size'])
def test_create_keras_inputs(self):
keras_inputs = feature.create_keras_inputs(
context_feature_columns=self._context_feature_columns,
example_feature_columns=self._example_feature_columns,
size_feature_name='example_list_size')
self.assertCountEqual(
keras_inputs.keys(),
list(self._context_feature_columns.keys()) +
list(self._example_feature_columns.keys()) + ['example_list_size'])
def test_keras_inputs_dynamic_list_shape(self):
keras_inputs = feature.create_keras_inputs(
context_feature_columns=self._context_feature_columns,
example_feature_columns=self._example_feature_columns,
size_feature_name=None)
self.assertEqual(keras_inputs['query_length'].shape.as_list(), [None, 1])
self.assertEqual(keras_inputs['query_length'].dtype, tf.int64)
self.assertEqual(keras_inputs['utility'].shape.as_list(), [None, None, 1])
self.assertEqual(keras_inputs['utility'].dtype, tf.float32)
self.assertEqual(keras_inputs['unigrams'].dtype, tf.string)
def create_keras_model(network,
loss,
metrics,
optimizer,
size_feature_name,
list_size=None):
"""Creates a Functional Keras ranking model.
A mask is inferred from size_feature_name and passed to the network, along
with feature dictionary as inputs.
Args:
network: (`tfr.keras.network.RankingNetwork`) A ranking network which
generates a list of scores.
loss: (`tfr.keras.losses._RankingLoss`) A ranking loss.
metrics: (list) List of ranking metrics, `tfr.keras.metrics._RankingMetric`
instances.
optimizer: (`tf.keras.optimizer.Optimizer`) Optimizer to minimize ranking
loss.
size_feature_name: (str) Name of feature for example list sizes. If not
None, this feature name corresponds to a `tf.int32` Tensor of size
[batch_size] corresponding to sizes of example lists. If `None`, all
examples are treated as valid.
list_size: (int) The list size for example features. If None, use dynamic
list size. A fixed list size is required for TPU training.
Returns:
A compiled ranking Keras model, a `tf.keras.Model` instance.
"""
# TODO: Support compatibility with TPUs.
keras_inputs = feature.create_keras_inputs(
context_feature_columns=network.context_feature_columns,
example_feature_columns=network.example_feature_columns,
size_feature_name=size_feature_name,
list_size=list_size)
# Create mask from sizes and list_size.
mask = None
if size_feature_name is not None:
mask = feature.GenerateMask(network.example_feature_columns,
size_feature_name)(
keras_inputs)
logits = network(inputs=keras_inputs, mask=mask)
ranker = tf.keras.Model(inputs=keras_inputs, outputs=logits)
ranker.compile(optimizer=optimizer, loss=loss, metrics=metrics)
return ranker
