





import argparse
import json
from pathlib import Path

import requests
import timm
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import DeiTImageProcessor, ViTConfig, ViTForImageClassification, ViTImageProcessor, ViTModel
from transformers.utils import logging


logging.set_verbosity_info()
_UpperCAmelCase				:       Dict	   =				logging.get_logger(__name__)


def        SCREAMING_SNAKE_CASE      (       _UpperCAmelCase						,	_UpperCAmelCase=False )			->     Union[str, Any]:
    lowerCamelCase__    :      Tuple					 =		[]
    for i in range(config.num_hidden_layers ):
        # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
        rename_keys.append((F"""blocks.{i}.norm1.weight""", F"""vit.encoder.layer.{i}.layernorm_before.weight""") )
        rename_keys.append((F"""blocks.{i}.norm1.bias""", F"""vit.encoder.layer.{i}.layernorm_before.bias""") )
        rename_keys.append((F"""blocks.{i}.attn.proj.weight""", F"""vit.encoder.layer.{i}.attention.output.dense.weight""") )
        rename_keys.append((F"""blocks.{i}.attn.proj.bias""", F"""vit.encoder.layer.{i}.attention.output.dense.bias""") )
        rename_keys.append((F"""blocks.{i}.norm2.weight""", F"""vit.encoder.layer.{i}.layernorm_after.weight""") )
        rename_keys.append((F"""blocks.{i}.norm2.bias""", F"""vit.encoder.layer.{i}.layernorm_after.bias""") )
        rename_keys.append((F"""blocks.{i}.mlp.fc1.weight""", F"""vit.encoder.layer.{i}.intermediate.dense.weight""") )
        rename_keys.append((F"""blocks.{i}.mlp.fc1.bias""", F"""vit.encoder.layer.{i}.intermediate.dense.bias""") )
        rename_keys.append((F"""blocks.{i}.mlp.fc2.weight""", F"""vit.encoder.layer.{i}.output.dense.weight""") )
        rename_keys.append((F"""blocks.{i}.mlp.fc2.bias""", F"""vit.encoder.layer.{i}.output.dense.bias""") )

    # projection layer + position embeddings
    rename_keys.extend(
        [
            ('cls_token', 'vit.embeddings.cls_token'),
            ('patch_embed.proj.weight', 'vit.embeddings.patch_embeddings.projection.weight'),
            ('patch_embed.proj.bias', 'vit.embeddings.patch_embeddings.projection.bias'),
            ('pos_embed', 'vit.embeddings.position_embeddings'),
        ] )

    if base_model:
        # layernorm + pooler
        rename_keys.extend(
            [
                ('norm.weight', 'layernorm.weight'),
                ('norm.bias', 'layernorm.bias'),
                ('pre_logits.fc.weight', 'pooler.dense.weight'),
                ('pre_logits.fc.bias', 'pooler.dense.bias'),
            ] )

        # if just the base model, we should remove "vit" from all keys that start with "vit"
        lowerCamelCase__    :      Any					 =		[(pair[0], pair[1][4:]) if pair[1].startswith('vit' ) else pair for pair in rename_keys]
    else:
        # layernorm + classification head
        rename_keys.extend(
            [
                ('norm.weight', 'vit.layernorm.weight'),
                ('norm.bias', 'vit.layernorm.bias'),
                ('head.weight', 'classifier.weight'),
                ('head.bias', 'classifier.bias'),
            ] )

    return rename_keys


