

"""simple docstring"""


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    (			__lowerCAmelCase							)				-> List[Any]:
      SCREAMING_SNAKE_CASE__      :					Optional[Any]											=		torch.exp(SCREAMING_SNAKE_CASE__							)
      SCREAMING_SNAKE_CASE__      :					Optional[Any]											=		torch.sum(SCREAMING_SNAKE_CASE__     ,    dim=1							)  # sum of exp(x_i)
      SCREAMING_SNAKE_CASE__      :					Optional[Any]											=		torch.sum(x * exp_x     ,    dim=1							)  # sum of x_i * exp(x_i)
      return torch.log(SCREAMING_SNAKE_CASE__							) - B / A

class    __a							(nn.Module):






     '''simple docstring'''

     def __init__(  self		,			_a     )     ->					Tuple:


           """simple docstring"""


           super().__init__()
           SCREAMING_SNAKE_CASE__      :					Optional[int]											=		config.output_attentions
           SCREAMING_SNAKE_CASE__      :					List[Any]											=		config.output_hidden_states
           SCREAMING_SNAKE_CASE__      :					Tuple											=		nn.ModuleList([BertLayer(lowercase_     ) for _ in range(config.num_hidden_layers     )]     )
           SCREAMING_SNAKE_CASE__      :					Union[str, Any]											=		nn.ModuleList([BertHighway(lowercase_     ) for _ in range(config.num_hidden_layers     )]     )

           SCREAMING_SNAKE_CASE__      :					str											=		[-1 for _ in range(config.num_hidden_layers     )]

     def     _a	(  self		,			_a     )     ->					str:


           """simple docstring"""


           if (type(lowercase_     ) is float) or (type(lowercase_     ) is int):
                 for i in range(len(self.early_exit_entropy     )     ):
                       SCREAMING_SNAKE_CASE__      :					Dict											=		x
           else:
                 SCREAMING_SNAKE_CASE__      :					List[str]											=		x

     def     _a	(  self		,			_a     )     ->					str:


           """simple docstring"""


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





     def     _a	(  self		,			_a		,			_a=None		,			_a=None		,			_a=None		,			_a=None		,			)     ->					Union[str, Any]:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__      :					List[Any]											=		()
           SCREAMING_SNAKE_CASE__      :					List[str]											=		()
           SCREAMING_SNAKE_CASE__      :					int											=		()
           for i, layer_module in enumerate(self.layer     ):
                 if self.output_hidden_states:
                       SCREAMING_SNAKE_CASE__      :					List[Any]											=		all_hidden_states + (hidden_states,)

                 SCREAMING_SNAKE_CASE__      :					int											=		layer_module(
                     lowercase_		,			lowercase_		,			head_mask[i]		,			lowercase_		,			lowercase_     )
                 SCREAMING_SNAKE_CASE__      :					Union[str, Any]											=		layer_outputs[0]

                 if self.output_attentions:
                       SCREAMING_SNAKE_CASE__      :					Union[str, Any]											=		all_attentions + (layer_outputs[1],)

                 SCREAMING_SNAKE_CASE__      :					int											=		(hidden_states,)
                 if self.output_hidden_states:
                       SCREAMING_SNAKE_CASE__      :					List[str]											=		current_outputs + (all_hidden_states,)
                 if self.output_attentions:
                       SCREAMING_SNAKE_CASE__      :					Optional[int]											=		current_outputs + (all_attentions,)

                 SCREAMING_SNAKE_CASE__      :					Tuple											=		self.highway[i](lowercase_     )
                 # logits, pooled_output

                 if not self.training:
                       SCREAMING_SNAKE_CASE__      :					Any											=		highway_exit[0]
                       SCREAMING_SNAKE_CASE__      :					str											=		entropy(lowercase_     )
                       SCREAMING_SNAKE_CASE__      :					Optional[Any]											=		highway_exit + (highway_entropy,)  # logits, hidden_states(?), entropy
                       SCREAMING_SNAKE_CASE__      :					Optional[int]											=		all_highway_exits + (highway_exit,)

                       if highway_entropy < self.early_exit_entropy[i]:
                             SCREAMING_SNAKE_CASE__      :					str											=		(highway_logits,) + current_outputs[1:] + (all_highway_exits,)
                             raise HighwayException(lowercase_		,			i + 1     )
                 else:
                       SCREAMING_SNAKE_CASE__      :					Any											=		all_highway_exits + (highway_exit,)

        # Add last layer
           if self.output_hidden_states:
                 SCREAMING_SNAKE_CASE__      :					Optional[int]											=		all_hidden_states + (hidden_states,)

           SCREAMING_SNAKE_CASE__      :					Dict											=		(hidden_states,)
           if self.output_hidden_states:
                 SCREAMING_SNAKE_CASE__      :					Dict											=		outputs + (all_hidden_states,)
           if self.output_attentions:
                 SCREAMING_SNAKE_CASE__      :					List[Any]											=		outputs + (all_attentions,)

           SCREAMING_SNAKE_CASE__      :					int											=		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). """      ,		UpperCamelCase_      ,		)
class    __a							(UpperCamelCase_):






     '''simple docstring'''

     def __init__(  self		,			_a     )     ->					List[Any]:


           """simple docstring"""


           super().__init__(lowercase_     )
           SCREAMING_SNAKE_CASE__      :					str											=		config

           SCREAMING_SNAKE_CASE__      :					Any											=		BertEmbeddings(lowercase_     )
           SCREAMING_SNAKE_CASE__      :					Tuple											=		DeeBertEncoder(lowercase_     )
           SCREAMING_SNAKE_CASE__      :					str											=		BertPooler(lowercase_     )

           self.init_weights()

     def     _a	(  self     )     ->					Dict:


           """simple docstring"""


           self.encoder.init_highway_pooler(self.pooler     )

     def     _a	(  self     )     ->					str:


           """simple docstring"""


           return self.embeddings.word_embeddings

     def     _a	(  self		,			_a     )     ->					Optional[Any]:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__      :					int											=		value

     def     _a	(  self		,			_a     )     ->					List[Any]:


           """simple docstring"""


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





     @add_start_docstrings_to_model_forward(lowercase_     )
     def     _a	(  self		,			_a=None		,			_a=None		,			_a=None		,			_a=None		,			_a=None		,			_a=None		,			_a=None		,			_a=None		,			)     ->					Any:


           """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__      :					List[str]											=		inputs_embeds.size()[:-1]
           else:
                 raise ValueError("""You have to specify either input_ids or inputs_embeds"""     )

           SCREAMING_SNAKE_CASE__      :					List[Any]											=		input_ids.device if input_ids is not None else inputs_embeds.device

           if attention_mask is None:
                 SCREAMING_SNAKE_CASE__      :					int											=		torch.ones(lowercase_		,			device=lowercase_     )
           if encoder_attention_mask is None:
                 SCREAMING_SNAKE_CASE__      :					Dict											=		torch.ones(lowercase_		,			device=lowercase_     )
           if token_type_ids is None:
                 SCREAMING_SNAKE_CASE__      :					Any											=		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 encoder_attention_mask.dim() == 3:
                 SCREAMING_SNAKE_CASE__      :					Dict											=		encoder_attention_mask[:, None, :, :]
           if encoder_attention_mask.dim() == 2:
                 SCREAMING_SNAKE_CASE__      :					str											=		encoder_attention_mask[:, None, None, :]

           SCREAMING_SNAKE_CASE__      :					Optional[int]											=		encoder_extended_attention_mask.to(
               dtype=next(self.parameters()     ).dtype     )  # fp16 compatibility
           SCREAMING_SNAKE_CASE__      :					int											=		(1.0 - encoder_extended_attention_mask) * -10_000.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]
           SCREAMING_SNAKE_CASE__      :					List[Any]											=		self.get_head_mask(lowercase_		,			self.config.num_hidden_layers     )

           SCREAMING_SNAKE_CASE__      :					Tuple											=		self.embeddings(
               input_ids=lowercase_		,			position_ids=lowercase_		,			token_type_ids=lowercase_		,			inputs_embeds=lowercase_     )
           SCREAMING_SNAKE_CASE__      :					List[str]											=		self.encoder(
               lowercase_		,			attention_mask=lowercase_		,			head_mask=lowercase_		,			encoder_hidden_states=lowercase_		,			encoder_attention_mask=lowercase_		,			)
           SCREAMING_SNAKE_CASE__      :					Optional[Any]											=		encoder_outputs[0]
           SCREAMING_SNAKE_CASE__      :					Optional[int]											=		self.pooler(lowercase_     )

           SCREAMING_SNAKE_CASE__      :					Tuple											=		(
               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							(UpperCamelCase_):






     '''simple docstring'''

     def __init__(  self		,			_a		,			_a     )     ->					List[str]:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__      :					Any											=		message
           SCREAMING_SNAKE_CASE__      :					Union[str, Any]											=		exit_layer  # start from 1!





class    __a							(nn.Module):






     '''simple docstring'''

     def __init__(  self		,			_a     )     ->					Optional[int]:


           """simple docstring"""


           super().__init__()
           SCREAMING_SNAKE_CASE__      :					Any											=		BertPooler(lowercase_     )
           SCREAMING_SNAKE_CASE__      :					Optional[Any]											=		nn.Dropout(config.hidden_dropout_prob     )
           SCREAMING_SNAKE_CASE__      :					List[Any]											=		nn.Linear(config.hidden_size		,			config.num_labels     )





     def     _a	(  self		,			_a     )     ->					List[str]:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__      :					Any											=		encoder_outputs[0]
           SCREAMING_SNAKE_CASE__      :					List[str]											=		self.pooler(lowercase_     )
           # "return" pooler_output

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

           # Dropout and classification
           SCREAMING_SNAKE_CASE__      :					Dict											=		bmodel_output[1]

           SCREAMING_SNAKE_CASE__      :					Tuple											=		self.dropout(lowercase_     )
           SCREAMING_SNAKE_CASE__      :					Optional[Any]											=		self.classifier(lowercase_     )

           return logits, pooled_output





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






     '''simple docstring'''

     def __init__(  self		,			_a     )     ->					Any:


           """simple docstring"""


           super().__init__(lowercase_     )
           SCREAMING_SNAKE_CASE__      :					int											=		config.num_labels
           SCREAMING_SNAKE_CASE__      :					Tuple											=		config.num_hidden_layers

           SCREAMING_SNAKE_CASE__      :					List[str]											=		DeeBertModel(lowercase_     )
           SCREAMING_SNAKE_CASE__      :					Union[str, Any]											=		nn.Dropout(config.hidden_dropout_prob     )
           SCREAMING_SNAKE_CASE__      :					List[Any]											=		nn.Linear(config.hidden_size		,			self.config.num_labels     )

           self.init_weights()





     @add_start_docstrings_to_model_forward(lowercase_     )
     def     _a	(  self		,			_a=None		,			_a=None		,			_a=None		,			_a=None		,			_a=None		,			_a=None		,			_a=None		,			_a=-1		,			_a=False		,			)     ->					Tuple:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__      :					List[Any]											=		self.num_layers
           try:
                 SCREAMING_SNAKE_CASE__      :					Union[str, Any]											=		self.bert(
                     lowercase_		,			attention_mask=lowercase_		,			token_type_ids=lowercase_		,			position_ids=lowercase_		,			head_mask=lowercase_		,			inputs_embeds=lowercase_		,			)
                 # sequence_output, pooled_output, (hidden_states), (attentions), highway exits

                 SCREAMING_SNAKE_CASE__      :					Optional[Any]											=		outputs[1]

                 SCREAMING_SNAKE_CASE__      :					int											=		self.dropout(lowercase_     )
                 SCREAMING_SNAKE_CASE__      :					int											=		self.classifier(lowercase_     )
                 SCREAMING_SNAKE_CASE__      :					Dict											=		(logits,) + outputs[2:]  # add hidden states and attention if they are here
           except HighwayException as e:
                 SCREAMING_SNAKE_CASE__      :					Optional[Any]											=		e.message
                 SCREAMING_SNAKE_CASE__      :					str											=		e.exit_layer
                 SCREAMING_SNAKE_CASE__      :					int											=		outputs[0]

           if not self.training:
                 SCREAMING_SNAKE_CASE__      :					List[str]											=		entropy(lowercase_     )
                 SCREAMING_SNAKE_CASE__      :					Union[str, Any]											=		[]
                 SCREAMING_SNAKE_CASE__      :					Any											=		[]
           if labels is not None:
                 if self.num_labels == 1:
                       #  We are doing regression
                       SCREAMING_SNAKE_CASE__      :					Optional[Any]											=		MSELoss()
                       SCREAMING_SNAKE_CASE__      :					Dict											=		loss_fct(logits.view(-1     )		,			labels.view(-1     )     )
                 else:
                       SCREAMING_SNAKE_CASE__      :					int											=		CrossEntropyLoss()
                       SCREAMING_SNAKE_CASE__      :					List[str]											=		loss_fct(logits.view(-1		,			self.num_labels     )		,			labels.view(-1     )     )

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

                 if train_highway:
                       SCREAMING_SNAKE_CASE__      :					List[str]											=		(sum(highway_losses[:-1]     ),) + outputs
                       # exclude the final highway, of course
                 else:
                       SCREAMING_SNAKE_CASE__      :					Union[str, Any]											=		(loss,) + outputs
           if not self.training:
                 SCREAMING_SNAKE_CASE__      :					str											=		outputs + ((original_entropy, highway_entropy), exit_layer)
                 if output_layer >= 0:
                       SCREAMING_SNAKE_CASE__      :					Union[str, Any]											=		(
                           (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)








import os

import pytest
from attr import dataclass


lowerCamelCase											=       'us-east-1'  # defaults region




@dataclass
class A  :
      UpperCamelCase__							:  str
      UpperCamelCase__							:  Dict      ='arn:aws:iam::558105141721:role/sagemaker_execution_role'
      UpperCamelCase__							:  int      ={
          'task_name': 'mnli',
          'per_device_train_batch_size': 16,
          'per_device_eval_batch_size': 16,
          'do_train': True,
          'do_eval': True,
          'do_predict': True,
          'output_dir': '/opt/ml/model',
          'overwrite_output_dir': True,
          'max_steps': 500,
          'save_steps': 5500,
      }
      UpperCamelCase__							:  Optional[Any]      ={**hyperparameters, 'max_steps': 1000}






      @property
      def        lowerCamelCase       (					self						:     List[Any] )					->						str:
             """simple docstring"""




             if self.framework == "pytorch":
                    return [
                        {"Name": "train_runtime", "Regex": r"train_runtime.*=\D*(.*?)$"},
                        {"Name": "eval_accuracy", "Regex": r"eval_accuracy.*=\D*(.*?)$"},
                        {"Name": "eval_loss", "Regex": r"eval_loss.*=\D*(.*?)$"},
                    ]
             else:
                    return [
                        {"Name": "train_runtime", "Regex": r"train_runtime.*=\D*(.*?)$"},
                        {"Name": "eval_accuracy", "Regex": r"loss.*=\D*(.*?)]?$"},
                        {"Name": "eval_loss", "Regex": r"sparse_categorical_accuracy.*=\D*(.*?)]?$"},
                    ]






      @property
      def        lowerCamelCase       (					self						:     List[Any] )					->						str:
             """simple docstring"""




             return F'''{self.framework}-transfromers-test'''






      @property
      def        lowerCamelCase       (					self						:     Union[str, Any] )					->						str:
             """simple docstring"""




             return F'''./tests/sagemaker/scripts/{self.framework}'''






      @property
      def        lowerCamelCase       (					self						:     Union[str, Any] )					->						str:
             """simple docstring"""




             if self.framework == "pytorch":
                    return "763104351884.dkr.ecr.us-east-1.amazonaws.com/huggingface-pytorch-training:1.7.1-transformers4.6.1-gpu-py36-cu110-ubuntu18.04"
             else:
                    return "763104351884.dkr.ecr.us-east-1.amazonaws.com/huggingface-tensorflow-training:2.4.1-transformers4.6.1-gpu-py37-cu110-ubuntu18.04"




@pytest.fixture(scope='class'     )
def 	a_      (    SCREAMING_SNAKE_CASE__					:      Tuple     ):
       '''simple docstring'''





       _lowerCamelCase       :		List[Any]		=SageMakerTestEnvironment(framework=request.cls.framework     )


