





from __future__ import annotations

import unittest
import warnings

from transformers import AutoTokenizer, MarianConfig, MarianTokenizer, TranslationPipeline, is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow
from transformers.utils import cached_property

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


if is_tf_available():
						import tensorflow as tf

						from transformers import TFAutoModelForSeqaSeqLM, TFMarianModel, TFMarianMTModel
@require_tf
class SCREAMING_SNAKE_CASE_    :
						__magic_name__:      Union[str, Any]			   =							MarianConfig
						__magic_name__:      List[Any]			   =							{}
						__magic_name__:      Any			   =							"gelu"






						def __init__(   self   :   List[Any]		, _A   :   Dict		, _A   :   Dict=13		, _A   :   Dict=7		, _A   :   Optional[Any]=True		, _A   :   Union[str, Any]=False		, _A   :   Union[str, Any]=99		, _A   :   Optional[Any]=32		, _A   :   str=2		, _A   :   Optional[Any]=4		, _A   :   Tuple=37		, _A   :   List[Any]=0.1		, _A   :   List[Any]=0.1		, _A   :   Tuple=20		, _A   :   Optional[Any]=2		, _A   :   Optional[int]=1		, _A   :   Tuple=0		, )				->     List[Any]:







							"""simple docstring"""


							snake_case_		:      Optional[int]           =						parent
							snake_case_		:      Tuple           =						batch_size
							snake_case_		:      int           =						seq_length
							snake_case_		:      Dict           =						is_training
							snake_case_		:      Any           =						use_labels
							snake_case_		:      List[Any]           =						vocab_size
							snake_case_		:      List[Any]           =						hidden_size
							snake_case_		:      Optional[Any]           =						num_hidden_layers
							snake_case_		:      Dict           =						num_attention_heads
							snake_case_		:      str           =						intermediate_size

							snake_case_		:      Optional[int]           =						hidden_dropout_prob
							snake_case_		:      List[str]           =						attention_probs_dropout_prob
							snake_case_		:      List[str]           =						max_position_embeddings
							snake_case_		:      Dict           =						eos_token_id
							snake_case_		:      Optional[int]           =						pad_token_id
							snake_case_		:      List[str]           =						bos_token_id






						def UpperCAmelCase_					(   self   :   Dict						)				->     Dict:







							"""simple docstring"""


							snake_case_		:      int           =						ids_tensor([self.batch_size, self.seq_length - 1]		, self.vocab_size						)
							snake_case_		:      List[str]           =						tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size						)		, 1						)
							snake_case_		:      Dict           =						tf.concat([input_ids, eos_tensor]		, axis=1						)

							snake_case_		:      Optional[int]           =						ids_tensor([self.batch_size, self.seq_length]		, self.vocab_size						)

							snake_case_		:      Tuple           =						self.config_cls(
							    vocab_size=self.vocab_size		, d_model=self.hidden_size		, encoder_layers=self.num_hidden_layers		, decoder_layers=self.num_hidden_layers		, encoder_attention_heads=self.num_attention_heads		, decoder_attention_heads=self.num_attention_heads		, encoder_ffn_dim=self.intermediate_size		, decoder_ffn_dim=self.intermediate_size		, dropout=self.hidden_dropout_prob		, attention_dropout=self.attention_probs_dropout_prob		, max_position_embeddings=self.max_position_embeddings		, eos_token_ids=[2]		, bos_token_id=self.bos_token_id		, pad_token_id=self.pad_token_id		, decoder_start_token_id=self.pad_token_id		, **self.config_updates		, )
							snake_case_		:      Optional[Any]           =						prepare_marian_inputs_dict(_A		, _A		, _A						)
							return config, inputs_dict






						def UpperCAmelCase_					(   self   :   Dict		, _A   :   int		, _A   :   Optional[Any]						)				->     Optional[Any]:







							"""simple docstring"""


							snake_case_		:      Dict           =						TFMarianModel(config=_A						).get_decoder()
							snake_case_		:      str           =						inputs_dict['input_ids']

							snake_case_		:      List[str]           =						input_ids[:1, :]
							snake_case_		:      int           =						inputs_dict['attention_mask'][:1, :]
							snake_case_		:      str           =						inputs_dict['head_mask']
							snake_case_		:      Optional[int]           =						1

							# first forward pass
							snake_case_		:      Optional[int]           =						model(_A		, attention_mask=_A		, head_mask=_A		, use_cache=_A						)

							snake_case_    ,snake_case_		:      str           =						outputs.to_tuple()

							# create hypothetical next token and extent to next_input_ids
							snake_case_		:      Union[str, Any]           =						ids_tensor((self.batch_size, 3)		, config.vocab_size						)
							snake_case_		:      int           =						tf.cast(ids_tensor((self.batch_size, 3)		, 2						)		, tf.inta						)

							# append to next input_ids and
							snake_case_		:      Optional[Any]           =						tf.concat([input_ids, next_tokens]		, axis=-1						)
							snake_case_		:      Optional[int]           =						tf.concat([attention_mask, next_attn_mask]		, axis=-1						)

							snake_case_		:      Optional[Any]           =						model(_A		, attention_mask=_A						)[0]
							snake_case_		:      Union[str, Any]           =						model(_A		, attention_mask=_A		, past_key_values=_A						)[0]

							self.parent.assertEqual(next_tokens.shape[1]		, output_from_past.shape[1]						)

							# select random slice
							snake_case_		:      str           =						int(ids_tensor((1,)		, output_from_past.shape[-1]						)						)
							snake_case_		:      Optional[Any]           =						output_from_no_past[:, -3:, random_slice_idx]
							snake_case_		:      Union[str, Any]           =						output_from_past[:, :, random_slice_idx]

							# test that outputs are equal for slice
							tf.debugging.assert_near(_A		, _A		, rtol=1E-3						)


