
'''simple docstring'''
import json
import os
import pickle
import subprocess
from hashlib import mda
from pathlib import Path
from tempfile import gettempdir
from textwrap import dedent
from types import CodeType, FunctionType
from unittest import TestCase
from unittest.mock import patch

import pytest
from multiprocess import Pool
from packaging import version

import datasets
from datasets.fingerprint import Hasher, fingerprint_transform
from datasets.table import InMemoryTable

from .utils import (
    require_regex,
    require_spacy,
    require_spacy_model,
    require_tiktoken,
    require_torch,
    require_transformers,
)


class        SCREAMING_SNAKE_CASE__	:







				def __init__(    self		,							A_			)->						Optional[int]:







											'''simple docstring'''






											UpperCamelCase					       =			foo







				def __call__(    self			)->						Tuple:







											'''simple docstring'''






											return self.foo


class        SCREAMING_SNAKE_CASE__	(			datasets.Dataset):







				@fingerprint_transform(inplace=A_			)
				def 		UpperCAmelCase_  (    self		,							A_		,							*A_		,							**A_			)->						int:







											'''simple docstring'''






											return DatasetChild(self.data		,							fingerprint=A_			)







				@fingerprint_transform(inplace=A_			)
				def 		UpperCAmelCase_  (    self		,							A_		,							*A_		,							**A_			)->						List[Any]:







											'''simple docstring'''






											return DatasetChild(self.data		,							fingerprint=A_			)


class        SCREAMING_SNAKE_CASE__	:







				def __init__(    self		,							A_			)->						Optional[Any]:







											'''simple docstring'''






											UpperCamelCase					       =			callable







				def __call__(    self		,							*A_		,							**A_			)->						List[str]:







											'''simple docstring'''






											if self.callable is not None:
																		return self.callable(*A_		,							**A_			)







				def __getstate__(    self			)->						Optional[int]:







											'''simple docstring'''






											raise pickle.PicklingError()


class        SCREAMING_SNAKE_CASE__	(			snake_case_):







				@require_transformers
				@pytest.mark.integration
				def 		UpperCAmelCase_  (    self			)->						Any:







											'''simple docstring'''






											from transformers import AutoTokenizer

											def encode(A_			):
																		return tokenizer(A_			)

											# TODO: add hash consistency tests across sessions
											UpperCamelCase					       =			AutoTokenizer.from_pretrained('bert-base-uncased'			)
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(lambda A_			: tokenizer(A_			)			)			).hexdigest()
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
											UpperCamelCase					       =			AutoTokenizer.from_pretrained('bert-base-cased'			)
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(lambda A_			: tokenizer(A_			)			)			).hexdigest()
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
											UpperCamelCase					       =			AutoTokenizer.from_pretrained('bert-base-uncased'			)
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(lambda A_			: tokenizer(A_			)			)			).hexdigest()
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
											self.assertEqual(A_		,							A_			)
											self.assertNotEqual(A_		,							A_			)
											self.assertEqual(A_		,							A_			)
											self.assertNotEqual(A_		,							A_			)
											self.assertEqual(A_		,							A_			)
											self.assertNotEqual(A_		,							A_			)







				@require_transformers
				@pytest.mark.integration
				def 		UpperCAmelCase_  (    self			)->						Any:







											'''simple docstring'''






											from transformers import AutoTokenizer

											UpperCamelCase					       =			AutoTokenizer.from_pretrained('gpt2'			)
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
											tokenizer('Hello world !'			)  # call once to change the tokenizer's cache
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
											self.assertEqual(A_		,							A_			)







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







											'''simple docstring'''






											import regex

											UpperCamelCase					       =			regex.Regex('foo'			)
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
											UpperCamelCase					       =			regex.Regex('bar'			)
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
											UpperCamelCase					       =			regex.Regex('foo'			)
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
											self.assertEqual(A_		,							A_			)
											self.assertNotEqual(A_		,							A_			)


