jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def update_state(self, y_true, y_pred, sample_weight=None):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
    """Update state of the metric.

    Args:
      y_true: the true labels with shape [batch_size, 1].
      y_pred: model output, which is the similarity matrix with shape
        [batch_size, label_embedding_vocab_size] between context and full vocab
        label embeddings.
      sample_weight: Optional weighting of each example. Defaults to 1.
    """
    similarities = y_pred
    label_indices = tf.expand_dims(tf.cast(y_true, dtype=tf.int32), -1)

    # Get the indices of similar labels in sorted order for each context.
    sorted_similarities = tf.argsort(
        similarities, axis=-1, direction='DESCENDING')

    ranks = tf.where(tf.equal(sorted_similarities, label_indices))[:, -1]
    ranks = tf.cast(ranks, 'float32')
    super().update_state(ranks, sample_weight=sample_weight)
//  [python]
def update_state(self, y_true, y_pred, sample_weight=None):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
    """Update state of the metric.

    Args:
      y_true: the true labels with shape [batch_size, 1].
      y_pred: model output, which is the similarity matrix with shape
        [batch_size, label_embedding_vocab_size] between context and full vocab
        label embeddings.
      sample_weight: Optional weighting of each example. Defaults to 1.
    """
    similarities = y_pred
    label_indices = tf.expand_dims(tf.cast(y_true, dtype=tf.int32), -1)

    # Get the indices of similar labels in sorted order for each context.
    sorted_similarities = tf.argsort(
        similarities, axis=-1, direction='DESCENDING')

    ranks = tf.where(tf.equal(sorted_similarities, label_indices))[:, -1]
    ranks = tf.keras.backend.cast(ranks, 'float32')
    super().update_state(ranks, sample_weight=sample_weight)
//  [python]
def _get_labels_embed(
        label_ids: tf.Tensor, all_labels_embed: tf.Tensor
    ) -> tf.Tensor:
        # instead of processing labels again, gather embeddings from
        # all_labels_embed using label ids

        indices = tf.to_int32(label_ids[:, :, 0])
        labels_embed = tf.gather(all_labels_embed, indices)

        return labels_embed
//  [python]
def _get_labels_embed(
        label_ids: tf.Tensor, all_labels_embed: tf.Tensor
    ) -> tf.Tensor:
        # instead of processing labels again, gather embeddings from
        # all_labels_embed using label ids

        indices = tf.cast(label_ids[:, :, 0], tf.int32)
        labels_embed = tf.gather(all_labels_embed, indices)

        return labels_embed
//  [python]
def call(self, attention_mask, past_key_values_length: int = 0):
        """`input_ids_shape` is expected to be [bsz x seqlen]."""
        attention_mask = t.cast(attention_mask, t.int64)

        # create positions depending on attention_mask
        positions = tf.math.cumsum(attention_mask, axis=1) * attention_mask - 1

        # cut positions if `past_key_values_length` is > 0
        positions = positions[:, past_key_values_length:]

        return super().call(positions + self.offset)
//  [python]
def call(self, attention_mask, past_key_values_length: int = 0):
        """`input_ids_shape` is expected to be [bsz x seqlen]."""
        attention_mask = tf.cast(attention_mask, tf.int64)

        # create positions depending on attention_mask
        positions = tf.math.cumsum(attention_mask, axis=1) * attention_mask - 1

        # cut positions if `past_key_values_length` is > 0
        positions = positions[:, past_key_values_length:]

        return super().call(positions + self.offset)
//  [python]
def __call__(self, shape, dtype=tf.float32):
        maxval = tf.sqrt(tf.constant(1.) / tf.cast(shape[0], tf.float32))
        maxval = t.cast(maxval, dtype=dtype)
        minval = -maxval

        weights = tf.random.uniform(shape, minval=minval, maxval=maxval, dtype=dtype)
        weights = t.cast(weights, dtype=tf.float32)
        return weights
//  [python]
def __call__(self, shape, dtype=tf.float32):
        maxval = tf.sqrt(tf.constant(1.) / tf.cast(shape[0], tf.float32))
        maxval = tf.cast(maxval, dtype=dtype)
        minval = -maxval

        weights = tf.random.uniform(shape, minval=minval, maxval=maxval, dtype=dtype)
        weights = tf.cast(weights, dtype=tf.float32)
        return weights
//  [python]
def _make_causal_mask(input_ids_shape: tf.TensorShape, past_key_values_length: int = 0):
    """
    Make causal mask used for bi-directional self-attention.
    """
    bsz = input_ids_shape[0]
    tgt_len = input_ids_shape[1]
    # We need triu with k = 1 but TF expects known compile-time dims for that, so we hack around it
    mask = tf.fill((tgt_len, tgt_len), tf.to_float(LARGE_NEGATIVE))
    mask = tf.linalg.band_part(mask, 0, -1) - tf.linalg.band_part(mask, 0, 0)

    if past_key_values_length > 0:
        mask = tf.concat([tf.zeros((tgt_len, past_key_values_length)), mask], axis=-1)

    return tf.tile(mask[None, None, :, :], (bsz, 1, 1, 1))
//  [python]
def _make_causal_mask(input_ids_shape: tf.TensorShape, past_key_values_length: int = 0):
    """
    Make causal mask used for bi-directional self-attention.
    """
    bsz = input_ids_shape[0]
    tgt_len = input_ids_shape[1]
    # We need triu with k = 1 but TF expects known compile-time dims for that, so we hack around it
    mask = tf.fill((tgt_len, tgt_len), tf.cast(LARGE_NEGATIVE, tf.float32))
    mask = tf.linalg.band_part(mask, 0, -1) - tf.linalg.band_part(mask, 0, 0)

    if past_key_values_length > 0:
        mask = tf.concat([tf.zeros((tgt_len, past_key_values_length)), mask], axis=-1)

    return tf.tile(mask[None, None, :, :], (bsz, 1, 1, 1))
