<gh_stars>1-10
from typing import Callable, Tuple, Union, List

import tensorflow as tf

from .. import PreTrainModel
from ..bert import BertEmbedding
from ..configs import DebertaConfig
from ..helper import get_initializer
from ..helper import shape_list
from ..transformer import TransformerEncoderLayer, TransformerEncoder, TransformerAttention, MultiHeadSelfAttention, \
    split_multi_heads


class DeBerta(PreTrainModel):
    """
    DeBerta
    """
    config_cls = DebertaConfig

    def __init__(self, use_emd: bool = True, name='deberta', **kwargs):
        """

        :param use_emd: whether use enhanced masked decoder
        """
        super(DeBerta, self).__init__(name=name, **kwargs)

        self.embeddings = DeBertaEmbedding(embed_size=self.config.hidden_size,
                                           vocab_size=self.config.vocab_size,
                                           initializer_range=self.config.initializer_range,
                                           hidden_dropout_prob=self.config.hidden_dropout_prob,
                                           type_vocab_size=self.config.type_vocab_size,
                                           name='embeddings')
        self.encoder = DeBertaEncoder(hidden_size=self.config.hidden_size,
                                      max_relative_positions=self.config.max_relative_positions,
                                      num_hidden_layers=self.config.num_hidden_layers,
                                      num_attention_heads=self.config.num_attention_heads,
                                      attention_probs_dropout_prob=self.config.attention_probs_dropout_prob,
                                      initializer_range=self.config.initializer_range,
                                      hidden_dropout_prob=self.config.hidden_dropout_prob,
                                      hidden_act=self.config.hidden_act,
                                      intermediate_size=self.config.intermediate_size,
                                      name='encoder')

        if use_emd:
            self.emd = DeBertaEMD(hidden_size=self.config.hidden_size,
                                  max_absolute_positions=self.config.max_absolute_positions,
                                  max_relative_positions=self.config.max_relative_positions,
                                  num_hidden_layers=self.config.num_emd_layers,
                                  num_attention_heads=self.config.num_attention_heads,
                                  attention_probs_dropout_prob=self.config.attention_probs_dropout_prob,
                                  initializer_range=self.config.initializer_range,
                                  hidden_dropout_prob=self.config.hidden_dropout_prob,
                                  hidden_act=self.config.hidden_act,
                                  intermediate_size=self.config.intermediate_size,
                                  name='emd')

    def call(self,
             inputs: Union[tf.Tensor, List[tf.Tensor], Tuple[tf.Tensor]],
             attention_mask: tf.Tensor = None,
             token_type_ids: tf.Tensor = None,
             num_layer_hidden: int = -1,
             **kwargs) -> tf.Tensor:
        if isinstance(inputs, (list, tuple)):
            input_ids = inputs[0]
            attention_mask = inputs[1] if len(inputs) > 1 else attention_mask
            token_type_ids = inputs[2] if len(inputs) > 1 else token_type_ids
        else:
            input_ids = inputs

        embeds = self.embeddings(inputs=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask)
        all_hidden_states = self.encoder(hidden_states=embeds, attention_mask=attention_mask)

        hidden_states = all_hidden_states[num_layer_hidden]

        if hasattr(self, 'emd'):
            hidden_states = self.emd(hidden_states=hidden_states,
                                     rel_embeddings=self.encoder.relative_position_embeddings)

        return hidden_states


class DeBertaEncoder(TransformerEncoder):
    """
    DeBerta Encoder
    """

    def __init__(self,
                 max_relative_positions: int,
                 hidden_size: int,
                 num_hidden_layers: int,
                 initializer_range: float,
                 rel_pos_embeds: bool = True,
                 name: str = 'encoder',
                 **kwargs):
        """

        :param max_relative_positions: max relative position
        :param hidden_size: hidden size
        :param num_hidden_layers: number of transformer layer
        :param initializer_range: weights initial range
        """
        super(DeBertaEncoder, self).__init__(hidden_size=hidden_size,
                                             initializer_range=initializer_range,
                                             num_hidden_layers=num_hidden_layers,
                                             name=name,
                                             **kwargs)

        self.max_relative_positions = max_relative_positions
        self.hidden_size = hidden_size
        self.initializer_range = initializer_range
        self.rel_pos_embeds = rel_pos_embeds

        self.layers = [DeBertaEncoderLayer(max_relative_positions=max_relative_positions,
                                           hidden_size=hidden_size,
                                           initializer_range=initializer_range,
                                           name=f'layer_{idx}',
                                           **kwargs)
                       for idx in range(num_hidden_layers)]

    def build(self, input_shape):
        if self.rel_pos_embeds:
            self.relative_position_embeddings = self.add_weight(
                shape=[self.max_relative_positions * 2, self.hidden_size],
                initializer=get_initializer(self.initializer_range),
                name='relative_positions_embeddings')

        return super(DeBertaEncoder, self).build(input_shape)

    def call(self, hidden_states: tf.Tensor, **kwargs) -> tf.Tensor:
        return super(DeBertaEncoder, self).call(hidden_states=hidden_states,
                                                rel_embeddings=self.relative_position_embeddings,
                                                **kwargs)


