


import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, BatchEncoding, PreTrainedTokenizer
from ...utils import logging


A				:       List[Any]      =  logging.get_logger(__name__)

A				:       List[Any]      =  '▁'

A				:       Any      =  {'vocab_file': 'sentencepiece.bpe.model'}

A				:       str      =  {
    'vocab_file': {
        'facebook/nllb-200-distilled-600M': (
            'https://huggingface.co/facebook/nllb-200-distilled-600M/blob/main/sentencepiece.bpe.model'
        ),
    }
}

A				:       Tuple      =  {
    'facebook/nllb-200-distilled-600M': 1024,
}

# fmt: off
A				:       Optional[int]      =  ['ace_Arab', 'ace_Latn', 'acm_Arab', 'acq_Arab', 'aeb_Arab', 'afr_Latn', 'ajp_Arab', 'aka_Latn', 'amh_Ethi', 'apc_Arab', 'arb_Arab', 'ars_Arab', 'ary_Arab', 'arz_Arab', 'asm_Beng', 'ast_Latn', 'awa_Deva', 'ayr_Latn', 'azb_Arab', 'azj_Latn', 'bak_Cyrl', 'bam_Latn', 'ban_Latn', 'bel_Cyrl', 'bem_Latn', 'ben_Beng', 'bho_Deva', 'bjn_Arab', 'bjn_Latn', 'bod_Tibt', 'bos_Latn', 'bug_Latn', 'bul_Cyrl', 'cat_Latn', 'ceb_Latn', 'ces_Latn', 'cjk_Latn', 'ckb_Arab', 'crh_Latn', 'cym_Latn', 'dan_Latn', 'deu_Latn', 'dik_Latn', 'dyu_Latn', 'dzo_Tibt', 'ell_Grek', 'eng_Latn', 'epo_Latn', 'est_Latn', 'eus_Latn', 'ewe_Latn', 'fao_Latn', 'pes_Arab', 'fij_Latn', 'fin_Latn', 'fon_Latn', 'fra_Latn', 'fur_Latn', 'fuv_Latn', 'gla_Latn', 'gle_Latn', 'glg_Latn', 'grn_Latn', 'guj_Gujr', 'hat_Latn', 'hau_Latn', 'heb_Hebr', 'hin_Deva', 'hne_Deva', 'hrv_Latn', 'hun_Latn', 'hye_Armn', 'ibo_Latn', 'ilo_Latn', 'ind_Latn', 'isl_Latn', 'ita_Latn', 'jav_Latn', 'jpn_Jpan', 'kab_Latn', 'kac_Latn', 'kam_Latn', 'kan_Knda', 'kas_Arab', 'kas_Deva', 'kat_Geor', 'knc_Arab', 'knc_Latn', 'kaz_Cyrl', 'kbp_Latn', 'kea_Latn', 'khm_Khmr', 'kik_Latn', 'kin_Latn', 'kir_Cyrl', 'kmb_Latn', 'kon_Latn', 'kor_Hang', 'kmr_Latn', 'lao_Laoo', 'lvs_Latn', 'lij_Latn', 'lim_Latn', 'lin_Latn', 'lit_Latn', 'lmo_Latn', 'ltg_Latn', 'ltz_Latn', 'lua_Latn', 'lug_Latn', 'luo_Latn', 'lus_Latn', 'mag_Deva', 'mai_Deva', 'mal_Mlym', 'mar_Deva', 'min_Latn', 'mkd_Cyrl', 'plt_Latn', 'mlt_Latn', 'mni_Beng', 'khk_Cyrl', 'mos_Latn', 'mri_Latn', 'zsm_Latn', 'mya_Mymr', 'nld_Latn', 'nno_Latn', 'nob_Latn', 'npi_Deva', 'nso_Latn', 'nus_Latn', 'nya_Latn', 'oci_Latn', 'gaz_Latn', 'ory_Orya', 'pag_Latn', 'pan_Guru', 'pap_Latn', 'pol_Latn', 'por_Latn', 'prs_Arab', 'pbt_Arab', 'quy_Latn', 'ron_Latn', 'run_Latn', 'rus_Cyrl', 'sag_Latn', 'san_Deva', 'sat_Beng', 'scn_Latn', 'shn_Mymr', 'sin_Sinh', 'slk_Latn', 'slv_Latn', 'smo_Latn', 'sna_Latn', 'snd_Arab', 'som_Latn', 'sot_Latn', 'spa_Latn', 'als_Latn', 'srd_Latn', 'srp_Cyrl', 'ssw_Latn', 'sun_Latn', 'swe_Latn', 'swh_Latn', 'szl_Latn', 'tam_Taml', 'tat_Cyrl', 'tel_Telu', 'tgk_Cyrl', 'tgl_Latn', 'tha_Thai', 'tir_Ethi', 'taq_Latn', 'taq_Tfng', 'tpi_Latn', 'tsn_Latn', 'tso_Latn', 'tuk_Latn', 'tum_Latn', 'tur_Latn', 'twi_Latn', 'tzm_Tfng', 'uig_Arab', 'ukr_Cyrl', 'umb_Latn', 'urd_Arab', 'uzn_Latn', 'vec_Latn', 'vie_Latn', 'war_Latn', 'wol_Latn', 'xho_Latn', 'ydd_Hebr', 'yor_Latn', 'yue_Hant', 'zho_Hans', 'zho_Hant', 'zul_Latn']






