jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def scaled_dot_product_attention(q, k, v, mask, outputs, name):
    """Calculate the attention weights.
    q, k, v must have matching leading dimensions.
    k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.
    The mask has different shapes depending on its type(padding or look ahead)
    but it must be broadcastable for addition.

    Args:
      q: query shape == (..., seq_len_q, depth)
      k: key shape == (..., seq_len_k, depth)
      v: value shape == (..., seq_len_v, depth_v)
      mask: Float tensor with shape broadcastable
            to (..., seq_len_q, seq_len_k). Defaults to None.

    Returns:
      output, attention_weights
    """

    matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)
    outputs[f'{name}_sat_matmul_qk'] = matmul_qk.numpy()
    
    # scale matmul_qk
    dk = tf.cast(tf.shape(k)[-1], tf.float32)
    scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)
    outputs[f'{name}_sat_logits'] = scaled_attention_logits.numpy()
    
    # add the mask to the scaled tensor.
    if mask is not None:
        scaled_attention_logits += (mask * -1e9)

    outputs[f'{name}_sat_mask'] = scaled_attention_logits.numpy()
    # softmax is normalized on the last axis (seq_len_k) so that the scores
    # add up to 1.
    attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., seq_len_q, seq_len_k)
    outputs[f'{name}_sat_softmax'] = attention_weights.numpy()
    
    output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)
    outputs[f'{name}_sat_matmul_atten'] = output.numpy()
    
    return output, attention_weights
//  [python]
def __call__(self, batch_size, **kwargs):
    """Sample a batch of context."""
    spec = self._context_spec
    contexts = tf.random_uniform(
        shape=[
            batch_size,
        ] + spec.shape.as_list(), dtype=tf.float32)
    contexts = tf.to_float(tf.greater(contexts, self._probs), dtype=spec.dtype)
    return contexts, contexts
//  [python]
def __call__(self, batch_size, **kwargs):
    """Sample a batch of context."""
    spec = self._context_spec
    contexts = tf.random_uniform(
        shape=[
            batch_size,
        ] + spec.shape.as_list(), dtype=tf.float32)
    contexts = tf.cast(tf.greater(contexts, self._probs), dtype=spec.dtype)
    return contexts, contexts
//  [python]
def _quantize(self, inputs, mode):
        if mode == "noise":
            if not self.sub_mean or self.mean is None:
                medians = None
            else:
                medians = self.mean
            return quantize(inputs, medians, method=self.approx, tau=self.tau, tau2=self.tau2, T=self.T, k=self.k)
        
        if self.mean is not None:
            inputs -= self.mean
        outputs = tf.math.round(inputs)
        outputs = t.cast(outputs, self.dtype)
        if mode == "dequantize":
            if self.mean is not None:
                outputs += self.mean
            return outputs
        else:
            assert mode == "symbols", mode
            outputs = t.cast(outputs, tf.int32)
            return outputs
//  [python]
def _quantize(self, inputs, mode):
        if mode == "noise":
            if not self.sub_mean or self.mean is None:
                medians = None
            else:
                medians = self.mean
            return quantize(inputs, medians, method=self.approx, tau=self.tau, tau2=self.tau2, T=self.T, k=self.k)
        
        if self.mean is not None:
            inputs -= self.mean
        outputs = tf.math.round(inputs)
        outputs = tf.cast(outputs, self.dtype)
        if mode == "dequantize":
            if self.mean is not None:
                outputs += self.mean
            return outputs
        else:
            assert mode == "symbols", mode
            outputs = tf.cast(outputs, tf.int32)
            return outputs
//  [python]
def _quantize(self, inputs, mode):
        _, _, _, input_slices = self._get_input_dims()
        medians = self._medians[input_slices]

        if mode == "noise":
            if not self.sub_mean:
                medians = None
            if self.approx == "DSl-Q":
                assert medians is None
                self.quantizer(inputs)
            else:
                return quantize(inputs, medians, method=self.approx, tau=self.tau, tau2=self.tau2, T=self.T, k=self.k)

        outputs = tf.math.round(inputs - medians)
        outputs = t.cast(outputs, self.dtype)
        if mode == "dequantize":
            return outputs + medians
        else:
            assert mode == "symbols", mode
            outputs = t.cast(outputs, tf.int32)
            return outputs
//  [python]
def _quantize(self, inputs, mode):
        _, _, _, input_slices = self._get_input_dims()
        medians = self._medians[input_slices]

        if mode == "noise":
            if not self.sub_mean:
                medians = None
            if self.approx == "DSl-Q":
                assert medians is None
                self.quantizer(inputs)
            else:
                return quantize(inputs, medians, method=self.approx, tau=self.tau, tau2=self.tau2, T=self.T, k=self.k)

        outputs = tf.math.round(inputs - medians)
        outputs = tf.cast(outputs, self.dtype)
        if mode == "dequantize":
            return outputs + medians
        else:
            assert mode == "symbols", mode
            outputs = tf.cast(outputs, tf.int32)
            return outputs
