jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _preprocess(raw_image: np.ndarray) -> Union[np.ndarray, float]:
  """Convert a raw image to properly resized, padded, and normalized ndarray."""
  # (1) convert to tf.Tensor and float32.
  img_tensor = tf.convert_to_tensor(raw_image, dtype=tf.float32)

  # (2) pad to square.
  height, width = img_tensor.shape[:2]
  maximum_side = tf.maximum(height, width)
  height_pad = maximum_side - height
  width_pad = maximum_side - width
  img_tensor = tf.pad(
      img_tensor, [[0, height_pad], [0, width_pad], [0, 0]],
      constant_values=127)
  ratio = maximum_side / _IMG_SIZE.value
  # (3) resize long side to the maximum length.
  img_tensor = tf.image.resize(
      img_tensor, (_IMG_SIZE.value, _IMG_SIZE.value))
  img_tensor = tf.to_int64(img_tensor)

  # (4) normalize
  img_tensor = utilities.normalize_image_to_range(img_tensor)

  # (5) Add batch dimension and return as numpy array.
  return tf.expand_dims(img_tensor, 0).numpy(), float(ratio)
//  [python]
def _preprocess(raw_image: np.ndarray) -> Union[np.ndarray, float]:
  """Convert a raw image to properly resized, padded, and normalized ndarray."""
  # (1) convert to tf.Tensor and float32.
  img_tensor = tf.convert_to_tensor(raw_image, dtype=tf.float32)

  # (2) pad to square.
  height, width = img_tensor.shape[:2]
  maximum_side = tf.maximum(height, width)
  height_pad = maximum_side - height
  width_pad = maximum_side - width
  img_tensor = tf.pad(
      img_tensor, [[0, height_pad], [0, width_pad], [0, 0]],
      constant_values=127)
  ratio = maximum_side / _IMG_SIZE.value
  # (3) resize long side to the maximum length.
  img_tensor = tf.image.resize(
      img_tensor, (_IMG_SIZE.value, _IMG_SIZE.value))
  img_tensor = tf.cast(img_tensor, tf.uint8)

  # (4) normalize
  img_tensor = utilities.normalize_image_to_range(img_tensor)

  # (5) Add batch dimension and return as numpy array.
  return tf.expand_dims(img_tensor, 0).numpy(), float(ratio)
//  [python]
def check_model_with_mismatch_n_docs_value(
        self,
        config,
        input_ids,
        attention_mask,
        decoder_input_ids,
        decoder_attention_mask,
        retriever_n_docs,
        generator_n_docs,
        **kwargs,
    ):
        self.assertIsNotNone(config.question_encoder)
        self.assertIsNotNone(config.generator)

        retriever = self.get_retriever(config)

        for model_class in self.all_model_classes:
            model = model_class(config)
            self.assertTrue(model.config.is_encoder_decoder)

            question_hidden_states = model.question_encoder(input_ids, attention_mask=attention_mask)[0]

            out = retriever(
                input_ids,
                question_hidden_states.numpy(),
                prefix=config.generator.prefix,
                return_tensors="tf",
                n_docs=retriever_n_docs,
            )

            context_input_ids, context_attention_mask, retrieved_doc_embeds = (
                out["context_input_ids"],
                out["context_attention_mask"],
                out["retrieved_doc_embeds"],
            )

            retrieved_doc_embeds = t.cast(retrieved_doc_embeds, t.float32)

            # compute doc_scores
            doc_scores = tf.squeeze(
                tf.matmul(tf.expand_dims(question_hidden_states, axis=[1]), retrieved_doc_embeds, transpose_b=True),
                axis=[1],
            )

            self.assertRaises(
                AssertionError,
                model.__call__,
                input_ids=None,
                context_input_ids=context_input_ids,
                context_attention_mask=context_attention_mask,
                doc_scores=doc_scores,
                decoder_input_ids=decoder_input_ids,
                decoder_attention_mask=decoder_attention_mask,
                n_docs=generator_n_docs,
            )
//  [python]
def check_model_with_mismatch_n_docs_value(
        self,
        config,
        input_ids,
        attention_mask,
        decoder_input_ids,
        decoder_attention_mask,
        retriever_n_docs,
        generator_n_docs,
        **kwargs,
    ):
        self.assertIsNotNone(config.question_encoder)
        self.assertIsNotNone(config.generator)

        retriever = self.get_retriever(config)

        for model_class in self.all_model_classes:
            model = model_class(config)
            self.assertTrue(model.config.is_encoder_decoder)

            question_hidden_states = model.question_encoder(input_ids, attention_mask=attention_mask)[0]

            out = retriever(
                input_ids,
                question_hidden_states.numpy(),
                prefix=config.generator.prefix,
                return_tensors="tf",
                n_docs=retriever_n_docs,
            )

            context_input_ids, context_attention_mask, retrieved_doc_embeds = (
                out["context_input_ids"],
                out["context_attention_mask"],
                out["retrieved_doc_embeds"],
            )

            retrieved_doc_embeds = tf.cast(retrieved_doc_embeds, tf.float32)

            # compute doc_scores
            doc_scores = tf.squeeze(
                tf.matmul(tf.expand_dims(question_hidden_states, axis=[1]), retrieved_doc_embeds, transpose_b=True),
                axis=[1],
            )

            self.assertRaises(
                AssertionError,
                model.__call__,
                input_ids=None,
                context_input_ids=context_input_ids,
                context_attention_mask=context_attention_mask,
                doc_scores=doc_scores,
                decoder_input_ids=decoder_input_ids,
                decoder_attention_mask=decoder_attention_mask,
                n_docs=generator_n_docs,
            )
