code: 
  
  
  '''simple docstring'''
  
  import unittest
  
  from transformers import is_torch_available
  from transformers.testing_utils import require_torch, slow, torch_device
  
  from ...test_configuration_common import ConfigTester
  from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
  from ...test_pipeline_mixin import PipelineTesterMixin
  
  
  if is_torch_available():
         import torch
  
         from transformers import (
             RemBertConfig,
             RemBertForCausalLM,
             RemBertForMaskedLM,
             RemBertForMultipleChoice,
             RemBertForQuestionAnswering,
             RemBertForSequenceClassification,
             RemBertForTokenClassification,
             RemBertModel,
         )
         from transformers.models.rembert.modeling_rembert import REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST
  class 						__lowerCAmelCase		:
  
  
  
  
  
        '''simple docstring'''
  
  
  
  
  
  
        def __init__(self					:					List[str]     ,		UpperCamelCase					:					str     ,		UpperCamelCase					:					List[Any]=13     ,		UpperCamelCase					:					Dict=7     ,		UpperCamelCase					:					Any=True     ,		UpperCamelCase					:					int=True     ,		UpperCamelCase					:					Optional[int]=True     ,		UpperCamelCase					:					Dict=True     ,		UpperCamelCase					:					str=99     ,		UpperCamelCase					:					List[str]=32     ,		UpperCamelCase					:					Optional[int]=18     ,		UpperCamelCase					:					Tuple=43     ,		UpperCamelCase					:					List[str]=5     ,		UpperCamelCase					:					Optional[Any]=4     ,		UpperCamelCase					:					List[Any]=37     ,		UpperCamelCase					:					Union[str, Any]="gelu"     ,		UpperCamelCase					:					Optional[int]=0.1     ,		UpperCamelCase					:					int=0.1     ,		UpperCamelCase					:					str=512     ,		UpperCamelCase					:					Dict=16     ,		UpperCamelCase					:					Union[str, Any]=2     ,		UpperCamelCase					:					Any=0.02     ,		UpperCamelCase					:					str=3     ,		UpperCamelCase					:					Union[str, Any]=4     ,		UpperCamelCase					:					List[Any]=None     ,		):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               lowercase__       =				parent
               lowercase__       =				batch_size
               lowercase__       =				seq_length
               lowercase__       =				is_training
               lowercase__       =				use_input_mask
               lowercase__       =				use_token_type_ids
               lowercase__       =				use_labels
               lowercase__       =				vocab_size
               lowercase__       =				hidden_size
               lowercase__       =				input_embedding_size
               lowercase__       =				output_embedding_size
               lowercase__       =				num_hidden_layers
               lowercase__       =				num_attention_heads
               lowercase__       =				intermediate_size
               lowercase__       =				hidden_act
               lowercase__       =				hidden_dropout_prob
               lowercase__       =				attention_probs_dropout_prob
               lowercase__       =				max_position_embeddings
               lowercase__       =				type_vocab_size
               lowercase__       =				type_sequence_label_size
               lowercase__       =				initializer_range
               lowercase__       =				num_labels
               lowercase__       =				num_choices
               lowercase__       =				scope
  
  
  
        def    UpperCamelCase__	(self					:					Dict				):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               lowercase__       =				ids_tensor([self.batch_size, self.seq_length]     ,		self.vocab_size				)
  
               lowercase__       =				None
               if self.use_input_mask:
                      lowercase__       =				random_attention_mask([self.batch_size, self.seq_length]				)
  
               lowercase__       =				None
               if self.use_token_type_ids:
                      lowercase__       =				ids_tensor([self.batch_size, self.seq_length]     ,		self.type_vocab_size				)
  
               lowercase__       =				None
               lowercase__       =				None
               lowercase__       =				None
               if self.use_labels:
                      lowercase__       =				ids_tensor([self.batch_size]     ,		self.type_sequence_label_size				)
                      lowercase__       =				ids_tensor([self.batch_size, self.seq_length]     ,		self.num_labels				)
                      lowercase__       =				ids_tensor([self.batch_size]     ,		self.num_choices				)
  
               lowercase__       =				RemBertConfig(
                   vocab_size=self.vocab_size     ,		hidden_size=self.hidden_size     ,		input_embedding_size=self.input_embedding_size     ,		output_embedding_size=self.output_embedding_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     ,		is_decoder=UpperCamelCase     ,		initializer_range=self.initializer_range     ,		)
  
               return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
  
  
  
        def    UpperCamelCase__	(self					:					Tuple				):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               (
                   (
                   lowercase__       
               )    ,(
                   lowercase__       
               )    ,(
                   lowercase__       
               )    ,(
                   lowercase__       
               )    ,(
                   lowercase__       
               )    ,(
                   lowercase__       
               )    ,(
                   lowercase__       
               )    ,       
               )       =				self.prepare_config_and_inputs()
  
               lowercase__       =				True
               lowercase__       =				floats_tensor([self.batch_size, self.seq_length, self.hidden_size]				)
               lowercase__       =				ids_tensor([self.batch_size, self.seq_length]     ,		vocab_size=2				)
  
               return (
                   config,
                   input_ids,
                   token_type_ids,
                   input_mask,
                   sequence_labels,
                   token_labels,
                   choice_labels,
                   encoder_hidden_states,
                   encoder_attention_mask,
               )
  
  
  
        def    UpperCamelCase__	(self					:					List[str]     ,		UpperCamelCase					:					Optional[Any]     ,		UpperCamelCase					:					Optional[int]     ,		UpperCamelCase					:					List[str]     ,		UpperCamelCase					:					Optional[int]     ,		UpperCamelCase					:					Dict     ,		UpperCamelCase					:					List[str]     ,		UpperCamelCase					:					Optional[Any]				):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               lowercase__       =				RemBertModel(config=UpperCamelCase				)
               model.to(UpperCamelCase				)
               model.eval()
               lowercase__       =				model(UpperCamelCase     ,		attention_mask=UpperCamelCase     ,		token_type_ids=UpperCamelCase				)
               lowercase__       =				model(UpperCamelCase     ,		token_type_ids=UpperCamelCase				)
               lowercase__       =				model(UpperCamelCase				)
               self.parent.assertEqual(result.last_hidden_state.shape     ,		(self.batch_size, self.seq_length, self.hidden_size)				)
  
  
  
        def    UpperCamelCase__	(self					:					str     ,		UpperCamelCase					:					List[str]     ,		UpperCamelCase					:					Tuple     ,		UpperCamelCase					:					str     ,		UpperCamelCase					:					Optional[int]     ,		UpperCamelCase					:					Tuple     ,		UpperCamelCase					:					Dict     ,		UpperCamelCase					:					List[str]     ,		UpperCamelCase					:					str     ,		UpperCamelCase					:					List[str]     ,		):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               lowercase__       =				True
               lowercase__       =				RemBertModel(UpperCamelCase				)
               model.to(UpperCamelCase				)
               model.eval()
               lowercase__       =				model(
                   UpperCamelCase     ,		attention_mask=UpperCamelCase     ,		token_type_ids=UpperCamelCase     ,		encoder_hidden_states=UpperCamelCase     ,		encoder_attention_mask=UpperCamelCase     ,		)
               lowercase__       =				model(
                   UpperCamelCase     ,		attention_mask=UpperCamelCase     ,		token_type_ids=UpperCamelCase     ,		encoder_hidden_states=UpperCamelCase     ,		)
               lowercase__       =				model(UpperCamelCase     ,		attention_mask=UpperCamelCase     ,		token_type_ids=UpperCamelCase				)
               self.parent.assertEqual(result.last_hidden_state.shape     ,		(self.batch_size, self.seq_length, self.hidden_size)				)
  
  
  
        def    UpperCamelCase__	(self					:					Optional[int]     ,		UpperCamelCase					:					int     ,		UpperCamelCase					:					Tuple     ,		UpperCamelCase					:					Optional[Any]     ,		UpperCamelCase					:					int     ,		UpperCamelCase					:					Optional[int]     ,		UpperCamelCase					:					Optional[int]     ,		UpperCamelCase					:					Dict     ,		UpperCamelCase					:					Optional[int]     ,		UpperCamelCase					:					int     ,		):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               lowercase__       =				RemBertForCausalLM(config=UpperCamelCase				)
               model.to(UpperCamelCase				)
               model.eval()
               lowercase__       =				model(UpperCamelCase     ,		attention_mask=UpperCamelCase     ,		token_type_ids=UpperCamelCase     ,		labels=UpperCamelCase				)
               self.parent.assertEqual(result.logits.shape     ,		(self.batch_size, self.seq_length, self.vocab_size)				)
  
  
  
        def    UpperCamelCase__	(self					:					Tuple     ,		UpperCamelCase					:					Optional[Any]     ,		UpperCamelCase					:					Union[str, Any]     ,		UpperCamelCase					:					Dict     ,		UpperCamelCase					:					Dict     ,		UpperCamelCase					:					Dict     ,		UpperCamelCase					:					List[str]     ,		UpperCamelCase					:					Union[str, Any]				):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               lowercase__       =				RemBertForMaskedLM(config=UpperCamelCase				)
               model.to(UpperCamelCase				)
               model.eval()
               lowercase__       =				model(UpperCamelCase     ,		attention_mask=UpperCamelCase     ,		token_type_ids=UpperCamelCase     ,		labels=UpperCamelCase				)
               self.parent.assertEqual(result.logits.shape     ,		(self.batch_size, self.seq_length, self.vocab_size)				)
  
  
  
        def    UpperCamelCase__	(self					:					Optional[int]     ,		UpperCamelCase					:					Optional[Any]     ,		UpperCamelCase					:					Tuple     ,		UpperCamelCase					:					Optional[int]     ,		UpperCamelCase					:					Any     ,		UpperCamelCase					:					Optional[int]     ,		UpperCamelCase					:					List[Any]     ,		UpperCamelCase					:					str     ,		UpperCamelCase					:					Any     ,		UpperCamelCase					:					int     ,		):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               lowercase__       =				True
               lowercase__       =				True
               lowercase__       =				RemBertForCausalLM(config=UpperCamelCase				)
               model.to(UpperCamelCase				)
               model.eval()
  
               # first forward pass
               lowercase__       =				model(
                   UpperCamelCase     ,		attention_mask=UpperCamelCase     ,		encoder_hidden_states=UpperCamelCase     ,		encoder_attention_mask=UpperCamelCase     ,		use_cache=UpperCamelCase     ,		)
               lowercase__       =				outputs.past_key_values
  
               # create hypothetical multiple next token and extent to next_input_ids
               lowercase__       =				ids_tensor((self.batch_size, 3)     ,		config.vocab_size				)
               lowercase__       =				ids_tensor((self.batch_size, 3)     ,		vocab_size=2				)
  
               # append to next input_ids and
               lowercase__       =				torch.cat([input_ids, next_tokens]     ,		dim=-1				)
               lowercase__       =				torch.cat([input_mask, next_mask]     ,		dim=-1				)
  
               lowercase__       =				model(
                   UpperCamelCase     ,		attention_mask=UpperCamelCase     ,		encoder_hidden_states=UpperCamelCase     ,		encoder_attention_mask=UpperCamelCase     ,		output_hidden_states=UpperCamelCase     ,		)['''hidden_states'''][0]
               lowercase__       =				model(
                   UpperCamelCase     ,		attention_mask=UpperCamelCase     ,		encoder_hidden_states=UpperCamelCase     ,		encoder_attention_mask=UpperCamelCase     ,		past_key_values=UpperCamelCase     ,		output_hidden_states=UpperCamelCase     ,		)['''hidden_states'''][0]
  
               # select random slice
               lowercase__       =				ids_tensor((1,)     ,		output_from_past.shape[-1]				).item()
               lowercase__       =				output_from_no_past[:, -3:, random_slice_idx].detach()
               lowercase__       =				output_from_past[:, :, random_slice_idx].detach()
  
               self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]				)
  
               # test that outputs are equal for slice
               self.parent.assertTrue(torch.allclose(UpperCamelCase     ,		UpperCamelCase     ,		atol=1E-3				)				)
  
  
  
        def    UpperCamelCase__	(self					:					Any     ,		UpperCamelCase					:					Any     ,		UpperCamelCase					:					Union[str, Any]     ,		UpperCamelCase					:					Dict     ,		UpperCamelCase					:					Dict     ,		UpperCamelCase					:					Union[str, Any]     ,		UpperCamelCase					:					Any     ,		UpperCamelCase					:					Tuple				):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               lowercase__       =				RemBertForQuestionAnswering(config=UpperCamelCase				)
               model.to(UpperCamelCase				)
               model.eval()
               lowercase__       =				model(
                   UpperCamelCase     ,		attention_mask=UpperCamelCase     ,		token_type_ids=UpperCamelCase     ,		start_positions=UpperCamelCase     ,		end_positions=UpperCamelCase     ,		)
               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    UpperCamelCase__	(self					:					int     ,		UpperCamelCase					:					Union[str, Any]     ,		UpperCamelCase					:					int     ,		UpperCamelCase					:					List[str]     ,		UpperCamelCase					:					Union[str, Any]     ,		UpperCamelCase					:					Optional[int]     ,		UpperCamelCase					:					Union[str, Any]     ,		UpperCamelCase					:					List[str]				):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               lowercase__       =				self.num_labels
               lowercase__       =				RemBertForSequenceClassification(UpperCamelCase				)
               model.to(UpperCamelCase				)
               model.eval()
               lowercase__       =				model(UpperCamelCase     ,		attention_mask=UpperCamelCase     ,		token_type_ids=UpperCamelCase     ,		labels=UpperCamelCase				)
               self.parent.assertEqual(result.logits.shape     ,		(self.batch_size, self.num_labels)				)
  
  
  
        def    UpperCamelCase__	(self					:					Union[str, Any]     ,		UpperCamelCase					:					Tuple     ,		UpperCamelCase					:					List[Any]     ,		UpperCamelCase					:					Tuple     ,		UpperCamelCase					:					Optional[Any]     ,		UpperCamelCase					:					str     ,		UpperCamelCase					:					Dict     ,		UpperCamelCase					:					int				):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               lowercase__       =				self.num_labels
               lowercase__       =				RemBertForTokenClassification(config=UpperCamelCase				)
               model.to(UpperCamelCase				)
               model.eval()
               lowercase__       =				model(UpperCamelCase     ,		attention_mask=UpperCamelCase     ,		token_type_ids=UpperCamelCase     ,		labels=UpperCamelCase				)
               self.parent.assertEqual(result.logits.shape     ,		(self.batch_size, self.seq_length, self.num_labels)				)
  
  
  
        def    UpperCamelCase__	(self					:					int     ,		UpperCamelCase					:					int     ,		UpperCamelCase					:					int     ,		UpperCamelCase					:					Tuple     ,		UpperCamelCase					:					List[str]     ,		UpperCamelCase					:					Tuple     ,		UpperCamelCase					:					List[Any]     ,		UpperCamelCase					:					Dict				):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               lowercase__       =				self.num_choices
               lowercase__       =				RemBertForMultipleChoice(config=UpperCamelCase				)
               model.to(UpperCamelCase				)
               model.eval()
               lowercase__       =				input_ids.unsqueeze(1				).expand(-1     ,		self.num_choices     ,		-1				).contiguous()
               lowercase__       =				token_type_ids.unsqueeze(1				).expand(-1     ,		self.num_choices     ,		-1				).contiguous()
               lowercase__       =				input_mask.unsqueeze(1				).expand(-1     ,		self.num_choices     ,		-1				).contiguous()
               lowercase__       =				model(
                   UpperCamelCase     ,		attention_mask=UpperCamelCase     ,		token_type_ids=UpperCamelCase     ,		labels=UpperCamelCase     ,		)
               self.parent.assertEqual(result.logits.shape     ,		(self.batch_size, self.num_choices)				)
  
  
  
  
  
  
        def    UpperCamelCase__	(self					:					int				):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               lowercase__       =				self.prepare_config_and_inputs()
               (
                   (
                   lowercase__       
               )    ,(
                   lowercase__       
               )    ,(
                   lowercase__       
               )    ,(
                   lowercase__       
               )    ,(
                   lowercase__       
               )    ,(
                   lowercase__       
               )    ,(
                   lowercase__       
               )    ,       
               )       =				config_and_inputs
               lowercase__       =				{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
               return config, inputs_dict
  
  
  
  
  
  
  @require_torch
  class 						__lowerCAmelCase		(__lowerCamelCase						,  __lowerCamelCase						,  unittest.TestCase  ):
  
  
  
  
  
        '''simple docstring'''
  
  
  
  
  
  
        lowerCAmelCase__      :  str						      =							(
            (
                RemBertModel,
                RemBertForMaskedLM,
                RemBertForCausalLM,
                RemBertForMultipleChoice,
                RemBertForQuestionAnswering,
                RemBertForSequenceClassification,
                RemBertForTokenClassification,
            )
            if is_torch_available()
            else ()
        )
        lowerCAmelCase__      :  Optional[Any]						      =							(RemBertForCausalLM,) if is_torch_available() else ()
        lowerCAmelCase__      :  Dict						      =							(
            {
                'feature-extraction': RemBertModel,
                'fill-mask': RemBertForMaskedLM,
                'question-answering': RemBertForQuestionAnswering,
                'text-classification': RemBertForSequenceClassification,
                'text-generation': RemBertForCausalLM,
                'token-classification': RemBertForTokenClassification,
                'zero-shot': RemBertForSequenceClassification,
            }
            if is_torch_available()
            else {}
        )
  
  
  
        def    UpperCamelCase__	(self					:					Optional[Any]				):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               lowercase__       =				RemBertModelTester(self				)
               lowercase__       =				ConfigTester(self     ,		config_class=UpperCamelCase     ,		hidden_size=37				)
  
  
  
        def    UpperCamelCase__	(self					:					Dict				):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               self.config_tester.run_common_tests()
  
  
  
        def    UpperCamelCase__	(self					:					Dict				):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               lowercase__       =				self.model_tester.prepare_config_and_inputs()
               self.model_tester.create_and_check_model(*UpperCamelCase				)
  
  
  
        def    UpperCamelCase__	(self					:					Dict				):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               lowercase__       =				self.model_tester.prepare_config_and_inputs()
               for type in ["absolute", "relative_key", "relative_key_query"]:
                      lowercase__       =				type
                      self.model_tester.create_and_check_model(*UpperCamelCase				)
  
  
  
        def    UpperCamelCase__	(self					:					Union[str, Any]				):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               lowercase__       =				self.model_tester.prepare_config_and_inputs()
               self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase				)
  
  
  
        def    UpperCamelCase__	(self					:					Tuple				):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               lowercase__       =				self.model_tester.prepare_config_and_inputs()
               self.model_tester.create_and_check_for_multiple_choice(*UpperCamelCase				)
  
  
  
        def    UpperCamelCase__	(self					:					List[Any]				):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               lowercase__       =				self.model_tester.prepare_config_and_inputs_for_decoder()
               self.model_tester.create_and_check_decoder_model_past_large_inputs(*UpperCamelCase				)
  
  
  
        def    UpperCamelCase__	(self					:					str				):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               lowercase__       =				self.model_tester.prepare_config_and_inputs()
               self.model_tester.create_and_check_for_question_answering(*UpperCamelCase				)
  
  
  
        def    UpperCamelCase__	(self					:					Dict				):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               lowercase__       =				self.model_tester.prepare_config_and_inputs()
               self.model_tester.create_and_check_for_sequence_classification(*UpperCamelCase				)
  
  
  
        def    UpperCamelCase__	(self					:					Any				):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               lowercase__       =				self.model_tester.prepare_config_and_inputs()
               self.model_tester.create_and_check_for_token_classification(*UpperCamelCase				)
  
  
  
        def    UpperCamelCase__	(self					:					List[Any]				):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               lowercase__       =				self.model_tester.prepare_config_and_inputs_for_decoder()
               self.model_tester.create_and_check_model_as_decoder(*UpperCamelCase				)
  
  
  
        def    UpperCamelCase__	(self					:					List[str]				):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               (
                   (
                   lowercase__       
               )    ,(
                   lowercase__       
               )    ,(
                   lowercase__       
               )    ,(
                   lowercase__       
               )    ,(
                   lowercase__       
               )    ,(
                   lowercase__       
               )    ,(
                   lowercase__       
               )    ,(
                   lowercase__       
               )    ,(
                   lowercase__       
               )    ,       
               )       =				self.model_tester.prepare_config_and_inputs_for_decoder()
  
               lowercase__       =				None
  
               self.model_tester.create_and_check_model_as_decoder(
                   UpperCamelCase     ,		UpperCamelCase     ,		UpperCamelCase     ,		UpperCamelCase     ,		UpperCamelCase     ,		UpperCamelCase     ,		UpperCamelCase     ,		UpperCamelCase     ,		UpperCamelCase     ,		)
  
  
  
  
  
  
        @slow
        def    UpperCamelCase__	(self					:					Tuple				):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               for model_name in REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                      lowercase__       =				RemBertModel.from_pretrained(UpperCamelCase				)
                      self.assertIsNotNone(UpperCamelCase				)
  
  
  
  
  
  
  @require_torch
  class 						__lowerCAmelCase		(unittest.TestCase  ):
  
  
  
  
  
        '''simple docstring'''
  
  
  
  
  
  
        @slow
        def    UpperCamelCase__	(self					:					List[Any]				):
  
               '''simple docstring'''
  
  
  
  
  
  
  
               lowercase__       =				RemBertModel.from_pretrained('''google/rembert'''				)
               lowercase__       =				torch.tensor([[312, 56498, 313, 2125, 313]]				)
               lowercase__       =				torch.tensor([[0, 0, 0, 1, 1]]				)
               with torch.no_grad():
                      lowercase__       =				model(UpperCamelCase     ,		token_type_ids=UpperCamelCase     ,		output_hidden_states=UpperCamelCase				)
  
               lowercase__       =				1152
  
               lowercase__       =				torch.Size((1, 5, hidden_size)				)
               self.assertEqual(output['''last_hidden_state'''].shape     ,		UpperCamelCase				)
  
               lowercase__       =				torch.tensor(
                   [
                       [
                           [0.07_54, -0.20_22, 0.19_04],
                           [-0.33_54, -0.36_92, -0.47_91],
                           [-0.23_14, -0.67_29, -0.07_49],
                           [-0.03_96, -0.31_05, -0.42_34],
                           [-0.15_71, -0.05_25, 0.53_53],
                       ]
                   ]				)
  
               # Running on the original tf implementation gives slightly different results here.
               # Not clear why this variations is present
               # TODO: Find reason for discrepancy
               # expected_original_implementation = [[
               #     [0.07630594074726105, -0.20146065950393677, 0.19107051193714142],
               #     [-0.3405614495277405, -0.36971670389175415, -0.4808273911476135],
               #     [-0.22587086260318756, -0.6656315922737122, -0.07844287157058716],
               #     [-0.04145475849509239, -0.3077218234539032, -0.42316967248916626],
               #     [-0.15887849032878876, -0.054529931396245956, 0.5356100797653198]
               # ]]
  
               self.assertTrue(torch.allclose(output['''last_hidden_state'''][:, :, :3]     ,		UpperCamelCase     ,		atol=1E-4				)				)
  
  
  
  
  