//  [python]
def evaluate(sentence):
    sentence[0] = '<SOS> ' + sentence[0] + ' <EOS>'
    sentence = pad(token_English.texts_to_sequences(sentence) , length = len(English[0]))
    sentence = tf.convert_to_tensor(np.array(sentence),dtype=tf.int64)

    decoder_input = tf.convert_to_tensor(np.array(token_Japanese.texts_to_sequences(['SOS'])),dtype=tf.int64)

    outputs = None
    for i in range(len(English[0])):
        enc_padding_mask, combined_mask, dec_padding_mask = create_masks(
            sentence, decoder_input)
        #tf.print(enc_padding_mask.shape, combined_mask.shape, dec_padding_mask.shape)
        # predictions.shape == (batch_size, seq_len, vocab_size)
        predictions, attention_weights, out = transformer([sentence,
                                                     decoder_input,
                                                     enc_padding_mask,
                                                     combined_mask,
                                                     dec_padding_mask],
                                                     False)

        # select the last word from the seq_len dimension
        predictions = predictions[: ,-1:, :]  # (batch_size, 1, vocab_size)

        predicted_id = tf.to_int64(tf.argmax(predictions, axis=-1))

        #tf.print(predicted_id,decoder_input)
        # return the result if the predicted_id is equal to the end token
        if predicted_id == token_Japanese.texts_to_sequences(['EOS']):
            out['input_sentence'] = sentence.numpy()
            return tf.squeeze(decoder_input, axis=0), attention_weights , sentence, out

        # concatentate the predicted_id to the output which is given to the decoder
        # as its input.
        decoder_input = tf.concat([decoder_input, predicted_id], axis=1)
        outputs = out

    outputs['input_sentence'] = sentence.numpy()
    return tf.squeeze(decoder_input, axis=0), attention_weights, sentence, outputs
//  [python]
def evaluate(sentence):
    sentence[0] = '<SOS> ' + sentence[0] + ' <EOS>'
    sentence = pad(token_English.texts_to_sequences(sentence) , length = len(English[0]))
    sentence = tf.convert_to_tensor(np.array(sentence),dtype=tf.int64)

    decoder_input = tf.convert_to_tensor(np.array(token_Japanese.texts_to_sequences(['SOS'])),dtype=tf.int64)

    outputs = None
    for i in range(len(English[0])):
        enc_padding_mask, combined_mask, dec_padding_mask = create_masks(
            sentence, decoder_input)
        #tf.print(enc_padding_mask.shape, combined_mask.shape, dec_padding_mask.shape)
        # predictions.shape == (batch_size, seq_len, vocab_size)
        predictions, attention_weights, out = transformer([sentence,
                                                     decoder_input,
                                                     enc_padding_mask,
                                                     combined_mask,
                                                     dec_padding_mask],
                                                     False)

        # select the last word from the seq_len dimension
        predictions = predictions[: ,-1:, :]  # (batch_size, 1, vocab_size)

        predicted_id = tf.cast(tf.argmax(predictions, axis=-1), tf.int64)

        #tf.print(predicted_id,decoder_input)
        # return the result if the predicted_id is equal to the end token
        if predicted_id == token_Japanese.texts_to_sequences(['EOS']):
            out['input_sentence'] = sentence.numpy()
            return tf.squeeze(decoder_input, axis=0), attention_weights , sentence, out

        # concatentate the predicted_id to the output which is given to the decoder
        # as its input.
        decoder_input = tf.concat([decoder_input, predicted_id], axis=1)
        outputs = out

    outputs['input_sentence'] = sentence.numpy()
    return tf.squeeze(decoder_input, axis=0), attention_weights, sentence, outputs
//  [python]
def __call__(self, batch_size, **kwargs):
    """Sample a batch of context.

    Args:
      batch_size: Batch size.
    Returns:
      Two [batch_size, num_context_dims] tensors.
    """
    if self._maxval == self._minval:
      contexts = tf.constant(
          self._maxval, shape=[batch_size, 1], dtype=tf.int32)
    else:
      contexts = tf.random_uniform(
          shape=[batch_size, 1],
          dtype=tf.int32,
          maxval=self._maxval,
          minval=self._minval)
    next_contexts = tf.maximum(contexts + self._timestep, 0)

    return tf.to_int32(
        contexts), tf.to_int32(
            next_contexts)
//  [python]
def __call__(self, batch_size, **kwargs):
    """Sample a batch of context.

    Args:
      batch_size: Batch size.
    Returns:
      Two [batch_size, num_context_dims] tensors.
    """
    if self._maxval == self._minval:
      contexts = tf.constant(
          self._maxval, shape=[batch_size, 1], dtype=tf.int32)
    else:
      contexts = tf.random_uniform(
          shape=[batch_size, 1],
          dtype=tf.int32,
          maxval=self._maxval,
          minval=self._minval)
    next_contexts = tf.maximum(contexts + self._timestep, 0)

    return tf.cast(
        contexts, dtype=self._context_spec.dtype), tf.cast(
            next_contexts, dtype=self._context_spec.dtype)
//  [python]
def _get_current_sparsities(self):
        current_sparsities = []
        for i in range(len(self._all_weight_variable_sorted)):
            weight = self._all_weight_variable_sorted[i]
            weight_abs = tf.abs(weight)
            sparsity = tf.reduce_mean(tf.to_float(weight_abs < 1e-8)) * tf.constant([1.])
            current_sparsities.append(sparsity)
        current_sparsities = tf.concat(current_sparsities, axis=0)

        return current_sparsities
//  [python]
def _get_current_sparsities(self):
        current_sparsities = []
        for i in range(len(self._all_weight_variable_sorted)):
            weight = self._all_weight_variable_sorted[i]
            weight_abs = tf.abs(weight)
            sparsity = tf.reduce_mean(tf.cast(weight_abs < 1e-8, tf.float32)) * tf.constant([1.])
            current_sparsities.append(sparsity)
        current_sparsities = tf.concat(current_sparsities, axis=0)

        return current_sparsities
//  [python]
def _decode(self, record: tf.Tensor):
    """Decodes a serialized tf.Example."""
    name_to_features = {
        self._params.src_lang: tf.io.FixedLenFeature([], tf.string),
        self._params.tgt_lang: tf.io.FixedLenFeature([], tf.string),
    }
    if self._params.has_unique_id:
      name_to_features['unique_id'] = tf.io.FixedLenFeature([], tf.int64)
    example = tf.io.parse_single_example(record, name_to_features)
    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in example:
      t = example[name]
      if t.dtype == tf.int64:
        t = t.astype(tf.int32)
      example[name] = t
    return example
