



import unittest

from parameterized import parameterized

from transformers import OpenLlamaConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
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 OpenLlamaForCausalLM, OpenLlamaForSequenceClassification, OpenLlamaModel





class 					__lowerCAmelCase :




				def __init__(							self						,				_snake_case						,				_snake_case=13						,				_snake_case=7						,				_snake_case=True						,				_snake_case=True						,				_snake_case=False						,				_snake_case=True						,				_snake_case=99						,				_snake_case=32						,				_snake_case=5						,				_snake_case=4						,				_snake_case=37						,				_snake_case="gelu"						,				_snake_case=0.1						,				_snake_case=0.1						,				_snake_case=512						,				_snake_case=16						,				_snake_case=2						,				_snake_case=0.02						,				_snake_case=3						,				_snake_case=4						,				_snake_case=None						,				):



									"""simple docstring"""





									_lowerCAmelCase						=    parent
									_lowerCAmelCase						=    batch_size
									_lowerCAmelCase						=    seq_length
									_lowerCAmelCase						=    is_training
									_lowerCAmelCase						=    use_input_mask
									_lowerCAmelCase						=    use_token_type_ids
									_lowerCAmelCase						=    use_labels
									_lowerCAmelCase						=    vocab_size
									_lowerCAmelCase						=    hidden_size
									_lowerCAmelCase						=    num_hidden_layers
									_lowerCAmelCase						=    num_attention_heads
									_lowerCAmelCase						=    intermediate_size
									_lowerCAmelCase						=    hidden_act
									_lowerCAmelCase						=    hidden_dropout_prob
									_lowerCAmelCase						=    attention_probs_dropout_prob
									_lowerCAmelCase						=    max_position_embeddings
									_lowerCAmelCase						=    type_vocab_size
									_lowerCAmelCase						=    type_sequence_label_size
									_lowerCAmelCase						=    initializer_range
									_lowerCAmelCase						=    num_labels
									_lowerCAmelCase						=    num_choices
									_lowerCAmelCase						=    scope




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    ids_tensor([self.batch_size, self.seq_length]						,				self.vocab_size		)

									_lowerCAmelCase						=    None
									if self.use_input_mask:
														_lowerCAmelCase						=    random_attention_mask([self.batch_size, self.seq_length]		)

									_lowerCAmelCase						=    None
									if self.use_token_type_ids:
														_lowerCAmelCase						=    ids_tensor([self.batch_size, self.seq_length]						,				self.type_vocab_size		)

									_lowerCAmelCase						=    None
									_lowerCAmelCase						=    None
									_lowerCAmelCase						=    None
									if self.use_labels:
														_lowerCAmelCase						=    ids_tensor([self.batch_size]						,				self.type_sequence_label_size		)
														_lowerCAmelCase						=    ids_tensor([self.batch_size, self.seq_length]						,				self.num_labels		)
														_lowerCAmelCase						=    ids_tensor([self.batch_size]						,				self.num_choices		)

									_lowerCAmelCase						=    self.get_config()

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




				def 							snake_case   (							self		):



									"""simple docstring"""





									return OpenLlamaConfig(
									    vocab_size=self.vocab_size						,				hidden_size=self.hidden_size						,				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						,				is_decoder=_snake_case						,				initializer_range=self.initializer_range						,				use_stable_embedding=_snake_case						,				)




				def 							snake_case   (							self						,				_snake_case						,				_snake_case						,				_snake_case						,				_snake_case						,				_snake_case						,				_snake_case						,				_snake_case		):



									"""simple docstring"""





									_lowerCAmelCase						=    OpenLlamaModel(config=_snake_case		)
									model.to(_snake_case		)
									model.eval()
									_lowerCAmelCase						=    model(_snake_case						,				attention_mask=_snake_case		)
									_lowerCAmelCase						=    model(_snake_case		)
									self.parent.assertEqual(result.last_hidden_state.shape						,				(self.batch_size, self.seq_length, self.hidden_size)		)




				def 							snake_case   (							self						,				_snake_case						,				_snake_case						,				_snake_case						,				_snake_case						,				_snake_case						,				_snake_case						,				_snake_case						,				_snake_case						,				_snake_case						,				):



									"""simple docstring"""





									_lowerCAmelCase						=    True
									_lowerCAmelCase						=    OpenLlamaModel(_snake_case		)
									model.to(_snake_case		)
									model.eval()
									_lowerCAmelCase						=    model(
									    _snake_case						,				attention_mask=_snake_case						,				encoder_hidden_states=_snake_case						,				encoder_attention_mask=_snake_case						,				)
									_lowerCAmelCase						=    model(
									    _snake_case						,				attention_mask=_snake_case						,				encoder_hidden_states=_snake_case						,				)
									_lowerCAmelCase						=    model(_snake_case						,				attention_mask=_snake_case		)
									self.parent.assertEqual(result.last_hidden_state.shape						,				(self.batch_size, self.seq_length, self.hidden_size)		)




				def 							snake_case   (							self						,				_snake_case						,				_snake_case						,				_snake_case						,				_snake_case						,				_snake_case						,				_snake_case						,				_snake_case						,				_snake_case						,				_snake_case						,				):



									"""simple docstring"""





									_lowerCAmelCase						=    OpenLlamaForCausalLM(config=_snake_case		)
									model.to(_snake_case		)
									model.eval()
									_lowerCAmelCase						=    model(_snake_case						,				attention_mask=_snake_case						,				labels=_snake_case		)
									self.parent.assertEqual(result.logits.shape						,				(self.batch_size, self.seq_length, self.vocab_size)		)




				def 							snake_case   (							self						,				_snake_case						,				_snake_case						,				_snake_case						,				_snake_case						,				_snake_case						,				_snake_case						,				_snake_case						,				_snake_case						,				_snake_case						,				):



									"""simple docstring"""





									_lowerCAmelCase						=    True
									_lowerCAmelCase						=    True
									_lowerCAmelCase						=    OpenLlamaForCausalLM(config=_snake_case		)
									model.to(_snake_case		)
									model.eval()

									# first forward pass
									_lowerCAmelCase						=    model(
									    _snake_case						,				attention_mask=_snake_case						,				encoder_hidden_states=_snake_case						,				encoder_attention_mask=_snake_case						,				use_cache=_snake_case						,				)
									_lowerCAmelCase						=    outputs.past_key_values

									# create hypothetical multiple next token and extent to next_input_ids
									_lowerCAmelCase						=    ids_tensor((self.batch_size, 3)						,				config.vocab_size		)
									_lowerCAmelCase						=    ids_tensor((self.batch_size, 3)						,				vocab_size=2		)

									# append to next input_ids and
									_lowerCAmelCase						=    torch.cat([input_ids, next_tokens]						,				dim=-1		)
									_lowerCAmelCase						=    torch.cat([input_mask, next_mask]						,				dim=-1		)

									_lowerCAmelCase						=    model(
									    _snake_case						,				attention_mask=_snake_case						,				encoder_hidden_states=_snake_case						,				encoder_attention_mask=_snake_case						,				output_hidden_states=_snake_case						,				)["""hidden_states"""][0]
									_lowerCAmelCase						=    model(
									    _snake_case						,				attention_mask=_snake_case						,				encoder_hidden_states=_snake_case						,				encoder_attention_mask=_snake_case						,				past_key_values=_snake_case						,				output_hidden_states=_snake_case						,				)["""hidden_states"""][0]

									# select random slice
									_lowerCAmelCase						=    ids_tensor((1,)						,				output_from_past.shape[-1]		).item()
									_lowerCAmelCase						=    output_from_no_past[:, -3:, random_slice_idx].detach()
									_lowerCAmelCase						=    output_from_past[:, :, random_slice_idx].detach()

									self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]		)

									# test that outputs are equal for slice
									self.parent.assertTrue(torch.allclose(_snake_case						,				_snake_case						,				atol=1e-3		)		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    self.prepare_config_and_inputs()
									(
									    (
									    _lowerCAmelCase
									)	,	(
									    _lowerCAmelCase
									)	,	(
									    _lowerCAmelCase
									)	,	(
									    _lowerCAmelCase
									)	,	(
									    _lowerCAmelCase
									)	,	(
									    _lowerCAmelCase
									)	,	(
									    _lowerCAmelCase
									)	,	
									)						=    config_and_inputs
									_lowerCAmelCase						=    {"""input_ids""": input_ids, """attention_mask""": input_mask}
									return config, inputs_dict





@require_torch
class 					__lowerCAmelCase (     lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						unittest.TestCase       ):
				__lowerCamelCase =				(
				    (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
				)
				__lowerCamelCase =				(OpenLlamaForCausalLM,) if is_torch_available() else ()
				__lowerCamelCase =				(
				    {
				        '''feature-extraction''': OpenLlamaModel,
				        '''text-classification''': OpenLlamaForSequenceClassification,
				        '''text-generation''': OpenLlamaForCausalLM,
				        '''zero-shot''': OpenLlamaForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)
				__lowerCamelCase =				False
				__lowerCamelCase =				False




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    OpenLlamaModelTester(self		)
									_lowerCAmelCase						=    ConfigTester(self						,				config_class=_snake_case						,				hidden_size=37		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									self.config_tester.run_common_tests()




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*_snake_case		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    self.model_tester.prepare_config_and_inputs()
									for type in ["absolute", "relative_key", "relative_key_query"]:
														_lowerCAmelCase						=    type
														self.model_tester.create_and_check_model(*_snake_case		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase	,	_lowerCAmelCase						=    self.model_tester.prepare_config_and_inputs_for_common()
									_lowerCAmelCase						=    3
									_lowerCAmelCase						=    input_dict["""input_ids"""]
									_lowerCAmelCase						=    input_ids.ne(1		).to(_snake_case		)
									_lowerCAmelCase						=    ids_tensor([self.model_tester.batch_size]						,				self.model_tester.type_sequence_label_size		)
									_lowerCAmelCase						=    OpenLlamaForSequenceClassification(_snake_case		)
									model.to(_snake_case		)
									model.eval()
									_lowerCAmelCase						=    model(_snake_case						,				attention_mask=_snake_case						,				labels=_snake_case		)
									self.assertEqual(result.logits.shape						,				(self.model_tester.batch_size, self.model_tester.num_labels)		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase	,	_lowerCAmelCase						=    self.model_tester.prepare_config_and_inputs_for_common()
									_lowerCAmelCase						=    3
									_lowerCAmelCase						=    """single_label_classification"""
									_lowerCAmelCase						=    input_dict["""input_ids"""]
									_lowerCAmelCase						=    input_ids.ne(1		).to(_snake_case		)
									_lowerCAmelCase						=    ids_tensor([self.model_tester.batch_size]						,				self.model_tester.type_sequence_label_size		)
									_lowerCAmelCase						=    OpenLlamaForSequenceClassification(_snake_case		)
									model.to(_snake_case		)
									model.eval()
									_lowerCAmelCase						=    model(_snake_case						,				attention_mask=_snake_case						,				labels=_snake_case		)
									self.assertEqual(result.logits.shape						,				(self.model_tester.batch_size, self.model_tester.num_labels)		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase	,	_lowerCAmelCase						=    self.model_tester.prepare_config_and_inputs_for_common()
									_lowerCAmelCase						=    3
									_lowerCAmelCase						=    """multi_label_classification"""
									_lowerCAmelCase						=    input_dict["""input_ids"""]
									_lowerCAmelCase						=    input_ids.ne(1		).to(_snake_case		)
									_lowerCAmelCase						=    ids_tensor(
									    [self.model_tester.batch_size, config.num_labels]						,				self.model_tester.type_sequence_label_size		).to(torch.float		)
									_lowerCAmelCase						=    OpenLlamaForSequenceClassification(_snake_case		)
									model.to(_snake_case		)
									model.eval()
									_lowerCAmelCase						=    model(_snake_case						,				attention_mask=_snake_case						,				labels=_snake_case		)
									self.assertEqual(result.logits.shape						,				(self.model_tester.batch_size, self.model_tester.num_labels)		)




				@unittest.skip("""Open-Llama buffers include complex numbers, which breaks this test"""		)
				def 							snake_case   (							self		):



									"""simple docstring"""





									pass




				@parameterized.expand([("""linear""",), ("""dynamic""",)]		)
				def 							snake_case   (							self						,				_snake_case		):



									"""simple docstring"""





									_lowerCAmelCase	,	_lowerCAmelCase						=    self.model_tester.prepare_config_and_inputs_for_common()
									_lowerCAmelCase						=    ids_tensor([1, 10]						,				config.vocab_size		)
									_lowerCAmelCase						=    ids_tensor([1, int(config.max_position_embeddings * 1.5		)]						,				config.vocab_size		)

									set_seed(42		)  # Fixed seed at init time so the two models get the same random weights
									_lowerCAmelCase						=    OpenLlamaModel(_snake_case		)
									original_model.to(_snake_case		)
									original_model.eval()
									_lowerCAmelCase						=    original_model(_snake_case		).last_hidden_state
									_lowerCAmelCase						=    original_model(_snake_case		).last_hidden_state

									set_seed(42		)  # Fixed seed at init time so the two models get the same random weights
									_lowerCAmelCase						=    {"""type""": scaling_type, """factor""": 10.0}
									_lowerCAmelCase						=    OpenLlamaModel(_snake_case		)
									scaled_model.to(_snake_case		)
									scaled_model.eval()
									_lowerCAmelCase						=    scaled_model(_snake_case		).last_hidden_state
									_lowerCAmelCase						=    scaled_model(_snake_case		).last_hidden_state

									# Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
									# maximum sequence length, so the outputs for the short input should match.
									if scaling_type == "dynamic":
														self.assertTrue(torch.allclose(_snake_case						,				_snake_case						,				atol=1e-5		)		)
									else:
														self.assertFalse(torch.allclose(_snake_case						,				_snake_case						,				atol=1e-5		)		)

									# The output should be different for long inputs
									self.assertFalse(torch.allclose(_snake_case						,				_snake_case						,				atol=1e-5		)		)











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      	=  logging.get_logger(__name__)


def __A    (      __lowerCAmelCase ,		__lowerCAmelCase=False	)->  Union[str, Any]:





						"""simple docstring"""



						_UpperCAmelCase									=			[]
						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"
												_UpperCAmelCase									=			[(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	)->  List[str]:





						"""simple docstring"""



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


def __A    (      __lowerCAmelCase	)->  Optional[Any]:





						"""simple docstring"""



						_UpperCAmelCase									=			['head.weight', 'head.bias']
						for k in ignore_keys:
												state_dict.pop(__lowerCAmelCase ,		__lowerCAmelCase	)


def __A    (      __lowerCAmelCase ,		__lowerCAmelCase ,		__lowerCAmelCase	)->  int:





						"""simple docstring"""



						_UpperCAmelCase									=			dct.pop(__lowerCAmelCase	)
						_UpperCAmelCase									=			val


def __A    (      )->  str:





						"""simple docstring"""



						_UpperCAmelCase									=			'http://images.cocodataset.org/val2017/000000039769.jpg'
						_UpperCAmelCase									=			Image.open(requests.get(__lowerCAmelCase ,		stream=__lowerCAmelCase	).raw	)
						return im


@torch.no_grad()
def __A    (      __lowerCAmelCase ,		__lowerCAmelCase ,		__lowerCAmelCase=True	)->  List[str]:





						"""simple docstring"""



						_UpperCAmelCase									=			ViTConfig()
						# patch_size
						if model_name[-1] == "8":
												_UpperCAmelCase									=			8
						# set labels if required
						if not base_model:
												_UpperCAmelCase									=			1_000
												_UpperCAmelCase									=			'huggingface/label-files'
												_UpperCAmelCase									=			'imagenet-1k-id2label.json'
												_UpperCAmelCase									=			json.load(open(hf_hub_download(__lowerCAmelCase ,		__lowerCAmelCase ,		repo_type='dataset'	) ,		'r'	)	)
												_UpperCAmelCase									=			{int(__lowerCAmelCase	): v for k, v in idalabel.items()}
												_UpperCAmelCase									=			idalabel
												_UpperCAmelCase									=			{v: k for k, v in idalabel.items()}
						# size of the architecture
						if model_name in ["dino_vits8", "dino_vits16"]:
												_UpperCAmelCase									=			384
												_UpperCAmelCase									=			1_536
												_UpperCAmelCase									=			12
												_UpperCAmelCase									=			6

						# load original model from torch hub
						_UpperCAmelCase									=			torch.hub.load('facebookresearch/dino:main' ,		__lowerCAmelCase	)
						original_model.eval()

						# load state_dict of original model, remove and rename some keys
						_UpperCAmelCase									=			original_model.state_dict()
						if base_model:
												remove_classification_head_(__lowerCAmelCase	)
						_UpperCAmelCase									=			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:
												_UpperCAmelCase									=			ViTModel(__lowerCAmelCase ,		add_pooling_layer=__lowerCAmelCase	).eval()
						else:
												_UpperCAmelCase									=			ViTForImageClassification(__lowerCAmelCase	).eval()
						model.load_state_dict(__lowerCAmelCase	)

						# Check outputs on an image, prepared by ViTImageProcessor
						_UpperCAmelCase									=			ViTImageProcessor()
						_UpperCAmelCase									=			image_processor(images=prepare_img() ,		return_tensors='pt'	)
						_UpperCAmelCase									=			encoding['pixel_values']
						_UpperCAmelCase									=			model(__lowerCAmelCase	)

						if base_model:
												_UpperCAmelCase									=			original_model(__lowerCAmelCase	)
												assert torch.allclose(__lowerCAmelCase ,		outputs.last_hidden_state[:, 0, :] ,		atol=1E-1	)
						else:
												_UpperCAmelCase									=			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      	=  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      	=  parser.parse_args()
		convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)
