



import unittest

from transformers import ElectraConfig, is_torch_available
from transformers.models.auto import get_values
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 (
          MODEL_FOR_PRETRAINING_MAPPING,
          ElectraForCausalLM,
          ElectraForMaskedLM,
          ElectraForMultipleChoice,
          ElectraForPreTraining,
          ElectraForQuestionAnswering,
          ElectraForSequenceClassification,
          ElectraForTokenClassification,
          ElectraModel,
      )
      from transformers.models.electra.modeling_electra import ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST

class    _SCREAMING_SNAKE_CASE						:





     def __init__(							self			:				Dict  ,	UpperCamelCase			:				List[str]  ,	UpperCamelCase			:				List[Any]=1_3  ,	UpperCamelCase			:				str=7  ,	UpperCamelCase			:				Union[str, Any]=True  ,	UpperCamelCase			:				List[str]=True  ,	UpperCamelCase			:				Optional[Any]=True  ,	UpperCamelCase			:				Optional[int]=True  ,	UpperCamelCase			:				str=9_9  ,	UpperCamelCase			:				Dict=3_2  ,	UpperCamelCase			:				List[str]=5  ,	UpperCamelCase			:				str=4  ,	UpperCamelCase			:				List[str]=3_7  ,	UpperCamelCase			:				str="gelu"  ,	UpperCamelCase			:				Dict=0.1  ,	UpperCamelCase			:				Dict=0.1  ,	UpperCamelCase			:				Optional[int]=5_1_2  ,	UpperCamelCase			:				str=1_6  ,	UpperCamelCase			:				Optional[int]=2  ,	UpperCamelCase			:				List[Any]=0.0_2  ,	UpperCamelCase			:				List[Any]=3  ,	UpperCamelCase			:				int=4  ,	UpperCamelCase			:				str=None  ,	)->List[str]:
      __SCREAMING_SNAKE_CASE							:     List[str]							=						parent
      __SCREAMING_SNAKE_CASE							:     List[str]							=						batch_size
      __SCREAMING_SNAKE_CASE							:     Tuple							=						seq_length
      __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						is_training
      __SCREAMING_SNAKE_CASE							:     Dict							=						use_input_mask
      __SCREAMING_SNAKE_CASE							:     int							=						use_token_type_ids
      __SCREAMING_SNAKE_CASE							:     List[Any]							=						use_labels
      __SCREAMING_SNAKE_CASE							:     List[str]							=						vocab_size
      __SCREAMING_SNAKE_CASE							:     Dict							=						hidden_size
      __SCREAMING_SNAKE_CASE							:     int							=						num_hidden_layers
      __SCREAMING_SNAKE_CASE							:     str							=						num_attention_heads
      __SCREAMING_SNAKE_CASE							:     Any							=						intermediate_size
      __SCREAMING_SNAKE_CASE							:     Any							=						hidden_act
      __SCREAMING_SNAKE_CASE							:     Any							=						hidden_dropout_prob
      __SCREAMING_SNAKE_CASE							:     int							=						attention_probs_dropout_prob
      __SCREAMING_SNAKE_CASE							:     Optional[int]							=						max_position_embeddings
      __SCREAMING_SNAKE_CASE							:     List[Any]							=						type_vocab_size
      __SCREAMING_SNAKE_CASE							:     Any							=						type_sequence_label_size
      __SCREAMING_SNAKE_CASE							:     Optional[int]							=						initializer_range
      __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						num_labels
      __SCREAMING_SNAKE_CASE							:     Dict							=						num_choices
      __SCREAMING_SNAKE_CASE							:     Any							=						scope





     def 						__snake_case     (							self			:				Any					)->Tuple:
      __SCREAMING_SNAKE_CASE							:     List[Any]							=						ids_tensor([self.batch_size, self.seq_length]  ,	self.vocab_size					)

      __SCREAMING_SNAKE_CASE							:     int							=						None
      if self.use_input_mask:
       __SCREAMING_SNAKE_CASE							:     str							=						random_attention_mask([self.batch_size, self.seq_length]					)

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

      __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						None
      __SCREAMING_SNAKE_CASE							:     Any							=						None
      __SCREAMING_SNAKE_CASE							:     str							=						None
      if self.use_labels:
       __SCREAMING_SNAKE_CASE							:     List[Any]							=						ids_tensor([self.batch_size]  ,	self.type_sequence_label_size					)
       __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						ids_tensor([self.batch_size, self.seq_length]  ,	self.num_labels					)
       __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						ids_tensor([self.batch_size]  ,	self.num_choices					)
       __SCREAMING_SNAKE_CASE							:     str							=						ids_tensor([self.batch_size, self.seq_length]  ,	1					)

      __SCREAMING_SNAKE_CASE							:     Dict							=						self.get_config()

      return (
          config,
          input_ids,
          token_type_ids,
          input_mask,
          sequence_labels,
          token_labels,
          choice_labels,
          fake_token_labels,
      )





     def 						__snake_case     (							self			:				List[Any]					)->Tuple:
      return ElectraConfig(
          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  ,	is_decoder=UpperCamelCase  ,	initializer_range=self.initializer_range  ,	)





     def 						__snake_case     (							self			:				List[str]					)->List[str]:
      (
          (
          __SCREAMING_SNAKE_CASE
      ),		(
          __SCREAMING_SNAKE_CASE
      ),		(
          __SCREAMING_SNAKE_CASE
      ),		(
          __SCREAMING_SNAKE_CASE
      ),		(
          __SCREAMING_SNAKE_CASE
      ),		(
          __SCREAMING_SNAKE_CASE
      ),		(
          __SCREAMING_SNAKE_CASE
      ),		(
          __SCREAMING_SNAKE_CASE
      ),		
      )							:     List[Any]							=						self.prepare_config_and_inputs()

      __SCREAMING_SNAKE_CASE							:     List[Any]							=						True
      __SCREAMING_SNAKE_CASE							:     List[str]							=						floats_tensor([self.batch_size, self.seq_length, self.hidden_size]					)
      __SCREAMING_SNAKE_CASE							:     int							=						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 						__snake_case     (							self			:				List[str]  ,	UpperCamelCase			:				Union[str, Any]  ,	UpperCamelCase			:				int  ,	UpperCamelCase			:				str  ,	UpperCamelCase			:				Tuple  ,	UpperCamelCase			:				Optional[Any]  ,	UpperCamelCase			:				Optional[Any]  ,	UpperCamelCase			:				Optional[int]  ,	UpperCamelCase			:				Union[str, Any]  ,	)->Union[str, Any]:
      __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						ElectraModel(config=UpperCamelCase					)
      model.to(UpperCamelCase					)
      model.eval()
      __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						model(UpperCamelCase  ,	attention_mask=UpperCamelCase  ,	token_type_ids=UpperCamelCase					)
      __SCREAMING_SNAKE_CASE							:     Optional[int]							=						model(UpperCamelCase  ,	token_type_ids=UpperCamelCase					)
      __SCREAMING_SNAKE_CASE							:     List[str]							=						model(UpperCamelCase					)
      self.parent.assertEqual(result.last_hidden_state.shape  ,	(self.batch_size, self.seq_length, self.hidden_size)					)





     def 						__snake_case     (							self			:				List[Any]  ,	UpperCamelCase			:				Optional[Any]  ,	UpperCamelCase			:				Tuple  ,	UpperCamelCase			:				int  ,	UpperCamelCase			:				List[Any]  ,	UpperCamelCase			:				List[Any]  ,	UpperCamelCase			:				str  ,	UpperCamelCase			:				Dict  ,	UpperCamelCase			:				List[str]  ,	UpperCamelCase			:				Dict  ,	)->Dict:
      __SCREAMING_SNAKE_CASE							:     Any							=						True
      __SCREAMING_SNAKE_CASE							:     List[Any]							=						ElectraModel(UpperCamelCase					)
      model.to(UpperCamelCase					)
      model.eval()
      __SCREAMING_SNAKE_CASE							:     Dict							=						model(
          UpperCamelCase  ,	attention_mask=UpperCamelCase  ,	token_type_ids=UpperCamelCase  ,	encoder_hidden_states=UpperCamelCase  ,	encoder_attention_mask=UpperCamelCase  ,	)
      __SCREAMING_SNAKE_CASE							:     Tuple							=						model(
          UpperCamelCase  ,	attention_mask=UpperCamelCase  ,	token_type_ids=UpperCamelCase  ,	encoder_hidden_states=UpperCamelCase  ,	)
      __SCREAMING_SNAKE_CASE							:     Optional[int]							=						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 						__snake_case     (							self			:				int  ,	UpperCamelCase			:				List[Any]  ,	UpperCamelCase			:				str  ,	UpperCamelCase			:				Any  ,	UpperCamelCase			:				int  ,	UpperCamelCase			:				Union[str, Any]  ,	UpperCamelCase			:				Optional[Any]  ,	UpperCamelCase			:				Any  ,	UpperCamelCase			:				int  ,	)->Union[str, Any]:
      __SCREAMING_SNAKE_CASE							:     Dict							=						ElectraForMaskedLM(config=UpperCamelCase					)
      model.to(UpperCamelCase					)
      model.eval()
      __SCREAMING_SNAKE_CASE							:     List[str]							=						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 						__snake_case     (							self			:				Optional[int]  ,	UpperCamelCase			:				Optional[Any]  ,	UpperCamelCase			:				Any  ,	UpperCamelCase			:				List[Any]  ,	UpperCamelCase			:				str  ,	UpperCamelCase			:				Union[str, Any]  ,	UpperCamelCase			:				int  ,	UpperCamelCase			:				List[Any]  ,	UpperCamelCase			:				Union[str, Any]  ,	UpperCamelCase			:				Any  ,	)->int:
      __SCREAMING_SNAKE_CASE							:     List[Any]							=						ElectraForCausalLM(config=UpperCamelCase					)
      model.to(UpperCamelCase					)
      model.eval()
      __SCREAMING_SNAKE_CASE							:     Any							=						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 						__snake_case     (							self			:				List[str]  ,	UpperCamelCase			:				List[str]  ,	UpperCamelCase			:				Union[str, Any]  ,	UpperCamelCase			:				Optional[Any]  ,	UpperCamelCase			:				Optional[int]  ,	UpperCamelCase			:				Dict  ,	UpperCamelCase			:				str  ,	UpperCamelCase			:				Optional[int]  ,	UpperCamelCase			:				Any  ,	)->Optional[Any]:
      __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						self.num_labels
      __SCREAMING_SNAKE_CASE							:     Tuple							=						ElectraForTokenClassification(config=UpperCamelCase					)
      model.to(UpperCamelCase					)
      model.eval()
      __SCREAMING_SNAKE_CASE							:     Tuple							=						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 						__snake_case     (							self			:				Optional[Any]  ,	UpperCamelCase			:				str  ,	UpperCamelCase			:				Dict  ,	UpperCamelCase			:				Optional[Any]  ,	UpperCamelCase			:				int  ,	UpperCamelCase			:				Tuple  ,	UpperCamelCase			:				List[str]  ,	UpperCamelCase			:				List[str]  ,	UpperCamelCase			:				Dict  ,	)->str:
      __SCREAMING_SNAKE_CASE							:     Optional[int]							=						self.num_labels
      __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						ElectraForPreTraining(config=UpperCamelCase					)
      model.to(UpperCamelCase					)
      model.eval()
      __SCREAMING_SNAKE_CASE							:     List[Any]							=						model(UpperCamelCase  ,	attention_mask=UpperCamelCase  ,	token_type_ids=UpperCamelCase  ,	labels=UpperCamelCase					)
      self.parent.assertEqual(result.logits.shape  ,	(self.batch_size, self.seq_length)					)





     def 						__snake_case     (							self			:				Union[str, Any]  ,	UpperCamelCase			:				Optional[Any]  ,	UpperCamelCase			:				Optional[int]  ,	UpperCamelCase			:				List[Any]  ,	UpperCamelCase			:				int  ,	UpperCamelCase			:				List[Any]  ,	UpperCamelCase			:				Optional[int]  ,	UpperCamelCase			:				Any  ,	UpperCamelCase			:				List[str]  ,	)->Any:
      __SCREAMING_SNAKE_CASE							:     int							=						self.num_labels
      __SCREAMING_SNAKE_CASE							:     List[Any]							=						ElectraForSequenceClassification(UpperCamelCase					)
      model.to(UpperCamelCase					)
      model.eval()
      __SCREAMING_SNAKE_CASE							:     str							=						model(UpperCamelCase  ,	attention_mask=UpperCamelCase  ,	token_type_ids=UpperCamelCase  ,	labels=UpperCamelCase					)
      self.parent.assertEqual(result.logits.shape  ,	(self.batch_size, self.num_labels)					)





     def 						__snake_case     (							self			:				Any  ,	UpperCamelCase			:				Union[str, Any]  ,	UpperCamelCase			:				Union[str, Any]  ,	UpperCamelCase			:				List[Any]  ,	UpperCamelCase			:				Optional[Any]  ,	UpperCamelCase			:				Tuple  ,	UpperCamelCase			:				List[Any]  ,	UpperCamelCase			:				List[Any]  ,	UpperCamelCase			:				Union[str, Any]  ,	)->Optional[Any]:
      __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						ElectraForQuestionAnswering(config=UpperCamelCase					)
      model.to(UpperCamelCase					)
      model.eval()
      __SCREAMING_SNAKE_CASE							:     List[str]							=						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 						__snake_case     (							self			:				int  ,	UpperCamelCase			:				int  ,	UpperCamelCase			:				Optional[int]  ,	UpperCamelCase			:				Optional[int]  ,	UpperCamelCase			:				List[Any]  ,	UpperCamelCase			:				List[Any]  ,	UpperCamelCase			:				str  ,	UpperCamelCase			:				Optional[Any]  ,	UpperCamelCase			:				Optional[int]  ,	)->int:
      __SCREAMING_SNAKE_CASE							:     Optional[int]							=						self.num_choices
      __SCREAMING_SNAKE_CASE							:     List[str]							=						ElectraForMultipleChoice(config=UpperCamelCase					)
      model.to(UpperCamelCase					)
      model.eval()
      __SCREAMING_SNAKE_CASE							:     Optional[int]							=						input_ids.unsqueeze(1					).expand(-1  ,	self.num_choices  ,	-1					).contiguous()
      __SCREAMING_SNAKE_CASE							:     Optional[int]							=						token_type_ids.unsqueeze(1					).expand(-1  ,	self.num_choices  ,	-1					).contiguous()
      __SCREAMING_SNAKE_CASE							:     Any							=						input_mask.unsqueeze(1					).expand(-1  ,	self.num_choices  ,	-1					).contiguous()
      __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						model(
          UpperCamelCase  ,	attention_mask=UpperCamelCase  ,	token_type_ids=UpperCamelCase  ,	labels=UpperCamelCase  ,	)
      self.parent.assertEqual(result.logits.shape  ,	(self.batch_size, self.num_choices)					)






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