//  [python]
def _embed_dialogue(
        self,
        dialogue_in: tf.Tensor,
        tf_batch_data: Dict[Text, Dict[Text, List[tf.Tensor]]],
    ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor, Optional[tf.Tensor]]:
        """Creates dialogue level embedding and mask.

        Args:
            dialogue_in: The encoded dialogue.
            tf_batch_data: Batch in model data format.

        Returns:
            The dialogue embedding, the mask, and (for diagnostic purposes)
            also the attention weights.
        """
        dialogue_lengths = t.cast(tf_batch_data[DIALOGUE][LENGTH][0], t.int32)
        mask = rasa_layers.compute_mask(dialogue_lengths)

        if self.max_history_featurizer_is_used:
            # invert dialogue sequence so that the last turn would always have
            # exactly the same positional encoding
            dialogue_in = tf.reverse_sequence(dialogue_in, dialogue_lengths, seq_axis=1)

        dialogue_transformed, attention_weights = self._tf_layers[
            f"transformer.{DIALOGUE}"
        ](dialogue_in, 1 - mask, self._training)
        dialogue_transformed = tf.nn.gelu(dialogue_transformed)

        if self.max_history_featurizer_is_used:
            # pick last vector if max history featurizer is used, since we inverted
            # dialogue sequence, the last vector is actually the first one
            dialogue_transformed = dialogue_transformed[:, :1, :]
            mask = tf.expand_dims(self._last_token(mask, dialogue_lengths), 1)
        elif not self._training:
            # during prediction we don't care about previous dialogue turns,
            # so to save computation time, use only the last one
            dialogue_transformed = tf.expand_dims(
                self._last_token(dialogue_transformed, dialogue_lengths), 1
            )
            mask = tf.expand_dims(self._last_token(mask, dialogue_lengths), 1)

        dialogue_embed = self._tf_layers[f"embed.{DIALOGUE}"](dialogue_transformed)

        return dialogue_embed, mask, dialogue_transformed, attention_weights
//  [python]
def _embed_dialogue(
        self,
        dialogue_in: tf.Tensor,
        tf_batch_data: Dict[Text, Dict[Text, List[tf.Tensor]]],
    ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor, Optional[tf.Tensor]]:
        """Creates dialogue level embedding and mask.

        Args:
            dialogue_in: The encoded dialogue.
            tf_batch_data: Batch in model data format.

        Returns:
            The dialogue embedding, the mask, and (for diagnostic purposes)
            also the attention weights.
        """
        dialogue_lengths = tf.cast(tf_batch_data[DIALOGUE][LENGTH][0], tf.int32)
        mask = rasa_layers.compute_mask(dialogue_lengths)

        if self.max_history_featurizer_is_used:
            # invert dialogue sequence so that the last turn would always have
            # exactly the same positional encoding
            dialogue_in = tf.reverse_sequence(dialogue_in, dialogue_lengths, seq_axis=1)

        dialogue_transformed, attention_weights = self._tf_layers[
            f"transformer.{DIALOGUE}"
        ](dialogue_in, 1 - mask, self._training)
        dialogue_transformed = tf.nn.gelu(dialogue_transformed)

        if self.max_history_featurizer_is_used:
            # pick last vector if max history featurizer is used, since we inverted
            # dialogue sequence, the last vector is actually the first one
            dialogue_transformed = dialogue_transformed[:, :1, :]
            mask = tf.expand_dims(self._last_token(mask, dialogue_lengths), 1)
        elif not self._training:
            # during prediction we don't care about previous dialogue turns,
            # so to save computation time, use only the last one
            dialogue_transformed = tf.expand_dims(
                self._last_token(dialogue_transformed, dialogue_lengths), 1
            )
            mask = tf.expand_dims(self._last_token(mask, dialogue_lengths), 1)

        dialogue_embed = self._tf_layers[f"embed.{DIALOGUE}"](dialogue_transformed)

        return dialogue_embed, mask, dialogue_transformed, attention_weights
//  [python]
def _create_last_dialogue_turns_mask(
        tf_batch_data: Dict[Text, Dict[Text, List[tf.Tensor]]], attribute: Text
    ) -> tf.Tensor:
        # Since max_history_featurizer_is_used is True,
        # we need to find the locations of last dialogue turns in
        # (combined batch dimension and dialogue length,) dimension,
        # so that we can use `_sequence_lengths` as a boolean  mask to pick
        # which ones are "real" textual input in these last dialogue turns.

        # In order to do that we can use given `dialogue_lengths`.
        # For example:
        # If we have `dialogue_lengths = [2, 1, 3]`, than
        # `dialogue_indices = [0, 1, 0, 0, 1, 2]` here we can spot that `0`
        # always indicates the first dialogue turn,
        # which means that previous dialogue turn is the last dialogue turn.
        # Combining this with the fact that the last element in
        # `dialogue_indices` is always the last dialogue turn, we can add
        # a `0` to the end, getting
        # `_dialogue_indices = [0, 1, 0, 0, 1, 2, 0]`.
        # Then removing the first element
        # `_last_dialogue_turn_inverse_indicator = [1, 0, 0, 1, 2, 0]`
        # we see that `0` points to the last dialogue turn.
        # We convert all positive numbers to `True` and take
        # the inverse mask to get
        # `last_dialogue_mask = [0, 1, 1, 0, 0, 1],
        # which precisely corresponds to the fact that first dialogue is of
        # length 2, the second 1 and the third 3.
        last_dialogue_turn_mask = tf.math.logical_not(
            tf.to_int64(
                tf.concat(
                    [
                        tf_batch_data[DIALOGUE][INDICES][0],
                        tf.zeros((1,), dtype=tf.int32),
                    ],
                    axis=0,
                )[1:],
            )
        )
        # get only the indices of real inputs
        return tf.boolean_mask(
            last_dialogue_turn_mask,
            tf.reshape(tf_batch_data[attribute][SEQUENCE_LENGTH][0], (-1,)),
        )