style_context: 
  
  
  '''simple docstring'''
  
  import argparse
  
  import torch
  
  from transformers import ElectraConfig, ElectraForMaskedLM, ElectraForPreTraining, load_tf_weights_in_electra
  from transformers.utils import logging
  
  
  logging.set_verbosity_info()
  
  def _SCREAMING_SNAKE_CASE   (A						,				A						,				A						,				A ) -> Dict:
  
  
  
  
  
  
         """simple docstring"""
  
  
  
  
  
  
  
         lowercase__       =				ElectraConfig.from_json_file(A )
         print(f"Building PyTorch model from configuration: {config}" )
  
         if discriminator_or_generator == "discriminator":
                lowercase__       =				ElectraForPreTraining(A )
         elif discriminator_or_generator == "generator":
                lowercase__       =				ElectraForMaskedLM(A )
         else:
                raise ValueError('''The discriminator_or_generator argument should be either \'discriminator\' or \'generator\'''' )
  
         # Load weights from tf checkpoint
         load_tf_weights_in_electra(
             A						,				A						,				A						,				discriminator_or_generator=A )
  
         # Save pytorch-model
         print(f"Save PyTorch model to {pytorch_dump_path}" )
         torch.save(model.state_dict()						,				A )
  
  
  if __name__ == "__main__":
   lowerCamelCase							:     Tuple								=			argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       '--tf_checkpoint_path', default=None, type=str, required=True, help='Path to the TensorFlow checkpoint path.'
   )
   parser.add_argument(
       '--config_file',
       default=None,
       type=str,
       required=True,
       help='The config json file corresponding to the pre-trained model. \nThis specifies the model architecture.',
   )
   parser.add_argument(
       '--pytorch_dump_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
   )
   parser.add_argument(
       '--discriminator_or_generator',
       default=None,
       type=str,
       required=True,
       help=(
           'Whether to export the generator or the discriminator. Should be a string, either \'discriminator\' or '
           '\'generator\'.'
       ),
   )
   lowerCamelCase							:     Optional[int]								=			parser.parse_args()
   convert_tf_checkpoint_to_pytorch(
       args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path, args.discriminator_or_generator
   )
  
  
  
  
  