@require_torch
class    _SCREAMING_SNAKE_CASE						(UpperCamelCase      ,     UpperCamelCase      ,     unittest.TestCase			):
     lowerCAmelCase       =      (
         (
             ElectraModel,
             ElectraForPreTraining,
             ElectraForMaskedLM,
             ElectraForCausalLM,
             ElectraForMultipleChoice,
             ElectraForTokenClassification,
             ElectraForSequenceClassification,
             ElectraForQuestionAnswering,
         )
         if is_torch_available()
         else ()
     )
     lowerCAmelCase       =      (
         {
             """feature-extraction""": ElectraModel,
             """fill-mask""": ElectraForMaskedLM,
             """question-answering""": ElectraForQuestionAnswering,
             """text-classification""": ElectraForSequenceClassification,
             """text-generation""": ElectraForCausalLM,
             """token-classification""": ElectraForTokenClassification,
             """zero-shot""": ElectraForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     lowerCAmelCase       =      True





     def 						__snake_case     (							self			:				Dict  ,	UpperCamelCase			:				str  ,	UpperCamelCase			:				Optional[int]  ,	UpperCamelCase			:				Dict=False					)->Dict:
      __SCREAMING_SNAKE_CASE							:     Tuple							=						super()._prepare_for_class(UpperCamelCase  ,	UpperCamelCase  ,	return_labels=UpperCamelCase					)

      if return_labels:
       if model_class in get_values(UpperCamelCase					):
        __SCREAMING_SNAKE_CASE							:     str							=						torch.zeros(
            (self.model_tester.batch_size, self.model_tester.seq_length)  ,	dtype=torch.long  ,	device=UpperCamelCase					)
      return inputs_dict





     def 						__snake_case     (							self			:				List[Any]					)->Optional[Any]:
      __SCREAMING_SNAKE_CASE							:     str							=						ElectraModelTester(self					)
      __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						ConfigTester(self  ,	config_class=UpperCamelCase  ,	hidden_size=3_7					)





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





     def 						__snake_case     (							self			:				Dict					)->Dict:
      __SCREAMING_SNAKE_CASE							:     Any							=						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_electra_model(*UpperCamelCase					)





     def 						__snake_case     (							self			:				Union[str, Any]					)->Any:
      __SCREAMING_SNAKE_CASE							:     Dict							=						self.model_tester.prepare_config_and_inputs_for_decoder()
      self.model_tester.create_and_check_electra_model_as_decoder(*UpperCamelCase					)





     def 						__snake_case     (							self			:				Union[str, Any]					)->int:
      __SCREAMING_SNAKE_CASE							:     Tuple							=						self.model_tester.prepare_config_and_inputs()
      for type in ["absolute", "relative_key", "relative_key_query"]:
       __SCREAMING_SNAKE_CASE							:     Any							=						type
       self.model_tester.create_and_check_electra_model(*UpperCamelCase					)





     def 						__snake_case     (							self			:				int					)->Optional[Any]:
      __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_electra_for_masked_lm(*UpperCamelCase					)





     def 						__snake_case     (							self			:				List[str]					)->Optional[int]:
      __SCREAMING_SNAKE_CASE							:     List[Any]							=						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_electra_for_token_classification(*UpperCamelCase					)





     def 						__snake_case     (							self			:				List[Any]					)->List[Any]:
      __SCREAMING_SNAKE_CASE							:     List[Any]							=						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_electra_for_pretraining(*UpperCamelCase					)





     def 						__snake_case     (							self			:				Union[str, Any]					)->Optional[int]:
      __SCREAMING_SNAKE_CASE							:     int							=						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_electra_for_sequence_classification(*UpperCamelCase					)





     def 						__snake_case     (							self			:				Dict					)->List[str]:
      __SCREAMING_SNAKE_CASE							:     int							=						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_electra_for_question_answering(*UpperCamelCase					)





     def 						__snake_case     (							self			:				List[str]					)->Optional[int]:
      __SCREAMING_SNAKE_CASE							:     Optional[int]							=						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_electra_for_multiple_choice(*UpperCamelCase					)





     @slow
     def 						__snake_case     (							self			:				str					)->Optional[int]:
      for model_name in ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       __SCREAMING_SNAKE_CASE							:     Dict							=						ElectraModel.from_pretrained(UpperCamelCase					)
       self.assertIsNotNone(UpperCamelCase					)






     def 						__snake_case     (							self			:				Tuple					)->Dict:
      __SCREAMING_SNAKE_CASE							:     Tuple							=						self.model_tester.prepare_config_and_inputs_for_decoder()
      self.model_tester.create_and_check_electra_for_causal_lm(*UpperCamelCase					)






@require_torch
class    _SCREAMING_SNAKE_CASE						(unittest.TestCase			):





     @slow
     def 						__snake_case     (							self			:				Any					)->Optional[Any]:
      __SCREAMING_SNAKE_CASE							:     str							=						ElectraModel.from_pretrained("google/electra-small-discriminator"					)
      __SCREAMING_SNAKE_CASE							:     List[str]							=						torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]					)
      __SCREAMING_SNAKE_CASE							:     int							=						torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]					)
      __SCREAMING_SNAKE_CASE							:     int							=						model(UpperCamelCase  ,	attention_mask=UpperCamelCase					)[0]
      __SCREAMING_SNAKE_CASE							:     Optional[int]							=						torch.Size((1, 1_1, 2_5_6)					)
      self.assertEqual(output.shape  ,	UpperCamelCase					)
      __SCREAMING_SNAKE_CASE							:     Tuple							=						torch.tensor(
          [[[0.4_4_7_1, 0.6_8_2_1, -0.3_2_6_5], [0.4_6_2_7, 0.5_2_5_5, -0.3_6_6_8], [0.4_5_3_2, 0.3_3_1_3, -0.4_3_4_4]]]					)

      self.assertTrue(torch.allclose(output[:, 1:4, 1:4]  ,	UpperCamelCase  ,	atol=1E-4					)					)











