
"""simple docstring"""

import unittest

from parameterized import parameterized

from transformers import AutoTokenizer, GPTNeoXConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch

      from transformers import (
          GPTNeoXForCausalLM,
          GPTNeoXForQuestionAnswering,
          GPTNeoXForSequenceClassification,
          GPTNeoXForTokenClassification,
          GPTNeoXModel,
      )

class __lowerCAmelCase	:





     def __init__(			self			,  __UpperCAmelCase			,  __UpperCAmelCase=13			,  __UpperCAmelCase=7			,  __UpperCAmelCase=True			,  __UpperCAmelCase=True			,  __UpperCAmelCase=True			,  __UpperCAmelCase=True			,  __UpperCAmelCase=99			,  __UpperCAmelCase=64			,  __UpperCAmelCase=5			,  __UpperCAmelCase=4			,  __UpperCAmelCase=37			,  __UpperCAmelCase="gelu"			,  __UpperCAmelCase=0.1			,  __UpperCAmelCase=0.1			,  __UpperCAmelCase=512			,  __UpperCAmelCase=16			,  __UpperCAmelCase=2			,  __UpperCAmelCase=0.0_2			,  __UpperCAmelCase=3			,  __UpperCAmelCase=4			,  __UpperCAmelCase=None			,  ):

      '''simple docstring'''




      __UpperCamelCase       						= parent
      __UpperCamelCase       						= batch_size
      __UpperCamelCase       						= seq_length
      __UpperCamelCase       						= is_training
      __UpperCamelCase       						= use_input_mask
      __UpperCamelCase       						= use_token_type_ids
      __UpperCamelCase       						= use_labels
      __UpperCamelCase       						= vocab_size
      __UpperCamelCase       						= hidden_size
      __UpperCamelCase       						= num_hidden_layers
      __UpperCamelCase       						= num_attention_heads
      __UpperCamelCase       						= intermediate_size
      __UpperCamelCase       						= hidden_act
      __UpperCamelCase       						= hidden_dropout_prob
      __UpperCamelCase       						= attention_probs_dropout_prob
      __UpperCamelCase       						= max_position_embeddings
      __UpperCamelCase       						= type_vocab_size
      __UpperCamelCase       						= type_sequence_label_size
      __UpperCamelCase       						= initializer_range
      __UpperCamelCase       						= num_labels
      __UpperCamelCase       						= num_choices
      __UpperCamelCase       						= scope
      __UpperCamelCase       						= vocab_size - 1





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      __UpperCamelCase       						= ids_tensor([self.batch_size, self.seq_length]			,  self.vocab_size   )

      __UpperCamelCase       						= None
      if self.use_input_mask:
       __UpperCamelCase       						= random_attention_mask([self.batch_size, self.seq_length]   )

      __UpperCamelCase       						= None
      if self.use_labels:
       __UpperCamelCase       						= ids_tensor([self.batch_size, self.seq_length]			,  self.num_labels   )

      __UpperCamelCase       						= self.get_config()

      return config, input_ids, input_mask, token_labels





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      return GPTNeoXConfig(
          vocab_size=self.vocab_size			,  hidden_size=self.hidden_size			,  num_hidden_layers=self.num_hidden_layers			,  num_attention_heads=self.num_attention_heads			,  intermediate_size=self.intermediate_size			,  hidden_act=self.hidden_act			,  hidden_dropout_prob=self.hidden_dropout_prob			,  attention_probs_dropout_prob=self.attention_probs_dropout_prob			,  max_position_embeddings=self.max_position_embeddings			,  type_vocab_size=self.type_vocab_size			,  is_decoder=__UpperCAmelCase			,  initializer_range=self.initializer_range			,  pad_token_id=self.pad_token_id			,  )





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      __UpperCamelCase						,			__UpperCamelCase						,			__UpperCamelCase						,			__UpperCamelCase       						= self.prepare_config_and_inputs()

      __UpperCamelCase       						= True

      return config, input_ids, input_mask, token_labels





     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

      '''simple docstring'''




      __UpperCamelCase       						= GPTNeoXModel(config=__UpperCAmelCase   )
      model.to(__UpperCAmelCase   )
      model.eval()
      __UpperCamelCase       						= model(__UpperCAmelCase			,  attention_mask=__UpperCAmelCase   )
      __UpperCamelCase       						= model(__UpperCAmelCase   )
      self.parent.assertEqual(result.last_hidden_state.shape			,  (self.batch_size, self.seq_length, self.hidden_size)   )





     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

      '''simple docstring'''




      __UpperCamelCase       						= True
      __UpperCamelCase       						= GPTNeoXModel(__UpperCAmelCase   )
      model.to(__UpperCAmelCase   )
      model.eval()
      __UpperCamelCase       						= model(__UpperCAmelCase			,  attention_mask=__UpperCAmelCase   )
      self.parent.assertEqual(result.last_hidden_state.shape			,  (self.batch_size, self.seq_length, self.hidden_size)   )





     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

      '''simple docstring'''




      __UpperCamelCase       						= GPTNeoXForCausalLM(config=__UpperCAmelCase   )
      model.to(__UpperCAmelCase   )
      model.eval()
      __UpperCamelCase       						= model(__UpperCAmelCase			,  attention_mask=__UpperCAmelCase			,  labels=__UpperCAmelCase   )
      self.parent.assertEqual(result.logits.shape			,  (self.batch_size, self.seq_length, self.vocab_size)   )





     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

      '''simple docstring'''




      __UpperCamelCase       						= self.num_labels
      __UpperCamelCase       						= GPTNeoXForQuestionAnswering(__UpperCAmelCase   )
      model.to(__UpperCAmelCase   )
      model.eval()
      __UpperCamelCase       						= model(__UpperCAmelCase			,  attention_mask=__UpperCAmelCase   )
      self.parent.assertEqual(result.start_logits.shape			,  (self.batch_size, self.seq_length)   )
      self.parent.assertEqual(result.end_logits.shape			,  (self.batch_size, self.seq_length)   )





     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

      '''simple docstring'''




      __UpperCamelCase       						= self.num_labels
      __UpperCamelCase       						= GPTNeoXForSequenceClassification(__UpperCAmelCase   )
      model.to(__UpperCAmelCase   )
      model.eval()
      __UpperCamelCase       						= ids_tensor([self.batch_size]			,  self.type_sequence_label_size   )
      __UpperCamelCase       						= model(__UpperCAmelCase			,  attention_mask=__UpperCAmelCase			,  labels=__UpperCAmelCase   )
      self.parent.assertEqual(result.logits.shape			,  (self.batch_size, self.num_labels)   )





     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

      '''simple docstring'''




      __UpperCamelCase       						= self.num_labels
      __UpperCamelCase       						= GPTNeoXForTokenClassification(__UpperCAmelCase   )
      model.to(__UpperCAmelCase   )
      model.eval()
      __UpperCamelCase       						= model(__UpperCAmelCase			,  attention_mask=__UpperCAmelCase			,  labels=__UpperCAmelCase   )
      self.parent.assertEqual(result.logits.shape			,  (self.batch_size, self.seq_length, self.num_labels)   )





     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

      '''simple docstring'''




      __UpperCamelCase       						= True
      __UpperCamelCase       						= GPTNeoXForCausalLM(config=__UpperCAmelCase   )
      model.to(__UpperCAmelCase   )
      model.eval()

      # first forward pass
      __UpperCamelCase       						= model(__UpperCAmelCase			,  attention_mask=__UpperCAmelCase			,  use_cache=__UpperCAmelCase   )
      __UpperCamelCase       						= outputs.past_key_values

      # create hypothetical multiple next token and extent to next_input_ids
      __UpperCamelCase       						= ids_tensor((self.batch_size, 3)			,  config.vocab_size   )
      __UpperCamelCase       						= ids_tensor((self.batch_size, 3)			,  vocab_size=2   )

      # append to next input_ids and
      __UpperCamelCase       						= torch.cat([input_ids, next_tokens]			,  dim=-1   )
      __UpperCamelCase       						= torch.cat([input_mask, next_mask]			,  dim=-1   )

      __UpperCamelCase       						= model(__UpperCAmelCase			,  attention_mask=__UpperCAmelCase			,  output_hidden_states=__UpperCAmelCase   )
      __UpperCamelCase       						= output_from_no_past['hidden_states'][0]
      __UpperCamelCase       						= model(
          __UpperCAmelCase			,  attention_mask=__UpperCAmelCase			,  past_key_values=__UpperCAmelCase			,  output_hidden_states=__UpperCAmelCase			,  )['hidden_states'][0]

      # select random slice
      __UpperCamelCase       						= ids_tensor((1,)			,  output_from_past.shape[-1]   ).item()
      __UpperCamelCase       						= output_from_no_past[:, -3:, random_slice_idx].detach()
      __UpperCamelCase       						= output_from_past[:, :, random_slice_idx].detach()

      self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]   )

      # test that outputs are equal for slice
      self.parent.assertTrue(torch.allclose(__UpperCAmelCase			,  __UpperCAmelCase			,  atol=1E-3   )   )







     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      __UpperCamelCase       						= self.prepare_config_and_inputs()
      __UpperCamelCase						,			__UpperCamelCase						,			__UpperCamelCase						,			__UpperCamelCase       						= config_and_inputs
      __UpperCamelCase       						= {'input_ids': input_ids, 'attention_mask': input_mask}
      return config, inputs_dict


