"""simple docstring"""







from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
     import tensorflow as tf

     from transformers import (
         TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
         TFXLMForMultipleChoice,
         TFXLMForQuestionAnsweringSimple,
         TFXLMForSequenceClassification,
         TFXLMForTokenClassification,
         TFXLMModel,
         TFXLMWithLMHeadModel,
         XLMConfig,
     )






class      _SCREAMING_SNAKE_CASE							:



   '''simple docstring'''



   def __init__(	self ,				_SCREAMING_SNAKE_CASE ,				):


      '''simple docstring'''

      __lowerCamelCase    		=					parent
      __lowerCamelCase    		=					13
      __lowerCamelCase    		=					7
      __lowerCamelCase    		=					True
      __lowerCamelCase    		=					True
      __lowerCamelCase    		=					True
      __lowerCamelCase    		=					True
      __lowerCamelCase    		=					True
      __lowerCamelCase    		=					False
      __lowerCamelCase    		=					False
      __lowerCamelCase    		=					False
      __lowerCamelCase    		=					2
      __lowerCamelCase    		=					99
      __lowerCamelCase    		=					0
      __lowerCamelCase    		=					32
      __lowerCamelCase    		=					2
      __lowerCamelCase    		=					4
      __lowerCamelCase    		=					0.1
      __lowerCamelCase    		=					0.1
      __lowerCamelCase    		=					512
      __lowerCamelCase    		=					16
      __lowerCamelCase    		=					2
      __lowerCamelCase    		=					0.02
      __lowerCamelCase    		=					3
      __lowerCamelCase    		=					4
      __lowerCamelCase    		=					'last'
      __lowerCamelCase    		=					True
      __lowerCamelCase    		=					None
      __lowerCamelCase    		=					0







   def    _A				(	self	):


      '''simple docstring'''

      __lowerCamelCase    		=					ids_tensor([self.batch_size, self.seq_length] ,				self.vocab_size	)
      __lowerCamelCase    		=					random_attention_mask([self.batch_size, self.seq_length] ,				dtype=tf.floataa	)

      __lowerCamelCase    		=					None
      if self.use_input_lengths:
         __lowerCamelCase    		=					(
             ids_tensor([self.batch_size] ,				vocab_size=2	) + self.seq_length - 2
         )  # small variation of seq_length

      __lowerCamelCase    		=					None
      if self.use_token_type_ids:
         __lowerCamelCase    		=					ids_tensor([self.batch_size, self.seq_length] ,				self.n_langs	)

      __lowerCamelCase    		=					None
      __lowerCamelCase    		=					None
      __lowerCamelCase    		=					None
      if self.use_labels:
         __lowerCamelCase    		=					ids_tensor([self.batch_size] ,				self.type_sequence_label_size	)
         __lowerCamelCase    		=					ids_tensor([self.batch_size, self.seq_length] ,				self.num_labels	)
         __lowerCamelCase    		=					ids_tensor([self.batch_size] ,				2 ,				dtype=tf.floataa	)
         __lowerCamelCase    		=					ids_tensor([self.batch_size] ,				self.num_choices	)

      __lowerCamelCase    		=					XLMConfig(
          vocab_size=self.vocab_size ,				n_special=self.n_special ,				emb_dim=self.hidden_size ,				n_layers=self.num_hidden_layers ,				n_heads=self.num_attention_heads ,				dropout=self.hidden_dropout_prob ,				attention_dropout=self.attention_probs_dropout_prob ,				gelu_activation=self.gelu_activation ,				sinusoidal_embeddings=self.sinusoidal_embeddings ,				asm=self.asm ,				causal=self.causal ,				n_langs=self.n_langs ,				max_position_embeddings=self.max_position_embeddings ,				initializer_range=self.initializer_range ,				summary_type=self.summary_type ,				use_proj=self.use_proj ,				bos_token_id=self.bos_token_id ,				)

      return (
          config,
          input_ids,
          token_type_ids,
          input_lengths,
          sequence_labels,
          token_labels,
          is_impossible_labels,
          choice_labels,
          input_mask,
      )







   def    _A				(	self ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				):


      '''simple docstring'''

      __lowerCamelCase    		=					TFXLMModel(config=_SCREAMING_SNAKE_CASE	)
      __lowerCamelCase    		=					{'input_ids': input_ids, 'lengths': input_lengths, 'langs': token_type_ids}
      __lowerCamelCase    		=					model(_SCREAMING_SNAKE_CASE	)

      __lowerCamelCase    		=					[input_ids, input_mask]
      __lowerCamelCase    		=					model(_SCREAMING_SNAKE_CASE	)
      self.parent.assertEqual(result.last_hidden_state.shape ,				(self.batch_size, self.seq_length, self.hidden_size)	)







   def    _A				(	self ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				):


      '''simple docstring'''

      __lowerCamelCase    		=					TFXLMWithLMHeadModel(_SCREAMING_SNAKE_CASE	)

      __lowerCamelCase    		=					{'input_ids': input_ids, 'lengths': input_lengths, 'langs': token_type_ids}
      __lowerCamelCase    		=					model(_SCREAMING_SNAKE_CASE	)

      __lowerCamelCase    		=					outputs

      self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.seq_length, self.vocab_size)	)







   def    _A				(	self ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				):


      '''simple docstring'''

      __lowerCamelCase    		=					TFXLMForQuestionAnsweringSimple(_SCREAMING_SNAKE_CASE	)

      __lowerCamelCase    		=					{'input_ids': input_ids, 'lengths': input_lengths}

      __lowerCamelCase    		=					model(_SCREAMING_SNAKE_CASE	)

      self.parent.assertEqual(result.start_logits.shape ,				(self.batch_size, self.seq_length)	)
      self.parent.assertEqual(result.end_logits.shape ,				(self.batch_size, self.seq_length)	)







   def    _A				(	self ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				):


      '''simple docstring'''

      __lowerCamelCase    		=					TFXLMForSequenceClassification(_SCREAMING_SNAKE_CASE	)

      __lowerCamelCase    		=					{'input_ids': input_ids, 'lengths': input_lengths}

      __lowerCamelCase    		=					model(_SCREAMING_SNAKE_CASE	)

      self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.type_sequence_label_size)	)







   def    _A				(	self ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				):


      '''simple docstring'''

      __lowerCamelCase    		=					self.num_labels
      __lowerCamelCase    		=					TFXLMForTokenClassification(config=_SCREAMING_SNAKE_CASE	)
      __lowerCamelCase    		=					{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
      __lowerCamelCase    		=					model(_SCREAMING_SNAKE_CASE	)
      self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.seq_length, self.num_labels)	)







   def    _A				(	self ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				):


      '''simple docstring'''

      __lowerCamelCase    		=					self.num_choices
      __lowerCamelCase    		=					TFXLMForMultipleChoice(config=_SCREAMING_SNAKE_CASE	)
      __lowerCamelCase    		=					tf.tile(tf.expand_dims(_SCREAMING_SNAKE_CASE ,				1	) ,				(1, self.num_choices, 1)	)
      __lowerCamelCase    		=					tf.tile(tf.expand_dims(_SCREAMING_SNAKE_CASE ,				1	) ,				(1, self.num_choices, 1)	)
      __lowerCamelCase    		=					tf.tile(tf.expand_dims(_SCREAMING_SNAKE_CASE ,				1	) ,				(1, self.num_choices, 1)	)
      __lowerCamelCase    		=					{
          'input_ids': multiple_choice_inputs_ids,
          'attention_mask': multiple_choice_input_mask,
          'token_type_ids': multiple_choice_token_type_ids,
      }
      __lowerCamelCase    		=					model(_SCREAMING_SNAKE_CASE	)
      self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.num_choices)	)







   def    _A				(	self	):


      '''simple docstring'''

      __lowerCamelCase    		=					self.prepare_config_and_inputs()
      (
          (
          __lowerCamelCase
      )		,(
          __lowerCamelCase
      )		,(
          __lowerCamelCase
      )		,(
          __lowerCamelCase
      )		,(
          __lowerCamelCase
      )		,(
          __lowerCamelCase
      )		,(
          __lowerCamelCase
      )		,(
          __lowerCamelCase
      )		,(
          __lowerCamelCase
      )		,
      )    		=					config_and_inputs
      __lowerCamelCase    		=					{
          'input_ids': input_ids,
          'token_type_ids': token_type_ids,
          'langs': token_type_ids,
          'lengths': input_lengths,
      }
      return config, inputs_dict






