






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

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

UpperCamelCase_	:				Dict										= {
    '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 :				Optional[int] )					->					Optional[Any]:
      __snake_case      :    Optional[Any]  			=EfficientNetConfig()
      __snake_case      :    str  			=CONFIG_MAP[model_name]['''hidden_dim''']
      __snake_case      :    Any  			=CONFIG_MAP[model_name]['''width_coef''']
      __snake_case      :    List[str]  			=CONFIG_MAP[model_name]['''depth_coef''']
      __snake_case      :    str  			=CONFIG_MAP[model_name]['''image_size''']
      __snake_case      :    int  			=CONFIG_MAP[model_name]['''dropout_rate''']
      __snake_case      :    Optional[int]  			=CONFIG_MAP[model_name]['''dw_padding''']

      __snake_case      :    Union[str, Any]  			='''huggingface/label-files'''
      __snake_case      :    List[Any]  			='''imagenet-1k-id2label.json'''
      __snake_case      :    Dict  			=1_000
      __snake_case      :    Union[str, Any]  			=json.load(open(hf_hub_download(snake_case_       ,		snake_case_       ,		repo_type='''dataset''' )       ,		'''r''' ) )
      __snake_case      :    List[str]  			={int(snake_case_ ): v for k, v in idalabel.items()}

      __snake_case      :    Optional[int]  			=idalabel
      __snake_case      :    Union[str, Any]  			={v: k for k, v in idalabel.items()}
      return config



def     __lowercase					(    )					->					int:
      __snake_case      :    Optional[Any]  			='''http://images.cocodataset.org/val2017/000000039769.jpg'''
      __snake_case      :    List[Any]  			=Image.open(requests.get(snake_case_       ,		stream=snake_case_ ).raw )
      return im



def     __lowercase					(    a :				str )					->					List[Any]:
      __snake_case      :    Union[str, Any]  			=CONFIG_MAP[model_name]['''image_size''']
      __snake_case      :    int  			=EfficientNetImageProcessor(
          size={'''height''': size, '''width''': size}       ,		image_mean=[0.4_8_5, 0.4_5_6, 0.4_0_6]       ,		image_std=[0.4_7_8_5_3_9_4_4, 0.4_7_3_2_8_6_4, 0.4_7_4_3_4_1_6_3]       ,		do_center_crop=snake_case_       ,		)
      return preprocessor



def     __lowercase					(    a :				Optional[int] )					->					Union[str, Any]:
      __snake_case      :    Any  			=[v.split('''_''' )[0].split('''block''' )[1] for v in original_param_names if v.startswith('''block''' )]
      __snake_case      :    Any  			=sorted(set(snake_case_ ) )
      __snake_case      :    int  			=len(snake_case_ )
      __snake_case      :    Dict  			={b: str(snake_case_ ) for b, i in zip(snake_case_       ,		range(snake_case_ ) )}

      __snake_case      :    Tuple  			=[]
      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      :    Optional[int]  			=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      :    Union[str, Any]  			={}
      for item in rename_keys:
            if item[0] in original_param_names:
                  __snake_case      :    Optional[int]  			='''efficientnet.''' + item[1]

      __snake_case      :    List[str]  			='''classifier.weight'''
      __snake_case      :    Tuple  			='''classifier.bias'''
      return key_mapping



def     __lowercase					(    a :				Optional[Any]       ,		a :				int       ,		a :				Optional[int] )					->					Any:
      for key, value in tf_params.items():
            if "normalization" in key:
                  continue

            __snake_case      :    Optional[Any]  			=key_mapping[key]
            if "_conv" in key and "kernel" in key:
                  __snake_case      :    Union[str, Any]  			=torch.from_numpy(snake_case_ ).permute(3       ,		2       ,		0       ,		1 )
            elif "depthwise_kernel" in key:
                  __snake_case      :    List[str]  			=torch.from_numpy(snake_case_ ).permute(2       ,		3       ,		0       ,		1 )
            elif "kernel" in key:
                  __snake_case      :    Optional[Any]  			=torch.from_numpy(np.transpose(snake_case_ ) )
            else:
                  __snake_case      :    Tuple  			=torch.from_numpy(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_(snake_case_ )



@torch.no_grad()
def     __lowercase					(    a :				Union[str, Any]       ,		a :				Optional[int]       ,		a :				List[Any]       ,		a :				Union[str, Any] )					->					List[str]:
      __snake_case      :    List[str]  			=model_classes[model_name](
          include_top=snake_case_       ,		weights='''imagenet'''       ,		input_tensor=snake_case_       ,		input_shape=snake_case_       ,		pooling=snake_case_       ,		classes=1_000       ,		classifier_activation='''softmax'''       ,		)

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

      # Load HuggingFace model
      __snake_case      :    int  			=get_efficientnet_config(snake_case_ )
      __snake_case      :    Optional[int]  			=EfficientNetForImageClassification(snake_case_ ).eval()
      __snake_case      :    Any  			=hf_model.state_dict()

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

      # Initialize preprocessor and preprocess input image
      __snake_case      :    Union[str, Any]  			=convert_image_processor(snake_case_ )
      __snake_case      :    int  			=preprocessor(images=prepare_img()       ,		return_tensors='''pt''' )

      # HF model inference
      hf_model.eval()
      with torch.no_grad():
            __snake_case      :    Union[str, Any]  			=hf_model(**snake_case_ )
      __snake_case      :    Optional[int]  			=outputs.logits.detach().numpy()

      # Original model inference
      __snake_case      :    int  			=False
      __snake_case      :    int  			=CONFIG_MAP[model_name]['''image_size''']
      __snake_case      :    Dict  			=prepare_img().resize((image_size, image_size)       ,		resample=PIL.Image.NEAREST )
      __snake_case      :    List[str]  			=image.img_to_array(snake_case_ )
      __snake_case      :    Union[str, Any]  			=np.expand_dims(snake_case_       ,		axis=0 )
      __snake_case      :    int  			=original_model.predict(snake_case_ )

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

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


if __name__ == "__main__":
      UpperCamelCase_	:				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""")

      UpperCamelCase_	:				str										= parser.parse_args()
      convert_efficientnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)










"""simple docstring"""





from typing import TYPE_CHECKING

from ...utils import _LazyModule


UpperCamelCase_	:				Union[str, Any]										= {"""tokenization_wav2vec2_phoneme""": ["""Wav2Vec2PhonemeCTCTokenizer"""]}


if TYPE_CHECKING:
    from .tokenization_wavaveca_phoneme import WavaVecaPhonemeCTCTokenizer
else:
    import sys

    UpperCamelCase_	:				List[Any]										= _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