//  [python]
def _create_last_dialogue_turns_mask(
        tf_batch_data: Dict[Text, Dict[Text, List[tf.Tensor]]], attribute: Text
    ) -> tf.Tensor:
        # Since max_history_featurizer_is_used is True,
        # we need to find the locations of last dialogue turns in
        # (combined batch dimension and dialogue length,) dimension,
        # so that we can use `_sequence_lengths` as a boolean  mask to pick
        # which ones are "real" textual input in these last dialogue turns.

        # In order to do that we can use given `dialogue_lengths`.
        # For example:
        # If we have `dialogue_lengths = [2, 1, 3]`, than
        # `dialogue_indices = [0, 1, 0, 0, 1, 2]` here we can spot that `0`
        # always indicates the first dialogue turn,
        # which means that previous dialogue turn is the last dialogue turn.
        # Combining this with the fact that the last element in
        # `dialogue_indices` is always the last dialogue turn, we can add
        # a `0` to the end, getting
        # `_dialogue_indices = [0, 1, 0, 0, 1, 2, 0]`.
        # Then removing the first element
        # `_last_dialogue_turn_inverse_indicator = [1, 0, 0, 1, 2, 0]`
        # we see that `0` points to the last dialogue turn.
        # We convert all positive numbers to `True` and take
        # the inverse mask to get
        # `last_dialogue_mask = [0, 1, 1, 0, 0, 1],
        # which precisely corresponds to the fact that first dialogue is of
        # length 2, the second 1 and the third 3.
        last_dialogue_turn_mask = tf.math.logical_not(
            tf.cast(
                tf.concat(
                    [
                        tf_batch_data[DIALOGUE][INDICES][0],
                        tf.zeros((1,), dtype=tf.int32),
                    ],
                    axis=0,
                )[1:],
                dtype=tf.bool,
            )
        )
        # get only the indices of real inputs
        return tf.boolean_mask(
            last_dialogue_turn_mask,
            tf.reshape(tf_batch_data[attribute][SEQUENCE_LENGTH][0], (-1,)),
        )
//  [python]
def creat_mask(self, seq):
        seq_pad = tf.to_float(tf.equal(seq, 0))
        return seq_pad[:, tf.newaxis, tf.newaxis, :]
//  [python]
def creat_mask(self, seq):
        seq_pad = tf.cast(tf.equal(seq, 0), tf.float32)
        return seq_pad[:, tf.newaxis, tf.newaxis, :]
//  [python]
def call(self, inputs, training=None):
    if training is None:
      training = tf_keras.backend.learning_phase()

    backbone_output = inputs[0]
    decoder_output = inputs[1][str(self._config_dict['level'])]

    x = decoder_output
    for i in range(len(self._config_dict['low_level'])):
      feature = backbone_output[str(self._config_dict['low_level'][i])]
      feature = self._projection_convs[i](feature)
      feature = self._projection_norms[i](feature, training=training)
      feature = self._activation(feature)

      shape = tf.shape(feature)
      x = tf.image.resize(
          x, size=[shape[1], shape[2]],
          method=self._config_dict['interpolation'])
      x = tf.cast(x, dtype=feature.dtype)
      x = tf.concat([x, feature], axis=self._channel_axis)

      x = self._fusion_convs[i](x)
      x = self._fusion_norms[i](x, training=training)
      x = self._activation(x)
    return x
//  [python]
def bert_feature_loss(self, real, pred):
        T1=tf.shape(real)[1]
        T2=tf.shape(pred)[1]
        T=tf.reduce_min([T1,T2])
        real=real[:,:T]
        pred=pred[:,:T]
        mask = tf.cast(tf.not_equal(real, -10.), 1.)
        loss = tf.square(real - pred)
        loss *= mask
        return tf.reduce_mean(tf.reduce_sum(loss,-1) / (tf.reduce_sum(mask,-1) + 1e-6),-1,True)
//  [python]
def classes_acc(self, real, pred):
        mask = tf.math.logical_not(tf.math.equal(real, 0))
        accs = tf.keras.metrics.sparse_categorical_accuracy(real,pred)

        mask = tf.to_float(mask)
        accs *= mask
        final=tf.reduce_sum(accs,-1)/tf.reduce_sum(mask,-1)

        return tf.reduce_mean(final)
//  [python]
def classes_acc(self, real, pred):
        mask = tf.math.logical_not(tf.math.equal(real, 0))
        accs = tf.keras.metrics.sparse_categorical_accuracy(real,pred)

        mask = tf.cast(mask, dtype=accs.dtype)
        accs *= mask
        final=tf.reduce_sum(accs,-1)/tf.reduce_sum(mask,-1)

        return tf.reduce_mean(final)
//  [python]
def classes_loss(self, real, pred):
        mask = tf.math.logical_not(tf.math.equal(real, 0))
        loss = tf.keras.losses.sparse_categorical_crossentropy(real, pred,True)
        mask_one= tf.math.logical_not(tf.math.equal(real, 1))
        mask = tf.to_float(mask)
        mask_one = tf.to_float(mask_one)
        mask_one*=mask
        loss_all =loss* mask
        final=tf.reduce_sum(loss_all,-1,True)/(tf.reduce_sum(mask,-1,True)+1e-6)
        loss_other = loss_all*mask_one
        final2 = tf.reduce_sum(loss_other, -1, True) / (tf.reduce_sum(mask_one, -1, True)+1e-6)
        return final+final2
//  [python]
def classes_loss(self, real, pred):
        mask = tf.math.logical_not(tf.math.equal(real, 0))
        loss = tf.keras.losses.sparse_categorical_crossentropy(real, pred,True)
        mask_one= tf.math.logical_not(tf.math.equal(real, 1))
        mask = tf.cast(mask, dtype=loss.dtype)
        mask_one = tf.cast(mask_one, dtype=loss.dtype)
        mask_one*=mask
        loss_all =loss* mask
        final=tf.reduce_sum(loss_all,-1,True)/(tf.reduce_sum(mask,-1,True)+1e-6)
        loss_other = loss_all*mask_one
        final2 = tf.reduce_sum(loss_other, -1, True) / (tf.reduce_sum(mask_one, -1, True)+1e-6)
        return final+final2