class        SCREAMING_SNAKE_CASE__	(			snake_case_):







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







											'''simple docstring'''






											def func():
																		return foo

											UpperCamelCase					       =			[0]
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
											UpperCamelCase					       =			[1]
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
											UpperCamelCase					       =			[0]
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
											self.assertEqual(A_		,							A_			)
											self.assertNotEqual(A_		,							A_			)







				def 		UpperCAmelCase_  (    self			)->						str:







											'''simple docstring'''






											def func():
																		pass

											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()

											def func():
																		pass

											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
											self.assertEqual(A_		,							A_			)







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







											'''simple docstring'''






											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(Foo([0]			)			)			).hexdigest()
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(Foo([1]			)			)			).hexdigest()
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(Foo([0]			)			)			).hexdigest()
											self.assertEqual(A_		,							A_			)
											self.assertNotEqual(A_		,							A_			)







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







											'''simple docstring'''






											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(Foo([0]			).__call__			)			).hexdigest()
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(Foo([1]			).__call__			)			).hexdigest()
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(Foo([0]			).__call__			)			).hexdigest()
											self.assertEqual(A_		,							A_			)
											self.assertNotEqual(A_		,							A_			)







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







											'''simple docstring'''






											UpperCamelCase					       =			(
											    'co_argcount',
											    'co_kwonlyargcount',
											    'co_nlocals',
											    'co_stacksize',
											    'co_flags',
											    'co_code',
											    'co_consts',
											    'co_names',
											    'co_varnames',
											    'co_filename',
											    'co_name',
											    'co_firstlineno',
											    'co_lnotab',
											    'co_freevars',
											    'co_cellvars',
											)

											def _create_code_pyaa(*A_			):
																		if hasattr(args[-3]		,							'encode'			):
																									UpperCamelCase					       =			list(A_			)
																									UpperCamelCase					       =			args[-3].encode()  # co_lnotab
																									UpperCamelCase					       =			args[-10].encode()  # co_code
																		if hasattr(A_		,							'co_posonlyargcount'			):
																									if len(A_			) == 16:
																																return CodeType(*A_			)
																									elif len(A_			) == 15:
																																return CodeType(args[0]		,							0		,							*args[1:]			)
																									return CodeType(args[0]		,							0		,							0		,							*args[1:]			)
																		elif hasattr(A_		,							'co_kwonlyargcount'			):
																									if len(A_			) == 16:
																																return CodeType(args[0]		,							*args[2:]			)
																									elif len(A_			) == 15:
																																return CodeType(*A_			)
																									return CodeType(args[0]		,							0		,							*args[1:]			)
																		if len(A_			) == 16:
																									return CodeType(args[0]		,							*args[3:]			)
																		elif len(A_			) == 15:
																									return CodeType(args[0]		,							*args[2:]			)
																		return CodeType(*A_			)

											def create_ipython_func(A_		,							A_			):
																		def func():
																									return returned_obj

																		UpperCamelCase					       =			func.__code__
																		# Use _create_code from dill in order to make it work for different python versions
																		if datasets.config.PY_VERSION < version.parse('3.8'			):
																									UpperCamelCase					       =			_create_code_pyaa(
																									    *[getattr(A_		,							A_			) if k != 'co_filename' else co_filename for k in code_args_pyaa]			)
																		else:
																									UpperCamelCase					       =			code.replace(co_filename=A_			)
																		return FunctionType(A_		,							func.__globals__		,							func.__name__		,							func.__defaults__		,							func.__closure__			)

											UpperCamelCase       ,     UpperCamelCase					       =			'<ipython-input-2-e0383a102aae>', [0]
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(create_ipython_func(A_		,							A_			)			)			).hexdigest()
											UpperCamelCase       ,     UpperCamelCase					       =			'<ipython-input-2-e0383a102aae>', [1]
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(create_ipython_func(A_		,							A_			)			)			).hexdigest()
											UpperCamelCase       ,     UpperCamelCase					       =			'<ipython-input-5-713f6613acf3>', [0]
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(create_ipython_func(A_		,							A_			)			)			).hexdigest()
											self.assertEqual(A_		,							A_			)
											self.assertNotEqual(A_		,							A_			)

											UpperCamelCase       ,     UpperCamelCase					       =			os.path.join(gettempdir()		,							'ipykernel_12345'		,							'321456789.py'			), [0]
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(create_ipython_func(A_		,							A_			)			)			).hexdigest()
											UpperCamelCase       ,     UpperCamelCase					       =			os.path.join(gettempdir()		,							'ipykernel_12345'		,							'321456789.py'			), [1]
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(create_ipython_func(A_		,							A_			)			)			).hexdigest()
											UpperCamelCase       ,     UpperCamelCase					       =			os.path.join(gettempdir()		,							'ipykernel_12345'		,							'654123987.py'			), [0]
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(create_ipython_func(A_		,							A_			)			)			).hexdigest()
											self.assertEqual(A_		,							A_			)
											self.assertNotEqual(A_		,							A_			)







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







											'''simple docstring'''






											UpperCamelCase       ,     UpperCamelCase					       =			[0], [1]

											def func():
																		return foo, bar

											UpperCamelCase					       =			'__main__'

											def globalvars_mocka_side_effect(A_		,							*A_		,							**A_			):
																		return {"foo": foo, "bar": bar}

											def globalvars_mocka_side_effect(A_		,							*A_		,							**A_			):
																		return {"bar": bar, "foo": foo}

											with patch('dill.detect.globalvars'		,							side_effect=A_			) as globalvars_mocka:
																		UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
																		self.assertGreater(globalvars_mocka.call_count		,							0			)
											with patch('dill.detect.globalvars'		,							side_effect=A_			) as globalvars_mocka:
																		UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
																		self.assertGreater(globalvars_mocka.call_count		,							0			)
											self.assertEqual(A_		,							A_			)