//  [python]
def _decode(self, record: tf.Tensor):
    """Decodes a serialized tf.Example."""
    name_to_features = {
        self._params.src_lang: tf.io.FixedLenFeature([], tf.string),
        self._params.tgt_lang: tf.io.FixedLenFeature([], tf.string),
    }
    if self._params.has_unique_id:
      name_to_features['unique_id'] = tf.io.FixedLenFeature([], tf.int64)
    example = tf.io.parse_single_example(record, name_to_features)
    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in example:
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      example[name] = t
    return example
//  [python]
def _generate_level_pruner_masks(self):
        apply_masks = []
        for i in range(len(self._all_weight_variable_sorted)):
            v = self._all_weight_variable_sorted[i]
            size = self._sorted_weight_variable_sizes[i]
            sparsity = self._perturbed_sparsities[i]
            
            sparse_num = size * sparsity
            sparse_num = tf.to_int64(sparse_num)
            weight_abs = tf.abs(v)
            weight_abs_sorted = tf.contrib.framework.sort(tf.reshape(weight_abs, [-1]))
            threshold = weight_abs_sorted[sparse_num]
            mask = tf.to_float(tf.cast(weight_abs >= threshold, tf.int32))
            if self._debug:
                print_op = tf.print("mask sparsity:", 1-tf.reduce_mean(mask), sparsity)
                with tf.control_dependencies([print_op]):
                    mask = tf.identity(mask)
            apply_mask = tf.assign(v, v * mask)
            apply_masks.append(apply_mask)
        
        return apply_masks
//  [python]
def _generate_level_pruner_masks(self):
        apply_masks = []
        for i in range(len(self._all_weight_variable_sorted)):
            v = self._all_weight_variable_sorted[i]
            size = self._sorted_weight_variable_sizes[i]
            sparsity = self._perturbed_sparsities[i]
            
            sparse_num = size * sparsity
            sparse_num = tf.cast(sparse_num, tf.int64)
            weight_abs = tf.abs(v)
            weight_abs_sorted = tf.contrib.framework.sort(tf.reshape(weight_abs, [-1]))
            threshold = weight_abs_sorted[sparse_num]
            mask = tf.cast(weight_abs >= threshold, tf.float32)
            if self._debug:
                print_op = tf.print("mask sparsity:", 1-tf.reduce_mean(mask), sparsity)
                with tf.control_dependencies([print_op]):
                    mask = tf.identity(mask)
            apply_mask = tf.assign(v, v * mask)
            apply_masks.append(apply_mask)
        
        return apply_masks
//  [python]
def _set_values_using_indicator(self, x, indicator, val):
    """Set the indicated fields of x to val.

    Args:
      x: tensor.
      indicator: boolean with same shape as x.
      val: scalar with value to set.

    Returns:
      modified tensor.
    """
    indicator = t.cast(indicator, x.dtype)
    return tf.add(tf.multiply(x, 1 - indicator), val * indicator)
//  [python]
def _set_values_using_indicator(self, x, indicator, val):
    """Set the indicated fields of x to val.

    Args:
      x: tensor.
      indicator: boolean with same shape as x.
      val: scalar with value to set.

    Returns:
      modified tensor.
    """
    indicator = tf.cast(indicator, x.dtype)
    return tf.add(tf.multiply(x, 1 - indicator), val * indicator)
//  [python]
def _generate_perturbed_sparsities(self):
        magnitude = self._current_temperature / self._start_temperature * self._perturbation_magnitude
        perturbation = tf.random.uniform([len(self._all_weight_variable_sorted)], -magnitude, magnitude)
        if self._debug:
            print_op = tf.print("perturbation, sparsities, sparsities + perturbation:", perturbation, self._sparsities, self._sparsities + perturbation)
            with tf.control_dependencies([print_op]):
                perturbation = tf.identity(perturbation)
        sparsities = tf.clip_by_value(self._sparsities + perturbation, 0, tf.reduce_max(self._sparsities + perturbation))
        sparsities = tf.contrib.framework.sort(sparsities)

        sparse_num = tf.reduce_sum(sparsities * self._sorted_weight_variable_sizes)
        overall_sparsity = sparse_num / tf.reduce_sum(self._sorted_weight_variable_sizes)
        scale_factor = self._target_sparsity / overall_sparsity
        perturbed_sparsities = scale_factor * sparsities

        if self._debug:
            print_op = tf.print("after sparsity rescale:", perturbed_sparsities)
            with tf.control_dependencies([print_op]):
                perturbed_sparsities = tf.identity(perturbed_sparsities)

        perturbed_sparsities_mask = tf.to_float(perturbed_sparsities < 1.0)
        perturbed_sparsities = perturbed_sparsities * perturbed_sparsities_mask + (1.0 - perturbed_sparsities_mask) * self._target_sparsity

        if self._debug:
            print_op = tf.print("before, after perturbation:", self._sparsities, perturbed_sparsities)
            with tf.control_dependencies([print_op]):
                return tf.identity(perturbed_sparsities)
        else:
            return perturbed_sparsities
