






from __future__ import annotations

import unittest

from transformers import DistilBertConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
       import tensorflow as tf

       from transformers.models.distilbert.modeling_tf_distilbert import (
           TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
           TFDistilBertForMaskedLM,
           TFDistilBertForMultipleChoice,
           TFDistilBertForQuestionAnswering,
           TFDistilBertForSequenceClassification,
           TFDistilBertForTokenClassification,
           TFDistilBertModel,
       )





class 					A__     :

     '''simple docstring'''






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






           """simple docstring"""




           UpperCamelCase					       =  parent
           UpperCamelCase					       =  13
           UpperCamelCase					       =  7
           UpperCamelCase					       =  True
           UpperCamelCase					       =  True
           UpperCamelCase					       =  False
           UpperCamelCase					       =  True
           UpperCamelCase					       =  99
           UpperCamelCase					       =  32
           UpperCamelCase					       =  2
           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







     def   _SCREAMING_SNAKE_CASE     ( self      :    str ):






           """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
           UpperCamelCase					       =  None
           UpperCamelCase					       =  None
           if self.use_labels:
                 UpperCamelCase					       =  ids_tensor([self.batch_size]   ,						self.type_sequence_label_size )
                 UpperCamelCase					       =  ids_tensor([self.batch_size, self.seq_length]   ,						self.num_labels )
                 UpperCamelCase					       =  ids_tensor([self.batch_size]   ,						self.num_choices )

           UpperCamelCase					       =  DistilBertConfig(
               vocab_size=self.vocab_size   ,						dim=self.hidden_size   ,						n_layers=self.num_hidden_layers   ,						n_heads=self.num_attention_heads   ,						hidden_dim=self.intermediate_size   ,						hidden_act=self.hidden_act   ,						dropout=self.hidden_dropout_prob   ,						attention_dropout=self.attention_probs_dropout_prob   ,						max_position_embeddings=self.max_position_embeddings   ,						initializer_range=self.initializer_range   ,						)

           return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels







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






           """simple docstring"""




           UpperCamelCase					       =  TFDistilBertModel(config=_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  {'input_ids': input_ids, 'attention_mask': input_mask}

           UpperCamelCase					       =  model(_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  [input_ids, input_mask]

           UpperCamelCase					       =  model(_SCREAMING_SNAKE_CASE )

           self.parent.assertEqual(result.last_hidden_state.shape   ,						(self.batch_size, self.seq_length, self.hidden_size) )







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






           """simple docstring"""




           UpperCamelCase					       =  TFDistilBertForMaskedLM(config=_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  {'input_ids': input_ids, 'attention_mask': input_mask}
           UpperCamelCase					       =  model(_SCREAMING_SNAKE_CASE )
           self.parent.assertEqual(result.logits.shape   ,						(self.batch_size, self.seq_length, self.vocab_size) )







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






           """simple docstring"""




           UpperCamelCase					       =  TFDistilBertForQuestionAnswering(config=_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  {
               'input_ids': input_ids,
               'attention_mask': input_mask,
           }
           UpperCamelCase					       =  model(_SCREAMING_SNAKE_CASE )
           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   _SCREAMING_SNAKE_CASE     ( self      :    int   ,						_SCREAMING_SNAKE_CASE      :    str   ,						_SCREAMING_SNAKE_CASE      :    str   ,						_SCREAMING_SNAKE_CASE      :    str   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    int   ,						_SCREAMING_SNAKE_CASE      :    List[Any] ):






           """simple docstring"""




           UpperCamelCase					       =  self.num_labels
           UpperCamelCase					       =  TFDistilBertForSequenceClassification(_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  {'input_ids': input_ids, 'attention_mask': input_mask}
           UpperCamelCase					       =  model(_SCREAMING_SNAKE_CASE )
           self.parent.assertEqual(result.logits.shape   ,						(self.batch_size, self.num_labels) )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Union[str, Any]   ,						_SCREAMING_SNAKE_CASE      :    int   ,						_SCREAMING_SNAKE_CASE      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    int   ,						_SCREAMING_SNAKE_CASE      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    str   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any] ):






           """simple docstring"""




           UpperCamelCase					       =  self.num_choices
           UpperCamelCase					       =  TFDistilBertForMultipleChoice(_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  tf.tile(tf.expand_dims(_SCREAMING_SNAKE_CASE   ,						1 )   ,						(1, self.num_choices, 1) )
           UpperCamelCase					       =  tf.tile(tf.expand_dims(_SCREAMING_SNAKE_CASE   ,						1 )   ,						(1, self.num_choices, 1) )
           UpperCamelCase					       =  {
               'input_ids': multiple_choice_inputs_ids,
               'attention_mask': multiple_choice_input_mask,
           }
           UpperCamelCase					       =  model(_SCREAMING_SNAKE_CASE )
           self.parent.assertEqual(result.logits.shape   ,						(self.batch_size, self.num_choices) )







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






           """simple docstring"""




           UpperCamelCase					       =  self.num_labels
           UpperCamelCase					       =  TFDistilBertForTokenClassification(_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  {'input_ids': input_ids, 'attention_mask': input_mask}
           UpperCamelCase					       =  model(_SCREAMING_SNAKE_CASE )
           self.parent.assertEqual(result.logits.shape   ,						(self.batch_size, self.seq_length, self.num_labels) )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Optional[Any] ):






           """simple docstring"""




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





@require_tf
class 					A__     (     __snake_case   ,    __snake_case   ,    unittest.TestCase ):

     '''simple docstring'''






     snake_case__					=					(
         (
             TFDistilBertModel,
             TFDistilBertForMaskedLM,
             TFDistilBertForQuestionAnswering,
             TFDistilBertForSequenceClassification,
             TFDistilBertForTokenClassification,
             TFDistilBertForMultipleChoice,
         )
         if is_tf_available()
         else None
     )
     snake_case__					=					(
         {
             """feature-extraction""": TFDistilBertModel,
             """fill-mask""": TFDistilBertForMaskedLM,
             """question-answering""": TFDistilBertForQuestionAnswering,
             """text-classification""": TFDistilBertForSequenceClassification,
             """token-classification""": TFDistilBertForTokenClassification,
             """zero-shot""": TFDistilBertForSequenceClassification,
         }
         if is_tf_available()
         else {}
     )
     snake_case__					=					False
     snake_case__					=					False







     def   _SCREAMING_SNAKE_CASE     ( self      :    Tuple ):






           """simple docstring"""




           UpperCamelCase					       =  TFDistilBertModelTester(self )
           UpperCamelCase					       =  ConfigTester(self   ,						config_class=_SCREAMING_SNAKE_CASE   ,						dim=37 )







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






           """simple docstring"""




           self.config_tester.run_common_tests()







     def   _SCREAMING_SNAKE_CASE     ( self      :    str ):






           """simple docstring"""




           UpperCamelCase					       =  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_model(*_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    str ):






           """simple docstring"""




           UpperCamelCase					       =  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_masked_lm(*_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Tuple ):






           """simple docstring"""




           UpperCamelCase					       =  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_question_answering(*_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Dict ):






           """simple docstring"""




           UpperCamelCase					       =  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_sequence_classification(*_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    int ):






           """simple docstring"""




           UpperCamelCase					       =  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_multiple_choice(*_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Optional[Any] ):






           """simple docstring"""




           UpperCamelCase					       =  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_token_classification(*_SCREAMING_SNAKE_CASE )







     @slow
     def   _SCREAMING_SNAKE_CASE     ( self      :    List[str] ):






           """simple docstring"""




           for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1] ):
                 UpperCamelCase					       =  TFDistilBertModel.from_pretrained(_SCREAMING_SNAKE_CASE )
                 self.assertIsNotNone(_SCREAMING_SNAKE_CASE )





@require_tf
class 					A__     (     unittest.TestCase ):

     '''simple docstring'''






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






           """simple docstring"""




           UpperCamelCase					       =  TFDistilBertModel.from_pretrained('distilbert-base-uncased' )
           UpperCamelCase					       =  tf.constant([[0, 1, 2, 3, 4, 5]] )
           UpperCamelCase					       =  model(_SCREAMING_SNAKE_CASE )[0]

           UpperCamelCase					       =  [1, 6, 768]
           self.assertEqual(output.shape   ,						_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  tf.constant(
               [
                   [
                       [0.1_9_2_6_1_8_8_5, -0.1_3_7_3_2_9_5_5, 0.4_1_1_9_7_9_9],
                       [0.2_2_1_5_0_1_5_6, -0.0_7_4_2_2_6_6_1, 0.3_9_0_3_7_2_0_4],
                       [0.2_2_7_5_6_0_1_8, -0.0_8_9_6_4_1_4, 0.3_7_0_1_4_6_7],
                   ]
               ] )
           tf.debugging.assert_near(output[:, :3, :3]   ,						_SCREAMING_SNAKE_CASE   ,						atol=1E-4 )











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)



