content: # coding=utf-8
  # Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
  # Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  """PyTorch BERT model. """
  
  import math
  import os
  import warnings
  from dataclasses import dataclass
  from typing import Optional, Tuple
  
  import torch
  import torch.nn.functional as F
  import torch.utils.checkpoint
  import transformers
  from torch import Tensor, device, dtype, nn
  from torch.nn import CrossEntropyLoss, MSELoss
  from transformers.activations import ACT2FN
  # from transformers.models.bert.configuration_bert import BertConfig
  from transformers.configuration_utils import PretrainedConfig
  from transformers.file_utils import (ModelOutput, add_start_docstrings,
                                       add_start_docstrings_to_model_forward,
                                       replace_return_docstrings)
  from transformers.modeling_outputs import (
      BaseModelOutputWithPastAndCrossAttentions,
      BaseModelOutputWithPoolingAndCrossAttentions,
      CausalLMOutputWithCrossAttentions, MaskedLMOutput,
      MultipleChoiceModelOutput, NextSentencePredictorOutput,
      QuestionAnsweringModelOutput, SequenceClassifierOutput,
      TokenClassifierOutput)
  from transformers.modeling_utils import (PreTrainedModel,
                                           apply_chunking_to_forward,
                                           find_pruneable_heads_and_indices,
                                           prune_linear_layer)
  from transformers.utils import logging
  
  transformers.logging.set_verbosity_error()
  
  logger = logging.get_logger(__name__)
  
  _CONFIG_FOR_DOC = "BertConfig"
  _TOKENIZER_FOR_DOC = "BertTokenizer"
  
  BERT_PRETRAINED_MODEL_ARCHIVE_LIST = [
      "bert-base-uncased",
      "bert-large-uncased",
      "bert-base-cased",
      "bert-large-cased",
      "bert-base-multilingual-uncased",
      "bert-base-multilingual-cased",
      "bert-base-chinese",
      "bert-base-german-cased",
      "bert-large-uncased-whole-word-masking",
      "bert-large-cased-whole-word-masking",
      "bert-large-uncased-whole-word-masking-finetuned-squad",
      "bert-large-cased-whole-word-masking-finetuned-squad",
      "bert-base-cased-finetuned-mrpc",
      "bert-base-german-dbmdz-cased",
      "bert-base-german-dbmdz-uncased",
      "cl-tohoku/bert-base-japanese",
      "cl-tohoku/bert-base-japanese-whole-word-masking",
      "cl-tohoku/bert-base-japanese-char",
      "cl-tohoku/bert-base-japanese-char-whole-word-masking",
      "TurkuNLP/bert-base-finnish-cased-v1",
      "TurkuNLP/bert-base-finnish-uncased-v1",
      "wietsedv/bert-base-dutch-cased",
      # See all BERT models at https://huggingface.co/models?filter=bert
  ]
  
  
  class BertConfig(PretrainedConfig):
      r"""
      This is the configuration class to store the configuration of a [`BertModel`] or a [`TFBertModel`]. It is used to
      instantiate a BERT model according to the specified arguments, defining the model architecture. Instantiating a
      configuration with the defaults will yield a similar configuration to that of the BERT
      [bert-base-uncased](https://huggingface.co/bert-base-uncased) architecture.
  
      Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
      documentation from [`PretrainedConfig`] for more information.
  
  
      Args:
          vocab_size (`int`, *optional*, defaults to 30522):
              Vocabulary size of the BERT model. Defines the number of different tokens that can be represented by the
              `inputs_ids` passed when calling [`BertModel`] or [`TFBertModel`].
          hidden_size (`int`, *optional*, defaults to 768):
              Dimensionality of the encoder layers and the pooler layer.
          num_hidden_layers (`int`, *optional*, defaults to 12):
              Number of hidden layers in the Transformer encoder.
          num_attention_heads (`int`, *optional*, defaults to 12):
              Number of attention heads for each attention layer in the Transformer encoder.
          intermediate_size (`int`, *optional*, defaults to 3072):
              Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
          hidden_act (`str` or `Callable`, *optional*, defaults to `"gelu"`):
              The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
              `"relu"`, `"silu"` and `"gelu_new"` are supported.
          hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
              The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
          attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
              The dropout ratio for the attention probabilities.
          max_position_embeddings (`int`, *optional*, defaults to 512):
              The maximum sequence length that this model might ever be used with. Typically set this to something large
              just in case (e.g., 512 or 1024 or 2048).
          type_vocab_size (`int`, *optional*, defaults to 2):
              The vocabulary size of the `token_type_ids` passed when calling [`BertModel`] or [`TFBertModel`].
          initializer_range (`float`, *optional*, defaults to 0.02):
              The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
          layer_norm_eps (`float`, *optional*, defaults to 1e-12):
              The epsilon used by the layer normalization layers.
          position_embedding_type (`str`, *optional*, defaults to `"absolute"`):
              Type of position embedding. Choose one of `"absolute"`, `"relative_key"`, `"relative_key_query"`. For
              positional embeddings use `"absolute"`. For more information on `"relative_key"`, please refer to
              [Self-Attention with Relative Position Representations (Shaw et al.)](https://arxiv.org/abs/1803.02155).
              For more information on `"relative_key_query"`, please refer to *Method 4* in [Improve Transformer Models
              with Better Relative Position Embeddings (Huang et al.)](https://arxiv.org/abs/2009.13658).
          use_cache (`bool`, *optional*, defaults to `True`):
              Whether or not the model should return the last key/values attentions (not used by all models). Only
              relevant if `config.is_decoder=True`.
          classifier_dropout (`float`, *optional*):
              The dropout ratio for the classification head.
  
      Examples:
  
      ```python
      >>> from transformers import BertModel, BertConfig
  
      >>> # Initializing a BERT bert-base-uncased style configuration
      >>> configuration = BertConfig()
  
      >>> # Initializing a model from the bert-base-uncased style configuration
      >>> model = BertModel(configuration)
  
      >>> # Accessing the model configuration
      >>> configuration = model.config
      ```"""
      model_type = "bert"
  
      def __init__(
          self,
          vocab_size=30522,
          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,
          pad_token_id=0,
          position_embedding_type="absolute",
          use_cache=True,
          classifier_dropout=None,
          cross_module="ca",
          **kwargs,
      ):
          super().__init__(pad_token_id=pad_token_id, **kwargs)
  
          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
          self.position_embedding_type = position_embedding_type
          self.use_cache = use_cache
          self.classifier_dropout = classifier_dropout
          self.cross_module = cross_module
  
  
  def load_tf_weights_in_bert(model, config, tf_checkpoint_path):
      """Load tf checkpoints in a pytorch model."""
      try:
          import re
  
          import numpy as np
          import tensorflow as tf
      except ImportError:
          logger.error(
              "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "
              "https://www.tensorflow.org/install/ for installation instructions."
          )
          raise
      tf_path = os.path.abspath(tf_checkpoint_path)
      logger.info("Converting TensorFlow checkpoint from {}".format(tf_path))
      # Load weights from TF model
      init_vars = tf.train.list_variables(tf_path)
      names = []
      arrays = []
      for name, shape in init_vars:
          logger.info("Loading TF weight {} with shape {}".format(name, shape))
          array = tf.train.load_variable(tf_path, name)
          names.append(name)
          arrays.append(array)
  
      for name, array in zip(names, arrays):
          name = name.split("/")
          # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v
          # which are not required for using pretrained model
          if any(
              n
              in [
                  "adam_v",
                  "adam_m",
                  "AdamWeightDecayOptimizer",
                  "AdamWeightDecayOptimizer_1",
                  "global_step",
              ]
              for n in name
          ):
              logger.info("Skipping {}".format("/".join(name)))
              continue
          pointer = model
          for m_name in name:
              if re.fullmatch(r"[A-Za-z]+_\d+", m_name):
                  scope_names = re.split(r"_(\d+)", m_name)
              else:
                  scope_names = [m_name]
              if scope_names[0] == "kernel" or scope_names[0] == "gamma":
                  pointer = getattr(pointer, "weight")
              elif scope_names[0] == "output_bias" or scope_names[0] == "beta":
                  pointer = getattr(pointer, "bias")
              elif scope_names[0] == "output_weights":
                  pointer = getattr(pointer, "weight")
              elif scope_names[0] == "squad":
                  pointer = getattr(pointer, "classifier")
              else:
                  try:
                      pointer = getattr(pointer, scope_names[0])
                  except AttributeError:
                      logger.info("Skipping {}".format("/".join(name)))
                      continue
              if len(scope_names) >= 2:
                  num = int(scope_names[1])
                  pointer = pointer[num]
          if m_name[-11:] == "_embeddings":
              pointer = getattr(pointer, "weight")
          elif m_name == "kernel":
              array = np.transpose(array)
          try:
              assert (
                  pointer.shape == array.shape
              ), f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched"
          except AssertionError as e:
              e.args += (pointer.shape, array.shape)
              raise
          logger.info("Initialize PyTorch weight {}".format(name))
          pointer.data = torch.from_numpy(array)
      return model
  
  
  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=config.pad_token_id
          )
          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 is not snake-cased to stick with TensorFlow model variable name and be able to load
          # any TensorFlow checkpoint file
          self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
          self.dropout = nn.Dropout(config.hidden_dropout_prob)
  
          # position_ids (1, len position emb) is contiguous in memory and exported when serialized
          self.register_buffer(
              "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1))
          )
          self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
  
          self.config = config
  
      def forward(
          self,
          input_ids=None,
          token_type_ids=None,
          position_ids=None,
          inputs_embeds=None,
          past_key_values_length=0,
      ):
          if input_ids is not None:
              input_shape = input_ids.size()
          else:
              input_shape = inputs_embeds.size()[:-1]
  
          seq_length = input_shape[1]
  
          if position_ids is None:
              position_ids = self.position_ids[
                  :, past_key_values_length : seq_length + past_key_values_length
              ]
  
          if token_type_ids is None:
              token_type_ids = torch.zeros(
                  input_shape, dtype=torch.long, device=self.position_ids.device
              )
  
          if inputs_embeds is None:
              inputs_embeds = self.word_embeddings(input_ids)
  
          token_type_embeddings = self.token_type_embeddings(token_type_ids)
  
          embeddings = inputs_embeds + token_type_embeddings
          if self.position_embedding_type == "absolute":
              position_embeddings = self.position_embeddings(position_ids)
              embeddings += position_embeddings
          embeddings = self.LayerNorm(embeddings)
          embeddings = self.dropout(embeddings)
          return embeddings
  
  
  class BertSelfAttention(nn.Module):
      def __init__(self, config, is_cross_attention):
          super().__init__()
          self.config = config
          if config.hidden_size % config.num_attention_heads != 0 and not hasattr(
              config, "embedding_size"
          ):
              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)
          if is_cross_attention:
              self.key = nn.Linear(config.encoder_width, self.all_head_size)
              self.value = nn.Linear(config.encoder_width, self.all_head_size)
          else:
              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)
          self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
          if (
              self.position_embedding_type == "relative_key"
              or self.position_embedding_type == "relative_key_query"
          ):
              self.max_position_embeddings = config.max_position_embeddings
              self.distance_embedding = nn.Embedding(
                  2 * config.max_position_embeddings - 1, self.attention_head_size
              )
          self.save_attention = False
  
      def save_attn_gradients(self, attn_gradients):
          self.attn_gradients = attn_gradients
  
      def get_attn_gradients(self):
          return self.attn_gradients
  
      def save_attention_map(self, attention_map):
          self.attention_map = attention_map
  
      def get_attention_map(self):
          return self.attention_map
  
      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)
          return x.permute(0, 2, 1, 3)
  
      def forward(
          self,
          hidden_states,
          attention_mask=None,
          head_mask=None,
          encoder_hidden_states=None,
          encoder_attention_mask=None,
          past_key_value=None,
          output_attentions=False,
      ):
          mixed_query_layer = self.query(hidden_states)
  
          # If this is instantiated as a cross-attention module, the keys
          # and values come from an encoder; the attention mask needs to be
          # such that the encoder's padding tokens are not attended to.
          is_cross_attention = encoder_hidden_states is not None
  
          if is_cross_attention:
              key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
              value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
              attention_mask = encoder_attention_mask
          elif past_key_value is not None:
              key_layer = self.transpose_for_scores(self.key(hidden_states))
              value_layer = self.transpose_for_scores(self.value(hidden_states))
              key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
              value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
          else:
              key_layer = self.transpose_for_scores(self.key(hidden_states))
              value_layer = self.transpose_for_scores(self.value(hidden_states))
  
          query_layer = self.transpose_for_scores(mixed_query_layer)
  
          past_key_value = (key_layer, 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))
  
          if (
              self.position_embedding_type == "relative_key"
              or self.position_embedding_type == "relative_key_query"
          ):
              seq_length = hidden_states.size()[1]
              position_ids_l = torch.arange(
                  seq_length, dtype=torch.long, device=hidden_states.device
              ).view(-1, 1)
              position_ids_r = torch.arange(
                  seq_length, dtype=torch.long, device=hidden_states.device
              ).view(1, -1)
              distance = position_ids_l - position_ids_r
              positional_embedding = self.distance_embedding(
                  distance + self.max_position_embeddings - 1
              )
              positional_embedding = positional_embedding.to(
                  dtype=query_layer.dtype
              )  # fp16 compatibility
  
              if self.position_embedding_type == "relative_key":
                  relative_position_scores = torch.einsum(
                      "bhld,lrd->bhlr", query_layer, positional_embedding
                  )
                  attention_scores = attention_scores + relative_position_scores
              elif self.position_embedding_type == "relative_key_query":
                  relative_position_scores_query = torch.einsum(
                      "bhld,lrd->bhlr", query_layer, positional_embedding
                  )
                  relative_position_scores_key = torch.einsum(
                      "bhrd,lrd->bhlr", key_layer, positional_embedding
                  )
                  attention_scores = (
                      attention_scores
                      + relative_position_scores_query
                      + relative_position_scores_key
                  )
  
          attention_scores = attention_scores / math.sqrt(self.attention_head_size)
          if attention_mask is not None:
              # 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)
  
          if is_cross_attention and self.save_attention:
              self.save_attention_map(attention_probs)
              attention_probs.register_hook(self.save_attn_gradients)
  
          # 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_dropped = self.dropout(attention_probs)
  
          # Mask heads if we want to
          if head_mask is not None:
              attention_probs_dropped = attention_probs_dropped * head_mask
  
          context_layer = torch.matmul(attention_probs_dropped, value_layer)
  
          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)
  
          # added `attention_scores` to return tuple
          outputs = (
              (context_layer, attention_probs, attention_scores)
              if output_attentions
              else (context_layer,)
          )
  
          outputs = outputs + (past_key_value,)
          return outputs
  
  
  class BertSelfOutput(nn.Module):
      def __init__(self, config):
          super().__init__()
          self.dense = nn.Linear(config.hidden_size, config.hidden_size)
          self.LayerNorm = nn.LayerNorm(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, is_cross_attention=False):
          super().__init__()
  
          self.self = BertSelfAttention(config, is_cross_attention)
  
          self.output = BertSelfOutput(config)
          self.pruned_heads = set()
  
      def prune_heads(self, heads):
          if len(heads) == 0:
              return
          heads, index = find_pruneable_heads_and_indices(
              heads,
              self.self.num_attention_heads,
              self.self.attention_head_size,
              self.pruned_heads,
          )
  
          # Prune linear layers
          self.self.query = prune_linear_layer(self.self.query, index)
          self.self.key = prune_linear_layer(self.self.key, index)
          self.self.value = prune_linear_layer(self.self.value, index)
          self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
  
          # Update hyper params and store pruned heads
          self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
          self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
          self.pruned_heads = self.pruned_heads.union(heads)
  
      def forward(
          self,
          hidden_states,
          attention_mask=None,
          head_mask=None,
          encoder_hidden_states=None,
          encoder_attention_mask=None,
          past_key_value=None,
          output_attentions=False,
      ):
          self_outputs = self.self(
              hidden_states,
              attention_mask,
              head_mask,
              encoder_hidden_states,
              encoder_attention_mask,
              past_key_value,
              output_attentions,
          )
          attention_output = self.output(self_outputs[0], hidden_states)
          # add attentions if we output them
          outputs = (attention_output,) + self_outputs[1:]
          return outputs  # (context_layer, attention_probs, attention_scores, past_key_value,)
  
  
  class BertIntermediate(nn.Module):
      def __init__(self, config):
          super().__init__()
          self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
          if isinstance(config.hidden_act, str):
              self.intermediate_act_fn = ACT2FN[config.hidden_act]
          else:
              self.intermediate_act_fn = 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 = nn.LayerNorm(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, layer_num):
          super().__init__()
          self.config = config
          self.chunk_size_feed_forward = config.chunk_size_feed_forward
          self.seq_len_dim = 1
          self.attention = BertAttention(config)
  
          self.has_cross_attention = layer_num >= config.fusion_layer
          if self.has_cross_attention:
              self.crossattention = BertAttention(config, is_cross_attention=True)
          self.intermediate = BertIntermediate(config)
          self.output = BertOutput(config)
  
      def forward(
          self,
          hidden_states,
          attention_mask=None,
          head_mask=None,
          encoder_hidden_states=None,
          encoder_attention_mask=None,
          past_key_value=None,
          output_attentions=False,
      ):
          # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
          self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
          self_attention_outputs = self.attention(
              hidden_states,
              attention_mask,
              head_mask,
              output_attentions=output_attentions,
              past_key_value=self_attn_past_key_value,
          )  # (context_layer, attention_probs, attention_scores, past_key_value,)
          attention_output = self_attention_outputs[0]
  
          outputs = self_attention_outputs[1:-1]
          present_key_value = self_attention_outputs[-1]
  
          if self.has_cross_attention:
              assert (
                  encoder_hidden_states is not None
              ), "encoder_hidden_states must be given for cross-attention layers"
  
              if type(encoder_hidden_states) == list:
                  cross_attention_outputs = self.crossattention(
                      attention_output,
                      attention_mask,
                      head_mask,
                      encoder_hidden_states[
                          (self.layer_num - self.config.fusion_layer)
                          % len(encoder_hidden_states)
                      ],
                      encoder_attention_mask[
                          (self.layer_num - self.config.fusion_layer)
                          % len(encoder_hidden_states)
                      ],
                      output_attentions=output_attentions,
                  )
                  attention_output = cross_attention_outputs[0]
                  outputs = outputs + cross_attention_outputs[1:-1]
  
              else:
                  cross_attention_outputs = self.crossattention(
                      attention_output,
                      attention_mask,
                      head_mask,
                      encoder_hidden_states,
                      encoder_attention_mask,
                      output_attentions=output_attentions,
                  )  # (context_layer, attention_probs, attention_scores, past_key_value,)
                  attention_output = cross_attention_outputs[0]
                  # add cross attentions if we output attention weights
                  outputs = outputs + cross_attention_outputs[1:-1]
          layer_output = apply_chunking_to_forward(
              self.feed_forward_chunk,
              self.chunk_size_feed_forward,
              self.seq_len_dim,
              attention_output,
          )
          outputs = (layer_output,) + outputs
  
          outputs = outputs + (present_key_value,)
  
          return outputs
  
      def feed_forward_chunk(self, attention_output):
          intermediate_output = self.intermediate(attention_output)
          layer_output = self.output(intermediate_output, attention_output)
          return layer_output
  
  
  class BertEncoder(nn.Module):
      def __init__(self, config):
          super().__init__()
          self.config = config
          self.layer = nn.ModuleList(
              [BertLayer(config, i) for i in range(config.num_hidden_layers)]
          )
          logger.info(f"build bert with cross_module: {config.cross_module}")
  
      def forward(
          self,
          hidden_states,
          attention_mask=None,
          head_mask=None,
          encoder_hidden_states=None,
          encoder_attention_mask=None,
          past_key_values=None,
          use_cache=None,
          output_attentions=False,
          output_hidden_states=False,
          return_dict=True,
          mode="multi_modal",
          normalize_attention=True,
      ):
          all_hidden_states = () if output_hidden_states else None
          all_self_attentions = () if output_attentions else None
          # all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
          all_cross_attentions = () if output_attentions else None
  
          next_decoder_cache = () if use_cache else None
  
          if (
              mode == "text" or mode == "temporal"
          ):  # temporal is added and used for temporal att module.
              start_layer = 0
              output_layer = self.config.fusion_layer
  
          elif mode == "fusion":
              start_layer = self.config.fusion_layer
              output_layer = self.config.num_hidden_layers
  
          elif mode == "multi_modal":
              start_layer = 0
              output_layer = self.config.num_hidden_layers
  
          for i in range(start_layer, output_layer):
              layer_module = self.layer[i]
              if output_hidden_states:
                  all_hidden_states = all_hidden_states + (hidden_states,)
  
              layer_head_mask = head_mask[i] if head_mask is not None else None
              past_key_value = past_key_values[i] if past_key_values is not None else None
  
              if getattr(self.config, "gradient_checkpointing", False) and self.training:
  
                  if use_cache:
                      logger.warn(
                          "`use_cache=True` is incompatible with `config.gradient_checkpointing=True`. Setting "
                          "`use_cache=False`..."
                      )
                      use_cache = False
  
                  def create_custom_forward(module):
                      def custom_forward(*inputs):
                          return module(*inputs, past_key_value, output_attentions)
  
                      return custom_forward
  
                  layer_outputs = torch.utils.checkpoint.checkpoint(
                      create_custom_forward(layer_module),
                      hidden_states,
                      attention_mask,
                      layer_head_mask,
                      encoder_hidden_states,
                      encoder_attention_mask,
                      use_reentrant=False,
                  )
              else:
                  layer_outputs = layer_module(
                      hidden_states,
                      attention_mask,
                      layer_head_mask,
                      encoder_hidden_states,
                      encoder_attention_mask,
                      past_key_value,
                      output_attentions,
                  )  # (context_layer, attention_probs, attention_scores, past_key_value,)
              hidden_states = layer_outputs[0]
              if use_cache:
                  next_decoder_cache += (layer_outputs[-1],)
              if output_attentions:
                  # whether to output normalized attention,
                  # note for unnormalized attention, there is a mask added
                  offset = int(normalize_attention)
                  # all_self_attentions = all_self_attentions + (layer_outputs[1], )
                  all_self_attentions = all_self_attentions + (layer_outputs[2 - offset],)
                  if hasattr(layer_module, "crossattention"):
                      # all_cross_attentions = all_cross_attentions + (layer_outputs[3], )
                      all_cross_attentions = all_cross_attentions + (layer_outputs[4 - offset],)
  
          if output_hidden_states:
  