


import argparse
import json
from pathlib import Path

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

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


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





def      UpperCAmelCase      (		lowerCamelCase_  :Union[str, Any]				,      lowerCamelCase_  :Any=False ):







		'''simple docstring'''
		snake_case_		:							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'''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"""),
				    ] )

				# if just the base model, we should remove "vit" from all keys that start with "vit"
				snake_case_		:							Optional[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      UpperCAmelCase      (		lowerCamelCase_  :Any				,      lowerCamelCase_  :Any				,      lowerCamelCase_  :int=False ):







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





def      UpperCAmelCase      (		lowerCamelCase_  :Optional[Any] ):







		'''simple docstring'''
		snake_case_		:							Any			       =						["""head.weight""", """head.bias"""]
		for k in ignore_keys:
				state_dict.pop(lowerCamelCase_				,      lowerCamelCase_ )





def      UpperCAmelCase      (		lowerCamelCase_  :str				,      lowerCamelCase_  :Optional[Any]				,      lowerCamelCase_  :Optional[int] ):







		'''simple docstring'''
		snake_case_		:							Optional[int]			       =						dct.pop(lowerCamelCase_ )
		snake_case_		:							int			       =						val





def      UpperCAmelCase      (		):







		'''simple docstring'''
		snake_case_		:							Union[str, Any]			       =						"""http://images.cocodataset.org/val2017/000000039769.jpg"""
		snake_case_		:							Union[str, Any]			       =						Image.open(requests.get(lowerCamelCase_				,      stream=lowerCamelCase_ ).raw )
		return im





@torch.no_grad()
def      UpperCAmelCase      (		lowerCamelCase_  :List[Any]				,      lowerCamelCase_  :str				,      lowerCamelCase_  :Dict=True ):







		'''simple docstring'''
		snake_case_		:							List[Any]			       =						ViTConfig()
		# patch_size
		if model_name[-1] == "8":
				snake_case_		:							Optional[int]			       =						8
		# set labels if required
		if not base_model:
				snake_case_		:							str			       =						10_00
				snake_case_		:							Tuple			       =						"""huggingface/label-files"""
				snake_case_		:							Union[str, Any]			       =						"""imagenet-1k-id2label.json"""
				snake_case_		:							int			       =						json.load(open(hf_hub_download(lowerCamelCase_				,      lowerCamelCase_				,      repo_type="""dataset""" )				,      """r""" ) )
				snake_case_		:							Tuple			       =						{int(lowerCamelCase_ ): v for k, v in idalabel.items()}
				snake_case_		:							Union[str, Any]			       =						idalabel
				snake_case_		:							int			       =						{v: k for k, v in idalabel.items()}
		# size of the architecture
		if model_name in ["dino_vits8", "dino_vits16"]:
				snake_case_		:							int			       =						3_84
				snake_case_		:							Dict			       =						15_36
				snake_case_		:							Any			       =						12
				snake_case_		:							Optional[Any]			       =						6

		# load original model from torch hub
		snake_case_		:							List[Any]			       =						torch.hub.load("""facebookresearch/dino:main"""				,      lowerCamelCase_ )
		original_model.eval()

		# load state_dict of original model, remove and rename some keys
		snake_case_		:							str			       =						original_model.state_dict()
		if base_model:
				remove_classification_head_(lowerCamelCase_ )
		snake_case_		:							Union[str, Any]			       =						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_				,      lowerCamelCase_ )

		# load HuggingFace model
		if base_model:
				snake_case_		:							Optional[int]			       =						ViTModel(lowerCamelCase_				,      add_pooling_layer=lowerCamelCase_ ).eval()
		else:
				snake_case_		:							int			       =						ViTForImageClassification(lowerCamelCase_ ).eval()
		model.load_state_dict(lowerCamelCase_ )

		# Check outputs on an image, prepared by ViTImageProcessor
		snake_case_		:							str			       =						ViTImageProcessor()
		snake_case_		:							str			       =						image_processor(images=prepare_img()				,      return_tensors="""pt""" )
		snake_case_		:							Union[str, Any]			       =						encoding["""pixel_values"""]
		snake_case_		:							Optional[int]			       =						model(lowerCamelCase_ )

		if base_model:
				snake_case_		:							List[str]			       =						original_model(lowerCamelCase_ )
				assert torch.allclose(lowerCamelCase_				,      outputs.last_hidden_state[:, 0, :]				,      atol=1E-1 )
		else:
				snake_case_		:							str			       =						original_model(lowerCamelCase_ )
				assert logits.shape == outputs.logits.shape
				assert torch.allclose(lowerCamelCase_				,      outputs.logits				,      atol=1E-3 )

		Path(lowerCamelCase_ ).mkdir(exist_ok=lowerCamelCase_ )
		print(F'''Saving model {model_name} 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__":
							__A	:	Dict     =     argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    '--model_name',
							    default='dino_vitb16',
							    type=str,
							    help='Name of the model trained with DINO you\'d like to convert.',
							)
							parser.add_argument(
							    '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
							)
							parser.add_argument(
							    '--base_model',
							    action='store_true',
							    help='Whether to only convert the base model (no projection head weights).',
							)

							parser.set_defaults(base_model=True)
							__A	:	Optional[Any]     =     parser.parse_args()
							convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)