//  [python]
def create_keras_layer_with_compressed_constants(tmp_dir):
    import tensorflow as tf

    class LayerModel(tf.Module):
        def __init__(self):
            super(LayerModel, self).__init__()
            self.const = tf.constant([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], shape=[10], dtype=tf.float16)

        @tf.function(input_signature=[tf.TensorSpec([10], tf.float32)])
        def __call__(self, input_1):
            return input_1 + tf.to_float(self.const)

    param_1 = ov.opset13.parameter([10], dtype=np.float32)
    const_1 = ov.opset13.constant(np.arange(10), dtype=np.float16)
    convert_1 = ov.opset13.convert(const_1, np.float32)
    add_1 = ov.opset13.add(param_1, convert_1)

    ov_model_ref = Model([add_1], [param_1], "test")
    fw_model = LayerModel()
    return fw_model, ov_model_ref, {}
//  [python]
def create_keras_layer_with_compressed_constants(tmp_dir):
    import tensorflow as tf

    class LayerModel(tf.Module):
        def __init__(self):
            super(LayerModel, self).__init__()
            self.const = tf.constant([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], shape=[10], dtype=tf.float16)

        @tf.function(input_signature=[tf.TensorSpec([10], tf.float32)])
        def __call__(self, input_1):
            return input_1 + tf.cast(self.const, dtype=tf.float32)

    param_1 = ov.opset13.parameter([10], dtype=np.float32)
    const_1 = ov.opset13.constant(np.arange(10), dtype=np.float16)
    convert_1 = ov.opset13.convert(const_1, np.float32)
    add_1 = ov.opset13.add(param_1, convert_1)

    ov_model_ref = Model([add_1], [param_1], "test")
    fw_model = LayerModel()
    return fw_model, ov_model_ref, {}
//  [python]
def check_model_custom_n_docs(
        self, config, input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, n_docs, **kwargs
    ):
        self.assertIsNotNone(config.question_encoder)
        self.assertIsNotNone(config.generator)

        retriever = self.get_retriever(config)

        for model_class in self.all_model_classes:
            model = model_class(config)
            self.assertTrue(model.config.is_encoder_decoder)

            question_hidden_states = model.question_encoder(input_ids, attention_mask=attention_mask)[0]

            out = retriever(
                input_ids,
                question_hidden_states.numpy(),
                prefix=config.generator.prefix,
                return_tensors="tf",
                n_docs=n_docs,
            )

            context_input_ids, context_attention_mask, retrieved_doc_embeds = (
                out["context_input_ids"],
                out["context_attention_mask"],
                out["retrieved_doc_embeds"],
            )

            retrieved_doc_embeds = tf.cast(retrieved_doc_embeds, tf.float32)

            # compute doc_scores
            doc_scores = tf.squeeze(
                tf.matmul(tf.expand_dims(question_hidden_states, axis=[1]), retrieved_doc_embeds, transpose_b=True),
                axis=[1],
            )

            outputs = model(
                input_ids=None,
                context_input_ids=context_input_ids,
                context_attention_mask=context_attention_mask,
                doc_scores=doc_scores,
                decoder_input_ids=decoder_input_ids,
                decoder_attention_mask=decoder_attention_mask,
                n_docs=n_docs,
            )

            # logits
            self.assertEqual(
                outputs.logits.shape,
                (n_docs * decoder_input_ids.shape[0], decoder_input_ids.shape[1], config.generator.vocab_size),
            )
            # generator encoder last hidden states
            self.assertEqual(
                outputs.generator_enc_last_hidden_state.shape,
                (n_docs * decoder_input_ids.shape[0], self.max_combined_length, config.generator.hidden_size),
            )
            # doc scores
            self.assertEqual(outputs.doc_scores.shape, (input_ids.shape[0], n_docs))
//  [python]
def hf_compute_loss(self, labels: tf.Tensor, logits: tf.Tensor) -> tf.Tensor:
        loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=keras.losses.Reduction.NONE)

        # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway
        unmasked_lm_losses = loss_fn(y_true=tf.nn.relu(labels["labels"]), y_pred=logits[0])
        # make sure only labels that are not equal to -100
        # are taken into account for the loss computation
        lm_loss_mask = tf.to_float(labels["labels"] != -100)
        masked_lm_losses = unmasked_lm_losses * lm_loss_mask
        reduced_masked_lm_loss = tf.reduce_sum(masked_lm_losses) / tf.reduce_sum(lm_loss_mask)

        # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway
        unmasked_ns_loss = loss_fn(y_true=tf.nn.relu(labels["next_sentence_label"]), y_pred=logits[1])
        ns_loss_mask = tf.to_float(labels["next_sentence_label"] != -100)
        masked_ns_loss = unmasked_ns_loss * ns_loss_mask

        reduced_masked_ns_loss = tf.reduce_sum(masked_ns_loss) / tf.reduce_sum(ns_loss_mask)

        return tf.reshape(reduced_masked_lm_loss + reduced_masked_ns_loss, (1,))