@require_torch
class __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ,						__SCREAMING_SNAKE_CASE    ,						__SCREAMING_SNAKE_CASE    ,						unittest.TestCase    ):
     lowercase		     =		(
         (
             GPTNeoXModel,
             GPTNeoXForCausalLM,
             GPTNeoXForQuestionAnswering,
             GPTNeoXForSequenceClassification,
             GPTNeoXForTokenClassification,
         )
         if is_torch_available()
         else ()
     )
     lowercase		     =		(GPTNeoXForCausalLM,) if is_torch_available() else ()
     lowercase		     =		(
         {
             "feature-extraction": GPTNeoXModel,
             "question-answering": GPTNeoXForQuestionAnswering,
             "text-classification": GPTNeoXForSequenceClassification,
             "text-generation": GPTNeoXForCausalLM,
             "token-classification": GPTNeoXForTokenClassification,
             "zero-shot": GPTNeoXForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     lowercase		     =		False
     lowercase		     =		False
     lowercase		     =		False
     lowercase		     =		False





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      __UpperCamelCase       						= GPTNeoXModelTester(self   )
      __UpperCamelCase       						= ConfigTester(self			,  config_class=__UpperCAmelCase			,  hidden_size=64			,  num_attention_heads=8   )





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      self.config_tester.run_common_tests()





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      __UpperCamelCase						,			__UpperCamelCase						,			__UpperCamelCase						,			__UpperCamelCase       						= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(__UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase   )





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      __UpperCamelCase						,			__UpperCamelCase						,			__UpperCamelCase						,			__UpperCamelCase       						= self.model_tester.prepare_config_and_inputs_for_decoder()
      self.model_tester.create_and_check_model_as_decoder(__UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase   )





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      __UpperCamelCase						,			__UpperCamelCase						,			__UpperCamelCase						,			__UpperCamelCase       						= self.model_tester.prepare_config_and_inputs_for_decoder()

      __UpperCamelCase       						= None

      self.model_tester.create_and_check_model_as_decoder(__UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase   )





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      __UpperCamelCase						,			__UpperCamelCase						,			__UpperCamelCase						,			__UpperCamelCase       						= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_decoder_model_past_large_inputs(__UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase   )





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      __UpperCamelCase       						= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_causal_lm(*__UpperCAmelCase   )





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      __UpperCamelCase       						= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_question_answering(*__UpperCAmelCase   )





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      __UpperCamelCase       						= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_sequence_classification(*__UpperCAmelCase   )





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      __UpperCamelCase       						= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_token_classification(*__UpperCAmelCase   )





     @unittest.skip(reason='Feed forward chunking is not implemented'   )
     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      pass







     @parameterized.expand([('linear',), ('dynamic',)]   )
     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      __UpperCamelCase						,			__UpperCamelCase       						= self.model_tester.prepare_config_and_inputs_for_common()
      __UpperCamelCase       						= ids_tensor([1, 10]			,  config.vocab_size   )
      __UpperCamelCase       						= ids_tensor([1, int(config.max_position_embeddings * 1.5   )]			,  config.vocab_size   )

      set_seed(42   )  # Fixed seed at init time so the two models get the same random weights
      __UpperCamelCase       						= GPTNeoXModel(__UpperCAmelCase   )
      original_model.to(__UpperCAmelCase   )
      original_model.eval()
      __UpperCamelCase       						= original_model(__UpperCAmelCase   ).last_hidden_state
      __UpperCamelCase       						= original_model(__UpperCAmelCase   ).last_hidden_state

      set_seed(42   )  # Fixed seed at init time so the two models get the same random weights
      __UpperCamelCase       						= {'type': scaling_type, 'factor': 10.0}
      __UpperCamelCase       						= GPTNeoXModel(__UpperCAmelCase   )
      scaled_model.to(__UpperCAmelCase   )
      scaled_model.eval()
      __UpperCamelCase       						= scaled_model(__UpperCAmelCase   ).last_hidden_state
      __UpperCamelCase       						= scaled_model(__UpperCAmelCase   ).last_hidden_state

      # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
      # maximum sequence length, so the outputs for the short input should match.
      if scaling_type == "dynamic":
       self.assertTrue(torch.allclose(__UpperCAmelCase			,  __UpperCAmelCase			,  atol=1E-5   )   )
      else:
       self.assertFalse(torch.allclose(__UpperCAmelCase			,  __UpperCAmelCase			,  atol=1E-5   )   )

      # The output should be different for long inputs
      self.assertFalse(torch.allclose(__UpperCAmelCase			,  __UpperCAmelCase			,  atol=1E-5   )   )


@require_torch
class __lowerCAmelCase	(   unittest.TestCase    ):





     @slow
     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      __UpperCamelCase       						= AutoTokenizer.from_pretrained('EleutherAI/pythia-410m-deduped'   )
      for checkpointing in [True, False]:
       __UpperCamelCase       						= GPTNeoXForCausalLM.from_pretrained('EleutherAI/pythia-410m-deduped'   )

       if checkpointing:
        model.gradient_checkpointing_enable()
       else:
        model.gradient_checkpointing_disable()
       model.to(__UpperCAmelCase   )

       __UpperCamelCase       						= tokenizer('My favorite food is'			,  return_tensors='pt'   ).to(__UpperCAmelCase   )
       # The hub repo. is updated on 2023-04-04, resulting in poor outputs.
       # See: https://github.com/huggingface/transformers/pull/24193
       __UpperCamelCase       						= 'My favorite food is a good old-fashioned, old-fashioned, old-fashioned.\n\nI\'m not sure'

       __UpperCamelCase       						= model.generate(**__UpperCAmelCase			,  do_sample=__UpperCAmelCase			,  max_new_tokens=20   )
       __UpperCamelCase       						= tokenizer.batch_decode(__UpperCAmelCase   )[0]

       self.assertEqual(__UpperCAmelCase			,  __UpperCAmelCase   )





"""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 							A		(       snake_case   :Optional[int]	)   ->			str:
 __UpperCamelCase       						= torch.exp(snake_case	)
 __UpperCamelCase       						= torch.sum(snake_case       ,      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(snake_case	) - B / A

class __lowerCAmelCase	(   nn.Module    ):





     def __init__(			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




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

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





     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




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





     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase   ):

      '''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 		UpperCAmelCase      (			self			,  __UpperCAmelCase			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=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(
           __UpperCAmelCase			,  __UpperCAmelCase			,  head_mask[i]			,  __UpperCAmelCase			,  __UpperCAmelCase   )
       __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](__UpperCAmelCase   )
       # logits, pooled_output

       if not self.training:
        __UpperCamelCase       						= highway_exit[0]
        __UpperCamelCase       						= entropy(__UpperCAmelCase   )
        __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(__UpperCAmelCase			,  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). "    ,						__SCREAMING_SNAKE_CASE    ,						)
class __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ):





     def __init__(			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      super().__init__(__UpperCAmelCase   )
      __UpperCamelCase       						= config

      __UpperCamelCase       						= BertEmbeddings(__UpperCAmelCase   )
      __UpperCamelCase       						= DeeBertEncoder(__UpperCAmelCase   )
      __UpperCamelCase       						= BertPooler(__UpperCAmelCase   )

      self.init_weights()





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      self.encoder.init_highway_pooler(self.pooler   )





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      return self.embeddings.word_embeddings





     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      __UpperCamelCase       						= value





     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




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







     @add_start_docstrings_to_model_forward(__UpperCAmelCase   )
     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=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(__UpperCAmelCase			,  device=__UpperCAmelCase   )
      if encoder_attention_mask is None:
       __UpperCamelCase       						= torch.ones(__UpperCAmelCase			,  device=__UpperCAmelCase   )
      if token_type_ids is None:
       __UpperCamelCase       						= torch.zeros(__UpperCAmelCase			,  dtype=torch.long			,  device=__UpperCAmelCase   )

      # 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(__UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase   )

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

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

      __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 __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ):





     def __init__(			self			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

      '''simple docstring'''




      __UpperCamelCase       						= message
      __UpperCamelCase       						= exit_layer  # start from 1!


class __lowerCAmelCase	(   nn.Module    ):





     def __init__(			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      super().__init__()
      __UpperCamelCase       						= BertPooler(__UpperCAmelCase   )
      __UpperCamelCase       						= nn.Dropout(config.hidden_dropout_prob   )
      __UpperCamelCase       						= nn.Linear(config.hidden_size			,  config.num_labels   )







     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      __UpperCamelCase       						= encoder_outputs[0]
      __UpperCamelCase       						= self.pooler(__UpperCAmelCase   )
      # "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(__UpperCAmelCase   )
      __UpperCamelCase       						= self.classifier(__UpperCAmelCase   )

      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. "    ,						__SCREAMING_SNAKE_CASE    ,						)
class __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ):





     def __init__(			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      super().__init__(__UpperCAmelCase   )
      __UpperCamelCase       						= config.num_labels
      __UpperCamelCase       						= config.num_hidden_layers

      __UpperCamelCase       						= DeeBertModel(__UpperCAmelCase   )
      __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(__UpperCAmelCase   )
     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=-1			,  __UpperCAmelCase=False			,  ):

      '''simple docstring'''




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

       __UpperCamelCase       						= outputs[1]

       __UpperCamelCase       						= self.dropout(__UpperCAmelCase   )
       __UpperCamelCase       						= self.classifier(__UpperCAmelCase   )
       __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(__UpperCAmelCase   )
       __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(__UpperCAmelCase   )
         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(__UpperCAmelCase   )

       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)



