




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




def 		lowerCAmelCase_ (						lowerCamelCase ):
				__magic_name__						:						int										=["""head.weight""", """head.bias"""]
				for k in ignore_keys:
								state_dict.pop(lowerCamelCase						,  lowerCamelCase )




def 		lowerCAmelCase_ (						lowerCamelCase ):
				# projection head is used in the self-supervised pre-training in MSN,
				# for downstream task it's not needed.
				__magic_name__						:						List[Any]										=[
				    """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 		lowerCAmelCase_ (						lowerCamelCase						,  lowerCamelCase						,  lowerCamelCase ):
				__magic_name__						:						Optional[Any]										=dct.pop(lowerCamelCase )
				__magic_name__						:						List[str]										=val




def 		lowerCAmelCase_ (						lowerCamelCase						,  lowerCamelCase ):
				__magic_name__						:						str										=ViTMSNConfig()
				__magic_name__						:						str										=1000

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

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

				__magic_name__						:						List[str]										=ViTMSNModel(lowerCamelCase )

				__magic_name__						:						str										=torch.hub.load_state_dict_from_url(lowerCamelCase						,  map_location="""cpu""" )["""target_encoder"""]

				__magic_name__						:						Dict										=ViTImageProcessor(size=config.image_size )

				remove_projection_head(lowerCamelCase )
				__magic_name__						:						Dict										=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()

				__magic_name__						:						List[Any]										="""http://images.cocodataset.org/val2017/000000039769.jpg"""

				__magic_name__						:						Dict										=Image.open(requests.get(lowerCamelCase						,  stream=lowerCamelCase ).raw )
				__magic_name__						:						Tuple										=ViTImageProcessor(
				    size=config.image_size						,  image_mean=lowerCamelCase						,  image_std=lowerCamelCase )
				__magic_name__						:						Tuple										=image_processor(images=lowerCamelCase						,  return_tensors="""pt""" )

				# forward pass
				torch.manual_seed(2 )
				__magic_name__						:						Dict										=model(**lowerCamelCase )
				__magic_name__						:						List[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:
								__magic_name__						:						Any										=torch.tensor([[-1.0_9_1_5, -1.4_8_7_6, -1.1_8_0_9]] )
				elif "b16" in checkpoint_url:
								__magic_name__						:						Optional[int]										=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:
								__magic_name__						:						Tuple										=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:
								__magic_name__						:						Tuple										=torch.tensor([[-4.3_8_6_8, 5.2_9_3_2, -0.4_1_3_7]] )
				else:
								__magic_name__						:						List[Any]										=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]						,  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_		:					Optional[Any]          =  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_		:					int          =  parser.parse_args()
							convert_vit_msn_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)










from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


UpperCAmelCase_		:					Tuple          =  logging.get_logger(__name__)

UpperCAmelCase_		:					List[str]          =  {
    "facebook/xlm-roberta-xl": "https://huggingface.co/facebook/xlm-roberta-xl/resolve/main/config.json",
    "facebook/xlm-roberta-xxl": "https://huggingface.co/facebook/xlm-roberta-xxl/resolve/main/config.json",
    # See all XLM-RoBERTa-XL models at https://huggingface.co/models?filter=xlm-roberta-xl
}




class 		__A   (      UpperCamelCase__      ):
						UpperCamelCase      =       """xlm-roberta-xl"""

						def __init__(		self :Dict     ,			__snake_case :Optional[Any]=25_08_80     ,			__snake_case :List[Any]=25_60     ,			__snake_case :Optional[Any]=36     ,			__snake_case :Any=32     ,			__snake_case :int=1_02_40     ,			__snake_case :List[Any]="gelu"     ,			__snake_case :Union[str, Any]=0.1     ,			__snake_case :Optional[Any]=0.1     ,			__snake_case :str=5_14     ,			__snake_case :Union[str, Any]=1     ,			__snake_case :Optional[int]=0.02     ,			__snake_case :str=1E-05     ,			__snake_case :str=1     ,			__snake_case :int=0     ,			__snake_case :Tuple=2     ,			__snake_case :Optional[int]="absolute"     ,			__snake_case :str=True     ,			__snake_case :Any=None     ,			**__snake_case :Dict     ,			):
										'''simple docstring'''

										super().__init__(pad_token_id=__snake_case     ,			bos_token_id=__snake_case     ,			eos_token_id=__snake_case     ,			**__snake_case			)
										__magic_name__						:						List[str]										=vocab_size
										__magic_name__						:						List[str]										=hidden_size
										__magic_name__						:						Union[str, Any]										=num_hidden_layers
										__magic_name__						:						Any										=num_attention_heads
										__magic_name__						:						Any										=hidden_act
										__magic_name__						:						List[str]										=intermediate_size
										__magic_name__						:						Any										=hidden_dropout_prob
										__magic_name__						:						Union[str, Any]										=attention_probs_dropout_prob
										__magic_name__						:						Any										=max_position_embeddings
										__magic_name__						:						Any										=type_vocab_size
										__magic_name__						:						List[str]										=initializer_range
										__magic_name__						:						Optional[int]										=layer_norm_eps
										__magic_name__						:						Dict										=position_embedding_type
										__magic_name__						:						Any										=use_cache
										__magic_name__						:						Dict										=classifier_dropout




class 		__A   (      UpperCamelCase__      ):

						@property
						def      A__    (		self :Dict			):
										'''simple docstring'''

										if self.task == "multiple-choice":
														__magic_name__						:						str										={0: """batch""", 1: """choice""", 2: """sequence"""}
										else:
														__magic_name__						:						Optional[Any]										={0: """batch""", 1: """sequence"""}
										return OrderedDict(
										    [
										        ("""input_ids""", dynamic_axis),
										        ("""attention_mask""", dynamic_axis),
										    ]			)




