
"""simple docstring"""





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()
__SCREAMING_SNAKE_CASE					:						Union[str, Any]				= logging.get_logger(__name__)

__SCREAMING_SNAKE_CASE					:						List[str]				= {
    '''b0''': efficientnet.EfficientNetBa,
    '''b1''': efficientnet.EfficientNetBa,
    '''b2''': efficientnet.EfficientNetBa,
    '''b3''': efficientnet.EfficientNetBa,
    '''b4''': efficientnet.EfficientNetBa,
    '''b5''': efficientnet.EfficientNetBa,
    '''b6''': efficientnet.EfficientNetBa,
    '''b7''': efficientnet.EfficientNetBa,
}

__SCREAMING_SNAKE_CASE					:						Optional[Any]				= {
    '''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 							A_     (__a				):

					'''simple docstring'''





					A_				 =						EfficientNetConfig()
					A_				 =						CONFIG_MAP[model_name]["hidden_dim"]
					A_				 =						CONFIG_MAP[model_name]["width_coef"]
					A_				 =						CONFIG_MAP[model_name]["depth_coef"]
					A_				 =						CONFIG_MAP[model_name]["image_size"]
					A_				 =						CONFIG_MAP[model_name]["dropout_rate"]
					A_				 =						CONFIG_MAP[model_name]["dw_padding"]

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

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

					'''simple docstring'''





					A_				 =						"http://images.cocodataset.org/val2017/000000039769.jpg"
					A_				 =						Image.open(requests.get(__a				,   stream=__a				).raw				)
					return im
def 							A_     (__a				):

					'''simple docstring'''





					A_				 =						CONFIG_MAP[model_name]["image_size"]
					A_				 =						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 							A_     (__a				):

					'''simple docstring'''





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

					A_				 =						[]
					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:
										A_				 =						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")				)

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

					A_				 =						"classifier.weight"
					A_				 =						"classifier.bias"
					return key_mapping
def 							A_     (__a				,   __a				,   __a				):

					'''simple docstring'''





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

										A_				 =						key_mapping[key]
										if "_conv" in key and "kernel" in key:
															A_				 =						torch.from_numpy(__a				).permute(3				,   2				,   0				,   1				)
										elif "depthwise_kernel" in key:
															A_				 =						torch.from_numpy(__a				).permute(2				,   3				,   0				,   1				)
										elif "kernel" in key:
															A_				 =						torch.from_numpy(np.transpose(__a				)				)
										else:
															A_				 =						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 							A_     (__a				,   __a				,   __a				,   __a				):

					'''simple docstring'''





					A_				 =						model_classes[model_name](
					    include_top=__a				,   weights="imagenet"				,   input_tensor=__a				,   input_shape=__a				,   pooling=__a				,   classes=1000				,   classifier_activation="softmax"				,   )

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

					# Load HuggingFace model
					A_				 =						get_efficientnet_config(__a				)
					A_				 =						EfficientNetForImageClassification(__a				).eval()
					A_				 =						hf_model.state_dict()

					# Create src-to-dst parameter name mapping dictionary
					print("Converting parameters..."				)
					A_				 =						rename_keys(__a				)
					replace_params(__a				,   __a				,   __a				)

					# Initialize preprocessor and preprocess input image
					A_				 =						convert_image_processor(__a				)
					A_				 =						preprocessor(images=prepare_img()				,   return_tensors="pt"				)

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

					# Original model inference
					A_				 =						False
					A_				 =						CONFIG_MAP[model_name]["image_size"]
					A_				 =						prepare_img().resize((image_size, image_size)				,   resample=PIL.Image.NEAREST				)
					A_				 =						image.img_to_array(__a				)
					A_				 =						np.expand_dims(__a				,   axis=0				)
					A_				 =						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...'				)
										A_				 =						f'efficientnet-{model_name}'
										preprocessor.push_to_hub(__a				)
										hf_model.push_to_hub(__a				)


if __name__ == "__main__":
		__SCREAMING_SNAKE_CASE					:						List[str]				= 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''')

		__SCREAMING_SNAKE_CASE					:						Union[str, Any]				= parser.parse_args()
		convert_efficientnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)






"""simple docstring"""





import inspect
import re

from transformers.utils import direct_transformers_import


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_config_docstrings.py
UpperCamelCase_					:						Optional[Any]				= '''src/transformers'''


# This is to make sure the transformers module imported is the one in the repo.
UpperCamelCase_					:						Union[str, Any]				= direct_transformers_import(PATH_TO_TRANSFORMERS)

UpperCamelCase_					:						List[Any]				= transformers.models.auto.configuration_auto.CONFIG_MAPPING

# Regex pattern used to find the checkpoint mentioned in the docstring of `config_class`.
# For example, `[bert-base-uncased](https://huggingface.co/bert-base-uncased)`
UpperCamelCase_					:						str				= re.compile(R'''\[(.+?)\]\((https://huggingface\.co/.+?)\)''')


UpperCamelCase_					:						Dict				= {
    '''DecisionTransformerConfig''',
    '''EncoderDecoderConfig''',
    '''MusicgenConfig''',
    '''RagConfig''',
    '''SpeechEncoderDecoderConfig''',
    '''TimmBackboneConfig''',
    '''VisionEncoderDecoderConfig''',
    '''VisionTextDualEncoderConfig''',
    '''LlamaConfig''',
}
def 							A_     (__a				):

					'''simple docstring'''





					A_				 =						None

					# source code of `config_class`
					A_				 =						inspect.getsource(__a				)
					A_				 =						_re_checkpoint.findall(__a				)

					# Each `checkpoint` is a tuple of a checkpoint name and a checkpoint link.
					# For example, `('bert-base-uncased', 'https://huggingface.co/bert-base-uncased')`
					for ckpt_name, ckpt_link in checkpoints:
										# allow the link to end with `/`
										if ckpt_link.endswith("/"				):
															A_				 =						ckpt_link[:-1]

										# verify the checkpoint name corresponds to the checkpoint link
										A_				 =						f'https://huggingface.co/{ckpt_name}'
										if ckpt_link == ckpt_link_from_name:
															A_				 =						ckpt_name
															break

					return checkpoint
def 							A_     ():

					'''simple docstring'''





					A_				 =						[]

					for config_class in list(CONFIG_MAPPING.values()				):
										# Skip deprecated models
										if "models.deprecated" in config_class.__module__:
															continue
										A_				 =						get_checkpoint_from_config_class(__a				)

										A_				 =						config_class.__name__
										if checkpoint is None and name not in CONFIG_CLASSES_TO_IGNORE_FOR_DOCSTRING_CHECKPOINT_CHECK:
															configs_without_checkpoint.append(__a				)

					if len(__a				) > 0:
										A_				 =						"\n".join(sorted(__a				)				)
										raise ValueError(f'The following configurations don\'t contain any valid checkpoint:\n{message}'				)


if __name__ == "__main__":
		check_config_docstrings_have_checkpoints()