class DeBertaEMD(DeBertaEncoder):
    """
    DeBerta Enhanced Mask Decoder
    """

    def __init__(self, max_absolute_positions: int, initializer_range: float, **kwargs):
        """

        :param max_absolute_positions: max absolute positions
        :param initializer_range: absolute position embedding weights initial range
        """
        super(DeBertaEMD, self).__init__(initializer_range=initializer_range, rel_pos_embeds=False, **kwargs)

        self.absolute_position_embeddings = tf.keras.layers.Embedding(max_absolute_positions,
                                                                      self.hidden_size,
                                                                      embeddings_initializer=get_initializer(
                                                                          initializer_range),
                                                                      name='absolute_position_embeddings')

    def call(self, hidden_states: tf.Tensor, **kwargs) -> tf.Tensor:
        position_ids = tf.range(shape_list(hidden_states)[-2], dtype=tf.int32)[tf.newaxis, :]
        abs_pos_embeds = self.absolute_position_embeddings(position_ids)

        layer_output = hidden_states + abs_pos_embeds
        for layer in self.layers:
            layer_output = layer(query=layer_output,
                                 key=hidden_states,
                                 value=hidden_states,
                                 **kwargs)

        return layer_output


class DeBertaEncoderLayer(TransformerEncoderLayer):
    def __init__(self, max_relative_positions: int, hidden_act: str, intermediate_size: int, **kwargs):
        super(DeBertaEncoderLayer, self).__init__(hidden_act=hidden_act, intermediate_size=intermediate_size, **kwargs)

        kwargs.pop('name')
        self.attention = DeBertaAttention(max_relative_positions=max_relative_positions, name='attention', **kwargs)


class DeBertaAttention(TransformerAttention):
    def __init__(self, max_relative_positions: int, **kwargs):
        super(DeBertaAttention, self).__init__(**kwargs)
        kwargs.pop('name')
        self.self_attention = DeBertaDisentangledSelfAttention(max_relative_positions=max_relative_positions,
                                                               name='self',
                                                               **kwargs)


class DeBertaDisentangledSelfAttention(MultiHeadSelfAttention):
    """
    DeBerta Disentangled Attention
    """

    def __init__(self,
                 hidden_size: int,
                 max_relative_positions: int,
                 attention_probs_dropout_prob: float,
                 hidden_dropout_prob: float,
                 initializer_range: float,
                 **kwargs):
        """

        :param hidden_size: hidden size
        :param max_relative_positions: max relative positions
        :param attention_probs_dropout_prob: attention dropout probability
        :param hidden_dropout_prob: hidden dropout probability
        :param initializer_range: weights initial range
        """
        super(DeBertaDisentangledSelfAttention, self).__init__(hidden_size=hidden_size,
                                                               attention_probs_dropout_prob=attention_probs_dropout_prob,
                                                               initializer_range=initializer_range,
                                                               **kwargs)

        self.scale_factor = tf.sqrt(tf.convert_to_tensor(3 * self.head_size, tf.float32))

        self.key_dense = tf.keras.layers.Dense(self.all_head_size,
                                               kernel_initializer=get_initializer(self.initializer_range),
                                               use_bias=False,
                                               name='key')
        self.dropout = DeBertaStableDropout(attention_probs_dropout_prob)

        self.relative_position = DeBertaRelativePositionEmbedding(max_relative_positions=max_relative_positions,
                                                                  all_head_size=hidden_size,
                                                                  num_attention_heads=self.num_attention_heads,
                                                                  hidden_dropout_prob=hidden_dropout_prob,
                                                                  initializer_range=initializer_range,
                                                                  name='relative_position')

    def relative_position_attention(self,
                                    query: tf.Tensor,
                                    key: tf.Tensor,
                                    rel_embeddings: tf.Tensor,
                                    training: tf.Tensor = False,
                                    **kwargs) -> tf.Tensor:
        """
        DeBerta relative positive attention
        :param query: query hidden states
        :param key: key hidden states
        :param rel_pos: relative position ids
        :param rel_embeddings: relative position embeddings
        :param training: training flag
        :return: relative position attention scores
        """
        return self.relative_position(query=query,
                                      key=key,
                                      rel_embeddings=rel_embeddings,
                                      training=training,
                                      **kwargs)


