





'''simple docstring'''




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

_SCREAMING_SNAKE_CASE							     =			logging.get_logger(__name__)





def 	__lowerCamelCase			(      __lowerCAmelCase       :							Optional[Any]      ,    __lowerCAmelCase       :							List[Any]		)		->  Union[str, Any]:
   def run_func(__lowerCAmelCase       :							int		):
      @wraps(__lowerCAmelCase		)
      def run_in_eager_mode(*__lowerCAmelCase       :							List[str]      ,    **__lowerCAmelCase       :							Tuple		):
         return func(*__lowerCAmelCase      ,    **__lowerCAmelCase		)

      @wraps(__lowerCAmelCase		)
      @tf.function(experimental_compile=__lowerCAmelCase		)
      def run_in_graph_mode(*__lowerCAmelCase       :							List[Any]      ,    **__lowerCAmelCase       :							Dict		):
         return func(*__lowerCAmelCase      ,    **__lowerCAmelCase		)

      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 	__lowerCamelCase			(      __lowerCAmelCase       :							Optional[int]      ,    __lowerCAmelCase       :							List[Any]      ,    __lowerCAmelCase       :							List[str]		)		->  int:
   snake_case 						=   random.Random()
   snake_case 						=   [rng.randint(0      ,    vocab_size - 1		) for i in range(batch_size * sequence_length		)]
   return tf.constant(__lowerCAmelCase      ,    shape=(batch_size, sequence_length)      ,    dtype=tf.intaa		)