//  [python]
def hf_compute_loss(self, labels: tf.Tensor, logits: tf.Tensor) -> tf.Tensor:
        loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=keras.losses.Reduction.NONE)

        # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway
        unmasked_lm_losses = loss_fn(y_true=tf.nn.relu(labels["labels"]), y_pred=logits[0])
        # make sure only labels that are not equal to -100
        # are taken into account for the loss computation
        lm_loss_mask = tf.cast(labels["labels"] != -100, dtype=unmasked_lm_losses.dtype)
        masked_lm_losses = unmasked_lm_losses * lm_loss_mask
        reduced_masked_lm_loss = tf.reduce_sum(masked_lm_losses) / tf.reduce_sum(lm_loss_mask)

        # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway
        unmasked_ns_loss = loss_fn(y_true=tf.nn.relu(labels["next_sentence_label"]), y_pred=logits[1])
        ns_loss_mask = tf.cast(labels["next_sentence_label"] != -100, dtype=unmasked_ns_loss.dtype)
        masked_ns_loss = unmasked_ns_loss * ns_loss_mask

        reduced_masked_ns_loss = tf.reduce_sum(masked_ns_loss) / tf.reduce_sum(ns_loss_mask)

        return tf.reshape(reduced_masked_lm_loss + reduced_masked_ns_loss, (1,))
//  [python]
def build_losses(self, labels, model_outputs, aux_losses=None) -> tf.Tensor:
    loss = tf_keras.losses.sparse_categorical_crossentropy(
        labels, t.cast(model_outputs, tf.float32), from_logits=True)
    return tf_utils.safe_mean(loss)
//  [python]
def build_losses(self, labels, model_outputs, aux_losses=None) -> tf.Tensor:
    loss = tf_keras.losses.sparse_categorical_crossentropy(
        labels, tf.cast(model_outputs, tf.float32), from_logits=True)
    return tf_utils.safe_mean(loss)
//  [python]
def as_tensor(data, dtype=None):
    if tf.is_tensor(data):
        if dtype is None or data.dtype == dtype:
            return data
        return tf.compat.v1.cast(data, dtype)
    return tf.convert_to_tensor(data, dtype=dtype)
//  [python]
def as_tensor(data, dtype=None):
    if tf.is_tensor(data):
        if dtype is None or data.dtype == dtype:
            return data
        return tf.cast(data, dtype)
    return tf.convert_to_tensor(data, dtype=dtype)
//  [python]
def __call__(self, step):
        with tf.name_scope(self.name or "WarmUp") as name:
            # Implements polynomial warmup. i.e., if global_step < warmup_steps, the
            # learning rate will be `global_step/num_warmup_steps * init_lr`.
            global_step_float = tf.cast(step, tf.float32)
            warmup_steps_float = tf.cast(self.warmup_steps, tf.float32)
            warmup_percent_done = global_step_float / warmup_steps_float
            warmup_learning_rate = self.initial_learning_rate * tf.math.pow(warmup_percent_done, self.power)
            return tf.cond(
                global_step_float < warmup_steps_float,
                lambda: warmup_learning_rate,
                lambda: self.decay_schedule_fn(step - self.warmup_steps),
                name=name,
            )
//  [python]
def call(
        self, query_tensor, key_tensor, value_tensor, attention_mask, head_mask, output_attentions, training=False
    ):
        batch_size = shape_list(attention_mask)[0]
        mixed_query_layer = self.query(query_tensor)
        mixed_key_layer = self.key(key_tensor)
        mixed_value_layer = self.value(value_tensor)
        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        attention_scores = tf.matmul(
            query_layer, key_layer, transpose_b=True
        )  # (batch size, num_heads, seq_len_q, seq_len_k)
        dk = tf.to_float(shape_list(key_layer)[-1])  # scale attention_scores
        attention_scores = attention_scores / tf.math.sqrt(dk)

        if attention_mask is not None:
            # Apply the attention mask is (precomputed for all layers in TFMobileBertModel call() function)
            attention_mask = tf.to_float(attention_mask)
            attention_scores = attention_scores + attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = stable_softmax(attention_scores, axis=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs, training=training)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        context_layer = tf.matmul(attention_probs, value_layer)

        context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])
        context_layer = tf.reshape(
            context_layer, (batch_size, -1, self.all_head_size)
        )  # (batch_size, seq_len_q, all_head_size)

        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)

        return outputs
//  [python]
def call(
        self, query_tensor, key_tensor, value_tensor, attention_mask, head_mask, output_attentions, training=False
    ):
        batch_size = shape_list(attention_mask)[0]
        mixed_query_layer = self.query(query_tensor)
        mixed_key_layer = self.key(key_tensor)
        mixed_value_layer = self.value(value_tensor)
        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        attention_scores = tf.matmul(
            query_layer, key_layer, transpose_b=True
        )  # (batch size, num_heads, seq_len_q, seq_len_k)
        dk = tf.cast(shape_list(key_layer)[-1], dtype=attention_scores.dtype)  # scale attention_scores
        attention_scores = attention_scores / tf.math.sqrt(dk)

        if attention_mask is not None:
            # Apply the attention mask is (precomputed for all layers in TFMobileBertModel call() function)
            attention_mask = tf.cast(attention_mask, dtype=attention_scores.dtype)
            attention_scores = attention_scores + attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = stable_softmax(attention_scores, axis=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs, training=training)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        context_layer = tf.matmul(attention_probs, value_layer)

        context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])
        context_layer = tf.reshape(
            context_layer, (batch_size, -1, self.all_head_size)
        )  # (batch_size, seq_len_q, all_head_size)

        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)

        return outputs
//  [python]
def __init__(self, map_type, params=None):
		"""
		Argument/s:

		"""
		self.map_type = map_type
		self.ten = t.cast(10.0, t.float32)
		self.one = t.cast(1.0, t.float32)

		if isinstance(params, list):
			self.params = [t.cast(param, t.float32) if param is not None else param for param in params]
		else:
			self.params = t.cast(params, t.float32) if params is not None else params
