



'''simple docstring'''





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 ConvNextConfig, ConvNextForImageClassification, ConvNextImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
lowercase_			=  logging.get_logger(__name__)




def lowerCAmelCase  (__A):



     """simple docstring"""



     _a     =  ConvNextConfig()

     if "tiny" in checkpoint_url:
          _a     =  [3, 3, 9, 3]
          _a     =  [96, 192, 384, 768]
     if "small" in checkpoint_url:
          _a     =  [3, 3, 27, 3]
          _a     =  [96, 192, 384, 768]
     if "base" in checkpoint_url:
          _a     =  [3, 3, 27, 3]
          _a     =  [128, 256, 512, 1_024]
     if "large" in checkpoint_url:
          _a     =  [3, 3, 27, 3]
          _a     =  [192, 384, 768, 1_536]
     if "xlarge" in checkpoint_url:
          _a     =  [3, 3, 27, 3]
          _a     =  [256, 512, 1_024, 2_048]

     if "1k" in checkpoint_url:
          _a     =  1_000
          _a     =  '''imagenet-1k-id2label.json'''
          _a     =  (1, 1_000)
     else:
          _a     =  21_841
          _a     =  '''imagenet-22k-id2label.json'''
          _a     =  (1, 21_841)

     _a     =  '''huggingface/label-files'''
     _a     =  num_labels
     _a     =  json.load(open(hf_hub_download(_lowerCAmelCase  ,     _lowerCAmelCase  ,     repo_type='''dataset''')  ,     '''r'''))
     _a     =  {int(_lowerCAmelCase): v for k, v in idalabel.items()}
     if "1k" not in checkpoint_url:
          # this dataset contains 21843 labels but the model only has 21841
          # we delete the classes as mentioned in https://github.com/google-research/big_transfer/issues/18
          del idalabel[9_205]
          del idalabel[15_027]
     _a     =  idalabel
     _a     =  {v: k for k, v in idalabel.items()}
     _a     =  hidden_sizes
     _a     =  depths

     return config, expected_shape




def lowerCAmelCase  (__A):



     """simple docstring"""



     if "downsample_layers.0.0" in name:
          _a     =  name.replace('''downsample_layers.0.0'''  ,     '''embeddings.patch_embeddings''')
     if "downsample_layers.0.1" in name:
          _a     =  name.replace('''downsample_layers.0.1'''  ,     '''embeddings.norm''')  # we rename to layernorm later on
     if "downsample_layers.1.0" in name:
          _a     =  name.replace('''downsample_layers.1.0'''  ,     '''stages.1.downsampling_layer.0''')
     if "downsample_layers.1.1" in name:
          _a     =  name.replace('''downsample_layers.1.1'''  ,     '''stages.1.downsampling_layer.1''')
     if "downsample_layers.2.0" in name:
          _a     =  name.replace('''downsample_layers.2.0'''  ,     '''stages.2.downsampling_layer.0''')
     if "downsample_layers.2.1" in name:
          _a     =  name.replace('''downsample_layers.2.1'''  ,     '''stages.2.downsampling_layer.1''')
     if "downsample_layers.3.0" in name:
          _a     =  name.replace('''downsample_layers.3.0'''  ,     '''stages.3.downsampling_layer.0''')
     if "downsample_layers.3.1" in name:
          _a     =  name.replace('''downsample_layers.3.1'''  ,     '''stages.3.downsampling_layer.1''')
     if "stages" in name and "downsampling_layer" not in name:
          # stages.0.0. for instance should be renamed to stages.0.layers.0.
          _a     =  name[: len('''stages.0''')] + '''.layers''' + name[len('''stages.0''') :]
     if "stages" in name:
          _a     =  name.replace('''stages'''  ,     '''encoder.stages''')
     if "norm" in name:
          _a     =  name.replace('''norm'''  ,     '''layernorm''')
     if "gamma" in name:
          _a     =  name.replace('''gamma'''  ,     '''layer_scale_parameter''')
     if "head" in name:
          _a     =  name.replace('''head'''  ,     '''classifier''')

     return name




def lowerCAmelCase  ():



     """simple docstring"""



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




