


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()
lowerCamelCase : Optional[int]						   =       logging.get_logger(__name__)

lowerCamelCase : int						   =       {
    '''b0''': efficientnet.EfficientNetBa,
    '''b1''': efficientnet.EfficientNetBa,
    '''b2''': efficientnet.EfficientNetBa,
    '''b3''': efficientnet.EfficientNetBa,
    '''b4''': efficientnet.EfficientNetBa,
    '''b5''': efficientnet.EfficientNetBa,
    '''b6''': efficientnet.EfficientNetBa,
    '''b7''': efficientnet.EfficientNetBa,
}

lowerCamelCase : Dict						   =       {
    '''b0''': {
        '''hidden_dim''': 12_80,
        '''width_coef''': 1.0,
        '''depth_coef''': 1.0,
        '''image_size''': 2_24,
        '''dropout_rate''': 0.2,
        '''dw_padding''': [],
    },
    '''b1''': {
        '''hidden_dim''': 12_80,
        '''width_coef''': 1.0,
        '''depth_coef''': 1.1,
        '''image_size''': 2_40,
        '''dropout_rate''': 0.2,
        '''dw_padding''': [16],
    },
    '''b2''': {
        '''hidden_dim''': 14_08,
        '''width_coef''': 1.1,
        '''depth_coef''': 1.2,
        '''image_size''': 2_60,
        '''dropout_rate''': 0.3,
        '''dw_padding''': [5, 8, 16],
    },
    '''b3''': {
        '''hidden_dim''': 15_36,
        '''width_coef''': 1.2,
        '''depth_coef''': 1.4,
        '''image_size''': 3_00,
        '''dropout_rate''': 0.3,
        '''dw_padding''': [5, 18],
    },
    '''b4''': {
        '''hidden_dim''': 17_92,
        '''width_coef''': 1.4,
        '''depth_coef''': 1.8,
        '''image_size''': 3_80,
        '''dropout_rate''': 0.4,
        '''dw_padding''': [6],
    },
    '''b5''': {
        '''hidden_dim''': 20_48,
        '''width_coef''': 1.6,
        '''depth_coef''': 2.2,
        '''image_size''': 4_56,
        '''dropout_rate''': 0.4,
        '''dw_padding''': [13, 27],
    },
    '''b6''': {
        '''hidden_dim''': 23_04,
        '''width_coef''': 1.8,
        '''depth_coef''': 2.6,
        '''image_size''': 5_28,
        '''dropout_rate''': 0.5,
        '''dw_padding''': [31],
    },
    '''b7''': {
        '''hidden_dim''': 25_60,
        '''width_coef''': 2.0,
        '''depth_coef''': 3.1,
        '''image_size''': 6_00,
        '''dropout_rate''': 0.5,
        '''dw_padding''': [18],
    },
}
def       snake_case_		(     lowerCAmelCase_       :						Tuple      ):
							__lowercase      :							Union[str, Any]				=     EfficientNetConfig()
							__lowercase      :							int				=     CONFIG_MAP[model_name]["""hidden_dim"""]
							__lowercase      :							List[Any]				=     CONFIG_MAP[model_name]["""width_coef"""]
							__lowercase      :							Union[str, Any]				=     CONFIG_MAP[model_name]["""depth_coef"""]
							__lowercase      :							List[Any]				=     CONFIG_MAP[model_name]["""image_size"""]
							__lowercase      :							List[Any]				=     CONFIG_MAP[model_name]["""dropout_rate"""]
							__lowercase      :							Dict				=     CONFIG_MAP[model_name]["""dw_padding"""]

							__lowercase      :							int				=     """huggingface/label-files"""
							__lowercase      :							int				=     """imagenet-1k-id2label.json"""
							__lowercase      :							Dict				=     1000
							__lowercase      :							str				=     json.load(open(hf_hub_download(lowerCAmelCase_       ,       lowerCAmelCase_       ,       repo_type="""dataset"""      )       ,       """r"""      )      )
							__lowercase      :							Tuple				=     {int(lowerCAmelCase_      ): v for k, v in idalabel.items()}

							__lowercase      :							Optional[Any]				=     idalabel
							__lowercase      :							str				=     {v: k for k, v in idalabel.items()}
							return config
def       snake_case_		(     ):
							__lowercase      :							Optional[Any]				=     """http://images.cocodataset.org/val2017/000000039769.jpg"""
							__lowercase      :							List[Any]				=     Image.open(requests.get(lowerCAmelCase_       ,       stream=lowerCAmelCase_      ).raw      )
							return im
