


import argparse
import os
import re


lowerCamelCase__							:						Optional[int]							     =   """src/transformers/models/auto"""


# re pattern that matches mapping introductions:
#    SUPER_MODEL_MAPPING_NAMES = OrderedDict or SUPER_MODEL_MAPPING = OrderedDict
lowerCamelCase__							:						List[Any]							     =   re.compile(r"""[A-Z_]+_MAPPING(\s+|_[A-Z_]+\s+)=\s+OrderedDict""")
# re pattern that matches identifiers in mappings
lowerCamelCase__							:						Dict							     =   re.compile(r"""\s*\(\s*\"(\S[^\"]+)\"""")







def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase							,   __lowerCAmelCase = False				)    ->				Union[str, Any]:
  with open(__lowerCAmelCase							,   '''r'''							,   encoding='''utf-8'''				) as f:
    snake_case__						= f.read()

  snake_case__						= content.split('''\n'''				)
  snake_case__						= []
  snake_case__						= 0
  while line_idx < len(__lowerCAmelCase				):
    if _re_intro_mapping.search(lines[line_idx]				) is not None:
      snake_case__						= len(re.search(r'''^(\s*)\S'''							,   lines[line_idx]				).groups()[0]				) + 8
      # Start of a new mapping!
      while not lines[line_idx].startswith(''' ''' * indent + '''('''				):
        new_lines.append(lines[line_idx]				)
        line_idx += 1

      snake_case__						= []
      while lines[line_idx].strip() != "]":
        # Blocks either fit in one line or not
        if lines[line_idx].strip() == "(":
          snake_case__						= line_idx
          while not lines[line_idx].startswith(''' ''' * indent + ''')'''				):
            line_idx += 1
          blocks.append('''\n'''.join(lines[start_idx : line_idx + 1]				)				)
        else:
          blocks.append(lines[line_idx]				)
        line_idx += 1

      # Sort blocks by their identifiers
      snake_case__						= sorted(__lowerCAmelCase							,   key=lambda __lowerCAmelCase				: _re_identifier.search(__lowerCAmelCase				).groups()[0]				)
      new_lines += blocks
    else:
      new_lines.append(lines[line_idx]				)
      line_idx += 1

  if overwrite:
    with open(__lowerCAmelCase							,   '''w'''							,   encoding='''utf-8'''				) as f:
      f.write('''\n'''.join(__lowerCAmelCase				)				)
  elif "\n".join(__lowerCAmelCase				) != content:
    return True







def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase = False				)    ->				Tuple:
  snake_case__						= [os.path.join(__lowerCAmelCase							,   __lowerCAmelCase				) for f in os.listdir(__lowerCAmelCase				) if f.endswith('''.py'''				)]
  snake_case__						= [sort_auto_mapping(__lowerCAmelCase							,   overwrite=__lowerCAmelCase				) for fname in fnames]

  if not overwrite and any(__lowerCAmelCase				):
    snake_case__						= [f for f, d in zip(__lowerCAmelCase							,   __lowerCAmelCase				) if d]
    raise ValueError(
        F"""The following files have auto mappings that need sorting: {', '.join(__lowerCAmelCase				)}. Run `make style` to fix"""
        ''' this.'''				)


if __name__ == "__main__":
   lowerCamelCase__							:						Any							     =   argparse.ArgumentParser()
   parser.add_argument("""--check_only""", action="""store_true""", help="""Whether to only check or fix style.""")
   lowerCamelCase__							:						Union[str, Any]							     =   parser.parse_args()

   sort_all_auto_mappings(not args.check_only)









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,
    BertEmbeddings,
    BertLayer,
    BertPooler,
    BertPreTrainedModel,
)
def 		lowercase__  (				_UpperCamelCase) ->			List[str]:
      """simple docstring"""

      UpperCamelCase					       =  torch.exp(_UpperCamelCase)
      UpperCamelCase					       =  torch.sum(_UpperCamelCase       ,							dim=1)  # sum of exp(x_i)
      UpperCamelCase					       =  torch.sum(x * exp_x       ,							dim=1)  # sum of x_i * exp(x_i)
      return torch.log(_UpperCamelCase) - B / A