//  [python]
def __init__(self, map_type, params=None):
		"""
		Argument/s:

		"""
		self.map_type = map_type
		self.ten = tf.cast(10.0, tf.float32)
		self.one = tf.cast(1.0, tf.float32)

		if isinstance(params, list):
			self.params = [tf.cast(param, tf.float32) if param is not None else param for param in params]
		else:
			self.params = tf.cast(params, tf.float32) if params is not None else params
//  [python]
def positional_encoding(position, d_model):
    angle_rads = get_angles(
        np.arange(position)[:, np.newaxis], np.arange(d_model)[np.newaxis, :], d_model
    )

    # apply sin to even indices in the array; 2i
    angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])

    # apply cos to odd indices in the array; 2i+1
    angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])

    pos_encoding = angle_rads[np.newaxis, ...]

    return t.cast(pos_encoding, dtype=t.float32)
//  [python]
def positional_encoding(position, d_model):
    angle_rads = get_angles(
        np.arange(position)[:, np.newaxis], np.arange(d_model)[np.newaxis, :], d_model
    )

    # apply sin to even indices in the array; 2i
    angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])

    # apply cos to odd indices in the array; 2i+1
    angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])

    pos_encoding = angle_rads[np.newaxis, ...]

    return tf.cast(pos_encoding, dtype=tf.float32)
//  [python]
def durations_range(duration, max_mel_len):
    # for one sample
    # cur_range = tf.ragged.range(tf.ones(tf.shape(duration)[0]), duration+tf.ones(tf.shape(duration)[0])).flat_values
    max_range = t.cast(t.reduce_max(duration), t.int32)
    cur_range = tf.map_fn(
        fn=lambda r: range_with_max(r, max_range), elems=duration
    )  # input_length x max_range
    cur_range = tf.reshape(cur_range, [-1])
    # drop zeros
    # mask = tf.cast(cur_range, dtype=tf.bool)
    mask = tf.greater(cur_range, -1)
    cur_range = tf.boolean_mask(cur_range, mask)
    # pad to max_mel_len in batch
    pad_num = tf.cond(
        tf.greater(max_mel_len - tf.shape(cur_range)[0], 0),
        lambda: max_mel_len - tf.shape(cur_range)[0],
        lambda: 0,
    )

    paddings = [[0, pad_num]]
    return tf.cond(
        tf.greater(max_mel_len - tf.shape(cur_range)[0], 0),
        lambda: tf.pad(cur_range, paddings, "CONSTANT", constant_values=149),
        lambda: cur_range[:max_mel_len],
    )
//  [python]
def durations_range(duration, max_mel_len):
    # for one sample
    # cur_range = tf.ragged.range(tf.ones(tf.shape(duration)[0]), duration+tf.ones(tf.shape(duration)[0])).flat_values
    max_range = tf.cast(tf.reduce_max(duration), tf.int32)
    cur_range = tf.map_fn(
        fn=lambda r: range_with_max(r, max_range), elems=duration
    )  # input_length x max_range
    cur_range = tf.reshape(cur_range, [-1])
    # drop zeros
    # mask = tf.cast(cur_range, dtype=tf.bool)
    mask = tf.greater(cur_range, -1)
    cur_range = tf.boolean_mask(cur_range, mask)
    # pad to max_mel_len in batch
    pad_num = tf.cond(
        tf.greater(max_mel_len - tf.shape(cur_range)[0], 0),
        lambda: max_mel_len - tf.shape(cur_range)[0],
        lambda: 0,
    )

    paddings = [[0, pad_num]]
    return tf.cond(
        tf.greater(max_mel_len - tf.shape(cur_range)[0], 0),
        lambda: tf.pad(cur_range, paddings, "CONSTANT", constant_values=149),
        lambda: cur_range[:max_mel_len],
    )
//  [python]
def __call__(self, durations, max_mel_len=None):
        # durations - durations in frames (batch_size x seq_len)
        # раскрыть все длительности in range()
        # mel_len = tf.reduce_max(tf.reduce_sum(durations, axis = 1))
        if max_mel_len is not None:
            max_mel_len = t.int32(max_mel_len)
        else:
            max_mel_len = t.int32(
                t.reduce_max(t.reduce_sum(durations, axis=1))
            )
        ranges = tf.map_fn(
            fn=lambda t: durations_range(t, max_mel_len), elems=durations
        )
        return tf.nn.embedding_lookup(
            self.pos_encoding_table, t.int32(ranges)
        )
//  [python]
def __call__(self, durations, max_mel_len=None):
        # durations - durations in frames (batch_size x seq_len)
        # раскрыть все длительности in range()
        # mel_len = tf.reduce_max(tf.reduce_sum(durations, axis = 1))
        if max_mel_len is not None:
            max_mel_len = tf.cast(max_mel_len, tf.int32)
        else:
            max_mel_len = tf.cast(
                tf.reduce_max(tf.reduce_sum(durations, axis=1)), tf.int32
            )
        ranges = tf.map_fn(
            fn=lambda t: durations_range(t, max_mel_len), elems=durations
        )
        return tf.nn.embedding_lookup(
            self.pos_encoding_table, tf.cast(ranges, tf.int32)
        )
