

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_       	=	logging.get_logger(__name__)

lowerCAmelCase_       	=	{
    '''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_       	=	{
    '''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   __SCREAMING_SNAKE_CASE					(SCREAMING_SNAKE_CASE__       ):
						snake_case_   =	EfficientNetConfig()
						snake_case_   =	CONFIG_MAP[model_name]['''hidden_dim''']
						snake_case_   =	CONFIG_MAP[model_name]['''width_coef''']
						snake_case_   =	CONFIG_MAP[model_name]['''depth_coef''']
						snake_case_   =	CONFIG_MAP[model_name]['''image_size''']
						snake_case_   =	CONFIG_MAP[model_name]['''dropout_rate''']
						snake_case_   =	CONFIG_MAP[model_name]['''dw_padding''']

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

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


def   __SCREAMING_SNAKE_CASE					():
						snake_case_   =	'''http://images.cocodataset.org/val2017/000000039769.jpg'''
						snake_case_   =	Image.open(requests.get(SCREAMING_SNAKE_CASE__  ,		stream=SCREAMING_SNAKE_CASE__       ).raw       )
						return im


def   __SCREAMING_SNAKE_CASE					(SCREAMING_SNAKE_CASE__       ):
						snake_case_   =	CONFIG_MAP[model_name]['''image_size''']
						snake_case_   =	EfficientNetImageProcessor(
						    size={'''height''': size, '''width''': size}  ,		image_mean=[0.485, 0.456, 0.406]  ,		image_std=[0.4785_3944, 0.473_2864, 0.4743_4163]  ,		do_center_crop=SCREAMING_SNAKE_CASE__  ,		)
						return preprocessor


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

						snake_case_   =	[]
						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:
												snake_case_   =	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''')       )

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

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


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

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

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


@torch.no_grad()
def   __SCREAMING_SNAKE_CASE					(SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__       ):
						snake_case_   =	model_classes[model_name](
						    include_top=SCREAMING_SNAKE_CASE__  ,		weights='''imagenet'''  ,		input_tensor=SCREAMING_SNAKE_CASE__  ,		input_shape=SCREAMING_SNAKE_CASE__  ,		pooling=SCREAMING_SNAKE_CASE__  ,		classes=1000  ,		classifier_activation='''softmax'''  ,		)

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

						# Load HuggingFace model
						snake_case_   =	get_efficientnet_config(SCREAMING_SNAKE_CASE__       )
						snake_case_   =	EfficientNetForImageClassification(SCREAMING_SNAKE_CASE__       ).eval()
						snake_case_   =	hf_model.state_dict()

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

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

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

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

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

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


if __name__ == "__main__":
					lowerCAmelCase_       	=	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_       	=	parser.parse_args()
					convert_efficientnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)




'''simple docstring'''





import argparse
import os
import pickle
import sys

import torch

from transformers import TransfoXLConfig, TransfoXLLMHeadModel, load_tf_weights_in_transfo_xl
from transformers.models.transfo_xl import tokenization_transfo_xl as data_utils
from transformers.models.transfo_xl.tokenization_transfo_xl import CORPUS_NAME, VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()

# We do this to be able to load python 2 datasets pickles
# See e.g. https://stackoverflow.com/questions/2121874/python-pickling-after-changing-a-modules-directory/2121918#2121918
_UpperCamelCase											=						data_utils.TransfoXLTokenizer
_UpperCamelCase											=						data_utils.TransfoXLCorpus
_UpperCamelCase											=						data_utils
_UpperCamelCase											=						data_utils

def        a_     (							_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       )      ->      Optional[Any]:
  if transfo_xl_dataset_file:
    # Convert a pre-processed corpus (see original TensorFlow repo)
    with open(_lowerCAmelCase       ,'rb'       ) as fp:
      __lowerCamelCase       :       Optional[Any]           = pickle.load(_lowerCAmelCase       ,encoding='latin1'       )
    # Save vocabulary and dataset cache as Dictionaries (should be better than pickles for the long-term)
    __lowerCamelCase       :       Tuple           = pytorch_dump_folder_path + '/' + VOCAB_FILES_NAMES['pretrained_vocab_file']
    print(F'Save vocabulary to {pytorch_vocab_dump_path}'       )
    __lowerCamelCase       :       str           = corpus.vocab.__dict__
    torch.save(_lowerCAmelCase       ,_lowerCAmelCase       )

    __lowerCamelCase       :       Optional[Any]           = corpus.__dict__
    corpus_dict_no_vocab.pop('vocab'       ,_lowerCAmelCase       )
    __lowerCamelCase       :       Optional[Any]           = pytorch_dump_folder_path + '/' + CORPUS_NAME
    print(F'Save dataset to {pytorch_dataset_dump_path}'       )
    torch.save(_lowerCAmelCase       ,_lowerCAmelCase       )

  if tf_checkpoint_path:
    # Convert a pre-trained TensorFlow model
    __lowerCamelCase       :       int           = os.path.abspath(_lowerCAmelCase       )
    __lowerCamelCase       :       Any           = os.path.abspath(_lowerCAmelCase       )

    print(F'Converting Transformer XL checkpoint from {tf_path} with config at {config_path}.'       )
    # Initialise PyTorch model
    if transfo_xl_config_file == "":
      __lowerCamelCase       :       Optional[int]           = TransfoXLConfig()
    else:
      __lowerCamelCase       :       Optional[int]           = TransfoXLConfig.from_json_file(_lowerCAmelCase       )
    print(F'Building PyTorch model from configuration: {config}'       )
    __lowerCamelCase       :       List[str]           = TransfoXLLMHeadModel(_lowerCAmelCase       )

    __lowerCamelCase       :       Dict           = load_tf_weights_in_transfo_xl(_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       )
    # Save pytorch-model
    __lowerCamelCase       :       List[str]           = os.path.join(_lowerCAmelCase       ,_lowerCAmelCase       )
    __lowerCamelCase       :       int           = os.path.join(_lowerCAmelCase       ,_lowerCAmelCase       )
    print(F'Save PyTorch model to {os.path.abspath(_lowerCAmelCase       )}'       )
    torch.save(model.state_dict()       ,_lowerCAmelCase       )
    print(F'Save configuration file to {os.path.abspath(_lowerCAmelCase       )}'       )
    with open(_lowerCAmelCase       ,'w'       ,encoding='utf-8'       ) as f:
      f.write(config.to_json_string()       )


if __name__ == "__main__":
   _UpperCamelCase											=						argparse.ArgumentParser()
   parser.add_argument(
       '--pytorch_dump_folder_path',
       default=None,
       type=str,
       required=True,
       help='Path to the folder to store the PyTorch model or dataset/vocab.',
   )
   parser.add_argument(
       '--tf_checkpoint_path',
       default='',
       type=str,
       help='An optional path to a TensorFlow checkpoint path to be converted.',
   )
   parser.add_argument(
       '--transfo_xl_config_file',
       default='',
       type=str,
       help=(
           'An optional config json file corresponding to the pre-trained BERT model. \n'
           'This specifies the model architecture.'
       ),
   )
   parser.add_argument(
       '--transfo_xl_dataset_file',
       default='',
       type=str,
       help='An optional dataset file to be converted in a vocabulary.',
   )
   _UpperCamelCase											=						parser.parse_args()
   convert_transfo_xl_checkpoint_to_pytorch(
       args.tf_checkpoint_path,
       args.transfo_xl_config_file,
       args.pytorch_dump_folder_path,
       args.transfo_xl_dataset_file,
   )






