




import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

from transformers import (
    WavaVecaConfig,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaForCTC,
    WavaVecaForPreTraining,
    WavaVecaProcessor,
    logging,
)
from transformers.models.wavaveca.modeling_wavaveca import WavaVecaForSequenceClassification


logging.set_verbosity_info()
lowerCamelCase__ :   Any					=							logging.get_logger(__name__)

lowerCamelCase__ :   int					=							{
    'post_extract_proj': 'feature_projection.projection',
    'encoder.pos_conv.0': 'encoder.pos_conv_embed.conv',
    'self_attn.k_proj': 'encoder.layers.*.attention.k_proj',
    'self_attn.v_proj': 'encoder.layers.*.attention.v_proj',
    'self_attn.q_proj': 'encoder.layers.*.attention.q_proj',
    'self_attn.out_proj': 'encoder.layers.*.attention.out_proj',
    'self_attn_layer_norm': 'encoder.layers.*.layer_norm',
    'fc1': 'encoder.layers.*.feed_forward.intermediate_dense',
    'fc2': 'encoder.layers.*.feed_forward.output_dense',
    'final_layer_norm': 'encoder.layers.*.final_layer_norm',
    'encoder.layer_norm': 'encoder.layer_norm',
    'adapter_layer': 'encoder.layers.*.adapter_layer',
    'w2v_model.layer_norm': 'feature_projection.layer_norm',
    'quantizer.weight_proj': 'quantizer.weight_proj',
    'quantizer.vars': 'quantizer.codevectors',
    'project_q': 'project_q',
    'final_proj': 'project_hid',
    'w2v_encoder.proj': 'lm_head',
    'mask_emb': 'masked_spec_embed',
    'pooling_layer.linear': 'projector',
    'pooling_layer.projection': 'classifier',
}
lowerCamelCase__ :   int					=							[
    'lm_head',
    'quantizer.weight_proj',
    'quantizer.codevectors',
    'project_q',
    'project_hid',
    'projector',
    'classifier',
]

def 					UpperCAmelCase_							(      __UpperCAmelCase						:	List[str]							)					-> Union[str, Any]:
      SCREAMING_SNAKE_CASE_							=			{}
      with open(__UpperCAmelCase			,			'r'							) as file:
            for line_number, line in enumerate(__UpperCAmelCase							):
                  SCREAMING_SNAKE_CASE_							=			line.strip()
                  if line:
                        SCREAMING_SNAKE_CASE_							=			line.split()
                        SCREAMING_SNAKE_CASE_							=			line_number
                        SCREAMING_SNAKE_CASE_							=			words[0]
                        SCREAMING_SNAKE_CASE_							=			value
      return result

def 					UpperCAmelCase_							(      __UpperCAmelCase						:	Dict			,			__UpperCAmelCase						:	Tuple			,			__UpperCAmelCase						:	List[Any]			,			__UpperCAmelCase						:	Optional[int]			,			__UpperCAmelCase						:	Optional[int]							)					-> int:
      for attribute in key.split('.'							):
            SCREAMING_SNAKE_CASE_							=			getattr(__UpperCAmelCase			,			__UpperCAmelCase							)

      SCREAMING_SNAKE_CASE_							=			None
      for param_key in PARAM_MAPPING.keys():
            if full_name.endswith(__UpperCAmelCase							):
                  SCREAMING_SNAKE_CASE_							=			PARAM_MAPPING[full_name.split('.'							)[-1]]
                  SCREAMING_SNAKE_CASE_							=			'param'

      if weight_type is not None and weight_type != "param":
            SCREAMING_SNAKE_CASE_							=			getattr(__UpperCAmelCase			,			__UpperCAmelCase							).shape
      elif weight_type is not None and weight_type == "param":
            SCREAMING_SNAKE_CASE_							=			hf_pointer
            for attribute in hf_param_name.split('.'							):
                  SCREAMING_SNAKE_CASE_							=			getattr(__UpperCAmelCase			,			__UpperCAmelCase							)
            SCREAMING_SNAKE_CASE_							=			shape_pointer.shape

            # let's reduce dimension
            SCREAMING_SNAKE_CASE_							=			value[0]
      else:
            SCREAMING_SNAKE_CASE_							=			hf_pointer.shape

      if hf_shape != value.shape:
            raise ValueError(
                f"Shape of hf {key + '.' + weight_type if weight_type is not None else ''} is {hf_shape}, but should be"
                f" {value.shape} for {full_name}"							)

      if weight_type == "weight":
            SCREAMING_SNAKE_CASE_							=			value
      elif weight_type == "weight_g":
            SCREAMING_SNAKE_CASE_							=			value
      elif weight_type == "weight_v":
            SCREAMING_SNAKE_CASE_							=			value
      elif weight_type == "bias":
            SCREAMING_SNAKE_CASE_							=			value
      elif weight_type == "param":
            for attribute in hf_param_name.split('.'							):
                  SCREAMING_SNAKE_CASE_							=			getattr(__UpperCAmelCase			,			__UpperCAmelCase							)
            SCREAMING_SNAKE_CASE_							=			value
      else:
            SCREAMING_SNAKE_CASE_							=			value

      logger.info(f"{key + '.' + weight_type if weight_type is not None else ''} was initialized from {full_name}."							)