//  [python]
def reward(
      self, observation: types.NestedTensor, t: types.Int
  ) -> types.NestedTensor:
    def same_interval_parameters():
      """Returns the parameters of the current piece.

      Returns:
        The pair of `tf.Tensor` `(observation_to_reward, additive_reward)`.
      """
      return [
          self._current_observation_to_reward,
          self._current_additive_reward,
      ]

    def new_interval_parameters():
      """Update and returns the piece parameters.

      Returns:
        The pair of `tf.Tensor` `(observation_to_reward, additive_reward)`.
      """
      tf.compat.v1.assign_add(
          self._current_interval,
          t(self._interval_distribution.sample(), dtype=tf.int64),
      )
      tf.compat.v1.assign(
          self._current_observation_to_reward,
          self._observation_to_reward_distribution.sample(),
      )
      tf.compat.v1.assign(
          self._current_additive_reward,
          self._additive_reward_distribution.sample(),
      )

      return [
          self._current_observation_to_reward,
          self._current_additive_reward,
      ]

    observation_to_reward, additive_reward = tf.cond(
        t < self._current_interval,
        same_interval_parameters,
        new_interval_parameters,
    )

    reward = tf.matmul(observation, observation_to_reward) + tf.reshape(
        additive_reward, [1, -1]
    )
    return reward
//  [python]
def reward(
      self, observation: types.NestedTensor, t: types.Int
  ) -> types.NestedTensor:
    def same_interval_parameters():
      """Returns the parameters of the current piece.

      Returns:
        The pair of `tf.Tensor` `(observation_to_reward, additive_reward)`.
      """
      return [
          self._current_observation_to_reward,
          self._current_additive_reward,
      ]

    def new_interval_parameters():
      """Update and returns the piece parameters.

      Returns:
        The pair of `tf.Tensor` `(observation_to_reward, additive_reward)`.
      """
      tf.compat.v1.assign_add(
          self._current_interval,
          tf.cast(self._interval_distribution.sample(), dtype=tf.int64),
      )
      tf.compat.v1.assign(
          self._current_observation_to_reward,
          self._observation_to_reward_distribution.sample(),
      )
      tf.compat.v1.assign(
          self._current_additive_reward,
          self._additive_reward_distribution.sample(),
      )

      return [
          self._current_observation_to_reward,
          self._current_additive_reward,
      ]

    observation_to_reward, additive_reward = tf.cond(
        t < self._current_interval,
        same_interval_parameters,
        new_interval_parameters,
    )

    reward = tf.matmul(observation, observation_to_reward) + tf.reshape(
        additive_reward, [1, -1]
    )
    return reward
//  [python]
def _length_normalization(alpha, length, dtype=tf.float32):
  """Return length normalization factor."""
  return tf.pow(((5. + tf.cast(length, dtype)) / 6.), alpha)
//  [python]
def _get_global_attn_indices(is_index_global_attn):
        """ compute global attn indices required throughout forward pass """
        # helper variable
        num_global_attn_indices = tf.reduce_sum(tf.to_int32(is_index_global_attn), axis=1)

        # max number of global attn indices in batch
        max_num_global_attn_indices = tf.reduce_max(num_global_attn_indices)

        # indices of global attn
        is_index_global_attn_nonzero = tf.where(is_index_global_attn)

        # helper variable
        is_local_index_global_attn = tf.range(max_num_global_attn_indices) < tf.expand_dims(
            num_global_attn_indices, axis=-1
        )

        # location of the non-padding values within global attention indices
        is_local_index_global_attn_nonzero = tf.where(is_local_index_global_attn)

        # location of the padding values within global attention indices
        is_local_index_no_global_attn_nonzero = tf.where(tf.math.logical_not(is_local_index_global_attn))

        return (
            max_num_global_attn_indices,
            is_index_global_attn_nonzero,
            is_local_index_global_attn_nonzero,
            is_local_index_no_global_attn_nonzero,
        )
//  [python]
def _get_global_attn_indices(is_index_global_attn):
        """ compute global attn indices required throughout forward pass """
        # helper variable
        num_global_attn_indices = tf.reduce_sum(tf.cast(is_index_global_attn, dtype=tf.dtypes.int32), axis=1)

        # max number of global attn indices in batch
        max_num_global_attn_indices = tf.reduce_max(num_global_attn_indices)

        # indices of global attn
        is_index_global_attn_nonzero = tf.where(is_index_global_attn)

        # helper variable
        is_local_index_global_attn = tf.range(max_num_global_attn_indices) < tf.expand_dims(
            num_global_attn_indices, axis=-1
        )

        # location of the non-padding values within global attention indices
        is_local_index_global_attn_nonzero = tf.where(is_local_index_global_attn)

        # location of the padding values within global attention indices
        is_local_index_no_global_attn_nonzero = tf.where(tf.math.logical_not(is_local_index_global_attn))

        return (
            max_num_global_attn_indices,
            is_index_global_attn_nonzero,
            is_local_index_global_attn_nonzero,
            is_local_index_no_global_attn_nonzero,
        )
//  [python]
def _compute_global_attention_mask(input_ids_shape, sep_token_indices, before_sep_token=True):
    """
        Computes global attention mask by putting attention on all tokens
        before `sep_token_id` if `before_sep_token is True` else after
        `sep_token_id`.
    """

    assert sep_token_indices.shape[1] == 2, "`input_ids` should have two dimensions"
    question_end_index = tf.reshape(sep_token_indices, (input_ids_shape[0], 3, 2))[:, 0, 1]
    question_end_index = tf.cast(question_end_index[:, None], tf.dtypes.int32)  # size: batch_size x 1
    # bool attention mask with True in locations of global attention
    attention_mask = tf.range(input_ids_shape[1])
    if before_sep_token is True:
        attention_mask = tf.to_int32(
            tf.broadcast_to(attention_mask, input_ids_shape) < tf.broadcast_to(question_end_index, input_ids_shape),
        )
    else:
        # last token is separation token and should not be counted and in the middle are two separation tokens
        attention_mask = tf.to_int32(
            tf.broadcast_to(attention_mask, input_ids_shape)
            > tf.broadcast_to(question_end_index + 1, input_ids_shape),
        ) * tf.to_int32(tf.broadcast_to(attention_mask, input_ids_shape) < input_ids_shape[-1])

    return attention_mask
