code: 
  
  
  
  
  
  '''simple docstring'''
  
  
  
  
  
  import os
  from shutil import copyfile
  from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
  
  import sentencepiece as spm
  
  from ...tokenization_utils import AddedToken, PreTrainedTokenizer
  from ...utils import logging
  
  
  if TYPE_CHECKING:
       from transformers.pipelines.conversational import Conversation
  
  UpperCAmelCase			:			str	  = logging.get_logger(__name__)
  
  UpperCAmelCase			:			List[Any]	  = {"""vocab_file""": """tokenizer.model"""}
  
  UpperCAmelCase			:			Dict	  = {
      """vocab_file""": {
          """hf-internal-testing/llama-tokenizer""": """https://huggingface.co/hf-internal-testing/llama-tokenizer/resolve/main/tokenizer.model""",
      },
      """tokenizer_file""": {
          """hf-internal-testing/llama-tokenizer""": """https://huggingface.co/hf-internal-testing/llama-tokenizer/resolve/main/tokenizer_config.json""",
      },
  }
  UpperCAmelCase			:			Optional[int]	  = {
      """hf-internal-testing/llama-tokenizer""": 2048,
  }
  UpperCAmelCase			:			Any	  = """▁"""
  
  UpperCAmelCase			:			int	  = """[INST]""", """[/INST]"""
  UpperCAmelCase			:			List[str]	  = """<<SYS>>\n""", """\n<</SYS>>\n\n"""
  
  # fmt: off
  UpperCAmelCase			:			List[str]	  = """You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your\
  answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure\
  that your responses are socially unbiased and positive in nature.
  
  If a question does not make any sense, or is not factually coherent, explain why instead of answering something not\
  correct. If you don't know the answer to a question, please don't share false information."""
  
  
  
  
  
  
  
  class        __lowerCAmelCase (					UpperCamelCase__):
        _lowercase       :      Dict        =    VOCAB_FILES_NAMES
        _lowercase       :      Optional[Any]        =    PRETRAINED_VOCAB_FILES_MAP
        _lowercase       :      Optional[Any]        =    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
        _lowercase       :      Dict        =    ["""input_ids""", """attention_mask"""]
  
  
  
  
  
  
  
        def __init__(					self	,   lowerCAmelCase__	,   lowerCAmelCase__="<unk>"	,   lowerCAmelCase__="<s>"	,   lowerCAmelCase__="</s>"	,   lowerCAmelCase__=None	,   lowerCAmelCase__ = None	,   lowerCAmelCase__=True	,   lowerCAmelCase__=False	,   lowerCAmelCase__=False	,   lowerCAmelCase__=True	,   **lowerCAmelCase__	,   )     ->		Optional[Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      Union[str, Any]            ={} if sp_model_kwargs is None else sp_model_kwargs
          a__			:      Union[str, Any]            =AddedToken(lowerCAmelCase__	,   lstrip=lowerCAmelCase__	,   rstrip=lowerCAmelCase__							) if isinstance(lowerCAmelCase__	,   lowerCAmelCase__							) else bos_token
          a__			:      Any            =AddedToken(lowerCAmelCase__	,   lstrip=lowerCAmelCase__	,   rstrip=lowerCAmelCase__							) if isinstance(lowerCAmelCase__	,   lowerCAmelCase__							) else eos_token
          a__			:      int            =AddedToken(lowerCAmelCase__	,   lstrip=lowerCAmelCase__	,   rstrip=lowerCAmelCase__							) if isinstance(lowerCAmelCase__	,   lowerCAmelCase__							) else unk_token
          a__			:      int            =AddedToken(lowerCAmelCase__	,   lstrip=lowerCAmelCase__	,   rstrip=lowerCAmelCase__							) if isinstance(lowerCAmelCase__	,   lowerCAmelCase__							) else pad_token
          super().__init__(
              bos_token=lowerCAmelCase__	,   eos_token=lowerCAmelCase__	,   unk_token=lowerCAmelCase__	,   pad_token=lowerCAmelCase__	,   add_bos_token=lowerCAmelCase__	,   add_eos_token=lowerCAmelCase__	,   sp_model_kwargs=self.sp_model_kwargs	,   clean_up_tokenization_spaces=lowerCAmelCase__	,   legacy=lowerCAmelCase__	,   **lowerCAmelCase__	,   )
          if legacy:
            logger.warning_once(
                F'''You are using the legacy behaviour of the {self.__class__}. This means that tokens that come after special tokens will not be properly handled. We recommend you to'''
                " read the related pull request available at https://github.com/huggingface/transformers/pull/24565"							)
          a__			:      Optional[Any]            =legacy
          a__			:      Optional[int]            =vocab_file
          a__			:      Optional[int]            =add_bos_token
          a__			:      Dict            =add_eos_token
          a__			:      Optional[Any]            =spm.SentencePieceProcessor(**self.sp_model_kwargs							)
          self.sp_model.Load(lowerCAmelCase__							)
  
  
  
  
  
  
  
        def __getstate__(					self							)     ->		int:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      Optional[int]            =self.__dict__.copy()
          a__			:      int            =None
          a__			:      Union[str, Any]            =self.sp_model.serialized_model_proto()
          return state
  
  
  
  
  
  
  
        def __setstate__(					self	,   lowerCAmelCase__							)     ->		Union[str, Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      Optional[Any]            =d
          a__			:      Optional[int]            =spm.SentencePieceProcessor(**self.sp_model_kwargs							)
          self.sp_model.LoadFromSerializedProto(self.sp_model_proto							)
  
  
  
  
  
  
  
        @property
        def      _lowercase						(					self							)     ->		str:
  
  
  
          '''simple docstring'''
  
  
  
  
          return self.sp_model.get_piece_size()
  
  
  
  
  
  
  
        def      _lowercase						(					self							)     ->		Any:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      Union[str, Any]            ={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	,   lowerCAmelCase__	,   **lowerCAmelCase__							)     ->		List[str]:
  
  
  
          '''simple docstring'''
  
  
  
  
          if not self.legacy:
            a__			:      Union[str, Any]            =SPIECE_UNDERLINE + text.replace(lowerCAmelCase__	,   " "							)
          return super().tokenize(lowerCAmelCase__	,   **lowerCAmelCase__							)
  
  
  
  
  
  
  
        def      _lowercase						(					self	,   lowerCAmelCase__							)     ->		List[Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          if not self.legacy:
            a__			:      Union[str, Any]            =text.startswith(lowerCAmelCase__							)
            if is_first:
              a__			:      Optional[int]            =text[1:]
  
          a__			:      Dict            =self.sp_model.encode(lowerCAmelCase__	,   out_type=lowerCAmelCase__							)
  
          if not self.legacy and not is_first and not text.startswith(" "							) and tokens[0].startswith(lowerCAmelCase__							):
            a__			:      Optional[int]            =([tokens[0][1:]] if len(tokens[0]							) > 1 else []) + tokens[1:]
          return tokens
  
  
  
  
  
  
  
        def      _lowercase						(					self	,   lowerCAmelCase__							)     ->		Optional[Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          return self.sp_model.piece_to_id(lowerCAmelCase__							)
  
  
  
  
  
  
  
        def      _lowercase						(					self	,   lowerCAmelCase__							)     ->		Optional[Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      Optional[Any]            =self.sp_model.IdToPiece(lowerCAmelCase__							)
          return token
  
  
  
  
  
  
  
        def      _lowercase						(					self	,   lowerCAmelCase__							)     ->		Union[str, Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      Union[str, Any]            =[]
          a__			:      Tuple            =""
          a__			:      int            =False
          for i, token in enumerate(lowerCAmelCase__							):
            # make sure that special tokens are not decoded using sentencepiece model
            if token in self.all_special_tokens:
              if not prev_is_special and i != 0:
                out_string += " "
              out_string += self.sp_model.decode(lowerCAmelCase__							) + token
              a__			:      Optional[int]            =True
              a__			:      List[str]            =[]
            else:
              current_sub_tokens.append(lowerCAmelCase__							)
              a__			:      List[Any]            =False
          out_string += self.sp_model.decode(lowerCAmelCase__							)
          return out_string
  
  
  
  
  
  
  
        def      _lowercase						(					self	,   lowerCAmelCase__	,   lowerCAmelCase__ = None							)     ->		Tuple[str]:
  
  
  
          '''simple docstring'''
  
  
  
  
          if not os.path.isdir(lowerCAmelCase__							):
            logger.error(F'''Vocabulary path ({save_directory}) should be a directory'''							)
            return
          a__			:      Optional[Any]            =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:
              a__			:      Union[str, Any]            =self.sp_model.serialized_model_proto()
              fi.write(lowerCAmelCase__							)
  
          return (out_vocab_file,)
  
  
  
  
  
  
  
        def      _lowercase						(					self	,   lowerCAmelCase__	,   lowerCAmelCase__=None							)     ->		str:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      Any            =[self.bos_token_id] if self.add_bos_token else []
          a__			:      List[Any]            =[self.eos_token_id] if self.add_eos_token else []
  
          a__			:      Dict            =bos_token_id + token_ids_a + eos_token_id
  
          if token_ids_a is not None:
            a__			:      Dict            =output + bos_token_id + token_ids_a + eos_token_id
  
          return output
  
  
  
  
  
  
  
        def      _lowercase						(					self	,   lowerCAmelCase__	,   lowerCAmelCase__ = None	,   lowerCAmelCase__ = False							)     ->		List[int]:
  
  
  
          '''simple docstring'''
  
  
  
  
          if already_has_special_tokens:
            return super().get_special_tokens_mask(
                token_ids_a=lowerCAmelCase__	,   token_ids_a=lowerCAmelCase__	,   already_has_special_tokens=lowerCAmelCase__							)
  
          a__			:      Any            =[1] if self.add_bos_token else []
          a__			:      Any            =[1] if self.add_eos_token else []
  
          if token_ids_a is None:
            return bos_token_id + ([0] * len(lowerCAmelCase__							)) + eos_token_id
          return (
              bos_token_id
              + ([0] * len(lowerCAmelCase__							))
              + eos_token_id
              + bos_token_id
              + ([0] * len(lowerCAmelCase__							))
              + eos_token_id
          )
  
  
  
  
  
  
  
        def      _lowercase						(					self	,   lowerCAmelCase__	,   lowerCAmelCase__ = None							)     ->		List[int]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      str            =[self.bos_token_id] if self.add_bos_token else []
          a__			:      Dict            =[self.eos_token_id] if self.add_eos_token else []
  
          a__			:      int            =[0] * len(bos_token_id + token_ids_a + eos_token_id							)
  
          if token_ids_a is not None:
            output += [1] * len(bos_token_id + token_ids_a + eos_token_id							)
  
          return output
  
  
  
  
  
  
  
        def      _lowercase						(					self	,   lowerCAmelCase__							)     ->		List[int]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      Tuple            =list(conversation.iter_texts()							)
          if not all([is_user for is_user, msg in dialogue[::2]]							) or not all(
              [not is_user for is_user, msg in dialogue[1::2]]							):
            raise ValueError(
                "The model only supports 'user' and 'assistant' roles, starting with user and alternating (u/a/u/a/u...)"							)
  
          a__			:      List[int]            =[]
          if len(conversation.past_user_inputs							) > 0:
            if not conversation.past_user_inputs[0].startswith(lowerCAmelCase__							) or E_SYS not in conversation.past_user_inputs[0]:
              a__			:      Dict            =(
                  B_SYS + DEFAULT_SYSTEM_PROMPT + E_SYS + conversation.past_user_inputs[0]
              )
          elif not dialogue[0][1].startswith(lowerCAmelCase__							) or E_SYS not in dialogue[0][1]:
            a__			:      Optional[Any]            =(dialogue[0][0], B_SYS + DEFAULT_SYSTEM_PROMPT + E_SYS + dialogue[0][1])
  
          dialog_tokens += sum(
              [
                  [self.bos_token_id]
                  + self.encode(
                      F'''{B_INST} {(prompt[1]).strip()} {E_INST} {(answer[1]).strip()} '''	,   add_special_tokens=lowerCAmelCase__							)
                  + [self.eos_token_id]
                  for prompt, answer in zip(dialogue[::2]	,   dialogue[1::2]							)
              ]	,   []	,   )
          if not (dialogue[-1][0]):
            raise ValueError(F'''Last message must be from user, got {dialogue[-1]["role"]}'''							)
          dialog_tokens += [self.bos_token_id] + self.encode(
              F'''{B_INST} {(dialogue[-1][1]).strip()} {E_INST}'''	,   add_special_tokens=lowerCAmelCase__							)
          return dialog_tokens
  
  
  
  
  
  
  
style_context: 
  
  
  
  
  
  import argparse
  import json
  import os
  import re
  from collections import OrderedDict
  from dataclasses import dataclass, field
  from functools import partial
  from pathlib import Path
  from pprint import pprint
  from typing import Dict, List, Tuple
  
  import torch
  import torch.nn as nn
  from classy_vision.models.regnet import RegNet, RegNetParams
  from huggingface_hub import cached_download, hf_hub_url
  from torch import Tensor
  from vissl.models.model_helpers import get_trunk_forward_outputs
  
  from transformers import AutoImageProcessor, RegNetConfig, RegNetForImageClassification, RegNetModel
  from transformers.modeling_utils import PreTrainedModel
  from transformers.utils import logging
  
  
  logging.set_verbosity_info()
  UpperCAmelCase			:			Optional[Any]	  = logging.get_logger()
  
  
  
  
  
  
  
  @dataclass
  class        __lowerCAmelCase :
        _lowercase       :      nn.Module
        _lowercase       :      List[nn.Module]        =    field(default_factory=UpperCamelCase__)
        _lowercase       :      list        =    field(default_factory=UpperCamelCase__)
        _lowercase       :      Dict[str, nn.Module]        =    field(default_factory=UpperCamelCase__)
  
  
  
  
  
  
  
        def      _lowercase						(					self	,   lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__							)     ->		str:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      Dict            =len(list(m.modules()							)							) == 1 or isinstance(lowerCAmelCase__	,   nn.Convad							) or isinstance(lowerCAmelCase__	,   nn.BatchNormad							)
          if has_not_submodules:
            self.traced.append(lowerCAmelCase__							)
            a__			:      List[Any]            =m
  
  
  
  
  
  
  
        def __call__(					self	,   lowerCAmelCase__							)     ->		Optional[Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          for name, m in self.module.named_modules():
            self.handles.append(m.register_forward_hook(partial(self._forward_hook	,   name=lowerCAmelCase__							)							)							)
          self.module(lowerCAmelCase__							)
          [x.remove() for x in self.handles]
          return self
  
  
  
  
  
  
  
        @property
        def      _lowercase						(					self							)     ->		Optional[Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          return {k: v for k, v in self.nameamodule.items() if len(list(v.state_dict().keys()							)							) > 0}
  
  
  
  
  
  
  
  class        __lowerCAmelCase (					nn.Module):
  
  
  
  
  
  
  
        def __init__(					self	,   lowerCAmelCase__							)     ->		List[Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          super().__init__()
  
          a__			:      List[Tuple[str, nn.Module]]            =[]
          # - get the stem
          feature_blocks.append(("conv1", model.stem)							)
          # - get all the feature blocks
          for k, v in model.trunk_output.named_children():
            assert k.startswith("block"							), F'''Unexpected layer name {k}'''
            a__			:      int            =len(lowerCAmelCase__							) + 1
            feature_blocks.append((F'''res{block_index}''', v)							)
  
          a__			:      List[str]            =nn.ModuleDict(lowerCAmelCase__							)
  
  
  
  
  
  
  
        def      _lowercase						(					self	,   lowerCAmelCase__							)     ->		Union[str, Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          return get_trunk_forward_outputs(
              lowerCAmelCase__	,   out_feat_keys=lowerCAmelCase__	,   feature_blocks=self._feature_blocks	,   )
  
  
  
  
  
  
  
  class        __lowerCAmelCase (					UpperCamelCase__):
  
  
  
  
  
  
  
        def      _lowercase						(					self							)     ->		Tuple:
  
  
  
          '''simple docstring'''
  
  
  
  
          return [(8, 2, 2, 8, 1.0), (8, 2, 7, 8, 1.0), (8, 2, 1_7, 8, 1.0), (8, 2, 1, 8, 1.0)]
  def   _A (	SCREAMING_SNAKE_CASE  :     str		):
  
  
  
  
  
    """simple docstring"""
  
  
  
    a__			:      Union[str, Any]            =RegNetConfig(depths=[2, 7, 17, 1] , hidden_sizes=[8, 8, 8, 8] , groups_width=8		)
    if "in1k" in model_name:
      a__			:      Dict            =RegNetForImageClassification(SCREAMING_SNAKE_CASE		)
    else:
      a__			:      List[Any]            =RegNetModel(SCREAMING_SNAKE_CASE		)
    # create from model (with small weights)
    a__			:      Tuple            =FakeRegNetVisslWrapper(
        RegNet(FakeRegNetParams(depth=27 , group_width=1_010 , w_a=1_744 , w_a=6_2_0.8_3 , w_m=2.5_2		)		)		)
  
    with torch.no_grad():
      a__			:      int            =from_model.eval()
      a__			:      int            =our_model.eval()
  
      a__			:      Dict            =torch.randn((1, 3, 32, 32)		)
      # trace both
      a__			:      Tuple            =Tracker(SCREAMING_SNAKE_CASE		)
      a__			:      List[Any]            =dest_tracker(SCREAMING_SNAKE_CASE		).parametrized
  
      pprint(dest_tracker.nameamodule		)
      a__			:      Optional[Any]            =Tracker(SCREAMING_SNAKE_CASE		)
      a__			:      Tuple            =src_tracker(SCREAMING_SNAKE_CASE		).parametrized
  
    # convert the keys -> module dict to keys -> params
    def to_params_dict(SCREAMING_SNAKE_CASE  :     Any		):
      a__			:      List[Any]            =OrderedDict()
      for name, module in dict_with_modules.items():
        for param_name, param in module.state_dict().items():
          a__			:      Optional[Any]            =param
      return params_dict
  
    a__			:      Union[str, Any]            ={}
  
    a__			:      Optional[int]            =to_params_dict(SCREAMING_SNAKE_CASE		)
    a__			:      List[Any]            =to_params_dict(SCREAMING_SNAKE_CASE		)
  
    for (src_key, src_param), (dest_key, dest_param) in zip(src_state_dict.items() , dst_state_dict.items()		):
      a__			:      List[str]            =dest_key
      logger.info(f'''{src_key} -> {dest_key}'''		)
    # if "in1k" was in the model_name it means it must have a classification head (was finetuned)
    if "in1k" in model_name:
      a__			:      str            ="classifier.1.weight"
      a__			:      Optional[Any]            ="classifier.1.bias"
  
    return from_to_ours_keys
  def   _A (	SCREAMING_SNAKE_CASE  :     Path , SCREAMING_SNAKE_CASE  :     str = None , SCREAMING_SNAKE_CASE  :     bool = True		):
  
  
  
  
  
    """simple docstring"""
  
  
  
    a__			:      int            ="imagenet-1k-id2label.json"
    a__			:      int            =1_000
  
    a__			:      Optional[int]            ="huggingface/label-files"
    a__			:      Union[str, Any]            =num_labels
    a__			:      Optional[int]            =json.load(open(cached_download(hf_hub_url(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , repo_type="dataset"		)		) , "r"		)		)
    a__			:      Tuple            ={int(SCREAMING_SNAKE_CASE		): v for k, v in idalabel.items()}
  
    a__			:      Any            =idalabel
    a__			:      Optional[int]            ={v: k for k, v in idalabel.items()}
  
    a__			:      str            =partial(SCREAMING_SNAKE_CASE , num_labels=SCREAMING_SNAKE_CASE , idalabel=SCREAMING_SNAKE_CASE , labelaid=SCREAMING_SNAKE_CASE		)
  
    a__			:      Tuple            ={
        "regnet-y-10b-seer": ImageNetPreTrainedConfig(
            depths=[2, 7, 17, 1] , hidden_sizes=[2_020, 4_040, 11_110, 28_280] , groups_width=1_010		),
        # finetuned on imagenet
        "regnet-y-10b-seer-in1k": ImageNetPreTrainedConfig(
            depths=[2, 7, 17, 1] , hidden_sizes=[2_020, 4_040, 11_110, 28_280] , groups_width=1_010		),
    }
  
    # add seer weights logic
    def load_using_classy_vision(SCREAMING_SNAKE_CASE  :     str		) -> Tuple[Dict, Dict]:
      a__			:      Optional[int]            =torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE , model_dir=str(SCREAMING_SNAKE_CASE		) , map_location="cpu"		)
      # check if we have a head, if yes add it
      a__			:      Optional[Any]            =files["classy_state_dict"]["base_model"]["model"]
      return model_state_dict["trunk"], model_state_dict["heads"]
  
    a__			:      int            ={
        "regnet-y-10b-seer": partial(
            SCREAMING_SNAKE_CASE , "https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_regnet10B/model_iteration124500_conso.torch" , ),
        "regnet-y-10b-seer-in1k": partial(
            SCREAMING_SNAKE_CASE , "https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_finetuned/seer_10b_finetuned_in1k_model_phase28_conso.torch" , ),
    }
  
    a__			:      List[Any]            =get_from_to_our_keys(SCREAMING_SNAKE_CASE		)
  
    if not (save_directory / f'''{model_name}.pth''').exists():
      logger.info("Loading original state_dict."		)
      a__      ,							a__			:      Optional[int]            =names_to_from_model[model_name]()
      a__			:      Dict            =from_state_dict_trunk
      if "in1k" in model_name:
        # add the head
        a__			:      Optional[int]            ={**from_state_dict_trunk, **from_state_dict_head}
      logger.info("Done!"		)
  
      a__			:      List[str]            ={}
  
      a__			:      Any            =list(from_state_dict.keys()		)
      a__			:      Union[str, Any]            =r"\.block.-part."
      # this is "interesting", so the original checkpoints have `block[0,1]-part` in each key name, we remove it
      for key in from_state_dict.keys():
        # remove the weird "block[0,1]-part" from the key
        a__			:      int            =re.sub(SCREAMING_SNAKE_CASE , "" , SCREAMING_SNAKE_CASE		)
        # now src_key from the model checkpoints is the one we got from the original model after tracing, so use it to get the correct destination key
        a__			:      Optional[int]            =from_to_ours_keys[src_key]
        # store the parameter with our key
        a__			:      Tuple            =from_state_dict[key]
        not_used_keys.remove(SCREAMING_SNAKE_CASE		)
      # check that all keys have been updated
      assert len(SCREAMING_SNAKE_CASE		) == 0, f'''Some keys where not used {",".join(SCREAMING_SNAKE_CASE		)}'''
  
      logger.info(f'''The following keys were not used: {",".join(SCREAMING_SNAKE_CASE		)}'''		)
  
      # save our state dict to disk
      torch.save(SCREAMING_SNAKE_CASE , save_directory / f'''{model_name}.pth'''		)
  
      del converted_state_dict
    else:
      logger.info("The state_dict was already stored on disk."		)
    if push_to_hub:
      logger.info(f'''Token is {os.environ["HF_TOKEN"]}'''		)
      logger.info("Loading our model."		)
      # create our model
      a__			:      Tuple            =names_to_config[model_name]
      a__			:      Tuple            =RegNetModel
      if "in1k" in model_name:
        a__			:      Any            =RegNetForImageClassification
      a__			:      Tuple            =our_model_func(SCREAMING_SNAKE_CASE		)
      # place our model to the meta device (so remove all the weights)
      our_model.to(torch.device("meta"		)		)
      logger.info("Loading state_dict in our model."		)
      # load state dict
      a__			:      int            =our_model.state_dict().keys()
      PreTrainedModel._load_pretrained_model_low_mem(
          SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , [save_directory / f'''{model_name}.pth''']		)
      logger.info("Finally, pushing!"		)
      # push it to hub
      our_model.push_to_hub(
          repo_path_or_name=save_directory / model_name , commit_message="Add model" , output_dir=save_directory / model_name , )
      a__			:      Dict            =384
      # we can use the convnext one
      a__			:      Any            =AutoImageProcessor.from_pretrained("facebook/convnext-base-224-22k-1k" , size=SCREAMING_SNAKE_CASE		)
      image_processor.push_to_hub(
          repo_path_or_name=save_directory / model_name , commit_message="Add image processor" , output_dir=save_directory / model_name , )
  
  
  if __name__ == "__main__":
       UpperCAmelCase			:			Tuple	  = argparse.ArgumentParser()
       # Required parameters
       parser.add_argument(
           """--model_name""",
           default=None,
           type=str,
           help=(
               """The name of the model you wish to convert, it must be one of the supported regnet* architecture,"""
               """ currently: regnetx-*, regnety-*. If `None`, all of them will the converted."""
           ),
       )
       parser.add_argument(
           """--pytorch_dump_folder_path""",
           default=None,
           type=Path,
           required=True,
           help="""Path to the output PyTorch model directory.""",
       )
       parser.add_argument(
           """--push_to_hub""",
           default=True,
           type=bool,
           required=False,
           help="""If True, push model and image processor to the hub.""",
       )
  
       UpperCAmelCase			:			int	  = parser.parse_args()
  
       UpperCAmelCase			:			Path	  = args.pytorch_dump_folder_path
       pytorch_dump_folder_path.mkdir(exist_ok=True, parents=True)
       convert_weights_and_push(pytorch_dump_folder_path, args.model_name, args.push_to_hub)
  
  
  
  
  
  
  