def 					UpperCAmelCase_							(      __UpperCAmelCase						:	List[str]			,			__UpperCAmelCase						:	List[str]			,			__UpperCAmelCase						:	Dict			,			__UpperCAmelCase						:	List[str]			,			__UpperCAmelCase						:	List[Any]							)					-> List[Any]:
      SCREAMING_SNAKE_CASE_							=			None
      for param_key in PARAM_MAPPING.keys():
            if full_name.endswith(__UpperCAmelCase							):
                  SCREAMING_SNAKE_CASE_							=			PARAM_MAPPING[full_name.split('.'							)[-1]]
                  SCREAMING_SNAKE_CASE_							=			'param'

      if weight_type is not None and weight_type != "param":
            SCREAMING_SNAKE_CASE_							=			'.'.join([key, weight_type]							)
      elif weight_type is not None and weight_type == "param":
            SCREAMING_SNAKE_CASE_							=			'.'.join([key, hf_param_name]							)
      else:
            SCREAMING_SNAKE_CASE_							=			key

      SCREAMING_SNAKE_CASE_							=			value if 'lm_head' in full_key else value[0]


lowerCamelCase__ :   Dict					=							{
    'W_a': 'linear_1.weight',
    'W_b': 'linear_2.weight',
    'b_a': 'linear_1.bias',
    'b_b': 'linear_2.bias',
    'ln_W': 'norm.weight',
    'ln_b': 'norm.bias',
}

def 					UpperCAmelCase_							(      __UpperCAmelCase						:	List[Any]			,			__UpperCAmelCase						:	Dict			,			__UpperCAmelCase						:	str=None			,			__UpperCAmelCase						:	str=None							)					-> Optional[Any]:
      SCREAMING_SNAKE_CASE_							=			False
      for key, mapped_key in MAPPING.items():
            SCREAMING_SNAKE_CASE_							=			'wav2vec2.' + mapped_key if mapped_key not in TOP_LEVEL_KEYS else mapped_key
            if key in name or key.split('w2v_model.'							)[-1] == name.split('.'							)[0]:
                  SCREAMING_SNAKE_CASE_							=			True
                  if "*" in mapped_key:
                        SCREAMING_SNAKE_CASE_							=			name.split(__UpperCAmelCase							)[0].split('.'							)[-2]
                        SCREAMING_SNAKE_CASE_							=			mapped_key.replace('*'			,			__UpperCAmelCase							)
                  if "weight_g" in name:
                        SCREAMING_SNAKE_CASE_							=			'weight_g'
                  elif "weight_v" in name:
                        SCREAMING_SNAKE_CASE_							=			'weight_v'
                  elif "bias" in name:
                        SCREAMING_SNAKE_CASE_							=			'bias'
                  elif "weight" in name:
                        # TODO: don't match quantizer.weight_proj
                        SCREAMING_SNAKE_CASE_							=			'weight'
                  else:
                        SCREAMING_SNAKE_CASE_							=			None
                  if hf_dict is not None:
                        rename_dict(__UpperCAmelCase			,			__UpperCAmelCase			,			__UpperCAmelCase			,			__UpperCAmelCase			,			__UpperCAmelCase							)
                  else:
                        set_recursively(__UpperCAmelCase			,			__UpperCAmelCase			,			__UpperCAmelCase			,			__UpperCAmelCase			,			__UpperCAmelCase							)
                  return is_used
      return is_used