class DeBertaEmbedding(BertEmbedding):
    """
    DeBerta Embedding Layer
    """

    def __init__(self, hidden_dropout_prob: float, **kwargs):
        super(DeBertaEmbedding, self).__init__(hidden_dropout_prob=hidden_dropout_prob,
                                               max_position_embeddings=None,
                                               **kwargs)
        self.dropout = DeBertaStableDropout(hidden_dropout_prob)

    def embedding(self,
                  inputs: tf.Tensor,
                  token_type_ids: tf.Tensor = None,
                  attention_mask: tf.Tensor = None,
                  **kwargs) -> tf.Tensor:
        embeds = self.get_embedding([inputs, None, token_type_ids], **kwargs)
        embeds = self.LayerNorm(embeds)

        if attention_mask is not None:
            if len(shape_list(attention_mask)) != len(shape_list(embeds)):
                attention_mask = tf.cast(tf.expand_dims(attention_mask, axis=2), tf.float32)

            embeds = embeds * attention_mask

        embeds = self.dropout(embeds)

        return embeds


class DeBertaRelativePositionEmbedding(tf.keras.layers.Layer):
    """
    DeBerta Relative Position Attention Embedding
    """

    def __init__(self,
                 max_relative_positions: int,
                 all_head_size: int,
                 num_attention_heads: int,
                 hidden_dropout_prob: float,
                 initializer_range: float,
                 **kwargs):
        if all_head_size % num_attention_heads != 0:
            raise ValueError(f'head size must be an integer, got {all_head_size / num_attention_heads}')

        super(DeBertaRelativePositionEmbedding, self).__init__(**kwargs)

        self.max_relative_positions = max_relative_positions
        self.num_attention_heads = num_attention_heads

        self.scale_factor = tf.sqrt(tf.convert_to_tensor(3 * all_head_size / num_attention_heads, tf.float32))

        self.pos_proj = tf.keras.layers.Dense(all_head_size,
                                              kernel_initializer=get_initializer(initializer_range),
                                              name='pos_proj',
                                              use_bias=False)
        self.pos_q_proj = tf.keras.layers.Dense(all_head_size,
                                                kernel_initializer=get_initializer(initializer_range),
                                                name='pos_q_proj')

        self.pos_dropout = DeBertaStableDropout(hidden_dropout_prob)

    def call(self,
             query: tf.Tensor,
             key: tf.Tensor,
             rel_embeddings: tf.Tensor,
             training: tf.Tensor = False,
             **kwargs) -> tf.Tensor:
        """
        :param query: query hidden states
        :param key: key hidden states
        :param rel_embeddings: relative position embeddings
        :param training: training flag
        :return: relative position attention scores
        """
        rel_embeddings = self.pos_dropout(rel_embeddings, training=training)

        rel_pos = build_relative_position(shape_list(query)[-2], shape_list(key)[-2])
        rel_pos = tf.expand_dims(rel_pos, axis=1)

        multi_head_query = tf.transpose(split_multi_heads(query, self.num_attention_heads), perm=[0, 2, 1, 3])
        multi_head_key = tf.transpose(split_multi_heads(key, self.num_attention_heads), perm=[0, 2, 1, 3])

        attention_span = tf.cast(
            tf.minimum(
                tf.maximum(shape_list(multi_head_query)[-2], shape_list(multi_head_key)[-2]),
                self.max_relative_positions
            ),
            tf.int64
        )

        rel_embeddings = tf.expand_dims(
            rel_embeddings[self.max_relative_positions - attention_span: self.max_relative_positions + attention_span],
            axis=0
        )

        # content to position
        pos_key_layer = self.pos_proj(rel_embeddings)
        pos_key_layer = tf.transpose(split_multi_heads(pos_key_layer, self.num_attention_heads), perm=[0, 2, 1, 3])
        c2p_attn = tf.matmul(multi_head_query, pos_key_layer, transpose_b=True)
        c2p_pos = tf.clip_by_value(rel_pos + attention_span, 0, attention_span * 2 - 1)
        c2p_indices = tf.broadcast_to(c2p_pos,
                                      shape=shape_list(multi_head_query)[:-1] + shape_list(multi_head_key)[-2:-1])
        c2p_attn = torch_gather(c2p_attn, c2p_indices)

        # position to content
        pos_query_layer = self.pos_q_proj(rel_embeddings)
        pos_query_layer = tf.transpose(split_multi_heads(pos_query_layer, self.num_attention_heads), perm=[0, 2, 1, 3])
        p2c_attn = tf.matmul(multi_head_key, pos_query_layer, transpose_b=True)
        p2c_pos = tf.clip_by_value(-rel_pos + attention_span, 0, attention_span * 2 - 1)
        p2c_indices = tf.broadcast_to(p2c_pos,
                                      shape=shape_list(multi_head_query)[:2] + shape_list(multi_head_key)[
                                                                               -2:-1] + shape_list(
                                          multi_head_query)[-2: -1])
        # [batch_size, head, key_sequence_length, query_sequence_length]
        p2c_attn = torch_gather(p2c_attn, p2c_indices)
        # [batch_size, head, query_sequence_length, key_sequence_length]
        p2c_attn = tf.transpose(p2c_attn, perm=[0, 1, 3, 2])

        # add c2p with p2c and scale
        rel_pos_attn_scores = tf.divide(c2p_attn + p2c_attn, self.scale_factor)
        return rel_pos_attn_scores


