


import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

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


if is_torch_available():
   import torch

   from transformers import (
       MODEL_FOR_PRETRAINING_MAPPING,
       MegatronBertForCausalLM,
       MegatronBertForMaskedLM,
       MegatronBertForMultipleChoice,
       MegatronBertForNextSentencePrediction,
       MegatronBertForPreTraining,
       MegatronBertForQuestionAnswering,
       MegatronBertForSequenceClassification,
       MegatronBertForTokenClassification,
       MegatronBertModel,
   )

class 			UpperCamelCase__			:






      def __init__(  self       : List[Any],		__lowerCamelCase       : str,		__lowerCamelCase       : Any=13,		__lowerCamelCase       : List[Any]=7,		__lowerCamelCase       : str=True,		__lowerCamelCase       : Optional[int]=True,		__lowerCamelCase       : Any=True,		__lowerCamelCase       : int=True,		__lowerCamelCase       : Optional[Any]=99,		__lowerCamelCase       : List[Any]=64,		__lowerCamelCase       : Tuple=32,		__lowerCamelCase       : str=5,		__lowerCamelCase       : Optional[Any]=4,		__lowerCamelCase       : Optional[Any]=37,		__lowerCamelCase       : Optional[int]="gelu",		__lowerCamelCase       : Tuple=0.1,		__lowerCamelCase       : int=0.1,		__lowerCamelCase       : Optional[int]=5_12,		__lowerCamelCase       : Tuple=16,		__lowerCamelCase       : List[Any]=2,		__lowerCamelCase       : Dict=0.02,		__lowerCamelCase       : Any=3,		__lowerCamelCase       : Optional[Any]=4,		__lowerCamelCase       : str=None,		)	->			List[str]:
           UpperCamelCase__				:	Any										=  parent
           UpperCamelCase__				:	Tuple										=  batch_size
           UpperCamelCase__				:	Dict										=  seq_length
           UpperCamelCase__				:	List[Any]										=  is_training
           UpperCamelCase__				:	str										=  use_input_mask
           UpperCamelCase__				:	List[Any]										=  use_token_type_ids
           UpperCamelCase__				:	List[Any]										=  use_labels
           UpperCamelCase__				:	Union[str, Any]										=  vocab_size
           UpperCamelCase__				:	Union[str, Any]										=  hidden_size
           UpperCamelCase__				:	Optional[Any]										=  embedding_size
           UpperCamelCase__				:	Dict										=  num_hidden_layers
           UpperCamelCase__				:	Optional[int]										=  num_attention_heads
           UpperCamelCase__				:	List[str]										=  intermediate_size
           UpperCamelCase__				:	Optional[Any]										=  hidden_act
           UpperCamelCase__				:	Any										=  hidden_dropout_prob
           UpperCamelCase__				:	str										=  attention_probs_dropout_prob
           UpperCamelCase__				:	Union[str, Any]										=  max_position_embeddings
           UpperCamelCase__				:	int										=  type_vocab_size
           UpperCamelCase__				:	Any										=  type_sequence_label_size
           UpperCamelCase__				:	str										=  initializer_range
           UpperCamelCase__				:	int										=  num_labels
           UpperCamelCase__				:	Union[str, Any]										=  num_choices
           UpperCamelCase__				:	int										=  scope






      def 						__lowercase(  self       : Dict )	->			Any:
           UpperCamelCase__				:	str										=  ids_tensor([self.batch_size, self.seq_length],		self.vocab_size )

           UpperCamelCase__				:	Optional[int]										=  None
           if self.use_input_mask:
                UpperCamelCase__				:	Optional[Any]										=  random_attention_mask([self.batch_size, self.seq_length] )

           UpperCamelCase__				:	str										=  None
           if self.use_token_type_ids:
                UpperCamelCase__				:	Optional[Any]										=  ids_tensor([self.batch_size, self.seq_length],		self.type_vocab_size )

           UpperCamelCase__				:	List[str]										=  None
           UpperCamelCase__				:	List[str]										=  None
           UpperCamelCase__				:	Dict										=  None
           if self.use_labels:
                UpperCamelCase__				:	Tuple										=  ids_tensor([self.batch_size],		self.type_sequence_label_size )
                UpperCamelCase__				:	str										=  ids_tensor([self.batch_size, self.seq_length],		self.num_labels )
                UpperCamelCase__				:	str										=  ids_tensor([self.batch_size],		self.num_choices )

           UpperCamelCase__				:	List[str]										=  self.get_config()

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






      def 						__lowercase(  self       : int )	->			Tuple:
           return MegatronBertConfig(
               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,		embedding_size=self.embedding_size,		hidden_act=self.hidden_act,		hidden_dropout_prob=self.hidden_dropout_prob,		attention_probs_dropout_prob=self.attention_probs_dropout_prob,		max_position_embeddings=self.max_position_embeddings,		type_vocab_size=self.type_vocab_size,		is_decoder=lowerCamelCase__,		initializer_range=self.initializer_range,		)






      def 						__lowercase(  self       : Dict,		__lowerCamelCase       : Any,		__lowerCamelCase       : Any,		__lowerCamelCase       : Tuple,		__lowerCamelCase       : str,		__lowerCamelCase       : int,		__lowerCamelCase       : Tuple,		__lowerCamelCase       : int )	->			Tuple:
           UpperCamelCase__				:	Tuple										=  MegatronBertModel(config=lowerCamelCase__ )
           model.to(lowerCamelCase__ )
           model.eval()
           UpperCamelCase__				:	str										=  model(lowerCamelCase__,		attention_mask=lowerCamelCase__,		token_type_ids=lowerCamelCase__ )
           UpperCamelCase__				:	str										=  model(lowerCamelCase__,		token_type_ids=lowerCamelCase__ )
           UpperCamelCase__				:	List[Any]										=  model(lowerCamelCase__ )

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






      def 						__lowercase(  self       : Tuple,		__lowerCamelCase       : str,		__lowerCamelCase       : Tuple,		__lowerCamelCase       : str,		__lowerCamelCase       : List[Any],		__lowerCamelCase       : List[Any],		__lowerCamelCase       : Union[str, Any],		__lowerCamelCase       : str )	->			Union[str, Any]:
           UpperCamelCase__				:	Optional[int]										=  MegatronBertForMaskedLM(config=lowerCamelCase__ )
           model.to(lowerCamelCase__ )
           model.eval()
           UpperCamelCase__				:	Union[str, Any]										=  model(lowerCamelCase__,		attention_mask=lowerCamelCase__,		token_type_ids=lowerCamelCase__,		labels=lowerCamelCase__ )
           self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.seq_length, self.vocab_size) )






      def 						__lowercase(  self       : int,		__lowerCamelCase       : Dict,		__lowerCamelCase       : Optional[int],		__lowerCamelCase       : Dict,		__lowerCamelCase       : Union[str, Any],		__lowerCamelCase       : str,		__lowerCamelCase       : Tuple,		__lowerCamelCase       : Optional[Any] )	->			Union[str, Any]:
           UpperCamelCase__				:	List[Any]										=  MegatronBertForCausalLM(config=lowerCamelCase__ )
           model.to(lowerCamelCase__ )
           model.eval()
           UpperCamelCase__				:	str										=  model(lowerCamelCase__,		attention_mask=lowerCamelCase__,		token_type_ids=lowerCamelCase__,		labels=lowerCamelCase__ )
           self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.seq_length, self.vocab_size) )






      def 						__lowercase(  self       : Tuple,		__lowerCamelCase       : Any,		__lowerCamelCase       : Tuple,		__lowerCamelCase       : List[Any],		__lowerCamelCase       : Optional[Any],		__lowerCamelCase       : Tuple,		__lowerCamelCase       : int,		__lowerCamelCase       : Tuple )	->			str:
           UpperCamelCase__				:	Dict										=  MegatronBertForNextSentencePrediction(config=lowerCamelCase__ )
           model.to(lowerCamelCase__ )
           model.eval()
           UpperCamelCase__				:	str										=  model(
               lowerCamelCase__,		attention_mask=lowerCamelCase__,		token_type_ids=lowerCamelCase__,		labels=lowerCamelCase__,		)
           self.parent.assertEqual(result.logits.shape,		(self.batch_size, 2) )






      def 						__lowercase(  self       : Any,		__lowerCamelCase       : Union[str, Any],		__lowerCamelCase       : Optional[int],		__lowerCamelCase       : int,		__lowerCamelCase       : Dict,		__lowerCamelCase       : int,		__lowerCamelCase       : Dict,		__lowerCamelCase       : Tuple )	->			str:
           UpperCamelCase__				:	Dict										=  MegatronBertForPreTraining(config=lowerCamelCase__ )
           model.to(lowerCamelCase__ )
           model.eval()
           UpperCamelCase__				:	Optional[int]										=  model(
               lowerCamelCase__,		attention_mask=lowerCamelCase__,		token_type_ids=lowerCamelCase__,		labels=lowerCamelCase__,		next_sentence_label=lowerCamelCase__,		)
           self.parent.assertEqual(result.prediction_logits.shape,		(self.batch_size, self.seq_length, self.vocab_size) )
           self.parent.assertEqual(result.seq_relationship_logits.shape,		(self.batch_size, 2) )






      def 						__lowercase(  self       : Tuple,		__lowerCamelCase       : Any,		__lowerCamelCase       : str,		__lowerCamelCase       : str,		__lowerCamelCase       : Any,		__lowerCamelCase       : Any,		__lowerCamelCase       : str,		__lowerCamelCase       : Optional[int] )	->			Dict:
           UpperCamelCase__				:	str										=  MegatronBertForQuestionAnswering(config=lowerCamelCase__ )
           model.to(lowerCamelCase__ )
           model.eval()
           UpperCamelCase__				:	Optional[int]										=  model(
               lowerCamelCase__,		attention_mask=lowerCamelCase__,		token_type_ids=lowerCamelCase__,		start_positions=lowerCamelCase__,		end_positions=lowerCamelCase__,		)
           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 						__lowercase(  self       : Any,		__lowerCamelCase       : List[Any],		__lowerCamelCase       : Optional[int],		__lowerCamelCase       : Any,		__lowerCamelCase       : List[str],		__lowerCamelCase       : List[Any],		__lowerCamelCase       : Optional[Any],		__lowerCamelCase       : str )	->			int:
           UpperCamelCase__				:	List[Any]										=  self.num_labels
           UpperCamelCase__				:	Optional[int]										=  MegatronBertForSequenceClassification(lowerCamelCase__ )
           model.to(lowerCamelCase__ )
           model.eval()
           UpperCamelCase__				:	Dict										=  model(lowerCamelCase__,		attention_mask=lowerCamelCase__,		token_type_ids=lowerCamelCase__,		labels=lowerCamelCase__ )
           self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.num_labels) )






      def 						__lowercase(  self       : Union[str, Any],		__lowerCamelCase       : int,		__lowerCamelCase       : Union[str, Any],		__lowerCamelCase       : Optional[int],		__lowerCamelCase       : List[Any],		__lowerCamelCase       : List[str],		__lowerCamelCase       : Dict,		__lowerCamelCase       : Union[str, Any] )	->			Dict:
           UpperCamelCase__				:	Any										=  self.num_labels
           UpperCamelCase__				:	Optional[int]										=  MegatronBertForTokenClassification(config=lowerCamelCase__ )
           model.to(lowerCamelCase__ )
           model.eval()
           UpperCamelCase__				:	Optional[int]										=  model(lowerCamelCase__,		attention_mask=lowerCamelCase__,		token_type_ids=lowerCamelCase__,		labels=lowerCamelCase__ )
           self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.seq_length, self.num_labels) )






      def 						__lowercase(  self       : List[Any],		__lowerCamelCase       : Union[str, Any],		__lowerCamelCase       : List[Any],		__lowerCamelCase       : str,		__lowerCamelCase       : Union[str, Any],		__lowerCamelCase       : Tuple,		__lowerCamelCase       : Union[str, Any],		__lowerCamelCase       : List[str] )	->			Optional[Any]:
           UpperCamelCase__				:	Optional[int]										=  self.num_choices
           UpperCamelCase__				:	List[str]										=  MegatronBertForMultipleChoice(config=lowerCamelCase__ )
           model.to(lowerCamelCase__ )
           model.eval()
           UpperCamelCase__				:	Any										=  input_ids.unsqueeze(1 ).expand(-1,		self.num_choices,		-1 ).contiguous()
           UpperCamelCase__				:	List[Any]										=  token_type_ids.unsqueeze(1 ).expand(-1,		self.num_choices,		-1 ).contiguous()
           UpperCamelCase__				:	int										=  input_mask.unsqueeze(1 ).expand(-1,		self.num_choices,		-1 ).contiguous()
           UpperCamelCase__				:	Dict										=  model(
               lowerCamelCase__,		attention_mask=lowerCamelCase__,		token_type_ids=lowerCamelCase__,		labels=lowerCamelCase__,		)
           self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.num_choices) )






      def 						__lowercase(  self       : List[Any] )	->			List[str]:
           UpperCamelCase__				:	List[str]										=  self.prepare_config_and_inputs()
           (
               UpperCamelCase__     
           )				:	str										=  config_and_inputs
           UpperCamelCase__				:	List[str]										=  {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
           return config, inputs_dict

@require_torch
class 			UpperCamelCase__			(			UpperCAmelCase_				,    UpperCAmelCase_				,    unittest.TestCase      ):
      a__ :  Union[str, Any]		    =    (
          (
              MegatronBertModel,
              MegatronBertForMaskedLM,
              MegatronBertForCausalLM,
              MegatronBertForMultipleChoice,
              MegatronBertForNextSentencePrediction,
              MegatronBertForPreTraining,
              MegatronBertForQuestionAnswering,
              MegatronBertForSequenceClassification,
              MegatronBertForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      a__ :  Union[str, Any]		    =    (
          {
              '''feature-extraction''': MegatronBertModel,
              '''fill-mask''': MegatronBertForMaskedLM,
              '''question-answering''': MegatronBertForQuestionAnswering,
              '''text-classification''': MegatronBertForSequenceClassification,
              '''text-generation''': MegatronBertForCausalLM,
              '''token-classification''': MegatronBertForTokenClassification,
              '''zero-shot''': MegatronBertForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      a__ :  str		    =    True
      # test_resize_embeddings = False
      a__ :  Optional[Any]		    =    False






      def 						__lowercase(  self       : List[str],		__lowerCamelCase       : Optional[int],		__lowerCamelCase       : List[str],		__lowerCamelCase       : List[Any]=False )	->			Union[str, Any]:
           UpperCamelCase__				:	List[str]										=  super()._prepare_for_class(lowerCamelCase__,		lowerCamelCase__,		return_labels=lowerCamelCase__ )

           if return_labels:
                if model_class in get_values(lowerCamelCase__ ):
                     UpperCamelCase__				:	int										=  torch.zeros(
                         (self.model_tester.batch_size, self.model_tester.seq_length),		dtype=torch.long,		device=lowerCamelCase__ )
                     UpperCamelCase__				:	Union[str, Any]										=  torch.zeros(
                         self.model_tester.batch_size,		dtype=torch.long,		device=lowerCamelCase__ )
           return inputs_dict






      def 						__lowercase(  self       : str )	->			Any:
           UpperCamelCase__				:	Optional[Any]										=  MegatronBertModelTester(self )
           UpperCamelCase__				:	Tuple										=  ConfigTester(self,		config_class=lowerCamelCase__,		hidden_size=37 )






      def 						__lowercase(  self       : int )	->			Union[str, Any]:
           self.config_tester.run_common_tests()






      def 						__lowercase(  self       : Optional[int] )	->			Tuple:
           UpperCamelCase__				:	Any										=  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_megatron_bert_model(*lowerCamelCase__ )






      def 						__lowercase(  self       : Optional[int] )	->			List[str]:
           UpperCamelCase__				:	Union[str, Any]										=  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_megatron_bert_for_masked_lm(*lowerCamelCase__ )






      def 						__lowercase(  self       : str )	->			Optional[Any]:
           UpperCamelCase__				:	Optional[int]										=  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*lowerCamelCase__ )






      def 						__lowercase(  self       : Any )	->			Dict:
           UpperCamelCase__				:	Union[str, Any]										=  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*lowerCamelCase__ )






      def 						__lowercase(  self       : Tuple )	->			Tuple:
           UpperCamelCase__				:	List[Any]										=  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_megatron_bert_for_pretraining(*lowerCamelCase__ )






      def 						__lowercase(  self       : Union[str, Any] )	->			Dict:
           UpperCamelCase__				:	List[Any]										=  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_megatron_bert_for_question_answering(*lowerCamelCase__ )






      def 						__lowercase(  self       : int )	->			Tuple:
           UpperCamelCase__				:	List[Any]										=  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*lowerCamelCase__ )






      def 						__lowercase(  self       : Optional[Any] )	->			Optional[Any]:
           UpperCamelCase__				:	int										=  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_megatron_bert_for_token_classification(*lowerCamelCase__ )







def     _lowercase							(  __lowerCamelCase      :       str      )	-> Any:





     '''simple docstring'''





     return torch.tensor(
         __lowerCAmelCase		,dtype=torch.long		,device=__lowerCAmelCase		,)


_SCREAMING_SNAKE_CASE			:       Dict      =						1e-4







@require_torch
@require_sentencepiece
@require_tokenizers
class 			UpperCamelCase__			(			unittest.TestCase      ):






      @slow
      @unittest.skip('''Model is not available.''' )
      def 						__lowercase(  self       : Optional[Any] )	->			str:
           UpperCamelCase__				:	Any										=  """nvidia/megatron-bert-uncased-345m"""
           if "MYDIR" in os.environ:
                UpperCamelCase__				:	List[str]										=  os.path.join(os.environ['''MYDIR'''],		lowerCamelCase__ )
           UpperCamelCase__				:	int										=  MegatronBertModel.from_pretrained(lowerCamelCase__ )
           model.to(lowerCamelCase__ )
           model.half()
           UpperCamelCase__				:	int										=  _long_tensor([[1_01, 71_10, 10_05, 10_56, 20_23, 1_13_33, 1_74_13, 10_29, 1_02]] )
           with torch.no_grad():
                UpperCamelCase__				:	Tuple										=  model(lowerCamelCase__ )[0]
           UpperCamelCase__				:	Tuple										=  torch.Size((1, 9, 10_24) )
           self.assertEqual(output.shape,		lowerCamelCase__ )

           UpperCamelCase__				:	Any										=  [-0.6040, -0.2517, -0.1025, 0.3420, -0.6758, -0.0017, -0.1089, -0.1990, 0.5728]
           for ii in range(3 ):
                for jj in range(3 ):
                     UpperCamelCase__				:	Any										=  output[0, ii, jj]
                     UpperCamelCase__				:	Dict										=  expected[3 * ii + jj]
                     UpperCamelCase__				:	str										=  """ii={} jj={} a={} b={}""".format(lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__ )
                     self.assertTrue(math.isclose(lowerCamelCase__,		lowerCamelCase__,		rel_tol=lowerCamelCase__,		abs_tol=lowerCamelCase__ ),		msg=lowerCamelCase__ )












import logging

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,
    BertEncoder,
    BertModel,
    BertPreTrainedModel,
)


UpperCamelCase    							=      logging.getLogger(__name__)







class 		_A						( UpperCAmelCase_ ):

  def    a					(    self						:	Tuple   ,				lowerCamelCase__						:	Optional[int]   ,				lowerCamelCase__						:	Union[str, Any]   ,				lowerCamelCase__						:	str=None   ,				lowerCamelCase__						:	int=None			):
     """simple docstring"""

     __UpperCamelCase	:			Optional[Any]					=					self.layer[current_layer](lowerCamelCase__   ,				lowerCamelCase__   ,				head_mask[current_layer]			)

     __UpperCamelCase	:			Optional[Any]					=					layer_outputs[0]

     return hidden_states







@add_start_docstrings(
    '''The bare Bert Model transformer with PABEE outputting raw hidden-states without any specific head on top.'''	,				UpperCAmelCase_	,				)
class 		_A						( UpperCAmelCase_ ):

  def __init__(    self						:	Tuple   ,				lowerCamelCase__						:	Tuple			):
     """simple docstring"""

     super().__init__(lowerCamelCase__			)

     __UpperCamelCase	:			List[str]					=					BertEncoderWithPabee(lowerCamelCase__			)

     self.init_weights()
     __UpperCamelCase	:			Optional[int]					=					0
     __UpperCamelCase	:			Tuple					=					0
     __UpperCamelCase	:			Any					=					0

     __UpperCamelCase	:			Optional[Any]					=					0

  def    a					(    self						:	Union[str, Any]   ,				lowerCamelCase__						:	Any			):
     """simple docstring"""

     __UpperCamelCase	:			List[Any]					=					threshold

  def    a					(    self						:	Optional[int]   ,				lowerCamelCase__						:	Union[str, Any]			):
     """simple docstring"""

     __UpperCamelCase	:			int					=					patience

  def    a					(    self						:	Optional[int]			):
     """simple docstring"""

     __UpperCamelCase	:			int					=					0
     __UpperCamelCase	:			Tuple					=					0

  def    a					(    self						:	int			):
     """simple docstring"""

     __UpperCamelCase	:			int					=					self.inference_layers_num / self.inference_instances_num
     __UpperCamelCase	:			str					=					(
         f'*** Patience = {self.patience} Avg. Inference Layers = {avg_inf_layers:.2f} Speed Up ='
         f' {1 - avg_inf_layers / self.config.num_hidden_layers:.2f} ***'
     )
     print(lowerCamelCase__			)

  @add_start_docstrings_to_model_forward(lowerCamelCase__			)
  def    a					(    self						:	Any   ,				lowerCamelCase__						:	List[Any]=None   ,				lowerCamelCase__						:	List[str]=None   ,				lowerCamelCase__						:	List[Any]=None   ,				lowerCamelCase__						:	Optional[Any]=None   ,				lowerCamelCase__						:	List[str]=None   ,				lowerCamelCase__						:	str=None   ,				lowerCamelCase__						:	List[Any]=None   ,				lowerCamelCase__						:	Any=None   ,				lowerCamelCase__						:	List[str]=None   ,				lowerCamelCase__						:	List[str]=None   ,				lowerCamelCase__						:	str=False   ,				):
     """simple docstring"""

     if input_ids is not None and inputs_embeds is not None:
        raise ValueError("""You cannot specify both input_ids and inputs_embeds at the same time"""			)
     elif input_ids is not None:
        __UpperCamelCase	:			Optional[Any]					=					input_ids.size()
     elif inputs_embeds is not None:
        __UpperCamelCase	:			Tuple					=					inputs_embeds.size()[:-1]
     else:
        raise ValueError("""You have to specify either input_ids or inputs_embeds"""			)

     __UpperCamelCase	:			Tuple					=					input_ids.device if input_ids is not None else inputs_embeds.device

     if attention_mask is None:
        __UpperCamelCase	:			Optional[Any]					=					torch.ones(lowerCamelCase__   ,				device=lowerCamelCase__			)
     if token_type_ids is None:
        __UpperCamelCase	:			str					=					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.
     __UpperCamelCase	:			torch.Tensor					=					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 self.config.is_decoder and encoder_hidden_states is not None:
        __UpperCamelCase	,     __UpperCamelCase	,     __UpperCamelCase	:			Optional[Any]					=					encoder_hidden_states.size()
        __UpperCamelCase	:			Union[str, Any]					=					(encoder_batch_size, encoder_sequence_length)
        if encoder_attention_mask is None:
           __UpperCamelCase	:			Union[str, Any]					=					torch.ones(lowerCamelCase__   ,				device=lowerCamelCase__			)
        __UpperCamelCase	:			Optional[int]					=					self.invert_attention_mask(lowerCamelCase__			)
     else:
        __UpperCamelCase	:			str					=					None

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

     __UpperCamelCase	:			Any					=					self.embeddings(
         input_ids=lowerCamelCase__   ,				position_ids=lowerCamelCase__   ,				token_type_ids=lowerCamelCase__   ,				inputs_embeds=lowerCamelCase__			)
     __UpperCamelCase	:			Optional[int]					=					embedding_output

     if self.training:
        __UpperCamelCase	:			List[Any]					=					[]
        for i in range(self.config.num_hidden_layers			):
           __UpperCamelCase	:			List[Any]					=					self.encoder.adaptive_forward(
               lowerCamelCase__   ,				current_layer=lowerCamelCase__   ,				attention_mask=lowerCamelCase__   ,				head_mask=lowerCamelCase__			)

           __UpperCamelCase	:			Tuple					=					self.pooler(lowerCamelCase__			)
           __UpperCamelCase	:			str					=					output_layers[i](output_dropout(lowerCamelCase__			)			)
           res.append(lowerCamelCase__			)
     elif self.patience == 0:  # Use all layers for inference
        __UpperCamelCase	:			Union[str, Any]					=					self.encoder(
            lowerCamelCase__   ,				attention_mask=lowerCamelCase__   ,				head_mask=lowerCamelCase__   ,				encoder_hidden_states=lowerCamelCase__   ,				encoder_attention_mask=lowerCamelCase__   ,				)
        __UpperCamelCase	:			Optional[int]					=					self.pooler(encoder_outputs[0]			)
        __UpperCamelCase	:			Union[str, Any]					=					[output_layers[self.config.num_hidden_layers - 1](lowerCamelCase__			)]
     else:
        __UpperCamelCase	:			Union[str, Any]					=					0
        __UpperCamelCase	:			Optional[Any]					=					None
        __UpperCamelCase	:			List[Any]					=					0
        for i in range(self.config.num_hidden_layers			):
           calculated_layer_num += 1
           __UpperCamelCase	:			Union[str, Any]					=					self.encoder.adaptive_forward(
               lowerCamelCase__   ,				current_layer=lowerCamelCase__   ,				attention_mask=lowerCamelCase__   ,				head_mask=lowerCamelCase__			)

           __UpperCamelCase	:			str					=					self.pooler(lowerCamelCase__			)
           __UpperCamelCase	:			Any					=					output_layers[i](lowerCamelCase__			)
           if regression:
              __UpperCamelCase	:			Optional[int]					=					logits.detach()
              if patient_result is not None:
                 __UpperCamelCase	:			Dict					=					patient_result.detach()
              if (patient_result is not None) and torch.abs(patient_result - labels			) < self.regression_threshold:
                 patient_counter += 1
              else:
                 __UpperCamelCase	:			List[str]					=					0
           else:
              __UpperCamelCase	:			Optional[int]					=					logits.detach().argmax(dim=1			)
              if patient_result is not None:
                 __UpperCamelCase	:			Union[str, Any]					=					patient_result.detach().argmax(dim=1			)
              if (patient_result is not None) and torch.all(labels.eq(lowerCamelCase__			)			):
                 patient_counter += 1
              else:
                 __UpperCamelCase	:			str					=					0

           __UpperCamelCase	:			Union[str, Any]					=					logits
           if patient_counter == self.patience:
              break
        __UpperCamelCase	:			List[str]					=					[patient_result]
        self.inference_layers_num += calculated_layer_num
        self.inference_instances_num += 1

     return res







@add_start_docstrings(
    '''Bert Model transformer with PABEE and a sequence classification/regression head on top (a linear layer on top of
    the pooled output) e.g. for GLUE tasks. '''	,				UpperCAmelCase_	,				)
class 		_A						( UpperCAmelCase_ ):

  def __init__(    self						:	Any   ,				lowerCamelCase__						:	Union[str, Any]			):
     """simple docstring"""

     super().__init__(lowerCamelCase__			)
     __UpperCamelCase	:			Dict					=					config.num_labels

     __UpperCamelCase	:			List[Any]					=					BertModelWithPabee(lowerCamelCase__			)
     __UpperCamelCase	:			Optional[int]					=					nn.Dropout(config.hidden_dropout_prob			)
     __UpperCamelCase	:			List[Any]					=					nn.ModuleList(
         [nn.Linear(config.hidden_size   ,				self.config.num_labels			) for _ in range(config.num_hidden_layers			)]			)

     self.init_weights()

  @add_start_docstrings_to_model_forward(lowerCamelCase__			)
  def    a					(    self						:	List[Any]   ,				lowerCamelCase__						:	Union[str, Any]=None   ,				lowerCamelCase__						:	Dict=None   ,				lowerCamelCase__						:	Union[str, Any]=None   ,				lowerCamelCase__						:	List[Any]=None   ,				lowerCamelCase__						:	Optional[Any]=None   ,				lowerCamelCase__						:	Optional[int]=None   ,				lowerCamelCase__						:	Dict=None   ,				):
     """simple docstring"""

     __UpperCamelCase	:			str					=					self.bert(
         input_ids=lowerCamelCase__   ,				attention_mask=lowerCamelCase__   ,				token_type_ids=lowerCamelCase__   ,				position_ids=lowerCamelCase__   ,				head_mask=lowerCamelCase__   ,				inputs_embeds=lowerCamelCase__   ,				output_dropout=self.dropout   ,				output_layers=self.classifiers   ,				regression=self.num_labels == 1   ,				)

     __UpperCamelCase	:			List[str]					=					(logits[-1],)

     if labels is not None:
        __UpperCamelCase	:			Tuple					=					None
        __UpperCamelCase	:			Union[str, Any]					=					0
        for ix, logits_item in enumerate(lowerCamelCase__			):
           if self.num_labels == 1:
              #  We are doing regression
              __UpperCamelCase	:			Union[str, Any]					=					MSELoss()
              __UpperCamelCase	:			Optional[Any]					=					loss_fct(logits_item.view(-1			)   ,				labels.view(-1			)			)
           else:
              __UpperCamelCase	:			Dict					=					CrossEntropyLoss()
              __UpperCamelCase	:			Dict					=					loss_fct(logits_item.view(-1   ,				self.num_labels			)   ,				labels.view(-1			)			)
           if total_loss is None:
              __UpperCamelCase	:			Dict					=					loss
           else:
              total_loss += loss * (ix + 1)
           total_weights += ix + 1
        __UpperCamelCase	:			Tuple					=					(total_loss / total_weights,) + outputs

     return outputs



