def setUp(self):
super(GenerateMaskTest, self).setUp()
self.context_feature_columns = _context_feature_columns()
self.example_feature_columns = _example_feature_columns()
# Batch size = 2, list_size = 2.
self.features = _features()
self.mask_generator_layer = feature.GenerateMask(
example_feature_columns=self.example_feature_columns,
size_feature_name='example_feature_size')
def setUp(self):
super(GenerateMaskTest, self).setUp()
(context_feature_columns, example_feature_columns,
custom_objects) = _get_feature_columns()
self._context_feature_columns = context_feature_columns
self._example_feature_columns = example_feature_columns
self._custom_objects = custom_objects
# Batch size = 2, list_size = 2.
self._features = _features()
self._mask_generator_layer = feature.GenerateMask(
example_feature_columns=self._example_feature_columns,
size_feature_name='example_feature_size')
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