@torch.no_grad()
def lowerCAmelCase  (__A  ,     __A):



     """simple docstring"""



     _a  ,   _a     =  get_convnext_config(_lowerCAmelCase)
     # load original state_dict from URL
     _a     =  torch.hub.load_state_dict_from_url(_lowerCAmelCase)['''model''']
     # rename keys
     for key in state_dict.copy().keys():
          _a     =  state_dict.pop(_lowerCAmelCase)
          _a     =  val
     # add prefix to all keys expect classifier head
     for key in state_dict.copy().keys():
          _a     =  state_dict.pop(_lowerCAmelCase)
          if not key.startswith('''classifier'''):
               _a     =  '''convnext.''' + key
          _a     =  val

     # load HuggingFace model
     _a     =  ConvNextForImageClassification(_lowerCAmelCase)
     model.load_state_dict(_lowerCAmelCase)
     model.eval()

     # Check outputs on an image, prepared by ConvNextImageProcessor
     _a     =  224 if '''224''' in checkpoint_url else 384
     _a     =  ConvNextImageProcessor(size=_lowerCAmelCase)
     _a     =  image_processor(images=prepare_img()  ,     return_tensors='''pt''').pixel_values

     _a     =  model(_lowerCAmelCase).logits

     # note: the logits below were obtained without center cropping
     if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth":
          _a     =  torch.tensor([-0.12_10, -0.66_05, 0.19_18])
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
          _a     =  torch.tensor([-0.44_73, -0.18_47, -0.63_65])
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
          _a     =  torch.tensor([0.45_25, 0.75_39, 0.03_08])
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
          _a     =  torch.tensor([0.35_61, 0.63_50, -0.03_84])
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
          _a     =  torch.tensor([0.41_74, -0.09_89, 0.14_89])
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
          _a     =  torch.tensor([0.25_13, -0.13_49, -0.16_13])
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
          _a     =  torch.tensor([1.29_80, 0.36_31, -0.11_98])
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
          _a     =  torch.tensor([1.29_63, 0.12_27, 0.17_23])
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
          _a     =  torch.tensor([1.79_56, 0.83_90, 0.28_20])
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
          _a     =  torch.tensor([-0.28_22, -0.05_02, -0.08_78])
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
          _a     =  torch.tensor([-0.56_72, -0.07_30, -0.43_48])
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
          _a     =  torch.tensor([0.26_81, 0.23_65, 0.62_46])
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
          _a     =  torch.tensor([-0.26_42, 0.39_31, 0.51_16])
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
          _a     =  torch.tensor([-0.66_77, -0.18_73, -0.83_79])
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
          _a     =  torch.tensor([-0.77_49, -0.29_67, -0.64_44])
     else:
          raise ValueError(F'''Unknown URL: {checkpoint_url}''')

     assert torch.allclose(logits[0, :3]  ,     _lowerCAmelCase  ,     atol=1e-3)
     assert logits.shape == expected_shape

     Path(_lowerCAmelCase).mkdir(exist_ok=_lowerCAmelCase)
     print(F'''Saving model to {pytorch_dump_folder_path}''')
     model.save_pretrained(_lowerCAmelCase)
     print(F'''Saving image processor to {pytorch_dump_folder_path}''')
     image_processor.save_pretrained(_lowerCAmelCase)

     print('''Pushing model to the hub...''')
     _a     =  '''convnext'''
     if "tiny" in checkpoint_url:
          model_name += "-tiny"
     elif "small" in checkpoint_url:
          model_name += "-small"
     elif "base" in checkpoint_url:
          model_name += "-base"
     elif "xlarge" in checkpoint_url:
          model_name += "-xlarge"
     elif "large" in checkpoint_url:
          model_name += "-large"
     if "224" in checkpoint_url:
          model_name += "-224"
     elif "384" in checkpoint_url:
          model_name += "-384"
     if "22k" in checkpoint_url and "1k" not in checkpoint_url:
          model_name += "-22k"
     if "22k" in checkpoint_url and "1k" in checkpoint_url:
          model_name += "-22k-1k"

     model.push_to_hub(
         repo_path_or_name=Path(_lowerCAmelCase  ,     _lowerCAmelCase)  ,     organization='''nielsr'''  ,     commit_message='''Add model'''  ,     )


if __name__ == "__main__":
  lowercase_			=  argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      "--checkpoint_url",
      default="https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth",
      type=str,
      help="URL of the original ConvNeXT checkpoint you'd like to convert.",
  )
  parser.add_argument(
      "--pytorch_dump_folder_path",
      default=None,
      type=str,
      required=True,
      help="Path to the output PyTorch model directory.",
  )

  lowercase_			=  parser.parse_args()
  convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)








'''simple docstring'''





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

from tokenizers import processors

from ...tokenization_utils import AddedToken, BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
  from .tokenization_mbartaa import MBartaaTokenizer
else:
  lowercase_			=  None


lowercase_			=  logging.get_logger(__name__)

lowercase_			=  {"vocab_file": "sentencepiece.bpe.model", "tokenizer_file": "tokenizer.json"}

lowercase_			=  {
    "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"
        ),
    },
    "tokenizer_file": {
        "facebook/mbart-large-50-one-to-many-mmt": (
            "https://huggingface.co/facebook/mbart-large-50-one-to-many-mmt/resolve/main/tokenizer.json"
        ),
    },
}