@require_tf
class      _SCREAMING_SNAKE_CASE							(			UpperCAmelCase    ,				UpperCAmelCase    ,				unittest.TestCase				):



   '''simple docstring'''



   __UpperCAmelCase	:							List[Any]    =    (
       (
           TFXLMModel,
           TFXLMWithLMHeadModel,
           TFXLMForSequenceClassification,
           TFXLMForQuestionAnsweringSimple,
           TFXLMForTokenClassification,
           TFXLMForMultipleChoice,
       )
       if is_tf_available()
       else ()
   )
   __UpperCAmelCase	:							List[str]    =    (
       (TFXLMWithLMHeadModel,) if is_tf_available() else ()
   )  # TODO (PVP): Check other models whether language generation is also applicable
   __UpperCAmelCase	:							List[str]    =    (
       {
           "feature-extraction": TFXLMModel,
           "fill-mask": TFXLMWithLMHeadModel,
           "question-answering": TFXLMForQuestionAnsweringSimple,
           "text-classification": TFXLMForSequenceClassification,
           "text-generation": TFXLMWithLMHeadModel,
           "token-classification": TFXLMForTokenClassification,
           "zero-shot": TFXLMForSequenceClassification,
       }
       if is_tf_available()
       else {}
   )
   __UpperCAmelCase	:							Union[str, Any]    =    False
   __UpperCAmelCase	:							Optional[Any]    =    False







   def    _A				(	self ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE	):


      '''simple docstring'''

      if (
          pipeline_test_casse_name == "QAPipelineTests"
          and tokenizer_name is not None
          and not tokenizer_name.endswith('Fast'	)
      ):
         # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
         # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
         # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
         return True

      return False







   def    _A				(	self	):


      '''simple docstring'''

      __lowerCamelCase    		=					TFXLMModelTester(self	)
      __lowerCamelCase    		=					ConfigTester(self ,				config_class=_SCREAMING_SNAKE_CASE ,				emb_dim=37	)







   def    _A				(	self	):


      '''simple docstring'''

      self.config_tester.run_common_tests()







   def    _A				(	self	):


      '''simple docstring'''

      __lowerCamelCase    		=					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_xlm_model(*_SCREAMING_SNAKE_CASE	)







   def    _A				(	self	):


      '''simple docstring'''

      __lowerCamelCase    		=					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_xlm_lm_head(*_SCREAMING_SNAKE_CASE	)







   def    _A				(	self	):


      '''simple docstring'''

      __lowerCamelCase    		=					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_xlm_qa(*_SCREAMING_SNAKE_CASE	)







   def    _A				(	self	):


      '''simple docstring'''

      __lowerCamelCase    		=					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_xlm_sequence_classif(*_SCREAMING_SNAKE_CASE	)







   def    _A				(	self	):


      '''simple docstring'''

      __lowerCamelCase    		=					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_xlm_for_token_classification(*_SCREAMING_SNAKE_CASE	)







   def    _A				(	self	):


      '''simple docstring'''

      __lowerCamelCase    		=					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_xlm_for_multiple_choice(*_SCREAMING_SNAKE_CASE	)







   @slow
   def    _A				(	self	):


      '''simple docstring'''

      for model_name in TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         __lowerCamelCase    		=					TFXLMModel.from_pretrained(_SCREAMING_SNAKE_CASE	)
         self.assertIsNotNone(_SCREAMING_SNAKE_CASE	)