def        SCREAMING_SNAKE_CASE      (       _UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase=False )			->     Tuple:


    for i in range(config.num_hidden_layers ):
        if base_model:
            lowerCamelCase__    :      Optional[int]					 =		''
        else:
            lowerCamelCase__    :      List[str]					 =		'vit.'
        # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
        lowerCamelCase__    :      str					 =		state_dict.pop(F"""blocks.{i}.attn.qkv.weight""" )
        lowerCamelCase__    :      List[str]					 =		state_dict.pop(F"""blocks.{i}.attn.qkv.bias""" )
        # next, add query, keys and values (in that order) to the state dict
        lowerCamelCase__    :      int					 =		in_proj_weight[
            : config.hidden_size, :
        ]
        lowerCamelCase__    :      str					 =		in_proj_bias[: config.hidden_size]
        lowerCamelCase__    :      int					 =		in_proj_weight[
            config.hidden_size : config.hidden_size * 2, :
        ]
        lowerCamelCase__    :      Any					 =		in_proj_bias[
            config.hidden_size : config.hidden_size * 2
        ]
        lowerCamelCase__    :      str					 =		in_proj_weight[
            -config.hidden_size :, :
        ]
        lowerCamelCase__    :      List[Any]					 =		in_proj_bias[-config.hidden_size :]


def        SCREAMING_SNAKE_CASE      (       _UpperCAmelCase )			->     Union[str, Any]:
    lowerCamelCase__    :      int					 =		['head.weight', 'head.bias']
    for k in ignore_keys:
        state_dict.pop(__A						,	__A )


def        SCREAMING_SNAKE_CASE      (       _UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase )			->     int:
    lowerCamelCase__    :      Union[str, Any]					 =		dct.pop(__A )
    lowerCamelCase__    :      List[str]					 =		val


def        SCREAMING_SNAKE_CASE      (       )			->     List[Any]:
    lowerCamelCase__    :      List[str]					 =		'http://images.cocodataset.org/val2017/000000039769.jpg'
    lowerCamelCase__    :      Optional[Any]					 =		Image.open(requests.get(__A						,	stream=__A ).raw )
    return im


