



"""simple docstring"""




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


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

_lowerCAmelCase							:   Any				=				"""▁"""

_lowerCAmelCase							:   Optional[Any]				=				{"""vocab_file""": """sentencepiece.bpe.model"""}

_lowerCAmelCase							:   List[str]				=				{
    """vocab_file""": {
        """facebook/nllb-200-distilled-600M""": (
            """https://huggingface.co/facebook/nllb-200-distilled-600M/blob/main/sentencepiece.bpe.model"""
        ),
    }
}

_lowerCAmelCase							:   Dict				=				{
    """facebook/nllb-200-distilled-600M""": 10_24,
}

# fmt: off
_lowerCAmelCase							:   Union[str, Any]				=				["""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_    (							_UpperCamelCase		):
      __SCREAMING_SNAKE_CASE		:  Union[str, Any]						=					VOCAB_FILES_NAMES
      __SCREAMING_SNAKE_CASE		:  Tuple						=					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      __SCREAMING_SNAKE_CASE		:  Any						=					PRETRAINED_VOCAB_FILES_MAP
      __SCREAMING_SNAKE_CASE		:  List[str]						=					["""input_ids""", """attention_mask"""]

      __SCREAMING_SNAKE_CASE		:  List[int]						=					[]
      __SCREAMING_SNAKE_CASE		:  List[int]						=					[]



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

           _UpperCAmelCase     :      List[str]														=       {} if sp_model_kwargs is None else sp_model_kwargs
           _UpperCAmelCase     :      List[Any]														=       legacy_behaviour

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

           _UpperCAmelCase     :      Dict														=       spm.SentencePieceProcessor(**self.sp_model_kwargs		)
           self.sp_model.Load(str(_lowercase		)		)
           _UpperCAmelCase     :      List[Any]														=       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
           _UpperCAmelCase     :      Optional[Any]														=       {"<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
           _UpperCAmelCase     :      List[str]														=       1

           _UpperCAmelCase     :      Union[str, Any]														=       len(self.sp_model		)
           _UpperCAmelCase     :      Dict														=       {
               code: self.sp_model_size + i + self.fairseq_offset for i, code in enumerate(_lowercase		)
           }
           _UpperCAmelCase     :      Dict														=       {v: k for k, v in self.lang_code_to_id.items()}
           _UpperCAmelCase     :      Any														=       len(self.sp_model		) + len(self.lang_code_to_id		) + self.fairseq_offset

           self.fairseq_tokens_to_ids.update(self.lang_code_to_id		)
           _UpperCAmelCase     :      Any														=       {v: k for k, v in self.fairseq_tokens_to_ids.items()}
           _UpperCAmelCase     :      int														=       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]		)

           _UpperCAmelCase     :      List[str]														=       src_lang if src_lang is not None else "eng_Latn"
           _UpperCAmelCase     :      Union[str, Any]														=       self.lang_code_to_id[self._src_lang]
           _UpperCAmelCase     :      Tuple														=       tgt_lang
           self.set_src_lang_special_tokens(self._src_lang		)



      def __getstate__(	self	:      Optional[int]		):
           _UpperCAmelCase     :      Optional[int]														=       self.__dict__.copy()
           _UpperCAmelCase     :      Optional[int]														=       None
           _UpperCAmelCase     :      int														=       self.sp_model.serialized_model_proto()
           return state



      def __setstate__(	self	:      Optional[int]	,      A	:      Optional[Any]		):
           _UpperCAmelCase     :      List[str]														=       d

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

           _UpperCAmelCase     :      List[str]														=       spm.SentencePieceProcessor(**self.sp_model_kwargs		)
           self.sp_model.LoadFromSerializedProto(self.sp_model_proto		)



      @property
      def 		snake_case_   (	self	:      Tuple		):
           return len(self.sp_model		) + len(self.lang_code_to_id		) + self.fairseq_offset + 1  # Plus 1 for the mask token



      @property
      def 		snake_case_   (	self	:      Union[str, Any]		):
           return self._src_lang



      @src_lang.setter
      def 		snake_case_   (	self	:      Optional[int]	,      A	:      Optional[Any]		):
           _UpperCAmelCase     :      int														=       new_src_lang
           self.set_src_lang_special_tokens(self._src_lang		)



      def 		snake_case_   (	self	:      Tuple	,      A	:      Union[str, Any]	,      A	:      Any = None	,      A	:      Union[str, Any] = False		):

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

           _UpperCAmelCase     :      List[str]														=       [1] * len(self.prefix_tokens		)
           _UpperCAmelCase     :      List[Any]														=       [1] * len(self.suffix_tokens		)
           if token_ids_a is None:
                return prefix_ones + ([0] * len(_lowercase		)) + suffix_ones
           return prefix_ones + ([0] * len(_lowercase		)) + ([0] * len(_lowercase		)) + suffix_ones



      def 		snake_case_   (	self	:      List[Any]	,      A	:      Any	,      A	:      Tuple = None		):
           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 		snake_case_   (	self	:      int	,      A	:      str	,      A	:      Any = None		):
           _UpperCAmelCase     :      Dict														=       [self.sep_token_id]
           _UpperCAmelCase     :      Tuple														=       [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 		snake_case_   (	self	:      int	,      A	:      Any	,      A	:      List[str]	,      A	:      Dict	,      A	:      Union[str, Any]	,      **A	:      Any		):
           if src_lang is None or tgt_lang is None:
                raise ValueError("Translation requires a `src_lang` and a `tgt_lang` for this model"		)
           _UpperCAmelCase     :      Union[str, Any]														=       src_lang
           _UpperCAmelCase     :      Tuple														=       self(_lowercase	,      add_special_tokens=_lowercase	,      return_tensors=_lowercase	,      **_lowercase		)
           _UpperCAmelCase     :      Tuple														=       self.convert_tokens_to_ids(_lowercase		)
           _UpperCAmelCase     :      Optional[int]														=       tgt_lang_id
           return inputs



      def 		snake_case_   (	self	:      Optional[Any]		):
           _UpperCAmelCase     :      Optional[Any]														=       {self.convert_ids_to_tokens(_lowercase		): i for i in range(self.vocab_size		)}
           vocab.update(self.added_tokens_encoder		)
           return vocab



      def 		snake_case_   (	self	:      Tuple	,      A	:      Tuple		):
           return self.sp_model.encode(_lowercase	,      out_type=_lowercase		)



      def 		snake_case_   (	self	:      Dict	,      A	:      Tuple		):
           if token in self.fairseq_tokens_to_ids:
                return self.fairseq_tokens_to_ids[token]
           _UpperCAmelCase     :      Optional[int]														=       self.sp_model.PieceToId(_lowercase		)

           # 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 		snake_case_   (	self	:      Optional[Any]	,      A	:      List[Any]		):
           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 		snake_case_   (	self	:      Optional[int]	,      A	:      Tuple		):
           _UpperCAmelCase     :      Dict														=       "".join(_lowercase		).replace(_lowercase	,      " "		).strip()
           return out_string



      def 		snake_case_   (	self	:      Optional[Any]	,      A	:      Any	,      A	:      Tuple = None		):
           if not os.path.isdir(_lowercase		):
                logger.error(f'Vocabulary path ({save_directory}) should be a directory'		)
                return
           _UpperCAmelCase     :      Any														=       os.path.join(
               _lowercase	,      (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]		)

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

           return (out_vocab_file,)



      def 		snake_case_   (	self	:      int	,      A	:      str	,      A	:      int = "eng_Latn"	,      A	:      Tuple = None	,      A	:      List[Any] = "fra_Latn"	,      **A	:      int	,      ):
           _UpperCAmelCase     :      Optional[Any]														=       src_lang
           _UpperCAmelCase     :      Optional[Any]														=       tgt_lang
           return super().prepare_seqaseq_batch(_lowercase	,      _lowercase	,      **_lowercase		)



      def 		snake_case_   (	self	:      Optional[Any]		):
           return self.set_src_lang_special_tokens(self.src_lang		)



      def 		snake_case_   (	self	:      Optional[int]		):
           return self.set_tgt_lang_special_tokens(self.tgt_lang		)



      def 		snake_case_   (	self	:      Dict	,      A	:      Dict		):
           _UpperCAmelCase     :      Optional[int]														=       self.lang_code_to_id[src_lang]
           if self.legacy_behaviour:
                _UpperCAmelCase     :      Tuple														=       []
                _UpperCAmelCase     :      Tuple														=       [self.eos_token_id, self.cur_lang_code]
           else:
                _UpperCAmelCase     :      Optional[int]														=       [self.cur_lang_code]
                _UpperCAmelCase     :      Dict														=       [self.eos_token_id]



      def 		snake_case_   (	self	:      Union[str, Any]	,      A	:      List[Any]		):
           _UpperCAmelCase     :      Any														=       self.lang_code_to_id[lang]
           if self.legacy_behaviour:
                _UpperCAmelCase     :      str														=       []
                _UpperCAmelCase     :      Tuple														=       [self.eos_token_id, self.cur_lang_code]
           else:
                _UpperCAmelCase     :      int														=       [self.cur_lang_code]
                _UpperCAmelCase     :      int														=       [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 timm import create_model
from timm.data import resolve_data_config
from timm.data.transforms_factory import create_transform

from transformers import BitConfig, BitForImageClassification, BitImageProcessor
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


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




def 			lowerCAmelCase(      SCREAMING_SNAKE_CASE_			)->						int:



    """simple docstring"""






    UpperCamelCase_			  =						"huggingface/label-files"
    UpperCamelCase_			  =						"imagenet-1k-id2label.json"
    UpperCamelCase_			  =						json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE_			,						SCREAMING_SNAKE_CASE_			,						repo_type="dataset"			)			,						"r"			)			)
    UpperCamelCase_			  =						{int(SCREAMING_SNAKE_CASE_			): v for k, v in idalabel.items()}
    UpperCamelCase_			  =						{v: k for k, v in idalabel.items()}

    UpperCamelCase_			  =						"std_conv" if "bit" in model_name else False

    # note that when using BiT as backbone for ViT-hybrid checkpoints,
    # one needs to additionally set config.layer_type = "bottleneck", config.stem_type = "same",
    # config.conv_layer = "std_conv_same"
    UpperCamelCase_			  =						BitConfig(
        conv_layer=SCREAMING_SNAKE_CASE_			,						num_labels=1_0_0_0			,						idalabel=SCREAMING_SNAKE_CASE_			,						labelaid=SCREAMING_SNAKE_CASE_			,						)

    return config




def 			lowerCAmelCase(      SCREAMING_SNAKE_CASE_			)->						Optional[Any]:



    """simple docstring"""






    if "stem.conv" in name:
        UpperCamelCase_			  =						name.replace("stem.conv"			,						"bit.embedder.convolution"			)
    if "blocks" in name:
        UpperCamelCase_			  =						name.replace("blocks"			,						"layers"			)
    if "head.fc" in name:
        UpperCamelCase_			  =						name.replace("head.fc"			,						"classifier.1"			)
    if name.startswith("norm"			):
        UpperCamelCase_			  =						"bit." + name
    if "bit" not in name and "classifier" not in name:
        UpperCamelCase_			  =						"bit.encoder." + name

    return name




def 			lowerCAmelCase(      )->						int:



    """simple docstring"""






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




@torch.no_grad()
def 			lowerCAmelCase(      SCREAMING_SNAKE_CASE_			,						SCREAMING_SNAKE_CASE_			,						SCREAMING_SNAKE_CASE_=False			)->						Tuple:



    """simple docstring"""






    UpperCamelCase_			  =						get_config(SCREAMING_SNAKE_CASE_			)

    # load original model from timm
    UpperCamelCase_			  =						create_model(SCREAMING_SNAKE_CASE_			,						pretrained=SCREAMING_SNAKE_CASE_			)
    timm_model.eval()

    # load state_dict of original model
    UpperCamelCase_			  =						timm_model.state_dict()
    for key in state_dict.copy().keys():
        UpperCamelCase_			  =						state_dict.pop(SCREAMING_SNAKE_CASE_			)
        UpperCamelCase_			  =						val.squeeze() if "head" in key else val

    # load HuggingFace model
    UpperCamelCase_			  =						BitForImageClassification(SCREAMING_SNAKE_CASE_			)
    model.eval()
    model.load_state_dict(SCREAMING_SNAKE_CASE_			)

    # create image processor
    UpperCamelCase_			  =						create_transform(**resolve_data_config({}			,						model=SCREAMING_SNAKE_CASE_			)			)
    UpperCamelCase_			  =						transform.transforms

    UpperCamelCase_			  =						{
        "bilinear": PILImageResampling.BILINEAR,
        "bicubic": PILImageResampling.BICUBIC,
        "nearest": PILImageResampling.NEAREST,
    }

    UpperCamelCase_			  =						BitImageProcessor(
        do_resize=SCREAMING_SNAKE_CASE_			,						size={"shortest_edge": timm_transforms[0].size}			,						resample=pillow_resamplings[timm_transforms[0].interpolation.value]			,						do_center_crop=SCREAMING_SNAKE_CASE_			,						crop_size={"height": timm_transforms[1].size[0], "width": timm_transforms[1].size[1]}			,						do_normalize=SCREAMING_SNAKE_CASE_			,						image_mean=timm_transforms[-1].mean.tolist()			,						image_std=timm_transforms[-1].std.tolist()			,						)

    UpperCamelCase_			  =						prepare_img()
    UpperCamelCase_			  =						transform(SCREAMING_SNAKE_CASE_			).unsqueeze(0			)
    UpperCamelCase_			  =						processor(SCREAMING_SNAKE_CASE_			,						return_tensors="pt"			).pixel_values

    # verify pixel values
    assert torch.allclose(SCREAMING_SNAKE_CASE_			,						SCREAMING_SNAKE_CASE_			)

    # verify logits
    with torch.no_grad():
        UpperCamelCase_			  =						model(SCREAMING_SNAKE_CASE_			)
        UpperCamelCase_			  =						outputs.logits

    print("Logits:"			,						logits[0, :3]			)
    print("Predicted class:"			,						model.config.idalabel[logits.argmax(-1			).item()]			)
    UpperCamelCase_			  =						timm_model(SCREAMING_SNAKE_CASE_			)
    assert timm_logits.shape == outputs.logits.shape
    assert torch.allclose(SCREAMING_SNAKE_CASE_			,						outputs.logits			,						atol=1E-3			)
    print("Looks ok!"			)

    if pytorch_dump_folder_path is not None:
        Path(SCREAMING_SNAKE_CASE_			).mkdir(exist_ok=SCREAMING_SNAKE_CASE_			)
        print(f"Saving model {model_name} and processor to {pytorch_dump_folder_path}"			)
        model.save_pretrained(SCREAMING_SNAKE_CASE_			)
        processor.save_pretrained(SCREAMING_SNAKE_CASE_			)

    if push_to_hub:
        print(f"Pushing model {model_name} and processor to the hub"			)
        model.push_to_hub(f"ybelkada/{model_name}"			)
        processor.push_to_hub(f"ybelkada/{model_name}"			)


if __name__ == "__main__":
 SCREAMING_SNAKE_CASE		:Optional[Any]						=     argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     """--model_name""",
     default="""resnetv2_50x1_bitm""",
     type=str,
     help="""Name of the BiT 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."""
 )
 parser.add_argument(
     """--push_to_hub""",
     action="""store_true""",
     help="""Whether to push the model to the hub.""",
 )

 SCREAMING_SNAKE_CASE		:Union[str, Any]						=     parser.parse_args()
 convert_bit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)
