#
# Copyright (c) 2021, NVIDIA CORPORATION.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
from typing import Optional, Tuple, Union
import tensorflow as tf
from tensorflow.keras.backend import random_bernoulli
from merlin.models.tf.core import combinators
from merlin.models.tf.core.base import Block, BlockType, PredictionOutput
from merlin.models.tf.core.combinators import TabularBlock
from merlin.models.tf.transforms.features import PrepareFeatures
from merlin.models.tf.typing import TabularData
from merlin.models.tf.utils import tf_utils
from merlin.models.utils import schema_utils
from merlin.models.utils.dependencies import is_transformers_available
from merlin.schema import ColumnSchema, Schema, Tags
@Block.registry.register_with_multiple_names("remove_pad_3d")
@tf.keras.utils.register_keras_serializable(package="merlin_models")
class RemovePad3D(Block):
"""
Flatten the sequence of labels and filter out non-targets positions
Parameters
----------
padding_idx: int
The padding index value.
Defaults to 0.
Returns
-------
targets: tf.Tensor
The flattened vector of true targets positions
flatten_predictions: tf.Tensor
If the predictions are 3-D vectors (sequential task),
flatten the predictions vectors to keep only the ones related to target positions.
"""
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.padding_idx = 0
def compute_output_shape(self, input_shape):
return input_shape
def call_outputs(
self, outputs: PredictionOutput, training=True, **kwargs
) -> "PredictionOutput":
targets, predictions = outputs.targets, outputs.predictions
targets = tf.reshape(targets, (-1,))
non_pad_mask = targets != self.padding_idx
targets = tf.boolean_mask(targets, non_pad_mask)
assert isinstance(predictions, tf.Tensor), "Predictions must be a tensor"
if len(tuple(predictions.get_shape())) == 3:
predictions = tf.reshape(predictions, (-1, predictions.shape[-1]))
predictions = tf.boolean_mask(
predictions, tf.broadcast_to(tf.expand_dims(non_pad_mask, 1), tf.shape(predictions))
)
return outputs.copy_with_updates(predictions=predictions, targets=targets)
@tf.keras.utils.register_keras_serializable(package="merlin_models")
class SequenceTransform(TabularBlock):
"""Base class for preparing sequential inputs and targets.
Parameters
----------
schema : Schema
The schema with the sequential columns to be truncated
target : Union[str, Tags, ColumnSchema]
The sequential input column that will be used to extract the target
pre : Optional[BlockType], optional
A block that is called before this method call().
P.s. The PrepareFeatures() block is applied to convert
the tuple representation of sequential features to RaggedTensors,
so that the tensors sequences can be shifted/truncated
transformer: Optional[TransformerBlock]
The transformer block that leverages the group of sequences returned
by the given SequenceTransform, by default None.
"""
def __init__(
self,
schema: Schema,
target: Union[str, Tags, ColumnSchema],
pre: Optional[BlockType] = None,
transformer=None,
**kwargs,
):
_pre = PrepareFeatures(schema)
if pre:
_pre = _pre.connect(pre)
super().__init__(pre=_pre, schema=schema, **kwargs)
self.target = target
self.target_name = self._get_target(target)
self.transformer = transformer
def _get_target(self, target):
if (
(isinstance(target, str) and target not in self.schema.column_names)
or (isinstance(target, Tags) and len(self.schema.select_by_tag(target)) > 0)
or (isinstance(target, ColumnSchema) and target not in self.schema)
):
raise ValueError("The target column needs to be part of the sequential schema")
target_name = target
if isinstance(target, ColumnSchema):
target_name = target.name
if isinstance(target, Tags):
if len(self.schema.select_by_tag(target)) > 1:
raise ValueError(
"Only 1 column should the Tag ({target}) provided for target, but"
f"the following columns have that tag: "
f"{self.schema.select_by_tag(target).column_names}"
)
target_name = self.schema.select_by_tag(target).column_names[0]
return target_name
def call(
self, inputs: TabularData, targets=None, training=False, testing=False, **kwargs
) -> Tuple:
raise NotImplementedError()
def _check_seq_inputs_targets(self, inputs: TabularData):
if self.target_name not in inputs:
raise ValueError(
f"The inputs provided does contain the target column ({self.target_name})"
)
target_shape = inputs[self.target_name].get_shape().as_list()
if len(target_shape) < 2:
raise ValueError(
f"The sequential target column ({self.target_name}) cannot be a 1D tensor,"
f" but the shape is {target_shape}"
)
if target_shape[1] == 1:
raise ValueError(
f"The 2nd dim of the target column ({self.target_name}) should be greater"
" than 1, so that the sequential input can be shifted as target"
)
seq_inputs_shapes = {
col: inputs[col].get_shape().as_list() for col in self.schema.column_names
}
seq_shapes = list(seq_inputs_shapes.values())
if not all(x[:2] == seq_shapes[0][:2] for x in seq_shapes):
raise ValueError(
"The sequential inputs must have the first two dims equal, but they "
f"are different: {seq_inputs_shapes}"
)
def compute_output_shape(self, input_shape):
new_input_shapes = dict()
for k, v in input_shape.items():
new_input_shapes[k] = v
if k in self.schema.column_names:
# If it is a list/sparse feature (in tuple representation), uses the offset as shape
if isinstance(v, tuple) and isinstance(v[1], tf.TensorShape):
new_input_shapes[k] = tf.TensorShape([v[1][0], None])
else:
new_input_shapes[k] = v
return new_input_shapes
def get_config(self):
"""Returns the config of the layer as a Python dictionary."""
config = super().get_config()
target = self.target
if isinstance(target, ColumnSchema):
target = schema_utils.schema_to_tensorflow_metadata_json(Schema([target]))
config["target"] = target
return config
@classmethod
def from_config(cls, config):
"""Creates layer from its config. Returning the instance."""
config = tf_utils.maybe_deserialize_keras_objects(config, ["pre", "post", "aggregation"])
config["schema"] = schema_utils.tensorflow_metadata_json_to_schema(config["schema"])
if config["target"].startswith("{"): # we have a schema
config["target"] = [
col for col in schema_utils.tensorflow_metadata_json_to_schema(config["target"])
][0]
schema = config.pop("schema")
target = config.pop("target")
return cls(schema, target, **config)
def configure_for_train(self):
"""Method called by the model.fit() to set additional model's
configuration before calling keras parent class `fit()`
"""
pass
def configure_for_test(self):
"""Method called by the model.evaluate() to check any custom model's
configuration before calling keras parent class `evaluate()`
"""
pass
[docs]@Block.registry.register_with_multiple_names("seq_predict_next")
@tf.keras.utils.register_keras_serializable(package="merlin_models")
class SequencePredictNext(SequenceTransform):
"""Prepares sequential inputs and targets for next-item prediction.
The target is extracted from the shifted sequence of item ids and
the sequential input features are truncated in the last position.
Parameters
----------
schema : Schema
The schema with the sequential columns to be truncated
target : Union[str, Tags, ColumnSchema]
The sequential input column that will be used to extract the target
pre : Optional[BlockType], optional
A block that is called before this method call().
P.s. The PrepareFeatures() is called automatically before the pre
to convert the list features representation
"""
[docs] def call(
self, inputs: TabularData, targets=None, training=False, testing=False, **kwargs
) -> Tuple:
self._check_seq_inputs_targets(inputs)
# Shifts the target column to be the next item of corresponding input column
new_target = inputs[self.target_name][:, 1:]
if targets is None:
targets = dict({self.target_name: new_target})
elif isinstance(targets, dict):
targets[self.target_name] = new_target
else:
raise ValueError("Targets should be None or a dict of tensors")
new_inputs = dict()
for k, v in inputs.items():
if k in self.schema.column_names:
# Removes the last item of the sequence, as it belongs to the target
new_inputs[k] = v[:, :-1]
else:
new_inputs[k] = v
return (new_inputs, targets)
[docs] def compute_output_shape(self, input_shape):
new_input_shapes = dict()
for k, v in input_shape.items():
new_input_shapes[k] = v
if k in self.schema.column_names:
# If it is a list/sparse feature (in tuple representation), uses the offset as shape
if isinstance(v, tuple) and isinstance(v[1], tf.TensorShape):
new_input_shapes[k] = tf.TensorShape([v[1][0], None])
else:
# Reducing 1 position of the seq length
new_input_shapes[k] = tf.TensorShape([v[0], v[1] - 1])
return new_input_shapes
[docs] def compute_mask(self, inputs, mask=None):
new_item_id_seq = inputs[self.target_name][:, :-1]
target_mask = tf.RaggedTensor.from_row_lengths(
values=tf.ones_like(new_item_id_seq.flat_values, dtype=tf.bool),
row_lengths=new_item_id_seq.row_lengths(),
)
self.target_mask = tf.squeeze(target_mask, axis=-1)
targets_mask = dict({self.target_name: self.target_mask})
inputs_mask = dict()
for k, v in inputs.items():
if k in self.schema.column_names:
inputs_mask[k] = self.target_mask
else:
inputs_mask[k] = None
return (inputs_mask, targets_mask)
[docs]@Block.registry.register_with_multiple_names("seq_predict_last")
@tf.keras.utils.register_keras_serializable(package="merlin_models")
class SequencePredictLast(SequenceTransform):
"""Prepares sequential inputs and targets for last-item prediction.
The target is extracted from the last element of sequence of item ids and
the sequential input features are truncated before the last position.
Parameters
----------
schema : Schema
The schema with the sequential columns to be truncated
target : Union[str, Tags, ColumnSchema]
The sequential input column that will be used to extract the target
pre : Optional[BlockType], optional
A block that is called before this method call().
P.s. The PrepareFeatures() is called automatically before the pre
to convert the list features representation
"""
[docs] @tf.function
def call(
self, inputs: TabularData, targets=None, training=False, testing=False, **kwargs
) -> Tuple:
self._check_seq_inputs_targets(inputs)
# Shifts the target column to be the next item of corresponding input column
new_target = tf.squeeze(inputs[self.target_name][:, -1:], axis=1)
if targets is None:
targets = dict({self.target_name: new_target})
elif isinstance(targets, dict):
targets[self.target_name] = new_target
else:
raise ValueError("Targets should be None or a dict of tensors")
new_inputs = dict()
for k, v in inputs.items():
if k in self.schema.column_names:
# Removes the last item of the sequence, as it belongs to the target
new_inputs[k] = v[:, :-1]
else:
new_inputs[k] = v
return (new_inputs, targets)
[docs] def compute_output_shape(self, input_shape):
new_input_shapes = dict()
for k, v in input_shape.items():
new_input_shapes[k] = v
if k in self.schema.column_names:
# If it is a list/sparse feature (in tuple representation), uses the offset as shape
if isinstance(v, tuple) and isinstance(v[1], tf.TensorShape):
new_input_shapes[k] = tf.TensorShape([v[1][0], None])
else:
# Reducing 1 position of the seq length
new_input_shapes[k] = tf.TensorShape([v[0], v[1] - 1])
return new_input_shapes
[docs] def compute_mask(self, inputs, mask=None):
new_item_id_seq = inputs[self.target_name][:, :-1]
self.target_mask = self._generate_target_mask(new_item_id_seq)
inputs_mask = dict()
for k, v in inputs.items():
if k in self.schema.column_names:
inputs_mask[k] = self.target_mask
else:
inputs_mask[k] = None
return (inputs_mask, self.target_mask)
def _generate_target_mask(self, ids_seq: tf.RaggedTensor) -> tf.RaggedTensor:
"""Returns a bool ragged tensor with the last positions of the sequence masked
Parameters
----------
ids_seq : tf.RaggedTensor
Sequence of ids, which are used to infer how many values
each sequence contains
Returns
-------
tf.RaggedTensor
Mask tensor, with True at the last positions
"""
row_lengths = ids_seq.row_lengths(1)
max_seq_length = tf.cast(tf.reduce_max(row_lengths), tf.int32)
padding_mask = tf.sequence_mask(row_lengths)
targets_mask = tf.ragged.boolean_mask(
tf.cast(tf.one_hot(row_lengths - 1, max_seq_length), tf.bool), padding_mask
)
return targets_mask
[docs]@Block.registry.register_with_multiple_names("seq_predict_random")
@tf.keras.utils.register_keras_serializable(package="merlin_models")
class SequencePredictRandom(SequenceTransform):
"""Prepares sequential inputs and targets for random-item prediction.
A random element in the sequence (except the first one) is selected
as target and all elements before the selected target as used as
input features.
Parameters
----------
schema : Schema
The schema with the sequential columns to be truncated
target : Union[str, Tags, ColumnSchema]
The sequential input column that will be used to extract the target
pre : Optional[BlockType], optional
A block that is called before this method call().
P.s. The PrepareFeatures() is called automatically before the pre
to convert the list features representation
"""
[docs] def call(
self, inputs: TabularData, targets=None, training=False, testing=False, **kwargs
) -> Tuple:
self._check_seq_inputs_targets(inputs)
batch_size = inputs[self.target_name].nrows()
seq_length = inputs[self.target_name].row_lengths(1)
max_length = tf.reduce_max(seq_length)
random_targets_indices = tf.expand_dims(
tf.cast(
tf.math.floor(
(
tf.random.uniform(
shape=[batch_size], minval=0.0, maxval=1.0, dtype=tf.float32
)
* tf.cast(seq_length - 1, tf.float32)
)
+ 1
),
tf.int32,
),
1,
)
positions_matrix = tf.tile(
tf.expand_dims(tf.range(0, max_length, dtype=tf.int32), 0), [batch_size, 1]
)
self.random_mask = positions_matrix < random_targets_indices
target_mask = positions_matrix == random_targets_indices
new_target = tf.squeeze(
tf.ragged.boolean_mask(inputs[self.target_name], target_mask), axis=1
)
if targets is None:
targets = dict({self.target_name: new_target})
elif isinstance(targets, dict):
targets[self.target_name] = new_target
else:
raise ValueError("Targets should be None or a dict of tensors")
new_inputs = dict()
for k, v in inputs.items():
if k in self.schema.column_names:
new_inputs[k] = tf.ragged.boolean_mask(v, self.random_mask)
else:
new_inputs[k] = v
return (new_inputs, targets)
[docs] def compute_mask(self, inputs, mask=None):
new_item_id_seq = tf.ragged.boolean_mask(inputs[self.target_name], self.random_mask)
self.target_mask = self._generate_target_mask(new_item_id_seq)
inputs_mask = dict()
for k, v in inputs.items():
if k in self.schema.column_names:
inputs_mask[k] = self.target_mask
else:
inputs_mask[k] = None
return (inputs_mask, self.target_mask)
def _generate_target_mask(self, ids_seq: tf.RaggedTensor) -> tf.RaggedTensor:
"""Returns a bool ragged tensor with the last positions of the sequence masked
Parameters
----------
ids_seq : tf.RaggedTensor
Sequence of ids, which are used to infer how many values
each sequence contains
Returns
-------
tf.RaggedTensor
Mask tensor, with True at the last positions
"""
row_lengths = ids_seq.row_lengths(1)
max_seq_length = tf.cast(tf.reduce_max(row_lengths), tf.int32)
padding_mask = tf.sequence_mask(row_lengths)
targets_mask = tf.ragged.boolean_mask(
tf.cast(tf.one_hot(row_lengths - 1, max_seq_length), tf.bool), padding_mask
)
return targets_mask
[docs]@tf.keras.utils.register_keras_serializable(package="merlin.models")
class SequenceMaskRandom(SequenceTargetAsInput):
"""This block implements the Masked Language Modeling (MLM) training approach
introduced in BERT (NLP) and later adapted to RecSys by BERT4Rec [1].
Given an input tf.RaggedTensor with sequences of embeddings
and the corresponding sequence of item ids, some positions are randomly selected (masked)
to be the targets for prediction.
The targets are output being the same as the input ids sequence.
The target masks are returned by using Keras Masking
(._keras_mask), which is set by the compute_mask() method.
Note: The SequenceMaskRandom is meant to be used as a `pre` of model.fit(),
e.g. model.fit(..., pre=SequenceMaskRandom(...)).
Note: Typically during model.evaluate() you want to evaluate the model
to predict the last position of the sequence, as it mimics the next-item predicion
task. In that case, you should use model.evaluate(..., pre=SequenceMaskLast(...))
instead of SequenceMaskRandom(...).
References
----------
.. [1] Sun, Fei, et al. "BERT4Rec: Sequential recommendation with bidirectional encoder
representations from transformer." Proceedings of the 28th ACM international
conference on information and knowledge management. 2019.
Parameters
----------
schema : Schema
The input schema, that will be used to discover the name
of the item id column
target : Union[str, Tags, ColumnSchema]
The sequential input column that will be used to extract the target
masking_prob : float, optional
Probability of an item to be selected (masked) as a label of the given sequence.
Note: We enforce that at least one item is masked for each sequence, so that it
is useful for training, by default 0.2
"""
[docs] def __init__(
self,
schema: Schema,
target: Union[str, Tags, ColumnSchema],
masking_prob: float = 0.2,
**kwargs,
):
self.masking_prob = masking_prob
super().__init__(schema, target, **kwargs)
[docs] @tf.function
def compute_mask(self, inputs, mask=None):
"""Selects (masks) some positions of the targets to be predicted.
This method is called by Keras after call()
and returns the targets mask that will
be assigned to the input tensors and targets, being accessible
by tensor._keras_mask
"""
item_id_seq = inputs[self.target_name]
self.target_mask = self._generate_target_mask(item_id_seq)
inputs_mask = dict()
for k in inputs:
if k in self.schema.column_names:
inputs_mask[k] = self.target_mask
else:
inputs_mask[k] = None
targets_mask = dict({self.target_name: self.target_mask})
return (inputs_mask, targets_mask)
def _generate_target_mask(self, ids_seq: tf.RaggedTensor) -> tf.RaggedTensor:
"""Generates a target mask according to the defined probability and
to the constraints for Masked Language Modeling training (i.e., each
sequence might have between 1 and length-1 masked positions.)
Parameters
----------
ids_seq : tf.RaggedTensor
Sequence of ids, which are used to infer how many values
each sequence contains
Returns
-------
tf.RaggedTensor
Mask tensor, with True at the positions where targets were
selected to be predicted
"""
row_lengths = ids_seq.row_lengths(1)
assertion_min_seq_length = tf.Assert(tf.reduce_all(row_lengths > 1), [row_lengths])
with tf.control_dependencies([assertion_min_seq_length]):
# Targets are masked according to a probability
target_mask_by_prob = self._get_masked_by_prob(row_lengths, prob=self.masking_prob)
# Exactly one target is masked per row
one_target_mask = self._get_one_masked(row_lengths)
# For sequences (rows) with either all or none elements sampled (masked) as targets
# as those sequences would be invalid for training
# the row mask is replaced by a row mask that contains exactly one masked target
replacement_cond = tf.logical_or(
tf.logical_not(tf.reduce_any(target_mask_by_prob, axis=1)),
tf.reduce_all(target_mask_by_prob, axis=1),
)
target_mask = tf.where(
tf.expand_dims(replacement_cond, -1), one_target_mask, target_mask_by_prob
)
padding_mask = tf.sequence_mask(row_lengths)
target_mask_ragged = tf.ragged.boolean_mask(target_mask, padding_mask)
return target_mask_ragged
@staticmethod
def _get_masked_by_prob(row_lengths: tf.Tensor, prob: float) -> tf.Tensor:
"""Generates a dense mask boolean tensor with True values
for randomly selected targets
"""
max_seq_length = tf.cast(tf.reduce_max(row_lengths), tf.int32)
batch_size = tf.shape(row_lengths)[0]
output = tf.cast(random_bernoulli([batch_size, max_seq_length], p=prob), tf.bool)
padding_mask = tf.sequence_mask(row_lengths)
# Ignoring masked items in the padding positions
output = tf.logical_and(output, padding_mask)
return output
@staticmethod
def _get_one_masked(row_lengths: tf.Tensor):
"""Generates a dense mask boolean tensor where for each tensor (row)
there is exactly one True value (selected target)
"""
max_seq_length = tf.cast(tf.reduce_max(row_lengths), tf.int32)
random_targets_indices = tf.cast(
tf.math.floor(
(
tf.random.uniform(
shape=tf.shape(row_lengths), minval=0.0, maxval=1.0, dtype=tf.float32
)
* tf.cast(row_lengths, tf.float32)
)
),
tf.int32,
)
one_target_mask = tf.cast(tf.one_hot(random_targets_indices, max_seq_length), tf.bool)
return one_target_mask
[docs] def get_config(self):
config = super().get_config()
config["masking_prob"] = self.masking_prob
return config
[docs] @classmethod
def from_config(cls, config):
config = tf_utils.maybe_deserialize_keras_objects(config, ["pre", "post", "aggregation"])
schema = schema_utils.tensorflow_metadata_json_to_schema(config.pop("schema"))
target = config.pop("target")
masking_prob = config.pop("masking_prob")
return cls(schema, target, masking_prob, **config)
[docs]@tf.keras.utils.register_keras_serializable(package="merlin.models")
class SequenceMaskLast(SequenceTargetAsInput):
"""This block copies one of the sequence input features to be
the target feature. The last item of the target (and corresponding
sequences) is selected (masked) to be predicted
The input and target masks are returned by using Keras Masking
(._keras_mask), which is set by the compute_mask() method.
Parameters
----------
schema : Schema
The input schema, that will be used to discover the name
of the item id column
target : Union[str, Tags, ColumnSchema]
The sequential input column that will be used to extract the target
"""
[docs] def compute_mask(self, inputs, mask=None):
"""Selects (masks) the last position of the sequential targets
to be predicted.
This method is called by Keras after call()
and returns the targets mask that will
be assigned to the input tensors and targets, being accessible
by tensor._keras_mask
"""
item_id_seq = inputs[self.target_name]
self.target_mask = self._generate_target_mask(item_id_seq)
inputs_mask = dict()
for k, v in inputs.items():
if k in self.schema.column_names:
inputs_mask[k] = self.target_mask
else:
inputs_mask[k] = None
return (inputs_mask, self.target_mask)
def _generate_target_mask(self, ids_seq: tf.RaggedTensor) -> tf.RaggedTensor:
"""Returns a bool ragged tensor with the last positions of the sequence masked
Parameters
----------
ids_seq : tf.RaggedTensor
Sequence of ids, which are used to infer how many values
each sequence contains
Returns
-------
tf.RaggedTensor
Mask tensor, with True at the last positions
"""
row_lengths = ids_seq.row_lengths(1)
max_seq_length = tf.cast(tf.reduce_max(row_lengths), tf.int32)
assertion_min_seq_length = tf.Assert(tf.reduce_all(row_lengths > 1), [row_lengths])
with tf.control_dependencies([assertion_min_seq_length]):
padding_mask = tf.sequence_mask(row_lengths)
targets_mask = tf.ragged.boolean_mask(
tf.cast(tf.one_hot(row_lengths - 1, max_seq_length), tf.bool), padding_mask
)
return targets_mask
[docs] @classmethod
def from_config(cls, config):
config = tf_utils.maybe_deserialize_keras_objects(config, ["pre", "post", "aggregation"])
schema = schema_utils.tensorflow_metadata_json_to_schema(config.pop("schema"))
target = config.pop("target")
return cls(schema, target, **config)
@tf.keras.utils.register_keras_serializable(package="merlin.models")
class SequenceMaskLastInference(Block):
def call(self, inputs, training=False, testing=False):
self.inference_mode = not training and not testing
if self.inference_mode:
# Extending sequences in one position by copying the last embedding
repeat = inputs[:, -1:, :]
# repeat = tf.expand_dims(repeat, 1)
inputs = tf.concat([inputs, repeat], axis=1)
return inputs
def compute_mask(self, inputs, mask=None):
"""Selects (masks) the next position after the
last valid (non-padded) position of the sequential targets
to be predicted.
This method is called by Keras after call()
and returns the mask that is going to be assigned
to the input tensors, being accessible
by tensor._keras_mask
"""
targets_mask = None
if self.inference_mode:
if isinstance(inputs, tf.RaggedTensor):
row_lengths = inputs.row_lengths(1) + 1
max_seq_length = tf.cast(tf.reduce_max(row_lengths), tf.int32)
padding_mask = tf.sequence_mask(row_lengths)
targets_mask = tf.ragged.boolean_mask(
tf.cast(tf.one_hot(row_lengths - 1, max_seq_length), tf.bool), padding_mask
)
return targets_mask
[docs]@tf.keras.utils.register_keras_serializable(package="merlin.models")
class ReplaceMaskedEmbeddings(Block):
"""Takes a 3D input tensor (batch size x seq. length x embedding dim) and replaces
by a dummy trainable single embedding at the positions to be masked.
This is useful to be used when PredictMasked() transformation is used in
the fit()/eval() methods, which randomly selects some targets to be predicted and uses
Keras Masking to cascade the `_keras_mask`. By replacing input embeddings
at masked positions we avoid target leakage when training models with
Masked Language Modeling (BERT-like).
To support masked training approach in Transformer-based model,
SequenceMaskRandom and SequenceLastRandom implements `configure_for_train` method
that sets `ReplaceMaskedEmbeddings` as part of the `masking_pre` of
the transformer block.
"""
[docs] def __init__(self, **kwargs):
"""Initializes the block."""
super().__init__(**kwargs)
self.supports_masking = True
[docs] def build(self, input_shape):
"""Builds the block's internal variables.
This method creates a trainable embedding to replace masked interactions in the input.
Parameters
----------
input_shape : tf.TensorShape
Shape of the input tensor.
Returns
-------
None
Raises
------
ValueError
If the last dimension of the input shape is None.
"""
self.hidden_size = input_shape[-1]
if self.hidden_size is None:
raise ValueError("The last dim of inputs cannot be None")
# Create a trainable embedding to replace masked interactions
initializer = tf.random_normal_initializer(mean=0.0, stddev=0.001)
self.masked_embedding = tf.Variable(initializer(shape=[self.hidden_size], dtype=tf.float32))
return super().build(input_shape)
[docs] def call(
self,
inputs: Union[tf.Tensor, tf.RaggedTensor],
) -> Union[tf.Tensor, tf.RaggedTensor]:
"""If the sequence of input embeddings is masked (with `tensor._keras_mask` defined),
replaces the input embeddings for masked elements
Parameters
----------
inputs : Union[tf.Tensor, tf.RaggedTensor]
A tensor with sequences of vectors.
Needs to be 3D (batch_size, sequence_length, embeddings dim).
If inputs._keras_mask is defined uses it to infer the mask
Returns
-------
Union[tf.Tensor, tf.RaggedTensor]
returns a tensor with the masked inputs replaced by the dummy embedding
"""
outputs = inputs
if getattr(inputs, "_keras_mask", None) is not None:
# Replaces the embeddings at masked positions by a dummy trainable embedding
outputs = self._replace_masked_embeddings(inputs, inputs._keras_mask)
return outputs
def _replace_masked_embeddings(
self, inputs: Union[tf.Tensor, tf.RaggedTensor], mask: Union[tf.Tensor, tf.RaggedTensor]
) -> tf.RaggedTensor:
"""
Replaces in the inputs tensors the values masked as targets by a common trainable
embedding
"""
tf.Assert(
tf_utils.check_inputs_mask_compatible_shape(inputs, mask),
[
"The inputs and mask need to be compatible: have the same dtype "
"(tf.Tensor or tf.RaggedTensor) and the tf.rank(mask) == tf.rank(inputs)-1"
],
)
if isinstance(mask, tf.RaggedTensor):
mask = mask.with_row_splits_dtype(inputs.row_splits.dtype)
output = tf.where(
tf.cast(tf.expand_dims(mask, -1), tf.bool),
tf.cast(self.masked_embedding, dtype=inputs.dtype),
inputs,
)
return output
@tf.keras.utils.register_keras_serializable(package="merlin.models")
class ExtractMaskFromTargets(Block):
"""
Recovers the mask information for the inputs from the mask information
stored in the targets.
This block looks for the Keras mask (`._keras_mask`) in the following order:
1. Checks if the input tensor has a mask.
2. Checks if there is a single target and if it has a mask.
3. If there are multiple targets (dictionary), returns the mask of the target
that matches the first two dimensions of the input.
This is useful to use when the mask information for the inputs may be lost in
previous non-mask-aware Merlin blocks.
"""
def call(
self,
inputs: Union[tf.Tensor, tf.RaggedTensor],
targets: Optional[Union[tf.Tensor, tf.RaggedTensor, TabularData]] = None,
) -> Union[tf.Tensor, tf.RaggedTensor]:
mask = self._infer_mask_from_inputs_or_targets(inputs, targets)
inputs._keras_mask = mask
return inputs
def _infer_mask_from_inputs_or_targets(
self,
inputs: Union[tf.Tensor, tf.RaggedTensor],
targets: Optional[Union[tf.Tensor, tf.RaggedTensor]] = None,
):
mask = None
if getattr(inputs, "_keras_mask", None) is not None:
mask = inputs._keras_mask
elif targets is not None:
if isinstance(targets, dict):
if len(targets) == 1:
single_target = list(targets.values())[0]
if getattr(single_target, "_keras_mask", None) is not None:
mask = single_target._keras_mask
elif len(targets) > 1:
# If there is more than 1 target, checks if only one
# target matches the shape sequence of the inputs
for _, v in targets.items():
if getattr(
v, "_keras_mask", None
) is not None and tf_utils.check_inputs_mask_compatible_shape(
inputs, v._keras_mask
):
if mask is None:
mask = v._keras_mask
else:
raise ValueError(
"It is not possible to infer the mask "
"from the targets because there are more than "
"one target with the expected shape that matches "
"the inputs shape (batch_size x seq_length)"
)
elif getattr(targets, "_keras_mask", None) is not None:
mask = targets._keras_mask
return mask
@tf.keras.utils.register_keras_serializable(package="merlin.models")
class SequenceCausalLastInference(Block):
def call(self, inputs, training=False, testing=False):
self.inference_mode = not training and not testing
return inputs
def compute_mask(self, inputs, mask=None):
"""Selects (masks) the last non padded position of the
input sequence to be predicted.
This method is called by Keras after call()
and returns the mask that is going to be assigned
to the input tensors, being accessible
by tensor._keras_mask
"""
if self.inference_mode:
if isinstance(inputs, tf.RaggedTensor):
row_lengths = inputs.row_lengths(1)
max_seq_length = tf.cast(tf.reduce_max(row_lengths), tf.int32)
padding_mask = tf.sequence_mask(row_lengths)
mask = tf.ragged.boolean_mask(
tf.cast(tf.one_hot(row_lengths - 1, max_seq_length), tf.bool), padding_mask
)
return mask