//  [python]
def _compute_global_attention_mask(input_ids_shape, sep_token_indices, before_sep_token=True):
    """
        Computes global attention mask by putting attention on all tokens
        before `sep_token_id` if `before_sep_token is True` else after
        `sep_token_id`.
    """

    assert sep_token_indices.shape[1] == 2, "`input_ids` should have two dimensions"
    question_end_index = tf.reshape(sep_token_indices, (input_ids_shape[0], 3, 2))[:, 0, 1]
    question_end_index = tf.cast(question_end_index[:, None], tf.dtypes.int32)  # size: batch_size x 1
    # bool attention mask with True in locations of global attention
    attention_mask = tf.range(input_ids_shape[1])
    if before_sep_token is True:
        attention_mask = tf.cast(
            tf.broadcast_to(attention_mask, input_ids_shape) < tf.broadcast_to(question_end_index, input_ids_shape),
            tf.dtypes.int32,
        )
    else:
        # last token is separation token and should not be counted and in the middle are two separation tokens
        attention_mask = tf.cast(
            tf.broadcast_to(attention_mask, input_ids_shape)
            > tf.broadcast_to(question_end_index + 1, input_ids_shape),
            tf.dtypes.int32,
        ) * tf.cast(tf.broadcast_to(attention_mask, input_ids_shape) < input_ids_shape[-1], tf.dtypes.int32)

    return attention_mask
//  [python]
def _get_triplet_mask(labels):
        '''
        得到一个3D的mask [a, p, n], 对应triplet（a, p, n）是valid的位置是True
        ----------------------------------
        Args:
            labels: 对应训练数据的labels, shape = (batch_size,)
        
        Returns:
            mask: 3D,shape = (batch_size, batch_size, batch_size)
        
        '''

        # 初始化一个二维矩阵，坐标(i, j)不相等置为1，得到indices_not_equal
        indices_equal = tf.equal(tf.eye(tf.shape(labels)[0]), 1)
        indices_not_equal = tf.logical_not(indices_equal)
        # 因为最后得到一个3D的mask矩阵(i, j, k)，增加一个维度，则 i_not_equal_j 在第三个维度增加一个即，(batch_size, batch_size, 1), 其他同理
        i_not_equal_j = tf.expand_dims(indices_not_equal, 2)
        i_not_equal_k = tf.expand_dims(indices_not_equal, 1)
        j_not_equal_k = tf.expand_dims(indices_not_equal, 0)
        # 想得到i!=j!=k, 三个不等取and即可, 最后可以得到当下标（i, j, k）不相等时才取True
        distinct_indices = tf.logical_and(tf.logical_and(
            i_not_equal_j, i_not_equal_k), j_not_equal_k)

        # 同样根据labels得到对应i=j, i!=k
        label_equal = tf.equal(tf.expand_dims(labels, 0),
                            tf.expand_dims(labels, 1))
        i_equal_j = tf.expand_dims(label_equal, 2)
        i_equal_k = tf.expand_dims(label_equal, 1)
        valid_labels = tf.logical_and(i_equal_j, tf.logical_not(i_equal_k))
        # mask即为满足上面两个约束，所以两个3D取and
        mask = tf.logical_and(distinct_indices, valid_labels)
        return mask
//  [python]
def _get_triplet_mask(labels):
        '''
        得到一个3D的mask [a, p, n], 对应triplet（a, p, n）是valid的位置是True
        ----------------------------------
        Args:
            labels: 对应训练数据的labels, shape = (batch_size,)
        
        Returns:
            mask: 3D,shape = (batch_size, batch_size, batch_size)
        
        '''

        # 初始化一个二维矩阵，坐标(i, j)不相等置为1，得到indices_not_equal
        indices_equal = tf.cast(tf.eye(tf.shape(labels)[0]), tf.bool)
        indices_not_equal = tf.logical_not(indices_equal)
        # 因为最后得到一个3D的mask矩阵(i, j, k)，增加一个维度，则 i_not_equal_j 在第三个维度增加一个即，(batch_size, batch_size, 1), 其他同理
        i_not_equal_j = tf.expand_dims(indices_not_equal, 2)
        i_not_equal_k = tf.expand_dims(indices_not_equal, 1)
        j_not_equal_k = tf.expand_dims(indices_not_equal, 0)
        # 想得到i!=j!=k, 三个不等取and即可, 最后可以得到当下标（i, j, k）不相等时才取True
        distinct_indices = tf.logical_and(tf.logical_and(
            i_not_equal_j, i_not_equal_k), j_not_equal_k)

        # 同样根据labels得到对应i=j, i!=k
        label_equal = tf.equal(tf.expand_dims(labels, 0),
                            tf.expand_dims(labels, 1))
        i_equal_j = tf.expand_dims(label_equal, 2)
        i_equal_k = tf.expand_dims(label_equal, 1)
        valid_labels = tf.logical_and(i_equal_j, tf.logical_not(i_equal_k))
        # mask即为满足上面两个约束，所以两个3D取and
        mask = tf.logical_and(distinct_indices, valid_labels)
        return mask
//  [python]
def _continue_search(self, state):
    """Return whether to continue the search loop.

    The loops should terminate when
      1) when decode length has been reached, or
      2) when the worst score in the finished sequences is better than the best
         score in the alive sequences (i.e. the finished sequences are provably
         unchanging)

    Args:
      state: A dictionary with the current loop state.

    Returns:
      Bool tensor with value True if loop should continue, False if loop should
      terminate.
    """
    i = state[_StateKeys.CUR_INDEX]
    alive_log_probs = state[_StateKeys.ALIVE_LOG_PROBS]
    finished_scores = state[_StateKeys.FINISHED_SCORES]
    finished_flags = state[_StateKeys.FINISHED_FLAGS]

    not_at_max_decode_length = tf.less(i, self.max_decode_length)

    # Calculate largest length penalty (the larger penalty, the better score).
    max_length_norm = _length_normalization(
        self.alpha, self.max_decode_length, dtype=self.dtype)
    # Get the best possible scores from alive sequences.
    # This tf.slice/tf.squeeze is equivalent to alive_log_probs[:, 0] which
    # emits a tf.strided_slice. tf.slice is easier to reason about as we aren't
    # actually taking a non trivial stride.
    best_alive_scores = tf.squeeze(tf.slice(alive_log_probs, [0, 0], [-1, 1]),
                                   axis=1) / max_length_norm

    # Compute worst score in finished sequences for each batch element
    finished_scores *= t.cast(finished_flags,
                               self.dtype)  # set filler scores to zero
    lowest_finished_scores = tf.reduce_min(finished_scores, axis=1)

    # If there are no finished sequences in a batch element, then set the lowest
    # finished score to -INF for that element.
    finished_batches = tf.reduce_any(finished_flags, 1)
    lowest_finished_scores += ((1.0 - t.cast(finished_batches, self.dtype)) *
                               -inf(self.dtype))

    worst_finished_score_better_than_best_alive_score = tf.reduce_all(
        tf.greater(lowest_finished_scores, best_alive_scores))

    return tf.logical_and(
        not_at_max_decode_length,
        tf.logical_not(worst_finished_score_better_than_best_alive_score))