class        SCREAMING_SNAKE_CASE__	(			snake_case_):







				def 		UpperCAmelCase_  (    self			)->						Tuple:







											'''simple docstring'''






											UpperCamelCase					       =			Hasher.hash('hello'			)
											UpperCamelCase					       =			Hasher.hash('hello'			)
											UpperCamelCase					       =			Hasher.hash('there'			)
											self.assertEqual(A_		,							A_			)
											self.assertNotEqual(A_		,							A_			)







				def 		UpperCAmelCase_  (    self			)->						Dict:







											'''simple docstring'''






											UpperCamelCase					       =			Hasher.hash(Foo('hello'			)			)
											UpperCamelCase					       =			Hasher.hash(Foo('hello'			)			)
											UpperCamelCase					       =			Hasher.hash(Foo('there'			)			)
											self.assertEqual(A_		,							A_			)
											self.assertNotEqual(A_		,							A_			)







				def 		UpperCAmelCase_  (    self			)->						int:







											'''simple docstring'''






											UpperCamelCase					       =			Hasher()
											for x in ["hello", Foo('hello'			)]:
																		hasher.update(A_			)
											UpperCamelCase					       =			hasher.hexdigest()
											UpperCamelCase					       =			Hasher()
											for x in ["hello", Foo('hello'			)]:
																		hasher.update(A_			)
											UpperCamelCase					       =			hasher.hexdigest()
											UpperCamelCase					       =			Hasher()
											for x in ["there", Foo('there'			)]:
																		hasher.update(A_			)
											UpperCamelCase					       =			hasher.hexdigest()
											self.assertEqual(A_		,							A_			)
											self.assertNotEqual(A_		,							A_			)







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







											'''simple docstring'''






											with self.assertRaises(pickle.PicklingError			):
																		Hasher.hash(UnpicklableCallable(Foo('hello'			)			)			)







				def 		UpperCAmelCase_  (    self			)->						str:







											'''simple docstring'''






											UpperCamelCase					       =			'abc'
											UpperCamelCase					       =			[string, string]  # two strings have the same ids
											UpperCamelCase					       =			[string, string]
											UpperCamelCase					       =			json.loads(F'''["{string}", "{string}"]'''			)  # two strings have different ids
											self.assertIs(obja[0]		,							A_			)
											self.assertIs(obja[0]		,							obja[1]			)
											self.assertIs(obja[0]		,							A_			)
											self.assertIs(obja[0]		,							obja[1]			)
											self.assertIsNot(obja[0]		,							A_			)
											self.assertIsNot(obja[0]		,							obja[1]			)
											UpperCamelCase					       =			Hasher.hash(A_			)
											UpperCamelCase					       =			Hasher.hash(A_			)
											UpperCamelCase					       =			Hasher.hash(A_			)
											self.assertEqual(A_		,							A_			)
											self.assertEqual(A_		,							A_			)







				@require_tiktoken
				def 		UpperCAmelCase_  (    self			)->						Dict:







											'''simple docstring'''






											import tiktoken

											UpperCamelCase					       =			tiktoken.get_encoding('gpt2'			)
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
											UpperCamelCase					       =			tiktoken.get_encoding('r50k_base'			)
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
											UpperCamelCase					       =			tiktoken.get_encoding('gpt2'			)
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
											self.assertEqual(A_		,							A_			)
											self.assertNotEqual(A_		,							A_			)







				@require_torch
				def 		UpperCAmelCase_  (    self			)->						int:







											'''simple docstring'''






											import torch

											UpperCamelCase					       =			torch.tensor([1.0]			)
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
											UpperCamelCase					       =			torch.tensor([2.0]			)
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
											UpperCamelCase					       =			torch.tensor([1.0]			)
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
											self.assertEqual(A_		,							A_			)
											self.assertNotEqual(A_		,							A_			)







				@require_spacy
				@require_spacy_model('en_core_web_sm'			)
				@require_spacy_model('fr_core_news_sm'			)
				@pytest.mark.integration
				def 		UpperCAmelCase_  (    self			)->						Optional[Any]:







											'''simple docstring'''






											import spacy

											UpperCamelCase					       =			spacy.load('en_core_web_sm'			)
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
											UpperCamelCase					       =			spacy.load('fr_core_news_sm'			)
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
											UpperCamelCase					       =			spacy.load('en_core_web_sm'			)
											UpperCamelCase					       =			mda(datasets.utils.py_utils.dumps(A_			)			).hexdigest()
											self.assertEqual(A_		,							A_			)
											self.assertNotEqual(A_		,							A_			)