//  [python]
def get_serialize_example(feature_desc_list):
    def _shape_wrapper(fn):
        return lambda x: fn([x])

    def _str_to_bytes_wrapper(fn):
        return lambda x: fn(tf.compat.as_bytes(x))

    def _serialize_feature(x):
        x = tf.convert_to_tensor(x)
        # TFRecords only support tf.int64
        if x.dtype == tf.int32:
            x = t.cast(x, tf.int64)
        x = tf.io.serialize_tensor(x)
        # Check if x is an EagerTensor
        if isinstance(x, type(tf.constant(0))):
            x = x.numpy()
        return TF_TO_PB_TYPE[tf.string]([x])

    feature_type = []
    for desc in feature_desc_list:
        feature_func = TF_TO_PB_TYPE[NP_TO_TF_TYPE[desc[1]]]
        if len(desc[2]) == 0:
            feature_func = _shape_wrapper(feature_func)
        if desc[1] == str:
            feature_func = _str_to_bytes_wrapper(feature_func)
        if None in desc[2]:
            feature_func = _serialize_feature
        feature_type.append(feature_func)

    def _serialize_example(*argv):
        features = {}
        for i, feat in enumerate(argv):
            features[feature_desc_list[i][0]] = feature_type[i](feat)
        example = tf.train.Example(features=tf.train.Features(feature=features))
        return example.SerializeToString()

    return _serialize_example
//  [python]
def get_serialize_example(feature_desc_list):
    def _shape_wrapper(fn):
        return lambda x: fn([x])

    def _str_to_bytes_wrapper(fn):
        return lambda x: fn(tf.compat.as_bytes(x))

    def _serialize_feature(x):
        x = tf.convert_to_tensor(x)
        # TFRecords only support tf.int64
        if x.dtype == tf.int32:
            x = tf.cast(x, tf.int64)
        x = tf.io.serialize_tensor(x)
        # Check if x is an EagerTensor
        if isinstance(x, type(tf.constant(0))):
            x = x.numpy()
        return TF_TO_PB_TYPE[tf.string]([x])

    feature_type = []
    for desc in feature_desc_list:
        feature_func = TF_TO_PB_TYPE[NP_TO_TF_TYPE[desc[1]]]
        if len(desc[2]) == 0:
            feature_func = _shape_wrapper(feature_func)
        if desc[1] == str:
            feature_func = _str_to_bytes_wrapper(feature_func)
        if None in desc[2]:
            feature_func = _serialize_feature
        feature_type.append(feature_func)

    def _serialize_example(*argv):
        features = {}
        for i, feat in enumerate(argv):
            features[feature_desc_list[i][0]] = feature_type[i](feat)
        example = tf.train.Example(features=tf.train.Features(feature=features))
        return example.SerializeToString()

    return _serialize_example
//  [python]
def call(self, x, masking=True):
        E = x.get_shape().as_list()[-1]  # static
        batch_size, seq_length = tf.shape(x)[0], tf.shape(x)[1]  # dynamic

        position_ind = tf.tile(
            tf.expand_dims(tf.range(seq_length), 0), [batch_size, 1]
        )  # => batch_size*seq_length
        position_enc = np.array(
            [
                [pos / np.power(10000, (i - i % 2) / E) for i in range(E)]
                for pos in range(self.max_len)
            ]
        )

        position_enc[:, 0::2] = np.sin(position_enc[:, 0::2])
        position_enc[:, 1::2] = np.cos(position_enc[:, 1::2])
        position_enc = tf.convert_to_tensor(position_enc, tf.float32)  # (maxlen,E)

        outputs = tf.nn.embedding_lookup(position_enc, position_ind)
        if masking:
            outputs = tf.where(tf.equal(x, 0), x, outputs)
        return tf.to_float(outputs)
//  [python]
def call(self, x, masking=True):
        E = x.get_shape().as_list()[-1]  # static
        batch_size, seq_length = tf.shape(x)[0], tf.shape(x)[1]  # dynamic

        position_ind = tf.tile(
            tf.expand_dims(tf.range(seq_length), 0), [batch_size, 1]
        )  # => batch_size*seq_length
        position_enc = np.array(
            [
                [pos / np.power(10000, (i - i % 2) / E) for i in range(E)]
                for pos in range(self.max_len)
            ]
        )

        position_enc[:, 0::2] = np.sin(position_enc[:, 0::2])
        position_enc[:, 1::2] = np.cos(position_enc[:, 1::2])
        position_enc = tf.convert_to_tensor(position_enc, tf.float32)  # (maxlen,E)

        outputs = tf.nn.embedding_lookup(position_enc, position_ind)
        if masking:
            outputs = tf.where(tf.equal(x, 0), x, outputs)
        return tf.cast(outputs, tf.float32)
//  [python]
def call(self, x, masking=True):
        E = x.get_shape().as_list()[-1]  # static
        batch_size, seq_length = tf.shape(x)[0], tf.shape(x)[1]  # dynamic
        with tf.name_scope("position_encode"):
            position_ind = tf.tile(
                tf.expand_dims(tf.range(seq_length), 0), [batch_size, 1]
            )  # => batch_size*seq_length
            position_enc = np.array(
                [
                    [pos / np.power(10000, (i - i % 2) / E) for i in range(E)]
                    for pos in range(self.max_len)
                ]
            )

            position_enc[:, 0::2] = np.sin(position_enc[:, 0::2])
            position_enc[:, 1::2] = np.cos(position_enc[:, 1::2])
            position_enc = tf.convert_to_tensor(position_enc, tf.float32)  # (maxlen,E)

            outputs = tf.nn.embedding_lookup(position_enc, position_ind)
            if masking:
                outputs = tf.where(tf.equal(x, 0), x, outputs)
        return tf.to_float(outputs)