class DeBertaStableDropout(tf.keras.layers.Layer):
    """
    DeBerta Custom Dropout, it use to bernoulli to generate dropout distribution
    """

    def __init__(self, drop_prob: float, **kwargs):
        """
        :param drop_prob: dropout probability
        """
        super().__init__(**kwargs)
        self.drop_prob = tf.convert_to_tensor(drop_prob, dtype=tf.float32)

    def call(self, inputs: tf.Tensor, training: tf.Tensor = False, **kwargs) -> Union[
        tf.Tensor, Tuple[tf.Tensor, Callable]]:
        if training and self.drop_prob > 0:
            return DeBertaXDropout(inputs, self.drop_prob)

        return inputs


@tf.custom_gradient
def DeBertaXDropout(inputs: tf.Tensor, drop_prob: tf.Tensor) -> Tuple[tf.Tensor, Callable]:
    """
    DeBerta custom dropout main layer
    :param inputs: inputs
    :param drop_prob: dropout probability
    :return: dropped out inputs and custom gradient function
    """
    mask = tf.cast(
        1 - tf.compat.v1.distributions.Bernoulli(probs=1 - drop_prob).sample(sample_shape=shape_list(inputs)), tf.bool
    )

    scale = tf.convert_to_tensor(1.0 / (1 - drop_prob), dtype=tf.float32)
    inputs = tf.cond(drop_prob > 0, lambda: tf.where(mask, 0.0, inputs) * scale, lambda: inputs)

    def custom_grad(upstream_grad):
        return tf.cond(
            scale > 1, lambda: (tf.where(mask, 0.0, upstream_grad) * scale, None), lambda: (upstream_grad, None)
        )

    return inputs, custom_grad


def build_relative_position(query_size: int, key_size: int) -> tf.Tensor:
    """
    Build the position representation of the key relative to the query.
    note:
        1. if the key position is 4, the query position is 1, and then the relative position is -3
        2. if the key position is 1, the query position is 4, and then the relative position is 3
    :param query_size: self attention query length
    :param key_size: self attention key length
    :return: relative position ids
    """
    q_ids = tf.range(query_size, dtype=tf.int32)
    k_ids = tf.range(key_size, dtype=tf.int32)

    rel_pos_ids = q_ids[:, None] - tf.tile(tf.reshape(k_ids, [1, -1]), [query_size, 1])[:query_size, :]
    rel_pos_ids = tf.expand_dims(rel_pos_ids, axis=0)

    return tf.cast(rel_pos_ids, dtype=tf.int64)


def torch_gather(x: tf.Tensor, indices: tf.Tensor) -> tf.Tensor:
    """
    Gather relative position attentions scores from given indices
    :param x: inputs
    :param indices: gathered indices
    :return: gathered inputs
    """
    flat_x = tf.reshape(x, [-1, shape_list(x)[-1]])
    flat_indices = tf.reshape(indices, [-1, shape_list(indices)[-1]])

    return tf.reshape(tf.gather(flat_x, flat_indices, batch_dims=1), shape_list(indices))


from typing import Callable, Tuple, Union, List

import tensorflow as tf

from .. import PreTrainModel
from ..bert import BertEmbedding
from ..configs import DebertaConfig
from ..helper import get_initializer
from ..helper import shape_list
from ..transformer import TransformerEncoderLayer, TransformerEncoder, TransformerAttention, MultiHeadSelfAttention, \
    split_multi_heads


