






"""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	     =    logging.getLogger(__name__)







class    lowerCamelCase__   (			snake_case					):






    def 							_UpperCamelCase				(					self      ,A      ,A      ,A=None      ,A=None   ):
          UpperCAmelCase							     =			self.layer[current_layer](snake_case__      ,snake_case__      ,head_mask[current_layer]   )

          UpperCAmelCase							     =			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.'''       ,	snake_case       ,	)
class    lowerCamelCase__   (			snake_case					):






    def __init__(					self      ,A   ):
          super().__init__(snake_case__   )

          UpperCAmelCase							     =			BertEncoderWithPabee(snake_case__   )

          self.init_weights()
          UpperCAmelCase							     =			0
          UpperCAmelCase							     =			0
          UpperCAmelCase							     =			0

          UpperCAmelCase							     =			0






    def 							_UpperCamelCase				(					self      ,A   ):
          UpperCAmelCase							     =			threshold






    def 							_UpperCamelCase				(					self      ,A   ):
          UpperCAmelCase							     =			patience






    def 							_UpperCamelCase				(					self   ):
          UpperCAmelCase							     =			0
          UpperCAmelCase							     =			0






    def 							_UpperCamelCase				(					self   ):
          UpperCAmelCase							     =			self.inference_layers_num / self.inference_instances_num
          UpperCAmelCase							     =			(
              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(snake_case__   )






    @add_start_docstrings_to_model_forward(snake_case__   )
    def 							_UpperCamelCase				(					self      ,A=None      ,A=None      ,A=None      ,A=None      ,A=None      ,A=None      ,A=None      ,A=None      ,A=None      ,A=None      ,A=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:
                UpperCAmelCase							     =			input_ids.size()
          elif inputs_embeds is not None:
                UpperCAmelCase							     =			inputs_embeds.size()[:-1]
          else:
                raise ValueError("""You have to specify either input_ids or inputs_embeds"""   )

          UpperCAmelCase							     =			input_ids.device if input_ids is not None else inputs_embeds.device

          if attention_mask is None:
                UpperCAmelCase							     =			torch.ones(snake_case__      ,device=snake_case__   )
          if token_type_ids is None:
                UpperCAmelCase							     =			torch.zeros(snake_case__      ,dtype=torch.long      ,device=snake_case__   )

          # 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.
          UpperCAmelCase							     =			self.get_extended_attention_mask(snake_case__      ,snake_case__      ,snake_case__   )

          # 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:
                UpperCAmelCase							     =			encoder_hidden_states.size()
                UpperCAmelCase							     =			(encoder_batch_size, encoder_sequence_length)
                if encoder_attention_mask is None:
                      UpperCAmelCase							     =			torch.ones(snake_case__      ,device=snake_case__   )
                UpperCAmelCase							     =			self.invert_attention_mask(snake_case__   )
          else:
                UpperCAmelCase							     =			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]
          UpperCAmelCase							     =			self.get_head_mask(snake_case__      ,self.config.num_hidden_layers   )

          UpperCAmelCase							     =			self.embeddings(
              input_ids=snake_case__      ,position_ids=snake_case__      ,token_type_ids=snake_case__      ,inputs_embeds=snake_case__   )
          UpperCAmelCase							     =			embedding_output

          if self.training:
                UpperCAmelCase							     =			[]
                for i in range(self.config.num_hidden_layers   ):
                      UpperCAmelCase							     =			self.encoder.adaptive_forward(
                          snake_case__      ,current_layer=snake_case__      ,attention_mask=snake_case__      ,head_mask=snake_case__   )

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

                      UpperCAmelCase							     =			self.pooler(snake_case__   )
                      UpperCAmelCase							     =			output_layers[i](snake_case__   )
                      if regression:
                            UpperCAmelCase							     =			logits.detach()
                            if patient_result is not None:
                                  UpperCAmelCase							     =			patient_result.detach()
                            if (patient_result is not None) and torch.abs(patient_result - labels   ) < self.regression_threshold:
                                  patient_counter += 1
                            else:
                                  UpperCAmelCase							     =			0
                      else:
                            UpperCAmelCase							     =			logits.detach().argmax(dim=1   )
                            if patient_result is not None:
                                  UpperCAmelCase							     =			patient_result.detach().argmax(dim=1   )
                            if (patient_result is not None) and torch.all(labels.eq(snake_case__   )   ):
                                  patient_counter += 1
                            else:
                                  UpperCAmelCase							     =			0

                      UpperCAmelCase							     =			logits
                      if patient_counter == self.patience:
                            break
                UpperCAmelCase							     =			[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. '''       ,	snake_case       ,	)
class    lowerCamelCase__   (			snake_case					):






    def __init__(					self      ,A   ):
          super().__init__(snake_case__   )
          UpperCAmelCase							     =			config.num_labels

          UpperCAmelCase							     =			BertModelWithPabee(snake_case__   )
          UpperCAmelCase							     =			nn.Dropout(config.hidden_dropout_prob   )
          UpperCAmelCase							     =			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(snake_case__   )
    def 							_UpperCamelCase				(					self      ,A=None      ,A=None      ,A=None      ,A=None      ,A=None      ,A=None      ,A=None      ,):
          UpperCAmelCase							     =			self.bert(
              input_ids=snake_case__      ,attention_mask=snake_case__      ,token_type_ids=snake_case__      ,position_ids=snake_case__      ,head_mask=snake_case__      ,inputs_embeds=snake_case__      ,output_dropout=self.dropout      ,output_layers=self.classifiers      ,regression=self.num_labels == 1      ,)

          UpperCAmelCase							     =			(logits[-1],)

          if labels is not None:
                UpperCAmelCase							     =			None
                UpperCAmelCase							     =			0
                for ix, logits_item in enumerate(snake_case__   ):
                      if self.num_labels == 1:
                            #  We are doing regression
                            UpperCAmelCase							     =			MSELoss()
                            UpperCAmelCase							     =			loss_fct(logits_item.view(-1   )      ,labels.view(-1   )   )
                      else:
                            UpperCAmelCase							     =			CrossEntropyLoss()
                            UpperCAmelCase							     =			loss_fct(logits_item.view(-1      ,self.num_labels   )      ,labels.view(-1   )   )
                      if total_loss is None:
                            UpperCAmelCase							     =			loss
                      else:
                            total_loss += loss * (ix + 1)
                      total_weights += ix + 1
                UpperCAmelCase							     =			(total_loss / total_weights,) + outputs

          return outputs





'''simple docstring'''





SCREAMING_SNAKE_CASE     =       {
    """meter""": """m""",
    """kilometer""": """km""",
    """megametre""": """Mm""",
    """gigametre""": """Gm""",
    """terametre""": """Tm""",
    """petametre""": """Pm""",
    """exametre""": """Em""",
    """zettametre""": """Zm""",
    """yottametre""": """Ym""",
}
# Exponent of the factor(meter)
SCREAMING_SNAKE_CASE     =       {
    """m""": 0,
    """km""": 3,
    """Mm""": 6,
    """Gm""": 9,
    """Tm""": 1_2,
    """Pm""": 1_5,
    """Em""": 1_8,
    """Zm""": 2_1,
    """Ym""": 2_4,
}







def snake_case_     (					lowercase__   ,       lowercase__   ,       lowercase__		):
  UpperCAmelCase__      :					str				       =							from_type.lower().strip("s"		)
  UpperCAmelCase__      :					Any				       =							to_type.lower().strip("s"		)

  UpperCAmelCase__      :					Union[str, Any]				       =							UNIT_SYMBOL.get(lowercase__   ,       lowercase__		)
  UpperCAmelCase__      :					Any				       =							UNIT_SYMBOL.get(lowercase__   ,       lowercase__		)

  if from_sanitized not in METRIC_CONVERSION:
    UpperCAmelCase__      :					str				       =							(
        F"""Invalid 'from_type' value: {from_type!r}.\n"""
        F"""Conversion abbreviations are: {", ".join(lowercase__		)}"""
    )
    raise ValueError(lowercase__		)
  if to_sanitized not in METRIC_CONVERSION:
    UpperCAmelCase__      :					List[str]				       =							(
        F"""Invalid 'to_type' value: {to_type!r}.\n"""
        F"""Conversion abbreviations are: {", ".join(lowercase__		)}"""
    )
    raise ValueError(lowercase__		)
  UpperCAmelCase__      :					Optional[int]				       =							METRIC_CONVERSION[from_sanitized]
  UpperCAmelCase__      :					str				       =							METRIC_CONVERSION[to_sanitized]
  UpperCAmelCase__      :					Tuple				       =							1

  if from_exponent > to_exponent:
    UpperCAmelCase__      :					Tuple				       =							from_exponent - to_exponent
  else:
    UpperCAmelCase__      :					Optional[int]				       =							-(to_exponent - from_exponent)

  return value * pow(1_0   ,       lowercase__		)


if __name__ == "__main__":
      from doctest import testmod

      testmod()