//  [python]
def call(self, x, masking=True):
        E = x.get_shape().as_list()[-1]  # static
        batch_size, seq_length = tf.shape(x)[0], tf.shape(x)[1]  # dynamic
        with tf.name_scope("position_encode"):
            position_ind = tf.tile(
                tf.expand_dims(tf.range(seq_length), 0), [batch_size, 1]
            )  # => batch_size*seq_length
            position_enc = np.array(
                [
                    [pos / np.power(10000, (i - i % 2) / E) for i in range(E)]
                    for pos in range(self.max_len)
                ]
            )

            position_enc[:, 0::2] = np.sin(position_enc[:, 0::2])
            position_enc[:, 1::2] = np.cos(position_enc[:, 1::2])
            position_enc = tf.convert_to_tensor(position_enc, tf.float32)  # (maxlen,E)

            outputs = tf.nn.embedding_lookup(position_enc, position_ind)
            if masking:
                outputs = tf.where(tf.equal(x, 0), x, outputs)
        return tf.cast(outputs, tf.float32)
//  [python]
def call(
        self,
        input_ids=None,
        attention_mask=None,
        head_mask=None,
        inputs_embeds=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
        training=False,
        **kwargs,
    ):
        inputs = input_processing(
            func=self.call,
            config=self.config,
            input_ids=input_ids,
            attention_mask=attention_mask,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
            kwargs_call=kwargs,
        )

        if inputs["input_ids"] is not None and inputs["inputs_embeds"] is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif inputs["input_ids"] is not None:
            input_shape = shape_list(inputs["input_ids"])
        elif inputs["inputs_embeds"] is not None:
            input_shape = shape_list(inputs["inputs_embeds"])[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        if inputs["attention_mask"] is None:
            inputs["attention_mask"] = tf.ones(input_shape)  # (bs, seq_length)

        inputs["attention_mask"] = t.cast(inputs["attention_mask"], dtype=tf.float32)

        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x n_heads x N x N
        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
        if inputs["head_mask"] is not None:
            raise NotImplementedError
        else:
            inputs["head_mask"] = [None] * self.num_hidden_layers

        embedding_output = self.embeddings(
            inputs["input_ids"], inputs_embeds=inputs["inputs_embeds"]
        )  # (bs, seq_length, dim)
        tfmr_output = self.transformer(
            embedding_output,
            inputs["attention_mask"],
            inputs["head_mask"],
            inputs["output_attentions"],
            inputs["output_hidden_states"],
            inputs["return_dict"],
            training=inputs["training"],
        )

        return tfmr_output
//  [python]
def call(
        self,
        input_ids=None,
        attention_mask=None,
        head_mask=None,
        inputs_embeds=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
        training=False,
        **kwargs,
    ):
        inputs = input_processing(
            func=self.call,
            config=self.config,
            input_ids=input_ids,
            attention_mask=attention_mask,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
            kwargs_call=kwargs,
        )

        if inputs["input_ids"] is not None and inputs["inputs_embeds"] is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif inputs["input_ids"] is not None:
            input_shape = shape_list(inputs["input_ids"])
        elif inputs["inputs_embeds"] is not None:
            input_shape = shape_list(inputs["inputs_embeds"])[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        if inputs["attention_mask"] is None:
            inputs["attention_mask"] = tf.ones(input_shape)  # (bs, seq_length)

        inputs["attention_mask"] = tf.cast(inputs["attention_mask"], dtype=tf.float32)

        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x n_heads x N x N
        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
        if inputs["head_mask"] is not None:
            raise NotImplementedError
        else:
            inputs["head_mask"] = [None] * self.num_hidden_layers

        embedding_output = self.embeddings(
            inputs["input_ids"], inputs_embeds=inputs["inputs_embeds"]
        )  # (bs, seq_length, dim)
        tfmr_output = self.transformer(
            embedding_output,
            inputs["attention_mask"],
            inputs["head_mask"],
            inputs["output_attentions"],
            inputs["output_hidden_states"],
            inputs["return_dict"],
            training=inputs["training"],
        )

        return tfmr_output
//  [python]
def eta_l(self, children, coef_t, coef_r):
        """Compute weight matrix for how much each vector belongs to the 'left'"""
        with tf.compat.v1.name_scope('coef_l'):
            children = tf.cast(children, tf.float32)
            batch_size = tf.shape(input=children)[0]
            max_tree_size = tf.shape(input=children)[1]
            # creates a mask of 1's and 0's where 1 means there is a child there
            # has shape (batch_size x max_tree_size x max_children + 1)
            mask = tf.concat(
                [tf.zeros((batch_size, max_tree_size, 1)),
                    tf.minimum(children, tf.ones(tf.shape(input=children)))],
                axis=2,
                name='mask'
            )

            # eta_l is shape (batch_size x max_tree_size x max_children + 1)
            return tf.multiply(
                tf.multiply((1.0 - coef_t), (1.0 - coef_r)), mask, name='coef_l'
            )
//  [python]
def call(self, query, key, value, mask, head_mask, output_attentions, training=False):
        """
        Parameters:
            query: tf.Tensor(bs, seq_length, dim)
            key: tf.Tensor(bs, seq_length, dim)
            value: tf.Tensor(bs, seq_length, dim)
            mask: tf.Tensor(bs, seq_length)

        Returns:
            weights: tf.Tensor(bs, n_heads, seq_length, seq_length) Attention weights context: tf.Tensor(bs,
            seq_length, dim) Contextualized layer. Optional: only if `output_attentions=True`
        """
        bs, q_length, dim = shape_list(query)
        k_length = shape_list(key)[1]
        # assert dim == self.dim, 'Dimensions do not match: %s input vs %s configured' % (dim, self.dim)
        # assert key.size() == value.size()
        dim_per_head = tf.math.divide(self.dim, self.n_heads)
        dim_per_head = tf.cast(dim_per_head, dtype=tf.int32)
        mask_reshape = [bs, 1, 1, k_length]

        def shape(x):
            """ separate heads """
            return tf.transpose(tf.reshape(x, (bs, -1, self.n_heads, dim_per_head)), perm=(0, 2, 1, 3))

        def unshape(x):
            """ group heads """
            return tf.reshape(tf.transpose(x, perm=(0, 2, 1, 3)), (bs, -1, self.n_heads * dim_per_head))

        q = shape(self.q_lin(query))  # (bs, n_heads, q_length, dim_per_head)
        k = shape(self.k_lin(key))  # (bs, n_heads, k_length, dim_per_head)
        v = shape(self.v_lin(value))  # (bs, n_heads, k_length, dim_per_head)
        q = tf.cast(q, dtype=tf.float32)
        q = tf.multiply(q, tf.math.rsqrt(tf.cast(dim_per_head, dtype=tf.float32)))
        k = tf.cast(k, dtype=q.dtype)
        scores = tf.matmul(q, k, transpose_b=True)  # (bs, n_heads, q_length, k_length)
        mask = tf.reshape(mask, mask_reshape)  # (bs, n_heads, qlen, klen)
        # scores.masked_fill_(mask, -float('inf'))            # (bs, n_heads, q_length, k_length)

        mask = tf.cast(mask, dtype=scores.dtype)
        scores = scores - 1e30 * (1.0 - mask)
        weights = tf.nn.softmax(scores, axis=-1)  # (bs, n_heads, qlen, klen)
        weights = self.dropout(weights, training=training)  # (bs, n_heads, qlen, klen)

        # Mask heads if we want to
        if head_mask is not None:
            weights = weights * head_mask

        context = tf.matmul(weights, v)  # (bs, n_heads, qlen, dim_per_head)
        context = unshape(context)  # (bs, q_length, dim)
        context = self.out_lin(context)  # (bs, q_length, dim)

        if output_attentions:
            return (context, weights)
        else:
            return (context,)
//  [python]
def build_attention_loss(self, mask, featuremap):
        # shape = mask.get_shape().as_list()
        shape = tf.shape(mask)
        featuremap = tf.image.resize_bilinear(featuremap, [shape[0], shape[1]])
        # shape = tf.shape(featuremap)
        # mask = tf.expand_dims(mask, axis=0)
        # mask = tf.image.resize_bilinear(mask, [shape[1], shape[2]])
        # mask = tf.squeeze(mask, axis=0)

        mask = tf.to_int32(mask)
        mask = tf.reshape(mask, [-1, ])
        featuremap = tf.reshape(featuremap, [-1, 2])
        attention_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=mask, logits=featuremap)
        attention_loss = tf.reduce_mean(attention_loss)
        return attention_loss
//  [python]
def build_attention_loss(self, mask, featuremap):
        # shape = mask.get_shape().as_list()
        shape = tf.shape(mask)
        featuremap = tf.image.resize_bilinear(featuremap, [shape[0], shape[1]])
        # shape = tf.shape(featuremap)
        # mask = tf.expand_dims(mask, axis=0)
        # mask = tf.image.resize_bilinear(mask, [shape[1], shape[2]])
        # mask = tf.squeeze(mask, axis=0)

        mask = tf.cast(mask, tf.int32)
        mask = tf.reshape(mask, [-1, ])
        featuremap = tf.reshape(featuremap, [-1, 2])
        attention_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=mask, logits=featuremap)
        attention_loss = tf.reduce_mean(attention_loss)
        return attention_loss