class DeBerta(PreTrainModel):
    """
    DeBerta
    """
    config_cls = DebertaConfig

    def __init__(self, use_emd: bool = True, name='deberta', **kwargs):
        """

        :param use_emd: whether use enhanced masked decoder
        """
        super(DeBerta, self).__init__(name=name, **kwargs)

        self.embeddings = DeBertaEmbedding(embed_size=self.config.hidden_size,
                                           vocab_size=self.config.vocab_size,
                                           initializer_range=self.config.initializer_range,
                                           hidden_dropout_prob=self.config.hidden_dropout_prob,
                                           type_vocab_size=self.config.type_vocab_size,
                                           name='embeddings')
        self.encoder = DeBertaEncoder(hidden_size=self.config.hidden_size,
                                      max_relative_positions=self.config.max_relative_positions,
                                      num_hidden_layers=self.config.num_hidden_layers,
                                      num_attention_heads=self.config.num_attention_heads,
                                      attention_probs_dropout_prob=self.config.attention_probs_dropout_prob,
                                      initializer_range=self.config.initializer_range,
                                      hidden_dropout_prob=self.config.hidden_dropout_prob,
                                      hidden_act=self.config.hidden_act,
                                      intermediate_size=self.config.intermediate_size,
                                      name='encoder')

        if use_emd:
            self.emd = DeBertaEMD(hidden_size=self.config.hidden_size,
                                  max_absolute_positions=self.config.max_absolute_positions,
                                  max_relative_positions=self.config.max_relative_positions,
                                  num_hidden_layers=self.config.num_emd_layers,
                                  num_attention_heads=self.config.num_attention_heads,
                                  attention_probs_dropout_prob=self.config.attention_probs_dropout_prob,
                                  initializer_range=self.config.initializer_range,
                                  hidden_dropout_prob=self.config.hidden_dropout_prob,
                                  hidden_act=self.config.hidden_act,
                                  intermediate_size=self.config.intermediate_size,
                                  name='emd')

    def call(self,
             inputs: Union[tf.Tensor, List[tf.Tensor], Tuple[tf.Tensor]],
             attention_mask: tf.Tensor = None,
             token_type_ids: tf.Tensor = None,
             num_layer_hidden: int = -1,
             **kwargs) -> tf.Tensor:
        if isinstance(inputs, (list, tuple)):
            input_ids = inputs[0]
            attention_mask = inputs[1] if len(inputs) > 1 else attention_mask
            token_type_ids = inputs[2] if len(inputs) > 1 else token_type_ids
        else:
            input_ids = inputs

        embeds = self.embeddings(inputs=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask)
        all_hidden_states = self.encoder(hidden_states=embeds, attention_mask=attention_mask)

        hidden_states = all_hidden_states[num_layer_hidden]

        if hasattr(self, 'emd'):
            hidden_states = self.emd(hidden_states=hidden_states,
                                     rel_embeddings=self.encoder.relative_position_embeddings)

        return hidden_states


class DeBertaEncoder(TransformerEncoder):
    """
    DeBerta Encoder
    """

    def __init__(self,
                 max_relative_positions: int,
                 hidden_size: int,
                 num_hidden_layers: int,
                 initializer_range: float,
                 rel_pos_embeds: bool = True,
                 name: str = 'encoder',
                 **kwargs):
        """

        :param max_relative_positions: max relative position
        :param hidden_size: hidden size
        :param num_hidden_layers: number of transformer layer
        :param initializer_range: weights initial range
        """
        super(DeBertaEncoder, self).__init__(hidden_size=hidden_size,
                                             initializer_range=initializer_range,
                                             num_hidden_layers=num_hidden_layers,
                                             name=name,
                                             **kwargs)

        self.max_relative_positions = max_relative_positions
        self.hidden_size = hidden_size
        self.initializer_range = initializer_range
        self.rel_pos_embeds = rel_pos_embeds

        self.layers = [DeBertaEncoderLayer(max_relative_positions=max_relative_positions,
                                           hidden_size=hidden_size,
                                           initializer_range=initializer_range,
                                           name=f'layer_{idx}',
                                           **kwargs)
                       for idx in range(num_hidden_layers)]

    def build(self, input_shape):
        if self.rel_pos_embeds:
            self.relative_position_embeddings = self.add_weight(
                shape=[self.max_relative_positions * 2, self.hidden_size],
                initializer=get_initializer(self.initializer_range),
                name='relative_positions_embeddings')

        return super(DeBertaEncoder, self).build(input_shape)

    def call(self, hidden_states: tf.Tensor, **kwargs) -> tf.Tensor:
        return super(DeBertaEncoder, self).call(hidden_states=hidden_states,
                                                rel_embeddings=self.relative_position_embeddings,
                                                **kwargs)