@pytest.mark.integration
def 						A_(			A		:						Path):
							UpperCamelCase					       =			tmp_path / 'dir1'
							UpperCamelCase					       =			tmp_path / 'dir2'
							dira.mkdir()
							dira.mkdir()
							UpperCamelCase					       =			'script.py'
							UpperCamelCase					       =			dedent(
							    '\n    from datasets.fingerprint import Hasher\n    def foo():\n        pass\n    print(Hasher.hash(foo))\n    ')
							UpperCamelCase					       =			dira / script_filename
							UpperCamelCase					       =			dira / script_filename
							with script_patha.open('w') as f:
														f.write(A)
							with script_patha.open('w') as f:
														f.write(A)
							UpperCamelCase					       =			subprocess.check_output(['python', str(A)])
							UpperCamelCase					       =			subprocess.check_output(['python', str(A)])
							assert fingerprinta == fingerprinta







def 						A_(			):
							UpperCamelCase					       =			{'a': [0, 1, 2]}
							UpperCamelCase					       =			DatasetChild(InMemoryTable.from_pydict(A))
							UpperCamelCase					       =			dataset.funca()._fingerprint
							assert expected_fingerprint == dataset.funca()._fingerprint
							assert expected_fingerprint != dataset.funca()._fingerprint

							with Pool(2) as p:
														assert expected_fingerprint == p.apply_async(dataset.funca).get()._fingerprint
														assert expected_fingerprint != p.apply_async(dataset.funca).get()._fingerprint







