def __init__(
self,
embedding_layer: tf.keras.layers.Layer,
bottom_stack: Optional[tf.keras.layers.Layer] = None,
feature_interaction: Optional[tf.keras.layers.Layer] = None,
top_stack: Optional[tf.keras.layers.Layer] = None,
task: Optional[tasks.Task] = None) -> None:
"""Initializes the model.
Args:
embedding_layer: The embedding layer is applied to categorical features.
It expects a string-to-tensor (or SparseTensor/RaggedTensor) dict as
an input, and outputs a dictionary of string-to-tensor of feature_name,
embedded_value pairs.
{feature_name_i: tensor_i} -> {feature_name_i: emb(tensor_i)}.
bottom_stack: The `bottom_stack` layer is applied to dense features before
feature interaction. If None, an MLP with layer sizes [256, 64, 16] is
used. For DLRM model, the output of bottom_stack should be of shape
(batch_size, embedding dimension).
feature_interaction: Feature interaction layer is applied to the
`bottom_stack` output and sparse feature embeddings. If it is None,
DotInteraction layer is used.
top_stack: The `top_stack` layer is applied to the `feature_interaction`
output. The output of top_stack should be in the range [0, 1]. If it is
None, MLP with layer sizes [512, 256, 1] is used.
task: The task which the model should optimize for. Defaults to a
`tfrs.tasks.Ranking` task with a binary cross-entropy loss, suitable
for tasks like click prediction.
"""
super().__init__()
self._embedding_layer = embedding_layer
self._bottom_stack = bottom_stack if bottom_stack else layers.blocks.MLP(
units=[256, 64, 16], final_activation="relu")
self._top_stack = top_stack if top_stack else layers.blocks.MLP(
units=[512, 256, 1], final_activation="sigmoid")
self._feature_interaction = (feature_interaction if feature_interaction
else feature_interaction_lib.DotInteraction())
if task is not None:
self._task = task
else:
self._task = tasks.Ranking(
loss=tf.keras.losses.BinaryCrossentropy(
reduction=tf.keras.losses.Reduction.NONE
),
metrics=[
tf.keras.metrics.AUC(name="auc"),
tf.keras.metrics.BinaryAccuracy(name="accuracy"),
],
prediction_metrics=[
tf.keras.metrics.Mean("prediction_mean"),
],
label_metrics=[
tf.keras.metrics.Mean("label_mean")
]
)
def __init__(self, candidate_dataset):
super().__init__()
self.query_model = tf.keras.layers.Dense(16)
self.candidate_model = tf.keras.layers.Dense(16)
self.ctr_model = tf.keras.layers.Dense(1, activation="sigmoid")
self.retrieval_task = tasks.Retrieval(
metrics=metrics.FactorizedTopK(
candidates=candidate_dataset.map(self.candidate_model),
ks=[5]))
self.ctr_task = tasks.Ranking(
metrics=[tf.keras.metrics.AUC(name="ctr_auc")])
def __init__(self,
vocab_sizes: List[int],
embedding_dim: int = 16,
emb_optimizer: Optional[tf.keras.optimizers.Optimizer] = None,
bottom_stack: Optional[tf.keras.layers.Layer] = None,
feature_interaction: Optional[tf.keras.layers.Layer] = None,
top_stack: Optional[tf.keras.layers.Layer] = None,
task: Optional[tasks.Task] = None) -> None:
super().__init__()
emb_feature_config = _get_tpu_embedding_feature_config(
vocab_sizes=vocab_sizes, embedding_dim=embedding_dim)
if not emb_optimizer:
emb_optimizer = tf.keras.optimizers.Adam()
self._tpu_embeddings = embedding.TPUEmbedding(emb_feature_config,
emb_optimizer)
self._bottom_stack = bottom_stack if bottom_stack else MlpBlock(
units=[256, 64, embedding_dim], out_activation="relu")
self._top_stack = top_stack if top_stack else MlpBlock(
units=[512, 256, 1], out_activation="sigmoid")
self._feature_interaction = (feature_interaction if feature_interaction
else DotInteraction())
if task is not None:
self._task = task
else:
self._task = tasks.Ranking(
loss=tf.keras.losses.BinaryCrossentropy(
reduction=tf.keras.losses.Reduction.NONE),
metrics=[
tf.keras.metrics.AUC(name="auc"),
tf.keras.metrics.BinaryAccuracy(name="accuracy"),
],
prediction_metrics=[
tf.keras.metrics.Mean("prediction_mean"),
],
label_metrics=[tf.keras.metrics.Mean("label_mean")])
def __init__(
self,
embedding_layer: tf.keras.layers.Layer,
bottom_stack: Optional[tf.keras.layers.Layer] = None,
feature_interaction: Optional[tf.keras.layers.Layer] = None,
top_stack: Optional[tf.keras.layers.Layer] = None,
task: Optional[tasks.Task] = None) -> None:
super().__init__()
self._embedding_layer = embedding_layer
self._bottom_stack = bottom_stack if bottom_stack else MlpBlock(
units=[256, 64, 16], out_activation="relu")
self._top_stack = top_stack if top_stack else MlpBlock(
units=[512, 256, 1], out_activation="sigmoid")
self._feature_interaction = (feature_interaction if feature_interaction
else feature_interaction_lib.DotInteraction())
if task is not None:
self._task = task
else:
self._task = tasks.Ranking(
loss=tf.keras.losses.BinaryCrossentropy(
reduction=tf.keras.losses.Reduction.NONE
),
metrics=[
tf.keras.metrics.AUC(name="auc"),
tf.keras.metrics.BinaryAccuracy(name="accuracy"),
],
prediction_metrics=[
tf.keras.metrics.Mean("prediction_mean"),
],
label_metrics=[
tf.keras.metrics.Mean("label_mean")
]
)
def __init__(self):
super().__init__()
self._dense = tf.keras.layers.Dense(1)
self.task = tasks.Ranking(
loss=tf.keras.losses.BinaryCrossentropy(),
metrics=[tf.keras.metrics.BinaryAccuracy(name="accuracy")])
