





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


_snake_case  :   Optional[int]						   =  logging.get_logger(__name__)

_snake_case  :   Any						   =  '▁'

_snake_case  :   int						   =  {'vocab_file': 'sentencepiece.bpe.model'}

_snake_case  :   Optional[Any]						   =  {
    'vocab_file': {
        'facebook/mbart-large-50-one-to-many-mmt': (
            'https://huggingface.co/facebook/mbart-large-50-one-to-many-mmt/resolve/main/sentencepiece.bpe.model'
        ),
    }
}

_snake_case  :   List[str]						   =  {
    'facebook/mbart-large-50-one-to-many-mmt': 1024,
}

# fmt: off
_snake_case  :   Optional[Any]						   =  ['ar_AR', 'cs_CZ', 'de_DE', 'en_XX', 'es_XX', 'et_EE', 'fi_FI', 'fr_XX', 'gu_IN', 'hi_IN', 'it_IT', 'ja_XX', 'kk_KZ', 'ko_KR', 'lt_LT', 'lv_LV', 'my_MM', 'ne_NP', 'nl_XX', 'ro_RO', 'ru_RU', 'si_LK', 'tr_TR', 'vi_VN', 'zh_CN', 'af_ZA', 'az_AZ', 'bn_IN', 'fa_IR', 'he_IL', 'hr_HR', 'id_ID', 'ka_GE', 'km_KH', 'mk_MK', 'ml_IN', 'mn_MN', 'mr_IN', 'pl_PL', 'ps_AF', 'pt_XX', 'sv_SE', 'sw_KE', 'ta_IN', 'te_IN', 'th_TH', 'tl_XX', 'uk_UA', 'ur_PK', 'xh_ZA', 'gl_ES', 'sl_SI']



