rasa-pro 3.12.0.dev1__py3-none-any.whl

rasa/utils/tensorflow/layers.py

@@ -0,0 +1,1565 @@
+import logging
+from typing import List, Optional, Text, Tuple, Callable, Union, Any
+import tensorflow as tf
+
+# TODO: The following is not (yet) available via tf.keras
+from keras.src.utils.control_flow_util import smart_cond
+import tensorflow.keras.backend as K
+
+import rasa.utils.tensorflow.crf
+from rasa.utils.tensorflow.constants import (
+    SOFTMAX,
+    MARGIN,
+    COSINE,
+    INNER,
+    CROSS_ENTROPY,
+    LABEL,
+    LABEL_PAD_ID,
+)
+from rasa.core.constants import DIALOGUE
+from rasa.shared.nlu.constants import FEATURE_TYPE_SENTENCE, FEATURE_TYPE_SEQUENCE
+from rasa.shared.nlu.constants import TEXT, INTENT, ACTION_NAME, ACTION_TEXT
+
+from rasa.utils.tensorflow.metrics import F1Score
+from rasa.utils.tensorflow.exceptions import TFLayerConfigException
+import rasa.utils.tensorflow.layers_utils as layers_utils
+from rasa.utils.tensorflow.crf import crf_log_likelihood
+
+logger = logging.getLogger(__name__)
+
+
+POSSIBLE_ATTRIBUTES = [
+    TEXT,
+    INTENT,
+    LABEL,
+    DIALOGUE,
+    ACTION_NAME,
+    ACTION_TEXT,
+    f"{LABEL}_{ACTION_NAME}",
+    f"{LABEL}_{ACTION_TEXT}",
+]
+
+
+class SparseDropout(tf.keras.layers.Dropout):
+    """Applies Dropout to the input.
+
+    Dropout consists in randomly setting
+    a fraction `rate` of input units to 0 at each update during training time,
+    which helps prevent overfitting.
+
+    Arguments:
+        rate: Fraction of the input units to drop (between 0 and 1).
+    """
+
+    def call(
+        self, inputs: tf.SparseTensor, training: Optional[Union[tf.Tensor, bool]] = None
+    ) -> tf.SparseTensor:
+        """Apply dropout to sparse inputs.
+
+        Arguments:
+            inputs: Input sparse tensor (of any rank).
+            training: Indicates whether the layer should behave in
+                training mode (adding dropout) or in inference mode (doing nothing).
+
+        Returns:
+            Output of dropout layer.
+
+        Raises:
+            A ValueError if inputs is not a sparse tensor
+        """
+        if not isinstance(inputs, tf.SparseTensor):
+            raise ValueError("Input tensor should be sparse.")
+
+        if training is None:
+            training = K.learning_phase()
+
+        def dropped_inputs() -> tf.SparseTensor:
+            to_retain_prob = tf.random.uniform(
+                tf.shape(inputs.values), 0, 1, inputs.values.dtype
+            )
+            to_retain = tf.greater_equal(to_retain_prob, self.rate)
+            return tf.sparse.retain(inputs, to_retain)
+
+        outputs = smart_cond(training, dropped_inputs, lambda: tf.identity(inputs))
+        # need to explicitly recreate sparse tensor, because otherwise the shape
+        # information will be lost after `retain`
+        # noinspection PyProtectedMember
+        return tf.SparseTensor(outputs.indices, outputs.values, inputs._dense_shape)
+
+
+class DenseForSparse(tf.keras.layers.Dense):
+    """Dense layer for sparse input tensor.
+
+    Just your regular densely-connected NN layer but for sparse tensors.
+
+    `Dense` implements the operation:
+    `output = activation(dot(input, kernel) + bias)`
+    where `activation` is the element-wise activation function
+    passed as the `activation` argument, `kernel` is a weights matrix
+    created by the layer, and `bias` is a bias vector created by the layer
+    (only applicable if `use_bias` is `True`).
+
+    Note: If the input to the layer has a rank greater than 2, then
+    it is flattened prior to the initial dot product with `kernel`.
+
+    Arguments:
+        units: Positive integer, dimensionality of the output space.
+            activation: Activation function to use.
+            If you don't specify anything, no activation is applied
+            (ie. "linear" activation: `a(x) = x`).
+        use_bias: Indicates whether the layer uses a bias vector.
+        kernel_initializer: Initializer for the `kernel` weights matrix.
+        bias_initializer: Initializer for the bias vector.
+        reg_lambda: regularization factor
+        bias_regularizer: Regularizer function applied to the bias vector.
+        activity_regularizer: Regularizer function applied to
+            the output of the layer (its "activation")..
+        kernel_constraint: Constraint function applied to
+            the `kernel` weights matrix.
+        bias_constraint: Constraint function applied to the bias vector.
+
+    Input shape:
+        N-D tensor with shape: `(batch_size, ..., input_dim)`.
+        The most common situation would be
+        a 2D input with shape `(batch_size, input_dim)`.
+
+    Output shape:
+        N-D tensor with shape: `(batch_size, ..., units)`.
+        For instance, for a 2D input with shape `(batch_size, input_dim)`,
+        the output would have shape `(batch_size, units)`.
+    """
+
+    def __init__(self, reg_lambda: float = 0, **kwargs: Any) -> None:
+        if reg_lambda > 0:
+            regularizer = tf.keras.regularizers.l2(reg_lambda)
+        else:
+            regularizer = None
+
+        super().__init__(kernel_regularizer=regularizer, **kwargs)
+
+    def get_units(self) -> int:
+        """Returns number of output units."""
+        return self.units
+
+    def get_kernel(self) -> tf.Tensor:
+        """Returns kernel tensor."""
+        return self.kernel
+
+    def get_bias(self) -> Union[tf.Tensor, None]:
+        """Returns bias tensor."""
+        if self.use_bias:
+            return self.bias
+        return None
+
+    def get_feature_type(self) -> Union[Text, None]:
+        """Returns a feature type of the data that's fed to the layer.
+
+        In order to correctly return a feature type, the function heavily relies
+        on the name of `DenseForSparse` layer to contain the feature type.
+        Acceptable values of feature types are `FEATURE_TYPE_SENTENCE`
+        and `FEATURE_TYPE_SEQUENCE`.
+
+        Returns:
+            feature type of dense layer.
+        """
+        for feature_type in [FEATURE_TYPE_SENTENCE, FEATURE_TYPE_SEQUENCE]:
+            if feature_type in self.name:
+                return feature_type
+        return None
+
+    def get_attribute(self) -> Union[Text, None]:
+        """Returns the attribute for which this layer was constructed.
+
+        For example: TEXT, LABEL, etc.
+
+        In order to correctly return an attribute, the function heavily relies
+        on the name of `DenseForSparse` layer being in the following format:
+        f"sparse_to_dense.{attribute}_{feature_type}".
+
+        Returns:
+            attribute of the layer.
+        """
+        metadata = self.name.split(".")
+        if len(metadata) > 1:
+            attribute_splits = metadata[1].split("_")[:-1]
+            attribute = "_".join(attribute_splits)
+            if attribute in POSSIBLE_ATTRIBUTES:
+                return attribute
+        return None
+
+    def call(self, inputs: tf.SparseTensor) -> tf.Tensor:
+        """Apply dense layer to sparse inputs.
+
+        Arguments:
+            inputs: Input sparse tensor (of any rank).
+
+        Returns:
+            Output of dense layer.
+
+        Raises:
+            A ValueError if inputs is not a sparse tensor
+        """
+        if not isinstance(inputs, tf.SparseTensor):
+            raise ValueError("Input tensor should be sparse.")
+
+        # outputs will be 2D
+        outputs = tf.sparse.sparse_dense_matmul(
+            tf.sparse.reshape(inputs, [-1, tf.shape(inputs)[-1]]), self.kernel
+        )
+
+        if len(inputs.shape) == 3:
+            # reshape back
+            outputs = tf.reshape(
+                outputs, (tf.shape(inputs)[0], tf.shape(inputs)[1], self.units)
+            )
+
+        if self.use_bias:
+            outputs = tf.nn.bias_add(outputs, self.bias)
+        if self.activation is not None:
+            return self.activation(outputs)
+        return outputs
+
+
+class RandomlyConnectedDense(tf.keras.layers.Dense):
+    """Layer with dense ouputs that are connected to a random subset of inputs.
+
+    `RandomlyConnectedDense` implements the operation:
+    `output = activation(dot(input, kernel) + bias)`
+    where `activation` is the element-wise activation function
+    passed as the `activation` argument, `kernel` is a weights matrix
+    created by the layer, and `bias` is a bias vector created by the layer
+    (only applicable if `use_bias` is `True`).
+    It creates `kernel_mask` to set a fraction of the `kernel` weights to zero.
+
+    Note: If the input to the layer has a rank greater than 2, then
+    it is flattened prior to the initial dot product with `kernel`.
+
+    The output is guaranteed to be dense (each output is connected to at least one
+    input), and no input is disconnected (each input is connected to at least one
+    output).
+
+    At `density = 0.0` the number of trainable weights is `max(input_size, units)`. At
+    `density = 1.0` this layer is equivalent to `tf.keras.layers.Dense`.
+
+    Input shape:
+        N-D tensor with shape: `(batch_size, ..., input_dim)`.
+        The most common situation would be
+        a 2D input with shape `(batch_size, input_dim)`.
+
+    Output shape:
+        N-D tensor with shape: `(batch_size, ..., units)`.
+        For instance, for a 2D input with shape `(batch_size, input_dim)`,
+        the output would have shape `(batch_size, units)`.
+    """
+
+    def __init__(self, density: float = 0.2, **kwargs: Any) -> None:
+        """Declares instance variables with default values.
+
+        Args:
+            density: Approximate fraction of trainable weights (between 0 and 1).
+            units: Positive integer, dimensionality of the output space.
+            activation: Activation function to use.
+                If you don't specify anything, no activation is applied
+                (ie. "linear" activation: `a(x) = x`).
+            use_bias: Indicates whether the layer uses a bias vector.
+            kernel_initializer: Initializer for the `kernel` weights matrix.
+            bias_initializer: Initializer for the bias vector.
+            kernel_regularizer: Regularizer function applied to
+                the `kernel` weights matrix.
+            bias_regularizer: Regularizer function applied to the bias vector.
+            activity_regularizer: Regularizer function applied to
+                the output of the layer (its "activation")..
+            kernel_constraint: Constraint function applied to
+                the `kernel` weights matrix.
+            bias_constraint: Constraint function applied to the bias vector.
+        """
+        super().__init__(**kwargs)
+
+        if density < 0.0 or density > 1.0:
+            raise TFLayerConfigException("Layer density must be in [0, 1].")
+
+        self.density = density
+
+    def build(self, input_shape: tf.TensorShape) -> None:
+        """Prepares the kernel mask.
+
+        Args:
+            input_shape: Shape of the inputs to this layer
+        """
+        super().build(input_shape)
+
+        if self.density == 1.0:
+            self.kernel_mask = None
+            return
+
+        # Construct mask with given density and guarantee that every output is
+        # connected to at least one input
+        kernel_mask = self._minimal_mask() + self._random_mask()
+
+        # We might accidently have added a random connection on top of
+        # a fixed connection
+        kernel_mask = tf.clip_by_value(kernel_mask, 0, 1)
+
+        self.kernel_mask = tf.Variable(
+            initial_value=kernel_mask, trainable=False, name="kernel_mask"
+        )
+
+    def _random_mask(self) -> tf.Tensor:
+        """Creates a random matrix with `num_ones` 1s and 0s otherwise.
+
+        Returns:
+            A random mask matrix
+        """
+        mask = tf.random.uniform(tf.shape(self.kernel), 0, 1)
+        mask = tf.cast(tf.math.less(mask, self.density), self.kernel.dtype)
+        return mask
+
+    def _minimal_mask(self) -> tf.Tensor:
+        """Creates a matrix with a minimal number of 1s to connect everythinig.
+
+        If num_rows == num_cols, this creates the identity matrix.
+        If num_rows > num_cols, this creates
+            1 0 0 0
+            0 1 0 0
+            0 0 1 0
+            0 0 0 1
+            1 0 0 0
+            0 1 0 0
+            0 0 1 0
+            . . . .
+            . . . .
+            . . . .
+        If num_rows < num_cols, this creates
+            1 0 0 1 0 0 1 ...
+            0 1 0 0 1 0 0 ...
+            0 0 1 0 0 1 0 ...
+
+        Returns:
+            A tiled and croped identity matrix.
+        """
+        kernel_shape = tf.shape(self.kernel)
+        num_rows = kernel_shape[0]
+        num_cols = kernel_shape[1]
+        short_dimension = tf.minimum(num_rows, num_cols)
+
+        mask = tf.tile(
+            tf.eye(short_dimension, dtype=self.kernel.dtype),
+            [
+                tf.math.ceil(num_rows / short_dimension),
+                tf.math.ceil(num_cols / short_dimension),
+            ],
+        )[:num_rows, :num_cols]
+
+        return mask
+
+    def call(self, inputs: tf.Tensor) -> tf.Tensor:
+        """Processes the given inputs.
+
+        Args:
+            inputs: What goes into this layer
+
+        Returns:
+            The processed inputs.
+        """
+        if self.density < 1.0:
+            # Set fraction of the `kernel` weights to zero according to precomputed mask
+            self.kernel.assign(self.kernel * self.kernel_mask)
+        return super().call(inputs)
+
+
+class Ffnn(tf.keras.layers.Layer):
+    """Feed-forward network layer.
+
+    Arguments:
+        layer_sizes: List of integers with dimensionality of the layers.
+        dropout_rate: Fraction of the input units to drop (between 0 and 1).
+        reg_lambda: regularization factor.
+        density: Approximate fraction of trainable weights (between 0 and 1).
+        layer_name_suffix: Text added to the name of the layers.
+
+    Input shape:
+        N-D tensor with shape: `(batch_size, ..., input_dim)`.
+        The most common situation would be
+        a 2D input with shape `(batch_size, input_dim)`.
+
+    Output shape:
+        N-D tensor with shape: `(batch_size, ..., layer_sizes[-1])`.
+        For instance, for a 2D input with shape `(batch_size, input_dim)`,
+        the output would have shape `(batch_size, layer_sizes[-1])`.
+    """
+
+    def __init__(
+        self,
+        layer_sizes: List[int],
+        dropout_rate: float,
+        reg_lambda: float,
+        density: float,
+        layer_name_suffix: Text,
+    ) -> None:
+        super().__init__(name=f"ffnn_{layer_name_suffix}")
+
+        l2_regularizer = tf.keras.regularizers.l2(reg_lambda)
+        self._ffn_layers = []
+        for i, layer_size in enumerate(layer_sizes):
+            self._ffn_layers.append(
+                RandomlyConnectedDense(
+                    units=layer_size,
+                    density=density,
+                    activation=tf.nn.gelu,
+                    kernel_regularizer=l2_regularizer,
+                    name=f"hidden_layer_{layer_name_suffix}_{i}",
+                )
+            )
+            self._ffn_layers.append(tf.keras.layers.Dropout(dropout_rate))
+
+    def call(
+        self, x: tf.Tensor, training: Optional[Union[tf.Tensor, bool]] = None
+    ) -> tf.Tensor:
+        """Apply feed-forward network layer."""
+        for layer in self._ffn_layers:
+            x = layer(x, training=training)
+
+        return x
+
+
+class Embed(tf.keras.layers.Layer):
+    """Dense embedding layer.
+
+    Input shape:
+        N-D tensor with shape: `(batch_size, ..., input_dim)`.
+        The most common situation would be
+        a 2D input with shape `(batch_size, input_dim)`.
+
+    Output shape:
+        N-D tensor with shape: `(batch_size, ..., embed_dim)`.
+        For instance, for a 2D input with shape `(batch_size, input_dim)`,
+        the output would have shape `(batch_size, embed_dim)`.
+    """
+
+    def __init__(
+        self, embed_dim: int, reg_lambda: float, layer_name_suffix: Text
+    ) -> None:
+        """Initialize layer.
+
+        Args:
+            embed_dim: Dimensionality of the output space.
+            reg_lambda: Regularization factor.
+            layer_name_suffix: Text added to the name of the layers.
+        """
+        super().__init__(name=f"embed_{layer_name_suffix}")
+
+        regularizer = tf.keras.regularizers.l2(reg_lambda)
+        self._dense = tf.keras.layers.Dense(
+            units=embed_dim,
+            activation=None,
+            kernel_regularizer=regularizer,
+            name=f"embed_layer_{layer_name_suffix}",
+        )
+
+    # noinspection PyMethodOverriding
+    def call(self, x: tf.Tensor) -> tf.Tensor:
+        """Apply dense layer."""
+        x = self._dense(x)
+        return x
+
+
+class InputMask(tf.keras.layers.Layer):
+    """The layer that masks 15% of the input.
+
+    Input shape:
+        N-D tensor with shape: `(batch_size, ..., input_dim)`.
+        The most common situation would be
+        a 2D input with shape `(batch_size, input_dim)`.
+
+    Output shape:
+        N-D tensor with shape: `(batch_size, ..., input_dim)`.
+        For instance, for a 2D input with shape `(batch_size, input_dim)`,
+        the output would have shape `(batch_size, input_dim)`.
+    """
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+
+        self._masking_prob = 0.85
+        self._mask_vector_prob = 0.7
+        self._random_vector_prob = 0.1
+
+    def build(self, input_shape: tf.TensorShape) -> None:
+        self.mask_vector = self.add_weight(
+            shape=(1, 1, input_shape[-1]), name="mask_vector"
+        )
+        self.built = True
+
+    # noinspection PyMethodOverriding
+    def call(
+        self,
+        x: tf.Tensor,
+        mask: tf.Tensor,
+        training: Optional[Union[tf.Tensor, bool]] = None,
+    ) -> Tuple[tf.Tensor, tf.Tensor]:
+        """Randomly mask input sequences.
+
+        Arguments:
+            x: Input sequence tensor of rank 3.
+            mask: A tensor representing sequence mask,
+                contains `1` for inputs and `0` for padding.
+            training: Indicates whether the layer should run in
+                training mode (mask inputs) or in inference mode (doing nothing).
+
+        Returns:
+            A tuple of masked inputs and boolean mask.
+        """
+        if training is None:
+            training = K.learning_phase()
+
+        lm_mask_prob = tf.random.uniform(tf.shape(mask), 0, 1, mask.dtype) * mask
+        lm_mask_bool = tf.greater_equal(lm_mask_prob, self._masking_prob)
+
+        def x_masked() -> tf.Tensor:
+            x_random_pad = tf.random.uniform(
+                tf.shape(x), tf.reduce_min(x), tf.reduce_max(x), x.dtype
+            ) * (1 - mask)
+            # shuffle over batch dim
+            x_shuffle = tf.random.shuffle(x * mask + x_random_pad)
+
+            # shuffle over sequence dim
+            x_shuffle = tf.transpose(x_shuffle, [1, 0, 2])
+            x_shuffle = tf.random.shuffle(x_shuffle)
+            x_shuffle = tf.transpose(x_shuffle, [1, 0, 2])
+
+            # shuffle doesn't support backprop
+            x_shuffle = tf.stop_gradient(x_shuffle)
+
+            mask_vector = tf.tile(self.mask_vector, (tf.shape(x)[0], tf.shape(x)[1], 1))
+
+            other_prob = tf.random.uniform(tf.shape(mask), 0, 1, mask.dtype)
+            other_prob = tf.tile(other_prob, (1, 1, x.shape[-1]))
+            x_other = tf.where(
+                other_prob < self._mask_vector_prob,
+                mask_vector,
+                tf.where(
+                    other_prob < self._mask_vector_prob + self._random_vector_prob,
+                    x_shuffle,
+                    x,
+                ),
+            )
+
+            return tf.where(tf.tile(lm_mask_bool, (1, 1, x.shape[-1])), x_other, x)
+
+        return (smart_cond(training, x_masked, lambda: tf.identity(x)), lm_mask_bool)
+
+
+def _scale_loss(log_likelihood: tf.Tensor) -> tf.Tensor:
+    """Creates scaling loss coefficient depending on the prediction probability.
+
+    Arguments:
+        log_likelihood: a tensor, log-likelihood of prediction
+
+    Returns:
+        Scaling tensor.
+    """
+    p = tf.math.exp(log_likelihood)
+    # only scale loss if some examples are already learned
+    return tf.cond(
+        tf.reduce_max(p) > 0.5,
+        lambda: tf.stop_gradient(tf.pow((1 - p) / 0.5, 4)),
+        lambda: tf.ones_like(p),
+    )
+
+
+class CRF(tf.keras.layers.Layer):
+    """CRF layer.
+
+    Arguments:
+        num_tags: Positive integer, number of tags.
+        reg_lambda: regularization factor.
+        name: Optional name of the layer.
+    """
+
+    def __init__(
+        self,
+        num_tags: int,
+        reg_lambda: float,
+        scale_loss: bool,
+        name: Optional[Text] = None,
+    ) -> None:
+        super().__init__(name=name)
+        self.num_tags = num_tags
+        self.scale_loss = scale_loss
+        self.transition_regularizer = tf.keras.regularizers.l2(reg_lambda)
+        self.f1_score_metric = F1Score(
+            num_classes=num_tags - 1,  # `0` prediction is not a prediction
+            average="micro",
+        )
+
+    def build(self, input_shape: tf.TensorShape) -> None:
+        # the weights should be created in `build` to apply random_seed
+        self.transition_params = self.add_weight(
+            shape=(self.num_tags, self.num_tags),
+            regularizer=self.transition_regularizer,
+            name="transitions",
+        )
+        self.built = True
+
+    # noinspection PyMethodOverriding
+    def call(
+        self, logits: tf.Tensor, sequence_lengths: tf.Tensor
+    ) -> Tuple[tf.Tensor, tf.Tensor]:
+        """Decodes the highest scoring sequence of tags.
+
+        Arguments:
+            logits: A [batch_size, max_seq_len, num_tags] tensor of
+                unary potentials.
+            sequence_lengths: A [batch_size] vector of true sequence lengths.
+
+        Returns:
+            A [batch_size, max_seq_len] matrix, with dtype `tf.int32`.
+            Contains the highest scoring tag indices.
+            A [batch_size, max_seq_len] matrix, with dtype `tf.float32`.
+            Contains the confidence values of the highest scoring tag indices.
+        """
+        predicted_ids, scores, _ = rasa.utils.tensorflow.crf.crf_decode(
+            logits, self.transition_params, sequence_lengths
+        )
+        # set prediction index for padding to `0`
+        mask = tf.sequence_mask(
+            sequence_lengths,
+            maxlen=tf.shape(predicted_ids)[1],
+            dtype=predicted_ids.dtype,
+        )
+
+        confidence_values = scores * tf.cast(mask, tf.float32)
+        predicted_ids = predicted_ids * mask
+
+        return predicted_ids, confidence_values
+
+    def loss(
+        self, logits: tf.Tensor, tag_indices: tf.Tensor, sequence_lengths: tf.Tensor
+    ) -> tf.Tensor:
+        """Computes the log-likelihood of tag sequences in a CRF.
+
+        Arguments:
+            logits: A [batch_size, max_seq_len, num_tags] tensor of unary potentials
+                to use as input to the CRF layer.
+            tag_indices: A [batch_size, max_seq_len] matrix of tag indices for which
+                we compute the log-likelihood.
+            sequence_lengths: A [batch_size] vector of true sequence lengths.
+
+        Returns:
+            Negative mean log-likelihood of all examples,
+            given the sequence of tag indices.
+        """
+        log_likelihood, _ = crf_log_likelihood(
+            logits, tag_indices, sequence_lengths, self.transition_params
+        )
+        loss = -log_likelihood
+        if self.scale_loss:
+            loss *= _scale_loss(log_likelihood)
+
+        return tf.reduce_mean(loss)
+
+    def f1_score(
+        self, tag_ids: tf.Tensor, pred_ids: tf.Tensor, mask: tf.Tensor
+    ) -> tf.Tensor:
+        """Calculates f1 score for train predictions."""
+        mask_bool = tf.cast(mask[:, :, 0], tf.bool)
+
+        # pick only non padding values and flatten sequences
+        tag_ids_flat = tf.boolean_mask(tag_ids, mask_bool)
+        pred_ids_flat = tf.boolean_mask(pred_ids, mask_bool)
+
+        # set `0` prediction to not a prediction
+        num_tags = self.num_tags - 1
+
+        tag_ids_flat_one_hot = tf.one_hot(tag_ids_flat - 1, num_tags)
+        pred_ids_flat_one_hot = tf.one_hot(pred_ids_flat - 1, num_tags)
+
+        return self.f1_score_metric(tag_ids_flat_one_hot, pred_ids_flat_one_hot)
+
+
+class DotProductLoss(tf.keras.layers.Layer):
+    """Abstract dot-product loss layer class.
+
+    Idea based on StarSpace paper: http://arxiv.org/abs/1709.03856
+
+    Implements similarity methods
+    * `sim` (computes a similarity between vectors)
+    * `get_similarities_and_confidences_from_embeddings` (calls `sim` and also computes
+        confidence values)
+
+    Specific loss functions (single- or multi-label) must be implemented in child
+    classes.
+    """
+
+    def __init__(
+        self,
+        num_candidates: int,
+        scale_loss: bool = False,
+        constrain_similarities: bool = True,
+        model_confidence: Text = SOFTMAX,
+        similarity_type: Text = INNER,
+        name: Optional[Text] = None,
+        **kwargs: Any,
+    ):
+        """Declares instance variables with default values.
+
+        Args:
+            num_candidates: Number of labels besides the positive one. Depending on
+                whether single- or multi-label loss is implemented (done in
+                sub-classes), these can be all negative example labels, or a mixture of
+                negative and further positive labels, respectively.
+            scale_loss: Boolean, if `True` scale loss inverse proportionally to
+                the confidence of the correct prediction.
+            constrain_similarities: Boolean, if `True` applies sigmoid on all
+                similarity terms and adds to the loss function to
+                ensure that similarity values are approximately bounded.
+                Used inside _loss_cross_entropy() only.
+            model_confidence: Normalization of confidence values during inference.
+                Currently, the only possible value is `SOFTMAX`.
+            similarity_type: Similarity measure to use, either `cosine` or `inner`.
+            name: Optional name of the layer.
+
+        Raises:
+            TFLayerConfigException: When `similarity_type` is not one of `COSINE` or
+                `INNER`.
+        """
+        super().__init__(name=name)
+        self.num_neg = num_candidates
+        self.scale_loss = scale_loss
+        self.constrain_similarities = constrain_similarities
+        self.model_confidence = model_confidence
+        self.similarity_type = similarity_type
+        if self.similarity_type not in {COSINE, INNER}:
+            raise TFLayerConfigException(
+                f"Unsupported similarity type '{self.similarity_type}', "
+                f"should be '{COSINE}' or '{INNER}'."
+            )
+
+    def sim(
+        self, a: tf.Tensor, b: tf.Tensor, mask: Optional[tf.Tensor] = None
+    ) -> tf.Tensor:
+        """Calculates similarity between `a` and `b`.
+
+        Operates on the last dimension. When `a` and `b` are vectors, then `sim`
+        computes either the dot-product, or the cosine of the angle between `a` and `b`,
+        depending on `self.similarity_type`.
+        Specifically, when the similarity type is `INNER`, then we compute the scalar
+        product `a . b`. When the similarity type is `COSINE`, we compute
+        `a . b / (|a| |b|)`, i.e. the cosine of the angle between `a` and `b`.
+
+        Args:
+            a: Any float tensor
+            b: Any tensor of the same shape and type as `a`
+            mask: Mask (should contain 1s for inputs and 0s for padding). Note, that
+                `len(mask.shape) == len(a.shape) - 1` should hold.
+
+        Returns:
+            Similarities between vectors in `a` and `b`.
+        """
+        if self.similarity_type == COSINE:
+            a = tf.nn.l2_normalize(a, axis=-1)
+            b = tf.nn.l2_normalize(b, axis=-1)
+        sim = tf.reduce_sum(a * b, axis=-1)
+        if mask is not None:
+            sim *= tf.expand_dims(mask, 2)
+
+        return sim
+
+    def get_similarities_and_confidences_from_embeddings(
+        self,
+        input_embeddings: tf.Tensor,
+        label_embeddings: tf.Tensor,
+        mask: Optional[tf.Tensor] = None,
+    ) -> Tuple[tf.Tensor, tf.Tensor]:
+        """Computes similary between input and label embeddings and model's confidence.
+
+        First compute the similarity from embeddings and then apply an activation
+        function if needed to get the confidence.
+
+        Args:
+            input_embeddings: Embeddings of input.
+            label_embeddings: Embeddings of labels.
+            mask: Mask (should contain 1s for inputs and 0s for padding). Note, that
+                `len(mask.shape) == len(a.shape) - 1` should hold.
+
+        Returns:
+            similarity between input and label embeddings and model's prediction
+            confidence for each label.
+        """
+        similarities = self.sim(input_embeddings, label_embeddings, mask)
+        confidences = similarities
+        if self.model_confidence == SOFTMAX:
+            confidences = tf.nn.softmax(similarities)
+        return similarities, confidences
+
+    def call(self, *args: Any, **kwargs: Any) -> Tuple[tf.Tensor, tf.Tensor]:
+        """Layer's logic - to be implemented in child class."""
+        raise NotImplementedError
+
+    def apply_mask_and_scaling(
+        self, loss: tf.Tensor, mask: Optional[tf.Tensor]
+    ) -> tf.Tensor:
+        """Scales the loss and applies the mask if necessary.
+
+        Args:
+            loss: The loss tensor
+            mask: (Optional) A mask to multiply with the loss
+
+        Returns:
+            The scaled loss, potentially averaged over the sequence
+            dimension.
+        """
+        if self.scale_loss:
+            # in case of cross entropy log_likelihood = -loss
+            loss *= _scale_loss(-loss)
+
+        if mask is not None:
+            loss *= mask
+
+        if len(loss.shape) == 2:
+            # average over the sequence
+            if mask is not None:
+                loss = tf.reduce_sum(loss, axis=-1) / tf.reduce_sum(mask, axis=-1)
+            else:
+                loss = tf.reduce_mean(loss, axis=-1)
+
+        return loss
+
+
+class SingleLabelDotProductLoss(DotProductLoss):
+    """Single-label dot-product loss layer.
+
+    This loss layer assumes that only one output (label) is correct for any given input.
+    """
+
+    def __init__(
+        self,
+        num_candidates: int,
+        scale_loss: bool = False,
+        constrain_similarities: bool = True,
+        model_confidence: Text = SOFTMAX,
+        similarity_type: Text = INNER,
+        name: Optional[Text] = None,
+        loss_type: Text = CROSS_ENTROPY,
+        mu_pos: float = 0.8,
+        mu_neg: float = -0.2,
+        use_max_sim_neg: bool = True,
+        neg_lambda: float = 0.5,
+        same_sampling: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        """Declares instance variables with default values.
+
+        Args:
+            num_candidates: Positive integer, the number of incorrect labels;
+                the algorithm will minimize their similarity to the input.
+            loss_type: The type of the loss function, either `cross_entropy` or
+                `margin`.
+            mu_pos: Indicates how similar the algorithm should
+                try to make embedding vectors for correct labels;
+                should be 0.0 < ... < 1.0 for `cosine` similarity type.
+            mu_neg: Maximum negative similarity for incorrect labels,
+                should be -1.0 < ... < 1.0 for `cosine` similarity type.
+            use_max_sim_neg: If `True` the algorithm only minimizes
+                maximum similarity over incorrect intent labels,
+                used only if `loss_type` is set to `margin`.
+            neg_lambda: The scale of how important it is to minimize
+                the maximum similarity between embeddings of different labels,
+                used only if `loss_type` is set to `margin`.
+            scale_loss: If `True` scale loss inverse proportionally to
+                the confidence of the correct prediction.
+            similarity_type: Similarity measure to use, either `cosine` or `inner`.
+            name: Optional name of the layer.
+            same_sampling: If `True` sample same negative labels
+                for the whole batch.
+            constrain_similarities: If `True` and loss_type is `cross_entropy`, a
+                sigmoid loss term is added to the total loss to ensure that similarity
+                values are approximately bounded.
+            model_confidence: Normalization of confidence values during inference.
+                Currently, the only possible value is `SOFTMAX`.
+        """
+        super().__init__(
+            num_candidates,
+            scale_loss=scale_loss,
+            constrain_similarities=constrain_similarities,
+            model_confidence=model_confidence,
+            similarity_type=similarity_type,
+            name=name,
+        )
+        self.loss_type = loss_type
+        self.mu_pos = mu_pos
+        self.mu_neg = mu_neg
+        self.use_max_sim_neg = use_max_sim_neg
+        self.neg_lambda = neg_lambda
+        self.same_sampling = same_sampling
+
+    def _get_bad_mask(
+        self, labels: tf.Tensor, target_labels: tf.Tensor, idxs: tf.Tensor
+    ) -> tf.Tensor:
+        """Calculate bad mask for given indices.
+
+        Checks that input features are different for positive negative samples.
+        """
+        pos_labels = tf.expand_dims(target_labels, axis=-2)
+        neg_labels = layers_utils.get_candidate_values(labels, idxs)
+
+        return tf.cast(
+            tf.reduce_all(tf.equal(neg_labels, pos_labels), axis=-1), pos_labels.dtype
+        )
+
+    def _get_negs(
+        self, embeds: tf.Tensor, labels: tf.Tensor, target_labels: tf.Tensor
+    ) -> Tuple[tf.Tensor, tf.Tensor]:
+        """Gets negative examples from given tensor."""
+        embeds_flat = layers_utils.batch_flatten(embeds)
+        labels_flat = layers_utils.batch_flatten(labels)
+        target_labels_flat = layers_utils.batch_flatten(target_labels)
+
+        total_candidates = tf.shape(embeds_flat)[0]
+        target_size = tf.shape(target_labels_flat)[0]
+
+        neg_ids = layers_utils.random_indices(
+            target_size, self.num_neg, total_candidates
+        )
+
+        neg_embeds = layers_utils.get_candidate_values(embeds_flat, neg_ids)
+        bad_negs = self._get_bad_mask(labels_flat, target_labels_flat, neg_ids)
+
+        # check if inputs have sequence dimension
+        if len(target_labels.shape) == 3:
+            # tensors were flattened for sampling, reshape back
+            # add sequence dimension if it was present in the inputs
+            target_shape = tf.shape(target_labels)
+            neg_embeds = tf.reshape(
+                neg_embeds, (target_shape[0], target_shape[1], -1, embeds.shape[-1])
+            )
+            bad_negs = tf.reshape(bad_negs, (target_shape[0], target_shape[1], -1))
+
+        return neg_embeds, bad_negs
+
+    def _sample_negatives(
+        self,
+        inputs_embed: tf.Tensor,
+        labels_embed: tf.Tensor,
+        labels: tf.Tensor,
+        all_labels_embed: tf.Tensor,
+        all_labels: tf.Tensor,
+    ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor]:
+        """Sample negative examples."""
+        pos_inputs_embed = tf.expand_dims(inputs_embed, axis=-2)
+        pos_labels_embed = tf.expand_dims(labels_embed, axis=-2)
+
+        # sample negative inputs
+        neg_inputs_embed, inputs_bad_negs = self._get_negs(inputs_embed, labels, labels)
+        # sample negative labels
+        neg_labels_embed, labels_bad_negs = self._get_negs(
+            all_labels_embed, all_labels, labels
+        )
+        return (
+            pos_inputs_embed,
+            pos_labels_embed,
+            neg_inputs_embed,
+            neg_labels_embed,
+            inputs_bad_negs,
+            labels_bad_negs,
+        )
+
+    def _train_sim(
+        self,
+        pos_inputs_embed: tf.Tensor,
+        pos_labels_embed: tf.Tensor,
+        neg_inputs_embed: tf.Tensor,
+        neg_labels_embed: tf.Tensor,
+        inputs_bad_negs: tf.Tensor,
+        labels_bad_negs: tf.Tensor,
+        mask: Optional[tf.Tensor],
+    ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor]:
+        """Define similarity."""
+        # calculate similarity with several
+        # embedded actions for the loss
+        neg_inf = tf.constant(-1e9)
+
+        sim_pos = self.sim(pos_inputs_embed, pos_labels_embed, mask)
+        sim_neg_il = (
+            self.sim(pos_inputs_embed, neg_labels_embed, mask)
+            + neg_inf * labels_bad_negs
+        )
+        sim_neg_ll = (
+            self.sim(pos_labels_embed, neg_labels_embed, mask)
+            + neg_inf * labels_bad_negs
+        )
+        sim_neg_ii = (
+            self.sim(pos_inputs_embed, neg_inputs_embed, mask)
+            + neg_inf * inputs_bad_negs
+        )
+        sim_neg_li = (
+            self.sim(pos_labels_embed, neg_inputs_embed, mask)
+            + neg_inf * inputs_bad_negs
+        )
+
+        # output similarities between user input and bot actions
+        # and similarities between bot actions and similarities between user inputs
+        return sim_pos, sim_neg_il, sim_neg_ll, sim_neg_ii, sim_neg_li
+
+    @staticmethod
+    def _calc_accuracy(sim_pos: tf.Tensor, sim_neg: tf.Tensor) -> tf.Tensor:
+        """Calculate accuracy."""
+        max_all_sim = tf.reduce_max(tf.concat([sim_pos, sim_neg], axis=-1), axis=-1)
+        sim_pos = tf.squeeze(sim_pos, axis=-1)
+        return layers_utils.reduce_mean_equal(max_all_sim, sim_pos)
+
+    def _loss_margin(
+        self,
+        sim_pos: tf.Tensor,
+        sim_neg_il: tf.Tensor,
+        sim_neg_ll: tf.Tensor,
+        sim_neg_ii: tf.Tensor,
+        sim_neg_li: tf.Tensor,
+        mask: Optional[tf.Tensor],
+    ) -> tf.Tensor:
+        """Define max margin loss."""
+        # loss for maximizing similarity with correct action
+        loss = tf.maximum(0.0, self.mu_pos - tf.squeeze(sim_pos, axis=-1))
+
+        # loss for minimizing similarity with `num_neg` incorrect actions
+        if self.use_max_sim_neg:
+            # minimize only maximum similarity over incorrect actions
+            max_sim_neg_il = tf.reduce_max(sim_neg_il, axis=-1)
+            loss += tf.maximum(0.0, self.mu_neg + max_sim_neg_il)
+        else:
+            # minimize all similarities with incorrect actions
+            max_margin = tf.maximum(0.0, self.mu_neg + sim_neg_il)
+            loss += tf.reduce_sum(max_margin, axis=-1)
+
+        # penalize max similarity between pos bot and neg bot embeddings
+        max_sim_neg_ll = tf.maximum(
+            0.0, self.mu_neg + tf.reduce_max(sim_neg_ll, axis=-1)
+        )
+        loss += max_sim_neg_ll * self.neg_lambda
+
+        # penalize max similarity between pos dial and neg dial embeddings
+        max_sim_neg_ii = tf.maximum(
+            0.0, self.mu_neg + tf.reduce_max(sim_neg_ii, axis=-1)
+        )
+        loss += max_sim_neg_ii * self.neg_lambda
+
+        # penalize max similarity between pos bot and neg dial embeddings
+        max_sim_neg_li = tf.maximum(
+            0.0, self.mu_neg + tf.reduce_max(sim_neg_li, axis=-1)
+        )
+        loss += max_sim_neg_li * self.neg_lambda
+
+        if mask is not None:
+            # mask loss for different length sequences
+            loss *= mask
+            # average the loss over sequence length
+            loss = tf.reduce_sum(loss, axis=-1) / tf.reduce_sum(mask, axis=1)
+
+        # average the loss over the batch
+        loss = tf.reduce_mean(loss)
+
+        return loss
+
+    def _loss_cross_entropy(
+        self,
+        sim_pos: tf.Tensor,
+        sim_neg_il: tf.Tensor,
+        sim_neg_ll: tf.Tensor,
+        sim_neg_ii: tf.Tensor,
+        sim_neg_li: tf.Tensor,
+        mask: Optional[tf.Tensor],
+    ) -> tf.Tensor:
+        """Defines cross entropy loss."""
+        loss = self._compute_softmax_loss(
+            sim_pos, sim_neg_il, sim_neg_ll, sim_neg_ii, sim_neg_li
+        )
+
+        if self.constrain_similarities:
+            loss += self._compute_sigmoid_loss(
+                sim_pos, sim_neg_il, sim_neg_ll, sim_neg_ii, sim_neg_li
+            )
+
+        loss = self.apply_mask_and_scaling(loss, mask)
+
+        # average the loss over the batch
+        return tf.reduce_mean(loss)
+
+    @staticmethod
+    def _compute_sigmoid_loss(
+        sim_pos: tf.Tensor,
+        sim_neg_il: tf.Tensor,
+        sim_neg_ll: tf.Tensor,
+        sim_neg_ii: tf.Tensor,
+        sim_neg_li: tf.Tensor,
+    ) -> tf.Tensor:
+        # Constrain similarity values in a range by applying sigmoid
+        # on them individually so that they saturate at extreme values.
+        sigmoid_logits = tf.concat(
+            [sim_pos, sim_neg_il, sim_neg_ll, sim_neg_ii, sim_neg_li], axis=-1
+        )
+        sigmoid_labels = tf.concat(
+            [
+                tf.ones_like(sigmoid_logits[..., :1]),
+                tf.zeros_like(sigmoid_logits[..., 1:]),
+            ],
+            axis=-1,
+        )
+        sigmoid_loss = tf.nn.sigmoid_cross_entropy_with_logits(
+            labels=sigmoid_labels, logits=sigmoid_logits
+        )
+        # average over logits axis
+        return tf.reduce_mean(sigmoid_loss, axis=-1)
+
+    def _compute_softmax_loss(
+        self,
+        sim_pos: tf.Tensor,
+        sim_neg_il: tf.Tensor,
+        sim_neg_ll: tf.Tensor,
+        sim_neg_ii: tf.Tensor,
+        sim_neg_li: tf.Tensor,
+    ) -> tf.Tensor:
+        # Similarity terms between input and label should be optimized relative
+        # to each other and hence use them as logits for softmax term
+        softmax_logits = tf.concat([sim_pos, sim_neg_il, sim_neg_li], axis=-1)
+        if not self.constrain_similarities:
+            # Concatenate other similarity terms as well. Due to this,
+            # similarity values between input and label may not be
+            # approximately bounded in a defined range.
+            softmax_logits = tf.concat(
+                [softmax_logits, sim_neg_ii, sim_neg_ll], axis=-1
+            )
+        # create label_ids for softmax
+        softmax_label_ids = tf.zeros_like(softmax_logits[..., 0], tf.int32)
+        softmax_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
+            labels=softmax_label_ids, logits=softmax_logits
+        )
+        return softmax_loss
+
+    @property
+    def _chosen_loss(self) -> Callable:
+        """Use loss depending on given option."""
+        if self.loss_type == MARGIN:
+            return self._loss_margin
+        elif self.loss_type == CROSS_ENTROPY:
+            return self._loss_cross_entropy
+        else:
+            raise TFLayerConfigException(
+                f"Wrong loss type '{self.loss_type}', "
+                f"should be '{MARGIN}' or '{CROSS_ENTROPY}'"
+            )
+
+    # noinspection PyMethodOverriding
+    def call(
+        self,
+        inputs_embed: tf.Tensor,
+        labels_embed: tf.Tensor,
+        labels: tf.Tensor,
+        all_labels_embed: tf.Tensor,
+        all_labels: tf.Tensor,
+        mask: Optional[tf.Tensor] = None,
+    ) -> Tuple[tf.Tensor, tf.Tensor]:
+        """Calculate loss and accuracy.
+
+        Args:
+            inputs_embed: Embedding tensor for the batch inputs;
+                shape `(batch_size, ..., num_features)`
+            labels_embed: Embedding tensor for the batch labels;
+                shape `(batch_size, ..., num_features)`
+            labels: Tensor representing batch labels; shape `(batch_size, ..., 1)`
+            all_labels_embed: Embedding tensor for the all labels;
+                shape `(num_labels, num_features)`
+            all_labels: Tensor representing all labels; shape `(num_labels, 1)`
+            mask: Optional mask, contains `1` for inputs and `0` for padding;
+                shape `(batch_size, 1)`
+
+        Returns:
+            loss: Total loss.
+            accuracy: Training accuracy.
+        """
+        (
+            pos_inputs_embed,
+            pos_labels_embed,
+            neg_inputs_embed,
+            neg_labels_embed,
+            inputs_bad_negs,
+            labels_bad_negs,
+        ) = self._sample_negatives(
+            inputs_embed, labels_embed, labels, all_labels_embed, all_labels
+        )
+
+        # calculate similarities
+        sim_pos, sim_neg_il, sim_neg_ll, sim_neg_ii, sim_neg_li = self._train_sim(
+            pos_inputs_embed,
+            pos_labels_embed,
+            neg_inputs_embed,
+            neg_labels_embed,
+            inputs_bad_negs,
+            labels_bad_negs,
+            mask,
+        )
+
+        accuracy = self._calc_accuracy(sim_pos, sim_neg_il)
+
+        loss = self._chosen_loss(
+            sim_pos, sim_neg_il, sim_neg_ll, sim_neg_ii, sim_neg_li, mask
+        )
+
+        return loss, accuracy
+
+
+class MultiLabelDotProductLoss(DotProductLoss):
+    """Multi-label dot-product loss layer.
+
+    This loss layer assumes that multiple outputs (labels) can be correct for any given
+    input. To accomodate for this, we use a sigmoid cross-entropy loss here.
+    """
+
+    def __init__(
+        self,
+        num_candidates: int,
+        scale_loss: bool = False,
+        constrain_similarities: bool = True,
+        model_confidence: Text = SOFTMAX,
+        similarity_type: Text = INNER,
+        name: Optional[Text] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Declares instance variables with default values.
+
+        Args:
+            num_candidates: Positive integer, the number of candidate labels.
+            scale_loss: If `True` scale loss inverse proportionally to
+                the confidence of the correct prediction.
+            similarity_type: Similarity measure to use, either `cosine` or `inner`.
+            name: Optional name of the layer.
+            constrain_similarities: Boolean, if `True` applies sigmoid on all
+                similarity terms and adds to the loss function to
+                ensure that similarity values are approximately bounded.
+                Used inside _loss_cross_entropy() only.
+            model_confidence: Normalization of confidence values during inference.
+                Currently, the only possible value is `SOFTMAX`.
+        """
+        super().__init__(
+            num_candidates,
+            scale_loss=scale_loss,
+            similarity_type=similarity_type,
+            name=name,
+            constrain_similarities=constrain_similarities,
+            model_confidence=model_confidence,
+        )
+
+    def call(
+        self,
+        batch_inputs_embed: tf.Tensor,
+        batch_labels_embed: tf.Tensor,
+        batch_labels_ids: tf.Tensor,
+        all_labels_embed: tf.Tensor,
+        all_labels_ids: tf.Tensor,
+        mask: Optional[tf.Tensor] = None,
+    ) -> Tuple[tf.Tensor, tf.Tensor]:
+        """Calculates loss and accuracy.
+
+        Args:
+            batch_inputs_embed: Embeddings of the batch inputs (e.g. featurized
+                trackers); shape `(batch_size, 1, num_features)`
+            batch_labels_embed: Embeddings of the batch labels (e.g. featurized intents
+                for IntentTED);
+                shape `(batch_size, max_num_labels_per_input, num_features)`
+            batch_labels_ids: Batch label indices (e.g. indices of the intents). We
+                assume that indices are integers that run from `0` to
+                `(number of labels) - 1`.
+                shape `(batch_size, max_num_labels_per_input, 1)`
+            all_labels_embed: Embeddings for all labels in the domain;
+                shape `(batch_size, num_features)`
+            all_labels_ids: Indices for all labels in the domain;
+                shape `(num_labels, 1)`
+            mask: Optional sequence mask, which contains `1` for inputs and `0` for
+                padding.
+
+        Returns:
+            loss: Total loss (based on StarSpace http://arxiv.org/abs/1709.03856);
+                scalar
+            accuracy: Training accuracy; scalar
+        """
+        (
+            pos_inputs_embed,  # (batch_size, 1, 1, num_features)
+            pos_labels_embed,  # (batch_size, 1, max_num_labels_per_input, num_features)
+            candidate_labels_embed,  # (batch_size, 1, num_candidates, num_features)
+            pos_neg_labels,  # (batch_size, num_candidates)
+        ) = self._sample_candidates(
+            batch_inputs_embed,
+            batch_labels_embed,
+            batch_labels_ids,
+            all_labels_embed,
+            all_labels_ids,
+        )
+
+        # Calculate similarities
+        sim_pos, sim_candidate_il = self._train_sim(
+            pos_inputs_embed, pos_labels_embed, candidate_labels_embed, mask
+        )
+
+        label_padding_mask = self._construct_mask_for_label_padding(
+            batch_labels_ids, tf.shape(pos_neg_labels)[-1]
+        )
+
+        # Repurpose the `mask` argument of `_accuracy` and `_loss_sigmoid`
+        # to pass the `label_padding_mask`. We can do this right now because
+        # we don't use `MultiLabelDotProductLoss` for sequence tagging tasks
+        # yet. Hence, the `mask` argument passed to this function will always
+        # be empty. Whenever, we come across a use case where `mask` is
+        # non-empty we'll have to refactor the `_accuracy` and `_loss_sigmoid`
+        # functions to take into consideration both, sequence level masks as
+        # well as label padding masks.
+
+        accuracy = self._accuracy(
+            sim_pos, sim_candidate_il, pos_neg_labels, label_padding_mask
+        )
+        loss = self._loss_sigmoid(
+            sim_pos, sim_candidate_il, pos_neg_labels, mask=label_padding_mask
+        )
+
+        return loss, accuracy
+
+    @staticmethod
+    def _construct_mask_for_label_padding(
+        batch_labels_ids: tf.Tensor, num_candidates: tf.Tensor
+    ) -> tf.Tensor:
+        """Constructs a mask which indicates indices for valid label ids.
+
+        Indices corresponding to valid label ids have a
+        `1` and indices corresponding to `LABEL_PAD_ID`
+        have a `0`.
+
+        Args:
+            batch_labels_ids: Batch label indices (e.g. indices of the intents). We
+                assume that indices are integers that run from `0` to
+                `(number of labels) - 1` with a special
+                value for padding which is set to `LABEL_PAD_ID`.
+                shape `(batch_size, max_num_labels_per_input, 1)`
+            num_candidates: Number of candidates sampled.
+
+        Returns:
+            Constructed mask.
+        """
+        pos_label_pad_indices = tf.cast(
+            tf.equal(tf.squeeze(batch_labels_ids, -1), LABEL_PAD_ID), dtype=tf.float32
+        )
+
+        # Flip 1 and 0 to 0 and 1 respectively
+        pos_label_pad_mask = 1 - pos_label_pad_indices
+
+        # `pos_label_pad_mask` only contains the mask for label ids
+        # seen in the batch. For sampled candidate label ids, the mask
+        # should be a tensor of `1`s since all candidate label ids
+        # are valid. From this, we construct the padding mask for
+        # all label ids: label ids seen in the batch + label ids sampled.
+        all_label_pad_mask = tf.concat(
+            [
+                pos_label_pad_mask,
+                tf.ones(
+                    (tf.shape(batch_labels_ids)[0], num_candidates), dtype=tf.float32
+                ),
+            ],
+            axis=-1,
+        )
+
+        return all_label_pad_mask
+
+    def _train_sim(
+        self,
+        pos_inputs_embed: tf.Tensor,
+        pos_labels_embed: tf.Tensor,
+        candidate_labels_embed: tf.Tensor,
+        mask: tf.Tensor,
+    ) -> Tuple[tf.Tensor, tf.Tensor]:
+        sim_pos = self.sim(
+            pos_inputs_embed, pos_labels_embed, mask
+        )  # (batch_size, 1, max_labels_per_input)
+        sim_candidate_il = self.sim(
+            pos_inputs_embed, candidate_labels_embed, mask
+        )  # (batch_size, 1, num_candidates)
+
+        return sim_pos, sim_candidate_il
+
+    def _sample_candidates(
+        self,
+        batch_inputs_embed: tf.Tensor,
+        batch_labels_embed: tf.Tensor,
+        batch_labels_ids: tf.Tensor,
+        all_labels_embed: tf.Tensor,
+        all_labels_ids: tf.Tensor,
+    ) -> Tuple[
+        tf.Tensor,  # (batch_size, 1, 1, num_features)
+        tf.Tensor,  # (batch_size, 1, num_features)
+        tf.Tensor,  # (batch_size, 1, num_candidates, num_features)
+        tf.Tensor,  # (batch_size, num_candidates)
+    ]:
+        """Samples candidate examples.
+
+        Args:
+            batch_inputs_embed: Embeddings of the batch inputs (e.g. featurized
+                trackers) # (batch_size, 1, num_features)
+            batch_labels_embed: Embeddings of the batch labels (e.g. featurized intents
+                for IntentTED) # (batch_size, max_num_labels_per_input, num_features)
+            batch_labels_ids: Batch label indices (e.g. indices of the
+                intents) # (batch_size, max_num_labels_per_input, 1)
+            all_labels_embed: Embeddings for all labels in
+                the domain # (num_labels, num_features)
+            all_labels_ids: Indices for all labels in the
+                domain # (num_labels, 1)
+
+        Returns:
+            pos_inputs_embed: Embeddings of the batch inputs
+            pos_labels_embed: Embeddings of the batch labels with an extra
+                dimension inserted.
+            candidate_labels_embed: More examples of embeddings of labels, some positive
+                some negative
+            pos_neg_indicators: Indicator for which candidates are positives and which
+                are negatives
+        """
+        pos_inputs_embed = tf.expand_dims(
+            batch_inputs_embed, axis=-2, name="expand_pos_input"
+        )
+
+        pos_labels_embed = tf.expand_dims(
+            batch_labels_embed, axis=1, name="expand_pos_labels"
+        )
+
+        # Pick random examples from the batch
+        candidate_ids = layers_utils.random_indices(
+            batch_size=tf.shape(batch_inputs_embed)[0],
+            n=self.num_neg,
+            n_max=tf.shape(all_labels_embed)[0],
+        )
+
+        # Get the label embeddings corresponding to candidate indices
+        candidate_labels_embed = layers_utils.get_candidate_values(
+            all_labels_embed, candidate_ids
+        )
+        candidate_labels_embed = tf.expand_dims(candidate_labels_embed, axis=1)
+
+        # Get binary indicators of whether a candidate is positive or not
+        pos_neg_indicators = self._get_pos_neg_indicators(
+            all_labels_ids, batch_labels_ids, candidate_ids
+        )
+
+        return (
+            pos_inputs_embed,
+            pos_labels_embed,
+            candidate_labels_embed,
+            pos_neg_indicators,
+        )
+
+    def _get_pos_neg_indicators(
+        self,
+        all_labels_ids: tf.Tensor,
+        batch_labels_ids: tf.Tensor,
+        candidate_ids: tf.Tensor,
+    ) -> tf.Tensor:
+        """Computes indicators for which candidates are positive labels.
+
+        Args:
+            all_labels_ids: Indices of all the labels
+            batch_labels_ids: Indices of the labels in the examples
+            candidate_ids: Indices of labels that may or may not appear in the examples
+
+        Returns:
+            Binary indicators of whether or not a label is positive
+        """
+        candidate_labels_ids = layers_utils.get_candidate_values(
+            all_labels_ids, candidate_ids
+        )
+        candidate_labels_ids = tf.expand_dims(candidate_labels_ids, axis=1)
+
+        # Determine how many distinct labels exist (highest label index)
+        max_label_id = tf.cast(tf.math.reduce_max(all_labels_ids), dtype=tf.int32)
+
+        # Convert the positive label ids to their one_hot representation.
+        # Note: -1 indices yield a zeros-only vector. We use -1 as a padding token,
+        # as the number of positive labels in each example can differ. The padding is
+        # added in the TrackerFeaturizer.
+        batch_labels_one_hot = tf.one_hot(
+            tf.cast(tf.squeeze(batch_labels_ids, axis=-1), tf.int32),
+            max_label_id + 1,
+            axis=-1,
+        )  # (batch_size, max_num_labels_per_input, max_label_id)
+
+        # Collapse the extra dimension and convert to a multi-hot representation
+        # by aggregating all ones in the one-hot representation.
+        # We use tf.reduce_any instead of tf.reduce_sum because several examples can
+        # have the same postivie label.
+        batch_labels_multi_hot = tf.cast(
+            tf.math.reduce_any(tf.cast(batch_labels_one_hot, dtype=tf.bool), axis=-2),
+            tf.float32,
+        )  # (batch_size, max_label_id)
+
+        # Remove extra dimensions for gather
+        candidate_labels_ids = tf.squeeze(tf.squeeze(candidate_labels_ids, 1), -1)
+
+        # Collect binary indicators of whether or not a label is positive
+        return tf.gather(
+            batch_labels_multi_hot,
+            tf.cast(candidate_labels_ids, tf.int32),
+            batch_dims=1,
+            name="gather_labels",
+        )
+
+    def _loss_sigmoid(
+        self,
+        sim_pos: tf.Tensor,  # (batch_size, 1, max_num_labels_per_input)
+        sim_candidates_il: tf.Tensor,  # (batch_size, 1, num_candidates)
+        pos_neg_labels: tf.Tensor,  # (batch_size, num_candidates)
+        mask: Optional[
+            tf.Tensor
+        ] = None,  # (batch_size, max_num_labels_per_input + num_candidates)
+    ) -> tf.Tensor:  # ()
+        """Computes the sigmoid loss."""
+        # Concatenate the guaranteed positive examples with the candidate examples,
+        # some of which are positives and others are negatives. Which are which
+        # is stored in `pos_neg_labels`.
+        logits = tf.concat([sim_pos, sim_candidates_il], axis=-1, name="logit_concat")
+        logits = tf.squeeze(logits, 1)
+
+        # Create label_ids for sigmoid. `mask` will take care of the
+        # extra 1s we create as label ids for indices corresponding
+        # to padding ids.
+        pos_label_ids = tf.squeeze(tf.ones_like(sim_pos, tf.float32), 1)
+        label_ids = tf.concat(
+            [pos_label_ids, pos_neg_labels], axis=-1, name="gt_concat"
+        )
+
+        # Compute the sigmoid cross-entropy loss. When minimized, the embeddings
+        # for the two classes (positive and negative) are pushed away from each
+        # other in the embedding space, while it is allowed that any input embedding
+        # corresponds to more than one label.
+        loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=label_ids, logits=logits)
+
+        loss = self.apply_mask_and_scaling(loss, mask)
+
+        # Average the loss over the batch
+        return tf.reduce_mean(loss)
+
+    @staticmethod
+    def _accuracy(
+        sim_pos: tf.Tensor,  # (batch_size, 1, max_num_labels_per_input)
+        sim_candidates: tf.Tensor,  # (batch_size, 1, num_candidates)
+        pos_neg_indicators: tf.Tensor,  # (batch_size, num_candidates)
+        mask: tf.Tensor,  # (batch_size, max_num_labels_per_input + num_candidates)
+    ) -> tf.Tensor:  # ()
+        """Calculates the accuracy."""
+        all_preds = tf.concat(
+            [sim_pos, sim_candidates], axis=-1, name="acc_concat_preds"
+        )
+        all_preds_sigmoid = tf.nn.sigmoid(all_preds)
+        all_pred_labels = tf.squeeze(tf.math.round(all_preds_sigmoid), 1)
+
+        # Create an indicator for the positive labels by concatenating the 1 for all
+        # guaranteed positive labels and the `pos_neg_indicators`
+        all_positives = tf.concat(
+            [tf.squeeze(tf.ones_like(sim_pos), axis=1), pos_neg_indicators],
+            axis=-1,
+            name="acc_concat_gt",
+        )
+
+        return layers_utils.reduce_mean_equal(all_pred_labels, all_positives, mask=mask)