def __init__()

in tensorflow_recommenders/experimental/models/ranking.py [0:0]

• 33 lines of code
• 5 McCabe index (conditional complexity)

  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")
          ]
      )