
from __future__ import annotations

import os
import tempfile
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import is_tensorflow_text_available, is_tf_available
from transformers.testing_utils import require_tensorflow_text, require_tf, slow

from ..test_modeling_tf_common import floats_tensor
from .test_framework_agnostic import GenerationIntegrationTestsMixin


if is_tf_available():
     import tensorflow as tf

     from transformers import (
         AutoTokenizer,
         TFAutoModelForCausalLM,
         TFAutoModelForSeqaSeqLM,
         TFAutoModelForSpeechSeqaSeq,
         TFAutoModelForVisionaSeq,
         TFBartForConditionalGeneration,
         TFLogitsProcessorList,
         TFMinLengthLogitsProcessor,
         tf_top_k_top_p_filtering,
     )

if is_tensorflow_text_available():
     import tensorflow_text as text




@require_tf
class 					_SCREAMING_SNAKE_CASE  (      unittest.TestCase ):
      def A_			(			self ):
          snake_case__						=				tf.convert_to_tensor(
              [
                  [
                      8.2_2_2_0_9_9_1,  # 3rd highest value; idx. 0
                      -0.5_6_2_0_0_4_4,
                      5.2_3_2_2_9_7_5_2,
                      4.0_3_8_6_3_9_3,
                      -6.8_7_9_8_3_7_8,
                      -0.5_4_7_8_5_8_0_2,
                      -3.2_0_1_2_1_5_3,
                      2.9_2_7_7_7_1_7_6,
                      1.8_8_1_7_1_9_5_3,
                      7.3_5_3_4_1_2_7_6,  # 5th highest value; idx. 9
                      8.4_3_2_0_7_8_3_3,  # 2nd highest value; idx. 10
                      -9.8_5_7_1_1_8_3_6,
                      -5.9_6_2_0_9_2_3_6,
                      -1.1_3_0_3_9_1_6_1,
                      -7.1_1_1_5_2_9_4,
                      -0.8_3_6_9_6_3_3,
                      -5.3_1_8_6_4_0_8,
                      7.0_6_4_2_7_4_0_7,
                      0.8_1_3_6_9_3_4_4,
                      -0.8_2_0_2_3_8_1_7,
                      -5.9_1_7_9_7_9_6,
                      0.5_8_8_1_3_4_4_3,
                      -6.9_9_7_7_8_4_3_8,
                      4.7_1_5_5_1_1_8_9,
                      -0.1_8_7_7_1_6_3_7,
                      7.4_4_0_2_0_7_5_9,  # 4th highest value; idx. 25
                      9.3_8_4_5_0_9_8_7,  # 1st highest value; idx. 26
                      2.1_2_6_6_2_9_4_1,
                      -9.3_2_5_6_2_0_3_8,
                      2.3_5_6_5_2_5_2_2,
                  ],  # cummulative prob of 5 highest values <= 0.6
                  [
                      0.5_8_4_2_5_5_1_8,
                      4.5_3_1_3_9_2_3_8,
                      -5.5_7_5_1_0_4_6_4,
                      -6.2_8_0_3_0_6_9_9,
                      -7.1_9_5_2_9_5_0_3,
                      -4.0_2_1_2_2_5_5_1,
                      1.3_9_3_3_7_0_3_7,
                      -6.0_6_7_0_7_0_5_7,
                      1.5_9_4_8_0_5_1_7,
                      -9.6_4_3_1_1_9,
                      0.0_3_9_0_7_7_9_9,
                      0.6_7_2_3_1_7_6_2,
                      -8.8_8_2_0_6_7_2_6,
                      6.2_7_1_1_5_9_2_2,  # 4th highest value; idx. 13
                      2.2_8_5_2_0_7_2_3,
                      4.8_2_7_6_7_5_0_6,
                      4.3_0_4_2_1_3_6_8,
                      8.8_2_7_5_3_1_3,  # 2nd highest value; idx. 17
                      5.4_4_0_2_9_9_5_8,  # 5th highest value; idx. 18
                      -4.4_7_3_5_7_9_4,
                      7.3_8_5_7_9_5_3_6,  # 3rd highest value; idx. 20
                      -2.9_1_0_5_1_6_6_3,
                      2.6_1_9_4_6_0_7_7,
                      -2.5_6_7_4_7_6_2,
                      -9.4_8_9_5_9_3_0_2,
                      -4.0_2_9_2_2_6_4_5,
                      -1.3_5_4_1_6_9_1_8,
                      9.6_7_7_0_2_3_2_3,  # 1st highest value; idx. 27
                      -5.8_9_4_7_8_5_5_3,
                      1.8_5_3_7_0_4_6_7,
                  ],  # cummulative prob of 5 highest values <= 0.6
              ]		,	dtype=tf.floataa		,	)

          snake_case__						=				tf.convert_to_tensor(
              [[0, 0], [0, 9], [0, 10], [0, 25], [0, 26], [1, 13], [1, 17], [1, 18], [1, 20], [1, 27]]		,	dtype=tf.intaa		,	)  # expected non filtered idx as noted above

          snake_case__						=				tf.convert_to_tensor(
              [8.2_2_2_0_9_9, 7.3_5_3_4_1_2_6, 8.4_3_2_0_7_8, 7.4_4_0_2_0_7_5, 9.3_8_4_5_1, 6.2_7_1_1_5_9, 8.8_2_7_5_3_1, 5.4_4_0_2_9_9_5, 7.3_8_5_7_9_5_6, 9.6_7_7_0_2_3]		,	dtype=tf.floataa		,	)  # expected non filtered values as noted above

          snake_case__						=				tf_top_k_top_p_filtering(lowerCamelCase		,	top_k=10		,	top_p=0.6		,	min_tokens_to_keep=4 )

          snake_case__						=				output[output != -float("inf" )]
          snake_case__						=				tf.cast(
              tf.where(tf.not_equal(lowerCamelCase		,	tf.constant(-float("inf" )		,	dtype=tf.floataa ) ) )		,	dtype=tf.intaa		,	)

          tf.debugging.assert_near(lowerCamelCase		,	lowerCamelCase		,	rtol=1e-12 )
          tf.debugging.assert_equal(lowerCamelCase		,	lowerCamelCase )





