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.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)