text: import logging
  import math
  import os
  
  import torch
  from torch import nn
  from torch.nn import CrossEntropyLoss, MSELoss
  
  def swish(x):
      return x * torch.sigmoid(x)
  
  def gelu(x):
      """ 
      """
      return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))
  
  def mish(x):
      return x * torch.tanh(nn.functional.softplus(x))
  
  ACT2FN = {"gelu": gelu, "relu": torch.nn.functional.relu, "swish": swish, "mish": mish}
  
  class BertConfig(object):
      
      def __init__(
          self,
          vocab_size,
          hidden_size=768,
          num_hidden_layers=12,
          num_attention_heads=12,
          intermediate_size=3072,
          hidden_act="gelu",
          hidden_dropout_prob=0.1,
          attention_probs_dropout_prob=0.1,
          max_position_embeddings=512,
          type_vocab_size=2,
          initializer_range=0.02,
          layer_norm_eps=1e-12,
      ):
  
          self.vocab_size = vocab_size
          self.hidden_size = hidden_size
          self.num_hidden_layers = num_hidden_layers
          self.num_attention_heads = num_attention_heads
          self.hidden_act = hidden_act
          self.intermediate_size = intermediate_size
          self.hidden_dropout_prob = hidden_dropout_prob
          self.attention_probs_dropout_prob = attention_probs_dropout_prob
          self.max_position_embeddings = max_position_embeddings
          self.type_vocab_size = type_vocab_size
          self.initializer_range = initializer_range
          self.layer_norm_eps = layer_norm_eps
  
  class RobertaLargeConfig(object):
      
      def __init__(
          self,
          vocab_size,
          hidden_size=1024,
          num_hidden_layers=24,
          num_attention_heads=16,
          intermediate_size=4096,
          hidden_act="gelu",
          hidden_dropout_prob=0.1,
          attention_probs_dropout_prob=0.1,
          max_position_embeddings=512,
          type_vocab_size=2,
          initializer_range=0.02,
          layer_norm_eps=1e-12,
      ):
  
          self.vocab_size = vocab_size
          self.hidden_size = hidden_size
          self.num_hidden_layers = num_hidden_layers
          self.num_attention_heads = num_attention_heads
          self.hidden_act = hidden_act
          self.intermediate_size = intermediate_size
          self.hidden_dropout_prob = hidden_dropout_prob
          self.attention_probs_dropout_prob = attention_probs_dropout_prob
          self.max_position_embeddings = max_position_embeddings
          self.type_vocab_size = type_vocab_size
          self.initializer_range = initializer_range
          self.layer_norm_eps = layer_norm_eps
  
  
  class BertLayerNorm(nn.Module):
      """LayerNorm层, 见Transformer(一), 讲编码器(encoder)的第3部分"""
      def __init__(self, hidden_size, eps=1e-12, conditional=False):
          """Construct a layernorm module in the TF style (epsilon inside the square root).
          """
          super(BertLayerNorm, self).__init__()
          self.weight = nn.Parameter(torch.ones(hidden_size))
          self.bias = nn.Parameter(torch.zeros(hidden_size))
          self.variance_epsilon = eps
          self.conditional = conditional
          if conditional == True:
              #说明是条件 ln
              self.weight_dense = nn.Linear(2 * hidden_size, hidden_size, bias=False)
              self.weight_dense.weight.data.uniform_(0, 0)
              self.bias_dense = nn.Linear(2 * hidden_size, hidden_size, bias=False)
              self.bias_dense.weight.data.uniform_(0, 0)
  
      def forward(self, x):
          if self.conditional == False:
              u = x.mean(-1, keepdim=True)
              s = (x - u).pow(2).mean(-1, keepdim=True)
              x = (x - u) / torch.sqrt(s + self.variance_epsilon)
              return self.weight * x + self.bias
          else :
              inputs = x[0]
              cond = x[1]
              for _ in range(len(inputs.shape) - len(cond.shape)):
                  cond = cond.unsqueeze(dim=1)
             
              weight = self.weight + self.weight_dense(cond)
              bias = self.bias + self.bias_dense(cond)
              u = inputs.mean(-1, keepdim=True)
              s = (inputs - u).pow(2).mean(-1, keepdim=True)
              x = (inputs - u) / torch.sqrt(s + self.variance_epsilon)
             
              return weight * x + bias
  
  
  
  class BertEmbeddings(nn.Module):
      """Construct the embeddings from word, position and token_type embeddings.
      """
  
      def __init__(self, config):
          super().__init__()
          self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=0)
          self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
          self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
  
          self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
          self.dropout = nn.Dropout(config.hidden_dropout_prob)
  
      def forward(self, input_ids, token_type_ids=None, position_ids=None):
       
          input_shape = input_ids.size()
          
          seq_length = input_shape[1]
          device = input_ids.device
          if position_ids is None:
              position_ids = torch.arange(seq_length, dtype=torch.long, device=device)
              position_ids = position_ids.unsqueeze(0).expand(input_shape)
          if token_type_ids is None:
              token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
  
          inputs_embeds = self.word_embeddings(input_ids)
          position_embeddings = self.position_embeddings(position_ids)
          token_type_embeddings = self.token_type_embeddings(token_type_ids)
  
          embeddings = inputs_embeds + position_embeddings + token_type_embeddings
          embeddings = self.LayerNorm(embeddings)
          embeddings = self.dropout(embeddings)
          return embeddings
  
  
  class BertSelfAttention(nn.Module):
      def __init__(self, config: BertConfig):
          super().__init__()
          if config.hidden_size % config.num_attention_heads != 0:
              raise ValueError(
                  "The hidden size (%d) is not a multiple of the number of attention "
                  "heads (%d)" % (config.hidden_size, config.num_attention_heads)
              )
          
          self.num_attention_heads = config.num_attention_heads
          self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
          self.all_head_size = self.num_attention_heads * self.attention_head_size
  
          self.query = nn.Linear(config.hidden_size, self.all_head_size)
          self.key = nn.Linear(config.hidden_size, self.all_head_size)
          self.value = nn.Linear(config.hidden_size, self.all_head_size)
  
          self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
  
      def transpose_for_scores(self, x):
          new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
          x = x.view(*new_x_shape)
  
          ## 最后xshape (batch_size, num_attention_heads, seq_len, head_size)
          return x.permute(0, 2, 1, 3)
  
      def forward(
          self,
          hidden_states,
          attention_mask,
          output_attentions=False
      ):
          mixed_query_layer = self.query(hidden_states)
          mixed_key_layer = self.key(hidden_states)
          mixed_value_layer = self.value(hidden_states)
  
          query_layer = self.transpose_for_scores(mixed_query_layer)
          key_layer = self.transpose_for_scores(mixed_key_layer)
          value_layer = self.transpose_for_scores(mixed_value_layer)
  
          # Take the dot product between "query" and "key" to get the raw attention scores.
          attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
          attention_scores = attention_scores / math.sqrt(self.attention_head_size)
          
          # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
          attention_scores = attention_scores + attention_mask
  
          # Normalize the attention scores to probabilities.
          attention_probs = nn.Softmax(dim=-1)(attention_scores)
  
          # 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)
  
          # 注意力加权
          context_layer = torch.matmul(attention_probs, value_layer)
          # 把加权后的V reshape, 得到[batch_size, length, embedding_dimension]
          context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
          new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
          context_layer = context_layer.view(*new_context_layer_shape)
  
          # 得到输出
          if output_attentions:
              return context_layer, attention_probs
          return context_layer, None
  
  
  class BertSelfOutput(nn.Module):
      def __init__(self, config):
          super().__init__()
          self.dense = nn.Linear(config.hidden_size, config.hidden_size)
          self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
          self.dropout = nn.Dropout(config.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 BertAttention(nn.Module):
      def __init__(self, config):
          super().__init__()
          self.self = BertSelfAttention(config)
          self.output = BertSelfOutput(config)
  
      def forward(
          self,
          hidden_states,
          attention_mask,
          output_attentions=False
      ):
          self_outputs, attention_metrix = self.self(hidden_states, attention_mask, output_attentions=output_attentions)
          attention_output = self.output(self_outputs, hidden_states)
          
          return attention_output, attention_metrix
  
  
  class BertIntermediate(nn.Module):
      def __init__(self, config):
          super().__init__()
          self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
          self.intermediate_act_fn = ACT2FN[config.hidden_act]
  
      def forward(self, hidden_states):
          hidden_states = self.dense(hidden_states)
          hidden_states = self.intermediate_act_fn(hidden_states)
          return hidden_states
  
  class BertOutput(nn.Module):
      def __init__(self, config):
          super().__init__()
          self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
          self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
          self.dropout = nn.Dropout(config.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 BertLayer(nn.Module):
      def __init__(self, config):
          super().__init__()
          self.attention = BertAttention(config)
          self.intermediate = BertIntermediate(config)
          self.output = BertOutput(config)
  
      def forward(
          self,
          hidden_states,
          attention_mask,
          output_attentions=False
      ):
          attention_output, attention_matrix = self.attention(hidden_states, attention_mask, output_attentions=output_attentions)
          intermediate_output = self.intermediate(attention_output)
          layer_output = self.output(intermediate_output, attention_output)
          return layer_output, attention_matrix
  
  
  class BertEncoder(nn.Module):
      def __init__(self, config):
          super().__init__()
          self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.num_hidden_layers)])
  
      def forward(
          self,
          hidden_states,
          attention_mask,
          output_all_encoded_layers=True,
          output_attentions=False
      ):
          all_encoder_layers = []
          all_attention_matrices = []
          for i, layer_module in enumerate(self.layer):
              
              layer_output, attention_matrix = layer_module(
                  hidden_states, attention_mask, output_attentions=output_attentions
              )
              hidden_states = layer_output
              if output_all_encoded_layers:
                  all_encoder_layers.append(hidden_states)
                  all_attention_matrices.append(attention_matrix)
          if not output_all_encoded_layers:
              all_encoder_layers.append(hidden_states)
              all_attention_matrices.append(attention_matrix)
              
          return all_encoder_layers, all_attention_matrices
  
  
  class BertPooler(nn.Module):
      def __init__(self, config):
          super().__init__()
          self.dense = nn.Linear(config.hidden_size, config.hidden_size)
          self.activation = nn.Tanh()
  
      def forward(self, hidden_states):
          # We "pool" the model by simply taking the hidden state corresponding
          # to the first token.
          first_token_tensor = hidden_states[:, 0]
          pooled_output = self.dense(first_token_tensor)
          pooled_output = self.activation(pooled_output)
          return pooled_output
  
  
  class BertPredictionHeadTransform(nn.Module):
      def __init__(self, config):
          super().__init__()
          self.dense = nn.Linear(config.hidden_size, config.hidden_size)
          self.transform_act_fn = ACT2FN[config.hidden_act]
          self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
  
      def forward(self, hidden_states):
          hidden_states = self.dense(hidden_states)
          hidden_states = self.transform_act_fn(hidden_states)
          hidden_states = self.LayerNorm(hidden_states)
          return hidden_states
  
  class Predictions(nn.Module):
      def __init__(self, config):
          super().__init__()
          self.transform = BertPredictionHeadTransform(config)
          self.decoder = nn.Linear(config.hidden_size, config.vocab_size)  
          self.bias = nn.Parameter(torch.zeros(config.vocab_size))
          # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
          self.decoder.bias = self.bias
  
      def forward(self, x):
          x = self.transform(x)
  
          return x, self.decoder(x)
  
  class CLS(nn.Module):
      def __init__(self, config):
          super().__init__()
          self.predictions = Predictions(config)
  
      def forward(self, x):
          return self.predictions(x)
  
  class BertLMPredictionHead(nn.Module):
      def __init__(self, config):
          super().__init__()
          self.transform = BertPredictionHeadTransform(config)
  
          # The output weights are the same as the input embeddings, but there is
          # an output-only bias for each token.
          self.decoder = nn.Linear(config.hidden_size, config.vocab_size)
  
          # self.decoder.weight = bert_model_embedding_weights
          self.bias = nn.Parameter(torch.zeros(config.vocab_size))
  
          # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
          self.decoder.bias = self.bias
  
      def forward(self, hidden_states):
          hidden_states = self.transform(hidden_states)
          hidden_states = self.decoder(hidden_states)
          return hidden_states
  
  
  class BertOnlyMLMHead(nn.Module):
      def __init__(self, config):
          super().__init__()
          self.predictions = BertLMPredictionHead(config)
  
      def forward(self, sequence_output):
          prediction_scores = self.predictions(sequence_output)
          return prediction_scores
  
  
  class BertOnlyNSPHead(nn.Module):
      def __init__(self, config):
          super().__init__()
          self.seq_relationship = nn.Linear(config.hidden_size, 2)
  
      def forward(self, pooled_output):
          seq_relationship_score = self.seq_relationship(pooled_output)
          return seq_relationship_score
  
  
  class BertPreTrainingHeads(nn.Module):
      def __init__(self, config):
          super().__init__()
          self.predictions = BertLMPredictionHead(config)
          self.seq_relationship = nn.Linear(config.hidden_size, 2)
  
      def forward(self, sequence_output, pooled_output):
          prediction_scores = self.predictions(sequence_output)
          seq_relationship_score = self.seq_relationship(pooled_output)
          return prediction_scores, seq_relationship_score
  
  
  class BertPreTrainedModel(nn.Module):
      """ An abstract class to handle weights initialization and
          a simple interface for downloading and loading pretrained models.
      """
  
      def __init__(self, config, *inputs, **kwargs):
          super(BertPreTrainedModel, self).__init__()
          if not isinstance(config, BertConfig) and not isinstance(config, RobertaLargeConfig):
              raise ValueError(
                  "Parameter config in `{}(config)` should be an instance of class `BertConfig`. "
                  "To create a model from a Google pretrained model use "
                  "`model = {}.from_pretrained(PRETRAINED_MODEL_NAME)`".format(
                      self.__class__.__name__, self.__class__.__name__
                  ))
          self.config = config
  
      def init_bert_weights(self, module):
          """ Initialize the weights.
          """
          if isinstance(module, (nn.Linear)):
              # 初始线性映射层的参数为正态分布
              module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
          elif isinstance(module, BertLayerNorm):
              # 初始化LayerNorm中的alpha为全1, beta为全0
              module.bias.data.zero_()
              module.weight.data.fill_(1.0)
          if isinstance(module, nn.Linear) and module.bias is not None:
              # 初始化偏置为0
              module.bias.data.zero_()
  
  
  class BertModel(BertPreTrainedModel):
      """
      The model can behave as an encoder (with only self-attention) as well
      as a decoder, in which case a layer of cross-attention is added between
      the self-attention layers, following the architecture described in `Attention is all you need`_ by Ashish Vaswani,
      Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
      To behave as an decoder the model needs to be initialized with the
      :obj:`is_decoder` argument of the configuration set to :obj:`True`; an
      :obj:`encoder_hidden_states` is expected as an input to the forward pass.
      .. _`Attention is all you need`:
          https://arxiv.org/abs/1706.03762
      """
  
      def __init__(self, config):
          super().__init__(config)
          self.config = config
  
          self.embeddings = BertEmbeddings(config)
          self.encoder = BertEncoder(config)
          self.pooler = BertPooler(config)
  
          self.apply(self.init_bert_weights)
  
      def forward(
          self,
          input_ids,
          attention_mask=None,
          token_type_ids=None,
          position_ids=None,
          output_all_encoded_layers=True,
          output_attentions=False
      ):  
          
          extended_attention_mask = (input_ids > 0).float()
          # print(extended_attention_mask)
          # 注意力矩阵mask: [batch_size, 1, 1, seq_length]
          extended_attention_mask = extended_attention_mask.unsqueeze(1).unsqueeze(2)
          if attention_mask is not None :
              ## 如果传进来的注意力mask不是null，那就直接用传进来的注意力mask 乘 原始mask
              # 注意 原始mask是extended_attention_mask，这个是用来把pad部分置为0，去掉pad部分影响
              extended_attention_mask = attention_mask * extended_attention_mask
  
          if token_type_ids is None:
              token_type_ids = torch.zeros_like(input_ids, dtype=torch.long, device=input_ids.device)
  
          # print(extended_attention_mask)
          # extended_attention_mask = extended_attention_mask.to(dtype=next(self.parameters()).dtype)  # fp16 compatibility
          extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
  
          embedding_output = self.embeddings(
              input_ids=input_ids, position_ids=position_ids, token_type_ids=token_type_ids
          )
          encoder_layers, all_attention_matrices = self.encoder(
              embedding_output,
              attention_mask=extended_attention_mask,
              output_all_encoded_layers=output_all_encoded_layers,
              output_attentions=output_attentions
          )
          sequence_output = encoder_layers[-1]
          pooled_output = self.pooler(sequence_output)
  
          if output_attentions:
              return all_attention_matrices
          if not output_all_encoded_layers:
              # 如果不用输出所有encoder层
              encoder_layers = encoder_layers[-1]
          return encoder_layers, pooled_output
  