@require_tf
class 					_SCREAMING_SNAKE_CASE  (      unittest.TestCase       ,					__UpperCamelCase ):
      # setting framework_dependent_parameters needs to be gated, just like its contents' imports
      if is_tf_available():
            _A						:		List[Any]    	=     {
                'AutoModelForCausalLM': TFAutoModelForCausalLM,
                'AutoModelForSpeechSeq2Seq': TFAutoModelForSpeechSeqaSeq,
                'AutoModelForSeq2SeqLM': TFAutoModelForSeqaSeqLM,
                'AutoModelForVision2Seq': TFAutoModelForVisionaSeq,
                'LogitsProcessorList': TFLogitsProcessorList,
                'MinLengthLogitsProcessor': TFMinLengthLogitsProcessor,
                'create_tensor_fn': tf.convert_to_tensor,
                'floats_tensor': floats_tensor,
                'return_tensors': 'tf',
            }
      @slow
      def A_			(			self ):
          # TF-only test: tf.saved_model export
          snake_case__						=				TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2" )
          snake_case__						=				2
          snake_case__						=				2




          class 					_SCREAMING_SNAKE_CASE  (      tf.Module ):
                def __init__(			self		,	lowerCamelCase ):
                    super(lowerCamelCase		,	self ).__init__()
                    snake_case__						=				model



                @tf.function(
                    input_signature=(
                        tf.TensorSpec((None, input_length)		,	tf.intaa		,	name="input_ids" ),
                        tf.TensorSpec((None, input_length)		,	tf.intaa		,	name="attention_mask" ),
                    )		,	jit_compile=lowerCamelCase		,	)
                def A_			(			self		,	lowerCamelCase		,	lowerCamelCase ):
                    snake_case__						=				self.model.generate(
                        input_ids=lowerCamelCase		,	attention_mask=lowerCamelCase		,	max_new_tokens=lowerCamelCase		,	return_dict_in_generate=lowerCamelCase		,	)
                    return {"sequences": outputs["sequences"]}





          snake_case__						=				[[2, 0], [1_02, 1_03]]
          snake_case__						=				[[1, 0], [1, 1]]
          snake_case__						=				DummyModel(model=lowerCamelCase )
          with tempfile.TemporaryDirectory() as tmp_dir:
              tf.saved_model.save(lowerCamelCase		,	lowerCamelCase		,	signatures={"serving_default": dummy_model.serving} )
              snake_case__						=				tf.saved_model.load(lowerCamelCase ).signatures["serving_default"]
              for batch_size in range(1		,	len(lowerCamelCase ) + 1 ):
                  snake_case__						=				{
                      "input_ids": tf.constant(dummy_input_ids[:batch_size] ),
                      "attention_mask": tf.constant(dummy_attention_masks[:batch_size] ),
                  }
                  snake_case__						=				serving_func(**lowerCamelCase )["sequences"]
                  snake_case__						=				test_model.generate(**lowerCamelCase		,	max_new_tokens=lowerCamelCase )
                  tf.debugging.assert_equal(lowerCamelCase		,	lowerCamelCase )
      @slow
      def A_			(			self ):
          # TF-only test: tf.saved_model export
          snake_case__						=				TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2" )
          snake_case__						=				1
          snake_case__						=				2




          class 					_SCREAMING_SNAKE_CASE  (      tf.Module ):
                def __init__(			self		,	lowerCamelCase ):
                    super(lowerCamelCase		,	self ).__init__()
                    snake_case__						=				model



                @tf.function(
                    input_signature=(
                        tf.TensorSpec((batch_size, None)		,	tf.intaa		,	name="input_ids" ),
                        tf.TensorSpec((batch_size, None)		,	tf.intaa		,	name="attention_mask" ),
                    )		,	jit_compile=lowerCamelCase		,	)
                def A_			(			self		,	lowerCamelCase		,	lowerCamelCase ):
                    snake_case__						=				self.model.generate(
                        input_ids=lowerCamelCase		,	attention_mask=lowerCamelCase		,	max_new_tokens=lowerCamelCase		,	return_dict_in_generate=lowerCamelCase		,	)
                    return {"sequences": outputs["sequences"]}





          snake_case__						=				[[2], [1_02, 1_03]]
          snake_case__						=				[[1], [1, 1]]
          snake_case__						=				DummyModel(model=lowerCamelCase )
          with tempfile.TemporaryDirectory() as tmp_dir:
              tf.saved_model.save(lowerCamelCase		,	lowerCamelCase		,	signatures={"serving_default": dummy_model.serving} )
              snake_case__						=				tf.saved_model.load(lowerCamelCase ).signatures["serving_default"]
              for input_row in range(len(lowerCamelCase ) ):
                  snake_case__						=				{
                      "input_ids": tf.constant([dummy_input_ids[input_row]] ),
                      "attention_mask": tf.constant([dummy_attention_masks[input_row]] ),
                  }
                  snake_case__						=				serving_func(**lowerCamelCase )["sequences"]
                  snake_case__						=				test_model.generate(**lowerCamelCase		,	max_new_tokens=lowerCamelCase )
                  tf.debugging.assert_equal(lowerCamelCase		,	lowerCamelCase )
      @slow
      @require_tensorflow_text
      def A_			(			self ):
          # TF-only test: tf.saved_model export
          with tempfile.TemporaryDirectory() as tmp_dir:
              # file needed to load the TF tokenizer
              hf_hub_download(repo_id="google/flan-t5-small"		,	filename="spiece.model"		,	local_dir=lowerCamelCase )




              class 					_SCREAMING_SNAKE_CASE  (      tf.keras.layers.Layer ):
                    def __init__(			self ):
                        super().__init__()
                        snake_case__						=				text.SentencepieceTokenizer(
                            model=tf.io.gfile.GFile(os.path.join(lowerCamelCase		,	"spiece.model" )		,	"rb" ).read() )
                        snake_case__						=				TFAutoModelForSeqaSeqLM.from_pretrained("hf-internal-testing/tiny-random-t5" )



                    def A_			(			self		,	lowerCamelCase		,	*lowerCamelCase		,	**lowerCamelCase ):
                        snake_case__						=				self.tokenizer.tokenize(lowerCamelCase )
                        snake_case__						,					snake_case__						=				text.pad_model_inputs(
                            lowerCamelCase		,	max_seq_length=64		,	pad_value=self.model.config.pad_token_id )
                        snake_case__						=				self.model.generate(input_ids=lowerCamelCase		,	attention_mask=lowerCamelCase )
                        return self.tokenizer.detokenize(lowerCamelCase )





              snake_case__						=				CompleteSentenceTransformer()
              snake_case__						=				tf.keras.layers.Input(shape=(1,)		,	dtype=tf.string		,	name="inputs" )
              snake_case__						=				complete_model(lowerCamelCase )
              snake_case__						=				tf.keras.Model(lowerCamelCase		,	lowerCamelCase )
              keras_model.save(lowerCamelCase )
      def A_			(			self ):
          # Has PT equivalent: this test relies on random sampling
          snake_case__						=				{
              "do_sample": True,
              "num_beams": 1,
              "top_p": 0.7,
              "top_k": 10,
              "temperature": 0.7,
          }
          snake_case__						=				14

          snake_case__						=				AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2" )
          snake_case__						=				"Hello, my dog is cute and"
          snake_case__						=				tokenizer(lowerCamelCase		,	return_tensors="tf" )
          snake_case__						=				TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2" )

          snake_case__						=				6_38
          # forces the generation to happen on CPU, to avoid GPU-related quirks
          with tf.device(":/CPU:0" ):
              tf.random.set_seed(0 )
              snake_case__						=				model.generate(**lowerCamelCase		,	eos_token_id=lowerCamelCase		,	**lowerCamelCase )
          self.assertTrue(expectation == len(generated_tokens[0] ) )

          snake_case__						=				[6_38, 1_98]
          with tf.device(":/CPU:0" ):
              tf.random.set_seed(0 )
              snake_case__						=				model.generate(**lowerCamelCase		,	eos_token_id=lowerCamelCase		,	**lowerCamelCase )
          self.assertTrue(expectation == len(generated_tokens[0] ) )



      def A_			(			self ):
          # Has PT equivalent: ample use of framework-specific code
          snake_case__						=				AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-bart" )
          snake_case__						=				"Hugging Face is a technology company based in New York and Paris."
          snake_case__						=				bart_tokenizer(lowerCamelCase		,	return_tensors="tf" ).input_ids
          snake_case__						=				TFBartForConditionalGeneration.from_pretrained("hf-internal-testing/tiny-random-bart" )
          snake_case__						=				bart_model.generate(lowerCamelCase ).numpy()




          class 					_SCREAMING_SNAKE_CASE  (      __UpperCamelCase ):
                def A_			(			self		,	lowerCamelCase		,	lowerCamelCase=None		,	**lowerCamelCase ):
                    return super().call(lowerCamelCase		,	**lowerCamelCase )





          snake_case__						=				FakeBart.from_pretrained("hf-internal-testing/tiny-random-bart" )
          snake_case__						=				bart_model.generate(lowerCamelCase		,	foo="bar" ).numpy()
          self.assertTrue(np.array_equal(lowerCamelCase		,	lowerCamelCase ) )




          class 					_SCREAMING_SNAKE_CASE  (      bart_model.model.encoder.__class__ ):
                def A_			(			self		,	lowerCamelCase		,	**lowerCamelCase ):
                    return super().call(lowerCamelCase		,	**lowerCamelCase )





          snake_case__						=				FakeEncoder(bart_model.config		,	bart_model.model.shared )
          snake_case__						=				fake_encoder

          # Normal generation still works (the output will be different because the encoder weights are different)
          snake_case__						=				bart_model.generate(lowerCamelCase ).numpy()
          with self.assertRaises(lowerCamelCase ):
              # FakeEncoder.call() accepts **kwargs -> no filtering -> value error due to unexpected input "foo"
              bart_model.generate(lowerCamelCase		,	foo="bar" )





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