class     _UpperCAmelCase   (   _UpperCamelCase   ):

  """simple docstring"""


  a_      =			VOCAB_FILES_NAMES
  a_      =			PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
  a_      =			PRETRAINED_VOCAB_FILES_MAP
  a_      =			["""input_ids""", """attention_mask"""]

  a_      =			[]
  a_      =			[]






  def __init__(    self			:		Dict ,  lowerCAmelCase_			:		Optional[int] ,  lowerCAmelCase_			:		List[str]=None ,  lowerCAmelCase_			:		int=None ,  lowerCAmelCase_			:		Union[str, Any]="</s>" ,  lowerCAmelCase_			:		Dict="</s>" ,  lowerCAmelCase_			:		str="<s>" ,  lowerCAmelCase_			:		Optional[int]="<unk>" ,  lowerCAmelCase_			:		Optional[int]="<pad>" ,  lowerCAmelCase_			:		List[str]="<mask>" ,  lowerCAmelCase_			:		Optional[Dict[str, Any]] = None ,  **lowerCAmelCase_			:		List[Any] ,  )							->						None:
     # Mask token behave like a normal word, i.e. include the space before it
     __lowerCAmelCase							     =     AddedToken(lowerCAmelCase_ ,  lstrip=lowerCAmelCase_ ,  rstrip=lowerCAmelCase_					) if isinstance(lowerCAmelCase_ ,  lowerCAmelCase_					) else mask_token

     __lowerCAmelCase							     =     {} if sp_model_kwargs is None else sp_model_kwargs

     __lowerCAmelCase							     =     kwargs.get('additional_special_tokens' ,  []					)
     kwargs["additional_special_tokens"] += [
         code for code in FAIRSEQ_LANGUAGE_CODES if code not in kwargs["additional_special_tokens"]
     ]

     super().__init__(
         src_lang=lowerCAmelCase_ ,  tgt_lang=lowerCAmelCase_ ,  eos_token=lowerCAmelCase_ ,  unk_token=lowerCAmelCase_ ,  sep_token=lowerCAmelCase_ ,  cls_token=lowerCAmelCase_ ,  pad_token=lowerCAmelCase_ ,  mask_token=lowerCAmelCase_ ,  sp_model_kwargs=self.sp_model_kwargs ,  **lowerCAmelCase_ ,  )

     __lowerCAmelCase							     =     spm.SentencePieceProcessor(**self.sp_model_kwargs					)
     self.sp_model.Load(str(lowerCAmelCase_					)					)
     __lowerCAmelCase							     =     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>' | ',' | '.' | '▁' | 's'   | '▁de' | '-'
     # spm      | '<unk>' | '<s>'   | '</s>' | ','     | '.' | '▁' | 's' | '▁de' | '-'   | '▁a'

     # Mimic fairseq token-to-id alignment for the first 4 token
     __lowerCAmelCase							     =     {'<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
     __lowerCAmelCase							     =     1

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

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

     __lowerCAmelCase							     =     src_lang if src_lang is not None else 'en_XX'
     __lowerCAmelCase							     =     self.lang_code_to_id[self._src_lang]
     __lowerCAmelCase							     =     tgt_lang
     self.set_src_lang_special_tokens(self._src_lang					)






  @property
  def     lowercase		(    self			:		str					)							->						int:
     return len(self.sp_model					) + len(self.lang_code_to_id					) + self.fairseq_offset + 1  # Plus 1 for the mask token






  @property
  def     lowercase		(    self			:		Optional[Any]					)							->						str:
     return self._src_lang






  @src_lang.setter
  def     lowercase		(    self			:		Optional[int] ,  lowerCAmelCase_			:		str					)							->						None:
     __lowerCAmelCase							     =     new_src_lang
     self.set_src_lang_special_tokens(self._src_lang					)






  def __getstate__(    self			:		Union[str, Any]					)							->						Dict:
     __lowerCAmelCase							     =     self.__dict__.copy()
     __lowerCAmelCase							     =     None
     return state






  def __setstate__(    self			:		Any ,  lowerCAmelCase_			:		Dict					)							->						None:
     __lowerCAmelCase							     =     d

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

     __lowerCAmelCase							     =     spm.SentencePieceProcessor(**self.sp_model_kwargs					)
     self.sp_model.Load(self.vocab_file					)






  def     lowercase		(    self			:		List[Any]					)							->						Dict:
     __lowerCAmelCase							     =     {self.convert_ids_to_tokens(lowerCAmelCase_					): i for i in range(self.vocab_size					)}
     vocab.update(self.added_tokens_encoder					)
     return vocab






  def     lowercase		(    self			:		Tuple ,  lowerCAmelCase_			:		str					)							->						List[str]:
     return self.sp_model.encode(lowerCAmelCase_ ,  out_type=lowerCAmelCase_					)






  def     lowercase		(    self			:		str ,  lowerCAmelCase_			:		str					)							->						int:
     if token in self.fairseq_tokens_to_ids:
        return self.fairseq_tokens_to_ids[token]
     __lowerCAmelCase							     =     self.sp_model.PieceToId(lowerCAmelCase_					)

     # 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     lowercase		(    self			:		Dict ,  lowerCAmelCase_			:		int					)							->						str:
     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     lowercase		(    self			:		Any ,  lowerCAmelCase_			:		Union[str, Any]					)							->						int:
     __lowerCAmelCase							     =     []
     __lowerCAmelCase							     =     ''
     __lowerCAmelCase							     =     False
     for token in tokens:
        # make sure that special tokens are not decoded using sentencepiece model
        if token in self.all_special_tokens:
           if not prev_is_special:
              out_string += " "
           out_string += self.sp_model.decode(lowerCAmelCase_					) + token
           __lowerCAmelCase							     =     True
           __lowerCAmelCase							     =     []
        else:
           current_sub_tokens.append(lowerCAmelCase_					)
           __lowerCAmelCase							     =     False
     out_string += self.sp_model.decode(lowerCAmelCase_					)
     return out_string.strip()






  def     lowercase		(    self			:		str ,  lowerCAmelCase_			:		str ,  lowerCAmelCase_			:		Optional[str] = None					)							->						Tuple[str]:
     if not os.path.isdir(lowerCAmelCase_					):
        logger.error(f"""Vocabulary path ({save_directory}) should be a directory"""					)
        return
     __lowerCAmelCase							     =     os.path.join(
         lowerCAmelCase_ ,  (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file']					)

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

     return (out_vocab_file,)






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

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

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






  def     lowercase		(    self			:		Optional[Any] ,  lowerCAmelCase_			:		List[int] ,  lowerCAmelCase_			:		Optional[List[int]] = None					)							->						List[int]:
     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     lowercase		(    self			:		str ,  lowerCAmelCase_			:		Optional[int] ,  lowerCAmelCase_			:		str ,  lowerCAmelCase_			:		Optional[str] ,  lowerCAmelCase_			:		Optional[str] ,  **lowerCAmelCase_			:		Optional[int]					)							->						Tuple:
     if src_lang is None or tgt_lang is None:
        raise ValueError('Translation requires a `src_lang` and a `tgt_lang` for this model'					)
     __lowerCAmelCase							     =     src_lang
     __lowerCAmelCase							     =     self(lowerCAmelCase_ ,  add_special_tokens=lowerCAmelCase_ ,  return_tensors=lowerCAmelCase_ ,  **lowerCAmelCase_					)
     __lowerCAmelCase							     =     self.convert_tokens_to_ids(lowerCAmelCase_					)
     __lowerCAmelCase							     =     tgt_lang_id
     return inputs






  def     lowercase		(    self			:		Dict ,  lowerCAmelCase_			:		List[str] ,  lowerCAmelCase_			:		str = "en_XX" ,  lowerCAmelCase_			:		Optional[List[str]] = None ,  lowerCAmelCase_			:		str = "ro_RO" ,  **lowerCAmelCase_			:		Optional[Any] ,  )							->						BatchEncoding:
     __lowerCAmelCase							     =     src_lang
     __lowerCAmelCase							     =     tgt_lang
     return super().prepare_seqaseq_batch(lowerCAmelCase_ ,  lowerCAmelCase_ ,  **lowerCAmelCase_					)






  def     lowercase		(    self			:		Optional[int]					)							->						List[str]:
     return self.set_src_lang_special_tokens(self.src_lang					)






  def     lowercase		(    self			:		str					)							->						List[str]:
     return self.set_tgt_lang_special_tokens(self.tgt_lang					)






  def     lowercase		(    self			:		List[str] ,  lowerCAmelCase_			:		str					)							->						None:
     __lowerCAmelCase							     =     self.lang_code_to_id[src_lang]
     __lowerCAmelCase							     =     [self.cur_lang_code_id]
     __lowerCAmelCase							     =     [self.eos_token_id]






  def     lowercase		(    self			:		int ,  lowerCAmelCase_			:		str					)							->						None:
     __lowerCAmelCase							     =     self.lang_code_to_id[tgt_lang]
     __lowerCAmelCase							     =     [self.cur_lang_code_id]
     __lowerCAmelCase							     =     [self.eos_token_id]








import argparse
import json
from pathlib import Path

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

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


logging.set_verbosity_info()
_snake_case  :   List[str]						   =  logging.get_logger(__name__)

def   a_   (			lowerCAmelCase_				:      Optional[Any],     lowerCAmelCase_				:      Union[str, Any]=False     ):
   __lowerCAmelCase							     =     []
   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'),
              ('pre_logits.fc.weight', 'pooler.dense.weight'),
              ('pre_logits.fc.bias', 'pooler.dense.bias'),
          ]     )

      # if just the base model, we should remove "vit" from all keys that start with "vit"
      __lowerCAmelCase							     =     [(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_				:      List[Any],     lowerCAmelCase_				:      Dict,     lowerCAmelCase_				:      int=False     ):
   for i in range(config.num_hidden_layers     ):
      if base_model:
         __lowerCAmelCase							     =     ''
      else:
         __lowerCAmelCase							     =     'vit.'
      # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
      __lowerCAmelCase							     =     state_dict.pop(F"""blocks.{i}.attn.qkv.weight"""     )
      __lowerCAmelCase							     =     state_dict.pop(F"""blocks.{i}.attn.qkv.bias"""     )
      # next, add query, keys and values (in that order) to the state dict
      __lowerCAmelCase							     =     in_proj_weight[
          : config.hidden_size, :
      ]
      __lowerCAmelCase							     =     in_proj_bias[: config.hidden_size]
      __lowerCAmelCase							     =     in_proj_weight[
          config.hidden_size : config.hidden_size * 2, :
      ]
      __lowerCAmelCase							     =     in_proj_bias[
          config.hidden_size : config.hidden_size * 2
      ]
      __lowerCAmelCase							     =     in_proj_weight[
          -config.hidden_size :, :
      ]
      __lowerCAmelCase							     =     in_proj_bias[-config.hidden_size :]

def   a_   (			lowerCAmelCase_				:      Optional[int]     ):
   __lowerCAmelCase							     =     ['head.weight', 'head.bias']
   for k in ignore_keys:
      state_dict.pop(lowerCAmelCase_,     lowerCAmelCase_     )

def   a_   (			lowerCAmelCase_				:      Optional[Any],     lowerCAmelCase_				:      Tuple,     lowerCAmelCase_				:      Optional[int]     ):
   __lowerCAmelCase							     =     dct.pop(lowerCAmelCase_     )
   __lowerCAmelCase							     =     val

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

@torch.no_grad()
def   a_   (			lowerCAmelCase_				:      List[str],     lowerCAmelCase_				:      Union[str, Any]     ):
   __lowerCAmelCase							     =     ViTConfig()
   __lowerCAmelCase							     =     False
   # dataset (ImageNet-21k only or also fine-tuned on ImageNet 2012), patch_size and image_size
   if vit_name[-5:] == "in21k":
      __lowerCAmelCase							     =     True
      __lowerCAmelCase							     =     int(vit_name[-12:-10]     )
      __lowerCAmelCase							     =     int(vit_name[-9:-6]     )
   else:
      __lowerCAmelCase							     =     1000
      __lowerCAmelCase							     =     'huggingface/label-files'
      __lowerCAmelCase							     =     'imagenet-1k-id2label.json'
      __lowerCAmelCase							     =     json.load(open(hf_hub_download(lowerCAmelCase_,     lowerCAmelCase_,     repo_type='dataset'     ),     'r'     )     )
      __lowerCAmelCase							     =     {int(lowerCAmelCase_     ): v for k, v in idalabel.items()}
      __lowerCAmelCase							     =     idalabel
      __lowerCAmelCase							     =     {v: k for k, v in idalabel.items()}
      __lowerCAmelCase							     =     int(vit_name[-6:-4]     )
      __lowerCAmelCase							     =     int(vit_name[-3:]     )
   # size of the architecture
   if "deit" in vit_name:
      if vit_name[9:].startswith('tiny'     ):
         __lowerCAmelCase							     =     192
         __lowerCAmelCase							     =     768
         __lowerCAmelCase							     =     12
         __lowerCAmelCase							     =     3
      elif vit_name[9:].startswith('small'     ):
         __lowerCAmelCase							     =     384
         __lowerCAmelCase							     =     1536
         __lowerCAmelCase							     =     12
         __lowerCAmelCase							     =     6
      else:
         pass
   else:
      if vit_name[4:].startswith('small'     ):
         __lowerCAmelCase							     =     768
         __lowerCAmelCase							     =     2304
         __lowerCAmelCase							     =     8
         __lowerCAmelCase							     =     8
      elif vit_name[4:].startswith('base'     ):
         pass
      elif vit_name[4:].startswith('large'     ):
         __lowerCAmelCase							     =     1024
         __lowerCAmelCase							     =     4096
         __lowerCAmelCase							     =     24
         __lowerCAmelCase							     =     16
      elif vit_name[4:].startswith('huge'     ):
         __lowerCAmelCase							     =     1280
         __lowerCAmelCase							     =     5120
         __lowerCAmelCase							     =     32
         __lowerCAmelCase							     =     16

    # load original model from timm
   __lowerCAmelCase							     =     timm.create_model(lowerCAmelCase_,     pretrained=lowerCAmelCase_     )
   timm_model.eval()

   # load state_dict of original model, remove and rename some keys
   __lowerCAmelCase							     =     timm_model.state_dict()
   if base_model:
      remove_classification_head_(lowerCAmelCase_     )
   __lowerCAmelCase							     =     create_rename_keys(lowerCAmelCase_,     lowerCAmelCase_     )
   for src, dest in rename_keys:
      rename_key(lowerCAmelCase_,     lowerCAmelCase_,     lowerCAmelCase_     )
   read_in_q_k_v(lowerCAmelCase_,     lowerCAmelCase_,     lowerCAmelCase_     )

   # load HuggingFace model
   if vit_name[-5:] == "in21k":
      __lowerCAmelCase							     =     ViTModel(lowerCAmelCase_     ).eval()
   else:
      __lowerCAmelCase							     =     ViTForImageClassification(lowerCAmelCase_     ).eval()
   model.load_state_dict(lowerCAmelCase_     )

   # Check outputs on an image, prepared by ViTImageProcessor/DeiTImageProcessor
   if "deit" in vit_name:
      __lowerCAmelCase							     =     DeiTImageProcessor(size=config.image_size     )
   else:
      __lowerCAmelCase							     =     ViTImageProcessor(size=config.image_size     )
   __lowerCAmelCase							     =     image_processor(images=prepare_img(),     return_tensors='pt'     )
   __lowerCAmelCase							     =     encoding['pixel_values']
   __lowerCAmelCase							     =     model(lowerCAmelCase_     )

   if base_model:
      __lowerCAmelCase							     =     timm_model.forward_features(lowerCAmelCase_     )
      assert timm_pooled_output.shape == outputs.pooler_output.shape
      assert torch.allclose(lowerCAmelCase_,     outputs.pooler_output,     atol=1E-3     )
   else:
      __lowerCAmelCase							     =     timm_model(lowerCAmelCase_     )
      assert timm_logits.shape == outputs.logits.shape
      assert torch.allclose(lowerCAmelCase_,     outputs.logits,     atol=1E-3     )

   Path(lowerCAmelCase_     ).mkdir(exist_ok=lowerCAmelCase_     )
   print(F"""Saving model {vit_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__":
     _snake_case  :   List[Any]						   =  argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         '--vit_name',
         default='vit_base_patch16_224',
         type=str,
         help='Name of the ViT timm model you\'d like to convert.',
     )
     parser.add_argument(
         '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
     )

     _snake_case  :   Optional[int]						   =  parser.parse_args()
     convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path)

