


from __future__ import annotations

import copy
import inspect
import json
import math
import os
import tempfile
import unittest
from importlib import import_module

import numpy as np

from transformers import ViTMAEConfig
from transformers.file_utils import cached_property, is_tf_available, is_vision_available
from transformers.testing_utils import require_tf, require_vision, slow

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


if is_tf_available():
   import tensorflow as tf

   from transformers import TFViTMAEForPreTraining, TFViTMAEModel


if is_vision_available():
   from PIL import Image

   from transformers import ViTImageProcessor







class lowerCAmelCase_						:

 '''simple docstring'''

 def __init__(     self		,				_UpperCAmelCase		,				_UpperCAmelCase=13		,				_UpperCAmelCase=30		,				_UpperCAmelCase=2		,				_UpperCAmelCase=3		,				_UpperCAmelCase=True		,				_UpperCAmelCase=True		,				_UpperCAmelCase=32		,				_UpperCAmelCase=2		,				_UpperCAmelCase=4		,				_UpperCAmelCase=37		,				_UpperCAmelCase="gelu"		,				_UpperCAmelCase=0.1		,				_UpperCAmelCase=0.1		,				_UpperCAmelCase=10		,				_UpperCAmelCase=0.02		,				_UpperCAmelCase=3		,				_UpperCAmelCase=0.6		,				_UpperCAmelCase=None		,				):
       snake_case_	=				parent
       snake_case_	=				batch_size
       snake_case_	=				image_size
       snake_case_	=				patch_size
       snake_case_	=				num_channels
       snake_case_	=				is_training
       snake_case_	=				use_labels
       snake_case_	=				hidden_size
       snake_case_	=				num_hidden_layers
       snake_case_	=				num_attention_heads
       snake_case_	=				intermediate_size
       snake_case_	=				hidden_act
       snake_case_	=				hidden_dropout_prob
       snake_case_	=				attention_probs_dropout_prob
       snake_case_	=				type_sequence_label_size
       snake_case_	=				initializer_range
       snake_case_	=				mask_ratio
       snake_case_	=				scope

       # in ViTMAE, the expected sequence length = (num_patches + 1) * (1 - config.mask_ratio), rounded above
       # (we add 1 for the [CLS] token)
       snake_case_	=				(image_size // patch_size) ** 2
       snake_case_	=				int(math.ceil((1 - mask_ratio) * (num_patches + 1) ) )



 def 						UpperCamelCase__					(     self ):
       snake_case_	=				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

       snake_case_	=				None
       if self.use_labels:
             snake_case_	=				ids_tensor([self.batch_size]		,				self.type_sequence_label_size )

       snake_case_	=				self.get_config()

       return config, pixel_values, labels



 def 						UpperCamelCase__					(     self ):
       return ViTMAEConfig(
           image_size=self.image_size		,				patch_size=self.patch_size		,				num_channels=self.num_channels		,				hidden_size=self.hidden_size		,				num_hidden_layers=self.num_hidden_layers		,				num_attention_heads=self.num_attention_heads		,				intermediate_size=self.intermediate_size		,				decoder_hidden_size=self.hidden_size		,				decoder_num_hidden_layers=self.num_hidden_layers		,				decoder_num_attention_heads=self.num_attention_heads		,				decoder_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		,				is_decoder=UpperCamelCase__		,				initializer_range=self.initializer_range		,				mask_ratio=self.mask_ratio		,				)



 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase ):
       snake_case_	=				TFViTMAEModel(config=UpperCamelCase__ )
       snake_case_	=				model(UpperCamelCase__		,				training=UpperCamelCase__ )
       self.parent.assertEqual(result.last_hidden_state.shape		,				(self.batch_size, self.seq_length, self.hidden_size) )



 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase ):
       snake_case_	=				TFViTMAEForPreTraining(UpperCamelCase__ )
       snake_case_	=				model(UpperCamelCase__		,				training=UpperCamelCase__ )
       # expected sequence length = num_patches
       snake_case_	=				(self.image_size // self.patch_size) ** 2
       snake_case_	=				self.patch_size**2 * self.num_channels
       self.parent.assertEqual(result.logits.shape		,				(self.batch_size, num_patches, expected_num_channels) )

       # test greyscale images
       snake_case_	=				1
       snake_case_	=				TFViTMAEForPreTraining(UpperCamelCase__ )

       snake_case_	=				floats_tensor([self.batch_size, 1, self.image_size, self.image_size] )
       snake_case_	=				model(UpperCamelCase__		,				training=UpperCamelCase__ )
       snake_case_	=				self.patch_size**2
       self.parent.assertEqual(result.logits.shape		,				(self.batch_size, num_patches, expected_num_channels) )





 def 						UpperCamelCase__					(     self ):
       snake_case_	=				self.prepare_config_and_inputs()
       ((snake_case_)	,   (snake_case_)	,   (snake_case_))	=				config_and_inputs
       snake_case_	=				{'''pixel_values''': pixel_values}
       return config, inputs_dict







@require_tf
class lowerCAmelCase_						(			__a ,						__a ,						unittest.TestCase	):

 '''simple docstring'''

 __snake_case												=  (TFViTMAEModel, TFViTMAEForPreTraining) if is_tf_available() else ()
 __snake_case												=  {"feature-extraction": TFViTMAEModel} if is_tf_available() else {}

 __snake_case												=  False
 __snake_case												=  False
 __snake_case												=  False
 __snake_case												=  False



 def 						UpperCamelCase__					(     self ):
       snake_case_	=				TFViTMAEModelTester(self )
       snake_case_	=				ConfigTester(self		,				config_class=UpperCamelCase__		,				has_text_modality=UpperCamelCase__		,				hidden_size=37 )



 def 						UpperCamelCase__					(     self ):
       self.config_tester.run_common_tests()



 @unittest.skip(reason='''ViTMAE does not use inputs_embeds''' )
 def 						UpperCamelCase__					(     self ):
       pass



 def 						UpperCamelCase__					(     self ):
       snake_case_	,   snake_case_	=				self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
             snake_case_	=				model_class(UpperCamelCase__ )
             self.assertIsInstance(model.get_input_embeddings()		,				(tf.keras.layers.Layer) )
             snake_case_	=				model.get_output_embeddings()
             self.assertTrue(x is None or isinstance(UpperCamelCase__		,				tf.keras.layers.Layer ) )



 def 						UpperCamelCase__					(     self ):
       snake_case_	,   snake_case_	=				self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
             snake_case_	=				model_class(UpperCamelCase__ )
             snake_case_	=				inspect.signature(model.call )
             # signature.parameters is an OrderedDict => so arg_names order is deterministic
             snake_case_	=				[*signature.parameters.keys()]

             snake_case_	=				['''pixel_values''']
             self.assertListEqual(arg_names[:1]		,				UpperCamelCase__ )



 def 						UpperCamelCase__					(     self ):
       snake_case_	=				self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*UpperCamelCase__ )



 def 						UpperCamelCase__					(     self ):
       snake_case_	=				self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_pretraining(*UpperCamelCase__ )



 def 						UpperCamelCase__					(     self ):
       np.random.seed(2 )

       snake_case_	,   snake_case_	=				self.model_tester.prepare_config_and_inputs_for_common()

       snake_case_	=				int((config.image_size // config.patch_size) ** 2 )
       snake_case_	=				np.random.uniform(size=(self.model_tester.batch_size, num_patches) )

       for model_class in self.all_model_classes:
             snake_case_	=				model_class(UpperCamelCase__ )
             snake_case_	=				self._prepare_for_class(UpperCamelCase__		,				UpperCamelCase__ )

             snake_case_	=				model(UpperCamelCase__		,				noise=UpperCamelCase__ )

             snake_case_	=				copy.deepcopy(self._prepare_for_class(UpperCamelCase__		,				UpperCamelCase__ ) )
             snake_case_	=				model(**UpperCamelCase__		,				noise=UpperCamelCase__ )
             snake_case_	=				outputs_dict[0].numpy()
             snake_case_	=				outputs_keywords[0].numpy()

             self.assertLess(np.sum(np.abs(output_dict - output_keywords ) )		,				1E-6 )



 def 						UpperCamelCase__					(     self ):
       np.random.seed(2 )

       snake_case_	,   snake_case_	=				self.model_tester.prepare_config_and_inputs_for_common()

       snake_case_	=				int((config.image_size // config.patch_size) ** 2 )
       snake_case_	=				np.random.uniform(size=(self.model_tester.batch_size, num_patches) )

       def prepare_numpy_arrays(_UpperCAmelCase ):
             snake_case_	=				{}
             for k, v in inputs_dict.items():
                   if tf.is_tensor(UpperCamelCase__ ):
                         snake_case_	=				v.numpy()
                   else:
                         snake_case_	=				np.array(UpperCamelCase__ )

             return inputs_np_dict

       for model_class in self.all_model_classes:
             snake_case_	=				model_class(UpperCamelCase__ )

             snake_case_	=				self._prepare_for_class(UpperCamelCase__		,				UpperCamelCase__ )
             snake_case_	=				prepare_numpy_arrays(UpperCamelCase__ )

             snake_case_	=				model(UpperCamelCase__		,				noise=UpperCamelCase__ )
             snake_case_	=				model(**UpperCamelCase__		,				noise=UpperCamelCase__ )
             self.assert_outputs_same(UpperCamelCase__		,				UpperCamelCase__ )



 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase ):
       np.random.seed(2 )

       snake_case_	=				int((tf_model.config.image_size // tf_model.config.patch_size) ** 2 )
       snake_case_	=				np.random.uniform(size=(self.model_tester.batch_size, num_patches) )
       snake_case_	=				tf.constant(UpperCamelCase__ )

       # Add `noise` argument.
       # PT inputs will be prepared in `super().check_pt_tf_models()` with this added `noise` argument
       snake_case_	=				tf_noise

       super().check_pt_tf_models(UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__ )



 def 						UpperCamelCase__					(     self ):
       np.random.seed(2 )

       snake_case_	,   snake_case_	=				self.model_tester.prepare_config_and_inputs_for_common()

       snake_case_	=				{
           module_member
           for model_class in self.all_model_classes
           for module in (import_module(model_class.__module__ ),)
           for module_member_name in dir(UpperCamelCase__ )
           if module_member_name.endswith('''MainLayer''' )
           # This condition is required, since `modeling_tf_clip.py` has 3 classes whose names end with `MainLayer`.
           and module_member_name[: -len('''MainLayer''' )] == model_class.__name__[: -len('''Model''' )]
           for module_member in (getattr(UpperCamelCase__		,				UpperCamelCase__ ),)
           if isinstance(UpperCamelCase__		,				UpperCamelCase__ )
           and tf.keras.layers.Layer in module_member.__bases__
           and getattr(UpperCamelCase__		,				'''_keras_serializable'''		,				UpperCamelCase__ )
       }

       snake_case_	=				int((config.image_size // config.patch_size) ** 2 )
       snake_case_	=				np.random.uniform(size=(self.model_tester.batch_size, num_patches) )
       snake_case_	=				tf.convert_to_tensor(UpperCamelCase__ )
       inputs_dict.update({'''noise''': noise} )

       for main_layer_class in tf_main_layer_classes:
             snake_case_	=				main_layer_class(UpperCamelCase__ )

             snake_case_	=				{
                 name: tf.keras.Input(tensor.shape[1:]		,				dtype=tensor.dtype ) for name, tensor in inputs_dict.items()
             }

             snake_case_	=				tf.keras.Model(UpperCamelCase__		,				outputs=main_layer(UpperCamelCase__ ) )
             snake_case_	=				model(UpperCamelCase__ )

             with tempfile.TemporaryDirectory() as tmpdirname:
                   snake_case_	=				os.path.join(UpperCamelCase__		,				'''keras_model.h5''' )
                   model.save(UpperCamelCase__ )
                   snake_case_	=				tf.keras.models.load_model(
                       UpperCamelCase__		,				custom_objects={main_layer_class.__name__: main_layer_class} )
                   assert isinstance(UpperCamelCase__		,				tf.keras.Model )
                   snake_case_	=				model(UpperCamelCase__ )
                   self.assert_outputs_same(UpperCamelCase__		,				UpperCamelCase__ )



 @slow
 def 						UpperCamelCase__					(     self ):
       np.random.seed(2 )

       snake_case_	,   snake_case_	=				self.model_tester.prepare_config_and_inputs_for_common()

       snake_case_	=				int((config.image_size // config.patch_size) ** 2 )
       snake_case_	=				np.random.uniform(size=(self.model_tester.batch_size, num_patches) )

       for model_class in self.all_model_classes:
             snake_case_	=				model_class(UpperCamelCase__ )
             snake_case_	=				self._prepare_for_class(UpperCamelCase__		,				UpperCamelCase__ )
             snake_case_	=				model(UpperCamelCase__		,				noise=UpperCamelCase__ )

             if model_class.__name__ == "TFViTMAEModel":
                   snake_case_	=				outputs.last_hidden_state.numpy()
                   snake_case_	=				0
             else:
                   snake_case_	=				outputs.logits.numpy()
                   snake_case_	=				0

             with tempfile.TemporaryDirectory() as tmpdirname:
                   model.save_pretrained(UpperCamelCase__		,				saved_model=UpperCamelCase__ )
                   snake_case_	=				model_class.from_pretrained(UpperCamelCase__ )
                   snake_case_	=				model(UpperCamelCase__		,				noise=UpperCamelCase__ )

                   if model_class.__name__ == "TFViTMAEModel":
                         snake_case_	=				after_outputs['''last_hidden_state'''].numpy()
                         snake_case_	=				0
                   else:
                         snake_case_	=				after_outputs['''logits'''].numpy()
                         snake_case_	=				0

                   snake_case_	=				np.amax(np.abs(out_a - out_a ) )
                   self.assertLessEqual(UpperCamelCase__		,				1E-5 )



 def 						UpperCamelCase__					(     self ):
       np.random.seed(2 )

       snake_case_	,   snake_case_	=				self.model_tester.prepare_config_and_inputs_for_common()

       snake_case_	=				int((config.image_size // config.patch_size) ** 2 )
       snake_case_	=				np.random.uniform(size=(self.model_tester.batch_size, num_patches) )

       for model_class in self.all_model_classes:
             snake_case_	=				model_class(UpperCamelCase__ )
             snake_case_	=				self._prepare_for_class(UpperCamelCase__		,				UpperCamelCase__ )

             snake_case_	=				model(UpperCamelCase__		,				noise=UpperCamelCase__ )
             snake_case_	=				model.get_config()
             # make sure that returned config is jsonifiable, which is required by keras
             json.dumps(UpperCamelCase__ )
             snake_case_	=				model_class.from_config(model.get_config() )
             # make sure it also accepts a normal config
             snake_case_	=				model_class.from_config(model.config )
             snake_case_	=				new_model(UpperCamelCase__ )  # Build model
             new_model.set_weights(model.get_weights() )
             snake_case_	=				new_model(UpperCamelCase__		,				noise=UpperCamelCase__ )

             self.assert_outputs_same(UpperCamelCase__		,				UpperCamelCase__ )



 @unittest.skip(
     reason='''ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load
    to get deterministic results.''' )
 def 						UpperCamelCase__					(     self ):
       pass



 @unittest.skip(reason='''ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load''' )
 def 						UpperCamelCase__					(     self ):
       pass





 @slow
 def 						UpperCamelCase__					(     self ):
       snake_case_	=				TFViTMAEModel.from_pretrained('''google/vit-base-patch16-224''' )
       self.assertIsNotNone(UpperCamelCase__ )






def 	__lowerCAmelCase   ()->  Optional[Any]:




      """simple docstring"""

      snake_case_	=				Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png''' )
      return image







@require_tf
@require_vision
class lowerCAmelCase_						(			unittest.TestCase	):

 '''simple docstring'''

 @cached_property
 def 						UpperCamelCase__					(     self ):
       return ViTImageProcessor.from_pretrained('''facebook/vit-mae-base''' ) if is_vision_available() else None





 @slow
 def 						UpperCamelCase__					(     self ):
       np.random.seed(2 )

       snake_case_	=				TFViTMAEForPreTraining.from_pretrained('''facebook/vit-mae-base''' )

       snake_case_	=				self.default_image_processor
       snake_case_	=				prepare_img()
       snake_case_	=				image_processor(images=UpperCamelCase__		,				return_tensors='''tf''' )

       # prepare a noise vector that will be also used for testing the TF model
       # (this way we can ensure that the PT and TF models operate on the same inputs)
       snake_case_	=				ViTMAEConfig()
       snake_case_	=				int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2 )
       snake_case_	=				np.random.uniform(size=(1, num_patches) )

       # forward pass
       snake_case_	=				model(**UpperCamelCase__		,				noise=UpperCamelCase__ )

       # verify the logits
       snake_case_	=				tf.convert_to_tensor([1, 1_96, 7_68] )
       self.assertEqual(outputs.logits.shape		,				UpperCamelCase__ )

       snake_case_	=				tf.convert_to_tensor(
           [[-0.0_548, -1.7_023, -0.9_325], [0.3_721, -0.5_670, -0.2_233], [0.8_235, -1.3_878, -0.3_524]] )

       tf.debugging.assert_near(outputs.logits[0, :3, :3]		,				UpperCamelCase__		,				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 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE )->  List[Any]:




      """simple docstring"""

      snake_case_	=				torch.exp(SCREAMING_SNAKE_CASE )
      snake_case_	=				torch.sum(SCREAMING_SNAKE_CASE , 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(SCREAMING_SNAKE_CASE ) - B / A







class lowerCAmelCase_						(			nn.Module	):

 '''simple docstring'''

 def __init__(     self		,				_UpperCAmelCase ):
       super().__init__()
       snake_case_	=				config.output_attentions
       snake_case_	=				config.output_hidden_states
       snake_case_	=				nn.ModuleList([BertLayer(_UpperCAmelCase ) for _ in range(config.num_hidden_layers )] )
       snake_case_	=				nn.ModuleList([BertHighway(_UpperCAmelCase ) for _ in range(config.num_hidden_layers )] )

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



 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase ):
       if (type(_UpperCAmelCase ) is float) or (type(_UpperCAmelCase ) is int):
             for i in range(len(self.early_exit_entropy ) ):
                   snake_case_	=				x
       else:
             snake_case_	=				x



 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase ):
       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 						UpperCamelCase__					(     self		,				_UpperCAmelCase		,				_UpperCAmelCase=None		,				_UpperCAmelCase=None		,				_UpperCAmelCase=None		,				_UpperCAmelCase=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(
                 _UpperCAmelCase		,				_UpperCAmelCase		,				head_mask[i]		,				_UpperCAmelCase		,				_UpperCAmelCase )
             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](_UpperCAmelCase )
             # logits, pooled_output

             if not self.training:
                   snake_case_	=				highway_exit[0]
                   snake_case_	=				entropy(_UpperCAmelCase )
                   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(_UpperCAmelCase		,				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). " ,						lowerCamelCase__ ,						)
class lowerCAmelCase_						(			lowerCamelCase__	):

 '''simple docstring'''

 def __init__(     self		,				_UpperCAmelCase ):
       super().__init__(_UpperCAmelCase )
       snake_case_	=				config

       snake_case_	=				BertEmbeddings(_UpperCAmelCase )
       snake_case_	=				DeeBertEncoder(_UpperCAmelCase )
       snake_case_	=				BertPooler(_UpperCAmelCase )

       self.init_weights()



 def 						UpperCamelCase__					(     self ):
       self.encoder.init_highway_pooler(self.pooler )



 def 						UpperCamelCase__					(     self ):
       return self.embeddings.word_embeddings



 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase ):
       snake_case_	=				value



 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase ):
       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		,				):
       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(_UpperCAmelCase		,				device=_UpperCAmelCase )
       if encoder_attention_mask is None:
             snake_case_	=				torch.ones(_UpperCAmelCase		,				device=_UpperCAmelCase )
       if token_type_ids is None:
             snake_case_	=				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.
       snake_case_	=				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:
             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) * -10_000.0

       # Prepare head mask if needed
       # 1.0 in head_mask indicate we keep the head
       # attention_probs has shape bsz x n_heads x N x N
       # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
       # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
       snake_case_	=				self.get_head_mask(_UpperCAmelCase		,				self.config.num_hidden_layers )

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

       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 lowerCAmelCase_						(			lowerCamelCase__	):

 '''simple docstring'''

 def __init__(     self		,				_UpperCAmelCase		,				_UpperCAmelCase ):
       snake_case_	=				message
       snake_case_	=				exit_layer  # start from 1!







class lowerCAmelCase_						(			nn.Module	):

 '''simple docstring'''

 def __init__(     self		,				_UpperCAmelCase ):
       super().__init__()
       snake_case_	=				BertPooler(_UpperCAmelCase )
       snake_case_	=				nn.Dropout(config.hidden_dropout_prob )
       snake_case_	=				nn.Linear(config.hidden_size		,				config.num_labels )





 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase ):
       # Pooler
       snake_case_	=				encoder_outputs[0]
       snake_case_	=				self.pooler(_UpperCAmelCase )
       # "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(_UpperCAmelCase )
       snake_case_	=				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. " ,						lowerCamelCase__ ,						)
class lowerCAmelCase_						(			lowerCamelCase__	):

 '''simple docstring'''

 def __init__(     self		,				_UpperCAmelCase ):
       super().__init__(_UpperCAmelCase )
       snake_case_	=				config.num_labels
       snake_case_	=				config.num_hidden_layers

       snake_case_	=				DeeBertModel(_UpperCAmelCase )
       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(_UpperCAmelCase )
 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase=None		,				_UpperCAmelCase=None		,				_UpperCAmelCase=None		,				_UpperCAmelCase=None		,				_UpperCAmelCase=None		,				_UpperCAmelCase=None		,				_UpperCAmelCase=None		,				_UpperCAmelCase=-1		,				_UpperCAmelCase=False		,				):
       snake_case_	=				self.num_layers
       try:
             snake_case_	=				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

             snake_case_	=				outputs[1]

             snake_case_	=				self.dropout(_UpperCAmelCase )
             snake_case_	=				self.classifier(_UpperCAmelCase )
             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(_UpperCAmelCase )
             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(_UpperCAmelCase )
                         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(_UpperCAmelCase )

             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)