更新模型版本并修复问题。

QiDeBERTa.py

@@ -1,68 +1,473 @@
+from dataclasses import dataclass
 from typing import Optional, Tuple
 
 import torch
-from torch.nn import Module, Linear, Parameter, LayerNorm, Embedding, Dropout, ModuleList, MultiheadAttention
+from torch.nn import Module, Linear, Parameter, LayerNorm, Embedding, Dropout, ModuleList, MultiheadAttention, functional
 from torch.nn.functional import gelu
 from transformers import DebertaV2PreTrainedModel
-from transformers.modeling_outputs import BaseModelOutput, TokenClassifierOutput, MaskedLMOutput, SequenceClassifierOutput
-from transformers.models.deberta_v2.modeling_deberta_v2 import build_relative_position, DebertaV2Layer
+from transformers.modeling_outputs import BaseModelOutput as EncoderOutput, MaskedLMOutput
+from transformers.models.deberta_v2.modeling_deberta_v2 import build_relative_position, scaled_size_sqrt, build_rpos
+from transformers.utils import ModelOutput
 
-from .Configuration import QiDeBERTaConfig
+from .configuration import QiDeBERTaConfig
+
+
+@dataclass
+class BaseModelOutput(ModelOutput):
+    last_hidden_state: Optional[torch.FloatTensor] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+    embedding_output: Optional[torch.FloatTensor] = None
+    token_embeddings: Optional[torch.FloatTensor] = None
 
 
 class QiDeBERTaEmbeddings(Module):
     """Construct the embeddings from word, position and token_type embeddings."""
-
     def __init__(
         self,
         pad_token_id: int,
-        hidden_size: int,
+        d_model: int,
         vocab_size: int,
+        layer_norm_eps: float,
+        hidden_dropout_prob: float,
+    ):
+        super().__init__()
+        self.word_embeddings = Embedding(num_embeddings=vocab_size, embedding_dim=d_model, padding_idx=pad_token_id)
+        self.LayerNorm = LayerNorm(normalized_shape=d_model, eps=layer_norm_eps)
+        self.dropout = Dropout(p=hidden_dropout_prob)
+
+    def forward(self, input_ids: torch.Tensor, mask: torch.Tensor):
+        inputs_embeds = self.word_embeddings(input_ids)
+        embeddings = self.LayerNorm(inputs_embeds)
+        if mask.dim() != embeddings.dim():
+            if mask.dim() == 4:
+                mask = mask.squeeze(1).squeeze(1)
+            mask = mask.unsqueeze(2)
+        mask = mask.to(embeddings.dtype)
+        return self.dropout(embeddings * mask), inputs_embeds
+
+
+class QiDeBERTaDisentangledSelfAttention(Module):
+    """
+    Disentangled self-attention module
+    """
+    def __init__(
+        self,
+        num_heads: int,
+        d_model: int,
+        share_att_key: bool,
+        relative_attention: bool,
+        max_position_embeddings: int,
+        hidden_dropout_prob: float,
+        attention_probs_dropout_prob: float,
+        pos_att_type: Optional[list] = None,
+        position_buckets: int = -1,
+        max_relative_positions: int = -1,
     ):
         super().__init__()
-        self.word_embeddings = Embedding(vocab_size, hidden_size, padding_idx=pad_token_id)
+        self.num_attention_heads = num_heads
+        self.attention_head_size = d_model // num_heads
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+        self.query_proj = Linear(in_features=d_model, out_features=self.all_head_size, bias=True)
+        self.key_proj = Linear(in_features=d_model, out_features=self.all_head_size, bias=True)
+        self.value_proj = Linear(in_features=d_model, out_features=self.all_head_size, bias=True)
+
+        self.share_att_key = share_att_key
+        self.pos_att_type = pos_att_type if pos_att_type is not None else []
+        self.relative_attention = relative_attention
+
+        if self.relative_attention:
+            self.position_buckets = position_buckets
+            self.max_relative_positions = max_relative_positions
+            if self.max_relative_positions < 1:
+                self.max_relative_positions = max_position_embeddings
+            self.pos_ebd_size = self.max_relative_positions
+            if self.position_buckets > 0:
+                self.pos_ebd_size = self.position_buckets
+
+            self.pos_dropout = Dropout(p=hidden_dropout_prob)
+
+            if not self.share_att_key:
+                if "c2p" in self.pos_att_type:
+                    self.pos_key_proj = Linear(in_features=d_model, out_features=self.all_head_size, bias=True)
+                if "p2c" in self.pos_att_type:
+                    self.pos_query_proj = Linear(in_features=d_model, out_features=self.all_head_size)
+
+        self.dropout = Dropout(p=attention_probs_dropout_prob)
 
-        self.LayerNorm = LayerNorm(normalized_shape=hidden_size, eps=1e-7)
-        self.dropout = Dropout(p=0.1)
+    @staticmethod
+    def transpose_for_scores(x, attention_heads) -> torch.Tensor:
+        new_x_shape = x.size()[:-1] + (attention_heads, -1)
+        x = x.view(new_x_shape)
+        return x.permute(0, 2, 1, 3).contiguous().view(-1, x.size(1), x.size(-1))
 