class DeBertaEMD(DeBertaEncoder):
    """
    DeBerta Enhanced Mask Decoder
    """

    def __init__(self, max_absolute_positions: int, initializer_range: float, **kwargs):
        """

        :param max_absolute_positions: max absolute positions
        :param initializer_range: absolute position embedding weights initial range
        """
        super(DeBertaEMD, self).__init__(initializer_range=initializer_range, rel_pos_embeds=False, **kwargs)

        self.absolute_position_embeddings = tf.keras.layers.Embedding(max_absolute_positions,
                                                                      self.hidden_size,
                                                                      embeddings_initializer=get_initializer(
                                                                          initializer_range),
                                                                      name='absolute_position_embeddings')

    def call(self, hidden_states: tf.Tensor, **kwargs) -> tf.Tensor:
        position_ids = tf.range(shape_list(hidden_states)[-2], dtype=tf.int32)[tf.newaxis, :]
        abs_pos_embeds = self.absolute_position_embeddings(position_ids)

        layer_output = hidden_states + abs_pos_embeds
        for layer in self.layers:
            layer_output = layer(query=layer_output,
                                 key=hidden_states,
                                 value=hidden_states,
                                 **kwargs)

        return layer_output


class DeBertaEncoderLayer(TransformerEncoderLayer):
    def __init__(self, max_relative_positions: int, hidden_act: str, intermediate_size: int, **kwargs):
        super(DeBertaEncoderLayer, self).__init__(hidden_act=hidden_act, intermediate_size=intermediate_size, **kwargs)

        kwargs.pop('name')
        self.attention = DeBertaAttention(max_relative_positions=max_relative_positions, name='attention', **kwargs)


class DeBertaAttention(TransformerAttention):
    def __init__(self, max_relative_positions: int, **kwargs):
        super(DeBertaAttention, self).__init__(**kwargs)
        kwargs.pop('name')
        self.self_attention = DeBertaDisentangledSelfAttention(max_relative_positions=max_relative_positions,
                                                               name='self',
                                                               **kwargs)


class DeBertaDisentangledSelfAttention(MultiHeadSelfAttention):
    """
    DeBerta Disentangled Attention
    """

    def __init__(self,
                 hidden_size: int,
                 max_relative_positions: int,
                 attention_probs_dropout_prob: float,
                 hidden_dropout_prob: float,
                 initializer_range: float,
                 **kwargs):
        """

        :param hidden_size: hidden size
        :param max_relative_positions: max relative positions
        :param attention_probs_dropout_prob: attention dropout probability
        :param hidden_dropout_prob: hidden dropout probability
        :param initializer_range: weights initial range
        """
        super(DeBertaDisentangledSelfAttention, self).__init__(hidden_size=hidden_size,
                                                               attention_probs_dropout_prob=attention_probs_dropout_prob,
                                                               initializer_range=initializer_range,
                                                               **kwargs)

        self.scale_factor = tf.sqrt(tf.convert_to_tensor(3 * self.head_size, tf.float32))

        self.key_dense = tf.keras.layers.Dense(self.all_head_size,
                                               kernel_initializer=get_initializer(self.initializer_range),
                                               use_bias=False,
                                               name='key')
        self.dropout = DeBertaStableDropout(attention_probs_dropout_prob)

        self.relative_position = DeBertaRelativePositionEmbedding(max_relative_positions=max_relative_positions,
                                                                  all_head_size=hidden_size,
                                                                  num_attention_heads=self.num_attention_heads,
                                                                  hidden_dropout_prob=hidden_dropout_prob,
                                                                  initializer_range=initializer_range,
                                                                  name='relative_position')

    def relative_position_attention(self,
                                    query: tf.Tensor,
                                    key: tf.Tensor,
                                    rel_embeddings: tf.Tensor,
                                    training: tf.Tensor = False,
                                    **kwargs) -> tf.Tensor:
        """
        DeBerta relative positive attention
        :param query: query hidden states
        :param key: key hidden states
        :param rel_pos: relative position ids
        :param rel_embeddings: relative position embeddings
        :param training: training flag
        :return: relative position attention scores
        """
        return self.relative_position(query=query,
                                      key=key,
                                      rel_embeddings=rel_embeddings,
                                      training=training,
                                      **kwargs)


class DeBertaEmbedding(BertEmbedding):
    """
    DeBerta Embedding Layer
    """

    def __init__(self, hidden_dropout_prob: float, **kwargs):
        super(DeBertaEmbedding, self).__init__(hidden_dropout_prob=hidden_dropout_prob,
                                               max_position_embeddings=None,
                                               **kwargs)
        self.dropout = DeBertaStableDropout(hidden_dropout_prob)

    def embedding(self,
                  inputs: tf.Tensor,
                  token_type_ids: tf.Tensor = None,
                  attention_mask: tf.Tensor = None,
                  **kwargs) -> tf.Tensor:
        embeds = self.get_embedding([inputs, None, token_type_ids], **kwargs)
        embeds = self.LayerNorm(embeds)

        if attention_mask is not None:
            if len(shape_list(attention_mask)) != len(shape_list(embeds)):
                attention_mask = tf.cast(tf.expand_dims(attention_mask, axis=2), tf.float32)

            embeds = embeds * attention_mask

        embeds = self.dropout(embeds)

        return embeds