def       snake_case_		(     lowerCAmelCase_       :						Optional[int]      ):
							__lowercase      :							Tuple				=     CONFIG_MAP[model_name]["""image_size"""]
							__lowercase      :							Any				=     EfficientNetImageProcessor(
							    size={"""height""": size, """width""": size}       ,       image_mean=[0.485, 0.456, 0.406]       ,       image_std=[0.47_853_944, 0.4_732_864, 0.47_434_163]       ,       do_center_crop=lowerCAmelCase_       ,       )
							return preprocessor
def       snake_case_		(     lowerCAmelCase_       :						Optional[Any]      ):
							__lowercase      :							Any				=     [v.split("""_"""      )[0].split("""block"""      )[1] for v in original_param_names if v.startswith("""block"""      )]
							__lowercase      :							Optional[int]				=     sorted(set(lowerCAmelCase_      )      )
							__lowercase      :							Optional[int]				=     len(lowerCAmelCase_      )
							__lowercase      :							Union[str, Any]				=     {b: str(lowerCAmelCase_      ) for b, i in zip(lowerCAmelCase_       ,       range(lowerCAmelCase_      )      )}

							__lowercase      :							Any				=     []
							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:
														__lowercase      :							List[Any]				=     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""")      )

							__lowercase      :							List[str]				=     {}
							for item in rename_keys:
														if item[0] in original_param_names:
																					__lowercase      :							List[str]				=     """efficientnet.""" + item[1]

							__lowercase      :							Optional[Any]				=     """classifier.weight"""
							__lowercase      :							Optional[int]				=     """classifier.bias"""
							return key_mapping
def       snake_case_		(     lowerCAmelCase_       :						Optional[int]       ,       lowerCAmelCase_       :						Tuple       ,       lowerCAmelCase_       :						Dict      ):
							for key, value in tf_params.items():
														if "normalization" in key:
																					continue

														__lowercase      :							str				=     key_mapping[key]
														if "_conv" in key and "kernel" in key:
																					__lowercase      :							List[str]				=     torch.from_numpy(lowerCAmelCase_      ).permute(3       ,       2       ,       0       ,       1      )
														elif "depthwise_kernel" in key:
																					__lowercase      :							int				=     torch.from_numpy(lowerCAmelCase_      ).permute(2       ,       3       ,       0       ,       1      )
														elif "kernel" in key:
																					__lowercase      :							Tuple				=     torch.from_numpy(np.transpose(lowerCAmelCase_      )      )
														else:
																					__lowercase      :							List[Any]				=     torch.from_numpy(lowerCAmelCase_      )

														# Replace HF parameters with original TF model parameters
														assert hf_params[hf_key].shape == new_hf_value.shape
														hf_params[hf_key].copy_(lowerCAmelCase_      )
@torch.no_grad()
def       snake_case_		(     lowerCAmelCase_       :						Optional[Any]       ,       lowerCAmelCase_       :						str       ,       lowerCAmelCase_       :						str       ,       lowerCAmelCase_       :						Dict      ):
							__lowercase      :							int				=     model_classes[model_name](
							    include_top=lowerCAmelCase_       ,       weights="""imagenet"""       ,       input_tensor=lowerCAmelCase_       ,       input_shape=lowerCAmelCase_       ,       pooling=lowerCAmelCase_       ,       classes=1000       ,       classifier_activation="""softmax"""       ,       )

							__lowercase      :							Tuple				=     original_model.trainable_variables
							__lowercase      :							str				=     original_model.non_trainable_variables
							__lowercase      :							Optional[int]				=     {param.name: param.numpy() for param in tf_params}
							for param in tf_non_train_params:
														__lowercase      :							Optional[Any]				=     param.numpy()
							__lowercase      :							Optional[Any]				=     list(tf_params.keys()      )

							# Load HuggingFace model
							__lowercase      :							str				=     get_efficientnet_config(lowerCAmelCase_      )
							__lowercase      :							List[Any]				=     EfficientNetForImageClassification(lowerCAmelCase_      ).eval()
							__lowercase      :							Union[str, Any]				=     hf_model.state_dict()

							# Create src-to-dst parameter name mapping dictionary
							print("""Converting parameters..."""      )
							__lowercase      :							List[str]				=     rename_keys(lowerCAmelCase_      )
							replace_params(lowerCAmelCase_       ,       lowerCAmelCase_       ,       lowerCAmelCase_      )

							# Initialize preprocessor and preprocess input image
							__lowercase      :							Optional[Any]				=     convert_image_processor(lowerCAmelCase_      )
							__lowercase      :							Any				=     preprocessor(images=prepare_img()       ,       return_tensors="""pt"""      )

							# HF model inference
							hf_model.eval()
							with torch.no_grad():
														__lowercase      :							Tuple				=     hf_model(**lowerCAmelCase_      )
							__lowercase      :							List[Any]				=     outputs.logits.detach().numpy()

							# Original model inference
							__lowercase      :							Any				=     False
							__lowercase      :							Optional[Any]				=     CONFIG_MAP[model_name]["""image_size"""]
							__lowercase      :							int				=     prepare_img().resize((image_size, image_size)       ,       resample=PIL.Image.NEAREST      )
							__lowercase      :							Tuple				=     image.img_to_array(lowerCAmelCase_      )
							__lowercase      :							str				=     np.expand_dims(lowerCAmelCase_       ,       axis=0      )
							__lowercase      :							Optional[Any]				=     original_model.predict(lowerCAmelCase_      )

							# Check whether original and HF model outputs match  -> np.allclose
							assert np.allclose(lowerCAmelCase_       ,       lowerCAmelCase_       ,       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(lowerCAmelCase_      ):
																					os.mkdir(lowerCAmelCase_      )
														# Save converted model and image processor
														hf_model.save_pretrained(lowerCAmelCase_      )
														preprocessor.save_pretrained(lowerCAmelCase_      )

							if push_to_hub:
														# Push model and image processor to hub
														print(F"Pushing converted {model_name} to the hub..."      )
														__lowercase      :							List[str]				=     F"efficientnet-{model_name}"
														preprocessor.push_to_hub(lowerCAmelCase_      )
														hf_model.push_to_hub(lowerCAmelCase_      )


if __name__ == "__main__":
	lowerCamelCase : Optional[Any]						   =       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''')

	lowerCamelCase : Any						   =       parser.parse_args()
	convert_efficientnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)







from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


lowerCamelCase : int						   =       {
    '''configuration_funnel''': ['''FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''FunnelConfig'''],
    '''convert_funnel_original_tf_checkpoint_to_pytorch''': [],
    '''tokenization_funnel''': ['''FunnelTokenizer'''],
}

try:
	if not is_tokenizers_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	lowerCamelCase : List[str]						   =       ['''FunnelTokenizerFast''']

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	lowerCamelCase : Union[str, Any]						   =       [
	    '''FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''FunnelBaseModel''',
	    '''FunnelForMaskedLM''',
	    '''FunnelForMultipleChoice''',
	    '''FunnelForPreTraining''',
	    '''FunnelForQuestionAnswering''',
	    '''FunnelForSequenceClassification''',
	    '''FunnelForTokenClassification''',
	    '''FunnelModel''',
	    '''FunnelPreTrainedModel''',
	    '''load_tf_weights_in_funnel''',
	]

try:
	if not is_tf_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	lowerCamelCase : str						   =       [
	    '''TF_FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''TFFunnelBaseModel''',
	    '''TFFunnelForMaskedLM''',
	    '''TFFunnelForMultipleChoice''',
	    '''TFFunnelForPreTraining''',
	    '''TFFunnelForQuestionAnswering''',
	    '''TFFunnelForSequenceClassification''',
	    '''TFFunnelForTokenClassification''',
	    '''TFFunnelModel''',
	    '''TFFunnelPreTrainedModel''',
	]


if TYPE_CHECKING:
	from .configuration_funnel import FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP, FunnelConfig
	from .tokenization_funnel import FunnelTokenizer

	try:
		if not is_tokenizers_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .tokenization_funnel_fast import FunnelTokenizerFast

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_funnel import (
		    FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST,
		    FunnelBaseModel,
		    FunnelForMaskedLM,
		    FunnelForMultipleChoice,
		    FunnelForPreTraining,
		    FunnelForQuestionAnswering,
		    FunnelForSequenceClassification,
		    FunnelForTokenClassification,
		    FunnelModel,
		    FunnelPreTrainedModel,
		    load_tf_weights_in_funnel,
		)

	try:
		if not is_tf_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_tf_funnel import (
		    TF_FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFFunnelBaseModel,
		    TFFunnelForMaskedLM,
		    TFFunnelForMultipleChoice,
		    TFFunnelForPreTraining,
		    TFFunnelForQuestionAnswering,
		    TFFunnelForSequenceClassification,
		    TFFunnelForTokenClassification,
		    TFFunnelModel,
		    TFFunnelPreTrainedModel,
		)

else:
	import sys

	lowerCamelCase : Any						   =       _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)