@torch.no_grad()
def        SCREAMING_SNAKE_CASE      (       _UpperCAmelCase						,	_UpperCAmelCase )			->     List[str]:
    lowerCamelCase__    :      Dict					 =		ViTConfig()
    lowerCamelCase__    :      Union[str, Any]					 =		False
    # dataset (ImageNet-21k only or also fine-tuned on ImageNet 2012), patch_size and image_size
    if vit_name[-5:] == "in21k":
        lowerCamelCase__    :      Dict					 =		True
        lowerCamelCase__    :      Optional[Any]					 =		int(vit_name[-12:-10] )
        lowerCamelCase__    :      Optional[Any]					 =		int(vit_name[-9:-6] )
    else:
        lowerCamelCase__    :      Optional[int]					 =		1000
        lowerCamelCase__    :      str					 =		'huggingface/label-files'
        lowerCamelCase__    :      List[Any]					 =		'imagenet-1k-id2label.json'
        lowerCamelCase__    :      List[Any]					 =		json.load(open(hf_hub_download(__A						,	__A						,	repo_type='dataset' )						,	'r' ) )
        lowerCamelCase__    :      Dict					 =		{int(__A ): v for k, v in idalabel.items()}
        lowerCamelCase__    :      Optional[int]					 =		idalabel
        lowerCamelCase__    :      Optional[Any]					 =		{v: k for k, v in idalabel.items()}
        lowerCamelCase__    :      Optional[int]					 =		int(vit_name[-6:-4] )
        lowerCamelCase__    :      Optional[int]					 =		int(vit_name[-3:] )
    # size of the architecture
    if "deit" in vit_name:
        if vit_name[9:].startswith('tiny' ):
            lowerCamelCase__    :      Any					 =		192
            lowerCamelCase__    :      str					 =		768
            lowerCamelCase__    :      Tuple					 =		12
            lowerCamelCase__    :      Any					 =		3
        elif vit_name[9:].startswith('small' ):
            lowerCamelCase__    :      Dict					 =		384
            lowerCamelCase__    :      int					 =		1536
            lowerCamelCase__    :      Tuple					 =		12
            lowerCamelCase__    :      List[Any]					 =		6
        else:
            pass
    else:
        if vit_name[4:].startswith('small' ):
            lowerCamelCase__    :      Optional[int]					 =		768
            lowerCamelCase__    :      Tuple					 =		2304
            lowerCamelCase__    :      List[str]					 =		8
            lowerCamelCase__    :      List[str]					 =		8
        elif vit_name[4:].startswith('base' ):
            pass
        elif vit_name[4:].startswith('large' ):
            lowerCamelCase__    :      int					 =		1024
            lowerCamelCase__    :      Any					 =		4096
            lowerCamelCase__    :      Tuple					 =		24
            lowerCamelCase__    :      int					 =		16
        elif vit_name[4:].startswith('huge' ):
            lowerCamelCase__    :      Any					 =		1280
            lowerCamelCase__    :      int					 =		5120
            lowerCamelCase__    :      str					 =		32
            lowerCamelCase__    :      List[Any]					 =		16

    # load original model from timm
    lowerCamelCase__    :      List[str]					 =		timm.create_model(__A						,	pretrained=__A )
    timm_model.eval()

    # load state_dict of original model, remove and rename some keys
    lowerCamelCase__    :      int					 =		timm_model.state_dict()
    if base_model:
        remove_classification_head_(__A )
    lowerCamelCase__    :      Optional[int]					 =		create_rename_keys(__A						,	__A )
    for src, dest in rename_keys:
        rename_key(__A						,	__A						,	__A )
    read_in_q_k_v(__A						,	__A						,	__A )

    # load HuggingFace model
    if vit_name[-5:] == "in21k":
        lowerCamelCase__    :      Tuple					 =		ViTModel(__A ).eval()
    else:
        lowerCamelCase__    :      List[str]					 =		ViTForImageClassification(__A ).eval()
    model.load_state_dict(__A )

    # Check outputs on an image, prepared by ViTImageProcessor/DeiTImageProcessor
    if "deit" in vit_name:
        lowerCamelCase__    :      Union[str, Any]					 =		DeiTImageProcessor(size=config.image_size )
    else:
        lowerCamelCase__    :      Dict					 =		ViTImageProcessor(size=config.image_size )
    lowerCamelCase__    :      Dict					 =		image_processor(images=prepare_img()						,	return_tensors='pt' )
    lowerCamelCase__    :      Optional[int]					 =		encoding['pixel_values']
    lowerCamelCase__    :      List[Any]					 =		model(__A )

    if base_model:
        lowerCamelCase__    :      Union[str, Any]					 =		timm_model.forward_features(__A )
        assert timm_pooled_output.shape == outputs.pooler_output.shape
        assert torch.allclose(__A						,	outputs.pooler_output						,	atol=1e-3 )
    else:
        lowerCamelCase__    :      Any					 =		timm_model(__A )
        assert timm_logits.shape == outputs.logits.shape
        assert torch.allclose(__A						,	outputs.logits						,	atol=1e-3 )

    Path(__A ).mkdir(exist_ok=__A )
    print(F"""Saving model {vit_name} to {pytorch_dump_folder_path}""" )
    model.save_pretrained(__A )
    print(F"""Saving image processor to {pytorch_dump_folder_path}""" )
    image_processor.save_pretrained(__A )


if __name__ == "__main__":
    _UpperCAmelCase				:       Tuple	   =				argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        """--vit_name""",
        default="""vit_base_patch16_224""",
        type=str,
        help="""Name of the ViT timm model you'd like to convert.""",
    )
    parser.add_argument(
        """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model directory."""
    )

    _UpperCAmelCase				:       List[str]	   =				parser.parse_args()
    convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path)











import random
import timeit
from functools import wraps
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_tf_auto import TF_MODEL_MAPPING, TF_MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_tf_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_tf_available():
    import tensorflow as tf
    from tensorflow.python.framework.errors_impl import ResourceExhaustedError

    from .benchmark_args_tf import TensorFlowBenchmarkArguments

if is_pyanvml_available():
    import pyanvml.pyanvml as nvml

snake_case_ : Tuple    		=       logging.get_logger(__name__)