//  [python]
def _continue_search(self, state):
    """Return whether to continue the search loop.

    The loops should terminate when
      1) when decode length has been reached, or
      2) when the worst score in the finished sequences is better than the best
         score in the alive sequences (i.e. the finished sequences are provably
         unchanging)

    Args:
      state: A dictionary with the current loop state.

    Returns:
      Bool tensor with value True if loop should continue, False if loop should
      terminate.
    """
    i = state[_StateKeys.CUR_INDEX]
    alive_log_probs = state[_StateKeys.ALIVE_LOG_PROBS]
    finished_scores = state[_StateKeys.FINISHED_SCORES]
    finished_flags = state[_StateKeys.FINISHED_FLAGS]

    not_at_max_decode_length = tf.less(i, self.max_decode_length)

    # Calculate largest length penalty (the larger penalty, the better score).
    max_length_norm = _length_normalization(
        self.alpha, self.max_decode_length, dtype=self.dtype)
    # Get the best possible scores from alive sequences.
    # This tf.slice/tf.squeeze is equivalent to alive_log_probs[:, 0] which
    # emits a tf.strided_slice. tf.slice is easier to reason about as we aren't
    # actually taking a non trivial stride.
    best_alive_scores = tf.squeeze(tf.slice(alive_log_probs, [0, 0], [-1, 1]),
                                   axis=1) / max_length_norm

    # Compute worst score in finished sequences for each batch element
    finished_scores *= tf.cast(finished_flags,
                               self.dtype)  # set filler scores to zero
    lowest_finished_scores = tf.reduce_min(finished_scores, axis=1)

    # If there are no finished sequences in a batch element, then set the lowest
    # finished score to -INF for that element.
    finished_batches = tf.reduce_any(finished_flags, 1)
    lowest_finished_scores += ((1.0 - tf.cast(finished_batches, self.dtype)) *
                               -inf(self.dtype))

    worst_finished_score_better_than_best_alive_score = tf.reduce_all(
        tf.greater(lowest_finished_scores, best_alive_scores))

    return tf.logical_and(
        not_at_max_decode_length,
        tf.logical_not(worst_finished_score_better_than_best_alive_score))
//  [python]
def share_bert_layer(self, is_train_place, query_ids, mask_ids, seg_ids, use_bert_pre=1):
        self.bert_config = modeling.BertConfig.from_json_file(
            self.cfg["bert_dir"] + self.cfg["bert_config"])
        bert_model = modeling.BertModel(
            config=self.bert_config,
            is_training=is_train_place,
            input_ids=query_ids,
            input_mask=mask_ids,
            token_type_ids=seg_ids,
            use_one_hot_embeddings=False,
            scope="bert")
        if use_bert_pre:
            tvars = tf.trainable_variables()
            bert_init_dir = self.cfg["bert_dir"] + \
                self.cfg["bert_init_checkpoint"]
            (assignment, initialized_variable_names) = modeling.get_assignment_map_from_checkpoint(tvars,
                                                                                                   bert_init_dir)
            tf.train.init_from_checkpoint(bert_init_dir, assignment)
        bert_output_seq = bert_model.get_sequence_output()

        # 默认使用cls输出
        pooled = bert_model.get_pooled_output()
        embedding_table = bert_model.embedding_table
        input_mask_ = tf.cast(tf.expand_dims(mask_ids, axis=-1), dtype=tf.float32)
        if self.cfg['sentence_embedding_type'] == "avg":
            # 最后一层avg pooling
            pooled = tf.reduce_sum(bert_output_seq * input_mask_, axis=1) / tf.reduce_sum(input_mask_, axis=1)
        elif self.cfg['sentence_embedding_type'].startswith("avg-last-last-"):
            # 使用最后的第n层 avg pooling
            n_last = int(self.cfg['sentence_embedding_type'][-1])
            sequence = bert_model.all_encoder_layers[-n_last] # [batch_size, seq_length, hidden_size]
            pooled = tf.reduce_sum(sequence * input_mask_, axis=1) / tf.reduce_sum(input_mask_, axis=1)
        elif self.cfg['sentence_embedding_type'].startswith("avg-last-"):
            # 使用最后的n层 avg pooling
            pooled = 0
            n_last = int(self.cfg['sentence_embedding_type'][-1])
            for i in range(n_last):
                sequence = bert_model.all_encoder_layers[-i]
                pooled += tf.reduce_sum(sequence * input_mask_, axis=1) / tf.reduce_sum(input_mask_, axis=1)
            pooled /= float(n_last)
        elif self.cfg['sentence_embedding_type'].startswith("avg-last-concat-"):
            pooled = []
            n_last = int(self.cfg['sentence_embedding_type'][-1])
            for i in range(n_last):
                sequence = bert_model.all_encoder_layers[-i]
                pooled += [tf.reduce_sum(sequence * input_mask_, axis=1) / tf.reduce_sum(input_mask_, axis=1)]
            pooled = tf.concat(pooled, axis=-1)
        return pooled, bert_output_seq, embedding_table