def 				SCREAMING_SNAKE_CASE__    (     __a ,			__a ,			__a ,			__a=None ,			__a=None ,			__a=None ,			__a=None ,			__a=None ,			):
	if attention_mask is None:
		snake_case_		:      Optional[Any]           =						tf.cast(tf.math.not_equal(__a ,			config.pad_token_id ) ,			tf.inta )
	if decoder_attention_mask is None:
		snake_case_		:      Dict           =						tf.concat(
		    [
		        tf.ones(decoder_input_ids[:, :1].shape ,			dtype=tf.inta ),
		        tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:] ,			config.pad_token_id ) ,			tf.inta ),
		    ] ,			axis=-1 ,			)
	if head_mask is None:
		snake_case_		:      Dict           =						tf.ones((config.encoder_layers, config.encoder_attention_heads) )
	if decoder_head_mask is None:
		snake_case_		:      Optional[int]           =						tf.ones((config.decoder_layers, config.decoder_attention_heads) )
	if cross_attn_head_mask is None:
		snake_case_		:      Tuple           =						tf.ones((config.decoder_layers, config.decoder_attention_heads) )
	return {
	    "input_ids": input_ids,
	    "decoder_input_ids": decoder_input_ids,
	    "attention_mask": attention_mask,
	    "decoder_attention_mask": decoder_attention_mask,
	    "head_mask": head_mask,
	    "decoder_head_mask": decoder_head_mask,
	    "cross_attn_head_mask": cross_attn_head_mask,
	}
@require_tf
class SCREAMING_SNAKE_CASE_    (       snake_case_							,					snake_case_							,					unittest.TestCase       ):
						__magic_name__:      Dict			   =							(TFMarianMTModel, TFMarianModel) if is_tf_available() else ()
						__magic_name__:      Union[str, Any]			   =							(TFMarianMTModel,) if is_tf_available() else ()
						__magic_name__:      int			   =							(
						    {
						        "conversational": TFMarianMTModel,
						        "feature-extraction": TFMarianModel,
						        "summarization": TFMarianMTModel,
						        "text2text-generation": TFMarianMTModel,
						        "translation": TFMarianMTModel,
						    }
						    if is_tf_available()
						    else {}
						)
						__magic_name__:      List[str]			   =							True
						__magic_name__:      Union[str, Any]			   =							False
						__magic_name__:      int			   =							False






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







							"""simple docstring"""


							snake_case_		:      Optional[int]           =						TFMarianModelTester(self						)
							snake_case_		:      Optional[Any]           =						ConfigTester(self		, config_class=_A						)






						def UpperCAmelCase_					(   self   :   List[Any]						)				->     List[str]:







							"""simple docstring"""


							self.config_tester.run_common_tests()






						def UpperCAmelCase_					(   self   :   Optional[Any]						)				->     Optional[int]:







							"""simple docstring"""


							snake_case_		:      int           =						self.model_tester.prepare_config_and_inputs_for_common()
							self.model_tester.check_decoder_model_past_large_inputs(*_A						)