class      UpperCamelCase(	_a ):
					snake_case_		:					int    					=						VOCAB_FILES_NAMES
					snake_case_		:					List[Any]    					=						PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
					snake_case_		:					Optional[int]    					=						PRETRAINED_VOCAB_FILES_MAP
					snake_case_		:					List[Any]    					=						["""input_ids""", """attention_mask"""]

					snake_case_		:					List[int]    					=						[]
					snake_case_		:					List[int]    					=						[]


					def __init__(    self						:		Optional[int]  ,     SCREAMING_SNAKE_CASE						:		Optional[int]  ,     SCREAMING_SNAKE_CASE						:		Union[str, Any]="<s>"  ,     SCREAMING_SNAKE_CASE						:		Optional[int]="</s>"  ,     SCREAMING_SNAKE_CASE						:		str="</s>"  ,     SCREAMING_SNAKE_CASE						:		int="<s>"  ,     SCREAMING_SNAKE_CASE						:		str="<unk>"  ,     SCREAMING_SNAKE_CASE						:		str="<pad>"  ,     SCREAMING_SNAKE_CASE						:		str="<mask>"  ,     SCREAMING_SNAKE_CASE						:		List[str]=None  ,     SCREAMING_SNAKE_CASE						:		Optional[int]=None  ,     SCREAMING_SNAKE_CASE						:		List[Any]=None  ,     SCREAMING_SNAKE_CASE						:		Optional[Dict[str, Any]] = None  ,     SCREAMING_SNAKE_CASE						:		str=None  ,     SCREAMING_SNAKE_CASE						:		List[Any]=False  ,     **SCREAMING_SNAKE_CASE						:		Optional[int]  ,     )      ->						Optional[int]:




											'''simple docstring'''

											__snake_case				=		AddedToken(SCREAMING_SNAKE_CASE  ,     lstrip=SCREAMING_SNAKE_CASE  ,     rstrip=SCREAMING_SNAKE_CASE    ) if isinstance(SCREAMING_SNAKE_CASE  ,     SCREAMING_SNAKE_CASE    ) else mask_token

											__snake_case				=		{} if sp_model_kwargs is None else sp_model_kwargs
											__snake_case				=		legacy_behaviour

											super().__init__(
											    bos_token=SCREAMING_SNAKE_CASE  ,     eos_token=SCREAMING_SNAKE_CASE  ,     unk_token=SCREAMING_SNAKE_CASE  ,     sep_token=SCREAMING_SNAKE_CASE  ,     cls_token=SCREAMING_SNAKE_CASE  ,     pad_token=SCREAMING_SNAKE_CASE  ,     mask_token=SCREAMING_SNAKE_CASE  ,     tokenizer_file=SCREAMING_SNAKE_CASE  ,     src_lang=SCREAMING_SNAKE_CASE  ,     tgt_lang=SCREAMING_SNAKE_CASE  ,     additional_special_tokens=SCREAMING_SNAKE_CASE  ,     sp_model_kwargs=self.sp_model_kwargs  ,     legacy_behaviour=SCREAMING_SNAKE_CASE  ,     **SCREAMING_SNAKE_CASE  ,     )

											__snake_case				=		spm.SentencePieceProcessor(**self.sp_model_kwargs    )
											self.sp_model.Load(str(SCREAMING_SNAKE_CASE    )    )
											__snake_case				=		vocab_file
											# Original fairseq vocab and spm vocab must be "aligned":
											# Vocab    |    0    |    1    |   2    |    3    |  4   |  5   |  6   |   7  |   8  |  9
											# -------- | ------- | ------- | ------ | ------- | ---- | ---- | ---- | ---- | ---- | ----
											# fairseq  | '<s>'   | '<pad>' | '</s>' | '<unk>' | 'an' | '▁n' | '▁m' | '▁t' | '▁k' | '▁a'
											# spm      | '<unk>' | '<s>'   | '</s>' | 'an'    | '▁n' | '▁m' | '▁t' | '▁k' | '▁a' | '▁s'

											# Mimic fairseq token-to-id alignment for the first 4 token
											__snake_case				=		{"<s>": 0, "<pad>": 1, "</s>": 2, "<unk>": 3}

											# The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
											__snake_case				=		1

											__snake_case				=		len(self.sp_model    )
											__snake_case				=		{
											    code: self.sp_model_size + i + self.fairseq_offset for i, code in enumerate(SCREAMING_SNAKE_CASE    )
											}
											__snake_case				=		{v: k for k, v in self.lang_code_to_id.items()}
											__snake_case				=		len(self.sp_model    ) + len(self.lang_code_to_id    ) + self.fairseq_offset

											self.fairseq_tokens_to_ids.update(self.lang_code_to_id    )
											__snake_case				=		{v: k for k, v in self.fairseq_tokens_to_ids.items()}
											__snake_case				=		list(self.lang_code_to_id.keys()    )

											if additional_special_tokens is not None:
																	# Only add those special tokens if they are not already there.
																	self._additional_special_tokens.extend(
																	    [t for t in additional_special_tokens if t not in self._additional_special_tokens]    )

											__snake_case				=		src_lang if src_lang is not None else "eng_Latn"
											__snake_case				=		self.lang_code_to_id[self._src_lang]
											__snake_case				=		tgt_lang
											self.set_src_lang_special_tokens(self._src_lang    )


					def __getstate__(    self						:		Union[str, Any]    )      ->						Optional[int]:




											'''simple docstring'''

											__snake_case				=		self.__dict__.copy()
											__snake_case				=		None
											__snake_case				=		self.sp_model.serialized_model_proto()
											return state


					def __setstate__(    self						:		Optional[int]  ,     SCREAMING_SNAKE_CASE						:		Optional[int]    )      ->						Optional[Any]:




											'''simple docstring'''

											__snake_case				=		d

											# for backward compatibility
											if not hasattr(self  ,     "sp_model_kwargs"    ):
																	__snake_case				=		{}

											__snake_case				=		spm.SentencePieceProcessor(**self.sp_model_kwargs    )
											self.sp_model.LoadFromSerializedProto(self.sp_model_proto    )


					@property
					def 				SCREAMING_SNAKE_CASE_	(    self						:		List[Any]    )      ->						Any:




											'''simple docstring'''

											return len(self.sp_model    ) + len(self.lang_code_to_id    ) + self.fairseq_offset + 1  # Plus 1 for the mask token


					@property
					def 				SCREAMING_SNAKE_CASE_	(    self						:		Union[str, Any]    )      ->						str:




											'''simple docstring'''

											return self._src_lang


					@src_lang.setter
					def 				SCREAMING_SNAKE_CASE_	(    self						:		List[str]  ,     SCREAMING_SNAKE_CASE						:		str    )      ->						None:




											'''simple docstring'''

											__snake_case				=		new_src_lang
											self.set_src_lang_special_tokens(self._src_lang    )


					def 				SCREAMING_SNAKE_CASE_	(    self						:		Tuple  ,     SCREAMING_SNAKE_CASE						:		List[int]  ,     SCREAMING_SNAKE_CASE						:		Optional[List[int]] = None  ,     SCREAMING_SNAKE_CASE						:		bool = False    )      ->						List[int]:




											'''simple docstring'''

											if already_has_special_tokens:
																	return super().get_special_tokens_mask(
																	    token_ids_a=SCREAMING_SNAKE_CASE  ,     token_ids_a=SCREAMING_SNAKE_CASE  ,     already_has_special_tokens=SCREAMING_SNAKE_CASE    )

											__snake_case				=		[1] * len(self.prefix_tokens    )
											__snake_case				=		[1] * len(self.suffix_tokens    )
											if token_ids_a is None:
																	return prefix_ones + ([0] * len(SCREAMING_SNAKE_CASE    )) + suffix_ones
											return prefix_ones + ([0] * len(SCREAMING_SNAKE_CASE    )) + ([0] * len(SCREAMING_SNAKE_CASE    )) + suffix_ones


					def 				SCREAMING_SNAKE_CASE_	(    self						:		List[str]  ,     SCREAMING_SNAKE_CASE						:		List[int]  ,     SCREAMING_SNAKE_CASE						:		Optional[List[int]] = None    )      ->						List[int]:




											'''simple docstring'''

											if token_ids_a is None:
																	return self.prefix_tokens + token_ids_a + self.suffix_tokens
											# We don't expect to process pairs, but leave the pair logic for API consistency
											return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens


					def 				SCREAMING_SNAKE_CASE_	(    self						:		Tuple  ,     SCREAMING_SNAKE_CASE						:		List[int]  ,     SCREAMING_SNAKE_CASE						:		Optional[List[int]] = None    )      ->						List[int]:




											'''simple docstring'''

											__snake_case				=		[self.sep_token_id]
											__snake_case				=		[self.cls_token_id]

											if token_ids_a is None:
																	return len(cls + token_ids_a + sep    ) * [0]
											return len(cls + token_ids_a + sep + sep + token_ids_a + sep    ) * [0]


					def 				SCREAMING_SNAKE_CASE_	(    self						:		str  ,     SCREAMING_SNAKE_CASE						:		Optional[int]  ,     SCREAMING_SNAKE_CASE						:		str  ,     SCREAMING_SNAKE_CASE						:		Optional[str]  ,     SCREAMING_SNAKE_CASE						:		Optional[str]  ,     **SCREAMING_SNAKE_CASE						:		List[Any]    )      ->						List[Any]:




											'''simple docstring'''

											if src_lang is None or tgt_lang is None:
																	raise ValueError("Translation requires a `src_lang` and a `tgt_lang` for this model"    )
											__snake_case				=		src_lang
											__snake_case				=		self(SCREAMING_SNAKE_CASE  ,     add_special_tokens=SCREAMING_SNAKE_CASE  ,     return_tensors=SCREAMING_SNAKE_CASE  ,     **SCREAMING_SNAKE_CASE    )
											__snake_case				=		self.convert_tokens_to_ids(SCREAMING_SNAKE_CASE    )
											__snake_case				=		tgt_lang_id
											return inputs


					def 				SCREAMING_SNAKE_CASE_	(    self						:		str    )      ->						Union[str, Any]:




											'''simple docstring'''

											__snake_case				=		{self.convert_ids_to_tokens(SCREAMING_SNAKE_CASE    ): i for i in range(self.vocab_size    )}
											vocab.update(self.added_tokens_encoder    )
											return vocab


					def 				SCREAMING_SNAKE_CASE_	(    self						:		List[str]  ,     SCREAMING_SNAKE_CASE						:		str    )      ->						List[str]:




											'''simple docstring'''

											return self.sp_model.encode(SCREAMING_SNAKE_CASE  ,     out_type=SCREAMING_SNAKE_CASE    )


					def 				SCREAMING_SNAKE_CASE_	(    self						:		str  ,     SCREAMING_SNAKE_CASE						:		List[str]    )      ->						Union[str, Any]:




											'''simple docstring'''

											if token in self.fairseq_tokens_to_ids:
																	return self.fairseq_tokens_to_ids[token]
											__snake_case				=		self.sp_model.PieceToId(SCREAMING_SNAKE_CASE    )

											# Need to return unknown token if the SP model returned 0
											return spm_id + self.fairseq_offset if spm_id else self.unk_token_id


					def 				SCREAMING_SNAKE_CASE_	(    self						:		List[str]  ,     SCREAMING_SNAKE_CASE						:		Any    )      ->						List[Any]:




											'''simple docstring'''

											if index in self.fairseq_ids_to_tokens:
																	return self.fairseq_ids_to_tokens[index]
											return self.sp_model.IdToPiece(index - self.fairseq_offset    )


					def 				SCREAMING_SNAKE_CASE_	(    self						:		List[str]  ,     SCREAMING_SNAKE_CASE						:		List[Any]    )      ->						Optional[int]:




											'''simple docstring'''

											__snake_case				=		"".join(SCREAMING_SNAKE_CASE    ).replace(SCREAMING_SNAKE_CASE  ,     " "    ).strip()
											return out_string


					def 				SCREAMING_SNAKE_CASE_	(    self						:		int  ,     SCREAMING_SNAKE_CASE						:		str  ,     SCREAMING_SNAKE_CASE						:		Optional[str] = None    )      ->						Tuple[str]:




											'''simple docstring'''

											if not os.path.isdir(SCREAMING_SNAKE_CASE    ):
																	logger.error(f'''Vocabulary path ({save_directory}) should be a directory'''    )
																	return
											__snake_case				=		os.path.join(
											    SCREAMING_SNAKE_CASE  ,     (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]    )

											if os.path.abspath(self.vocab_file    ) != os.path.abspath(SCREAMING_SNAKE_CASE    ) and os.path.isfile(self.vocab_file    ):
																	copyfile(self.vocab_file  ,     SCREAMING_SNAKE_CASE    )
											elif not os.path.isfile(self.vocab_file    ):
																	with open(SCREAMING_SNAKE_CASE  ,     "wb"    ) as fi:
																							__snake_case				=		self.sp_model.serialized_model_proto()
																							fi.write(SCREAMING_SNAKE_CASE    )

											return (out_vocab_file,)


					def 				SCREAMING_SNAKE_CASE_	(    self						:		Any  ,     SCREAMING_SNAKE_CASE						:		List[str]  ,     SCREAMING_SNAKE_CASE						:		str = "eng_Latn"  ,     SCREAMING_SNAKE_CASE						:		Optional[List[str]] = None  ,     SCREAMING_SNAKE_CASE						:		str = "fra_Latn"  ,     **SCREAMING_SNAKE_CASE						:		Optional[Any]  ,     )      ->						BatchEncoding:




											'''simple docstring'''

											__snake_case				=		src_lang
											__snake_case				=		tgt_lang
											return super().prepare_seqaseq_batch(SCREAMING_SNAKE_CASE  ,     SCREAMING_SNAKE_CASE  ,     **SCREAMING_SNAKE_CASE    )


					def 				SCREAMING_SNAKE_CASE_	(    self						:		Tuple    )      ->						Optional[Any]:




											'''simple docstring'''

											return self.set_src_lang_special_tokens(self.src_lang    )


					def 				SCREAMING_SNAKE_CASE_	(    self						:		int    )      ->						str:




											'''simple docstring'''

											return self.set_tgt_lang_special_tokens(self.tgt_lang    )


					def 				SCREAMING_SNAKE_CASE_	(    self						:		List[str]  ,     SCREAMING_SNAKE_CASE						:		Union[str, Any]    )      ->						None:




											'''simple docstring'''

											__snake_case				=		self.lang_code_to_id[src_lang]
											if self.legacy_behaviour:
																	__snake_case				=		[]
																	__snake_case				=		[self.eos_token_id, self.cur_lang_code]
											else:
																	__snake_case				=		[self.cur_lang_code]
																	__snake_case				=		[self.eos_token_id]






					def 				SCREAMING_SNAKE_CASE_	(    self						:		Dict  ,     SCREAMING_SNAKE_CASE						:		str    )      ->						None:




											'''simple docstring'''

											__snake_case				=		self.lang_code_to_id[lang]
											if self.legacy_behaviour:
																	__snake_case				=		[]
																	__snake_case				=		[self.eos_token_id, self.cur_lang_code]
											else:
																	__snake_case				=		[self.cur_lang_code]
																	__snake_case				=		[self.eos_token_id]









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