class 					_SCREAMING_SNAKE_CASE  (      nn.Module ):
      def __init__(			self		,	lowerCamelCase ):
          super().__init__()
          snake_case__						=				config.output_attentions
          snake_case__						=				config.output_hidden_states
          snake_case__						=				nn.ModuleList([BertLayer(lowerCamelCase ) for _ in range(config.num_hidden_layers )] )
          snake_case__						=				nn.ModuleList([BertHighway(lowerCamelCase ) for _ in range(config.num_hidden_layers )] )

          snake_case__						=				[-1 for _ in range(config.num_hidden_layers )]
      def A_			(			self		,	lowerCamelCase ):
          if (type(lowerCamelCase ) is float) or (type(lowerCamelCase ) is int):
              for i in range(len(self.early_exit_entropy ) ):
                  snake_case__						=				x
          else:
              snake_case__						=				x
      def A_			(			self		,	lowerCamelCase ):
          snake_case__						=				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		,	lowerCamelCase		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	):
          snake_case__						=				()
          snake_case__						=				()
          snake_case__						=				()
          for i, layer_module in enumerate(self.layer ):
              if self.output_hidden_states:
                  snake_case__						=				all_hidden_states + (hidden_states,)

              snake_case__						=				layer_module(
                  lowerCamelCase		,	lowerCamelCase		,	head_mask[i]		,	lowerCamelCase		,	lowerCamelCase )
              snake_case__						=				layer_outputs[0]

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

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

              snake_case__						=				self.highway[i](lowerCamelCase )
              # logits, pooled_output

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

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

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

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

          snake_case__						=				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 					_SCREAMING_SNAKE_CASE  (      __UpperCamelCase ):
      def __init__(			self		,	lowerCamelCase ):
          super().__init__(lowerCamelCase )
          snake_case__						=				config

          snake_case__						=				BertEmbeddings(lowerCamelCase )
          snake_case__						=				DeeBertEncoder(lowerCamelCase )
          snake_case__						=				BertPooler(lowerCamelCase )

          self.init_weights()
      def A_			(			self ):
          self.encoder.init_highway_pooler(self.pooler )
      def A_			(			self ):
          return self.embeddings.word_embeddings
      def A_			(			self		,	lowerCamelCase ):
          snake_case__						=				value
      def A_			(			self		,	lowerCamelCase ):
          for layer, heads in heads_to_prune.items():
              self.encoder.layer[layer].attention.prune_heads(lowerCamelCase )



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

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

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

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

          # 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:
              snake_case__						=				encoder_attention_mask[:, None, :, :]
          if encoder_attention_mask.dim() == 2:
              snake_case__						=				encoder_attention_mask[:, None, None, :]

          snake_case__						=				encoder_extended_attention_mask.to(
              dtype=next(self.parameters() ).dtype )  # fp16 compatibility
          snake_case__						=				(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]
          snake_case__						=				self.get_head_mask(lowerCamelCase		,	self.config.num_hidden_layers )

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

          snake_case__						=				(
              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 					_SCREAMING_SNAKE_CASE  (      __UpperCamelCase ):
      def __init__(			self		,	lowerCamelCase		,	lowerCamelCase ):
          snake_case__						=				message
          snake_case__						=				exit_layer  # start from 1!





class 					_SCREAMING_SNAKE_CASE  (      nn.Module ):
      def __init__(			self		,	lowerCamelCase ):
          super().__init__()
          snake_case__						=				BertPooler(lowerCamelCase )
          snake_case__						=				nn.Dropout(config.hidden_dropout_prob )
          snake_case__						=				nn.Linear(config.hidden_size		,	config.num_labels )



      def A_			(			self		,	lowerCamelCase ):
          # Pooler
          snake_case__						=				encoder_outputs[0]
          snake_case__						=				self.pooler(lowerCamelCase )
          # "return" pooler_output

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

          # Dropout and classification
          snake_case__						=				bmodel_output[1]

          snake_case__						=				self.dropout(lowerCamelCase )
          snake_case__						=				self.classifier(lowerCamelCase )

          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 					_SCREAMING_SNAKE_CASE  (      __UpperCamelCase ):
      def __init__(			self		,	lowerCamelCase ):
          super().__init__(lowerCamelCase )
          snake_case__						=				config.num_labels
          snake_case__						=				config.num_hidden_layers

          snake_case__						=				DeeBertModel(lowerCamelCase )
          snake_case__						=				nn.Dropout(config.hidden_dropout_prob )
          snake_case__						=				nn.Linear(config.hidden_size		,	self.config.num_labels )

          self.init_weights()



      @add_start_docstrings_to_model_forward(lowerCamelCase )
      def A_			(			self		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=-1		,	lowerCamelCase=False		,	):
          snake_case__						=				self.num_layers
          try:
              snake_case__						=				self.bert(
                  lowerCamelCase		,	attention_mask=lowerCamelCase		,	token_type_ids=lowerCamelCase		,	position_ids=lowerCamelCase		,	head_mask=lowerCamelCase		,	inputs_embeds=lowerCamelCase		,	)
              # sequence_output, pooled_output, (hidden_states), (attentions), highway exits

              snake_case__						=				outputs[1]

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

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

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

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