'''simple docstring'''






import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch

							from transformers import EsmForMaskedLM, EsmForSequenceClassification, EsmForTokenClassification, EsmModel
							from transformers.models.esm.modeling_esm import (
							    ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
							    EsmEmbeddings,
							    create_position_ids_from_input_ids,
							)
class 						__UpperCamelCase	:





							def __init__(  self		:Union[str, Any]  ,_UpperCamelCase		:int  ,_UpperCamelCase		:str=1_3  ,_UpperCamelCase		:Tuple=7  ,_UpperCamelCase		:Tuple=False  ,_UpperCamelCase		:Tuple=True  ,_UpperCamelCase		:Union[str, Any]=False  ,_UpperCamelCase		:int=True  ,_UpperCamelCase		:List[str]=3_3  ,_UpperCamelCase		:Any=3_2  ,_UpperCamelCase		:Any=5  ,_UpperCamelCase		:List[str]=4  ,_UpperCamelCase		:Tuple=3_7  ,_UpperCamelCase		:Optional[Any]="gelu"  ,_UpperCamelCase		:Any=0.1  ,_UpperCamelCase		:List[Any]=0.1  ,_UpperCamelCase		:Any=5_1_2  ,_UpperCamelCase		:Tuple=1_6  ,_UpperCamelCase		:Any=2  ,_UpperCamelCase		:Optional[Any]=0.02  ,_UpperCamelCase		:List[str]=3  ,_UpperCamelCase		:Union[str, Any]=4  ,_UpperCamelCase		:Dict=None  ,):
									snake_case_		:							Tuple			       =						parent
									snake_case_		:							List[str]			       =						batch_size
									snake_case_		:							List[str]			       =						seq_length
									snake_case_		:							Any			       =						is_training
									snake_case_		:							List[Any]			       =						use_input_mask
									snake_case_		:							int			       =						use_token_type_ids
									snake_case_		:							Optional[int]			       =						use_labels
									snake_case_		:							List[str]			       =						vocab_size
									snake_case_		:							Dict			       =						hidden_size
									snake_case_		:							Optional[int]			       =						num_hidden_layers
									snake_case_		:							Optional[Any]			       =						num_attention_heads
									snake_case_		:							int			       =						intermediate_size
									snake_case_		:							Optional[Any]			       =						hidden_act
									snake_case_		:							int			       =						hidden_dropout_prob
									snake_case_		:							Optional[int]			       =						attention_probs_dropout_prob
									snake_case_		:							Optional[int]			       =						max_position_embeddings
									snake_case_		:							List[str]			       =						type_vocab_size
									snake_case_		:							int			       =						type_sequence_label_size
									snake_case_		:							Dict			       =						initializer_range
									snake_case_		:							Any			       =						num_labels
									snake_case_		:							Any			       =						num_choices
									snake_case_		:							Tuple			       =						scope





							def 		a__   (  self		:str			):
									snake_case_		:							Dict			       =						ids_tensor([self.batch_size, self.seq_length]  ,self.vocab_size			)

									snake_case_		:							str			       =						None
									if self.use_input_mask:
											snake_case_		:							Optional[int]			       =						random_attention_mask([self.batch_size, self.seq_length]			)

									snake_case_		:							Dict			       =						None
									snake_case_		:							List[str]			       =						None
									snake_case_		:							Tuple			       =						None
									if self.use_labels:
											snake_case_		:							Dict			       =						ids_tensor([self.batch_size]  ,self.type_sequence_label_size			)
											snake_case_		:							Any			       =						ids_tensor([self.batch_size, self.seq_length]  ,self.num_labels			)
											snake_case_		:							Union[str, Any]			       =						ids_tensor([self.batch_size]  ,self.num_choices			)

									snake_case_		:							Dict			       =						self.get_config()

									return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels





							def 		a__   (  self		:Optional[Any]			):
									return EsmConfig(
									    vocab_size=self.vocab_size  ,hidden_size=self.hidden_size  ,pad_token_id=1  ,num_hidden_layers=self.num_hidden_layers  ,num_attention_heads=self.num_attention_heads  ,intermediate_size=self.intermediate_size  ,hidden_act=self.hidden_act  ,hidden_dropout_prob=self.hidden_dropout_prob  ,attention_probs_dropout_prob=self.attention_probs_dropout_prob  ,max_position_embeddings=self.max_position_embeddings  ,type_vocab_size=self.type_vocab_size  ,initializer_range=self.initializer_range  ,)





							def 		a__   (  self		:str  ,_UpperCamelCase		:List[Any]  ,_UpperCamelCase		:List[str]  ,_UpperCamelCase		:Any  ,_UpperCamelCase		:Optional[int]  ,_UpperCamelCase		:List[Any]  ,_UpperCamelCase		:Optional[int]			):
									snake_case_		:							Union[str, Any]			       =						EsmModel(config=_UpperCamelCase			)
									model.to(_UpperCamelCase			)
									model.eval()
									snake_case_		:							int			       =						model(_UpperCamelCase  ,attention_mask=_UpperCamelCase			)
									snake_case_		:							str			       =						model(_UpperCamelCase			)
									snake_case_		:							Optional[int]			       =						model(_UpperCamelCase			)

									self.parent.assertEqual(result.last_hidden_state.shape  ,(self.batch_size, self.seq_length, self.hidden_size)			)
									self.parent.assertEqual(result.pooler_output.shape  ,(self.batch_size, self.hidden_size)			)





							def 		a__   (  self		:Any  ,_UpperCamelCase		:Optional[Any]  ,_UpperCamelCase		:Optional[Any]  ,_UpperCamelCase		:Any  ,_UpperCamelCase		:Dict  ,_UpperCamelCase		:Union[str, Any]  ,_UpperCamelCase		:Optional[int]			):
									snake_case_		:							str			       =						EsmForMaskedLM(config=_UpperCamelCase			)
									model.to(_UpperCamelCase			)
									model.eval()
									snake_case_		:							Dict			       =						model(_UpperCamelCase  ,attention_mask=_UpperCamelCase  ,labels=_UpperCamelCase			)
									self.parent.assertEqual(result.logits.shape  ,(self.batch_size, self.seq_length, self.vocab_size)			)





							def 		a__   (  self		:Optional[int]  ,_UpperCamelCase		:Optional[Any]  ,_UpperCamelCase		:List[str]  ,_UpperCamelCase		:Union[str, Any]  ,_UpperCamelCase		:Union[str, Any]  ,_UpperCamelCase		:Union[str, Any]  ,_UpperCamelCase		:Any			):
									snake_case_		:							List[Any]			       =						self.num_labels
									snake_case_		:							int			       =						EsmForTokenClassification(config=_UpperCamelCase			)
									model.to(_UpperCamelCase			)
									model.eval()
									snake_case_		:							int			       =						model(_UpperCamelCase  ,attention_mask=_UpperCamelCase  ,labels=_UpperCamelCase			)
									self.parent.assertEqual(result.logits.shape  ,(self.batch_size, self.seq_length, self.num_labels)			)





							def 		a__   (  self		:Union[str, Any]			):
									snake_case_		:							Dict			       =						self.prepare_config_and_inputs()
									(
									    (
									    snake_case_
									)       ,      (
									    snake_case_
									)       ,      (
									    snake_case_
									)       ,      (
									    snake_case_
									)       ,      (
									    snake_case_
									)       ,      (
									    snake_case_
									)       ,      
									)		:							Optional[int]			       =						config_and_inputs
									snake_case_		:							List[str]			       =						{"""input_ids""": input_ids, """attention_mask""": input_mask}
									return config, inputs_dict