def 					UpperCAmelCase_							(      __UpperCAmelCase						:	Dict			,			__UpperCAmelCase						:	List[str]			,			__UpperCAmelCase						:	Union[str, Any]							)					-> List[Any]:
      SCREAMING_SNAKE_CASE_							=			[]
      SCREAMING_SNAKE_CASE_							=			fairseq_model.state_dict()

      SCREAMING_SNAKE_CASE_							=			hf_model.wavaveca.feature_extractor

      for name, value in fairseq_dict.items():
            SCREAMING_SNAKE_CASE_							=			False
            if "conv_layers" in name:
                  load_conv_layer(
                      __UpperCAmelCase			,			__UpperCAmelCase			,			__UpperCAmelCase			,			__UpperCAmelCase			,			hf_model.config.feat_extract_norm == 'group'			,			)
                  SCREAMING_SNAKE_CASE_							=			True
            else:
                  SCREAMING_SNAKE_CASE_							=			load_wavaveca_layer(__UpperCAmelCase			,			__UpperCAmelCase			,			__UpperCAmelCase							)
            if not is_used:
                  unused_weights.append(__UpperCAmelCase							)

      logger.warning(f"Unused weights: {unused_weights}"							)

def 					UpperCAmelCase_							(      __UpperCAmelCase						:	List[str]			,			__UpperCAmelCase						:	Any			,			__UpperCAmelCase						:	Optional[int]			,			__UpperCAmelCase						:	Dict			,			__UpperCAmelCase						:	List[str]							)					-> Tuple:
      SCREAMING_SNAKE_CASE_							=			full_name.split('conv_layers.'							)[-1]
      SCREAMING_SNAKE_CASE_							=			name.split('.'							)
      SCREAMING_SNAKE_CASE_							=			int(items[0]							)
      SCREAMING_SNAKE_CASE_							=			int(items[1]							)

      if type_id == 0:
            if "bias" in name:
                  if value.shape != feature_extractor.conv_layers[layer_id].conv.bias.data.shape:
                        raise ValueError(
                            f"{full_name} has size {value.shape}, but"
                            f" {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found."							)
                  SCREAMING_SNAKE_CASE_							=			value
                  logger.info(f"Feat extract conv layer {layer_id} was initialized from {full_name}."							)
            elif "weight" in name:
                  if value.shape != feature_extractor.conv_layers[layer_id].conv.weight.data.shape:
                        raise ValueError(
                            f"{full_name} has size {value.shape}, but"
                            f" {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found."							)
                  SCREAMING_SNAKE_CASE_							=			value
                  logger.info(f"Feat extract conv layer {layer_id} was initialized from {full_name}."							)
      elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
            if "bias" in name:
                  if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape:
                        raise ValueError(
                            f"{full_name} has size {value.shape}, but"
                            f" {feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape} was found."							)
                  SCREAMING_SNAKE_CASE_							=			value
                  logger.info(f"Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."							)
            elif "weight" in name:
                  if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape:
                        raise ValueError(
                            f"{full_name} has size {value.shape}, but"
                            f" {feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape} was found."							)
                  SCREAMING_SNAKE_CASE_							=			value
                  logger.info(f"Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."							)
      else:
            unused_weights.append(__UpperCAmelCase							)