def 						_lowerCAmelCase							(     _lowerCAmelCase   ,						_lowerCAmelCase=False			)    ->						Optional[int]:


						'''simple docstring'''






						__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'''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				=		[(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			)    ->						int:


						'''simple docstring'''






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





def 						_lowerCAmelCase							(     _lowerCAmelCase			)    ->						List[str]:


						'''simple docstring'''






						__snake_case				=		["head.weight", "head.bias"]
						for k in ignore_keys:
												state_dict.pop(_lowerCAmelCase   ,						_lowerCAmelCase			)





def 						_lowerCAmelCase							(     _lowerCAmelCase   ,						_lowerCAmelCase   ,						_lowerCAmelCase			)    ->						Dict:


						'''simple docstring'''






						__snake_case				=		dct.pop(_lowerCAmelCase			)
						__snake_case				=		val





def 						_lowerCAmelCase							(     )    ->						Any:


						'''simple docstring'''






						__snake_case				=		"http://images.cocodataset.org/val2017/000000039769.jpg"
						__snake_case				=		Image.open(requests.get(_lowerCAmelCase   ,						stream=_lowerCAmelCase			).raw			)
						return im





@torch.no_grad()
def 						_lowerCAmelCase							(     _lowerCAmelCase   ,						_lowerCAmelCase   ,						_lowerCAmelCase=True			)    ->						int:


						'''simple docstring'''






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

						# load original model from torch hub
						__snake_case				=		torch.hub.load("facebookresearch/dino:main"   ,						_lowerCAmelCase			)
						original_model.eval()

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

						# Check outputs on an image, prepared by ViTImageProcessor
						__snake_case				=		ViTImageProcessor()
						__snake_case				=		image_processor(images=prepare_img()   ,						return_tensors="pt"			)
						__snake_case				=		encoding["pixel_values"]
						__snake_case				=		model(_lowerCAmelCase			)

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





