

import os
import tempfile
import unittest

from transformers import NezhaConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
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,
				    NezhaForMaskedLM,
				    NezhaForMultipleChoice,
				    NezhaForNextSentencePrediction,
				    NezhaForPreTraining,
				    NezhaForQuestionAnswering,
				    NezhaForSequenceClassification,
				    NezhaForTokenClassification,
				    NezhaModel,
				)
				from transformers.models.nezha.modeling_nezha import NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST

class 				__UpperCamelCase				:







	def __init__(   self: int  , __UpperCamelCase: List[Any]  , __UpperCamelCase: Optional[Any]=13  , __UpperCamelCase: List[Any]=7  , __UpperCamelCase: List[str]=True  , __UpperCamelCase: Any=True  , __UpperCamelCase: List[Any]=True  , __UpperCamelCase: Dict=True  , __UpperCamelCase: str=99  , __UpperCamelCase: List[str]=32  , __UpperCamelCase: Union[str, Any]=5  , __UpperCamelCase: Dict=4  , __UpperCamelCase: List[str]=37  , __UpperCamelCase: Tuple="gelu"  , __UpperCamelCase: Optional[int]=0.1  , __UpperCamelCase: Union[str, Any]=0.1  , __UpperCamelCase: Any=1_28  , __UpperCamelCase: int=32  , __UpperCamelCase: Dict=16  , __UpperCamelCase: Union[str, Any]=2  , __UpperCamelCase: List[str]=0.02  , __UpperCamelCase: Union[str, Any]=3  , __UpperCamelCase: int=4  , __UpperCamelCase: Tuple=None  , ):

						'''simple docstring'''





						__magic_name__							=    parent
						__magic_name__							=    batch_size
						__magic_name__							=    seq_length
						__magic_name__							=    is_training
						__magic_name__							=    use_input_mask
						__magic_name__							=    use_token_type_ids
						__magic_name__							=    use_labels
						__magic_name__							=    vocab_size
						__magic_name__							=    hidden_size
						__magic_name__							=    num_hidden_layers
						__magic_name__							=    num_attention_heads
						__magic_name__							=    intermediate_size
						__magic_name__							=    hidden_act
						__magic_name__							=    hidden_dropout_prob
						__magic_name__							=    attention_probs_dropout_prob
						__magic_name__							=    max_position_embeddings
						__magic_name__							=    type_vocab_size
						__magic_name__							=    type_sequence_label_size
						__magic_name__							=    initializer_range
						__magic_name__							=    num_labels
						__magic_name__							=    num_choices
						__magic_name__							=    scope







	def 	_SCREAMING_SNAKE_CASE							(   self: List[Any]      ):

						'''simple docstring'''





						__magic_name__							=    ids_tensor([self.batch_size, self.seq_length]  , self.vocab_size      )

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

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

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

						__magic_name__							=    self.get_config()

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







	def 	_SCREAMING_SNAKE_CASE							(   self: Any      ):

						'''simple docstring'''





						return NezhaConfig(
						    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=lowerCAmelCase_  , initializer_range=self.initializer_range  , )







	def 	_SCREAMING_SNAKE_CASE							(   self: Tuple      ):

						'''simple docstring'''





						(
						    (
						    __magic_name__    
						),	(
						    __magic_name__    
						),	(
						    __magic_name__    
						),	(
						    __magic_name__    
						),	(
						    __magic_name__    
						),	(
						    __magic_name__    
						),	(
						    __magic_name__    
						),	    
						)							=    self.prepare_config_and_inputs()

						__magic_name__							=    True
						__magic_name__							=    floats_tensor([self.batch_size, self.seq_length, self.hidden_size]      )
						__magic_name__							=    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 	_SCREAMING_SNAKE_CASE							(   self: Dict  , __UpperCamelCase: Dict  , __UpperCamelCase: int  , __UpperCamelCase: Union[str, Any]  , __UpperCamelCase: str  , __UpperCamelCase: str  , __UpperCamelCase: int  , __UpperCamelCase: int      ):

						'''simple docstring'''





						__magic_name__							=    NezhaModel(config=lowerCAmelCase_      )
						model.to(lowerCAmelCase_      )
						model.eval()
						__magic_name__							=    model(lowerCAmelCase_  , attention_mask=lowerCAmelCase_  , token_type_ids=lowerCAmelCase_      )
						__magic_name__							=    model(lowerCAmelCase_  , token_type_ids=lowerCAmelCase_      )
						__magic_name__							=    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 	_SCREAMING_SNAKE_CASE							(   self: Optional[int]  , __UpperCamelCase: Dict  , __UpperCamelCase: Union[str, Any]  , __UpperCamelCase: List[Any]  , __UpperCamelCase: Optional[Any]  , __UpperCamelCase: str  , __UpperCamelCase: int  , __UpperCamelCase: List[Any]  , __UpperCamelCase: Union[str, Any]  , __UpperCamelCase: Tuple  , ):

						'''simple docstring'''





						__magic_name__							=    True
						__magic_name__							=    NezhaModel(lowerCAmelCase_      )
						model.to(lowerCAmelCase_      )
						model.eval()
						__magic_name__							=    model(
						    lowerCAmelCase_  , attention_mask=lowerCAmelCase_  , token_type_ids=lowerCAmelCase_  , encoder_hidden_states=lowerCAmelCase_  , encoder_attention_mask=lowerCAmelCase_  , )
						__magic_name__							=    model(
						    lowerCAmelCase_  , attention_mask=lowerCAmelCase_  , token_type_ids=lowerCAmelCase_  , encoder_hidden_states=lowerCAmelCase_  , )
						__magic_name__							=    model(lowerCAmelCase_  , attention_mask=lowerCAmelCase_  , token_type_ids=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 	_SCREAMING_SNAKE_CASE							(   self: Optional[Any]  , __UpperCamelCase: Any  , __UpperCamelCase: List[str]  , __UpperCamelCase: Any  , __UpperCamelCase: int  , __UpperCamelCase: Union[str, Any]  , __UpperCamelCase: Any  , __UpperCamelCase: int      ):

						'''simple docstring'''





						__magic_name__							=    NezhaForMaskedLM(config=lowerCAmelCase_      )
						model.to(lowerCAmelCase_      )
						model.eval()
						__magic_name__							=    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 	_SCREAMING_SNAKE_CASE							(   self: Any  , __UpperCamelCase: Union[str, Any]  , __UpperCamelCase: Dict  , __UpperCamelCase: Any  , __UpperCamelCase: str  , __UpperCamelCase: Optional[Any]  , __UpperCamelCase: Tuple  , __UpperCamelCase: Any      ):

						'''simple docstring'''





						__magic_name__							=    NezhaForNextSentencePrediction(config=lowerCAmelCase_      )
						model.to(lowerCAmelCase_      )
						model.eval()
						__magic_name__							=    model(
						    lowerCAmelCase_  , attention_mask=lowerCAmelCase_  , token_type_ids=lowerCAmelCase_  , labels=lowerCAmelCase_  , )
						self.parent.assertEqual(result.logits.shape  , (self.batch_size, 2)      )







	def 	_SCREAMING_SNAKE_CASE							(   self: Any  , __UpperCamelCase: List[str]  , __UpperCamelCase: int  , __UpperCamelCase: Any  , __UpperCamelCase: Union[str, Any]  , __UpperCamelCase: List[Any]  , __UpperCamelCase: Union[str, Any]  , __UpperCamelCase: Union[str, Any]      ):

						'''simple docstring'''





						__magic_name__							=    NezhaForPreTraining(config=lowerCAmelCase_      )
						model.to(lowerCAmelCase_      )
						model.eval()
						__magic_name__							=    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 	_SCREAMING_SNAKE_CASE							(   self: Dict  , __UpperCamelCase: Dict  , __UpperCamelCase: Dict  , __UpperCamelCase: Any  , __UpperCamelCase: Optional[int]  , __UpperCamelCase: Any  , __UpperCamelCase: int  , __UpperCamelCase: Any      ):

						'''simple docstring'''





						__magic_name__							=    NezhaForQuestionAnswering(config=lowerCAmelCase_      )
						model.to(lowerCAmelCase_      )
						model.eval()
						__magic_name__							=    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 	_SCREAMING_SNAKE_CASE							(   self: Optional[int]  , __UpperCamelCase: Tuple  , __UpperCamelCase: Dict  , __UpperCamelCase: Tuple  , __UpperCamelCase: List[str]  , __UpperCamelCase: Dict  , __UpperCamelCase: str  , __UpperCamelCase: int      ):

						'''simple docstring'''





						__magic_name__							=    self.num_labels
						__magic_name__							=    NezhaForSequenceClassification(lowerCAmelCase_      )
						model.to(lowerCAmelCase_      )
						model.eval()
						__magic_name__							=    model(lowerCAmelCase_  , attention_mask=lowerCAmelCase_  , token_type_ids=lowerCAmelCase_  , labels=lowerCAmelCase_      )
						self.parent.assertEqual(result.logits.shape  , (self.batch_size, self.num_labels)      )







	def 	_SCREAMING_SNAKE_CASE							(   self: Tuple  , __UpperCamelCase: int  , __UpperCamelCase: List[str]  , __UpperCamelCase: List[Any]  , __UpperCamelCase: Tuple  , __UpperCamelCase: List[str]  , __UpperCamelCase: int  , __UpperCamelCase: Optional[int]      ):

						'''simple docstring'''





						__magic_name__							=    self.num_labels
						__magic_name__							=    NezhaForTokenClassification(config=lowerCAmelCase_      )
						model.to(lowerCAmelCase_      )
						model.eval()
						__magic_name__							=    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 	_SCREAMING_SNAKE_CASE							(   self: Tuple  , __UpperCamelCase: Optional[Any]  , __UpperCamelCase: Tuple  , __UpperCamelCase: Optional[Any]  , __UpperCamelCase: Optional[Any]  , __UpperCamelCase: Dict  , __UpperCamelCase: Tuple  , __UpperCamelCase: Optional[Any]      ):

						'''simple docstring'''





						__magic_name__							=    self.num_choices
						__magic_name__							=    NezhaForMultipleChoice(config=lowerCAmelCase_      )
						model.to(lowerCAmelCase_      )
						model.eval()
						__magic_name__							=    input_ids.unsqueeze(1      ).expand(-1  , self.num_choices  , -1      ).contiguous()
						__magic_name__							=    token_type_ids.unsqueeze(1      ).expand(-1  , self.num_choices  , -1      ).contiguous()
						__magic_name__							=    input_mask.unsqueeze(1      ).expand(-1  , self.num_choices  , -1      ).contiguous()
						__magic_name__							=    model(
						    lowerCAmelCase_  , attention_mask=lowerCAmelCase_  , token_type_ids=lowerCAmelCase_  , labels=lowerCAmelCase_  , )
						self.parent.assertEqual(result.logits.shape  , (self.batch_size, self.num_choices)      )







	def 	_SCREAMING_SNAKE_CASE							(   self: Optional[Any]      ):

						'''simple docstring'''





						__magic_name__							=    self.prepare_config_and_inputs()
						(
						    (
						    __magic_name__    
						),	(
						    __magic_name__    
						),	(
						    __magic_name__    
						),	(
						    __magic_name__    
						),	(
						    __magic_name__    
						),	(
						    __magic_name__    
						),	(
						    __magic_name__    
						),	    
						)							=    config_and_inputs
						__magic_name__							=    {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
						return config, inputs_dict




@require_torch
class 				__UpperCamelCase				(					__UpperCAmelCase  ,						__UpperCAmelCase  ,						__UpperCAmelCase  ,						unittest.TestCase			):
	_lowercase  :   Any       				=  (
	    (
	        NezhaModel,
	        NezhaForMaskedLM,
	        NezhaForMultipleChoice,
	        NezhaForNextSentencePrediction,
	        NezhaForPreTraining,
	        NezhaForQuestionAnswering,
	        NezhaForSequenceClassification,
	        NezhaForTokenClassification,
	    )
	    if is_torch_available()
	    else ()
	)
	_lowercase  :   str       				=  (
	    {
	        "feature-extraction": NezhaModel,
	        "fill-mask": NezhaForMaskedLM,
	        "question-answering": NezhaForQuestionAnswering,
	        "text-classification": NezhaForSequenceClassification,
	        "token-classification": NezhaForTokenClassification,
	        "zero-shot": NezhaForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	_lowercase  :   Optional[Any]       				=  True







	def 	_SCREAMING_SNAKE_CASE							(   self: int  , __UpperCamelCase: str  , __UpperCamelCase: int  , __UpperCamelCase: int=False      ):

						'''simple docstring'''





						__magic_name__							=    super()._prepare_for_class(lowerCAmelCase_  , lowerCAmelCase_  , return_labels=lowerCAmelCase_      )

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







	def 	_SCREAMING_SNAKE_CASE							(   self: Optional[int]      ):

						'''simple docstring'''





						__magic_name__							=    NezhaModelTester(self      )
						__magic_name__							=    ConfigTester(self  , config_class=lowerCAmelCase_  , hidden_size=37      )







	def 	_SCREAMING_SNAKE_CASE							(   self: Union[str, Any]      ):

						'''simple docstring'''





						self.config_tester.run_common_tests()







	def 	_SCREAMING_SNAKE_CASE							(   self: List[str]      ):

						'''simple docstring'''





						__magic_name__							=    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*lowerCAmelCase_      )







	def 	_SCREAMING_SNAKE_CASE							(   self: str      ):

						'''simple docstring'''





						__magic_name__							=    self.model_tester.prepare_config_and_inputs_for_decoder()
						self.model_tester.create_and_check_model_as_decoder(*lowerCAmelCase_      )







	def 	_SCREAMING_SNAKE_CASE							(   self: str      ):

						'''simple docstring'''





						(
						    (
						    __magic_name__    
						),	(
						    __magic_name__    
						),	(
						    __magic_name__    
						),	(
						    __magic_name__    
						),	(
						    __magic_name__    
						),	(
						    __magic_name__    
						),	(
						    __magic_name__    
						),	(
						    __magic_name__    
						),	(
						    __magic_name__    
						),	    
						)							=    self.model_tester.prepare_config_and_inputs_for_decoder()

						__magic_name__							=    None

						self.model_tester.create_and_check_model_as_decoder(
						    lowerCAmelCase_  , lowerCAmelCase_  , lowerCAmelCase_  , lowerCAmelCase_  , lowerCAmelCase_  , lowerCAmelCase_  , lowerCAmelCase_  , lowerCAmelCase_  , lowerCAmelCase_  , )







	def 	_SCREAMING_SNAKE_CASE							(   self: List[Any]      ):

						'''simple docstring'''





						__magic_name__							=    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*lowerCAmelCase_      )







	def 	_SCREAMING_SNAKE_CASE							(   self: int      ):

						'''simple docstring'''





						__magic_name__							=    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_multiple_choice(*lowerCAmelCase_      )







	def 	_SCREAMING_SNAKE_CASE							(   self: int      ):

						'''simple docstring'''





						__magic_name__							=    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_next_sequence_prediction(*lowerCAmelCase_      )







	def 	_SCREAMING_SNAKE_CASE							(   self: Optional[Any]      ):

						'''simple docstring'''





						__magic_name__							=    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_pretraining(*lowerCAmelCase_      )







	def 	_SCREAMING_SNAKE_CASE							(   self: Any      ):

						'''simple docstring'''





						__magic_name__							=    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_question_answering(*lowerCAmelCase_      )







	def 	_SCREAMING_SNAKE_CASE							(   self: Tuple      ):

						'''simple docstring'''





						__magic_name__							=    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_sequence_classification(*lowerCAmelCase_      )







	def 	_SCREAMING_SNAKE_CASE							(   self: Union[str, Any]      ):

						'''simple docstring'''





						__magic_name__							=    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_token_classification(*lowerCAmelCase_      )







	@slow
	def 	_SCREAMING_SNAKE_CASE							(   self: int      ):

						'''simple docstring'''





						for model_name in NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											__magic_name__							=    NezhaModel.from_pretrained(lowerCAmelCase_      )
											self.assertIsNotNone(lowerCAmelCase_      )







	@slow
	@require_torch_gpu
	def 	_SCREAMING_SNAKE_CASE							(   self: Union[str, Any]      ):

						'''simple docstring'''





						__magic_name__,	__magic_name__							=    self.model_tester.prepare_config_and_inputs_for_common()
						for model_class in self.all_model_classes:
											# NezhaForMultipleChoice behaves incorrectly in JIT environments.
											if model_class == NezhaForMultipleChoice:
																return

											__magic_name__							=    True
											__magic_name__							=    model_class(config=lowerCAmelCase_      )

											__magic_name__							=    self._prepare_for_class(lowerCAmelCase_  , lowerCAmelCase_      )
											__magic_name__							=    torch.jit.trace(
											    lowerCAmelCase_  , (inputs_dict['input_ids'].to('cpu'      ), inputs_dict['attention_mask'].to('cpu'      ))      )

											with tempfile.TemporaryDirectory() as tmp:
																torch.jit.save(lowerCAmelCase_  , os.path.join(lowerCAmelCase_  , 'bert.pt'      )      )
																__magic_name__							=    torch.jit.load(os.path.join(lowerCAmelCase_  , 'bert.pt'      )  , map_location=lowerCAmelCase_      )
																loaded(inputs_dict['input_ids'].to(lowerCAmelCase_      )  , inputs_dict['attention_mask'].to(lowerCAmelCase_      )      )




@require_torch
class 				__UpperCamelCase				(					unittest.TestCase			):







	@slow
	def 	_SCREAMING_SNAKE_CASE							(   self: Any      ):

						'''simple docstring'''





						__magic_name__							=    NezhaModel.from_pretrained('sijunhe/nezha-cn-base'      )
						__magic_name__							=    torch.tensor([[0, 1, 2, 3, 4, 5]]      )
						__magic_name__							=    torch.tensor([[0, 1, 1, 1, 1, 1]]      )
						with torch.no_grad():
											__magic_name__							=    model(lowerCAmelCase_  , attention_mask=lowerCAmelCase_      )[0]
						__magic_name__							=    torch.Size((1, 6, 7_68)      )
						self.assertEqual(output.shape  , lowerCAmelCase_      )
						__magic_name__							=    torch.tensor([[[0.0685, 0.2441, 0.1102], [0.0600, 0.1906, 0.1349], [0.0221, 0.0819, 0.0586]]]      )

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







	@slow
	def 	_SCREAMING_SNAKE_CASE							(   self: int      ):

						'''simple docstring'''





						__magic_name__							=    NezhaForMaskedLM.from_pretrained('sijunhe/nezha-cn-base'      )
						__magic_name__							=    torch.tensor([[0, 1, 2, 3, 4, 5]]      )
						__magic_name__							=    torch.tensor([[1, 1, 1, 1, 1, 1]]      )
						with torch.no_grad():
											__magic_name__							=    model(lowerCAmelCase_  , attention_mask=lowerCAmelCase_      )[0]
						__magic_name__							=    torch.Size((1, 6, 2_11_28)      )
						self.assertEqual(output.shape  , lowerCAmelCase_      )
						__magic_name__							=    torch.tensor(
						    [[-2.7939, -1.7902, -2.2189], [-2.8585, -1.8908, -2.3723], [-2.6499, -1.7750, -2.2558]]      )

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







import argparse
import json
import os

import numpy as np
import PIL
import requests
import tensorflow.keras.applications.efficientnet as efficientnet
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from tensorflow.keras.preprocessing import image

from transformers import (
    EfficientNetConfig,
    EfficientNetForImageClassification,
    EfficientNetImageProcessor,
)
from transformers.utils import logging


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

A__						=							{
    "b0": efficientnet.EfficientNetBa,
    "b1": efficientnet.EfficientNetBa,
    "b2": efficientnet.EfficientNetBa,
    "b3": efficientnet.EfficientNetBa,
    "b4": efficientnet.EfficientNetBa,
    "b5": efficientnet.EfficientNetBa,
    "b6": efficientnet.EfficientNetBa,
    "b7": efficientnet.EfficientNetBa,
}

A__						=							{
    "b0": {
        "hidden_dim": 1280,
        "width_coef": 1.0,
        "depth_coef": 1.0,
        "image_size": 224,
        "dropout_rate": 0.2,
        "dw_padding": [],
    },
    "b1": {
        "hidden_dim": 1280,
        "width_coef": 1.0,
        "depth_coef": 1.1,
        "image_size": 240,
        "dropout_rate": 0.2,
        "dw_padding": [16],
    },
    "b2": {
        "hidden_dim": 1408,
        "width_coef": 1.1,
        "depth_coef": 1.2,
        "image_size": 260,
        "dropout_rate": 0.3,
        "dw_padding": [5, 8, 16],
    },
    "b3": {
        "hidden_dim": 1536,
        "width_coef": 1.2,
        "depth_coef": 1.4,
        "image_size": 300,
        "dropout_rate": 0.3,
        "dw_padding": [5, 18],
    },
    "b4": {
        "hidden_dim": 1792,
        "width_coef": 1.4,
        "depth_coef": 1.8,
        "image_size": 380,
        "dropout_rate": 0.4,
        "dw_padding": [6],
    },
    "b5": {
        "hidden_dim": 2048,
        "width_coef": 1.6,
        "depth_coef": 2.2,
        "image_size": 456,
        "dropout_rate": 0.4,
        "dw_padding": [13, 27],
    },
    "b6": {
        "hidden_dim": 2304,
        "width_coef": 1.8,
        "depth_coef": 2.6,
        "image_size": 528,
        "dropout_rate": 0.5,
        "dw_padding": [31],
    },
    "b7": {
        "hidden_dim": 2560,
        "width_coef": 2.0,
        "depth_coef": 3.1,
        "image_size": 600,
        "dropout_rate": 0.5,
        "dw_padding": [18],
    },
}




def 					_lowercase     (				a_				:					List[Any]				)				->	Any:




					'''simple docstring'''







					__magic_name__							=    EfficientNetConfig()
					__magic_name__							=    CONFIG_MAP[model_name]['hidden_dim']
					__magic_name__							=    CONFIG_MAP[model_name]['width_coef']
					__magic_name__							=    CONFIG_MAP[model_name]['depth_coef']
					__magic_name__							=    CONFIG_MAP[model_name]['image_size']
					__magic_name__							=    CONFIG_MAP[model_name]['dropout_rate']
					__magic_name__							=    CONFIG_MAP[model_name]['dw_padding']

					__magic_name__							=    'huggingface/label-files'
					__magic_name__							=    'imagenet-1k-id2label.json'
					__magic_name__							=    1_0_0_0
					__magic_name__							=    json.load(open(hf_hub_download(a_					,a_					,repo_type='dataset'				)					,'r'				)				)
					__magic_name__							=    {int(a_				): v for k, v in idalabel.items()}

					__magic_name__							=    idalabel
					__magic_name__							=    {v: k for k, v in idalabel.items()}
					return config




def 					_lowercase     (				)				->	int:




					'''simple docstring'''







					__magic_name__							=    'http://images.cocodataset.org/val2017/000000039769.jpg'
					__magic_name__							=    Image.open(requests.get(a_					,stream=a_				).raw				)
					return im




def 					_lowercase     (				a_				:					Optional[int]				)				->	Optional[Any]:




					'''simple docstring'''







					__magic_name__							=    CONFIG_MAP[model_name]['image_size']
					__magic_name__							=    EfficientNetImageProcessor(
					    size={'height': size, 'width': size}					,image_mean=[0.485, 0.456, 0.406]					,image_std=[0.47853944, 0.4732864, 0.47434163]					,do_center_crop=a_					,)
					return preprocessor




def 					_lowercase     (				a_				:					List[str]				)				->	List[Any]:




					'''simple docstring'''







					__magic_name__							=    [v.split('_'				)[0].split('block'				)[1] for v in original_param_names if v.startswith('block'				)]
					__magic_name__							=    sorted(set(a_				)				)
					__magic_name__							=    len(a_				)
					__magic_name__							=    {b: str(a_				) for b, i in zip(a_					,range(a_				)				)}

					__magic_name__							=    []
					rename_keys.append(('stem_conv/kernel:0', 'embeddings.convolution.weight')				)
					rename_keys.append(('stem_bn/gamma:0', 'embeddings.batchnorm.weight')				)
					rename_keys.append(('stem_bn/beta:0', 'embeddings.batchnorm.bias')				)
					rename_keys.append(('stem_bn/moving_mean:0', 'embeddings.batchnorm.running_mean')				)
					rename_keys.append(('stem_bn/moving_variance:0', 'embeddings.batchnorm.running_var')				)

					for b in block_names:
										__magic_name__							=    block_name_mapping[b]
										rename_keys.append((F'block{b}_expand_conv/kernel:0', F'encoder.blocks.{hf_b}.expansion.expand_conv.weight')				)
										rename_keys.append((F'block{b}_expand_bn/gamma:0', F'encoder.blocks.{hf_b}.expansion.expand_bn.weight')				)
										rename_keys.append((F'block{b}_expand_bn/beta:0', F'encoder.blocks.{hf_b}.expansion.expand_bn.bias')				)
										rename_keys.append(
										    (F'block{b}_expand_bn/moving_mean:0', F'encoder.blocks.{hf_b}.expansion.expand_bn.running_mean')				)
										rename_keys.append(
										    (F'block{b}_expand_bn/moving_variance:0', F'encoder.blocks.{hf_b}.expansion.expand_bn.running_var')				)
										rename_keys.append(
										    (F'block{b}_dwconv/depthwise_kernel:0', F'encoder.blocks.{hf_b}.depthwise_conv.depthwise_conv.weight')				)
										rename_keys.append((F'block{b}_bn/gamma:0', F'encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.weight')				)
										rename_keys.append((F'block{b}_bn/beta:0', F'encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.bias')				)
										rename_keys.append(
										    (F'block{b}_bn/moving_mean:0', F'encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_mean')				)
										rename_keys.append(
										    (F'block{b}_bn/moving_variance:0', F'encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_var')				)

										rename_keys.append((F'block{b}_se_reduce/kernel:0', F'encoder.blocks.{hf_b}.squeeze_excite.reduce.weight')				)
										rename_keys.append((F'block{b}_se_reduce/bias:0', F'encoder.blocks.{hf_b}.squeeze_excite.reduce.bias')				)
										rename_keys.append((F'block{b}_se_expand/kernel:0', F'encoder.blocks.{hf_b}.squeeze_excite.expand.weight')				)
										rename_keys.append((F'block{b}_se_expand/bias:0', F'encoder.blocks.{hf_b}.squeeze_excite.expand.bias')				)
										rename_keys.append(
										    (F'block{b}_project_conv/kernel:0', F'encoder.blocks.{hf_b}.projection.project_conv.weight')				)
										rename_keys.append((F'block{b}_project_bn/gamma:0', F'encoder.blocks.{hf_b}.projection.project_bn.weight')				)
										rename_keys.append((F'block{b}_project_bn/beta:0', F'encoder.blocks.{hf_b}.projection.project_bn.bias')				)
										rename_keys.append(
										    (F'block{b}_project_bn/moving_mean:0', F'encoder.blocks.{hf_b}.projection.project_bn.running_mean')				)
										rename_keys.append(
										    (F'block{b}_project_bn/moving_variance:0', F'encoder.blocks.{hf_b}.projection.project_bn.running_var')				)

					rename_keys.append(('top_conv/kernel:0', 'encoder.top_conv.weight')				)
					rename_keys.append(('top_bn/gamma:0', 'encoder.top_bn.weight')				)
					rename_keys.append(('top_bn/beta:0', 'encoder.top_bn.bias')				)
					rename_keys.append(('top_bn/moving_mean:0', 'encoder.top_bn.running_mean')				)
					rename_keys.append(('top_bn/moving_variance:0', 'encoder.top_bn.running_var')				)

					__magic_name__							=    {}
					for item in rename_keys:
										if item[0] in original_param_names:
															__magic_name__							=    'efficientnet.' + item[1]

					__magic_name__							=    'classifier.weight'
					__magic_name__							=    'classifier.bias'
					return key_mapping




def 					_lowercase     (				a_				:					Optional[Any]					,a_				:					Optional[Any]					,a_				:					List[Any]				)				->	Any:




					'''simple docstring'''







					for key, value in tf_params.items():
										if "normalization" in key:
															continue

										__magic_name__							=    key_mapping[key]
										if "_conv" in key and "kernel" in key:
															__magic_name__							=    torch.from_numpy(a_				).permute(3					,2					,0					,1				)
										elif "depthwise_kernel" in key:
															__magic_name__							=    torch.from_numpy(a_				).permute(2					,3					,0					,1				)
										elif "kernel" in key:
															__magic_name__							=    torch.from_numpy(np.transpose(a_				)				)
										else:
															__magic_name__							=    torch.from_numpy(a_				)

										# Replace HF parameters with original TF model parameters
										assert hf_params[hf_key].shape == new_hf_value.shape
										hf_params[hf_key].copy_(a_				)




@torch.no_grad()
def 					_lowercase     (				a_				:					Dict					,a_				:					Optional[Any]					,a_				:					Any					,a_				:					Dict				)				->	List[Any]:




					'''simple docstring'''







					__magic_name__							=    model_classes[model_name](
					    include_top=a_					,weights='imagenet'					,input_tensor=a_					,input_shape=a_					,pooling=a_					,classes=1_0_0_0					,classifier_activation='softmax'					,)

					__magic_name__							=    original_model.trainable_variables
					__magic_name__							=    original_model.non_trainable_variables
					__magic_name__							=    {param.name: param.numpy() for param in tf_params}
					for param in tf_non_train_params:
										__magic_name__							=    param.numpy()
					__magic_name__							=    list(tf_params.keys()				)

					# Load HuggingFace model
					__magic_name__							=    get_efficientnet_config(a_				)
					__magic_name__							=    EfficientNetForImageClassification(a_				).eval()
					__magic_name__							=    hf_model.state_dict()

					# Create src-to-dst parameter name mapping dictionary
					print('Converting parameters...'				)
					__magic_name__							=    rename_keys(a_				)
					replace_params(a_					,a_					,a_				)

					# Initialize preprocessor and preprocess input image
					__magic_name__							=    convert_image_processor(a_				)
					__magic_name__							=    preprocessor(images=prepare_img()					,return_tensors='pt'				)

					# HF model inference
					hf_model.eval()
					with torch.no_grad():
										__magic_name__							=    hf_model(**a_				)
					__magic_name__							=    outputs.logits.detach().numpy()

					# Original model inference
					__magic_name__							=    False
					__magic_name__							=    CONFIG_MAP[model_name]['image_size']
					__magic_name__							=    prepare_img().resize((image_size, image_size)					,resample=PIL.Image.NEAREST				)
					__magic_name__							=    image.img_to_array(a_				)
					__magic_name__							=    np.expand_dims(a_					,axis=0				)
					__magic_name__							=    original_model.predict(a_				)

					# Check whether original and HF model outputs match  -> np.allclose
					assert np.allclose(a_					,a_					,atol=1e-3				), "The predicted logits are not the same."
					print('Model outputs match!'				)

					if save_model:
										# Create folder to save model
										if not os.path.isdir(a_				):
															os.mkdir(a_				)
										# Save converted model and image processor
										hf_model.save_pretrained(a_				)
										preprocessor.save_pretrained(a_				)

					if push_to_hub:
										# Push model and image processor to hub
										print(F'Pushing converted {model_name} to the hub...'				)
										__magic_name__							=    F'efficientnet-{model_name}'
										preprocessor.push_to_hub(a_				)
										hf_model.push_to_hub(a_				)


if __name__ == "__main__":
						A__						=							argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    "--model_name",
						    default="b0",
						    type=str,
						    help="Version name of the EfficientNet model you want to convert, select from [b0, b1, b2, b3, b4, b5, b6, b7].",
						)
						parser.add_argument(
						    "--pytorch_dump_folder_path",
						    default="hf_model",
						    type=str,
						    help="Path to the output PyTorch model directory.",
						)
						parser.add_argument("--save_model", action="store_true", help="Save model to local")
						parser.add_argument("--push_to_hub", action="store_true", help="Push model and image processor to the hub")

						A__						=							parser.parse_args()
						convert_efficientnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)