//  [python]
def _generate_perturbed_sparsities(self):
        magnitude = self._current_temperature / self._start_temperature * self._perturbation_magnitude
        perturbation = tf.random.uniform([len(self._all_weight_variable_sorted)], -magnitude, magnitude)
        if self._debug:
            print_op = tf.print("perturbation, sparsities, sparsities + perturbation:", perturbation, self._sparsities, self._sparsities + perturbation)
            with tf.control_dependencies([print_op]):
                perturbation = tf.identity(perturbation)
        sparsities = tf.clip_by_value(self._sparsities + perturbation, 0, tf.reduce_max(self._sparsities + perturbation))
        sparsities = tf.contrib.framework.sort(sparsities)

        sparse_num = tf.reduce_sum(sparsities * self._sorted_weight_variable_sizes)
        overall_sparsity = sparse_num / tf.reduce_sum(self._sorted_weight_variable_sizes)
        scale_factor = self._target_sparsity / overall_sparsity
        perturbed_sparsities = scale_factor * sparsities

        if self._debug:
            print_op = tf.print("after sparsity rescale:", perturbed_sparsities)
            with tf.control_dependencies([print_op]):
                perturbed_sparsities = tf.identity(perturbed_sparsities)

        perturbed_sparsities_mask = tf.cast(perturbed_sparsities < 1.0, tf.float32)
        perturbed_sparsities = perturbed_sparsities * perturbed_sparsities_mask + (1.0 - perturbed_sparsities_mask) * self._target_sparsity

        if self._debug:
            print_op = tf.print("before, after perturbation:", self._sparsities, perturbed_sparsities)
            with tf.control_dependencies([print_op]):
                return tf.identity(perturbed_sparsities)
        else:
            return perturbed_sparsities
//  [python]
def __call__(self, inputs):
        # Split string into unicode characters
        characters = t.cast(tf.io.decode_raw(inputs, out_type=tf.uint8), tf.uint32)
        # Apply hashing function recursively to each character of string
        h = tf.foldl(
            lambda a, x: tf.bitwise.bitwise_xor(a, x)
            * tf.constant(16777619, dtype="uint32"),
            characters,
            initializer=tf.constant(2166136261, dtype="uint32"),
        )
        return h
//  [python]
def __call__(self, inputs):
        # Split string into unicode characters
        characters = tf.cast(tf.io.decode_raw(inputs, out_type=tf.uint8), tf.uint32)
        # Apply hashing function recursively to each character of string
        h = tf.foldl(
            lambda a, x: tf.bitwise.bitwise_xor(a, x)
            * tf.constant(16777619, dtype="uint32"),
            characters,
            initializer=tf.constant(2166136261, dtype="uint32"),
        )
        return h
//  [python]
def hard_mining_mse(k):
    """
    Compute MSE for steering evaluation and hard-mining for the current batch.

    # Arguments
        k: number of samples for hard-mining.

    # Returns
        custom_mse: average MSE for the current batch.
    """

    def custom_mse(y_true, y_pred):
        # Parameter t indicates the type of experiment
        t = y_true[:,0]

        # Number of steering samples
        samples_steer = t.equal(1)
        n_samples_steer = tf.reduce_sum(tf.cast(samples_steer, tf.int32))

        if n_samples_steer == 0:
            return 0.0
        else:
            # Predicted and real steerings
            pred_steer = tf.squeeze(y_pred, squeeze_dims=-1)
            true_steer = y_true[:,1]

            # Steering loss
            l_steer = tf.multiply(t, K.square(pred_steer - true_steer))

            # Hard mining
            k_min = tf.minimum(k, n_samples_steer)
            _, indices = tf.nn.top_k(l_steer, k=k_min)
            max_l_steer = tf.gather(l_steer, indices)
            hard_l_steer = tf.divide(tf.reduce_sum(max_l_steer), tf.cast(k,tf.float32))

            return hard_l_steer

    return custom_mse
//  [python]
def hard_mining_mse(k):
    """
    Compute MSE for steering evaluation and hard-mining for the current batch.

    # Arguments
        k: number of samples for hard-mining.

    # Returns
        custom_mse: average MSE for the current batch.
    """

    def custom_mse(y_true, y_pred):
        # Parameter t indicates the type of experiment
        t = y_true[:,0]

        # Number of steering samples
        samples_steer = tf.cast(tf.equal(t,1), tf.int32)
        n_samples_steer = tf.reduce_sum(samples_steer)

        if n_samples_steer == 0:
            return 0.0
        else:
            # Predicted and real steerings
            pred_steer = tf.squeeze(y_pred, squeeze_dims=-1)
            true_steer = y_true[:,1]

            # Steering loss
            l_steer = tf.multiply(t, K.square(pred_steer - true_steer))

            # Hard mining
            k_min = tf.minimum(k, n_samples_steer)
            _, indices = tf.nn.top_k(l_steer, k=k_min)
            max_l_steer = tf.gather(l_steer, indices)
            hard_l_steer = tf.divide(tf.reduce_sum(max_l_steer), tf.cast(k,tf.float32))

            return hard_l_steer

    return custom_mse
//  [python]
def hard_mining_entropy(k):
    """
    Compute binary cross-entropy for collision evaluation and hard-mining.

    # Arguments
        k: Number of samples for hard-mining.

    # Returns
        custom_bin_crossentropy: average binary cross-entropy for the current batch.
    """

    def custom_bin_crossentropy(y_true, y_pred):
        # Parameter t indicates the type of experiment
        t = y_true[:,0]

        # Number of collision samples
        samples_coll = t.equal(0)
        n_samples_coll = tf.reduce_sum(tf.cast(samples_coll, tf.int32))

        if n_samples_coll == 0:
            return 0.0
        else:
            # Predicted and real labels
            pred_coll = tf.squeeze(y_pred, squeeze_dims=-1)
            true_coll = y_true[:,1]

            # Collision loss
            l_coll = tf.multiply((1-t), K.binary_crossentropy(true_coll, pred_coll))

            # Hard mining
            k_min = tf.minimum(k, n_samples_coll)
            _, indices = tf.nn.top_k(l_coll, k=k_min)
            max_l_coll = tf.gather(l_coll, indices)
            hard_l_coll = tf.divide(tf.reduce_sum(max_l_coll), tf.cast(k, tf.float32))

            return hard_l_coll

    return custom_bin_crossentropy