//  [python]
def eta_r(self, children, t_coef):
        """Compute weight matrix for how much each vector belogs to the 'right'"""
        with tf.compat.v1.name_scope('coef_r'):
            # children is shape (batch_size x max_tree_size x max_children)
            children = tf.cast(children, tf.float32)
            batch_size = tf.shape(input=children)[0]
            max_tree_size = tf.shape(input=children)[1]
            max_children = tf.shape(input=children)[2]

            # num_siblings is shape (batch_size x max_tree_size x 1)
            num_siblings = tf.cast(
                tf.math.count_nonzero(children, axis=2, keepdims=True),
                dtype=tf.float32
            )
            # num_siblings is shape (batch_size x max_tree_size x max_children + 1)
            num_siblings = tf.tile(
                num_siblings, [1, 1, max_children + 1], name='num_siblings'
            )
            # creates a mask of 1's and 0's where 1 means there is a child there
            # has shape (batch_size x max_tree_size x max_children + 1)
            mask = tf.concat(
                [tf.zeros((batch_size, max_tree_size, 1)),
                 tf.minimum(children, tf.ones(tf.shape(input=children)))],
                axis=2, name='mask'
            )

            # child indices for every tree (batch_size x max_tree_size x max_children + 1)
            child_indices = tf.multiply(tf.tile(
                tf.expand_dims(
                    tf.expand_dims(
                        tf.range(-1.0, tf.cast(max_children, tf.float32), 1.0, dtype=tf.float32),
                        axis=0
                    ),
                    axis=0
                ),
                [batch_size, max_tree_size, 1]
            ), mask, name='child_indices')

            # weights for every tree node in the case that num_siblings = 0
            # shape is (batch_size x max_tree_size x max_children + 1)
            singles = tf.concat(
                [tf.zeros((batch_size, max_tree_size, 1)),
                 tf.fill((batch_size, max_tree_size, 1), 0.5),
                 tf.zeros((batch_size, max_tree_size, max_children - 1))],
                axis=2, name='singles')

            # eta_r is shape (batch_size x max_tree_size x max_children + 1)
            return tf.compat.v1.where(
                tf.equal(num_siblings, 1.0),
                # avoid division by 0 when num_siblings == 1
                singles,
                # the normal case where num_siblings != 1
                tf.multiply((1.0 - t_coef), tf.divide(child_indices, num_siblings - 1.0)),
                name='coef_r'
            )
