tfr.keras.model.ModelBuilderWithMask

Interface to build a tf.keras.Model for ranking with a mask Tensor.

Inherits From: AbstractModelBuilder

tfr.keras.model.ModelBuilderWithMask(
    mask_feature_name: str, name: Optional[str] = None
)

The ModelBuilderWithMask class is an abstract class to build a ranking model based on dense Tensors and a mask Tensor to indicate the padded ones. All the boilerplate codes related to constructing a tf.keras.Model are integrated in the ModelBuilder class.

To be implemented by subclasses:

  • create_inputs(): Contains the logic to create tf.keras.Input for context and example inputs and mask for valid list items.
  • preprocess(): Contains the logic to preprocess context and example inputs.
  • score(): Contains the logic to score examples in list and return outputs.

Example subclass implementation:

class SimpleModelBuilder(ModelBuilderWithMask):

  def __init__(self, context_feature_spec, example_feature_spec,
               mask_feature_name, name=None):
    self._context_feature_spec = context_feature_spec
    self._example_feature_spec = example_feature_spec
    self._mask_feature_name = mask_feature_name
    self._name = name

  def create_inputs(self):
    context_inputs = {
        name: tf.keras.Input(
            shape=tuple(spec.shape),
            name=name,
            dtype=spec.dtype
        ) for name, spec in self._context_feature_spec.items()
    }
    example_inputs = {
        name: tf.keras.Input(
            shape=(None,) + tuple(spec.shape),
            name=name,
            dtype=spec.dtype
        ) for name, spec in self._example_feature_spec.items()
    }
    mask = tf.keras.Input(
        name=self._mask_feature_name, shape=(None,), dtype=tf.bool)
    return context_inputs, example_inputs, mask

  def preprocess(self, context_inputs, example_inputs, mask):
    context_features = {
        name: tf.math.log1p(
            tf.abs(tensor)) for name, tensor in context_inputs.items()
    }
    example_features = {
        name: tf.math.log1p(
            tf.abs(tensor)) for name, tensor in example_inputs.items()
    }
    return context_features, example_features

  def score(self, context_features, example_features, mask):
    x = tf.concat([tensor for tensor in example_features.values()], -1)
    return tf.keras.layers.Dense(1)(x)

mask_feature_name name of 2D mask boolean feature.
name (optional) name of the Model.

Methods

build

View source

build() -> tf.keras.Model

Builds a Keras Model for Ranking Pipeline.

Example usage:

model_builder = SimpleModelBuilder(
    {},
    {"example_feature_1": tf.io.FixedLenFeature(
        shape=(1,), dtype=tf.float32, default_value=0.0)},
    "list_mask", "model_builder")
model = model_builder.build()

Returns
A tf.keras.Model.

create_inputs

View source

@abc.abstractmethod
create_inputs() -> Tuple[tfr.keras.model.TensorDict, tfr.keras.model.TensorDict, tf.Tensor]

Creates context and example inputs.

Example usage:

model_builder = SimpleModelBuilder(
    {},
    {"example_feature_1": tf.io.FixedLenFeature(
        shape=(1,), dtype=tf.float32, default_value=0.0)},
    "list_mask", "model_builder")
context_inputs, example_inputs, mask = model_builder.create_inputs()

Returns
A tuple of
context_inputs maps from context feature keys to Keras Input.
example_inputs maps from example feature keys to Keras Input.
mask Keras Input for the mask feature.

preprocess

View source

@abc.abstractmethod
preprocess(
    context_inputs: tfr.keras.model.TensorDict,
    example_inputs: tfr.keras.model.TensorDict,
    mask: tf.Tensor
) -> Tuple[tfr.keras.model.TensorDict, tfr.keras.model.TensorDict]

Preprocesses context and example inputs.

Example usage:

model_builder = SimpleModelBuilder(
    {},
    {"example_feature_1": tf.io.FixedLenFeature(
        shape=(1,), dtype=tf.float32, default_value=0.0)},
    "list_mask", "model_builder")
context_inputs, example_inputs, mask = model_builder.create_inputs()
context_features, example_features = model_builder.preprocess(
    context_inputs, example_inputs, mask)

Args
context_inputs maps context feature keys to tf.keras.Input.
example_inputs maps example feature keys to tf.keras.Input.
mask [batch_size, list_size]-tensor of mask for valid examples.

Returns
A tuple of
context_features maps from context feature keys to [batch_size, feature_dims]-tensors of preprocessed context features.
example_features maps from example feature keys to [batch_size, list_size, feature_dims]-tensors of preprocessed example features.

score

View source

@abc.abstractmethod
score(
    context_features: tfr.keras.model.TensorDict,
    example_features: tfr.keras.model.TensorDict,
    mask: tf.Tensor
) -> Union[TensorLike, TensorDict]

Scores all examples and returns outputs.

Example usage:

model_builder = SimpleModelBuilder(
    {},
    {"example_feature_1": tf.io.FixedLenFeature(
        shape=(1,), dtype=tf.float32, default_value=0.0)},
    "list_mask", "model_builder")
context_inputs, example_inputs, mask = model_builder.create_inputs()
context_features, example_features = model_builder.preprocess(
    context_inputs, example_inputs, mask)
scores = model_builder.score(context_features, example_features)

Args
context_features maps from context feature keys to [batch_size, feature_dims]-tensors of preprocessed context features.
example_features maps from example feature keys to [batch_size, list_size, feature_dims]-tensors of preprocessed example features.
mask [batch_size, list_size]-tensor of mask for valid examples.

Returns
A [batch_size, list_size]-tensor of logits or a dict mapping task name to logits in the multi-task setting.