//  [python]
def generate_outputs(
      self,
      raw_scores: Dict[str, tf.Tensor],
      raw_boxes: Dict[str, tf.Tensor],
      raw_attributes: Dict[str, Dict[str, tf.Tensor]],
      image_shape: Optional[tf.Tensor] = None,
      anchor_boxes: Optional[Mapping[str, tf.Tensor]] = None,
      generate_detections: bool = False) -> Mapping[str, Any]:
    if not raw_attributes:
      raise ValueError('PointPillars model needs attribute heads.')
    # Clap heading to [-pi, pi]
    if 'heading' in raw_attributes:
      raw_attributes['heading'] = utils.clip_heading(raw_attributes['heading'])

    outputs = {
        'cls_outputs': raw_scores,
        'box_outputs': raw_boxes,
        'attribute_outputs': raw_attributes,
    }
    # Cast raw prediction to float32 for loss calculation.
    outputs = tf.nest.map_structure(lambda x: t.cast(x, tf.float32), outputs)
    if not generate_detections:
      return outputs

    if image_shape is None:
      raise ValueError('Image_shape should not be None for evaluation.')
    if anchor_boxes is None:
      # Generate anchors if needed.
      anchor_boxes = utils.generate_anchors(
          self._min_level,
          self._max_level,
          self._image_size,
          self._anchor_sizes,
      )
      for l in anchor_boxes:
        anchor_boxes[l] = tf.tile(
            tf.expand_dims(anchor_boxes[l], axis=0),
            [tf.shape(image_shape)[0], 1, 1, 1])

    # Generate detected boxes.
    if not self._detection_generator.get_config()['apply_nms']:
      raise ValueError('An NMS algorithm is required for detection generator')
    detections = self._detection_generator(raw_boxes, raw_scores,
                                           anchor_boxes, image_shape,
                                           raw_attributes)
    outputs.update({
        'boxes': detections['detection_boxes'],
        'scores': detections['detection_scores'],
        'classes': detections['detection_classes'],
        'num_detections': detections['num_detections'],
        'attributes': detections['detection_attributes'],
    })
    return outputs
//  [python]
def generate_outputs(
      self,
      raw_scores: Dict[str, tf.Tensor],
      raw_boxes: Dict[str, tf.Tensor],
      raw_attributes: Dict[str, Dict[str, tf.Tensor]],
      image_shape: Optional[tf.Tensor] = None,
      anchor_boxes: Optional[Mapping[str, tf.Tensor]] = None,
      generate_detections: bool = False) -> Mapping[str, Any]:
    if not raw_attributes:
      raise ValueError('PointPillars model needs attribute heads.')
    # Clap heading to [-pi, pi]
    if 'heading' in raw_attributes:
      raw_attributes['heading'] = utils.clip_heading(raw_attributes['heading'])

    outputs = {
        'cls_outputs': raw_scores,
        'box_outputs': raw_boxes,
        'attribute_outputs': raw_attributes,
    }
    # Cast raw prediction to float32 for loss calculation.
    outputs = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), outputs)
    if not generate_detections:
      return outputs

    if image_shape is None:
      raise ValueError('Image_shape should not be None for evaluation.')
    if anchor_boxes is None:
      # Generate anchors if needed.
      anchor_boxes = utils.generate_anchors(
          self._min_level,
          self._max_level,
          self._image_size,
          self._anchor_sizes,
      )
      for l in anchor_boxes:
        anchor_boxes[l] = tf.tile(
            tf.expand_dims(anchor_boxes[l], axis=0),
            [tf.shape(image_shape)[0], 1, 1, 1])

    # Generate detected boxes.
    if not self._detection_generator.get_config()['apply_nms']:
      raise ValueError('An NMS algorithm is required for detection generator')
    detections = self._detection_generator(raw_boxes, raw_scores,
                                           anchor_boxes, image_shape,
                                           raw_attributes)
    outputs.update({
        'boxes': detections['detection_boxes'],
        'scores': detections['detection_scores'],
        'classes': detections['detection_classes'],
        'num_detections': detections['num_detections'],
        'attributes': detections['detection_attributes'],
    })
    return outputs
//  [python]
def create_products_of_givens_rotations(dim, seed):
  r"""Constructs a 2D-tensor which is a product of Givens random rotations.

  Constructs a 2D-tensor of the form G_1 * ... * G_k, where G_i is a Givens
  random rotation. The resulting tensor mimics a matrix taken uniformly at
  random form the orthogonal group.

  Args:
    dim: number of rows/columns of the resulting 2D-tensor.
    seed: random seed.

  Returns:
    The product of Givens random rotations.
  """
  nb_givens_rotations = dim * int(math.ceil(math.log(float(dim))))
  q = np.eye(dim, dim)
  np.random.seed(seed)
  for _ in range(nb_givens_rotations):
    random_angle = math.pi * np.random.uniform()
    random_indices = np.random.choice(dim, 2)
    index_i = min(random_indices[0], random_indices[1])
    index_j = max(random_indices[0], random_indices[1])
    slice_i = q[index_i]
    slice_j = q[index_j]
    new_slice_i = math.cos(random_angle) * slice_i + math.sin(
        random_angle) * slice_j
    new_slice_j = -math.sin(random_angle) * slice_i + math.cos(
        random_angle) * slice_j
    q[index_i] = new_slice_i
    q[index_j] = new_slice_j
  return t.cast(t.constant(q), dtype=t.float32)
//  [python]
def create_products_of_givens_rotations(dim, seed):
  r"""Constructs a 2D-tensor which is a product of Givens random rotations.

  Constructs a 2D-tensor of the form G_1 * ... * G_k, where G_i is a Givens
  random rotation. The resulting tensor mimics a matrix taken uniformly at
  random form the orthogonal group.

  Args:
    dim: number of rows/columns of the resulting 2D-tensor.
    seed: random seed.

  Returns:
    The product of Givens random rotations.
  """
  nb_givens_rotations = dim * int(math.ceil(math.log(float(dim))))
  q = np.eye(dim, dim)
  np.random.seed(seed)
  for _ in range(nb_givens_rotations):
    random_angle = math.pi * np.random.uniform()
    random_indices = np.random.choice(dim, 2)
    index_i = min(random_indices[0], random_indices[1])
    index_j = max(random_indices[0], random_indices[1])
    slice_i = q[index_i]
    slice_j = q[index_j]
    new_slice_i = math.cos(random_angle) * slice_i + math.sin(
        random_angle) * slice_j
    new_slice_j = -math.sin(random_angle) * slice_i + math.cos(
        random_angle) * slice_j
    q[index_i] = new_slice_i
    q[index_j] = new_slice_j
  return tf.cast(tf.constant(q), dtype=tf.float32)
