






"""simple docstring"""






import logging

import torch
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.bert.modeling_bert import (
    BERT_INPUTS_DOCSTRING,
    BERT_START_DOCSTRING,
    BertEncoder,
    BertModel,
    BertPreTrainedModel,
)


UpperCAmelCase_     :				Optional[int]         =				logging.getLogger(__name__)




class 							lowerCAmelCase__       (    UpperCAmelCase__				):






      '''simple docstring'''

      def _SCREAMING_SNAKE_CASE				(		self  :						Dict		,			lowercase_  :						int		,			lowercase_  :						Union[str, Any]		,			lowercase_  :						Dict=None		,			lowercase_  :						List[str]=None):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Tuple							=	self.layer[current_layer](lowercase_		,			lowercase_		,			head_mask[current_layer])

          SCREAMING_SNAKE_CASE_   :						str							=	layer_outputs[0]

          return hidden_states




@add_start_docstrings(
    "The bare Bert Model transformer with PABEE outputting raw hidden-states without any specific head on top."       ,	UpperCAmelCase__       ,	)
class 							lowerCAmelCase__       (    UpperCAmelCase__				):






      '''simple docstring'''

      def __init__(		self  :						Any		,			lowercase_  :						Optional[Any]):




          '''simple docstring'''


          super().__init__(lowercase_)

          SCREAMING_SNAKE_CASE_   :						int							=	BertEncoderWithPabee(lowercase_)

          self.init_weights()
          SCREAMING_SNAKE_CASE_   :						Dict							=	0
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	0
          SCREAMING_SNAKE_CASE_   :						Tuple							=	0

          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	0
      def _SCREAMING_SNAKE_CASE				(		self  :						List[Any]		,			lowercase_  :						Union[str, Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	threshold
      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[Any]		,			lowercase_  :						str):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	patience
      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[int]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	0
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	0
      def _SCREAMING_SNAKE_CASE				(		self  :						List[str]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						str							=	self.inference_layers_num / self.inference_instances_num
          SCREAMING_SNAKE_CASE_   :						Any							=	(
              F'*** Patience = {self.patience} Avg. Inference Layers = {avg_inf_layers:.2f} Speed Up ='
              F' {1 - avg_inf_layers / self.config.num_hidden_layers:.2f} ***'
          )
          print(lowercase_)



      @add_start_docstrings_to_model_forward(lowercase_)
      def _SCREAMING_SNAKE_CASE				(		self  :						List[Any]		,			lowercase_  :						int=None		,			lowercase_  :						Union[str, Any]=None		,			lowercase_  :						List[Any]=None		,			lowercase_  :						List[str]=None		,			lowercase_  :						Optional[Any]=None		,			lowercase_  :						List[Any]=None		,			lowercase_  :						Dict=None		,			lowercase_  :						Any=None		,			lowercase_  :						List[str]=None		,			lowercase_  :						List[str]=None		,			lowercase_  :						Any=False		,			):




          '''simple docstring'''


          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:
              SCREAMING_SNAKE_CASE_   :						Tuple							=	input_ids.size()
          elif inputs_embeds is not None:
              SCREAMING_SNAKE_CASE_   :						Optional[int]							=	inputs_embeds.size()[:-1]
          else:
              raise ValueError('''You have to specify either input_ids or inputs_embeds''')

          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	input_ids.device if input_ids is not None else inputs_embeds.device

          if attention_mask is None:
              SCREAMING_SNAKE_CASE_   :						Any							=	torch.ones(lowercase_		,			device=lowercase_)
          if token_type_ids is None:
              SCREAMING_SNAKE_CASE_   :						int							=	torch.zeros(lowercase_		,			dtype=torch.long		,			device=lowercase_)

          # 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.
          SCREAMING_SNAKE_CASE_   :						torch.Tensor							=	self.get_extended_attention_mask(lowercase_		,			lowercase_		,			lowercase_)

          # If a 2D ou 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:
              SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	encoder_hidden_states.size()
              SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	(encoder_batch_size, encoder_sequence_length)
              if encoder_attention_mask is None:
                  SCREAMING_SNAKE_CASE_   :						List[Any]							=	torch.ones(lowercase_		,			device=lowercase_)
              SCREAMING_SNAKE_CASE_   :						str							=	self.invert_attention_mask(lowercase_)
          else:
              SCREAMING_SNAKE_CASE_   :						str							=	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]
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	self.get_head_mask(lowercase_		,			self.config.num_hidden_layers)

          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	self.embeddings(
              input_ids=lowercase_		,			position_ids=lowercase_		,			token_type_ids=lowercase_		,			inputs_embeds=lowercase_)
          SCREAMING_SNAKE_CASE_   :						int							=	embedding_output

          if self.training:
              SCREAMING_SNAKE_CASE_   :						Optional[int]							=	[]
              for i in range(self.config.num_hidden_layers):
                  SCREAMING_SNAKE_CASE_   :						Tuple							=	self.encoder.adaptive_forward(
                      lowercase_		,			current_layer=lowercase_		,			attention_mask=lowercase_		,			head_mask=lowercase_)

                  SCREAMING_SNAKE_CASE_   :						List[str]							=	self.pooler(lowercase_)
                  SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	output_layers[i](output_dropout(lowercase_))
                  res.append(lowercase_)
          elif self.patience == 0:  # Use all layers for inference
              SCREAMING_SNAKE_CASE_   :						Dict							=	self.encoder(
                  lowercase_		,			attention_mask=lowercase_		,			head_mask=lowercase_		,			encoder_hidden_states=lowercase_		,			encoder_attention_mask=lowercase_		,			)
              SCREAMING_SNAKE_CASE_   :						Tuple							=	self.pooler(encoder_outputs[0])
              SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	[output_layers[self.config.num_hidden_layers - 1](lowercase_)]
          else:
              SCREAMING_SNAKE_CASE_   :						Any							=	0
              SCREAMING_SNAKE_CASE_   :						Any							=	None
              SCREAMING_SNAKE_CASE_   :						Dict							=	0
              for i in range(self.config.num_hidden_layers):
                  calculated_layer_num += 1
                  SCREAMING_SNAKE_CASE_   :						Dict							=	self.encoder.adaptive_forward(
                      lowercase_		,			current_layer=lowercase_		,			attention_mask=lowercase_		,			head_mask=lowercase_)

                  SCREAMING_SNAKE_CASE_   :						List[str]							=	self.pooler(lowercase_)
                  SCREAMING_SNAKE_CASE_   :						Any							=	output_layers[i](lowercase_)
                  if regression:
                      SCREAMING_SNAKE_CASE_   :						Optional[int]							=	logits.detach()
                      if patient_result is not None:
                          SCREAMING_SNAKE_CASE_   :						Dict							=	patient_result.detach()
                      if (patient_result is not None) and torch.abs(patient_result - labels) < self.regression_threshold:
                          patient_counter += 1
                      else:
                          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	0
                  else:
                      SCREAMING_SNAKE_CASE_   :						Tuple							=	logits.detach().argmax(dim=1)
                      if patient_result is not None:
                          SCREAMING_SNAKE_CASE_   :						List[str]							=	patient_result.detach().argmax(dim=1)
                      if (patient_result is not None) and torch.all(labels.eq(lowercase_)):
                          patient_counter += 1
                      else:
                          SCREAMING_SNAKE_CASE_   :						Tuple							=	0

                  SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	logits
                  if patient_counter == self.patience:
                      break
              SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	[patient_result]
              self.inference_layers_num += calculated_layer_num
              self.inference_instances_num += 1

          return res




@add_start_docstrings(
    "Bert Model transformer with PABEE and a sequence classification/regression head on top (a linear layer on top of\n    the pooled output) e.g. for GLUE tasks. "       ,	UpperCAmelCase__       ,	)
class 							lowerCAmelCase__       (    UpperCAmelCase__				):






      '''simple docstring'''

      def __init__(		self  :						str		,			lowercase_  :						Any):




          '''simple docstring'''


          super().__init__(lowercase_)
          SCREAMING_SNAKE_CASE_   :						Any							=	config.num_labels

          SCREAMING_SNAKE_CASE_   :						List[str]							=	BertModelWithPabee(lowercase_)
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	nn.Dropout(config.hidden_dropout_prob)
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	nn.ModuleList(
              [nn.Linear(config.hidden_size		,			self.config.num_labels) for _ in range(config.num_hidden_layers)])

          self.init_weights()



      @add_start_docstrings_to_model_forward(lowercase_)
      def _SCREAMING_SNAKE_CASE				(		self  :						int		,			lowercase_  :						Any=None		,			lowercase_  :						int=None		,			lowercase_  :						List[str]=None		,			lowercase_  :						Any=None		,			lowercase_  :						Optional[int]=None		,			lowercase_  :						str=None		,			lowercase_  :						int=None		,			):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	self.bert(
              input_ids=lowercase_		,			attention_mask=lowercase_		,			token_type_ids=lowercase_		,			position_ids=lowercase_		,			head_mask=lowercase_		,			inputs_embeds=lowercase_		,			output_dropout=self.dropout		,			output_layers=self.classifiers		,			regression=self.num_labels == 1		,			)

          SCREAMING_SNAKE_CASE_   :						Tuple							=	(logits[-1],)

          if labels is not None:
              SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	None
              SCREAMING_SNAKE_CASE_   :						Tuple							=	0
              for ix, logits_item in enumerate(lowercase_):
                  if self.num_labels == 1:
                      #  We are doing regression
                      SCREAMING_SNAKE_CASE_   :						Optional[int]							=	MSELoss()
                      SCREAMING_SNAKE_CASE_   :						int							=	loss_fct(logits_item.view(-1)		,			labels.view(-1))
                  else:
                      SCREAMING_SNAKE_CASE_   :						int							=	CrossEntropyLoss()
                      SCREAMING_SNAKE_CASE_   :						Tuple							=	loss_fct(logits_item.view(-1		,			self.num_labels)		,			labels.view(-1))
                  if total_loss is None:
                      SCREAMING_SNAKE_CASE_   :						int							=	loss
                  else:
                      total_loss += loss * (ix + 1)
                  total_weights += ix + 1
              SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	(total_loss / total_weights,) + outputs

          return outputs










"""simple docstring"""






from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


UpperCAmelCase_     :				List[str]         =				logging.get_logger(__name__)

UpperCAmelCase_     :				Dict         =				{
    """facebook/data2vec-text-base""": """https://huggingface.co/data2vec/resolve/main/config.json""",
}




class 							lowerCAmelCase__       (    UpperCAmelCase__				):






      '''simple docstring'''

      __UpperCamelCase							=						"data2vec-text"
      def __init__(		self  :						Any		,			lowercase_  :						Any=30522		,			lowercase_  :						Any=768		,			lowercase_  :						Union[str, Any]=12		,			lowercase_  :						Dict=12		,			lowercase_  :						List[Any]=3072		,			lowercase_  :						str="gelu"		,			lowercase_  :						int=0.1		,			lowercase_  :						Dict=0.1		,			lowercase_  :						str=512		,			lowercase_  :						Optional[int]=2		,			lowercase_  :						int=0.02		,			lowercase_  :						int=1e-12		,			lowercase_  :						Any=1		,			lowercase_  :						Any=0		,			lowercase_  :						List[Any]=2		,			lowercase_  :						Tuple="absolute"		,			lowercase_  :						Optional[int]=True		,			lowercase_  :						int=None		,			**lowercase_  :						Tuple		,			):




          '''simple docstring'''


          super().__init__(pad_token_id=lowercase_		,			bos_token_id=lowercase_		,			eos_token_id=lowercase_		,			**lowercase_)

          SCREAMING_SNAKE_CASE_   :						List[Any]							=	vocab_size
          SCREAMING_SNAKE_CASE_   :						Tuple							=	hidden_size
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	num_hidden_layers
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	num_attention_heads
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	hidden_act
          SCREAMING_SNAKE_CASE_   :						Any							=	intermediate_size
          SCREAMING_SNAKE_CASE_   :						Tuple							=	hidden_dropout_prob
          SCREAMING_SNAKE_CASE_   :						Tuple							=	attention_probs_dropout_prob
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	max_position_embeddings
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	type_vocab_size
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	initializer_range
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	layer_norm_eps
          SCREAMING_SNAKE_CASE_   :						str							=	position_embedding_type
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	use_cache
          SCREAMING_SNAKE_CASE_   :						str							=	classifier_dropout




class 							lowerCAmelCase__       (    UpperCAmelCase__				):






      '''simple docstring'''

      @property
      def _SCREAMING_SNAKE_CASE				(		self  :						int):




          '''simple docstring'''


          if self.task == "multiple-choice":
              SCREAMING_SNAKE_CASE_   :						List[Any]							=	{0: '''batch''', 1: '''choice''', 2: '''sequence'''}
          else:
              SCREAMING_SNAKE_CASE_   :						Any							=	{0: '''batch''', 1: '''sequence'''}
          return OrderedDict(
              [
                  ('''input_ids''', dynamic_axis),
                  ('''attention_mask''', dynamic_axis),
              ])