def    A						(__A   :    bool   ,				__A   :    bool       )   ->      Optional[Any]:







     """simple docstring"""


     def run_func(__A   :    Optional[Any]       ):
          @wraps(__A       )
          def run_in_eager_mode(*__A   :    Dict   ,				**__A   :    List[Any]       ):
               return func(*__A   ,				**__A       )

          @wraps(__A       )
          @tf.function(experimental_compile=__A       )
          def run_in_graph_mode(*__A   :    Optional[Any]   ,				**__A   :    Any       ):
               return func(*__A   ,				**__A       )

          if do_eager_mode is True:
               if use_xla is not False:
                    raise ValueError(
                        '''Cannot run model in XLA, if `args.eager_mode` is set to `True`. Please set `args.eager_mode=False`.'''       )
               return run_in_eager_mode
          else:
               return run_in_graph_mode

     return run_func


def    A						(__A   :    int   ,				__A   :    int   ,				__A   :    int       )   ->      ["tf.Tensor"]:







     """simple docstring"""


     UpperCAmelCase_          =							random.Random()
     UpperCAmelCase_          =							[rng.randint(0   ,				vocab_size - 1       ) for i in range(batch_size * sequence_length       )]
     return tf.constant(__A   ,				shape=(batch_size, sequence_length)   ,				dtype=tf.intaa       )