@require_tf
class SCREAMING_SNAKE_CASE_    (       unittest.TestCase       ):
						__magic_name__:      str			   =							1000  # show more chars for failing integration tests






						@classmethod
						def UpperCAmelCase_					(   cls   :   List[str]						)				->     None:







							"""simple docstring"""


							snake_case_		:      List[str]           =						F"""Helsinki-NLP/opus-mt-{cls.src}-{cls.tgt}"""
							return cls






						@cached_property
						def UpperCAmelCase_					(   self   :   List[Any]						)				->     MarianTokenizer:







							"""simple docstring"""


							return AutoTokenizer.from_pretrained(self.model_name						)






						@property
						def UpperCAmelCase_					(   self   :   List[Any]						)				->     int:







							"""simple docstring"""


							return self.tokenizer.eos_token_id






						@cached_property
						def UpperCAmelCase_					(   self   :   str						)				->     Optional[Any]:







							"""simple docstring"""


							warnings.simplefilter('error'						)
							snake_case_		:      TFMarianMTModel           =						TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name						)
							assert isinstance(_A		, _A						)
							snake_case_		:      List[str]           =						model.config
							self.assertListEqual(c.bad_words_ids		, [[c.pad_token_id]]						)
							self.assertEqual(c.max_length		, 512						)
							self.assertEqual(c.decoder_start_token_id		, c.pad_token_id						)
							return model






						def UpperCAmelCase_					(   self   :   Any		, **_A   :   Tuple						)				->     List[Any]:







							"""simple docstring"""


							snake_case_		:      Optional[Any]           =						self.translate_src_text(**_A						)
							self.assertListEqual(self.expected_text		, _A						)






						def UpperCAmelCase_					(   self   :   List[Any]		, **_A   :   int						)				->     Any:







							"""simple docstring"""


							snake_case_		:      List[str]           =						self.tokenizer(self.src_text		, **_A		, padding=_A		, return_tensors='tf'						)
							snake_case_		:      Optional[int]           =						self.model.generate(
							    model_inputs.input_ids		, attention_mask=model_inputs.attention_mask		, num_beams=2		, max_length=128						)
							snake_case_		:      Dict           =						self.tokenizer.batch_decode(generated_ids.numpy()		, skip_special_tokens=_A						)
							return generated_words
@require_sentencepiece
@require_tokenizers
@require_tf
class SCREAMING_SNAKE_CASE_    (       snake_case_       ):
						__magic_name__:      str			   =							"mt"
						__magic_name__:      Dict			   =							"en"
						__magic_name__:      Optional[Any]			   =							["Billi messu b'mod ġentili, Ġesù fejjaq raġel li kien milqut bil - marda kerha tal - ġdiem."]
						__magic_name__:      Optional[int]			   =							["Touching gently, Jesus healed a man who was affected by the sad disease of leprosy."]






						@unittest.skip('Skipping until #12647 is resolved.'						)
						@slow
						def UpperCAmelCase_					(   self   :   int						)				->     List[Any]:







							"""simple docstring"""


							self._assert_generated_batch_equal_expected()