-    def forward(self, input_ids=None, mask=None, inputs_embeds=None):
-        if inputs_embeds is None:
-            inputs_embeds = self.word_embeddings(input_ids)
+    def forward(
+        self,
+        hidden_states,
+        attention_mask,
+        output_attentions=False,
+        relative_pos=None,
+        rel_embeddings=None,
+    ):
+        """
+        Call the module
 
-        embeddings = inputs_embeds
+        Args:
+            hidden_states (`torch.FloatTensor`):
+                Input states to the module usually the output from previous layer, it will be the Q,K and V in
+                *Attention(Q,K,V)*
 
-        embeddings = self.LayerNorm(embeddings)
+            attention_mask (`torch.BoolTensor`):
+                An attention mask matrix of shape [*B*, *N*, *N*] where *B* is the batch size, *N* is the maximum
+                sequence length in which element [i,j] = *1* means the *i* th token in the input can attend to the *j*
+                th token.
 
-        if mask is not None:
-            if mask.dim() != embeddings.dim():
-                if mask.dim() == 4:
-                    mask = mask.squeeze(1).squeeze(1)
-                mask = mask.unsqueeze(2)
-            mask = mask.to(embeddings.dtype)
+            output_attentions (`bool`, *optional*):
+                Whether return the attention matrix.
 
-            embeddings = embeddings * mask
+            relative_pos (`torch.LongTensor`):
+                The relative position encoding between the tokens in the sequence. It's of shape [*B*, *N*, *N*] with
+                values ranging in [*-max_relative_positions*, *max_relative_positions*].
 
-        embeddings = self.dropout(embeddings)
-        return embeddings
+            rel_embeddings (`torch.FloatTensor`):
+                The embedding of relative distances. It's a tensor of shape [\\(2 \\times
+                \\text{max_relative_positions}\\), *hidden_size*].
 
 
-class QiDeBERTaEncoder(Module):
-    """Modified BertEncoder with relative position bias support"""
+        """
+        query_layer = self.transpose_for_scores(self.query_proj(hidden_states), self.num_attention_heads)
+        key_layer = self.transpose_for_scores(self.key_proj(hidden_states), self.num_attention_heads)
+        value_layer = self.transpose_for_scores(self.value_proj(hidden_states), self.num_attention_heads)
+
+        rel_att = None
+        # Take the dot product between "query" and "key" to get the raw attention scores.
+        scale_factor = 1
+        if "c2p" in self.pos_att_type:
+            scale_factor += 1
+        if "p2c" in self.pos_att_type:
+            scale_factor += 1
+        scale = scaled_size_sqrt(query_layer, scale_factor)
+        attention_scores = torch.bmm(query_layer, key_layer.transpose(-1, -2) / scale.to(dtype=query_layer.dtype))
+        if self.relative_attention:
+            rel_embeddings = self.pos_dropout(rel_embeddings)
+            rel_att = self.disentangled_attention_bias(
+                query_layer, key_layer, relative_pos, rel_embeddings, scale_factor
+            )
+
+        if rel_att is not None:
+            attention_scores = attention_scores + rel_att
+        attention_scores = attention_scores
+        attention_scores = attention_scores.view(
+            -1, self.num_attention_heads, attention_scores.size(-2), attention_scores.size(-1)
+        )
 