def 						A_(			):
							UpperCamelCase					       =			{'a': [0, 1, 2]}


							class        SCREAMING_SNAKE_CASE__	(			datasets.Dataset):







											@fingerprint_transform(inplace=A_			)
											def 		UpperCAmelCase_  (    self		,							A_			)->						List[Any]:







																		'''simple docstring'''






																		return DummyDatasetChild(self.data		,							fingerprint=A_			)


							UpperCamelCase					       =			DummyDatasetChild(InMemoryTable.from_pydict(A)).func()


							class        SCREAMING_SNAKE_CASE__	(			datasets.Dataset):







											@fingerprint_transform(inplace=A_		,							version='1.0.0'			)
											def 		UpperCAmelCase_  (    self		,							A_			)->						str:







																		'''simple docstring'''






																		return DummyDatasetChild(self.data		,							fingerprint=A_			)


							UpperCamelCase					       =			DummyDatasetChild(InMemoryTable.from_pydict(A)).func()


							class        SCREAMING_SNAKE_CASE__	(			datasets.Dataset):







											@fingerprint_transform(inplace=A_		,							version='2.0.0'			)
											def 		UpperCAmelCase_  (    self		,							A_			)->						List[str]:







																		'''simple docstring'''






																		return DummyDatasetChild(self.data		,							fingerprint=A_			)


							UpperCamelCase					       =			DummyDatasetChild(InMemoryTable.from_pydict(A)).func()

							assert len({fingeprint_no_version, fingeprint_a, fingeprint_a}) == 3







def 						A_(			):
							# AttributeError: module 'dill._dill' has no attribute 'stack'
							UpperCamelCase					       =			Hasher()
							hasher.update(lambda A: x)