@torch.no_grad()
def 					UpperCAmelCase_							(      __UpperCAmelCase						:	List[str]			,			__UpperCAmelCase						:	List[Any]			,			__UpperCAmelCase						:	Any=None			,			__UpperCAmelCase						:	Optional[Any]=None			,			__UpperCAmelCase						:	List[Any]=True			,			__UpperCAmelCase						:	List[str]=False							)					-> Tuple:
      if config_path is not None:
            SCREAMING_SNAKE_CASE_							=			WavaVecaConfig.from_pretrained(__UpperCAmelCase							)
      else:
            SCREAMING_SNAKE_CASE_							=			WavaVecaConfig()

      if is_seq_class:
            SCREAMING_SNAKE_CASE_							=			read_txt_into_dict(__UpperCAmelCase							)
            SCREAMING_SNAKE_CASE_							=			idalabel
            SCREAMING_SNAKE_CASE_							=			WavaVecaForSequenceClassification(__UpperCAmelCase							)
            SCREAMING_SNAKE_CASE_							=			WavaVecaFeatureExtractor(
                feature_size=1			,			sampling_rate=1_60_00			,			padding_value=0			,			do_normalize=__UpperCAmelCase			,			return_attention_mask=__UpperCAmelCase			,			)
            feature_extractor.save_pretrained(__UpperCAmelCase							)

      elif is_finetuned:
            if dict_path:
                  SCREAMING_SNAKE_CASE_							=			Dictionary.load(__UpperCAmelCase							)

                  # important change bos & pad token id since CTC symbol is <pad> and
                  # not <s> as in fairseq
                  SCREAMING_SNAKE_CASE_							=			target_dict.pad_index
                  SCREAMING_SNAKE_CASE_							=			target_dict.bos_index
                  SCREAMING_SNAKE_CASE_							=			target_dict.eos_index
                  SCREAMING_SNAKE_CASE_							=			len(target_dict.symbols							)
                  SCREAMING_SNAKE_CASE_							=			os.path.join(__UpperCAmelCase			,			'vocab.json'							)
                  if not os.path.isdir(__UpperCAmelCase							):
                        logger.error('--pytorch_dump_folder_path ({}) should be a directory'.format(__UpperCAmelCase							)							)
                        return
                  os.makedirs(__UpperCAmelCase			,			exist_ok=__UpperCAmelCase							)
                  SCREAMING_SNAKE_CASE_							=			target_dict.indices

                  # fairseq has the <pad> and <s> switched
                  SCREAMING_SNAKE_CASE_							=			0
                  SCREAMING_SNAKE_CASE_							=			1
                  with open(__UpperCAmelCase			,			'w'			,			encoding='utf-8'							) as vocab_handle:
                        json.dump(__UpperCAmelCase			,			__UpperCAmelCase							)
                  SCREAMING_SNAKE_CASE_							=			WavaVecaCTCTokenizer(
                      __UpperCAmelCase			,			unk_token=target_dict.unk_word			,			pad_token=target_dict.pad_word			,			bos_token=target_dict.bos_word			,			eos_token=target_dict.eos_word			,			word_delimiter_token='|'			,			do_lower_case=__UpperCAmelCase			,			)
                  SCREAMING_SNAKE_CASE_							=			True if config.feat_extract_norm == 'layer' else False
                  SCREAMING_SNAKE_CASE_							=			WavaVecaFeatureExtractor(
                      feature_size=1			,			sampling_rate=1_60_00			,			padding_value=0			,			do_normalize=__UpperCAmelCase			,			return_attention_mask=__UpperCAmelCase			,			)
                  SCREAMING_SNAKE_CASE_							=			WavaVecaProcessor(feature_extractor=__UpperCAmelCase			,			tokenizer=__UpperCAmelCase							)
                  processor.save_pretrained(__UpperCAmelCase							)

            SCREAMING_SNAKE_CASE_							=			WavaVecaForCTC(__UpperCAmelCase							)
      else:
            SCREAMING_SNAKE_CASE_							=			WavaVecaForPreTraining(__UpperCAmelCase							)

      if is_finetuned or is_seq_class:
            SCREAMING_SNAKE_CASE_						, SCREAMING_SNAKE_CASE_						, SCREAMING_SNAKE_CASE_							=			fairseq.checkpoint_utils.load_model_ensemble_and_task(
                [checkpoint_path]			,			arg_overrides={'data': '/'.join(dict_path.split('/'							)[:-1]							)}							)
      else:
            SCREAMING_SNAKE_CASE_							=			argparse.Namespace(task='audio_pretraining'							)
            SCREAMING_SNAKE_CASE_							=			fairseq.tasks.setup_task(__UpperCAmelCase							)

            SCREAMING_SNAKE_CASE_						, SCREAMING_SNAKE_CASE_						, SCREAMING_SNAKE_CASE_							=			fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]			,			task=__UpperCAmelCase							)

      SCREAMING_SNAKE_CASE_							=			model[0].eval()

      recursively_load_weights(__UpperCAmelCase			,			__UpperCAmelCase			,			not is_finetuned							)

      hf_wavavec.save_pretrained(__UpperCAmelCase							)