class DeBertaRelativePositionEmbedding(tf.keras.layers.Layer):
    """
    DeBerta Relative Position Attention Embedding
    """

    def __init__(self,
                 max_relative_positions: int,
                 all_head_size: int,
                 num_attention_heads: int,
                 hidden_dropout_prob: float,
                 initializer_range: float,
                 **kwargs):
        if all_head_size % num_attention_heads != 0:
            raise ValueError(f'head size must be an integer, got {all_head_size / num_attention_heads}')

        super(DeBertaRelativePositionEmbedding, self).__init__(**kwargs)

        self.max_relative_positions = max_relative_positions
        self.num_attention_heads = num_attention_heads

        self.scale_factor = tf.sqrt(tf.convert_to_tensor(3 * all_head_size / num_attention_heads, tf.float32))

        self.pos_proj = tf.keras.layers.Dense(all_head_size,
                                              kernel_initializer=get_initializer(initializer_range),
                                              name='pos_proj',
                                              use_bias=False)
        self.pos_q_proj = tf.keras.layers.Dense(all_head_size,
                                                kernel_initializer=get_initializer(initializer_range),
                                                name='pos_q_proj')

        self.pos_dropout = DeBertaStableDropout(hidden_dropout_prob)

    def call(self,
             query: tf.Tensor,
             key: tf.Tensor,
             rel_embeddings: tf.Tensor,
             training: tf.Tensor = False,
             **kwargs) -> tf.Tensor:
        """
        :param query: query hidden states
        :param key: key hidden states
        :param rel_embeddings: relative position embeddings
        :param training: training flag
        :return: relative position attention scores
        """
        rel_embeddings = self.pos_dropout(rel_embeddings, training=training)

        rel_pos = build_relative_position(shape_list(query)[-2], shape_list(key)[-2])
        rel_pos = tf.expand_dims(rel_pos, axis=1)

        multi_head_query = tf.transpose(split_multi_heads(query, self.num_attention_heads), perm=[0, 2, 1, 3])
        multi_head_key = tf.transpose(split_multi_heads(key, self.num_attention_heads), perm=[0, 2, 1, 3])

        attention_span = tf.cast(
            tf.minimum(
                tf.maximum(shape_list(multi_head_query)[-2], shape_list(multi_head_key)[-2]),
                self.max_relative_positions
            ),
            tf.int64
        )

        rel_embeddings = tf.expand_dims(
            rel_embeddings[self.max_relative_positions - attention_span: self.max_relative_positions + attention_span],
            axis=0
        )

        # content to position
        pos_key_layer = self.pos_proj(rel_embeddings)
        pos_key_layer = tf.transpose(split_multi_heads(pos_key_layer, self.num_attention_heads), perm=[0, 2, 1, 3])
        c2p_attn = tf.matmul(multi_head_query, pos_key_layer, transpose_b=True)
        c2p_pos = tf.clip_by_value(rel_pos + attention_span, 0, attention_span * 2 - 1)
        c2p_indices = tf.broadcast_to(c2p_pos,
                                      shape=shape_list(multi_head_query)[:-1] + shape_list(multi_head_key)[-2:-1])
        c2p_attn = torch_gather(c2p_attn, c2p_indices)

        # position to content
        pos_query_layer = self.pos_q_proj(rel_embeddings)
        pos_query_layer = tf.transpose(split_multi_heads(pos_query_layer, self.num_attention_heads), perm=[0, 2, 1, 3])
        p2c_attn = tf.matmul(multi_head_key, pos_query_layer, transpose_b=True)
        p2c_pos = tf.clip_by_value(-rel_pos + attention_span, 0, attention_span * 2 - 1)
        p2c_indices = tf.broadcast_to(p2c_pos,
                                      shape=shape_list(multi_head_query)[:2] + shape_list(multi_head_key)[
                                                                               -2:-1] + shape_list(
                                          multi_head_query)[-2: -1])
        # [batch_size, head, key_sequence_length, query_sequence_length]
        p2c_attn = torch_gather(p2c_attn, p2c_indices)
        # [batch_size, head, query_sequence_length, key_sequence_length]
        p2c_attn = tf.transpose(p2c_attn, perm=[0, 1, 3, 2])

        # add c2p with p2c and scale
        rel_pos_attn_scores = tf.divide(c2p_attn + p2c_attn, self.scale_factor)
        return rel_pos_attn_scores