-    def __init__(self, config):
+        attention_mask = attention_mask.bool()
+        attention_scores = attention_scores.masked_fill(~(attention_mask), torch.finfo(query_layer.dtype).min)
+        # bsz x height x length x dimension
+        attention_probs = functional.softmax(attention_scores, dim=-1)
+
+        attention_probs = self.dropout(attention_probs)
+        context_layer = torch.bmm(
+            attention_probs.view(-1, attention_probs.size(-2), attention_probs.size(-1)), value_layer
+        )
+        context_layer = (
+            context_layer.view(-1, self.num_attention_heads, context_layer.size(-2), context_layer.size(-1))
+            .permute(0, 2, 1, 3)
+            .contiguous()
+        )
+        new_context_layer_shape = context_layer.size()[:-2] + (-1,)
+        context_layer = context_layer.view(new_context_layer_shape)
+
+        return (context_layer, attention_probs) if output_attentions else (context_layer, None)
+
+    def disentangled_attention_bias(self, query_layer, key_layer, relative_pos, rel_embeddings, scale_factor):
+        if relative_pos is None:
+            relative_pos = build_relative_position(
+                query_layer,
+                key_layer,
+                bucket_size=self.position_buckets,
+                max_position=self.max_relative_positions,
+            )
+        if relative_pos.dim() == 2:
+            relative_pos = relative_pos.unsqueeze(0).unsqueeze(0)
+        elif relative_pos.dim() == 3:
+            relative_pos = relative_pos.unsqueeze(1)
+        # bsz x height x query x key
+        elif relative_pos.dim() != 4:
+            raise ValueError(f"Relative position ids must be of dim 2 or 3 or 4. {relative_pos.dim()}")
+
+        att_span = self.pos_ebd_size
+        relative_pos = relative_pos.to(device=query_layer.device, dtype=torch.long)
+
+        rel_embeddings = rel_embeddings[0 : att_span * 2, :].unsqueeze(0)
+        if self.share_att_key:
+            pos_query_layer = self.transpose_for_scores(
+                self.query_proj(rel_embeddings), self.num_attention_heads
+            ).repeat(query_layer.size(0) // self.num_attention_heads, 1, 1)
+            pos_key_layer = self.transpose_for_scores(self.key_proj(rel_embeddings), self.num_attention_heads).repeat(
+                query_layer.size(0) // self.num_attention_heads, 1, 1
+            )
+        else:
+            if "c2p" in self.pos_att_type:
+                pos_key_layer = self.transpose_for_scores(
+                    self.pos_key_proj(rel_embeddings), self.num_attention_heads
+                ).repeat(query_layer.size(0) // self.num_attention_heads, 1, 1)  # .split(self.all_head_size, dim=-1)
+            if "p2c" in self.pos_att_type:
+                pos_query_layer = self.transpose_for_scores(
+                    self.pos_query_proj(rel_embeddings), self.num_attention_heads
+                ).repeat(query_layer.size(0) // self.num_attention_heads, 1, 1)  # .split(self.all_head_size, dim=-1)
+
+        score = 0
+        # content->position
+        if "c2p" in self.pos_att_type:
+            scale = scaled_size_sqrt(pos_key_layer, scale_factor)
+            c2p_att = torch.bmm(query_layer, pos_key_layer.transpose(-1, -2))
+            c2p_pos = torch.clamp(relative_pos + att_span, 0, att_span * 2 - 1)
+            c2p_att = torch.gather(
+                c2p_att,
+                dim=-1,
+                index=c2p_pos.squeeze(0).expand([query_layer.size(0), query_layer.size(1), relative_pos.size(-1)]),
+            )
+            score += c2p_att / scale.to(dtype=c2p_att.dtype)
+
+        # position->content
+        if "p2c" in self.pos_att_type:
+            scale = scaled_size_sqrt(pos_query_layer, scale_factor)
+            r_pos = build_rpos(
+                query_layer,
+                key_layer,
+                relative_pos,
+                self.max_relative_positions,
+                self.position_buckets,
+            )
+            p2c_pos = torch.clamp(-r_pos + att_span, 0, att_span * 2 - 1)
+            p2c_att = torch.bmm(key_layer, pos_query_layer.transpose(-1, -2))
+            p2c_att = torch.gather(
+                p2c_att,
+                dim=-1,
+                index=p2c_pos.squeeze(0).expand([query_layer.size(0), key_layer.size(-2), key_layer.size(-2)]),
+            ).transpose(-1, -2)
+            score += p2c_att / scale.to(dtype=p2c_att.dtype)
+
+        return score
+
+
+class QiDeBERTaSelfOutput(Module):
+    def __init__(
+        self,
+        d_model: int,
+        layer_norm_eps: float,
+        hidden_dropout_prob: float,
+    ):
         super().__init__()
+        self.dense = Linear(in_features=d_model, out_features=d_model)
+        self.LayerNorm = LayerNorm(normalized_shape=d_model, eps=layer_norm_eps)
+        self.dropout = Dropout(p=hidden_dropout_prob)
+
+    def forward(self, hidden_states, input_tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+class QiDeBERTaAttention(Module):
+    def __init__(
+        self,
+        num_heads: int,
+        d_model: int,
+        share_att_key: bool,
+        relative_attention: bool,
+        max_position_embeddings: int,
+        hidden_dropout_prob: float,
+        attention_probs_dropout_prob: float,
+        layer_norm_eps: float,
+        pos_att_type: Optional[list] = None,
+        position_buckets: int = -1,
+        max_relative_positions: int = -1,
+    ):
+        super().__init__()
+        self.self = QiDeBERTaDisentangledSelfAttention(
+            num_heads=num_heads,
+            d_model=d_model,
+            share_att_key=share_att_key,
+            relative_attention=relative_attention,
+            max_position_embeddings=max_position_embeddings,
+            hidden_dropout_prob=hidden_dropout_prob,
+            attention_probs_dropout_prob=attention_probs_dropout_prob,
+            pos_att_type=pos_att_type,
+            position_buckets=position_buckets,
+            max_relative_positions=max_relative_positions,
+        )
+        self.output = QiDeBERTaSelfOutput(
+            d_model=d_model,
+            layer_norm_eps=layer_norm_eps,
+            hidden_dropout_prob=hidden_dropout_prob,
+        )
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask,
+        output_attentions: bool = False,
+        relative_pos=None,
+        rel_embeddings=None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        self_output, att_matrix = self.self(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            relative_pos=relative_pos,
+            rel_embeddings=rel_embeddings,
+        )
+        attention_output = self.output(hidden_states=self_output, input_tensor=hidden_states)
+
+        return (attention_output, att_matrix) if output_attentions else (attention_output, None)
+
+
+class QiDeBERTaIntermediate(Module):
+    def __init__(
+        self,
+        d_model: int,
+        d_ff: int,
+    ):
+        super().__init__()
+        self.dense = Linear(in_features=d_model, out_features=d_ff)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.dense(hidden_states)
+        hidden_states = functional.gelu(hidden_states)
+        return hidden_states
+
+
+class QiDeBERTaOutput(Module):
+    def __init__(
+        self,
+        d_ff: int,
+        d_model: int,
+        layer_norm_eps: float,
+        hidden_dropout_prob: float,
+    ):
+        super().__init__()
+        self.dense = Linear(in_features=d_ff, out_features=d_model)
+        self.LayerNorm = LayerNorm(normalized_shape=d_model, eps=layer_norm_eps)
+        self.dropout = Dropout(p=hidden_dropout_prob)
+
+    def forward(self, hidden_states, input_tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
 
-        self.layer = ModuleList([DebertaV2Layer(config) for _ in range(config.num_hidden_layers)])
 
-        self.max_relative_positions = config.max_position_embeddings
+class QiDeBERTaLayer(Module):
+    def __init__(
+        self,
+        num_heads: int,
+        d_model: int,
+        d_ff: int,
+        share_att_key: bool,
+        relative_attention: bool,
+        max_position_embeddings: int,
+        hidden_dropout_prob: float,
+        attention_probs_dropout_prob: float,
+        layer_norm_eps: float,
+        pos_att_type: Optional[list] = None,
+        position_buckets: int = -1,
+        max_relative_positions: int = -1,
+    ):
+        super().__init__()
+        self.attention = QiDeBERTaAttention(
+            num_heads=num_heads,
+            d_model=d_model,
+            share_att_key=share_att_key,
+            relative_attention=relative_attention,
+            max_position_embeddings=max_position_embeddings,
+            hidden_dropout_prob=hidden_dropout_prob,
+            attention_probs_dropout_prob=attention_probs_dropout_prob,
+            layer_norm_eps=layer_norm_eps,
+            pos_att_type=pos_att_type,
+            position_buckets=position_buckets,
+            max_relative_positions=max_relative_positions,
+        )
+        self.intermediate = QiDeBERTaIntermediate(
+            d_model=d_model,
+            d_ff=d_ff
+        )
+        self.output = QiDeBERTaOutput(
+            d_ff=d_ff,
+            d_model=d_model,
+            layer_norm_eps=layer_norm_eps,
+            hidden_dropout_prob=hidden_dropout_prob,
+        )
 
-        self.position_buckets = config.position_buckets
+    def forward(
+        self,
+        hidden_states,
+        attention_mask,
+        relative_pos=None,
+        rel_embeddings=None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        attention_output, att_matrix = self.attention(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            relative_pos=relative_pos,
+            rel_embeddings=rel_embeddings,
+        )
+        intermediate_output = self.intermediate(attention_output)
+        layer_output = self.output(intermediate_output, attention_output)
 
-        pos_ebd_size = self.position_buckets * 2
+        return (layer_output, att_matrix) if output_attentions else (layer_output, None)
 
-        self.rel_embeddings = Embedding(pos_ebd_size, config.hidden_size)
 
-        self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps, elementwise_affine=True)
+class QiDeBERTaEncoder(Module):
+    """Modified BertEncoder with relative position bias support"""
+    def __init__(
+        self,
+        num_layers: int,
+        num_heads: int,
+        d_model: int,
+        d_ff: int,
+        share_att_key: bool,
+        relative_attention: bool,
+        max_position_embeddings: int,
+        hidden_dropout_prob: float,
+        attention_probs_dropout_prob: float,
+        layer_norm_eps: float,
+        pos_att_type: Optional[list] = None,
+        position_buckets: int = -1,
+        max_relative_positions: int = -1,
+    ):
+        super().__init__()
+
+        self.layer = ModuleList([
+            QiDeBERTaLayer(
+                num_heads=num_heads,
+                d_model=d_model,
+                d_ff=d_ff,
+                share_att_key=share_att_key,
+                relative_attention=relative_attention,
+                max_position_embeddings=max_position_embeddings,
+                hidden_dropout_prob=hidden_dropout_prob,
+                attention_probs_dropout_prob=attention_probs_dropout_prob,
+                layer_norm_eps=layer_norm_eps,
+                pos_att_type=pos_att_type,
+                position_buckets=position_buckets,
+                max_relative_positions=max_relative_positions,
+            )
+            for _ in range(num_layers)
+        ])
+
+        self.max_relative_positions = max_position_embeddings
+
+        self.position_buckets = position_buckets
+
+        pos_ebd_size = position_buckets * 2
+        self.rel_embeddings = Embedding(num_embeddings=pos_ebd_size, embedding_dim=d_model)
+
+        self.LayerNorm = LayerNorm(normalized_shape=d_model, eps=layer_norm_eps, elementwise_affine=True)
 
         self.gradient_checkpointing = False
 
@@ -72,12 +477,10 @@ class QiDeBERTaEncoder(Module):
             rel_embeddings = self.LayerNorm(rel_embeddings)
         return rel_embeddings
 
-    def get_attention_mask(self, attention_mask):
-        if attention_mask.dim() <= 2:
-            extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
-            attention_mask = extended_attention_mask * extended_attention_mask.squeeze(-2).unsqueeze(-1)
-        elif attention_mask.dim() == 3:
-            attention_mask = attention_mask.unsqueeze(1)
+    @staticmethod
+    def get_attention_mask(attention_mask):
+        extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+        attention_mask = extended_attention_mask * extended_attention_mask.squeeze(-2).unsqueeze(-1)
 
         return attention_mask
 
@@ -94,6 +497,7 @@ class QiDeBERTaEncoder(Module):
         self,
         hidden_states,
         attention_mask,
+        output_attentions: bool = True,
     ):
         attention_mask = self.get_attention_mask(attention_mask)
         relative_pos = self.get_rel_pos(hidden_states)
@@ -109,34 +513,35 @@ class QiDeBERTaEncoder(Module):
                     layer_module.__call__,
                     next_kv,
                     attention_mask,
-                    None,
                     relative_pos,
                     rel_embeddings,
-                    True,
+                    output_attentions,
                 )
             else:
                 output_states, attn_weights = layer_module(
-                    next_kv,
-                    attention_mask,
-                    query_states=None,
+                    hidden_states=next_kv,
+                    attention_mask=attention_mask,
                     relative_pos=relative_pos,
                     rel_embeddings=rel_embeddings,
-                    output_attentions=True,
+                    output_attentions=output_attentions,
                 )
 
-            all_attentions = all_attentions + (attn_weights,)
+            if output_attentions:
+                all_attentions = all_attentions + (attn_weights,)
 
             all_hidden_states = all_hidden_states + (output_states,)
 
             next_kv = output_states
 
-        return BaseModelOutput(
-            last_hidden_state=output_states, hidden_states=all_hidden_states, attentions=all_attentions
+        return EncoderOutput(
+            last_hidden_state=output_states,
+            hidden_states=all_hidden_states,
+            attentions=all_attentions if output_attentions else None
         )
 
 
 class QiDeBERTaBase(DebertaV2PreTrainedModel):
-    VERSION = '1.0.1'
+    VERSION = '1.1.0'
     config_class = QiDeBERTaConfig
     base_model_prefix = 'qideberta'
     _encoder_layer_path = ''
@@ -171,7 +576,7 @@ class QiDeBERTaBase(DebertaV2PreTrainedModel):
         else:
             encoder_layer.requires_grad_(requires_grad=True)
 
-    def freeze_embed_layer(self, freeze: bool = True):
+    def freeze_encoder_embed_layer(self, freeze: bool = True):
         """
         Freeze the embedding layer
         :param freeze:
@@ -198,10 +603,26 @@ class QiDeBERTa(QiDeBERTaBase):
 
         self.embeddings = QiDeBERTaEmbeddings(
             pad_token_id=config.pad_token_id,
-            hidden_size=config.hidden_size,
+            d_model=config.d_model,
             vocab_size=config.vocab_size,
+            layer_norm_eps=config.layer_norm_eps,
+            hidden_dropout_prob=config.hidden_dropout_prob,
+        )
+        self.encoder = QiDeBERTaEncoder(
+            num_layers=config.num_layers,
+            num_heads=config.num_heads,
+            max_position_embeddings=config.max_position_embeddings,
+            position_buckets=config.position_buckets,
+            d_model=config.d_model,
+            d_ff=config.d_ff,
+            layer_norm_eps=config.layer_norm_eps,
+            share_att_key=config.share_att_key,
+            relative_attention=config.relative_attention,
+            hidden_dropout_prob=config.hidden_dropout_prob,
+            attention_probs_dropout_prob=config.attention_probs_dropout_prob,
+            pos_att_type=config.pos_att_type,
+            max_relative_positions=config.max_relative_positions,
         )
-        self.encoder = QiDeBERTaEncoder(config)
         # Initialize weights and apply final processing
         self.post_init()
 
@@ -213,84 +634,59 @@ class QiDeBERTa(QiDeBERTaBase):
 
     def forward(
         self,
-        input_ids: Optional[torch.Tensor] = None,
+        input_ids: torch.Tensor,
         attention_mask: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        deep_recurrent_refinement_steps: int = 0,
+        output_attentions: bool = True,
     ) -> BaseModelOutput:
         """
         Forward pass of the model
 
-        :param input_ids:
+        :param input_ids: Token indices of input sequence tokens in the vocabulary. (batch_size, sequence_length)
         :param attention_mask: Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
             - 1 for tokens that are **not masked**,
             - 0 for tokens that are **masked**.
-        :param inputs_embeds:
-        :param deep_recurrent_refinement_steps:
+            (batch_size, sequence_length)
+        :param output_attentions:
         :return:
         """
-        if input_ids is not None and inputs_embeds is not None:
-            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
-        if input_ids is not None:
-            self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
-            input_shape = input_ids.size()
-        elif inputs_embeds is not None:
-            input_shape = inputs_embeds.size()[:-1]
-        else:
-            raise ValueError("You have to specify either input_ids or inputs_embeds")
-
-        device = input_ids.device if input_ids is not None else inputs_embeds.device
+        self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
+        input_shape = input_ids.size()
+        device = input_ids.device
 
         if attention_mask is None:
             attention_mask = torch.ones(input_shape, device=device)
 
-        embedding_output = self.embeddings(
+        embedding_output, token_embeddings = self.embeddings(
             input_ids=input_ids,
             mask=attention_mask,
-            inputs_embeds=inputs_embeds,
         )
 
         encoder_outputs = self.encoder(
             hidden_states=embedding_output,
             attention_mask=attention_mask,
+            output_attentions=output_attentions,
         )
-        encoded_layers = encoder_outputs.hidden_states
-
-        if deep_recurrent_refinement_steps > 1:
-            hidden_states = encoded_layers[-2]
-            layers = [self.encoder.layer[-1] for _ in range(self.z_steps)]
-            query_states = encoded_layers[-1]
-            rel_embeddings = self.encoder.get_rel_embedding()
-            attention_mask = self.encoder.get_attention_mask(attention_mask)
-            rel_pos = self.encoder.get_rel_pos(embedding_output)
-            for layer in layers[1:]:
-                query_states = layer(
-                    hidden_states,
-                    attention_mask,
-                    output_attentions=False,
-                    query_states=query_states,
-                    relative_pos=rel_pos,
-                    rel_embeddings=rel_embeddings,
-                )
-                encoded_layers.append(query_states)
 
         return BaseModelOutput(
             last_hidden_state=encoder_outputs.last_hidden_state,
             hidden_states=encoder_outputs.hidden_states,
-            attentions=encoder_outputs.attentions,
+            attentions=encoder_outputs.attentions if output_attentions else None,
+            embedding_output=embedding_output,
+            token_embeddings=token_embeddings,  # [B, L, H]
         )
 
 
 class QiDeBERTaMLMHead(Module):
     def __init__(
         self,
-        hidden_size: int,
-        vocab_size: int
+        d_model: int,
+        vocab_size: int,
+        layer_norm_eps: float,
     ):
         super().__init__()
-        self.dense = Linear(in_features=hidden_size, out_features=hidden_size)
+        self.dense = Linear(in_features=d_model, out_features=d_model)
 
-        self.LayerNorm = LayerNorm(normalized_shape=hidden_size, eps=1e-7, elementwise_affine=True)
+        self.LayerNorm = LayerNorm(normalized_shape=d_model, eps=layer_norm_eps, elementwise_affine=True)
 
         self.bias = Parameter(torch.zeros(vocab_size))
 
@@ -302,26 +698,101 @@ class QiDeBERTaMLMHead(Module):
             if module.bias is not None:
                 module.bias.data.zero_()
 
+    def _initialize_weights(self, module):
+        if getattr(module, "_is_hf_initialized", False):
+            return
+        self._init_weights(module)
+        module._is_hf_initialized = True
+
     def forward(self, hidden_states: torch.Tensor, word_embeddings: Embedding):
         hidden_states = self.dense(hidden_states)
-        hidden_states = gelu(hidden_states)
+        hidden_states = functional.gelu(hidden_states)
         hidden_states = self.LayerNorm(hidden_states)
         hidden_states = torch.matmul(hidden_states, word_embeddings.weight.t()) + self.bias
         return hidden_states
 
 
+class QiDeBERTaClassificationHead(Module):
+    def __init__(
+        self,
+        d_model: int,
+        num_labels: int,
+        hidden_dropout_prob: float,
+    ):
+        super().__init__()
+
+        self.dropout = Dropout(p=hidden_dropout_prob)
+        self.classifier = Linear(in_features=d_model, out_features=num_labels)
+
+    @staticmethod
+    def _init_weights(module):
+        """Initialize the weights."""
+        if isinstance(module, Linear):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+            if module.bias is not None:
+                module.bias.data.zero_()
+
+    def _initialize_weights(self, module):
+        if getattr(module, "_is_hf_initialized", False):
+            return
+        self._init_weights(module)
+        module._is_hf_initialized = True
+
+    def forward(self, hidden_states: torch.Tensor):
+        dropped = self.dropout(hidden_states)
+        logits = self.classifier(dropped)
+
+        return logits
+
+
+class ContextPooler(Module):
+    def __init__(
+        self,
+        pooler_hidden_size: int = 1024,
+        pooler_mode: str = 'token',  # mean, max, attn, token
+    ):
+        super().__init__()
+        if pooler_mode not in ['mean', 'max', 'attn', 'token']:
+            raise ValueError(f'Invalid pooler mode: {pooler_mode}')
+
+        self.dense = Linear(in_features=pooler_hidden_size, out_features=pooler_hidden_size)
+        self.pooler_mode = pooler_mode
+        if self.pooler_mode == 'attn':
+            self.attn = MultiheadAttention(embed_dim=pooler_hidden_size, num_heads=pooler_hidden_size, batch_first=True)
+            self.LayerNorm = LayerNorm(normalized_shape=pooler_hidden_size, eps=1e-7)
+
+    def forward(self, hidden_states):
+        if self.pooler_mode == 'attn':
+            query = hidden_states[:, 0:1, :]
+            attn_output, _ = self.attn(query, hidden_states, hidden_states)  # [batch_size, 1, hidden_size]
+            attn_output = attn_output.squeeze(1)  # [batch_size, hidden_size]
+            context_token = attn_output + hidden_states[:, 0]  # 残差连接
+            context_token = self.LayerNorm(context_token)  # 仅对 attn 方式归一化
+        elif self.pooler_mode == 'mean':
+            context_token = hidden_states.mean(dim=1)  # 计算所有 token 的平均表示
+        elif self.pooler_mode == 'max':
+            context_token = hidden_states.max(dim=1).values  # 计算所有 token 的最大表示
+        elif self.pooler_mode == 'token':
+            context_token = hidden_states[:, 0]  # 取第一个 token 的表示
+
+        pooled_output = self.dense(context_token)
+        pooled_output = gelu(pooled_output)
+        return pooled_output
+
+
 class QiDeBERTaForMaskedLM(QiDeBERTaBase):
     _tied_weights_keys = ["mlm_head.weight", "qideberta.embeddings.word_embeddings.weight"]
     _encoder_layer_path = 'qideberta.encoder'
     _embedding_layer_path = 'qideberta.embeddings'
-    task_head = 'mlm_head'
+    task_head = ['mlm_head']
 
     def __init__(self, config: QiDeBERTaConfig):
         super().__init__(config)
         self.qideberta = QiDeBERTa(config=config)
         self.mlm_head = QiDeBERTaMLMHead(
-            hidden_size=config.hidden_size,
-            vocab_size=config.vocab_size
+            d_model=config.d_model,
+            vocab_size=config.vocab_size,
+            layer_norm_eps=config.layer_norm_eps,
         )
 
         self.post_init()
@@ -346,7 +817,7 @@ class QiDeBERTaForMaskedLM(QiDeBERTaBase):
             deep_recurrent_refinement_steps=0,
         )
 
-        prediction_scores = self.mlm_head(hidden_states=outputs.last_hidden_state, word_embeddings=self.qideberta.embeddings.word_embeddings)
+        prediction_scores = self.mlm_head(hidden_states=outputs.last_hidden_state, word_embeddings=self.get_output_embeddings())
 
         return MaskedLMOutput(
             logits=prediction_scores,

README.md

@@ -98,7 +98,7 @@ texts = [
   "我爱北京天安门，天安门上太阳升。"
 ]
 
-outputs = model(**tokenizer(texts, padding=True))  # BaseModelOutput[last_hidden_state, hidden_states, attentions]
+outputs = model(**tokenizer(texts, padding=True, return_tensors='pt'))  # BaseModelOutput[last_hidden_state, hidden_states, attentions]
 ```
 
 ## Citation

config.json

@@ -5,25 +5,26 @@
   ],
   "attention_probs_dropout_prob": 0.1,
   "auto_map": {
-    "AutoConfig": "Configuration.QiDeBERTaConfig",
+    "AutoConfig": "configuration.QiDeBERTaConfig",
     "AutoModel": "QiDeBERTa.QiDeBERTa",
     "AutoModelForMaskedLM": "QiDeBERTa.QiDeBERTaForMaskedLM"
   },
   "bos_token_id": 1,
+  "classifier_num_labels": -1,
+  "d_ff": 3072,
+  "d_model": 768,
   "eos_token_id": 2,
   "hidden_act": "gelu",
   "hidden_dropout_prob": 0.1,
-  "hidden_size": 768,
   "initializer_range": 0.02,
-  "intermediate_size": 3072,
   "layer_norm_eps": 1e-07,
   "mask_token_id": 4,
   "max_position_embeddings": 512,
   "max_relative_positions": -1,
   "model_type": "QiDeBERTa",
   "norm_rel_ebd": "layer_norm",
-  "num_attention_heads": 12,
-  "num_hidden_layers": 12,
+  "num_heads": 12,
+  "num_layers": 12,
   "pad_token_id": 3,
   "pooler_hidden_size": 768,
   "pooler_mode": "token",
@@ -36,7 +37,7 @@
   "relative_attention": true,
   "share_att_key": true,
   "torch_dtype": "float32",
-  "transformers_version": "4.50.0",
+  "transformers_version": "4.52.4",
   "unk_token_id": 0,
   "vocab_size": 25500
 }

Configuration.py → configuration.py

@@ -3,14 +3,20 @@ from transformers import PretrainedConfig
 
 class QiDeBERTaConfig(PretrainedConfig):
     model_type = "QiDeBERTa"
+    attribute_map = {
+        "hidden_size": "d_model",
+        "num_attention_heads": "num_heads",
+        "num_hidden_layers": "num_layers",
+        "intermediate_size": "d_ff",
+    }
 
     def __init__(
         self,
         vocab_size=25500,
-        hidden_size=1024,
-        num_hidden_layers=24,
-        num_attention_heads=16,
-        intermediate_size=4096,
+        d_model=1024,
+        num_layers=24,
+        num_heads=16,
+        d_ff=4096,
         hidden_act="gelu",
         hidden_dropout_prob=0.1,
         attention_probs_dropout_prob=0.1,
@@ -20,6 +26,7 @@ class QiDeBERTaConfig(PretrainedConfig):
         relative_attention=True,
         max_relative_positions=-1,
         norm_rel_ebd='layer_norm',
+        classifier_num_labels=-1,
         unk_token_id=0,
         bos_token_id=1,
         eos_token_id=2,
@@ -34,10 +41,10 @@ class QiDeBERTaConfig(PretrainedConfig):
     ):
         super().__init__(**kwargs)
 
-        self.hidden_size = hidden_size
-        self.num_hidden_layers = num_hidden_layers
-        self.num_attention_heads = num_attention_heads
-        self.intermediate_size = intermediate_size
+        self.d_model = d_model
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.d_ff = d_ff
         self.hidden_act = hidden_act
         self.hidden_dropout_prob = hidden_dropout_prob
         self.attention_probs_dropout_prob = attention_probs_dropout_prob
@@ -45,6 +52,7 @@ class QiDeBERTaConfig(PretrainedConfig):
         self.initializer_range = initializer_range
         self.relative_attention = relative_attention
         self.max_relative_positions = max_relative_positions
+        self.classifier_num_labels = classifier_num_labels
         self.unk_token_id = unk_token_id
         self.bos_token_id = bos_token_id
         self.eos_token_id = eos_token_id
@@ -63,5 +71,5 @@ class QiDeBERTaConfig(PretrainedConfig):
         self.vocab_size = vocab_size
         self.layer_norm_eps = layer_norm_eps
 
-        self.pooler_hidden_size = kwargs.get("pooler_hidden_size", hidden_size)
+        self.pooler_hidden_size = kwargs.get("pooler_hidden_size", d_model)
         self.pooler_mode = pooler_mode

tokenizer.py

@@ -1,12 +1,12 @@
 import os
-from typing import Optional, Dict, Any, List, Tuple, Union
+from typing import Optional, Dict, Any, Tuple
 
 import sentencepiece
-import torch
-from torch.nn.utils.rnn import pad_sequence
+from torch import TensorType
 from transformers import PreTrainedTokenizer
 from transformers.models.deberta_v2.tokenization_deberta_v2 import SPMTokenizer
-from transformers.tokenization_utils_base import TextInput, BatchEncoding
+from transformers.tokenization_utils_base import TextInput, PreTokenizedInput, TruncationStrategy
+from transformers.utils import PaddingStrategy
 
 
 class QiDeBERTaTokenizer(PreTrainedTokenizer):
@@ -55,6 +55,7 @@ class QiDeBERTaTokenizer(PreTrainedTokenizer):
             **kwargs,
         )
         self._tokenizer.special_tokens = self.all_special_tokens
+        self.space_token_id = self._tokenizer.spm.PieceToId('▁')
 
     @property
     def vocab_size(self):
@@ -66,63 +67,66 @@ class QiDeBERTaTokenizer(PreTrainedTokenizer):
 
     def __call__(
         self,
-        texts: Union[str, list[str]],
+        text: TextInput|PreTokenizedInput|list[TextInput]|list[PreTokenizedInput],
+        text_pair: Optional[TextInput|PreTokenizedInput|list[TextInput]|list[PreTokenizedInput]] = None,
+        text_target: Optional[TextInput|PreTokenizedInput|list[TextInput]|list[PreTokenizedInput]] = None,
+        text_pair_target: Optional[TextInput|PreTokenizedInput|list[TextInput]|list[PreTokenizedInput]] = None,
         add_special_tokens: bool = True,
-        padding: bool = False,
+        padding: bool|str|PaddingStrategy = False,
+        truncation: Optional[bool|str|TruncationStrategy] = None,
+        max_length: Optional[int] = None,
+        stride: int = 0,
+        is_split_into_words: bool = False,
+        pad_to_multiple_of: Optional[int] = None,
+        padding_side: Optional[str] = None,
+        return_tensors: str|TensorType = 'pt',
+        return_token_type_ids: bool = False,
         return_attention_mask: bool = True,
-    ) -> BatchEncoding:
-        """
-        对输入文本进行编码，返回token ids和attention mask
-        :return:
-        """
-        if isinstance(texts, str):
-            texts = [texts]
-        if (isinstance(texts, list) and all(isinstance(text, str) for text in texts)) is not True:
-            raise ValueError(
-                f"Input must be a string or a list of strings, but got {type(texts)}"
-            )
-
-        if len(texts) > 1 and all(len(text) == len(texts[0]) for text in texts) is False and padding is False:
-            # Torch不支持可变长度的tensor，所以需要padding
-            print(f'[Warning] The input texts are not the same length, padding is required.')
-            padding = True
-
-        if padding:
-            token_ids = pad_sequence(
-                sequences=[torch.LongTensor(input_id) for input_id in self.encode(texts=texts, add_special_tokens=add_special_tokens)],
-                batch_first=True,
-                padding_value=self.processor().pad_id()
-            )
-        else:
-            token_ids = torch.LongTensor(self.encode(texts=texts, add_special_tokens=add_special_tokens))
-
-        if return_attention_mask:
-            return BatchEncoding(
-                data={
-                    'input_ids': token_ids,
-                    'attention_mask': token_ids != self.processor().pad_id()
-                }
-            )
-        else:
-            return BatchEncoding(
-                data={
-                    'input_ids': token_ids
-                }
-            )
+        return_overflowing_tokens: bool = False,
+        return_special_tokens_mask: bool = False,
+        return_offsets_mapping: bool = False,
+        return_length: bool = False,
+        verbose: bool = True,
+        **kwargs,
+    ):
+        return super().__call__(
+            text=text,
+            text_pair=text_pair,
+            text_target=text_target,
+            text_pair_target=text_pair_target,
+            add_special_tokens=add_special_tokens,
+            padding=padding,
+            truncation=truncation,
+            max_length=max_length,
+            stride=stride,
+            is_split_into_words=is_split_into_words,
+            pad_to_multiple_of=pad_to_multiple_of,
+            padding_side=padding_side,
+            return_tensors=return_tensors,
+            return_token_type_ids=return_token_type_ids,
+            return_attention_mask=return_attention_mask,
+            return_overflowing_tokens=return_overflowing_tokens,
+            return_special_tokens_mask=return_special_tokens_mask,
+            return_offsets_mapping=return_offsets_mapping,
+            return_length=return_length,
+            verbose=verbose,
+            **kwargs,
+        )
 
     def get_vocab(self):
         vocab = self.vocab.copy()
         vocab.update(self.get_added_vocab())
         return vocab
 
-    def _tokenize(self, text: str) -> List[str]:
+    def _tokenize(self, text: str) -> list[str]:
         """Take as input a string and return a list of strings (tokens) for words/sub-words"""
         if self.do_lower_case:
             text = text.lower()
         return self._tokenizer.tokenize(text)
 
-    def tokenize(self, text: TextInput, **kwargs) -> List[str]:
-        return super().tokenize(text, **kwargs)[1:]
+    def tokenize(self, text: TextInput, **kwargs) -> list[str]:
+        result = super().tokenize(text, **kwargs)
+        return result[1:] if result[0] == '▁' else result
 
     def _convert_token_to_id(self, token: str):
         """Converts a token (str) in an id using the vocab."""
@@ -229,34 +233,6 @@ class QiDeBERTaTokenizer(PreTrainedTokenizer):
     def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
         return self._tokenizer.save_pretrained(save_directory, filename_prefix=filename_prefix)
 
-    def encode(self, texts: str|list[str], add_special_tokens: bool = True) -> list[int]|list[list[int]]:
-        """
-        编码文本
-        :param texts:
-        :return:
-        """
-        if isinstance(texts, str):
-            return [self._tokenizer.spm.bos_id()] + self._tokenizer.spm.encode_as_ids(texts)[1:] + [self._tokenizer.spm.eos_id()] if add_special_tokens else self._tokenizer.spm.encode_as_ids(texts)[1:]
-        elif isinstance(texts, list):
-            return [
-                [self._tokenizer.spm.bos_id()] + ids[1:] + [self._tokenizer.spm.eos_id()] if add_special_tokens else ids[1:]
-                for ids in self._tokenizer.spm.encode_as_ids(texts)
-            ]
-
-    def decode(self, token_ids: list[int] or list[list[int]]) -> list[str] or list[list[str]]:
-        """
-        解码文本
-        :param token_ids:
-        :return:
-        """
-        if token_ids and isinstance(token_ids[0], list):
-            return [
-                self._tokenizer.spm.DecodeIds(input=ids)
-                for ids in token_ids
-            ]
-        elif token_ids and isinstance(token_ids[0], int):
-            return self._tokenizer.spm.DecodeIds(input=token_ids)
-
     def _get_bos_piece(self) -> str:
         """
         获取BOS Piece