//  [python]
def average_precision_voc12(precision, recall, name=None):
    """Compute (interpolated) average precision from precision and recall Tensors.

    The implementation follows Pascal 2012 and ILSVRC guidelines.
    See also: https://sanchom.wordpress.com/tag/average-precision/
    """
    with tf.name_scope(name, 'average_precision_voc12', [precision, recall]):
        # Convert to float64 to decrease error on Riemann sums.
        precision = tf.cast(precision, dtype=tf.float64)
        recall = tf.cast(recall, dtype=tf.float64)

        # Add bounds values to precision and recall.
        precision = tf.concat([[0.], precision, [0.]], axis=0)
        recall = tf.concat([[0.], recall, [1.]], axis=0)
        # Ensures precision is increasing in reverse order.
        precision = tfe_math.cummax(precision, reverse=True)

        # Riemann sums for estimating the integral.
        # mean_pre = (precision[1:] + precision[:-1]) / 2.
        mean_pre = precision[1:]
        diff_rec = recall[1:] - recall[:-1]
        ap = tf.reduce_sum(mean_pre * diff_rec)
        return ap
//  [python]
def average_precision_voc07(precision, recall, name=None):
    """Compute (interpolated) average precision from precision and recall Tensors.

    The implementation follows Pascal 2007 guidelines.
    See also: https://sanchom.wordpress.com/tag/average-precision/
    """
    with tf.name_scope(name, 'average_precision_voc07', [precision, recall]):
        # Convert to float64 to decrease error on cumulated sums.
        precision = tf.to_float(precision)
        recall = tf.to_float(recall)
        # Add zero-limit value to avoid any boundary problem...
        precision = tf.concat([precision, [0.]], axis=0)
        recall = tf.concat([recall, [np.inf]], axis=0)

        # Split the integral into 10 bins.
        l_aps = []
        for t in np.arange(0., 1.1, 0.1):
            mask = tf.greater_equal(recall, t)
            v = tf.reduce_max(tf.boolean_mask(precision, mask))
            l_aps.append(v / 11.)
        ap = tf.add_n(l_aps)
        return ap
//  [python]
def average_precision_voc07(precision, recall, name=None):
    """Compute (interpolated) average precision from precision and recall Tensors.

    The implementation follows Pascal 2007 guidelines.
    See also: https://sanchom.wordpress.com/tag/average-precision/
    """
    with tf.name_scope(name, 'average_precision_voc07', [precision, recall]):
        # Convert to float64 to decrease error on cumulated sums.
        precision = tf.cast(precision, dtype=tf.float64)
        recall = tf.cast(recall, dtype=tf.float64)
        # Add zero-limit value to avoid any boundary problem...
        precision = tf.concat([precision, [0.]], axis=0)
        recall = tf.concat([recall, [np.inf]], axis=0)

        # Split the integral into 10 bins.
        l_aps = []
        for t in np.arange(0., 1.1, 0.1):
            mask = tf.greater_equal(recall, t)
            v = tf.reduce_max(tf.boolean_mask(precision, mask))
            l_aps.append(v / 11.)
        ap = tf.add_n(l_aps)
        return ap
//  [python]
def num_correct_prediction(logits, labels):
  """Evaluate the quality of the logits at predicting the label.
  Return:
      the number of correct predictions
  """
  correct = tf.equal(tf.arg_max(logits, 1), tf.arg_max(labels, 1))
  correct = tf.to_int32(correct)
  n_correct = tf.reduce_sum(correct)
  return n_correct
//  [python]
def num_correct_prediction(logits, labels):
  """Evaluate the quality of the logits at predicting the label.
  Return:
      the number of correct predictions
  """
  correct = tf.equal(tf.arg_max(logits, 1), tf.arg_max(labels, 1))
  correct = tf.cast(correct, tf.int32)
  n_correct = tf.reduce_sum(correct)
  return n_correct
//  [python]
def accuracy(logits, labels):
  """Evaluate the quality of the logits at predicting the label.
  Args:
    logits: Logits tensor, float - [batch_size, NUM_CLASSES].
    labels: Labels tensor, 
  """
  with tf.name_scope('accuracy') as scope:
      correct = tf.equal(tf.arg_max(logits, 1), tf.arg_max(labels, 1))
      correct = tf.to_float(correct)
      accuracy = tf.reduce_mean(correct)*100.0
      tf.summary.scalar(scope+'/accuracy', accuracy)
  return accuracy
//  [python]
def accuracy(logits, labels):
  """Evaluate the quality of the logits at predicting the label.
  Args:
    logits: Logits tensor, float - [batch_size, NUM_CLASSES].
    labels: Labels tensor, 
  """
  with tf.name_scope('accuracy') as scope:
      correct = tf.equal(tf.arg_max(logits, 1), tf.arg_max(labels, 1))
      correct = tf.cast(correct, tf.float32)
      accuracy = tf.reduce_mean(correct)*100.0
      tf.summary.scalar(scope+'/accuracy', accuracy)
  return accuracy
//  [python]
def test_rot_pos(self):
        """Test 3GPP channel coefficient calculation: Rotate position according
        to orientation"""
        batch_size = TestChannelCoefficientsGenerator.BATCH_SIZE
        orientations = tf.random.normal(shape=[batch_size,3]).numpy()
        positions = tf.random.normal(shape=[batch_size,3, 1]).numpy()

        pos_r_ref = self.rot_pos_ref(orientations, positions)
        pos_r = self.ccg._rot_pos(  t.cast(orientations, tf.float64),
                                    t.cast(positions, tf.float64)).numpy()
        max_err = self.max_rel_err(pos_r_ref, pos_r)
        err_tol = TestChannelCoefficientsGenerator.MAX_ERR
        self.assertLessEqual(max_err, err_tol)
//  [python]
def test_rot_pos(self):
        """Test 3GPP channel coefficient calculation: Rotate position according
        to orientation"""
        batch_size = TestChannelCoefficientsGenerator.BATCH_SIZE
        orientations = tf.random.normal(shape=[batch_size,3]).numpy()
        positions = tf.random.normal(shape=[batch_size,3, 1]).numpy()

        pos_r_ref = self.rot_pos_ref(orientations, positions)
        pos_r = self.ccg._rot_pos(  tf.cast(orientations, tf.float64),
                                    tf.cast(positions, tf.float64)).numpy()
        max_err = self.max_rel_err(pos_r_ref, pos_r)
        err_tol = TestChannelCoefficientsGenerator.MAX_ERR
        self.assertLessEqual(max_err, err_tol)