//  [python]
def hard_mining_entropy(k):
    """
    Compute binary cross-entropy for collision evaluation and hard-mining.

    # Arguments
        k: Number of samples for hard-mining.

    # Returns
        custom_bin_crossentropy: average binary cross-entropy for the current batch.
    """

    def custom_bin_crossentropy(y_true, y_pred):
        # Parameter t indicates the type of experiment
        t = y_true[:,0]

        # Number of collision samples
        samples_coll = tf.cast(tf.equal(t,0), tf.int32)
        n_samples_coll = tf.reduce_sum(samples_coll)

        if n_samples_coll == 0:
            return 0.0
        else:
            # Predicted and real labels
            pred_coll = tf.squeeze(y_pred, squeeze_dims=-1)
            true_coll = y_true[:,1]

            # Collision loss
            l_coll = tf.multiply((1-t), K.binary_crossentropy(true_coll, pred_coll))

            # Hard mining
            k_min = tf.minimum(k, n_samples_coll)
            _, indices = tf.nn.top_k(l_coll, k=k_min)
            max_l_coll = tf.gather(l_coll, indices)
            hard_l_coll = tf.divide(tf.reduce_sum(max_l_coll), tf.cast(k, tf.float32))

            return hard_l_coll

    return custom_bin_crossentropy
//  [python]
def call(self, inputs_embeds, attention_mask, training=False):
    """Forward pass for the PositionlessTFGPT2MainLayer.

    Args:
      inputs_embeds: A tf.Tensor of input embeddings.
      attention_mask: A tf.Tensor of booleans indicating which inputs can be
        attended to by the layer.
      training: A bool for whether training or evaluating.

    Returns:
      hidden_states: A tf.Tensor representing the final hidden states.
    """
    attention_mask = attention_mask[:, tf.newaxis, tf.newaxis, :]

    # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
    # masked positions, this operation will create a tensor which is 0.0 for
    # positions we want to attend and -10000.0 for masked positions.
    # Since we are adding it to the raw scores before the softmax, this is
    # effectively the same as removing these entirely.
    attention_mask = tf.cast(attention_mask, tf.float32)
    attention_mask = (1.0 - attention_mask) * -10000.0

    hidden_states = inputs_embeds
    hidden_states = self._drop(hidden_states, training=training)

    input_shape = modeling_tf_utils.shape_list(inputs_embeds)[:-1]
    output_shape = input_shape + [
        modeling_tf_utils.shape_list(hidden_states)[-1]
    ]

    for block in self._h:
      outputs = block(
          hidden_states,
          layer_past=None,
          encoder_attention_mask=attention_mask,
          head_mask=None,
          encoder_hidden_states=None,
          encoder_attention_mask=None,
          use_cache=False,
          output_attentions=False,
          training=training)
      hidden_states = outputs[0]

    hidden_states = self._ln_f(hidden_states)
    hidden_states = tf.reshape(hidden_states, output_shape)

    return hidden_states
//  [python]
def call(self, inputs):
        # Get the position in the vector for largest probability
        pred = tf.to_int32(tf.math.argmax(inputs, axis=-1))
        # Lookup the vector position corresponding label
        return self.table.lookup(pred)
//  [python]
def call(self, inputs):
        # Get the position in the vector for largest probability
        pred = tf.cast(tf.math.argmax(inputs, axis=-1), tf.int32)
        # Lookup the vector position corresponding label
        return self.table.lookup(pred)
//  [python]
def __call__(self, inputs):
        words = tf.strings.split(inputs)
        # Map every wordngram tensor to a hashes tensor
        hashes = tf.map_fn(
            HashWord(),
            words,
            fn_output_signature=tf.uint32,
        )
        hashes = tf.to_int32(tf.to_int32(hashes))
        hashes = tf.to_int64(hashes)
        # Recursively hash again each word hash
        h = tf.foldl(
            lambda a, x: tf.to_int64(
                tf.constant(116049371, dtype="uint64") * a
                + tf.to_int64(tf.to_int32(x))
            ),
            hashes[1:],
            initializer=tf.gather(hashes, 0),
        )
        # Modulo by bucket size
        wordngram_index = tf.math.floormod(
            h, tf.constant(self.bucket_size, dtype="uint64")
        )
        # Add vocabulary length to find index in input_matrix
        return wordngram_index + self.vocab_size
//  [python]
def __call__(self, inputs):
        words = tf.strings.split(inputs)
        # Map every wordngram tensor to a hashes tensor
        hashes = tf.map_fn(
            HashWord(),
            words,
            fn_output_signature=tf.uint32,
        )
        hashes = tf.cast(tf.cast(hashes, "int32"), "uint64")
        # Recursively hash again each word hash
        h = tf.foldl(
            lambda a, x: tf.cast(
                tf.constant(116049371, dtype="uint64") * a
                + tf.cast(tf.cast(x, "int32"), "uint64"),
                "uint64",
            ),
            hashes[1:],
            initializer=tf.gather(hashes, 0),
        )
        # Modulo by bucket size
        wordngram_index = tf.math.floormod(
            h, tf.constant(self.bucket_size, dtype="uint64")
        )
        # Add vocabulary length to find index in input_matrix
        return wordngram_index + self.vocab_size