class __snake_case			(			a		):
   UpperCAmelCase__  :    TensorFlowBenchmarkArguments
   UpperCAmelCase__  :    PretrainedConfig
   UpperCAmelCase__  :    str          =  "TensorFlow"







   @property
   def lowerCamelCase						(      self		:						List[str]):

        """simple docstring"""


        return tf.__version__







   def lowerCamelCase						(      self		:						Dict					,						_snake_case		:						str					,						_snake_case		:						int					,						_snake_case		:						int):

        """simple docstring"""


        UpperCAmelCase_          =							self.args.strategy
        if strategy is None:
             raise ValueError('''A device strategy has to be initialized before using TensorFlow.''')
        UpperCAmelCase_          =							self._prepare_inference_func(_snake_case					,						_snake_case					,						_snake_case)
        return self._measure_speed(_inference)







   def lowerCamelCase						(      self		:						Any					,						_snake_case		:						str					,						_snake_case		:						int					,						_snake_case		:						int):

        """simple docstring"""


        UpperCAmelCase_          =							self.args.strategy
        if strategy is None:
             raise ValueError('''A device strategy has to be initialized before using TensorFlow.''')
        UpperCAmelCase_          =							self._prepare_train_func(_snake_case					,						_snake_case					,						_snake_case)
        return self._measure_speed(_train)







   def lowerCamelCase						(      self		:						Any					,						_snake_case		:						str					,						_snake_case		:						int					,						_snake_case		:						int):

        """simple docstring"""


        if self.args.is_gpu:
             tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]					,						_snake_case)
        UpperCAmelCase_          =							self.args.strategy
        if strategy is None:
             raise ValueError('''A device strategy has to be initialized before using TensorFlow.''')
        UpperCAmelCase_          =							self._prepare_inference_func(_snake_case					,						_snake_case					,						_snake_case)
        return self._measure_memory(_inference)







   def lowerCamelCase						(      self		:						Optional[Any]					,						_snake_case		:						str					,						_snake_case		:						int					,						_snake_case		:						int):

        """simple docstring"""


        if self.args.is_gpu:
             tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]					,						_snake_case)
        UpperCAmelCase_          =							self.args.strategy
        if strategy is None:
             raise ValueError('''A device strategy has to be initialized before using TensorFlow.''')

        UpperCAmelCase_          =							self._prepare_train_func(_snake_case					,						_snake_case					,						_snake_case)
        return self._measure_memory(_train)







   def lowerCamelCase						(      self		:						Optional[int]					,						_snake_case		:						str					,						_snake_case		:						int					,						_snake_case		:						int):

        """simple docstring"""


        UpperCAmelCase_          =							self.config_dict[model_name]

        if self.args.fpaa:
             raise NotImplementedError('''Mixed precision is currently not supported.''')

        UpperCAmelCase_          =							(
            hasattr(_snake_case					,						'''architectures''')
            and isinstance(config.architectures					,						_snake_case)
            and len(config.architectures) > 0
        )
        if not self.args.only_pretrain_model and has_model_class_in_config:
             try:
                  UpperCAmelCase_          =							'''TF''' + config.architectures[0]  # prepend 'TF' for tensorflow model
                  UpperCAmelCase_          =							__import__('''transformers'''					,						fromlist=[model_class])
                  UpperCAmelCase_          =							getattr(_snake_case					,						_snake_case)
                  UpperCAmelCase_          =							model_cls(_snake_case)
             except ImportError:
                  raise ImportError(
                      F"""{model_class} does not exist. If you just want to test the pretrained model, you might want to"""
                      ''' set `--only_pretrain_model` or `args.only_pretrain_model=True`.''')
        else:
             UpperCAmelCase_          =							TF_MODEL_MAPPING[config.__class__](_snake_case)

        # encoder-decoder has vocab size saved differently
        UpperCAmelCase_          =							config.vocab_size if hasattr(_snake_case					,						'''vocab_size''') else config.encoder.vocab_size
        UpperCAmelCase_          =							random_input_ids(_snake_case					,						_snake_case					,						_snake_case)

        @run_with_tf_optimizations(self.args.eager_mode					,						self.args.use_xla)
        def encoder_decoder_forward():
             return model(_snake_case					,						decoder_input_ids=_snake_case					,						training=_snake_case)

        @run_with_tf_optimizations(self.args.eager_mode					,						self.args.use_xla)
        def encoder_forward():
             return model(_snake_case					,						training=_snake_case)

        UpperCAmelCase_          =							encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

        return _inference







   def lowerCamelCase						(      self		:						Optional[Any]					,						_snake_case		:						str					,						_snake_case		:						int					,						_snake_case		:						int):

        """simple docstring"""


        UpperCAmelCase_          =							self.config_dict[model_name]

        if self.args.eager_mode is not False:
             raise ValueError('''Training cannot be done in eager mode. Please make sure that `args.eager_mode = False`.''')

        if self.args.fpaa:
             raise NotImplementedError('''Mixed precision is currently not supported.''')

        UpperCAmelCase_          =							(
            hasattr(_snake_case					,						'''architectures''')
            and isinstance(config.architectures					,						_snake_case)
            and len(config.architectures) > 0
        )
        if not self.args.only_pretrain_model and has_model_class_in_config:
             try:
                  UpperCAmelCase_          =							'''TF''' + config.architectures[0]  # prepend 'TF' for tensorflow model
                  UpperCAmelCase_          =							__import__('''transformers'''					,						fromlist=[model_class])
                  UpperCAmelCase_          =							getattr(_snake_case					,						_snake_case)
                  UpperCAmelCase_          =							model_cls(_snake_case)
             except ImportError:
                  raise ImportError(
                      F"""{model_class} does not exist. If you just want to test the pretrained model, you might want to"""
                      ''' set `--only_pretrain_model` or `args.only_pretrain_model=True`.''')
        else:
             UpperCAmelCase_          =							TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](_snake_case)

        # encoder-decoder has vocab size saved differently
        UpperCAmelCase_          =							config.vocab_size if hasattr(_snake_case					,						'''vocab_size''') else config.encoder.vocab_size
        UpperCAmelCase_          =							random_input_ids(_snake_case					,						_snake_case					,						_snake_case)

        @run_with_tf_optimizations(self.args.eager_mode					,						self.args.use_xla)
        def encoder_decoder_train():
             UpperCAmelCase_          =							model(_snake_case					,						decoder_input_ids=_snake_case					,						labels=_snake_case					,						training=_snake_case)[0]
             UpperCAmelCase_          =							tf.gradients(_snake_case					,						model.trainable_variables)
             return gradients

        @run_with_tf_optimizations(self.args.eager_mode					,						self.args.use_xla)
        def encoder_train():
             UpperCAmelCase_          =							model(_snake_case					,						labels=_snake_case					,						training=_snake_case)[0]
             UpperCAmelCase_          =							tf.gradients(_snake_case					,						model.trainable_variables)
             return gradients

        UpperCAmelCase_          =							encoder_decoder_train if config.is_encoder_decoder else encoder_train

        return _train







   def lowerCamelCase						(      self		:						Any					,						_snake_case		:						Optional[Any]):

        """simple docstring"""


        with self.args.strategy.scope():
             try:
                  if self.args.is_tpu or self.args.use_xla:
                       # run additional 10 times to stabilize compilation for tpu
                       logger.info('''Do inference on TPU. Running model 5 times to stabilize compilation''')
                       timeit.repeat(_snake_case					,						repeat=1					,						number=5)

                  # as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
                  UpperCAmelCase_          =							timeit.repeat(
                      _snake_case					,						repeat=self.args.repeat					,						number=10					,						)

                  return min(_snake_case) / 1_0.0
             except ResourceExhaustedError as e:
                  self.print_fn(F"""Doesn't fit on GPU. {e}""")







   def lowerCamelCase						(      self		:						Dict					,						_snake_case		:						Callable[[], None]):

        """simple docstring"""


        logger.info(
            '''Note that TensorFlow allocates more memory than '''
            '''it might need to speed up computation. '''
            '''The memory reported here corresponds to the memory '''
            '''reported by `nvidia-smi`, which can vary depending '''
            '''on total available memory on the GPU that is used.''')
        with self.args.strategy.scope():
             try:
                  if self.args.trace_memory_line_by_line:
                       if not self.args.eager_mode:
                            raise ValueError(
                                '''`args.eager_mode` is set to `False`. Make sure to run model in eager mode to measure memory'''
                                ''' consumption line by line.''')
                       UpperCAmelCase_          =							start_memory_tracing('''transformers''')

                  if self.args.is_tpu:
                       # tpu
                       raise NotImplementedError(
                           '''Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking'''
                           ''' with `args.memory=False`''')
                  elif self.args.is_gpu:
                       # gpu
                       if not is_pyanvml_available():
                            logger.warning(
                                '''py3nvml not installed, we won\'t log GPU memory usage. '''
                                '''Install py3nvml (pip install py3nvml) to log information about GPU.''')
                            UpperCAmelCase_          =							'''N/A'''
                       else:
                            logger.info(
                                '''Measuring total GPU usage on GPU device. Make sure to not have additional processes'''
                                ''' running on the same GPU.''')
                            # init nvml
                            nvml.nvmlInit()
                            func()
                            UpperCAmelCase_          =							nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx)
                            UpperCAmelCase_          =							nvml.nvmlDeviceGetMemoryInfo(_snake_case)
                            UpperCAmelCase_          =							meminfo.used
                            UpperCAmelCase_          =							Memory(_snake_case)
                            # shutdown nvml
                            nvml.nvmlShutdown()
                  else:
                       # cpu
                       if self.args.trace_memory_line_by_line:
                            logger.info(
                                '''When enabling line by line tracing, the max peak memory for CPU is inaccurate in'''
                                ''' TensorFlow.''')
                            UpperCAmelCase_          =							None
                       else:
                            UpperCAmelCase_          =							measure_peak_memory_cpu(_snake_case)
                            UpperCAmelCase_          =							Memory(_snake_case) if isinstance(_snake_case					,						_snake_case) else memory_bytes
                  if self.args.trace_memory_line_by_line:
                       UpperCAmelCase_          =							stop_memory_tracing(_snake_case)
                       if memory is None:
                            UpperCAmelCase_          =							summary.total
                  else:
                       UpperCAmelCase_          =							None

                  return memory, summary
             except ResourceExhaustedError as e:
                  self.print_fn(F"""Doesn't fit on GPU. {e}""")
                  return "N/A", None