lowercase_			=  {
    "facebook/mbart-large-50-one-to-many-mmt": 1_024,
}

# fmt: off
lowercase_			=  ["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 __A      (  A						):



 '''simple docstring'''







 __lowerCamelCase   :	List[str]										=	VOCAB_FILES_NAMES
 __lowerCamelCase   :	int										=	PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 __lowerCamelCase   :	Optional[int]										=	PRETRAINED_VOCAB_FILES_MAP
 __lowerCamelCase   :	Any										=	['input_ids', 'attention_mask']
 __lowerCamelCase   :	Optional[Any]										=	MBartaaTokenizer

 __lowerCamelCase   :	List[int]										=	[]
 __lowerCamelCase   :	List[int]										=	[]



 def __init__(self			,       A=None			,       A=None			,       A=None			,       A=None			,       A="</s>"			,       A="</s>"			,       A="<s>"			,       A="<unk>"			,       A="<pad>"			,       A="<mask>"			,       **A			,       )   ->					List[str]:


      """simple docstring"""



      _a     =  AddedToken(A			,       lstrip=A			,       rstrip=A	) if isinstance(A			,       A	) else mask_token

      _a     =  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__(
          A			,       src_lang=A			,       tgt_lang=A			,       tokenizer_file=A			,       eos_token=A			,       sep_token=A			,       cls_token=A			,       unk_token=A			,       pad_token=A			,       mask_token=A			,       **A			,       )

      _a     =  vocab_file
      _a     =  False if not self.vocab_file else True

      _a     =  {
          lang_code: self.convert_tokens_to_ids(A	) for lang_code in FAIRSEQ_LANGUAGE_CODES
      }

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



 @property
 def  a__  (self	)   ->					str:


      """simple docstring"""



      return self._src_lang



 @src_lang.setter
 def  a__  (self			,       A	)   ->					None:


      """simple docstring"""



      _a     =  new_src_lang
      self.set_src_lang_special_tokens(self._src_lang	)



 def  a__  (self			,       A			,       A = 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  a__  (self			,       A			,       A = "en_XX"			,       A = None			,       A = "ro_RO"			,       **A			,       )   ->					BatchEncoding:


      """simple docstring"""



      _a     =  src_lang
      _a     =  tgt_lang
      return super().prepare_seqaseq_batch(A			,       A			,       **A	)



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


      """simple docstring"""



      return self.set_src_lang_special_tokens(self.src_lang	)



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


      """simple docstring"""



      return self.set_tgt_lang_special_tokens(self.tgt_lang	)



 def  a__  (self			,       A	)   ->					None:


      """simple docstring"""



      _a     =  self.convert_tokens_to_ids(A	)
      _a     =  [self.cur_lang_code_id]
      _a     =  [self.eos_token_id]

      _a     =  self.convert_ids_to_tokens(self.prefix_tokens	)
      _a     =  self.convert_ids_to_tokens(self.suffix_tokens	)

      _a     =  processors.TemplateProcessing(
          single=prefix_tokens_str + ['''$A'''] + suffix_tokens_str			,       pair=prefix_tokens_str + ['''$A''', '''$B'''] + suffix_tokens_str			,       special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str			,       self.prefix_tokens + self.suffix_tokens	)	)			,       )



 def  a__  (self			,       A	)   ->					None:


      """simple docstring"""



      _a     =  self.convert_tokens_to_ids(A	)
      _a     =  [self.cur_lang_code_id]
      _a     =  [self.eos_token_id]

      _a     =  self.convert_ids_to_tokens(self.prefix_tokens	)
      _a     =  self.convert_ids_to_tokens(self.suffix_tokens	)

      _a     =  processors.TemplateProcessing(
          single=prefix_tokens_str + ['''$A'''] + suffix_tokens_str			,       pair=prefix_tokens_str + ['''$A''', '''$B'''] + suffix_tokens_str			,       special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str			,       self.prefix_tokens + self.suffix_tokens	)	)			,       )



 def  a__  (self			,       A			,       A			,       A			,       A			,       **A	)   ->					Optional[int]:


      """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'''	)
      _a     =  src_lang
      _a     =  self(A			,       add_special_tokens=A			,       return_tensors=A			,       **A	)
      _a     =  self.convert_tokens_to_ids(A	)
      _a     =  tgt_lang_id
      return inputs



 def  a__  (self			,       A			,       A = None	)   ->					Tuple[str]:


      """simple docstring"""



      if not self.can_save_slow_tokenizer:
           raise ValueError(
               '''Your fast tokenizer does not have the necessary information to save the vocabulary for a slow '''
               '''tokenizer.'''	)

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

      if os.path.abspath(self.vocab_file	) != os.path.abspath(A	):
           copyfile(self.vocab_file			,       A	)

      return (out_vocab_file,)