//  [python]
def call(self, inputs):
        # Create ragged tensor with words per sentence
        sentence_tokens = tf.strings.split(inputs)
        # Make ragged tensor of word ngrams per sentence
        sentence_word_ngrams = tf.strings.ngrams(
            sentence_tokens, [i for i in range(2, self.max_ngram_length + 1)]
        )
        # Get the indexes for each of the ngrams
        ngram_indexes = tf.ragged.map_flat_values(
            GetWordNgramIndexes(self.vocab_size, self.bucket_size),
            sentence_word_ngrams,
        )
        return t.cast(ngram_indexes, "int64")
//  [python]
def call(self, inputs):
        # Create ragged tensor with words per sentence
        sentence_tokens = tf.strings.split(inputs)
        # Make ragged tensor of word ngrams per sentence
        sentence_word_ngrams = tf.strings.ngrams(
            sentence_tokens, [i for i in range(2, self.max_ngram_length + 1)]
        )
        # Get the indexes for each of the ngrams
        ngram_indexes = tf.ragged.map_flat_values(
            GetWordNgramIndexes(self.vocab_size, self.bucket_size),
            sentence_word_ngrams,
        )
        return tf.cast(ngram_indexes, "int64")
//  [python]
def call(self, inputs):
        # Split sentences into a ragged tensor with words per sentence
        words = tf.strings.split(inputs)
        # Add BOW and EOW chars to all words
        words_ = tf.ragged.map_flat_values(
            tf.strings.reduce_join,
            [
                tf.broadcast_to(self.bow, shape=(tf.size(words),)),
                words,
                tf.broadcast_to(self.eow, shape=(tf.size(words),)),
            ],
            separator="",
            axis=0,
        )
        # Make ngrams of characters
        subwords = tf.strings.ngrams(
            tf.strings.unicode_split(words_, "UTF-8"),  # Split into characters
            [i for i in range(self.minn, self.maxn + 1)],  # For all sizes
            separator="",
        )
        # Map each ngram to an index of input_matrix
        subwords_indexes = tf.ragged.map_flat_values(
            GetSubwordsIndexes(self.vocab_size, self.bucket_size), subwords
        )
        # Cast needed for embedding layer
        return tf.to_int64(subwords_indexes)
//  [python]
def call(self, inputs):
        # Split sentences into a ragged tensor with words per sentence
        words = tf.strings.split(inputs)
        # Add BOW and EOW chars to all words
        words_ = tf.ragged.map_flat_values(
            tf.strings.reduce_join,
            [
                tf.broadcast_to(self.bow, shape=(tf.size(words),)),
                words,
                tf.broadcast_to(self.eow, shape=(tf.size(words),)),
            ],
            separator="",
            axis=0,
        )
        # Make ngrams of characters
        subwords = tf.strings.ngrams(
            tf.strings.unicode_split(words_, "UTF-8"),  # Split into characters
            [i for i in range(self.minn, self.maxn + 1)],  # For all sizes
            separator="",
        )
        # Map each ngram to an index of input_matrix
        subwords_indexes = tf.ragged.map_flat_values(
            GetSubwordsIndexes(self.vocab_size, self.bucket_size), subwords
        )
        # Cast needed for embedding layer
        return tf.cast(subwords_indexes, dtype="int64")
//  [python]
def call(
        self, pixel_values: tf.Tensor, bool_masked_pos: Optional[tf.Tensor] = None, training: bool = False
    ) -> tf.Tensor:
        embeddings = self.patch_embeddings(pixel_values)
        batch_size, seq_length, _ = shape_list(embeddings)

        if bool_masked_pos is not None:
            mask_tokens = tf.tile(self.mask_token, [batch_size, seq_length, 1])
            # replace the masked visual tokens by mask_tokens
            mask = tf.expand_dims(bool_masked_pos, axis=-1)
            mask = tf.cast(mask, dtype=mask_tokens.dtype)
            embeddings = embeddings * (1.0 - mask) + mask_tokens * mask

        cls_tokens = tf.repeat(self.cls_token, repeats=batch_size, axis=0)
        distillation_tokens = tf.repeat(self.distillation_token, repeats=batch_size, axis=0)
        embeddings = tf.concat((cls_tokens, distillation_tokens, embeddings), axis=1)
        embeddings = embeddings + self.position_embeddings
        embeddings = self.dropout(embeddings, training=training)
        return embeddings
//  [python]
def _is_box_near_crop_edge_tf(boxes, crop_box, orig_box, atol=20.0):
    """Filter masks at the edge of a crop, but not at the edge of the original image."""
    crop_box_tf = tf.convert_to_tensor(crop_box, dtype=tf.float32)
    orig_box_tf = tf.convert_to_tensor(orig_box, dtype=tf.float32)

    left, top, _, _ = crop_box
    offset = tf.convert_to_tensor([[left, top, left, top]])
    # Check if boxes has a channel dimension
    if len(boxes.shape) == 3:
        offset = tf.expand_dims(offset, 1)
    boxes = t.cast(boxes + offset, tf.float32)

    near_crop_edge = tnp.isclose(boxes, crop_box_tf[None, :], atol=atol, rtol=0)
    near_image_edge = tnp.isclose(boxes, orig_box_tf[None, :], atol=atol, rtol=0)
    near_crop_edge = tf.math.logical_and(near_crop_edge, ~near_image_edge)
    return tf.reduce_any(near_crop_edge, axis=1)
