code: # 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 CUTOFFBERT model. """
  
  
  import math
  import os
  import warnings
  import numpy as np
  from dataclasses import dataclass
  from typing import Optional, Tuple
  
  import torch
  import torch.utils.checkpoint
  import torch.nn.functional as F
  from packaging import version
  from torch import nn
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss, KLDivLoss
  from torch.distributions.beta import Beta
  
  from ...activations import ACT2FN
  from ...file_utils import (
      ModelOutput,
      add_code_sample_docstrings,
      add_start_docstrings,
      add_start_docstrings_to_model_forward,
      replace_return_docstrings,
  )
  from ...modeling_outputs import (
      BaseModelOutputWithPastAndCrossAttentions,
      BaseModelOutputWithPoolingAndCrossAttentions,
      CausalLMOutputWithCrossAttentions,
      MaskedLMOutput,
      MultipleChoiceModelOutput,
      NextSentencePredictorOutput,
      QuestionAnsweringModelOutput,
      SequenceClassifierOutput,
      TokenClassifierOutput,
      DualPassageEncoderModelOutput,
  )
  from ...modeling_utils import (
      PreTrainedModel,
      apply_chunking_to_forward,
      find_pruneable_heads_and_indices,
      prune_linear_layer,
  )
  from ...utils import logging
  from .configuration_cutoffbert import CutoffBertConfig
  from ..bert.modeling_bert import BertEmbeddings as CutoffBertEmbeddings
  from ..bert.modeling_bert import BertEncoder as CutoffBertEncoder
  from ..bert.modeling_bert import BertPooler as CutoffBertPooler
  
  
  logger = logging.get_logger(__name__)
  
  _CHECKPOINT_FOR_DOC = "bert-base-uncased"
  _CONFIG_FOR_DOC = "CutoffBertConfig"
  _TOKENIZER_FOR_DOC = "CutoffBertTokenizer"
  
  CUTOFFBERT_PRETRAINED_MODEL_ARCHIVE_LIST = [
      "bert-base-uncased",
      "bert-large-uncased",
      "bert-base-cased",
      "bert-large-cased",
      "bert-base-multilingual-uncased",
      "bert-base-multilingual-cased",
      # See all BERT models at https://huggingface.co/models?filter=bert
  ]
  
  
  def load_tf_weights_in_cutoffbert(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(f"Converting TensorFlow checkpoint from {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(f"Loading TF weight {name} with shape {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(f"Skipping {'/'.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(f"Skipping {'/'.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(f"Initialize PyTorch weight {name}")
          pointer.data = torch.from_numpy(array)
      return model
  
  
  class CutoffBertPreTrainedModel(PreTrainedModel):
      """
      An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
      models.
      """
  
      config_class = CutoffBertConfig
      load_tf_weights = load_tf_weights_in_cutoffbert
      base_model_prefix = "bert"
      _keys_to_ignore_on_load_missing = [r"position_ids"]
  
      def _init_weights(self, module):
          """Initialize the weights"""
          if isinstance(module, nn.Linear):
              # Slightly different from the TF version which uses truncated_normal for initialization
              # cf https://github.com/pytorch/pytorch/pull/5617
              module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
              if module.bias is not None:
                  module.bias.data.zero_()
          elif isinstance(module, nn.Embedding):
              module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
              if module.padding_idx is not None:
                  module.weight.data[module.padding_idx].zero_()
          elif isinstance(module, nn.LayerNorm):
              module.bias.data.zero_()
              module.weight.data.fill_(1.0)
  
  
  CUTOFFBERT_START_DOCSTRING = r"""
  
      This model inherits from :class:`~transformers.PreTrainedModel`. Check the superclass documentation for the generic
      methods the library implements for all its model (such as downloading or saving, resizing the input embeddings,
      pruning heads etc.)
  
      This model is also a PyTorch `torch.nn.Module <https://pytorch.org/docs/stable/nn.html#torch.nn.Module>`__
      subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to
      general usage and behavior.
  
      Parameters:
          config (:class:`~transformers.BertConfig`): Model configuration class with all the parameters of the model.
              Initializing with a config file does not load the weights associated with the model, only the
              configuration. Check out the :meth:`~transformers.PreTrainedModel.from_pretrained` method to load the model
              weights.
  """
  
  CUTOFFBERT_INPUTS_DOCSTRING = r"""
      Args:
          input_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`):
              Indices of input sequence tokens in the vocabulary.
  
              Indices can be obtained using :class:`~transformers.BertTokenizer`. See
              :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for
              details.
  
              `What are input IDs? <../glossary.html#input-ids>`__
          attention_mask (:obj:`torch.FloatTensor` of shape :obj:`({0})`, `optional`):
              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**.
  
              `What are attention masks? <../glossary.html#attention-mask>`__
          token_type_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`, `optional`):
              Segment token indices to indicate first and second portions of the inputs. Indices are selected in ``[0,
              1]``:
  
              - 0 corresponds to a `sentence A` token,
              - 1 corresponds to a `sentence B` token.
  
              `What are token type IDs? <../glossary.html#token-type-ids>`_
          position_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`, `optional`):
              Indices of positions of each input sequence tokens in the position embeddings. Selected in the range ``[0,
              config.max_position_embeddings - 1]``.
  
              `What are position IDs? <../glossary.html#position-ids>`_
          head_mask (:obj:`torch.FloatTensor` of shape :obj:`(num_heads,)` or :obj:`(num_layers, num_heads)`, `optional`):
              Mask to nullify selected heads of the self-attention modules. Mask values selected in ``[0, 1]``:
  
              - 1 indicates the head is **not masked**,
              - 0 indicates the head is **masked**.
  
          inputs_embeds (:obj:`torch.FloatTensor` of shape :obj:`({0}, hidden_size)`, `optional`):
              Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded representation.
              This is useful if you want more control over how to convert :obj:`input_ids` indices into associated
              vectors than the model's internal embedding lookup matrix.
          output_attentions (:obj:`bool`, `optional`):
              Whether or not to return the attentions tensors of all attention layers. See ``attentions`` under returned
              tensors for more detail.
          output_hidden_states (:obj:`bool`, `optional`):
              Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for
              more detail.
          return_dict (:obj:`bool`, `optional`):
              Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple.
  """
  
  
  @add_start_docstrings(
      "The bare CutoffBert Model transformer outputting raw hidden-states without any specific head on top.",
      CUTOFFBERT_START_DOCSTRING,
  )
  class CutoffBertModel(CutoffBertPreTrainedModel):
      """
  
      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 <https://arxiv.org/abs/1706.03762>`__ 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`. To be used in a Seq2Seq model, the model needs to initialized with both :obj:`is_decoder`
      argument and :obj:`add_cross_attention` set to :obj:`True`; an :obj:`encoder_hidden_states` is then expected as an
      input to the forward pass.
      """
  
      def __init__(self, config, add_pooling_layer=True):
          super().__init__(config)
          self.config = config
  
          self.embeddings = CutoffBertEmbeddings(config)
          self.encoder = CutoffBertEncoder(config)
  
          self.pooler = CutoffBertPooler(config) if add_pooling_layer else None
  
          self.init_weights()
  
      def get_input_embeddings(self):
          return self.embeddings.word_embeddings
  
      def set_input_embeddings(self, value):
          self.embeddings.word_embeddings = value
  
      def _prune_heads(self, heads_to_prune):
          """
          Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
          class PreTrainedModel
          """
          for layer, heads in heads_to_prune.items():
              self.encoder.layer[layer].attention.prune_heads(heads)
  
      @add_start_docstrings_to_model_forward(CUTOFFBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
      @add_code_sample_docstrings(
          tokenizer_class=_TOKENIZER_FOR_DOC,
          checkpoint=_CHECKPOINT_FOR_DOC,
          output_type=BaseModelOutputWithPoolingAndCrossAttentions,
          config_class=_CONFIG_FOR_DOC,
      )
      def forward(
          self,
          input_ids=None,
          attention_mask=None,
          token_type_ids=None,
          position_ids=None,
          head_mask=None,
          inputs_embeds=None,
          encoder_hidden_states=None,
          encoder_attention_mask=None,
          past_key_values=None,
          use_cache=None,
          output_attentions=None,
          output_hidden_states=None,
          return_dict=None,
      ):
          r"""
          encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
              Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
              the model is configured as a decoder.
          encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
              Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
              the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:
  
              - 1 for tokens that are **not masked**,
              - 0 for tokens that are **masked**.
          past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
              Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
  
              If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`
              (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`
              instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.
          use_cache (:obj:`bool`, `optional`):
              If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up
              decoding (see :obj:`past_key_values`).
          """
          output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
          output_hidden_states = (
              output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
          )
          return_dict = return_dict if return_dict is not None else self.config.use_return_dict
  
          if self.config.is_decoder:
              use_cache = use_cache if use_cache is not None else self.config.use_cache
          else:
              use_cache = False
  
          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")
          elif input_ids is not None:
              input_shape = input_ids.size()
              batch_size, seq_length = input_shape
          elif inputs_embeds is not None:
              input_shape = inputs_embeds.size()[:-1]
              batch_size, seq_length = input_shape
          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
  
          # past_key_values_length
          past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
  
          if attention_mask is None:
              attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)
  
          if token_type_ids is None:
              if hasattr(self.embeddings, "token_type_ids"):
                  buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length]
                  buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length)
                  token_type_ids = buffered_token_type_ids_expanded
              else:
                  token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
  
          # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
          # ourselves in which case we just need to make it broadcastable to all heads.
          extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape, device)
  
          # If a 2D or 3D attention mask is provided for the cross-attention
          # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
          if self.config.is_decoder and encoder_hidden_states is not None:
              encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
              encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
              if encoder_attention_mask is None:
                  encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
              encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
          else:
              encoder_extended_attention_mask = None
  
          # Prepare head mask if needed
          # 1.0 in head_mask indicate we keep the head
          # attention_probs has shape bsz x n_heads x N x N
          # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
          # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
          head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
  
          embedding_output = self.embeddings(
              input_ids=input_ids,
              position_ids=position_ids,
              token_type_ids=token_type_ids,
              inputs_embeds=inputs_embeds,
              past_key_values_length=past_key_values_length,
          )
          encoder_outputs = self.encoder(
              embedding_output,
              attention_mask=extended_attention_mask,
              head_mask=head_mask,
              encoder_hidden_states=encoder_hidden_states,
              encoder_attention_mask=encoder_extended_attention_mask,
              past_key_values=past_key_values,
              use_cache=use_cache,
              output_attentions=output_attentions,
              output_hidden_states=output_hidden_states,
              return_dict=return_dict,
          )
          sequence_output = encoder_outputs[0]
          pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
  
          if not return_dict:
              return (sequence_output, pooled_output) + encoder_outputs[1:]
  
          return BaseModelOutputWithPoolingAndCrossAttentions(
              last_hidden_state=sequence_output,
              pooler_output=pooled_output,
              past_key_values=encoder_outputs.past_key_values,
              hidden_states=encoder_outputs.hidden_states,
              attentions=encoder_outputs.attentions,
              cross_attentions=encoder_outputs.cross_attentions,
          )
  
  
  @add_start_docstrings(
      """
      CutoffBert Model transformer with a sequence classification head on top (a linear layer on top of the pooled
      output) + Cut-off data augmentation support.
      """,
      CUTOFFBERT_START_DOCSTRING,
  )
  class CutoffBertForSequenceClassification(CutoffBertPreTrainedModel):
      def __init__(self, config):
          super().__init__(config)
          self.num_labels = config.num_labels
          self.cls_token_id = config.cls_token_id
          self.sep_token_id = config.sep_token_id
          self.mask_token_id = config.mask_token_id
          self.masking_prob = config.cutoff_masking_prob
          self.temperature = config.cutoff_temperature
          self.mask_loss_wgt = config.cutoff_mask_loss_wgt
          self.js_loss_wgt = config.cutoff_js_loss_wgt
          self.config = config
  
          self.bert = CutoffBertModel(config)
          classifier_dropout = (
              config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
          )
          self.dropout = nn.Dropout(classifier_dropout)
          self.classifier = nn.Linear(config.hidden_size, config.num_labels)
  
          self.init_weights()
      
      def _apply_cutoff(self, inputs):
          masked_inputs = inputs.clone()
          valid_masking_indices = (inputs != self.cls_token_id) & (inputs != self.sep_token_id)
          random_masking_indices = torch.bernoulli(torch.full(inputs.shape, self.masking_prob, device=inputs.device)).bool()
          masking_indices = random_masking_indices & valid_masking_indices
          masked_inputs[masking_indices] = self.mask_token_id
          return masked_inputs
  
      @add_start_docstrings_to_model_forward(CUTOFFBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
      @add_code_sample_docstrings(
          tokenizer_class=_TOKENIZER_FOR_DOC,
          checkpoint=_CHECKPOINT_FOR_DOC,
          output_type=SequenceClassifierOutput,
          config_class=_CONFIG_FOR_DOC,
      )
      def forward(
          self,
          input_ids=None,
          attention_mask=None,
          token_type_ids=None,
          position_ids=None,
          head_mask=None,
          inputs_embeds=None,
          labels=None,
          output_attentions=None,
          output_hidden_states=None,
          return_dict=None,
      ):
          r"""
          labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
              Labels for computing the sequence classification/regression loss. Indices should be in :obj:`[0, ...,
              config.num_labels - 1]`. If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
              If :obj:`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
          """
          return_dict = return_dict if return_dict is not None else self.config.use_return_dict
  
          if labels is None:           
              outputs = self.bert(
                  input_ids,
                  attention_mask=attention_mask,
                  token_type_ids=token_type_ids,
                  position_ids=position_ids,
                  head_mask=head_mask,
                  inputs_embeds=inputs_embeds,
                  output_attentions=output_attentions,
                  output_hidden_states=output_hidden_states,
                  return_dict=return_dict,
              )
  
              pooled_output = self.dropout(outputs[1])
              logits = self.classifier(pooled_output)
  
              if not return_dict:
                  return (logits,) + outputs[2:]
  
              return SequenceClassifierOutput(
                  logits=logits,
                  hidden_states=outputs.hidden_states,
                  attentions=outputs.attentions,
              )
  
          b, l = input_ids.size()
          masked_input_ids = self._apply_cutoff(input_ids.clone())
          flatten_input_ids = torch.stack((input_ids, masked_input_ids), dim=1).reshape(-1, l)
          flatten_attention_mask = attention_mask.unsqueeze(1).expand(-1, 2, -1).reshape(-1, l) if attention_mask is not None else None
          flatten_token_type_ids = token_type_ids.unsqueeze(1).expand(-1, 2, -1).reshape(-1, l) if token_type_ids is not None else None
          flatten_position_ids = position_ids.unsqueeze(1).expand(-1, 2, -1).reshape(-1, l) if position_ids is not None else None
          flatten_inputs_embeds = inputs_embeds.unsqueeze(1).expand(-1, 2, -1, -1).reshape(-1, l, self.config.hidden_size) if inputs_embeds is not None else None
  
          flatten_outputs = self.bert(
              flatten_input_ids,
              attention_mask=flatten_attention_mask,
              token_type_ids=flatten_token_type_ids,
              position_ids=flatten_position_ids,
              head_mask=head_mask,
              inputs_embeds=flatten_inputs_embeds,
              output_attentions=output_attentions,
              output_hidden_states=output_hidden_states,
              return_dict=return_dict,
          )
  
          flatten_pooled_output = self.dropout(flatten_outputs[1])
          flatten_logits = self.classifier(flatten_pooled_output)
  
          logits, masked_logits = flatten_logits.reshape(b, 2, self.config.num_labels).chunk(2, dim=1)
          logits, masked_logits = logits.squeeze(dim=1).contiguous(), masked_logits.squeeze(dim=1).contiguous()
  
          loss_fct = CrossEntropyLoss()
          loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
  
          if self.mask_loss_wgt is not None and self.mask_loss_wgt > 0.0:
              mask_loss = loss_fct(masked_logits.view(-1, self.num_labels), labels.view(-1))
              loss += mask_loss * self.mask_loss_wgt
  
          if self.js_loss_wgt is not None and self.js_loss_wgt > 0.0:
              kl_loss_fct = KLDivLoss(reduction="batchmean")
              src_logits, trg_logits = logits, masked_logits
              mean_logits = (src_logits + trg_logits) * 0.5
              src_loss = kl_loss_fct(
                  F.log_softmax(src_logits / self.temperature, dim=-1),
                  F.softmax(mean_logits / self.temperature, dim=-1)
              ) * (self.temperature ** 2)
              trg_loss = kl_loss_fct(
                  F.log_softmax(trg_logits / self.temperature, dim=-1),
                  F.softmax(mean_logits / self.temperature, dim=-1)
              ) * (self.temperature ** 2)
              js_loss = (src_loss + trg_loss) * 0.5
              loss += js_loss * self.js_loss_wgt
  
          if not return_dict:
              return (loss, logits)
  
          return SequenceClassifierOutput(
              loss=loss,
              logits=logits,
          )
   # 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 PROMPTBERT model. """
  
  
  import math
  import os
  import warnings
  import numpy as np
  from dataclasses import dataclass
  from typing import Optional, Tuple
  
  import torch
  import torch.utils.checkpoint
  import torch.nn.functional as F
  from packaging import version
  from torch import nn
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss, KLDivLoss
  from torch.distributions.beta import Beta
  
  from ...activations import ACT2FN
  from ...file_utils import (
      ModelOutput,
      add_code_sample_docstrings,
      add_start_docstrings,
      add_start_docstrings_to_model_forward,
      replace_return_docstrings,
  )
  from ...modeling_outputs import (
      BaseModelOutputWithPastAndCrossAttentions,
      BaseModelOutputWithPoolingAndCrossAttentions,
      CausalLMOutputWithCrossAttentions,
      MaskedLMOutput,
      MultipleChoiceModelOutput,
      NextSentencePredictorOutput,
      QuestionAnsweringModelOutput,
      SequenceClassifierOutput,
      TokenClassifierOutput,
      DualPassageEncoderModelOutput,
  )
  from ...modeling_utils import (
      PreTrainedModel,
      apply_chunking_to_forward,
      find_pruneable_heads_and_indices,
      prune_linear_layer,
  )
  from ...utils import logging
  from .configuration_promptbert import PromptBertConfig
  from ..bert.modeling_bert import BertEmbeddings as PromptBertEmbeddings
  from ..bert.modeling_bert import BertEncoder as PromptBertEncoder
  from ..bert.modeling_bert import BertPooler as PromptBertPooler
  
  
  logger = logging.get_logger(__name__)
  
  _CHECKPOINT_FOR_DOC = "bert-base-uncased"
  _CONFIG_FOR_DOC = "PromptBertConfig"
  _TOKENIZER_FOR_DOC = "PromptBertTokenizer"
  
  PROMPTBERT_PRETRAINED_MODEL_ARCHIVE_LIST = [
      "bert-base-uncased",
      "bert-large-uncased",
      "bert-base-cased",
      "bert-large-cased",
      "bert-base-multilingual-uncased",
      "bert-base-multilingual-cased",
      # See all BERT models at https://huggingface.co/models?filter=bert
  ]
  
  
  def load_tf_weights_in_promptbert(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(f"Converting TensorFlow checkpoint from {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(f"Loading TF weight {name} with shape {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(f"Skipping {'/'.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(f"Skipping {'/'.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(f"Initialize PyTorch weight {name}")
          pointer.data = torch.from_numpy(array)
      return model
  
  
  class PromptBertPreTrainedModel(PreTrainedModel):
      """
      An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
      models.
      """
  
      config_class = PromptBertConfig
      load_tf_weights = load_tf_weights_in_promptbert
      base_model_prefix = "bert"
      _keys_to_ignore_on_load_missing = [r"position_ids"]
  
      def _init_weights(self, module):
          """Initialize the weights"""
          if isinstance(module, nn.Linear):
              # Slightly different from the TF version which uses truncated_normal for initialization
              # cf https://github.com/pytorch/pytorch/pull/5617
              module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
              if module.bias is not None:
                  module.bias.data.zero_()
          elif isinstance(module, nn.Embedding):
              module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
              if module.padding_idx is not None:
                  module.weight.data[module.padding_idx].zero_()
          elif isinstance(module, nn.LayerNorm):
              module.bias.data.zero_()
              module.weight.data.fill_(1.0)
  
  
  PROMPTBERT_START_DOCSTRING = r"""
  
      This model inherits from :class:`~transformers.PreTrainedModel`. Check the superclass documentation for the generic
      methods the library implements for all its model (such as downloading or saving, resizing the input embeddings,
      pruning heads etc.)
  
      This model is also a PyTorch `torch.nn.Module <https://pytorch.org/docs/stable/nn.html#torch.nn.Module>`__
      subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to
      general usage and behavior.
  
      Parameters:
          config (:class:`~transformers.BertConfig`): Model configuration class with all the parameters of the model.
              Initializing with a config file does not load the weights associated with the model, only the
              configuration. Check out the :meth:`~transformers.PreTrainedModel.from_pretrained` method to load the model
              weights.
  """
  
  PROMPTBERT_INPUTS_DOCSTRING = r"""
      Args:
          input_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`):
              Indices of input sequence tokens in the vocabulary.
  
              Indices can be obtained using :class:`~transformers.BertTokenizer`. See
              :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for
              details.
  
              `What are input IDs? <../glossary.html#input-ids>`__
          attention_mask (:obj:`torch.FloatTensor` of shape :obj:`({0})`, `optional`):
              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**.
  
              `What are attention masks? <../glossary.html#attention-mask>`__
          token_type_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`, `optional`):
              Segment token indices to indicate first and second portions of the inputs. Indices are selected in ``[0,
              1]``:
  
              - 0 corresponds to a `sentence A` token,
              - 1 corresponds to a `sentence B` token.
  
              `What are token type IDs? <../glossary.html#token-type-ids>`_
          position_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`, `optional`):
              Indices of positions of each input sequence tokens in the position embeddings. Selected in the range ``[0,
              config.max_position_embeddings - 1]``.
  
              `What are position IDs? <../glossary.html#position-ids>`_
          head_mask (:obj:`torch.FloatTensor` of shape :obj:`(num_heads,)` or :obj:`(num_layers, num_heads)`, `optional`):
              Mask to nullify selected heads of the self-attention modules. Mask values selected in ``[0, 1]``:
  
              - 1 indicates the head is **not masked**,
              - 0 indicates the head is **masked**.
  
          inputs_embeds (:obj:`torch.FloatTensor` of shape :obj:`({0}, hidden_size)`, `optional`):
              Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded representation.
              This is useful if you want more control over how to convert :obj:`input_ids` indices into associated
              vectors than the model's internal embedding lookup matrix.
          output_attentions (:obj:`bool`, `optional`):
              Whether or not to return the attentions tensors of all attention layers. See ``attentions`` under returned
              tensors for more detail.
          output_hidden_states (:obj:`bool`, `optional`):
              Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for
              more detail.
          return_dict (:obj:`bool`, `optional`):
              Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple.
  """
  
  
  @add_start_docstrings(
      "The bare PromptBert Model transformer outputting raw hidden-states without any specific head on top.",
      PROMPTBERT_START_DOCSTRING,
  )
  class PromptBertModel(PromptBertPreTrainedModel):
      """
  
      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 <https://arxiv.org/abs/1706.03762>`__ 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`. To be used in a Seq2Seq model, the model needs to initialized with both :obj:`is_decoder`
      argument and :obj:`add_cross_attention` set to :obj:`True`; an :obj:`encoder_hidden_states` is then expected as an
      input to the forward pass.
      """
  
      def __init__(self, config, add_pooling_layer=True):
          super().__init__(config)
          self.config = config
  
          self.embeddings = PromptBertEmbeddings(config)
          self.encoder = PromptBertEncoder(config)
  
          self.pooler = PromptBertPooler(config) if add_pooling_layer else None
  
          self.init_weights()
  
      def get_input_embeddings(self):
          return self.embeddings.word_embeddings
  
      def set_input_embeddings(self, value):
          self.embeddings.word_embeddings = value
  
      def _prune_heads(self, heads_to_prune):
          """
          Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
          class PreTrainedModel
          """
          for layer, heads in heads_to_prune.items():
              self.encoder.layer[layer].attention.prune_heads(heads)
  
      @add_start_docstrings_to_model_forward(PROMPTBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
      @add_code_sample_docstrings(
          tokenizer_class=_TOKENIZER_FOR_DOC,
          checkpoint=_CHECKPOINT_FOR_DOC,
          output_type=BaseModelOutputWithPoolingAndCrossAttentions,
          config_class=_CONFIG_FOR_DOC,
      )
      def forward(
          self,
          input_ids=None,
          attention_mask=None,
          token_type_ids=None,
          position_ids=None,
          head_mask=None,
          inputs_embeds=None,
          encoder_hidden_states=None,
          encoder_attention_mask=None,
          past_key_values=None,
          use_cache=None,
          output_attentions=None,
          output_hidden_states=None,
          return_dict=None,
      ):
          r"""
          encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
              Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
              the model is configured as a decoder.
          encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
              Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
              the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:
  
              - 1 for tokens that are **not masked**,
              - 0 for tokens that are **masked**.
          past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
              Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
  
              If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`
              (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`
              instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.
          use_cache (:obj:`bool`, `optional`):
              If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up
              decoding (see :obj:`past_key_values`).
          """
          output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
          output_hidden_states = (
              output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
          )
          return_dict = return_dict if return_dict is not None else self.config.use_return_dict
  
          if self.config.is_decoder:
              use_cache = use_cache if use_cache is not None else self.config.use_cache
          else:
              use_cache = False
  
          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")
          elif input_ids is not None:
              input_shape = input_ids.size()
              batch_size, seq_length = input_shape
          elif inputs_embeds is not None:
              input_shape = inputs_embeds.size()[:-1]
              batch_size, seq_length = input_shape
          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
  
          # past_key_values_length
          past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
  
          if attention_mask is None:
              attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)
  
          if token_type_ids is None:
              if hasattr(self.embeddings, "token_type_ids"):
                  buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length]
                  buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length)
                  token_type_ids = buffered_token_type_ids_expanded
              else:
                  token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
  
          # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
          # ourselves in which case we just need to make it broadcastable to all heads.
          extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape, device)
  
          # If a 2D or 3D attention mask is provided for the cross-attention
          # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
          if self.config.is_decoder and encoder_hidden_states is not None:
              encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
              encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
              if encoder_attention_mask is None:
                  encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
              encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
          else:
              encoder_extended_attention_mask = None
  
          # Prepare head mask if needed
          # 1.0 in head_mask indicate we keep the head
          # attention_probs has shape bsz x n_heads x N x N
          # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
          # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
          head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
  
          embedding_output = self.embeddings(
              input_ids=input_ids,
              position_ids=position_ids,
              token_type_ids=token_type_ids,
              inputs_embeds=inputs_embeds,
              past_key_values_length=past_key_values_length,
          )
          encoder_outputs = self.encoder(
              embedding_output,
              attention_mask=extended_attention_mask,
              head_mask=head_mask,
              encoder_hidden_states=encoder_hidden_states,
              encoder_attention_mask=encoder_extended_attention_mask,
              past_key_values=past_key_values,
              use_cache=use_cache,
              output_attentions=output_attentions,
              output_hidden_states=output_hidden_states,
              return_dict=return_dict,
          )
          sequence_output = encoder_outputs[0]
          pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
  
          if not return_dict:
              return (sequence_output, pooled_output) + encoder_outputs[1:]
  
          return BaseModelOutputWithPoolingAndCrossAttentions(
              last_hidden_state=sequence_output,
              pooler_output=pooled_output,
              past_key_values=encoder_outputs.past_key_values,
              hidden_states=encoder_outputs.hidden_states,
              attentions=encoder_outputs.attentions,
              cross_attentions=encoder_outputs.cross_attentions,
          )
  
  
  @add_start_docstrings(
      """
      PromptBert Model transformer with a sequence classification head on top (a linear layer on top of the pooled output).
      """,
      PROMPTBERT_START_DOCSTRING,
  )
  class PromptBertForSequenceClassification(PromptBertPreTrainedModel):
      def __init__(self, config):
          super().__init__(config)
          self.num_labels = config.num_labels
          self.config = config
  
          self.bert = PromptBertModel(config)
          classifier_dropout = (
              config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
          )
          self.dropout = nn.Dropout(classifier_dropout)
          self.classifier = nn.Linear(config.hidden_size, config.num_labels)
  
          self.init_weights()
  
      @add_start_docstrings_to_model_forward(PROMPTBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
      @add_code_sample_docstrings(
          tokenizer_class=_TOKENIZER_FOR_DOC,
          checkpoint=_CHECKPOINT_FOR_DOC,
          output_type=SequenceClassifierOutput,
          config_class=_CONFIG_FOR_DOC,
      )
      def forward(
          self,
          input_ids=None,
          attention_mask=None,
          token_type_ids=None,
          position_ids=None,
          head_mask=None,
          inputs_embeds=None,
          labels=None,
          output_attentions=None,
          output_hidden_states=None,
          return_dict=None,
      ):
          r"""
          labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
              Labels for computing the sequence classification/regression loss. Indices should be in :obj:`[0, ...,
              config.num_labels - 1]`. If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
              If :obj:`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
          """
          return_dict = return_dict if return_dict is not None else self.config.use_return_dict
  
          outputs = self.bert(
              input_ids,
              attention_mask=attention_mask,
              token_type_ids=token_type_ids,
              position_ids=position_ids,
              head_mask=head_mask,
              inputs_embeds=inputs_embeds,
              output_attentions=output_attentions,
              output_hidden_states=output_hidden_states,
              return_dict=return_dict,
          )
  
          pooled_output = outputs[1]
  
          pooled_output = self.dropout(pooled_output)
          logits = self.classifier(pooled_output)
  
          loss = None
          if labels is not None:
              if self.config.problem_type is None:
                  if self.num_labels == 1:
                      self.config.problem_type = "regression"
                  elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
                      self.config.problem_type = "single_label_classification"
                  else:
                      self.config.problem_type = "multi_label_classification"
  
              if self.config.problem_type == "regression":
                  loss_fct = MSELoss()
                  if self.num_labels == 1:
                      loss = loss_fct(logits.squeeze(), labels.squeeze())
                  else:
                      loss = loss_fct(logits, labels)
              elif self.config.problem_type == "single_label_classification":
                  loss_fct = CrossEntropyLoss()
                  loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
              elif self.config.problem_type == "multi_label_classification":
                  loss_fct = BCEWithLogitsLoss()
                  loss = loss_fct(logits, labels)
          if not return_dict:
              output = (logits,) + outputs[2:]
              return ((loss,) + output) if loss is not None else output
  
          return SequenceClassifierOutput(
              loss=loss,
              logits=logits,
              hidden_states=outputs.hidden_states,
              attentions=outputs.attentions,
          )
  
  
  @add_start_docstrings(
      """
      Bert Model with a dual encoder head on top for passage retrieval tasks (a linear layer on top of the pooled output
      for computing source-target similarity).
      """,
      PROMPTBERT_START_DOCSTRING,
  )
  class PromptBertForDualPassageEncoder(PromptBertPreTrainedModel):
      def __init__(self, config, cls_loss_wgt=None):
          super().__init__(config)
          self.num_labels = config.num_labels
          self.cls_loss_wgt = cls_loss_wgt
  
          self.bert = PromptBertModel(config)
          self.pooler = PromptBertPooler(config)
          self.dropout = nn.Dropout(config.hidden_dropout_prob)
          if self.cls_loss_wgt is not None and cls_loss_wgt > 0.0:
              self.classifier = nn.Linear(config.hidden_size, config.num_labels)
  
          self.init_weights()
  
      @add_start_docstrings_to_model_forward(PROMPTBERT_INPUTS_DOCSTRING.format("batch_size, 2, sequence_length"))
      @add_code_sample_docstrings(
          tokenizer_class=_TOKENIZER_FOR_DOC,
          checkpoint=_CHECKPOINT_FOR_DOC,
          output_type=DualPassageEncoderModelOutput,
          config_class=_CONFIG_FOR_DOC,
      )
      def forward(
          self,
          input_ids=None,
          attention_mask=None,
          token_type_ids=None,
          position_ids=None,
          head_mask=None,
          inputs_embeds=None,
          labels=None,
          output_attentions=None,
          output_hidden_states=None,
          return_dict=None,
      ):
          r"""
          labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
              Labels for computing the sequence classification/regression loss. Indices should be in :obj:`[0, ...,
              config.num_labels - 1]`.
          """
          return_dict = return_dict if return_dict is not None else self.config.use_return_dict
  
          if labels is None or len(input_ids.size()) < 3:
              outputs = self.bert(
                  input_ids,
                  attention_mask=attention_mask,
                  token_type_ids=token_type_ids,
                  position_ids=position_ids,
                  head_mask=head_mask,
                  inputs_embeds=inputs_embeds,
                  output_attentions=output_attentions,
                  output_hidden_states=output_hidden_states,
                  return_dict=return_dict,
              )
  
              pooled_output = self.pooler(outputs[0])
              pooled_output = self.dropout(pooled_output)
  
              if not return_dict:
                  return (pooled_output,) + outputs[2:]
  
              return DualPassageEncoderModelOutput(
                  pooled_output=pooled_output,
                  hidden_states=outputs.hidden_states,
                  attentions=outputs.attentions,
              )
  
          b, _, l = input_ids.size()
          flatten_input_ids = input_ids.reshape(-1, l)
          flatten_attention_mask = attention_mask.reshape(-1, l) if attention_mask is not None else None
          flatten_token_type_ids = token_type_ids.reshape(-1, l) if token_type_ids is not None else None
          flatten_position_ids = position_ids.reshape(-1, l) if position_ids is not None else None
          flatten_inputs_embeds = inputs_embeds.reshape(-1, l, self.config.hidden_size) if inputs_embeds is not None else None
  
          flatten_outputs = self.bert(
              flatten_input_ids,
              attention_mask=flatten_attention_mask,
              token_type_ids=flatten_token_type_ids,
              position_ids=flatten_position_ids,
              head_mask=head_mask,
              inputs_embeds=flatten_inputs_embeds,
              output_attentions=output_attentions,
              output_hidden_states=output_hidden_states,
              return_dict=return_dict,
          )
  
          flatten_pooled_output = self.pooler(flatten_outputs[0])
          src_pooled_output, trg_pooled_output = flatten_pooled_output.reshape(b, 2, self.config.hidden_size).chunk(2, dim=1)
          src_pooled_output, trg_pooled_output = src_pooled_output.squeeze(dim=1).contiguous(), trg_pooled_output.squeeze(dim=1).contiguous()
  
          mask = (labels.unsqueeze(-1).expand(-1, b) == labels.unsqueeze(0).expand(b, -1)) & (1 - torch.eye(b)).to(labels.device).bool()
          cl_logits = torch.einsum('ik,jk->ij', src_pooled_output, trg_pooled_output).masked_fill(mask, float('-inf'))
          cl_labels = torch.arange(b).to(labels.device)
          
          loss_fct = CrossEntropyLoss()
          cl_loss = loss_fct(cl_logits.view(-1, labels.size(0)), cl_labels.view(-1))
  
          if self.cls_loss_wgt is not None and self.cls_loss_wgt > 0.0:
              flatten_logits = self.classifier(self.dropout(flatten_outputs[1]))
              src_logits, trg_logits = flatten_logits.reshape(b, 2, self.num_labels).chunk(2, dim=1)
              src_logits, trg_logits = src_logits.squeeze(dim=1).contiguous(), trg_logits.squeeze(dim=1).contiguous()
              src_loss = loss_fct(src_logits.view(-1, self.num_labels), labels.view(-1))
              trg_loss = loss_fct(trg_logits.view(-1, self.num_labels), labels.view(-1))
              cls_loss = src_loss + trg_loss
              cls_logits = src_logits + trg_logits
              loss = cl_loss + cls_loss * self.cls_loss_wgt
              logits = cls_logits
          else:
              loss = cl_loss
              logits = cl_logits
  
          if not return_dict:
              return (loss, logits,)
  
          return DualPassageEncoderModelOutput(
              loss=loss,
              logits=logits,
          )
  
apis: torch.ones torch.stack torch.nn.Linear numpy.transpose torch.nn.functional.log_softmax tensorflow.train.list_variables torch.nn.functional.softmax torch.full torch.nn.CrossEntropyLoss torch.nn.KLDivLoss torch.from_numpy torch.zeros tensorflow.train.load_variable torch.nn.Dropout
apis: torch.ones torch.nn.Linear numpy.transpose torch.nn.MSELoss torch.einsum tensorflow.train.list_variables torch.nn.CrossEntropyLoss torch.from_numpy torch.arange torch.nn.BCEWithLogitsLoss torch.zeros tensorflow.train.load_variable torch.eye torch.nn.Dropout