@require_torch
class 						__UpperCamelCase	(						lowercase__						,  lowercase__						,  unittest.TestCase							):
							lowercase			:		Optional[int]   		=  False

							lowercase			:		List[str]   		=  (
							    (
							        EsmForMaskedLM,
							        EsmModel,
							        EsmForSequenceClassification,
							        EsmForTokenClassification,
							    )
							    if is_torch_available()
							    else ()
							)
							lowercase			:		int   		=  ()
							lowercase			:		List[str]   		=  (
							    {
							        'feature-extraction': EsmModel,
							        'fill-mask': EsmForMaskedLM,
							        'text-classification': EsmForSequenceClassification,
							        'token-classification': EsmForTokenClassification,
							        'zero-shot': EsmForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)
							lowercase			:		str   		=  True





							def 		a__   (  self		:Any			):
									snake_case_		:							Any			       =						EsmModelTester(self			)
									snake_case_		:							str			       =						ConfigTester(self  ,config_class=_UpperCamelCase  ,hidden_size=3_7			)





							def 		a__   (  self		:Optional[Any]			):
									self.config_tester.run_common_tests()





							def 		a__   (  self		:Optional[Any]			):
									snake_case_		:							str			       =						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*_UpperCamelCase			)





							def 		a__   (  self		:Dict			):
									snake_case_		:							Dict			       =						self.model_tester.prepare_config_and_inputs()
									for type in ["absolute", "relative_key", "relative_key_query"]:
											snake_case_		:							int			       =						type
											self.model_tester.create_and_check_model(*_UpperCamelCase			)





							def 		a__   (  self		:Dict			):
									snake_case_		:							str			       =						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_lm(*_UpperCamelCase			)





							def 		a__   (  self		:List[str]			):
									snake_case_		:							List[str]			       =						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_token_classification(*_UpperCamelCase			)





							@slow
							def 		a__   (  self		:Union[str, Any]			):
									for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											snake_case_		:							int			       =						EsmModel.from_pretrained(_UpperCamelCase			)
											self.assertIsNotNone(_UpperCamelCase			)





							def 		a__   (  self		:Tuple			):
									snake_case_		:							Dict			       =						self.model_tester.prepare_config_and_inputs()[0]
									snake_case_		:							Optional[int]			       =						EsmEmbeddings(config=_UpperCamelCase			)

									snake_case_		:							List[Any]			       =						torch.as_tensor([[1_2, 3_1, 1_3, model.padding_idx]]			)
									snake_case_		:							Union[str, Any]			       =						torch.as_tensor(
									    [
									        [
									            0 + model.padding_idx + 1,
									            1 + model.padding_idx + 1,
									            2 + model.padding_idx + 1,
									            model.padding_idx,
									        ]
									    ]			)
									snake_case_		:							Optional[Any]			       =						create_position_ids_from_input_ids(_UpperCamelCase  ,model.padding_idx			)
									self.assertEqual(position_ids.shape  ,expected_positions.shape			)
									self.assertTrue(torch.all(torch.eq(_UpperCamelCase  ,_UpperCamelCase			)			)			)





							def 		a__   (  self		:List[Any]			):
									snake_case_		:							Union[str, Any]			       =						self.model_tester.prepare_config_and_inputs()[0]
									snake_case_		:							List[Any]			       =						EsmEmbeddings(config=_UpperCamelCase			)

									snake_case_		:							Dict			       =						torch.empty(2  ,4  ,3_0			)
									snake_case_		:							List[Any]			       =						[
									    0 + embeddings.padding_idx + 1,
									    1 + embeddings.padding_idx + 1,
									    2 + embeddings.padding_idx + 1,
									    3 + embeddings.padding_idx + 1,
									]
									snake_case_		:							Any			       =						torch.as_tensor([expected_single_positions, expected_single_positions]			)
									snake_case_		:							str			       =						embeddings.create_position_ids_from_inputs_embeds(_UpperCamelCase			)
									self.assertEqual(position_ids.shape  ,expected_positions.shape			)
									self.assertTrue(torch.all(torch.eq(_UpperCamelCase  ,_UpperCamelCase			)			)			)





							@unittest.skip("""Esm does not support embedding resizing"""			)
							def 		a__   (  self		:Optional[Any]			):
									pass





							@unittest.skip("""Esm does not support embedding resizing"""			)
							def 		a__   (  self		:Optional[int]			):
									pass





							@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""			)
							def 		a__   (  self		:Optional[int]			):
									pass