if __name__ == "__main__":
   lowerCamelCase__ :   Tuple					=							argparse.ArgumentParser()
   parser.add_argument('--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model.')
   parser.add_argument('--checkpoint_path', default=None, type=str, help='Path to fairseq checkpoint')
   parser.add_argument('--dict_path', default=None, type=str, help='Path to dict of fine-tuned model')
   parser.add_argument('--config_path', default=None, type=str, help='Path to hf config.json of model to convert')
   parser.add_argument(
       '--not_finetuned', action='store_true', help='Whether the model to convert is a fine-tuned model or not'
   )
   parser.add_argument(
       '--is_seq_class',
       action='store_true',
       help='Whether the model to convert is a fine-tuned sequence classification model or not',
   )
   lowerCamelCase__ :   Any					=							parser.parse_args()

   lowerCamelCase__ :   Optional[Any]					=							not args.not_finetuned and not args.is_seq_class
   convert_wavaveca_checkpoint(
       args.checkpoint_path,
       args.pytorch_dump_folder_path,
       args.config_path,
       args.dict_path,
       is_finetuned,
       args.is_seq_class,
   )





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 					UpperCAmelCase_							(      __UpperCAmelCase						:	Any			,			__UpperCAmelCase						:	str=False							)					-> List[str]:
      SCREAMING_SNAKE_CASE_							=			[]
      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"
            SCREAMING_SNAKE_CASE_							=			[(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 					UpperCAmelCase_							(      __UpperCAmelCase						:	Union[str, Any]			,			__UpperCAmelCase						:	Union[str, Any]			,			__UpperCAmelCase						:	Optional[Any]=False							)					-> Any:
      for i in range(config.num_hidden_layers							):
            if base_model:
                  SCREAMING_SNAKE_CASE_							=			''
            else:
                  SCREAMING_SNAKE_CASE_							=			'vit.'
            # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
            SCREAMING_SNAKE_CASE_							=			state_dict.pop(f"module.blocks.{i}.attn.qkv.weight"							)
            SCREAMING_SNAKE_CASE_							=			state_dict.pop(f"module.blocks.{i}.attn.qkv.bias"							)
            # next, add query, keys and values (in that order) to the state dict
            SCREAMING_SNAKE_CASE_							=			in_proj_weight[
                : config.hidden_size, :
            ]
            SCREAMING_SNAKE_CASE_							=			in_proj_bias[: config.hidden_size]
            SCREAMING_SNAKE_CASE_							=			in_proj_weight[
                config.hidden_size : config.hidden_size * 2, :
            ]
            SCREAMING_SNAKE_CASE_							=			in_proj_bias[
                config.hidden_size : config.hidden_size * 2
            ]
            SCREAMING_SNAKE_CASE_							=			in_proj_weight[
                -config.hidden_size :, :
            ]
            SCREAMING_SNAKE_CASE_							=			in_proj_bias[-config.hidden_size :]

def 					UpperCAmelCase_							(      __UpperCAmelCase						:	List[Any]							)					-> Union[str, Any]:
      SCREAMING_SNAKE_CASE_							=			['head.weight', 'head.bias']
      for k in ignore_keys:
            state_dict.pop(__UpperCAmelCase			,			__UpperCAmelCase							)

def 					UpperCAmelCase_							(      __UpperCAmelCase						:	Optional[int]							)					-> str:
      # projection head is used in the self-supervised pre-training in MSN,
      # for downstream task it's not needed.
      SCREAMING_SNAKE_CASE_							=			[
          '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(__UpperCAmelCase			,			__UpperCAmelCase							)

def 					UpperCAmelCase_							(      __UpperCAmelCase						:	Dict			,			__UpperCAmelCase						:	Any			,			__UpperCAmelCase						:	Any							)					-> Dict:
      SCREAMING_SNAKE_CASE_							=			dct.pop(__UpperCAmelCase							)
      SCREAMING_SNAKE_CASE_							=			val

def 					UpperCAmelCase_							(      __UpperCAmelCase						:	Optional[int]			,			__UpperCAmelCase						:	Tuple							)					-> List[str]:
      SCREAMING_SNAKE_CASE_							=			ViTMSNConfig()
      SCREAMING_SNAKE_CASE_							=			10_00

      SCREAMING_SNAKE_CASE_							=			'datasets/huggingface/label-files'
      SCREAMING_SNAKE_CASE_							=			'imagenet-1k-id2label.json'
      SCREAMING_SNAKE_CASE_							=			json.load(open(hf_hub_download(__UpperCAmelCase			,			__UpperCAmelCase							)			,			'r'							)							)
      SCREAMING_SNAKE_CASE_							=			{int(__UpperCAmelCase							): v for k, v in idalabel.items()}
      SCREAMING_SNAKE_CASE_							=			idalabel
      SCREAMING_SNAKE_CASE_							=			{v: k for k, v in idalabel.items()}

      if "s16" in checkpoint_url:
            SCREAMING_SNAKE_CASE_							=			3_84
            SCREAMING_SNAKE_CASE_							=			15_36
            SCREAMING_SNAKE_CASE_							=			6
      elif "l16" in checkpoint_url:
            SCREAMING_SNAKE_CASE_							=			10_24
            SCREAMING_SNAKE_CASE_							=			40_96
            SCREAMING_SNAKE_CASE_							=			24
            SCREAMING_SNAKE_CASE_							=			16
            SCREAMING_SNAKE_CASE_							=			0.1
      elif "b4" in checkpoint_url:
            SCREAMING_SNAKE_CASE_							=			4
      elif "l7" in checkpoint_url:
            SCREAMING_SNAKE_CASE_							=			7
            SCREAMING_SNAKE_CASE_							=			10_24
            SCREAMING_SNAKE_CASE_							=			40_96
            SCREAMING_SNAKE_CASE_							=			24
            SCREAMING_SNAKE_CASE_							=			16
            SCREAMING_SNAKE_CASE_							=			0.1

      SCREAMING_SNAKE_CASE_							=			ViTMSNModel(__UpperCAmelCase							)

      SCREAMING_SNAKE_CASE_							=			torch.hub.load_state_dict_from_url(__UpperCAmelCase			,			map_location='cpu'							)['target_encoder']

      SCREAMING_SNAKE_CASE_							=			ViTImageProcessor(size=config.image_size							)

      remove_projection_head(__UpperCAmelCase							)
      SCREAMING_SNAKE_CASE_							=			create_rename_keys(__UpperCAmelCase			,			base_model=__UpperCAmelCase							)

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

      model.load_state_dict(__UpperCAmelCase							)
      model.eval()

      SCREAMING_SNAKE_CASE_							=			'http://images.cocodataset.org/val2017/000000039769.jpg'

      SCREAMING_SNAKE_CASE_							=			Image.open(requests.get(__UpperCAmelCase			,			stream=__UpperCAmelCase							).raw							)
      SCREAMING_SNAKE_CASE_							=			ViTImageProcessor(
          size=config.image_size			,			image_mean=__UpperCAmelCase			,			image_std=__UpperCAmelCase							)
      SCREAMING_SNAKE_CASE_							=			image_processor(images=__UpperCAmelCase			,			return_tensors='pt'							)

      # forward pass
      torch.manual_seed(2							)
      SCREAMING_SNAKE_CASE_							=			model(**__UpperCAmelCase							)
      SCREAMING_SNAKE_CASE_							=			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:
            SCREAMING_SNAKE_CASE_							=			torch.tensor([[-1.0_9_1_5, -1.4_8_7_6, -1.1_8_0_9]]							)
      elif "b16" in checkpoint_url:
            SCREAMING_SNAKE_CASE_							=			torch.tensor([[1_4.2_8_8_9, -1_8.9_0_4_5, 1_1.7_2_8_1]]							)
      elif "l16" in checkpoint_url:
            SCREAMING_SNAKE_CASE_							=			torch.tensor([[4_1.5_0_2_8, -2_2.8_6_8_1, 4_5.6_4_7_5]]							)
      elif "b4" in checkpoint_url:
            SCREAMING_SNAKE_CASE_							=			torch.tensor([[-4.3_8_6_8, 5.2_9_3_2, -0.4_1_3_7]]							)
      else:
            SCREAMING_SNAKE_CASE_							=			torch.tensor([[-0.1_7_9_2, -0.6_4_6_5, 2.4_2_6_3]]							)

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

      print(f"Saving model to {pytorch_dump_folder_path}"							)
      model.save_pretrained(__UpperCAmelCase							)

      print(f"Saving image processor to {pytorch_dump_folder_path}"							)
      image_processor.save_pretrained(__UpperCAmelCase							)


if __name__ == "__main__":
   lowerCamelCase__ :   int					=							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.'
   )

   lowerCamelCase__ :   Union[str, Any]					=							parser.parse_args()
   convert_vit_msn_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)