class 					A__     (     nn.Module ):

     '''simple docstring'''






     def __init__( self      :    Union[str, Any]   ,						_SCREAMING_SNAKE_CASE      :    Any ):






           """simple docstring"""




           super().__init__()
           UpperCamelCase					       =  config.output_attentions
           UpperCamelCase					       =  config.output_hidden_states
           UpperCamelCase					       =  nn.ModuleList([BertLayer(_SCREAMING_SNAKE_CASE ) for _ in range(config.num_hidden_layers )] )
           UpperCamelCase					       =  nn.ModuleList([BertHighway(_SCREAMING_SNAKE_CASE ) for _ in range(config.num_hidden_layers )] )

           UpperCamelCase					       =  [-1 for _ in range(config.num_hidden_layers )]







     def   _SCREAMING_SNAKE_CASE     ( self      :    Optional[int]   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any] ):






           """simple docstring"""




           if (type(_SCREAMING_SNAKE_CASE ) is float) or (type(_SCREAMING_SNAKE_CASE ) is int):
                 for i in range(len(self.early_exit_entropy ) ):
                       UpperCamelCase					       =  x
           else:
                 UpperCamelCase					       =  x







     def   _SCREAMING_SNAKE_CASE     ( self      :    str   ,						_SCREAMING_SNAKE_CASE      :    List[Any] ):






           """simple docstring"""




           UpperCamelCase					       =  pooler.state_dict()
           for highway in self.highway:
                 for name, param in highway.pooler.state_dict().items():
                       param.copy_(loaded_model[name] )







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






           """simple docstring"""




           UpperCamelCase					       =  ()
           UpperCamelCase					       =  ()
           UpperCamelCase					       =  ()
           for i, layer_module in enumerate(self.layer ):
                 if self.output_hidden_states:
                       UpperCamelCase					       =  all_hidden_states + (hidden_states,)

                 UpperCamelCase					       =  layer_module(
                     _SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						head_mask[i]   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
                 UpperCamelCase					       =  layer_outputs[0]

                 if self.output_attentions:
                       UpperCamelCase					       =  all_attentions + (layer_outputs[1],)

                 UpperCamelCase					       =  (hidden_states,)
                 if self.output_hidden_states:
                       UpperCamelCase					       =  current_outputs + (all_hidden_states,)
                 if self.output_attentions:
                       UpperCamelCase					       =  current_outputs + (all_attentions,)

                 UpperCamelCase					       =  self.highway[i](_SCREAMING_SNAKE_CASE )
                 # logits, pooled_output

                 if not self.training:
                       UpperCamelCase					       =  highway_exit[0]
                       UpperCamelCase					       =  entropy(_SCREAMING_SNAKE_CASE )
                       UpperCamelCase					       =  highway_exit + (highway_entropy,)  # logits, hidden_states(?), entropy
                       UpperCamelCase					       =  all_highway_exits + (highway_exit,)

                       if highway_entropy < self.early_exit_entropy[i]:
                             UpperCamelCase					       =  (highway_logits,) + current_outputs[1:] + (all_highway_exits,)
                             raise HighwayException(_SCREAMING_SNAKE_CASE   ,						i + 1 )
                 else:
                       UpperCamelCase					       =  all_highway_exits + (highway_exit,)

        # Add last layer
           if self.output_hidden_states:
                 UpperCamelCase					       =  all_hidden_states + (hidden_states,)

           UpperCamelCase					       =  (hidden_states,)
           if self.output_hidden_states:
                 UpperCamelCase					       =  outputs + (all_hidden_states,)
           if self.output_attentions:
                 UpperCamelCase					       =  outputs + (all_attentions,)

           UpperCamelCase					       =  outputs + (all_highway_exits,)
           return outputs  # last-layer hidden state, (all hidden states), (all attentions), all highway exits





@add_start_docstrings(
    """The Bert Model transformer with early exiting (DeeBERT). """   ,    __snake_case   ,    )
class 					A__     (     __snake_case ):

     '''simple docstring'''






     def __init__( self      :    Any   ,						_SCREAMING_SNAKE_CASE      :    int ):






           """simple docstring"""




           super().__init__(_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  config

           UpperCamelCase					       =  BertEmbeddings(_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  DeeBertEncoder(_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  BertPooler(_SCREAMING_SNAKE_CASE )

           self.init_weights()







     def   _SCREAMING_SNAKE_CASE     ( self      :    int ):






           """simple docstring"""




           self.encoder.init_highway_pooler(self.pooler )







     def   _SCREAMING_SNAKE_CASE     ( self      :    str ):






           """simple docstring"""




           return self.embeddings.word_embeddings







     def   _SCREAMING_SNAKE_CASE     ( self      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    Dict ):






           """simple docstring"""




           UpperCamelCase					       =  value







     def   _SCREAMING_SNAKE_CASE     ( self      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    Tuple ):






           """simple docstring"""




           for layer, heads in heads_to_prune.items():
                 self.encoder.layer[layer].attention.prune_heads(_SCREAMING_SNAKE_CASE )







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






           """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:
                 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(_SCREAMING_SNAKE_CASE   ,						device=_SCREAMING_SNAKE_CASE )
           if encoder_attention_mask is None:
                 UpperCamelCase					       =  torch.ones(_SCREAMING_SNAKE_CASE   ,						device=_SCREAMING_SNAKE_CASE )
           if token_type_ids is None:
                 UpperCamelCase					       =  torch.zeros(_SCREAMING_SNAKE_CASE   ,						dtype=torch.long   ,						device=_SCREAMING_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(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_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 encoder_attention_mask.dim() == 3:
                 UpperCamelCase					       =  encoder_attention_mask[:, None, :, :]
           if encoder_attention_mask.dim() == 2:
                 UpperCamelCase					       =  encoder_attention_mask[:, None, None, :]

           UpperCamelCase					       =  encoder_extended_attention_mask.to(
               dtype=next(self.parameters() ).dtype )  # fp16 compatibility
           UpperCamelCase					       =  (1.0 - encoder_extended_attention_mask) * -1_0_0_0_0.0

           # 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(_SCREAMING_SNAKE_CASE   ,						self.config.num_hidden_layers )

           UpperCamelCase					       =  self.embeddings(
               input_ids=_SCREAMING_SNAKE_CASE   ,						position_ids=_SCREAMING_SNAKE_CASE   ,						token_type_ids=_SCREAMING_SNAKE_CASE   ,						inputs_embeds=_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  self.encoder(
               _SCREAMING_SNAKE_CASE   ,						attention_mask=_SCREAMING_SNAKE_CASE   ,						head_mask=_SCREAMING_SNAKE_CASE   ,						encoder_hidden_states=_SCREAMING_SNAKE_CASE   ,						encoder_attention_mask=_SCREAMING_SNAKE_CASE   ,						)
           UpperCamelCase					       =  encoder_outputs[0]
           UpperCamelCase					       =  self.pooler(_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  (
               sequence_output,
               pooled_output,
           ) + encoder_outputs[
               1:
           ]  # add hidden_states and attentions if they are here
           return outputs  # sequence_output, pooled_output, (hidden_states), (attentions), highway exits





class 					A__     (     __snake_case ):

     '''simple docstring'''






     def __init__( self      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    int ):






           """simple docstring"""




           UpperCamelCase					       =  message
           UpperCamelCase					       =  exit_layer  # start from 1!





class 					A__     (     nn.Module ):

     '''simple docstring'''






     def __init__( self      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    str ):






           """simple docstring"""




           super().__init__()
           UpperCamelCase					       =  BertPooler(_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  nn.Dropout(config.hidden_dropout_prob )
           UpperCamelCase					       =  nn.Linear(config.hidden_size   ,						config.num_labels )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any] ):






           """simple docstring"""




           UpperCamelCase					       =  encoder_outputs[0]
           UpperCamelCase					       =  self.pooler(_SCREAMING_SNAKE_CASE )
           # "return" pooler_output

           # BertModel
           UpperCamelCase					       =  (pooler_input, pooler_output) + encoder_outputs[1:]
           # "return" bmodel_output

           # Dropout and classification
           UpperCamelCase					       =  bmodel_output[1]

           UpperCamelCase					       =  self.dropout(_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  self.classifier(_SCREAMING_SNAKE_CASE )

           return logits, pooled_output





@add_start_docstrings(
    """Bert Model (with early exiting - DeeBERT) with a classifier on top,
    also takes care of multi-layer training. """   ,    __snake_case   ,    )
class 					A__     (     __snake_case ):

     '''simple docstring'''






     def __init__( self      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    Dict ):






           """simple docstring"""




           super().__init__(_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  config.num_labels
           UpperCamelCase					       =  config.num_hidden_layers

           UpperCamelCase					       =  DeeBertModel(_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  nn.Dropout(config.hidden_dropout_prob )
           UpperCamelCase					       =  nn.Linear(config.hidden_size   ,						self.config.num_labels )

           self.init_weights()







     @add_start_docstrings_to_model_forward(_SCREAMING_SNAKE_CASE )
     def   _SCREAMING_SNAKE_CASE     ( self      :    str   ,						_SCREAMING_SNAKE_CASE      :    Any=None   ,						_SCREAMING_SNAKE_CASE      :    int=None   ,						_SCREAMING_SNAKE_CASE      :    Tuple=None   ,						_SCREAMING_SNAKE_CASE      :    List[Any]=None   ,						_SCREAMING_SNAKE_CASE      :    List[str]=None   ,						_SCREAMING_SNAKE_CASE      :    Dict=None   ,						_SCREAMING_SNAKE_CASE      :    List[str]=None   ,						_SCREAMING_SNAKE_CASE      :    str=-1   ,						_SCREAMING_SNAKE_CASE      :    str=False   ,						):






           """simple docstring"""




           UpperCamelCase					       =  self.num_layers
           try:
                 UpperCamelCase					       =  self.bert(
                     _SCREAMING_SNAKE_CASE   ,						attention_mask=_SCREAMING_SNAKE_CASE   ,						token_type_ids=_SCREAMING_SNAKE_CASE   ,						position_ids=_SCREAMING_SNAKE_CASE   ,						head_mask=_SCREAMING_SNAKE_CASE   ,						inputs_embeds=_SCREAMING_SNAKE_CASE   ,						)
                 # sequence_output, pooled_output, (hidden_states), (attentions), highway exits

                 UpperCamelCase					       =  outputs[1]

                 UpperCamelCase					       =  self.dropout(_SCREAMING_SNAKE_CASE )
                 UpperCamelCase					       =  self.classifier(_SCREAMING_SNAKE_CASE )
                 UpperCamelCase					       =  (logits,) + outputs[2:]  # add hidden states and attention if they are here
           except HighwayException as e:
                 UpperCamelCase					       =  e.message
                 UpperCamelCase					       =  e.exit_layer
                 UpperCamelCase					       =  outputs[0]

           if not self.training:
                 UpperCamelCase					       =  entropy(_SCREAMING_SNAKE_CASE )
                 UpperCamelCase					       =  []
                 UpperCamelCase					       =  []
           if labels is not None:
                 if self.num_labels == 1:
                       #  We are doing regression
                       UpperCamelCase					       =  MSELoss()
                       UpperCamelCase					       =  loss_fct(logits.view(-1 )   ,						labels.view(-1 ) )
                 else:
                       UpperCamelCase					       =  CrossEntropyLoss()
                       UpperCamelCase					       =  loss_fct(logits.view(-1   ,						self.num_labels )   ,						labels.view(-1 ) )

                 # work with highway exits
                 UpperCamelCase					       =  []
                 for highway_exit in outputs[-1]:
                       UpperCamelCase					       =  highway_exit[0]
                       if not self.training:
                             highway_logits_all.append(_SCREAMING_SNAKE_CASE )
                             highway_entropy.append(highway_exit[2] )
                       if self.num_labels == 1:
                             #  We are doing regression
                             UpperCamelCase					       =  MSELoss()
                             UpperCamelCase					       =  loss_fct(highway_logits.view(-1 )   ,						labels.view(-1 ) )
                       else:
                             UpperCamelCase					       =  CrossEntropyLoss()
                             UpperCamelCase					       =  loss_fct(highway_logits.view(-1   ,						self.num_labels )   ,						labels.view(-1 ) )
                       highway_losses.append(_SCREAMING_SNAKE_CASE )

                 if train_highway:
                       UpperCamelCase					       =  (sum(highway_losses[:-1] ),) + outputs
                       # exclude the final highway, of course
                 else:
                       UpperCamelCase					       =  (loss,) + outputs
           if not self.training:
                 UpperCamelCase					       =  outputs + ((original_entropy, highway_entropy), exit_layer)
                 if output_layer >= 0:
                       UpperCamelCase					       =  (
                           (outputs[0],) + (highway_logits_all[output_layer],) + outputs[2:]
                       )  # use the highway of the last layer

           return outputs  # (loss), logits, (hidden_states), (attentions), (highway_exits)