@require_torch
class 						__UpperCamelCase	(						lowercase__							):





							@slow
							def 		a__   (  self		:Any			):
									with torch.no_grad():
											snake_case_		:							Optional[Any]			       =						EsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D"""			)
											model.eval()
											snake_case_		:							List[Any]			       =						torch.tensor([[0, 1, 2, 3, 4, 5]]			)
											snake_case_		:							Dict			       =						model(_UpperCamelCase			)[0]

											snake_case_		:							Optional[int]			       =						3_3

											snake_case_		:							List[Any]			       =						torch.Size((1, 6, vocab_size)			)
											self.assertEqual(output.shape  ,_UpperCamelCase			)

											snake_case_		:							Union[str, Any]			       =						torch.tensor(
											    [[[8.92_15, -10.58_98, -6.46_71], [-6.39_67, -13.91_14, -1.12_12], [-7.78_12, -13.95_16, -3.74_06]]]			)
											self.assertTrue(torch.allclose(output[:, :3, :3]  ,_UpperCamelCase  ,atol=1E-4			)			)





							@slow
							def 		a__   (  self		:List[Any]			):
									with torch.no_grad():
											snake_case_		:							List[Any]			       =						EsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D"""			)
											model.eval()

											snake_case_		:							str			       =						torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]			)
											snake_case_		:							Optional[Any]			       =						model(_UpperCamelCase			)[0]
											# compare the actual values for a slice.
											snake_case_		:							List[str]			       =						torch.tensor(
											    [[[0.14_44, 0.54_13, 0.32_48], [0.30_34, 0.00_53, 0.31_08], [0.32_28, -0.24_99, 0.34_15]]]			)
											self.assertTrue(torch.allclose(output[:, :3, :3]  ,_UpperCamelCase  ,atol=1E-4			)			)