'''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()
lowerCAmelCase :  Union[str, Any]      =      logging.get_logger(__name__)







def 						A_(			A		:						int):
							UpperCamelCase					       =			ConvNextVaConfig()

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

							UpperCamelCase					       =			1000
							UpperCamelCase					       =			'imagenet-1k-id2label.json'
							UpperCamelCase					       =			(1, 1000)

							UpperCamelCase					       =			'huggingface/label-files'
							UpperCamelCase					       =			num_labels
							UpperCamelCase					       =			json.load(open(hf_hub_download(A						,			A						,			repo_type='dataset')						,			'r'))
							UpperCamelCase					       =			{int(A): v for k, v in idalabel.items()}

							UpperCamelCase					       =			idalabel
							UpperCamelCase					       =			{v: k for k, v in idalabel.items()}
							UpperCamelCase					       =			hidden_sizes
							UpperCamelCase					       =			depths

							return config, expected_shape







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

							return name







def 						A_(			):
							UpperCamelCase					       =			'http://images.cocodataset.org/val2017/000000039769.jpg'
							UpperCamelCase					       =			Image.open(requests.get(A						,			stream=A).raw)
							return im







def 						A_(			A		:						Tuple):
							if "224" in checkpoint_url:
														UpperCamelCase					       =			224
														UpperCamelCase					       =			224 / 256
							elif "384" in checkpoint_url:
														UpperCamelCase					       =			384
														UpperCamelCase					       =			None
							else:
														UpperCamelCase					       =			512
														UpperCamelCase					       =			None

							return ConvNextImageProcessor(
							    size=A						,			crop_pct=A						,			image_mean=[0.485, 0.456, 0.406]						,			image_std=[0.229, 0.224, 0.225]						,			resample=PILImageResampling.BICUBIC						,			)







@torch.no_grad()
def 						A_(			A		:						List[str]						,			A		:						List[str]						,			A		:						Any						,			A		:						int):
							print('Downloading original model from checkpoint...')
							# define ConvNeXTV2 configuration based on URL
							UpperCamelCase       ,     UpperCamelCase					       =			get_convnextva_config(A)
							# load original state_dict from URL
							UpperCamelCase					       =			torch.hub.load_state_dict_from_url(A)['model']

							print('Converting model parameters...')
							# rename keys
							for key in state_dict.copy().keys():
														UpperCamelCase					       =			state_dict.pop(A)
														UpperCamelCase					       =			val
							# add prefix to all keys expect classifier head
							for key in state_dict.copy().keys():
														UpperCamelCase					       =			state_dict.pop(A)
														if not key.startswith('classifier'):
																					UpperCamelCase					       =			'convnextv2.' + key
														UpperCamelCase					       =			val

							# load HuggingFace model
							UpperCamelCase					       =			ConvNextVaForImageClassification(A)
							model.load_state_dict(A)
							model.eval()

							# Check outputs on an image, prepared by ConvNextImageProcessor
							UpperCamelCase					       =			convert_preprocessor(A)
							UpperCamelCase					       =			preprocessor(images=prepare_img()						,			return_tensors='pt')
							UpperCamelCase					       =			model(**A).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					       =			torch.tensor([-0.3_930, 0.1_747, -0.5_246, 0.4_177, 0.4_295])
							elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
														UpperCamelCase					       =			torch.tensor([-0.1_727, -0.5_341, -0.7_818, -0.4_745, -0.6_566])
							elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
														UpperCamelCase					       =			torch.tensor([-0.0_333, 0.1_563, -0.9_137, 0.1_054, 0.0_381])
							elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
														UpperCamelCase					       =			torch.tensor([-0.1_744, -0.1_555, -0.0_713, 0.0_950, -0.1_431])
							elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
														UpperCamelCase					       =			torch.tensor([0.9_996, 0.1_966, -0.4_386, -0.3_472, 0.6_661])
							elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
														UpperCamelCase					       =			torch.tensor([-0.2_553, -0.6_708, -0.1_359, 0.2_518, -0.2_488])
							elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
														UpperCamelCase					       =			torch.tensor([-0.0_673, -0.5_627, -0.3_753, -0.2_722, 0.0_178])
							elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
														UpperCamelCase					       =			torch.tensor([-0.6_377, -0.7_458, -0.2_150, 0.1_184, -0.0_597])
							elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
														UpperCamelCase					       =			torch.tensor([1.0_799, 0.2_322, -0.8_860, 1.0_219, 0.6_231])
							elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
														UpperCamelCase					       =			torch.tensor([0.3_766, 0.4_917, -1.1_426, 0.9_942, 0.6_024])
							elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
														UpperCamelCase					       =			torch.tensor([0.4_220, -0.6_919, -0.4_317, -0.2_881, -0.6_609])
							elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
														UpperCamelCase					       =			torch.tensor([0.1_082, -0.8_286, -0.5_095, 0.4_681, -0.8_085])
							elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
														UpperCamelCase					       =			torch.tensor([-0.2_419, -0.6_221, 0.2_176, -0.0_980, -0.7_527])
							elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
														UpperCamelCase					       =			torch.tensor([0.0_391, -0.4_371, 0.3_786, 0.1_251, -0.2_784])
							elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
														UpperCamelCase					       =			torch.tensor([-0.0_504, 0.5_636, -0.1_729, -0.6_507, -0.3_949])
							elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
														UpperCamelCase					       =			torch.tensor([0.3_560, 0.9_486, 0.3_149, -0.2_667, -0.5_138])
							elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
														UpperCamelCase					       =			torch.tensor([-0.2_469, -0.4_550, -0.5_853, -0.0_810, 0.0_309])
							elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
														UpperCamelCase					       =			torch.tensor([-0.3_090, 0.0_802, -0.0_682, -0.1_979, -0.2_826])
							else:
														raise ValueError(f'''Unknown URL: {checkpoint_url}''')

							assert torch.allclose(logits[0, :5]						,			A						,			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(A):
																					os.mkdir(A)

														model.save_pretrained(A)
														preprocessor.save_pretrained(A)

							UpperCamelCase					       =			'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(A)
														preprocessor.push_to_hub(A)


if __name__ == "__main__":
						lowerCAmelCase :  List[Any]      =      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')

						lowerCAmelCase :  int      =      parser.parse_args()
						convert_convnextva_checkpoint(
						    args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
						)