//  [python]
def _is_box_near_crop_edge_tf(boxes, crop_box, orig_box, atol=20.0):
    """Filter masks at the edge of a crop, but not at the edge of the original image."""
    crop_box_tf = tf.convert_to_tensor(crop_box, dtype=tf.float32)
    orig_box_tf = tf.convert_to_tensor(orig_box, dtype=tf.float32)

    left, top, _, _ = crop_box
    offset = tf.convert_to_tensor([[left, top, left, top]])
    # Check if boxes has a channel dimension
    if len(boxes.shape) == 3:
        offset = tf.expand_dims(offset, 1)
    boxes = tf.cast(boxes + offset, tf.float32)

    near_crop_edge = tnp.isclose(boxes, crop_box_tf[None, :], atol=atol, rtol=0)
    near_image_edge = tnp.isclose(boxes, orig_box_tf[None, :], atol=atol, rtol=0)
    near_crop_edge = tf.math.logical_and(near_crop_edge, ~near_image_edge)
    return tf.reduce_any(near_crop_edge, axis=1)
//  [python]
def _get_seq_err_ops( batch_num_sequence_errors, label_length ):
    """Calculates the sequence error by accumulating for batches and returns
    the average"""

    var_collections=[tf.GraphKeys.LOCAL_VARIABLES]

    # Variables to tally across batches (all initially zero)
    total_num_sequence_errors =  tf.Variable( 0, trainable=False,
                                              name='total_num_sequence_errors',
                                              dtype=tf.int64,
                                              collections=var_collections )

    total_num_sequences =  tf.Variable( 0, trainable=False,
                                        name='total_num_sequences',
                                        dtype=tf.int64,
                                        collections=var_collections )

    # Get the batch size and cast it appropriately
    batch_size = tf.shape( label_length )[0]
    batch_size = tf.to_int64( batch_size )

    # Create the "+=" update ops and group together as one
    update_sequence_errors = tf.assign_add( total_num_sequence_errors,
                                            batch_num_sequence_errors )
    update_num_sequences   = tf.assign_add( total_num_sequences,
                                            batch_size )

    update_op = tf.group(update_sequence_errors, update_num_sequences)

    # Get the average sequence error across all inputs
    sequence_error = tf.truediv( total_num_sequence_errors,
                                 total_num_sequences,
                                 name='sequence_error' )
    
    return sequence_error, update_op, total_num_sequence_errors,\
        total_num_sequences
//  [python]
def _get_seq_err_ops( batch_num_sequence_errors, label_length ):
    """Calculates the sequence error by accumulating for batches and returns
    the average"""

    var_collections=[tf.GraphKeys.LOCAL_VARIABLES]

    # Variables to tally across batches (all initially zero)
    total_num_sequence_errors =  tf.Variable( 0, trainable=False,
                                              name='total_num_sequence_errors',
                                              dtype=tf.int64,
                                              collections=var_collections )

    total_num_sequences =  tf.Variable( 0, trainable=False,
                                        name='total_num_sequences',
                                        dtype=tf.int64,
                                        collections=var_collections )

    # Get the batch size and cast it appropriately
    batch_size = tf.shape( label_length )[0]
    batch_size = tf.cast( batch_size, tf.int64 )

    # Create the "+=" update ops and group together as one
    update_sequence_errors = tf.assign_add( total_num_sequence_errors,
                                            batch_num_sequence_errors )
    update_num_sequences   = tf.assign_add( total_num_sequences,
                                            batch_size )

    update_op = tf.group(update_sequence_errors, update_num_sequences)

    # Get the average sequence error across all inputs
    sequence_error = tf.truediv( total_num_sequence_errors,
                                 total_num_sequences,
                                 name='sequence_error' )
    
    return sequence_error, update_op, total_num_sequence_errors,\
        total_num_sequences
//  [python]
def _get_testing( rnn_logits,sequence_length,label,label_length,
                  continuous_eval, lexicon, lexicon_prior ):
    """Create ops for testing (all scalars): 
       loss: CTC loss function value, 
       label_error:   batch level edit distance on beam search max
       sequence_error: batch level sequence error rate
    """

    with tf.name_scope( "train" ):
        # Reduce by mean (rather than sum) if doing continuous evaluation
        batch_loss = model.ctc_loss_layer( rnn_logits,label,sequence_length,
                                           reduce_mean=continuous_eval) 
    with tf.name_scope( "test" ):
        predictions,_ = _get_output( rnn_logits, sequence_length,
                                     lexicon, lexicon_prior )

        hypothesis = tf.cast( predictions[0], tf.int32 ) # for edit_distance

        # Per-sequence statistic
        num_label_errors = tf.edit_distance( hypothesis, label, 
                                             normalize=False )

        # Per-batch summary counts
        batch_num_label_errors = tf.reduce_sum( num_label_errors)
        batch_num_sequence_errors = tf.count_nonzero( num_label_errors, axis=0 )
        batch_num_labels = tf.reduce_sum( label_length )
        
        # Wide integer type casts (prefer unsigned, but truediv dislikes those)
        batch_num_label_errors = tf.to_int64( batch_num_label_errors )
        batch_num_sequence_errors = tf.to_int64( batch_num_sequence_errors )
        batch_num_labels = tf.to_int64( batch_num_labels)
        
    return batch_loss, batch_num_label_errors, batch_num_sequence_errors, \
        batch_num_labels, predictions