class 			_lowerCAmelCase		(			__A	):






     """simple docstring"""






     snake_case_						      =   42
     snake_case_						      =   42
     snake_case_						      =   "TensorFlow"



     @property
     def      lowerCAmelCase			(   self							:			Optional[Any]						)->	Tuple:
        return tf.__version__



     def      lowerCAmelCase			(   self							:			List[Any]   ,       __snake_case							:			Tuple   ,       __snake_case							:			Any   ,       __snake_case							:			str						)->	List[str]:
        snake_case 						=   self.args.strategy
        if strategy is None:
           raise ValueError("""A device strategy has to be initialized before using TensorFlow."""						)
        snake_case 						=   self._prepare_inference_func(snake_case_   ,       snake_case_   ,       snake_case_						)
        return self._measure_speed(_inference						)



     def      lowerCAmelCase			(   self							:			Optional[int]   ,       __snake_case							:			Tuple   ,       __snake_case							:			int   ,       __snake_case							:			List[str]						)->	List[Any]:
        snake_case 						=   self.args.strategy
        if strategy is None:
           raise ValueError("""A device strategy has to be initialized before using TensorFlow."""						)
        snake_case 						=   self._prepare_train_func(snake_case_   ,       snake_case_   ,       snake_case_						)
        return self._measure_speed(_train						)



     def      lowerCAmelCase			(   self							:			str   ,       __snake_case							:			Dict   ,       __snake_case							:			List[Any]   ,       __snake_case							:			List[Any]						)->	List[Any]:


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



     def      lowerCAmelCase			(   self							:			List[Any]   ,       __snake_case							:			Optional[int]   ,       __snake_case							:			str   ,       __snake_case							:			List[str]						)->	str:


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

        snake_case 						=   self._prepare_train_func(snake_case_   ,       snake_case_   ,       snake_case_						)
        return self._measure_memory(_train						)



     def      lowerCAmelCase			(   self							:			List[str]   ,       __snake_case							:			Dict   ,       __snake_case							:			str   ,       __snake_case							:			Optional[Any]						)->	Optional[int]:
        snake_case 						=   self.config_dict[model_name]

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

        snake_case 						=   (
            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:
              snake_case 						=   """TF""" + config.architectures[0]  # prepend 'TF' for tensorflow model
              snake_case 						=   __import__("""transformers"""   ,       fromlist=[model_class]						)
              snake_case 						=   getattr(snake_case_   ,       snake_case_						)
              snake_case 						=   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:
           snake_case 						=   TF_MODEL_MAPPING[config.__class__](snake_case_						)

        # encoder-decoder has vocab size saved differently
        snake_case 						=   config.vocab_size if hasattr(snake_case_   ,       """vocab_size"""						) else config.encoder.vocab_size
        snake_case 						=   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_						)

        snake_case 						=   encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

        return _inference



     def      lowerCAmelCase			(   self							:			List[str]   ,       __snake_case							:			Tuple   ,       __snake_case							:			Dict   ,       __snake_case							:			int						)->	Any:
        snake_case 						=   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."""						)

        snake_case 						=   (
            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:
              snake_case 						=   """TF""" + config.architectures[0]  # prepend 'TF' for tensorflow model
              snake_case 						=   __import__("""transformers"""   ,       fromlist=[model_class]						)
              snake_case 						=   getattr(snake_case_   ,       snake_case_						)
              snake_case 						=   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:
           snake_case 						=   TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](snake_case_						)

        # encoder-decoder has vocab size saved differently
        snake_case 						=   config.vocab_size if hasattr(snake_case_   ,       """vocab_size"""						) else config.encoder.vocab_size
        snake_case 						=   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():
           snake_case 						=   model(snake_case_   ,       decoder_input_ids=snake_case_   ,       labels=snake_case_   ,       training=snake_case_						)[0]
           snake_case 						=   tf.gradients(snake_case_   ,       model.trainable_variables						)
           return gradients

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

        snake_case 						=   encoder_decoder_train if config.is_encoder_decoder else encoder_train

        return _train



     def      lowerCAmelCase			(   self							:			str   ,       __snake_case							:			List[str]						)->	Optional[int]:


        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
              snake_case 						=   timeit.repeat(
                  snake_case_   ,       repeat=self.args.repeat   ,       number=10   ,       )

              return min(snake_case_						) / 10.0
           except ResourceExhaustedError as e:
              self.print_fn(f'''Doesn\'t fit on GPU. {e}'''						)



     def      lowerCAmelCase			(   self							:			Optional[int]   ,       __snake_case							:			Dict						)->	Any:
        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."""						)
                 snake_case 						=   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."""						)
                    snake_case 						=   """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()
                    snake_case 						=   nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx						)
                    snake_case 						=   nvml.nvmlDeviceGetMemoryInfo(snake_case_						)
                    snake_case 						=   meminfo.used
                    snake_case 						=   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."""						)
                    snake_case 						=   None
                 else:
                    snake_case 						=   measure_peak_memory_cpu(snake_case_						)
                    snake_case 						=   Memory(snake_case_						) if isinstance(snake_case_   ,       snake_case_						) else memory_bytes
              if self.args.trace_memory_line_by_line:
                 snake_case 						=   stop_memory_tracing(snake_case_						)
                 if memory is None:
                    snake_case 						=   summary.total
              else:
                 snake_case 						=   None

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


"""simple docstring"""



import argparse
import json

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

from transformers import ViTImageProcessor, ViTMSNConfig, ViTMSNModel
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD


torch.set_grad_enabled(False)

def  _A(		lowerCAmelCase			,  lowerCAmelCase=False							):
							A__       :     Dict								=			[]
							for i in range(config.num_hidden_layers							):
														# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
														rename_keys.append((F'''module.blocks.{i}.norm1.weight''', F'''vit.encoder.layer.{i}.layernorm_before.weight''')							)
														rename_keys.append((F'''module.blocks.{i}.norm1.bias''', F'''vit.encoder.layer.{i}.layernorm_before.bias''')							)
														rename_keys.append(
														    (F'''module.blocks.{i}.attn.proj.weight''', F'''vit.encoder.layer.{i}.attention.output.dense.weight''')							)
														rename_keys.append((F'''module.blocks.{i}.attn.proj.bias''', F'''vit.encoder.layer.{i}.attention.output.dense.bias''')							)
														rename_keys.append((F'''module.blocks.{i}.norm2.weight''', F'''vit.encoder.layer.{i}.layernorm_after.weight''')							)
														rename_keys.append((F'''module.blocks.{i}.norm2.bias''', F'''vit.encoder.layer.{i}.layernorm_after.bias''')							)
														rename_keys.append((F'''module.blocks.{i}.mlp.fc1.weight''', F'''vit.encoder.layer.{i}.intermediate.dense.weight''')							)
														rename_keys.append((F'''module.blocks.{i}.mlp.fc1.bias''', F'''vit.encoder.layer.{i}.intermediate.dense.bias''')							)
														rename_keys.append((F'''module.blocks.{i}.mlp.fc2.weight''', F'''vit.encoder.layer.{i}.output.dense.weight''')							)
														rename_keys.append((F'''module.blocks.{i}.mlp.fc2.bias''', F'''vit.encoder.layer.{i}.output.dense.bias''')							)

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

							if base_model:
														# layernorm + pooler
														rename_keys.extend(
														    [
														        ("""module.norm.weight""", """layernorm.weight"""),
														        ("""module.norm.bias""", """layernorm.bias"""),
														    ]							)

														# if just the base model, we should remove "vit" from all keys that start with "vit"
														A__       :     Dict								=			[(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  _A(		lowerCAmelCase			,  lowerCAmelCase			,  lowerCAmelCase=False							):
							for i in range(config.num_hidden_layers							):
														if base_model:
																					A__       :     List[str]								=			""""""
														else:
																					A__       :     Dict								=			"""vit."""
														# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
														A__       :     List[str]								=			state_dict.pop(F'''module.blocks.{i}.attn.qkv.weight'''							)
														A__       :     Tuple								=			state_dict.pop(F'''module.blocks.{i}.attn.qkv.bias'''							)
														# next, add query, keys and values (in that order) to the state dict
														A__       :     Tuple								=			in_proj_weight[
														    : config.hidden_size, :
														]
														A__       :     Union[str, Any]								=			in_proj_bias[: config.hidden_size]
														A__       :     List[str]								=			in_proj_weight[
														    config.hidden_size : config.hidden_size * 2, :
														]
														A__       :     int								=			in_proj_bias[
														    config.hidden_size : config.hidden_size * 2
														]
														A__       :     int								=			in_proj_weight[
														    -config.hidden_size :, :
														]
														A__       :     List[Any]								=			in_proj_bias[-config.hidden_size :]

def  _A(		lowerCAmelCase							):
							A__       :     str								=			["""head.weight""", """head.bias"""]
							for k in ignore_keys:
														state_dict.pop(lowerCAmelCase			,  lowerCAmelCase							)

def  _A(		lowerCAmelCase							):
							# projection head is used in the self-supervised pre-training in MSN,
							# for downstream task it's not needed.
							A__       :     Dict								=			[
							    """module.fc.fc1.weight""",
							    """module.fc.fc1.bias""",
							    """module.fc.bn1.weight""",
							    """module.fc.bn1.bias""",
							    """module.fc.bn1.running_mean""",
							    """module.fc.bn1.running_var""",
							    """module.fc.bn1.num_batches_tracked""",
							    """module.fc.fc2.weight""",
							    """module.fc.fc2.bias""",
							    """module.fc.bn2.weight""",
							    """module.fc.bn2.bias""",
							    """module.fc.bn2.running_mean""",
							    """module.fc.bn2.running_var""",
							    """module.fc.bn2.num_batches_tracked""",
							    """module.fc.fc3.weight""",
							    """module.fc.fc3.bias""",
							]
							for k in ignore_keys:
														state_dict.pop(lowerCAmelCase			,  lowerCAmelCase							)

def  _A(		lowerCAmelCase			,  lowerCAmelCase			,  lowerCAmelCase							):
							A__       :     Optional[Any]								=			dct.pop(lowerCAmelCase							)
							A__       :     List[Any]								=			val

def  _A(		lowerCAmelCase			,  lowerCAmelCase							):
							A__       :     str								=			ViTMSNConfig()
							A__       :     List[str]								=			1000

							A__       :     Optional[int]								=			"""datasets/huggingface/label-files"""
							A__       :     Optional[int]								=			"""imagenet-1k-id2label.json"""
							A__       :     str								=			json.load(open(hf_hub_download(lowerCAmelCase			,  lowerCAmelCase							)			,  """r"""							)							)
							A__       :     Tuple								=			{int(lowerCAmelCase							): v for k, v in idalabel.items()}
							A__       :     Optional[Any]								=			idalabel
							A__       :     List[Any]								=			{v: k for k, v in idalabel.items()}

							if "s16" in checkpoint_url:
														A__       :     Tuple								=			384
														A__       :     int								=			1536
														A__       :     Union[str, Any]								=			6
							elif "l16" in checkpoint_url:
														A__       :     int								=			1024
														A__       :     int								=			4096
														A__       :     Any								=			24
														A__       :     int								=			16
														A__       :     Union[str, Any]								=			0.1
							elif "b4" in checkpoint_url:
														A__       :     Dict								=			4
							elif "l7" in checkpoint_url:
														A__       :     List[str]								=			7
														A__       :     Optional[int]								=			1024
														A__       :     Optional[int]								=			4096
														A__       :     Any								=			24
														A__       :     Union[str, Any]								=			16
														A__       :     Tuple								=			0.1

							A__       :     List[str]								=			ViTMSNModel(lowerCAmelCase							)

							A__       :     Any								=			torch.hub.load_state_dict_from_url(lowerCAmelCase			,  map_location="""cpu"""							)["""target_encoder"""]

							A__       :     List[str]								=			ViTImageProcessor(size=config.image_size							)

							remove_projection_head(lowerCAmelCase							)
							A__       :     List[str]								=			create_rename_keys(lowerCAmelCase			,  base_model=lowerCAmelCase							)

							for src, dest in rename_keys:
														rename_key(lowerCAmelCase			,  lowerCAmelCase			,  lowerCAmelCase							)
							read_in_q_k_v(lowerCAmelCase			,  lowerCAmelCase			,  base_model=lowerCAmelCase							)

							model.load_state_dict(lowerCAmelCase							)
							model.eval()

							A__       :     str								=			"""http://images.cocodataset.org/val2017/000000039769.jpg"""

							A__       :     int								=			Image.open(requests.get(lowerCAmelCase			,  stream=lowerCAmelCase							).raw							)
							A__       :     List[Any]								=			ViTImageProcessor(
							    size=config.image_size			,  image_mean=lowerCAmelCase			,  image_std=lowerCAmelCase							)
							A__       :     Any								=			image_processor(images=lowerCAmelCase			,  return_tensors="""pt"""							)

							# forward pass
							torch.manual_seed(2							)
							A__       :     Tuple								=			model(**lowerCAmelCase							)
							A__       :     Union[str, Any]								=			outputs.last_hidden_state

							# The following Colab Notebook was used to generate these outputs:
							# https://colab.research.google.com/gist/sayakpaul/3672419a04f5997827503fd84079bdd1/scratchpad.ipynb
							if "s16" in checkpoint_url:
														A__       :     Union[str, Any]								=			torch.tensor([[-1.0915, -1.4876, -1.1809]]							)
							elif "b16" in checkpoint_url:
														A__       :     int								=			torch.tensor([[14.2889, -18.9045, 11.7281]]							)
							elif "l16" in checkpoint_url:
														A__       :     List[Any]								=			torch.tensor([[41.5028, -22.8681, 45.6475]]							)
							elif "b4" in checkpoint_url:
														A__       :     Optional[Any]								=			torch.tensor([[-4.3868, 5.2932, -0.4137]]							)
							else:
														A__       :     str								=			torch.tensor([[-0.1792, -0.6465, 2.4263]]							)

							# verify logits
							assert torch.allclose(last_hidden_state[:, 0, :3]			,  lowerCAmelCase			,  atol=1E-4							)

							print(F'''Saving model to {pytorch_dump_folder_path}'''							)
							model.save_pretrained(lowerCAmelCase							)

							print(F'''Saving image processor to {pytorch_dump_folder_path}'''							)
							image_processor.save_pretrained(lowerCAmelCase							)


if __name__ == "__main__":
				_UpperCamelCase								=		argparse.ArgumentParser()
				# Required parameters
				parser.add_argument(
				    "--checkpoint_url",
				    default="https://dl.fbaipublicfiles.com/msn/vits16_800ep.pth.tar",
				    type=str,
				    help="URL of the checkpoint 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								=		parser.parse_args()
				convert_vit_msn_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)