@require_tf
class      _SCREAMING_SNAKE_CASE							(			unittest.TestCase				):



   '''simple docstring'''



   @slow
   def    _A				(	self	):


      '''simple docstring'''

      __lowerCamelCase    		=					TFXLMWithLMHeadModel.from_pretrained('xlm-mlm-en-2048'	)
      __lowerCamelCase    		=					tf.convert_to_tensor([[14, 447]] ,				dtype=tf.intaa	)  # the president
      __lowerCamelCase    		=					[
          14,
          447,
          14,
          447,
          14,
          447,
          14,
          447,
          14,
          447,
          14,
          447,
          14,
          447,
          14,
          447,
          14,
          447,
          14,
          447,
      ]  # the president the president the president the president the president the president the president the president the president the president
      # TODO(PVP): this and other input_ids I tried for generation give pretty bad results. Not sure why. Model might just not be made for auto-regressive inference
      __lowerCamelCase    		=					model.generate(_SCREAMING_SNAKE_CASE ,				do_sample=_SCREAMING_SNAKE_CASE	)
      self.assertListEqual(output_ids[0].numpy().tolist() ,				_SCREAMING_SNAKE_CASE	)










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






def 			__lowerCAmelCase       (SCREAMING_SNAKE_CASE   )->    Tuple:




      """simple docstring"""

      snake_case_			= ConvNextConfig()

      if "tiny" in checkpoint_url:
            snake_case_			= [3, 3, 9, 3]
            snake_case_			= [96, 192, 384, 768]
      if "small" in checkpoint_url:
            snake_case_			= [3, 3, 27, 3]
            snake_case_			= [96, 192, 384, 768]
      if "base" in checkpoint_url:
            snake_case_			= [3, 3, 27, 3]
            snake_case_			= [128, 256, 512, 1024]
      if "large" in checkpoint_url:
            snake_case_			= [3, 3, 27, 3]
            snake_case_			= [192, 384, 768, 1536]
      if "xlarge" in checkpoint_url:
            snake_case_			= [3, 3, 27, 3]
            snake_case_			= [256, 512, 1024, 2048]

      if "1k" in checkpoint_url:
            snake_case_			= 1000
            snake_case_			= '''imagenet-1k-id2label.json'''
            snake_case_			= (1, 1000)
      else:
            snake_case_			= 2_1841
            snake_case_			= '''imagenet-22k-id2label.json'''
            snake_case_			= (1, 2_1841)

      snake_case_			= '''huggingface/label-files'''
      snake_case_			= num_labels
      snake_case_			= json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   repo_type='''dataset'''   )   ,   '''r'''   )   )
      snake_case_			= {int(SCREAMING_SNAKE_CASE   ): 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[9205]
            del idalabel[1_5027]
      snake_case_			= idalabel
      snake_case_			= {v: k for k, v in idalabel.items()}
      snake_case_			= hidden_sizes
      snake_case_			= depths

      return config, expected_shape






def 			__lowerCAmelCase       (SCREAMING_SNAKE_CASE   )->    Tuple:




      """simple docstring"""

      if "downsample_layers.0.0" in name:
            snake_case_			= name.replace('''downsample_layers.0.0'''   ,   '''embeddings.patch_embeddings'''   )
      if "downsample_layers.0.1" in name:
            snake_case_			= name.replace('''downsample_layers.0.1'''   ,   '''embeddings.norm'''   )  # we rename to layernorm later on
      if "downsample_layers.1.0" in name:
            snake_case_			= name.replace('''downsample_layers.1.0'''   ,   '''stages.1.downsampling_layer.0'''   )
      if "downsample_layers.1.1" in name:
            snake_case_			= name.replace('''downsample_layers.1.1'''   ,   '''stages.1.downsampling_layer.1'''   )
      if "downsample_layers.2.0" in name:
            snake_case_			= name.replace('''downsample_layers.2.0'''   ,   '''stages.2.downsampling_layer.0'''   )
      if "downsample_layers.2.1" in name:
            snake_case_			= name.replace('''downsample_layers.2.1'''   ,   '''stages.2.downsampling_layer.1'''   )
      if "downsample_layers.3.0" in name:
            snake_case_			= name.replace('''downsample_layers.3.0'''   ,   '''stages.3.downsampling_layer.0'''   )
      if "downsample_layers.3.1" in name:
            snake_case_			= 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.
            snake_case_			= name[: len('''stages.0'''   )] + '''.layers''' + name[len('''stages.0'''   ) :]
      if "stages" in name:
            snake_case_			= name.replace('''stages'''   ,   '''encoder.stages'''   )
      if "norm" in name:
            snake_case_			= name.replace('''norm'''   ,   '''layernorm'''   )
      if "gamma" in name:
            snake_case_			= name.replace('''gamma'''   ,   '''layer_scale_parameter'''   )
      if "head" in name:
            snake_case_			= name.replace('''head'''   ,   '''classifier'''   )

      return name






def 			__lowerCAmelCase       ()->    Optional[Any]:




      """simple docstring"""

      snake_case_			= '''http://images.cocodataset.org/val2017/000000039769.jpg'''
      snake_case_			= 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   )->    str:




      """simple docstring"""

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

      # load HuggingFace model
      snake_case_			= ConvNextForImageClassification(SCREAMING_SNAKE_CASE   )
      model.load_state_dict(SCREAMING_SNAKE_CASE   )
      model.eval()

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

      snake_case_			= model(SCREAMING_SNAKE_CASE   ).logits

      # note: the logits below were obtained without center cropping
      if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth":
            snake_case_			= torch.tensor([-0.1_2_1_0, -0.6_6_0_5, 0.1_9_1_8]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
            snake_case_			= torch.tensor([-0.4_4_7_3, -0.1_8_4_7, -0.6_3_6_5]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
            snake_case_			= torch.tensor([0.4_5_2_5, 0.7_5_3_9, 0.0_3_0_8]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
            snake_case_			= torch.tensor([0.3_5_6_1, 0.6_3_5_0, -0.0_3_8_4]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
            snake_case_			= torch.tensor([0.4_1_7_4, -0.0_9_8_9, 0.1_4_8_9]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
            snake_case_			= torch.tensor([0.2_5_1_3, -0.1_3_4_9, -0.1_6_1_3]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
            snake_case_			= torch.tensor([1.2_9_8_0, 0.3_6_3_1, -0.1_1_9_8]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
            snake_case_			= torch.tensor([1.2_9_6_3, 0.1_2_2_7, 0.1_7_2_3]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
            snake_case_			= torch.tensor([1.7_9_5_6, 0.8_3_9_0, 0.2_8_2_0]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
            snake_case_			= torch.tensor([-0.2_8_2_2, -0.0_5_0_2, -0.0_8_7_8]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
            snake_case_			= torch.tensor([-0.5_6_7_2, -0.0_7_3_0, -0.4_3_4_8]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
            snake_case_			= torch.tensor([0.2_6_8_1, 0.2_3_6_5, 0.6_2_4_6]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
            snake_case_			= torch.tensor([-0.2_6_4_2, 0.3_9_3_1, 0.5_1_1_6]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
            snake_case_			= torch.tensor([-0.6_6_7_7, -0.1_8_7_3, -0.8_3_7_9]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
            snake_case_			= torch.tensor([-0.7_7_4_9, -0.2_9_6_7, -0.6_4_4_4]   )
      else:
            raise ValueError(f'''Unknown URL: {checkpoint_url}'''   )

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

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

      print('''Pushing model to the hub...'''   )
      snake_case_			= '''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(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )   ,   organization='''nielsr'''   ,   commit_message='''Add model'''   ,   )


if __name__ == "__main__":
   UpperCAmelCase     =       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.""",
   )

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