import argparse
import json
import re
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    MobileNetVaConfig,
    MobileNetVaForImageClassification,
    MobileNetVaForSemanticSegmentation,
    MobileNetVaImageProcessor,
    load_tf_weights_in_mobilenet_va,
)
from transformers.utils import logging


logging.set_verbosity_info()
_lowerCamelCase         =	logging.get_logger(__name__)






def 						_lowerCAmelCase    (					__lowerCamelCase :  Union[str, Any]  ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE							:     Dict							=						MobileNetVaConfig(layer_norm_eps=0.001  )

 if "quant" in model_name:
  raise ValueError("Quantized models are not supported."  )

 __SCREAMING_SNAKE_CASE							:     Tuple							=						re.match(r"^.*mobilenet_v2_([^_]*)_([^_]*)$"		,						__lowerCamelCase  )
 if matches:
  __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						float(matches[1]  )
  __SCREAMING_SNAKE_CASE							:     str							=						int(matches[2]  )

 if model_name.startswith("deeplabv3_"  ):
  __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						8
  __SCREAMING_SNAKE_CASE							:     Dict							=						21
  __SCREAMING_SNAKE_CASE							:     Tuple							=						"pascal-voc-id2label.json"
 else:
  # The TensorFlow version of MobileNetV2 predicts 1001 classes instead
  # of the usual 1000. The first class (index 0) is "background".
  __SCREAMING_SNAKE_CASE							:     int							=						1001
  __SCREAMING_SNAKE_CASE							:     Dict							=						"imagenet-1k-id2label.json"

 __SCREAMING_SNAKE_CASE							:     List[Any]							=						"huggingface/label-files"
 __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						json.load(open(hf_hub_download(__lowerCamelCase		,						__lowerCamelCase		,						repo_type="dataset"  )		,						"r"  )  )

 if config.num_labels == 1001:
  __SCREAMING_SNAKE_CASE							:     int							=						{int(__lowerCamelCase  ) + 1: v for k, v in idalabel.items()}
  __SCREAMING_SNAKE_CASE							:     int							=						"background"
 else:
  __SCREAMING_SNAKE_CASE							:     Any							=						{int(__lowerCamelCase  ): v for k, v in idalabel.items()}

 __SCREAMING_SNAKE_CASE							:     str							=						idalabel
 __SCREAMING_SNAKE_CASE							:     Tuple							=						{v: k for k, v in idalabel.items()}

 return config






def 						_lowerCAmelCase    (					):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE							:     Dict							=						"http://images.cocodataset.org/val2017/000000039769.jpg"
 __SCREAMING_SNAKE_CASE							:     Any							=						Image.open(requests.get(__lowerCamelCase		,						stream=__lowerCamelCase  ).raw  )
 return im






@torch.no_grad()
def 						_lowerCAmelCase    (					__lowerCamelCase :  str		,						__lowerCamelCase :  Optional[Any]		,						__lowerCamelCase :  Optional[Any]		,						__lowerCamelCase :  Tuple=False  ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE							:     Any							=						get_mobilenet_va_config(__lowerCamelCase  )

 # Load 🤗 model
 if model_name.startswith("deeplabv3_"  ):
  __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						MobileNetVaForSemanticSegmentation(__lowerCamelCase  ).eval()
 else:
  __SCREAMING_SNAKE_CASE							:     str							=						MobileNetVaForImageClassification(__lowerCamelCase  ).eval()

 # Load weights from TensorFlow checkpoint
 load_tf_weights_in_mobilenet_va(__lowerCamelCase		,						__lowerCamelCase		,						__lowerCamelCase  )

 # Check outputs on an image, prepared by MobileNetV2ImageProcessor
 __SCREAMING_SNAKE_CASE							:     Tuple							=						MobileNetVaImageProcessor(
     crop_size={"width": config.image_size, "height": config.image_size}		,						size={"shortest_edge": config.image_size + 32}		,						)
 __SCREAMING_SNAKE_CASE							:     Any							=						image_processor(images=prepare_img()		,						return_tensors="pt"  )
 __SCREAMING_SNAKE_CASE							:     List[str]							=						model(**__lowerCamelCase  )
 __SCREAMING_SNAKE_CASE							:     Any							=						outputs.logits

 if model_name.startswith("deeplabv3_"  ):
  assert logits.shape == (1, 21, 65, 65)

  if model_name == "deeplabv3_mobilenet_v2_1.0_513":
   __SCREAMING_SNAKE_CASE							:     List[Any]							=						torch.tensor(
       [
           [[17.5790, 17.7581, 18.3355], [18.3257, 18.4230, 18.8973], [18.6169, 18.8650, 19.2187]],
           [[-2.1595, -2.0977, -2.3741], [-2.4226, -2.3028, -2.6835], [-2.7819, -2.5991, -2.7706]],
           [[4.2058, 4.8317, 4.7638], [4.4136, 5.0361, 4.9383], [4.5028, 4.9644, 4.8734]],
       ]  )

  else:
   raise ValueError(F"""Unknown model name: {model_name}"""  )

  assert torch.allclose(logits[0, :3, :3, :3]		,						__lowerCamelCase		,						atol=1E-4  )
 else:
  assert logits.shape == (1, 1001)

  if model_name == "mobilenet_v2_1.4_224":
   __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						torch.tensor([0.0181, -1.0015, 0.4688]  )
  elif model_name == "mobilenet_v2_1.0_224":
   __SCREAMING_SNAKE_CASE							:     str							=						torch.tensor([0.2445, -1.1993, 0.1905]  )
  elif model_name == "mobilenet_v2_0.75_160":
   __SCREAMING_SNAKE_CASE							:     int							=						torch.tensor([0.2482, 0.4136, 0.6669]  )
  elif model_name == "mobilenet_v2_0.35_96":
   __SCREAMING_SNAKE_CASE							:     Any							=						torch.tensor([0.1451, -0.4624, 0.7192]  )
  else:
   __SCREAMING_SNAKE_CASE							:     Optional[int]							=						None

  if expected_logits is not None:
   assert torch.allclose(logits[0, :3]		,						__lowerCamelCase		,						atol=1E-4  )

 Path(__lowerCamelCase  ).mkdir(exist_ok=__lowerCamelCase  )
 print(F"""Saving model {model_name} to {pytorch_dump_folder_path}"""  )
 model.save_pretrained(__lowerCamelCase  )
 print(F"""Saving image processor to {pytorch_dump_folder_path}"""  )
 image_processor.save_pretrained(__lowerCamelCase  )

 if push_to_hub:
  print("Pushing to the hub..."  )
  __SCREAMING_SNAKE_CASE							:     Tuple							=						"google/" + model_name
  image_processor.push_to_hub(__lowerCamelCase  )
  model.push_to_hub(__lowerCamelCase  )


if __name__ == "__main__":
      _lowerCamelCase         =	argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          """--model_name""",
          default="""mobilenet_v2_1.0_224""",
          type=str,
          help="""Name of the MobileNetV2 model you'd like to convert. Should in the form 'mobilenet_v2_<depth>_<size>'.""",
      )
      parser.add_argument(
          """--checkpoint_path""", required=True, type=str, help="""Path to the original TensorFlow checkpoint (.ckpt file)."""
      )
      parser.add_argument(
          """--pytorch_dump_folder_path""", required=True, type=str, help="""Path to the output PyTorch model directory."""
      )
      parser.add_argument(
          """--push_to_hub""", action="""store_true""", help="""Whether or not to push the converted model to the 🤗 hub."""
      )

      _lowerCamelCase         =	parser.parse_args()
      convert_movilevit_checkpoint(
          args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub
      )