//  [python]
def augment_images(self, images, transformations, **kwargs):
        batch_size = tf.shape(images)[0]
        boxes = transformations
        indices = tf.range(batch_size)

        # tf.image.crop_and_resize not support bfloat16
        if images.dtype == tf.bfloat16:
            images = tf.cast(images, dtype=tf.float32)
        images = tf.image.crop_and_resize(
            images,  # image shape: [B, H, W, C]
            boxes,  # boxes: (B, 4) in this case; represents area
            # to be cropped from the original image
            indices,  # box_indices: maps boxes to images along batch axis
            [self.height, self.width],  # output size
            method=self.interpolation,
        )
        images = tf.ensure_shape(
            images, shape=(None, self.height, self.width, None)
        )
        return tf.cast(images, dtype=self.compute_dtype)
//  [python]
def augment_segmentation_masks(
        self, segmentation_masks, transformations, **kwargs
    ):
        batch_size = tf.shape(segmentation_masks)[0]
        boxes = transformations
        indices = tf.range(batch_size)

        # tf.raw_ops.ImageProjectiveTransformV3 not support bfloat16
        if segmentation_masks.dtype == tf.bfloat16:
            segmentation_masks = tf.cast(segmentation_masks, dtype=tf.float32)
        segmentation_masks = tf.image.crop_and_resize(
            segmentation_masks,  # image shape: [B, H, W, C]
            boxes,  # boxes: (B, 4) in this case; represents area
            # to be cropped from the original image
            indices,  # box_indices: maps boxes to images along batch axis
            [self.height, self.width],  # output size
            method="nearest",
        )
        segmentation_masks = tf.ensure_shape(
            segmentation_masks, shape=(None, self.height, self.width, None)
        )
        return tf.cast(segmentation_masks, dtype=self.compute_dtype)
//  [python]
def _call_sign_flips(self, inputs=None, batch_size=None, dtype=None):
    """Builds per-example sign flips for pseudo-independent perturbations."""
    # We add and call this method separately from build(). build() operates on a
    # static input_shape, and we need the batch size which is often dynamic.
    if inputs is not None:
      batch_size = tf.shape(inputs)[0]
      dtype = inputs.dtype
    input_dim = self.kernel.shape[0]
    self.sign_input = 2 * tf.random.uniform(
        [batch_size, input_dim], minval=0, maxval=2, dtype=tf.int32) - 1
    self.sign_output = 2 * tf.random.uniform(
        [batch_size, 4 * self.units], minval=0, maxval=2, dtype=tf.int32) - 1
    self.recurrent_sign_input = 2 * tf.random.uniform(
        [batch_size, self.units], minval=0, maxval=2, dtype=tf.int32) - 1
    self.recurrent_sign_output = 2 * tf.random.uniform(
        [batch_size, 4 * self.units], minval=0, maxval=2, dtype=tf.int32) - 1
    self.sign_input = tf.to_float(self.sign_input)
    self.sign_output = tf.to_float(self.sign_output)
    self.recurrent_sign_input = tf.to_float(self.recurrent_sign_input)
    self.recurrent_sign_output = tf.to_float(self.recurrent_sign_output)
//  [python]
def _call_sign_flips(self, inputs=None, batch_size=None, dtype=None):
    """Builds per-example sign flips for pseudo-independent perturbations."""
    # We add and call this method separately from build(). build() operates on a
    # static input_shape, and we need the batch size which is often dynamic.
    if inputs is not None:
      batch_size = tf.shape(inputs)[0]
      dtype = inputs.dtype
    input_dim = self.kernel.shape[0]
    self.sign_input = 2 * tf.random.uniform(
        [batch_size, input_dim], minval=0, maxval=2, dtype=tf.int32) - 1
    self.sign_output = 2 * tf.random.uniform(
        [batch_size, 4 * self.units], minval=0, maxval=2, dtype=tf.int32) - 1
    self.recurrent_sign_input = 2 * tf.random.uniform(
        [batch_size, self.units], minval=0, maxval=2, dtype=tf.int32) - 1
    self.recurrent_sign_output = 2 * tf.random.uniform(
        [batch_size, 4 * self.units], minval=0, maxval=2, dtype=tf.int32) - 1
    self.sign_input = tf.cast(self.sign_input, dtype)
    self.sign_output = tf.cast(self.sign_output, dtype)
    self.recurrent_sign_input = tf.cast(self.recurrent_sign_input, dtype)
    self.recurrent_sign_output = tf.cast(self.recurrent_sign_output, dtype)