@require_sentencepiece
@require_tokenizers
@require_tf
class SCREAMING_SNAKE_CASE_    (       snake_case_       ):
						__magic_name__:      List[str]			   =							"en"
						__magic_name__:      int			   =							"zh"
						__magic_name__:      List[Any]			   =							["My name is Wolfgang and I live in Berlin"]
						__magic_name__:      int			   =							["我叫沃尔夫冈 我住在柏林"]






						@unittest.skip('Skipping until #12647 is resolved.'						)
						@slow
						def UpperCAmelCase_					(   self   :   List[Any]						)				->     Dict:







							"""simple docstring"""


							self._assert_generated_batch_equal_expected()
@require_sentencepiece
@require_tokenizers
@require_tf
class SCREAMING_SNAKE_CASE_    (       snake_case_       ):
						__magic_name__:      Any			   =							"en"
						__magic_name__:      Any			   =							"ROMANCE"
						__magic_name__:      List[Any]			   =							[
						    ">>fr<< Don't spend so much time watching TV.",
						    ">>pt<< Your message has been sent.",
						    ">>es<< He's two years older than me.",
						]
						__magic_name__:      List[Any]			   =							[
						    "Ne passez pas autant de temps à regarder la télé.",
						    "A sua mensagem foi enviada.",
						    "Es dos años más viejo que yo.",
						]






						@unittest.skip('Skipping until #12647 is resolved.'						)
						@slow
						def UpperCAmelCase_					(   self   :   Optional[Any]						)				->     List[str]:







							"""simple docstring"""


							self._assert_generated_batch_equal_expected()






						@unittest.skip('Skipping until #12647 is resolved.'						)
						@slow
						def UpperCAmelCase_					(   self   :   Union[str, Any]						)				->     int:







							"""simple docstring"""


							snake_case_		:      Dict           =						TranslationPipeline(self.model		, self.tokenizer		, framework='tf'						)
							snake_case_		:      List[str]           =						pipeline(self.src_text						)
							self.assertEqual(self.expected_text		, [x['translation_text'] for x in output]						)








'''simple docstring'''







import argparse
import json
import os

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

from transformers import ConvNextImageProcessor, ConvNextVaConfig, ConvNextVaForImageClassification
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


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





def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      Optional[int] )			->						Dict:
						UpperCAmelCase_      :   int      =	ConvNextVaConfig()

						if "atto" in checkpoint_url:
												UpperCAmelCase_      :   Tuple      =	[2, 2, 6, 2]
												UpperCAmelCase_      :   Tuple      =	[40, 80, 160, 320]
						if "femto" in checkpoint_url:
												UpperCAmelCase_      :   List[Any]      =	[2, 2, 6, 2]
												UpperCAmelCase_      :   Any      =	[48, 96, 192, 384]
						if "pico" in checkpoint_url:
												UpperCAmelCase_      :   Union[str, Any]      =	[2, 2, 6, 2]
												UpperCAmelCase_      :   int      =	[64, 128, 256, 512]
						if "nano" in checkpoint_url:
												UpperCAmelCase_      :   str      =	[2, 2, 8, 2]
												UpperCAmelCase_      :   List[str]      =	[80, 160, 320, 640]
						if "tiny" in checkpoint_url:
												UpperCAmelCase_      :   Optional[Any]      =	[3, 3, 9, 3]
												UpperCAmelCase_      :   Any      =	[96, 192, 384, 768]
						if "base" in checkpoint_url:
												UpperCAmelCase_      :   Union[str, Any]      =	[3, 3, 27, 3]
												UpperCAmelCase_      :   Optional[int]      =	[128, 256, 512, 1024]
						if "large" in checkpoint_url:
												UpperCAmelCase_      :   int      =	[3, 3, 27, 3]
												UpperCAmelCase_      :   List[Any]      =	[192, 384, 768, 1536]
						if "huge" in checkpoint_url:
												UpperCAmelCase_      :   Tuple      =	[3, 3, 27, 3]
												UpperCAmelCase_      :   Dict      =	[352, 704, 1408, 2816]

						UpperCAmelCase_      :   Union[str, Any]      =	1000
						UpperCAmelCase_      :   Optional[Any]      =	'''imagenet-1k-id2label.json'''
						UpperCAmelCase_      :   int      =	(1, 1000)

						UpperCAmelCase_      :   Any      =	'''huggingface/label-files'''
						UpperCAmelCase_      :   List[str]      =	num_labels
						UpperCAmelCase_      :   Optional[int]      =	json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__,      repo_type='''dataset''' ),      '''r''' ) )
						UpperCAmelCase_      :   Tuple      =	{int(SCREAMING_SNAKE_CASE__ ): v for k, v in idalabel.items()}

						UpperCAmelCase_      :   Union[str, Any]      =	idalabel
						UpperCAmelCase_      :   List[Any]      =	{v: k for k, v in idalabel.items()}
						UpperCAmelCase_      :   str      =	hidden_sizes
						UpperCAmelCase_      :   Dict      =	depths

						return config, expected_shape