class DeBertaStableDropout(tf.keras.layers.Layer):
    """
    DeBerta Custom Dropout, it use to bernoulli to generate dropout distribution
    """

    def __init__(self, drop_prob: float, **kwargs):
        """
        :param drop_prob: dropout probability
        """
        super().__init__(**kwargs)
        self.drop_prob = tf.convert_to_tensor(drop_prob, dtype=tf.float32)

    def call(self, inputs: tf.Tensor, training: tf.Tensor = False, **kwargs) -> Union[
        tf.Tensor, Tuple[tf.Tensor, Callable]]:
        if training and self.drop_prob > 0:
            return DeBertaXDropout(inputs, self.drop_prob)

        return inputs


@tf.custom_gradient
def DeBertaXDropout(inputs: tf.Tensor, drop_prob: tf.Tensor) -> Tuple[tf.Tensor, Callable]:
    """
    DeBerta custom dropout main layer
    :param inputs: inputs
    :param drop_prob: dropout probability
    :return: dropped out inputs and custom gradient function
    """
    mask = tf.cast(
        1 - tf.compat.v1.distributions.Bernoulli(probs=1 - drop_prob).sample(sample_shape=shape_list(inputs)), tf.bool
    )

    scale = tf.convert_to_tensor(1.0 / (1 - drop_prob), dtype=tf.float32)
    inputs = tf.cond(drop_prob > 0, lambda: tf.where(mask, 0.0, inputs) * scale, lambda: inputs)

    def custom_grad(upstream_grad):
        return tf.cond(
            scale > 1, lambda: (tf.where(mask, 0.0, upstream_grad) * scale, None), lambda: (upstream_grad, None)
        )

    return inputs, custom_grad


def build_relative_position(query_size: int, key_size: int) -> tf.Tensor:
    """
    Build the position representation of the key relative to the query.
    note:
        1. if the key position is 4, the query position is 1, and then the relative position is -3
        2. if the key position is 1, the query position is 4, and then the relative position is 3
    :param query_size: self attention query length
    :param key_size: self attention key length
    :return: relative position ids
    """
    q_ids = tf.range(query_size, dtype=tf.int32)
    k_ids = tf.range(key_size, dtype=tf.int32)

    rel_pos_ids = q_ids[:, None] - tf.tile(tf.reshape(k_ids, [1, -1]), [query_size, 1])[:query_size, :]
    rel_pos_ids = tf.expand_dims(rel_pos_ids, axis=0)

    return tf.cast(rel_pos_ids, dtype=tf.int64)


def torch_gather(x: tf.Tensor, indices: tf.Tensor) -> tf.Tensor:
    """
    Gather relative position attentions scores from given indices
    :param x: inputs
    :param indices: gathered indices
    :return: gathered inputs
    """
    flat_x = tf.reshape(x, [-1, shape_list(x)[-1]])
    flat_indices = tf.reshape(indices, [-1, shape_list(indices)[-1]])

    return tf.reshape(tf.gather(flat_x, flat_indices, batch_dims=1), shape_list(indices))
DeBerta :param use_emd: whether use enhanced masked decoder DeBerta Encoder :param max_relative_positions: max relative position         :param hidden_size: hidden size         :param num_hidden_layers: number of transformer layer         :param initializer_range: weights initial range DeBerta Enhanced Mask Decoder :param max_absolute_positions: max absolute positions         :param initializer_range: absolute position embedding weights initial range DeBerta Disentangled Attention :param hidden_size: hidden size         :param max_relative_positions: max relative positions         :param attention_probs_dropout_prob: attention dropout probability         :param hidden_dropout_prob: hidden dropout probability         :param initializer_range: weights initial range DeBerta relative positive attention         :param query: query hidden states         :param key: key hidden states         :param rel_pos: relative position ids         :param rel_embeddings: relative position embeddings         :param training: training flag         :return: relative position attention scores DeBerta Embedding Layer DeBerta Relative Position Attention Embedding :param query: query hidden states         :param key: key hidden states         :param rel_embeddings: relative position embeddings         :param training: training flag         :return: relative position attention scores # content to position # position to content # [batch_size, head, key_sequence_length, query_sequence_length] # [batch_size, head, query_sequence_length, key_sequence_length] # add c2p with p2c and scale DeBerta Custom Dropout, it use to bernoulli to generate dropout distribution :param drop_prob: dropout probability DeBerta custom dropout main layer     :param inputs: inputs     :param drop_prob: dropout probability     :return: dropped out inputs and custom gradient function Build the position representation of the key relative to the query.     note:         1. if the key position is 4, the query position is 1, and then the relative position is -3         2. if the key position is 1, the query position is 4, and then the relative position is 3     :param query_size: self attention query length     :param key_size: self attention key length     :return: relative position ids Gather relative position attentions scores from given indices     :param x: inputs     :param indices: gathered indices     :return: gathered inputs