
"""simple docstring"""
from __future__ import annotations

import unittest

from transformers import RoFormerConfig, 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 import (
			    TFRoFormerForCausalLM,
			    TFRoFormerForMaskedLM,
			    TFRoFormerForMultipleChoice,
			    TFRoFormerForQuestionAnswering,
			    TFRoFormerForSequenceClassification,
			    TFRoFormerForTokenClassification,
			    TFRoFormerModel,
			)
			from transformers.models.roformer.modeling_tf_roformer import (
			    TFRoFormerSelfAttention,
			    TFRoFormerSinusoidalPositionalEmbedding,
			)
class       lowerCAmelCase_		:



							'''simple docstring'''





							def __init__( self    ,				snake_case_    ,				snake_case_=13    ,				snake_case_=7    ,				snake_case_=True    ,				snake_case_=True    ,				snake_case_=True    ,				snake_case_=True    ,				snake_case_=99    ,				snake_case_=32    ,				snake_case_=2    ,				snake_case_=4    ,				snake_case_=37    ,				snake_case_="gelu"    ,				snake_case_=0.1    ,				snake_case_=0.1    ,				snake_case_=512    ,				snake_case_=16    ,				snake_case_=2    ,				snake_case_=0.02    ,				snake_case_=3    ,				snake_case_=4    ,				snake_case_=None    ,				)       ->  List[str]:
													__lowerCAmelCase     =  parent
													__lowerCAmelCase     =  13
													__lowerCAmelCase     =  7
													__lowerCAmelCase     =  True
													__lowerCAmelCase     =  True
													__lowerCAmelCase     =  True
													__lowerCAmelCase     =  True
													__lowerCAmelCase     =  99
													__lowerCAmelCase     =  32
													__lowerCAmelCase     =  2
													__lowerCAmelCase     =  4
													__lowerCAmelCase     =  37
													__lowerCAmelCase     =  """gelu"""
													__lowerCAmelCase     =  0.1
													__lowerCAmelCase     =  0.1
													__lowerCAmelCase     =  512
													__lowerCAmelCase     =  16
													__lowerCAmelCase     =  2
													__lowerCAmelCase     =  0.02
													__lowerCAmelCase     =  3
													__lowerCAmelCase     =  4
													__lowerCAmelCase     =  None




							def    A__						( self				)       ->  List[str]:
													__lowerCAmelCase     =  ids_tensor([self.batch_size, self.seq_length]    ,				self.vocab_size				)

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

													__lowerCAmelCase     =  None
													if self.use_token_type_ids:
																			__lowerCAmelCase     =  ids_tensor([self.batch_size, self.seq_length]    ,				self.type_vocab_size				)

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

													__lowerCAmelCase     =  RoFormerConfig(
													    vocab_size=self.vocab_size    ,				hidden_size=self.hidden_size    ,				num_hidden_layers=self.num_hidden_layers    ,				num_attention_heads=self.num_attention_heads    ,				intermediate_size=self.intermediate_size    ,				hidden_act=self.hidden_act    ,				hidden_dropout_prob=self.hidden_dropout_prob    ,				attention_probs_dropout_prob=self.attention_probs_dropout_prob    ,				max_position_embeddings=self.max_position_embeddings    ,				type_vocab_size=self.type_vocab_size    ,				initializer_range=self.initializer_range    ,				return_dict=snake_case_    ,				)

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




							def    A__						( self    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_				)       ->  List[str]:
													__lowerCAmelCase     =  TFRoFormerModel(config=snake_case_				)
													__lowerCAmelCase     =  {"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}

													__lowerCAmelCase     =  [input_ids, input_mask]
													__lowerCAmelCase     =  model(snake_case_				)

													__lowerCAmelCase     =  model(snake_case_				)

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




							def    A__						( self    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_				)       ->  Optional[int]:
													__lowerCAmelCase     =  True
													__lowerCAmelCase     =  TFRoFormerForCausalLM(config=snake_case_				)
													__lowerCAmelCase     =  {
													    """input_ids""": input_ids,
													    """attention_mask""": input_mask,
													    """token_type_ids""": token_type_ids,
													}
													__lowerCAmelCase     =  model(snake_case_				)["""logits"""]
													self.parent.assertListEqual(
													    list(prediction_scores.numpy().shape				)    ,				[self.batch_size, self.seq_length, self.vocab_size]				)




							def    A__						( self    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_				)       ->  Union[str, Any]:
													__lowerCAmelCase     =  TFRoFormerForMaskedLM(config=snake_case_				)
													__lowerCAmelCase     =  {
													    """input_ids""": input_ids,
													    """attention_mask""": input_mask,
													    """token_type_ids""": token_type_ids,
													}
													__lowerCAmelCase     =  model(snake_case_				)
													self.parent.assertEqual(result.logits.shape    ,				(self.batch_size, self.seq_length, self.vocab_size)				)




							def    A__						( self    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_				)       ->  Optional[int]:
													__lowerCAmelCase     =  self.num_labels
													__lowerCAmelCase     =  TFRoFormerForSequenceClassification(config=snake_case_				)
													__lowerCAmelCase     =  {
													    """input_ids""": input_ids,
													    """attention_mask""": input_mask,
													    """token_type_ids""": token_type_ids,
													}

													__lowerCAmelCase     =  model(snake_case_				)
													self.parent.assertEqual(result.logits.shape    ,				(self.batch_size, self.num_labels)				)




							def    A__						( self    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_				)       ->  Union[str, Any]:
													__lowerCAmelCase     =  self.num_choices
													__lowerCAmelCase     =  TFRoFormerForMultipleChoice(config=snake_case_				)
													__lowerCAmelCase     =  tf.tile(tf.expand_dims(snake_case_    ,				1				)    ,				(1, self.num_choices, 1)				)
													__lowerCAmelCase     =  tf.tile(tf.expand_dims(snake_case_    ,				1				)    ,				(1, self.num_choices, 1)				)
													__lowerCAmelCase     =  tf.tile(tf.expand_dims(snake_case_    ,				1				)    ,				(1, self.num_choices, 1)				)
													__lowerCAmelCase     =  {
													    """input_ids""": multiple_choice_inputs_ids,
													    """attention_mask""": multiple_choice_input_mask,
													    """token_type_ids""": multiple_choice_token_type_ids,
													}
													__lowerCAmelCase     =  model(snake_case_				)
													self.parent.assertEqual(result.logits.shape    ,				(self.batch_size, self.num_choices)				)




							def    A__						( self    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_				)       ->  Any:
													__lowerCAmelCase     =  self.num_labels
													__lowerCAmelCase     =  TFRoFormerForTokenClassification(config=snake_case_				)
													__lowerCAmelCase     =  {
													    """input_ids""": input_ids,
													    """attention_mask""": input_mask,
													    """token_type_ids""": token_type_ids,
													}
													__lowerCAmelCase     =  model(snake_case_				)
													self.parent.assertEqual(result.logits.shape    ,				(self.batch_size, self.seq_length, self.num_labels)				)




							def    A__						( self    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_				)       ->  Optional[int]:
													__lowerCAmelCase     =  TFRoFormerForQuestionAnswering(config=snake_case_				)
													__lowerCAmelCase     =  {
													    """input_ids""": input_ids,
													    """attention_mask""": input_mask,
													    """token_type_ids""": token_type_ids,
													}

													__lowerCAmelCase     =  model(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    A__						( self				)       ->  List[str]:
													__lowerCAmelCase     =  self.prepare_config_and_inputs()
													(
													    (
													    __lowerCAmelCase
													)   ,     (
													    __lowerCAmelCase
													)   ,     (
													    __lowerCAmelCase
													)   ,     (
													    __lowerCAmelCase
													)   ,     (
													    __lowerCAmelCase
													)   ,     (
													    __lowerCAmelCase
													)   ,     (
													    __lowerCAmelCase
													)   ,     
													)     =  config_and_inputs
													__lowerCAmelCase     =  {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
													return config, inputs_dict
@require_tf
class       lowerCAmelCase_		(				A__				,				A__				,				unittest.TestCase ):



							'''simple docstring'''





							_snake_case						      =		(
							    (
							        TFRoFormerModel,
							        TFRoFormerForCausalLM,
							        TFRoFormerForMaskedLM,
							        TFRoFormerForQuestionAnswering,
							        TFRoFormerForSequenceClassification,
							        TFRoFormerForTokenClassification,
							        TFRoFormerForMultipleChoice,
							    )
							    if is_tf_available()
							    else ()
							)
							_snake_case						      =		(
							    {
							        '''feature-extraction''': TFRoFormerModel,
							        '''fill-mask''': TFRoFormerForMaskedLM,
							        '''question-answering''': TFRoFormerForQuestionAnswering,
							        '''text-classification''': TFRoFormerForSequenceClassification,
							        '''text-generation''': TFRoFormerForCausalLM,
							        '''token-classification''': TFRoFormerForTokenClassification,
							        '''zero-shot''': TFRoFormerForSequenceClassification,
							    }
							    if is_tf_available()
							    else {}
							)

							_snake_case						      =		False
							_snake_case						      =		False




							def    A__						( self    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_				)       ->  Tuple:
													if pipeline_test_casse_name == "TextGenerationPipelineTests":
																			return True

													return False




							def    A__						( self				)       ->  Tuple:
													__lowerCAmelCase     =  TFRoFormerModelTester(self				)
													__lowerCAmelCase     =  ConfigTester(self    ,				config_class=snake_case_    ,				hidden_size=37				)




							def    A__						( self				)       ->  Dict:
													self.config_tester.run_common_tests()




							def    A__						( self				)       ->  List[Any]:
													__lowerCAmelCase     =  self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_model(*snake_case_				)




							def    A__						( self				)       ->  List[Any]:
													__lowerCAmelCase     =  self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_masked_lm(*snake_case_				)




							def    A__						( self				)       ->  Tuple:
													__lowerCAmelCase     =  self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_lm_head(*snake_case_				)




							def    A__						( self				)       ->  Tuple:
													__lowerCAmelCase     =  self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_multiple_choice(*snake_case_				)




							def    A__						( self				)       ->  Any:
													__lowerCAmelCase     =  self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_question_answering(*snake_case_				)




							def    A__						( self				)       ->  Tuple:
													__lowerCAmelCase     =  self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_sequence_classification(*snake_case_				)




							def    A__						( self				)       ->  Optional[int]:
													__lowerCAmelCase     =  self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_token_classification(*snake_case_				)




							@slow
							def    A__						( self				)       ->  int:
													__lowerCAmelCase     =  TFRoFormerModel.from_pretrained("""junnyu/roformer_chinese_base"""				)
													self.assertIsNotNone(snake_case_				)
@require_tf
class       lowerCAmelCase_		(				unittest.TestCase ):



							'''simple docstring'''





							@slow
							def    A__						( self				)       ->  Any:
													__lowerCAmelCase     =  TFRoFormerForMaskedLM.from_pretrained("""junnyu/roformer_chinese_base"""				)
													__lowerCAmelCase     =  tf.constant([[0, 1, 2, 3, 4, 5]]				)
													__lowerCAmelCase     =  model(snake_case_				)[0]

													# TODO Replace vocab size
													__lowerCAmelCase     =  50_000

													__lowerCAmelCase     =  [1, 6, vocab_size]
													self.assertEqual(output.shape    ,				snake_case_				)

													print(output[:, :3, :3]				)

													# TODO Replace values below with what was printed above.
													__lowerCAmelCase     =  tf.constant(
													    [
													        [
													            [-0.12_053_341, -1.0_264_901, 0.29_221_946],
													            [-1.5_133_783, 0.197_433, 0.15_190_607],
													            [-5.0_135_403, -3.900_256, -0.84_038_764],
													        ]
													    ]				)
													tf.debugging.assert_near(output[:, :3, :3]    ,				snake_case_    ,				atol=1e-4				)
@require_tf
class       lowerCAmelCase_		(				unittest.TestCase ):



							'''simple docstring'''





							_snake_case						      =		1e-4




							def    A__						( self				)       ->  Optional[Any]:
													__lowerCAmelCase     =  tf.constant([[4, 10]]				)
													__lowerCAmelCase     =  TFRoFormerSinusoidalPositionalEmbedding(num_positions=6    ,				embedding_dim=6				)

													__lowerCAmelCase     =  emba(input_ids.shape				)
													__lowerCAmelCase     =  tf.constant(
													    [[0.0_000, 0.0_000, 0.0_000, 1.0_000, 1.0_000, 1.0_000], [0.8_415, 0.0_464, 0.0_022, 0.5_403, 0.9_989, 1.0_000]]				)

													tf.debugging.assert_near(snake_case_    ,				snake_case_    ,				atol=self.tolerance				)




							def    A__						( self				)       ->  str:
													__lowerCAmelCase     =  tf.constant(
													    [
													        [0.0_000, 0.0_000, 0.0_000, 0.0_000, 0.0_000],
													        [0.8_415, 0.8_219, 0.8_020, 0.7_819, 0.7_617],
													        [0.9_093, 0.9_364, 0.9_581, 0.9_749, 0.9_870],
													    ]				)
													__lowerCAmelCase     =  TFRoFormerSinusoidalPositionalEmbedding(num_positions=512    ,				embedding_dim=512				)
													emba([2, 16, 512]				)
													__lowerCAmelCase     =  emba.weight[:3, :5]

													tf.debugging.assert_near(snake_case_    ,				snake_case_    ,				atol=self.tolerance				)
@require_tf
class       lowerCAmelCase_		(				unittest.TestCase ):



							'''simple docstring'''





							_snake_case						      =		1e-4




							def    A__						( self				)       ->  Union[str, Any]:
													# 2,12,16,64
													__lowerCAmelCase     =  tf.reshape(tf.range(2 * 12 * 16 * 64    ,				dtype=tf.floataa				)    ,				shape=(2, 12, 16, 64)				) / 100

													__lowerCAmelCase     =  -tf.reshape(tf.range(2 * 12 * 16 * 64    ,				dtype=tf.floataa				)    ,				shape=(2, 12, 16, 64)				) / 100

													__lowerCAmelCase     =  TFRoFormerSinusoidalPositionalEmbedding(num_positions=32    ,				embedding_dim=64				)
													__lowerCAmelCase     =  embed_positions([2, 16, 768]				)[None, None, :, :]

													__lowerCAmelCase   ,     __lowerCAmelCase     =  TFRoFormerSelfAttention.apply_rotary_position_embeddings(
													    snake_case_    ,				snake_case_    ,				snake_case_				)

													__lowerCAmelCase     =  tf.constant(
													    [
													        [0.0_000, 0.0_100, 0.0_200, 0.0_300, 0.0_400, 0.0_500, 0.0_600, 0.0_700],
													        [-0.2_012, 0.8_897, 0.0_263, 0.9_401, 0.2_074, 0.9_463, 0.3_481, 0.9_343],
													        [-1.7_057, 0.6_271, -1.2_145, 1.3_897, -0.6_303, 1.7_647, -0.1_173, 1.8_985],
													        [-2.1_731, -1.6_397, -2.7_358, 0.2_854, -2.1_840, 1.7_183, -1.3_018, 2.4_871],
													        [0.2_717, -3.6_173, -2.9_206, -2.1_988, -3.6_638, 0.3_858, -2.9_155, 2.2_980],
													        [3.9_859, -2.1_580, -0.7_984, -4.4_904, -4.1_181, -2.0_252, -4.4_782, 1.1_253],
													    ]				)
													__lowerCAmelCase     =  tf.constant(
													    [
													        [0.0_000, -0.0_100, -0.0_200, -0.0_300, -0.0_400, -0.0_500, -0.0_600, -0.0_700],
													        [0.2_012, -0.8_897, -0.0_263, -0.9_401, -0.2_074, -0.9_463, -0.3_481, -0.9_343],
													        [1.7_057, -0.6_271, 1.2_145, -1.3_897, 0.6_303, -1.7_647, 0.1_173, -1.8_985],
													        [2.1_731, 1.6_397, 2.7_358, -0.2_854, 2.1_840, -1.7_183, 1.3_018, -2.4_871],
													        [-0.2_717, 3.6_173, 2.9_206, 2.1_988, 3.6_638, -0.3_858, 2.9_155, -2.2_980],
													        [-3.9_859, 2.1_580, 0.7_984, 4.4_904, 4.1_181, 2.0_252, 4.4_782, -1.1_253],
													    ]				)

													tf.debugging.assert_near(query_layer[0, 0, :6, :8]    ,				snake_case_    ,				atol=self.tolerance				)
													tf.debugging.assert_near(key_layer[0, 0, :6, :8]    ,				snake_case_    ,				atol=self.tolerance				)


"""simple docstring"""
import torch
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.bert.modeling_bert import (
    BERT_INPUTS_DOCSTRING,
    BERT_START_DOCSTRING,
    BertEmbeddings,
    BertLayer,
    BertPooler,
    BertPreTrainedModel,
)

def  lowercase     (_lowerCAmelCase						):
						__lowerCAmelCase     =  torch.exp(_lowerCAmelCase						)
						__lowerCAmelCase     =  torch.sum(_lowerCAmelCase ,    dim=1						)  # sum of exp(x_i)
						__lowerCAmelCase     =  torch.sum(x * exp_x ,    dim=1						)  # sum of x_i * exp(x_i)
						return torch.log(_lowerCAmelCase						) - B / A
class       lowerCAmelCase_		(				nn.Module ):



							'''simple docstring'''





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

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




							def    A__						( self    ,				snake_case_				)       ->  List[Any]:
													if (type(snake_case_				) is float) or (type(snake_case_				) is int):
																			for i in range(len(self.early_exit_entropy				)				):
																									__lowerCAmelCase     =  x
													else:
																			__lowerCAmelCase     =  x




							def    A__						( self    ,				snake_case_				)       ->  Dict:
													__lowerCAmelCase     =  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    ,				snake_case_    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				)       ->  List[str]:
													__lowerCAmelCase     =  ()
													__lowerCAmelCase     =  ()
													__lowerCAmelCase     =  ()
													for i, layer_module in enumerate(self.layer				):
																			if self.output_hidden_states:
																									__lowerCAmelCase     =  all_hidden_states + (hidden_states,)

																			__lowerCAmelCase     =  layer_module(
																			    snake_case_    ,				snake_case_    ,				head_mask[i]    ,				snake_case_    ,				snake_case_				)
																			__lowerCAmelCase     =  layer_outputs[0]

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

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

																			__lowerCAmelCase     =  self.highway[i](snake_case_				)
																			# logits, pooled_output

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

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

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

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

													__lowerCAmelCase     =  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). '''				,				A__				,				)
class       lowerCAmelCase_		(				A__ ):



							'''simple docstring'''





							def __init__( self    ,				snake_case_				)       ->  Tuple:
													super().__init__(snake_case_				)
													__lowerCAmelCase     =  config

													__lowerCAmelCase     =  BertEmbeddings(snake_case_				)
													__lowerCAmelCase     =  DeeBertEncoder(snake_case_				)
													__lowerCAmelCase     =  BertPooler(snake_case_				)

													self.init_weights()




							def    A__						( self				)       ->  Dict:
													self.encoder.init_highway_pooler(self.pooler				)




							def    A__						( self				)       ->  List[Any]:
													return self.embeddings.word_embeddings




							def    A__						( self    ,				snake_case_				)       ->  Tuple:
													__lowerCAmelCase     =  value




							def    A__						( self    ,				snake_case_				)       ->  List[str]:
													for layer, heads in heads_to_prune.items():
																			self.encoder.layer[layer].attention.prune_heads(snake_case_				)




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

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

													if attention_mask is None:
																			__lowerCAmelCase     =  torch.ones(snake_case_    ,				device=snake_case_				)
													if encoder_attention_mask is None:
																			__lowerCAmelCase     =  torch.ones(snake_case_    ,				device=snake_case_				)
													if token_type_ids is None:
																			__lowerCAmelCase     =  torch.zeros(snake_case_    ,				dtype=torch.long    ,				device=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.
													__lowerCAmelCase     =  self.get_extended_attention_mask(snake_case_    ,				snake_case_    ,				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:
																			__lowerCAmelCase     =  encoder_attention_mask[:, None, :, :]
													if encoder_attention_mask.dim() == 2:
																			__lowerCAmelCase     =  encoder_attention_mask[:, None, None, :]

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

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

													__lowerCAmelCase     =  (
													    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_		(				A__ ):



							'''simple docstring'''





							def __init__( self    ,				snake_case_    ,				snake_case_				)       ->  Optional[int]:
													__lowerCAmelCase     =  message
													__lowerCAmelCase     =  exit_layer  # start from 1!
class       lowerCAmelCase_		(				nn.Module ):



							'''simple docstring'''





							def __init__( self    ,				snake_case_				)       ->  List[Any]:
													super().__init__()
													__lowerCAmelCase     =  BertPooler(snake_case_				)
													__lowerCAmelCase     =  nn.Dropout(config.hidden_dropout_prob				)
													__lowerCAmelCase     =  nn.Linear(config.hidden_size    ,				config.num_labels				)




							def    A__						( self    ,				snake_case_				)       ->  Optional[int]:
													# Pooler
													__lowerCAmelCase     =  encoder_outputs[0]
													__lowerCAmelCase     =  self.pooler(snake_case_				)
													# "return" pooler_output

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

													# Dropout and classification
													__lowerCAmelCase     =  bmodel_output[1]

													__lowerCAmelCase     =  self.dropout(snake_case_				)
													__lowerCAmelCase     =  self.classifier(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. '''				,				A__				,				)
class       lowerCAmelCase_		(				A__ ):



							'''simple docstring'''





							def __init__( self    ,				snake_case_				)       ->  Optional[Any]:
													super().__init__(snake_case_				)
													__lowerCAmelCase     =  config.num_labels
													__lowerCAmelCase     =  config.num_hidden_layers

													__lowerCAmelCase     =  DeeBertModel(snake_case_				)
													__lowerCAmelCase     =  nn.Dropout(config.hidden_dropout_prob				)
													__lowerCAmelCase     =  nn.Linear(config.hidden_size    ,				self.config.num_labels				)

													self.init_weights()




							@add_start_docstrings_to_model_forward(snake_case_				)
							def    A__						( self    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=-1    ,				snake_case_=False    ,				)       ->  Optional[int]:
													__lowerCAmelCase     =  self.num_layers
													try:
																			__lowerCAmelCase     =  self.bert(
																			    snake_case_    ,				attention_mask=snake_case_    ,				token_type_ids=snake_case_    ,				position_ids=snake_case_    ,				head_mask=snake_case_    ,				inputs_embeds=snake_case_    ,				)
																			# sequence_output, pooled_output, (hidden_states), (attentions), highway exits

																			__lowerCAmelCase     =  outputs[1]

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

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

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

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