def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      Dict )			->						Optional[Any]:
						if "downsample_layers.0.0" in name:
												UpperCAmelCase_      :   str      =	name.replace('''downsample_layers.0.0''',      '''embeddings.patch_embeddings''' )
						if "downsample_layers.0.1" in name:
												UpperCAmelCase_      :   Optional[Any]      =	name.replace('''downsample_layers.0.1''',      '''embeddings.norm''' )  # we rename to layernorm later on
						if "downsample_layers.1.0" in name:
												UpperCAmelCase_      :   Dict      =	name.replace('''downsample_layers.1.0''',      '''stages.1.downsampling_layer.0''' )
						if "downsample_layers.1.1" in name:
												UpperCAmelCase_      :   Any      =	name.replace('''downsample_layers.1.1''',      '''stages.1.downsampling_layer.1''' )
						if "downsample_layers.2.0" in name:
												UpperCAmelCase_      :   Optional[int]      =	name.replace('''downsample_layers.2.0''',      '''stages.2.downsampling_layer.0''' )
						if "downsample_layers.2.1" in name:
												UpperCAmelCase_      :   List[Any]      =	name.replace('''downsample_layers.2.1''',      '''stages.2.downsampling_layer.1''' )
						if "downsample_layers.3.0" in name:
												UpperCAmelCase_      :   Optional[int]      =	name.replace('''downsample_layers.3.0''',      '''stages.3.downsampling_layer.0''' )
						if "downsample_layers.3.1" in name:
												UpperCAmelCase_      :   Optional[int]      =	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.
												UpperCAmelCase_      :   Any      =	name[: len('''stages.0''' )] + '''.layers''' + name[len('''stages.0''' ) :]
						if "gamma" in name:
												UpperCAmelCase_      :   Optional[Any]      =	name.replace('''gamma''',      '''weight''' )
						if "beta" in name:
												UpperCAmelCase_      :   Dict      =	name.replace('''beta''',      '''bias''' )
						if "stages" in name:
												UpperCAmelCase_      :   str      =	name.replace('''stages''',      '''encoder.stages''' )
						if "norm" in name:
												UpperCAmelCase_      :   Union[str, Any]      =	name.replace('''norm''',      '''layernorm''' )
						if "head" in name:
												UpperCAmelCase_      :   Union[str, Any]      =	name.replace('''head''',      '''classifier''' )

						return name





def 							lowerCamelCase_   (	)			->						Union[str, Any]:
						UpperCAmelCase_      :   int      =	'''http://images.cocodataset.org/val2017/000000039769.jpg'''
						UpperCAmelCase_      :   Optional[int]      =	Image.open(requests.get(SCREAMING_SNAKE_CASE__,      stream=SCREAMING_SNAKE_CASE__ ).raw )
						return im





def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      Optional[int] )			->						int:
						if "224" in checkpoint_url:
												UpperCAmelCase_      :   Optional[Any]      =	224
												UpperCAmelCase_      :   Any      =	224 / 256
						elif "384" in checkpoint_url:
												UpperCAmelCase_      :   int      =	384
												UpperCAmelCase_      :   Optional[Any]      =	None
						else:
												UpperCAmelCase_      :   List[Any]      =	512
												UpperCAmelCase_      :   Union[str, Any]      =	None

						return ConvNextImageProcessor(
						    size=SCREAMING_SNAKE_CASE__,      crop_pct=SCREAMING_SNAKE_CASE__,      image_mean=[0.4_85, 0.4_56, 0.4_06],      image_std=[0.2_29, 0.2_24, 0.2_25],      resample=PILImageResampling.BICUBIC,      )