//  [python]
def _get_testing( rnn_logits,sequence_length,label,label_length,
                  continuous_eval, lexicon, lexicon_prior ):
    """Create ops for testing (all scalars): 
       loss: CTC loss function value, 
       label_error:   batch level edit distance on beam search max
       sequence_error: batch level sequence error rate
    """

    with tf.name_scope( "train" ):
        # Reduce by mean (rather than sum) if doing continuous evaluation
        batch_loss = model.ctc_loss_layer( rnn_logits,label,sequence_length,
                                           reduce_mean=continuous_eval) 
    with tf.name_scope( "test" ):
        predictions,_ = _get_output( rnn_logits, sequence_length,
                                     lexicon, lexicon_prior )

        hypothesis = tf.cast( predictions[0], tf.int32 ) # for edit_distance

        # Per-sequence statistic
        num_label_errors = tf.edit_distance( hypothesis, label, 
                                             normalize=False )

        # Per-batch summary counts
        batch_num_label_errors = tf.reduce_sum( num_label_errors)
        batch_num_sequence_errors = tf.count_nonzero( num_label_errors, axis=0 )
        batch_num_labels = tf.reduce_sum( label_length )
        
        # Wide integer type casts (prefer unsigned, but truediv dislikes those)
        batch_num_label_errors = tf.cast( batch_num_label_errors, tf.int64 )
        batch_num_sequence_errors = tf.cast( batch_num_sequence_errors, 
                                             tf.int64 )
        batch_num_labels = tf.cast( batch_num_labels, tf.int64)
        
    return batch_loss, batch_num_label_errors, batch_num_sequence_errors, \
        batch_num_labels, predictions
//  [python]
def inference(model, A, B, sess, args, src_test_iterator, src_test_next, src_vocab_rev, result_dir=None, step=None):
    sess.run([src_test_iterator.initializer])

    n_batch = 0
    t0 = time.time()
    probs = []
    if result_dir is not None:
        if not os.path.exists(result_dir):
            os.makedirs(result_dir)
        src_test_file = args.test_data[A].split("/")[-1]
        step = str(step) + "_" if step is not None else ''
        result_save_path1 = result_dir + "/" + step + src_test_file + '-' + str(B) + '.tsf'
        result_save_path2 = result_dir + "/" + step + src_test_file + '.tsf'
        print("Result save path:" + result_save_path1 + ", " + result_save_path2)
        dst_f1 = open(result_save_path1, "w")
        dst_f2 = open(result_save_path2, "w")

    while True:
        try:
            n_batch += 1
            src = sess.run(src_test_next)
            feed_dict = {
                model.input_ids: src["ids"],
                model.input_length: src["length"],
            }
            predictions = sess.run(model.predictions, feed_dict=feed_dict)
            probs.extend(predictions["log_probs"].reshape(-1))

            if result_dir is not None:
                def save_prediction():
                    batch_size = predictions["tokens"].shape[0]
                    decode_width = predictions["tokens"].shape[1]
                    src_tokens = src_vocab_rev.lookup(tf.to_int64(src["ids"]))
                    src_tokens = sess.run(src_tokens)
                    for i in range(batch_size):
                        src_tokens_ = src_tokens[i][:src["length"][i]]
                        src_sent = " ".join(src_tokens_)
                        for j in range(decode_width):
                            log_probs = predictions["log_probs"][i][j]
                            pred_tokens = predictions["tokens"][i][j][:predictions["length"][i][j] - 1]  # ignore 

                            pred_sent = " ".join(pred_tokens)
                            dst_f1.write("%s\t%s\t%s\n" % (src_sent, pred_sent, log_probs))
                            if j == 0:
                                dst_f2.write("%s\n" % pred_sent)

                save_prediction()

        except tf.errors.OutOfRangeError:  # next epoch
            print("INFERENCE---Total N batch:{}\tGenerate probs:{}\tCost time:{}".format(
                n_batch, np.mean(probs), time.time() - t0))
            break
    if result_dir is not None:
        dst_f1.close()
        dst_f2.close()
    return np.mean(probs), result_save_path2
//  [python]
def inference(model, A, B, sess, args, src_test_iterator, src_test_next, src_vocab_rev, result_dir=None, step=None):
    sess.run([src_test_iterator.initializer])

    n_batch = 0
    t0 = time.time()
    probs = []
    if result_dir is not None:
        if not os.path.exists(result_dir):
            os.makedirs(result_dir)
        src_test_file = args.test_data[A].split("/")[-1]
        step = str(step) + "_" if step is not None else ''
        result_save_path1 = result_dir + "/" + step + src_test_file + '-' + str(B) + '.tsf'
        result_save_path2 = result_dir + "/" + step + src_test_file + '.tsf'
        print("Result save path:" + result_save_path1 + ", " + result_save_path2)
        dst_f1 = open(result_save_path1, "w")
        dst_f2 = open(result_save_path2, "w")

    while True:
        try:
            n_batch += 1
            src = sess.run(src_test_next)
            feed_dict = {
                model.input_ids: src["ids"],
                model.input_length: src["length"],
            }
            predictions = sess.run(model.predictions, feed_dict=feed_dict)
            probs.extend(predictions["log_probs"].reshape(-1))

            if result_dir is not None:
                def save_prediction():
                    batch_size = predictions["tokens"].shape[0]
                    decode_width = predictions["tokens"].shape[1]
                    src_tokens = src_vocab_rev.lookup(tf.cast(src["ids"], tf.int64))
                    src_tokens = sess.run(src_tokens)
                    for i in range(batch_size):
                        src_tokens_ = src_tokens[i][:src["length"][i]]
                        src_sent = " ".join(src_tokens_)
                        for j in range(decode_width):
                            log_probs = predictions["log_probs"][i][j]
                            pred_tokens = predictions["tokens"][i][j][:predictions["length"][i][j] - 1]  # ignore 

                            pred_sent = " ".join(pred_tokens)
                            dst_f1.write("%s\t%s\t%s\n" % (src_sent, pred_sent, log_probs))
                            if j == 0:
                                dst_f2.write("%s\n" % pred_sent)

                save_prediction()

        except tf.errors.OutOfRangeError:  # next epoch
            print("INFERENCE---Total N batch:{}\tGenerate probs:{}\tCost time:{}".format(
                n_batch, np.mean(probs), time.time() - t0))
            break
    if result_dir is not None:
        dst_f1.close()
        dst_f2.close()
    return np.mean(probs), result_save_path2