@torch.no_grad()
def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      List[Any],      SCREAMING_SNAKE_CASE__   :      List[Any],      SCREAMING_SNAKE_CASE__   :      List[Any],      SCREAMING_SNAKE_CASE__   :      str )			->						Optional[Any]:
						print('''Downloading original model from checkpoint...''' )
						# define ConvNeXTV2 configuration based on URL
						UpperCAmelCase_    ,							UpperCAmelCase_      :   Tuple      =	get_convnextva_config(SCREAMING_SNAKE_CASE__ )
						# load original state_dict from URL
						UpperCAmelCase_      :   int      =	torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE__ )['''model''']

						print('''Converting model parameters...''' )
						# rename keys
						for key in state_dict.copy().keys():
												UpperCAmelCase_      :   Optional[Any]      =	state_dict.pop(SCREAMING_SNAKE_CASE__ )
												UpperCAmelCase_      :   List[str]      =	val
						# add prefix to all keys expect classifier head
						for key in state_dict.copy().keys():
												UpperCAmelCase_      :   List[str]      =	state_dict.pop(SCREAMING_SNAKE_CASE__ )
												if not key.startswith('''classifier''' ):
																		UpperCAmelCase_      :   Optional[Any]      =	'''convnextv2.''' + key
												UpperCAmelCase_      :   Union[str, Any]      =	val

						# load HuggingFace model
						UpperCAmelCase_      :   Optional[int]      =	ConvNextVaForImageClassification(SCREAMING_SNAKE_CASE__ )
						model.load_state_dict(SCREAMING_SNAKE_CASE__ )
						model.eval()

						# Check outputs on an image, prepared by ConvNextImageProcessor
						UpperCAmelCase_      :   Any      =	convert_preprocessor(SCREAMING_SNAKE_CASE__ )
						UpperCAmelCase_      :   Tuple      =	preprocessor(images=prepare_img(),      return_tensors='''pt''' )
						UpperCAmelCase_      :   Optional[Any]      =	model(**SCREAMING_SNAKE_CASE__ ).logits

						# note: the logits below were obtained without center cropping
						if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt":
												UpperCAmelCase_      :   Optional[Any]      =	torch.tensor([-0.39_30, 0.17_47, -0.52_46, 0.41_77, 0.42_95] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
												UpperCAmelCase_      :   List[Any]      =	torch.tensor([-0.17_27, -0.53_41, -0.78_18, -0.47_45, -0.65_66] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
												UpperCAmelCase_      :   Dict      =	torch.tensor([-0.03_33, 0.15_63, -0.91_37, 0.10_54, 0.03_81] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
												UpperCAmelCase_      :   Tuple      =	torch.tensor([-0.17_44, -0.15_55, -0.07_13, 0.09_50, -0.14_31] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
												UpperCAmelCase_      :   Optional[int]      =	torch.tensor([0.99_96, 0.19_66, -0.43_86, -0.34_72, 0.66_61] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
												UpperCAmelCase_      :   Optional[Any]      =	torch.tensor([-0.25_53, -0.67_08, -0.13_59, 0.25_18, -0.24_88] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
												UpperCAmelCase_      :   int      =	torch.tensor([-0.06_73, -0.56_27, -0.37_53, -0.27_22, 0.01_78] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
												UpperCAmelCase_      :   Optional[int]      =	torch.tensor([-0.63_77, -0.74_58, -0.21_50, 0.11_84, -0.05_97] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
												UpperCAmelCase_      :   List[Any]      =	torch.tensor([1.07_99, 0.23_22, -0.88_60, 1.02_19, 0.62_31] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
												UpperCAmelCase_      :   Optional[int]      =	torch.tensor([0.37_66, 0.49_17, -1.14_26, 0.99_42, 0.60_24] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
												UpperCAmelCase_      :   Dict      =	torch.tensor([0.42_20, -0.69_19, -0.43_17, -0.28_81, -0.66_09] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
												UpperCAmelCase_      :   Dict      =	torch.tensor([0.10_82, -0.82_86, -0.50_95, 0.46_81, -0.80_85] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
												UpperCAmelCase_      :   Optional[int]      =	torch.tensor([-0.24_19, -0.62_21, 0.21_76, -0.09_80, -0.75_27] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
												UpperCAmelCase_      :   Optional[int]      =	torch.tensor([0.03_91, -0.43_71, 0.37_86, 0.12_51, -0.27_84] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
												UpperCAmelCase_      :   Any      =	torch.tensor([-0.05_04, 0.56_36, -0.17_29, -0.65_07, -0.39_49] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
												UpperCAmelCase_      :   Optional[Any]      =	torch.tensor([0.35_60, 0.94_86, 0.31_49, -0.26_67, -0.51_38] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
												UpperCAmelCase_      :   List[Any]      =	torch.tensor([-0.24_69, -0.45_50, -0.58_53, -0.08_10, 0.03_09] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
												UpperCAmelCase_      :   Tuple      =	torch.tensor([-0.30_90, 0.08_02, -0.06_82, -0.19_79, -0.28_26] )
						else:
												raise ValueError(F"""Unknown URL: {checkpoint_url}""" )

						assert torch.allclose(logits[0, :5],      SCREAMING_SNAKE_CASE__,      atol=1E-3 )
						assert logits.shape == expected_shape
						print('''Model outputs match the original results!''' )

						if save_model:
												print('''Saving model to local...''' )
												# Create folder to save model
												if not os.path.isdir(SCREAMING_SNAKE_CASE__ ):
																		os.mkdir(SCREAMING_SNAKE_CASE__ )

												model.save_pretrained(SCREAMING_SNAKE_CASE__ )
												preprocessor.save_pretrained(SCREAMING_SNAKE_CASE__ )

						UpperCAmelCase_      :   Dict      =	'''convnextv2'''
						if "atto" in checkpoint_url:
												model_name += "-atto"
						if "femto" in checkpoint_url:
												model_name += "-femto"
						if "pico" in checkpoint_url:
												model_name += "-pico"
						if "nano" in checkpoint_url:
												model_name += "-nano"
						elif "tiny" in checkpoint_url:
												model_name += "-tiny"
						elif "base" in checkpoint_url:
												model_name += "-base"
						elif "large" in checkpoint_url:
												model_name += "-large"
						elif "huge" in checkpoint_url:
												model_name += "-huge"
						if "22k" in checkpoint_url and "1k" not in checkpoint_url:
												model_name += "-22k"
						elif "22k" in checkpoint_url and "1k" in checkpoint_url:
												model_name += "-22k-1k"
						elif "1k" in checkpoint_url:
												model_name += "-1k"
						if "224" in checkpoint_url:
												model_name += "-224"
						elif "384" in checkpoint_url:
												model_name += "-384"
						elif "512" in checkpoint_url:
												model_name += "-512"

						if push_to_hub:
												print(F"""Pushing {model_name} to the hub...""" )
												model.push_to_hub(SCREAMING_SNAKE_CASE__ )
												preprocessor.push_to_hub(SCREAMING_SNAKE_CASE__ )


if __name__ == "__main__":
					snake_case_  :							Tuple            =	argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    "--checkpoint_url",
					    default="https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt",
					    type=str,
					    help="URL of the original ConvNeXTV2 checkpoint you'd like to convert.",
					)
					parser.add_argument(
					    "--pytorch_dump_folder_path",
					    default="model",
					    type=str,
					    help="Path to the output PyTorch model directory.",
					)
					parser.add_argument("--save_model", action="store_true", help="Save model to local")
					parser.add_argument("--push_to_hub", action="store_true", help="Push model and image preprocessor to the hub")

					snake_case_  :							Optional[int]            =	parser.parse_args()
					convert_convnextva_checkpoint(
					    args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
					)





